from typing import Set

import numpy as np
from matplotlib import pyplot as plt

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 plot_v(v1: Vector, v2: Vector, *args, **kwargs):
  plt.plot(
    [v1.a[0, 0], v2.a[0, 0]],
    [v1.a[1, 0], v2.a[1, 0]],
    *args, **kwargs)


intersectables = []
f = plt.figure(figsize=(6, 6))
ax = plt.gca()

# declare a plane
plane = Plane(
  r0=Vector.fromargs(1, 0),
  n=Vector.fromargs(-2, 1),
)
intersectables.append((plane, 1.5))
plt.plot([-1, 3], [-4, 4], label='плоскость')

# declare a sphere
sphere = Sphere(
  r0=Vector.fromargs(0, 2),
  radius=0.8,
)
intersectables.append((sphere, 2))
circle = plt.Circle(
  (sphere.r0.a[0, 0], sphere.r0.a[1, 0]),
  sphere.radius, color='blue', fill=False)
ax.add_artist(circle)


def raytrace(
    ray: Ray,
    current_n: float,
    already_intersected: Set,
    depth: int):
  if depth == 0:
    return

  min_t = None
  surface = None
  surface_n = None

  for intersectable, intersectable_n in intersectables:
    t = intersectable.intersect(ray)
    if t is None or iszero(t):
      continue

    if min_t is None or t < min_t:
      min_t = t
      surface = intersectable
      surface_n = intersectable_n

  if min_t is None:
    plot_v(ray.r0, ray.prolong(5), color='red')
    return

  intersection = ray.prolong(min_t)
  plot_v(ray.r0, intersection, color='orange')

  normal = surface.normal_at(intersection)
  new_n = surface_n
  new_already_intersected = set(already_intersected)

  if id(surface) not in already_intersected:
    # moves_into
    new_already_intersected.add(id(surface))
  else:
    # moves outside
    new_already_intersected.remove(id(surface))
    normal = -1 * normal
    new_n = 1

  # reflection
  reflected = reflect(ray.e, normal)
  raytrace(Ray(intersection, reflected), current_n, already_intersected, depth - 1)

  # refraction
  refracted = refract(ray.e, normal, current_n, new_n)
  raytrace(Ray(intersection, refracted), new_n, new_already_intersected, depth - 1)


raytrace(Ray(
  r0=Vector.fromargs(-4, 2),
  e=Vector.fromargs(2, -0.5)
), 1, set(), 10)

# 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(-5, 5)
plt.ylim(-5, 5)
# plt.legend()
plt.grid(True)
plt.show()
