import numpy as np

from scm.plams.core.functions import config, delete_job, finish, init
from scm.plams.core.settings import Settings
from scm.plams.interfaces.adfsuite.ams import AMSJob
from scm.plams.interfaces.adfsuite.amsworker import AMSWorker
from scm.plams.mol.molecule import Molecule

__all__ = ["preoptimize", "refine_density", "refine_lattice"]


def preoptimize(
    molecule: Molecule, model: str = "UFF", settings: Settings = None, nproc: int = 1, maxiterations: int = 100
):
    """
    Returns an optimized Molecule (or list of optimized molecules)

    molecule: Molecule or list of Molecules
        Molecule to optimize
    model: str
        Shorthand for some model, e.g. 'UFF'
    settings: Settings
        Custom engine settings (overrides ``model``)
    nproc: int
        Number of processes
    maxiterations: int
        Maximum number of iterations for the geometry optimization.
    """
    my_settings = _get_quick_settings(model, settings, nproc)
    single_molecule = isinstance(molecule, Molecule)

    input_molecules = [molecule] if single_molecule else molecule
    output_molecules = []
    with AMSWorker(my_settings) as worker:
        for i, mol in enumerate(input_molecules):
            results = worker.GeometryOptimization(
                name=f"preoptimization_{i}", molecule=mol, maxiterations=maxiterations, pretendconverged=True
            )
            output_molecules.append(results.get_main_molecule())

    if single_molecule:
        return output_molecules[0]
    else:
        return output_molecules


def refine_density(
    molecule: Molecule,
    density: float,
    step_size=50,
    model: str = "UFF",
    settings: Settings = None,
    nproc: int = 1,
    maxiterations: int = 100,
):
    """

    Performs a series of geometry optimizations with densities approaching
    ``density``. This can be useful if you want to compress a system to a
    given density, but cannot just use apply_strain() (because
    apply_strain() also scales the bond lengths).

    This function can be useful if for example packmol does not succeed to
    pack molecules with the desired density. Packmol can then generate a
    structure with a lower density, and this function can be used to
    increase the density to the desired value.

    Returns: a Molecule with the requested density.

    molecule: Molecule
        The molecule must have a 3D lattice
    density: float
        Target density in kg/m^3   (1000x the density in g/cm^3)
    step_size: float
        Step size for the density (in kg/m^3). Set step_size to a large number to only use 1 step.
    model: str
        e.g. 'UFF'
    settings: Settings
        Engine settings (overrides ``model``)
    maxiterations: int
        maximum number of iterations for the geometry optimization.


    """

    assert len(molecule.lattice) == 3
    tolerance = 1e-3
    my_settings = _get_quick_settings(model, settings, nproc)
    current_density = molecule.get_density()  # kg/m^3
    output_molecule = molecule.copy()
    counter = 0
    with AMSWorker(my_settings) as worker:
        while current_density < density - tolerance or current_density > density + tolerance:
            counter += 1
            if current_density < density - step_size:
                new_density = current_density + step_size
            elif current_density > density + step_size:
                new_density = current_density - step_size
            else:
                new_density = density

            output_molecule.set_density(new_density)

            results = worker.GeometryOptimization(
                name=f"preoptimization_{counter}",
                molecule=output_molecule,
                maxiterations=maxiterations,
                pretendconverged=True,
            )

            output_molecule = results.get_main_molecule()
            current_density = output_molecule.get_density()  # kg/m^3

    return output_molecule


def refine_lattice(
    molecule: Molecule,
    lattice,
    n_points=None,
    max_strain=0.15,
    model: str = "UFF",
    settings: Settings = None,
    nproc: int = 1,
    maxiterations: int = 10,
):
    """

    Returns a ``Molecule`` for which the lattice of the ``molecule`` is
    transformed to ``lattice``, by performing short geometry optimizations
    (each for at most ``maxiterations``) on gradually distorted lattices
    (linearly interpolating from the original lattice to the new lattice
    using ``n_points`` points).

    This can be useful for transforming an orthorhombic box of a liquid
    into a non-orthorhombic box of a liquid, where the gradual transformation
    of the lattice ensures that the molecules do not become too distorted.

    If init() has been called before calling this function, the job will be run in
    the current PLAMS working directory and will be deleted when the job finishes.

    Returns: a Molecule with the requested lattice. If the refinement
    fails, ``None`` is returned.

    molecule: Molecule
        The initial molecule
    lattice: list of list of float
        List with 1, 2, or 3 elements. Each element is a list of float with 3 elements each. For example, ``lattice=[[10, 0, 0],[-5, 5, 0],[0, 0, 12]]``.
    n_points: None or int >=2
        Number of points used for the linear interpolation. If None, n_points will be chosen such that the maximum strain for any step is at most ``max_strain`` compared to the original lattice vector lengths.
    max_strain: float
        Only if ``n_points=None``, use this value to determine the maximum allowed strain from one step to the next (as a fraction of the length of the original lattice vectors).
    model: str
        e.g. 'UFF'
    settings: Settings
        Engine settings (overrides ``model``)
    nproc: int
        Number of processes used by the job
    maxiterations: int
        maximum number of iterations for the geometry optimizations.

    """
    assert len(lattice) == len(
        molecule.lattice
    ), f"Different number of lattice vectors: len(molecule.lattice) = {len(molecule.lattice)}, len(lattice) = {len(lattice)}"
    assert all(len(x) == 3 for x in lattice), f"Lattice vectors must have three components. Lattice: {lattice}"
    assert all(len(x) == 3 for x in molecule.lattice), f"Lattice vectors must have three components. Lattice: {lattice}"
    assert (
        len(lattice) >= 1 and len(lattice) <= 3
    ), f"{len(lattice)} lattice vectors given but must be between 1 and 3. Lattice: {lattice}"

    def lattice2str(latt):
        return "\n".join(" ".join(str(j) for j in i) for i in latt)

    if n_points is None:
        lat1 = np.array(molecule.lattice)
        lat2 = np.array(lattice)
        strain = np.linalg.norm(lat2, axis=1) / np.linalg.norm(lat1, axis=1)
        n_points = np.ceil(np.max(np.abs(strain - 1) / max_strain))
        n_points = int(max(n_points, 2))

    assert n_points >= 2, f"Expected n_points >=2, but received {n_points}"

    called_plams_init = _ensure_init()

    s = _get_quick_settings(model, settings, nproc)
    s.input.ams.task = "PESScan"
    s.input.ams.PESScan.ScanCoordinate.FromLattice._1 = lattice2str(molecule.lattice)
    s.input.ams.PESScan.ScanCoordinate.ToLattice._1 = lattice2str(lattice)
    s.input.ams.PESScan.ScanCoordinate.NPoints = n_points
    s.input.ams.GeometryOptimization.MaxIterations = maxiterations
    s.input.ams.GeometryOptimization.PretendConverged = "Yes"

    job = AMSJob(settings=s, molecule=molecule, name="refine_density")
    job.run()

    final_molecule = None
    if job.ok():
        results = job.results.get_pesscan_results(molecules=True)
        final_molecule = results["Molecules"][-1]

    delete_job(job)

    if called_plams_init:
        finish()

    return final_molecule


def _ensure_init():
    if config.init:
        called_plams_init = False
    else:
        s = Settings()
        s.erase_workdir = True
        s.log.stdout = 0
        init(config_settings=s)
        called_plams_init = True

    return called_plams_init


def model_to_settings(model: str):
    """
    Returns Settings
    """
    settings = Settings()
    model = model.lower()
    if model == "uff":
        settings.input.ForceField.Type = "UFF"
    elif model == "gaff":
        settings.input.ForceField.Type = "GAFF"
        settings.input.ForceField.AnteChamberIntegration = "Yes"
    elif model == "gfnff":
        settings.input.GFNFF
    elif model == "ani-2x" or model == "ani2x":
        settings.input.MLPotential.Model = "ANI-2x"
    elif model == "gfn1xtb" or model == "gfn1-xtb":
        settings.input.DFTB.Model = "GFN1-xTB"
    elif model == "m3gnet-up-2022" or model == "m3gnetup2022":
        settings.input.MLPotential.Model = "M3GNet-UP-2022"
    else:
        raise ValueError("Unknown model: {}".format(model))

    return settings


def _get_quick_settings(model, settings, nproc):
    if settings is None:
        my_settings = model_to_settings(model)
    else:
        my_settings = settings.copy()

    if nproc:
        my_settings.runscript.nproc = nproc

    return my_settings
