// Lean 4 → LambdaIR → MiniC → C → Rust
// Module: HeytingLean.PadicDecoupling.Hamiltonian.EnergyConservation
// Theorem preserved: step_conserves_energy

use crate::tripartite::TripartiteHamiltonian;

/// A single time step redistributing energy across the tripartite channels.
pub struct TimeStep {
    pub before: TripartiteHamiltonian,
    pub after: TripartiteHamiltonian,
    pub transfer: f64,
    pub signal_fraction: f64,
}

/// Apply a time step: redistribute transfer across channels.
pub fn apply_time_step(before: TripartiteHamiltonian, transfer: f64, signal_fraction: f64) -> TripartiteHamiltonian {
    TripartiteHamiltonian {
        e_worker: before.e_worker - transfer,
        e_signal: before.e_signal + transfer * signal_fraction,
        e_interaction: before.e_interaction + transfer * (1.0 - signal_fraction),
    }
}

/// Theorem preserved: step_conserves_energy
/// ts.after.total = ts.before.total
pub fn check_step_conserves_energy(ts: &TimeStep) -> bool {
    (ts.after.total() - ts.before.total()).abs() < 1e-12
}

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

    #[test]
    fn conservation_after_apply() {
        let before = TripartiteHamiltonian {
            e_worker: 10.0,
            e_signal: 2.0,
            e_interaction: 3.0,
        };
        let after = apply_time_step(before, 4.0, 0.7);
        let ts = TimeStep { before, after, transfer: 4.0, signal_fraction: 0.7 };
        assert!(check_step_conserves_energy(&ts));
    }

    #[test]
    fn total_preserved() {
        let before = TripartiteHamiltonian {
            e_worker: 1.0,
            e_signal: 2.0,
            e_interaction: 3.0,
        };
        let after = apply_time_step(before, 0.5, 0.6);
        assert!((before.total() - after.total()).abs() < 1e-12);
    }
}