"""Kd prediction module using Prodigy."""

import logging
import subprocess
import warnings
from collections import OrderedDict
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union

import joblib
import numpy as np
import pandas as pd
from joblib import Parallel, delayed
from prodigy_prot.modules.parsers import parse_structure
from prodigy_prot.modules.prodigy import Prodigy

logger = logging.getLogger(__name__)

# ##===============================================================
# ## this part is not ideal, but necessary to maintain compatibility

# # Compatibility shim for loading joblib/pickle files that reference
# # the old numpy RNG module path `numpy.random._mt19937.MT19937`.
# # Some models were pickled with a different numpy version where the
# # BitGenerator class lived under that module path. If the current
# # numpy exposes `MT19937` under a different path, the unpickler
# # fails with "is not a known BitGenerator module." To avoid that,
# # insert a small module alias into sys.modules mapping the old
# # module name to the current `MT19937` implementation.
# try:
#     import sys
#     import types

#     import numpy as _np

#     if "numpy.random._mt19937" not in sys.modules:
#         _mod = types.ModuleType("numpy.random._mt19937")
#         # new numpy exposes MT19937 as numpy.random.MT19937
#         try:
#             _mod.MT19937 = _np.random.MT19937
#         except AttributeError:
#             # Fallback: older numpy may not have this attribute; nothing to do
#             pass
#         sys.modules["numpy.random._mt19937"] = _mod
# except Exception:
#     # Don't block module import if the shim fails for any reason.
#     pass
# # Models are loaded lazily to avoid import-time errors when pickle files
# # were created with a different numpy/joblib setup. If loading fails,
# # an informative error will be raised when the models are actually needed.
# GP_SINGLE_MODEL = None
# GP_MULTIPLE_MODEL = None
# PCA_SINGLE_MODEL = None
# PCA_MULTIPLE_MODEL = None
# SCALER_MULTIPLE_MODEL = None
# SCALER_SINGLE_MODEL = None
# ##===============================================================


GP_SINGLE_MODEL = joblib.load(Path(__file__).parent / "gp_model_single_AF2.pkl")
GP_MULTIPLE_MODEL = joblib.load(Path(__file__).parent / "gp_model_multiple_AF2.pkl")
PCA_SINGLE_MODEL = joblib.load(Path(__file__).parent / "pca_single_AF2.pkl")
PCA_MULTIPLE_MODEL = joblib.load(Path(__file__).parent / "pca_multiple_AF2.pkl")
SCALER_MULTIPLE_MODEL = joblib.load(Path(__file__).parent / "scaler_multiple_AF2.pkl")
SCALER_SINGLE_MODEL = joblib.load(Path(__file__).parent / "scaler_single_AF2.pkl")


def _try_load_model(path: Path):
    """Attempt to load a model with joblib and return it.

    Any exception is propagated to the caller for handling.
    """
    return joblib.load(path)


class ProdigyPredictor:
    """Wrapper for Prodigy Kd prediction."""

    def __init__(
        self,
        temperature: float = 25.0,
        distance_cutoff: float = 5.5,
        acc_threshold: float = 0.05,
    ):
        """Initialize Prodigy predictor.

        Parameters
        ----------
        temperature : float, optional
            temperatureerature in Celsius, by default 25.0.
        distance_cutoff : float, optional
            Distance cutoff with default 5.5.
        acc_threshold : float, optional
            Acceptance threshold, by default 0.05.

        """
        self.temperature = temperature
        self.distance_cutoff = distance_cutoff
        self.acc_threshold = acc_threshold

    def feature_gen_single(self, pdb_path: Union[Path, str]) -> dict:
        """Predict Kd for a single PDB structure.

        Parameters
        ----------
        pdb_path : Path
            Path to PDB file. Supports ".pdb", ".cif", ".ent" formats.

        Returns
        -------
        prodigy_results : dict
            Dictionary of Prodigy features.

        """
        pdb_path = Path(pdb_path).expanduser().resolve()
        structures, _, _ = parse_structure(pdb_path)
        prodigy_model = Prodigy(
            model=structures[0],
            name=pdb_path.name,
            selection=None,
            temp=self.temperature,
        )  # type: ignore
        logger.info(f"Predicting Kd for {pdb_path.name}...")

        prodigy_model.predict(
            temp=self.temperature,
            distance_cutoff=self.distance_cutoff,
            acc_threshold=self.acc_threshold,
        )
        prodigy_results = prodigy_model.as_dict()
        # convert the unit of kd to µM
        prodigy_results["kd_val"] = prodigy_model.kd_val * 1.0e6
        prodigy_results.pop("model", None)
        prodigy_results.pop("selection", None)

        return prodigy_results

    def predict_kd_multiple(self, pdb_paths: List[Path]) -> List[float]:
        """Predict Kd for multiple PDB structures.

        Parameters
        ----------
        pdb_paths : List[Path]
            List of PDB file paths.

        """
        prodigy_results = Parallel(n_jobs=-1)(
            delayed(self.feature_gen_single)(pdb_path) for pdb_path in pdb_paths
        )
        kd_values = [
            float(result["kd_val"]) if result is not None else float("NaN")
            for result in prodigy_results
        ]
        logger.info(f"Predicted Kd values for {len(kd_values)} PDB structures")

        self.kd_values = kd_values

        return kd_values

    def aggregate_kd_values(self, mutation_name: str) -> dict:
        """Aggregate Kd values and return a dictionary.

        Returns
        -------
        float
            Aggregated Kd value in uM.
        """
        # build a dictionary so that we can convert to pandas dataframe in the future

        kd_dict = OrderedDict()
        kd_dict["Mutation"] = mutation_name
        for idx, kd in enumerate(self.kd_values):
            kd_dict[f"Kd{idx + 1} (µM)"] = kd

        kd_dict["avg Kd (µM)"] = np.nanmean(self.kd_values)
        kd_dict["SD (µM)"] = np.nanstd(self.kd_values)

        return kd_dict


class GaussianProcessPredictor:
    """Gaussian Process regression model for Kd prediction."""

    def __init__(
        self,
        wdir: Path,
        model_type: str = "gp_multiple",
        temperature: float = 25.0,
        distance_cutoff: float = 5.5,
        acc_threshold: float = 0.05,
    ):
        """Initialize Gaussian Process predictor."""
        self.wdir = Path(wdir).expanduser().resolve()
        self.model_type = model_type
        self.temperature = temperature
        self.distance_cutoff = distance_cutoff
        self.acc_threshold = acc_threshold
        # load the GP model
        self.gp_model, self.pca_model, self.scaler_model = self.load_model()

    # def load_model(self):
    #     """Load the Gaussian Process model."""
    #     models_dir = Path(__file__).parent / "models"

    #     def _load_or_global(global_var, fname):
    #         # if already loaded in memory, return it
    #         if global_var is not None:
    #             return global_var
    #         # otherwise attempt to load from file
    #         p = models_dir / fname
    #         if not p.exists():
    #             raise FileNotFoundError(f"Required model file not found: {p}")
    #         try:
    #             return _try_load_model(p)
    #         except Exception as e:
    #             raise RuntimeError(
    #                 f"Failed to load model file {p}: {e}\n"
    #                 "This often happens when the model pickle was created with a "
    #                 "different numpy/joblib version. Recreate the model with the "
    #                 "current environment or use a matching numpy version."
    #             )

    #     mt = self.model_type.lower()
    #     if mt == "gp_multiple":
    #         gp_model = _load_or_global(GP_MULTIPLE_MODEL, "gp_model_multiple_AF2.pkl")
    #         pca_model = _load_or_global(PCA_MULTIPLE_MODEL, "pca_multiple_AF2.pkl")
    #         scaler_model = _load_or_global(SCALER_MULTIPLE_MODEL, "scaler_multiple_AF2.pkl")
    #     elif mt == "gp_single":
    #         gp_model = _load_or_global(GP_SINGLE_MODEL, "gp_model_single_AF2.pkl")
    #         pca_model = _load_or_global(PCA_SINGLE_MODEL, "pca_single_AF2.pkl")
    #         scaler_model = _load_or_global(SCALER_SINGLE_MODEL, "scaler_single_AF2.pkl")
    #     else:
    #         raise ValueError("model_type must be either 'gp_multiple' or 'gp_single'")

    #     return gp_model, pca_model, scaler_model

    def load_model(self):
        """Load the Gaussian Process model."""
        if self.model_type == "gp_multiple":
            gp_model = GP_MULTIPLE_MODEL
            pca_model = PCA_MULTIPLE_MODEL
            scaler_model = SCALER_MULTIPLE_MODEL
        elif self.model_type == "gp_single":
            gp_model = GP_SINGLE_MODEL
            pca_model = PCA_SINGLE_MODEL
            scaler_model = SCALER_SINGLE_MODEL
        else:
            raise ValueError("model_type must be either 'gp_multiple' or 'gp_single'")

        return gp_model, pca_model, scaler_model

    def predict_kd(
        self,
        pdb_fpath: Optional[Path] = None,
    ) -> Tuple[Union[float, np.ndarray], Union[float, np.ndarray]]:
        """Predict Kd values for one or many PDBs in a single batch (one mutant type).

        Parameters
        ----------
        pdb_fpath : Optional[Path], optional
            Path to PDB file or directory containing PDB files. If None, all PDB files in
            the working directory will be used, by default None.

        Returns
        -------
        kd: Union[float, np.ndarray]
            Predicted Kd value(s) in µM.

        Notes
        -----
        - For 'gp_single' model type, exactly one PDB file is required. If a directory
          is provided, the function will look for files containing 'rank_001' in their names
          and use the first one found.
        - For 'gp_multiple' model type, exactly five PDB files are required; otherwise, it will use
          the first five PDB files sorted alphabetically.

        """
        # gather all pdb files in the working directory if pdb_fpath is None
        if pdb_fpath is None:
            pdb_paths = sorted(self.wdir.glob("*.pdb"))
        # when pdb_fpath is provided, check if it's a file or directory
        else:
            pdb_source = Path(pdb_fpath).expanduser().resolve()
            # if the provided path is a directory, gather all pdb files in it and get the top 1
            # based on "rank_001" tag
            if pdb_source.is_dir():
                pdb_paths = sorted(pdb_source.glob("*.pdb"))
                rank_001 = [p for p in pdb_paths if "rank_001" in p.name]
                pdb_paths = [rank_001[0] if rank_001 else pdb_paths[0]] if pdb_paths else []
            else:
                pdb_paths = [pdb_source]

        # check the compatibility between number of pdbs and model type
        if not pdb_paths:
            raise FileNotFoundError("No PDB files found for prediction.")
        if self.model_type == "gp_single" and len(pdb_paths) != 1:
            warnings.warn(
                "gp_single model requires exactly one PDB file. Taking the first one only."
            )
            pdb_paths = [pdb_paths[0]]
        if self.model_type == "gp_multiple":
            if len(pdb_paths) < 5:
                raise ValueError(
                    "gp_multiple model requires exactly five PDB files. "
                    f"Only {len(pdb_paths)} found."
                )
            elif len(pdb_paths) > 5:
                warnings.warn(
                    "gp_multiple model requires exactly five PDB files. Taking the first five only."
                )
                pdb_paths = pdb_paths[:5]
            else:
                pass  # exactly five pdbs provided

        prodigy_pred = ProdigyPredictor(
            temperature=self.temperature,
            distance_cutoff=self.distance_cutoff,
            acc_threshold=self.acc_threshold,
        )

        df_feat_list = []
        for idx, pdb_path in enumerate(pdb_paths):
            feat_dict = prodigy_pred.feature_gen_single(pdb_path)
            # print(
            #     f"logKd from prodigy: {np.log10(float(feat_dict['kd_val'])):.10f} "
            #     f"{'µM':>3}"
            # )
            feat_dict.pop("kd_val", None)
            df_features_tmp = pd.DataFrame(feat_dict, index=[0])
            # make sure the order of columns is correct so that the model can read them properly
            # this can avoid potential problems of data misalignment and shifting
            df_features_tmp = df_features_tmp[
                [
                    "ICs",
                    "nis_a",
                    "nis_c",
                    "ba_val",
                    "AA",
                    "PP",
                    "CC",
                    "AP",
                    "CP",
                    "AC",
                ]
            ]
            df_features_tmp.columns = [col + f"_{idx}" for col in df_features_tmp.columns]

            df_feat_list.append(df_features_tmp)

        df_feats = pd.concat(df_feat_list, axis=1)
        # print(f"df_feats.head():\n{df_feats.head()}")

        # check if df_feats has nan values
        if df_feats.isnull().values.any():
            raise ValueError("Feature matrix contains NaN values. Cannot proceed.")
        # scale the features
        feats = self.scaler_model.transform(df_feats.to_numpy())
        # perform PCA transformation
        feats = self.pca_model.transform(feats)
        # predict Kd using GP model
        kd_pred_mean, kd_pred_std = self.gp_model.predict(X=feats, return_std=True)

        return kd_pred_mean, kd_pred_std
