use nalgebra::{DMatrix, DVector, DVectorView, Vector3};
use rustc_hash::FxHashMap;
use slab::Slab;
use std::{collections::BTreeSet, rc::Rc, sync::atomic::{AtomicU64, Ordering}};
use sycamore::prelude::*;
use web_sys::{console, wasm_bindgen::JsValue}; /* DEBUG */

use crate::{
    engine::{
        Q,
        change_half_curvature,
        local_unif_to_std,
        realize_gram,
        sphere,
        ConfigSubspace,
        ConstraintProblem
    },
    outline::OutlineItem,
    specified::SpecifiedValue
};

// the types of the keys we use to access an assembly's elements and regulators
pub type ElementKey = usize;
pub type RegulatorKey = usize;

pub type ElementColor = [f32; 3];

/* KLUDGE */
// we should reconsider this design when we build a system for switching between
// assemblies. at that point, we might want to switch to hierarchical keys,
// where each each element has a key that identifies it within its assembly and
// each assembly has a key that identifies it within the sesssion
static NEXT_ELEMENT_SERIAL: AtomicU64 = AtomicU64::new(0);

pub trait ProblemPoser {
    fn pose(&self, problem: &mut ConstraintProblem, elts: &Slab<Element>);
}

#[derive(Clone, PartialEq)]
pub struct Element {
    pub id: String,
    pub label: String,
    pub color: ElementColor,
    pub representation: Signal<DVector<f64>>,
    
    // the regulators this element is subject to. the assembly that owns the
    // element is responsible for keeping this set up to date
    pub regulators: Signal<BTreeSet<RegulatorKey>>,
    
    // a serial number, assigned by `Element::new`, that uniquely identifies
    // each element
    pub serial: u64,
    
    // the configuration matrix column index that was assigned to this element
    // last time the assembly was realized, or `None` if the element has never
    // been through a realization
    column_index: Option<usize>
}

impl Element {
    pub fn new(
        id: String,
        label: String,
        color: ElementColor,
        representation: DVector<f64>
    ) -> Element {
        // take the next serial number, panicking if that was the last number we
        // had left. the technique we use to panic on overflow is taken from
        // _Rust Atomics and Locks_, by Mara Bos
        //
        //   https://marabos.nl/atomics/atomics.html#example-handle-overflow
        //
        let serial = NEXT_ELEMENT_SERIAL.fetch_update(
            Ordering::SeqCst, Ordering::SeqCst,
            |serial| serial.checked_add(1)
        ).expect("Out of serial numbers for elements");
        
        Element {
            id: id,
            label: label,
            color: color,
            representation: create_signal(representation),
            regulators: create_signal(BTreeSet::default()),
            serial: serial,
            column_index: None
        }
    }
    
    // the smallest positive depth, represented as a multiple of `dir`, where
    // the line generated by `dir` hits the element (which is assumed to be a
    // sphere). returns `None` if the line misses the sphere. this function
    // should be kept synchronized with `sphere_cast` in `inversive.frag`, which
    // does essentially the same thing on the GPU side
    pub fn cast(&self, dir: Vector3<f64>, assembly_to_world: &DMatrix<f64>) -> Option<f64> {
        // if `a/b` is less than this threshold, we approximate
        // `a*u^2 + b*u + c` by the linear function `b*u + c`
        const DEG_THRESHOLD: f64 = 1e-9;
        
        let rep = self.representation.with_untracked(|rep| assembly_to_world * rep);
        let a = -rep[3] * dir.norm_squared();
        let b = rep.rows_range(..3).dot(&dir);
        let c = -rep[4];
        
        let adjust = 4.0*a*c/(b*b);
        if adjust < 1.0 {
            // as long as `b` is non-zero, the linear approximation of
            //
            //   a*u^2 + b*u + c
            //
            // at `u = 0` will reach zero at a finite depth `u_lin`. the root of
            // the quadratic adjacent to `u_lin` is stored in `lin_root`. if
            // both roots have the same sign, `lin_root` will be the one closer
            // to `u = 0`
            let square_rect_ratio = 1.0 + (1.0 - adjust).sqrt();
            let lin_root = -(2.0*c)/b / square_rect_ratio;
            if a.abs() > DEG_THRESHOLD * b.abs() {
                if lin_root > 0.0 {
                    Some(lin_root)
                } else {
                    let other_root = -b/(2.*a) * square_rect_ratio;
                    (other_root > 0.0).then_some(other_root)
                }
            } else {
                (lin_root > 0.0).then_some(lin_root)
            }
        } else {
            // the line through `dir` misses the sphere completely
            None
        }
    }
}

impl ProblemPoser for Element {
    fn pose(&self, problem: &mut ConstraintProblem, _elts: &Slab<Element>) {
        if let Some(index) = self.column_index {
            problem.gram.push_sym(index, index, 1.0);
            problem.guess.set_column(index, &self.representation.get_clone_untracked());
        } else {
            panic!("Tried to write problem data from an unindexed element: \"{}\"", self.id);
        }
    }
}

pub trait Regulator: ProblemPoser + OutlineItem {
    fn subjects(&self) -> Vec<ElementKey>;
    fn measurement(&self) -> ReadSignal<f64>;
    fn set_point(&self) -> Signal<SpecifiedValue>;
    
    fn activate(&self, _assembly: &Assembly) {}
}

pub struct ProductRegulator {
    pub subjects: [ElementKey; 2],
    pub measurement: ReadSignal<f64>,
    pub set_point: Signal<SpecifiedValue>
}

impl ProductRegulator {
    pub fn new(subjects: [ElementKey; 2], assembly: &Assembly) -> ProductRegulator {
        let measurement = assembly.elements.map(
            move |elts| {
                let representations = subjects.map(|subj| elts[subj].representation);
                representations[0].with(|rep_0|
                    representations[1].with(|rep_1|
                        rep_0.dot(&(&*Q * rep_1))
                    )
                )
            }
        );
        
        let set_point = create_signal(SpecifiedValue::from_empty_spec());
        
        ProductRegulator {
            subjects: subjects,
            measurement: measurement,
            set_point: set_point
        }
    }
}

impl Regulator for ProductRegulator {
    fn subjects(&self) -> Vec<ElementKey> {
        self.subjects.into()
    }
    
    fn measurement(&self) -> ReadSignal<f64> {
        self.measurement
    }
    
    fn set_point(&self) -> Signal<SpecifiedValue> {
        self.set_point
    }
}

impl ProblemPoser for ProductRegulator {
    fn pose(&self, problem: &mut ConstraintProblem, elts: &Slab<Element>) {
        self.set_point.with_untracked(|set_pt| {
            if let Some(val) = set_pt.value {
                let subject_column_indices = self.subjects.map(
                    |subj| elts[subj].column_index
                );
                if let [Some(row), Some(col)] = subject_column_indices {
                    problem.gram.push_sym(row, col, val);
                } else {
                    panic!("Tried to write problem data from a regulator with an unindexed subject");
                }
            }
        });
    }
}

pub struct HalfCurvatureRegulator {
    pub subject: ElementKey,
    pub measurement: ReadSignal<f64>,
    pub set_point: Signal<SpecifiedValue>
}

impl HalfCurvatureRegulator {
    pub fn new(subject: ElementKey, assembly: &Assembly) -> HalfCurvatureRegulator {
        const CURVATURE_COMPONENT: usize = 3;
        let measurement = assembly.elements.map(
            move |elts| elts[subject].representation.with(|rep| rep[CURVATURE_COMPONENT])
        );
        
        let set_point = create_signal(SpecifiedValue::from_empty_spec());
        
        HalfCurvatureRegulator {
            subject: subject,
            measurement: measurement,
            set_point: set_point
        }
    }
}

impl Regulator for HalfCurvatureRegulator {
    fn subjects(&self) -> Vec<ElementKey> {
        vec![self.subject]
    }
    
    fn measurement(&self) -> ReadSignal<f64> {
        self.measurement
    }
    
    fn set_point(&self) -> Signal<SpecifiedValue> {
        self.set_point
    }
    
    fn activate(&self, assembly: &Assembly) {
        if let Some(half_curv) = self.set_point.with_untracked(|set_pt| set_pt.value) {
            let representation = assembly.elements.with_untracked(
                |elts| elts[self.subject].representation
            );
            representation.update(
                |rep| change_half_curvature(rep, half_curv)
            );
        }
    }
}

impl ProblemPoser for HalfCurvatureRegulator {
    fn pose(&self, problem: &mut ConstraintProblem, elts: &Slab<Element>) {
        self.set_point.with_untracked(|set_pt| {
            if let Some(val) = set_pt.value {
                if let Some(col) = elts[self.subject].column_index {
                    const CURVATURE_COMPONENT: usize = 3;
                    problem.frozen.push(CURVATURE_COMPONENT, col, val);
                } else {
                    panic!("Tried to write problem data from a regulator with an unindexed subject");
                }
            }
        });
    }
}

// the velocity is expressed in uniform coordinates
pub struct ElementMotion<'a> {
    pub key: ElementKey,
    pub velocity: DVectorView<'a, f64>
}

type AssemblyMotion<'a> = Vec<ElementMotion<'a>>;

// a complete, view-independent description of an assembly
#[derive(Clone)]
pub struct Assembly {
    // elements and regulators
    pub elements: Signal<Slab<Element>>,
    pub regulators: Signal<Slab<Rc<dyn Regulator>>>,
    
    // solution variety tangent space. the basis vectors are stored in
    // configuration matrix format, ordered according to the elements' column
    // indices. when you realize the assembly, every element that's present
    // during realization gets a column index and is reflected in the tangent
    // space. since the methods in this module never assign column indices
    // without later realizing the assembly, we get the following invariant:
    //
    //  (1) if an element has a column index, its tangent motions can be found
    //      in that column of the tangent space basis matrices
    //
    pub tangent: Signal<ConfigSubspace>,
    
    // indexing
    pub elements_by_id: Signal<FxHashMap<String, ElementKey>>
}

impl Assembly {
    pub fn new() -> Assembly {
        Assembly {
            elements: create_signal(Slab::new()),
            regulators: create_signal(Slab::new()),
            tangent: create_signal(ConfigSubspace::zero(0)),
            elements_by_id: create_signal(FxHashMap::default())
        }
    }
    
    // --- inserting elements and regulators ---
    
    // insert a sphere into the assembly without checking whether we already
    // have an element with the same identifier. any element that does have the
    // same identifier will get kicked out of the `elements_by_id` index
    fn insert_sphere_unchecked(&self, elt: Element) -> ElementKey {
        // insert the sphere
        let id = elt.id.clone();
        let key = self.elements.update(|elts| elts.insert(elt));
        self.elements_by_id.update(|elts_by_id| elts_by_id.insert(id, key));
        
        // regulate the sphere's curvature
        self.insert_regulator(HalfCurvatureRegulator::new(key, &self));
        
        key
    }
    
    pub fn try_insert_sphere(&self, elt: Element) -> Option<ElementKey> {
        let can_insert = self.elements_by_id.with_untracked(
            |elts_by_id| !elts_by_id.contains_key(&elt.id)
        );
        if can_insert {
            Some(self.insert_sphere_unchecked(elt))
        } else {
            None
        }
    }
    
    pub fn insert_new_sphere(&self) {
        // find the next unused identifier in the default sequence
        let mut id_num = 1;
        let mut id = format!("sphere{}", id_num);
        while self.elements_by_id.with_untracked(
            |elts_by_id| elts_by_id.contains_key(&id)
        ) {
            id_num += 1;
            id = format!("sphere{}", id_num);
        }
        
        // create and insert a sphere
        let _ = self.insert_sphere_unchecked(
            Element::new(
                id,
                format!("Sphere {}", id_num),
                [0.75_f32, 0.75_f32, 0.75_f32],
                sphere(0.0, 0.0, 0.0, 1.0)
            )
        );
    }
    
    pub fn insert_regulator<T: Regulator + 'static>(&self, regulator: T) {
        // add the regulator to the assembly's regulator list
        let regulator_rc = Rc::new(regulator);
        let key = self.regulators.update(
            |regs| regs.insert(regulator_rc.clone())
        );
        
        // add the regulator to each subject's regulator list
        let subjects = regulator_rc.subjects();
        let subject_regulators: Vec<_> = self.elements.with_untracked(
            |elts| subjects.into_iter().map(
                |subj| elts[subj].regulators
            ).collect()
        );
        for regulators in subject_regulators {
            regulators.update(|regs| regs.insert(key));
        }
        
        // update the realization when the regulator becomes a constraint, or is
        // edited while acting as a constraint
        let self_for_effect = self.clone();
        create_effect(move || {
            console::log_1(&JsValue::from(
                format!("Updated regulator with subjects {:?}", regulator_rc.subjects())
            ));
            if regulator_rc.set_point().with(|set_pt| set_pt.is_present()) {
                regulator_rc.activate(&self_for_effect);
                self_for_effect.realize();
            }
        });
        
        /* DEBUG */
        // print an updated list of regulators
        console::log_1(&JsValue::from("Regulators:"));
        self.regulators.with_untracked(|regs| {
            for (_, reg) in regs.into_iter() {
                console::log_1(&JsValue::from(format!(
                    "  {:?}: {}",
                    reg.subjects(),
                    reg.set_point().with_untracked(
                        |set_pt| {
                            let spec = &set_pt.spec;
                            if spec.is_empty() {
                                "__".to_string()
                            } else {
                                spec.clone()
                            }
                        }
                    )
                )));
            }
        });
    }
    
    // --- realization ---
    
    pub fn realize(&self) {
        // index the elements
        self.elements.update_silent(|elts| {
            for (index, (_, elt)) in elts.into_iter().enumerate() {
                elt.column_index = Some(index);
            }
        });
        
        // set up the constraint problem
        let problem = self.elements.with_untracked(|elts| {
            let mut problem = ConstraintProblem::new(elts.len());
            for (_, elt) in elts {
                elt.pose(&mut problem, elts);
            }
            self.regulators.with_untracked(|regs| {
                for (_, reg) in regs {
                    reg.pose(&mut problem, elts);
                }
            });
            problem
        });
        
        /* DEBUG */
        // log the Gram matrix
        console::log_1(&JsValue::from("Gram matrix:"));
        problem.gram.log_to_console();
        
        /* DEBUG */
        // log the initial configuration matrix
        console::log_1(&JsValue::from("Old configuration:"));
        for j in 0..problem.guess.nrows() {
            let mut row_str = String::new();
            for k in 0..problem.guess.ncols() {
                row_str.push_str(format!(" {:>8.3}", problem.guess[(j, k)]).as_str());
            }
            console::log_1(&JsValue::from(row_str));
        }
        
        // look for a configuration with the given Gram matrix
        let (config, tangent, success, history) = realize_gram(
            &problem, 1.0e-12, 0.5, 0.9, 1.1, 200, 110
        );
        
        /* DEBUG */
        // report the outcome of the search
        console::log_1(&JsValue::from(
            if success {
                "Target accuracy achieved!"
            } else {
                "Failed to reach target accuracy"
            }
        ));
        console::log_2(&JsValue::from("Steps:"), &JsValue::from(history.scaled_loss.len() - 1));
        console::log_2(&JsValue::from("Loss:"), &JsValue::from(*history.scaled_loss.last().unwrap()));
        console::log_2(&JsValue::from("Tangent dimension:"), &JsValue::from(tangent.dim()));
        
        if success {
            // read out the solution
            for (_, elt) in self.elements.get_clone_untracked() {
                elt.representation.update(
                    |rep| rep.set_column(0, &config.column(elt.column_index.unwrap()))
                );
            }
            
            // save the tangent space
            self.tangent.set_silent(tangent);
        }
    }
    
    // --- deformation ---
    
    // project the given motion to the tangent space of the solution variety and
    // move the assembly along it. the implementation is based on invariant (1)
    // from above and the following additional invariant:
    //
    //  (2) if an element is affected by a constraint, it has a column index
    //
    // we have this invariant because the assembly gets realized each time you
    // add a constraint
    pub fn deform(&self, motion: AssemblyMotion) {
        /* KLUDGE */
        // when the tangent space is zero, deformation won't do anything, but
        // the attempt to deform should be registered in the UI. this console
        // message will do for now
        if self.tangent.with(|tan| tan.dim() <= 0 && tan.assembly_dim() > 0) {
            console::log_1(&JsValue::from("The assembly is rigid"));
        }
        
        // give a column index to each moving element that doesn't have one yet.
        // this temporarily breaks invariant (1), but the invariant will be
        // restored when we realize the assembly at the end of the deformation.
        // in the process, we find out how many matrix columns we'll need to
        // hold the deformation
        let realized_dim = self.tangent.with(|tan| tan.assembly_dim());
        let motion_dim = self.elements.update_silent(|elts| {
            let mut next_column_index = realized_dim;
            for elt_motion in motion.iter() {
                let moving_elt = &mut elts[elt_motion.key];
                if moving_elt.column_index.is_none() {
                    moving_elt.column_index = Some(next_column_index);
                    next_column_index += 1;
                }
            }
            next_column_index
        });
        
        // project the element motions onto the tangent space of the solution
        // variety and sum them to get a deformation of the whole assembly. the
        // matrix `motion_proj` that holds the deformation has extra columns for
        // any moving elements that aren't reflected in the saved tangent space
        const ELEMENT_DIM: usize = 5;
        let mut motion_proj = DMatrix::zeros(ELEMENT_DIM, motion_dim);
        for elt_motion in motion {
            // we can unwrap the column index because we know that every moving
            // element has one at this point
            let column_index = self.elements.with_untracked(
                |elts| elts[elt_motion.key].column_index.unwrap()
            );
            
            if column_index < realized_dim {
                // this element had a column index when we started, so by
                // invariant (1), it's reflected in the tangent space
                let mut target_columns = motion_proj.columns_mut(0, realized_dim);
                target_columns += self.tangent.with(
                    |tan| tan.proj(&elt_motion.velocity, column_index)
                );
            } else {
                // this element didn't have a column index when we started, so
                // by invariant (2), it's unconstrained
                let mut target_column = motion_proj.column_mut(column_index);
                let unif_to_std = self.elements.with_untracked(
                    |elts| {
                        elts[elt_motion.key].representation.with_untracked(
                            |rep| local_unif_to_std(rep.as_view())
                        )
                    }
                );
                target_column += unif_to_std * elt_motion.velocity;
            }
        }
        
        // step the assembly along the deformation. this changes the elements'
        // normalizations, so we restore those afterward
        /* KLUDGE */
        // since our test assemblies only include spheres, we assume that every
        // element is on the 1 mass shell
        for (_, elt) in self.elements.get_clone_untracked() {
            elt.representation.update_silent(|rep| {
                match elt.column_index {
                    Some(column_index) => {
                        // step the assembly along the deformation
                        *rep += motion_proj.column(column_index);
                        
                        // restore normalization by contracting toward the last
                        // coordinate axis
                        let q_sp = rep.fixed_rows::<3>(0).norm_squared();
                        let half_q_lt = -2.0 * rep[3] * rep[4];
                        let half_q_lt_sq = half_q_lt * half_q_lt;
                        let scaling = half_q_lt + (q_sp + half_q_lt_sq).sqrt();
                        rep.fixed_rows_mut::<4>(0).scale_mut(1.0 / scaling);
                    },
                    None => {
                        console::log_1(&JsValue::from(
                            format!("No velocity to unpack for fresh element \"{}\"", elt.id)
                        ))
                    }
                };
            });
        }
        
        // bring the configuration back onto the solution variety. this also
        // gets the elements' column indices and the saved tangent space back in
        // sync
        self.realize();
    }
}

#[cfg(test)]
mod tests {
    use crate::engine;
    
    use super::*;
    
    #[test]
    #[should_panic(expected = "Tried to write problem data from an unindexed element: \"sphere\"")]
    fn unindexed_element_test() {
        let _ = create_root(|| {
            Element::new(
                "sphere".to_string(),
                "Sphere".to_string(),
                [1.0_f32, 1.0_f32, 1.0_f32],
                engine::sphere(0.0, 0.0, 0.0, 1.0)
            ).pose(&mut ConstraintProblem::new(1), &Slab::new());
        });
    }
    
    #[test]
    #[should_panic(expected = "Tried to write problem data from a regulator with an unindexed subject")]
    fn unindexed_subject_test() {
        let _ = create_root(|| {
            let mut elts = Slab::new();
            let subjects = [0, 1].map(|k| {
                elts.insert(
                    Element::new(
                        format!("sphere{k}"),
                        format!("Sphere {k}"),
                        [1.0_f32, 1.0_f32, 1.0_f32],
                        engine::sphere(0.0, 0.0, 0.0, 1.0)
                    )
                )
            });
            elts[subjects[0]].column_index = Some(0);
            ProductRegulator {
                subjects: subjects,
                measurement: create_memo(|| 0.0),
                set_point: create_signal(SpecifiedValue::try_from("0.0".to_string()).unwrap())
            }.pose(&mut ConstraintProblem::new(2), &elts);
        });
    }
}