import os
import shutil
import multiprocessing
import tempfile
from enum import Enum

from rdkit import Chem
from typing import Optional, Dict
from typing_extensions import Literal  # Required for Python 3.7. From 3.8 Literal is in typing.
from pydantic import PrivateAttr, BaseModel, Field

from dockstream.core.RDkit.RDkit_ligand_preparator import RDkitLigandPreparator

from copy import deepcopy
from dockstream.core.ligand.ligand import Ligand, get_enumerations_for_ligand
from dockstream.core.Schrodinger.license_token_guard import SchrodingerLicenseTokenGuard
from dockstream.loggers.blank_logger import BlankLogger

from dockstream.core.ligand_preparator import LigandPreparator, _LE
from dockstream.utils.general_utils import gen_temp_file

from dockstream.utils.parallelization.general_utils import split_into_sublists, get_progress_bar_string

from dockstream.utils.dockstream_exceptions import LigandPreparationFailed

from dockstream.utils.smiles import to_smiles
from dockstream.utils.execute_external.Schrodinger import SchrodingerExecutor
from dockstream.utils.enums.Schrodinger_enums import LigprepLigandPreparationEnum, SchrodingerExecutablesEnum, \
                                                 SchrodingerOutputEnum, SchrodingerDockingConfigurationEnum

_LP = LigprepLigandPreparationEnum()
_EE = SchrodingerExecutablesEnum()
_OE = SchrodingerOutputEnum()
_CE = SchrodingerDockingConfigurationEnum()


class Parallelization(BaseModel):
    number_cores: int = 4
    max_compounds_per_subjob: Optional[int] = 0


class EpikParameters(BaseModel):
    """Epik - Empirical pKa Prediction"""

    target_pH: float = 7.0
    pH_tolerance: float = 2.0


class ChiralityParameters(BaseModel):
    """Manipulate the stereogenic composition of a three-dimensional molecular structure.

    Sample some or all of the chiral centers,
    in one or multiple input structures.

    Stereoisomers are generated by varying the chiralities of chiral atoms
    whose chiralities are not specified.
    """

    max_number_stereoisomers: int = Field(default=10, description="Generate up to this many stereoisomers per input structure. This is -s option of the ligprep command.")


class LigprepForceField(str, Enum):
    """Force field."""

    OPLS_2005 = "OPLS_2005"
    OPLS3e = "OPLS3e"
    # OPLS4 = "OPLS4"


class LigprepLigandPreparatorParameters(BaseModel):
    prefix_execution: Optional[str] = None
    binary_location: Optional[str] = None
    token_guard: Optional[Dict] = None
    parallelization: Optional[Parallelization] = Parallelization()
    use_epik: Optional[EpikParameters] = None
    chirality: Optional[ChiralityParameters] = None
    IGNORE_CHIRALITIES: Optional[bool] = Field(default=False, title="Ignore chiralities and generate all combinations", description="Do not respect existing chirality properties and do not respect chiralities from the input geometry. Generate stereoisomers for all chiral centers up to the number permitted (specified using the -s option, Default: 32). This is equivalent to 'Generate all combinations' in the Ligand Preparation user interface, and -ac option of the ligprep command.")
    filter_file: Optional[Dict] = None
    force_field: LigprepForceField = LigprepForceField.OPLS_2005
    command_line_parameters: Optional[Dict[str, str]] = None

    class Config:
        use_enum_values = True


class LigprepLigandPreparator(LigandPreparator, BaseModel):
    """Class that acts as an interface to the "Ligprep" executable to prepare ligands."""

    type: Literal["Ligprep"] = "Ligprep"
    parameters: LigprepLigandPreparatorParameters = LigprepLigandPreparatorParameters()

    class Config:
        underscore_attrs_are_private = True

    _ligprep_executor: SchrodingerExecutor = PrivateAttr()
    _token_guard: SchrodingerLicenseTokenGuard = PrivateAttr(default=None)
    _logger_blank = PrivateAttr()

    def __init__(self, **data):
        super().__init__(**data)

        self._logger_blank = BlankLogger()

        # initialize the executor for all "Ligprep" related calls and also check if it is available
        self._ligprep_executor = SchrodingerExecutor(prefix_execution=self.parameters.prefix_execution,
                                                     binary_location=self.parameters.binary_location)
        if not self._ligprep_executor.is_available():
            raise LigandPreparationFailed("Cannot initialize Ligprep backend - abort.")
        self._logger.log(f"Checked Glide backend availability (prefix_execution={self.parameters.prefix_execution}, binary_location={self.parameters.binary_location}).",
                         _LE.DEBUG)

        # if necessary parameters are specified, initialize the token guard
        if self.parameters.token_guard is not None:
            self._token_guard = SchrodingerLicenseTokenGuard(**self.parameters.token_guard)

    def _load_references(self):
        references = []
        ref_format = self.align.reference_format.upper()
        for path in self.align.reference_paths:
            if ref_format == _LP.ALIGN_REFERENCE_FORMAT_PDB:
                ref_mol = Chem.MolFromPDBFile(path, sanitize=True)
                ref_mol.SetProp("_Name", os.path.basename(path))
                references.append(ref_mol)
            elif ref_format == _LP.ALIGN_REFERENCE_FORMAT_SDF:
                mol_supplier = Chem.SDMolSupplier(path)
                for mol in mol_supplier:
                    references.append(mol)
            else:
                raise IOError("Specified format not supported.")
        if len(references) == 0:
            raise LigandPreparationFailed("No reference molecules could be loaded with path(s) specified.")
        self._references = references
        self._logger.log(f"Stored {len(references)} reference molecules.", _LE.DEBUG)

    def _get_RDkit_aligner(self, conf, ligands):
        return RDkitLigandPreparator(ligands=ligands, **conf)

    def _prepare_ligprep_arguments(self) -> list:
        arguments_list = []

        # add user-specified command-line settings (if provided)
        if self.parameters.command_line_parameters is not None:
            for key in self.parameters.command_line_parameters.keys():
                arguments_list.append(key)
                if self.parameters.command_line_parameters[key] is not None and \
                   self.parameters.command_line_parameters[key] != "":
                    arguments_list.append(self.parameters.command_line_parameters[key])

        # add default settings, that are not exposed to the user yet
        if _EE.LIGPREP_HOST not in arguments_list:
            arguments_list.append(_EE.LIGPREP_HOST)
            arguments_list.append(_EE.LIGPREP_HOST_LOCALHOST)
        arguments_list.append(_EE.LIGPREP_WAIT)
        arguments_list = arguments_list + [_EE.LIGPREP_NJOBS, 1]

        # Epik: can be used to enumerate tautomers and ionization states, optional parameters set the target pH and the
        #       tolerance (in log units)
        if self.parameters.use_epik is not None:
            arguments_list.append(_EE.LIGPREP_EPIK)
            arguments_list = arguments_list + [_EE.LIGPREP_PH, self.parameters.use_epik.target_pH]
            arguments_list = arguments_list + [_EE.LIGPREP_PHT, self.parameters.use_epik.pH_tolerance]

        # Chirality: for now, pre-defined sterocenters are respected / kept (see definition above)
        if self.parameters.chirality is not None:
            arguments_list = arguments_list + [_EE.LIGPREP_S, self.parameters.chirality.max_number_stereoisomers]
        else:
            arguments_list.append(_EE.LIGPREP_G)   # respect chiralities that are pre-specified

        if self.parameters.IGNORE_CHIRALITIES:
            arguments_list.append(_EE.LIGPREP_AC)

        # Force field used: there are two implemented; while "OPLS_2005" seems to be the default for the command-line,
        #                   the "Ligprep" interface defaults to "OPLS3e"
        arguments_list.append(_EE.LIGPREP_BFF)
        ff_string = self.parameters.force_field
        if ff_string == _LP.FORCE_FIELD_OPLS_2005:
            arguments_list.append(_EE.LIGPREP_FF_OPLS_2005)
        elif ff_string == _LP.FORCE_FIELD_OPLS3e:
            arguments_list.append(_EE.LIGPREP_FF_OPLS3e)
        else:
            self._logger.log(f"Specified force field {ff_string} is not supported - defaulting to \"OPLS_2005\".",
                             _LE.WARNING)

        return arguments_list

    def _apply_token_guard(self):
        if self._token_guard is not None:
            self._token_guard.guard()

    def _get_sublists_for_embedding(self, number_cores, enforce_singletons=False):
        if enforce_singletons:
            return split_into_sublists(input_list=self.ligands, partitions=None, slice_size=1)

        # decide how to slice the ligand list depending on whether a maximum length is defined or not
        if self.parameters.parallelization.max_compounds_per_subjob > 0:
            slice_size = min(max(self.parameters.parallelization.max_compounds_per_subjob, 1),
                             len(self.ligands))
            return split_into_sublists(input_list=self.ligands, partitions=None, slice_size=slice_size)
        else:
            # split the ligands into as many cores as available
            partitions = min(number_cores, len(self.ligands))
            return split_into_sublists(input_list=self.ligands, partitions=partitions, slice_size=None)

    def _generate_temporary_input_output_files(self, start_indices, sublists):
        # in case singletons are handed over, wrap them in a list for "zipping" later
        if not isinstance(start_indices, list):
            start_indices = [start_indices]
        if not isinstance(sublists, list):
            sublists = [sublists]

        tmp_output_dirs = []
        tmp_input_smi_paths = []
        tmp_input_filter_paths = []
        tmp_output_sdf_paths = []
        for start_index, sublist in zip(start_indices, sublists):
            # generate temporary input files and output directory
            cur_tmp_output_dir = tempfile.mkdtemp()
            tmp_output_dirs.append(cur_tmp_output_dir)
            cur_tmp_smi = gen_temp_file(prefix=str(start_index), suffix=".smi", dir=cur_tmp_output_dir)
            cur_tmp_filter = gen_temp_file(prefix=str(start_index), suffix=".lff", dir=cur_tmp_output_dir)
            tmp_input_smi_paths.append(cur_tmp_smi)
            tmp_input_filter_paths.append(cur_tmp_filter)

            # write smiles to temporary file as "Ligprep" backend
            with open(cur_tmp_smi, 'w') as f:
                for lig in sublist:
                    f.write(lig.get_smile() + " " + lig.get_identifier() + "\n")

            # add the path to which "_dock_subjob()" will write the result SDF
            output_sdf_path = gen_temp_file(prefix=str(start_index), suffix="_result.sdf", dir=cur_tmp_output_dir)
            tmp_output_sdf_paths.append(output_sdf_path)
        return tmp_output_dirs, tmp_input_smi_paths, tmp_output_sdf_paths, tmp_input_filter_paths

    def _log_docking_progress(self, number_done, number_total):
        self._logger.log(get_progress_bar_string(number_done, number_total, length=65), _LE.INFO)

    def _get_number_cores(self):
        # prepare the parallelization and set the number of cores to be used
        number_cores = self.parameters.parallelization.number_cores
        if number_cores == 0:
            number_cores = 1
        elif number_cores < 0:
            # subtract the number of cores (neg. value, thus add up) from total number of cores, e.g. -1 will
            # use all available cores minus 1
            number_cores = multiprocessing.cpu_count() + number_cores
        return number_cores

    def generate3Dcoordinates(self):
        """Method to generate 3D coordinates, in case the molecules have to be built from SMILES."""

        number_cores = self._get_number_cores()
        start_indices, sublists = self._get_sublists_for_embedding(number_cores=number_cores)
        number_sublists = len(sublists)
        self._logger.log(f"Split ligands into {number_sublists} sublists for embedding.",
                         _LE.DEBUG)

        jobs_submitted = 0
        slices_per_iteration = min(number_cores, number_sublists)
        while jobs_submitted < len(sublists):
            upper_bound_slice = min((jobs_submitted + slices_per_iteration), len(sublists))
            cur_slice_start_indices = start_indices[jobs_submitted:upper_bound_slice]
            cur_slice_sublists = sublists[jobs_submitted:upper_bound_slice]

            # generate paths and initialize molecules (so that if they fail, this can be covered)
            tmp_output_dirs, tmp_input_smi_paths, \
            tmp_output_sdf_paths, tmp_input_filter_paths = self._generate_temporary_input_output_files(cur_slice_start_indices,
                                                                                                       cur_slice_sublists)

            # call "token guard" method (only executed, if block is specified in the configuration), which will wait
            # with the execution if not enough tokens are available at the moment
            self._apply_token_guard()

            # run in parallel; wait for all subjobs to finish before proceeding
            processes = []
            for chunk_index in range(len(tmp_output_dirs)):
                p = multiprocessing.Process(target=self._run_embedding_subjob, args=(tmp_input_smi_paths[chunk_index],
                                                                                     tmp_output_sdf_paths[chunk_index],
                                                                                     tmp_output_dirs[chunk_index],
                                                                                     tmp_input_filter_paths[chunk_index]))
                processes.append(p)
                p.start()
                jobs_submitted += 1
            for p in processes:
                p.join()

            # load and store the conformers; name it sequentially
            # note, that some backends require the H-coordinates (such as Glide) - so keep them!
            ligands_embedded = []
            for path_sdf_results in tmp_output_sdf_paths:
                if not os.path.isfile(path_sdf_results):
                    continue
                mol_supplier = Chem.SDMolSupplier(path_sdf_results, removeHs=False)
                for mol in mol_supplier:
                    # Ligprep adds a "-1" to "-[N]" to the names in the variants tag; this tag is always added
                    # alternatively, the "_Name" property could be loaded
                    if mol is not None and mol.HasProp(_OE.LIGPREP_VARIANTS):
                        identifier, new_enum = mol.GetProp(_OE.LIGPREP_VARIANTS).split('-')
                        lig_id, enum_id = identifier.split(':')
                        ligands_embedded.append(Ligand(smile="", ligand_number=int(lig_id), enumeration=int(enum_id),
                                                       molecule=mol, mol_type=_LP.TYPE_LIGPREP))
                    else:
                        self._logger.log(f"Skipped molecule when loading as specified property {_OE.LIGPREP_VARIANTS} could not be found - typically, this indicates that ligprep could not embed the molecule.", _LE.WARNING)

            # update internal (self.ligands) list of ligands with new molecules
            self._expand_enumerations(ligands_embedded)

            # remove temporary files
            for path in tmp_output_dirs:
                shutil.rmtree(path)
            self._log_docking_progress(number_done=jobs_submitted, number_total=number_sublists)

        # check success and failure with embedding
        failed = 0
        for ligand in self.ligands:
            if ligand.get_molecule() is None:
                failed += 1
                self._logger.log(f"Enumeration {ligand.get_identifier()} could not be embedded (smile: {ligand.get_smile()}).",
                                 _LE.DEBUG)
        if failed > 0:
            self._logger.log(f"It appears ligprep could not embed all {len(self.ligands)} ligands ({failed} were not found) - compounds not embedded might be ignored at the docking stage.",
                             _LE.WARNING)

    def _add_filtering(self, arguments: list, tmp_input_filter: str) -> list:
        if self.parameters.filter_file is not None:
            dicFilters = self.parameters.filter_file
            filter_file = open(tmp_input_filter, 'w')
            for key in dicFilters.keys():
                filter_file.write(f"{key}                             {dicFilters[key]}\n")
            filter_file.close()
            arguments = arguments + [_EE.LIGPREP_F,
                                     tmp_input_filter]
        return arguments

    def _run_embedding_subjob(self, smi_ligand_path, path_sdf_results, tmp_output_dir, tmp_input_filter):
        # 1) change to directory, to be able to use relative paths (to compensate for Schrodinger bug with AWS)
        working_dir = os.getcwd()
        os.chdir(tmp_output_dir)

        # 2) prepare "Ligprep" arguments
        arguments = self._prepare_ligprep_arguments()
        arguments = self._add_filtering(arguments, tmp_input_filter)
        arguments = arguments + [_EE.LIGPREP_INPUT_ISMI, os.path.basename(smi_ligand_path)]
        arguments = arguments + [_EE.LIGPREP_OUTPUT_OSD, os.path.basename(path_sdf_results)]

        # 3) run "Ligprep" backend and add log file to "debug" mode logging
        result = self._ligprep_executor.execute(command=_EE.LIGPREP,
                                                arguments=arguments,
                                                location=tmp_output_dir,
                                                check=False)
        self._logger.log(f"Executed Ligprep backend (output file: {path_sdf_results}).", _LE.DEBUG)
        path_tmp_log = os.path.join(tmp_output_dir,
                                    "".join([os.path.splitext(os.path.basename(path_sdf_results))[0],
                                             _OE.LIGPREP_LOG]))
        self._print_log_file(path=path_tmp_log)

        # 4) revert back to working directory
        os.chdir(working_dir)

    def _expand_enumerations(self, ligands_embedded):
        # store the generated conformations with the original ligands; if embedding failed, keep the old (emtpy) one
        new_ligand_list = []
        for ligand in self.ligands:
            enums = get_enumerations_for_ligand(ligands=ligands_embedded, ligand_id=ligand.get_ligand_number())
            if len(enums) == 0:
                new_ligand_list.append(ligand)
            else:
                for enum_id, enum in enumerate(enums):
                    buf_lig = deepcopy(ligand)
                    buf_lig.set_enumeration(enum_id)
                    buf_lig.set_molecule(enum.get_molecule())
                    buf_lig.set_mol_type(enum.get_mol_type())
                    buf_lig.set_smile(to_smiles(enum.get_molecule()))
                    new_ligand_list.append(buf_lig)
        self.ligands = new_ligand_list

    def _print_log_file(self, path):
        if os.path.isfile(path):
            with open(path, 'r') as log_file:
                log_file_raw = log_file.readlines()
                self._logger.log(f"Printing log file {path}:\n", _LE.DEBUG)
                for line in log_file_raw:
                    self._logger_blank.log(line.rstrip("\n"), _LE.DEBUG)
                self._logger_blank.log("", _LE.DEBUG)
                self._logger.log("--- End file", _LE.DEBUG)

    def align_ligands(self):
        self.ligands = self._align_ligands_with_RDkit_preparator(self.ligands)
