use std::{env, hint::black_box, process, sync::Arc, time::Instant};

use laddu::{
    amplitudes::{
        breit_wigner::BreitWigner,
        common::{ComplexScalar, Scalar},
        kmatrix::{KopfKMatrixA0, KopfKMatrixA2, KopfKMatrixF0, KopfKMatrixF2},
        parameter,
        zlm::Zlm,
    },
    data::{Dataset, DatasetReadOptions},
    extensions::NLL,
    io,
    traits::LikelihoodTerm,
    utils::{
        enums::{Frame, Sign},
        variables::{Angles, Mass, Polarization, Topology},
    },
    Evaluator, Expression, RngSubsetExtension,
};
use rayon::ThreadPoolBuilder;

const BENCH_DATASET_RELATIVE_PATH: &str = "benches/bench.parquet";
const P4_NAMES: [&str; 4] = ["beam", "proton", "kshort1", "kshort2"];
const AUX_NAMES: [&str; 2] = ["pol_magnitude", "pol_angle"];
const PARTIAL_WAVE_DATA_SAMPLE_EVENTS: usize = 512;
const PARTIAL_WAVE_ACCMC_SAMPLE_EVENTS: usize = 512;
const PARTIAL_WAVE_GENMC_SAMPLE_EVENTS: usize = 512;
const PARTIAL_WAVE_DATA_SEED: u64 = 11;
const PARTIAL_WAVE_ACCMC_SEED: u64 = 17;
const PARTIAL_WAVE_GENMC_SEED: u64 = 23;
const KMATRIX_DATASET_SEED: u64 = 71;
const DEFAULT_WARMUP_ITERS: usize = 64;
const DEFAULT_ITERS: usize = 4096;

#[derive(Clone, Copy, Debug)]
enum Mode {
    BreitWignerEvaluateLocal,
    BreitWignerNllValue,
    BreitWignerProjectionTotal,
    BreitWignerProjectionGeneratedMc,
    KmatrixValue,
    KmatrixDataTerm,
    KmatrixMcTerm,
}

impl Mode {
    fn parse(raw: &str) -> Option<Self> {
        match raw {
            "bw_evaluate_local" => Some(Self::BreitWignerEvaluateLocal),
            "bw_nll_value" => Some(Self::BreitWignerNllValue),
            "bw_projection_total" => Some(Self::BreitWignerProjectionTotal),
            "bw_projection_generated_mc" => Some(Self::BreitWignerProjectionGeneratedMc),
            "kmatrix_value" => Some(Self::KmatrixValue),
            "kmatrix_data_term" => Some(Self::KmatrixDataTerm),
            "kmatrix_mc_term" => Some(Self::KmatrixMcTerm),
            _ => None,
        }
    }

    fn as_str(self) -> &'static str {
        match self {
            Self::BreitWignerEvaluateLocal => "bw_evaluate_local",
            Self::BreitWignerNllValue => "bw_nll_value",
            Self::BreitWignerProjectionTotal => "bw_projection_total",
            Self::BreitWignerProjectionGeneratedMc => "bw_projection_generated_mc",
            Self::KmatrixValue => "kmatrix_value",
            Self::KmatrixDataTerm => "kmatrix_data_term",
            Self::KmatrixMcTerm => "kmatrix_mc_term",
        }
    }
}

#[derive(Debug)]
struct Args {
    mode: Mode,
    dataset_path: String,
    iterations: usize,
    warmup_iters: usize,
    thread_count: Option<usize>,
    max_events: Option<usize>,
}

impl Args {
    fn parse() -> Self {
        let mut mode = None;
        let mut dataset_path = benchmark_parquet_path();
        let mut iterations = DEFAULT_ITERS;
        let mut warmup_iters = DEFAULT_WARMUP_ITERS;
        let mut thread_count = None;
        let mut max_events = None;

        let mut args = env::args().skip(1);
        while let Some(arg) = args.next() {
            match arg.as_str() {
                "--mode" => mode = Some(parse_mode(args.next())),
                _ if arg.starts_with("--mode=") => {
                    mode = Some(parse_mode(Some(arg["--mode=".len()..].to_string())));
                }
                "--dataset" => {
                    dataset_path = args
                        .next()
                        .unwrap_or_else(|| usage_and_exit("missing dataset path"));
                }
                _ if arg.starts_with("--dataset=") => {
                    dataset_path = arg["--dataset=".len()..].to_string();
                }
                "--iterations" => {
                    iterations = parse_positive(args.next(), "iterations");
                }
                _ if arg.starts_with("--iterations=") => {
                    iterations = parse_positive(
                        Some(arg["--iterations=".len()..].to_string()),
                        "iterations",
                    );
                }
                "--warmup" => {
                    warmup_iters = parse_nonnegative(args.next(), "warmup");
                }
                _ if arg.starts_with("--warmup=") => {
                    warmup_iters =
                        parse_nonnegative(Some(arg["--warmup=".len()..].to_string()), "warmup");
                }
                "--threads" => {
                    thread_count = Some(parse_positive(args.next(), "threads"));
                }
                _ if arg.starts_with("--threads=") => {
                    thread_count = Some(parse_positive(
                        Some(arg["--threads=".len()..].to_string()),
                        "threads",
                    ));
                }
                "--max-events" => {
                    max_events = Some(parse_positive(args.next(), "max-events"));
                }
                _ if arg.starts_with("--max-events=") => {
                    max_events = Some(parse_positive(
                        Some(arg["--max-events=".len()..].to_string()),
                        "max-events",
                    ));
                }
                "--help" | "-h" => {
                    print_usage();
                    process::exit(0);
                }
                other => usage_and_exit(&format!("unknown argument: {other}")),
            }
        }

        Self {
            mode: mode.unwrap_or_else(|| usage_and_exit("missing required --mode")),
            dataset_path,
            iterations,
            warmup_iters,
            thread_count,
            max_events,
        }
    }
}

struct BreitWignerContext {
    nll: Box<NLL>,
    gen_evaluator: Evaluator,
    evaluator: Evaluator,
    params: Vec<f64>,
    n_data_events: usize,
    n_gen_events: usize,
}

fn main() {
    let args = Args::parse();
    match args.mode {
        Mode::BreitWignerEvaluateLocal
        | Mode::BreitWignerNllValue
        | Mode::BreitWignerProjectionTotal
        | Mode::BreitWignerProjectionGeneratedMc => {
            let ctx = build_breit_wigner_context(&args.dataset_path);
            run_breit_wigner_mode(&args, &ctx);
        }
        Mode::KmatrixValue | Mode::KmatrixDataTerm | Mode::KmatrixMcTerm => {
            let nll = build_kmatrix_nll(&args.dataset_path, args.max_events);
            let params = deterministic_parameter_vector(nll.n_free(), -100.0);
            run_kmatrix_mode(&args, &nll, &params);
        }
    }
}

fn print_usage() {
    eprintln!(
        "Usage: cargo run --release -p laddu --bin profile_value_hotspots -- \\\n+         --mode <mode> [--dataset PATH] [--iterations N] [--warmup N] [--threads N] [--max-events N]\n\
         modes:\n\
           bw_evaluate_local\n\
           bw_nll_value\n\
           bw_projection_total\n\
           bw_projection_generated_mc\n\
           kmatrix_value\n\
           kmatrix_data_term\n\
           kmatrix_mc_term"
    );
}

fn usage_and_exit(message: &str) -> ! {
    eprintln!("{message}");
    print_usage();
    process::exit(1);
}

fn parse_mode(value: Option<String>) -> Mode {
    let value = value.unwrap_or_else(|| usage_and_exit("missing value for --mode"));
    Mode::parse(&value).unwrap_or_else(|| usage_and_exit("invalid profiling mode"))
}

fn parse_positive(value: Option<String>, flag: &str) -> usize {
    let parsed = value
        .unwrap_or_else(|| usage_and_exit(&format!("missing value for --{flag}")))
        .parse::<usize>()
        .unwrap_or_else(|_| usage_and_exit(&format!("invalid integer for --{flag}")));
    if parsed == 0 {
        usage_and_exit(&format!("--{flag} must be greater than zero"));
    }
    parsed
}

fn parse_nonnegative(value: Option<String>, flag: &str) -> usize {
    value
        .unwrap_or_else(|| usage_and_exit(&format!("missing value for --{flag}")))
        .parse::<usize>()
        .unwrap_or_else(|_| usage_and_exit(&format!("invalid integer for --{flag}")))
}

fn benchmark_parquet_path() -> String {
    std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
        .join(BENCH_DATASET_RELATIVE_PATH)
        .to_string_lossy()
        .into_owned()
}

fn read_benchmark_dataset(dataset_path: &str) -> Arc<Dataset> {
    let options = DatasetReadOptions::default()
        .p4_names(P4_NAMES)
        .aux_names(AUX_NAMES);
    io::read_parquet(dataset_path, &options).expect("benchmark dataset should open")
}

fn sample_dataset(dataset: &Arc<Dataset>, seed: u64, max_events: usize) -> Arc<Dataset> {
    let mut rng = fastrand::Rng::with_seed(seed);
    let n_total = dataset.n_events();
    let n_take = max_events.min(n_total);
    let indices = rng.subset(n_total, n_take);
    let events = indices
        .into_iter()
        .map(|index| {
            dataset
                .event(index)
                .expect("subset index should be valid")
                .data_arc()
        })
        .collect();
    Arc::new(Dataset::new_with_metadata(events, dataset.metadata_arc()))
}

fn deterministic_parameter_vector(n_free: usize, offset: f64) -> Vec<f64> {
    (0..n_free)
        .map(|index| offset + (index as f64 + 1.0) * 0.25)
        .collect()
}

fn build_breit_wigner_partial_wave_model() -> Expression {
    let topology = Topology::missing_k2("beam", ["kshort1", "kshort2"], "proton");
    let angles = Angles::new(topology.clone(), "kshort1", Frame::Helicity);
    let polarization = Polarization::new(topology, "pol_magnitude", "pol_angle");
    let resonance_mass = Mass::new(["kshort1", "kshort2"]);
    let daughter_1_mass = Mass::new(["kshort1"]);
    let daughter_2_mass = Mass::new(["kshort2"]);

    let z00p = Zlm::new("Z00+", 0, 0, Sign::Positive, &angles, &polarization)
        .expect("z00 should construct");
    let z22p = Zlm::new("Z22+", 2, 2, Sign::Positive, &angles, &polarization)
        .expect("z22 should construct");
    let bw_f01500 = BreitWigner::new(
        "f0(1500)",
        laddu::constant("f0_mass", 1.506),
        parameter("f0_width"),
        0,
        &daughter_1_mass,
        &daughter_2_mass,
        &resonance_mass,
    )
    .expect("f0(1500) should construct");
    let bw_f21525 = BreitWigner::new(
        "f2(1525)",
        laddu::constant("f2_mass", 1.517),
        parameter("f2_width"),
        2,
        &daughter_1_mass,
        &daughter_2_mass,
        &resonance_mass,
    )
    .expect("f2(1525) should construct");
    let s0p = Scalar::new("S0+", parameter("S0+ re")).expect("S0+ should construct");
    let d2p = ComplexScalar::new("D2+", parameter("D2+ re"), parameter("D2+ im"))
        .expect("D2+ should construct");

    let pos_re = (&s0p * &bw_f01500 * z00p.real() + &d2p * &bw_f21525 * z22p.real()).norm_sqr();
    let pos_im = (&s0p * &bw_f01500 * z00p.imag() + &d2p * &bw_f21525 * z22p.imag()).norm_sqr();
    pos_re + pos_im
}

fn build_breit_wigner_context(dataset_path: &str) -> BreitWignerContext {
    let dataset = read_benchmark_dataset(dataset_path);
    let ds_data = sample_dataset(
        &dataset,
        PARTIAL_WAVE_DATA_SEED,
        PARTIAL_WAVE_DATA_SAMPLE_EVENTS,
    );
    let ds_accmc = sample_dataset(
        &dataset,
        PARTIAL_WAVE_ACCMC_SEED,
        PARTIAL_WAVE_ACCMC_SAMPLE_EVENTS,
    );
    let ds_genmc = sample_dataset(
        &dataset,
        PARTIAL_WAVE_GENMC_SEED,
        PARTIAL_WAVE_GENMC_SAMPLE_EVENTS,
    );
    let model = build_breit_wigner_partial_wave_model();
    let nll = NLL::new(&model, &ds_data, &ds_accmc, None).expect("unbinned NLL should build");
    let gen_evaluator = model
        .load(&ds_genmc)
        .expect("generated-mc evaluator should build");
    let evaluator = model.load(&ds_data).expect("data evaluator should build");
    let params = vec![100.0, 0.112, 50.0, 50.0, 0.086];
    BreitWignerContext {
        nll,
        gen_evaluator,
        evaluator,
        params,
        n_data_events: ds_data.n_events(),
        n_gen_events: ds_genmc.n_events(),
    }
}

fn build_kmatrix_nll(dataset_path: &str, max_events: Option<usize>) -> Box<NLL> {
    let dataset = read_benchmark_dataset(dataset_path);
    let (ds_data, ds_mc) = if let Some(max_events) = max_events {
        let sampled = sample_dataset(&dataset, KMATRIX_DATASET_SEED, max_events);
        (sampled.clone(), sampled)
    } else {
        (dataset.clone(), dataset)
    };
    let topology = Topology::missing_k2("beam", ["kshort1", "kshort2"], "proton");
    let angles = Angles::new(topology.clone(), "kshort1", Frame::Helicity);
    let polarization = Polarization::new(topology.clone(), "pol_magnitude", "pol_angle");
    let resonance_mass = Mass::new(["kshort1", "kshort2"]);
    let z00p = Zlm::new("Z00+", 0, 0, Sign::Positive, &angles, &polarization)
        .expect("z00+ should construct");
    let z00n = Zlm::new("Z00-", 0, 0, Sign::Negative, &angles, &polarization)
        .expect("z00- should construct");
    let z22p = Zlm::new("Z22+", 2, 2, Sign::Positive, &angles, &polarization)
        .expect("z22+ should construct");
    let f0p = KopfKMatrixF0::new(
        "f0+",
        [
            [
                laddu::constant("f0+ c00 re", 0.0),
                laddu::constant("f0+ c00 im", 0.0),
            ],
            [
                parameter("f0(980)+ re"),
                laddu::constant("f0(980)+ im_fix", 0.0),
            ],
            [parameter("f0(1370)+ re"), parameter("f0(1370)+ im")],
            [parameter("f0(1500)+ re"), parameter("f0(1500)+ im")],
            [parameter("f0(1710)+ re"), parameter("f0(1710)+ im")],
        ],
        0,
        &resonance_mass,
        None,
    )
    .expect("f0+ should construct");
    let a0p = KopfKMatrixA0::new(
        "a0+",
        [
            [parameter("a0(980)+ re"), parameter("a0(980)+ im")],
            [parameter("a0(1450)+ re"), parameter("a0(1450)+ im")],
        ],
        0,
        &resonance_mass,
        None,
    )
    .expect("a0+ should construct");
    let f0n = KopfKMatrixF0::new(
        "f0-",
        [
            [
                laddu::constant("f0- c00 re", 0.0),
                laddu::constant("f0- c00 im", 0.0),
            ],
            [
                parameter("f0(980)- re"),
                laddu::constant("f0(980)- im_fix", 0.0),
            ],
            [parameter("f0(1370)- re"), parameter("f0(1370)- im")],
            [parameter("f0(1500)- re"), parameter("f0(1500)- im")],
            [parameter("f0(1710)- re"), parameter("f0(1710)- im")],
        ],
        0,
        &resonance_mass,
        None,
    )
    .expect("f0- should construct");
    let a0n = KopfKMatrixA0::new(
        "a0-",
        [
            [parameter("a0(980)- re"), parameter("a0(980)- im")],
            [parameter("a0(1450)- re"), parameter("a0(1450)- im")],
        ],
        0,
        &resonance_mass,
        None,
    )
    .expect("a0- should construct");
    let f2 = KopfKMatrixF2::new(
        "f2",
        [
            [parameter("f2(1270) re"), parameter("f2(1270) im")],
            [parameter("f2(1525) re"), parameter("f2(1525) im")],
            [parameter("f2(1850) re"), parameter("f2(1850) im")],
            [parameter("f2(1910) re"), parameter("f2(1910) im")],
        ],
        2,
        &resonance_mass,
        None,
    )
    .expect("f2 should construct");
    let a2 = KopfKMatrixA2::new(
        "a2",
        [
            [parameter("a2(1320) re"), parameter("a2(1320) im")],
            [parameter("a2(1700) re"), parameter("a2(1700) im")],
        ],
        2,
        &resonance_mass,
        None,
    )
    .expect("a2 should construct");
    let s0p = f0p + a0p;
    let s0n = f0n + a0n;
    let d2p = f2 + a2;
    let pos_re = (&s0p * z00p.real() + &d2p * z22p.real()).norm_sqr();
    let pos_im = (&s0p * z00p.imag() + &d2p * z22p.imag()).norm_sqr();
    let neg_re = (&s0n * z00n.real()).norm_sqr();
    let neg_im = (&s0n * z00n.imag()).norm_sqr();
    let expr = pos_re + pos_im + neg_re + neg_im;
    NLL::new(&expr, &ds_data, &ds_mc, None).expect("k-matrix NLL should build")
}

fn run_breit_wigner_mode(args: &Args, ctx: &BreitWignerContext) {
    let warmup = args.warmup_iters;
    let iters = args.iterations;
    match args.mode {
        Mode::BreitWignerEvaluateLocal => {
            for _ in 0..warmup {
                let values = ctx.evaluator.evaluate(black_box(&ctx.params));
                black_box(values.iter().map(|value| value.re).sum::<f64>());
            }
            let mut sink = 0.0;
            let t0 = Instant::now();
            for _ in 0..iters {
                let values = ctx.evaluator.evaluate(black_box(&ctx.params));
                sink += black_box(values.iter().map(|value| value.re).sum::<f64>());
            }
            eprintln!(
                "mode={} iters={iters} warmup={warmup} n_events={} sink={sink} elapsed={:?}",
                args.mode.as_str(),
                ctx.n_data_events,
                t0.elapsed(),
            );
        }
        Mode::BreitWignerNllValue => {
            for _ in 0..warmup {
                black_box(
                    ctx.nll
                        .evaluate(black_box(&ctx.params))
                        .expect("breit-wigner nll should succeed"),
                );
            }
            let mut sink = 0.0;
            let t0 = Instant::now();
            for _ in 0..iters {
                sink += black_box(
                    ctx.nll
                        .evaluate(black_box(&ctx.params))
                        .expect("breit-wigner nll should succeed"),
                );
            }
            eprintln!(
                "mode={} iters={iters} warmup={warmup} n_events={} sink={sink} elapsed={:?}",
                args.mode.as_str(),
                ctx.n_data_events,
                t0.elapsed(),
            );
        }
        Mode::BreitWignerProjectionTotal => {
            for _ in 0..warmup {
                let weights = ctx
                    .nll
                    .project_weights(black_box(&ctx.params), None)
                    .expect("breit-wigner projection should succeed");
                black_box(weights.iter().sum::<f64>());
            }
            let mut sink = 0.0;
            let t0 = Instant::now();
            for _ in 0..iters {
                let weights = ctx
                    .nll
                    .project_weights(black_box(&ctx.params), None)
                    .expect("breit-wigner projection should succeed");
                sink += black_box(weights.iter().sum::<f64>());
            }
            eprintln!(
                "mode={} iters={iters} warmup={warmup} n_events={} sink={sink} elapsed={:?}",
                args.mode.as_str(),
                ctx.n_data_events,
                t0.elapsed(),
            );
        }
        Mode::BreitWignerProjectionGeneratedMc => {
            for _ in 0..warmup {
                let weights = ctx
                    .nll
                    .project_weights(black_box(&ctx.params), Some(ctx.gen_evaluator.clone()))
                    .expect("generated-mc projection should succeed");
                black_box(weights.iter().sum::<f64>());
            }
            let mut sink = 0.0;
            let t0 = Instant::now();
            for _ in 0..iters {
                let weights = ctx
                    .nll
                    .project_weights(black_box(&ctx.params), Some(ctx.gen_evaluator.clone()))
                    .expect("generated-mc projection should succeed");
                sink += black_box(weights.iter().sum::<f64>());
            }
            eprintln!(
                "mode={} iters={iters} warmup={warmup} n_events={} sink={sink} elapsed={:?}",
                args.mode.as_str(),
                ctx.n_gen_events,
                t0.elapsed(),
            );
        }
        Mode::KmatrixValue | Mode::KmatrixDataTerm | Mode::KmatrixMcTerm => {
            unreachable!("kmatrix mode handled separately")
        }
    }
}

fn run_kmatrix_mode(args: &Args, nll: &NLL, params: &[f64]) {
    let warmup = args.warmup_iters;
    let iters = args.iterations;
    let threads = args.thread_count.unwrap_or_else(|| {
        std::thread::available_parallelism()
            .map(usize::from)
            .unwrap_or(1)
    });
    let pool = ThreadPoolBuilder::new()
        .num_threads(threads)
        .build()
        .expect("rayon pool should build");

    let run_once = || match args.mode {
        Mode::KmatrixValue => nll
            .evaluate(black_box(params))
            .expect("k-matrix value should succeed"),
        Mode::KmatrixDataTerm => nll.profile_data_term_local_value(black_box(params)),
        Mode::KmatrixMcTerm => nll.profile_mc_term_local_value(black_box(params)),
        Mode::BreitWignerEvaluateLocal
        | Mode::BreitWignerNllValue
        | Mode::BreitWignerProjectionTotal
        | Mode::BreitWignerProjectionGeneratedMc => unreachable!("non-kmatrix mode"),
    };

    for _ in 0..warmup {
        black_box(pool.install(run_once));
    }
    let mut sink = 0.0;
    let t0 = Instant::now();
    for _ in 0..iters {
        sink += black_box(pool.install(run_once));
    }
    eprintln!(
        "mode={} iters={iters} warmup={warmup} threads={threads} n_events={} sink={sink} elapsed={:?}",
        args.mode.as_str(),
        nll.data_evaluator.dataset.n_events(),
        t0.elapsed(),
    );
}
