use nalgebra::DMatrix;
use std::{array, f64::consts::PI};

use dyna3::engine::{Q, point, realize_gram, PartialMatrix};

fn main() {
    // set up a kaleidocycle, made of points with fixed distances between them,
    // and find its tangent space
    const N_POINTS: usize = 12;
    let gram = {
        let mut gram_to_be = PartialMatrix::new();
        for block in (0..N_POINTS).step_by(2) {
            let block_next = (block + 2) % N_POINTS;
            for j in 0..2 {
                // diagonal and hinge edges
                for k in j..2 {
                    gram_to_be.push_sym(block + j, block + k, if j == k { 0.0 } else { -0.5 });
                }
                
                // non-hinge edges
                for k in 0..2 {
                    gram_to_be.push_sym(block + j, block_next + k, -0.625);
                }
            }
        }
        gram_to_be
    };
    let guess = {
        const N_HINGES: usize = 6;
        let guess_elts = (0..N_HINGES).step_by(2).flat_map(
            |n| {
                let ang_hor = (n as f64) * PI/3.0;
                let ang_vert = ((n + 1) as f64) * PI/3.0;
                let x_vert = ang_vert.cos();
                let y_vert = ang_vert.sin();
                [
                    point(0.0, 0.0, 0.0),
                    point(ang_hor.cos(), ang_hor.sin(), 0.0),
                    point(x_vert, y_vert, -0.5),
                    point(x_vert, y_vert, 0.5)
                ]
            }
        ).collect::<Vec<_>>();
        DMatrix::from_columns(&guess_elts)
    };
    let frozen: [_; N_POINTS] = array::from_fn(|k| (3, k));
    let (config, tangent, success, history) = realize_gram(
        &gram, guess, &frozen,
        1.0e-12, 0.5, 0.9, 1.1, 200, 110
    );
    print!("Completed Gram matrix:{}", config.tr_mul(&*Q) * &config);
    print!("Configuration:{}", config);
    if success {
        println!("Target accuracy achieved!");
    } else {
        println!("Failed to reach target accuracy");
    }
    println!("Steps: {}", history.scaled_loss.len() - 1);
    println!("Loss: {}\n", history.scaled_loss.last().unwrap());
    
    // find the kaleidocycle's twist motion
    let twist_motion: DMatrix<_> = (0..N_POINTS).step_by(4).flat_map(
        |n| {
            let up_field = {
                DMatrix::from_column_slice(5, 5, &[
                    0.0, 0.0, 0.0, 0.0, 0.0,
                    0.0, 0.0, 0.0, 0.0, 0.0,
                    0.0, 0.0, 0.0, 0.0, 1.0,
                    0.0, 0.0, 2.0, 0.0, 0.0,
                    0.0, 0.0, 0.0, 0.0, 0.0
                ])
            };
            [
                tangent.proj(&(&up_field * config.column(n)).as_view(), n),
                tangent.proj(&(-&up_field * config.column(n+1)).as_view(), n+1)
            ]
        }
    ).sum();
    let normalization = 5.0 / twist_motion[(2, 0)];
    print!("Twist motion:{}", normalization * twist_motion);
}