from typing import Type
import numpy as np
import sys
from scipy.spatial.transform import Rotation
from math import sin, cos, hypot


class Vector(np.ndarray):
    """"""

    ATOL = 1e-15

    def __new__(cls, x: float, y: float, z: float = 0.0, *, w: float = 0.0):
        if w == 1.0:
            return np.asarray((x, y, z, w), dtype=float).view(Point)
        elif w == 0.0:
            return np.asarray((x, y, z, w), dtype=float).view(cls)
        else:
            return np.asarray((x / w, y / w, z / w, 1.0), dtype=float).view(Point)

    @property
    def kind(self):
        if np.isclose(self.w, 0.0, atol=self.ATOL):
            return sys.intern("Vector")
        elif np.isclose(self.w, 1.0, atol=self.ATOL):
            return sys.intern("Point")
        else:
            return sys.intern("PseudoPoint")

    def __bool__(self) -> bool:
        if self.kind is sys.intern("Vector"):
            return not np.all(np.isclose(self, 0.0, atol=self.ATOL))
        else:
            return True

    def __repr__(self):
        if self.kind is sys.intern("PseudoPoint"):
            return f"{self.kind}({self.x}, {self.y}, {self.z}, {self.w})"
        else:
            return f"{self.kind}({self.x}, {self.y}, {self.z})"

    def __str__(self):
        if self.kind is sys.intern("Vector"):
            return f"[{self.x:+.2f}, {self.y:+.2f}, {self.z:+.2f}]"
        elif self.kind is sys.intern("Point"):
            return f"({self.x:+.2f}, {self.y:+.2f}, {self.z:+.2f})"
        elif self.kind is sys.intern("PseudoPoint"):
            return f"{{{self.x:+.2f}, {self.y:+.2f}, {self.z:+.2f}}}"
        else:
            raise TypeError(f"Invalid Vector kind {self.kind}")

    def __abs__(self) -> float:
        return np.linalg.norm(self)

    def rotate(self, t: float, a: np.ndarray = None, c: np.ndarray = None):
        """\
        Rotates Vector in 't' radians according to the 'a' axis.
        Rotates Point in 't' radians according to the 'a' axis.
        """
        if self.kind is sys.intern("Vector"):
            if c is not None:
                raise TypeError("Can't rotate 'Vector' with respect to 'Point'.")
            return Transform.rotation(t, a) @ self
        elif self.kind is sys.intern("Point"):
            return c + (Transform.rotation(t, a) @ (self - c))
        elif self.kind is sys.intern("PseudoPoint"):
            raise TypeError("Can't rotate 'PseudoPoint'.")
        else:
            raise TypeError(f"Can't rotate '{type(self)}'.")

    def cross(self, other):
        if isinstance(other, Vector):
            v = np.cross(self[:3], other[:3])
            return Vector(*v)
        else:
            raise TypeError(f'Cross Product must be between vectors.')

    def unit(self):
        return self / abs(self)

    @property
    def x(self):
        return self[0]

    @property
    def y(self):
        return self[1]

    @property
    def z(self):
        return self[2]

    @property
    def w(self):
        return self[3]

    @property
    def xyz(self):
        return np.array([self.x, self.y, self.z], dtype=float)

    def adjust(self):
        if self.kind is sys.intern("Vector"):
            raise TypeError("Can't adjust 'Vector' to its homogeneous coordinate.")
        elif self.kind is sys.intern("Point"):
            return self
        elif self.kind is sys.intern("PseudoPoint"):
            return Point(self.x/self.w, self.y/self.w, self.z/self.w)
        else:
            raise TypeError(f"Can't adjust '{type(self)}' to its homogeneous coordinate.")
        


class Point(Vector):
    """"""

    def __new__(cls, x: float, y: float, z: float = 0.0, *, w: float = 1.0):
        if w == 1.0:
            return np.asarray((x, y, z, w), dtype=float).view(cls)
        elif w == 0.0:
            return np.asarray((x, y, z, w), dtype=float).view(Vector)
        else:
            return np.asarray((x / w, y / w, z / w, 1.0), dtype=float).view(cls)


class PseudoPoint(Point):
    """"""

    def __new__(cls, x: float, y: float, z: float, w: float):
        return np.asarray((x, y, z, w), dtype=float).view(Vector)


class Transform(np.ndarray):
    def __new__(cls, T: np.ndarray = None):
        if T is None:
            return np.eye(4, dtype=float).view(cls)
        else:
            return np.asarray(T, dtype=float).view(cls)

    def __matmul__(self, other):
        if isinstance(other, Point):
            p = np.ndarray.__matmul__(self, other)
            return Point(p[0], p[1], p[2], w=p[3]).adjust()
        elif isinstance(other, Vector):
            v = np.ndarray.__matmul__(self, other)
            return Vector(v[0], v[1], v[2])
        else:
            return np.ndarray.__matmul__(self, other)

    @classmethod
    def rotation(cls, t: float, a: Vector = None, c: Point = None) -> np.ndarray:
        """Rotation Matrix (4 x 4) in 't' radians according to the 'a' axis."""
        if a is None:
            u = Vector(0.0, 0.0, 1.0)
        else:
            u = a / np.linalg.norm(a)

        if c is None:
            c = Point(0.0, 0.0, 0.0)

        M = np.eye(4, dtype=float)
        U = np.eye(4, dtype=float)
        V = np.eye(4, dtype=float)

        M[:3, :3] = Rotation.from_rotvec(t * u.xyz).as_matrix()
        U[:3, 3] = c[:3]
        V[:3, 3] = -c[:3]
        return Transform(U @ M @ V)

    @classmethod
    def scale(cls, a: float):
        """Scale Matrix"""
        M = np.eye(4, dtype=float)
        M[:3, :3] *= a
        return Transform(M)


class Domain(object):
    def __init__(self, X: tuple, Y: tuple):
        self.ax, self.bx = X
        self.ay, self.by = Y

    def __contains__(self, p: Point):
        return (self.ax <= p.x <= self.bx) and (self.ay <= p.y <= self.by)

    def __str__(self):
        return f"[{self.ax}, {self.bx}] × [{self.ay}, {self.by}]"

    def __repr__(self):
        return (
            f"{self.__class__.__name__}(({self.ax}, {self.bx}), ({self.ax}, {self.by}))"
        )

    def __call__(self, m: int, n: int):
        if not isinstance(m, int) or not isinstance(n, int) or m < 1 or n < 1:
            raise TypeError("'m', 'n' must be positive integers greater than 2.")
        else:
            for i in range(m):
                yield tuple(Point(
                        x=self.ax * (i / (m - 1)) + self.bx * (1.0 - (i / (m - 1))),
                        y=self.ay * (j / (n - 1)) + self.by * (1.0 - (j / (n - 1))),
                    ) for j in range(n))

            for j in range(n):
                yield tuple(Point(
                    x=self.ax * (i / (m - 1)) + self.bx * (1.0 - (i / (m - 1))),
                    y=self.ay * (j / (n - 1)) + self.by * (1.0 - (j / (n - 1))),
                ) for i in range(m))


class Map(object):
    def __init__(self, f, D: Domain):
        self.f = f
        self.D = D

    def __contains__(self, p: Point) -> bool:
        return p in self.D

    def __call__(self, m: int, n: int) -> tuple:
        for line in self.D(m, n):
            yield tuple(Point(p.x, p.y, self.f(p.x, p.y)) for p in line)


class Surface(object):
    def __init__(self, *maps: Map):
        self.maps = maps
        self.T = Transform()

    def __call__(self, m: int, n: int) -> tuple:
        for M in self.maps:
            for line in M(m, n):
                yield tuple((self.T @ q).adjust() for q in line)

    def transform(self, T: Transform):
        self.T = self.T @ T

Vector.X = Vector(1.0, 0.0, 0.0)
Vector.Y = Vector(0.0, 1.0, 0.0)
Vector.Z = Vector(0.0, 0.0, 1.0)

Point.O = Point(0.0, 0.0, 0.0)