"""
hbond module is part of DuIvyTools library, which is a tool for analysis and 
visualization of GROMACS result files. This module is written by CharlesHahn.

This module is designed to process hbond information.

This module contains:
    - hbond function

This file is provided to you by GPLv3 license."""


import os
import sys
import logging
import argparse
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors

# from pathlib import Path
# dit_path = Path(os.path.dirname(__file__))
# sys.path.append(dit_path)
from DuIvyTools.XPM import XPM


def hbond(
    xpmfile: str = "",
    ndxfile: str = "",
    grofile: str = "",
    select: list = [],
    noshow: bool = False,
    figout: str = None,
    csv: str = None,
    hnf: str = None,
) -> None:
    """
    hbond: a function to figure out hbond information, occupancy and occupancy
    map, occupancy table.
    This function will read one gro file, one hbmap.xpm, one hbond.ndx, and
    return some hbond information.

    :parameters:
        gro_file: the gro file which saved coordinates to pare atom names
        hbmap.xpm: the hbond occupancy map file generated by 'gmx hbond'
        hbond.ndx: the index file generated by 'gmx hbond'

    :return:
        None
    """

    for suffix, file in zip([".xpm", ".ndx", ".gro"], [xpmfile, ndxfile, grofile]):
        if file == None:
            logging.error("You have to specify the {} file for input".format(suffix))
            sys.exit()
        if not os.path.exists(file):
            logging.error("no {} in current dirrectory. ".format(file))
            sys.exit()
        if file[-4:] != suffix:
            logging.error(
                "only accept {} file with suffix {}".format(suffix.strip("."), suffix)
            )
            sys.exit()

    ## parse ndx file to get index of hydrogen bonds
    donor_ndxs, hydrogen_ndxs, acceptor_ndxs = [], [], []
    with open(ndxfile, "r") as fo:
        lines = [line for line in fo.readlines() if line.strip()]
    lines.reverse()
    for line in lines:
        if "hbonds_" in line:
            break
        items = line.split()
        donor_ndxs.append(int(items[0].strip()))
        hydrogen_ndxs.append(int(items[1].strip()))
        acceptor_ndxs.append(int(items[2].strip()))
    donor_ndxs.reverse()
    hydrogen_ndxs.reverse()
    acceptor_ndxs.reverse()
    if len(donor_ndxs) != len(hydrogen_ndxs) != len(acceptor_ndxs):
        logging.error("wrong length in donor, hydrogen, acceptor indexs")
        sys.exit()

    ## read the gro file and parse atom names
    if hnf == None:
        hnf = (
            "d_resname(d_resnum)@d_atomname(d_atomnum)->h_atomname("
            + "h_atomnum)...a_resname(a_resnum)@a_atomname(a_atomnum)"
        )
    features = [
        "d_resname",
        "d_resnum",
        "d_atomname",
        "d_atomnum",
        "h_resname",
        "h_resnum",
        "h_atomname",
        "h_atomnum",
        "a_resname",
        "a_resnum",
        "a_atomname",
        "a_atomnum",
    ]
    for feature in features:
        if feature in hnf or hnf == "number":
            break
    else:
        logging.warning(
            "no key feature detected in your specified hbond name "
            + "format, use the default format"
        )
        hnf = (
            "d_resname(d_resnum)@d_atomname(d_atomnum)->h_atomname("
            + "h_atomnum)...a_resname(a_resnum)@a_atomname(a_atomnum)"
        )

    donor_names, hydrogen_names, acceptor_names = [], [], []
    with open(grofile, "r") as fo:
        lines = fo.readlines()[2:-1]
    for ind in donor_ndxs:
        line = lines[ind - 1]
        res_num = line[:5].strip()
        res_name = line[5:10].strip()
        atom_name = line[10:15].strip()
        atom_num = line[15:20].strip()
        # name = f"{res_name}({res_num})@{atom_name}({atom_num})"
        # donor_names.append(name)
        donor_names.append([res_name, res_num, atom_name, atom_num])
    for ind in hydrogen_ndxs:
        line = lines[ind - 1]
        res_num = line[:5].strip()
        res_name = line[5:10].strip()
        atom_name = line[10:15].strip()
        atom_num = line[15:20].strip()
        # name = f"{res_name}({res_num})@{atom_name}({atom_num})"
        # hydrogen_names.append(name)
        hydrogen_names.append([res_name, res_num, atom_name, atom_num])
    for ind in acceptor_ndxs:
        line = lines[ind - 1]
        res_num = line[:5].strip()
        res_name = line[5:10].strip()
        atom_name = line[10:15].strip()
        atom_num = line[15:20].strip()
        # name = f"{res_name}({res_num})@{atom_name}({atom_num})"
        # acceptor_names.append(name)
        acceptor_names.append([res_name, res_num, atom_name, atom_num])
    if (
        len(donor_names)
        != len(hydrogen_names)
        != len(acceptor_names)
        != len(acceptor_ndxs)
    ):
        logging.error("wrong length in donor, hydrogen, acceptor names")
        sys.exit()
    # hbond_names = [
    #     f'{donor_names[i]}->{hydrogen_names[i].split("@")[1]}...{acceptor_names[i]}'
    #     for i in range(len(donor_names))]
    hbond_names = []
    for i in range(len(donor_names)):
        hbond_name = hnf.replace("d_resname", donor_names[i][0])
        hbond_name = hbond_name.replace("d_resnum", donor_names[i][1])
        hbond_name = hbond_name.replace("d_atomname", donor_names[i][2])
        hbond_name = hbond_name.replace("d_atomnum", donor_names[i][3])
        hbond_name = hbond_name.replace("h_resname", hydrogen_names[i][0])
        hbond_name = hbond_name.replace("h_resnum", hydrogen_names[i][1])
        hbond_name = hbond_name.replace("h_atomname", hydrogen_names[i][2])
        hbond_name = hbond_name.replace("h_atomnum", hydrogen_names[i][3])
        hbond_name = hbond_name.replace("a_resname", acceptor_names[i][0])
        hbond_name = hbond_name.replace("a_resnum", acceptor_names[i][1])
        hbond_name = hbond_name.replace("a_atomname", acceptor_names[i][2])
        hbond_name = hbond_name.replace("a_atomnum", acceptor_names[i][3])
        hbond_names.append(hbond_name)

    ## parse xpmfile
    xpm = XPM(xpmfile)
    if xpm.xpm_height < len(hbond_names):
        logging.warning(
            "height of xpm ({}) in {} is not equal to number of hbond ({}) in {}, removed the excess in xpm".format(
                xpm.xpm_height, xpmfile, len(hydrogen_names), ndxfile
            )
        )
        gap = xpm.xpm_height - len(hbond_names)
        xpm.xpm_height -= gap
        xpm.xpm_datalines = xpm.xpm_datalines[gap:]
        xpm.xpm_yaxis = xpm.xpm_yaxis[gap:]
    xpm.xpm_datalines.reverse()
    xpm.xpm_yaxis.reverse()

    ## get occupancy
    occupancy, xpm_datamatrix = [], []
    for dataline in xpm.xpm_datalines:
        dot_list = []
        for i in range(
            0, xpm.xpm_width * xpm.xpm_char_per_pixel, xpm.xpm_char_per_pixel
        ):
            dot_list.append(xpm.chars.index(dataline[i : i + xpm.xpm_char_per_pixel]))
        xpm_datamatrix.append(dot_list)
        occupancy.append(sum(dot_list) / (1.0 * len(dot_list)))

    ## deal with the selection
    # select = [5]
    if select == []:
        select = [i for i in range(len(hbond_names))]
    else:
        occupancy = [occupancy[i] for i in select]
        xpm_datamatrix = [xpm_datamatrix[i] for i in select]
        hbond_names = [hbond_names[i] for i in select]

    ## draw map
    plt.figure()
    cmap = mcolors.ListedColormap(["white", "#F94C66"])
    hbond = plt.pcolormesh(
        xpm.xpm_xaxis,
        [i for i in range(len(select))],
        xpm_datamatrix,
        cmap=cmap,
        shading="auto",
    )
    if len(select) == 1:
        hbond = plt.imshow(xpm_datamatrix, cmap=cmap, aspect="auto")
    cb = plt.colorbar(hbond, orientation="horizontal", fraction=0.03)
    cb.set_ticks([0.25, 0.75])
    cb.set_ticklabels(["None", "Present"])
    for i in range(1, len(select)):
        plt.hlines(i - 0.5, min(xpm.xpm_xaxis), max(xpm.xpm_xaxis), colors="white")
    plt.title(xpm.xpm_title)
    plt.xlabel(xpm.xpm_xlabel)
    plt.ylabel(xpm.xpm_ylabel)
    if xpm.xpm_height <= 20 and hnf != "number":
        plt.yticks([i for i in range(len(select))], hbond_names)
    plt.tight_layout()
    if figout != None:
        plt.savefig(figout, dpi=300)
    if not noshow:
        plt.show()

    ## show table
    print("-" * 79)
    print("{:<2} {:<60} {}".format("id", "donor->hydrogen...acceptor", "occupancy(%)"))
    print("-" * 79)
    for i in range(len(hbond_names)):
        print(
            "{:<2d} {:<60} {:.2f}".format(
                select[i], hbond_names[i], occupancy[i] * 100.0
            )
        )
    print("-" * 79)
    if csv != None:
        with open(csv, "w") as fo:
            fo.write(
                "{},{},{}\n".format("id", "donor->hydrogen...acceptor", "occupancy(%)")
            )
            for i in range(len(hbond_names)):
                fo.write(
                    "{},{},{:.2f}\n".format(
                        select[i], hbond_names[i], occupancy[i] * 100.0
                    )
                )


def hbond_call_functions(arguments: list = []):
    """call functions according to arguments"""

    if len(arguments) == 0:
        arguments = [argv for argv in sys.argv]

    ## parse the command parameters
    parser = argparse.ArgumentParser(
        description="hbond command process the hbond infos"
    )
    parser.add_argument("-f", "--input", help="gro file for input")
    parser.add_argument("-n", "--index", help="hbond ndx file for input")
    parser.add_argument("-m", "--map", help="hbond map file for input")
    parser.add_argument(
        "-c", "--select", nargs="+", help="to select row of data, like: 1 2-4 6"
    )
    parser.add_argument("-o", "--output", help="figure name for output")
    parser.add_argument("-csv", "--csv", help="store table info into csv file")
    parser.add_argument(
        "-ns",
        "--noshow",
        action="store_true",
        help="whether not to show picture, useful on computer without gui",
    )
    parser.add_argument(
        "-hnf",
        "--hbond_name_format",
        help="define the hbond name format by user! Each atom has four"
        + " features: resname, resnum, atomname, atomnum. Distinguish "
        + "donor, hydrogen, acceptor by adding one prefix to each feature,"
        + " like: d_resname, a_resnum, h_atomname. \nSo you may able to "
        + "define hbond name style by: 'd_resname(d_resnum)@d_atomname(d_"
        + "atomnum)->h_atomname(h_atomnum)...a_resname(a_resnum)@a_atomn"
        + "ame(a_atomnum)' which is the default style,  or also you could"
        + " specify 'd_atomname@h_atomname...a_atomname' or some format "
        + "you would like. \nOr you could just set the hnf to be 'number'",
    )

    if len(arguments) < 2:
        logging.error("no input parameters, -h or --help for help messages")
        exit()

    method = arguments[1]
    if method in ["-h", "--help"]:
        parser.parse_args(arguments[1:])
        exit()
    if len(arguments) == 2:
        logging.error("no parameters, type 'dit <command> -h' for more infos.")
        exit()

    args = parser.parse_args(arguments[2:])
    select = []
    if args.select != None:
        for c in args.select:
            if "-" in c.strip("-"):
                select += [i for i in range(int(c.split("-")[0]), int(c.split("-")[1]))]
            else:
                select.append(int(c))

    if method == "hbond":
        grofile = args.input
        ndxfile = args.index
        xpmfile = args.map
        noshow = args.noshow
        figout = args.output
        csv = args.csv
        hnf = args.hbond_name_format
        hbond(xpmfile, ndxfile, grofile, select, noshow, figout, csv, hnf)
    else:
        logging.error("unknown method {}".format(method))
        exit()

    logging.info("May you good day !")


def main():
    hbond_call_functions()


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(levelname)s -> %(message)s")
    logger = logging.getLogger(__name__)
    main()
