//! TRACES: FR-DEV-3g
//! The learned denoise in a develop session (docs/dev/denoise.md §7).
//!
//! The classical demosaic shows at once; when the photograph asks for the
//! learned one, it is computed off the UI thread from the mosaic the
//! session kept, and swapped in when it lands. Grain is a blend of the two
//! results, made once per slider value by [`dr_gpu::GrainBlend`] and handed
//! to the render as its source — the adjust pass never knows.
//!
//! `demosaiced` stays the classical result for the session's life: the raw
//! histogram, the white balance picker, masks and segmentation measure the
//! sensor data, and only [`DevelopSession::developed_source`] — what the
//! render draws — changes.

use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc;
use std::sync::Arc;

use dr_decode::RawImage;
use dr_gpu::{DemosaicedImage, GrainBlend};

use super::session::DevelopSession;

/// What the session holds for the learned denoise.
#[derive(Default)]
pub(crate) struct DenoiseState {
    /// The mosaic as decoded, kept only for a photograph that can take the
    /// learned stage. The job repairs a copy.
    mosaic: Option<Arc<RawImage>>,
    /// The file's `NoiseProfile` and ISO, read from the header at open.
    profile: Option<Vec<(f32, f32)>>,
    iso: Option<u32>,
    /// The network's result, once it has landed.
    result: Option<Arc<DemosaicedImage>>,
    /// The last grain blend made, and the grain it was made at.
    blended: Option<(f32, Arc<DemosaicedImage>)>,
    blend: Option<GrainBlend>,
    job: Option<Job>,
    /// Why the last attempt failed; not retried until the switch is
    /// toggled, so a photograph that cannot be denoised does not loop.
    failed: Option<String>,
    /// Where the noise figures came from, for the panel.
    source: Option<dr_denoise::Source>,
}

struct Job {
    rx: mpsc::Receiver<Msg>,
    cancel: Arc<AtomicBool>,
}

enum Msg {
    Progress(usize, usize),
    Done(Result<Finished, String>),
}

struct Finished {
    rgb: Vec<f32>,
    width: u32,
    height: u32,
    source: dr_denoise::Source,
    rung: String,
    seconds: f64,
}

/// What a poll found, for the develop view's status line and redraw.
#[derive(Debug, PartialEq)]
pub enum DenoiseStatus {
    /// Nothing running and nothing new.
    Idle,
    /// Tiles done of tiles.
    Running(usize, usize),
    /// The result landed this poll: redraw.
    Landed,
    Failed(String),
}

impl DevelopSession {
    /// Keep the mosaic for the learned denoise, if it can take this frame.
    pub(super) fn keep_mosaic(&mut self, raw: RawImage) {
        if dr_denoise::eligible(&raw) {
            self.denoise.mosaic = Some(Arc::new(raw));
        }
        self.graph
            .set_denoise_available(self.denoise.mosaic.is_some());
    }

    /// The header's part: the DNG's measured noise and the ISO.
    pub fn prepare_denoise(&mut self, bytes: &[u8], meta: &dr_decode::Metadata) {
        if self.denoise.mosaic.is_some() {
            self.denoise.profile = dr_decode::noise_profile(bytes);
            self.denoise.iso = meta.iso;
        }
    }

    /// Bring what is computed in line with what the edit asks for: start the
    /// network when it is wanted and has not run, stop it when it is not.
    /// Cheap when nothing changed; the develop view calls it after every
    /// change to the edit, whatever made it — a slider, undo, a version.
    pub fn reconcile_denoise(&mut self) {
        let wanted = self.graph.denoise_applied() && self.denoise.mosaic.is_some();
        if !wanted {
            if let Some(job) = self.denoise.job.take() {
                job.cancel.store(true, Ordering::Relaxed);
            }
            // Toggling off is how a failure is retried.
            self.denoise.failed = None;
            return;
        }
        if self.denoise.result.is_some()
            || self.denoise.job.is_some()
            || self.denoise.failed.is_some()
        {
            return;
        }
        let Some(model) = crate::library::denoise_model() else {
            self.denoise.failed = Some("the denoise model is not installed".into());
            return;
        };
        let (tx, rx) = mpsc::channel();
        let cancel = Arc::new(AtomicBool::new(false));
        let work = Work {
            ctx: self.ctx.clone(),
            raw: self.denoise.mosaic.clone().expect("checked above"),
            profile: self.denoise.profile.clone(),
            iso: self.denoise.iso,
            model,
            cancel: cancel.clone(),
        };
        crate::executors::spawn(crate::executors::Executor::Decode, "denoise", move || {
            let result = work.run(&mut |done, total| {
                let _ = tx.send(Msg::Progress(done, total));
            });
            let _ = tx.send(Msg::Done(result));
        });
        self.denoise.job = Some(Job { rx, cancel });
    }

    /// Collect what the job sent since the last poll.
    pub fn poll_denoise(&mut self) -> DenoiseStatus {
        let Some(job) = &self.denoise.job else {
            return DenoiseStatus::Idle;
        };
        let mut last = None;
        let mut done = None;
        while let Ok(msg) = job.rx.try_recv() {
            match msg {
                Msg::Progress(d, t) => last = Some((d, t)),
                Msg::Done(r) => done = Some(r),
            }
        }
        match done {
            Some(Ok(f)) => {
                self.denoise.job = None;
                match self.land(f) {
                    Ok(()) => DenoiseStatus::Landed,
                    Err(e) => self.fail(e),
                }
            }
            Some(Err(e)) => {
                self.denoise.job = None;
                self.fail(e)
            }
            None => last.map_or(DenoiseStatus::Running(0, 1), |(d, t)| {
                DenoiseStatus::Running(d, t)
            }),
        }
    }

    /// Compute the result now, on this thread, if the edit wants it and it
    /// is not here: for an export, which must not write the classical
    /// picture of a photograph that asks for the learned one (§7.1).
    pub fn denoise_blocking(&mut self) -> Result<(), String> {
        if !self.graph.denoise_applied() || self.denoise.result.is_some() {
            return Ok(());
        }
        let Some(raw) = self.denoise.mosaic.clone() else {
            return Ok(());
        };
        // Already under way: wait for it rather than start again.
        if let Some(job) = self.denoise.job.take() {
            for msg in job.rx.iter() {
                if let Msg::Done(result) = msg {
                    return result.and_then(|f| self.land(f));
                }
            }
        }
        let model = crate::library::denoise_model().ok_or("the denoise model is not installed")?;
        let work = Work {
            ctx: self.ctx.clone(),
            raw,
            profile: self.denoise.profile.clone(),
            iso: self.denoise.iso,
            model,
            cancel: Arc::new(AtomicBool::new(false)),
        };
        let finished = work.run(&mut |_, _| {})?;
        self.land(finished)
    }

    /// What the render draws: the learned result with the asked-for grain
    /// where there is one, else the classical demosaic.
    pub(super) fn developed_source(&mut self) -> Arc<DemosaicedImage> {
        if !self.graph.denoise_applied() {
            return self.demosaiced.clone();
        }
        let Some(result) = self.denoise.result.clone() else {
            return self.demosaiced.clone();
        };
        let grain = self.graph.denoise_grain();
        if grain <= 0.0 {
            return result;
        }
        if let Some((g, image)) = &self.denoise.blended {
            if (*g - grain).abs() < 1e-4 {
                return image.clone();
            }
        }
        let blend = self
            .denoise
            .blend
            .get_or_insert_with(|| GrainBlend::new(&self.ctx));
        match blend.blend(&result, &self.demosaiced, grain) {
            Ok(image) => {
                self.denoise.blended = Some((grain, image.clone()));
                image
            }
            Err(e) => {
                log::warn!("grain blend failed, showing the denoised result without grain: {e}");
                result
            }
        }
    }

    /// For the panel: where the noise figures came from, once computed.
    pub fn denoise_source(&self) -> Option<dr_denoise::Source> {
        self.denoise.source
    }

    pub fn denoise_failure(&self) -> Option<&str> {
        self.denoise.failed.as_deref()
    }

    fn land(&mut self, f: Finished) -> Result<(), String> {
        let image =
            DemosaicedImage::from_rgb_f32(&self.ctx, &self.demosaiced, f.width, f.height, &f.rgb)
                .map_err(|e| e.to_string())?;
        log::info!(
            "learned denoise: {}×{} on {} in {:.1} s, noise {}",
            f.width,
            f.height,
            f.rung,
            f.seconds,
            f.source.label()
        );
        self.denoise.result = Some(Arc::new(image));
        self.denoise.blended = None;
        self.denoise.source = Some(f.source);
        Ok(())
    }

    fn fail(&mut self, e: String) -> DenoiseStatus {
        log::warn!("learned denoise failed: {e}");
        self.denoise.failed = Some(e.clone());
        DenoiseStatus::Failed(e)
    }
}

/// Everything the job needs, owned, so it can leave the UI thread.
struct Work {
    ctx: dr_gpu::GpuContext,
    raw: Arc<RawImage>,
    profile: Option<Vec<(f32, f32)>>,
    iso: Option<u32>,
    model: std::path::PathBuf,
    cancel: Arc<AtomicBool>,
}

impl Work {
    fn run(self, progress: &mut dyn FnMut(usize, usize)) -> Result<Finished, String> {
        let started = std::time::Instant::now();
        // The app's own hot-pixel pass, on a copy: the classical source was
        // repaired by the same pass inside `Demosaicer::run`.
        let mut raw = (*self.raw).clone();
        dr_gpu::Demosaicer::new(&self.ctx)
            .and_then(|d| d.repair_hot_pixels(&mut raw))
            .map_err(|e| e.to_string())?;
        let noise = dr_denoise::noise::for_frame_with(&raw, self.profile.as_deref(), self.iso)
            .ok_or("this photograph gives no way to measure its noise")?;
        let mut net =
            dr_denoise::onnx::OnnxNet::from_path(&self.model).map_err(|e| e.to_string())?;
        let rung = net
            .rung()
            .map(|r| r.label().to_string())
            .unwrap_or_default();
        let cancel = self.cancel;
        let rgb = dr_denoise::denoise(&raw, &noise, &mut net, &mut |done, total| {
            progress(done, total);
            !cancel.load(Ordering::Relaxed)
        })
        .map_err(|e| e.to_string())?
        .ok_or("stopped")?;
        Ok(Finished {
            rgb,
            width: raw.crop.width,
            height: raw.crop.height,
            source: noise.source,
            rung,
            seconds: started.elapsed().as_secs_f64(),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::develop::test_support::headless;
    use dr_pipeline::learned_denoise;

    fn bayer(w: u32, h: u32) -> RawImage {
        RawImage {
            width: w,
            height: h,
            data: vec![800; (w * h) as usize],
            cfa_pattern: dr_decode::CfaPattern::Rggb,
            black_level: [0; 4],
            white_level: 4095,
            wb_coeffs: [2.0, 1.0, 1.5, 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: 1,
            profile: None,
            profile_tables: None,
            baseline_exposure: 0.0,
            make: String::new(),
            model: String::new(),
            crop: dr_decode::CropRect {
                x: 0,
                y: 0,
                width: w,
                height: h,
            },
        }
    }

    fn landed(session: &mut DevelopSession) {
        let (w, h) = (64, 64);
        session
            .land(Finished {
                rgb: vec![0.2; (w * h * 3) as usize],
                width: w,
                height: h,
                source: dr_denoise::Source::Measured,
                rung: "test".into(),
                seconds: 0.0,
            })
            .expect("upload");
    }

    #[test]
    fn a_bayer_raw_offers_the_switch_and_an_rgb_image_does_not() {
        let Some(ctx) = headless() else { return };
        let raw =
            DevelopSession::open_owned(&ctx, bayer(64, 64), dr_types::Orientation::NORMAL).unwrap();
        assert!(raw
            .graph
            .capabilities()
            .iter()
            .any(|c| c.id == learned_denoise::ID));
        let rgba: Vec<u8> = (0..64 * 64).flat_map(|_| [128u8, 128, 128, 255]).collect();
        let jpeg =
            DevelopSession::open_rgb(&ctx, &rgba, 64, 64, dr_types::Orientation::NORMAL).unwrap();
        assert!(!jpeg
            .graph
            .capabilities()
            .iter()
            .any(|c| c.id == learned_denoise::ID));
    }

    #[test]
    fn the_render_draws_what_the_edit_asks_for() {
        let Some(ctx) = headless() else { return };
        let mut s =
            DevelopSession::open_owned(&ctx, bayer(64, 64), dr_types::Orientation::NORMAL).unwrap();
        let classical = s.demosaiced.clone();
        landed(&mut s);
        let result = s.denoise.result.clone().unwrap();
        let same = |a: &Arc<DemosaicedImage>, b: &Arc<DemosaicedImage>| Arc::ptr_eq(a, b);

        // Off: the classical demosaic, result or no result.
        assert!(same(&s.developed_source(), &classical));
        // On, no grain: the network's result as it is.
        s.graph
            .set_param(learned_denoise::ID, learned_denoise::APPLY, 1.0);
        assert!(same(&s.developed_source(), &result));
        // Grain: a blend, made once per value and reused until it moves.
        s.graph
            .set_param(learned_denoise::ID, learned_denoise::GRAIN, 40.0);
        let blended = s.developed_source();
        assert!(!same(&blended, &result) && !same(&blended, &classical));
        assert!(
            same(&s.developed_source(), &blended),
            "the same grain must not blend again"
        );
        s.graph
            .set_param(learned_denoise::ID, learned_denoise::GRAIN, 60.0);
        assert!(!same(&s.developed_source(), &blended));
        // The sensor's own reading stays the classical one throughout.
        assert!(same(&s.demosaiced, &classical));
    }

    #[test]
    fn switching_off_stops_and_forgets_a_failure() {
        let Some(ctx) = headless() else { return };
        let mut s =
            DevelopSession::open_owned(&ctx, bayer(64, 64), dr_types::Orientation::NORMAL).unwrap();
        s.denoise.failed = Some("no model".into());
        s.graph
            .set_param(learned_denoise::ID, learned_denoise::APPLY, 1.0);
        s.reconcile_denoise();
        assert!(
            s.denoise.job.is_none(),
            "a failure is not retried while the switch stays on"
        );
        s.graph
            .set_param(learned_denoise::ID, learned_denoise::APPLY, 0.0);
        s.reconcile_denoise();
        assert!(
            s.denoise.failed.is_none(),
            "toggling off is how a failure is retried"
        );
    }
}