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

/// Basis function trait. Given a dimension k and index j, evaluates the
/// j-th basis function at fixed-point coordinate x.
pub trait BasisFn {
    fn eval(&self, dim: usize, index: usize, x: i64, fp: &FixedPoint) -> i64;
}

/// Point distribution trait. Returns the i-th grid point in
/// dimension k, as a fixed-point value.
pub trait PointDist {
    fn point(&self, dim: usize, index: usize, fp: &FixedPoint) -> i64;
}

/// Monomial basis: phi_j(x) = x^j in fixed-point.
pub struct MonomialBasis;

impl BasisFn for MonomialBasis {
    fn eval(&self, _dim: usize, index: usize, x: i64, fp: &FixedPoint) -> i64 {
        let mut result = fp.one();
        for _ in 0..index {
            result = fp.mul(result, x);
        }
        result
    }
}

/// Golden-ratio point distribution: point(i) = frac((i+1) * phi).
pub struct GoldenPoints;

impl PointDist for GoldenPoints {
    fn point(&self, _dim: usize, index: usize, fp: &FixedPoint) -> i64 {
        let phi = 0.5 * (1.0 + 5.0_f64.sqrt());
        let val = (index + 1) as f64 * phi;
        let frac = val - val.floor();
        fp.from_f64(frac)
    }
}

/// Uniform points on [0, 1]: point(i) = i / (n-1) for n points.
pub struct UniformPoints {
    pub n: usize,
}

impl PointDist for UniformPoints {
    fn point(&self, _dim: usize, index: usize, fp: &FixedPoint) -> i64 {
        if self.n <= 1 {
            return 0;
        }
        fp.div(
            fp.from_f64(index as f64),
            fp.from_f64((self.n - 1) as f64),
        )
    }
}

/// Build the interpolation matrix for one dimension: M_k(i, j) = phi_k(j)(x_k(i)).
pub fn build_1d_matrix(
    n: usize,
    dim: usize,
    basis: &dyn BasisFn,
    points: &dyn PointDist,
    fp: &FixedPoint,
) -> DenseMatrix {
    let mut m = DenseMatrix::zeros(n, n);
    for i in 0..n {
        let x = points.point(dim, i, fp);
        for j in 0..n {
            m.set(i, j, basis.eval(dim, j, x, fp));
        }
    }
    m
}

/// Create a SparseTPOperator for interpolation with the given basis and points.
/// Requires d >= 2.
pub fn create_interpolation_operator(
    fp: FixedPoint,
    iter: &BoundedSumIter,
    basis: &dyn BasisFn,
    points: &dyn PointDist,
) -> SparseTPOperator {
    assert!(iter.dim() >= 2, "sparse grid requires d >= 2");
    let d = iter.dim();
    let bounds = iter.index_bounds();

    let matrices: Vec<DenseMatrix> = (0..d)
        .map(|k| build_1d_matrix(bounds[k], k, basis, points, &fp))
        .collect();

    SparseTPOperator::new(fp, iter.clone(), matrices)
}

/// Evaluate a function at all sparse grid points.
/// The function takes a slice of fixed-point coordinates and returns
/// a fixed-point value. Requires d >= 2.
pub fn evaluate_function(
    iter: &BoundedSumIter,
    f: impl Fn(&[i64]) -> i64,
    points: &dyn PointDist,
    fp: &FixedPoint,
) -> MultiDimVec {
    assert!(iter.dim() >= 2, "sparse grid requires d >= 2");
    let d = iter.dim();
    let mut v = MultiDimVec::new(iter.clone());
    let mut indexes: Vec<usize> = vec![0; v.data.len()];

    let mut jump = iter.clone();
    jump.reset();
    let mut point = vec![0i64; d];

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

        for dim in 0..(d - 1) {
            point[dim] = points.point(dim, jump.index_at(dim), fp);
        }

        for last_idx in 0..last_count {
            point[d - 1] = points.point(d - 1, last_idx, fp);
            let val = f(&point);
            let first = jump.first_index();
            let idx = indexes[first];
            v.data[first][idx] = val;
            indexes[first] = idx + 1;
        }

        jump.next();
    }

    v
}

/// Full interpolation pipeline: evaluate function, build operator, solve.
/// Requires d >= 2.
pub fn interpolate(
    fp: FixedPoint,
    iter: &BoundedSumIter,
    f: impl Fn(&[i64]) -> i64,
    basis: &dyn BasisFn,
    points: &dyn PointDist,
) -> MultiDimVec {
    let rhs = evaluate_function(iter, f, points, &fp);
    let mut op = create_interpolation_operator(fp, iter, basis, points);
    op.solve(rhs)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn monomial_eval() {
        let fp = FixedPoint::new(24);
        let basis = MonomialBasis;
        let x = fp.from_f64(0.5);
        assert_eq!(basis.eval(0, 0, x, &fp), fp.one());
        let v = basis.eval(0, 1, x, &fp);
        assert!((fp.to_f64(v) - 0.5).abs() < 0.001);
        let v = basis.eval(0, 2, x, &fp);
        assert!((fp.to_f64(v) - 0.25).abs() < 0.001);
    }

    #[test]
    fn golden_points_in_unit_interval() {
        let fp = FixedPoint::new(24);
        let pts = GoldenPoints;
        for i in 0..10 {
            let p = fp.to_f64(pts.point(0, i, &fp));
            assert!(p >= 0.0 && p < 1.0, "point {i} = {p} out of [0,1)");
        }
    }

    #[test]
    fn interpolate_2d_product() {
        let fp = FixedPoint::new(24);
        let iter = BoundedSumIter::new(2, 3);
        let basis = MonomialBasis;
        let pts = GoldenPoints;

        // f(x, y) = x * y -- product of coordinates
        let coeffs = interpolate(
            fp,
            &iter,
            |pt| fp.mul(pt[0], pt[1]),
            &basis,
            &pts,
        );

        // Reconstruct: apply the operator to coefficients and compare to rhs
        let rhs = evaluate_function(&iter, |pt| fp.mul(pt[0], pt[1]), &pts, &fp);
        let mut op = create_interpolation_operator(fp, &iter, &basis, &pts);
        let reconstructed = op.apply(coeffs);

        // Compute reconstruction error
        let mut err_sq: i128 = 0;
        let mut norm_sq: i128 = 0;
        for (r_row, rhs_row) in reconstructed.data.iter().zip(rhs.data.iter()) {
            for (&r, &rh) in r_row.iter().zip(rhs_row.iter()) {
                let diff = r - rh;
                err_sq += (diff as i128) * (diff as i128);
                norm_sq += (rh as i128) * (rh as i128);
            }
        }

        let rel_err = if norm_sq > 0 {
            (err_sq as f64 / norm_sq as f64).sqrt()
        } else {
            0.0
        };
        assert!(rel_err < 0.01, "relative reconstruction error = {rel_err}");
    }

    #[test]
    fn interpolate_3d_product() {
        let fp = FixedPoint::new(24);
        let iter = BoundedSumIter::new(3, 3);
        let basis = MonomialBasis;
        let pts = GoldenPoints;

        // f(x, y, z) = x * y * z
        let coeffs = interpolate(
            fp,
            &iter,
            |pt| fp.mul(fp.mul(pt[0], pt[1]), pt[2]),
            &basis,
            &pts,
        );

        let rhs = evaluate_function(
            &iter,
            |pt| fp.mul(fp.mul(pt[0], pt[1]), pt[2]),
            &pts,
            &fp,
        );
        let mut op = create_interpolation_operator(fp, &iter, &basis, &pts);
        let reconstructed = op.apply(coeffs);

        let mut err_sq: i128 = 0;
        let mut norm_sq: i128 = 0;
        for (r_row, rhs_row) in reconstructed.data.iter().zip(rhs.data.iter()) {
            for (&r, &rh) in r_row.iter().zip(rhs_row.iter()) {
                let diff = r - rh;
                err_sq += (diff as i128) * (diff as i128);
                norm_sq += (rh as i128) * (rh as i128);
            }
        }

        let rel_err = if norm_sq > 0 {
            (err_sq as f64 / norm_sq as f64).sqrt()
        } else {
            0.0
        };
        assert!(rel_err < 0.01, "3D relative reconstruction error = {rel_err}");
    }
}
