# **************************************************************************
# * Authors: Rémi Vuillemot             (remi.vuillemot@upmc.fr)
# *
# * IMPMC, UPMC Sorbonne University
# *
# * This program is free software; you can redistribute it and/or modify
# * it under the terms of the GNU General Public License as published by
# * the Free Software Foundation; either version 2 of the License, or
# * (at your option) any later version.
# *
# * This program is distributed in the hope that it will be useful,
# * but WITHOUT ANY WARRANTY; without even the implied warranty of
# * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# * GNU General Public License for more details.
# *
# * You should have received a copy of the GNU General Public License
# * along with this program; if not, write to the Free Software
# * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
# * 02111-1307  USA
# *
# *  All comments concerning this program package may be sent to the
# *  e-mail address 'scipion@cnb.csic.es'
# **************************************************************************

from pyworkflow.viewer import (ProtocolViewer, DESKTOP_TKINTER, WEB_DJANGO)
from continuousflex.protocols.protocol_genesis import *
from continuousflex.protocols.utilities.genesis_utilities import *
from .plotter import FlexPlotter
from pwem.viewers import VmdView, ChimeraView
from pyworkflow.utils import getListFromRangeString
import numpy as np
import os
import glob
from matplotlib.pyplot import cm


class GenesisViewer(ProtocolViewer):
    """ Visualization of results from the GENESIS protocol
    """
    _label = 'viewer genesis'
    _targets = [ProtGenesis]
    _environments = [DESKTOP_TKINTER, WEB_DJANGO]

    def _defineParams(self, form):
        form.addSection(label='Visualization')

        if self.protocol.getNumberOfSimulation() >1:
            form.addParam('fitRange', params.NumericRangeParam,
                          label="Simulation selection",
                          default="1-%i"%self.protocol.getNumberOfSimulation(),
                          important = True,
                          help=' Select the simulation to display. Examples:'
                               ' "1,3-5" -> [1,3,4,5]'
                               ' "1, 2, 4" -> [1,2,4]')
        if self.protocol.simulationType.get() == SIMULATION_REMD\
                or self.protocol.simulationType.get() == SIMULATION_RENMMD:
            form.addParam('replicaRange', params.NumericRangeParam,
                          label="Replica selection",
                          default="1-%i"%self.protocol.nreplica.get(),
                          help=' Select the replicas to display. Examples:'
                               ' "1,3-5" -> [1,3,4,5]'
                               ' "1, 2, 4" -> [1,2,4]')

        form.addParam('compareToPDB', params.BooleanParam, default=False,
                        label="Compare to external PDB",
                        help='TODO')

        group = form.addGroup('External PDB',condition= "compareToPDB")
        group.addParam('targetPDB', params.PathParam, default=None,
                        label="Target PDB (s)", important=True,
                        help=' Target PDBs to compute RMSD against. Atom mathcing is performed between '
                             ' the output PDBs and the target PDBs. Use the file pattern as file location with /*.pdb',
                       condition= "compareToPDB")
        group.addParam('referencePDB', params.PathParam, default="",
                        label="Intial PDB",
                        help='Atom matching will ignore the output PDB and will use the initial PDB instead.',
                        expertLevel=params.LEVEL_ADVANCED,condition= "compareToPDB")

        group.addParam('alignTarget', params.BooleanParam, default=False,
                        label="Align Target PDB",
                        help='TODO',condition= "compareToPDB")

        group = form.addGroup('Chimera 3D view')
        group.addParam('displayChimera', params.LabelParam,
                      label='Display results in Chimera',
                      help='Show initial and final structures in Chimera')

        group = form.addGroup('VMD trajectory view')
        group.addParam('displayTrajVMD', params.LabelParam,
                      label='Display trajecory in VMD',
                      help='TODO')

        group = form.addGroup('Energy Analysis')
        group.addParam('displayEnergy', params.LabelParam,
                      label='Display Potential Energy',
                      help='Show time series of the potentials used in MD simulation/Minimization')

        group = form.addGroup('RMSD analysis', condition= "compareToPDB")
        group.addParam('displayRMSDts', params.LabelParam,
                      label='Display RMSD time series',
                      help='TODO',condition= "compareToPDB")

        group.addParam('displayRMSD', params.LabelParam,
                      label='Display final RMSD',
                      help='TODO',condition= "compareToPDB")


        if self.protocol.EMfitChoice.get() != EMFIT_NONE:
            group = form.addGroup('Cryo EM fitting')
            group.addParam('displayCC', params.LabelParam,
                          label='Display Correlation Coefficient',
                          help='Show C.C. time series during the simulation')

    def _getVisualizeDict(self):
        return {
            'displayChimera': self._plotChimera,
            'displayEnergy': self._plotEnergy,
            'displayCC': self._plotCC,
            'displayRMSDts': self._plotRMSDts,
            'displayRMSD': self._plotRMSD,
            'displayTrajVMD': self._plotTrajVMD,
                }

    def _plotChimera(self, paramName):
        tmpChimeraFile = self.protocol._getExtraPath("chimera.cxc")
        index = self.getSimulationList()[0]

        with open(tmpChimeraFile, "w") as f:
            f.write("open %s.pdb \n"% os.path.abspath(self.protocol.getInputPDBprefix(index)))
            f.write("color #1 magenta \n" )
            count = 1

            outpdbfile = self.getOutputPrefixAll(index)[0] +".pdb"
            if os.path.exists(outpdbfile):
                if os.path.getsize(outpdbfile) != 0:
                    f.write("open %s \n"% os.path.abspath(outpdbfile))
                    count+=1
                    f.write("color #%s lime \n"%count)

            if self.compareToPDB.get():
                f.write("open %s \n" % os.path.abspath(self.getTargetPDB(index)))
                count+=1
                f.write("color #%s orange \n"%count)

            if self.protocol.EMfitChoice.get() == EMFIT_VOLUMES:
                f.write("open %s.mrc \n" % os.path.abspath(self.protocol.getInputEMprefix(index)))
                count+=1
                f.write("volume #%i transparency 0.5\n"%count)

            f.write("hide atoms \n")
            f.write("show cartoons \n")
            f.write("lighting soft \n")

        cv = ChimeraView(tmpChimeraFile)
        cv.show()

    def _plotTrajVMD(self, paramName):
        tmpVmdFile = self.protocol._getExtraPath("vmd.tcl")
        index = self.getSimulationList()[0]
        with open(tmpVmdFile, "w") as f:
            f.write("mol new %s.pdb waitfor all\n" % self.protocol.getInputPDBprefix(index))
            f.write("mol addfile %s.dcd waitfor all\n" % self.getOutputPrefixAll(index)[0])

            if self.protocol.forcefield.get() == FORCEFIELD_CHARMM:
                f.write("mol modstyle 0 0 NewCartoon \n")
            else:
                f.write("mol modstyle 0 0 Tube \n")
            f.write("mol modcolor 0 0 Molecule\n")

            if self.protocol.EMfitChoice.get() == EMFIT_VOLUMES:
                f.write("mol addfile %s.mrc waitfor all\n" % self.protocol.getInputEMprefix(index))
                f.write("mol addrep 0 \n")
                f.write("mol modstyle 1 0 Isosurface 0.5 0 0 0 1 1  \n")
                f.write("mol modmaterial 1 0 Transparent \n")

            if self.compareToPDB.get():
                targetFile = self.getTargetPDB(index)
                f.write("set nf [molinfo top get numframes]\n")
                f.write("mol new %s waitfor all\n" %targetFile)
                f.write( "for {set i 1 } {$i < $nf} {incr i} {\n")
                f.write( "animate dup frame 0 1\n")
                f.write("}\n")
                if self.protocol.forcefield.get() == FORCEFIELD_CHARMM:
                    f.write("mol modstyle 0 1 NewCartoon \n")
                else:
                    f.write("mol modstyle 0 1 Tube \n")
                f.write("mol modcolor 0 1 Molecule\n")
            f.write("animate style Loop\n")
            f.write("display projection Orthographic\n")

        vmdviewer = VmdView(" -e " + tmpVmdFile)
        vmdviewer.show()

    def _plotEnergy(self, paramName):
        self._plotEnergyTotal()
        self._plotEnergyDetail()

    def _plotEnergyTotal(self):
        ene_default = ["TOTAL_ENE", "POTENTIAL_ENE", "KINETIC_ENE"]

        ene = {}
        for i in self.getSimulationList():
            outputPrefix = self.getOutputPrefixAll(i)
            for j in outputPrefix:
                log_file = readLogFile(j + ".log")
                for e in ene_default:
                    if e in log_file:
                        if e in ene :
                            ene[e].append(log_file[e])
                        else:
                            ene[e] = [log_file[e]]
        enelist =[]
        labels=[]
        for e in ene :
            labels.append(e)
            enelist.append(ene[e])


        self.genesisPlotter(title="Energy (kcal/mol)", data=enelist, ndata=len(enelist),
                            nrep=len(enelist[0]), labels=labels)

    def _plotEnergyDetail(self):
        ene_default = ["BOND", "ANGLE", "UREY-BRADLEY", "DIHEDRAL", "IMPROPER", "CMAP", "VDWAALS", "ELECT", "NATIVE_CONTACT",
               "NON-NATIVE_CONT", "RESTRAINT_TOTAL"]

        ene = {}
        for i in self.getSimulationList():
            outputPrefix = self.getOutputPrefixAll(i)
            for j in outputPrefix:
                log_file = readLogFile(j+".log")
                for e in ene_default:
                    if e in log_file:
                        if e in ene :
                            ene[e].append(log_file[e])
                        else:
                            ene[e] = [log_file[e]]
        enelist =[]
        labels=[]
        for e in ene :
            labels.append(e)
            enelist.append(ene[e])

        self.genesisPlotter(title="Energy (kcal/mol)", data=enelist, ndata=len(enelist),
                            nrep=len(enelist[0]), labels=labels)

    def _plotCC(self, paramName):
        # Get CC list
        cc = []
        labels=[]
        simlist = self.getSimulationList()
        for i in simlist:
            outputPrefix = self.getOutputPrefixAll(i)
            cc_rep = []
            if len(simlist) == 1:
                labels.append("CC")
            else:
                labels.append("CC %s" % str(i + 1))
            for j in outputPrefix:
                log_file = readLogFile(j + ".log")
                if 'RESTR_CVS001' in log_file:
                    cc_rep.append(log_file['RESTR_CVS001'])
                else:
                    raise RuntimeError("CC not present in the log file")
            cc.append(cc_rep)

        self.genesisPlotter(title="CC", data=cc, ndata=len(simlist),
                            nrep=len(self.getOutputPrefixAll()), labels=labels)


    def genesisPlotter(self, title, data, ndata, nrep, labels):
        plotter = FlexPlotter()
        ax = plotter.createSubPlot(title, "", title)
        nmax= 10
        colors = [cm.get_cmap("tab10", 10)(i) for i in range(nmax) ]

        for i in range(ndata):
            if ndata <= nmax and nrep > 1:
                try:
                    meandata = np.mean(data[i], axis=0)
                    stddata = np.std(data[i], axis=0)
                    x = self.getTimePeriod(len(meandata))
                    ax.errorbar(x=x, y=meandata, yerr=stddata,
                                capthick=1.7, capsize=5, elinewidth=1.7, color=colors[i] if ndata!= 1 else "black",
                                errorevery=np.max([len(meandata) // 10, 1]), label="%s Average" % labels[i])
                except TypeError:
                    x = self.getTimePeriod(len(data[i][0]))
                    ax.plot(x, data[i][0], color=colors[i], label=labels[i])

            for j in range(nrep):
                x = self.getTimePeriod(len(data[i][j]))
                if 1 < nrep <= nmax:
                    if ndata == 1 :
                        ax.plot(x, data[i][j], color= colors[j], alpha=0.5, label="#%i"%(j+1))
                    # else:
                    #     ax.plot(x, data[i][j], color= colors[i], alpha=0.5)
                if nrep == 1 and ndata <= 10:
                    ax.plot(x, data[i][j], color= colors[i],label=labels[i])
        if ndata > nmax :
            try:
                meandata = np.mean(data, axis=(0,1))
                stddata = np.std(data, axis=(0,1))
                x = self.getTimePeriod(len(meandata))
                ax.errorbar(x=x, y=meandata, yerr=stddata,
                            capthick=1.7, capsize=5, elinewidth=1.7, color="black",
                            errorevery=np.max([len(meandata) // 10, 1]), label="Global Average")
            except TypeError:
                x = self.getTimePeriod(len(data[0][0]))
                ax.plot(x, data[0][0], color=colors[0], label=labels[0])
        xlabel = "Number of iterations" if self.protocol.simulationType.get() == SIMULATION_MIN else "Time (ps)"
        ax.set_xlabel(xlabel)
        plotter.legend()
        plotter.show()

    def _plotRMSDts(self, paramName):

        # Get matching atoms
        if self.referencePDB.get() != "":
            ref_pdb = ContinuousFlexPDBHandler(self.referencePDB.get())
        else:
            ref_pdb = ContinuousFlexPDBHandler(self.protocol.getInputPDBprefix()+".pdb")
        target_pdb = ContinuousFlexPDBHandler(self.getTargetPDB())
        idx_matchin_atoms = ref_pdb.matchPDBatoms(reference_pdb=target_pdb, ca_only=True)

        # Get RMSD list
        rmsd = []
        labels=[]
        simlist = self.getSimulationList()
        for i in simlist:
            outputPrefix = self.getOutputPrefixAll(i)
            if len(simlist) == 1:
                labels.append("RMSD")
            else:
                labels.append("RMSD %s"%str(i+1))
            rmsd_rep=[]
            for outprf in outputPrefix:
                rmsd_curr = []

                inputPDB = ContinuousFlexPDBHandler(self.protocol.getInputPDBprefix(i)+".pdb")
                targetPDB = ContinuousFlexPDBHandler(self.getTargetPDB(i))
                rmsd_curr.append(inputPDB.getRMSD(reference_pdb=targetPDB, align=self.alignTarget.get(), idx_matching_atoms=idx_matchin_atoms))
                coord_arr = dcd2numpyArr(outprf + ".dcd")
                for i in range(len(coord_arr)):
                    inputPDB.coords[:, :] = coord_arr[i]
                    rmsd_curr.append(inputPDB.getRMSD(reference_pdb=targetPDB, align=self.alignTarget.get(), idx_matching_atoms=idx_matchin_atoms))

                rmsd_rep.append(rmsd_curr)
            rmsd.append(rmsd_rep)

        self.genesisPlotter(title="RMSD ($\AA$)", data=rmsd, ndata=len(simlist),
                            nrep=len(self.getOutputPrefixAll()), labels=labels)


    def _plotRMSD(self, paramName):
        plotter = FlexPlotter()
        ax = plotter.createSubPlot("RMSD ($\AA$)", "# Simulation", "RMSD ($\AA$)")

        # Get RMSD list
        initial_mols = []
        final_mols = []
        target_mols = []
        for i in self.getSimulationList():
            inputPDB = self.protocol.getInputPDBprefix(i)+".pdb"
            targetPDB = self.getTargetPDB(i)
            outputPrefs = self.getOutputPrefixAll(i)
            target_mols.append(ContinuousFlexPDBHandler(targetPDB))
            initial_mols.append(ContinuousFlexPDBHandler(inputPDB))
            for outputPrefix in outputPrefs:
                outputPDB = outputPrefix +".pdb"
                final_mols.append(ContinuousFlexPDBHandler(outputPDB))

        if self.referencePDB.get() != "":
            ref_mol = ContinuousFlexPDBHandler(self.referencePDB.get())
        else:
            ref_mol = initial_mols[0]
        idx_match = ref_mol.matchPDBatoms(reference_pdb=target_mols[0],ca_only=True)
        rmsdi=[]
        rmsdf=[]
        for i in range(len(self.getSimulationList())):
            for j in range(len(outputPrefs)):
                rmsdi.append(initial_mols[i].getRMSD(reference_pdb=target_mols[i],
                                                     idx_matching_atoms=idx_match,
                                                     align=self.alignTarget.get()))
                rmsdf.append(final_mols[i*len(outputPrefs) + j].getRMSD(reference_pdb=target_mols[i],
                                                    idx_matching_atoms=idx_match, align=self.alignTarget.get()))

        ax.plot(rmsdf, "o", color="tab:blue", label="Final RMSD", markeredgecolor='black')
        ax.plot(rmsdi, "o", color="tab:green", label="Initial RMSD", markeredgecolor='black')
        plotter.legend()
        plotter.show()

    def getSimulationList(self):
        if self.protocol.getNumberOfSimulation() > 1:
            return np.array(getListFromRangeString(self.fitRange.get())) -1
        else:
            return np.array([0])

    def getTargetPDB(self, index=0):
        targetPDBlist = [f for f in glob.glob(self.targetPDB.get())]
        targetPDBlist.sort()
        if index < len(targetPDBlist):
            return targetPDBlist[index]
        else:
            return targetPDBlist[0]

    def getOutputPrefixAll(self, index=0):
        outPrf = np.array(self.protocol.getOutputPrefixAll(index))
        if self.protocol.simulationType.get() == SIMULATION_REMD \
                or self.protocol.simulationType.get() == SIMULATION_RENMMD:
            return outPrf[np.array(getListFromRangeString(self.replicaRange.get())) - 1]
        else:
            return outPrf

    def getTimePeriod(self, length):
        if self.protocol.simulationType.get() == SIMULATION_MIN:
            timestep = 1.0
        else:
            timestep = float(self.protocol.time_step.get())
        eneperiod = int(self.protocol.eneout_period.get())

        return np.arange(length) * eneperiod * timestep

