import numpy as np
from abc import ABC, abstractmethod
import svg

rgen = np.random.default_rng()

# the following functions are taken from Ben Southgate:
# https://bsouthga.dev/posts/colour-gradients-with-python


def hex_to_RGB(hex):
    """ "#FFFFFF" -> [255,255,255]"""
    # Pass 16 to the integer function for change of base
    return [int(hex[i : i + 2], 16) for i in range(1, 6, 2)]


def RGB_to_hex(RGB):
    """[255,255,255] -> "#FFFFFF" """
    # Components need to be integers for hex to make sense
    RGB = [int(x) for x in RGB]
    return "#" + "".join(
        ["0{0:x}".format(v) if v < 16 else "{0:x}".format(v) for v in RGB]
    )


def colour_dict(gradient):
    """Takes in a list of RGB sub-lists and returns dictionary of
    colours in RGB and hex form for use in a graphing function
    defined later on."""
    return {
        "hex": [RGB_to_hex(RGB) for RGB in gradient],
        "r": [RGB[0] for RGB in gradient],
        "g": [RGB[1] for RGB in gradient],
        "b": [RGB[2] for RGB in gradient],
    }


def linear_gradient(start_hex, finish_hex="#FFFFFF", n=10):
    """returns a gradient list of (n) colours between
    two hex colours. start_hex and finish_hex
    should be the full six-digit colour string,
    inlcuding the number sign ("#FFFFFF")"""
    # Starting and ending colours in RGB form
    s = hex_to_RGB(start_hex)
    f = hex_to_RGB(finish_hex)
    # Initilize a list of the output colours with the starting colour
    RGB_list = [s]
    # Calcuate a colour at each evenly spaced value of t from 1 to n
    for t in range(0, n):
        # Interpolate RGB vector for colour at the current value of t
        curr_vector = [
            int(s[j] + (float(t) / (n - 1)) * (f[j] - s[j])) for j in range(3)
        ]
        # Add it to our list of output colours
        RGB_list.append(curr_vector)

    return colour_dict(RGB_list)


def rand_hex_colour(num=1):
    """Generate random hex colours, default is one,
    returning a string. If num is greater than
    1, an array of strings is returned."""
    colours = [RGB_to_hex([x * 255 for x in rgen.rand(3)]) for i in range(num)]
    if num == 1:
        return colours[0]
    else:
        return colours


def polylinear_gradient(colours, n):
    """returns a list of colours forming linear gradients between
    all sequential pairs of colours. "n" specifies the total
    number of desired output colours"""
    # The number of colours per individual linear gradient
    n_out = int(float(n) / (len(colours) - 1))
    # returns dictionary defined by colour_dict()
    gradient_dict = linear_gradient(colours[0], colours[1], n_out)

    if len(colours) > 1:
        for col in range(1, len(colours) - 1):
            next = linear_gradient(colours[col], colours[col + 1], n_out)
            for k in ("hex", "r", "g", "b"):
                # Exclude first point to avoid duplicates
                gradient_dict[k] += next[k][1:]

    return gradient_dict


class ColourMap(ABC):
    @abstractmethod
    def __call__(self, v: float): ...


class LinearGradientColourMap(ColourMap):
    def __init__(
        self,
        colours: list[str] | None = ["#ff0000", "#ffffff", "#0000ff"],
        min_value: float | None = 0,
        max_value: float | None = 1,
        bins: int = 100,
    ):
        self.colours = polylinear_gradient(colours, bins)
        self.min, self.max = min_value, max_value

    def __call__(self, v: float):
        v = max(0, int((v - self.min) / (self.max - self.min) * 100) - 1)
        if v >= len(self.colours["hex"]):
            breakpoint()
        return self.colours["hex"][v]


class RandomColourMap(ColourMap):
    def __init__(self, random_state: int | list[int] | None = [2, 3, 4, 5, 6]):
        self.rgen = np.random.default_rng(random_state)

    def __call__(self, v: float):
        return RGB_to_hex([x * 255 for x in self.rgen.random(3)])


class MatrixVisualisation:
    def __init__(
        self,
        matrix: np.typing.NDArray,
        cmap: ColourMap,
        text: bool = False,
        labels: int | list[str] | bool = False,
    ):
        self.m, self.n = matrix.shape
        width = 20
        height = 20
        gap = 1
        self.text = text
        self.total_width = (gap + width) * n + gap
        self.total_height = (gap + height) * m + gap
        self.cmap = cmap

        self.elements = []
        self.elements.append(
            svg.Style(text=".mono { font: monospace; text-align: center;}")
        )
        self.elements.append(svg.Style(text=".small { font-size: 25%; }"))
        self.elements.append(svg.Style(text=".normal { font-size: 12px; }"))

        for i, y in enumerate(range(gap, self.total_height, gap + height)):
            for j, x in enumerate(range(gap, self.total_width, gap + width)):
                self.elements.append(
                    svg.Rect(
                        x=x,
                        y=y,
                        width=width,
                        height=height,
                        stroke="transparent",
                        fill=cmap(matrix[i, j]),
                    )
                )
                if text:
                    self.elements.append(
                        svg.Text(
                            x=x + width / 5,
                            y=y + 3 * height / 4,
                            textLength=width / 2,
                            lengthAdjust="spacingAndGlyphs",
                            class_=["mono"],
                            text=f"{matrix[i, j]:.02f}",
                        )
                    )

    def colourbar(
        self,
        min_value: float = 0,
        max_value: float = 1,
        height: int | None = None,
        width: int = 20,
        resolution: int = 256,
        border: int | bool = 1,
        labels: int | list[str] | bool = False,
    ):
        if height is None:
            height = int(self.total_height * 2 / 3)

        lines = [
            svg.Rect(
                fill=self.cmap(v),
                x=0,
                y=y,
                width=width,
                height=1.1 * height / resolution,
                stroke="none",
            )
            for y, v in zip(
                np.linspace(0, height, resolution),
                np.linspace(min_value, max_value, resolution - 1),
            )
        ]

        if labels is None:
            label = []
        elif isinstance(labels, int):
            label = svg.G(
                id="colourbar labels",
                elements=[
                    svg.Text(
                        text=f"— {v:.02f}",
                        class_=["normal"],
                        x=width,
                        y=y,
                        dy=3,
                    )
                    for y, v in zip(
                        np.linspace(0, height, labels),
                        np.linspace(min_value, max_value, labels),
                    )
                ],
            )
        elif isinstance(labels, list):
            if all(isinstance(n, str) for n in labels):
                label = svg.G(
                    id="colourbar labels",
                    elements=[
                        svg.Text(
                            text=f"— {v}",
                            class_=["normal"],
                            x=width,
                            y=y,
                            dy=3,
                        )
                        for y, v in zip(np.linspace(0, height, len(labels)), labels)
                    ],
                )
            if all(isinstance(n, float) or isinstance(n, int) for n in labels):
                label = svg.G(
                    id="colourbar labels",
                    elements=[
                        svg.Text(
                            text=f"— {v:.02f}",
                            class_=["normal"],
                            x=width,
                            y=(v - min_value) / (max_value - min_value) * height,
                            dy=3,
                        )
                        for v in labels
                    ],
                )

        cbar = svg.G(
            id="colourbar",
            elements=[
                lines,
                label,
                svg.Rect(
                    x=0,
                    y=0,
                    width=width,
                    height=height,
                    fill="none",
                    stroke_width=border,
                    stroke="black",
                ),
            ],
            transform=[
                svg.Translate(
                    x=int(self.total_width + width / 2),
                    y=int((self.total_height - height) / 2),
                )
            ],
        )
        self.elements.append(cbar)
        self.total_width = self.total_width + 2 * width + 40 * bool(labels)

    @property
    def svg(self):
        return str(
            svg.SVG(
                width=self.total_width, height=self.total_height, elements=self.elements
            )
        )

    def __repr__(self):
        return f"""Matrix Visualisation: 
    shape: {matrix.shape}
    size:  {self.total_width}x{self.total_height}
            """


if __name__ == "__main__":
    m, n = 30, 20
    matrix = rgen.random(size=(m, n))

    colours = ["#f5d72a", "#ffffff", "#2182af"]
    # colours = ["#ff0000", "#00ff00", "#0000ff"]
    cmap = LinearGradientColourMap(colours, matrix.min(), matrix.max())
    # cmap = RandomColourMap()

    fig = MatrixVisualisation(matrix, cmap=cmap)
    fig.colourbar(labels=["yellow", "white", "blue"])
    fig.colourbar(labels=5)
    fig.colourbar(labels=[0.2, 0.5, 0.55, 0.66, 1])
    filename = "matrix.svg"
    print(fig)
    with open(filename, "w") as f:
        f.write(fig.svg)