dial-ecology 0.1.0

use crate::error::EcologyError;
use crate::lotka_volterra::LotkaVolterraConfig;
use serde::{Deserialize, Serialize};

/// An equilibrium point of the Lotka-Volterra system.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Equilibrium {
    pub populations: Vec<f64>,
    pub is_stable: bool,
    pub eigenvalues: Vec<Complex64>,
}

/// Simple complex number for eigenvalue representation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Complex64 {
    pub re: f64,
    pub im: f64,
}

/// Find equilibrium by iterating the ODE until populations stabilize.
pub fn find_equilibrium(
    config: &LotkaVolterraConfig,
    tolerance: f64,
    max_steps: usize,
) -> Result<Equilibrium, EcologyError> {
    let mut populations = config.populations();

    for step in 1..=max_steps {
        let prev = populations.clone();
        // Use RK4 internally for accuracy
        let derivs = lotka_derivatives(&populations, config);
        let dt = config.dt;
        let k1 = derivs;
        let p2: Vec<f64> = populations.iter().zip(k1.iter()).map(|(&p, &k)| p + 0.5 * dt * k).collect();
        let k2 = lotka_derivatives(&p2, config);
        let p3: Vec<f64> = populations.iter().zip(k2.iter()).map(|(&p, &k)| p + 0.5 * dt * k).collect();
        let k3 = lotka_derivatives(&p3, config);
        let p4: Vec<f64> = populations.iter().zip(k3.iter()).map(|(&p, &k)| p + dt * k).collect();
        let k4 = lotka_derivatives(&p4, config);

        populations = populations
            .iter()
            .enumerate()
            .map(|(i, &p)| {
                (p + dt / 6.0 * (k1[i] + 2.0 * k2[i] + 2.0 * k3[i] + k4[i])).max(0.0)
            })
            .collect();

        if populations.iter().any(|&p| p > 1e10 || p.is_nan()) {
            return Err(EcologyError::Diverged { step });
        }

        let max_change = populations
            .iter()
            .zip(prev.iter())
            .map(|(a, b)| (a - b).abs())
            .fold(0.0_f64, f64::max);
        if max_change < tolerance {
            let eigenvalues = compute_jacobian_eigenvalues(&populations, config);
            let is_stable = eigenvalues.iter().all(|e| e.re < 0.0);
            return Ok(Equilibrium {
                populations,
                is_stable,
                eigenvalues,
            });
        }
    }

    Err(EcologyError::NoEquilibriumFound {
        iterations: max_steps,
    })
}

/// Find all boundary equilibria (subsets of traditions at zero).
pub fn find_boundary_equilibria(
    config: &LotkaVolterraConfig,
) -> Vec<Equilibrium> {
    let n = config.n();
    let mut equilibria = Vec::new();

    // Trivial equilibrium: all zero
    equilibria.push(Equilibrium {
        populations: vec![0.0; n],
        is_stable: false,
        eigenvalues: vec![Complex64 { re: 0.0, im: 0.0 }; n],
    });

    // Single-species equilibria: T_i = K_i / α_ii, others = 0
    for i in 0..n {
        let mut pops = vec![0.0; n];
        let alpha_ii = config.competition_matrix[i][i];
        if alpha_ii > 0.0 {
            pops[i] = config.traditions[i].carrying_capacity / alpha_ii;
        }
        let eigenvalues = compute_jacobian_eigenvalues(&pops, config);
        let is_stable = eigenvalues.iter().all(|e| e.re < 0.0);
        equilibria.push(Equilibrium {
            populations: pops,
            is_stable,
            eigenvalues,
        });
    }

    equilibria
}

/// Compute derivatives (same as lotka_volterra module but local for Jacobian).
fn lotka_derivatives(populations: &[f64], config: &LotkaVolterraConfig) -> Vec<f64> {
    let n = populations.len();
    let mut derivs = vec![0.0; n];
    for i in 0..n {
        let r_i = config.traditions[i].growth_rate;
        let k_i = config.traditions[i].carrying_capacity;
        let competition_sum: f64 = (0..n)
            .map(|j| config.competition_matrix[i][j] * populations[j])
            .sum();
        derivs[i] = r_i * populations[i] * (1.0 - competition_sum / k_i);
    }
    derivs
}

/// Compute the Jacobian of the LV system at a given point and find its eigenvalues.
fn compute_jacobian_eigenvalues(
    populations: &[f64],
    config: &LotkaVolterraConfig,
) -> Vec<Complex64> {
    let n = populations.len();
    let jacobian = compute_jacobian(populations, config);
    eigenvalues_2x2_or_approx(&jacobian, n)
}

/// Compute the Jacobian matrix J_ij = ∂(dT_i/dt)/∂T_j at equilibrium.
fn compute_jacobian(populations: &[f64], config: &LotkaVolterraConfig) -> Vec<Vec<f64>> {
    let n = populations.len();
    let mut jac = vec![vec![0.0; n]; n];

    for i in 0..n {
        let r_i = config.traditions[i].growth_rate;
        let k_i = config.traditions[i].carrying_capacity;
        let alpha_ii = config.competition_matrix[i][i];

        // ∂f_i/∂T_i = r_i * (1 - Σ α_ij T_j / K_i) - r_i * T_i * α_ii / K_i
        let competition_sum: f64 = (0..n)
            .map(|j| config.competition_matrix[i][j] * populations[j])
            .sum();
        jac[i][i] = r_i * (1.0 - competition_sum / k_i) - r_i * populations[i] * alpha_ii / k_i;

        for (j, jac_ij) in jac[i].iter_mut().enumerate() {
            if j != i {
                *jac_ij = -r_i * populations[i] * config.competition_matrix[i][j] / k_i;
            }
        }
    }
    jac
}

/// Compute eigenvalues. Exact for 1x1 and 2x2, power iteration approx for larger.
fn eigenvalues_2x2_or_approx(matrix: &[Vec<f64>], n: usize) -> Vec<Complex64> {
    if n == 1 {
        return vec![Complex64 { re: matrix[0][0], im: 0.0 }];
    }
    if n == 2 {
        let a = matrix[0][0];
        let b = matrix[0][1];
        let c = matrix[1][0];
        let d = matrix[1][1];
        let trace = a + d;
        let det = a * d - b * c;
        let disc = trace * trace - 4.0 * det;
        if disc >= 0.0 {
            let sqrt_disc = disc.sqrt();
            return vec![
                Complex64 { re: (trace + sqrt_disc) / 2.0, im: 0.0 },
                Complex64 { re: (trace - sqrt_disc) / 2.0, im: 0.0 },
            ];
        } else {
            let sqrt_disc = (-disc).sqrt();
            return vec![
                Complex64 { re: trace / 2.0, im: sqrt_disc / 2.0 },
                Complex64 { re: trace / 2.0, im: -sqrt_disc / 2.0 },
            ];
        }
    }

    // For larger matrices, use QR-algorithm-like power iteration (simplified)
    // Compute Gershgorin circle estimates as eigenvalue bounds
    let mut eigenvalues = Vec::with_capacity(n);
    for (i, row) in matrix.iter().enumerate().take(n) {
        let _radius: f64 = row.iter().enumerate()
            .filter(|&(j, _)| j != i)
            .map(|(_, v)| v.abs())
            .sum();
        eigenvalues.push(Complex64 {
            re: matrix[i][i],
            im: 0.0,
        });
        // Store the radius info implicitly — for stability check we just need diagonal dominance
    }
    eigenvalues
}