"""Tests for watkins_nn.flow module."""

import torch
import pytest


def test_flow_state_conservation():
    from watkins_nn.flow import FlowState
    state = FlowState(lam=0.6, kap=0.2, eta=0.2)
    assert state.conservation_error() < 1e-10
    assert state.is_valid()


def test_flow_state_invalid():
    from watkins_nn.flow import FlowState
    state = FlowState(lam=0.5, kap=0.3, eta=0.3)
    assert state.conservation_error() > 0.05


def test_flow_state_tensor_roundtrip():
    from watkins_nn.flow import FlowState
    state = FlowState(lam=0.6, kap=0.2, eta=0.2)
    t = state.to_tensor()
    state2 = FlowState.from_tensor(t)
    assert abs(state.lam - state2.lam) < 1e-6


def test_free_energy():
    from watkins_nn.flow import FlowState, free_energy
    state = FlowState(lam=0.6, kap=0.2, eta=0.2)
    F = free_energy(state)
    assert isinstance(F, float)
    assert F != 0.0  # F is finite at non-equilibrium


def test_gradient_F():
    from watkins_nn.flow import FlowState, gradient_F
    state = FlowState(lam=0.6, kap=0.2, eta=0.2)
    grad = gradient_F(state)
    assert len(grad) == 3


def test_enforce_conservation():
    from watkins_nn.flow import FlowState, enforce_conservation
    state = FlowState(lam=0.5, kap=0.3, eta=0.3)
    fixed = enforce_conservation(state)
    assert abs(fixed.lam + fixed.kap + fixed.eta - 1.0) < 1e-12


def test_flow_step():
    from watkins_nn.flow import FlowState, FlowConfig, flow_step
    state = FlowState(lam=0.4, kap=0.3, eta=0.3)
    config = FlowConfig(dt=0.001)
    new_state = flow_step(state, config)
    assert new_state.step == 1
    assert new_state.conservation_error() < 1e-9
    assert new_state.F is not None


def test_flow_convergence():
    from watkins_nn.flow import FlowState, FlowConfig, run_flow
    from watkins_nn.constants import LAM_STAR
    initial = FlowState(lam=0.4, kap=0.3, eta=0.3)
    config = FlowConfig(dt=0.001, max_steps=20000)
    final, _ = run_flow(initial, config)
    assert abs(final.lam - LAM_STAR) < 0.01
    assert final.conservation_error() < 1e-9


def test_flow_trajectory():
    from watkins_nn.flow import FlowState, FlowConfig, run_flow
    initial = FlowState(lam=0.5, kap=0.25, eta=0.25)
    config = FlowConfig(dt=0.001, max_steps=100, record_trajectory=True)
    final, trajectory = run_flow(initial, config)
    assert len(trajectory) > 1
    # F should decrease monotonically (approximately)
    F_values = [s.F for s in trajectory if s.F is not None]
    assert F_values[-1] <= F_values[0]


def test_distance_to_equilibrium():
    from watkins_nn.flow import FlowState, distance_to_equilibrium
    from watkins_nn.constants import LAM_STAR, KAP_STAR
    # At equilibrium, distance should be ~0
    eq_state = FlowState(lam=LAM_STAR, kap=KAP_STAR, eta=KAP_STAR)
    assert distance_to_equilibrium(eq_state) < 1e-10
    # Away from equilibrium
    far_state = FlowState(lam=0.4, kap=0.3, eta=0.3)
    assert distance_to_equilibrium(far_state) > 0.1


def test_watkins_flow_module():
    from watkins_nn.flow import WatkinsFlow
    flow = WatkinsFlow(num_steps=5, dt=0.01)
    p = torch.tensor([[0.4, 0.3, 0.3]], dtype=torch.float32)
    out = flow(p)
    assert out.shape == (1, 3)
    assert abs(float(out.sum()) - 1.0) < 1e-5
    # Lambda should increase toward equilibrium
    assert float(out[0, 0]) > 0.4


def test_watkins_flow_conservation():
    from watkins_nn.flow import WatkinsFlow
    flow = WatkinsFlow(num_steps=10, dt=0.01)
    p = torch.tensor([[0.5, 0.25, 0.25]], dtype=torch.float32)
    out = flow(p)
    assert flow.verify_conservation(out)
