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

use core::marker::PhantomData;
use core::cmp::Ordering;

use mat4;
use mat32;
use hibitset::BitSet;
use number_traits::Float;
use specs_guided_join::GuidedJoin;

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

use super::{TransformGuide, Child, Init, LocalTransform2D, Transform2D, LocalTransform3D, Transform3D};


pub struct TransformSystem<T: 'static + Sync + Send + Copy + Float> {
    depth: BTreeMap<Entity, usize>,
    children: BTreeMap<Entity, Entity>,
    guide: BTreeSet<Entity>,
    new: Vec<Entity>,
    dead: BitSet,
    dirty: BitSet,
    _marker: PhantomData<T>,
}

impl<T> Default for TransformSystem<T>
    where T: 'static + Sync + Send + Copy + Float,
{
    #[inline(always)]
    fn default() -> Self {
        TransformSystem {
            depth: BTreeMap::default(),
            children: BTreeMap::default(),
            guide: BTreeSet::default(),
            new: Vec::default(),
            dead: BitSet::default(),
            dirty: BitSet::default(),
            _marker: PhantomData,
        }
    }
}

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

    #[inline(always)]
    pub fn new() -> Self {
        Self::default()
    }
}


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

    fn run(&mut self, (
        entities, mut guide, children, mut inits, locals_2d, locals_3d, mut globals_2d, mut globals_3d
    ): Self::SystemData) {

        for (entity, _, child, _) in (&*entities, &locals_2d, &children, !&inits).join() {
            self.children.insert(entity, child.parent());
            self.new.push(entity);
            self.guide.insert(entity);
        }
        for (entity, _, child, _) in (&*entities, &locals_3d, &children, !&inits).join() {
            self.children.insert(entity, child.parent());
            self.new.push(entity);
            self.guide.insert(entity);
        }

        for (entity, _, _) in (&*entities, &locals_2d, !&inits).join() {
            self.guide.insert(entity);
        }
        for (entity, _, _) in (&*entities, &locals_3d, !&inits).join() {
            self.guide.insert(entity);
        }

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

        let needs_sort = self.new.len() != 0;

        for entity in self.new.drain(..) {
            get_depth(
                &mut self.depth, &self.children,
                self.children.get(&entity).map(|p| *p),
                entity, 0
            );
            inits.insert(entity, Init);
        }

        for (entity, local, global, _) in (&*entities, &locals_2d, &mut globals_2d, !&children).join() {
            if local.is_dirty() {
                self.dirty.add(entity.id());
                global.0 = local.matrix();
                local.flag(false);
            }
        }
        for (entity, local, global, _) in (&*entities, &locals_3d, &mut globals_3d, !&children).join() {
            if local.is_dirty() {
                self.dirty.add(entity.id());
                global.0 = local.matrix();
                local.flag(false);
            }
        }

        if needs_sort {
            guide.set(self.guide.iter().map(|e| *e).collect());

            guide.as_slice_mut().sort_by(|a, b| {
                if let Some(a) = self.depth.get(a) {
                    if let Some(b) = self.depth.get(b) {
                        a.cmp(b)
                    } else {
                        Ordering::Equal
                    }
                } else {
                    Ordering::Equal
                }
            });
        }

        for (entity, local, child) in (&*entities, &locals_2d, &children).guided_join(guide.as_slice()) {
            let mut global_matrix = local.matrix();

            if let Some(parent_global) = globals_2d.get(child.parent()) {
                let local_matrix = global_matrix.clone();
                mat32::mul(&mut global_matrix, &parent_global.0, &local_matrix);
            } else if let Some(parent_global) = globals_3d.get(child.parent()) {
                let local_matrix = global_matrix.clone();
                let mut parent_global_2d = mat32::new_identity();
                mat32::from_mat4(&mut parent_global_2d, &parent_global.0);
                mat32::mul(&mut global_matrix, &parent_global_2d, &local_matrix);
            }

            if let Some(global_2d) = globals_2d.get_mut(entity) {
                global_2d.0 = global_matrix;
            }

            local.flag(false);
            child.flag(false);
        }

        for (entity, local, child) in (&*entities, &locals_3d, &children).guided_join(guide.as_slice()) {
            let mut global_matrix = local.matrix();

            if let Some(parent_global) = globals_3d.get(child.parent()) {
                let local_matrix = global_matrix.clone();
                mat4::mul(&mut global_matrix, &parent_global.0, &local_matrix);
            } else if let Some(parent_global) = globals_2d.get(child.parent()) {
                let local_matrix = global_matrix.clone();
                let mut parent_global_3d = mat4::new_identity();
                mat4::from_mat32(&mut parent_global_3d, &parent_global.0);
                mat4::mul(&mut global_matrix, &parent_global_3d, &local_matrix);
            }

            if let Some(global_3d) = globals_3d.get_mut(entity) {
                global_3d.0 = global_matrix;
            }

            local.flag(false);
            child.flag(false);
        }

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

#[inline]
fn get_depth<'a>(
    depth: &'a mut BTreeMap<Entity, usize>,
    children: &'a BTreeMap<Entity, Entity>,
    parent: Option<Entity>,
    entity: Entity,
    current: usize
) -> usize {
    if let Some(depth) = depth.get(&entity) {
        return *depth;
    }

    if let Some(parent) = parent {
        depth.insert(parent, current);
        get_depth(depth, children, children.get(&parent).map(|p| *p), entity, current + 1)
    } else {
        depth.insert(entity, current);
        current
    }
}