use rayon::iter::ParallelIterator;
pub mod bspline;
pub mod error;
pub mod filter;
pub mod metric;
pub mod optimize;
pub mod par_indexed_iter;
pub mod register;
pub mod transform;

use ndarray::prelude::*;
use num::Complex;
use thiserror::Error;

use crate::par_indexed_iter::ParallelIndexedIterMut;
use crate::transform::transform_point;
use error::Error;

/// An example of generating julia fractals, for testing purposes.
/// parameters: 2x2 flattened rotation matrix + xy translation
/// center: center of rotation
pub fn julia_image(
    shape: &[usize; 2],
    parameters: &[f64; 6],
    center: &[f64; 2],
    c: &[f64; 2],
) -> Array2<u8> {
    let c = Complex::new(c[0], c[1]);
    let scaley = 3.0 / shape[0] as f64;
    let scalex = 3.0 / shape[1] as f64;
    let mut im = Array2::<u8>::zeros([shape[0], shape[1]]);
    im.par_indexed_iter_mut().for_each(|(i, x)| {
        let cv = transform_point(i.as_slice(), center, parameters);
        let mut z = Complex::new(cv[1] * scalex - 1.5, cv[0] * scaley - 1.5);
        let mut i = 0;
        while i < 255 && z.norm() <= 2.0 {
            z = z * z + c;
            i += 1;
        }
        *x = i;
    });

    im
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::transform::Transform;
    use tempfile::NamedTempFile;
    use tiffwrite::IJTiffFile;

    #[test]
    fn serialization() -> Result<(), Error> {
        let file = NamedTempFile::new()?;
        let t = Transform::<Ix2>::new_with_center(
            vec![1.2, 0.3, -0.4, 0.9, 10.2, -9.5],
            vec![59.5, 49.5],
            vec![120, 100],
        );
        t.to_file(file.path().to_path_buf())?;
        let s = Transform::from_file(file.path().to_path_buf())?;
        assert_eq!(s, t);
        Ok(())
    }

    #[test]
    fn transform_point() -> Result<(), Box<dyn std::error::Error>> {
        let transform = Transform::<Ix2>::new_with_center(
            vec![1.0, 0.0, 0.0, 1.0, -120.0, 10.0],
            vec![299.5, 399.5],
            vec![600, 800],
        );
        let p = transform.transform_point(&[300, 400]);
        assert_eq!(p, [180.0, 410.0]);
        Ok(())
    }

    #[test]
    fn interpbs_f64() -> Result<(), Box<dyn std::error::Error>> {
        let shape = [1200, 1600];
        let transform_j = [1.0, 0.0, 0.0, 1.0, 0.0, 0.0];
        let transform_k = [1.0, 0.0, 0.0, 1.0, 0.0, 0.0];
        let center: [f64; 2] = shape
            .iter()
            .map(|&s| (s - 1) as f64 / 2.0)
            .collect::<Vec<_>>()
            .try_into()
            .unwrap();
        let transform =
            Transform::<Ix2>::new(transform_j.to_vec(), shape.to_vec()).with_rotation(1.0);
        let transform_j = transform.parameters.clone().try_into().unwrap();
        let c = [-0.8, 0.156];
        let j = julia_image(&shape, &transform_j, &center, &c).mapv(|x| x as f64);
        let k = julia_image(&shape, &transform_k, &center, &c).mapv(|x| x as f64);
        let n = transform.interpolate::<3, _, _>(k.view())?;
        let sj = j.iter().sum::<f64>();
        // let sk = k.iter().sum::<f64>();
        let sn = n.iter().sum::<f64>();
        // println!("sj: {}, sk: {}, sn: {}", sj, sk, sn);

        // let mut tiff = tiffwrite::IJTiffFile::new("interpbs_f64_0.tif")?;
        // tiff.save(j.view(), 0, 0, 0)?;
        // tiff.save(k.view(), 1, 0, 0)?;
        // tiff.save(n.view(), 2, 0, 0)?;

        let s = (sj.ln() - sn.ln()).abs();
        let d = (j.ln() - n.ln())
            .powi(2)
            .iter()
            .filter(|i| i.is_finite())
            .sum::<f64>()
            .sqrt();
        // println!("s: {}, d: {}", s, d);
        assert!(s < 1e-2);
        assert!(2000.0 * d <= (shape[0] * shape[1]) as f64);
        Ok(())
    }

    #[test]
    fn interpbs_f64_par() -> Result<(), Box<dyn std::error::Error>> {
        let shape = [1200, 1600];
        let transform_j = [1.0, 0.0, 0.0, 1.0, 0.0, 0.0];
        let transform_k = [1.0, 0.0, 0.0, 1.0, 0.0, 0.0];
        let center: [f64; 2] = shape
            .iter()
            .map(|&s| (s - 1) as f64 / 2.0)
            .collect::<Vec<_>>()
            .try_into()
            .unwrap();
        let transform =
            Transform::<Ix2>::new(transform_j.to_vec(), shape.to_vec()).with_rotation(1.0);
        let transform_j = transform.parameters.clone().try_into().unwrap();
        let c = [-0.8, 0.156];
        let j = julia_image(&shape, &transform_j, &center, &c).mapv(|x| x as f64);
        let k = julia_image(&shape, &transform_k, &center, &c).mapv(|x| x as f64);
        let n = transform.interpolate_par::<3, _, _>(k.view())?;
        let sj = j.iter().sum::<f64>();
        let sn = n.iter().sum::<f64>();

        let s = (sj.ln() - sn.ln()).abs();
        let d = (j.ln() - n.ln())
            .powi(2)
            .iter()
            .filter(|i| i.is_finite())
            .sum::<f64>()
            .sqrt();
        assert!(s < 1e-2);
        assert!(2000.0 * d <= (shape[0] * shape[1]) as f64);
        Ok(())
    }

    #[test]
    fn interpolate_par_unsafe() -> Result<(), Box<dyn std::error::Error>> {
        let image_a = julia_image(
            &[8000, 6000],
            &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
            &[3999.5, 2999.5],
            &[-0.8, 0.156],
        );
        let transform = Transform::from_rotation(1.0, &[3999.5, 2999.5]);
        let image_b = transform.interpolate_par::<3, _, _>(image_a.view())?;
        assert_ne!(image_b[[4000, 3000]], 0.0);
        Ok(())
    }

    #[test]
    fn julia() -> Result<(), Box<dyn std::error::Error>> {
        let image_a = julia_image(
            &[60, 80],
            &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
            &[29.5, 39.5],
            &[-0.8, 0.156],
        );
        let mut t = IJTiffFile::new(std::env::home_dir().unwrap().join("tmp/julia.tif"))?;
        t.save(image_a.view(), 0, 0, 0)?;
        Ok(())
    }
}