import dataclasses
from typing import List

from matplotlib.colors import hsv_to_rgb

from geometry.ray import Ray
from geometry.reflections import reflect, refract
from geometry.surface import Surface
from util import iszero


@dataclasses.dataclass
class RaytracingSettings:
  STARTING_LW: float
  REFLECT_LW_DECAY: float
  REFRACT_LW_DECAY: float

  STARTING_HUE: float
  REFLECT_HUE_SHIFT: float
  REFRACT_HUE_SHIFT: float

  MAX_DEPTH: int

  lw: float = None
  hue: float = None
  depth: float = None
  n: float = 1

  def __post_init__(self):
    if self.lw is None:
      self.lw = self.STARTING_LW

    if self.hue is None:
      self.hue = self.STARTING_HUE

    if self.depth is None:
      self.depth = self.MAX_DEPTH

  def copy(self):
    return dataclasses.replace(self)

  def reflect(self):
    return dataclasses.replace(
      self,
      lw=self.lw * self.REFLECT_LW_DECAY,
      hue=(self.hue + self.REFLECT_HUE_SHIFT) % 1,
      depth=self.depth - 1,
    )

  def refract(self, new_n):
    return dataclasses.replace(
      self,
      lw=self.lw * self.REFRACT_LW_DECAY,
      hue=(self.hue + self.REFRACT_HUE_SHIFT) % 1,
      depth=self.depth - 1,
      n=new_n,
    )


def raytrace(ray: Ray,
             settings: RaytracingSettings,
             axis,
             surfaces: List[Surface],
             line_drawer):
  if settings.depth == 0:
    return

  min_t = None
  surface = None

  for cur_surface in surfaces:
    t = cur_surface.intersect(ray)
    if t is None or iszero(t):
      continue

    if min_t is None or t < min_t:
      min_t = t
      surface = cur_surface

  if min_t is None:
    line_drawer(axis, ray.r0, ray.prolong(1),
                color=hsv_to_rgb([settings.hue, settings.lw / settings.STARTING_LW, 1]),
                zorder=settings.depth + 5,
                linestyle='--',
                lw=settings.lw)
    return

  intersection = ray.prolong(min_t)
  line_drawer(axis, ray.r0, intersection,
              color=hsv_to_rgb([settings.hue, settings.lw / settings.STARTING_LW, 1]),
              zorder=settings.depth + 5,
              linestyle='-',
              lw=settings.lw)

  if ray.r0 not in surface:
    # moves into
    normal = surface.normal_at(intersection)
    new_n = settings.n * surface.refraction_index
  else:
    # moves outside
    normal = -surface.normal_at(intersection)
    new_n = settings.n / surface.refraction_index

  # reflection
  reflected = reflect(ray.e, normal)
  raytrace(Ray(intersection, reflected), settings.reflect(), axis, surfaces, line_drawer)

  # refraction
  refracted = refract(ray.e, normal, settings.n, new_n)
  raytrace(Ray(intersection, refracted), settings.refract(new_n), axis, surfaces, line_drawer)
