/// Generate a synthetic PSG sample for testing.
///
/// Creates a safetensors file with random multichannel data
/// mimicking 4 modalities (BAS, RESP, EKG, EMG) at 128 Hz.
///
/// Usage:
///   cargo run --bin gen_sample_psg -- --output data/sample_psg.safetensors

use clap::Parser;
use safetensors::{Dtype, View};
use std::borrow::Cow;

#[derive(Parser)]
#[command(about = "Generate synthetic PSG sample data")]
struct Args {
    #[arg(long, default_value = "data/sample_psg.safetensors")]
    output: String,

    /// Sampling frequency in Hz.
    #[arg(long, default_value = "128")]
    sample_rate: usize,

    /// Duration in minutes.
    #[arg(long, default_value = "5")]
    duration_min: usize,

    /// Number of BAS (EEG/EOG) channels.
    #[arg(long, default_value = "10")]
    bas_channels: usize,

    /// Number of RESP channels.
    #[arg(long, default_value = "7")]
    resp_channels: usize,

    /// Number of EKG channels.
    #[arg(long, default_value = "2")]
    ekg_channels: usize,

    /// Number of EMG channels.
    #[arg(long, default_value = "4")]
    emg_channels: usize,
}

struct RawTensor {
    data: Vec<u8>,
    shape: Vec<usize>,
}

impl View for RawTensor {
    fn dtype(&self) -> Dtype { Dtype::F32 }
    fn shape(&self) -> &[usize] { &self.shape }
    fn data(&self) -> Cow<'_, [u8]> { Cow::Borrowed(&self.data) }
    fn data_len(&self) -> usize { self.data.len() }
}

fn gen_signal(channels: usize, samples: usize) -> RawTensor {
    // Simple synthetic: sine waves at different frequencies per channel
    let mut data = Vec::with_capacity(channels * samples * 4);
    for ch in 0..channels {
        let freq = 0.5 + ch as f32 * 0.3;
        for t in 0..samples {
            let val = (2.0 * std::f32::consts::PI * freq * t as f32 / 128.0).sin();
            data.extend_from_slice(&val.to_le_bytes());
        }
    }
    RawTensor { data, shape: vec![channels, samples] }
}

fn main() -> anyhow::Result<()> {
    let args = Args::parse();
    let samples = args.duration_min * 60 * args.sample_rate;

    let modalities = vec![
        ("BAS", args.bas_channels),
        ("RESP", args.resp_channels),
        ("EKG", args.ekg_channels),
        ("EMG", args.emg_channels),
    ];

    let mut keys: Vec<String> = Vec::new();
    let mut tensors: Vec<RawTensor> = Vec::new();

    for (name, channels) in &modalities {
        eprintln!("Generating {} signal: {} channels × {} samples", name, channels, samples);
        keys.push(name.to_string());
        tensors.push(gen_signal(*channels, samples));
    }

    // Save metadata
    let sr = args.sample_rate as f32;
    keys.push("sample_rate".into());
    tensors.push(RawTensor {
        data: sr.to_le_bytes().to_vec(),
        shape: vec![1],
    });

    let pairs: Vec<(&str, RawTensor)> = keys.iter()
        .map(|s| s.as_str())
        .zip(tensors)
        .collect();

    std::fs::create_dir_all(std::path::Path::new(&args.output).parent().unwrap_or(std::path::Path::new(".")))?;
    let bytes = safetensors::serialize(pairs, None)?;
    std::fs::write(&args.output, bytes)?;
    eprintln!("Wrote {}", args.output);

    Ok(())
}
