use crate::fixed::FixedPoint;
use crate::index::BoundedSumIter;
use crate::matrix::DenseMatrix;
use crate::vector::MultiDimVec;

/// Sparse tensor product operator.
///
/// Implicitly stores a matrix that is the restriction of a tensor product
/// M_1 x M_2 x ... x M_d to a downward-closed multi-index set.
///
/// Uses the unidirectional principle: multiply one dimension at a time,
/// cycling indices between multiplications. O(d * N * n) where N is the
/// total number of grid points and n is the 1D matrix size, compared to
/// O(N^2) for a naive matrix-vector product.
///
/// All arithmetic is fixed-point.
pub struct SparseTPOperator {
    fp: FixedPoint,
    iter: BoundedSumIter,
    d: usize,
    l_inv: Option<Vec<DenseMatrix>>,
    u_inv: Option<Vec<DenseMatrix>>,
    lower: Option<Vec<DenseMatrix>>,
    upper: Option<Vec<DenseMatrix>>,
    matrices: Vec<DenseMatrix>,
}

impl SparseTPOperator {
    pub fn new(fp: FixedPoint, iter: BoundedSumIter, matrices: Vec<DenseMatrix>) -> Self {
        let d = matrices.len();
        Self {
            fp,
            iter,
            d,
            l_inv: None,
            u_inv: None,
            lower: None,
            upper: None,
            matrices,
        }
    }

    fn ensure_lu(&mut self) {
        if self.lower.is_some() {
            return;
        }
        let fp = self.fp;
        let mut ls = Vec::with_capacity(self.d);
        let mut us = Vec::with_capacity(self.d);
        for m in &self.matrices {
            let (l, u) = m.lu_decompose(
                fp.one(),
                |a, b| fp.mul(a, b),
                |a, b| fp.div(a, b),
            );
            ls.push(l);
            us.push(u);
        }
        self.lower = Some(ls);
        self.upper = Some(us);
    }

    fn ensure_solve(&mut self) {
        if self.l_inv.is_some() {
            return;
        }
        self.ensure_lu();
        let fp = self.fp;
        let lower = self.lower.as_ref().unwrap();
        let upper = self.upper.as_ref().unwrap();

        let l_inv: Vec<DenseMatrix> = lower
            .iter()
            .map(|l| l.invert_unit_lower(fp.one(), |a, b| fp.mul(a, b)))
            .collect();

        let u_inv: Vec<DenseMatrix> = upper
            .iter()
            .map(|u| {
                u.invert_upper(
                    fp.one(),
                    |a, b| fp.mul(a, b),
                    |a, b| fp.div(a, b),
                )
            })
            .collect();

        self.l_inv = Some(l_inv);
        self.u_inv = Some(u_inv);
    }

    /// Matrix-vector product: y = A * x where A is the sparse tensor product.
    pub fn apply(&mut self, input: MultiDimVec) -> MultiDimVec {
        self.ensure_lu();
        let mut v = input;
        let mut buffer = MultiDimVec::new(self.iter.clone());

        let upper = self.upper.as_ref().unwrap();
        for k in (0..self.d).rev() {
            multiply_single_upper_triangular(&upper[k], &mut v, &mut buffer, &self.fp);
        }

        let lower = self.lower.as_ref().unwrap();
        for k in (0..self.d).rev() {
            multiply_single_lower_triangular(&lower[k], &mut v, &mut buffer, &self.fp);
        }

        v
    }

    /// Solve A * x = rhs for x, where A is the sparse tensor product.
    pub fn solve(&mut self, rhs: MultiDimVec) -> MultiDimVec {
        self.ensure_solve();
        let mut v = rhs;
        let mut buffer = MultiDimVec::new(self.iter.clone());

        let l_inv = self.l_inv.as_ref().unwrap();
        for k in (0..self.d).rev() {
            multiply_single_lower_triangular(&l_inv[k], &mut v, &mut buffer, &self.fp);
        }

        let u_inv = self.u_inv.as_ref().unwrap();
        for k in (0..self.d).rev() {
            multiply_single_upper_triangular(&u_inv[k], &mut v, &mut buffer, &self.fp);
        }

        v
    }
}

/// Below this inner-loop size, slice setup overhead exceeds the
/// savings from deferred-shift accumulation.
const SLICE_THRESHOLD: usize = 16;

/// Multiply v by \hat{I x ... x I x L} and cycle indices.
/// L is lower triangular. Result goes into buffer, then swap.
///
/// For count >= SLICE_THRESHOLD, uses deferred-shift i128 accumulation
/// with pre-sliced data to eliminate per-element shifts and bounds
/// checks. For small counts, uses the simpler per-element fp.mul path.
fn multiply_single_lower_triangular(
    l: &DenseMatrix,
    v: &mut MultiDimVec,
    buffer: &mut MultiDimVec,
    fp: &FixedPoint,
) {
    let it = v.iter().clone();
    let cycled = it.cycle();
    buffer.reset_with_iter(cycled);

    let mut indexes: Vec<usize> = vec![0; buffer.data.len()];
    let mut jump = it;
    jump.reset();

    let mut first_v_idx = 0usize;
    let mut second_v_idx = 0usize;

    let frac_bits = fp.frac_bits;
    let l_data = &l.data;
    let l_cols = l.cols;

    while jump.valid() {
        let count = jump.last_dim_count();

        if count >= SLICE_THRESHOLD {
            let src_slice = &v.data[first_v_idx][second_v_idx..second_v_idx + count];
            for i in 0..count {
                let row = &l_data[i * l_cols..i * l_cols + i + 1];
                let src = &src_slice[..i + 1];
                let n = row.len();
                let mut acc: i128 = 0;
                for j in 0..n {
                    acc += row[j] as i128 * src[j] as i128;
                }
                let bi = indexes[i];
                buffer.data[i][bi] = (acc >> frac_bits) as i64;
                indexes[i] = bi + 1;
            }
        } else {
            let src = &v.data[first_v_idx];
            for i in 0..count {
                let mut sum: i128 = 0;
                for j in 0..=i {
                    sum += fp.mul(l.get(i, j), src[second_v_idx + j]) as i128;
                }
                let bi = indexes[i];
                buffer.data[i][bi] = sum as i64;
                indexes[i] = bi + 1;
            }
        }

        second_v_idx += count;
        if second_v_idx >= v.data[first_v_idx].len() {
            second_v_idx = 0;
            first_v_idx += 1;
        }

        jump.next();
    }

    v.swap(buffer);
}

/// Multiply v by \hat{I x ... x I x U} and cycle indices.
/// U is upper triangular.
fn multiply_single_upper_triangular(
    u: &DenseMatrix,
    v: &mut MultiDimVec,
    buffer: &mut MultiDimVec,
    fp: &FixedPoint,
) {
    let it = v.iter().clone();
    let cycled = it.cycle();
    buffer.reset_with_iter(cycled);

    let mut indexes: Vec<usize> = vec![0; buffer.data.len()];
    let mut jump = it;
    jump.reset();

    let mut first_v_idx = 0usize;
    let mut second_v_idx = 0usize;

    let frac_bits = fp.frac_bits;
    let u_data = &u.data;
    let u_cols = u.cols;

    while jump.valid() {
        let count = jump.last_dim_count();

        if count >= SLICE_THRESHOLD {
            let src_slice = &v.data[first_v_idx][second_v_idx..second_v_idx + count];
            for i in 0..count {
                let row = &u_data[i * u_cols + i..i * u_cols + count];
                let sv = &src_slice[i..count];
                let n = row.len();
                let mut acc: i128 = 0;
                for j in 0..n {
                    acc += row[j] as i128 * sv[j] as i128;
                }
                let bi = indexes[i];
                buffer.data[i][bi] = (acc >> frac_bits) as i64;
                indexes[i] = bi + 1;
            }
        } else {
            let src = &v.data[first_v_idx];
            for i in 0..count {
                let mut sum: i128 = 0;
                for j in i..count {
                    sum += fp.mul(u.get(i, j), src[second_v_idx + j]) as i128;
                }
                let bi = indexes[i];
                buffer.data[i][bi] = sum as i64;
                indexes[i] = bi + 1;
            }
        }

        second_v_idx += count;
        if second_v_idx >= v.data[first_v_idx].len() {
            second_v_idx = 0;
            first_v_idx += 1;
        }

        jump.next();
    }

    v.swap(buffer);
}

/// Cycle indices without matrix multiplication (identity multiply).
pub fn cycle_identity(v: &mut MultiDimVec, buffer: &mut MultiDimVec) {
    let it = v.iter().clone();
    let cycled = it.cycle();
    buffer.reset_with_iter(cycled);

    let mut indexes: Vec<usize> = vec![0; buffer.data.len()];
    let mut jump = it;
    jump.reset();

    let mut first_v_idx = 0usize;
    let mut second_v_idx = 0usize;

    while jump.valid() {
        let count = jump.last_dim_count();
        let src = &v.data[first_v_idx];

        for i in 0..count {
            let bi = indexes[i];
            buffer.data[i][bi] = src[second_v_idx + i];
            indexes[i] = bi + 1;
        }

        second_v_idx += count;
        if second_v_idx >= v.data[first_v_idx].len() {
            second_v_idx = 0;
            first_v_idx += 1;
        }

        jump.next();
    }

    v.swap(buffer);
}
