"""Causal Tensor Network — tensor networks with retrocausal edges."""

from __future__ import annotations

from typing import Any

import networkx as nx
import numpy as np
from numpy.typing import NDArray

from tqce.network.node import CausalNode


class CausalTensorNetwork:
    """A tensor network where edges can point backward in time.

    Creates Closed Causal Loops (CCLs) — the computational analog of
    Closed Timelike Curves (CTCs) in physics.

    The network is solved not by a forward pass, but by fixed-point iteration
    across the entire temporal structure. The fixed point IS the self-consistent
    timeline.

    Attributes:
        nodes: Mapping of node_id to CausalNode.
        graph: NetworkX directed graph encoding the full topology.
        causal_edges: List of (from_id, to_id, weight) forward-time edges.
        retro_edges: List of (from_future_id, to_past_id, weight) backward-time edges.
        consistency_threshold: Convergence threshold for the fixed-point solver.
        iteration_history: Record of total inconsistency at each iteration.
    """

    def __init__(self, consistency_threshold: float = 1e-7) -> None:
        self.nodes: dict[str, CausalNode] = {}
        self.graph = nx.DiGraph()
        self.retro_edges: list[tuple[str, str, float]] = []
        self.causal_edges: list[tuple[str, str, float]] = []
        self.consistency_threshold = consistency_threshold
        self.iteration_history: list[float] = []
        self._projections: dict[tuple[str, str], NDArray[np.floating]] = {}

    def add_node(self, node: CausalNode) -> None:
        """Add a node to the network."""
        self.nodes[node.node_id] = node
        self.graph.add_node(node.node_id, time=node.time_coordinate)

    def _get_projection(self, from_id: str, to_id: str,
                        from_dim: int, to_dim: int) -> NDArray[np.floating]:
        """Get or create a deterministic projection matrix for mismatched dims."""
        key = (from_id, to_id)
        if key not in self._projections:
            rng = np.random.RandomState(hash(key) % (2**31))
            self._projections[key] = rng.randn(to_dim, from_dim) * 0.1
        return self._projections[key]

    def add_causal_edge(self, from_id: str, to_id: str,
                        weight: float = 1.0) -> None:
        """Add a forward-time causal connection.

        Args:
            from_id: Source node (must be earlier in time).
            to_id: Target node (must be later in time).
            weight: Connection strength.

        Raises:
            ValueError: If from_node is not strictly before to_node in time.
        """
        from_node = self.nodes[from_id]
        to_node = self.nodes[to_id]

        if from_node.time_coordinate >= to_node.time_coordinate:
            raise ValueError(
                f"Causal edge from {from_id} (t={from_node.time_coordinate}) "
                f"to {to_id} (t={to_node.time_coordinate}) violates causality. "
                f"Use add_retrocausal_edge() for backward-time connections."
            )

        to_node.causal_parents.append((from_id, weight))
        self.causal_edges.append((from_id, to_id, weight))
        self.graph.add_edge(from_id, to_id, weight=weight, type="causal")

    def add_retrocausal_edge(self, from_future_id: str, to_past_id: str,
                              weight: float = 1.0) -> None:
        """Add a backward-time (retrocausal) connection.

        This creates a Closed Causal Loop when combined with forward edges.

        Args:
            from_future_id: Source node (must be later in time).
            to_past_id: Target node (must be earlier in time).
            weight: Connection strength.

        Raises:
            ValueError: If from_future is not strictly after to_past in time.
        """
        future_node = self.nodes[from_future_id]
        past_node = self.nodes[to_past_id]

        if future_node.time_coordinate <= past_node.time_coordinate:
            raise ValueError(
                f"Retrocausal edge must go from future to past. "
                f"{from_future_id} (t={future_node.time_coordinate}) -> "
                f"{to_past_id} (t={past_node.time_coordinate}) is invalid."
            )

        past_node.retrocausal_parents.append((from_future_id, weight))
        self.retro_edges.append((from_future_id, to_past_id, weight))
        self.graph.add_edge(from_future_id, to_past_id,
                            weight=weight, type="retrocausal")

    def _apply_parent_contribution(
        self, node: CausalNode, parent_id: str, weight: float,
        is_retrocausal: bool = False
    ) -> NDArray[np.floating]:
        """Compute a single parent's contribution to a node's state."""
        parent_state = self.nodes[parent_id].state
        fn = node.retro_fn if is_retrocausal else node.forward_fn
        effective_weight = weight * 0.5 if is_retrocausal else weight

        if fn is not None:
            return effective_weight * fn(parent_state)

        if len(parent_state) == node.state_dim:
            return effective_weight * np.tanh(parent_state)

        proj = self._get_projection(parent_id, node.node_id,
                                    len(parent_state), node.state_dim)
        return effective_weight * np.tanh(proj @ parent_state)

    def compute_node_state(self, node_id: str) -> NDArray[np.floating]:
        """Compute a node's state from all its causal and retrocausal parents.

        If the node has no parents, returns its current state (bootstrap node).
        """
        node = self.nodes[node_id]
        contributions: list[NDArray] = []

        for parent_id, weight in node.causal_parents:
            contributions.append(
                self._apply_parent_contribution(node, parent_id, weight, False)
            )

        for parent_id, weight in node.retrocausal_parents:
            contributions.append(
                self._apply_parent_contribution(node, parent_id, weight, True)
            )

        if not contributions:
            return node.state

        return np.mean(contributions, axis=0)

    def run_temporal_fixed_point(
        self,
        max_iterations: int = 500,
        damping: float = 0.3,
        verbose: bool = False,
    ) -> dict[str, Any]:
        """Solve the self-consistency problem via fixed-point iteration.

        Repeatedly updates all nodes using their parents' current states until
        the total state change drops below ``consistency_threshold``.

        Args:
            max_iterations: Maximum solver iterations.
            damping: Blending factor in [0,1]. Higher = more aggressive updates.
            verbose: Print progress to stdout.

        Returns:
            Dict with keys: converged, iterations, final_inconsistency,
            node_states, paradox_pressures, iteration_history.
        """
        if verbose:
            print("+" + "-" * 44 + "+")
            print("|   TQCE Temporal Fixed-Point Solver         |")
            print("+" + "-" * 44 + "+")

        sorted_nodes = sorted(
            self.nodes.values(), key=lambda n: n.time_coordinate
        )

        iteration = 0
        total_delta = float("inf")

        while iteration < max_iterations and total_delta > self.consistency_threshold:
            total_delta = 0.0

            # Forward pass
            for node in sorted_nodes:
                new_state = self.compute_node_state(node.node_id)
                delta = float(np.linalg.norm(new_state - node.state))
                total_delta += delta
                node.state = (1 - damping) * node.state + damping * new_state
                node.paradox_pressure = delta

            # Backward pass (retrocausal sweep)
            for node in reversed(sorted_nodes):
                new_state = self.compute_node_state(node.node_id)
                delta = float(np.linalg.norm(new_state - node.state))
                total_delta += delta
                node.state = (1 - damping) * node.state + damping * new_state

            self.iteration_history.append(total_delta)

            if verbose and iteration % 50 == 0:
                pressure_max = max(
                    n.paradox_pressure for n in self.nodes.values()
                )
                print(
                    f"|  Iter {iteration:4d} | dS = {total_delta:.2e} | "
                    f"MaxP = {pressure_max:.4f}  |"
                )

            iteration += 1

        converged = total_delta <= self.consistency_threshold

        if verbose:
            print("+" + "-" * 44 + "+")
            status = "CONVERGED" if converged else "DIVERGED (Paradox?)"
            print(f"|  Status: {status:<34}|")
            print(f"|  Final dS: {total_delta:.2e}{' ' * 23}|")
            print(f"|  Iterations: {iteration:<29}|")
            print("+" + "-" * 44 + "+")

        return {
            "converged": converged,
            "iterations": iteration,
            "final_inconsistency": total_delta,
            "node_states": {
                nid: n.state.copy() for nid, n in self.nodes.items()
            },
            "paradox_pressures": {
                nid: n.paradox_pressure for nid, n in self.nodes.items()
            },
            "iteration_history": self.iteration_history.copy(),
        }
