from typing import Set

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

from plane import Plane
from ray import Ray
from reflections import reflect, refract
from sphere import Sphere
from vector import Vector
from util import iszero


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))

surfaces = []

np.random.seed(5)

for i in range(5):
  sphere = Sphere(
    refraction_index=3,
    r0=Vector.fromiterable(np.random.uniform(-2, 2, size=2)),
    radius=np.random.uniform(0.5, 2),
  )
  surfaces.append(sphere)

  plt.gca().add_artist(plt.Circle(
    (sphere.r0.a[0, 0], sphere.r0.a[1, 0]),
    sphere.radius, color='blue', fill=False, zorder=200))
  plt.gca().add_artist(plt.Circle(
    (sphere.r0.a[0, 0], sphere.r0.a[1, 0]),
    sphere.radius, color='blue', fill=True, zorder=0, alpha=0.05))


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(5), color=hsv_to_rgb([hue, 1, 1]), zorder=depth + 5, lw=lw)
    return

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

  # normal = surface.normal_at(intersection)
  # new_n = current_n *

  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 * 1,
    current_n=current_n,
    depth=depth - 1)

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


raytrace(Ray(
  r0=Vector.fromargs(-3.5, 0),
  e=Vector.fromargs(2, -0.5)
), hue=0, lw=5, current_n=1, depth=8)

# for i in range(30):
#   trace_ray(Ray(
#     r0=Vector.fromargs(-4, -2),
#     e=Vector.fromargs(np.cos(i / 20 / np.pi), np.sin(i / 20 / np.pi)),
#   ))

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