from dataclasses import dataclass
from typing import Optional

import numpy as np

from geometry.ray import Ray
from geometry.surface import Surface
from geometry.line_drawers import draw_line_3d
from geometry.vector import Vector
from geometry.rotations import rotate2d, rotate3d, rotate_rodrigues
from util import double_equality, minpos, solve_quadratic


@dataclass
class Ellipse(Surface):
  refraction_index: float
  f1: Vector
  f2: Vector
  distance: float

  def __post_init__(self):
    if abs(self.f1 - self.f2) > self.distance:
      raise ValueError()

  def distance_to(self, v):
    return abs(self.f1 - v) + abs(self.f2 - v)

  def __and__(self, v: Vector) -> bool:
    return double_equality(self.distance_to(v), self.distance)

  def __contains__(self, v: Vector) -> bool:
    return self & v or self.distance_to(v) < self.distance

  def normal_at(self, v: Vector) -> Vector:
    d1 = (v - self.f1).normalize()
    d2 = (v - self.f2).normalize()
    return (d1 + d2).normalize()

  def intersect(self, ray: Ray) -> Optional[float]:
    D = self.distance
    b1 = 2 * (ray.e @ (ray.r0 - self.f1))
    b2 = 2 * (ray.e @ (ray.r0 - self.f2))
    c1 = (ray.r0 - self.f1) @ (ray.r0 - self.f1)
    c2 = (ray.r0 - self.f2) @ (ray.r0 - self.f2)

    a = 4 * D ** 2 - (b1 - b2) ** 2
    b = 2 * D ** 2 * (b1 + b2) - 2 * b1 * c1 + 2 * b1 * c2 + 2 * b2 * c1 - 2 * b2 * c2
    c = 2 * D ** 2 * (c1 + c2) - (c1 - c2) ** 2 - D ** 4

    return minpos(solve_quadratic(a, b, c))

  def draw2d(self, plt, **kwargs):
    r0 = (self.f1 + self.f2) / 2
    e = Vector.fromargs(1, 0)

    xs = []
    ys = []

    for alpha in np.linspace(0, 2 * np.pi, 200):
      r = Ray(r0=r0, e=rotate2d(e, alpha))
      p = r.prolong(self.intersect(r))
      xs.append(p.a[0, 0])
      ys.append(p.a[1, 0])

      # n = self.normal_at(p)
      # p2 = p + n
      # plt.plot(
      #   [p.a[0, 0], p2.a[0, 0]],
      #   [p.a[1, 0], p2.a[1, 0]],
      # )

    plt.plot(xs, ys, **kwargs)

  def draw3d_old(self, axis, **kwargs):
    r0 = (self.f1 + self.f2) / 2
    e = Vector.fromargs(0, 0, 1)

    N1 = 20
    N2 = 100

    for alpha in np.linspace(0, 2 * np.pi, N1):
      xs = []
      ys = []
      zs = []

      for beta in np.linspace(0, 2 * np.pi, N2):
        d = Vector.fromargs(
          np.sin(alpha) * np.cos(beta),
          np.sin(alpha) * np.sin(beta),
          np.cos(alpha),
        )

        r = Ray(r0=r0, e=d)
        p = r.prolong(self.intersect(r))

        xs.append(p.a[0, 0])
        ys.append(p.a[1, 0])
        zs.append(p.a[2, 0])

      axis.plot(xs, ys, zs, color='blue', **kwargs)

      xs = []
      ys = []
      zs = []

      for beta in np.linspace(0, 2 * np.pi, N2):
        d = Vector.fromargs(
          np.sin(beta) * np.cos(alpha),
          np.sin(beta) * np.sin(alpha),
          np.cos(beta),
        )

        r = Ray(r0=r0, e=d)
        p = r.prolong(self.intersect(r))

        xs.append(p.a[0, 0])
        ys.append(p.a[1, 0])
        zs.append(p.a[2, 0])

      axis.plot(xs, ys, zs, color='blue', **kwargs)

  def draw3d(self, axis, **kwargs):
    r0 = (self.f1 + self.f2) / 2
    a1 = (self.f2 - r0).normalize()

    a2_candidates = [
      a1 | Vector.fromargs(1, 0, 0),
      a1 | Vector.fromargs(0, 1, 0),
      a1 | Vector.fromargs(0, 0, 1),
    ]
    a2 = a2_candidates[np.argmax(np.abs(a2_candidates))]
    a2 = a2.normalize()
    a3 = (a1 | a2).normalize()

    draw_line_3d(axis, r0, r0 + a1, lw=5, zorder=0, color='red')
    draw_line_3d(axis, r0, r0 + a2, lw=5, zorder=0, color='green')
    draw_line_3d(axis, r0, r0 + a3, lw=5, zorder=0, color='blue')

    N1 = 8
    N2 = 100

    for lat in np.linspace(0, np.pi, N1):
      vs = []
      e = rotate_rodrigues(a1, a2, lat)

      for long in np.linspace(0, 2 * np.pi, N2):
        r = Ray(r0=r0, e=rotate_rodrigues(e, a1, long))
        p = r.prolong(self.intersect(r))

        vs.append(p.a)

      V = np.hstack(vs)
      axis.plot(V[0, :], V[1, :], V[2, :], color='navy', **kwargs)

    for long in np.linspace(0, np.pi, N1):
      vs = []
      e = rotate_rodrigues(a2, a1, long)
      a3_e = rotate_rodrigues(a3, a1, long)

      for lat in np.linspace(0, 2 * np.pi, N2):
        r = Ray(r0=r0, e=rotate_rodrigues(e, a3_e, lat))
        p = r.prolong(self.intersect(r))

        vs.append(p.a)

      V = np.hstack(vs)
      axis.plot(V[0, :], V[1, :], V[2, :], color='indigo', **kwargs)
