"""Chrono-Entropic Regularization (CER).

Physical insight: Real closed timelike curves are constrained by the
second law of thermodynamics. Information in a causal loop cannot grow
unboundedly — entropy must be non-decreasing.

CER penalizes states that would imply entropy decrease across a causal loop:
    CER(S) = max(0, H(S_past) - H(S_future))
where H is the approximate entropy of the state distribution.
"""

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray

from tqce.network.node import CausalNode

from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from tqce.network.ctn import CausalTensorNetwork


def _approx_entropy(state: NDArray[np.floating]) -> float:
    """Approximate entropy of a state vector via normalized magnitude distribution."""
    if len(state) < 2:
        return 0.0
    probs = np.abs(state) / (np.sum(np.abs(state)) + 1e-10)
    return float(-np.sum(probs * np.log(probs + 1e-10)))


class ChronoEntropicRegularizer:
    """Thermodynamics-inspired regularizer for causal loop systems.

    Penalizes states where a past node has higher approximate entropy than
    its retrocausal (future) parents, which would imply information flowing
    backward in violation of the second law.

    Args:
        ctn: The Causal Tensor Network to regularize.
        eps: Finite difference step size for numerical gradients.
    """

    def __init__(self, ctn: CausalTensorNetwork, eps: float = 1e-5) -> None:
        self.ctn = ctn
        self.eps = eps

    def compute(self, node: CausalNode) -> float:
        """Compute the CER penalty for a single node.

        Returns 0.0 for nodes with state_dim < 2 or no retrocausal parents.
        """
        if node.state_dim < 2:
            return 0.0

        node_entropy = _approx_entropy(node.state)
        penalty = 0.0

        for future_id, weight in node.retrocausal_parents:
            future_node = self.ctn.nodes[future_id]
            future_entropy = _approx_entropy(future_node.state)
            if node_entropy > future_entropy:
                penalty += weight * (node_entropy - future_entropy)

        return penalty

    def gradient(self, node: CausalNode) -> NDArray[np.floating]:
        """Numerical gradient of CER with respect to node state."""
        grad = np.zeros_like(node.state)
        base = self.compute(node)

        for i in range(len(node.state)):
            node.state[i] += self.eps
            perturbed = self.compute(node)
            grad[i] = (perturbed - base) / self.eps
            node.state[i] -= self.eps

        return grad
