//! Train the DINOv3 feature → RGB decoder, using meganeura for both halves.
//!
//! ```text
//! cargo run --release --example train_decoder -- \
//!   <dataset-dir-or-manifest> <model.safetensors> \
//!   [steps] [images] [layers] [size] [seed] [output-dir]
//! ```
//!
//! Two phases. First every image is encoded once and its features cached in
//! memory, because the encoder is frozen and running it inside the training
//! loop would cost ~30× the decoder's own forward pass for no gradient.
//! Then the decoder trains on those pairs with Adam.
//!
//! Writes `decoder.bin` — raw f32, parameters in graph declaration order —
//! next to the working directory, plus a PNG strip of reconstructions so
//! the result can be judged by eye rather than by loss alone.

use std::path::{Path, PathBuf};
use std::time::{Instant, SystemTime, UNIX_EPOCH};

use dinovision::decoder;
use dinovision::dinov3::Config;
use meganeura::graph::Op;
use meganeura::train::{Mode, SessionConfig};
use meganeura::{Graph, Session};
use serde::Serialize;

mod common;

const BATCH: usize = 8;

fn random_u64(state: &mut u64) -> u64 {
    *state ^= *state >> 12;
    *state ^= *state << 25;
    *state ^= *state >> 27;
    state.wrapping_mul(0x2545_F491_4F6C_DD1D)
}

fn seeded_state(base: u64, seed: u64) -> u64 {
    let state = base ^ seed;
    if state == 0 {
        0xA076_1D64_78BD_642F
    } else {
        state
    }
}

fn shuffle(order: &mut [usize], state: &mut u64) {
    for i in (1..order.len()).rev() {
        order.swap(i, random_u64(state) as usize % (i + 1));
    }
}

/// Deterministic small init. Kaiming-ish: scaled by fan-in so activations
/// neither vanish nor explode through four stages.
fn init_parameters(session: &mut Session, graph: &Graph, seed: u64) {
    let mut state = seeded_state(0x9E37_79B9_7F4A_7C15, seed);
    let mut next = move || ((random_u64(&mut state) >> 40) as f32 / (1u32 << 24) as f32) - 0.5;
    for node in graph.nodes() {
        let Op::Parameter { name } = &node.op else {
            continue;
        };
        let n = node.ty.num_elements();
        let shape = &node.ty.shape;
        let data: Vec<f32> = if name.ends_with("norm.weight") {
            vec![1.0; n]
        } else if name.ends_with(".bias") {
            vec![0.0; n]
        } else {
            // shape is [out, in, kh, kw]; fan_in = in * kh * kw.
            let fan_in: usize = shape.iter().skip(1).product::<usize>().max(1);
            let scale = (2.0 / fan_in as f32).sqrt() * 2.0;
            (0..n).map(|_| next() * scale).collect()
        };
        session.set_parameter(name, &data);
    }
}

#[derive(Serialize)]
struct TrainingRecord {
    schema_version: u32,
    completed_unix_seconds: u64,
    dataset: String,
    dataset_manifest_sha256: Option<String>,
    requested_images: usize,
    training_images: usize,
    model: String,
    model_sha256: String,
    initialization: String,
    initialization_sha256: Option<String>,
    seed: u64,
    image_size: usize,
    encoder_layers: usize,
    batch_size: usize,
    steps: usize,
    data_order: &'static str,
    objective: &'static str,
    optimizer: &'static str,
    initial_learning_rate: f32,
    final_learning_rate: f32,
    decoder_parameters: usize,
    encoding_seconds: f64,
    training_seconds: f64,
    final_l1: f32,
    decoder: String,
    decoder_sha256: String,
    diagnostic_samples: String,
}

fn portable_path(path: &Path) -> String {
    path.to_string_lossy().replace('\\', "/")
}

fn save_parameters(session: &Session, graph: &Graph, path: &Path) -> std::io::Result<()> {
    let mut bytes = Vec::new();
    for node in graph.nodes() {
        let Op::Parameter { name } = &node.op else {
            continue;
        };
        let mut buf = vec![0.0f32; node.ty.num_elements()];
        session.read_param(name, &mut buf);
        bytes.extend_from_slice(bytemuck::cast_slice(&buf));
    }
    std::fs::write(path, bytes)
}

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 image_dir = PathBuf::from(args.next().expect(
        "usage: train_decoder <dataset-dir-or-manifest> <model.safetensors> \
         [steps] [images] [layers] [size] [seed] [output-dir]",
    ));
    let weights = PathBuf::from(args.next().expect("need the DINOv3 checkpoint"));
    let steps: usize = args.next().and_then(|s| s.parse().ok()).unwrap_or(3000);
    let max_images: usize = args.next().and_then(|s| s.parse().ok()).unwrap_or(2500);
    // Truncating the encoder is the cheapest way to speed it up on a
    // headset, and the later layers are where colour gets discarded, so a
    // shallower encoder may reconstruct better as well as faster.
    let layers: usize = args.next().and_then(|s| s.parse().ok()).unwrap_or(12);
    let size: usize = args.next().and_then(|s| s.parse().ok()).unwrap_or(224);
    let seed: u64 = args.next().and_then(|s| s.parse().ok()).unwrap_or(0);
    let output_dir = args
        .next()
        .map(PathBuf::from)
        .unwrap_or_else(|| PathBuf::from("."));
    assert!(steps > 0, "steps must be positive");
    std::fs::create_dir_all(&output_dir).expect("create output directory");

    let config = Config::vits16().at_resolution(size).with_layers(layers);
    let gpu = dinovision::init_context(None).expect("GPU context");

    // ---- Phase 1: encode the dataset once ----
    let paths = if image_dir.extension().and_then(|e| e.to_str()) == Some("json") {
        common::images_for_split(&image_dir, "train", max_images)
            .unwrap_or_else(|e| panic!("{}: {e}", image_dir.display()))
            .into_iter()
            .map(|(path, image)| {
                common::verify_image(&path, &image)
                    .unwrap_or_else(|e| panic!("dataset verification failed: {e}"));
                path
            })
            .collect()
    } else {
        common::find_images(&image_dir, max_images)
    };
    assert!(!paths.is_empty(), "no images under {}", image_dir.display());
    log::info!("encoding {} images", paths.len());

    let (mut encoder, _) = dinovision::bench::build_encoder_session(gpu.clone(), &config, None);
    let model = meganeura::data::safetensors::SafeTensorsModel::load(weights.clone())
        .expect("read weights");
    dinovision::weights::load_encoder(&mut encoder, &model, &config).expect("bind weights");

    let feat_len = config.hidden_size * config.num_patches();
    let img_len = 3 * config.image_size * config.image_size;
    let mut features: Vec<f32> = Vec::with_capacity(paths.len() * feat_len);
    // Targets stay u8 and are widened per batch. As f32 the cache would be
    // 900 KB an image, and a few thousand images then no longer fit in RAM.
    let mut targets: Vec<u8> = Vec::with_capacity(paths.len() * img_len);

    let start = Instant::now();
    let mut kept = 0usize;
    let mut scratch = vec![0.0f32; config.num_tokens() * config.hidden_size];
    for (i, path) in paths.iter().enumerate() {
        let Some(rgb) = common::load_frame(path, config.image_size as u32) else {
            continue;
        };
        let patches = dinovision::preprocess::patches_from_rgb8(&rgb, &config);
        encoder.set_input("patches", &patches);
        encoder.step();
        encoder.wait();
        encoder.read_output_by_index(0, &mut scratch);

        features.extend_from_slice(&decoder::patch_features_to_nchw(&scratch, &config));
        // Target is plain CHW in [0, 1] — the decoder predicts pixels, not
        // ImageNet-normalized values.
        let size = config.image_size;
        for c in 0..3 {
            for p in 0..size * size {
                targets.push(rgb[p * 3 + c]);
            }
        }
        kept += 1;
        if i % 500 == 0 {
            log::info!(
                "  {i}/{} ({:.0}s)",
                paths.len(),
                start.elapsed().as_secs_f64()
            );
        }
    }
    drop(encoder);
    assert!(kept > 0, "none of the selected images could be decoded");
    let encoding_seconds = start.elapsed().as_secs_f64();
    log::info!(
        "encoded {kept} images in {:.0}s ({:.0} MB cached)",
        encoding_seconds,
        (features.len() * 4 + targets.len()) as f64 / 1e6
    );

    // ---- Phase 2: train the decoder ----
    let mut g = Graph::new();
    let feat_in = g.input("feat", &[BATCH * feat_len]);
    let recon = decoder::build_decoder(&mut g, &config, feat_in, BATCH);
    let target_in = g.input("target", &[BATCH * img_len]);
    // L1 rather than MSE. Squared error optimises the conditional mean, so
    // wherever a feature is ambiguous the decoder hedges by averaging every
    // possibility, which is blur by construction. L1 optimises the median
    // and commits to one answer, usually looking markedly sharper at the
    // same PSNR.
    let loss = g.l1_loss(recon, target_in);
    g.set_outputs(vec![loss, recon]);

    let (mut session, _) = meganeura::train::build(
        &g,
        SessionConfig {
            mode: Mode::Training,
            gpu: Some(gpu),
            ..Default::default()
        },
    );
    // `DINOVISION_INIT=decoder.bin` continues from existing weights instead
    // of starting over. Necessary when adapting to a few hundred captured
    // frames: 2M parameters trained from scratch on that much data would
    // simply memorise it, where fine-tuning shifts an already-general
    // decoder onto the new distribution.
    let (initialization, initialization_sha256) = match std::env::var("DINOVISION_INIT") {
        Ok(path) => {
            let path = PathBuf::from(path);
            decoder::load_parameters(&mut session, &g, &path)
                .unwrap_or_else(|e| panic!("could not load {}: {e}", path.display()));
            log::info!("fine-tuning from {}", path.display());
            let digest = common::sha256(&path).expect("hash initial decoder");
            (portable_path(&path), Some(digest))
        }
        Err(_) => {
            init_parameters(&mut session, &g, seed);
            ("random".to_string(), None)
        }
    };
    session.set_adam(2e-3, 0.9, 0.999, 1e-8);
    log::info!(
        "training {} decoder parameters for {steps} steps, batch {BATCH}, encoder depth {layers}, \
         resolution {size}, seed {seed}",
        decoder::parameter_count(&config)
    );

    let mut feat_batch = vec![0.0f32; BATCH * feat_len];
    let mut target_batch = vec![0.0f32; BATCH * img_len];
    let mut order: Vec<usize> = (0..kept).collect();
    let mut shuffle_state = seeded_state(0xD1B5_4A32_D192_ED03, seed);
    shuffle(&mut order, &mut shuffle_state);
    let mut cursor = 0usize;
    let train_start = Instant::now();
    let mut final_l1 = f32::NAN;
    let mut final_learning_rate = 2e-3;

    for step in 0..steps {
        for b in 0..BATCH {
            if cursor == kept {
                shuffle(&mut order, &mut shuffle_state);
                cursor = 0;
            }
            let idx = order[cursor];
            cursor += 1;
            feat_batch[b * feat_len..(b + 1) * feat_len]
                .copy_from_slice(&features[idx * feat_len..(idx + 1) * feat_len]);
            for (dst, &src) in target_batch[b * img_len..(b + 1) * img_len]
                .iter_mut()
                .zip(&targets[idx * img_len..(idx + 1) * img_len])
            {
                *dst = src as f32 / 255.0;
            }
        }
        session.set_input("feat", &feat_batch);
        session.set_input("target", &target_batch);
        // Linear decay with a 5%-of-initial-rate floor. This is an
        // implementation choice, not an admitted optimizer ablation.
        let progress = step as f32 / steps as f32;
        final_learning_rate = 2e-3 * (1.0 - progress).max(0.05);
        session.set_adam(final_learning_rate, 0.9, 0.999, 1e-8);
        session.step();
        session.wait();

        if step % 100 == 0 || step == steps - 1 {
            final_l1 = session.read_loss();
            log::info!(
                "step {step:>5}  L1 {final_l1:.5}  ({:.0}s)",
                train_start.elapsed().as_secs_f64()
            );
        }
    }
    let training_seconds = train_start.elapsed().as_secs_f64();

    let decoder_path = output_dir.join("decoder.bin");
    save_parameters(&session, &g, &decoder_path).expect("write decoder.bin");
    log::info!("wrote {}", decoder_path.display());

    // ---- Sample strip: original above, reconstruction below ----
    let size = config.image_size;
    for b in 0..BATCH {
        let idx = b % kept;
        feat_batch[b * feat_len..(b + 1) * feat_len]
            .copy_from_slice(&features[idx * feat_len..(idx + 1) * feat_len]);
        for (dst, &src) in target_batch[b * img_len..(b + 1) * img_len]
            .iter_mut()
            .zip(&targets[idx * img_len..(idx + 1) * img_len])
        {
            *dst = src as f32 / 255.0;
        }
    }
    session.set_input("feat", &feat_batch);
    session.step();
    session.wait();
    let mut recon_out = vec![0.0f32; BATCH * img_len];
    session.read_output_by_index(1, &mut recon_out);

    let cols = 6.min(BATCH);
    let mut strip = image::RgbImage::new((cols * size) as u32, (2 * size) as u32);
    let mut total_psnr = 0.0;
    for b in 0..cols {
        let t = &target_batch[b * img_len..(b + 1) * img_len];
        let r = &recon_out[b * img_len..(b + 1) * img_len];
        total_psnr += decoder::psnr(r, t);
        for y in 0..size {
            for x in 0..size {
                let at = |src: &[f32], c: usize| {
                    (src[c * size * size + y * size + x].clamp(0.0, 1.0) * 255.0) as u8
                };
                strip.put_pixel(
                    (b * size + x) as u32,
                    y as u32,
                    image::Rgb([at(t, 0), at(t, 1), at(t, 2)]),
                );
                strip.put_pixel(
                    (b * size + x) as u32,
                    (size + y) as u32,
                    image::Rgb([at(r, 0), at(r, 1), at(r, 2)]),
                );
            }
        }
    }
    let samples_path = output_dir.join("decoder_samples.png");
    strip.save(&samples_path).expect("write samples");
    let manifest_sha256 = (image_dir.extension().and_then(|e| e.to_str()) == Some("json"))
        .then(|| common::sha256(&image_dir).expect("hash dataset manifest"));
    let record = TrainingRecord {
        schema_version: 1,
        completed_unix_seconds: SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("system time before Unix epoch")
            .as_secs(),
        dataset: portable_path(&image_dir),
        dataset_manifest_sha256: manifest_sha256,
        requested_images: max_images,
        training_images: kept,
        model: portable_path(&weights),
        model_sha256: common::sha256(&weights).expect("hash encoder model"),
        initialization,
        initialization_sha256,
        seed,
        image_size: config.image_size,
        encoder_layers: config.num_hidden_layers,
        batch_size: BATCH,
        steps,
        data_order: "seeded Fisher-Yates; reshuffled after every complete pass",
        objective: "mean absolute error (L1)",
        optimizer: "Adam(beta1=0.9,beta2=0.999,epsilon=1e-8), linear decay to a 0.0001 floor at 95% of updates",
        initial_learning_rate: 2e-3,
        final_learning_rate,
        decoder_parameters: decoder::parameter_count(&config),
        encoding_seconds,
        training_seconds,
        final_l1,
        decoder: portable_path(&decoder_path),
        decoder_sha256: common::sha256(&decoder_path).expect("hash trained decoder"),
        diagnostic_samples: portable_path(&samples_path),
    };
    let record_path = output_dir.join("training.json");
    std::fs::write(
        &record_path,
        serde_json::to_vec_pretty(&record).expect("serialize training record"),
    )
    .expect("write training record");
    println!(
        "wrote {} and {}\nin-sample diagnostic mean PSNR {:.2} dB",
        samples_path.display(),
        record_path.display(),
        total_psnr / cols as f32
    );
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn shuffle_is_seeded_and_preserves_every_index() {
        let run = |seed| {
            let mut order: Vec<usize> = (0..100).collect();
            let mut state = seeded_state(0xD1B5_4A32_D192_ED03, seed);
            shuffle(&mut order, &mut state);
            order
        };
        let a = run(7);
        let b = run(7);
        let c = run(8);
        assert_eq!(a, b);
        assert_ne!(a, c);
        let mut sorted = a;
        sorted.sort_unstable();
        assert_eq!(sorted, (0..100).collect::<Vec<_>>());
    }
}