use alloc::vec::Vec;
use alloc::btree_map::BTreeMap;
use alloc::btree_set::BTreeSet;

use core::marker::PhantomData;

use mat4;
use number_traits::Float;

use specs::{Entities, Entity, Join, ReadStorage, System, WriteStorage};

use super::super::{Child, Init};
use super::{LocalTransform3D, Transform3D};


pub struct Transform3DSystem<T: 'static + Sync + Send + Copy + Float> {
    indices: BTreeMap<Entity, usize>,
    sorted: Vec<(Entity, Entity)>,
    new: Vec<Entity>,
    dead: BTreeSet<Entity>,
    dirty: BTreeSet<Entity>,
    swapped: BTreeSet<Entity>,
    _marker: PhantomData<T>,
}


impl<T: 'static + Sync + Send + Copy + Float> Transform3DSystem<T> {

    #[inline(always)]
    pub fn new() -> Self {
        Transform3DSystem {
            indices: BTreeMap::default(),
            sorted: Vec::new(),
            new: Vec::new(),
            dead: BTreeSet::default(),
            dirty: BTreeSet::default(),
            swapped: BTreeSet::default(),
            _marker: PhantomData,
        }
    }
}


impl<'a, T: 'static + Sync + Send + Copy + Float> System<'a> for Transform3DSystem<T> {
    type SystemData = (
        Entities<'a>,
        ReadStorage<'a, LocalTransform3D<T>>,
        ReadStorage<'a, Child>,
        WriteStorage<'a, Init>,
        WriteStorage<'a, Transform3D<T>>
    );

    fn run(&mut self, (entities, locals, children, mut init, mut globals): Self::SystemData) {

        for (entity, _, child, _) in (&*entities, &locals, &children, !&init).join() {
            self.indices.insert(entity, self.sorted.len());
            self.sorted.push((entity, child.parent()));
            self.new.push(entity);
        }

        for &(entity, _) in &self.sorted {
            if let Some(child) = children.get(entity) {
                if !entities.is_alive(child.parent()) || self.dead.contains(&child.parent()) {
                    let _ = entities.delete(entity);
                    self.dead.insert(entity);
                }
            }
        }

        for entity in self.new.drain(..) {
            init.insert(entity, Init);
        }

        for (entity, local, global, _) in (&*entities, &locals, &mut globals, !&children).join() {
            if local.is_dirty() {
                self.dirty.insert(entity);
                global.0 = local.matrix();
                local.flag(false);
            }
        }

        let mut index = 0;
        while index < self.sorted.len() {
            let (entity, parent_entity) = self.sorted[index];

            match (
                children.get(entity),
                locals.get(entity),
                self.dead.contains(&entity),
            ) {
                (Some(child), Some(local), false) => {
                    if child.is_dirty() && !self.swapped.contains(&entity) {
                        if child.parent() != parent_entity {
                            self.sorted[index] = (entity, child.parent());
                        }

                        let mut swap = None;

                        if let Some(parent_index) = self.indices.get(&child.parent()) {
                            if parent_index > &index {
                                swap = Some(*parent_index);
                            }
                        }

                        if let Some(p) = swap {
                            self.sorted.swap(p, index);
                            self.indices.insert(child.parent(), index);
                            self.indices.insert(entity, p);
                            self.swapped.insert(entity);
                            continue;
                        }
                    }

                    if local.is_dirty() || child.is_dirty() ||
                        self.dirty.contains(&child.parent())
                    {
                        let mut global_matrix = local.matrix();

                        if let Some(parent_global) = globals.get(child.parent()) {
                            let local_matrix = global_matrix.clone();
                            mat4::mul(&mut global_matrix, &parent_global.0, &local_matrix);
                        }

                        if let Some(global) = globals.get_mut(entity) {
                            global.0 = global_matrix;
                        }

                        local.flag(false);
                        child.flag(false);
                        self.dirty.insert(entity);
                    }
                }
                _ => {
                    self.sorted.swap_remove(index);
                    if let Some(swapped) = self.sorted.get(index) {
                        self.indices.insert(swapped.0, index);

                        if let Some(parent) = children.get(swapped.0) {
                            parent.flag(true);
                        }
                    }
                    self.indices.remove(&entity);
                    init.remove(entity);

                    continue;
                }
            }

            index += 1;
        }

        self.dirty.clear();
        self.dead.clear();
        self.swapped.clear();
    }
}


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

    use core::marker::PhantomData;

    use super::super::super::register;
    use specs::{World, DispatcherBuilder};


    pub struct TestSystem<T: 'static + Sync + Send + Copy + Float>(PhantomData<T>);


    impl<T: 'static + Sync + Send + Copy + Float> TestSystem<T> {
        #[inline(always)]
        pub fn new() -> Self {
            TestSystem(PhantomData)
        }
    }


    impl<'a, T: 'static + Sync + Send + Copy + Float> System<'a> for TestSystem<T> {
        type SystemData = WriteStorage<'a, LocalTransform3D<T>>;

        fn run(&mut self, mut locals: Self::SystemData) {
            for local in (&mut locals).join() {
                local.position[0] = T::from_f32(1.0);
                local.position[1] = T::from_f32(1.0);
                local.position[2] = T::from_f32(1.0);
            }
        }
    }


    #[test]
    fn test() {
        let mut world = World::new();

        register::<f32>(&mut world);

        let parent = world.create_entity()
            .with(LocalTransform3D::<f32>::new())
            .with(Transform3D::<f32>::new())
            .build();

        world.create_entity()
            .with(Child::new(parent))
            .with(LocalTransform3D::<f32>::new())
            .with(Transform3D::<f32>::new())
            .build();

        let mut dispatcher = DispatcherBuilder::new()
            .add(Transform3DSystem::<f32>::new(), "transform_3d_system", &[])
            .add(TestSystem::<f32>::new(), "test_system", &["transform_3d_system"])
            .build();

        dispatcher.dispatch(&mut world.res); // first frame init
        dispatcher.dispatch(&mut world.res);

        let globals = world.read::<Transform3D<f32>>();
        let child = (&globals).join().nth(1).unwrap();

        assert_eq!(&child.0, &[
            1.0, 0.0, 0.0, 0.0,
            0.0, 1.0, 0.0, 0.0,
            0.0, 0.0, 1.0, 0.0,
            2.0, 2.0, 2.0, 1.0
        ]);
    }
}
