"""Tests for visualization module."""

import numpy as np
import pytest
import matplotlib
matplotlib.use("Agg")

from tqce.network.node import CausalNode
from tqce.network.ctn import CausalTensorNetwork
from tqce.viz.plots import (
    plot_convergence,
    plot_paradox_pressure_landscape,
    plot_ctn_graph,
    plot_tec_probabilities,
    plot_amplification_trace,
    animate_convergence,
)


@pytest.fixture
def simple_net():
    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


class TestVisualization:
    def test_plot_convergence(self):
        history = [10.0, 5.0, 2.0, 1.0, 0.5, 0.1]
        fig = plot_convergence(history)
        assert fig is not None
        import matplotlib.pyplot as plt
        plt.close(fig)

    def test_plot_convergence_with_loss(self):
        history = [10.0, 5.0, 2.0, 1.0]
        loss_history = [8.0, 4.0, 1.5, 0.5]
        fig = plot_convergence(history, loss_history=loss_history)
        assert fig is not None
        import matplotlib.pyplot as plt
        plt.close(fig)

    def test_plot_paradox_pressure_landscape(self):
        pressures = {"A": 0.5, "B": 0.2, "C": 0.8}
        fig = plot_paradox_pressure_landscape(pressures)
        assert fig is not None
        import matplotlib.pyplot as plt
        plt.close(fig)

    def test_plot_ctn_graph(self, simple_net):
        fig = plot_ctn_graph(simple_net)
        assert fig is not None
        import matplotlib.pyplot as plt
        plt.close(fig)

    def test_plot_tec_probabilities(self):
        probs = np.array([0.1, 0.3, 0.4, 0.15, 0.05])
        fig = plot_tec_probabilities(probs)
        assert fig is not None
        import matplotlib.pyplot as plt
        plt.close(fig)

    def test_plot_amplification_trace(self):
        trace = [1.0, 2.0, 4.0, 7.0, 9.5, 10.0, 10.0]
        fig = plot_amplification_trace(trace, target=10.0)
        assert fig is not None
        import matplotlib.pyplot as plt
        plt.close(fig)

    def test_animate_convergence(self):
        histories = [[10.0, 5.0, 2.0], [8.0, 3.0, 1.0]]
        anim = animate_convergence(histories)
        assert anim is not None
        import matplotlib.pyplot as plt
        plt.close("all")
