//! A floating-point renderer of the model `wai-splat-render:i1` pins (§12): the same camera,
//! projection, low-pass, cut-off, alpha limits, early stop and order, in `f32` or `f64` with
//! the platform's `exp` and `sqrt`. It is the reference the integer renderer's quality is
//! measured against, and the illustration of why a float renderer cannot be signed: its
//! frames depend on the width it computes in. It is not part of any conformance check.

use super::{Planes, Pose, PLANES};

/// The arithmetic the renderer is written over.
pub trait Real: Copy + PartialOrd + core::ops::Add<Output = Self> + core::ops::Sub<Output = Self> + core::ops::Mul<Output = Self> + core::ops::Div<Output = Self> + core::ops::Neg<Output = Self> {
    fn of(v: f64) -> Self;
    fn to(self) -> f64;
    fn sqrt(self) -> Self;
    fn exp(self) -> Self;
}

impl Real for f32 {
    fn of(v: f64) -> Self {
        v as f32
    }
    fn to(self) -> f64 {
        self as f64
    }
    fn sqrt(self) -> Self {
        f32::sqrt(self)
    }
    fn exp(self) -> Self {
        f32::exp(self)
    }
}

impl Real for f64 {
    fn of(v: f64) -> Self {
        v
    }
    fn to(self) -> f64 {
        self
    }
    fn sqrt(self) -> Self {
        f64::sqrt(self)
    }
    fn exp(self) -> Self {
        f64::exp(self)
    }
}

type V3<R> = [R; 3];

fn dot<R: Real>(a: &V3<R>, b: &V3<R>) -> R {
    a[0] * b[0] + a[1] * b[1] + a[2] * b[2]
}

fn cross<R: Real>(a: &V3<R>, b: &V3<R>) -> V3<R> {
    [a[1] * b[2] - a[2] * b[1], a[2] * b[0] - a[0] * b[2], a[0] * b[1] - a[1] * b[0]]
}

fn norm<R: Real>(a: &V3<R>) -> V3<R> {
    let n = dot(a, a).sqrt();
    [a[0] / n, a[1] / n, a[2] / n]
}

fn clamp<R: Real>(x: R, lo: R, hi: R) -> R {
    if x < lo { lo } else if x > hi { hi } else { x }
}

/// The attribute values as a float decode presents them: `lo + (hi − lo)·u/255` per plane.
pub fn attributes<R: Real>(p: &Planes) -> Option<Vec<[R; PLANES]>> {
    if p.planes as usize != PLANES {
        return None;
    }
    let cells = (p.s * p.s) as usize;
    Some(
        (0..p.n_gaussians as usize)
            .map(|k| {
                std::array::from_fn(|i| {
                    let (lo, hi) = (R::of(p.ranges[i].0 as f64), R::of(p.ranges[i].1 as f64));
                    lo + (hi - lo) * R::of(p.bytes[i * cells + k] as f64) / R::of(255.0)
                })
            })
            .collect(),
    )
}

struct Splat<R> {
    key: R,
    index: usize,
    mx: R,
    my: R,
    x0: i64,
    x1: i64,
    y0: i64,
    y1: i64,
    conic: [R; 3],
    opacity: R,
    color: V3<R>,
}

/// The frame at `pose`, RGB8, rows top to bottom.
pub fn render<R: Real>(attrs: &[[R; PLANES]], pose: &Pose) -> Vec<u8> {
    let q = |v: i64| R::of(v as f64 / 65536.0);
    let eye = pose.eye.map(q);
    let fw = norm(&[q(pose.target[0]) - eye[0], q(pose.target[1]) - eye[1], q(pose.target[2]) - eye[2]]);
    let right = norm(&cross(&fw, &pose.up.map(q)));
    let up = cross(&right, &fw);
    let (w, h, fl) = (R::of(pose.width as f64), R::of(pose.height as f64), R::of(pose.focal as f64));
    let (zero, one, half) = (R::of(0.0), R::of(1.0), R::of(0.5));
    let mut splats = Vec::new();
    for (index, a) in attrs.iter().enumerate() {
        let d = [a[0] - eye[0], a[1] - eye[1], a[2] - eye[2]];
        let (xc, yc, zc) = (dot(&right, &d), dot(&up, &d), dot(&fw, &d));
        if zc < R::of(0.2) {
            continue;
        }
        let mx = w * half + fl * xc / zc;
        let my = h * half - fl * yc / zc;
        let lx = R::of(1.3) * w * half / fl * zc;
        let ly = R::of(1.3) * h * half / fl * zc;
        let (xcl, ycl) = (clamp(xc, -lx, lx), clamp(yc, -ly, ly));
        let j00 = fl / zc;
        let j02 = -fl * xcl / (zc * zc);
        let j12 = fl * ycl / (zc * zc);
        let t = [
            [j00 * right[0] + j02 * fw[0], j00 * right[1] + j02 * fw[1], j00 * right[2] + j02 * fw[2]],
            [-j00 * up[0] + j12 * fw[0], -j00 * up[1] + j12 * fw[1], -j00 * up[2] + j12 * fw[2]],
        ];
        let (qw, qx, qy, qz) = (a[6], a[7], a[8], a[9]);
        let n = qw * qw + qx * qx + qy * qy + qz * qz;
        let r = if n == zero {
            [[one, zero, zero], [zero, one, zero], [zero, zero, one]]
        } else {
            let two = R::of(2.0);
            [
                [(qw * qw + qx * qx - qy * qy - qz * qz) / n, two * (qx * qy - qw * qz) / n, two * (qx * qz + qw * qy) / n],
                [two * (qx * qy + qw * qz) / n, (qw * qw - qx * qx + qy * qy - qz * qz) / n, two * (qy * qz - qw * qx) / n],
                [two * (qx * qz - qw * qy) / n, two * (qy * qz + qw * qx) / n, (qw * qw - qx * qx - qy * qy + qz * qz) / n],
            ]
        };
        let s = [a[3], a[4], a[5]];
        let m: [[R; 3]; 3] = std::array::from_fn(|i| std::array::from_fn(|j| r[i][j] * s[j]));
        let s3: [[R; 3]; 3] = std::array::from_fn(|i| std::array::from_fn(|j| m[i][0] * m[j][0] + m[i][1] * m[j][1] + m[i][2] * m[j][2]));
        let u: [[R; 3]; 2] = std::array::from_fn(|i| std::array::from_fn(|k| t[i][0] * s3[0][k] + t[i][1] * s3[1][k] + t[i][2] * s3[2][k]));
        let vv = |i: usize, l: usize| u[i][0] * t[l][0] + u[i][1] * t[l][1] + u[i][2] * t[l][2];
        let (ca, cb, cc) = (vv(0, 0) + R::of(0.3), vv(0, 1), vv(1, 1) + R::of(0.3));
        let rx = (R::of(3.0) * ca.sqrt()).to().ceil() as i64;
        let ry = (R::of(3.0) * cc.sqrt()).to().ceil() as i64;
        if rx > super::MAX_RADIUS || ry > super::MAX_RADIUS {
            continue;
        }
        // Pixel i's centre is i + 1/2.
        let x0 = ((mx - R::of(rx as f64) - half).to().ceil() as i64).max(0);
        let x1 = ((mx + R::of(rx as f64) - half).to().floor() as i64).min(pose.width as i64 - 1);
        let y0 = ((my - R::of(ry as f64) - half).to().ceil() as i64).max(0);
        let y1 = ((my + R::of(ry as f64) - half).to().floor() as i64).min(pose.height as i64 - 1);
        let det = ca * cc - cb * cb;
        if x0 > x1 || y0 > y1 || det.partial_cmp(&zero) != Some(core::cmp::Ordering::Greater) {
            continue;
        }
        splats.push(Splat {
            key: dot(&d, &d),
            index,
            mx,
            my,
            x0,
            x1,
            y0,
            y1,
            conic: [cc / det, -cb / det, ca / det],
            opacity: clamp(a[13], zero, one),
            color: [clamp(a[10], zero, one), clamp(a[11], zero, one), clamp(a[12], zero, one)],
        });
    }
    splats.sort_by(|a, b| a.key.partial_cmp(&b.key).expect("finite").then(a.index.cmp(&b.index)));
    let (wi, hi) = (pose.width as usize, pose.height as usize);
    let mut t = vec![one; wi * hi];
    let mut acc = vec![[zero; 3]; wi * hi];
    let mut done = vec![false; wi * hi];
    let (amax, amin, tmin) = (R::of(0.99), R::of(1.0 / 255.0), R::of(1e-4));
    for s in &splats {
        for j in s.y0..=s.y1 {
            for i in s.x0..=s.x1 {
                let p = j as usize * wi + i as usize;
                if done[p] {
                    continue;
                }
                let dx = R::of(i as f64) + half - s.mx;
                let dy = R::of(j as f64) + half - s.my;
                let d2 = s.conic[0] * dx * dx + R::of(2.0) * s.conic[1] * dx * dy + s.conic[2] * dy * dy;
                if d2 > R::of(9.0) {
                    continue;
                }
                let mut alpha = s.opacity * (-half * d2).exp();
                if alpha > amax {
                    alpha = amax;
                }
                if alpha < amin {
                    continue;
                }
                let tn = t[p] * (one - alpha);
                if tn < tmin {
                    done[p] = true;
                    continue;
                }
                let tp = t[p];
                for (a, &c) in acc[p].iter_mut().zip(&s.color) {
                    *a = *a + c * alpha * tp;
                }
                t[p] = tn;
            }
        }
    }
    acc.iter().flat_map(|c| c.map(|v| (v.to() * 255.0 + 0.5).floor().clamp(0.0, 255.0) as u8)).collect()
}

/// PSNR in decibels between two RGB8 frames of the same size; infinite where they are equal.
pub fn psnr(a: &[u8], b: &[u8]) -> f64 {
    assert_eq!(a.len(), b.len());
    let se: f64 = a.iter().zip(b).map(|(&x, &y)| (x as f64 - y as f64).powi(2)).sum();
    if se == 0.0 { f64::INFINITY } else { 10.0 * (255.0f64 * 255.0 * a.len() as f64 / se).log10() }
}
