import itertools
import numpy as np
import math
from sympy import N, Number, Rational, init_printing, latex, simplify, Expr
from sympy.vector import CoordSys3D, Vector

init_printing()

import sympy

unit = lambda v: v/np.linalg.norm(v)

def perspective_matrix(vfov, width, height, left, right, bottom, top, near, far):
  aspect = width / height
  return np.array([
    [1.0 / math.tan(vfov / 2.0) / aspect, 0, 0, 0],
    [0, 1.0 / math.tan(vfov / 2.0), 0, 0],
    [0, 0, -(far + near) / (far - near), -2.0 * far * near / (far - near)],
    [0, 0, -1, 0]
  ])
  # return np.array([
  #   [2.0 * near / (right - left), 0, (right + left) / (right - left), 0],
  #   [0, 2.0 * near / (top - bottom), (top + bottom) / (top - bottom), 0],
  #   [0, 0, -(far + near) / (far - near), -(2.0 * far * near) / (far - near)],
  #   [0, 0, -1, 0],
  # ])

def view_matrix(camera_pos, view_dir, up_dir):
  n = unit(-view_dir)
  u = unit(np.cross(up_dir, n))
  v = np.cross(n, u)
  return np.array([
    [u[0], u[1], u[2], -np.dot(camera_pos, u)],
    [v[0], v[1], v[2], -np.dot(camera_pos, v)],
    [n[0], n[1], n[2], -np.dot(camera_pos, n)],
    [0, 0, 0, 1],
  ])

def print_trans(before, after):
  def style_vec(v):
    start = "$[\\begin{matrix}"
    mid = str(v[0]) + " & " + str(v[1]) + " & " + str(v[2])
    end = "\\end{matrix}]$"
    return f"{start}{mid}{end}"
  return style_vec(before) + " $\\rightarrow$ " + style_vec(after)

def compute_view(near, vfov, hfov):
  width = 2.0 * near * math.tan(hfov / 2.0)
  height = 2.0 * near * math.tan(hfov / 2.0)

  left = -width / 2.0
  right = width / 2.0
  bottom = -height / 2.0
  top = height / 2.0
  return width, height, left, right, bottom, top

def print_bmatrix(arr):
  for row in arr:
    for j, col in enumerate(row):
      end = " " if j == len(row) - 1 else " & "
      print(col, end=end)
    print("\\\\")

def problem_1():
  C = CoordSys3D('C')
  p = 1 * C.i + 4 * C.j + 8 * C.k
  e = 0 * C.i + 0 * C.j + 0 * C.k
  s = 2 * C.i + 2 * C.j + 10 * C.k
  radius = 3

  v0 = p - e
  print("v0 (incoming ray) =", v0)
  print(" - |v0| =", v0.magnitude())
  print()

  n = p - s
  print("n (normal) =", n)
  n_norm = n.normalize()
  print(" - n_norm (normalized) =", n_norm)
  print()

  cos_theta_i = (-v0).dot(n) / (v0.magnitude() * n.magnitude())
  theta_i = sympy.acos(cos_theta_i)
  print("[1a] theta_i (angle of incidence) =", theta_i)
  print()

  proj_len = cos_theta_i * v0.magnitude()
  proj = n_norm * proj_len
  print("projection of v0 onto n =", proj)

  nx = p + proj
  print("nx (proj point) =", nx)
  print()

  v1 = nx - e
  print("v1 (from e to nx) =", v1)

  z = e + 2 * v1
  print("z =", z)

  v3 = z - p
  print("v3 =", v3)
  print()

  sin_theta_i = sympy.sin(theta_i)
  print("sin(theta_i) =", sin_theta_i)

  theta_t = sympy.asin(Rational(2, 3) * sin_theta_i)
  print("[1d] theta_t (angle of transmission) =", theta_t)
  print(" - approx =", theta_t.evalf())
  print()

  v5_len = sympy.tan(theta_t) * (s - p).magnitude()
  print("v5 len =", v5_len)
  v4 = (s - p) + v1.normalize() * v5_len
  print("v4 =", v4)
  print(" - [1e] normalized =", v4.normalize())
  print()

  print("s - p =", s - p)
  print("|s - p| =", (s - p).magnitude())
  print("|v1| =", v1.normalize())
  print("tan(theta_t) =", sympy.tan(theta_t))

def problem_4():
  camera_pos = np.array([2, 3, 5])
  view_dir = np.array([1, -1, -1])
  up_dir = np.array([0, 1, 0])
  V = view_matrix(camera_pos, view_dir, up_dir)
  print(V)

  f = np.vectorize(lambda c: (c * c))
  print(f(V))


def build_translation_matrix(vec):
  return np.array([
    [1, 0, 0, vec[0]],
    [0, 1, 0, vec[1]],
    [0, 0, 1, vec[2]],
    [0, 0, 0, 1],
  ])

def problem_5():
  b1 = build_translation_matrix(np.array([0, 0, 5]))
  theta = math.radians(-90)
  sin_theta = round( math.sin(theta), 5)
  cos_theta = round(math.cos(theta), 5)
  b2 = np.array([
    [cos_theta, 0, sin_theta, 0],
    [0, 1, 0, 0],
    [-sin_theta, 0, cos_theta, 0],
    [0, 0, 0, 1],
  ])
  b3 = build_translation_matrix(np.array([0, 0, -5]))

  print("b1", b1)
  print("b2", b2)
  print("b3", b3)
  M = b3 @ b2 @ b1
  print("M", M)

  ex1 = np.array([1, 1, -4, 1])
  print("ex1", ex1, M @ ex1)

  up = np.array([0, 1, 0])
  n = np.array([1, 0, 0])
  u = unit(np.cross(up, n))
  v = np.cross(n, u)

  print(f"{up = }, {n = }, {u = }, {v = }")

  eye = np.array([5, 0, -5])
  dx = -(np.dot(eye, u))
  dy = -(np.dot(eye, v))
  dz = -(np.dot(eye, n))
  print(f"{dx = }, {dy = }, {dz = }")

def problem_8():
  near = 0.5
  far = 20



  print("part 8a")
  vfov = hfov = math.radians(60)
  width, height, left, right, bottom, top = compute_view(near, vfov, hfov)
  print(perspective_matrix(vfov, width, height, left, right, bottom, top, near, far))
  print()

def problem_7():

  def solve(camera_pos, angle):
    angle_radians = math.radians(angle)
    near = 1
    far = 10
    view_dir = np.array([0, 0, -1])
    up_dir = np.array([0, 1, 0])
    width, height, left, right, bottom, top = compute_view(near, angle_radians, angle_radians)
    print("faces of the viewing frustum", left, right, bottom, top)
    P = perspective_matrix(angle_radians, width, height, left, right, bottom, top, near, far)
    V = view_matrix(camera_pos, view_dir, up_dir)
    print("P")
    print_bmatrix(np.around(P, 4))
    print("V")
    print_bmatrix(np.around(V, 4))
    m = P @ V

    points = [
      np.array([0.5, 0.5, -4]),
      np.array([0.5, -0.5, -4]),
      np.array([-0.5, -0.5, -4]),
      np.array([-0.5, 0.5, -4]),

      np.array([0.5, 0.5, -5]),
      np.array([0.5, -0.5, -5]),
      np.array([-0.5, -0.5, -5]),
      np.array([-0.5, 0.5, -5]),

      np.array([1, 1.5, -7]),
      np.array([1, -0.5, -7]),
      np.array([-1, -0.5, -7]),
      np.array([-1, 1.5, -7]),

      np.array([1, 1.5, -9]),
      np.array([1, -0.5, -9]),
      np.array([-1, -0.5, -9]),
      np.array([-1, 1.5, -9]),
    ]

    for point in points:
      point_ = np.r_[point, [1]]
      trans = m @ point_
      def style_vec(v):
        start = "$[\\begin{matrix}"
        mid = str(v[0]) + " & " + str(v[1]) + " & " + str(v[2])
        end = "\\end{matrix}]$"
        return f"{start}{mid}{end}"
      print("-", style_vec(point), "$\\rightarrow$", style_vec(np.around(trans[:3], 4)))

  print("Part A")
  camera_pos = np.array([0, 0, 0])
  angle = 90
  solve(camera_pos, angle)

  print("Part B")
  camera_pos = np.array([0, 0, -2])
  angle = 90
  solve(camera_pos, angle)

  print("Part C")
  camera_pos = np.array([0, 0, 0])
  angle = 45
  solve(camera_pos, angle)

def problem_6():
  # sqrt3 = math.sqrt(3)
  sqrt3 = sympy.sqrt(3)
  width = 2 * sqrt3

  def ortho_transform(left, right, bottom, top, near, far):
    step_1 = np.array([
      [1, 0, 0, 0],
      [0, 1, 0, 0],
      [0, 0, -1, 0],
      [0, 0, 0, 1],
    ])

    step_2 = np.array([
      [1, 0, 0, -((left + right) / 2)],
      [0, 1, 0, -((bottom + top) / 2)],
      [0, 0, 1, -((near + far) / 2)],
      [0, 0, 0, 1],
    ])

    step_3 = np.array([
      [(2 / (right - left)), 0, 0, 0],
      [0, (2 / (top - bottom)), 0, 0],
      [0, 0, (2 / (far - near)), 0],
      [0, 0, 0, 1],
    ])

    return step_3 @ step_2 @ step_1

  def oblique_transform(left, right, bottom, top, near, far):
    step_0 = np.array([
      [1, 0, (1 / sqrt3), 0],
      [0, 1, 0, 0],
      [0, 0, 1, 0],
      [0, 0, 0, 1],
    ])

    M_ortho = ortho_transform(left, right, bottom, top, near, far)

    return M_ortho @ step_0

  def calculate(M, points):
    left = min(map(lambda p: p[0], points))
    right = max(map(lambda p: p[0], points))
    bottom = min(map(lambda p: p[1], points))
    top = max(map(lambda p: p[1], points))
    near = min(map(lambda p: p[2], points))
    far = max(map(lambda p: p[2], points))

    for point in points:
      point_ = np.r_[point, [1]]
      trans_ = M(left, right, bottom, top, near, far) @ point_
      trans = trans_[:3]
      v = np.vectorize(lambda x: float(x.evalf()))
      l = np.vectorize(lambda x: latex(simplify(x)))
      point = l(point)
      trans = l(trans)
      print("-", print_trans(point, trans))

  cube_center = np.array([0, 0, -3 * sqrt3])

  points = []
  for (dz, dx, dy) in itertools.product([-1, 1], [-1, 1], [-1, 1]):
    point = cube_center + np.array([
      dx * width / 2,
      dy * width / 2,
      dz * width / 2,
    ])
    points.append(point)

  print("Helosu Ortho")
  calculate(ortho_transform, points)

  print("Helosu Oblique")
  calculate(oblique_transform, points)

def problem_9():
  pass

print("\nPROBLEM 8 -------------------------"); problem_8()
print("\nPROBLEM 5 -------------------------"); problem_5()
print("\nPROBLEM 9 -------------------------"); problem_9()
print("\nPROBLEM 7 -------------------------"); problem_7()
print("\nPROBLEM 4 -------------------------"); problem_4()
print("\nPROBLEM 6 -------------------------"); problem_6()
print("\nPROBLEM 1 -------------------------"); problem_1()