//! Run the full pipeline on desktop and write the result as an image.
//!
//! Exercises exactly what the headset does — inference worker, PCA fit,
//! colour projection — minus OpenXR, and writes a side-by-side PPM so the
//! output can actually be looked at. Rendering a feature view you cannot
//! see is a poor way to find out it is wrong.
//!
//! ```text
//! cargo run --release --example preview -- [out.ppm] [model.safetensors]
//! ```
//!
//! Without a checkpoint it runs on synthetic weights: the pipeline is
//! exercised end to end, but the colours mean nothing. With real weights,
//! regions that belong to the same object should share a colour.

use std::path::PathBuf;
use std::time::{Duration, Instant};

use dinovision::dinov3::Config;
use dinovision::inference::{self, Weights};
use dinovision::pca::COMPONENTS;
use dinovision::source::{FrameSource, TestPattern};

/// Nearest-neighbour upscale of the feature grid, so each patch reads as a
/// solid block. Bilinear would look prettier and hide the true resolution.
fn upscale(grid_rgb: &[f32], grid: usize, out: usize) -> Vec<u8> {
    let mut px = vec![0u8; out * out * 3];
    for y in 0..out {
        let gy = (y * grid) / out;
        for x in 0..out {
            let gx = (x * grid) / out;
            let src = (gy * grid + gx) * COMPONENTS;
            let dst = (y * out + x) * 3;
            for c in 0..3 {
                px[dst + c] = (grid_rgb[src + c].clamp(0.0, 1.0) * 255.0) as u8;
            }
        }
    }
    px
}

/// Source on the left, feature view on the right.
fn write_side_by_side(path: &PathBuf, left: &[u8], right: &[u8], size: usize) {
    let mut joined = vec![0u8; size * size * 6];
    for y in 0..size {
        let src = y * size * 3;
        let dst = y * size * 6;
        joined[dst..dst + size * 3].copy_from_slice(&left[src..src + size * 3]);
        joined[dst + size * 3..dst + size * 6].copy_from_slice(&right[src..src + size * 3]);
    }
    image::RgbImage::from_raw(size as u32 * 2, size as u32, joined)
        .expect("image dimensions do not match buffer")
        .save(path)
        .expect("failed to write image");
}

fn main() {
    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();

    let mut args = std::env::args().skip(1);
    let out_path = PathBuf::from(args.next().unwrap_or_else(|| "preview.png".into()));
    let weights = match args.next().map(PathBuf::from) {
        Some(p) if p.exists() => {
            log::info!("using weights from {}", p.display());
            Weights::SafeTensors(p)
        }
        Some(p) => panic!("{} does not exist", p.display()),
        None => Weights::Synthetic,
    };

    let config = Config::vits16().at_resolution(256);
    let gpu = dinovision::init_context(None).expect("failed to initialize GPU context");
    let worker = inference::spawn(gpu, config.clone(), weights, None, 1, dinovision::inference::Display::PcaColour);

    // `DINOVISION_SOURCE=screen` grabs the desktop instead of the synthetic
    // scene, which is the quickest way to see the colouring against real
    // imagery without opening a window.
    let mut source: Box<dyn FrameSource> = match std::env::var("DINOVISION_SOURCE").as_deref() {
        #[cfg(feature = "capture")]
        Ok("screen") => match dinovision::source::ScreenCapture::new(config.image_size, 0) {
            Ok(s) => Box::new(s),
            Err(e) => panic!("screen capture unavailable: {e}"),
        },
        #[cfg(not(feature = "capture"))]
        Ok("screen") => panic!("rebuild with --features capture for screen capture"),
        _ => Box::new(TestPattern::new(config.image_size)),
    };
    let mut last_frame: Vec<u8> = Vec::new();

    // Let the session build, then run a few frames so the PCA basis is
    // refitted from real features rather than the placeholder.
    let deadline = Instant::now() + Duration::from_secs(180);
    let mut completed = 0;
    while completed < 5 && Instant::now() < deadline {
        if worker.is_ready() {
            let rgb = source.next_frame().unwrap();
            last_frame = rgb.to_vec();
            let patches = dinovision::preprocess::patches_from_rgb8(rgb, &config);
            worker.submit(patches);
        }
        std::thread::sleep(Duration::from_millis(20));
        completed = worker.generation();
    }

    let grid = worker
        .latest()
        .expect("inference produced no result before the deadline");
    log::info!(
        "{} completed encodes, last took {:.1} ms ({}x{} grid)",
        completed,
        grid.latency_ms,
        grid.grid,
        grid.grid
    );

    let size = config.image_size;
    let features = upscale(grid.patch_rgb(), grid.grid, size);
    write_side_by_side(&out_path, &last_frame, &features, size);

    // A view that is one flat colour means the projection collapsed, which
    // is easy to miss by eye in a small image.
    let spread = {
        let mut lo = [255u8; 3];
        let mut hi = [0u8; 3];
        for px in features.chunks_exact(3) {
            for c in 0..3 {
                lo[c] = lo[c].min(px[c]);
                hi[c] = hi[c].max(px[c]);
            }
        }
        (0..3).map(|c| hi[c] as i32 - lo[c] as i32).max().unwrap()
    };
    println!("wrote {} (source | features)", out_path.display());
    println!("colour spread: {spread}/255");
    if spread < 16 {
        println!("WARNING: the feature view is nearly flat — projection may have collapsed");
    }
}