from iminuit.color import Gradient

__all__ = ["LatexTable"]


class LatexTable:
    """Latex table output.
    """

    float_format = "%10.5g"
    int_format = "%d"
    latex_kwd = [
        "alpha",
        "beta",
        "gamma",
        "delta",
        "epsilon",
        "zeta",
        "eta",
        "theta",
        "iota",
        "kappa",
        "lambda",
        "mu",
        "nu",
        "xi",
        "omicron",
        "pi",
        "rho",
        "sigma",
        "tau",
        "upsilon",
        "phi",
        "chi",
        "psi",
        "omega",
        "Alpha",
        "Beta",
        "Gamma",
        "Delta",
        "Epsilon",
        "Zeta",
        "Eta",
        "Theta",
        "Iota",
        "Kappa",
        "Lambda",
        "Mu",
        "Nu",
        "Xi",
        "Omicron",
        "Pi",
        "Rho",
        "Sigma",
        "Tau",
        "Upsilon",
        "Phi",
        "Chi",
        "Psi",
        "Omega",
    ]

    def __init__(
        self,
        data,
        headers=None,
        smart_latex=True,
        escape_under_score=True,
        alignment=None,
        rotate_header=False,
        latex_map=None,
    ):
        # Make sure #data columns matches #header columns, using any non-zero
        # length if data or headers are missing
        if len(data) > 0:
            num_col = len(data[0])
            if headers:
                assert num_col == len(headers)
        else:
            if headers is not None:
                num_col = len(headers)
            else:
                # LaTeX requires at least one column
                num_col = 1

        self.headers = headers
        self.data = data
        self.num_col = num_col
        self.smart_latex = smart_latex
        self.escape_under_score = escape_under_score
        self.alignment = self._auto_align() if alignment is None else alignment
        self.rotate_header = rotate_header
        self.latex_map = {} if latex_map is None else latex_map
        self.cell_color = {}  # map of tuple (i,j)=>(r, g, b) #i,j include header

    def _auto_align(self):
        return "|" + "c|" * self.num_col

    def _format(self, s):
        if s in self.latex_map:
            return self.latex_map[s]
        elif isinstance(s, float):
            return self.float_format % s
        elif isinstance(s, int):
            return self.int_format % s
        elif self.smart_latex:
            return self._convert_smart_latex(s)
        elif self.escape_under_score:
            return s.replace("_", r"\_")
        else:
            return s

    def _convert_smart_latex(self, s):
        """convert greek symbol to latex one
        transform
        a to $a$ if a is greek letter else just a
        a_xxx to $a_{xxx}$ and
        a_xxx_yyy_zzz to a xxx $yyy_{zzz}$
        """
        # FIXME: implement this

        parts = s.split("_")
        if len(parts) == 1:  # a to $a$ if a is greek letter else just a
            if parts[0] in self.latex_kwd:
                return r"$\%s$" % str(parts[0])
            else:
                return str(parts[0])
        elif len(parts) == 2:  # a_xxx to $a_{xxx}$ and
            first = "\\%s" % parts[0] if parts[0] in self.latex_kwd else parts[0]
            second = "\\%s" % parts[1] if parts[1] in self.latex_kwd else parts[1]
            return r"$%s_{%s}$" % (first, second)
        else:  # a_xxx_yyy_zzz to a xxx $yyy_{zzz}$
            textpart = map(self._convert_smart_latex, parts[:-2])
            textpart = " ".join(textpart)
            latexpart = self._convert_smart_latex("_".join(parts[-2:]))
            return textpart + " " + latexpart

    def set_cell_color(self, i, j, c):
        """colorize i,j cell with rgb color tuple c

        Note that i,j index includes header.
        i=0 is header if header is present. If header is not present then
        i=0 refer to first data row.
        """
        self.cell_color[(i, j)] = c

    def _prepare(self):  # return list of list
        ret = []
        if self.headers:
            tmp = list(map(self._format, self.headers))
            if self.rotate_header:
                tmp = list(map(lambda x: "\\rotatebox{90}{%s}" % x, tmp))

            ret.append(tmp)
        for x in self.data:
            ret.append(list(map(self._format, x)))
        return ret

    def __str__(self):
        hline = "\\hline\n"
        ret = ""
        if len(self.cell_color) != 0:
            ret += "%\\usepackage[table]{xcolor} % include this for color\n"
            ret += "%\\usepackage{rotating} % include this for rotate header\n"
            ret += "%\\documentclass[xcolor=table]{beamer} % for beamer\n"
        ret += "\\begin{tabular}{%s}\n" % self.alignment
        ret += hline
        tdata = self._prepare()
        # decorate it

        for (i, j), c in self.cell_color.items():
            xcolor = "[RGB]{%d,%d,%d}" % (c[0], c[1], c[2])
            tdata[i][j] = "\\cellcolor" + xcolor + " " + tdata[i][j]

        for line in tdata:
            ret += " & ".join(line) + "\\\\\n"
            ret += hline
        ret += "\\end{tabular}\n"

        return ret.strip()


class LatexFactory:
    @classmethod
    def build_matrix(cls, vnames, matrix, latex_map=None):
        """build latex correlation matrix"""
        # ret_link  = '<a onclick="$(\'#%s\').toggle()" href="#">Show Latex</a>'%uid
        headers = [""] + list(vnames)
        data = []
        color = {}
        grad = Gradient(
            (-1.0, 120.0, 120.0, 250.0),
            (0.0, 250.0, 250.0, 250.0),
            (1.0, 250.0, 100.0, 100.0),
        )
        for i, v1 in enumerate(vnames):
            tmp = [v1]
            for j in range(len(vnames)):
                m = matrix[i][j]
                tmp.append(m)
                color[(i + 1, j + 1)] = grad(m)
                # +1 for header on the side and top
            data.append(tmp)

        table = LatexTable(
            headers=headers, data=data, rotate_header=True, latex_map=latex_map
        )
        table.float_format = "%.2g"
        for (i, j), c in color.items():
            table.set_cell_color(i, j, c)
        return table

    @classmethod
    def build_param_table(
        cls, mps, merr=None, float_format="%5.3e", smart_latex=True, latex_map=None
    ):
        """build latex parameter table"""
        headers = [
            "",
            "Name",
            "Value",
            "Hesse Error",
            "Minos Error-",
            "Minos Error+",
            "Limit-",
            "Limit+",
            "Fixed?",
        ]

        data = []
        for i, mp in enumerate(mps):
            minos_p, minos_m = (
                ("", "")
                if merr is None or mp.name not in merr
                else ("%g" % merr[mp.name].upper, "%g" % merr[mp.name].lower)
            )
            limit_p = "" if mp.upper_limit is None else "%g" % mp.upper_limit
            limit_m = "" if mp.lower_limit is None else "%s" % mp.lower_limit
            fixed = "Yes" if mp.is_fixed else "No"
            tmp = [
                i,
                mp.name,
                "%g" % mp.value,
                "%g" % mp.error,
                minos_m,
                minos_p,
                limit_m,
                limit_p,
                fixed,
            ]
            data.append(tmp)
        alignment = "|c|r|r|r|r|r|r|r|c|"
        ret = LatexTable(
            data,
            headers=headers,
            alignment=alignment,
            smart_latex=smart_latex,
            latex_map=latex_map,
        )
        ret.float_format = float_format
        return ret
