use std::time::Instant;

use kira_spatial_core as kcore;
use kira_spatial_io::{Dataset, LoadConfig};

use crate::cli::{RunArgs, SignalKind};
use crate::error::OrchestratorError;
use crate::feature_slice::{
    FeatureSliceCache, detect_feature_slice_paths, feature_slice_cache_path,
    load_feature_slice_cache, load_feature_slice_values_from_h5, load_or_build_feature_index,
    open_h5, write_feature_slice_cache,
};
use crate::pipeline::grid::GridMapping;
use crate::pipeline::log_timing;
use crate::pipeline::signal::{compute_signal, find_gene_id};

use crate::pipeline::stats::{
    ColocalizeArtifacts, DomainArtifacts, SpatialStatsArtifacts, build_extras,
};

/// Result of the load+signal stage, including optional spatial-stats / domain /
/// co-localization artifacts computed on the normalized fields.
pub(crate) struct LoadedSignal {
    pub(crate) grid: GridMapping,
    pub(crate) signal_dense: Vec<f32>,
    pub(crate) per_gene_signal_dense: Vec<(String, Vec<f32>)>,
    pub(crate) stats: Option<SpatialStatsArtifacts>,
    pub(crate) domains: Option<DomainArtifacts>,
    pub(crate) coloc: Option<ColocalizeArtifacts>,
}

pub(crate) fn is_feature_slice_matrix_missing(err: &kira_spatial_io::SpatialIoError) -> bool {
    err.to_string().contains("missing /matrix/barcodes dataset")
}

/// Coarsen the grid by `factor` (N×N fine cells → one super-bin): per-gene
/// values are pooled by mean, the library size by sum, and a coarse bin is valid
/// if any of its fine bins was. Returns the coarse grid, core domain, validity
/// mask, (optional) library size, and pooled per-gene fields. This densifies
/// sparse HD bins and lifts the whole analysis to a compartment scale.
#[allow(clippy::type_complexity)]
fn apply_bin_aggregate(
    grid: &GridMapping,
    valid: &[bool],
    libsize: Option<&[f32]>,
    per_gene: &[Vec<f32>],
    factor: u32,
) -> Result<
    (
        GridMapping,
        kcore::SpatialDomain,
        Vec<bool>,
        Option<Vec<f32>>,
        Vec<Vec<f32>>,
    ),
    OrchestratorError,
> {
    let (coarse_grid, assign, n_coarse) = grid.coarsen_plan(factor)?;
    let mut counts = vec![0_u32; n_coarse];
    let mut coarse_valid = vec![false; n_coarse];
    for (i, &c) in assign.iter().enumerate() {
        counts[c as usize] += 1;
        if valid[i] {
            coarse_valid[c as usize] = true;
        }
    }
    let coarse_libsize = libsize.map(|ls| {
        let mut out = vec![0.0_f32; n_coarse];
        for (i, &c) in assign.iter().enumerate() {
            out[c as usize] += ls[i];
        }
        out
    });
    let mut coarse_per_gene = Vec::with_capacity(per_gene.len());
    for g in per_gene {
        let mut acc = vec![0.0_f32; n_coarse];
        for (i, &c) in assign.iter().enumerate() {
            acc[c as usize] += g[i];
        }
        for (c, a) in acc.iter_mut().enumerate() {
            if counts[c] > 0 {
                *a /= counts[c] as f32;
            }
        }
        coarse_per_gene.push(acc);
    }
    let coarse_domain = coarse_grid.to_core_domain_with_mask(&coarse_valid)?;
    eprintln!(
        "[bin-aggregate] factor {} -> {} coarse bins (from {} fine)",
        factor,
        n_coarse,
        assign.len()
    );
    Ok((
        coarse_grid,
        coarse_domain,
        coarse_valid,
        coarse_libsize,
        coarse_per_gene,
    ))
}

pub(crate) fn compute_signal_dense_from_dataset(
    args: &RunArgs,
    dataset: &Dataset,
) -> Result<LoadedSignal, OrchestratorError> {
    let t_total = Instant::now();
    let io_domain = dataset.spatial_domain();
    let io_expr = dataset.expression_csr();
    let features = &dataset.features().rows;

    let grid = GridMapping::from_io_domain(io_domain)?;
    let core_domain = grid.to_core_domain(io_domain)?;

    let mut core_indptr = Vec::with_capacity(io_expr.indptr.len());
    for i in 0..io_expr.indptr.len() {
        let p = io_expr.indptr.get(i);
        core_indptr.push(
            u32::try_from(p)
                .map_err(|_| OrchestratorError::InvalidInput("csr indptr does not fit u32"))?,
        );
    }
    let core_expr_view = kcore::BinMajorCsrView::new(
        io_expr.n_bins,
        io_expr.n_genes,
        &core_indptr,
        &io_expr.indices,
        &io_expr.data,
    );

    // Per-bin library sizes (total counts over ALL genes) for CPM normalization.
    let indptr_u64 = io_expr.indptr.to_u64_vec();
    let libsize = crate::pipeline::normalize::bin_library_sizes(&indptr_u64, &io_expr.data);
    let valid: Vec<bool> = io_domain.tissue_mask.iter().map(|b| *b).collect();

    // Build, normalize, and collect each requested gene's dense field.
    let mut per_gene_norm = Vec::<Vec<f32>>::with_capacity(args.genes.len());
    for gene in &args.genes {
        let gene_id = find_gene_id(features, gene)?;
        let gene_field = kcore::build_gene_field(
            &core_domain,
            &core_expr_view,
            gene_id,
            kcore::GeneFieldConfig::default(),
        )?;
        let mut vals: Vec<f32> = gene_field.values().iter().map(|&v| v as f32).collect();
        crate::pipeline::normalize::normalize_gene(
            &mut vals,
            Some(&libsize),
            args.normalize,
            args.target_sum,
        );
        per_gene_norm.push(vals);
    }

    // Optional bin aggregation: pool fine bins into N×N super-bins for the whole
    // downstream pipeline (operators, images, stats).
    let (grid, core_domain, valid, libsize, per_gene_norm) = if args.bin_aggregate > 1 {
        let (g, d, v, ls, pg) =
            apply_bin_aggregate(&grid, &valid, Some(&libsize), &per_gene_norm, args.bin_aggregate)?;
        (g, d, v, ls.unwrap_or_default(), pg)
    } else {
        (grid, core_domain, valid, libsize, per_gene_norm)
    };

    // Materialize normalized per-gene core fields (for per-gene output + cross-gradient).
    let mut per_gene_fields = Vec::<(String, kcore::Field<'_>)>::with_capacity(args.genes.len());
    for (gene, vals) in args.genes.iter().zip(per_gene_norm.iter()) {
        per_gene_fields.push((gene.clone(), kcore::Field::new(&core_domain, vals.clone())?));
    }

    // Aggregate the panel into a single field (module score by default).
    let core_values =
        crate::pipeline::normalize::aggregate_panel(&per_gene_norm, &valid, args.panel)?;
    let core_gene = kcore::Field::new(&core_domain, core_values)?;

    let signal_sparse = if args.signal == SignalKind::CrossGradient {
        if per_gene_fields.len() < 2 {
            return Err(OrchestratorError::InvalidInput(
                "cross-gradient requires at least two values in --genes",
            ));
        }
        compute_signal(
            args.signal,
            &core_domain,
            &per_gene_fields[0].1,
            Some(&per_gene_fields[1].1),
            args.pre_smooth_sigma,
            args.fractal_window,
        )?
    } else {
        compute_signal(
            args.signal,
            &core_domain,
            &core_gene,
            None,
            args.pre_smooth_sigma,
            args.fractal_window,
        )?
    };
    let signal_dense = grid.scatter_to_dense(&signal_sparse)?;

    let mut per_gene_signal_dense = Vec::<(String, Vec<f32>)>::new();
    for (gene, gene_field) in &per_gene_fields {
        let gene_signal_sparse = if args.signal == SignalKind::CrossGradient {
            compute_signal(
                args.signal,
                &core_domain,
                gene_field,
                Some(&core_gene),
                args.pre_smooth_sigma,
                args.fractal_window,
            )?
        } else {
            compute_signal(
                args.signal,
                &core_domain,
                gene_field,
                None,
                args.pre_smooth_sigma,
                args.fractal_window,
            )?
        };
        let gene_signal_dense = grid.scatter_to_dense(&gene_signal_sparse)?;
        per_gene_signal_dense.push((gene.clone(), gene_signal_dense));
    }
    let per_gene_ref: Vec<(&str, &[f32])> = args
        .genes
        .iter()
        .map(|g| g.as_str())
        .zip(per_gene_norm.iter().map(|v| v.as_slice()))
        .collect();
    let (stats, domains, coloc) =
        build_extras(&core_gene, &core_domain, &grid, args, &per_gene_ref, Some(&libsize))?;
    log_timing("dataset_path_total", t_total);
    Ok(LoadedSignal {
        grid,
        signal_dense,
        per_gene_signal_dense,
        stats,
        domains,
        coloc,
    })
}

pub(crate) fn compute_signal_dense_from_feature_slice(
    args: &RunArgs,
) -> Result<LoadedSignal, OrchestratorError> {
    let t_total = Instant::now();
    let Some(paths) = detect_feature_slice_paths(&args.h5) else {
        return Err(OrchestratorError::InvalidInputOwned(format!(
            "feature_slice.h5 layout not detected at {}",
            args.h5.display()
        )));
    };

    let mut genes = args.genes.iter();
    let first_gene = genes.next().ok_or(OrchestratorError::InvalidInput(
        "at least one gene is required in --genes",
    ))?;

    let cache_path = feature_slice_cache_path(&paths.parquet_path);
    eprintln!(
        "[timing] feature_slice_cache path={} exists={}",
        cache_path.display(),
        cache_path.is_file()
    );
    let cache = load_feature_slice_cache(&cache_path).ok();
    eprintln!("[timing] feature_slice_cache loaded={}", cache.is_some());

    let fallback_io_path = || compute_signal_dense_from_feature_slice_via_io(args);

    let (domain, summed, cache_for_rest, h5_and_index) = if let Some(cache) = cache {
        let t_fast = Instant::now();
        let h5 = open_h5(&paths.h5_path)?;
        let feature_index = match load_or_build_feature_index(&paths.h5_path, &h5) {
            Ok(v) => v,
            Err(_) => return fallback_io_path(),
        };
        let first_idx = *feature_index.get(first_gene).ok_or_else(|| {
            OrchestratorError::InvalidInputOwned(format!(
                "gene not found in /features/name: {first_gene}"
            ))
        })?;
        let domain = cache.domain.clone();
        let first_vals =
            load_feature_slice_values_from_h5(&h5, first_idx, domain.len(), &cache.coord_to_bin)?;
        log_timing("feature_slice_fast_first_gene", t_fast);
        (domain, first_vals, Some(cache), Some((h5, feature_index)))
    } else {
        let t_slow = Instant::now();
        let first =
            kira_spatial_io::load_feature_slice_gene(&args.h5, first_gene, LoadConfig::default())?;
        let built_cache = FeatureSliceCache::from_domain(&first.spatial_domain)?;
        if let Err(err) = write_feature_slice_cache(&cache_path, &built_cache) {
            eprintln!(
                "warning: failed writing feature-slice mmap cache {}: {}",
                cache_path.display(),
                err
            );
        }
        log_timing("feature_slice_io_first_gene", t_slow);
        let h5_and_index = match open_h5(&paths.h5_path)
            .and_then(|h5| load_or_build_feature_index(&paths.h5_path, &h5).map(|idx| (h5, idx)))
        {
            Ok(v) => Some(v),
            Err(err) => {
                eprintln!(
                    "warning: failed switching to fast path after cache build; staying io-only: {}",
                    err
                );
                None
            }
        };
        (
            first.spatial_domain,
            first.values,
            Some(built_cache),
            h5_and_index,
        )
    };
    let expected_len = summed.len();
    let mut per_gene_sparse = Vec::<(String, Vec<f32>)>::new();
    per_gene_sparse.push((first_gene.clone(), summed));

    // Per-bin library size (sum over all features) for CPM on the feature-slice
    // path; only the fast (H5 + coord-cache) branch can compute it.
    let mut libsize_opt: Option<Vec<f32>> = None;

    if let (Some(cache), Some((h5, feature_index))) = (cache_for_rest, h5_and_index) {
        for gene in genes {
            let t_gene = Instant::now();
            let idx = *feature_index.get(gene).ok_or_else(|| {
                OrchestratorError::InvalidInputOwned(format!(
                    "gene not found in /features/name: {gene}"
                ))
            })?;
            let values =
                load_feature_slice_values_from_h5(&h5, idx, domain.len(), &cache.coord_to_bin)?;
            per_gene_sparse.push((gene.clone(), values));
            eprintln!(
                "[timing] feature_slice_fast_gene {} {:?}",
                gene,
                t_gene.elapsed()
            );
        }
        // Library size is needed for CPM normalization and/or the --min-counts QC mask.
        if args.normalize == crate::cli::NormalizeMode::CpmLog1p || args.min_counts > 0.0 {
            match crate::feature_slice::compute_feature_slice_library_sizes(
                &h5,
                feature_index.len(),
                domain.len(),
                &cache.coord_to_bin,
                &cache_path,
            ) {
                Ok(ls) => libsize_opt = Some(ls),
                Err(e) => eprintln!(
                    "warning: feature-slice library-size computation failed (CPM falls back to \
                     log1p; QC mask skipped): {e}"
                ),
            }
        }
    } else {
        for gene in genes {
            let t_gene = Instant::now();
            let loaded =
                kira_spatial_io::load_feature_slice_gene(&args.h5, gene, LoadConfig::default())?;
            if loaded.spatial_domain.len() != domain.len() || loaded.values.len() != expected_len {
                return Err(OrchestratorError::InvalidInput(
                    "feature-slice domains mismatch across requested genes",
                ));
            }
            per_gene_sparse.push((gene.clone(), loaded.values));
            eprintln!(
                "[timing] feature_slice_io_gene {} {:?}",
                gene,
                t_gene.elapsed()
            );
        }
    }

    let t_post = Instant::now();
    let grid = GridMapping::from_io_domain(&domain)?;
    let core_domain = grid.to_core_domain(&domain)?;
    // CPM uses the per-bin library size when available (fast path); otherwise
    // CpmLog1p degrades to log1p inside normalize_gene.
    let valid: Vec<bool> = domain.tissue_mask.iter().map(|b| *b).collect();
    for (_, vals) in per_gene_sparse.iter_mut() {
        crate::pipeline::normalize::normalize_gene(
            vals,
            libsize_opt.as_deref(),
            args.normalize,
            args.target_sum,
        );
    }
    let (grid, core_domain, valid, libsize_opt) = if args.bin_aggregate > 1 {
        let fine_pg: Vec<Vec<f32>> = per_gene_sparse.iter().map(|(_, v)| v.clone()).collect();
        let (cg, cd, cv, cls, cpg) =
            apply_bin_aggregate(&grid, &valid, libsize_opt.as_deref(), &fine_pg, args.bin_aggregate)?;
        for (k, (_, v)) in per_gene_sparse.iter_mut().enumerate() {
            *v = cpg[k].clone();
        }
        (cg, cd, cv, cls)
    } else {
        (grid, core_domain, valid, libsize_opt)
    };
    let per_gene_vecs: Vec<Vec<f32>> = per_gene_sparse.iter().map(|(_, v)| v.clone()).collect();
    let core_values =
        crate::pipeline::normalize::aggregate_panel(&per_gene_vecs, &valid, args.panel)?;
    let core_gene = kcore::Field::new(&core_domain, core_values)?;

    let signal_sparse = if args.signal == SignalKind::CrossGradient {
        if per_gene_sparse.len() < 2 {
            return Err(OrchestratorError::InvalidInput(
                "cross-gradient requires at least two values in --genes",
            ));
        }
        let a = kcore::Field::new(&core_domain, per_gene_sparse[0].1.clone())?;
        let b = kcore::Field::new(&core_domain, per_gene_sparse[1].1.clone())?;
        compute_signal(
            args.signal,
            &core_domain,
            &a,
            Some(&b),
            args.pre_smooth_sigma,
            args.fractal_window,
        )?
    } else {
        compute_signal(
            args.signal,
            &core_domain,
            &core_gene,
            None,
            args.pre_smooth_sigma,
            args.fractal_window,
        )?
    };
    let signal_dense = grid.scatter_to_dense(&signal_sparse)?;
    let mut per_gene_signal_dense = Vec::<(String, Vec<f32>)>::with_capacity(per_gene_sparse.len());
    for (gene, sparse_vals) in per_gene_sparse {
        let gf = kcore::Field::new(&core_domain, sparse_vals)?;
        let sig_sparse = if args.signal == SignalKind::CrossGradient {
            compute_signal(
                args.signal,
                &core_domain,
                &gf,
                Some(&core_gene),
                args.pre_smooth_sigma,
                args.fractal_window,
            )?
        } else {
            compute_signal(
                args.signal,
                &core_domain,
                &gf,
                None,
                args.pre_smooth_sigma,
                args.fractal_window,
            )?
        };
        let sig_dense = grid.scatter_to_dense(&sig_sparse)?;
        per_gene_signal_dense.push((gene, sig_dense));
    }
    let per_gene_ref: Vec<(&str, &[f32])> = args
        .genes
        .iter()
        .map(|g| g.as_str())
        .zip(per_gene_vecs.iter().map(|v| v.as_slice()))
        .collect();
    let (stats, domains, coloc) = build_extras(
        &core_gene,
        &core_domain,
        &grid,
        args,
        &per_gene_ref,
        libsize_opt.as_deref(),
    )?;
    log_timing("feature_slice_postprocess", t_post);
    log_timing("feature_slice_total", t_total);
    Ok(LoadedSignal {
        grid,
        signal_dense,
        per_gene_signal_dense,
        stats,
        domains,
        coloc,
    })
}

pub(crate) fn compute_signal_dense_from_feature_slice_via_io(
    args: &RunArgs,
) -> Result<LoadedSignal, OrchestratorError> {
    let t_total = Instant::now();
    eprintln!("[timing] fallback=io_only");
    let mut genes = args.genes.iter();
    let first_gene = genes.next().ok_or(OrchestratorError::InvalidInput(
        "at least one gene is required in --genes",
    ))?;
    let first =
        kira_spatial_io::load_feature_slice_gene(&args.h5, first_gene, LoadConfig::default())?;

    let summed = first.values;
    let domain = first.spatial_domain;
    let expected_len = summed.len();
    let mut per_gene_sparse = Vec::<(String, Vec<f32>)>::new();
    per_gene_sparse.push((first_gene.clone(), summed));

    for gene in genes {
        let loaded =
            kira_spatial_io::load_feature_slice_gene(&args.h5, gene, LoadConfig::default())?;
        if loaded.spatial_domain.len() != domain.len() || loaded.values.len() != expected_len {
            return Err(OrchestratorError::InvalidInput(
                "feature-slice domains mismatch across requested genes",
            ));
        }
        per_gene_sparse.push((gene.clone(), loaded.values));
    }

    let grid = GridMapping::from_io_domain(&domain)?;
    let core_domain = grid.to_core_domain(&domain)?;
    // feature-slice inputs carry no library size; CpmLog1p degrades to log1p.
    let valid: Vec<bool> = domain.tissue_mask.iter().map(|b| *b).collect();
    for (_, vals) in per_gene_sparse.iter_mut() {
        crate::pipeline::normalize::normalize_gene(vals, None, args.normalize, args.target_sum);
    }
    let (grid, core_domain, valid) = if args.bin_aggregate > 1 {
        let fine_pg: Vec<Vec<f32>> = per_gene_sparse.iter().map(|(_, v)| v.clone()).collect();
        let (cg, cd, cv, _cls, cpg) =
            apply_bin_aggregate(&grid, &valid, None, &fine_pg, args.bin_aggregate)?;
        for (k, (_, v)) in per_gene_sparse.iter_mut().enumerate() {
            *v = cpg[k].clone();
        }
        (cg, cd, cv)
    } else {
        (grid, core_domain, valid)
    };
    let per_gene_vecs: Vec<Vec<f32>> = per_gene_sparse.iter().map(|(_, v)| v.clone()).collect();
    let core_values =
        crate::pipeline::normalize::aggregate_panel(&per_gene_vecs, &valid, args.panel)?;
    let core_gene = kcore::Field::new(&core_domain, core_values)?;

    let signal_sparse = if args.signal == SignalKind::CrossGradient {
        if per_gene_sparse.len() < 2 {
            return Err(OrchestratorError::InvalidInput(
                "cross-gradient requires at least two values in --genes",
            ));
        }
        let a = kcore::Field::new(&core_domain, per_gene_sparse[0].1.clone())?;
        let b = kcore::Field::new(&core_domain, per_gene_sparse[1].1.clone())?;
        compute_signal(
            args.signal,
            &core_domain,
            &a,
            Some(&b),
            args.pre_smooth_sigma,
            args.fractal_window,
        )?
    } else {
        compute_signal(
            args.signal,
            &core_domain,
            &core_gene,
            None,
            args.pre_smooth_sigma,
            args.fractal_window,
        )?
    };
    let signal_dense = grid.scatter_to_dense(&signal_sparse)?;
    let mut per_gene_signal_dense = Vec::<(String, Vec<f32>)>::with_capacity(per_gene_sparse.len());
    for (gene, sparse_vals) in per_gene_sparse {
        let gf = kcore::Field::new(&core_domain, sparse_vals)?;
        let sig_sparse = if args.signal == SignalKind::CrossGradient {
            compute_signal(
                args.signal,
                &core_domain,
                &gf,
                Some(&core_gene),
                args.pre_smooth_sigma,
                args.fractal_window,
            )?
        } else {
            compute_signal(
                args.signal,
                &core_domain,
                &gf,
                None,
                args.pre_smooth_sigma,
                args.fractal_window,
            )?
        };
        let sig_dense = grid.scatter_to_dense(&sig_sparse)?;
        per_gene_signal_dense.push((gene, sig_dense));
    }
    let per_gene_ref: Vec<(&str, &[f32])> = args
        .genes
        .iter()
        .map(|g| g.as_str())
        .zip(per_gene_vecs.iter().map(|v| v.as_slice()))
        .collect();
    let (stats, domains, coloc) =
        build_extras(&core_gene, &core_domain, &grid, args, &per_gene_ref, None)?;
    log_timing("feature_slice_io_total", t_total);
    Ok(LoadedSignal {
        grid,
        signal_dense,
        per_gene_signal_dense,
        stats,
        domains,
        coloc,
    })
}
