use std::collections::HashMap;

use kira_spatial_core as kcore;
use kira_spatial_field::gene_field;
use kira_spatial_io::Dataset;
use rayon::prelude::*;

use crate::error::OrchestratorError;

fn sanitize(v: f32) -> f32 {
    if v.is_finite() {
        if v == 0.0 { 0.0 } else { v }
    } else {
        0.0
    }
}

/// Deterministic domain id for the kira-spatial-field contract; stable across
/// releases so external `Field::from_gene` consumers can share caches.
pub(crate) fn domain_id(len: usize) -> u64 {
    0x4B49_5241_0000_0000_u64.wrapping_add(len as u64)
}

/// Exposes a `kira_spatial_io::Dataset` as an `ExpressionCsrView`, narrowing
/// `indptr` from `u64` to `u32` to match the field-crate contract.
pub(crate) struct FieldCsrAdapter {
    indptr: Vec<u32>,
    indices: Vec<u32>,
    data: Vec<f32>,
    gene_index: HashMap<String, usize>,
}

impl FieldCsrAdapter {
    #[allow(dead_code)]
    pub(crate) fn from_dataset(dataset: &Dataset) -> Result<Self, OrchestratorError> {
        let csr = dataset.expression_csr();
        let mut indptr = Vec::with_capacity(csr.indptr.len());
        for i in 0..csr.indptr.len() {
            let p = csr.indptr.get(i);
            indptr.push(
                u32::try_from(p)
                    .map_err(|_| OrchestratorError::InvalidInput("csr indptr does not fit u32"))?,
            );
        }
        let indices = csr.indices.clone();
        let data = csr.data.clone();

        let mut gene_index = HashMap::with_capacity(dataset.features().rows.len());
        for row in &dataset.features().rows {
            gene_index.insert(row.gene_name.clone(), row.gene_id as usize);
        }

        Ok(Self {
            indptr,
            indices,
            data,
            gene_index,
        })
    }
}

impl gene_field::ExpressionCsrView for FieldCsrAdapter {
    fn indptr(&self) -> &[u32] {
        &self.indptr
    }

    fn indices(&self) -> &[u32] {
        &self.indices
    }

    fn data(&self) -> &[f32] {
        &self.data
    }

    fn gene_index(&self) -> &HashMap<String, usize> {
        &self.gene_index
    }
}

/// Adapter exposing a flat bin count as a `kira_spatial_field::gene_field::SpatialDomainView`.
#[allow(dead_code)]
pub(crate) struct FieldDomainAdapter {
    pub(crate) id: u64,
    pub(crate) bins: usize,
}

impl gene_field::SpatialDomainView for FieldDomainAdapter {
    fn id(&self) -> u64 {
        self.id
    }

    fn bin_count(&self) -> usize {
        self.bins
    }
}

pub(crate) struct GridMapping {
    rows: Vec<u32>,
    cols: Vec<u32>,
    min_row: u32,
    min_col: u32,
    pub(crate) nx: usize,
    pub(crate) ny: usize,
    pub(crate) origin_x: f32,
    pub(crate) origin_y: f32,
}

impl GridMapping {
    pub(crate) fn from_io_domain(
        domain: &kira_spatial_io::SpatialDomain,
    ) -> Result<Self, OrchestratorError> {
        let rows = domain
            .grid_row
            .as_ref()
            .ok_or(OrchestratorError::InvalidInput(
                "grid_row is required for regular-grid projection",
            ))?;
        let cols = domain
            .grid_col
            .as_ref()
            .ok_or(OrchestratorError::InvalidInput(
                "grid_col is required for regular-grid projection",
            ))?;
        if rows.len() != cols.len() || rows.len() != domain.len() {
            return Err(OrchestratorError::InvalidInput(
                "grid_row/grid_col lengths are inconsistent",
            ));
        }

        let mut min_row = u32::MAX;
        let mut max_row = 0_u32;
        let mut min_col = u32::MAX;
        let mut max_col = 0_u32;
        for i in 0..rows.len() {
            min_row = min_row.min(rows[i]);
            max_row = max_row.max(rows[i]);
            min_col = min_col.min(cols[i]);
            max_col = max_col.max(cols[i]);
        }

        let nx_u32 = max_col
            .checked_sub(min_col)
            .and_then(|d| d.checked_add(1))
            .ok_or(OrchestratorError::InvalidInput("invalid grid_col range"))?;
        let ny_u32 = max_row
            .checked_sub(min_row)
            .and_then(|d| d.checked_add(1))
            .ok_or(OrchestratorError::InvalidInput("invalid grid_row range"))?;

        Ok(Self {
            rows: rows.clone(),
            cols: cols.clone(),
            min_row,
            min_col,
            nx: nx_u32 as usize,
            ny: ny_u32 as usize,
            origin_x: min_col as f32,
            origin_y: min_row as f32,
        })
    }
    pub(crate) fn to_core_domain(
        &self,
        io_domain: &kira_spatial_io::SpatialDomain,
    ) -> Result<kcore::SpatialDomain, OrchestratorError> {
        let width = u32::try_from(self.nx)
            .map_err(|_| OrchestratorError::InvalidInput("nx does not fit u32"))?;
        let height = u32::try_from(self.ny)
            .map_err(|_| OrchestratorError::InvalidInput("ny does not fit u32"))?;

        let mut x_q = Vec::with_capacity(self.cols.len());
        let mut y_q = Vec::with_capacity(self.rows.len());
        for i in 0..self.rows.len() {
            x_q.push(
                i32::try_from(self.cols[i])
                    .map_err(|_| OrchestratorError::InvalidInput("grid_col does not fit i32"))?,
            );
            y_q.push(
                i32::try_from(self.rows[i])
                    .map_err(|_| OrchestratorError::InvalidInput("grid_row does not fit i32"))?,
            );
        }

        let mut mask = Vec::with_capacity(io_domain.len());
        for i in 0..io_domain.len() {
            mask.push(u8::from(io_domain.tissue_mask[i]));
        }

        let grid = kcore::GridSpec::new(
            i32::try_from(self.min_col)
                .map_err(|_| OrchestratorError::InvalidInput("min_col does not fit i32"))?,
            i32::try_from(self.min_row)
                .map_err(|_| OrchestratorError::InvalidInput("min_row does not fit i32"))?,
            1,
            width,
            height,
        )?;

        Ok(kcore::SpatialDomain::new_quantized(
            x_q,
            y_q,
            Some(mask),
            Some(grid),
        )?)
    }

    /// Builds a core domain over the same geometry but with an explicit validity
    /// mask (`valid[i] == true` keeps bin `i`). Used to apply a per-bin QC mask
    /// (e.g. minimum library size) so spatial statistics ignore background bins.
    pub(crate) fn to_core_domain_with_mask(
        &self,
        valid: &[bool],
    ) -> Result<kcore::SpatialDomain, OrchestratorError> {
        if valid.len() != self.rows.len() {
            return Err(OrchestratorError::InvalidInput(
                "QC mask length mismatches spatial domain",
            ));
        }
        let width = u32::try_from(self.nx)
            .map_err(|_| OrchestratorError::InvalidInput("nx does not fit u32"))?;
        let height = u32::try_from(self.ny)
            .map_err(|_| OrchestratorError::InvalidInput("ny does not fit u32"))?;

        let mut x_q = Vec::with_capacity(self.cols.len());
        let mut y_q = Vec::with_capacity(self.rows.len());
        let mut mask = Vec::with_capacity(self.rows.len());
        for i in 0..self.rows.len() {
            x_q.push(
                i32::try_from(self.cols[i])
                    .map_err(|_| OrchestratorError::InvalidInput("grid_col does not fit i32"))?,
            );
            y_q.push(
                i32::try_from(self.rows[i])
                    .map_err(|_| OrchestratorError::InvalidInput("grid_row does not fit i32"))?,
            );
            mask.push(u8::from(valid[i]));
        }

        let grid = kcore::GridSpec::new(
            i32::try_from(self.min_col)
                .map_err(|_| OrchestratorError::InvalidInput("min_col does not fit i32"))?,
            i32::try_from(self.min_row)
                .map_err(|_| OrchestratorError::InvalidInput("min_row does not fit i32"))?,
            1,
            width,
            height,
        )?;

        Ok(kcore::SpatialDomain::new_quantized(
            x_q,
            y_q,
            Some(mask),
            Some(grid),
        )?)
    }

    /// Plans a coarsening of the grid by `factor` (NxN fine cells → one
    /// super-bin). Returns the coarse `GridMapping` (one entry per occupied
    /// coarse cell, in first-appearance order), an `assign` map from each fine
    /// bin to its coarse-bin index, and the coarse bin count. Deterministic.
    pub(crate) fn coarsen_plan(
        &self,
        factor: u32,
    ) -> Result<(GridMapping, Vec<u32>, usize), OrchestratorError> {
        if factor <= 1 {
            return Err(OrchestratorError::InvalidInput(
                "bin-aggregate factor must be >= 2",
            ));
        }
        let cnx = self.nx.div_ceil(factor as usize);
        let cny = self.ny.div_ceil(factor as usize);

        let mut cell_index: HashMap<(u32, u32), u32> = HashMap::new();
        let mut crows = Vec::new();
        let mut ccols = Vec::new();
        let mut assign = vec![0_u32; self.rows.len()];
        for i in 0..self.rows.len() {
            let ccol = (self.cols[i] - self.min_col) / factor;
            let crow = (self.rows[i] - self.min_row) / factor;
            let idx = *cell_index.entry((crow, ccol)).or_insert_with(|| {
                let id = crows.len() as u32;
                crows.push(crow);
                ccols.push(ccol);
                id
            });
            assign[i] = idx;
        }
        let n_coarse = crows.len();
        let coarse = GridMapping {
            rows: crows,
            cols: ccols,
            min_row: 0,
            min_col: 0,
            nx: cnx,
            ny: cny,
            origin_x: 0.0,
            origin_y: 0.0,
        };
        Ok((coarse, assign, n_coarse))
    }

    pub(crate) fn scatter_to_dense(&self, sparse: &[f32]) -> Result<Vec<f32>, OrchestratorError> {
        if sparse.len() != self.rows.len() {
            return Err(OrchestratorError::InvalidInput(
                "sparse field length mismatches spatial domain",
            ));
        }
        let mut mapped: Vec<(usize, f32)> = (0..sparse.len())
            .into_par_iter()
            .map(|i| {
                let gx = (self.cols[i] - self.min_col) as usize;
                let gy = (self.rows[i] - self.min_row) as usize;
                let idx = gy * self.nx + gx;
                (idx, sanitize(sparse[i]))
            })
            .collect();

        mapped.par_sort_unstable_by_key(|(idx, _)| *idx);
        let mut dense = vec![0.0_f32; self.nx * self.ny];
        let mut prev = None::<usize>;
        for (idx, value) in mapped {
            if prev == Some(idx) {
                return Err(OrchestratorError::InvalidInput(
                    "multiple bins map to the same grid cell",
                ));
            }
            dense[idx] = value;
            prev = Some(idx);
        }

        Ok(dense)
    }
}
