"""Fine-tuning helpers backed by the example's local MatterGen model package."""

import logging
from collections import OrderedDict
from copy import deepcopy
from pathlib import Path
from typing import Tuple

import hydra
import pytorch_lightning as pl
import torch
from omegaconf import DictConfig, open_dict

from .common.utils.data_classes import MatterGenCheckpointInfo, _rewrite_vendored_targets
from .common.utils.globals import get_device

logger = logging.getLogger(__name__)


def init_adapter_lightningmodule_from_pretrained(
    adapter_cfg: DictConfig, lightning_module_cfg: DictConfig
) -> Tuple[pl.LightningModule, DictConfig]:
    """Initialize the local adapter model from a MatterGen checkpoint."""
    if adapter_cfg.model_path is not None:
        if adapter_cfg.pretrained_name is not None:
            logger.warning(
                "pretrained_name is provided, but will be ignored since model_path is also provided."
            )
        model_path = Path(hydra.utils.to_absolute_path(adapter_cfg.model_path))
        checkpoint_info = MatterGenCheckpointInfo(model_path, adapter_cfg.load_epoch)
    elif adapter_cfg.pretrained_name is not None:
        checkpoint_info = MatterGenCheckpointInfo.from_hf_hub(
            adapter_cfg.pretrained_name
        )
    else:
        raise ValueError("Either adapter.model_path or adapter.pretrained_name is required.")

    checkpoint_path = checkpoint_info.checkpoint_path
    version_root_path = Path(checkpoint_path).relative_to(
        checkpoint_info.model_path
    ).parents[1]
    config_path = Path(checkpoint_info.model_path) / version_root_path
    pretrained_cfg_path = (
        config_path if (config_path / "config.yaml").exists() else config_path.parent.parent
    )

    hydra.core.global_hydra.GlobalHydra.instance().clear()
    with hydra.initialize_config_dir(
        str(pretrained_cfg_path.absolute()), version_base="1.1"
    ):
        pretrained_cfg = hydra.compose(config_name="config")
    pretrained_cfg = _rewrite_vendored_targets(pretrained_cfg)
    diffusion_module_cfg = deepcopy(pretrained_cfg.lightning_module.diffusion_module)
    denoiser_cfg = diffusion_module_cfg.model

    with open_dict(adapter_cfg.adapter):
        for key, value in denoiser_cfg.items():
            if key not in {"_target_", "property_embeddings_adapt"}:
                adapter_cfg.adapter[key] = value
            if key == "property_embeddings":
                for field in value:
                    if field in adapter_cfg.adapter.property_embeddings_adapt:
                        adapter_cfg.adapter.property_embeddings_adapt.remove(field)

        adapter_cfg.adapter.gemnet["_target_"] = (
            "model.common.gemnet.gemnet_ctrl.GemNetTCtrl"
        )
        adapter_cfg.adapter.gemnet.condition_on_adapt = list(
            adapter_cfg.adapter.property_embeddings_adapt
        )

    with open_dict(diffusion_module_cfg):
        diffusion_module_cfg.model = adapter_cfg.adapter
    with open_dict(lightning_module_cfg):
        lightning_module_cfg.diffusion_module = diffusion_module_cfg

    lightning_module = hydra.utils.instantiate(lightning_module_cfg)
    checkpoint = torch.load(checkpoint_path, map_location=get_device())
    pretrained_state: OrderedDict = checkpoint["state_dict"]
    local_state: OrderedDict = lightning_module.state_dict()
    local_state.update(
        (key, pretrained_state[key])
        for key in local_state.keys() & pretrained_state.keys()
    )
    lightning_module.load_state_dict(local_state, strict=True)

    if not adapter_cfg.full_finetuning:
        pretrained_keys = set(pretrained_state.keys())
        for name, parameter in lightning_module.named_parameters():
            if name in pretrained_keys:
                parameter.requires_grad_(False)

    return lightning_module, lightning_module_cfg