//! Mass-flow topology presolve for the system graph (CM1.4).
//!
//! This module identifies **series branches** — maximal sequences of flow edges
//! where every intermediate node has exactly one incoming and one outgoing edge
//! (no junction). All edges in such a branch share a single mass-flow Newton
//! unknown, reducing the state-vector size from `3|E|` to `|B| + 2|E|`.
//!
//! ## Algorithm
//!
//! 1. Iterate over all graph edges in petgraph iteration order.
//! 2. For each unvisited edge, trace the maximal series branch by walking
//!    forward (following the target node's single outgoing edge) and backward
//!    (following the source node's single incoming edge) while nodes have
//!    in-degree == 1 and out-degree == 1.
//! 3. Assign the same `mass_flow_branch_id` to every edge in the branch.
//! 4. Return the total branch count (`|B|`).
//!
//! ## Design reference
//!
//! Based on TESPy `presolve_massflow_topology()` in
//! `src/tespy/networks/network.py` (lines 983–1177).
//! See `complex-machine-design.md §1.6` and `epics-complex-machine.md` Story 1.4.

use entropyk_components::Component;
use petgraph::graph::{EdgeIndex, Graph};
use petgraph::visit::EdgeRef;
use petgraph::Directed;
use std::collections::HashSet;

use crate::system::FlowEdge;

/// Assigns `mass_flow_branch_id` to every edge in `graph` by grouping edges
/// into series branches (maximal paths with no junction node).
///
/// Returns the number of independent branches `|B|`. Every edge in the same
/// branch receives the same `mass_flow_branch_id`; edges in different branches
/// receive distinct ids.
///
/// # Series branch definition
///
/// Two adjacent edges belong to the same branch when their shared node has
/// **in-degree == 1 and out-degree == 1** in the directed graph. Nodes with
/// more than one incoming or outgoing edge are junction boundaries that start
/// a new branch.
///
/// For a pure series directed cycle (every node in-degree=1, out-degree=1), all
/// edges form one branch. The walk terminates when it revisits a branch-member
/// edge (cycle guard).
pub fn presolve_mass_flow_topology(
    graph: &mut Graph<Box<dyn Component>, FlowEdge, Directed>,
) -> usize {
    let mut visited: HashSet<EdgeIndex> = HashSet::new();
    let mut branch_id: usize = 0;

    // Collect edge indices first to avoid borrow conflict on the mutable graph.
    let edge_indices: Vec<EdgeIndex> = graph.edge_indices().collect();

    for start_edge in edge_indices {
        if visited.contains(&start_edge) {
            continue;
        }

        // Trace the maximal series branch containing start_edge.
        let branch = trace_series_branch(graph, start_edge, &visited);

        // Assign branch_id and mark all branch edges as visited.
        for &eid in &branch {
            if let Some(weight) = graph.edge_weight_mut(eid) {
                weight.mass_flow_branch_id = branch_id;
            }
            visited.insert(eid);
        }

        branch_id += 1;
    }

    branch_id
}

/// Traces the maximal series branch containing `start_edge`.
///
/// Walks forward from the target node and backward from the source node of
/// `start_edge`, collecting all edges reachable through 1-in-1-out nodes **or
/// through declared internal flow paths** ([`Component::flow_paths`]) of
/// multi-port components: a Modelica-style 4-port heat exchanger declares
/// `[(0, 1), (2, 3)]`, so the refrigerant inlet/outlet edges (ports 0→1) and
/// the secondary inlet/outlet edges (ports 2→3) each form one continuous
/// series branch, exactly like `port_a.m_flow + port_b.m_flow = 0`.
///
/// Stops when a genuine junction node (splitter/merger/drum without a matching
/// flow path) or a previously-visited/branch edge is encountered.
fn trace_series_branch(
    graph: &Graph<Box<dyn Component>, FlowEdge, Directed>,
    start_edge: EdgeIndex,
    visited: &HashSet<EdgeIndex>,
) -> Vec<EdgeIndex> {
    let mut branch: Vec<EdgeIndex> = vec![start_edge];

    let (src, tgt) = graph
        .edge_endpoints(start_edge)
        .expect("start_edge must be a valid edge");

    // ── Walk forward from the target node ──────────────────────────────────
    let mut cur = tgt;
    let mut cur_in = start_edge;
    loop {
        let next_eid = match forward_continuation(graph, cur, cur_in) {
            Some(e) => e,
            None => break,
        };

        // Cycle guard and already-visited guard.
        if branch.contains(&next_eid) || visited.contains(&next_eid) {
            break;
        }

        branch.push(next_eid);
        cur = graph
            .edge_endpoints(next_eid)
            .expect("next_eid must be valid")
            .1;
        cur_in = next_eid;
    }

    // ── Walk backward from the source node ────────────────────────────────
    let mut cur = src;
    let mut cur_out = start_edge;
    loop {
        let prev_eid = match backward_continuation(graph, cur, cur_out) {
            Some(e) => e,
            None => break,
        };

        if branch.contains(&prev_eid) || visited.contains(&prev_eid) {
            break;
        }

        branch.push(prev_eid);
        cur = graph
            .edge_endpoints(prev_eid)
            .expect("prev_eid must be valid")
            .0;
        cur_out = prev_eid;
    }

    branch
}

/// Returns the unique outgoing edge continuing the series branch through
/// `node`, entered via `in_edge` — either the plain 1-in-1-out rule or a
/// declared internal flow path matching the entry port.
fn forward_continuation(
    graph: &Graph<Box<dyn Component>, FlowEdge, Directed>,
    node: petgraph::graph::NodeIndex,
    in_edge: EdgeIndex,
) -> Option<EdgeIndex> {
    let in_edges: Vec<_> = graph
        .edges_directed(node, petgraph::Direction::Incoming)
        .collect();
    let out_edges: Vec<_> = graph
        .edges_directed(node, petgraph::Direction::Outgoing)
        .collect();

    // Plain series node.
    if in_edges.len() == 1 && out_edges.len() == 1 {
        return Some(out_edges[0].id());
    }

    // Multi-port component with declared internal flow paths (Modelica-style):
    // continue through the path whose inlet port matches the entry port.
    let entry_port = graph.edge_weight(in_edge).map(|w| w.target_port)?;
    let paths = graph.node_weight(node)?.flow_paths();
    let out_port = paths
        .iter()
        .find(|(p_in, _)| *p_in == entry_port)
        .map(|(_, p_out)| *p_out)?;
    let matching: Vec<EdgeIndex> = out_edges
        .iter()
        .filter(|e| e.weight().source_port == out_port)
        .map(|e| e.id())
        .collect();
    match matching.as_slice() {
        [only] => Some(*only),
        _ => None,
    }
}

/// Backward analogue of [`forward_continuation`]: unique incoming edge whose
/// target port feeds the flow path that exits through `out_edge`.
fn backward_continuation(
    graph: &Graph<Box<dyn Component>, FlowEdge, Directed>,
    node: petgraph::graph::NodeIndex,
    out_edge: EdgeIndex,
) -> Option<EdgeIndex> {
    let in_edges: Vec<_> = graph
        .edges_directed(node, petgraph::Direction::Incoming)
        .collect();
    let out_edges: Vec<_> = graph
        .edges_directed(node, petgraph::Direction::Outgoing)
        .collect();

    if in_edges.len() == 1 && out_edges.len() == 1 {
        return Some(in_edges[0].id());
    }

    let exit_port = graph.edge_weight(out_edge).map(|w| w.source_port)?;
    let paths = graph.node_weight(node)?.flow_paths();
    let in_port = paths
        .iter()
        .find(|(_, p_out)| *p_out == exit_port)
        .map(|(p_in, _)| *p_in)?;
    let matching: Vec<EdgeIndex> = in_edges
        .iter()
        .filter(|e| e.weight().target_port == in_port)
        .map(|e| e.id())
        .collect();
    match matching.as_slice() {
        [only] => Some(*only),
        _ => None,
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Unit tests
// ─────────────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use entropyk_components::{ComponentError, JacobianBuilder, ResidualVector, StateSlice};
    use petgraph::graph::Graph;

    // ── Minimal mock component for topology tests ─────────────────────────

    struct MockComponent;

    impl entropyk_components::Component for MockComponent {
        fn n_equations(&self) -> usize {
            2
        }

        fn compute_residuals(
            &self,
            _state: &StateSlice,
            _residuals: &mut ResidualVector,
        ) -> Result<(), ComponentError> {
            Ok(())
        }

        fn jacobian_entries(
            &self,
            _state: &StateSlice,
            _jacobian: &mut JacobianBuilder,
        ) -> Result<(), ComponentError> {
            Ok(())
        }

        fn get_ports(&self) -> &[entropyk_components::ConnectedPort] {
            &[]
        }

        fn signature(&self) -> String {
            "mock".to_string()
        }
    }

    /// Build a directed graph with `n_nodes` arranged as a directed cycle:
    /// node 0 → node 1 → ... → node n-1 → node 0.
    /// Returns (graph, edge_indices_in_order).
    fn build_series_cycle(
        n: usize,
    ) -> (
        Graph<Box<dyn Component>, FlowEdge, Directed>,
        Vec<EdgeIndex>,
    ) {
        let mut g: Graph<Box<dyn Component>, FlowEdge, Directed> = Graph::new();
        let nodes: Vec<_> = (0..n)
            .map(|_| g.add_node(Box::new(MockComponent) as Box<dyn Component>))
            .collect();
        let mut edges = Vec::new();
        for i in 0..n {
            let src = nodes[i];
            let tgt = nodes[(i + 1) % n];
            edges.push(g.add_edge(src, tgt, FlowEdge::new_unassigned()));
        }
        (g, edges)
    }

    /// Build a graph that represents a main loop (4 edges) plus a bypass
    /// branch: node 0 → (splitter) node 1 → two paths → (merger) node 4 → node 0.
    /// Topology: 0→1→2→4→0 (main) and 1→3→4 (bypass), so node 1 has out-degree 2
    /// and node 4 has in-degree 2.
    fn build_splitter_merger_topology() -> (
        Graph<Box<dyn Component>, FlowEdge, Directed>,
        Vec<EdgeIndex>,
    ) {
        // Nodes: 0 (source/sink), 1 (splitter), 2 (main path), 3 (bypass path), 4 (merger)
        let mut g: Graph<Box<dyn Component>, FlowEdge, Directed> = Graph::new();
        let n: Vec<_> = (0..5)
            .map(|_| g.add_node(Box::new(MockComponent) as Box<dyn Component>))
            .collect();

        // Main series edges: 0→1, 2→4, 4→0 (these are outside the junction nodes)
        let e01 = g.add_edge(n[0], n[1], FlowEdge::new_unassigned()); // branch A
        let e12 = g.add_edge(n[1], n[2], FlowEdge::new_unassigned()); // branch B (main)
        let e24 = g.add_edge(n[2], n[4], FlowEdge::new_unassigned()); // branch B (main)
        let e13 = g.add_edge(n[1], n[3], FlowEdge::new_unassigned()); // branch C (bypass)
        let e34 = g.add_edge(n[3], n[4], FlowEdge::new_unassigned()); // branch C (bypass)
        let e40 = g.add_edge(n[4], n[0], FlowEdge::new_unassigned()); // branch A

        (g, vec![e01, e12, e24, e13, e34, e40])
    }

    // ── Test: 4-edge series cycle → 1 branch ─────────────────────────────

    #[test]
    fn test_series_cycle_4_edges_one_branch() {
        let (mut g, edges) = build_series_cycle(4);

        let n_branches = presolve_mass_flow_topology(&mut g);

        assert_eq!(
            n_branches, 1,
            "pure series cycle must yield exactly 1 branch"
        );

        // All 4 edges must share branch_id == 0.
        let branch_ids: Vec<usize> = edges
            .iter()
            .map(|&e| g.edge_weight(e).unwrap().mass_flow_branch_id)
            .collect();
        assert!(
            branch_ids.iter().all(|&id| id == 0),
            "all edges must be in branch 0, got {:?}",
            branch_ids
        );
    }

    // ── Test: 2-edge series → 1 branch ───────────────────────────────────

    #[test]
    fn test_series_cycle_2_edges_one_branch() {
        let (mut g, edges) = build_series_cycle(2);
        let n_branches = presolve_mass_flow_topology(&mut g);
        assert_eq!(n_branches, 1);
        let ids: Vec<usize> = edges
            .iter()
            .map(|&e| g.edge_weight(e).unwrap().mass_flow_branch_id)
            .collect();
        assert!(ids.iter().all(|&id| id == 0));
    }

    // ── Test: single-node self-loop (degenerate) — 1 branch ──────────────

    #[test]
    fn test_single_edge_one_branch() {
        let (mut g, edges) = build_series_cycle(1);
        let n_branches = presolve_mass_flow_topology(&mut g);
        assert_eq!(n_branches, 1);
        assert_eq!(g.edge_weight(edges[0]).unwrap().mass_flow_branch_id, 0);
    }

    // ── Test: splitter-merger topology → 3 branches ───────────────────────

    #[test]
    fn test_splitter_merger_three_branches() {
        // Topology:
        //   node 0 →[e01]→ node 1 (splitter, out-degree 2)
        //   node 1 →[e12]→ node 2 →[e24]→ node 4 (merger, in-degree 2)
        //   node 1 →[e13]→ node 3 →[e34]→ node 4
        //   node 4 →[e40]→ node 0
        //
        // Expected branches:
        //   Branch A: e01, e40  (series path from node 0 through junction node 1 back via node 4)
        //   Wait — node 0 has in-degree 1 (from e40) and out-degree 1 (to e01), so it IS a 1-in-1-out node.
        //   node 1 has out-degree 2 → junction boundary.
        //   node 4 has in-degree 2 → junction boundary.
        //
        //   So tracing from e01:
        //     start: e01 (0→1)
        //     forward from node 1: out-degree=2 → STOP
        //     backward from node 0: in-edges = [e40], out-edges = [e01] → 1-in-1-out → add e40, move to node 4
        //     node 4: in-degree=2 → STOP
        //   Branch A: {e01, e40}
        //
        //   From e12 (unvisited):
        //     forward from node 2: in-degree=1, out-degree=1 → add e24, move to node 4; node 4 in-degree=2 → STOP
        //     backward from node 1: out-degree=2 → STOP
        //   Branch B: {e12, e24}
        //
        //   From e13 (unvisited):
        //     forward from node 3: in-degree=1, out-degree=1 → add e34, move to node 4; node 4 in-degree=2 → STOP
        //     backward from node 1: out-degree=2 → STOP
        //   Branch C: {e13, e34}
        //
        //   Total: 3 branches ✓

        let (mut g, edges) = build_splitter_merger_topology();
        let [e01, e12, e24, e13, e34, e40] = edges.as_slice() else {
            panic!()
        };

        let n_branches = presolve_mass_flow_topology(&mut g);

        assert_eq!(n_branches, 3, "splitter-merger must yield 3 branches");

        let id = |e: &EdgeIndex| g.edge_weight(*e).unwrap().mass_flow_branch_id;

        // e01 and e40 must share a branch.
        assert_eq!(id(e01), id(e40), "e01 and e40 must share branch");

        // e12 and e24 must share a branch (main path through node 2).
        assert_eq!(id(e12), id(e24), "e12 and e24 must share branch");

        // e13 and e34 must share a branch (bypass through node 3).
        assert_eq!(id(e13), id(e34), "e13 and e34 must share branch");

        // All three pairs must be in DISTINCT branches.
        let a = id(e01);
        let b = id(e12);
        let c = id(e13);
        assert_ne!(a, b, "branch A and B must differ");
        assert_ne!(a, c, "branch A and C must differ");
        assert_ne!(b, c, "branch B and C must differ");
    }

    // ── Test: no double-assignment on a 6-edge series ─────────────────────

    #[test]
    fn test_no_double_assignment_large_cycle() {
        let (mut g, edges) = build_series_cycle(6);
        let n_branches = presolve_mass_flow_topology(&mut g);
        assert_eq!(n_branches, 1, "6-edge series cycle is still 1 branch");
        for &e in &edges {
            assert_eq!(g.edge_weight(e).unwrap().mass_flow_branch_id, 0);
        }
    }
}
