"""Tests for Temporal Eigenstate Collapse (TEC) components."""

import numpy as np
import pytest

from tqce.quantum.wavefunction import TemporalWavefunction


class TestTemporalWavefunction:
    def test_initialization(self):
        wf = TemporalWavefunction(n_timelines=10, state_dim=5)
        assert wf.timeline_states.shape == (10, 5)
        assert wf.amplitudes.shape == (10,)

    def test_amplitudes_normalized(self):
        wf = TemporalWavefunction(n_timelines=10, state_dim=5)
        norm = np.sqrt(np.sum(np.abs(wf.amplitudes) ** 2))
        np.testing.assert_allclose(norm, 1.0, atol=1e-10)

    def test_interference_returns_probabilities(self):
        wf = TemporalWavefunction(n_timelines=5, state_dim=3)
        consistency_fn = lambda state: -np.linalg.norm(state)
        probs = wf.compute_causal_interference(consistency_fn)
        assert probs.shape == (5,)
        np.testing.assert_allclose(np.sum(probs), 1.0, atol=1e-6)
        assert np.all(probs >= 0)

    def test_collapse_returns_state_and_score(self):
        wf = TemporalWavefunction(n_timelines=10, state_dim=4)
        consistency_fn = lambda state: -np.linalg.norm(state)
        state, score = wf.collapse(consistency_fn, n_interference_steps=5)
        assert state.shape == (4,)
        assert isinstance(score, (float, np.floating))

    def test_consistent_timeline_wins(self):
        """The most consistent timeline should have highest probability."""
        np.random.seed(42)
        wf = TemporalWavefunction(n_timelines=5, state_dim=3)
        wf.timeline_states[2] = np.zeros(3)

        def consistency_fn(state):
            return -np.linalg.norm(state)

        state, score = wf.collapse(consistency_fn, n_interference_steps=10)
        assert np.linalg.norm(state) < np.mean(
            [np.linalg.norm(wf.timeline_states[i]) for i in [0, 1, 3, 4]]
        )

    def test_probabilities_sum_to_one_after_interference(self):
        wf = TemporalWavefunction(n_timelines=8, state_dim=4)
        fn = lambda s: float(-np.sum(s ** 2))
        for _ in range(5):
            probs = wf.compute_causal_interference(fn)
            np.testing.assert_allclose(np.sum(probs), 1.0, atol=1e-6)
