//! `LightTfrags`: the per-vertex lighting pass the game runs on tfrag geometry.
//!
//! Spec: `docs/formats/tfrag-lighting.md`.
//!
//! Every arithmetic step goes through [`crate::ps2`]. This is not
//! fastidiousness: the FPU truncates toward zero and has no denormals, and the
//! pass accumulates across every vertex of every patch. Computing it in `f32`
//! produces a visibly different level, not a subtly different one.

use gc_formats::tfrag::TfragGeometry;
use gc_formats::tfrag_light::{DirLightSet, LightBank, NormalTable};

use crate::ps2;

/// A vector of four PS2 float bit patterns.
type V4 = [u32; 4];

fn bits(v: [f32; 4]) -> V4 {
    v.map(f32::to_bits)
}

/// `vmul`: lane-wise product.
fn vmul(a: V4, b: V4) -> V4 {
    [0, 1, 2, 3].map(|k| ps2::mul(a[k], b[k]))
}

/// `vmulbc`: every lane times one scalar.
fn vmuls(a: V4, s: u32) -> V4 {
    a.map(|x| ps2::mul(x, s))
}

/// `vadda` / `vmadd`: `acc + a * s`, lane-wise.
fn vmadds(acc: V4, a: V4, s: u32) -> V4 {
    [0, 1, 2, 3].map(|k| ps2::add(acc[k], ps2::mul(a[k], s)))
}

/// `(ax*bx + ay*by) + 1.0*(az*bz)`.
///
/// The `1.0 *` is in the microprogram: `vmaddz.x r, vf21(1.0), t`. It matters
/// only for overflow behaviour, but dropping it changes the flush case.
fn dot3(a: V4, b: V4) -> u32 {
    let p = vmul(a, b);
    ps2::add(ps2::add(p[0], p[1]), ps2::mul(ps2::ONE, p[2]))
}

/// A colour byte as the float the microprogram builds: `65536 + c / 128`,
/// formed by `padduw` against `0x47800000`.
fn color_float(c: [u8; 4]) -> V4 {
    c.map(|b| 0x4780_0000 + b as u32)
}

/// Clamp to a byte with `vminibcx` against `0x478000ff`, then take the low
/// byte of each lane.
fn pack_color(v: V4) -> [u8; 4] {
    let clamp = 0x4780_00ff;
    [
        ps2::min(v[0], clamp) as u8,
        ps2::min(v[1], clamp) as u8,
        ps2::min(v[2], clamp) as u8,
        v[3] as u8,
    ]
}

/// PEXT5 of the 5:5:5:1 base colour.
pub fn pext5(c: u16) -> [u8; 4] {
    [
        ((c & 0x1f) << 3) as u8,
        ((c >> 5 & 0x1f) << 3) as u8,
        ((c >> 10 & 0x1f) << 3) as u8,
        ((c >> 15) << 7) as u8,
    ]
}

/// The light set a vertex uses, with the back-face factors pulled out of the
/// colour w lanes and those lanes cleared.
struct ActiveSet {
    ca: V4,
    da: V4,
    cb: V4,
    db: V4,
    wa: u32,
    wb: u32,
}

/// `select` names one light set, or blends two.
///
/// Bits 8..15 clear means "just set bits 0..3". Otherwise set bits 0..3 is
/// blended with set bits 4..7 by `t = bits 8..15 / 256`.
fn active_set(bank: &LightBank, select: u16) -> Option<ActiveSet> {
    let set = |i: u16| -> Option<DirLightSet> { bank.sets.get((i & 0xf) as usize).copied() };
    let (mut ca, da, mut cb, db);
    if select & 0xff00 == 0 {
        let s = set(select)?;
        (ca, da, cb, db) = (
            bits(s.color_a),
            bits(s.dir_a),
            bits(s.color_b),
            bits(s.dir_b),
        );
    } else {
        let t = ps2::itof12(((select as u32 >> 4) & 0xff0) as i32);
        let w = ps2::sub(ps2::ONE, t);
        let a = set(select)?;
        let b = set(select >> 4)?;
        let blend = |x: [f32; 4], y: [f32; 4]| -> V4 {
            let (p, q) = (vmuls(bits(x), w), vmuls(bits(y), t));
            [0, 1, 2, 3].map(|k| ps2::add(p[k], q[k]))
        };
        ca = blend(a.color_a, b.color_a);
        cb = blend(a.color_b, b.color_b);
        // Blended directions are re-normalised, which the unblended path does
        // not need because the stored ones are already unit.
        let norm = |v: V4| -> V4 {
            let q = ps2::rsqrt(ps2::ONE, dot3(v, v));
            [
                ps2::mul(v[0], q),
                ps2::mul(v[1], q),
                ps2::mul(v[2], q),
                v[3],
            ]
        };
        da = norm(blend(a.dir_a, b.dir_a));
        db = norm(blend(a.dir_b, b.dir_b));
    }
    // `vaddw.x/.y vf31 = 0 + colour.w`, then `vsubw.w colour = 1 - 1`.
    let (wa, wb) = (ps2::add(0, ca[3]), ps2::add(0, cb[3]));
    ca[3] = 0;
    cb[3] = 0;
    Some(ActiveSet {
        ca,
        da,
        cb,
        db,
        wa,
        wb,
    })
}

/// The pass's terms for one vertex.
///
/// This exists because a colour mismatch on its own cannot be attributed: the
/// final byte is a sum of five things, and the PCSX2 comparison -- the only
/// check on this pass there is -- is worthless without knowing *which* of them
/// the game disagrees with. It is a snapshot of what the pass computed, not a
/// second implementation: `light_vertex_n` and this share one body.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct VertexTerms {
    /// The record's azimuth and elevation indices.
    pub azimuth: u8,
    pub elevation: u8,
    /// The record's base colour, PEXT5-expanded.
    pub base: [u8; 4],
    /// The `select` word, the set it named, and whether it blended two.
    pub select: u16,
    pub set_index: Option<usize>,
    pub blended: bool,
    /// The vertex's normal.
    pub normal: [f32; 4],
    /// The active set's colours and directions, with the w lanes cleared as
    /// the pass clears them.
    pub color_a: [f32; 4],
    pub color_b: [f32; 4],
    pub dir_a: [f32; 4],
    pub dir_b: [f32; 4],
    /// The back-face factors taken from the colour w lanes.
    pub wa: f32,
    pub wb: f32,
    /// The two dot products, and each after its `max(d, d*w)`.
    pub dot_a: f32,
    pub dot_b: f32,
    pub max_a: f32,
    pub max_b: f32,
    /// The final colour.
    pub out: [u8; 4],
}

/// Light one vertex: `base + A.rgb * max(dA, dA*wA) + B.rgb * max(dB, dB*wB)`.
pub fn light_vertex(
    bank: &LightBank,
    n: V4,
    light: &gc_formats::tfrag::TfragLight,
) -> Option<[u8; 4]> {
    light_vertex_n(bank, n, light, light.select)
}

fn light_vertex_n(
    bank: &LightBank,
    n: V4,
    light: &gc_formats::tfrag::TfragLight,
    select: u16,
) -> Option<[u8; 4]> {
    Some(light_vertex_terms(bank, n, light, select)?.out)
}

/// Light one vertex, keeping the terms for the comparison.
pub fn light_vertex_terms(
    bank: &LightBank,
    n: V4,
    light: &gc_formats::tfrag::TfragLight,
    select: u16,
) -> Option<VertexTerms> {
    let s = active_set(bank, select)?;
    let da = dot3(n, s.da);
    let db = dot3(n, s.db);
    let da_max = ps2::max(da, ps2::mul(da, s.wa));
    let db_max = ps2::max(db, ps2::mul(db, s.wb));
    let base = pext5(light.color as u16);
    let acc = vmuls(color_float(base), ps2::ONE);
    let acc = vmadds(acc, s.ca, da_max);
    let out = pack_color(vmadds(acc, s.cb, db_max));
    let unbits = |v: V4| v.map(f32::from_bits);
    Some(VertexTerms {
        azimuth: light.azimuth_index(),
        elevation: light.elevation_index(),
        base,
        select,
        set_index: if select & 0xff00 == 0 {
            bank.sets
                .get((select & 0xf) as usize)
                .map(|_| (select & 0xf) as usize)
        } else {
            let low = bank.sets.get((select & 0xf) as usize);
            let high = bank.sets.get(((select >> 4) & 0xf) as usize);
            low.and(high).map(|_| (select & 0xf) as usize)
        },
        blended: select & 0xff00 != 0,
        normal: unbits(n),
        color_a: unbits(s.ca),
        color_b: unbits(s.cb),
        dir_a: unbits(s.da),
        dir_b: unbits(s.db),
        wa: f32::from_bits(s.wa),
        wb: f32::from_bits(s.wb),
        dot_a: f32::from_bits(da),
        dot_b: f32::from_bits(db),
        max_a: f32::from_bits(da_max),
        max_b: f32::from_bits(db_max),
        out,
    })
}

/// The normal a vertex lights with, as raw PS2 floats.
///
/// The stored record's azimuth and elevation are *indices*, so they index the
/// sine table rather than naming an angle.
///
/// All three components are negated: the stored normals point into the
/// surface. The z sign is what the PCSX2 comparison settled -- negating only x
/// and y still produces a plausible level, and it cost 48 percentage points of
/// agreement on level 0 (95.88% vs 100.00%) before it was measured.
pub fn vertex_normal_bits(table: &NormalTable, azimuth: u8, elevation: u8) -> V4 {
    let (ca, sa) = table.get(azimuth);
    let (ce, se) = table.get(elevation);
    let x = ps2::mul(ca.to_bits(), ce.to_bits());
    let y = ps2::mul(sa.to_bits(), ce.to_bits());
    [
        ps2::sub(0, x),
        ps2::sub(0, y),
        ps2::sub(0, se.to_bits()),
        se.to_bits(),
    ]
}

/// Run the lighting pass over a patch's light records.
///
/// Returns one colour per light record, in the same order. A record whose
/// `select` names a light set the bank does not have is left at its base
/// colour, which is what the game writes: the PCSX2 comparison shows those
/// records' RAM colours equal to their PEXT5 base exactly. Transparent black
/// was the first guess and it was wrong -- the vertex is unlit, not invisible.
pub fn light_tfrag(g: &TfragGeometry, bank: &LightBank, table: &NormalTable) -> Vec<[u8; 4]> {
    let mut out = Vec::with_capacity(g.lights.len());
    for l in &g.lights {
        let n = vertex_normal_bits(table, l.azimuth_index(), l.elevation_index());
        out.push(light_vertex_n(bank, n, l, l.select).unwrap_or(pext5(l.color as u16)));
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;
    use gc_formats::tfrag::TfragLight;

    fn table() -> NormalTable {
        let mut cos_sin = Vec::with_capacity(256);
        for i in 0..256usize {
            let a = i as f64 * std::f64::consts::PI / 128.0;
            cos_sin.push([(a.cos() as f32).to_bits(), (a.sin() as f32).to_bits()]);
        }
        NormalTable { cos_sin }
    }

    fn set(color: [f32; 4], dir: [f32; 4]) -> DirLightSet {
        DirLightSet {
            color_a: color,
            dir_a: dir,
            color_b: color,
            dir_b: dir,
        }
    }

    fn bank() -> LightBank {
        // One set: a white light pointing along -X. The table's azimuth 0
        // normal is -X too, so a vertex with that normal faces the light and
        // one with azimuth 128 faces away.
        LightBank {
            sets: vec![set([1.0, 1.0, 1.0, 1.0], [-1.0, 0.0, 0.0, 0.0])],
            point_lights: Vec::new(),
        }
    }

    fn light(azimuth: i8) -> TfragLight {
        TfragLight {
            unknown_0: 0,
            intensity: 0,
            azimuth,
            elevation: 0,
            color: 0,
            select: 0,
        }
    }

    #[test]
    fn pext5_expands_five_bits_to_eight() {
        assert_eq!(pext5(0x0000), [0, 0, 0, 0]);
        assert_eq!(pext5(0xffff), [248, 248, 248, 128]);
        assert_eq!(pext5(0x001f), [248, 0, 0, 0]);
    }

    #[test]
    fn a_facing_vertex_is_lit_and_a_back_facing_one_is_not() {
        let bank = bank();
        let t = table();
        // Azimuth 0 gives the normal -X, which faces the light. Azimuth 128
        // gives +X, which does not.
        let lit = light_tfrag(&patch(vec![light(0), light(-128)]), &bank, &t);
        assert_eq!(lit.len(), 2);
        assert!(
            lit[0][0] > 200,
            "facing vertex should be bright: {:?}",
            lit[0]
        );
        assert_eq!(
            lit[1][0], 0,
            "back-facing vertex should be dark: {:?}",
            lit[1]
        );
    }

    /// A patch with one light record, wired to `table` and `bank`.
    fn patch(lights: Vec<TfragLight>) -> TfragGeometry {
        TfragGeometry {
            lights,
            ..TfragGeometry::default()
        }
    }

    #[test]
    fn an_empty_bank_leaves_the_base_colour() {
        // The game's own answer, over PINE: a record whose select names a set
        // the bank does not have keeps its base colour. Transparent black was
        // the first guess, and the comparison read it as a 4.4% mismatch on
        // level 0 before this was measured.
        let bank = LightBank::default();
        let t = table();
        let mut l = light(0);
        // 5:5:5:1 with the top red step and the alpha bit set: (8, 0, 0,
        // 128). Stored as the i16 the record carries, expanded by PEXT5.
        l.color = 0x8001u16.cast_signed();
        let out = light_tfrag(&patch(vec![l]), &bank, &t);
        // PEXT5 of 1 is 1 << 3 = 8, with the alpha bit at 0x80.
        assert_eq!(out, vec![[8, 0, 0, 128]]);
    }

    #[test]
    fn normals_are_negated_and_unit_length() {
        let t = table();
        let n = vertex_normal_bits(&t, 0, 0);
        // Azimuth 0, elevation 0 is (1, 0, 0) before negation.
        assert_eq!(n[0], f32::to_bits(-1.0));
        assert_eq!(n[1], 0);
        assert_eq!(n[2], 0);
    }

    #[test]
    fn the_normal_s_z_is_negated_like_x_and_y() {
        // Elevation 64 is straight up: (0, 0, 1) before negation, so the
        // stored normal points down into the surface. Getting the z sign
        // wrong still produces a plausible level -- it flips which vertices
        // a light reaches -- and it cost 48 percentage points of agreement
        // before the PCSX2 comparison caught it.
        let t = table();
        let n = vertex_normal_bits(&t, 0, 64);
        let f = |b: u32| f32::from_bits(b);
        // The table's cos(pi/2) is 1.5e-8 rather than exactly zero, which is
        // what the game's own table holds, so x and y are ~0 rather than 0.
        assert!(f(n[0]).abs() < 1e-6, "x should be ~0, got {}", f(n[0]));
        assert!(f(n[1]).abs() < 1e-6, "y should be ~0, got {}", f(n[1]));
        assert_eq!(n[2], f32::to_bits(-1.0));
    }
}