//! TRACES: FR-DEV-3e
//! The camera profile's tables, end to end on a device (D20).
//!
//! `dr-pipeline` holds the lookup to the DNG SDK's algorithm on the CPU
//! (`ops::camera_profile::apply_reference`). Nothing there would notice a
//! shader that disagreed with it — a transposed constant matrix, an index
//! off by one column, a buffer bound in the wrong order — so this renders a
//! frame of 256 different colours through tables that move every one of them
//! a long way, and holds each pixel to the reference.
//!
//! The source is a linear three-sample frame, so the colours arrive exactly
//! as written with no demosaic between, and an identity stands in the view
//! transform's place so the readback is the scene colour, display-encoded.

use std::sync::Arc;

use dr_decode::{CfaPattern, CropRect, RawImage};
use dr_gpu::{AdjustPass, Demosaicer, GpuContext};
use dr_pipeline::descriptor::{Attribute, LocalizedKey, OpDescriptor, OpId, ParamId};
use dr_pipeline::operation::{Operation, Stage, Uniform};
use dr_pipeline::ops::camera_profile::{apply_reference, CameraProfile, APPLY, LOOK};
use dr_types::{HueSatTable, ProfileOrigin, ProfileTables, Transfer};

const SIZE: u32 = 16;

fn ctx() -> Option<GpuContext> {
    pollster::block_on(GpuContext::new_headless()).ok()
}

/// An identity in the view transform's place.
struct IdentityView;

impl Operation for IdentityView {
    fn descriptor(&self) -> Arc<OpDescriptor> {
        Arc::new(OpDescriptor {
            id: OpId("identity_view"),
            label: LocalizedKey("identity_view"),
            params: Vec::new(),
            attributes: vec![Attribute::Tone],
        })
    }
    fn set_param(&mut self, _: ParamId, _: f32) {}
    fn param(&self, _: ParamId) -> f32 {
        0.0
    }
    fn is_active(&self) -> bool {
        true
    }
    fn stage(&self) -> Stage {
        Stage::View
    }
    fn renders(&self) -> bool {
        true
    }
    fn wgsl_body(&self) -> String {
        String::new()
    }
    fn uniforms(&self) -> Vec<Uniform> {
        Vec::new()
    }
}

/// 256 colours across hue, saturation and value, kept under the prologue's
/// highlight desaturation and above black.
fn colours() -> Vec<[f32; 3]> {
    (0..SIZE * SIZE)
        .map(|i| {
            let f = |k: u32| {
                let x = (i.wrapping_mul(2_654_435_761).rotate_left(k * 7) >> 8) % 1000;
                0.04 + 0.86 * x as f32 / 1000.0
            };
            [f(1), f(2), f(3)]
        })
        .collect()
}

fn frame(tables: Option<ProfileTables>) -> RawImage {
    let data = colours()
        .iter()
        .flat_map(|c| c.map(|v| (v * 65535.0).round() as u16))
        .collect();
    RawImage {
        width: SIZE,
        height: SIZE,
        data,
        cfa_pattern: CfaPattern::Rggb,
        black_level: [0; 4],
        white_level: u16::MAX,
        wb_coeffs: [1.0, 1.0, 1.0, 1.0],
        color_matrix: Some([1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]),
        samples_per_pixel: 3,
        profile: None,
        profile_tables: tables.map(Arc::new),
        baseline_exposure: 0.0,
        make: String::new(),
        model: String::new(),
        crop: CropRect {
            x: 0,
            y: 0,
            width: SIZE,
            height: SIZE,
        },
    }
}

/// Tables that move every colour by a different amount: hue shifts of tens
/// of degrees, saturation scales either side of one, and a 3-D, sRGB-indexed
/// look whose value scale varies down the value axis.
fn strong_tables() -> ProfileTables {
    let (hd, sd) = (12u32, 5u32);
    let hue_sat = (0..hd * sd)
        .map(|i| {
            let (h, s) = (i / sd, i % sd);
            let a = h as f32 / hd as f32 * std::f32::consts::TAU;
            [25.0 * a.sin(), 1.0 + 0.3 * a.cos() * s as f32 / 4.0, 1.0]
        })
        .collect();
    let (lh, ls, lv) = (8u32, 4u32, 5u32);
    let look = (0..lh * ls * lv)
        .map(|i| {
            let v = i / (lh * ls);
            let h = (i / ls) % lh;
            [
                -15.0 + 4.0 * h as f32,
                1.25 - 0.05 * v as f32,
                0.85 + 0.06 * v as f32,
            ]
        })
        .collect();
    let mut look = HueSatTable::new(lh, ls, lv, true, look).unwrap();
    look.srgb_encoded = true;
    ProfileTables {
        name: "strong".into(),
        origin: ProfileOrigin::Embedded,
        hue_sat: Some(HueSatTable::new(hd, sd, 1, false, hue_sat).unwrap()),
        look: Some(look),
        tone_curve: None,
    }
}

fn render(ctx: &GpuContext, raw: &RawImage, op: CameraProfile) -> Vec<[u8; 3]> {
    let source = Demosaicer::new(ctx)
        .expect("demosaicer")
        .run(raw)
        .expect("upload");
    let ops: Vec<Box<dyn Operation>> = vec![Box::new(op), Box::new(IdentityView)];
    let shader = dr_pipeline::compose(&ops);
    let mut adjust = AdjustPass::new(ctx);
    adjust.render(&source, &shader, SIZE, SIZE).expect("render");
    let (pixels, _, _) = adjust.export_pixels().expect("readback");
    pixels.chunks_exact(4).map(|p| [p[0], p[1], p[2]]).collect()
}

fn encode(c: [f32; 3]) -> [i32; 3] {
    c.map(|v| (Transfer::Srgb.encode(v.clamp(0.0, 1.0)) * 255.0).round() as i32)
}

fn assert_agrees(got: &[[u8; 3]], expected: impl Fn([f32; 3]) -> [f32; 3], what: &str) {
    let mut moved = 0;
    for (i, (c, g)) in colours().into_iter().zip(got).enumerate() {
        let want = encode(expected(c));
        let g = g.map(i32::from);
        // Two 8-bit steps: the half-float source and intermediate, and the
        // rounding either side of the encode.
        assert!(
            want.iter().zip(g).all(|(w, g)| (w - g).abs() <= 2),
            "{what}: pixel {i} {c:?} rendered {g:?}, the reference says {want:?}"
        );
        if want != encode(c) {
            moved += 1;
        }
    }
    assert!(
        moved > 200,
        "{what}: only {moved} of 256 colours moved; the test proves little"
    );
}

#[test]
fn the_shader_agrees_with_the_cpu_reference() {
    let Some(ctx) = ctx() else {
        eprintln!("skipping: no GPU adapter");
        return;
    };
    let tables = strong_tables();
    let got = render(&ctx, &frame(Some(tables.clone())), CameraProfile::new());
    assert_agrees(&got, |c| apply_reference(&tables, c, 1.0), "at defaults");

    let mut doubled = CameraProfile::new();
    doubled.set_param(LOOK, 200.0);
    let got = render(&ctx, &frame(Some(tables.clone())), doubled);
    assert_agrees(&got, |c| apply_reference(&tables, c, 2.0), "look at 200%");
}

#[test]
fn camera_raw_tone_agrees_with_its_cpu_reference() {
    // TRACES: FR-DEV-3j
    // D21's rendering on 256 colours: the ProPhoto round trip, the clip, the
    // curve from the profile buffer's placeholder, and RGBTone's placement
    // of the middle channel, against `camera_raw::apply_reference`.
    let Some(ctx) = ctx() else {
        eprintln!("skipping: no GPU adapter");
        return;
    };
    let source = Demosaicer::new(&ctx)
        .expect("demosaicer")
        .run(&frame(None))
        .expect("upload");
    let mut view = dr_pipeline::ops::ViewTransform::new();
    view.set_param(
        dr_pipeline::ops::view_transform::CURVE,
        dr_pipeline::ops::view_transform::CAMERA_RAW,
    );
    let ops: Vec<Box<dyn Operation>> = vec![Box::new(view)];
    let shader = dr_pipeline::compose(&ops);
    let mut adjust = AdjustPass::new(&ctx);
    adjust.render(&source, &shader, SIZE, SIZE).expect("render");
    let (pixels, _, _) = adjust.export_pixels().expect("readback");
    let got: Vec<[u8; 3]> = pixels.chunks_exact(4).map(|p| [p[0], p[1], p[2]]).collect();
    let curve = &dr_types::tone::ACR3_DEFAULT;
    assert_agrees(
        &got,
        |c| {
            dr_pipeline::camera_raw::apply_reference(
                curve,
                c,
                dr_pipeline::view::DEFAULT_CONTRAST,
                dr_pipeline::view::DEFAULT_WHITE,
            )
        },
        "DNG reference tone",
    );
}

#[test]
fn switched_off_or_absent_the_render_is_unchanged() {
    let Some(ctx) = ctx() else {
        eprintln!("skipping: no GPU adapter");
        return;
    };
    let bare = render(&ctx, &frame(None), CameraProfile::new());
    let mut off = CameraProfile::new();
    off.set_param(APPLY, 0.0);
    let switched_off = render(&ctx, &frame(Some(strong_tables())), off);
    assert_eq!(bare, switched_off, "the switch off is the matrix alone");
    // Against the source colours, one 8-bit step for the half-float texture
    // the source is uploaded in; the exact comparison is the one above.
    for (c, g) in colours().into_iter().zip(&bare) {
        let want = encode(c);
        assert!(
            want.iter()
                .zip(g)
                .all(|(w, g)| (w - i32::from(*g)).abs() <= 1),
            "no tables, no change: {c:?} rendered {g:?}"
        );
    }
}

#[test]
fn the_libraries_adobe_standard_renders_as_the_reference_does() {
    // The real tables, when the library's 6D DNG is on this machine: a 90×30
    // HueSatMap and a 36×8×16 LookTable, at the sizes no synthetic test
    // reaches.
    let Some(ctx) = ctx() else {
        eprintln!("skipping: no GPU adapter");
        return;
    };
    let path = std::env::var_os("DR_DCP_SAMPLE")
        .map(std::path::PathBuf::from)
        .or_else(|| {
            std::env::var_os("HOME").map(|h| {
                std::path::PathBuf::from(h).join("Nextcloud/PhotosRaw/2017/2017-08-12/_MG_9080.dng")
            })
        });
    let Some(bytes) = path.and_then(|p| std::fs::read(p).ok()) else {
        eprintln!("skipping: no sample DNG");
        return;
    };
    let tables = dr_decode::dcp::embedded_in(&bytes)
        .expect("Adobe Standard")
        .tables(5000.0, ProfileOrigin::Embedded);
    let got = render(&ctx, &frame(Some(tables.clone())), CameraProfile::new());
    for (i, (c, g)) in colours().into_iter().zip(&got).enumerate() {
        let want = encode(apply_reference(&tables, c, 1.0));
        let g = g.map(i32::from);
        assert!(
            want.iter().zip(g).all(|(w, g)| (w - g).abs() <= 2),
            "pixel {i} {c:?} rendered {g:?}, the reference says {want:?}"
        );
    }
}