from __future__ import annotations

import logging
import math
import sys
from abc import ABC, abstractmethod
from collections import deque
from dataclasses import dataclass
from enum import Enum
from typing import (
    TYPE_CHECKING,
    Any,
    List,
    Optional,
    Sequence,
    Tuple,
    Union,
)

import numpy as np
from scipy.ndimage import label

from .grid import Grid
from .types import *

if TYPE_CHECKING:
    from . import Implicit


@dataclass(frozen=True)
class Lip:
    """
    Describes a vertical or horizontal line segment describing an interpolated point

    The point (ix1,iy1) must satisfy f(ix1,iy1) < altitude, while (ix2,iy2) must satisfy
    f(ix2,iy2) >= altitude.

    The lip must be horizontal (iy1 == iy2), or vertical (ix1 == ix2).

    The points (ix1,iy1) and (ix2,iy2) must lie on a grid generated by the lip delta.
    """

    #: Grid description
    grid: Grid

    #: Altitude of the isoline to plot
    altitude: float

    #: X coordinate of the point defining the lip, below the altitude
    ix1: int

    #: Y coordinate of the point defining the lip, below the altitude
    iy1: int

    #: Function value f(ix1,iy1)
    v1: float

    #: X coordinate of the point defining the lip, above or at altitude
    ix2: int

    #: Y coordinate of the point defining the lip, above or at altitude
    iy2: int

    #: Function value f(ix2,iy2)
    v2: float

    def __eq__(self, other: Any) -> bool:
        if not isinstance(other, Lip):
            return False
        return (
            self.ix1 == other.ix1
            and self.iy1 == other.iy1
            and self.ix2 == other.ix2
            and self.iy2 == other.iy2
        )

    def unpack(self) -> Tuple[int, int, float, int, int, float]:
        """
        Helper method the get information from the lip
        """
        return (self.ix1, self.iy1, self.v1, self.ix2, self.iy2, self.v2)

    def contains(self, other: Lip) -> bool:
        """
        Returns whether this lip contains another lip which refines it

        This method returns false if:

        - The lips differ in their orientation

        If both lips are vertical, this method returns false if:

        - The lips are not aligned on the same x coordinate.

        - The y interval of this lip does not include the y interval of the other lip

        If both lips are horizontal, this method returns false if:

        - The lips are not aligned on the same y coordinate.

        - The x interval of this lip does not include the x interval of the other lip

        Args:
            other: Other lip which may refine this lip
        """
        if self.is_horizontal():
            if not other.is_horizontal():
                return False
            if self.iy1 != other.iy1:
                return False
            x_min = min(self.ix1, self.ix2)
            x_max = max(self.ix1, self.ix2)
            return (x_min <= other.ix1 <= x_max) and (x_min <= other.ix2 <= x_max)
        else:
            assert self.is_vertical()
            if not other.is_vertical():
                return False
            if self.ix1 != other.ix1:
                return False
            y_min = min(self.iy1, self.iy2)
            y_max = max(self.iy1, self.iy2)
            return (y_min < other.iy1 <= y_max) and (y_min <= other.iy2 <= y_max)

    def on_border(self) -> bool:
        """
        Returns whether this lip is on the grid border
        """
        if self.is_vertical():
            return self.ix1 == 0 or self.ix1 == self.grid.x_divs
        else:
            assert self.is_horizontal()
            return self.iy1 == 0 or self.iy1 == self.grid.y_divs

    def is_vertical(self) -> bool:
        """
        Returns whether this lip is vertical
        """
        return self.ix1 == self.ix2

    def is_horizontal(self) -> bool:
        """
        Returns whether this lip is horizontal
        """
        return self.iy1 == self.iy2

    def delta(self) -> int:
        """
        Returns the length of the line segment supporting this lip
        """
        return max(abs(self.ix1 - self.ix2), abs(self.iy1 - self.iy2))

    def __post_init__(self) -> None:
        assert self.v1 < self.altitude
        assert self.v2 >= self.altitude
        assert self.is_horizontal() != self.is_vertical()
        delta = self.delta()
        ld = math.log2(delta)
        assert ld == round(ld)
        assert np.all(np.mod(np.array([self.ix1, self.iy1, self.ix2, self.iy2]), delta) == 0)

    def distance_squared(self, other: Union[Lip, Tuple[int, int]]) -> int:
        """
        Returns the distance squared of this lip to another lip or point
        """

        def norm_squared(dx: int, dy: int) -> int:
            return dx * dx + dy * dy

        if isinstance(other, Lip):
            d11 = norm_squared(self.ix1 - other.ix1, self.iy1 - other.iy1)
            d12 = norm_squared(self.ix1 - other.ix2, self.iy1 - other.iy2)
            d21 = norm_squared(self.ix2 - other.ix1, self.iy2 - other.iy1)
            d22 = norm_squared(self.ix2 - other.ix2, self.iy2 - other.iy2)
            return min(d11, d12, d21, d22)
        else:
            ix = other[0]
            iy = other[1]
            d1 = norm_squared(self.ix1 - ix, self.iy1 - iy)
            d2 = norm_squared(self.ix2 - ix, self.iy2 - iy)
            return min(d1, d2)

    def split(self, implicit: Implicit) -> Lip:
        """
        Returns this lip split in half

        Raises:
            ValueError: when this lip has :meth:`~.Lip.delta` == 1

        Args:
            implicit: Function to plot
        """
        ix1, iy1, v1, ix2, iy2, v2 = self.unpack()
        c1 = v1 >= self.altitude
        c2 = v2 >= self.altitude
        if self.delta() == 1:
            raise ValueError("Cannot split in half a lip with delta=1")
        assert self.delta() > 1
        ix = (ix1 + ix2) // 2
        iy = (iy1 + iy2) // 2
        v = implicit.eval(ix, iy)
        c = v >= self.altitude
        if c == c1:
            ix1 = ix
            iy1 = iy
            v1 = v
        else:
            ix2 = ix
            iy2 = iy
            v2 = v
        return Lip(self.grid, self.altitude, ix1, iy1, v1, ix2, iy2, v2)

    def split_until(self, delta: int, implicit: Implicit) -> Lip:
        """
        Returns the lip of specified delta inside the current lip

        Args:
            delta: Target size
            implicit: Function to plot
        """
        assert delta <= self.delta() and delta > 0
        res = self
        while delta < res.delta():
            res = res.split(implicit)
        return res

    def shared_point(self, other: Lip) -> Optional[Tuple[int, int]]:
        """
        Returns the point shared by two lips, if any

        Args:
            other: Other lip to check this lip with
        """
        if (self.ix1 == other.ix1 and self.iy1 == other.iy1) or (
            self.ix1 == other.ix2 and self.iy1 == other.iy2
        ):
            return (self.ix1, self.iy1)
        if (self.ix2 == other.ix1 and self.iy2 == other.iy1) or (
            self.ix2 == other.ix2 and self.iy2 == other.iy2
        ):
            return (self.ix2, self.iy2)
        return None

    def real_coordinates(self) -> Tuple[float, float]:
        """
        Returns the real coordinates of the linearly interpolated point
        """
        t = (self.altitude - self.v1) / (self.v2 - self.v1)
        x1, y1 = self.grid.real_coordinates(self.ix1, self.iy1)
        x2, y2 = self.grid.real_coordinates(self.ix2, self.iy2)
        x = (1 - t) * x1 + t * x2
        y = (1 - t) * y1 + t * y2
        return float(x), float(y)

    @staticmethod
    def from_below_above_real_pair(
        implicit: Implicit,
        below: Tuple[float, float],
        above: Tuple[float, float],
        delta: int,
    ) -> Tuple[Lip, int]:
        bx, by, delta = implicit.find_integer_coordinates(below[0], below[1], False, delta)
        ax, ay, delta = implicit.find_integer_coordinates(above[0], above[1], True, delta)
        return Lip.from_below_above_integer_pair((bx, by), (ax, ay), delta, implicit)

    @staticmethod
    def from_below_above_integer_pair(
        below: Tuple[int, int],
        above: Tuple[int, int],
        delta: int,
        implicit: Implicit,
    ) -> Tuple[Lip, int]:
        bx, by = below
        ax, ay = above
        assert (bx % delta) == 0 and (by % delta) == 0
        assert (ax % delta) == 0 and (ay % delta) == 0
        bv = implicit.eval(bx, by)
        av = implicit.eval(ax, ay)
        assert av >= implicit.altitude
        assert bv < implicit.altitude
        while abs(ax - bx) > delta or abs(ay - by) > delta:
            x = round((ax + bx) / 2 / delta) * delta
            y = round((ay + by) / 2 / delta) * delta
            v = implicit.eval(x, y)
            if v >= implicit.altitude:
                ax = x
                ay = y
                av = v
            else:
                bx = x
                by = y
                bv = v
        assert ax != bx or ay != by
        is_proper_lip = ax == bx or ay == by
        if not is_proper_lip:
            if implicit.is_above(ax, by):
                ay = by
            else:
                bx = ax
        lip = Lip(implicit.grid, implicit.altitude, bx, by, bv, ax, ay, av)
        return lip, lip.delta()

    @staticmethod
    def from_bisection(
        below: Optional[Tuple[int, int]],
        above: Optional[Tuple[int, int]],
        delta: int,
        implicit: Implicit,
    ) -> Tuple[Lip, int]:
        x_divs = implicit.grid.x_divs
        y_divs = implicit.grid.y_divs
        squares: deque[Tuple[int, int, int, int]] = deque([(0, 0, x_divs, y_divs)])
        while squares:
            ix1, iy1, ix2, iy2 = squares.popleft()
            for pt in [(ix1, iy1), (ix1, iy2), (ix2, iy1), (ix2, iy2)]:
                if below is not None and above is not None:
                    delta = np.gcd.reduce([delta, *below, *above])
                    return Lip.from_below_above_integer_pair(below, above, delta, implicit)
                x, y = pt
                if implicit.is_above(x, y):
                    above = (x, y)
                else:
                    below = (x, y)
            if abs(ix1 - ix2) > 1 and abs(iy1 - iy2) > 1:
                ix = (ix1 + ix2) // 2
                iy = (iy1 + iy2) // 2
                squares.append((ix1, iy1, ix, iy))
                squares.append((ix, iy1, ix2, iy))
                squares.append((ix1, iy, ix, iy2))
                squares.append((ix, iy, ix2, iy2))
        raise Exception("Could not find a lip from grid bisection")

    @staticmethod
    def from_grid(x_delta: int, y_delta: int, delta: int, implicit: Implicit) -> Sequence[Lip]:
        assert math.log2(x_delta).is_integer()
        assert math.log2(y_delta).is_integer()
        assert math.log2(delta).is_integer()
        assert delta <= x_delta and delta <= y_delta
        starting_lips: List[Lip] = []
        grid = implicit.grid
        x_divs = grid.x_divs
        y_divs = grid.y_divs
        n_x = x_divs // x_delta + 1
        n_y = y_divs // y_delta + 1
        features = np.zeros((n_x, n_y), dtype=i8)
        for x in range(n_x):
            for y in range(n_y):
                if implicit.is_above(x * x_delta, y * y_delta):
                    features[x, y] = 1
        xx, yy = np.where(features == 1)
        for i in range(len(xx)):
            ax = xx[i]
            ay = yy[i]
            assert implicit.is_above(ax * x_delta, ay * y_delta)

        n = label(features, output=features)
        for lbl in range(1, n + 1):
            start_lip: Optional[Lip] = None
            xx, yy = np.where(features == lbl)
            for i in range(len(xx)):
                ax = xx[i]
                ay = yy[i]
                assert implicit.is_above(ax * x_delta, ay * y_delta)
            for i in range(len(xx)):
                ax = xx[i]
                ay = yy[i]
                for bx, by in [(ax + 1, ay), (ax - 1, ay), (ax, ay + 1), (ax, ay - 1)]:
                    if 0 <= bx < n_x and 0 <= by < n_y and features[bx, by] == 0:
                        start_lip, _ = Lip.from_below_above_integer_pair(
                            (bx * x_delta, by * y_delta),
                            (ax * x_delta, ay * y_delta),
                            delta,
                            implicit,
                        )
                        break
                if start_lip is not None:
                    break
            assert start_lip is not None
            starting_lips.append(start_lip)
        return starting_lips
