"""Identifying titrable sites and integrity checking functions
"""
from config import Config
from constants import *
from formats import read_pdb_line, read_gro_line, new_pdb_line
from ffconverter import main_chains
from copy import copy


def identify_tit_sites(molecules, instanciate_sites=True):
    def SitesFileLine(resnumb, resname):
        for res in REGULARTITRATINGRES:
            if res[0:2] == resname[0:2]:
                resname = res
        if resname in TITRABLETAUTOMERS:
            ntautomers = TITRABLETAUTOMERS[resname]
        else:
            for res in REGULARTITRATINGRES:
                if res[0:2] == resname[0:2]:
                    ntautomers = TITRABLETAUTOMERS[res]
        # In debug mode, a .sites file is created
        text = "{0} ".format(resnumb)
        for i in range(ntautomers):
            text += "{0}tau{1} ".format(resname, i + 1)

        if res in ("NTR", "CTR"):
            resnumb += TERMINAL_OFFSET

        if instanciate_sites:
            instanciate_site(resnumb, resname, ntautomers)
        return text + "\n"

    def instanciate_site(resnumb, resname, ntautomers):
        sID = molecule.addSite(resnumb)
        molecule.addTautomers(sID, ntautomers, resname)

    def add2chain(chain, chain_res, resnumb, resname):
        chain_res[chain][resnumb] = resname

    sites_file = ""
    sites = {chain: [] for chain in molecules.keys()}
    chain_res = {chain: {} for chain in molecules.keys()}
    last_res = None
    with open(Config.pypka_params["f_in"]) as f:
        nline = 0
        resnumb = None
        resname = None
        for line in f:
            nline += 1
            if "ATOM " == line[0:5]:
                (aname, anumb, resname, chain, resnumb, x, y, z) = read_pdb_line(line)

                if line[26] != " ":
                    continue
                # if chain == ' ':
                #    chain = 'A'
                last_res = resnumb
                last_chain = chain
                chain_sites = []

                if chain in molecules.keys():
                    molecule = molecules[chain]
                    chain_sites = chain_res[chain]

                    if resname in PROTEIN_RESIDUES or resname in TITRABLERESIDUES:
                        if resnumb not in chain_sites:
                            if not chain_res[chain]:
                                sites_file += SitesFileLine(resnumb, "NTR")
                                add2chain(chain, chain_res, str(resnumb), "NTR")
                                if instanciate_sites:
                                    molecule.NTR = resnumb

                            if (
                                resname in TITRABLERESIDUES
                                and resname != "NTR"
                                and resname != "CTR"
                            ):
                                if (
                                    Config.pypka_params["ser_thr_titration"] == False
                                    and resname in ("SER", "THR")
                                ) or (
                                    Config.pypka_params["ffID"] == "CHARMM36m"
                                    and resname in ("SER", "THR")
                                ):
                                    continue
                                sites_file += SitesFileLine(resnumb, resname)
                                add2chain(chain, chain_res, resnumb, resname)

                        if "CTR" not in sites and aname in (
                            "CT",
                            "OT",
                            "OT1",
                            "OT2",
                            "O1",
                            "O2",
                            "OXT",
                        ):
                            sites_file += SitesFileLine(resnumb, "CTR")
                            add2chain(chain, chain_res, str(resnumb), "CTR")
                            if instanciate_sites:
                                molecule.CTR = resnumb

    if (
        "CTR" not in sites
        and aname in ("CT", "OT", "OT1", "OT2", "O1", "O2", "OXT")
        and molecule.sites == "all"
    ):
        sites_file += SitesFileLine(last_res, "CTR")
        add2chain(last_chain, chain_res, str(last_res), "CTR")
        if instanciate_sites:
            molecule = molecules[chain]
            molecule.CTR = resnumb

    if not chain_res and instanciate_sites:
        f_in = Config.pypka_params["f_in"]
        raise Exception("Not one titrable residue was found in {}".format(f_in))

    if instanciate_sites:
        for molecule in molecules.values():
            # Adding the reference tautomer to each site
            molecule.addReferenceTautomers()
            # Assigning a charge set to each tautomer
            molecule.addTautomersChargeSets()

    if Config.debug:
        with open("tmp.sites", "w") as f_new:
            f_new.write(sites_file)

    return chain_res


def check_sites_integrity(molecules, chains_res, useTMPpdb=False):
    """Identifies titrable residues and checks integrity of the residue blocks
    (excluding Hydrogens)
    """

    def check_site(prev_resname, cur_atoms, ter=None):
        def correctResName(resname):
            for res in REGULARTITRATINGRES:
                if res[0:2] == resname[0:2]:
                    return res
            return resname

        def makeSite(molecule, resnumb, resname, termini_resname=None):
            if resname in TITRABLETAUTOMERS:
                ntautomers = TITRABLETAUTOMERS[resname]
            else:
                for res in REGULARTITRATINGRES:
                    if res[0:2] == resname[0:2]:
                        ntautomers = TITRABLETAUTOMERS[res]
            sID = molecule.addSite(resnumb)
            molecule.addTautomers(
                sID, ntautomers, resname, termini_resname=termini_resname
            )
            # print('added', molecule.chain, resnumb, resname)

        prev_resname = correctResName(prev_resname)

        res_tits = True
        if ter:
            if not Config.pypka_params["ser_thr_titration"] and prev_resname in (
                "SER",
                "THR",
            ):
                res_tits = False
            else:
                res_tits = bool(prev_resname in TITRABLERESIDUES)

        res_atoms = copy(cur_atoms)
        (integrity_terminal, integrity_site) = check_integrity(
            prev_resname, res_atoms, ter=ter, site=res_tits
        )

        if integrity_terminal:
            ter_resnumb = prev_resnumb + TERMINAL_OFFSET
            makeSite(molecule, ter_resnumb, ter, termini_resname=prev_resname)
            if ter == "NTR":
                molecule.NTR = prev_resnumb
            elif ter == "CTR":
                molecule.CTR = prev_resnumb
        else:
            warning(molecule, prev_resnumb, ter, "")

        if prev_resnumb in sites:
            if integrity_site:
                makeSite(molecule, prev_resnumb, prev_resname)
            else:
                warning(molecule, prev_resnumb, prev_resname, cur_atoms)
        elif prev_resname == "CYS":  # dealing with a CYS that is not in sites
            if not integrity_site:
                warning(molecule, prev_resnumb, prev_resname, cur_atoms, mode="CYS")

    def warning(molecule, resnumb, resname, res_atoms, mode=None):
        if mode == "CYS" or resname == "CYS":
            CYS_atoms = ["N", "CA", "CB", "SG", "C", "O", "H"]
            if set(res_atoms).issubset(CYS_atoms) and set(CYS_atoms).issubset(
                res_atoms
            ):
                # no need to correct residue name
                warn = "{0} {1} is assumed to be participating " "in a SS-bond".format(
                    resnumb, resname
                )
                Config.log.report2log(warn)
                return
            CY0_atoms = ["N", "CA", "CB", "SG", "C", "O", "H", "HG1"]
            if set(res_atoms).issubset(CY0_atoms) and set(CY0_atoms).issubset(
                res_atoms
            ):
                molecule.correct_names[resnumb] = "CY0"
                return
            CY0_atoms = ["N", "CA", "CB", "SG", "C", "O", "H", "HG"]
            if set(res_atoms).issubset(CY0_atoms) and set(CY0_atoms).issubset(
                res_atoms
            ):
                molecule.correct_names[resnumb] = "CY0"
                molecule.correct_atoms[resnumb] = {"HG": "HG1"}
                return
            else:
                warn = "{0} {1} failed integrity check".format(resnumb, resname)
                Config.log.report2log(warn)
        elif resname not in TITRABLERESIDUES:
            return
        else:
            warn = "{0} {1} failed integrity check".format(resnumb, resname)
            Config.log.report2log(warn)

    if Config.pypka_params["f_in"] and not useTMPpdb:
        filename = Config.pypka_params["f_in"]
        filetype = "pdb"
    else:
        filename = "TMP.pdb"
        filetype = "pdb"

    resnumb = None
    cur_atoms = []
    prev_resnumb = None
    prev_resname = None
    last_chain = None
    with open(filename) as f:
        nline = 0
        f_lines = f.readlines()
        maxnlines = len(f_lines)
        for line in f_lines:
            resname = None
            nline += 1

            chain = None

            if "ATOM " == line[0:5]:
                (aname, anumb, resname, chain, resnumb, x, y, z) = read_pdb_line(line)

                if chain in molecules:
                    if not last_chain:
                        last_chain = chain
                    molecule = molecules[last_chain]
                    sites = chains_res[last_chain]
                    last_molecule = molecule
                    last_chain = chain

                if nline == maxnlines:
                    cur_atoms.append(aname)

            if line == "TER\n":
                resnumb += 1

            if (
                prev_resnumb != resnumb or nline == maxnlines
            ) and prev_resnumb is not None:
                if nline == maxnlines:
                    prev_resnumb = copy(resnumb)
                    resnumb = "None"

                if chain in molecules:
                    if (
                        prev_resname in TITRABLERESIDUES
                        or (prev_resnumb == molecule.NTR or resnumb == molecule.NTR)
                        or (prev_resnumb == molecule.CTR or resnumb == molecule.CTR)
                    ):
                        if prev_resnumb == molecule.NTR and resnumb != molecule.NTR:
                            check_site(prev_resname, cur_atoms, ter="NTR")
                            prev_resnumb = None
                        # Dealing with the last residue and CTR
                        elif prev_resnumb == molecule.CTR and resnumb != molecule.CTR:
                            check_site(prev_resname, cur_atoms, ter="CTR")
                        # Dealing with the previous residue
                        elif (
                            prev_resnumb is not None
                            and prev_resname in TITRABLERESIDUES
                        ):
                            if not (
                                not Config.pypka_params["ser_thr_titration"]
                                and prev_resname in ("SER", "THR")
                            ):
                                check_site(prev_resname, cur_atoms)

                    elif prev_resname == "ALA":
                        # TODO: check residue block integrity for other non titrating residues
                        pass
                elif (
                    last_molecule
                    and prev_resnumb == last_molecule.CTR
                    and resnumb != last_molecule.CTR
                ):
                    check_site(prev_resname, cur_atoms, ter="CTR")
                elif prev_resname in TITRABLERESIDUES and prev_resnumb is not None:
                    check_site(prev_resname, cur_atoms)

            # Dealing with the new residue
            if prev_resnumb != resnumb:
                cur_atoms = [aname]
                prev_resnumb = resnumb
                prev_resname = resname
            elif resnumb is not None:
                cur_atoms.append(aname)
                if prev_resname in ("NTR", "CTR") and prev_resname != resname:
                    prev_resname = resname

    for molecule in molecules.values():
        # Adding the reference tautomer to each site
        molecule.addReferenceTautomers()
        # Assigning a charge set to each tautomer
        molecule.addTautomersChargeSets()

    # TODO: report blocks that failed the check (in .log file with
    # numbering reference to stepwise scheme)
    # TODO: add lipid residues
    if Config.debug:
        print("exiting check_sites_integrity")


def check_integrity(resname, res_atoms, ter=False, site=True):
    def read_anames(filename):
        res_atoms_st = []
        with open(filename) as f:
            for line in f.readlines()[1:]:
                cols = line.strip().split()
                aname = cols[1]
                res_atoms_st.append(aname)
        return res_atoms_st

    def pop_atoms(anames1, anames2):
        trigger = False
        for aname in anames1:
            if aname not in anames2:
                trigger = True
                if Config.debug:
                    print((aname, "not in", resname))
            else:
                anames2.remove(aname)
        if trigger:
            return anames2, False
        else:
            return anames2, True

    if Config.debug:
        print("###### INTEGRITY CHECK ######")
        print(res_atoms)
    integrity = None
    integrity_ter = None

    if ter:
        st = "{0}/{1}/sts/{2}tau1.st".format(
            Config.pypka_params["script_dir"], Config.pypka_params["ffID"], ter
        )
        res_atoms_st = read_anames(st)
        res_atoms, integrity_ter = pop_atoms(res_atoms_st, res_atoms)

    if site:
        main_chain_atoms = main_chains[Config.pypka_params["ff_family"]]
        if ter == "NTR":
            main_chain = main_chain_atoms["NTR"]
        elif ter == "CTR":
            main_chain = main_chain_atoms["CTR"]
        else:
            main_chain = main_chain_atoms["main"]

        for aname in main_chain:
            if aname in res_atoms:
                res_atoms.remove(aname)
            elif not ter:
                integrity = False

        if Config.debug:
            print(("i", integrity))
        if integrity is not False:
            filename = "{0}/{1}/sts/{2}tau1.st".format(
                Config.pypka_params["script_dir"], Config.pypka_params["ffID"], resname
            )
            res_atoms_st = read_anames(filename)
            res_atoms, integrity = pop_atoms(res_atoms_st, res_atoms)
            if Config.debug:
                print((res_atoms_st, res_atoms))
                print(("i", integrity))

    if len(res_atoms) != 0 and integrity:
        raise Exception("Something is wrong")

    return integrity_ter, integrity


def make_delphi_inputfile(f_in, f_out, molecules):
    def getMaxCoords(coords, max_coords):
        x, y, z = coords
        max_x, max_y, max_z = max_coords
        if max_x < x:
            max_x = x
        if max_y < y:
            max_y = y
        if max_z < z:
            max_z = z
        return max_x, max_y, max_z

    def correct_termini(resnumb, resname, aname, ntr_res, ctr_res):
        if resnumb == ntr_res and aname in Config.pypka_params["NTR_atoms"]:
            resname = "NTR"
            resnumb += TERMINAL_OFFSET
        elif resnumb == ctr_res and aname in Config.pypka_params["CTR_atoms"]:
            resname = "CTR"
            resnumb += TERMINAL_OFFSET
            if aname == "C":
                aname = "CT"
        return resnumb, resname, aname

    def correct_res_names(molecule, resnumb, resname, aname):
        if resnumb in list(molecule.correct_names.keys()):
            resname = molecule.correct_names[resnumb]
        if (
            resnumb in list(molecule.correct_atoms.keys())
            and aname in molecule.correct_atoms[resnumb]
        ):
            aname = molecule.correct_atoms[resnumb][aname]

        return resnumb, resname, aname

    def assign_atoms(sites, resnumb, aname, site_Hs, site_positions):
        ref_tau_name = resname
        if resnumb in list(sites.keys()) and aname in list(
            sites[resnumb].getRefTautomer().charge_set.keys()
        ):
            # ( aname not in ('N', 'H', 'C', 'O', 'CA') or
            # (aname in ('N', 'H', 'C', 'O', 'CA') and resname == 'NTR')):
            # change res name to reference tautomer
            ref_tau_name = sites[resnumb].getRefTautomerName()

            # add atom to corresponding site
            sites[resnumb].addAtom(aname, anumb)

            if chain not in site_positions:
                site_positions[chain] = {}
                site_Hs[chain] = {}
            if resnumb not in site_positions[chain]:
                site_positions[chain][resnumb] = []
                site_Hs[chain][resnumb] = []

            if resnumb in site_positions[chain]:
                site_positions[chain][resnumb].append((x, y, z))
                if aname[0] == "H":
                    site_Hs[chain][resnumb].append((x, y, z))

        return site_Hs, ref_tau_name, site_positions

    new_pdb_content = ""
    site_positions = {}
    site_Hs = {}
    max_box = [0.0, 0.0, 0.0]
    aposition = -1
    with open(f_in) as f:
        for line in f:
            if line.startswith("ATOM"):
                aposition += 1
                (aname, anumb, resname, chain, resnumb, x, y, z) = read_pdb_line(line)

                max_box = getMaxCoords([x, y, z], max_box)

                if chain in molecules:
                    molecule = molecules[chain]
                    ntr_res = molecule.NTR
                    ctr_res = molecule.CTR
                    sites = molecule.sites

                    if (
                        resname == "HIS"
                        and aname == "HD1"
                        and resnumb not in sites.keys()
                    ):
                        aposition -= 1
                        continue

                    resnumb, resname, aname = correct_termini(
                        resnumb, resname, aname, ntr_res, ctr_res
                    )

                    resnumb, resname, aname = correct_res_names(
                        molecule, resnumb, resname, aname
                    )

                    titrable_res = False
                    if resnumb in sites.keys():
                        titrable_res = True

                    molecule.addAtom(aname, anumb, aposition, titrable_res)

                    (site_Hs, resname, site_positions) = assign_atoms(
                        sites, resnumb, aname, site_Hs, site_positions
                    )

                else:
                    if resname == "HIS" and aname == "HD1":
                        aposition -= 1
                        continue
                    resnumb, resname, aname = correct_res_names(
                        molecule, resnumb, resname, aname
                    )

                new_pdb_content += new_pdb_line(
                    aposition, aname, resname, resnumb, x, y, z
                )

            elif line.startswith("CRYST1"):
                parts = line.split()
                box = [float(i) for i in parts[1:4]]

    if box == [0.1, 0.1, 0.1]:
        box = max_box

    if Config.pypka_params["box"]:
        box = Config.pypka_params["box"]
    else:
        Config.pypka_params.setBox(box)

    if Config.delphi_params["pbc_dim"] == 2:
        Config.delphi_params.redefineScale()

    new_pdb_content += "TER\nENDMDL\n"
    with open(f_out, "w") as f_new:
        f_new.write(new_pdb_content)

    # TODO: check Terminal_offset has to be bigger than the total number of residues
    # TODO: delete terminal_offset and use another approach to distinguish between N- and C-ter
    # TODO: check size xy > config.cutoff * 2
    # if so, raise Exception, and ask to change cutoff value

    # TODO: check if pbc_dim -> set gsizes from pdb size xy and ignore perfil

    for chain in site_positions.keys():
        molecule = molecules[chain]
        for site in site_positions[chain]:
            if site in list(molecule.sites.keys()):
                pos_max = [-9999990, -999999, -999999]
                pos_min = [999999, 999999, 999999]
                focus_center = [0, 0, 0]
                for atom in site_positions[chain][site]:
                    for i in range(3):
                        if pos_max[i] < atom[i]:
                            pos_max[i] = atom[i]
                        if pos_min[i] > atom[i]:
                            pos_min[i] = atom[i]
                focus_center[0] = (pos_max[0] + pos_min[0]) / 2
                focus_center[1] = (pos_max[1] + pos_min[1]) / 2
                focus_center[2] = (pos_max[2] + pos_min[2]) / 2

                if Config.delphi_params["pbc_dim"] == 2:
                    molecule.sites[site].addCenter(
                        focus_center, boxsize=box[0], box_z=box[2]
                    )
                else:
                    molecule.sites[site].addCenter(focus_center)
                hx, hy, hz = 0, 0, 0
                nHs = len(site_Hs[chain][site])
                if nHs == 0:
                    sitename = molecule.sites[site].getName()
                    raise Exception(
                        "Site {1}{0} appears "
                        "to have no Hydrogen atoms".format(site, sitename)
                    )
                for h in site_Hs[chain][site]:
                    hx += h[0]
                    hy += h[1]
                    hz += h[2]
                hx /= nHs
                hy /= nHs
                hz /= nHs
                Hcenter = (hx, hy, hz)
                molecule.sites[site].addCenterH(Hcenter)


def get_chains_from_file(f_in):
    chain_list = []
    with open(f_in) as f:
        for line in f:
            if "ATOM " == line[0:5]:
                (aname, anumb, resname, chain, resnumb, x, y, z) = read_pdb_line(line)
                if chain not in chain_list:
                    chain_list.append(chain)

    return chain_list
