#!/usr/bin/env ipython
from __future__ import annotations
import torch
from random import uniform
import math

import pydiffvg

from util import rgb2hex

class Point:
    """2D point, optionally with control points;
        alternatively, 2-vector"""
    def __init__(self, x: float, y: float,
                 controls: list[Point] = None,
                 round=False):

        if isinstance(x, torch.Tensor):
            # Convert from tensor
            x = x.item()
            y = y.item()

        self.id = id(self)
        self.x = x
        self.y = y
        self.controls = controls or []

        if round:
            self.round()

    def add_control(self, control: Point):
        self.controls.append(control)

    def add_to_patch(self, patch: Patch):
        self.patches.append(patch)

    def as_xy(self):
        return [self.x, self.y]

    def round(self):
        self.x = int(self.x * 100) / 100.0
        self.y = int(self.y * 100) / 100.0

    def replace(self, pt: Point):
        """"Replace (x,y) coordinates of point while maintaining pointer."""
        self.x = pt.x
        self.y = pt.y

    def add(self, pt: Point):
        self.x += pt.x
        self.y += pt.y
        for cp in self.controls:
            cp.x += pt.x
            cp.y += pt.y

    def mult(self, pt: Point):
        self.x *= pt.x
        self.y *= pt.y
        for cp in self.controls:
            cp.x *= pt.x
            cp.y *= pt.y

    def equalize(self, other):
        self.id = other.id

    @classmethod
    def random(cls, rx=(0, 1), ry=(0, 1)):
        # Clamp to (0, 1)
        rx = (max(0, rx[0]), min(1, max(0, rx[1])))
        ry = (max(0, ry[0]), min(1, max(0, ry[1])))

        out = cls(uniform(*rx), uniform(*ry))
        out.round()
        return out

    def __eq__(self, other):
        return self.id == other.id

    def __hash__(self):
        return self.id

    def __repr__(self):
        return f"P<({self.x}, {self.y})[{len(self.controls)}]>"

    def __str__(self):
        return self.__repr__()


class Patch:
    """Cubic patch."""
    def __init__(self, points: list[Point], color=(0.2, 0.5, 0.7, 1.0)):
        self.points = points
        self.color = color

    def translate(self, pt: Point):
        for p in self.points:
            p.add(pt)

    def scale(self, pt: Point):
        for p in self.points:
            p.mult(pt)

    def as_path(self, width=256, height=256) -> pydiffvg.Path:
        ppoints = []
        for pt in self.points:
            ppoints.append([pt.x * width, pt.y * height])
            for cpt in pt.controls:
                ppoints.append([cpt.x * width, cpt.y * height])

        return pydiffvg.Path(
            num_control_points=torch.tensor(
                [len(p.controls) for p in self.points]
            ),
            points=torch.tensor(ppoints, requires_grad=True),
            is_closed=True
        )

    def as_shape_group(self, color=None) -> pydiffvg.ShapeGroup:
        # TODO proper id handling
        return pydiffvg.ShapeGroup(
            shape_ids=torch.tensor([0]),
            fill_color=torch.tensor(color or self.color, requires_grad=True)
        )

    def get_points(self):
        out = []
        for p in self.points:
            out.append([p.x, p.y])
            for cp in p.controls:
                out.append([cp.x, cp.y])

        return out

    @classmethod
    def random(cls, degree=4, num_control_points=2):
        num_control_points = [num_control_points] * degree

        # Random tweaks to regular polygon base
        angle = 2 * math.pi / degree
        angle = uniform(0.8 * angle, 1.2 * angle)
        points = []
        for i in range(degree):
            pt = Point(
                uniform(0.3, 0.7) + 0.5 * math.cos(i * angle),
                uniform(0.3, 0.7) + 0.5 * math.sin(i * angle)
            )

            # Stochastically clamp to (0,1)
            for c in ['x', 'y']:
                if (v := getattr(pt, c)) > 1:
                    diff = v - 1
                    setattr(pt, c, v - uniform(diff, 2 * diff))
                elif v < 0:
                    diff = -v
                    setattr(pt, c, v + uniform(diff, 2 * diff))

            points.append(pt)

        for i in range(len(num_control_points)):
            pt = points[i]
            npt = points[i+1 if i+1 < degree else 0]

            ncp = num_control_points[i]
            dx = (npt.x - pt.x) / (ncp + 1)
            dy = (npt.y - pt.y) / (ncp + 1)

            for j in range(1, ncp+1):
                midpoint = Point(
                    pt.x + j * dx * uniform(0.8, 1.2) + uniform(0, 0.2),
                    pt.y + j * dy * uniform(0.8, 1.2) + uniform(0, 0.2)
                )

                pt.add_control(midpoint)

        out = cls(points)
        out.color = (
            uniform(0, 1),
            uniform(0, 1),
            uniform(0, 1),
            0.7
        )

        return out

    def __repr__(self):
        out = f"F<({rgb2hex(*self.color)})"
        out += f"[{', '.join([str(x) for x in self.points])}]>"
        return out


class Quad(Patch):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.top, self.right, self.bottom, self.left = self.points
        by_y = sorted(self.points, key=lambda pt: pt.y)
        self.top, self.bottom = by_y[0], by_y[-1]
        by_x = sorted([pt for pt in self.points
                       if pt not in [self.top, self.bottom]],
                      key=lambda pt: pt.x)
        self.left, self.right = by_x

        for pt in by_y:
            pt.round()

    @classmethod
    def from_path_points(cls, pp, color=None):
        pt = None
        points = []
        for i in range(len(pp)):
            if i % (len(pp) // 4) == 0:
                if pt:
                    points.append(pt)
                pt = Point(*pp[i])
            else:
                pt.add_control(Point(*pp[i]))

        points.append(pt)
        return cls(points, color)

    def to_path_points(self):
        out = []
        for pt in self.points:
            out.append(pt.as_xy())
            for cp in pt.controls:
                out.append(cp.as_xy())

        return out


class GradientMesh:
    def __init__(self, *quads: Quad):
        self.quads = quads

    def as_shape_groups(self):
        sg = [quad.as_shape_group() for quad in self.quads]
        for i in range(len(sg)):
            sg[i].shape_ids = torch.tensor([i])
        return sg

    def as_shapes(self, width, height):
        return [quad.as_path(width, height) for quad in self.quads]

    def to_numbers(self):
        points = []
        controls = []
        for quad in self.quads:
            qp = []
            qcp = []
            for p in quad.points:
                qp.append([p.x, p.y])
                qcp.append([[cp.x, cp.y] for cp in p.controls])
            points.append(qp)
            controls.append(qcp)

        return [points, controls, [q.color for q in self.quads]]

    def from_numbers(self, numbers):
        points, controls, colors = numbers
        for q in range(4):
            self.quads[q].color = colors[q]
            pts = points[q]
            ctrls = controls[q]
            for p in range(4):
                self.quads[q].points[p].replace(Point(*pts[p]))
                for c in range(len(self.quads[q].points[p].controls)):
                    self.quads[q].points[p].controls[c].replace(
                        Point(*ctrls[p][c])
                    )

    @classmethod
    def from_path_points(cls, pp, colors):
        a, b, c, d = [Quad.from_path_points(pp[x], colors[x])
                      for x in range(len(pp))]
        join_quads(a, b, c, d, scale=False, translate=False)

        return cls(a, b, c, d)

    def to_path_points(self):
        out = []
        for q in self.quads:
            out.append(q.to_path_points())
        return out

    def to_point_map(self):
        # XXX this doesn't work
        # because of control points
        pts: list[Point] = []
        template: list[list[int]] = []

        for quad in self.quads:
            for pt in quad.points:
                pts.append(pt)

        pts = list(dict.fromkeys(pts))

        for idx in range(len(self.quads)):
            template.append([
                pts.index(pt) for pt in self.quads[idx].points
            ])

        return (pts, template)

    def from_point_map(self, pts, template):
        # XXX this not either
        mapped = []
        for q in template:
            mapped.append([
                pts[x] for x in q
            ])

        return mapped


def average_points(points: list[Point]) -> Point:
    x = sum([pt.x for pt in points]) / len(points)
    y = sum([pt.y for pt in points]) / len(points)
    return Point(x, y)


def equalize_points(points: list[Point]):
    first = points[0]
    for pt in points:
        pt.equalize(first)


def equalize_cp(points: list[Point]):
    first = points[0].controls[0]
    for pt in points:
        for cp in pt.controls:
            cp.equalize(first)


def merge_points(points: list[Point]):
    merged = average_points(points)
    for pt in points:
        pt.replace(merged)


def merge_cp(points: list[Point]):
    ct1 = points[0].controls
    ct2 = points[1].controls[::-1]
    for i in range(len(ct1)):
        merged = average_points([ct1[i], ct2[i]])
        ct1[i].replace(merged)
        ct2[i].replace(merged)


def join_quads(a: Quad, b: Quad, c: Quad, d: Quad, scale=True, translate=True):
    if translate:
        b.translate(a.top)
        c.translate(a.right)
        d.translate(a.bottom)

    merge_points([a.right, b.bottom, c.left, d.top])

    merge_points([a.top, b.left])
    merge_points([a.bottom, d.left])
    merge_points([b.right, c.top])
    merge_points([c.bottom, d.right])

    if scale:
        a.scale(Point(0.5, 0.5))
        b.scale(Point(0.5, 0.5))
        c.scale(Point(0.5, 0.5))
        d.scale(Point(0.5, 0.5))

    merge_cp([a.right, d.left])
    merge_cp([a.top, b.bottom])
    merge_cp([c.left, b.right])
    merge_cp([c.bottom, d.top])

    equalize_points([a.right, b.bottom, c.left, d.top])
    equalize_points([a.top, b.left])
    equalize_points([a.bottom, d.left])
    equalize_points([b.right, c.top])
    equalize_points([c.bottom, d.right])

    equalize_cp([a.right, d.left])
    equalize_cp([a.top, b.bottom])
    equalize_cp([c.left, b.right])
    equalize_cp([c.bottom, d.top])

