//! Machado, Oliveira & Fernandes (2009), severity 1.0 matrices in linear RGB.
//! https://www.inf.ufrgs.br/~oliveira/pubs_files/CVD_Simulation/CVD_Simulation.html
//! Model also used by Chromium: https://developer.chrome.com/docs/chromium/cvd
//! Matrix coefficients cross-checked against QGIS's published implementation:
//! https://api.qgis.org/api/qgsprevieweffect_8cpp_source.html
//! Pair warnings are our own heuristic, not a WCAG test or clinical prediction.
use super::Swatch;
use std::{collections::HashMap, sync::OnceLock};

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum Vision {
    #[default]
    Original,
    Protanopia,
    Deuteranopia,
    Tritanopia,
    Grayscale,
}

impl Vision {
    pub const CHECKS: [Self; 4] = [
        Self::Protanopia,
        Self::Deuteranopia,
        Self::Tritanopia,
        Self::Grayscale,
    ];
    pub fn label(self) -> &'static str {
        match self {
            Self::Original => "Original",
            Self::Protanopia => "Protanopia (red-cone loss)",
            Self::Deuteranopia => "Deuteranopia (green-cone loss)",
            Self::Tritanopia => "Tritanopia (blue-cone loss)",
            Self::Grayscale => "Grayscale (no color cues)",
        }
    }

    fn linear(self, rgb: [f64; 3]) -> [f64; 3] {
        let matrix = match self {
            Self::Original => return rgb,
            Self::Grayscale => return [luminance_linear(rgb); 3],
            Self::Protanopia => [
                [0.152286, 1.052583, -0.204868],
                [0.114503, 0.786281, 0.099216],
                [-0.003882, -0.048116, 1.051998],
            ],
            Self::Deuteranopia => [
                [0.367322, 0.860646, -0.227968],
                [0.280085, 0.672501, 0.047413],
                [-0.011820, 0.042940, 0.968881],
            ],
            Self::Tritanopia => [
                [1.255528, -0.076749, -0.178779],
                [-0.078411, 0.930809, 0.147602],
                [0.004733, 0.691367, 0.303900],
            ],
        };
        matrix.map(|row| {
            row.iter()
                .zip(rgb)
                .map(|(a, b)| a * b)
                .sum::<f64>()
                .clamp(0.0, 1.0)
        })
    }

    pub fn simulate(self, rgb: [u8; 3]) -> [u8; 3] {
        if self == Self::Original {
            return rgb;
        }
        self.linear(linear_rgb(rgb)).map(encode)
    }
}

fn decode(value: u8) -> f64 {
    let value = value as f64 / 255.0;
    if value <= 0.04045 {
        value / 12.92
    } else {
        ((value + 0.055) / 1.055).powf(2.4)
    }
}

fn encode(value: f64) -> u8 {
    let value = value.clamp(0.0, 1.0);
    let value = if value <= 0.0031308 {
        value * 12.92
    } else {
        1.055 * value.powf(1.0 / 2.4) - 0.055
    };
    (value * 255.0).round() as u8
}

pub fn linear_rgb(rgb: [u8; 3]) -> [f64; 3] {
    static LINEAR: OnceLock<[f64; 256]> = OnceLock::new();
    let table = LINEAR.get_or_init(|| std::array::from_fn(|i| decode(i as u8)));
    rgb.map(|v| table[v as usize])
}

fn luminance_linear([r, g, b]: [f64; 3]) -> f64 {
    0.2126 * r + 0.7152 * g + 0.0722 * b
}
pub fn luminance(rgb: [u8; 3]) -> f64 {
    luminance_linear(linear_rgb(rgb))
}

// OKLab reference transform: https://bottosson.github.io/posts/oklab/
fn oklab([r, g, b]: [f64; 3]) -> [f64; 3] {
    let l = (0.4122214708 * r + 0.5363325363 * g + 0.0514459929 * b).cbrt();
    let m = (0.2119034982 * r + 0.6806995451 * g + 0.1073969566 * b).cbrt();
    let s = (0.0883024619 * r + 0.2817188376 * g + 0.6299787005 * b).cbrt();
    [
        0.2104542553 * l + 0.7936177850 * m - 0.0040720468 * s,
        1.9779984951 * l - 2.4285922050 * m + 0.4505937099 * s,
        0.0259040371 * l + 0.7827717662 * m - 0.8086757660 * s,
    ]
}

fn distance(a: [f64; 3], b: [f64; 3]) -> f64 {
    a.iter()
        .zip(b)
        .map(|(a, b)| (a - b).powi(2))
        .sum::<f64>()
        .sqrt()
}

pub fn simulate_rgba(rgba: &[u8], mode: Vision) -> Vec<u8> {
    let mut output = rgba.to_vec();
    if mode == Vision::Original {
        return output;
    }
    let mut cache = HashMap::new();
    for pixel in output.as_chunks_mut::<4>().0 {
        if pixel[3] == 0 {
            continue;
        }
        let rgb = [pixel[0], pixel[1], pixel[2]];
        pixel[..3].copy_from_slice(cache.entry(rgb).or_insert_with(|| mode.simulate(rgb)));
    }
    output
}

pub const MAX_COLORS: usize = 64;
#[derive(Debug)]
pub struct PairWarning {
    pub a: [u8; 3],
    pub b: [u8; 3],
    pub before: f64,
    pub after: f64,
}

pub struct Scenario {
    pub mode: Vision,
    pub pairs: Vec<PairWarning>,
}

#[derive(Default)]
pub struct Report {
    pub scenarios: Vec<Scenario>,
    pub colors_checked: usize,
    pub opaque_colors: usize,
    pub pixels_checked: u64,
    pub opaque_pixels: u64,
}

pub fn check(palette: &[Swatch]) -> Report {
    // Only opaque colors: imported alpha pixels have an unknown backdrop.
    let mut colors: Vec<_> = palette.iter().filter(|s| s.opaque_count > 0).collect();
    colors.sort_unstable_by(|a, b| b.opaque_count.cmp(&a.opaque_count).then(a.rgb.cmp(&b.rgb)));
    let opaque_colors = colors.len();
    let opaque_pixels = colors.iter().map(|s| s.opaque_count).sum();
    colors.truncate(MAX_COLORS);
    let mut report = Report {
        colors_checked: colors.len(),
        opaque_colors,
        pixels_checked: colors.iter().map(|s| s.opaque_count).sum(),
        opaque_pixels,
        scenarios: Vec::new(),
    };
    let originals: Vec<_> = colors.iter().map(|s| oklab(linear_rgb(s.rgb))).collect();
    for mode in Vision::CHECKS {
        let simulated: Vec<_> = colors
            .iter()
            .map(|s| oklab(mode.linear(linear_rgb(s.rgb))))
            .collect();
        let mut pairs = Vec::new();
        for a in 0..colors.len() {
            for b in a + 1..colors.len() {
                let before = distance(originals[a], originals[b]);
                let after = distance(simulated[a], simulated[b]);
                // Deliberately exclude pairs already similar in the original.
                // These are tunable screening thresholds, not visibility limits.
                if before >= 0.08 && after <= 0.04 && after <= before * 0.5 {
                    pairs.push(PairWarning {
                        a: colors[a].rgb,
                        b: colors[b].rgb,
                        before,
                        after,
                    });
                }
            }
        }
        pairs.sort_by(|a, b| {
            a.after
                .total_cmp(&b.after)
                .then(a.a.cmp(&b.a))
                .then(a.b.cmp(&b.b))
        });
        report.scenarios.push(Scenario { mode, pairs });
    }
    report
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn gamma_roundtrip_luminance_and_neutrals() {
        for c in 0..=255 {
            assert_eq!(encode(decode(c)), c);
            for mode in Vision::CHECKS {
                assert_eq!(mode.simulate([c; 3]), [c; 3]);
            }
        }
        assert_eq!(luminance([0; 3]), 0.0);
        assert!((luminance([255; 3]) - 1.0).abs() < 1e-9);
        assert!((luminance([128; 3]) - 0.2158605).abs() < 1e-6);
        assert!(luminance([0, 255, 0]) > luminance([255, 0, 0]));
    }
    #[test]
    fn known_primary_simulations_use_linear_light() {
        // Independently calculated from severity-1 matrix columns + sRGB OETF.
        assert_eq!(Vision::Protanopia.simulate([255, 0, 0]), [109, 95, 0]);
        assert_eq!(Vision::Deuteranopia.simulate([255, 0, 0]), [163, 144, 0]);
        assert_eq!(Vision::Tritanopia.simulate([0, 0, 255]), [0, 107, 150]);
        assert_eq!(Vision::Grayscale.simulate([255, 0, 0]), [127; 3]);
    }
    #[test]
    fn preserves_alpha_and_original_pixels() {
        let pixels = [255, 0, 0, 128, 4, 5, 6, 0];
        assert_eq!(simulate_rgba(&pixels, Vision::Original), pixels);
        let simulated = simulate_rgba(&pixels, Vision::Deuteranopia);
        assert_eq!(simulated[3], 128);
        assert_eq!(&simulated[4..], &pixels[4..]);
    }
    #[test]
    fn warns_for_merging_colors_but_not_existing_gray_similarity() {
        // Red and this green have nearly equal luminance but different hues.
        let analysis = super::super::analyze(&[255, 0, 0, 255, 0, 148, 0, 255]);
        let report = check(&analysis.palette);
        assert_eq!(report.scenarios.last().unwrap().pairs.len(), 1);
        let gray = super::super::analyze(&[
            128, 128, 128, 255, 129, 129, 129, 255, 0, 0, 0, 255, 255, 255, 255, 255,
        ]);
        assert!(
            check(&gray.palette)
                .scenarios
                .iter()
                .all(|s| s.pairs.is_empty())
        );
    }
    #[test]
    fn catches_red_green_confusion_under_deuteranopia() {
        let a = super::super::analyze(&[255, 0, 0, 255, 0, 173, 0, 255]);
        let report = check(&a.palette);
        let deutan = report
            .scenarios
            .iter()
            .find(|s| s.mode == Vision::Deuteranopia)
            .unwrap();
        assert_eq!(deutan.pairs.len(), 1);
        assert!(deutan.pairs[0].before > deutan.pairs[0].after * 2.0);
    }

    #[test]
    fn coverage_is_bounded_and_alpha_is_not_a_pass() {
        let palette: Vec<_> = (0..100)
            .map(|i| Swatch {
                rgb: [i, 0, 0],
                count: 1,
                opaque_count: 1,
            })
            .collect();
        let report = check(&palette);
        assert_eq!(
            (
                report.colors_checked,
                report.opaque_colors,
                report.pixels_checked,
                report.opaque_pixels
            ),
            (64, 100, 64, 100)
        );
        let transparent = super::super::analyze(&[255, 0, 0, 128, 0, 148, 0, 0]);
        assert_eq!(check(&transparent.palette).colors_checked, 0);
    }
}