from pathlib import Path
from time import time
from typing import List, Union
from uuid import uuid4

import numpy as np

from poli.core.abstract_isolated_function import AbstractIsolatedFunction
from poli.core.util.proteins.foldx import FoldxInterface
from poli.core.util.proteins.pdb_parsing import (
    parse_pdb_as_residue_strings,
    parse_pdb_as_residues,
)

# This is the folder where all the files
# generated by FoldX will be stored.
# Feel free to change it if you want
# to keep the files somewhere else by
# passing tmp_folder to the black box.
# TODO: what happens if the user is on Windows?
DEFAULT_TMP_PATH = Path("/tmp").resolve()


class FoldxIsolatedFunction(AbstractIsolatedFunction):
    """
    A class representing the FoldxBlackBox, which is used for simulating protein mutations using FoldX.

    Parameters
    -----------
    wildtype_pdb_path : Union[Path, List[Path]], required
        The path(s) to the wildtype PDB file(s). (default: None)
    experiment_id : str, optional
        The experiment ID. (default: None)
    tmp_folder : Path, optional
        The temporary folder path. (default: None)
    eager_repair : bool, optional
        Flag indicating whether to eagerly repair the PDB files. (default: False)
    verbose : bool, optional
        Flag indicating whether we print the output from FoldX. (default: False)

    Attributes
    ----------
    experiment_id : str
        The experiment ID.
    tmp_folder : Path
        The temporary folder path.
    wildtype_pdb_paths : List[Path]
        The list of repaired wildtype PDB file paths.
    wildtype_residues : List[List[Residue]]
        The list of wildtype residues for each PDB file.
    wildtype_amino_acids : List[List[str]]
        The list of wildtype amino acids for each PDB file.
    wildtype_residue_strings : List[str]
        The list of wildtype residue strings for each PDB file.

    Methods
    -------
    create_working_directory() -> Path:
        Creates and returns the working directory path for the black box.

    """

    def __init__(
        self,
        wildtype_pdb_path: Union[Path, List[Path]],
        experiment_id: str = None,
        tmp_folder: Path = None,
        eager_repair: bool = False,
        verbose: bool = False,
    ):
        """
        Initialize the FoldxBlackBox.

        Parameters
        -----------
        wildtype_pdb_path : Union[Path, List[Path]]
            The path(s) to the wildtype PDB file(s).
        experiment_id : str, optional
            The experiment ID. (default: None)
        tmp_folder : Path, optional
            The temporary folder path. (default: None)
        eager_repair : bool, optional
            Flag indicating whether to eagerly repair the PDB files. (default: False)
        verbose : bool, optional
            Flag indicating whether we print the output from FoldX. (default: False)
        """
        # TODO: assert that wildtype_pdb_file is provided
        assert wildtype_pdb_path is not None, (
            "Missing required argument wildtype_pdb_file. "
            "Did you forget to pass it to create and into the black box?"
        )
        self.verbose = verbose

        # Defining the experiment id
        if experiment_id is None:
            experiment_id = f"{int(time())}_{str(uuid4())[:8]}"
        self.experiment_id = experiment_id

        self.tmp_folder = tmp_folder if tmp_folder is not None else DEFAULT_TMP_PATH

        if isinstance(wildtype_pdb_path, str):
            wildtype_pdb_path = Path(wildtype_pdb_path.strip())

        if isinstance(wildtype_pdb_path, Path):
            wildtype_pdb_path = [wildtype_pdb_path]

        if isinstance(wildtype_pdb_path, list):
            _wildtype_pdb_path = []
            for pdb_file in wildtype_pdb_path:
                if isinstance(pdb_file, str):
                    pdb_file = Path(pdb_file.strip())
                assert isinstance(
                    pdb_file, Path
                ), f"Expected a Path object or a string, but got {type(pdb_file)}."
                _wildtype_pdb_path.append(pdb_file)

            wildtype_pdb_path = _wildtype_pdb_path

        # At this point, wildtype_pdb_path is a list of Path objects.
        # We need to ensure that these are repaired pdb files.
        # We do this by creating a temporary folder and repairing
        # the pdbs there.
        if eager_repair:
            path_for_repairing_pdbs = self.tmp_folder / "foldx_tmp_files_for_repair"
            path_for_repairing_pdbs.mkdir(exist_ok=True, parents=True)
            foldx_interface_for_repairing = FoldxInterface(
                path_for_repairing_pdbs, verbose=self.verbose
            )

            # Re-writing wildtype_pdb_path to be the list of repaired pdb files.
            repaired_wildtype_pdb_files = [
                foldx_interface_for_repairing._repair_if_necessary_and_provide_path(
                    pdb_file
                )
                for pdb_file in wildtype_pdb_path
            ]

            # At this point, wildtype_pdb_path is a list of Path objects.
            self.wildtype_pdb_paths = repaired_wildtype_pdb_files
        else:
            self.wildtype_pdb_paths = wildtype_pdb_path

        self.wildtype_resiudes = [
            parse_pdb_as_residues(pdb_file) for pdb_file in self.wildtype_pdb_paths
        ]

        self.wildtype_amino_acids = [
            parse_pdb_as_residue_strings(pdb_file)
            for pdb_file in self.wildtype_pdb_paths
        ]

        self.wildtype_residue_strings = [
            "".join(amino_acids) for amino_acids in self.wildtype_amino_acids
        ]
        self.x0 = np.array(self.wildtype_residue_strings)

    def create_working_directory(self) -> Path:
        """
        Create and return the working directory path for the black box.

        Returns
        --------
        Path
            The path to the working directory.
        """
        sub_experiment_id = str(uuid4())[:8]

        working_dir = (
            self.tmp_folder / "foldx_tmp_files" / self.experiment_id / sub_experiment_id
        )
        working_dir.mkdir(exist_ok=True, parents=True)

        return working_dir

    def __call__(self, x: np.ndarray, context=None) -> np.ndarray:
        raise NotImplementedError
