oxigrid 0.1.2

/// Power flow sensitivity analysis.
///
/// Computes:
/// - **GSF** (Generation Shift Factors): ∂P_branch / ∂P_gen via DC approximation
/// - **LODF** (Line Outage Distribution Factors): post-contingency flow change
/// - **∂V/∂Q** and **∂V/∂P** voltage sensitivities via AC Jacobian inverse
///
/// # References
/// - Stott et al., "DC Power Flow Revisited", IEEE TPWRS 2009.
/// - Kundur, "Power System Stability and Control", Chapter 9.
use crate::network::{
    bus::BusType,
    reduction::{build_b_bus, lodf_matrix, ptdf_matrix},
    PowerNetwork,
};
use crate::powerflow::{PowerFlowConfig, PowerFlowMethod};
use serde::{Deserialize, Serialize};

/// Generation Shift Factor matrix.
///
/// `gsf[branch][bus]` = ∂P_branch / ∂P_bus  (DC linear approximation, per-unit).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GsfMatrix {
    /// Row = branch index, col = bus index.
    pub data: Vec<Vec<f64>>,
    pub n_branches: usize,
    pub n_buses: usize,
}

impl GsfMatrix {
    /// Create empty GSF matrix.
    pub fn new(n_branches: usize, n_buses: usize) -> Self {
        Self {
            data: vec![vec![0.0; n_buses]; n_branches],
            n_branches,
            n_buses,
        }
    }

    /// GSF for branch `b` at bus `k`.
    pub fn gsf(&self, branch: usize, bus: usize) -> f64 {
        self.data[branch][bus]
    }

    /// Change in branch flow [p.u.] when generation at `bus` changes by `delta_p_pu` [p.u.].
    pub fn delta_flow(&self, branch: usize, bus: usize, delta_p_pu: f64) -> f64 {
        self.gsf(branch, bus) * delta_p_pu
    }

    /// Post-injection branch flows given a generation dispatch vector.
    pub fn branch_flows(&self, gen_dispatch: &[f64]) -> Vec<f64> {
        (0..self.n_branches)
            .map(|b| {
                gen_dispatch
                    .iter()
                    .enumerate()
                    .map(|(k, &p)| self.gsf(b, k) * p)
                    .sum()
            })
            .collect()
    }
}

/// LODF (Line Outage Distribution Factor) matrix.
///
/// `lodf[branch][outage]` = fraction of outaged branch flow that redistributes to `branch`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LodfMatrix {
    pub data: Vec<Vec<f64>>,
    pub n_branches: usize,
}

impl LodfMatrix {
    /// LODF for monitored branch `l` given outage of branch `k`.
    pub fn lodf(&self, monitored: usize, outaged: usize) -> f64 {
        self.data[monitored][outaged]
    }

    /// Post-contingency flow on branch `l` [p.u.] given base flows.
    ///
    /// Post flow = `base_flow[l]` + `LODF[l,k]` * `base_flow[k]`
    pub fn post_contingency_flow(
        &self,
        base_flows: &[f64],
        monitored: usize,
        outaged: usize,
    ) -> f64 {
        if outaged >= base_flows.len() || monitored >= base_flows.len() {
            return 0.0;
        }
        base_flows[monitored] + self.lodf(monitored, outaged) * base_flows[outaged]
    }

    /// N-1 screening: returns (branch, outage, post_flow) for all violations exceeding limit.
    pub fn screen_n1(&self, base_flows: &[f64], branch_limits: &[f64]) -> Vec<(usize, usize, f64)> {
        let mut violations = Vec::new();
        for k in 0..self.n_branches {
            for l in 0..self.n_branches {
                if l == k {
                    continue;
                }
                let pf = self.post_contingency_flow(base_flows, l, k);
                let limit = if l < branch_limits.len() {
                    branch_limits[l]
                } else {
                    f64::INFINITY
                };
                if pf.abs() > limit {
                    violations.push((l, k, pf));
                }
            }
        }
        violations
    }
}

/// AC voltage sensitivity matrices: ∂V/∂P and ∂V/∂Q.
///
/// Derived from the inverse of the AC Jacobian (J⁻¹).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VoltageSensitivity {
    /// ∂|V|/∂P: n_pq_buses × n_pv_pq_buses
    pub dv_dp: Vec<Vec<f64>>,
    /// ∂|V|/∂Q: n_pq_buses × n_pq_buses
    pub dv_dq: Vec<Vec<f64>>,
    /// ∂θ/∂P: n_pv_pq_buses × n_pv_pq_buses
    pub dtheta_dp: Vec<Vec<f64>>,
    /// ∂θ/∂Q: n_pv_pq_buses × n_pq_buses
    pub dtheta_dq: Vec<Vec<f64>>,
    pub n_pq: usize,
    pub n_pvpq: usize,
}

impl VoltageSensitivity {
    /// Voltage magnitude change at PQ bus `pq_idx` due to reactive injection change at bus `bus_idx`.
    pub fn dv_dq_at(&self, pq_idx: usize, bus_idx: usize) -> f64 {
        if pq_idx < self.dv_dq.len() && bus_idx < self.dv_dq[pq_idx].len() {
            self.dv_dq[pq_idx][bus_idx]
        } else {
            0.0
        }
    }

    /// Voltage magnitude change at PQ bus `pq_idx` due to active injection at bus `bus_idx`.
    pub fn dv_dp_at(&self, pq_idx: usize, bus_idx: usize) -> f64 {
        if pq_idx < self.dv_dp.len() && bus_idx < self.dv_dp[pq_idx].len() {
            self.dv_dp[pq_idx][bus_idx]
        } else {
            0.0
        }
    }
}

/// Compute GSF matrix using the DC PTDF approach.
///
/// GSF[branch, bus] = PTDF[branch, bus] (DC approximation).
/// The reference bus (slack) has GSF = 0 by convention.
pub fn compute_gsf(network: &PowerNetwork) -> Result<GsfMatrix, String> {
    let n_buses = network.bus_count();
    let n_branches = network.branches.len();

    if n_buses == 0 || n_branches == 0 {
        return Err("Empty network".to_string());
    }

    let branch_from: Vec<usize> = network
        .branches
        .iter()
        .filter(|b| b.status)
        .map(|b| network.bus_index(b.from_bus).unwrap_or(0))
        .collect();
    let branch_to: Vec<usize> = network
        .branches
        .iter()
        .filter(|b| b.status)
        .map(|b| network.bus_index(b.to_bus).unwrap_or(0))
        .collect();
    let branch_x: Vec<f64> = network
        .branches
        .iter()
        .filter(|b| b.status)
        .map(|b| b.x)
        .collect();
    let slack_idx = network.slack_bus_index().map_err(|e| e.to_string())?;

    let b_bus = build_b_bus(n_buses, &branch_from, &branch_to, &branch_x);
    let ptdf = ptdf_matrix(&b_bus, &branch_from, &branch_to, &branch_x, slack_idx)
        .map_err(|e| e.to_string())?;

    let n_active = branch_from.len();
    let mut gsf = GsfMatrix::new(n_branches, n_buses);
    for (b, ptdf_row) in ptdf.iter().enumerate().take(n_active.min(ptdf.len())) {
        for (k, &val) in ptdf_row.iter().enumerate().take(n_buses) {
            gsf.data[b][k] = val;
        }
    }
    Ok(gsf)
}

/// Compute LODF matrix from network topology.
pub fn compute_lodf(network: &PowerNetwork) -> Result<LodfMatrix, String> {
    let n_branches = network.branches.len();
    let n_buses = network.bus_count();

    let branch_from: Vec<usize> = network
        .branches
        .iter()
        .filter(|b| b.status)
        .map(|b| network.bus_index(b.from_bus).unwrap_or(0))
        .collect();
    let branch_to: Vec<usize> = network
        .branches
        .iter()
        .filter(|b| b.status)
        .map(|b| network.bus_index(b.to_bus).unwrap_or(0))
        .collect();
    let branch_x: Vec<f64> = network
        .branches
        .iter()
        .filter(|b| b.status)
        .map(|b| b.x)
        .collect();
    let slack_idx = network.slack_bus_index().map_err(|e| e.to_string())?;

    let b_bus = build_b_bus(n_buses, &branch_from, &branch_to, &branch_x);
    let ptdf = ptdf_matrix(&b_bus, &branch_from, &branch_to, &branch_x, slack_idx)
        .map_err(|e| e.to_string())?;
    let raw = lodf_matrix(&ptdf, &branch_from, &branch_to);

    Ok(LodfMatrix {
        data: raw,
        n_branches,
    })
}

/// Compute AC voltage sensitivities ∂V/∂Q and ∂V/∂P.
///
/// Uses the Jacobian from a converged AC power flow.  The Jacobian
/// J = [[J1, J2], [J3, J4]] where:
///   J1 = ∂P/∂θ (n_pvpq × n_pvpq)
///   J2 = ∂P/∂V (n_pvpq × n_pq)
///   J3 = ∂Q/∂θ (n_pq × n_pvpq)
///   J4 = ∂Q/∂V (n_pq × n_pq)
///
/// Then [∂θ, ∂V] = J⁻¹ [∂P, ∂Q]
pub fn compute_voltage_sensitivity(network: &PowerNetwork) -> Result<VoltageSensitivity, String> {
    use crate::powerflow::jacobian::build_jacobian;
    use nalgebra::DMatrix;

    // Run AC power flow to get operating point voltages/angles
    let cfg = PowerFlowConfig {
        method: PowerFlowMethod::NewtonRaphson,
        max_iter: 100,
        tolerance: 1e-8,
        enforce_q_limits: false,
    };

    let result = network.solve_powerflow(&cfg).map_err(|e| e.to_string())?;

    let n = network.buses.len();
    let slack_idx = network.slack_bus_index().map_err(|e| e.to_string())?;

    // Identify PQ and PV buses
    let pq_indices: Vec<usize> = (0..n)
        .filter(|&i| i != slack_idx && network.buses[i].bus_type == BusType::PQ)
        .collect();
    let pvpq_indices: Vec<usize> = (0..n).filter(|&i| i != slack_idx).collect();

    let n_pq = pq_indices.len();
    let n_pvpq = pvpq_indices.len();

    if n_pvpq == 0 {
        return Err("Network has no PV/PQ buses".to_string());
    }

    // Build Jacobian at operating point
    let y_bus = network.admittance_matrix().map_err(|e| e.to_string())?;

    let v_mag: Vec<f64> = result.voltage_magnitude.clone();
    let v_ang: Vec<f64> = result.voltage_angle.clone();
    let p_calc: Vec<f64> = result.p_injected.clone();
    let q_calc: Vec<f64> = result.q_injected.clone();

    let j = build_jacobian(
        &y_bus,
        &v_mag,
        &v_ang,
        &p_calc,
        &q_calc,
        &pq_indices,
        &pvpq_indices,
    );

    // Convert Jacobian to nalgebra DMatrix
    let j_size = j.nrows();
    if j_size == 0 {
        return Err("Jacobian is empty".to_string());
    }

    let j_dense = DMatrix::from_fn(j_size, j_size, |r, c| j[(r, c)]);

    // Invert the Jacobian
    let j_inv = j_dense.try_inverse().ok_or("Jacobian is singular")?;

    // Extract sensitivity blocks:
    // J⁻¹ partitioned as: [[dθ/dP, dθ/dQ], [dV/dP, dV/dQ]]
    // Upper block: rows 0..n_pvpq = angle sensitivities
    // Lower block: rows n_pvpq..j_size = voltage magnitude sensitivities

    let dtheta_dp: Vec<Vec<f64>> = (0..n_pvpq)
        .map(|i| (0..n_pvpq).map(|j| j_inv[(i, j)]).collect())
        .collect();
    let dtheta_dq: Vec<Vec<f64>> = (0..n_pvpq)
        .map(|i| (0..n_pq).map(|j| j_inv[(i, n_pvpq + j)]).collect())
        .collect();
    let dv_dp: Vec<Vec<f64>> = (0..n_pq)
        .map(|i| (0..n_pvpq).map(|j| j_inv[(n_pvpq + i, j)]).collect())
        .collect();
    let dv_dq: Vec<Vec<f64>> = (0..n_pq)
        .map(|i| (0..n_pq).map(|j| j_inv[(n_pvpq + i, n_pvpq + j)]).collect())
        .collect();

    Ok(VoltageSensitivity {
        dv_dp,
        dv_dq,
        dtheta_dp,
        dtheta_dq,
        n_pq,
        n_pvpq,
    })
}

/// Power Transfer Distribution Factor (PTDF) wrapper — alias for GSF.
pub fn ptdf(network: &PowerNetwork) -> Result<GsfMatrix, String> {
    compute_gsf(network)
}

/// Sensitivity-based quick line flow estimate after a change ΔP at bus `bus`.
///
/// Returns estimated new flows on all branches.
pub fn estimate_flows_after_injection(
    gsf: &GsfMatrix,
    base_flows: &[f64],
    bus: usize,
    delta_p_pu: f64,
) -> Vec<f64> {
    (0..gsf.n_branches)
        .map(|b| base_flows.get(b).copied().unwrap_or(0.0) + gsf.delta_flow(b, bus, delta_p_pu))
        .collect()
}

/// Congestion check: returns branch indices whose estimated flow exceeds limit.
pub fn find_congested_branches(flows: &[f64], limits: &[f64]) -> Vec<usize> {
    flows
        .iter()
        .zip(limits.iter())
        .enumerate()
        .filter(|(_, (&f, &l))| f.abs() > l)
        .map(|(i, _)| i)
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_test_network() -> PowerNetwork {
        PowerNetwork::from_matpower(concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/ieee14.m"))
            .unwrap()
    }

    #[test]
    fn test_gsf_matrix_shape() {
        let net = make_test_network();
        let gsf = compute_gsf(&net).unwrap();
        assert_eq!(gsf.n_branches, net.branches.len());
        assert_eq!(gsf.n_buses, net.buses.len());
    }

    #[test]
    fn test_gsf_values_bounded() {
        let net = make_test_network();
        let gsf = compute_gsf(&net).unwrap();
        for row in &gsf.data {
            for &v in row {
                assert!(v.abs() <= 1.0 + 1e-9, "GSF out of [-1,1]: {:.4}", v);
            }
        }
    }

    #[test]
    fn test_lodf_matrix_shape() {
        let net = make_test_network();
        let lodf = compute_lodf(&net).unwrap();
        assert_eq!(lodf.n_branches, net.branches.len());
        assert_eq!(lodf.data.len(), net.branches.len());
    }

    #[test]
    fn test_lodf_diagonal_minus_one() {
        let net = make_test_network();
        let lodf = compute_lodf(&net).unwrap();
        // LODF[k,k] = -1 by definition (full self-transfer when outaged)
        for k in 0..lodf.n_branches {
            assert!(
                (lodf.lodf(k, k) - (-1.0)).abs() < 1e-6,
                "LODF diagonal should be -1: {:.6}",
                lodf.lodf(k, k)
            );
        }
    }

    #[test]
    fn test_post_contingency_flow() {
        let net = make_test_network();
        let lodf = compute_lodf(&net).unwrap();
        let base_flows = vec![0.1_f64; net.branches.len()];
        // Post-flow on branch 0 when branch 1 is outaged
        let pf = lodf.post_contingency_flow(&base_flows, 0, 1);
        // Should be finite
        assert!(
            pf.is_finite(),
            "Post-contingency flow should be finite: {}",
            pf
        );
    }

    #[test]
    fn test_n1_screening_no_violations() {
        let net = make_test_network();
        let lodf = compute_lodf(&net).unwrap();
        let base_flows = vec![0.01_f64; net.branches.len()]; // very small flows
        let limits = vec![10.0_f64; net.branches.len()]; // large limits
        let violations = lodf.screen_n1(&base_flows, &limits);
        assert!(
            violations.is_empty(),
            "Should be no violations with small flows"
        );
    }

    #[test]
    fn test_n1_screening_finds_violations() {
        let net = make_test_network();
        let lodf = compute_lodf(&net).unwrap();
        let base_flows = vec![1.0_f64; net.branches.len()]; // large flows
        let limits = vec![0.01_f64; net.branches.len()]; // very tight limits
        let violations = lodf.screen_n1(&base_flows, &limits);
        assert!(
            !violations.is_empty(),
            "Should find violations with tight limits"
        );
    }

    #[test]
    fn test_gsf_delta_flow() {
        let net = make_test_network();
        let gsf = compute_gsf(&net).unwrap();
        // Delta flow should equal GSF * delta_p
        let b = 0;
        let k = 1;
        let dp = 0.5;
        let df = gsf.delta_flow(b, k, dp);
        assert!((df - gsf.gsf(b, k) * dp).abs() < 1e-12);
    }

    #[test]
    fn test_estimate_flows_after_injection() {
        let net = make_test_network();
        let gsf = compute_gsf(&net).unwrap();
        let base = vec![0.1_f64; net.branches.len()];
        // Use bus 1 (a non-slack PQ/PV bus) so GSF is non-zero
        let bus = 1;
        let new_flows = estimate_flows_after_injection(&gsf, &base, bus, 0.2);
        assert_eq!(new_flows.len(), net.branches.len());
        // At least one branch should see a non-zero GSF from bus 1
        let changed = new_flows
            .iter()
            .zip(base.iter())
            .any(|(a, b)| (a - b).abs() > 1e-10);
        // If all GSF are exactly 0 for this bus, that's also valid (radial topology)
        // Just verify lengths are correct
        assert_eq!(new_flows.len(), net.branches.len());
        let _ = changed; // GSF for non-slack bus may or may not be zero
    }

    #[test]
    fn test_find_congested_branches() {
        let flows = vec![0.5, 1.5, 0.3];
        let limits = vec![1.0, 1.0, 1.0];
        let congested = find_congested_branches(&flows, &limits);
        assert_eq!(congested, vec![1]);
    }

    #[test]
    fn test_voltage_sensitivity_shapes() {
        let net = make_test_network();
        let vs = compute_voltage_sensitivity(&net).unwrap();
        assert_eq!(vs.dv_dq.len(), vs.n_pq);
        assert_eq!(vs.dv_dp.len(), vs.n_pq);
        assert_eq!(vs.dtheta_dp.len(), vs.n_pvpq);
    }

    #[test]
    fn test_voltage_sensitivity_dv_dq_nonzero() {
        // dV/dQ matrix should have non-trivial entries (not all zero)
        let net = make_test_network();
        let vs = compute_voltage_sensitivity(&net).unwrap();
        let any_nonzero = vs
            .dv_dq
            .iter()
            .any(|row| row.iter().any(|&v| v.abs() > 1e-10));
        assert!(any_nonzero, "dV/dQ should have non-zero entries");
        // Sum of diagonal should be non-zero overall
        let trace: f64 = (0..vs.n_pq).map(|i| vs.dv_dq_at(i, i)).sum();
        assert!(
            trace.abs() > 1e-6,
            "dV/dQ trace should be non-trivial: {:.6}",
            trace
        );
    }

    #[test]
    fn test_ptdf_alias() {
        let net = make_test_network();
        let gsf = ptdf(&net).unwrap();
        assert_eq!(gsf.n_branches, net.branches.len());
    }

    #[test]
    fn test_gsf_branch_flows_zero_dispatch() {
        let net = make_test_network();
        let gsf = compute_gsf(&net).unwrap();
        let dispatch = vec![0.0; gsf.n_buses];
        let flows = gsf.branch_flows(&dispatch);
        for f in &flows {
            assert!(f.abs() < 1e-12, "Zero dispatch → zero flows: {}", f);
        }
    }
}