"""Temporal Eigenstate Collapse (TEC).

Superposes multiple possible timelines as a 'temporal wavefunction'.
Each timeline is assigned a complex amplitude. Causal interference
amplifies consistent resolutions and cancels paradoxical ones.

This is inspired by quantum computing but runs on classical hardware.
The interference is CAUSAL interference, not quantum.
"""

from __future__ import annotations

from typing import Callable

import numpy as np
from numpy.typing import NDArray


class TemporalWavefunction:
    """A superposition of possible timeline states.

    Each timeline is assigned a complex amplitude alpha_k such that:
        |Psi> = sum_k alpha_k |Timeline_k>

    The probability of a timeline is |alpha_k|^2.
    Collapse selects the timeline with maximum probability after interference.

    Args:
        n_timelines: Number of candidate timelines to superpose.
        state_dim: Dimensionality of each timeline's state vector.
    """

    def __init__(self, n_timelines: int, state_dim: int) -> None:
        self.n_timelines = n_timelines
        self.state_dim = state_dim

        self.timeline_states = np.random.randn(n_timelines, state_dim) * 0.1

        phases = np.random.uniform(0, 2 * np.pi, n_timelines)
        magnitudes = np.ones(n_timelines) / np.sqrt(n_timelines)
        self.amplitudes: NDArray[np.complexfloating] = magnitudes * np.exp(
            1j * phases
        )

        self.consistency_scores = np.zeros(n_timelines)

    def compute_causal_interference(
        self, consistency_fn: Callable[[NDArray], float]
    ) -> NDArray[np.floating]:
        """Apply causal interference between timelines.

        Consistent timelines interfere constructively; paradoxical ones
        interfere destructively.

        Args:
            consistency_fn: Maps a state vector to a consistency score
                (higher = more consistent).

        Returns:
            Probability distribution over timelines (sums to 1).
        """
        new_amplitudes = np.zeros(self.n_timelines, dtype=complex)

        for i in range(self.n_timelines):
            self.consistency_scores[i] = consistency_fn(self.timeline_states[i])

            for j in range(self.n_timelines):
                similarity = np.exp(
                    -np.linalg.norm(
                        self.timeline_states[i] - self.timeline_states[j]
                    )
                )

                phase_diff = np.angle(self.amplitudes[j]) - np.angle(
                    self.amplitudes[i]
                )

                ci = self.consistency_scores[i]
                cj = consistency_fn(self.timeline_states[j])

                interference_sign = np.sign(ci * cj) if ci * cj != 0 else 1.0

                new_amplitudes[i] += (
                    self.amplitudes[j]
                    * similarity
                    * interference_sign
                    * np.exp(1j * phase_diff)
                )

        norm = np.sqrt(np.sum(np.abs(new_amplitudes) ** 2))
        if norm > 1e-10:
            new_amplitudes /= norm

        self.amplitudes = new_amplitudes
        return np.abs(self.amplitudes) ** 2

    def collapse(
        self,
        consistency_fn: Callable[[NDArray], float],
        n_interference_steps: int = 10,
    ) -> tuple[NDArray[np.floating], float]:
        """Run interference and collapse to the most consistent timeline.

        Args:
            consistency_fn: Consistency scoring function.
            n_interference_steps: Number of interference rounds before collapse.

        Returns:
            Tuple of (best_timeline_state, consistency_score).
        """
        probs = np.zeros(self.n_timelines)
        for _ in range(n_interference_steps):
            probs = self.compute_causal_interference(consistency_fn)

        best_idx = int(np.argmax(probs))
        return (
            self.timeline_states[best_idx].copy(),
            float(self.consistency_scores[best_idx]),
        )
