mod ast;
mod grammar;

#[allow(dead_code)]
mod selection_expr;
pub use selection_expr::*;

mod selection_error;
pub use selection_error::*;

mod selection_def;
pub use selection_def::*;

mod utils;

mod system;
pub use system::*;

mod sel;
pub use sel::*;

mod traits;
pub use traits::*;

mod par_split;
pub use par_split::*;


//############################################################
//#  Tests
//############################################################

#[cfg(test)]
mod tests {
    use crate::prelude::*;
    use rayon::iter::IndexedParallelIterator;
    pub use rayon::iter::ParallelIterator;
    use std::sync::LazyLock;

    static TOPST: LazyLock<(Topology, State)> = LazyLock::new(|| {
        let mut h = FileHandler::open("tests/protein.pdb").unwrap();
        h.read().unwrap()
    });

    #[test]
    fn builder_overlap() -> anyhow::Result<()> {
        let top = TOPST.0.clone();
        let st = TOPST.1.clone();
        let b = System::new(top, st)?;
        // Create two overlapping selections
        let _sel1 = b.select(0..10)?;
        let _sel2 = b.select(5..15)?;
        Ok(())
    }

    #[test]
    fn builder_par_no_overlap() {
        let top = TOPST.0.clone();
        let st = TOPST.1.clone();
        let b = System::new(top, st).unwrap();
        // Create two non-overlapping selections.
        let _sel1 = b.select(0..10).unwrap();
        let _sel2 = b.select(11..15).unwrap();
    }

    fn make_sel_all() -> anyhow::Result<(System, Sel)> {
        let top = TOPST.0.clone();
        let st = TOPST.1.clone();
        let b = System::new(top, st)?;
        let sel = b.select_all();
        Ok((b, sel))
    }

    fn make_sel_prot() -> anyhow::Result<(System, Sel)> {
        let top = TOPST.0.clone();
        let st = TOPST.1.clone();
        let b = System::new(top, st)?;
        let sel = b.select("not resname TIP3 POT CLA")?;
        Ok((b, sel))
    }

    #[test]
    fn test_measure() -> anyhow::Result<()> {
        let (sys, sel) = make_sel_all()?;
        let sel = sys.select_bound(sel)?;
        println!("before {}", sel.iter_pos().next().unwrap());

        let (minv, maxv) = sel.min_max();
        println!("{minv}:{maxv}");

        //sel.translate(&Vector3f::new(10.0,10.0,10.0));
        println!("after {}", sel.iter_pos().next().unwrap());
        println!("{:?}", sel.min_max());
        Ok(())
    }

    #[test]
    fn test_measure_pbc() -> anyhow::Result<()> {
        let (sys, sel) = make_sel_all()?;
        let sel = sys.select_bound(sel)?;
        let cm = sel.center_of_mass()?;
        println!("{cm}");
        Ok(())
    }

    #[test]
    fn test_translate() -> anyhow::Result<()> {
        let (mut sys, ind) = make_sel_all()?;
        println!("before {}", sys.try_bind(&ind)?.iter_pos().next().unwrap());
        sys.try_bind_mut(&ind)?
            .translate(&Vector3f::new(10.0, 10.0, 10.0));
        println!("after {}", sys.try_bind(&ind)?.iter_pos().next().unwrap());
        Ok(())
    }

    #[test]
    fn test_write_to_file() -> anyhow::Result<()> {
        let sys = System::from_file("tests/protein.pdb")?;
        let sel = sys.select("name CA")?;
        sys.try_bind(&sel)?
            .save(concat!(env!("OUT_DIR"), "/f.pdb"))?;

        // let mut h = FileHandler::create(concat!(env!("OUT_DIR"), "/f.pdb"))?;
        // h.write(&sel)?;

        println!("test {:p}", &sel);
        println!("out {}", env!("OUT_DIR"));

        //h.write(&sel)?;
        Ok(())
    }

    #[test]
    fn test_unwrap_connectivity_1() -> anyhow::Result<()> {
        let (mut sys, ind) = make_sel_prot()?;
        sys.try_bind_mut(&ind)?
            .unwrap_connectivity_dim(0.2, PBC_FULL)?;

        let mut h = FileHandler::create(concat!(env!("OUT_DIR"), "/unwrapped.pdb"))?;
        let sel = sys.select_bound(ind)?;
        h.write(&sel)?;
        Ok(())
    }

    #[test]
    fn eigen_test() -> anyhow::Result<()> {
        let (mut sys, ind1) = make_sel_prot()?;
        let (_, ind2) = make_sel_prot()?;

        sys.try_bind_mut(&ind2)?
            .rotate(&Vector3f::x_axis(), 80.0_f32.to_radians());

        let sel1 = sys.try_bind(&ind1)?;
        let sel2 = sys.try_bind(&ind2)?;
        sel1.save(concat!(env!("OUT_DIR"), "/sel2.pdb"))?;
        sel2.save(concat!(env!("OUT_DIR"), "/sel1_before.pdb"))?;
        println!("Initial RMSD:{}", rmsd_mw(&sel1, &sel2)?);

        let m = fit_transform(&sel1, &sel2)?;
        println!("{m}");

        sys.try_bind_mut(&ind1)?.apply_transform(&m);

        let sel1 = sys.try_bind(&ind1)?;
        let sel2 = sys.try_bind(&ind2)?;
        sel1.save(concat!(env!("OUT_DIR"), "/sel1_after.pdb"))?;
        println!("Final RMSD:{}", rmsd_mw(&sel1, &sel2)?);
        Ok(())
    }

    #[test]
    fn sasa_test() -> anyhow::Result<()> {
        let (sys, ind) = make_sel_all()?;
        let sel = sys.select_bound(ind)?;
        let s = sel.sasa()?;
        assert!(s.total_area() > 0.0);
        assert_eq!(s.areas().len(), sel.len());
        println!("Sasa area: {}", s.total_area());

        let sv = sel.sasa_vol()?;
        assert!(sv.total_volume() > 0.0);
        assert_eq!(sv.areas().len(), sel.len());
        println!("Sasa area: {}, volume: {}", sv.total_area(), sv.total_volume());
        Ok(())
    }

    #[test]
    fn tets_gyration() -> anyhow::Result<()> {
        let (sys, ind) = make_sel_prot()?;
        let g = sys.try_bind(&ind)?.gyration_pbc()?;
        println!("Gyration radius: {g}");
        Ok(())
    }

    #[test]
    fn test_inertia() -> anyhow::Result<()> {
        let (sys, ind) = make_sel_all()?;
        //sel1.rotate(
        //    &UnitVector3::new_normalize(Vector3f::new(1.0,2.0,3.0)),
        //    0.45
        //);
        let (moments, axes) = sys.try_bind(&ind)?.inertia_pbc()?;
        println!("Inertia moments: {moments}");
        println!("Inertia axes: {axes:?}");
        Ok(())
        // {0.7308828830718994 0.3332606256008148 -0.5956068634986877}
        // {0.6804488301277161 -0.28815552592277527 0.6737624406814575}
        // {-0.05291106179356575 0.8977214694023132 0.4373747706413269}
    }

    #[test]
    fn test_principal_transform() -> anyhow::Result<()> {
        let (mut sys, ind) = make_sel_prot()?;
        let sel1 = sys.try_bind(&ind)?;
        let tr = sel1.principal_transform()?;
        println!("Transform: {tr}");

        let (_, axes) = sel1.inertia()?;
        println!("Axes before: {axes}");

        sys.try_bind_mut(&ind)?.apply_transform(&tr);

        let sel1 = sys.try_bind(&ind)?;
        let (_, axes) = sel1.inertia()?;
        println!("Axes after: {axes}");

        sel1.save(concat!(env!("OUT_DIR"), "/oriented.pdb"))?;

        Ok(())
    }

    #[test]
    fn test_builder_append_from_self() -> anyhow::Result<()> {
        let (top, st) = FileHandler::open("tests/protein.pdb")?.read()?;
        let n = top.len();
        let mut builder = System::new(top, st)?;

        let sel = builder.select("resid 550:560")?;
        let added = sel.len();
        builder.append_from_self(&sel)?;
        let all = builder.select_all();
        assert_eq!(all.len(), n + added);
        Ok(())
    }

    #[test]
    fn test_select_from_vec() -> anyhow::Result<()> {
        let (top, st) = FileHandler::open("tests/protein.pdb")?.read()?;
        let src = System::new(top, st)?;
        let last = src.len() - 1;
        let _sel = src.select([1usize, 2, last].as_slice())?;
        Ok(())
    }

    // #[test]
    // #[should_panic]
    // fn test_builder_remove_from_self() {
    //     let (top, st) = FileHandler::open("tests/protein.pdb")
    //         .unwrap()
    //         .read()
    //         .unwrap();
    //     let n = top.len();
    //     let builder = System::new(top, st).unwrap();
    //     let sel = builder.select("resid 550:560").unwrap();
    //     let removed = sel.len();
    //     builder.remove(&sel).unwrap(); // Should break here
    //     let all = builder.select_all();
    //     assert_eq!(all.len(), n - removed);
    // }

    // #[test]
    // #[should_panic]
    // fn test_builder_fail_invalid_selection() {
    //     let (top, st) = FileHandler::open("tests/protein.pdb")
    //         .unwrap()
    //         .read()
    //         .unwrap();
    //     let builder = System::new(top, st).unwrap();
    //     let sel = builder.select("resid 809").unwrap(); // last residue
    //     builder.remove(&sel).unwrap();
    //     // Trying to call method on invalid selection
    //     let _cm = sel.center_of_mass();
    // }

    // #[test]
    // fn as_parrallel_test() -> anyhow::Result<()> {
    //     let top = TOPST.0.clone();
    //     let st = TOPST.1.clone();
    //     let ser = System::new(top, st)?;
    //     let ser_sel = ser.select_all()?;

    //     let mut res = vec![];

    //     ser_sel.as_parallel(|par_sel| {
    //         let fragments = par_sel.into_iter_fragments_resindex();
    //         res = fragments.par_iter().map(|sel| {
    //             let cm = sel.center_of_mass().unwrap();
    //             sel.translate(&cm.coords);
    //             println!("thread: {}", rayon::current_thread_index().unwrap());
    //             sel.center_of_mass().unwrap()
    //         }).collect::<Vec<_>>();
    //     })?;

    //     println!("cm = {:?}",res);

    //     Ok(())
    // }

    #[test]
    fn as_parrallel_test() -> anyhow::Result<()> {
        let top = TOPST.0.clone();
        let st = TOPST.1.clone();
        let mut sys = System::new(top, st)?;

        let parts = sys.split_par(|p| Some(p.atom.resindex))?;
        let coms = sys
            .iter_par_split_mut(&parts)
            .map(|sel| sel.center_of_mass())
            .collect::<Result<Vec<_>, _>>()?;

        sys.iter_par_split_mut(&parts)
            .enumerate()
            .for_each(|(i, mut sel)| sel.translate(&(-coms[i].coords)));

        println!("{:?}", coms);

        Ok(())
    }
}
