7#![allow(clippy::many_single_char_names)]

9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10struct ColorEmphasis {
11    r: bool,
12    g: bool,
13    b: bool,
14}
15
16impl ColorEmphasis {
17    const NONE: Self = Self { r: false, g: false, b: false };
18}
19
20use crate::NesPalette;
21use std::array;
22
23#[derive(Debug, Clone, Copy)]
24pub struct PaletteGenerationArgs {
25    pub brightness: f64,
26    pub saturation: f64,
27    pub contrast: f64,
28    pub gamma: f64,
29    pub hue_offset: f64,
30}
31
32impl Default for PaletteGenerationArgs {
33    fn default() -> Self {
34        Self { brightness: 1.0, saturation: 1.0, contrast: 1.0, gamma: 1.8, hue_offset: 0.0 }
35    }
36}
37
38fn generate_normalized_ntsc_signal(nes_color: u8, emphasis: ColorEmphasis, phase: u8) -> f64 {
39    const BLACK: f64 = 0.312;
40    const WHITE: f64 = 1.100;
41
42    const LOW: [f64; 4] = [0.228, 0.312, 0.552, 0.880];
43    const HIGH: [f64; 4] = [0.616, 0.840, 1.100, 1.100];
44    const ATTENUATED_LOW: [f64; 4] = [0.192, 0.256, 0.448, 0.712];
45    const ATTENUATED_HIGH: [f64; 4] = [0.500, 0.676, 0.896, 0.896];
46
47    let hue = nes_color & 0xF;
48
49    // Luma is forced to 1 for colors $xE and $xF
50    let luma = if hue < 0xE { (nes_color >> 4) & 3 } else { 1 };
51
52    let in_color_phase = |color| (color + phase) % 12 < 6;
53
54    // Color emphasis does not apply to colors $xE and $xF
55    let attenuate = hue < 0xE
56        && ((emphasis.r && in_color_phase(0))
57            || (emphasis.g && in_color_phase(4))
58            || (emphasis.b && in_color_phase(8)));
59
60    let (low, high) = if attenuate {
61        (ATTENUATED_LOW[luma as usize], ATTENUATED_HIGH[luma as usize])
62    } else {
63        (LOW[luma as usize], HIGH[luma as usize])
64    };
65
66    // Signal is always high for colors $x0 and always low for colors $xD, $xE, $xF
67    let signal = match hue {
68        0 => high,
69        1..=12 => {
70            if in_color_phase(hue) {
71                high
72            } else {
73                low
74            }
75        }
76        13..=15 => low,
77        _ => unreachable!("value & 0xF is always <= 15"),
78    };
79
80    // Normalize signal
81    (signal - BLACK) / (WHITE - BLACK)
82}

85const CHROMA_SATURATION_CORRECTION: f64 = 2.0;

87const SAMPLES: u8 = 12;
88const WEIGHT: f64 = 1.0 / (SAMPLES as f64);

92const BASE_HUE_OFFSET: f64 = 2.9;

94fn nes_to_yuv(nes_color: u8, emphasis: ColorEmphasis, hue_offset: f64) -> (f64, f64, f64) {
95    use std::f64::consts::PI;
96
97    let mut y = 0.0;
98    let mut u = 0.0;
99    let mut v = 0.0;
100
101    for phase in 0..SAMPLES {
102        let ntsc_signal = WEIGHT * generate_normalized_ntsc_signal(nes_color, emphasis, phase);
103        let wave_phase = f64::from(phase) + BASE_HUE_OFFSET + hue_offset;
104
105        y += ntsc_signal;
106        u += ntsc_signal * (wave_phase / 12.0 * 2.0 * PI).sin() * CHROMA_SATURATION_CORRECTION;
107        v += ntsc_signal * (wave_phase / 12.0 * 2.0 * PI).cos() * CHROMA_SATURATION_CORRECTION;
108    }
109
110    (y, u, v)
111}
112
113fn yuv_to_rgb(y: f64, u: f64, v: f64, gamma: f64) -> (u8, u8, u8) {
114    let apply_gamma = |c: f64| if c >= 0.0 { c.powf(2.2 / gamma) } else { 0.0 };
115    let clamp_to_u8 = |c: f64| (255.0 * c).clamp(0.0, 255.0).round() as u8;
116
117    let r = clamp_to_u8(apply_gamma(y + 1.139883 * v));
118    let g = clamp_to_u8(apply_gamma(y - 0.394642 * u - 0.580622 * v));
119    let b = clamp_to_u8(apply_gamma(y + 2.032062 * u));
120
121    (r, g, b)
122}
123
124fn emphasis_from_index(index: usize) -> ColorEmphasis {
125    ColorEmphasis { r: index & 0x040 != 0, g: index & 0x080 != 0, b: index & 0x100 != 0 }
126}
127
128#[must_use]
129pub fn generate(args: PaletteGenerationArgs) -> NesPalette {
130    NesPalette(array::from_fn(|color| {
131        let nes_color = (color & 0x03F) as u8;
132        let emphasis = emphasis_from_index(color);
133
134        let (mut y, mut u, mut v) = nes_to_yuv(nes_color, emphasis, args.hue_offset);
135
136        // Apply contrast
137        y = (y - 0.5) * args.contrast + 0.5;
138
139        // Apply brightness and saturation
140        y *= args.brightness;
141        u *= args.brightness * args.saturation;
142        v *= args.brightness * args.saturation;
143
144        yuv_to_rgb(y, u, v, args.gamma)
145    }))
146}
147
148fn rgb_to_yuv(r: u8, g: u8, b: u8) -> (f64, f64, f64) {
149    let u8_to_f64 = |c: u8| f64::from(c) / 255.0;
150
151    let r: f64 = u8_to_f64(r);
152    let g: f64 = u8_to_f64(g);
153    let b: f64 = u8_to_f64(b);
154
155    let y = r * 0.299 + g * 0.587 + b * 0.114;
156    let u = 0.492111 * (b - y);
157    let v = 0.877283 * (r - y);
158
159    (y, u, v)
160}
161
162#[must_use]
163pub fn extrapolate_64_to_512(palette: &[(u8, u8, u8); 64]) -> NesPalette {
164    use std::f64::consts::PI;
165
166    NesPalette(array::from_fn(|color| {
167        if color < 64 {
168            return palette[color];
169        }
170
171        let nes_color = (color & 0x03F) as u8;
172        let (r, g, b) = palette[nes_color as usize];
173        let emphasis = emphasis_from_index(color);
174
175        let (mut y, mut u, mut v) = rgb_to_yuv(r, g, b);
176
177        for phase in 0..SAMPLES {
178            let ntsc_without_emphasis =
179                generate_normalized_ntsc_signal(nes_color, ColorEmphasis::NONE, phase);
180            let ntsc_with_emphasis = generate_normalized_ntsc_signal(nes_color, emphasis, phase);
181            let difference = WEIGHT * (ntsc_without_emphasis - ntsc_with_emphasis);
182            debug_assert!(difference >= 0.0);
183
184            if difference < 1e-6 {
185                continue;
186            }
187
188            let wave_phase = (f64::from(phase) + BASE_HUE_OFFSET) / 12.0;
189
190            y -= difference;
191            u -= difference * (wave_phase * 2.0 * PI).sin() * CHROMA_SATURATION_CORRECTION;
192            v -= difference * (wave_phase * 2.0 * PI).cos() * CHROMA_SATURATION_CORRECTION;
193        }
194
195        // YUV to RGB conversion uses 2.2 as "base" gamma
196        yuv_to_rgb(y, u, v, 2.2)
197    }))
198}