//! Spatial-statistics and spatial-domain artifacts for the CLI pipeline.
//!
//! These run on the *normalized aggregate expression field* (not the operator
//! output), which is the scientifically meaningful target for spatial
//! autocorrelation (SVG-style) and domain detection.

use kira_spatial_core as kcore;
use serde::Serialize;

use crate::cli::RunArgs;
use crate::error::OrchestratorError;
use crate::pipeline::grid::GridMapping;

/// Global spatial-autocorrelation summary (serialized to `spatial_stats.json`).
#[derive(Debug, Serialize)]
pub(crate) struct SpatialStatsGlobal {
    pub global_moran_i: f64,
    pub global_moran_p_value: f64,
    pub global_moran_z_score: f64,
    pub global_geary_c: f64,
    pub n_valid: u32,
    pub weight_sum: f64,
    pub permutations: u32,
    pub seed: u64,
    pub neighbor_mode: &'static str,
    pub n_local_significant_q05: u32,
}

pub(crate) struct SpatialStatsArtifacts {
    pub global: SpatialStatsGlobal,
    pub nx: usize,
    pub ny: usize,
    pub local_moran_z: Vec<f32>,
    pub local_moran_q: Vec<f32>,
    pub getis_z: Vec<f32>,
    pub getis_q: Vec<f32>,
}

/// Spatial-domain (community) artifacts (serialized to `domains.json` + image).
#[derive(Debug, Serialize)]
pub(crate) struct DomainSummary {
    pub n_communities: u32,
    pub modularity: f64,
    pub passes: u32,
    pub resolution: f64,
    pub sizes: Vec<u32>,
}

pub(crate) struct DomainArtifacts {
    pub summary: DomainSummary,
    pub nx: usize,
    pub ny: usize,
    /// Dense per-cell labels shifted by +1 (0 = background / empty cell).
    pub labels_plus1: Vec<f32>,
}

/// Bivariate-Moran co-localization artifacts (serialized to `colocalize.json` + image).
#[derive(Debug, Serialize)]
pub(crate) struct ColocalizeSummary {
    pub gene_a: String,
    pub gene_b: String,
    pub global_bivariate_moran_r: f64,
    pub global_p_value: f64,
    pub global_z_score: f64,
    pub n_valid: u32,
    pub weight_sum: f64,
    pub permutations: u32,
    pub seed: u64,
    pub neighbor_mode: &'static str,
    pub neighbor_radius: u32,
    pub n_local_significant_q05: u32,
    pub interpretation: String,
}

pub(crate) struct ColocalizeArtifacts {
    pub summary: ColocalizeSummary,
    pub nx: usize,
    pub ny: usize,
    /// Dense local bivariate-Moran z-score (sign = co-cluster vs segregate).
    pub local_z: Vec<f32>,
    /// Dense local BH-FDR q-value.
    pub local_q: Vec<f32>,
}

/// Compute optional spatial-stats, spatial-domain, and co-localization artifacts
/// from the normalized fields, gated by CLI flags. `per_gene` holds the per-gene
/// normalized values aligned by gene name (used only for `--colocalize`).
pub(crate) fn build_extras(
    core_gene: &kcore::Field<'_>,
    core_domain: &kcore::SpatialDomain,
    grid: &GridMapping,
    args: &RunArgs,
    per_gene: &[(&str, &[f32])],
    qc_libsize: Option<&[f32]>,
) -> Result<
    (
        Option<SpatialStatsArtifacts>,
        Option<DomainArtifacts>,
        Option<ColocalizeArtifacts>,
    ),
    OrchestratorError,
> {
    if !(args.spatial_stats || args.domains || args.colocalize.is_some()) {
        return Ok((None, None, None));
    }

    // Optionally tighten the domain validity with a per-bin QC mask (min library
    // size), so spatial statistics ignore empty background bins.
    let qc_domain_owned: Option<kcore::SpatialDomain> = if args.min_counts > 0.0 {
        match qc_libsize {
            Some(libsize) => {
                let mut valid_qc = Vec::with_capacity(core_domain.len());
                for i in 0..core_domain.len() {
                    let lib = libsize.get(i).copied().unwrap_or(0.0);
                    valid_qc.push(core_domain.is_valid(i) && lib >= args.min_counts);
                }
                let kept = valid_qc.iter().filter(|&&v| v).count();
                eprintln!(
                    "[qc] min-counts={} -> {} / {} bins pass QC ({:.1}%)",
                    args.min_counts,
                    kept,
                    core_domain.len(),
                    100.0 * kept as f64 / core_domain.len().max(1) as f64
                );
                Some(grid.to_core_domain_with_mask(&valid_qc)?)
            }
            None => {
                eprintln!(
                    "warning: --min-counts requested but no per-bin library size is available on \
                     this input path; QC mask skipped"
                );
                None
            }
        }
    } else {
        None
    };
    let eff_domain: &kcore::SpatialDomain = qc_domain_owned.as_ref().unwrap_or(core_domain);

    // Re-materialize the aggregate field on the effective domain (same per-bin
    // values; only validity differs).
    let eff_gene = kcore::Field::new(eff_domain, core_gene.values().to_vec())?;

    let stats = if args.spatial_stats {
        Some(compute_spatial_stats(&eff_gene, eff_domain, grid, args)?)
    } else {
        None
    };
    let domains = if args.domains {
        Some(compute_domains(&eff_gene, eff_domain, grid, args)?)
    } else {
        None
    };
    let coloc = if args.colocalize.is_some() {
        Some(compute_colocalize(eff_domain, grid, args, per_gene)?)
    } else {
        None
    };
    Ok((stats, domains, coloc))
}

fn compute_colocalize(
    core_domain: &kcore::SpatialDomain,
    grid: &GridMapping,
    args: &RunArgs,
    per_gene: &[(&str, &[f32])],
) -> Result<ColocalizeArtifacts, OrchestratorError> {
    let (gene_a, gene_b) = args.colocalize_pair().ok_or(OrchestratorError::InvalidInput(
        "--colocalize expects exactly two comma-separated gene names: GENE_A,GENE_B",
    ))?;
    let find = |name: &str| -> Result<&[f32], OrchestratorError> {
        per_gene
            .iter()
            .find(|(n, _)| *n == name)
            .map(|(_, v)| *v)
            .ok_or_else(|| {
                OrchestratorError::InvalidInputOwned(format!(
                    "--colocalize gene '{name}' must also be listed in --genes"
                ))
            })
    };
    let va = find(&gene_a)?;
    let vb = find(&gene_b)?;

    let fa = kcore::Field::new(core_domain, va.to_vec())?;
    let fb = kcore::Field::new(core_domain, vb.to_vec())?;
    let gi = kcore::GridIndex::build(core_domain)?;
    let cfg = kcore::SpatialStatsConfig {
        neighbor_radius: args.neighbor_radius.max(1),
        ..kcore::SpatialStatsConfig::default()
    };
    let res = kcore::compute_bivariate_moran_inference(
        &fa,
        &fb,
        &gi,
        cfg,
        kcore::MoranInferenceConfig {
            permutations: args.stats_permutations,
            seed: args.stats_seed,
        },
    )?;

    let local_q_sparse = to_f32_vec(&res.local_q_value);
    let n_sig = local_q_sparse.iter().filter(|&&q| q < 0.05).count() as u32;

    // Effect-size-first interpretation. At HD scale (millions of bins) a
    // vanishingly small R still clears p<0.05, so |R| — not the p-value — is the
    // headline (see the "everything is significant" caveat for large N).
    let r = res.global_r;
    let mag = r.abs();
    let strength = if mag < 0.05 {
        "negligible"
    } else if mag < 0.15 {
        "weak"
    } else if mag < 0.30 {
        "moderate"
    } else {
        "strong"
    };
    let sig = res.global_p_value < 0.05;
    let interpretation = if mag < 0.05 {
        format!(
            "NEGLIGIBLE spatial association (R={r:.4}, |R|<0.05): '{gene_a}' and '{gene_b}' are \
             effectively spatially independent at this bin resolution. With {} valid bins even \
             trivial effects reach p<0.05 (p={:.3}) — judge by effect size, not the p-value.",
            res.n_valid, res.global_p_value
        )
    } else if sig {
        let dir = if r > 0.0 {
            "CO-CLUSTERED (spatially associated)"
        } else {
            "SEGREGATED / mutually excluded"
        };
        format!(
            "{strength} and significant: '{gene_a}' and '{gene_b}' are spatially {dir} \
             (R={r:.4}, p={:.3}, locally significant bins={n_sig}).",
            res.global_p_value
        )
    } else {
        format!(
            "no significant spatial association between '{gene_a}' and '{gene_b}' \
             (R={r:.4}, p={:.3}).",
            res.global_p_value
        )
    };

    Ok(ColocalizeArtifacts {
        summary: ColocalizeSummary {
            gene_a,
            gene_b,
            global_bivariate_moran_r: res.global_r,
            global_p_value: res.global_p_value,
            global_z_score: res.global_z_score,
            n_valid: res.n_valid,
            weight_sum: res.weight_sum,
            permutations: args.stats_permutations,
            seed: args.stats_seed,
            neighbor_mode: "rook",
            neighbor_radius: args.neighbor_radius.max(1),
            n_local_significant_q05: n_sig,
            interpretation,
        },
        nx: grid.nx,
        ny: grid.ny,
        local_z: grid.scatter_to_dense(&to_f32_vec(&res.local_z_score))?,
        local_q: grid.scatter_to_dense(&local_q_sparse)?,
    })
}

fn to_f32_vec(field: &kcore::Field<'_>) -> Vec<f32> {
    field.values().iter().map(|&v| v as f32).collect()
}

fn compute_spatial_stats(
    core_gene: &kcore::Field<'_>,
    core_domain: &kcore::SpatialDomain,
    grid: &GridMapping,
    args: &RunArgs,
) -> Result<SpatialStatsArtifacts, OrchestratorError> {
    let gi = kcore::GridIndex::build(core_domain)?;
    let cfg = kcore::SpatialStatsConfig::default();
    let inf = kcore::compute_moran_inference(
        core_gene,
        &gi,
        cfg,
        kcore::MoranInferenceConfig {
            permutations: args.stats_permutations,
            seed: args.stats_seed,
        },
    )?;
    let getis = kcore::compute_getis_ord_gstar_local(core_gene, &gi, cfg)?;

    // Count locally significant bins (BH-q < 0.05) across both statistics.
    let moran_q = to_f32_vec(&inf.local_q_value);
    let getis_q_sparse = to_f32_vec(&getis.q_value);
    let mut n_sig = 0_u32;
    for (&mq, &gq) in moran_q.iter().zip(getis_q_sparse.iter()) {
        if mq < 0.05 || gq < 0.05 {
            n_sig += 1;
        }
    }

    let global = SpatialStatsGlobal {
        global_moran_i: inf.stats.global_moran_i,
        global_moran_p_value: inf.global_p_value,
        global_moran_z_score: inf.global_z_score,
        global_geary_c: inf.stats.global_geary_c,
        n_valid: inf.stats.n_valid,
        weight_sum: inf.stats.weight_sum,
        permutations: args.stats_permutations,
        seed: args.stats_seed,
        neighbor_mode: "rook",
        n_local_significant_q05: n_sig,
    };

    Ok(SpatialStatsArtifacts {
        global,
        nx: grid.nx,
        ny: grid.ny,
        local_moran_z: grid.scatter_to_dense(&to_f32_vec(&inf.local_z_score))?,
        local_moran_q: grid.scatter_to_dense(&moran_q)?,
        getis_z: grid.scatter_to_dense(&to_f32_vec(&getis.z_score))?,
        getis_q: grid.scatter_to_dense(&getis_q_sparse)?,
    })
}

fn compute_domains(
    core_gene: &kcore::Field<'_>,
    core_domain: &kcore::SpatialDomain,
    grid: &GridMapping,
    args: &RunArgs,
) -> Result<DomainArtifacts, OrchestratorError> {
    let gi = kcore::GridIndex::build(core_domain)?;
    // Queen (8-neighbour) spatial adjacency, symmetric, tissue bins only.
    let graph = kcore::build_grid_adjacency(
        &gi,
        core_domain,
        kcore::GridAdjacency::Eight,
        false,
        true,
        None,
    )?;
    let values: Vec<f64> = core_gene.values().iter().map(|&v| v as f64).collect();
    let res = kcore::detect_domains(
        &graph,
        Some(&values),
        kcore::DomainConfig {
            resolution: args.domain_resolution as f64,
            ..kcore::DomainConfig::default()
        },
    )?;

    // Per-community sizes and a +1-shifted label field (0 = background).
    let mut sizes = vec![0_u32; res.n_communities as usize];
    let mut label_plus1_sparse = vec![0.0_f32; res.labels.len()];
    for (i, &lbl) in res.labels.iter().enumerate() {
        if lbl != kcore::BACKGROUND_LABEL {
            sizes[lbl as usize] += 1;
            label_plus1_sparse[i] = (lbl + 1) as f32;
        }
    }

    Ok(DomainArtifacts {
        summary: DomainSummary {
            n_communities: res.n_communities,
            modularity: res.modularity,
            passes: res.passes,
            resolution: args.domain_resolution as f64,
            sizes,
        },
        nx: grid.nx,
        ny: grid.ny,
        labels_plus1: grid.scatter_to_dense(&label_plus1_sparse)?,
    })
}
