use bone_types::{
    DegreesOfFreedom, Parameter, ParameterIndex, ResidualIndex, SolverResidual, SolverSeed,
};
use nalgebra::{DMatrix, DVector};

use crate::jacobian::{assemble_jacobian, evaluate_residuals};
use crate::system::ConstraintSystem;

#[derive(Clone, Debug, PartialEq)]
pub struct DofReport {
    dof: DegreesOfFreedom,
    under_constrained: Vec<ParameterIndex>,
    over_constrained: Vec<ResidualIndex>,
    redundant_consistent: Vec<ResidualIndex>,
}

impl DofReport {
    #[must_use]
    pub fn dof(&self) -> DegreesOfFreedom {
        self.dof
    }

    #[must_use]
    pub fn under_constrained(&self) -> &[ParameterIndex] {
        &self.under_constrained
    }

    #[must_use]
    pub fn over_constrained(&self) -> &[ResidualIndex] {
        &self.over_constrained
    }

    #[must_use]
    pub fn redundant_consistent(&self) -> &[ResidualIndex] {
        &self.redundant_consistent
    }

    #[must_use]
    pub fn display(&self) -> String {
        format!(
            "dof={} under=[{}] over=[{}] redundant=[{}]",
            self.dof.value(),
            self.under_constrained
                .iter()
                .map(ToString::to_string)
                .collect::<Vec<_>>()
                .join(","),
            self.over_constrained
                .iter()
                .map(ToString::to_string)
                .collect::<Vec<_>>()
                .join(","),
            self.redundant_consistent
                .iter()
                .map(ToString::to_string)
                .collect::<Vec<_>>()
                .join(","),
        )
    }
}

#[derive(Copy, Clone, Debug)]
pub struct DofConfig {
    pub seed: SolverSeed,
    pub witness_scale: f64,
    pub rank_tolerance: f64,
    pub residual_tolerance: SolverResidual,
}

impl DofConfig {
    pub const DEFAULT: Self = Self {
        seed: SolverSeed::DEFAULT,
        witness_scale: 0.5,
        rank_tolerance: 1e-9,
        residual_tolerance: SolverResidual::new(1e-7),
    };
}

#[must_use]
pub fn analyze_dof(system: &ConstraintSystem, cfg: DofConfig) -> DofReport {
    analyze_dof_at(system, system.parameters(), cfg)
}

#[must_use]
pub fn analyze_dof_at(
    system: &ConstraintSystem,
    params: &[Parameter],
    cfg: DofConfig,
) -> DofReport {
    let current_params: Vec<f64> = params.iter().map(|p| p.value()).collect();
    let n = system.parameter_count();
    let m = system.row_count();
    if m == 0 || n == 0 {
        let Ok(dof_n) = u32::try_from(n) else {
            unreachable!("parameter count fits in u32 via ParameterIndex")
        };
        return DofReport {
            dof: DegreesOfFreedom::new(dof_n),
            under_constrained: (0..dof_n).map(ParameterIndex::new).collect(),
            over_constrained: Vec::new(),
            redundant_consistent: Vec::new(),
        };
    }
    let witness_params = perturb(&current_params, cfg.seed, cfg.witness_scale);
    let jacobian = assemble_jacobian(system, &witness_params);
    let qr_cols = col_piv_qr(&jacobian, cfg.rank_tolerance);
    let transpose = jacobian.transpose();
    let qr_rows = col_piv_qr(&transpose, cfg.rank_tolerance);
    let rank = qr_cols.rank;
    let Ok(dof_n) = u32::try_from(n.saturating_sub(rank)) else {
        unreachable!("rank <= n <= ParameterIndex::MAX")
    };
    let under_constrained: Vec<ParameterIndex> = qr_cols.col_perm[rank..]
        .iter()
        .map(|&c| {
            let Ok(cv) = u32::try_from(c) else {
                unreachable!("pivoted column index fits in u32 via ParameterIndex")
            };
            ParameterIndex::new(cv)
        })
        .collect();
    let mut under_sorted = under_constrained;
    under_sorted.sort_by_key(|p| p.value());

    let residuals_at_current = evaluate_residuals(system, &current_params);
    let tol = cfg.residual_tolerance.value();
    let (over_constrained, redundant_consistent): (Vec<ResidualIndex>, Vec<ResidualIndex>) =
        qr_rows.col_perm[rank..]
            .iter()
            .map(|&row| {
                let Ok(rv) = u32::try_from(row) else {
                    unreachable!("dependent row index fits in u32 via ResidualIndex")
                };
                let value = residuals_at_current[row];
                (ResidualIndex::new(rv), value.abs() <= tol)
            })
            .fold(
                (Vec::new(), Vec::new()),
                |(mut over, mut red), (idx, consistent)| {
                    if consistent {
                        red.push(idx);
                    } else {
                        over.push(idx);
                    }
                    (over, red)
                },
            );
    let mut over_sorted = over_constrained;
    let mut red_sorted = redundant_consistent;
    over_sorted.sort_by_key(|r| r.value());
    red_sorted.sort_by_key(|r| r.value());
    DofReport {
        dof: DegreesOfFreedom::new(dof_n),
        under_constrained: under_sorted,
        over_constrained: over_sorted,
        redundant_consistent: red_sorted,
    }
}

fn perturb(params: &[f64], seed: SolverSeed, scale: f64) -> Vec<f64> {
    params
        .iter()
        .scan(WitnessRng::new(seed), |rng, &v| {
            let (u, next) = rng.step();
            *rng = next;
            Some(v + (u - 0.5) * scale)
        })
        .collect()
}

#[derive(Copy, Clone, Debug)]
struct WitnessRng {
    state: u64,
}

impl WitnessRng {
    fn new(seed: SolverSeed) -> Self {
        Self {
            state: seed.value(),
        }
    }

    fn step(self) -> (f64, Self) {
        let state = self.state.wrapping_add(0x9E37_79B9_7F4A_7C15);
        let mixed = mix64(state);
        let mantissa = mixed & 0x000F_FFFF_FFFF_FFFF;
        let uniform = f64::from_bits(0x3FF0_0000_0000_0000 | mantissa) - 1.0;
        (uniform, Self { state })
    }
}

fn mix64(x: u64) -> u64 {
    let a = (x ^ (x >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
    let b = (a ^ (a >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
    b ^ (b >> 31)
}

#[derive(Clone, Debug)]
struct ColPivQr {
    rank: usize,
    col_perm: Vec<usize>,
}

fn col_piv_qr(input: &DMatrix<f64>, tol: f64) -> ColPivQr {
    let m = input.nrows();
    let n = input.ncols();
    let initial_norms: Vec<f64> = (0..n).map(|j| input.column(j).norm_squared()).collect();
    let initial = QrState {
        a: input.clone(),
        col_perm: (0..n).collect(),
        col_norms: initial_norms,
        rank: 0,
    };
    let max_steps = m.min(n);
    let final_state = (0..max_steps).try_fold(initial, |state, k| state.step(k, tol));
    let state = match final_state {
        Ok(s) | Err(s) => s,
    };
    ColPivQr {
        rank: state.rank,
        col_perm: state.col_perm,
    }
}

#[derive(Clone, Debug)]
struct QrState {
    a: DMatrix<f64>,
    col_perm: Vec<usize>,
    col_norms: Vec<f64>,
    rank: usize,
}

impl QrState {
    fn step(self, pivot_row: usize, tol: f64) -> Result<Self, Self> {
        let QrState {
            a: matrix,
            mut col_perm,
            mut col_norms,
            rank,
        } = self;
        let rows = matrix.nrows();
        let cols = matrix.ncols();
        let (best_col, best_norm2) = col_norms.iter().enumerate().skip(pivot_row).fold(
            (pivot_row, f64::NEG_INFINITY),
            |(best, best_v), (j, &value)| {
                if value > best_v {
                    (j, value)
                } else {
                    (best, best_v)
                }
            },
        );
        if best_norm2 <= tol * tol {
            return Err(QrState {
                a: matrix,
                col_perm,
                col_norms,
                rank,
            });
        }
        let mut matrix = matrix;
        if best_col != pivot_row {
            matrix.swap_columns(pivot_row, best_col);
            col_perm.swap(pivot_row, best_col);
            col_norms.swap(pivot_row, best_col);
        }
        let tail = rows - pivot_row;
        let column: DVector<f64> = (0..tail)
            .map(|i| matrix[(pivot_row + i, pivot_row)])
            .collect::<Vec<_>>()
            .into();
        let column_norm = column.norm();
        if column_norm <= tol {
            return Err(QrState {
                a: matrix,
                col_perm,
                col_norms,
                rank,
            });
        }
        let alpha = if column[0] >= 0.0 {
            -column_norm
        } else {
            column_norm
        };
        let reflector: DVector<f64> = column
            .iter()
            .enumerate()
            .map(|(idx, &value)| if idx == 0 { value - alpha } else { value })
            .collect::<Vec<_>>()
            .into();
        let reflector_dot = reflector.dot(&reflector);
        let applied = if reflector_dot > 0.0 {
            let beta = 2.0 / reflector_dot;
            (pivot_row..cols).fold(matrix, |mut mat, col| {
                let projection: f64 = (0..tail)
                    .map(|i| reflector[i] * mat[(pivot_row + i, col)])
                    .sum();
                let coef = beta * projection;
                (0..tail).for_each(|i| mat[(pivot_row + i, col)] -= coef * reflector[i]);
                mat
            })
        } else {
            matrix
        };
        let new_norms: Vec<f64> = col_norms
            .iter()
            .enumerate()
            .map(|(col, &prev)| {
                if col <= pivot_row || pivot_row + 1 >= rows {
                    prev
                } else {
                    (0..(tail - 1))
                        .map(|i| {
                            let entry = applied[(pivot_row + 1 + i, col)];
                            entry * entry
                        })
                        .sum::<f64>()
                }
            })
            .collect();
        Ok(QrState {
            a: applied,
            col_perm,
            col_norms: new_norms,
            rank: rank + 1,
        })
    }
}
