#!/usr/bin/env python

from itertools import combinations, pairwise, product, count
from pathlib import Path
from typing import Sequence, Tuple, TypeAlias, Optional, Protocol
from numpy._typing import NDArray
from numpy.core import numerictypes
from numpy.core.function_base import linspace
from numpy.typing import ArrayLike
from math import cos, radians, sin, sqrt, tau, ceil, pi
from fractions import Fraction
from dataclasses import dataclass
from collections.abc import Iterable
import os

import numpy as np
import matplotlib.pyplot as plt
from skimage.io import imread


Number: TypeAlias = int | float
XY: TypeAlias = Tuple[Number, Number]


def PA(point: XY):
    return f"PA {round(point[0])},{round(point[1])};"


def SP(pen: int = 0):
    return f"SP {pen};"


def LB(string: str):
    # ASCII 3 = "\3" = "ETX" = (end of text)
    return f"LB{string}\3;"


def TEXT(point, label, run_over_rise=None, width=None, height=None):
    if run_over_rise:
        yield f"DI {round(run_over_rise[0])},{round(run_over_rise[1])};"

    if width and height:
        yield f"SI {width:.3f},{height:.3f}"

    yield from [PU, PA(point), LB(label)]


IN = "IN;"
PD = "PD;"
PU = "PU;"


class Plottable(Protocol):
    def __call__(self, offsets: XY = (0, 0), size: XY = (1, 1)) -> Iterable[str]: ...


GapUnit: TypeAlias = Number | Fraction
Gap: TypeAlias = GapUnit | Tuple[GapUnit, GapUnit]


class ZStack(Plottable):
    def __init__(self, children: Sequence[Plottable]):
        self.children = children

    def __call__(self, offsets: XY = (0, 0), size: XY = (1, 1)) -> Iterable[str]:
        for child in self.children:
            yield from child(offsets, size)


class Grid(Plottable):
    def __init__(
        self,
        children: Sequence[Plottable],
        columns=1,
        gap: Gap = (0, 0),
    ):
        self.children = children
        self.columns = columns
        if not isinstance(gap, tuple):
            gap = (gap, gap)
        self.gap = gap

    def __call__(self, offsets: XY = (0, 0), size: XY = (1, 1)) -> Iterable[str]:
        row_gap, column_gap = self.gap
        width = (1.0 - column_gap) / self.columns * size[0]
        row_count = ceil(len(self.children) / self.columns)
        height = (1.0 - row_gap) / row_count * size[1]
        child_size = (width, height)
        gap_width = column_gap / (self.columns - 1) * size[0] if self.columns > 1 else 0
        gap_height = row_gap / (row_count - 1) * size[1] if row_count > 1 else 0
        # print(locals())
        for i, child in enumerate(self.children):
            row, column = divmod(i, self.columns)
            child_offsets = (
                int(column * (width + gap_width) + offsets[0]),
                int(row * (height + gap_height) + offsets[1]),
            )
            yield from child(offsets=child_offsets, size=child_size)


class Page:
    default_height = 7650
    default_width = 10750
    default_size = (default_width, default_height)

    def __init__(
        self,
        child: Plottable,
        origin=(0, 0),
        size=default_size,
    ) -> None:
        self.child = child
        self.origin = origin
        self.size = size

    def __call__(self, number: Optional[int | str] = None):
        yield from [IN, SP(1)]
        yield from self.child(self.origin, self.size)
        if number and False:
            yield from TEXT(
                label=str(number),
                point=(Page.default_width + 200, Page.default_height / 2 - 62),
                run_over_rise=(0, 1),  # portrait bottom
            )
        yield from [PU, SP()]


class CalibratedPage(Page):
    """Magic values; calibrated with ruler"""

    def __init__(self, child: Plottable) -> None:
        # to equalize left and right margin...
        origin = (0, 220)
        # ...and top & bottom margin, too
        size = (Page.default_width - 80, Page.default_height - 220 - 10)
        super().__init__(child, origin, size)


def scaled(point: XY, offset: XY, size: XY):
    scaled_x = point[0] * size[0] + offset[0]
    scaled_y = point[1] * size[1] + offset[1]
    return (scaled_x, scaled_y)


class Path(Plottable):
    def __init__(self, vertices: Sequence[XY], close=False) -> None:
        self.vertices = vertices
        self.close = close

    def __call__(self, offsets: XY = (0, 0), size: XY = (1, 1)) -> Iterable[str]:
        start, rest = self.vertices[0], self.vertices[1:]
        scaled_start = scaled(start, offsets, size)
        yield from [PU, PA(scaled_start), PD]
        yield from [PA(scaled(point, offsets, size)) for point in rest]
        if self.close:
            yield PA(scaled_start)


def rgb2gray(rgb):
    r, g, b = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]
    gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
    return gray


class Postage(Plottable):
    def __init__(self, message=[], address=[]) -> None:
        self.message = message
        self.address = address
        self.line_height = 0.1

    def __call__(self, offsets: XY = (0, 0), size: XY = (1, 1)) -> Iterable[str]:
        cmds = []
        for i, line in enumerate(self.message):
            cmds += TEXT(
                label=line,
                point=scaled((-0.1, 0.1 + self.line_height * i), offsets, size),
                run_over_rise=(1, 0),  # portrait bottom
            )
        for i, line in enumerate(self.address):
            cmds += TEXT(
                label=line,
                point=scaled(
                    (0.4, self.line_height * 2 + self.line_height * i), offsets, size
                ),
                run_over_rise=(1, 0),  # portrait bottom
            )
        yield from cmds


class Cat(Plottable):
    def __init__(self, close=False) -> None:
        self.something = True

    def __call__(self, offsets: XY = (0, 0), size: XY = (1, 1)) -> Iterable[str]:

        filename = "mrchou.png"
        thresholds = [24]
        thresholdi_override = 1

        commands = [PU]
        for thresholdi, threshold in enumerate(thresholds):
            if thresholdi_override is not None:
                thresholdi = thresholdi_override
            # run `convert libby.png -colorspace Gray libby2.png`
            os.system(f"convert {filename} -colorspace Gray 1.png")
            os.system(f"convert 1.png -background white -alpha remove -alpha off 2.png")
            os.system(f"convert 2.png -threshold {threshold}% 3.png")
            im = imread("3.png")
            max_size = max(im.shape)
            os.system(
                f"convert -size {max_size}x{max_size} xc:white 3.png -gravity center -composite 4.png"
            )
            # reduce size of image
            im = imread("4.png")
            scaling_factor = 1
            spread_factor = 2
            os.system(
                f"convert 4.png -resize {im.shape[0]/scaling_factor}x{im.shape[1]/scaling_factor} 5.png"
            )
            os.system(f"convert 5.png -rotate 0 6.png")
            im = imread("6.png")

            # # # rotate image
            # im = np.rot90(im, 3)

            # get dimensions of image
            pen_down = False
            did_pen_down = False
            for i, row in enumerate(im):
                if i % spread_factor != 0 and spread_factor > 1:
                    continue
                for j, v in enumerate(row):
                    v = not v
                    x = (
                        offsets[0]
                        + scaling_factor * j * size[0] / im.shape[0]
                        + scaling_factor * size[0] / im.shape[0] * thresholdi / 3 * 2
                    )
                    y = (
                        offsets[1]
                        + scaling_factor * i * size[1] / im.shape[1]
                        + scaling_factor * size[1] / im.shape[1] * thresholdi / 3 * 2
                    )
                    if v and not pen_down:
                        commands.append(PA((x, y)))
                        commands.append(PD)
                        pen_down = True
                        did_pen_down = True
                    elif (not v) and pen_down:
                        commands.append(PA((x, y)))
                        commands.append(PU)
                        pen_down = False
                if did_pen_down:
                    commands.append(PU)
                    did_pen_down = False
        yield from commands


class UpTriangle(Plottable):
    def __call__(self, offsets: XY = (0, 0), size: XY = (1, 1)) -> Iterable[str]:
        path = Path([(0.0, 0.0), (0.5, 1.0), (1.0, 0.0)], close=True)
        yield from path(offsets, size)


class Outline(Plottable):
    def __init__(self, child: Plottable) -> None:
        self.child = child

    def __call__(self, offsets: XY = (0, 0), size: XY = (1, 1)) -> Iterable[str]:
        path = Path([(0, 0), (0, 1), (1, 1), (1, 0)], close=True)
        yield from path(offsets, size)
        yield from self.child(offsets, size)


class CenterSquare(Plottable):
    def __init__(self, child: Plottable) -> None:
        self.child = child

    def __call__(self, offsets: XY = (0, 0), size: XY = (1, 1)) -> Iterable[str]:
        major, minor = (0, 1) if size[0] >= size[1] else (1, 0)
        major_length, minor_length = size[major], size[minor]
        assert (major_length, minor_length) == (max(size), min(size))
        major_offset, minor_offset = offsets[major], offsets[minor]

        delta = major_length - minor_length
        child_offsets = ()
        major_minor_offsets = (offsets[major] + delta / 2, offsets[minor])
        child_offsets = (
            major_minor_offsets
            if size[0] >= size[1]
            else tuple(reversed(major_minor_offsets))
        )
        child_size = (minor_length, minor_length)
        yield from self.child(offsets=child_offsets, size=child_size)


# Function to parse HPGL commands and extract points
def parse_hpgl(hpgl):
    commands = hpgl.strip().split(";")
    lines = []
    texts = []
    pen_down = 0
    current_pos = (0, 0)
    for cmd in commands:
        if cmd.startswith("PU"):
            pen_down = 0
        elif cmd.startswith("PD"):
            pen_down = 1
        elif cmd.startswith("LB"):
            texts.append((current_pos, cmd[2:-1]))
        elif cmd.startswith("PA"):
            coords = cmd[2:].split(",")
            last_pos = current_pos
            current_pos = (int(coords[0]), int(coords[1]))
            if pen_down == 1:
                lines.append((last_pos, current_pos))

    return lines, texts


# Function to draw lines on a matplotlib plot
def draw_hpgl(hpgl):
    lines, texts = parse_hpgl(hpgl)
    fig, ax = plt.subplots()
    for line in lines:
        (x1, y1), (x2, y2) = line
        ax.plot([x1, x2], [y1, y2], "black")
    for text in texts:
        (x, y), t = text
        ax.text(x, y, t, fontsize=10)
    ax.set_aspect("equal", "box")
    plt.gca().invert_yaxis()  # Invert Y-axis to match HPGL coordinate system
    plt.axis("off")  # Turn off axes for a cleaner look
    plt.savefig("plot.png", bbox_inches="tight", pad_inches=0)  # Save plot to file
    # plt.show()


def main():
    cats = [
        Outline(
            CenterSquare(
                Postage(
                    message=["hello, world", "- zack"],
                    address=[
                        "hello",
                        "somewhere st",
                        "apt x",
                        "someplace",
                        "12345",
                    ],
                )
            )
        )
        for _ in range(4)
    ]
    # cats = [Outline(CenterSquare(Cat())) for _ in range(4)]
    grid = Grid(children=cats, columns=2, gap=(Fraction("1/64"), Fraction("1/64")))

    page = CalibratedPage(grid)
    hpgl_code = ""
    for line in page(number="testpage"):
        print(line)
        hpgl_code += line
    draw_hpgl(hpgl_code)


if __name__ == "__main__":
    main()