import numpy as np
from matplotlib import pyplot as plt
from matplotlib.colors import hsv_to_rgb

from ellipse import Ellipse
from ray import Ray
from reflections import reflect, refract
from util import iszero
from vector import Vector, rotate2d

MAX_DEPTH = 6

STARTING_LW = 5

REFRACT_HUE_SHIFT = 1 / 8

LW_DECAY = 0.7


def draw_line(v1: Vector, v2: Vector, **kwargs):
  plt.plot(
    [v1.a[0, 0], v2.a[0, 0]],
    [v1.a[1, 0], v2.a[1, 0]],
    **kwargs)


def draw_point(v: Vector, **kwargs):
  plt.gca().add_artist(
    plt.Circle(
      (v.a[0, 0], v.a[1, 0]), 0.1,
      fill=False, **kwargs))


f = plt.figure(figsize=(8, 8), dpi=200)

surfaces = []

np.random.seed(2)


# ellipse = Ellipse(
#   refraction_index=1.3,
#   f1=Vector.fromargs(1, 0),
#   f2=Vector.fromargs(-1, 0),
#   distance=3,
# )
# ellipse.draw2d(plt)
# surfaces.append(ellipse)


def raytrace(
    ray: Ray,
    hue: float,
    lw: float,
    current_n: float,
    depth: int):
  if 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:
    draw_line(ray.r0, ray.prolong(10),
              color=hsv_to_rgb([hue, lw / STARTING_LW, 1]),
              zorder=depth + 5, lw=lw)
    return

  intersection = ray.prolong(min_t)
  draw_line(ray.r0, intersection,
            color=hsv_to_rgb([hue, lw / STARTING_LW, 1]),
            zorder=depth + 5, lw=lw)

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

  # reflection
  reflected = reflect(ray.e, normal)
  raytrace(
    Ray(intersection, reflected),
    hue=hue,
    lw=lw * LW_DECAY,
    current_n=current_n,
    depth=depth - 1)

  # refraction
  refracted = refract(ray.e, normal, current_n, new_n)
  raytrace(
    Ray(intersection, refracted),
    hue=(hue + REFRACT_HUE_SHIFT) % 1,
    lw=lw * LW_DECAY,
    current_n=new_n,
    depth=depth - 1)


ellipse1 = Ellipse(
  refraction_index=3,
  f1=Vector.fromargs(2, -2),
  f2=Vector.fromargs(3, 2),
  distance=5,
)
ellipse1.draw2d(plt, color='blue')
surfaces.append(ellipse1)

ellipse2 = Ellipse(
  refraction_index=3,
  f1=Vector.fromargs(-1, -2),
  f2=Vector.fromargs(-1, 2),
  distance=4.2,
)
ellipse2.draw2d(plt, color='blue')
surfaces.append(ellipse2)

r0 = Vector.fromargs(-3, 0)
e = Vector.fromargs(1, 0)

for alpha in np.linspace(-np.pi / 4, np.pi / 4, 9):
  raytrace(Ray(
    r0=r0,
    e=rotate2d(e, alpha)
  ), hue=0, lw=STARTING_LW, current_n=1, depth=MAX_DEPTH)

plt.xlim(-4, 4)
plt.ylim(-4, 4)
# plt.legend()
plt.grid(True, zorder=1)
plt.show()
