forgewright 0.2.0

use glam::Vec3;

pub type LightVec = [f32; 3];

/// Solve surface normal for a single pixel using photometric stereo.
pub fn solve_pixel_normal(
    intensities: &[f32],
    light_vectors: &[LightVec],
    specular_threshold: f32,
) -> Vec3 {
    assert_eq!(intensities.len(), light_vectors.len());
    let mut valid_i: Vec<f32> = Vec::new();
    let mut valid_l: Vec<LightVec> = Vec::new();
    for (idx, &intensity) in intensities.iter().enumerate() {
        if intensity > 5.0 && intensity < specular_threshold {
            valid_i.push(intensity);
            valid_l.push(light_vectors[idx]);
        }
    }
    if valid_i.len() < 3 { return Vec3::Z; }
    let n = valid_i.len();
    let mut ltl = [[0.0f32; 3]; 3];
    for k in 0..n {
        let l = valid_l[k];
        for i in 0..3 { for j in 0..3 { ltl[i][j] += l[i] * l[j]; } }
    }
    let mut lti = [0.0f32; 3];
    for k in 0..n {
        let l = valid_l[k];
        let intensity = valid_i[k];
        for i in 0..3 { lti[i] += l[i] * intensity; }
    }
    let inv = match invert_3x3(&ltl) { Some(m) => m, None => return Vec3::Z };
    let nx = inv[0][0]*lti[0] + inv[0][1]*lti[1] + inv[0][2]*lti[2];
    let ny = inv[1][0]*lti[0] + inv[1][1]*lti[1] + inv[1][2]*lti[2];
    let nz = inv[2][0]*lti[0] + inv[2][1]*lti[1] + inv[2][2]*lti[2];
    let len = (nx*nx + ny*ny + nz*nz).sqrt();
    if len < 1e-8 { return Vec3::Z; }
    Vec3::new(nx/len, ny/len, nz/len)
}

fn invert_3x3(m: &[[f32; 3]; 3]) -> Option<[[f32; 3]; 3]> {
    let det = m[0][0]*(m[1][1]*m[2][2] - m[1][2]*m[2][1])
            - m[0][1]*(m[1][0]*m[2][2] - m[1][2]*m[2][0])
            + m[0][2]*(m[1][0]*m[2][1] - m[1][1]*m[2][0]);
    if det.abs() < 1e-10 { return None; }
    let d = 1.0 / det;
    Some([
        [(m[1][1]*m[2][2]-m[1][2]*m[2][1])*d, (m[0][2]*m[2][1]-m[0][1]*m[2][2])*d, (m[0][1]*m[1][2]-m[0][2]*m[1][1])*d],
        [(m[1][2]*m[2][0]-m[1][0]*m[2][2])*d, (m[0][0]*m[2][2]-m[0][2]*m[2][0])*d, (m[0][2]*m[1][0]-m[0][0]*m[1][2])*d],
        [(m[1][0]*m[2][1]-m[1][1]*m[2][0])*d, (m[0][1]*m[2][0]-m[0][0]*m[2][1])*d, (m[0][0]*m[1][1]-m[0][1]*m[1][0])*d],
    ])
}

pub const SYNTHETIC_LIGHTS: [LightVec; 4] = [
    [-0.577, 0.577, 0.577], [0.577, 0.577, 0.577],
    [-0.577, -0.577, 0.577], [0.577, -0.577, 0.577],
];

/// Estimate surface normals from a single grayscale image via Sobel gradient.
pub fn normals_from_gradient(gray: &[u8], width: u32, height: u32) -> Vec<Vec3> {
    normals_from_gradient_scaled(gray, width, height, 4.0)
}

/// Gradient-to-normal with configurable depth sensitivity.
pub fn normals_from_gradient_scaled(gray: &[u8], width: u32, height: u32, strength: f32) -> Vec<Vec3> {
    let w = width as i32;
    let h = height as i32;
    let get = |x: i32, y: i32| -> f32 {
        if x < 0 || y < 0 || x >= w || y >= h { return 0.0; }
        gray[(y * w + x) as usize] as f32 / 255.0
    };
    let mut normals = Vec::with_capacity((width * height) as usize);
    for y in 0..h {
        for x in 0..w {
            if x == 0 || y == 0 || x >= w - 1 || y >= h - 1 { normals.push(Vec3::Z); continue; }
            let gx = -get(x-1,y-1) - 2.0*get(x-1,y) - get(x-1,y+1) + get(x+1,y-1) + 2.0*get(x+1,y) + get(x+1,y+1);
            let gy = -get(x-1,y-1) - 2.0*get(x,y-1) - get(x+1,y-1) + get(x-1,y+1) + 2.0*get(x,y+1) + get(x+1,y+1);
            normals.push(Vec3::new(-gx * strength, -gy * strength, 1.0).normalize());
        }
    }
    normals
}

/// Synthetically relight a grayscale image using estimated normals (Lambertian).
pub fn virtual_relight(gray: &[u8], normals: &[Vec3], light: &LightVec) -> Vec<f32> {
    let light_vec = Vec3::new(light[0], light[1], light[2]).normalize();
    gray.iter().zip(normals.iter()).map(|(&pixel, normal)| {
        let albedo = pixel as f32 / 255.0;
        normal.dot(light_vec).max(0.0) * albedo * 255.0
    }).collect()
}

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

    #[test]
    fn solve_normal_basic() {
        let lights = [[-0.577f32, 0.577, 0.577], [0.577, 0.577, 0.577], [0.0, 0.0, 1.0]];
        let intensities = [150.0, 150.0, 200.0];
        let normal = solve_pixel_normal(&intensities, &lights, 245.0);
        assert!(normal.z > 0.5, "Normal should face camera, got {:?}", normal);
    }

    #[test]
    fn solve_normal_specular_rejection() {
        let lights = [[-0.577f32, 0.577, 0.577], [0.577, 0.577, 0.577], [0.0, -0.577, 0.577], [0.0, 0.0, 1.0]];
        let intensities = [100.0, 255.0, 120.0, 140.0];
        let normal = solve_pixel_normal(&intensities, &lights, 245.0);
        assert!(normal.z > 0.0);
    }

    #[test]
    fn solve_normal_insufficient_data() {
        let lights = [[0.0f32, 0.0, 1.0], [0.577, 0.577, 0.577]];
        let intensities = [250.0, 250.0];
        assert_eq!(solve_pixel_normal(&intensities, &lights, 245.0), Vec3::Z);
    }

    #[test]
    fn invert_identity() {
        let m = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];
        let inv = invert_3x3(&m).unwrap();
        for i in 0..3 { for j in 0..3 {
            let expected = if i == j { 1.0 } else { 0.0 };
            assert!((inv[i][j] - expected).abs() < 1e-5);
        }}
    }

    #[test]
    fn invert_singular_returns_none() {
        assert!(invert_3x3(&[[1.0, 2.0, 3.0], [2.0, 4.0, 6.0], [1.0, 1.0, 1.0]]).is_none());
    }

    #[test]
    fn gradient_normals_flat_surface() {
        let gray = vec![128u8; 64 * 64];
        let normals = normals_from_gradient(&gray, 64, 64);
        for n in &normals[65..normals.len() - 65] { assert!(n.z > 0.9); }
    }

    #[test]
    fn virtual_relight_basic() {
        let relit = virtual_relight(&vec![128u8; 16], &vec![Vec3::Z; 16], &[0.0, 0.0, 1.0]);
        for &val in &relit { assert!((val - 128.0).abs() < 1.0); }
    }

    #[test]
    fn scaled_normals_steeper_than_default() {
        let size = 32u32;
        let mut gray = vec![0u8; (size * size) as usize];
        for y in 0..size { for x in 0..size { gray[(y * size + x) as usize] = (x * 255 / (size - 1)) as u8; } }
        let n1 = normals_from_gradient_scaled(&gray, size, size, 1.0);
        let n4 = normals_from_gradient_scaled(&gray, size, size, 4.0);
        let c = (size / 2 * size + size / 2) as usize;
        assert!(n4[c].z < n1[c].z);
    }
}