from dataclasses import dataclass
from typing import Optional

import numpy as np

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


@dataclass
class TransformedSurface(Surface):
  refraction_index: float
  subsurface: Surface
  transform: AffineTransform

  def __contains__(self, v: Vector) -> bool:
    return self.transform.apply_shift(v) in self.subsurface

  def __and__(self, v: Vector) -> bool:
    return self.subsurface & self.transform.apply_shift(v)

  def normal_at(self, v: Vector) -> Vector:
    v_ = self.transform.apply_shift(v)
    result_ = self.subsurface.normal_at(v_)
    result = self.transform.invert_no_shift(result_)
    return result.normalize()

  def intersect(self, ray: Ray) -> Optional[float]:
    e_ = self.transform.apply_no_shift(ray.e)
    length_ratio = abs(e_)

    ray_ = Ray(
      r0=self.transform.apply_shift(ray.r0),
      e=e_.normalize(),
    )

    t_ = self.subsurface.intersect(ray_)
    if t_ is None:
      return None

    t = t_ / length_ratio

    return t

  def draw3d(self, painter: Painter):
    transformed_painter = TransformedPainter(painter, self.transform)

    self.subsurface.draw3d(transformed_painter)
