# (c) 2015-2022 Acellera Ltd http://www.acellera.com
# All Rights Reserved
# Distributed under HTMD Software License Agreement
# No redistribution in whole or part
#
from os import path
from htmd.adaptive.adaptivegoal import AdaptiveGoal
from htmd.model import macroAccumulate
from protocolinterface import val
import numpy as np
import os
import logging

logger = logging.getLogger(__name__)


class AdaptiveGoalEG(AdaptiveGoal):
    """Adaptive goal-directed sampling with epsilon-greedy state selection.

    Extends :class:`AdaptiveGoal` by splitting simulation budget between the
    undirected component (exploration) and the directed component (exploitation)
    using an epsilon-greedy policy rather than a fixed scale factor.

    Parameters
    ----------
    app : :class:`SimQueue <jobqueues.simqueue.SimQueue>` object, default=None
        A SimQueue class object used to retrieve and submit simulations
    project : str, default='adaptive'
        The name of the project
    nmin : int, default=0
        Minimum number of running simulations
    nmax : int, default=1
        Maximum number of running simulations
    nepochs : int, default=1000
        Stop adaptive once we have reached this number of epochs
    nframes : int, default=0
        Stop adaptive once we have simulated this number of aggregate simulation frames.
    inputpath : str, default='input'
        The directory used to store input folders
    generatorspath : str, default='generators'
        The directory containing the generators
    dryrun : boolean, default=False
        A dry run means that the adaptive will retrieve and generate a new epoch but not submit the simulations
    updateperiod : float, default=0
        When set to a value other than 0, the adaptive will run synchronously every `updateperiod` seconds
    coorname : str, default='input.coor'
        Name of the file containing the starting coordinates for the new simulations
    lock : bool, default=False
        Lock the folder while adaptive is ongoing
    datapath : str, default='data'
        The directory in which the completed simulations are stored
    filter : bool, default=True
        Enable or disable filtering of trajectories.
    filtersel : str, default='not water'
        Atom selection string for filtering.
        See more `here <http://www.ks.uiuc.edu/Research/vmd/vmd-1.9.2/ug/node89.html>`__
    filteredpath : str, default='filtered'
        The directory in which the filtered simulations will be stored
    projection : :class:`Projection <moleculekit.projections.projection.Projection>` object, default=None
        A Projection class object or a list of objects which will be used to project the simulation data before constructing a Markov model
    truncation : str, default=None
        Method for truncating the prob distribution (None, 'cumsum', 'statecut'
    statetype : ('micro', 'cluster', 'macro'), str, default='micro'
        What states (cluster, micro, macro) to use for calculations.
    macronum : int, default=8
        The number of macrostates to produce
    skip : int, default=1
        Allows skipping of simulation frames to reduce data. i.e. skip=3 will only keep every third frame
    lag : int, default=1
        The lagtime used to create the Markov model. Units are in frames.
    clustmethod : :class:`ClusterMixin <sklearn.base.ClusterMixin>` class, default=<class 'htmd.clustering.kcenters.KCenter'>
        Clustering algorithm used to cluster the contacts or distances
    method : str, default='1/Mc'
        Criteria used for choosing from which state to respawn from
    ticalag : int, default=20
        Lagtime to use for TICA in frames. When using `skip` remember to change this accordinly.
    ticadim : int, default=3
        Number of TICA dimensions to use. When set to 0 it disables TICA
    contactsym : str, default=None
        Contact symmetry
    save : bool, default=False
        Save the model generated
    goalfunction : function, default=None
        This function will be used to convert the goal-projected simulation data to a ranking whichcan be used for the directed component of FAST.
    ucscale : float, default=0.5
        Scaling factor for undirected component. Directed component scaling automatically calculated as (1-uscale)
    nosampledc : bool, default=False
        Spawn only from top DC conformations without sampling
    autoscale : bool, default=False
        Automatically scales exploration and exploitation ratios depending on how stuck the adaptive is at a given goal score.
    epsilon : float, default=0.5
        The epsilon for epsilon-greedy selection, where epsilon is the exploration ratio.

    Examples
    --------
    >>> crystalSS = MetricSecondaryStructure().project(Molecule('crystal.pdb'))[0]
    >>>
    >>> # First argument of a goal function always has to be a Molecule object
    >>> def ssGoal(mol):
    >>>     proj = MetricSecondaryStructure().project(mol)
    >>>     ss_score = np.sum(proj == crystalSS, axis=1) / proj.shape[1]  # How many predicted SS match
    >>>     return ss_score
    >>>
    >>> ag = AdaptiveGoal()
    >>> ag.generatorspath = '../generators/'
    >>> ag.nmin = 2
    >>> ag.nmax = 3
    >>> ag.projection = [MetricDistance('name CA', 'resname MOL', periodic='selections'), MetricDihedral()]
    >>> ag.goalfunction = ssGoal
    >>> ag.app = LocalGPUQueue()
    >>> ag.run()
    >>>
    >>> # Or alternatively if we have a multi-argument goal function
    >>> def ssGoalAlt(mol, ss):
    >>>     proj = MetricSecondaryStructure().project(mol)
    >>>     ss_score = np.sum(proj == ss, axis=1) / proj.shape[1]
    >>>     return ss_score
    >>> from joblib import delayed
    >>> ag.goalfunction = delayed(ssGoalAlt)(crystalSS)
    >>> ag.app = LocalGPUQueue()
    >>> ag.run()
    """

    def __init__(self):
        super().__init__()
        self._arg(
            "epsilon",
            "float",
            "The epsilon for epsilon greedy, epsilon being the exploration ratio",
            0.5,
            val.Number(float, "ANY"),
        )

    def _algorithm(self):
        sims = self._getSimlist()

        if self.nosampledc:
            print("Spawning only from top DC conformations without sampling")
            goaldata = self._getGoalData(sims)
            if not self._checkNFrames(goaldata):
                return False
            datconcat = np.concatenate(goaldata.dat).flatten()
            sortedabs = np.argsort(datconcat)[::-1]
            if self._debug:
                np.save("debug.npy", sortedabs[: self.nmax - self._running])
                return True
            self._writeInputs(goaldata.abs2sim(sortedabs[: self.nmax - self._running]))
            return True

        data = self._getData(sims)
        if self.save:
            if not path.exists("saveddata"):
                os.makedirs("saveddata")
            np.savetxt(
                path.join("saveddata", "e{}_report.npy".format(self._getEpoch())),
                [self._getEpoch(), data.numFrames, len(data.dat)],
            )

        if not self._checkNFrames(data):
            return False
        goaldata = self._getGoalData(
            data.simlist
        )  # Using the simlist of data in case some trajectories were dropped
        if len(data.simlist) != len(goaldata.simlist):
            logger.warning(
                "The goal function was not able to project all trajectories that the MSM projection could."
                "Check for possible errors in the goal function."
            )
        data.dropTraj(
            keepsims=goaldata.simlist
        )  # Ensuring that I use the intersection of projected simulations
        self._createMSM(data)

        model = self._model
        data = self._model.data

        # Undirected component
        uc = model.data.N  # Lower counts should give higher score
        if self.statetype == "micro":
            uc = uc[model.cluster_ofmicro]
        if self.statetype == "macro":
            uc = macroAccumulate(model, uc[model.cluster_ofmicro])
        uc = 1 / uc

        # Calculating the directed component
        dcmeans = dcstds = None
        if self.statetype == "micro":
            dcmeans, dcstds = self._calculateDirectedComponent(
                goaldata, model.data.St, model.micro_ofcluster
            )
        elif self.statetype == "macro":
            # TODO: Should we weigh by equilibrium population?
            dcmeans, dcstds = self._calculateDirectedComponent(
                goaldata, model.data.St, model.macro_ofcluster
            )

        ucunscaled = uc.copy()
        dcunscaled = dcmeans.copy()
        dc = dcmeans / np.sum(dcmeans)
        uc = uc / np.sum(uc)

        N = self.nmax - self._running
        Nuc, Ndc = np.random.multinomial(N, [self.epsilon, 1 - self.epsilon])
        rewardDC = self._truncate(uc * dc, Ndc)
        rewardUC = self._truncate(uc, Nuc)

        relFramesDC, spawncountsDC, truncprobDC = self._getSpawnFrames(
            rewardDC, self._model, data, Ndc
        )
        relFramesUC, spawncountsUC, truncprobUC = self._getSpawnFrames(
            rewardUC, self._model, data, Nuc
        )
        relFrames = np.concatenate(
            (np.concatenate(relFramesDC), np.concatenate(relFramesUC))
        )

        if self.save:
            if not path.exists("saveddata"):
                os.makedirs("saveddata")
            epoch = self._getEpoch()
            tosave = {
                "ucunscaled": -ucunscaled,
                "dcunscaled": dcunscaled,
                "uc": uc,
                "dc": dc,
                "epsilon": self.epsilon,
                "spawncountsDC": spawncountsDC,
                "spawncountsUC": spawncountsUC,
                "truncprobDC": truncprobDC,
                "truncprobUC": truncprobUC,
                "relFramesDC": relFramesDC,
                "relFramesUC": relFramesUC,
                "dcmeans": dcmeans,
                "dcstds": dcstds,
                "rewardDC": rewardDC,
                "rewardUC": rewardUC,
            }
            np.save(path.join("saveddata", "e{}_goalreport.npy".format(epoch)), tosave)
            np.save(
                path.join("saveddata", "e{}_spawnframes.npy".format(epoch)), relFrames
            )
            goaldata.save(path.join("saveddata", "e{}_goaldata.dat".format(epoch)))

        if self._debug:
            np.save("debug.npy", relFrames)
            return True

        self._writeInputs(data.rel2sim(relFrames))
        return True


if __name__ == "__main__":
    import htmd.home
    from moleculekit.projections.metricdistance import MetricDistance
    from jobqueues.localqueue import LocalGPUQueue
    from moleculekit.molecule import Molecule
    import shutil
    from htmd.util import tempname
    from joblib import delayed
    from htmd.home import home

    def rmsdgoal(proj):
        return -proj  # Lower RMSDs should give higher score

    tmpdir = tempname()
    shutil.copytree(htmd.home.home() + "/data/adaptive/", tmpdir)
    os.chdir(tmpdir)
    md = AdaptiveGoal()
    md.dryrun = True
    md.nmin = 1
    md.nmax = 2
    md.nepochs = 3
    md.ticalag = 2
    md.ticadim = 3
    md.updateperiod = 5
    md.projection = MetricDistance(
        "protein and name CA", "resname BEN and noh", periodic="selections"
    )
    # md.goalprojection = MetricRmsd(Molecule(htmd.home() + '/data/adaptive/generators/1/structure.pdb'),
    #                               'protein and name CA')
    md.goalfunction = rmsdgoal
    # md.app = LocalGPUQueue()
    # md.run()

    # Some real testing now
    from moleculekit.projections.metricsecondarystructure import (
        MetricSecondaryStructure,
    )
    from moleculekit.projections.metricdistance import MetricSelfDistance

    os.chdir(path.join(home(), "data", "test-adaptive"))

    goalProjectionDict = {
        "ss": MetricSecondaryStructure(),
        "contacts": MetricSelfDistance(
            "protein and name CA", metric="contacts", threshold=10
        ),
        "ss_contacts": [
            MetricSecondaryStructure(),
            MetricSelfDistance("protein and name CA", metric="contacts", threshold=10),
        ],
    }

    def getLongContacts(crystal, long=8):
        crystalMap = MetricSelfDistance(
            "protein and name CA", metric="contacts", threshold=10, pbc=False
        ).getMapping(crystal)
        indexes = np.vstack(crystalMap.atomIndexes.values)
        return crystal.resid[indexes[:, 1]] - crystal.resid[indexes[:, 0]] > long

    def getCrystalSS(crystal):
        return MetricSecondaryStructure().project(crystal)[0].flatten()

    def getCrystalCO(crystal):
        crystalCO = (
            MetricSelfDistance(
                "protein and name CA", metric="contacts", threshold=10, pbc=False
            )
            .project(crystal)
            .flatten()
        )
        longCO = getLongContacts(crystal)
        return crystalCO & longCO

    def ssContactGoal(mol, crystal, project=True, crystalSS=None, crystalCO=None):
        if crystalSS is None:
            crystalSS = getCrystalSS(crystal)
        if crystalCO is None:
            crystalCO = getCrystalCO(crystal)

        if project:
            projss = goalProjectionDict["ss"].copy().project(mol)
            projco = goalProjectionDict["contacts"].copy().project(mol)
        else:
            projss = mol[:, : len(crystalSS)]
            projco = mol[:, len(crystalSS) :]

        if len(crystalCO) != projco.shape[1]:
            raise RuntimeError(
                "Different lengths between crystal {} and traj {} contacts for fileloc {}".format(
                    len(crystalCO), projco.shape[1], mol.fileloc
                )
            )
        if len(crystalSS) != projss.shape[1]:
            raise RuntimeError(
                "Different lengths between crystal {} and traj {} SS for fileloc {}".format(
                    len(crystalSS), projss.shape[1], mol.fileloc
                )
            )

        ss_score = np.sum(projss == crystalSS, axis=1) / projss.shape[1]
        co_score = np.sum(projco[:, crystalCO] == 1, axis=1) / np.sum(
            crystalCO
        )  # Predicted conts are True?
        return 0.6 * ss_score + 0.4 * co_score

    refmol = Molecule("ntl9_2hbb.pdb")
    crystalSS = getCrystalSS(refmol)
    crystalCO = getCrystalCO(refmol)

    # TODO: Make this test work. Seems to ignore the random seed
    np.random.seed(10)
    ad = AdaptiveGoalEG()
    ad.app = LocalGPUQueue()
    ad.nmin = 10
    ad.nmax = 20
    ad.nepochs = 999999
    # ad.nframes = nframes['ntl9'] test that as well
    ad.generatorspath = "../../generators/"
    ad.projection = MetricSelfDistance("protein and name CA")
    ad.goalfunction = delayed(ssContactGoal)(refmol, True, crystalSS, crystalCO)
    ad.statetype = "micro"
    ad.truncation = "cumsum"
    ad.epsilon = 0.5
    ad._debug = True
    ad.run()
    # assert np.array_equal(np.load('debug.npy'), np.load('ref_adaptivegoaleg.npy'))
