linear-sim 0.10.3

use either::Either;
use stash::Stash;
use vec_map::VecMap;

use crate::{constraint, event, Constraint};
use crate::constraint::island::{Island, KeyType};
use crate::object::{self, InternalId, ObjectPair};
use crate::math::*;

#[cfg(feature = "derive_serdes")]
use serde::{Deserialize, Serialize};

/// Constraint island manager
#[cfg_attr(feature = "derive_serdes", derive(Deserialize, Serialize))]
#[derive(Clone, Debug)]
pub(crate) struct Manager {
  // NOTE: we put this field in an option so we can take and iterate over islands
  // without borrowing the manager
  pub islands                       : Option <Stash <Island>>,
  /// Pseudo-velocities.
  ///
  /// Post-impulse position corrections utilize this velocity which is reset to zero at
  /// the beginning of each collision step.
  pub pseudo_velocities             : VecMap <Vector3 <f64>>,
  /// Number of constraints for static objects.
  ///
  /// Entries are only occupied when the number of constraints is greater than zero.
  object_constraints_static         : VecMap <u32>,
  /// Island key and number of constraints for dynamic objects
  object_island_constraints_dynamic : VecMap <(KeyType, u32)>
}

impl Manager {
  pub(crate) fn get_island (&self, key : KeyType) -> Option <&Island> {
    self.islands.as_ref().unwrap().get (key as usize)
  }
  pub(crate) fn get_island_mut (&mut self, key : KeyType) -> Option <&mut Island> {
    self.islands.as_mut().unwrap().get_mut (key as usize)
  }
  pub(crate) fn remove_island (&mut self, key : KeyType) -> Option <Island> {
    self.islands.as_mut().unwrap().take (key as usize)
  }
  pub(crate) fn add_constraint (&mut self,
    object_pair : ObjectPair, constraint : constraint::Variant
  ) {
    let (object_id_a, object_id_b) = object_pair.into();
    let island = match (
      self.get_island_key (object_id_a), self.get_island_key (object_id_b)
    ) {
      (Some (island_a), None) => {
        // dynamic object b joins existing island a
        self.increment_constraint (object_id_a);
        self.assign_dynamic_island_key (object_id_b, island_a);
        self.get_island_mut (island_a).unwrap()
      }
      (None, Some (island_b)) => {
        // object a joins existing island b
        if object_id_a.kind() == object::Kind::Dynamic {
          self.assign_dynamic_island_key (object_id_a, island_b);
        } else {
          self.increment_constraint (object_id_a);
        }
        self.increment_constraint (object_id_b);
        self.get_island_mut (island_b).unwrap()
      }
      (Some (island_a), Some (island_b)) => {
        self.increment_constraint (object_id_a);
        self.increment_constraint (object_id_b);
        if island_a == island_b {
          self.get_island_mut (island_a).unwrap()
        } else {
          // merge islands
          let island_b = self.remove_island (island_b).unwrap();
          for (pair, _) in island_b.constraints.iter() {
            let (id_a, id_b) = (*pair).into();
            if id_a.kind() == object::Kind::Dynamic {
              self.change_dynamic_island_key (id_a, island_a);
            }
            self.change_dynamic_island_key (id_b, island_a);
          }
          let island_a = self.get_island_mut (island_a).unwrap();
          island_a.constraints.extend (island_b.constraints);
          island_a
        }
      }
      (None, None) => {
        // new island
        let new_island_index = self.islands.as_mut().unwrap()
          .put (Island::default());
        debug_assert!(new_island_index < KeyType::MAX as usize);
        let new_island_key = new_island_index as KeyType;
        if object_id_a.kind() == object::Kind::Dynamic {
          self.assign_dynamic_island_key (object_id_a, new_island_key);
        } else {
          self.increment_constraint (object_id_a);
        }
        self.assign_dynamic_island_key (object_id_b, new_island_key);
        self.get_island_mut (new_island_key).unwrap()
      }
    };
    island.constraints.push ((object_pair, constraint));
  }

  /// Remove the object if it belongs to a constraint island.
  ///
  /// Returns true if the object was found and removed, otherwise returns false if the
  /// object was not in a constraint island.
  #[must_use]
  pub(crate) fn remove_object (&mut self, object_id : InternalId) -> bool {
    if let Some (island_key) = self.get_island_key (object_id) {
      debug_assert_eq!(object_id.kind(), object::Kind::Dynamic);
      self.pseudo_velocities.remove (object_id.key().index()).unwrap();
      let island   = self.get_island_mut (island_key).unwrap();
      let removed = island.remove_object (object_id);
      let empty   = island.constraints.is_empty();
      if empty {
        let _ = self.remove_island (island_key).unwrap();
      }
      for id in removed {
        self.decrement_constraint (id);
      }
      self.object_island_constraints_dynamic.remove (object_id.key().index()).unwrap();
      true
    } else if object_id.kind() == object::Kind::Static {
      let index = object_id.key().index();
      if let Some (mut contact_count) = self.object_constraints_static.remove (index) {
        debug_assert!(contact_count > 0);
        let mut empty = vec![];
        let mut islands = self.islands.take().unwrap();
        for (i, island) in islands.iter_mut() {
          let removed = island.remove_object (object_id);
          if island.constraints.is_empty() {
            empty.push (i);
          }
          contact_count -= removed.len() as u32;
          for id in removed {
            self.decrement_constraint (id);
          }
          if contact_count == 0 {
            break
          }
        }
        for i in empty {
          islands.take (i);
        }
        self.islands = Some (islands);
        true
      } else {
        false
      }
    } else {
      false
    }
  }

  // TODO: move method ?
  pub(crate) fn output_contacts (&self, output : &mut Vec <event::Output>) {
    for island in self.islands.as_ref().unwrap().values() {
      for (object_pair, constraint) in island.constraints.iter() {
        #[expect(clippy::match_wildcard_for_single_variants)]
        match constraint {
          constraint::Variant::Contact (contact) => {
            let (id_a, id_b) = (*object_pair).into();
            output.push (event::Contact {
              object_id_a: id_a.into(),
              object_id_b: id_b.into(),
              contact:     contact.clone()
            }.into())
          }
          _ => {}
        }
      }
    }
  }

  /// Takes list of constraint indices to remove.
  ///
  /// List should be non-empty. Note that the modified island may no longer be connected
  /// and `island.partition()` should be called to get connected components.
  pub(crate) fn remove_constraints (&mut self,
    island : &mut Island, remove_list : &[u32]
  ) {
    debug_assert!(!remove_list.is_empty());
    island.constraints = island.constraints.drain (..).enumerate().filter_map (
      |(i, contact@(object_pair, _))|
      if remove_list.contains (&(i as u32)) {
        let (object_id_a, object_id_b) = object_pair.into();
        self.decrement_constraint (object_id_a);
        self.decrement_constraint (object_id_b);
        None
      } else {
        Some (contact)
      }
    ).collect();
  }

  /// Get the contact count for an object
  pub(crate) fn get_constraint_count (&self, object_id : InternalId) -> Option <u32> {
    let index = object_id.key().index();
    match object_id.kind() {
      object::Kind::Static  => self.object_constraints_static.get (index).copied(),
      object::Kind::Dynamic => self.object_island_constraints_dynamic.get (index)
        .map (|g| g.1),
      object::Kind::Nodetect => unreachable!()
    }
  }

  /// Get the island key for dynamic object, if it exists, otherwise returns None if the
  /// object is static
  pub(crate) fn get_island_key (&self, object_id : InternalId) -> Option <KeyType> {
    let index = object_id.key().index();
    match object_id.kind() {
      // static objects can be members of more than one island
      object::Kind::Static  => None,
      object::Kind::Dynamic => self.object_island_constraints_dynamic.get (index)
        .map (|g| g.0),
      object::Kind::Nodetect => unreachable!()
    }
  }

  /// Change existing island key
  pub(crate) fn change_dynamic_island_key (&mut self,
    object_id : InternalId, island_key : KeyType
  ) {
    match object_id.kind() {
      object::Kind::Static  => unreachable!(),
      object::Kind::Dynamic => {
        let index = object_id.key().index();
        self.object_island_constraints_dynamic[index].0 = island_key;
      }
      object::Kind::Nodetect => unimplemented!()
    }
  }

  /// Solve velocity constraints for persistent constraint islands
  pub(crate) fn constrain_velocities (&self,
    objects_static  : &VecMap <object::Static>,
    objects_dynamic : &mut VecMap <object::Dynamic>
  ) {
    for (island_index, island) in self.islands.as_ref().unwrap().iter() {
      log::debug!(island_index, contacts_count=island.constraints.len();
        "island constrain velocities");
      log::trace!(island_index, constraints:?=island.constraints;
        "island constrain velocities constraints");
      let mut iter = 0;
      loop {
        let mut satisfied = true;
        for (object_pair, constraint) in island.constraints.iter() {
          let (object_id_a, object_id_b) = (*object_pair).into();
          let index_a = object_id_a.key().index();
          let index_b = object_id_b.key().index();
          let mut object_a = match object_id_a.kind() {
            object::Kind::Static   => Either::Left  (objects_static[index_a].clone()),
            object::Kind::Dynamic  => Either::Right (objects_dynamic[index_a].clone()),
            object::Kind::Nodetect => unreachable!()
          };
          debug_assert_eq!(object_id_b.kind(), object::Kind::Dynamic);
          let mut object_b = objects_dynamic[index_b].clone();
          // NOTE: at this stage in the pipeline, pseudovelocities are zero
          let impulses = constraint.constrain_velocity (
            object_a.as_mut().map_left (|object_static| &*object_static),
            &mut object_b, None, Vector3::zero());
          if impulses.is_some() {
            match object_a {
              Either::Right (object_a) => objects_dynamic[index_a] = object_a,
              Either::Left  (_) => {}
            }
            objects_dynamic[index_b] = object_b;
            satisfied = false;
          }
        }
        iter += 1;
        if satisfied {
          // island satisfied
          log::debug!(island_index, iters=iter; "island velocity constraint satisfied");
          break
        }
      }
    }
  }

  /// Solve positions constraints for persistent constraint islands
  pub(crate) fn constrain_positions (&mut self,
    objects_static  : &VecMap <object::Static>,
    objects_dynamic : &mut VecMap <object::Dynamic>
  ) {
    let mut remove_islands = vec![];
    let mut new_islands    = vec![];
    let mut islands        = self.islands.take().unwrap();
    for (island_index, island) in islands.iter_mut() {
      log::debug!(island_index, constraint_count=island.constraints.len();
        "island constrain positions");
      log::trace!(island_index, contacts:?=island.constraints;
        "island constrain positions constraints");
      let mut iter = 0;
      loop {
        let mut satisfied = true;
        for (object_pair, constraint) in island.constraints.iter_mut() {
          let (object_id_a, object_id_b) = (*object_pair).into();
          let index_a = object_id_a.key().index();
          let index_b = object_id_b.key().index();
          let mut object_a = match object_id_a.kind() {
            object::Kind::Static   => Either::Left  (objects_static[index_a].clone()),
            object::Kind::Dynamic  => Either::Right (objects_dynamic[index_a].clone()),
            object::Kind::Nodetect => unreachable!()
          };
          debug_assert_eq!(object_id_b.kind(), object::Kind::Dynamic);
          let mut object_b = objects_dynamic[index_b].clone();
          let mut pseudovelocity_a = match object_a {
            Either::Left  (_) => None,
            Either::Right (_) => Some (self.pseudo_velocities[index_a])
          };
          let mut pseudovelocity_b = self.pseudo_velocities[index_b];
          let maybe_impulses       = constraint.constrain_position (
            -1.0, object_a.as_mut().map_left (|object_static| &*object_static),
            &mut object_b, pseudovelocity_a.as_mut(), &mut pseudovelocity_b);
          if maybe_impulses.is_some() {
            satisfied = false;
            match object_a {
              Either::Right (object_a) => {
                objects_dynamic[index_a]        = object_a;
                self.pseudo_velocities[index_a] = pseudovelocity_a.unwrap();
              }
              Either::Left  (_) => {}
            }
            objects_dynamic[index_b]        = object_b;
            self.pseudo_velocities[index_b] = pseudovelocity_b;
          }
        }
        iter += 1;
        if satisfied {
          log::debug!(island_index, iters=iter; "island position constraint satisfied");
          break
        }
      }
      // remove broken contacts
      let mut remove_contacts = vec![];
      for (constraint_index, contact) in island.constraints.iter()
        .enumerate().filter_map (|(index, (_, constraint))|
          if let constraint::Variant::Contact (contact) = constraint {
            Some ((index, contact))
          } else {
            None
          }
        )
      {
        if contact.distance >= constraint::contact::CONTACT_DISTANCE {
          remove_contacts.push (constraint_index as u32)
        }
      }
      if !remove_contacts.is_empty() {
        self.remove_constraints (island, remove_contacts.as_slice());
        if island.constraints.is_empty() {
          remove_islands.push (island_index as KeyType);
        } else {
          let partitions = island.partition();
          if !partitions.is_empty() {
            remove_islands.push (island_index as KeyType);
            new_islands.extend (partitions);
          }
        }
      }
    }
    for island_key in remove_islands {
      let _ = islands.take (island_key as usize).unwrap();
    }
    for new_island in new_islands {
      let dynamic_ids = new_island.dynamic_ids();
      let island_key = islands.put (new_island) as KeyType;
      for dynamic_id in dynamic_ids.into_vec() {
        self.change_dynamic_island_key (dynamic_id, island_key);
      }
    }
    self.islands = Some (islands);
  }

  /// Assign a new island key to a dynamic object
  fn assign_dynamic_island_key (&mut self,
    object_id : InternalId, island_key : KeyType
  ) {
    let index = object_id.key().index();
    match object_id.kind() {
      object::Kind::Static  => unreachable!("call increment_contact instead"),
      object::Kind::Dynamic =>
        assert!(self.object_island_constraints_dynamic.insert (index, (island_key, 1))
          .is_none()),
      object::Kind::Nodetect => unimplemented!()
    }
  }

  /// Increment the constraint counter.
  ///
  /// Adds the constraint counter if this is the first constraint of a static object.
  fn increment_constraint (&mut self, object_id : InternalId) {
    let index = object_id.key().index();
    match object_id.kind() {
      object::Kind::Static  =>
        *self.object_constraints_static.entry (index).or_insert (0) += 1,
      object::Kind::Dynamic => self.object_island_constraints_dynamic[index].1 += 1,
      object::Kind::Nodetect => unimplemented!()
    }
  }

  /// Reduce the constraint count for the given object ID
  fn decrement_constraint (&mut self, object_id : InternalId) {
    let index = object_id.key().index();
    match object_id.kind() {
      object::Kind::Dynamic => {
        let count =
          &mut self.object_island_constraints_dynamic.get_mut (index).unwrap().1;
        *count -= 1;
        if *count == 0 {
          self.object_island_constraints_dynamic.remove (index);
        }
      }
      object::Kind::Static  => {
        let count = self.object_constraints_static.get_mut (index).unwrap();
        *count -= 1;
        if *count == 0 {
          self.object_constraints_static.remove (index);
        }
      }
      object::Kind::Nodetect => unreachable!()
    }
  }
}

impl Default for Manager {
  fn default() -> Self {
    Manager {
      islands:                           Some (Stash::default()),
      pseudo_velocities:                 VecMap::default(),
      object_constraints_static:         VecMap::default(),
      object_island_constraints_dynamic: VecMap::default()
    }
  }
}