"""
End-to-end conservation tests: after a short run, total system mass
and total system energy must be conserved to machine precision. These
are the load-bearing tests for the "flux doubling" coupling mechanism.
"""
import numpy as np
import pytest
from tank_pipe.tank import Tank
from tank_pipe.pipe import Pipe
from tank_pipe.solver import run


GAMMA = 1.4
R_GAS = 287.0


def _build_scenario():
    """Default blowdown scenario at smaller scale — same physics, short run."""
    tank1 = Tank(V=5.0,  P_init=10e6,    T_init=300.0, gamma=GAMMA, R_gas=R_GAS)
    tank2 = Tank(V=10.0, P_init=101325.0, T_init=300.0, gamma=GAMMA, R_gas=R_GAS)
    pipe = Pipe(L=1.0, D=5e-3, N=20,
                P_init=101325.0, T_init=300.0, gamma=GAMMA, R_gas=R_GAS)
    return tank1, tank2, pipe


def _total_mass(tank1, tank2, pipe):
    return tank1.mass + tank2.mass + float(np.sum(pipe.W[0, :] * pipe.area * pipe.dx))


def _total_energy(tank1, tank2, pipe):
    return tank1.U + tank2.U + float(np.sum(pipe.W[2, :] * pipe.area * pipe.dx))


def test_total_mass_conserved_over_short_run():
    tank1, tank2, pipe = _build_scenario()
    m_init = _total_mass(tank1, tank2, pipe)
    run(tank1, tank2, pipe, t_end=1e-3, cfl=0.5)
    m_final = _total_mass(tank1, tank2, pipe)
    rel_err = abs(m_final - m_init) / m_init
    assert rel_err < 1e-10, (
        f"Total mass not conserved: m_init={m_init:.6e}, "
        f"m_final={m_final:.6e}, rel_err={rel_err:.2e}"
    )


def test_total_energy_conserved_over_short_run():
    tank1, tank2, pipe = _build_scenario()
    U_init = _total_energy(tank1, tank2, pipe)
    run(tank1, tank2, pipe, t_end=1e-3, cfl=0.5)
    U_final = _total_energy(tank1, tank2, pipe)
    rel_err = abs(U_final - U_init) / U_init
    assert rel_err < 1e-10, (
        f"Total energy not conserved: U_init={U_init:.6e}, "
        f"U_final={U_final:.6e}, rel_err={rel_err:.2e}"
    )