import numpy as np
from matplotlib import pyplot as plt
# noinspection PyUnresolvedReferences
from mpl_toolkits.mplot3d import Axes3D

from geometry.ellipse import Ellipse
from geometry.line_drawers import draw_line_3d
from geometry.ray import Ray
from geometry.rotations import rotate_rodrigues
from geometry.vector import Vector
from rendering import raytracing
from rendering.raytracing import RaytracingSettings

figure = plt.figure(figsize=(8, 8), dpi=200)
axis = figure.add_subplot(111, projection='3d')
axis.view_init(elev=10, azim=-80)

surfaces = []

ellipse = Ellipse(
  refraction_index=1.5,
  f1=Vector.fromargs(-2, 2, 1),
  f2=Vector.fromargs(1, 0, -2),
  distance=5.8,
)
surfaces.append(ellipse)
ellipse.draw3d(axis)

axis.set_xlim(-4, 4)
axis.set_ylim(-4, 4)
axis.set_zlim(-4, 4)

axis.set_xlabel('x')
axis.set_ylabel('y')
axis.set_zlabel('z')

axis.grid(True)

settings = RaytracingSettings(
  STARTING_LW=2,
  STARTING_HUE=0,

  REFLECT_HUE_SHIFT=0,
  REFLECT_LW_DECAY=0.7,

  REFRACT_HUE_SHIFT=1 / 8,
  REFRACT_LW_DECAY=0.8,

  MAX_DEPTH=3
)

for alpha in np.linspace(-np.pi / 32, np.pi / 32, 4):
  for beta in np.linspace(-np.pi / 32, np.pi / 32, 4):
    e = rotate_rodrigues(
      Vector.fromargs(-1, 0, 0),
      Vector.fromargs(0, 1, 0),
      alpha,
    )
    e = rotate_rodrigues(
      e,
      Vector.fromargs(0, 0, 1),
      beta,
    )

    raytracing.raytrace(
      Ray(
        r0=Vector.fromargs(4, 0, 0),
        e=e
      ),
      settings,
      axis,
      surfaces,
      draw_line_3d
    )

figure.show()

for i, azim in enumerate(np.linspace(0, 360, 16)):
  axis.view_init(elev=10, azim=azim)
  figure.savefig(f'movie/frame{i:03d}.png')
