"""
xpmParser module is part of DuIvyTools for parsing the xpm file generated by GROMACS.

Written by DuIvy and provided to you by GPLv3 license.
"""

import os
import sys
import string
from itertools import chain
from typing import List

base = os.path.dirname(os.path.realpath(os.path.join(__file__, "..")))
if base not in sys.path:
    sys.path.insert(0, base)

from utils import log


class XPM(log):
    """XPM class was designed to parse xpm file"""

    def __init__(
        self, xpmfile: str, is_file: bool = True, new_file: bool = False
    ) -> None:
        self.title: str = ""
        self.legend: str = ""
        self.type: str = ""
        self.xlabel: str = ""
        self.ylabel: str = ""
        self.width: int = 0
        self.height: int = 0
        self.color_num: int = 0
        self.char_per_pixel: int = 0
        self.chars: list[str] = []
        self.colors: list[str] = []
        self.notes: list[str] = []
        self.xaxis: list[float] = []
        self.yaxis: list[float] = []
        self.datalines: list[str] = []
        self.dot_matrix: list[list[str]] = []
        self.value_matrix: list[list[float]] = []

        if new_file:
            self.xpmfile = xpmfile
        else:
            if is_file:
                self.xpmfile: str = xpmfile
                if not os.path.exists(xpmfile):
                    self.error(f"No {xpmfile} detected ! check it !")
                if xpmfile[-4:] != ".xpm":
                    self.error(
                        f"you must specify a file with suffix .xpm, instead of {xpmfile}"
                    )
                with open(xpmfile, "r") as fo:
                    content = fo.read()
            else:
                content = xpmfile
            lines = [l.strip() for l in content.strip().split("\n")]
            self.parse_xpm(lines)
            if is_file:
                self.info(f"parsing data from {xpmfile} successfully !")

    def parse_xpm(self, lines: List[str]) -> None:
        """convert xpm content (lines) into XPM class

        Args:
            lines (List[str]): xpm file lines
        """
        flag_4_code: int = 0
        for line in lines:
            if flag_4_code == 1:
                flag_4_code = 2
                code4 = [int(c) for c in line.strip(' ,"').split()]
                self.width = code4[0]
                self.height = code4[1]
                self.color_num = code4[2]
                self.char_per_pixel = code4[3]
                continue
            elif (flag_4_code == 0) and line.startswith("static char"):
                flag_4_code = 1
                continue

            ## parsing comments and axis parts
            if line.startswith("/* x-axis"):
                self.xaxis += [float(n) for n in line.strip().split()[2:-1]]
                continue
            elif line.startswith("/* y-axis"):
                self.yaxis += [float(n) for n in line.strip().split()[2:-1]]
                continue
            elif line.startswith("/* title"):
                self.title = line.strip().split('"')[1]
                continue
            elif line.startswith("/* legend"):
                self.legend = line.strip().split('"')[1]
                continue
            elif line.startswith("/* x-label"):
                self.xlabel = line.strip().split('"')[1]
                continue
            elif line.startswith("/* y-label"):
                self.ylabel = line.strip().split('"')[1]
                continue
            elif line.startswith("/* type"):
                self.type = line.strip().split('"')[1]
                continue

            items = line.strip().split()
            ## for char-color-note part
            if len(items) == 7 and items[1] == "c":
                self.colors.append(items[2])
                self.notes.append(items[5].strip('"'))
                if len(items[0].strip('"')) == self.char_per_pixel:
                    self.chars.append(items[0].strip('"'))
                ## deal with blank char
                if len(items[0].strip('"')) < self.char_per_pixel:
                    char_item = line.strip('"')[: self.char_per_pixel]
                    self.chars.append(char_item)
                continue

            ## for figure content part
            if line.strip().startswith('"') and (
                len(line.strip().strip(",").strip('"'))
                == self.width * self.char_per_pixel
            ):
                self.datalines.append(line.strip().strip(",").strip('"'))

        ## check infos
        if self.color_num != len(self.chars):
            self.error(f"length of chars != color_num")
        if self.color_num != len(self.colors):
            self.error(f"length of colors != color_num")
        if self.color_num != len(self.notes):
            self.error(f"length of notes != color_num")
        if len(self.xaxis) != self.width and len(self.xaxis) != self.width + 1:
            self.error(
                f"length of xaxis ({len(self.xaxis)}) != xpm width ({self.width}) or xpm width +1"
            )
        if len(self.yaxis) != self.height and len(self.yaxis) != self.height + 1:
            self.error(
                f"length of yaxis ({len(self.yaxis)}) != xpm height ({self.height}) or xpm height +1"
            )
        if len(self.datalines) != self.height:
            self.error(
                f"rows of data ({len(self.datalines)}) is not equal to xpm height ({self.height}), check it!"
            )

        ## convert data
        if len(self.xaxis) == self.width + 1:
            self.xaxis = [
                (self.xaxis[i - 1] + self.xaxis[i]) / 2.0
                for i in range(1, len(self.xaxis))
            ]
            self.warn(
                "length of x-axis is 1 more than xpm width, use intermediate value for instead. "
            )
        if len(self.yaxis) == self.height + 1:
            self.yaxis = [
                (self.yaxis[i - 1] + self.yaxis[i]) / 2.0
                for i in range(1, len(self.yaxis))
            ]
            self.warn(
                "length of y-axis is 1 more than xpm height, use intermediate value for instead. "
            )
        if self.type == "Continuous":
            self.notes = [float(n) for n in self.notes]

        self.yaxis.reverse()  # IMPORTANT! from high to low now

        for dataline in self.datalines:
            dot_list: List[str] = []
            value_list: List[float] = []
            for i in range(0, self.width * self.char_per_pixel, self.char_per_pixel):
                dot = dataline[i : i + self.char_per_pixel]
                dot_list.append(dot)
                if self.type == "Continuous":
                    value = self.notes[self.chars.index(dot)]
                else:
                    # for Discrete, value store the index of chars|notes|colors
                    value = self.chars.index(dot)
                value_list.append(value)
            self.dot_matrix.append(dot_list)
            self.value_matrix.append(value_list)

        if (
            len(self.dot_matrix) != self.height
            or len(self.value_matrix) != self.height
            or len(self.dot_matrix[0]) != self.width
            or len(self.value_matrix[0]) != self.width
        ):
            self.error("Dimension error while parsing xpm file")

    def __sub__(self, xpm):  # diff_map
        """values of self - values of xpm correspondingly, return a new result XPM"""
        if self.type != "Continuous" or xpm.type != "Continuous":
            self.error("Only supported to apply xpm_diff to Continuous type of xpms")
        for key in ["title", "xlabel", "ylabel", "xaxis", "yaxis"]:
            if self.__dict__[key] != xpm.__dict__[key]:
                self.warn(
                    f"Detected different {key} in {self.xpmfile} and {xpm.xpmfile}. \nDIT strongly warns you that different type (meanings) of xpms should NOT be used to calculate difference. The results would NOT be reliable !!! "
                )
        if self.width != xpm.width or self.height != xpm.height:
            self.error(
                f"The shape of {self.xpmfile} ({self.width}, {self.height}) and {xpm.xpmfile} ({xpm.width}, {xpm.height}) are different, unable to calculate difference."
            )

        out = XPM("", is_file=False, new_file=True)
        for key, value in self.__dict__.items():
            out.__dict__[key] = value
        for h in range(self.height):
            for w in range(self.width):
                value = self.value_matrix[h][w] - xpm.value_matrix[h][w]
                value = float(f"{value:.6f}")
                out.value_matrix[h][w] = value
        out.refresh_by_value_matrix()
        return out

    def refresh_by_value_matrix(self, is_Continuous: bool = True) -> None:
        """generate the xpm class content by value_matrix

        Args:
            is_Continuous (bool, optional): set the type of xpm. Defaults to True.
        """
        out_value_list = sorted(list(set(chain(*self.value_matrix))))
        self.notes = out_value_list
        colors, l = [], len(out_value_list)
        for i in range(l):
            c = 16000000 // len(out_value_list) * i
            colors.append(f'#{hex(c)[-6:].replace("x", "0").upper():0>6}')
        self.color_num = l
        self.colors = colors

        letters = string.ascii_letters + "0123456789!@#$%^&*()-_=+{}|;"
        if l <= 0:
            self.error("no data detected in xpm.value_matrix")
        elif l > 0 and l <= len(letters):
            self.char_per_pixel = 1
            self.chars = [f"{x}" for x in letters][:l]
        elif l > len(letters) and l <= len(letters) * len(letters):
            self.char_per_pixel = 2
            self.chars = [f"{x}{y}" for x in letters for y in letters][:l]
        elif l > len(letters) * len(letters) and l <= len(letters) * len(letters) * len(
            letters
        ):
            self.warn(
                f"so many values ({l}) in xpm.value_matrix, may resulting in large xpm file"
            )
            self.char_per_pixel = 3
            self.chars = [
                f"{x}{y}{z}" for x in letters for y in letters for z in letters
            ][:l]
        else:
            self.error(
                f"too many values ({l}) in xpm.value_matrix, only able to construct XPM with {len(letters)*len(letters)*len(letters)} values or less"
            )

        if len(self.dot_matrix) == 0 or len(self.datalines) == 0:
            for h in range(self.height):
                self.dot_matrix.append(["" for _ in range(self.width)])
                self.datalines.append("")
        for h in range(self.height):
            dot_line: str = ""
            for w in range(self.width):
                dot = self.chars[out_value_list.index(self.value_matrix[h][w])]
                self.dot_matrix[h][w] = dot
                dot_line += dot
            self.datalines[h] = dot_line
        if not is_Continuous:  # refresh value_matrix from str to index
            for h in range(self.height):
                for w in range(self.width):
                    self.value_matrix[h][w] = out_value_list.index(
                        self.value_matrix[h][w]
                    )

    def save(self, outname: str) -> None:
        """dump XPM into xpm file"""
        outstr = "/* XPM */\n/* This file is generated by DuIvyTools */\n"
        outstr += f"""/* title: "{self.title}" */\n"""
        outstr += f"""/* legend: "{self.legend}" */\n"""
        outstr += f"""/* x-label: "{self.xlabel}" */\n"""
        outstr += f"""/* y-label: "{self.ylabel}" */\n"""
        outstr += f"""/* type: "{self.type}" */\n"""
        outstr += """static char *gromacs_xpm[] = {\n"""
        outstr += f""""{self.width} {self.height} {self.color_num} {self.char_per_pixel}",\n"""
        for char, color, note in zip(self.chars, self.colors, self.notes):
            outstr += f""""{char}  c {color} " /* "{note}" */,\n"""
        for x in range(0, len(self.xaxis), 50):
            outstr += (
                "/* x-axis: "
                + " ".join(str(v) for v in self.xaxis[x : x + 50])
                + " */\n"
            )
        # should not change original xpm data
        yaxis = [y for y in reversed(self.yaxis)]  # top to bottom -> bottom to top
        for y in range(0, len(yaxis), 50):
            outstr += (
                "/* y-axis: " + " ".join(str(v) for v in yaxis[y : y + 50]) + " */\n"
            )
        for line in self.datalines:
            outstr += f""""{line}",\n"""
        outstr = outstr.strip().strip(",") + "\n"
        with open(outname, "w") as fo:
            fo.write(outstr)
        self.info(f"Save results into {outname} successfully")


class XPMS(log):
    """XPMS class was designed to parse xpm file with multi-frames"""

    def __init__(self, xpmfile: str) -> None:
        self.xpmfile: str = xpmfile
        self.frames: list[XPM] = []

        with open(xpmfile, "r") as fo:
            contents = fo.read()
        contents = contents.split("/* XPM */")
        contents = [f"/* XPM */\n{c}" for c in contents if c.strip() != ""]
        for content in contents:
            xpm = XPM(content, is_file=False)
            self.frames.append(xpm)
        self.info(f"parsing multi-frames data from {xpmfile} successfully !")

    def __len__(self) -> int:
        """return number of frames"""
        return len(self.frames)

    def __getitem__(self, index: int) -> XPM:
        """get XPM of one frame"""
        return self.frames[index]

    def get_time_series(self) -> List[float]:
        """parsing time infos from xpm titles"""
        times: float = []
        for xpm in self:
            if xpm.title.startswith("t="):
                times.append(int(xpm.title[2:-2]))
            else:
                self.error(f"cannot parse time info from xpm title {xpm.title}")
        return times
