from dataclasses import dataclass
from typing import Optional

import numpy as np

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


@dataclass
class Sphere(Surface):
  refraction_index: float
  r0: Vector
  radius: float

  def __and__(self, v: Vector) -> bool:
    return double_equality(abs(self.r0 - v), self.radius)

  def __contains__(self, v: Vector) -> bool:
    return self & v or abs(self.r0 - v) < self.radius

  def normal_at(self, v: Vector) -> Vector:
    return (v - self.r0).normalize()

  def intersect(self, ray: Ray) -> Optional[float]:
    b = -2 * ((self.r0 - ray.r0) @ ray.e)
    c = (self.r0 - ray.r0) @ (self.r0 - ray.r0) - self.radius ** 2

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

  def draw3d(self, painter: Painter):
    a1 = Vector.fromargs(0, 0, 1)
    a2 = Vector.fromargs(1, 0, 0)
    a3 = Vector.fromargs(0, 1, 0)

    painter.line3d(self.r0, self.r0 + a1, lw=5, zorder=0, color='red')
    painter.line3d(self.r0, self.r0 + a2, lw=5, zorder=0, color='green')
    painter.line3d(self.r0, self.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=self.r0, e=rotate_rodrigues(e, a1, long))
        vs.append(r.prolong(self.intersect(r)))

      painter.plot3d(vs, color='navy')

    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=self.r0, e=rotate_rodrigues(e, a3_e, lat))
        vs.append(r.prolong(self.intersect(r)))

      painter.plot3d(vs, color='indigo')
