#!/usr/bin/env catnip
# Raymarcher géométrique : caméra orientée par rotor, intersections analytiques
#
# La caméra est orientée par un ROTOR R = s + B (scalaire + bivecteur). Le
# bivecteur B encode le PLAN de rotation ; son dual est l'axe. Chaque rayon de la
# base caméra est tourné par le produit sandwich v' = R v R̃, sans matrice de vue.
# Sous forme fermée v' = v + 2s(B×v) + 2 B×(B×v), et numpy applique ce sandwich à
# TOUT le tableau de rayons en une expression (le cross se diffuse sur (N,3)).
#
# La scène est résolue par intersection analytique rayon/sphère (racine du
# trinôme du second degré) et rayon/plan (t = (P0−O)·n / d·n), puis ombrée par un
# terme diffus de Lambert n·l. Catnip organise scène, matériaux et rendu ; numpy
# porte l'arithmétique vectorielle. Sorties : un PNG et deux erreurs numériques.
#
# DEPS: numpy, pillow

numpy = import('numpy')
math = import('math')
Image = import('PIL.Image')
import('pathlib', 'Path')

script_dir = Path(META.file).parent
output_dir = script_dir / 'output'
output_dir.mkdir(exist_ok=True)

EPS = 0.0001

struct Vec3 {
    x: float; y: float; z: float;

    arr(self) => { numpy.array(list(self.x, self.y, self.z)) }
}

from_arr = (a): Vec3 => { Vec3(float(a[0]), float(a[1]), float(a[2])) }

# Normalise chaque ligne d'un tableau (N,3) : divise par sa norme euclidienne.
normalize_rows = (v) => { v / numpy.linalg.norm(v, axis=1, keepdims=True) }

# Rotor = partie scalaire + bivecteur (composantes duales de l'axe de rotation).
struct Rotor {
    s: float; b: Vec3;

    # Sandwich R v R̃ vectorisé : appliqué à un tableau (N,3) de rayons d'un coup.
    # numpy.cross(bivecteur, rayons) diffuse le produit croisé sur toutes les lignes.
    apply_many(self, vs) => {
        bv = self.b.arr()
        t = 2.0 * numpy.cross(bv, vs)
        vs + self.s * t + numpy.cross(bv, t)
    }
}

rotor_from_axis_angle = (axis: Vec3, angle: float): Rotor => {
    k = axis.arr() / numpy.linalg.norm(axis.arr())
    half = angle / 2.0
    Rotor(math.cos(half), from_arr(math.sin(half) * k))
}

# Matrice de rotation de référence (Rodrigues), chemin indépendant du rotor pour
# le contrôle numérique : même (axe, angle), formule matricielle classique.
rotation_matrix = (axis: Vec3, angle: float) => {
    k = axis.arr() / numpy.linalg.norm(axis.arr())
    cross_k = numpy.array(list(
            list(0.0, -k[2], k[1]),
            list(k[2], 0.0, -k[0]),
            list(-k[1], k[0], 0.0),
        ))
    numpy.eye(3) * math.cos(angle) + math.sin(angle) * cross_k + (1.0 - math.cos(angle)) * numpy.outer(k, k)
}

struct Material {
    color: list[int]; ambient: float; diffuse: float;

    # Éclairage lambertien vectorisé : ambiante fixe + diffus n·l par pixel.
    shade(self, normals, light) => {
        diff = numpy.maximum(0.0, normals.dot(light))
        intensity = self.ambient + self.diffuse * diff
        base = numpy.array(self.color)
        numpy.clip(base * intensity.reshape(-1, 1), 0, 255)
    }
}

# Résultat d'une intersection sur tout le champ de rayons : distance t (inf où
# aucune intersection) et couleur ombrée par pixel.
struct Hit { t; color; }

# Contrat commun des objets de scène : chacun sait résoudre le champ de rayons
# entier et rendre un Hit. Le rendu (plus bas) dispatche dessus sans connaître le
# type concret — scene[i].render() suffit.
struct SceneObject {
    @abstract
    render(self, origin, dirs, light): Hit
}

struct Sphere extends(SceneObject) {
    center: Vec3; radius: float; material: Material;

    # Intersection analytique rayon/sphère : ‖O + t·d − C‖² = r² donne le trinôme
    # t² + b·t + c avec d unitaire (a=1). On garde la plus petite racine positive.
    render(self, origin, dirs, light): Hit => {
        oc = origin - self.center.arr()
        b = 2.0 * dirs.dot(oc)
        c = oc.dot(oc) - self.radius ** 2
        disc = b * b - 4.0 * c
        root = (-b - numpy.sqrt(numpy.maximum(disc, 0.0))) / 2.0
        hit = numpy.logical_and(disc > 0.0, root > EPS)
        t_geo = numpy.where(hit, root, 0.0)
        point = origin + t_geo.reshape(-1, 1) * dirs
        normal = normalize_rows(point - self.center.arr())
        Hit(numpy.where(hit, root, numpy.inf), self.material.shade(normal, light))
    }
}

struct Plane extends(SceneObject) {
    point: Vec3; normal: Vec3; material: Material;

    # Intersection rayon/plan : t = (P0 − O)·n / (d·n), défini si d·n ≠ 0. La
    # normale d'ombrage est retournée pour faire face au rayon incident.
    render(self, origin, dirs, light): Hit => {
        n = self.normal.arr()
        denom = dirs.dot(n)
        safe = numpy.where(numpy.abs(denom) < EPS, EPS, denom)
        t = (self.point.arr() - origin).dot(n) / safe
        hit = numpy.logical_and(numpy.abs(denom) > EPS, t > EPS)
        facing = numpy.where((denom < 0.0).reshape(-1, 1), n, -n)
        Hit(numpy.where(hit, t, numpy.inf), self.material.shade(facing, light))
    }
}

struct Camera {
    position: Vec3; rotor: Rotor; focal: float;

    # Rayons de la base caméra (avant rotation) : caméra en position, axe optique
    # vers −z, un rayon par pixel. Le rotor les oriente ensuite via apply_many.
    base_dirs(self, w: int, h: int) => {
        grid = numpy.meshgrid(numpy.arange(w), numpy.arange(h))
        px = grid[0].reshape(-1)
        py = grid[1].reshape(-1)
        dx = (px - w / 2.0 + 0.5) / self.focal
        dy = -(py - h / 2.0 + 0.5) / self.focal
        dz = numpy.full(w * h, -1.0)
        normalize_rows(numpy.column_stack(list(dx, dy, dz)))
    }
}

struct RenderConfig {
    width: int; height: int; background: list[int]; light: Vec3;
}

config = RenderConfig(320, 320, list(20, 22, 30), Vec3(-0.6, 1.0, 0.8))

# Le rotor oriente la caméra : léger basculement autour d'un axe proche de x pour
# plonger le regard vers le plan-sol. C'est le seul lien caméra → orientation.
rotor = rotor_from_axis_angle(Vec3(1.0, 0.3, 0.0), 0.18)
camera = Camera(Vec3(0.0, 0.6, 5.0), rotor, 320.0)

red = Material(list(220, 70, 70), 0.15, 0.85)
blue = Material(list(70, 120, 230), 0.15, 0.85)
yellow = Material(list(230, 200, 70), 0.15, 0.85)
floor = Material(list(120, 120, 130), 0.20, 0.60)

scene = list(
    Sphere(Vec3(-1.3, 0.0, 0.0), 0.9, red),
    Sphere(Vec3(1.3, 0.2, -0.5), 0.7, blue),
    Sphere(Vec3(0.0, -0.3, 1.2), 0.6, yellow),
    Plane(Vec3(0.0, -1.0, 0.0), Vec3(0.0, 1.0, 0.0), floor),
)

print(f"⇒ Rotor caméra : s={round(rotor.s, 4)}, " +
        f"bivecteur=({round(rotor.b.x, 4)}, {round(rotor.b.y, 4)}, {round(rotor.b.z, 4)})")

base = camera.base_dirs(config.width, config.height)
# Le sandwich rotor tourne le champ de rayons entier vers la direction visée.
rays = camera.rotor.apply_many(base)

# Contrôle 1 : une rotation est une isométrie, elle préserve la norme. Les rayons
# de base sont unitaires ; après rotor ils doivent le rester.
iso_error = float(numpy.max(numpy.abs(numpy.linalg.norm(rays, axis=1) - 1.0)))
print(f"  contrôle isométrie (‖rotor·d‖ − 1) : erreur max = {iso_error}")

# Contrôle 2 : le sandwich rotor vs la matrice de Rodrigues (même axe, même
# angle), deux chemins indépendants, sur tout le champ de rayons.
ref_matrix = rotation_matrix(Vec3(1.0, 0.3, 0.0), 0.18)
ref_rays = base.dot(ref_matrix.T)
rot_error = float(numpy.max(numpy.linalg.norm(ref_rays - rays, axis=1)))
print(f"  contrôle rotor vs matrice de Rodrigues : erreur max = {rot_error}")

# Rendu : pour chaque objet on résout tout le champ de rayons, puis on garde le
# plus proche par pixel (numpy.where sélectionne t et couleur là où c'est devant).
origin = camera.position.arr()
light = config.light.arr() / numpy.linalg.norm(config.light.arr())
n_pixels = config.width * config.height
t_best = numpy.full(n_pixels, numpy.inf)
color = numpy.zeros(list(n_pixels, 3)) + numpy.array(config.background)

i = 0
while i < len(scene) {
    hit = scene[i].render(origin, rays, light)
    closer = hit.t < t_best
    color = numpy.where(closer.reshape(-1, 1), hit.color, color)
    t_best = numpy.where(closer, hit.t, t_best)
    i = i + 1
}

covered = int(numpy.sum(numpy.array(t_best < numpy.inf)))
print(f"  pixels touchant la scène : {covered} / {n_pixels}")

pixels = color.reshape(config.height, config.width, 3).astype('uint8')
image = Image.fromarray(pixels, 'RGB')
output_path = output_dir / 'ga_raymarch.png'
image.save(str(output_path))
print()
print(f"⇒ Image → {output_path}")

# Aperçu navigateur : sert l'image une fois puis rend la main.
# --no-browser garde un chemin headless (le PNG reste écrit ci-dessus).
if '--no-browser' not in import('sys').argv {
    http = import('http')
    b64 = import('base64').b64encode(output_path.read_bytes()).decode('ascii')
    http.serve(f'<!doctype html><meta charset="utf-8"><title>Geometric algebra raymarch</title><body style="margin:0;background:#0d1117"><img style="max-width:100%;display:block;margin:0 auto" src="data:image/png;base64,{b64}">',
        0, 'text/html; charset=utf-8', True)
}