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

from util import rgb2hex, any_map


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

        self.x = x
        self.y = y

        if round:
            self.round()

    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

        return self

    def add(self, pt: Point):
        self.x += pt.x
        self.y += pt.y

        return self

    def mult(self, pt: Point):
        self.x *= pt.x
        self.y *= pt.y

        return self

    @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 __hash__(self):
        # Used for removing duplicate points
        return id(self)

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

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


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

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

        for q in self.controls:
            for p in q:
                p.add(pt)

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

        for q in self.controls:
            for p in q:
                p.mult(pt)

    @classmethod
    def random(cls, degree=4, num_control_points=2):
        num_control_points = [num_control_points] * degree
        """Returns a random Patch with `degree` vertices
        and `num_control_points` control points per edge."""

        # 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)

        control_points = []
        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)

            control_points.append([
                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)
                ) for j in range(1, ncp+1)
            ])

        out = cls(points, control_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):
    """Quadrilateral bicubic patch."""
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # Assign 4 vertices as self.{top, bototm, left, right}
        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

        # Assign 4 edges as self.{northeat, southeast, southwest, northwest}
        self.northeast = self.controls[self.points.index(self.top)]
        self.southeast = self.controls[self.points.index(self.right)]
        self.southwest = self.controls[self.points.index(self.bottom)]
        self.northwest = self.controls[self.points.index(self.left)]

        self.set_points()

    def set_points(self):
        """Reset self.points and self.controls; used after mutating
        e.g. self.top, self.southwest &c."""
        self.points = [self.left, self.top,
                       self.right, self.bottom]

        self.controls = [self.northwest, self.northeast,
                         self.southeast, self.southwest]


class PointMapping:
    """Mapping of unique points in a mesh to separate shapes for diffvg."""
    def __init__(self, points, controls, raw_points, colors):
        self.points = points
        self.controls = controls
        self.raw_points = raw_points
        self.data = torch.tensor(
            [pt.as_xy() for pt in raw_points], requires_grad=True
        )
        self.colors = colors

    def as_shapes(self):
        out = []

        for i in range(len(self.points)):
            quad = self.points[i]

            quadpoints = []
            for j in range(len(quad)):
                quadpoints.append(quad[j])
                quadpoints += self.controls[i][j]

            out.append(quadpoints)

        return [torch.stack(x) for x in
                any_map(lambda idx: self.data[idx], out)]


class GradientMesh:
    """Bicubic quadrilateral mesh."""
    def __init__(self, *quads: Quad):
        self.quads = quads

    def as_mapping(self):
        """Convert GradientMesh to PointMapping"""
        points = []
        controls = []
        raw_points = []

        for q in self.quads:
            points.append(q.points)
            controls.append(q.controls)
            raw_points += q.points
            for cp in q.controls:
                raw_points += cp

        raw_points = list(set(raw_points))

        points = any_map(lambda p: raw_points.index(p), points)
        controls = any_map(lambda p: raw_points.index(p), controls)

        colors = torch.tensor([q.color for q in self.quads],
                              requires_grad=True)

        return PointMapping(points, controls, raw_points, colors)

    def from_mapping(self, mapping: PointMapping):
        pass


def average_points(*points: list[Point]) -> Point:
    """Average (i.e. geometric center) of points."""
    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 join_quads(a: Quad, b: Quad, c: Quad, d: Quad,
               scale=True, translate=True,
               step=10):
    """Join 4 quadrilaterals so that they form a mesh with a center point."""

    def combine_cp(cpa, cpb):
        return [average_points(a, b) for a, b in zip(cpa, cpb)]

    # If each quad occupies full space, make it so that they occupy 1/4 of space
    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))

    # If quads are on top of each other, translate so they are not
    if translate:
        b.translate(a.top)
        c.translate(a.right)
        d.translate(a.bottom)

    # Equalize centerpoint
    a.right, b.bottom, c.left, d.top = [
        average_points(a.right, b.bottom, c.left, d.top)
    ] * 4

    # Equalize non-center shared points
    a.top, b.left = [average_points(a.top, b.left)] * 2
    a.bottom, d.left = [average_points(a.bottom, d.left)] * 2
    b.right, c.top = [average_points(b.right, c.top)] * 2
    c.bottom, d.right = [average_points(c.bottom, d.right)] * 2

    # Equalize edges
    a.northeast, b.southwest = (
        (comb := combine_cp(a.northeast, b.southwest)),
        comb[::-1]
    )
    b.southeast, c.northwest = (
        (comb := combine_cp(b.southeast, c.northwest)),
        comb[::-1]
    )
    c.southwest, d.northeast = (
        (comb := combine_cp(c.southwest, d.northeast)),
        comb[::-1]
    )
    d.northwest, a.southeast = (
        (comb := combine_cp(d.northwest, a.southeast)),
        comb[::-1]
    )

    # Update points
    for q in [a, b, c, d]:
        q.set_points()
