"""Tests for Causal Tensor Network components."""

import numpy as np
import pytest

from tqce.network.node import CausalNode
from tqce.network.ctn import CausalTensorNetwork


class TestCausalNode:
    def test_create_node_with_defaults(self):
        node = CausalNode("A", time_coordinate=0.0, state_dim=4)
        assert node.node_id == "A"
        assert node.time_coordinate == 0.0
        assert node.state_dim == 4
        assert node.state.shape == (4,)
        assert node.paradox_pressure == 0.0

    def test_create_node_with_explicit_state(self):
        state = np.array([1.0, 2.0, 3.0])
        node = CausalNode("B", time_coordinate=1.0, state_dim=3, state=state)
        np.testing.assert_array_equal(node.state, state)

    def test_default_state_is_small_random(self):
        node = CausalNode("C", time_coordinate=0.0, state_dim=100)
        assert np.abs(node.state).max() < 1.0

    def test_causal_parents_default_empty(self):
        node = CausalNode("D", time_coordinate=0.0, state_dim=4)
        assert node.causal_parents == []
        assert node.retrocausal_parents == []

    def test_node_with_custom_functions(self):
        fn = lambda x: x * 2
        node = CausalNode("E", time_coordinate=0.0, state_dim=4,
                          forward_fn=fn, retro_fn=fn)
        assert node.forward_fn is not None
        assert node.retro_fn is not None


class TestCausalTensorNetwork:
    def _make_simple_network(self) -> CausalTensorNetwork:
        """A -> B -> C with retrocausal C -> A."""
        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_add_nodes(self):
        net = CausalTensorNetwork()
        node = CausalNode("A", time_coordinate=0.0, state_dim=4)
        net.add_node(node)
        assert "A" in net.nodes
        assert "A" in net.graph

    def test_causal_edge_enforces_temporal_order(self):
        net = CausalTensorNetwork()
        net.add_node(CausalNode("A", time_coordinate=1.0, state_dim=4))
        net.add_node(CausalNode("B", time_coordinate=0.0, state_dim=4))
        with pytest.raises(ValueError, match="violates causality"):
            net.add_causal_edge("A", "B")

    def test_retrocausal_edge_enforces_backward_time(self):
        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))
        with pytest.raises(ValueError, match="must go from future to past"):
            net.add_retrocausal_edge("A", "B")

    def test_causal_edge_registers_parent(self):
        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", weight=0.7)
        assert ("A", 0.7) in net.nodes["B"].causal_parents

    def test_retrocausal_edge_registers_parent(self):
        net = self._make_simple_network()
        assert ("C", 0.4) in net.nodes["A"].retrocausal_parents

    def test_compute_node_state_no_parents(self):
        net = CausalTensorNetwork()
        node = CausalNode("X", time_coordinate=0.0, state_dim=4)
        net.add_node(node)
        result = net.compute_node_state("X")
        np.testing.assert_array_equal(result, node.state)

    def test_compute_node_state_with_parents(self):
        net = CausalTensorNetwork()
        a = CausalNode("A", time_coordinate=0.0, state_dim=4)
        b = CausalNode("B", time_coordinate=1.0, state_dim=4)
        a.state = np.ones(4)
        net.add_node(a)
        net.add_node(b)
        net.add_causal_edge("A", "B", weight=1.0)
        result = net.compute_node_state("B")
        expected = np.tanh(np.ones(4))
        np.testing.assert_allclose(result, expected, atol=1e-6)

    def test_compute_node_state_deterministic(self):
        """compute_node_state must produce the same result on repeated calls."""
        net = self._make_simple_network()
        r1 = net.compute_node_state("A")
        r2 = net.compute_node_state("A")
        np.testing.assert_array_equal(r1, r2)

    def test_mixed_dimension_projection(self):
        """Nodes with different state_dim should work via stored projections."""
        net = CausalTensorNetwork()
        net.add_node(CausalNode("A", time_coordinate=0.0, state_dim=3))
        net.add_node(CausalNode("B", time_coordinate=1.0, state_dim=5))
        net.add_causal_edge("A", "B", weight=1.0)
        result = net.compute_node_state("B")
        assert result.shape == (5,)


class TestFixedPointSolver:
    def test_simple_chain_converges(self):
        """A -> B -> C (no loops) converges immediately."""
        net = CausalTensorNetwork()
        net.add_node(CausalNode("A", time_coordinate=0.0, state_dim=3))
        net.add_node(CausalNode("B", time_coordinate=1.0, state_dim=3))
        net.add_node(CausalNode("C", time_coordinate=2.0, state_dim=3))
        net.add_causal_edge("A", "B")
        net.add_causal_edge("B", "C")
        result = net.run_temporal_fixed_point(verbose=False)
        assert result["converged"]

    def test_loop_converges(self):
        """A -> B -> C -> A (retro) should find a self-consistent fixed point."""
        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)
        result = net.run_temporal_fixed_point(verbose=False)
        assert result["converged"]
        assert result["final_inconsistency"] < 1e-7

    def test_iteration_history_decreasing(self):
        """Inconsistency should generally decrease over iterations."""
        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_retrocausal_edge("B", "A", weight=0.5)
        net.add_causal_edge("A", "B", weight=0.5)
        result = net.run_temporal_fixed_point(verbose=False)
        history = result["iteration_history"]
        assert history[-1] < history[0]

    def test_result_keys(self):
        net = CausalTensorNetwork()
        net.add_node(CausalNode("A", time_coordinate=0.0, state_dim=2))
        result = net.run_temporal_fixed_point(max_iterations=5, verbose=False)
        expected_keys = {
            "converged", "iterations", "final_inconsistency",
            "node_states", "paradox_pressures", "iteration_history"
        }
        assert set(result.keys()) == expected_keys
