"""Tests for optimization components (CER, PPGD)."""

import numpy as np
import pytest

from tqce.network.node import CausalNode
from tqce.network.ctn import CausalTensorNetwork
from tqce.optimization.regularizers import ChronoEntropicRegularizer
from tqce.optimization.ppgd import ParadoxPressureGradientDescent


class TestChronoEntropicRegularizer:
    def _make_loop(self) -> CausalTensorNetwork:
        net = CausalTensorNetwork()
        net.add_node(CausalNode("A", time_coordinate=0.0, state_dim=4))
        net.add_node(CausalNode("B", time_coordinate=1.0, state_dim=4))
        net.add_causal_edge("A", "B")
        net.add_retrocausal_edge("B", "A", weight=0.5)
        return net

    def test_penalty_nonnegative(self):
        net = self._make_loop()
        cer = ChronoEntropicRegularizer(net)
        penalty = cer.compute(net.nodes["A"])
        assert penalty >= 0.0

    def test_gradient_shape(self):
        net = self._make_loop()
        cer = ChronoEntropicRegularizer(net)
        grad = cer.gradient(net.nodes["A"])
        assert grad.shape == (4,)

    def test_zero_state_no_crash(self):
        net = self._make_loop()
        net.nodes["A"].state = np.zeros(4)
        net.nodes["B"].state = np.zeros(4)
        cer = ChronoEntropicRegularizer(net)
        penalty = cer.compute(net.nodes["A"])
        assert np.isfinite(penalty)

    def test_single_dim_returns_zero(self):
        net = CausalTensorNetwork()
        net.add_node(CausalNode("X", time_coordinate=0.0, state_dim=1))
        cer = ChronoEntropicRegularizer(net)
        penalty = cer.compute(net.nodes["X"])
        assert penalty == 0.0


class TestPPGD:
    def _make_loop_net(self) -> CausalTensorNetwork:
        np.random.seed(42)
        net = CausalTensorNetwork()
        net.add_node(CausalNode("A", time_coordinate=0.0, state_dim=4))
        net.add_node(CausalNode("B", time_coordinate=1.0, state_dim=4))
        net.add_node(CausalNode("C", time_coordinate=2.0, state_dim=4))
        net.add_causal_edge("A", "B", weight=0.8)
        net.add_causal_edge("B", "C", weight=0.6)
        net.add_retrocausal_edge("C", "A", weight=0.4)
        return net

    def test_single_step_returns_loss(self):
        net = self._make_loop_net()
        ppgd = ParadoxPressureGradientDescent(net)
        loss = ppgd.step()
        assert isinstance(loss, float)
        assert loss >= 0.0

    def test_loss_decreases(self):
        net = self._make_loop_net()
        ppgd = ParadoxPressureGradientDescent(net, lr=0.01)
        losses = [ppgd.step() for _ in range(50)]
        assert losses[-1] < losses[0]

    def test_optimize_returns_result_dict(self):
        net = self._make_loop_net()
        ppgd = ParadoxPressureGradientDescent(net)
        result = ppgd.optimize(max_steps=50, verbose=False)
        assert "final_loss" in result
        assert "steps" in result
        assert "loss_history" in result
        assert "converged" in result

    def test_optimize_with_cer(self):
        net = self._make_loop_net()
        ppgd = ParadoxPressureGradientDescent(net, cer_lambda=0.1)
        result = ppgd.optimize(max_steps=100, verbose=False)
        assert result["final_loss"] >= 0.0

    def test_momentum_updates_velocities(self):
        net = self._make_loop_net()
        ppgd = ParadoxPressureGradientDescent(net, momentum=0.9)
        ppgd.step()
        has_nonzero = any(
            np.any(v != 0) for v in ppgd.velocities.values()
        )
        assert has_nonzero
