use approx::*;
use bempp::assembly::boundary::BoundaryAssembler;
use bempp::assembly::{boundary, fmm_tools};
use bempp::function::SerialFunctionSpace;
use ndgrid::{shapes::regular_sphere, traits::Grid};
use bempp::traits::FunctionSpace;
use green_kernels::laplace_3d::Laplace3dKernel;
use green_kernels::{traits::Kernel, types::EvalType};
#[cfg(not(debug_assertions))]
use kifmm::traits::tree::Tree;
use kifmm::traits::{fmm::Fmm, tree::FmmTree};
use kifmm::FftFieldTranslation;
use kifmm::SingleNodeBuilder;
use ndelement::ciarlet::LagrangeElementFamily;
use ndelement::types::{Continuity, ReferenceCellType};
use rand::prelude::*;
use rlst::{
    empty_array, rlst_dynamic_array2, MultIntoResize, RandomAccessByRef, RandomAccessMut,
    RawAccess, RawAccessMut,
};

fn fmm_prototype<TestGrid: Grid<T = f64, EntityDescriptor = ReferenceCellType> + Sync, TrialGrid: Grid<T = f64, EntityDescriptor = ReferenceCellType> + Sync>(
    trial_space: &SerialFunctionSpace<f64, TrialGrid>,
    test_space: &SerialFunctionSpace<f64, TestGrid>,
) {
    let npts = 37;

    let test_grid = test_space.grid();
    let trial_grid = test_space.grid();
    if std::ptr::addr_of!(*test_grid) as usize != std::ptr::addr_of!(*trial_grid) as usize {
        panic!("Assembly on different grids not yet supported");
    }

    let grid = trial_space.grid();

    let test_ndofs = test_space.global_size();
    let trial_ndofs = trial_space.global_size();
    let nqpts = npts * grid.entity_types(2).iter().map(|&i| grid.entity_count(i)).sum::<usize>();
    let kernel = Laplace3dKernel::new();

    // Compute dense
    let mut matrix = rlst_dynamic_array2!(f64, [test_ndofs, trial_ndofs]);
    let mut a = boundary::LaplaceSingleLayerAssembler::<f64>::default();
    a.quadrature_degree(ReferenceCellType::Triangle, npts);
    a.assemble_into_dense(&mut matrix, trial_space, test_space);

    // Compute using FMM method
    let all_points = fmm_tools::get_all_quadrature_points::<f64, TrialGrid>(npts, grid);

    // k is the matrix that FMM will give us
    let mut k = rlst_dynamic_array2!(f64, [nqpts, nqpts]);
    kernel.assemble_st(
        EvalType::Value,
        all_points.data(),
        all_points.data(),
        k.data_mut(),
    );

    let mut p_t = rlst_dynamic_array2!(f64, [test_ndofs, nqpts]);
    fmm_tools::transpose_basis_to_quadrature_into_dense::<128, f64, TestGrid>(
        &mut p_t, npts, test_space,
    );

    let mut p = rlst_dynamic_array2!(f64, [nqpts, trial_ndofs]);
    fmm_tools::basis_to_quadrature_into_dense::<128, f64, TrialGrid>(&mut p, npts, trial_space);

    // matrix 2 = p_t @ k @ p - c + singular
    let mut matrix2 = rlst_dynamic_array2!(f64, [test_ndofs, trial_ndofs]);

    // matrix 2 = singular
    a.assemble_singular_into_dense(&mut matrix2, trial_space, test_space);

    let mut correction = rlst_dynamic_array2!(f64, [test_ndofs, trial_ndofs]);
    a.assemble_singular_correction_into_dense(&mut correction, trial_space, test_space);

    let temp = empty_array::<f64, 2>()
        .simple_mult_into_resize(empty_array::<f64, 2>().simple_mult_into_resize(p_t, k), p);
    for j in 0..trial_ndofs {
        for i in 0..test_ndofs {
            *matrix2.get_mut([i, j]).unwrap() +=
                *temp.get([i, j]).unwrap() - *correction.get([i, j]).unwrap();
        }
    }

    // Check two matrices are equal
    for i in 0..test_ndofs {
        for j in 0..trial_ndofs {
            assert_relative_eq!(
                *matrix.get([i, j]).unwrap(),
                *matrix2.get([i, j]).unwrap(),
                epsilon = 1e-8
            );
        }
    }
}

#[cfg(not(debug_assertions))]
fn fmm_matvec<TrialGrid: Grid<T = f64, EntityDescriptor = ReferenceCellType> + Sync, TestGrid: Grid<T = f64, EntityDescriptor = ReferenceCellType> + Sync>(
    trial_space: &SerialFunctionSpace<f64, TrialGrid>,
    test_space: &SerialFunctionSpace<f64, TestGrid>,
) {
    let npts = 37;

    let test_grid = test_space.grid();
    let trial_grid = test_space.grid();
    if std::ptr::addr_of!(*test_grid) as usize != std::ptr::addr_of!(*trial_grid) as usize {
        panic!("Assembly on different grids not yet supported");
    }

    let grid = trial_space.grid();

    let test_ndofs = test_space.global_size();
    let trial_ndofs = trial_space.global_size();
    let nqpts = npts * grid.entity_types(2).iter().map(|&i| grid.entity_count(i)).sum::<usize>();

    // Compute dense
    let mut matrix = rlst_dynamic_array2!(f64, [test_ndofs, trial_ndofs]);
    let mut a = boundary::LaplaceSingleLayerAssembler::<f64>::default();
    a.quadrature_degree(ReferenceCellType::Triangle, npts);
    a.assemble_into_dense(&mut matrix, trial_space, test_space);

    // Compute using FMM method
    let all_points = fmm_tools::get_all_quadrature_points::<f64, TrialGrid>(npts, grid);

    // FMM parameters
    let expansion_order = 6;
    let n_crit = Some(150);
    let sparse = true;

    let p_t =
        fmm_tools::transpose_basis_to_quadrature_into_csr::<128, f64, TestGrid>(npts, test_space);
    let p = fmm_tools::basis_to_quadrature_into_csr::<128, f64, TrialGrid>(npts, trial_space);
    let singular = a.assemble_singular_into_csr(trial_space, test_space);

    let correction = a.assemble_singular_correction_into_csr(trial_space, test_space);

    // matrix2 = p_t @ k @ p - c + singular
    let mut rng = rand::thread_rng();
    for _ in 0..10 {
        let mut vec = rlst_dynamic_array2!(f64, [trial_ndofs, 1]);
        for i in 0..trial_ndofs {
            *vec.get_mut([i, 0]).unwrap() = rng.gen();
        }
        let dense_result =
            empty_array::<f64, 2>().simple_mult_into_resize(matrix.view(), vec.view());

        let mut fmm_result = rlst_dynamic_array2!(f64, [test_ndofs, 1]);
        // (p_t @ k @ p - c + singular) @ vec
        let mut row = 0;
        for (i, (index, data)) in singular.indices().iter().zip(singular.data()).enumerate() {
            while i >= singular.indptr()[row + 1] {
                row += 1;
            }
            *fmm_result.get_mut([row, 0]).unwrap() += data * vec.get([*index, 0]).unwrap();
        }
        let mut row = 0;
        for (i, (index, data)) in correction
            .indices()
            .iter()
            .zip(correction.data())
            .enumerate()
        {
            while i >= correction.indptr()[row + 1] {
                row += 1;
            }
            *fmm_result.get_mut([row, 0]).unwrap() -= data * vec.get([*index, 0]).unwrap();
        }

        let mut temp0 = rlst_dynamic_array2!(f64, [nqpts, 1]);
        let mut row = 0;
        for (i, (index, data)) in p.indices().iter().zip(p.data()).enumerate() {
            while i >= p.indptr()[row + 1] {
                row += 1;
            }
            *temp0.get_mut([row, 0]).unwrap() += data * vec.get([*index, 0]).unwrap();
        }

        let fmm = SingleNodeBuilder::new()
            .tree(&all_points, &all_points, n_crit, sparse)
            .unwrap()
            .parameters(
                &temp0,
                expansion_order,
                Laplace3dKernel::new(),
                EvalType::Value,
                FftFieldTranslation::new(),
            )
            .unwrap()
            .build()
            .unwrap();

        let _ = fmm.evaluate(false);

        let mut temp1 = rlst_dynamic_array2!(f64, [nqpts, 1]);
        let indices = &fmm.tree().target_tree().all_global_indices().unwrap();

        for (i, j) in indices.iter().enumerate() {
            *temp1.get_mut([*j, 0]).unwrap() = fmm.potentials[i];
        }

        let mut row = 0;
        for (i, (index, data)) in p_t.indices().iter().zip(p_t.data()).enumerate() {
            while i >= p_t.indptr()[row + 1] {
                row += 1;
            }
            *fmm_result.get_mut([row, 0]).unwrap() += data * temp1.get([*index, 0]).unwrap();
        }

        for i in 0..test_ndofs {
            assert_relative_eq!(
                *dense_result.get([i, 0]).unwrap(),
                *fmm_result.get([i, 0]).unwrap(),
                epsilon = 1e-5
            );
        }
    }
}

#[test]
fn test_fmm_prototype_dp0_dp0() {
    #[cfg(debug_assertions)]
    let grid = regular_sphere(0);
    #[cfg(not(debug_assertions))]
    let grid = regular_sphere(2);

    let element = LagrangeElementFamily::<f64>::new(0, Continuity::Discontinuous);
    let space = SerialFunctionSpace::new(&grid, &element);

    fmm_prototype(&space, &space);
}

#[test]
fn test_fmm_prototype_p1_p1() {
    #[cfg(debug_assertions)]
    let grid = regular_sphere(0);
    #[cfg(not(debug_assertions))]
    let grid = regular_sphere(2);

    let element = LagrangeElementFamily::<f64>::new(1, Continuity::Standard);
    let space = SerialFunctionSpace::new(&grid, &element);

    fmm_prototype(&space, &space);
}

#[cfg(not(debug_assertions))]
#[test]
fn test_fmm_prototype_dp0_p1() {
    let grid = regular_sphere(2);

    let element0 = LagrangeElementFamily::<f64>::new(0, Continuity::Discontinuous);
    let element1 = LagrangeElementFamily::<f64>::new(1, Continuity::Standard);
    let space0 = SerialFunctionSpace::new(&grid, &element0);
    let space1 = SerialFunctionSpace::new(&grid, &element1);

    fmm_prototype(&space0, &space1);
}

#[cfg(not(debug_assertions))]
#[test]
fn test_fmm_dp0_dp0() {
    let grid = regular_sphere(2);

    let element = LagrangeElementFamily::<f64>::new(0, Continuity::Discontinuous);
    let space = SerialFunctionSpace::new(&grid, &element);

    fmm_matvec(&space, &space);
}

#[cfg(not(debug_assertions))]
#[test]
fn test_fmm_p1_p1() {
    let grid = regular_sphere(2);

    let element = LagrangeElementFamily::<f64>::new(1, Continuity::Standard);
    let space = SerialFunctionSpace::new(&grid, &element);

    fmm_matvec(&space, &space);
}

#[cfg(not(debug_assertions))]
#[test]
fn test_fmm_dp0_p1() {
    let grid = regular_sphere(2);

    let element0 = LagrangeElementFamily::<f64>::new(0, Continuity::Discontinuous);
    let element1 = LagrangeElementFamily::<f64>::new(1, Continuity::Standard);
    let space0 = SerialFunctionSpace::new(&grid, &element0);
    let space1 = SerialFunctionSpace::new(&grid, &element1);

    fmm_matvec(&space0, &space1);
}

#[test]
fn test_fmm_result() {
    let grid = regular_sphere(2);

    let npts = 1;

    let nqpts = npts * grid.entity_types(2).iter().map(|&i| grid.entity_count(i)).sum::<usize>();
    let kernel = Laplace3dKernel::new();

    let all_points = fmm_tools::get_all_quadrature_points::<f64, _>(npts, &grid);

    let expansion_order = 6;
    let n_crit = Some(1);
    let sparse = true;

    let mut k = rlst_dynamic_array2!(f64, [nqpts, nqpts]);
    kernel.assemble_st(
        EvalType::Value,
        all_points.data(),
        all_points.data(),
        k.data_mut(),
    );

    // let mut rng: ThreadRng = rand::thread_rng();
    let mut rng = StdRng::seed_from_u64(0);

    let mut vec = rlst_dynamic_array2!(f64, [nqpts, 1]);
    for i in 0..nqpts {
        *vec.get_mut([i, 0]).unwrap() = rng.gen();
    }
    let dense_result = empty_array::<f64, 2>().simple_mult_into_resize(k.view(), vec.view());

    let fmm = SingleNodeBuilder::new()
        .tree(&all_points, &all_points, n_crit, sparse)
        .unwrap()
        .parameters(
            &vec,
            expansion_order,
            Laplace3dKernel::new(),
            EvalType::Value,
            FftFieldTranslation::new(),
        )
        .unwrap()
        .build()
        .unwrap();

    let _ = fmm.evaluate(false);

    let indices = &fmm.tree().target_tree().global_indices;
    let mut fmm_result = rlst_dynamic_array2!(f64, [nqpts, 1]);
    for (i, j) in indices.iter().enumerate() {
        *fmm_result.get_mut([*j, 0]).unwrap() = fmm.potentials[i];
    }

    for i in 0..nqpts {
        assert_relative_eq!(
            *dense_result.get([i, 0]).unwrap(),
            *fmm_result.get([i, 0]).unwrap(),
            epsilon = 1e-5
        );
    }
}