gam-sae 0.3.149

use crate::manifold::{
    GraphCompressionKind, GraphEdge, LearnedGraphAtom, OccupancyLaw, graph_edge_rank_charge,
};
use crate::sparse_dict::{BlockSparseFit, harmonic_measure_coordinates};
use ndarray::Array2;

const WEEKDAY_ANCHORS: usize = 7;
const ROWS_PER_ANCHOR: usize = 24;
const KEEP_MARGIN: f64 = 1.4;
const DROP_MARGIN: f64 = 0.4;

fn unit_circle_anchor_embeddings(anchors: usize) -> Array2<f64> {
    Array2::<f64>::from_shape_fn((anchors, 2), |(i, j)| {
        let phase = std::f64::consts::TAU * i as f64 / anchors as f64;
        if j == 0 { phase.cos() } else { phase.sin() }
    })
}

fn arc_anchor_embeddings(anchors: usize) -> Array2<f64> {
    Array2::<f64>::from_shape_fn((anchors, 1), |(i, _)| i as f64 / (anchors - 1) as f64)
}

fn star_anchor_embeddings(leaves: usize) -> Array2<f64> {
    let anchors = leaves + 1;
    Array2::<f64>::from_shape_fn((anchors, 2), |(i, j)| {
        if i == 0 {
            0.0
        } else {
            let phase = std::f64::consts::TAU * (i - 1) as f64 / leaves as f64;
            if j == 0 { phase.cos() } else { phase.sin() }
        }
    })
}

fn two_cycle_anchor_embeddings() -> Array2<f64> {
    Array2::<f64>::from_shape_fn((8, 2), |(i, j)| {
        let cycle = i / 4;
        let local = i % 4;
        let center = if cycle == 0 { -4.0 } else { 4.0 };
        let phase = std::f64::consts::TAU * local as f64 / 4.0;
        if j == 0 {
            center + phase.cos()
        } else {
            phase.sin()
        }
    })
}

fn continuous_circle_rows(anchors: usize, rows_per_anchor: usize) -> Vec<f64> {
    let n = anchors * rows_per_anchor;
    (0..n).map(|i| i as f64 / n as f64).collect()
}

fn weekday_rows(anchors: usize, rows_per_anchor: usize) -> Vec<f64> {
    let mut rows = Vec::with_capacity(anchors * rows_per_anchor);
    for anchor in 0..anchors {
        for repeat in 0..rows_per_anchor {
            let jitter_rank = repeat as f64 - (rows_per_anchor - 1) as f64 * 0.5;
            let jitter =
                jitter_rank / (rows_per_anchor as f64 * anchors as f64 * rows_per_anchor as f64);
            rows.push((anchor as f64 / anchors as f64 + jitter).rem_euclid(1.0));
        }
    }
    rows
}

fn edge(a: usize, b: usize) -> GraphEdge {
    GraphEdge::new(a, b).expect("valid test edge")
}

fn path_edges(anchors: usize) -> Vec<GraphEdge> {
    (0..anchors - 1).map(|i| edge(i, i + 1)).collect()
}

fn cycle_edges(offset: usize, anchors: usize) -> Vec<GraphEdge> {
    (0..anchors)
        .map(|i| edge(offset + i, offset + ((i + 1) % anchors)))
        .collect()
}

fn keep_all(edges: usize, charge: f64) -> (Vec<f64>, Vec<f64>) {
    (vec![1.0; edges], vec![charge * KEEP_MARGIN; edges])
}

fn keep_with_extra_retired(
    edges: &[GraphEdge],
    keep: &[GraphEdge],
    charge: f64,
) -> (Vec<f64>, Vec<f64>) {
    let precisions = vec![1.0; edges.len()];
    let deltas = edges
        .iter()
        .map(|edge| {
            if keep.contains(edge) {
                charge * KEEP_MARGIN
            } else {
                charge * DROP_MARGIN
            }
        })
        .collect();
    (precisions, deltas)
}

fn circ_err(a: f64, b: f64) -> f64 {
    let d = (a - b).abs();
    d.min(1.0 - d)
}

fn harmonic_code(spikes: &[(f64, f32)], h_count: usize) -> Vec<f32> {
    let mut code = Vec::with_capacity(2 * h_count);
    for h in 1..=h_count {
        let mut cos_sum = 0.0_f32;
        let mut sin_sum = 0.0_f32;
        for &(t, amp) in spikes {
            let phase = std::f64::consts::TAU * h as f64 * t;
            cos_sum += amp * phase.cos() as f32;
            sin_sum += amp * phase.sin() as f32;
        }
        code.push(cos_sum);
        code.push(sin_sum);
    }
    code
}

fn harmonic_fit_from_codes(rows: &[Vec<f32>], block_size: usize) -> BlockSparseFit {
    let n = rows.len();
    let mut x = Array2::<f32>::zeros((n, block_size));
    let mut codes = ndarray::Array3::<f32>::zeros((n, 1, block_size));
    for row in 0..n {
        for col in 0..block_size {
            x[[row, col]] = rows[row][col];
            codes[[row, 0, col]] = rows[row][col];
        }
    }
    BlockSparseFit {
        decoder: Array2::<f32>::eye(block_size),
        blocks: Array2::<u32>::zeros((n, 1)),
        gates: Array2::<f32>::ones((n, 1)),
        codes,
        gamma: 1.0,
        block_utilization: vec![1.0],
        block_stable_rank: vec![2.0],
        matryoshka_prefix_losses: Vec::new(),
        explained_variance: 1.0,
        epochs: 0,
        convergence: crate::sparse_dict::BlockSparseConvergence::trivially_converged(),
        block_topk: 1,
        block_size,
    }
}

#[test]
fn graph_atom_reads_continuous_circle_as_one_loop_from_knn_edges() {
    let anchors = WEEKDAY_ANCHORS;
    let rows = continuous_circle_rows(anchors, ROWS_PER_ANCHOR);
    let embeddings = unit_circle_anchor_embeddings(anchors);
    let edges = LearnedGraphAtom::knn_candidate_edges(embeddings.view()).expect("knn edges");
    let n_eff = rows.len() as f64;
    let charge = graph_edge_rank_charge(n_eff, embeddings.ncols());
    let (precisions, deltas) = keep_all(edges.len(), charge);

    let atom = LearnedGraphAtom::from_reml_candidate_edges(
        embeddings.view(),
        &rows,
        n_eff,
        &edges,
        &precisions,
        &deltas,
    )
    .expect("continuous circle graph atom");
    let readout = atom.topology_readout();
    let smoothness = atom.surviving_smoothness_value();
    let selection = atom.structure_selection();

    assert_eq!(readout.b0, 1);
    assert_eq!(readout.b1, 1);
    assert_eq!(readout.surviving_edges, anchors);
    assert_eq!(selection.compression.kind, GraphCompressionKind::Circle);
    assert!(selection.compression.earns_standard_name());
    assert!(selection.compression.bits_saved > 0.0);
    assert!(selection.selected);
    assert!(
        matches!(
            atom.occupancy(),
            OccupancyLaw::Uniform | OccupancyLaw::Continuous
        ),
        "continuous circle occupancy should remain continuous/uniform, got {:?}",
        atom.occupancy()
    );
    assert!(smoothness > 0.0);
}

#[test]
fn graph_atom_reads_weekdays_as_atomic_cycle_without_fixed_menu_selection() {
    let anchors = WEEKDAY_ANCHORS;
    let rows = weekday_rows(anchors, ROWS_PER_ANCHOR);
    let embeddings = unit_circle_anchor_embeddings(anchors);
    let edges = LearnedGraphAtom::knn_candidate_edges(embeddings.view()).expect("knn edges");
    let n_eff = rows.len() as f64;
    let charge = graph_edge_rank_charge(n_eff, embeddings.ncols());
    let (precisions, deltas) = keep_all(edges.len(), charge);

    let atom = LearnedGraphAtom::from_reml_knn_edges(
        embeddings.view(),
        &rows,
        n_eff,
        &precisions,
        &deltas,
    )
    .expect("weekday graph atom");
    let readout = atom.topology_readout();
    let laplacian = atom.surviving_laplacian();
    let selection = atom.structure_selection();

    assert_eq!(readout.b0, 1);
    assert_eq!(readout.b1, 1);
    assert_eq!(readout.surviving_edges, anchors);
    assert_eq!(atom.occupancy(), OccupancyLaw::Discrete { anchors });
    assert_eq!(selection.compression.kind, GraphCompressionKind::Circle);
    assert!(selection.total_edge_charge > 0.0);
    assert!(selection.margin > 0.0);
    assert_eq!(laplacian.dim(), (anchors, anchors));
    for row in 0..anchors {
        assert!((laplacian.row(row).sum()).abs() < 1e-12);
        assert_eq!(laplacian[[row, row]], 2.0);
    }

    let offset = 3usize;
    let beta_dim = offset + anchors * atom.fiber_rank();
    let op = atom.surviving_penalty_op(offset, beta_dim);
    assert_eq!(op.dim(), beta_dim);
    assert_eq!(op.output_range(), Some(offset..beta_dim));

    let mut beta = vec![0.0; beta_dim];
    for anchor in 0..anchors {
        for channel in 0..atom.fiber_rank() {
            beta[offset + anchor * atom.fiber_rank() + channel] = embeddings[[anchor, channel]];
        }
    }
    let mut h_beta = vec![0.0; beta_dim];
    op.matvec(&beta, &mut h_beta);
    for row in 0..anchors {
        for channel in 0..atom.fiber_rank() {
            let mut expected = 0.0;
            for col in 0..anchors {
                expected += laplacian[[row, col]] * embeddings[[col, channel]];
            }
            let idx = offset + row * atom.fiber_rank() + channel;
            assert!((h_beta[idx] - expected).abs() < 1e-12);
        }
    }
}

#[test]
fn learned_graph_reads_path_as_interval() {
    let anchors = WEEKDAY_ANCHORS;
    let rows = continuous_circle_rows(anchors - 1, ROWS_PER_ANCHOR);
    let embeddings = arc_anchor_embeddings(anchors);
    let edges = path_edges(anchors);
    let n_eff = rows.len() as f64;
    let charge = graph_edge_rank_charge(n_eff, embeddings.ncols());
    let (precisions, deltas) = keep_all(edges.len(), charge);

    let atom = LearnedGraphAtom::from_reml_candidate_edges(
        embeddings.view(),
        &rows,
        n_eff,
        &edges,
        &precisions,
        &deltas,
    )
    .expect("path graph atom");
    let readout = atom.topology_readout();

    assert_eq!(readout.b0, 1);
    assert_eq!(readout.b1, 0);
    assert_eq!(readout.surviving_edges, anchors - 1);
    assert_eq!(
        atom.certified_compression().kind,
        GraphCompressionKind::Interval
    );
}

#[test]
fn learned_graph_reads_two_disconnected_cycles() {
    let rows = weekday_rows(8, ROWS_PER_ANCHOR);
    let embeddings = two_cycle_anchor_embeddings();
    let mut edges = cycle_edges(0, 4);
    edges.extend(cycle_edges(4, 4));
    let n_eff = rows.len() as f64;
    let charge = graph_edge_rank_charge(n_eff, embeddings.ncols());
    let (precisions, deltas) = keep_all(edges.len(), charge);

    let atom = LearnedGraphAtom::from_reml_candidate_edges(
        embeddings.view(),
        &rows,
        n_eff,
        &edges,
        &precisions,
        &deltas,
    )
    .expect("two-cycle graph atom");
    let readout = atom.topology_readout();

    assert_eq!(readout.b0, 2);
    assert_eq!(readout.b1, 2);
    assert_eq!(
        atom.certified_compression().kind,
        GraphCompressionKind::Graph
    );
    assert_eq!(
        atom.certified_compression().name,
        "structure without a standard name"
    );
}

#[test]
fn non_uniform_cycle_reports_no_standard_name() {
    let anchors = WEEKDAY_ANCHORS;
    let rows = continuous_circle_rows(anchors, ROWS_PER_ANCHOR);
    let embeddings = unit_circle_anchor_embeddings(anchors);
    let edges = cycle_edges(0, anchors);
    let n_eff = rows.len() as f64;
    let charge = graph_edge_rank_charge(n_eff, embeddings.ncols());
    let mut precisions = vec![1.0; edges.len()];
    precisions[0] = 2.0;
    let deltas = vec![charge * KEEP_MARGIN; edges.len()];

    let atom = LearnedGraphAtom::from_reml_candidate_edges(
        embeddings.view(),
        &rows,
        n_eff,
        &edges,
        &precisions,
        &deltas,
    )
    .expect("non-uniform cycle graph atom");
    let readout = atom.topology_readout();
    let compression = atom.certified_compression();

    assert_eq!(readout.b0, 1);
    assert_eq!(readout.b1, 1);
    assert_eq!(compression.kind, GraphCompressionKind::Graph);
    assert_eq!(compression.name, "structure without a standard name");
    assert_eq!(compression.bits_saved, 0.0);
}

#[test]
fn learned_graph_reads_branching_tree_and_detects_branch_vertex() {
    let leaves = 5usize;
    let embeddings = star_anchor_embeddings(leaves);
    let rows = continuous_circle_rows(leaves + 1, ROWS_PER_ANCHOR);
    let keep = (1..=leaves).map(|i| edge(0, i)).collect::<Vec<_>>();
    let mut edges = keep.clone();
    edges.extend((1..leaves).map(|i| edge(i, i + 1)));
    let n_eff = rows.len() as f64;
    let charge = graph_edge_rank_charge(n_eff, embeddings.ncols());
    let (precisions, deltas) = keep_with_extra_retired(&edges, &keep, charge);

    let atom = LearnedGraphAtom::from_reml_candidate_edges(
        embeddings.view(),
        &rows,
        n_eff,
        &edges,
        &precisions,
        &deltas,
    )
    .expect("branching tree graph atom");
    let readout = atom.topology_readout();
    let degrees = atom.surviving_degrees();

    assert_eq!(readout.b0, 1);
    assert_eq!(readout.b1, 0);
    assert!(
        degrees.iter().any(|&degree| degree > 2),
        "branching tree must expose a degree>2 vertex: {degrees:?}"
    );
    assert_eq!(
        atom.certified_compression().kind,
        GraphCompressionKind::Graph
    );
    assert_eq!(
        atom.certified_compression().name,
        "structure without a standard name"
    );
}

#[test]
fn two_date_modular_synthetic_binds_super_resolution_to_graph_base() {
    let anchors = WEEKDAY_ANCHORS;
    let rows = weekday_rows(anchors, ROWS_PER_ANCHOR);
    let embeddings = unit_circle_anchor_embeddings(anchors);
    let edges = LearnedGraphAtom::knn_candidate_edges(embeddings.view()).expect("knn edges");
    let n_eff = rows.len() as f64;
    let charge = graph_edge_rank_charge(n_eff, embeddings.ncols());
    let (precisions, deltas) = keep_all(edges.len(), charge);
    let atom = LearnedGraphAtom::from_reml_candidate_edges(
        embeddings.view(),
        &rows,
        n_eff,
        &edges,
        &precisions,
        &deltas,
    )
    .expect("weekday graph base");

    let selection = atom.structure_selection();
    assert!(
        selection.selected,
        "graph base must be selected before binding"
    );
    assert_eq!(selection.topology.b0, 1);
    assert_eq!(selection.topology.b1, 1);
    assert_eq!(selection.occupancy, OccupancyLaw::Discrete { anchors });

    let h_count = 8usize;
    let first_date = 1.0 / anchors as f64;
    let second_date = 4.0 / anchors as f64;
    assert!(
        circ_err(first_date, second_date) > crate::super_resolution::separation_limit(h_count),
        "two modular dates must clear the Prony separation guarantee"
    );
    let code_rows = vec![
        harmonic_code(&[(first_date, 1.0), (second_date, 0.75)], h_count),
        harmonic_code(&[(2.0 / anchors as f64, 0.9)], h_count),
    ];
    let fit = harmonic_fit_from_codes(&code_rows, 2 * h_count);
    let report = harmonic_measure_coordinates(&fit, 0).expect("measure readout");
    let row0 = report
        .firings
        .iter()
        .find(|firing| firing.row == 0)
        .expect("row 0 measure");

    assert!(
        row0.used_super_resolution,
        "two-date graph-base code must invoke Prony/matrix-pencil recovery"
    );
    assert_eq!(row0.spikes.len(), 2);
    assert!(
        row0.spikes
            .iter()
            .any(|spike| circ_err(spike.coordinate, first_date) < 1.0e-8),
        "missing first modular date"
    );
    assert!(
        row0.spikes
            .iter()
            .any(|spike| circ_err(spike.coordinate, second_date) < 1.0e-8),
        "missing second modular date"
    );
}

/// E1 graph-atom ENROLLMENT round-trip (#985): a discovered co-fire graph is
/// enrolled as a first-class graph atom and its Betti numbers read back.
///
/// Plant a ring of `V` atoms whose adjacent members co-fire (each row activates
/// one adjacent pair `(i, i+1 mod V)`), so the ONLY co-firing pairs are the `V`
/// ring edges. `coactive_pair_stats` discovers them; `enroll_from_coactivation`
/// enrolls the co-fire graph, the built-in ARD survival rule keeps the edges
/// whose co-fire evidence beats the one-edge charge, and `topology_readout`
/// reads the cycle Betti signature `(b0, b1) = (1, 1)`. Enroll → fit → read Betti,
/// end to end.
#[test]
fn coactivation_ring_enrolls_as_graph_atom_with_cycle_betti() {
    use crate::atom_codes::SparseAtomCodes;

    let vertices = 6usize;
    let reps = 120usize; // rows per ring edge
    let n = vertices * reps;
    let mut codes = SparseAtomCodes::empty(n, vertices);
    let mut row_coords = Vec::with_capacity(n);
    let mut row = 0usize;
    for edge in 0..vertices {
        let a = edge;
        let b = (edge + 1) % vertices;
        for _ in 0..reps {
            codes.row_mut(row).assign(a, 1.0);
            codes.row_mut(row).assign(b, 1.0);
            row_coords.push(edge as f64);
            row += 1;
        }
    }

    // Discovery: the sparse co-fire pass returns exactly the V ring edges (no
    // non-adjacent atoms ever co-fire), each at symmetric dependence 0.5.
    let pairs = codes.coactive_pair_stats();
    assert_eq!(pairs.len(), vertices, "only the ring edges co-fire");
    for (a, b, stats) in &pairs {
        assert!(
            (stats.dependence() - 0.5).abs() < 1e-9,
            "ring edge ({a},{b}) dependence {} != 0.5",
            stats.dependence()
        );
    }

    // Enrollment: vertices are atoms, candidate edges are the co-fire pairs above
    // the dependence floor. Vertices sit on a circle so the fiber embedding is
    // rank 2 and the surviving edge weights are uniform (a named circle).
    let embeddings = unit_circle_anchor_embeddings(vertices);
    let n_eff = n as f64;
    let atom = LearnedGraphAtom::enroll_from_coactivation(
        embeddings.view(),
        &row_coords,
        n_eff,
        &pairs,
        0.3, // keeps ring edges (0.5); would drop any incidental weak co-fire
    )
    .expect("ring co-fire graph enrolls as a graph atom");

    // Fit is the ARD survival applied inside the constructor; read Betti back.
    let topo = atom.topology_readout();
    assert_eq!(topo.vertices, vertices);
    assert_eq!(
        topo.surviving_edges, vertices,
        "the whole ring must clear the one-edge charge"
    );
    assert_eq!(topo.b0, 1, "a ring is connected: b0 = 1");
    assert_eq!(
        topo.b1, 1,
        "a ring has exactly one independent cycle: b1 = 1"
    );

    // The enrolled structure is SELECTED: summed co-fire evidence beats the summed
    // edge charge, and every edge's evidence is n_eff·MI over its one-edge charge.
    let selection = atom.structure_selection();
    assert!(
        selection.selected,
        "the enrolled ring must be selected (margin {})",
        selection.margin
    );
    let charge = graph_edge_rank_charge(n_eff, atom.fiber_rank());
    for (a, b, stats) in &pairs {
        let (_, delta) = crate::manifold::coactivation_edge_evidence(stats, n_eff);
        assert!(
            delta > charge,
            "ring edge ({a},{b}) evidence {delta} must exceed the one-edge charge {charge}"
        );
    }

    // A uniform-weight degree-2 ring compresses to the named circle.
    assert!(
        matches!(
            atom.certified_compression().kind,
            GraphCompressionKind::Circle
        ),
        "the enrolled ring compresses to the named circle"
    );
}