#!/usr/bin/env python


import argparse
import sys
import os
from pmx import ligand_alchemy
import glob


class PMXAssembleSystem:
    """
    Entrypoint for Icolos workflow manager, builds the system and merges topologies
    """

    def __init__(
        self,
        ligandPath: str,
        edges: str,
        workPath: str,
        ff: str = "amber99sb-star-ildn-mut.ff",
        water: str = "tip3p",
    ) -> None:
        super().__init__()
        self.ligandPath = ligandPath
        self.edges: list = edges.split(" ")
        self.workPath = workPath
        self.proteinPath = os.path.join(workPath, "input/protein/")
        self.protein = self._read_protein()
        self.ff = ff
        self.water = water

    def _read_protein(self):
        protein = {}
        # read protein folder
        protein["path"] = os.path.abspath(self.proteinPath)
        # get folder contents
        protein["posre"] = []
        protein["itp"] = []
        protein["mols"] = []  # mols to add to .top
        protein["str"] = ""
        for l in glob.glob("{0}/*".format(self.proteinPath)):
            fname = l.split("/")[-1]
            if ".itp" in fname:  # posre or top
                if "posre" in fname:
                    protein["posre"].append(os.path.abspath(l))
                else:
                    protein["itp"].append(os.path.abspath(l))
                    if fname.startswith("topol_"):
                        protein["mols"].append(fname[6:-4])
                    else:
                        protein["mols"].append(fname[:-4])
            if ".pdb" in fname:
                protein["str"] = fname

        protein["mols"].sort()
        return protein

    def assemble_systems(self):
        edges = self.edges
        for edge in edges:
            parts = edge.split("_")
            lig1 = parts[0]
            lig2 = parts[1]
            lig1path = "{0}/{1}".format(self.ligandPath, lig1)
            lig2path = "{0}/{1}".format(self.ligandPath, lig2)
            hybridStrTopPath = self._get_specific_path(
                workPath=self.workPath, edge=edge, bHybridStrTop=True
            )
            outLigPath = self._get_specific_path(
                workPath=self.workPath, edge=edge, wp="unbound"
            )
            outProtPath = self._get_specific_path(
                workPath=self.workPath, edge=edge, wp="bound"
            )

            # Ligand structure
            self._make_clean_pdb(
                "{0}/mergedA.pdb".format(hybridStrTopPath),
                "{0}/init.pdb".format(outLigPath),
            )

            # Ligand+Protein structure
            self._make_clean_pdb(
                "{0}/{1}".format(self.proteinPath, self.protein["str"]),
                "{0}/init.pdb".format(outProtPath),
            )
            self._make_clean_pdb(
                "{0}/mergedA.pdb".format(hybridStrTopPath),
                "{0}/init.pdb".format(outProtPath),
                bAppend=True,
            )

            # Ligand topology
            # ffitp; these are just the atomtypes for the forcefield.
            ffitpOut = "{0}/ffmerged.itp".format(hybridStrTopPath)
            ffitpIn1 = "{0}/ffMOL.itp".format(lig1path)
            ffitpIn2 = "{0}/ffMOL.itp".format(lig2path)
            ffitpIn3 = "{0}/ffmerged.itp".format(hybridStrTopPath)
            ligand_alchemy._merge_FF_files(
                ffitpOut, ffsIn=[ffitpIn1, ffitpIn2, ffitpIn3]
            )
            # top
            ligTopFname = "{0}/topol.top".format(outLigPath)
            ligFFitp = "{0}/ffmerged.itp".format(hybridStrTopPath)
            ligItp = "{0}/merged.itp".format(hybridStrTopPath)
            itps = [ligFFitp, ligItp]
            systemName = "ligand in water"
            self.create_top(fname=ligTopFname, itp=itps, systemName=systemName)

            # Ligand+Protein topology
            # top
            protTopFname = "{0}/topol.top".format(outProtPath)
            mols = []
            for m in self.protein["mols"]:
                mols.append([m, 1])
            mols.append(["MOL", 1])
            # require the protein to be parametrised with separate itp file, not included in the .top or it will be missed.
            itps.extend(self.protein["itp"])
            systemName = "protein and ligand in water"
            self.create_top(
                fname=protTopFname, itp=itps, mols=mols, systemName=systemName
            )

    def _get_specific_path(
        self,
        workPath=None,
        edge=None,
        bHybridStrTop=False,
        wp=None,
        state=None,
        r=None,
        sim=None,
    ):
        if edge == None:
            return workPath
        edgepath = "{0}/{1}".format(workPath, edge)

        if bHybridStrTop == True:
            hybridStrPath = "{0}/hybridStrTop".format(edgepath)
            return hybridStrPath

        if wp == None:
            return edgepath
        wppath = "{0}/{1}".format(edgepath, wp)

        if state == None:
            return wppath
        statepath = "{0}/{1}".format(wppath, state)

        if r == None:
            return statepath
        runpath = "{0}/run{1}".format(statepath, r)

        if sim == None:
            return runpath
        simpath = "{0}/{1}".format(runpath, sim)
        return simpath

    def create_top(
        self,
        fname="topol.top",
        itp=["merged.itp"],
        mols=[["MOL", 1]],
        systemName="simulation system",
        destination="",
        toppaths=[],
    ):

        fp = open(fname, "w")
        # ff itp
        fp.write('#include "%s/forcefield.itp"\n' % self.ff)
        # additional itp
        for i in itp:
            fp.write('#include "%s"\n' % i)
        # water itp
        fp.write('#include "%s/%s.itp"\n' % (self.ff, self.water))
        # ions
        fp.write('#include "%s/ions.itp"\n\n' % self.ff)
        # system
        fp.write("[ system ]\n")
        fp.write("{0}\n\n".format(systemName))
        # molecules
        fp.write("[ molecules ]\n")
        for mol in mols:
            fp.write("%s %s\n" % (mol[0], mol[1]))
        fp.close()

    def _clean_backup_files(self, path):
        toclean = glob.glob("{0}/*#".format(path))
        for clean in toclean:
            os.remove(clean)

    def _make_clean_pdb(self, fnameIn, fnameOut, bAppend=False):
        # read
        fp = open(fnameIn, "r")
        lines = fp.readlines()
        out = []
        for l in lines:
            if l.startswith("ATOM") or l.startswith("HETATM"):
                out.append(l)
        fp.close()

        # write
        if bAppend == True:
            fp = open(fnameOut, "a")
        else:
            fp = open(fnameOut, "w")
        for l in out:
            fp.write(l)
        fp.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Entrypoint for system preparation")
    parser.add_argument("-edges")
    parser.add_argument("-ligand_path")
    parser.add_argument("-workPath")

    args = parser.parse_args()

    assembler = PMXAssembleSystem(args.ligand_path, args.edges, args.workPath)
    assembler.assemble_systems()

    sys.exit(0)
