use lazy_static::lazy_static;
use nalgebra::{Const, DMatrix, DVector, Dyn};
use web_sys::{console, wasm_bindgen::JsValue}; /* DEBUG */

// --- elements ---

pub fn point(x: f64, y: f64, z: f64) -> DVector<f64> {
    DVector::from_column_slice(&[x, y, z, 0.5, 0.5*(x*x + y*y + z*z)])
}

// the sphere with the given center and radius, with inward-pointing normals
pub fn sphere(center_x: f64, center_y: f64, center_z: f64, radius: f64) -> DVector<f64> {
    let center_norm_sq = center_x * center_x + center_y * center_y + center_z * center_z;
    DVector::from_column_slice(&[
        center_x / radius,
        center_y / radius,
        center_z / radius,
        0.5 / radius,
        0.5 * (center_norm_sq / radius - radius)
    ])
}

// the sphere of curvature `curv` whose closest point to the origin has position
// `off * dir` and normal `dir`, where `dir` is a unit vector. setting the
// curvature to zero gives a plane
pub fn sphere_with_offset(dir_x: f64, dir_y: f64, dir_z: f64, off: f64, curv: f64) -> DVector<f64> {
    let norm_sp = 1.0 + off * curv;
    DVector::from_column_slice(&[
        norm_sp * dir_x,
        norm_sp * dir_y,
        norm_sp * dir_z,
        0.5 * curv,
        off * (1.0 + 0.5 * off * curv)
    ])
}

// --- partial matrices ---

struct MatrixEntry {
    index: (usize, usize),
    value: f64
}

pub struct PartialMatrix(Vec<MatrixEntry>);

impl PartialMatrix {
    pub fn new() -> PartialMatrix {
        PartialMatrix(Vec::<MatrixEntry>::new())
    }
    
    pub fn push_sym(&mut self, row: usize, col: usize, value: f64) {
        let PartialMatrix(entries) = self;
        entries.push(MatrixEntry { index: (row, col), value: value });
        if row != col {
            entries.push(MatrixEntry { index: (col, row), value: value });
        }
    }
    
    /* DEBUG */
    pub fn log_to_console(&self) {
        let PartialMatrix(entries) = self;
        for ent in entries {
            let ent_str = format!("  {} {} {}", ent.index.0, ent.index.1, ent.value);
            console::log_1(&JsValue::from(ent_str.as_str()));
        }
    }
    
    fn proj(&self, a: &DMatrix<f64>) -> DMatrix<f64> {
        let mut result = DMatrix::<f64>::zeros(a.nrows(), a.ncols());
        let PartialMatrix(entries) = self;
        for ent in entries {
            result[ent.index] = a[ent.index];
        }
        result
    }
    
    fn sub_proj(&self, rhs: &DMatrix<f64>) -> DMatrix<f64> {
        let mut result = DMatrix::<f64>::zeros(rhs.nrows(), rhs.ncols());
        let PartialMatrix(entries) = self;
        for ent in entries {
            result[ent.index] = ent.value - rhs[ent.index];
        }
        result
    }
}

// --- descent history ---

pub struct DescentHistory {
    pub config: Vec<DMatrix<f64>>,
    pub scaled_loss: Vec<f64>,
    pub neg_grad: Vec<DMatrix<f64>>,
    pub min_eigval: Vec<f64>,
    pub base_step: Vec<DMatrix<f64>>,
    pub backoff_steps: Vec<i32>
}

impl DescentHistory {
    fn new() -> DescentHistory {
        DescentHistory {
            config: Vec::<DMatrix<f64>>::new(),
            scaled_loss: Vec::<f64>::new(),
            neg_grad: Vec::<DMatrix<f64>>::new(),
            min_eigval: Vec::<f64>::new(),
            base_step: Vec::<DMatrix<f64>>::new(),
            backoff_steps: Vec::<i32>::new(),
        }
    }
}

// --- gram matrix realization ---

// the Lorentz form
lazy_static! {
    pub static ref Q: DMatrix<f64> = DMatrix::from_row_slice(5, 5, &[
        1.0, 0.0, 0.0,  0.0,  0.0,
        0.0, 1.0, 0.0,  0.0,  0.0,
        0.0, 0.0, 1.0,  0.0,  0.0,
        0.0, 0.0, 0.0,  0.0, -2.0,
        0.0, 0.0, 0.0, -2.0,  0.0
    ]);
}

struct SearchState {
    config: DMatrix<f64>,
    err_proj: DMatrix<f64>,
    loss: f64
}

impl SearchState {
    fn from_config(gram: &PartialMatrix, config: DMatrix<f64>) -> SearchState {
        let err_proj = gram.sub_proj(&(config.tr_mul(&*Q) * &config));
        let loss = err_proj.norm_squared();
        SearchState {
            config: config,
            err_proj: err_proj,
            loss: loss
        }
    }
}

fn basis_matrix(index: (usize, usize), nrows: usize, ncols: usize) -> DMatrix<f64> {
    let mut result = DMatrix::<f64>::zeros(nrows, ncols);
    result[index] = 1.0;
    result
}

// use backtracking line search to find a better configuration
fn seek_better_config(
    gram: &PartialMatrix,
    state: &SearchState,
    base_step: &DMatrix<f64>,
    base_target_improvement: f64,
    min_efficiency: f64,
    backoff: f64,
    max_backoff_steps: i32
) -> Option<(SearchState, i32)> {
    let mut rate = 1.0;
    for backoff_steps in 0..max_backoff_steps {
        let trial_config = &state.config + rate * base_step;
        let trial_state = SearchState::from_config(gram, trial_config);
        let improvement = state.loss - trial_state.loss;
        if improvement >= min_efficiency * rate * base_target_improvement {
            return Some((trial_state, backoff_steps));
        }
        rate *= backoff;
    }
    None
}

// seek a matrix `config` for which `config' * Q * config` matches the partial
// matrix `gram`. use gradient descent starting from `guess`
pub fn realize_gram(
    gram: &PartialMatrix,
    guess: DMatrix<f64>,
    frozen: &[(usize, usize)],
    scaled_tol: f64,
    min_efficiency: f64,
    backoff: f64,
    reg_scale: f64,
    max_descent_steps: i32,
    max_backoff_steps: i32
) -> (DMatrix<f64>, bool, DescentHistory) {
    // start the descent history
    let mut history = DescentHistory::new();
    
    // find the dimension of the search space
    let element_dim = guess.nrows();
    let assembly_dim = guess.ncols();
    let total_dim = element_dim * assembly_dim;
    
    // scale the tolerance
    let scale_adjustment = (gram.0.len() as f64).sqrt();
    let tol = scale_adjustment * scaled_tol;
    
    // convert the frozen indices to stacked format
    let frozen_stacked: Vec<usize> = frozen.into_iter().map(
        |index| index.1*element_dim + index.0
    ).collect();
    
    // use Newton's method with backtracking and gradient descent backup
    let mut state = SearchState::from_config(gram, guess);
    for _ in 0..max_descent_steps {
        // stop if the loss is tolerably low
        history.config.push(state.config.clone());
        history.scaled_loss.push(state.loss / scale_adjustment);
        if state.loss < tol { break; }
        
        // find the negative gradient of the loss function
        let neg_grad = 4.0 * &*Q * &state.config * &state.err_proj;
        let mut neg_grad_stacked = neg_grad.clone().reshape_generic(Dyn(total_dim), Const::<1>);
        history.neg_grad.push(neg_grad.clone());
        
        // find the negative Hessian of the loss function
        let mut hess_cols = Vec::<DVector<f64>>::with_capacity(total_dim);
        for col in 0..assembly_dim {
            for row in 0..element_dim {
                let index = (row, col);
                let basis_mat = basis_matrix(index, element_dim, assembly_dim);
                let neg_d_err =
                    basis_mat.tr_mul(&*Q) * &state.config
                    + state.config.tr_mul(&*Q) * &basis_mat;
                let neg_d_err_proj = gram.proj(&neg_d_err);
                let deriv_grad = 4.0 * &*Q * (
                    -&basis_mat * &state.err_proj
                    + &state.config * &neg_d_err_proj
                );
                hess_cols.push(deriv_grad.reshape_generic(Dyn(total_dim), Const::<1>));
            }
        }
        let mut hess = DMatrix::from_columns(hess_cols.as_slice());
        
        // regularize the Hessian
        let min_eigval = hess.symmetric_eigenvalues().min();
        if min_eigval <= 0.0 {
            hess -= reg_scale * min_eigval * DMatrix::identity(total_dim, total_dim);
        }
        history.min_eigval.push(min_eigval);
        
        // project the negative gradient and negative Hessian onto the
        // orthogonal complement of the frozen subspace
        let zero_col = DVector::zeros(total_dim);
        let zero_row = zero_col.transpose();
        for &k in &frozen_stacked {
            neg_grad_stacked[k] = 0.0;
            hess.set_row(k, &zero_row);
            hess.set_column(k, &zero_col);
            hess[(k, k)] = 1.0;
        }
        
        // compute the Newton step
        /*
          we need to either handle or eliminate the case where the minimum
          eigenvalue of the Hessian is zero, so the regularized Hessian is
          singular. right now, this causes the Cholesky decomposition to return
          `None`, leading to a panic when we unrap
        */
        let base_step_stacked = hess.cholesky().unwrap().solve(&neg_grad_stacked);
        let base_step = base_step_stacked.reshape_generic(Dyn(element_dim), Dyn(assembly_dim));
        history.base_step.push(base_step.clone());
        
        // use backtracking line search to find a better configuration
        match seek_better_config(
            gram, &state, &base_step, neg_grad.dot(&base_step),
            min_efficiency, backoff, max_backoff_steps
        ) {
            Some((better_state, backoff_steps)) => {
                state = better_state;
                history.backoff_steps.push(backoff_steps);
            },
            None => return (state.config, false, history)
        };
    }
    (state.config, state.loss < tol, history)
}

// --- tests ---

#[cfg(test)]
mod tests {
    use std::{array, f64::consts::PI};
    
    use super::*;
    
    #[test]
    fn sub_proj_test() {
        let target = PartialMatrix(vec![
            MatrixEntry { index: (0, 0), value: 19.0 },
            MatrixEntry { index: (0, 2), value: 39.0 },
            MatrixEntry { index: (1, 1), value: 59.0 },
            MatrixEntry { index: (1, 2), value: 69.0 }
        ]);
        let attempt = DMatrix::<f64>::from_row_slice(2, 3, &[
            1.0, 2.0, 3.0,
            4.0, 5.0, 6.0
        ]);
        let expected_result = DMatrix::<f64>::from_row_slice(2, 3, &[
            18.0, 0.0, 36.0,
            0.0, 54.0, 63.0
        ]);
        assert_eq!(target.sub_proj(&attempt), expected_result);
    }
    
    #[test]
    fn zero_loss_test() {
        let gram = PartialMatrix({
            let mut entries = Vec::<MatrixEntry>::new();
            for j in 0..3 {
                for k in 0..3 {
                    entries.push(MatrixEntry {
                        index: (j, k),
                        value: if j == k { 1.0 } else { -1.0 }
                    });
                }
            }
            entries
        });
        let config = {
            let a: f64 = 0.75_f64.sqrt();
            DMatrix::from_columns(&[
                sphere(1.0, 0.0, 0.0, a),
                sphere(-0.5, a, 0.0, a),
                sphere(-0.5, -a, 0.0, a)
            ])
        };
        let state = SearchState::from_config(&gram, config);
        assert!(state.loss.abs() < f64::EPSILON);
    }
    
    // this problem is from a sangaku by Irisawa Shintarō Hiroatsu. the article
    // below includes a nice translation of the problem statement, which was
    // recorded in Uchida Itsumi's book _Kokon sankan_ (_Mathematics, Past and
    // Present_)
    //
    //   "Japan's 'Wasan' Mathematical Tradition", by Abe Haruki
    //   https://www.nippon.com/en/japan-topics/c12801/
    //
    #[test]
    fn irisawa_hexlet_test() {
        let gram = PartialMatrix({
            let mut entries = Vec::<MatrixEntry>::new();
            for s in 0..9 {
                // each sphere is represented by a spacelike vector
                entries.push(MatrixEntry { index: (s, s), value: 1.0 });
                
                // the circumscribing sphere is tangent to all of the other
                // spheres, with matching orientation
                if s > 0 {
                    entries.push(MatrixEntry { index: (0, s), value: 1.0 });
                    entries.push(MatrixEntry { index: (s, 0), value: 1.0 });
                }
                
                if s > 2 {
                    // each chain sphere is tangent to the "sun" and "moon"
                    // spheres, with opposing orientation
                    for n in 1..3 {
                        entries.push(MatrixEntry { index: (s, n), value: -1.0 });
                        entries.push(MatrixEntry { index: (n, s), value: -1.0 });
                    }
                    
                    // each chain sphere is tangent to the next chain sphere,
                    // with opposing orientation
                    let s_next = 3 + (s-2) % 6;
                    entries.push(MatrixEntry { index: (s, s_next), value: -1.0 });
                    entries.push(MatrixEntry { index: (s_next, s), value: -1.0 });
                }
            }
            entries
        });
        let guess = DMatrix::from_columns(
            [
                sphere(0.0, 0.0, 0.0, 15.0),
                sphere(0.0, 0.0, -9.0, 5.0),
                sphere(0.0, 0.0, 11.0, 3.0)
            ].into_iter().chain(
                (1..=6).map(
                    |k| {
                        let ang = (k as f64) * PI/3.0;
                        sphere(9.0 * ang.cos(), 9.0 * ang.sin(), 0.0, 2.5)
                    }
                )
            ).collect::<Vec<_>>().as_slice()
        );
        let frozen: [(usize, usize); 4] = array::from_fn(|k| (3, k));
        const SCALED_TOL: f64 = 1.0e-12;
        let (config, success, history) = realize_gram(
            &gram, guess, &frozen,
            SCALED_TOL, 0.5, 0.9, 1.1, 200, 110
        );
        let entry_tol = SCALED_TOL.sqrt();
        let solution_diams = [30.0, 10.0, 6.0, 5.0, 15.0, 10.0, 3.75, 2.5, 2.0 + 8.0/11.0];
        for (k, diam) in solution_diams.into_iter().enumerate() {
            assert!((config[(3, k)] - 1.0 / diam).abs() < entry_tol);
        }
        print!("\nCompleted Gram matrix:{}", config.tr_mul(&*Q) * &config);
        if success {
            println!("Target accuracy achieved!");
        } else {
            println!("Failed to reach target accuracy");
        }
        println!("Steps: {}", history.scaled_loss.len() - 1);
        println!("Loss: {}", history.scaled_loss.last().unwrap());
        if success {
            println!("\nChain diameters:");
            println!("  {} sun (given)", 1.0 / config[(3, 3)]);
            for k in 4..9 {
                println!("  {} sun", 1.0 / config[(3, k)]);
            }
        }
        println!("\nStep │ Loss\n─────┼────────────────────────────────");
        for (step, scaled_loss) in history.scaled_loss.into_iter().enumerate() {
            println!("{:<4} │ {}", step, scaled_loss);
        }
    }
}