"""Paradox Pressure Gradient Descent (PPGD).

A gradient-based optimizer where the loss is causal inconsistency
(paradox pressure) across a temporal system.

Loss function:
    L(S) = sum_i ||S_i - T(S_i)||^2 + lambda * CER(S)

The gradient propagates not just through the network but ACROSS the
causal loop itself, making this a fundamentally new optimizer class.
"""

from __future__ import annotations

from typing import Any

import numpy as np
from numpy.typing import NDArray

from tqce.network.ctn import CausalTensorNetwork
from tqce.optimization.regularizers import ChronoEntropicRegularizer


class ParadoxPressureGradientDescent:
    """PPGD Optimizer — descend the paradox pressure loss surface.

    Treats causal inconsistency as a differentiable loss and uses
    momentum-based gradient descent to find self-consistent states.

    Args:
        ctn: The Causal Tensor Network to optimize.
        lr: Base learning rate.
        momentum: Momentum coefficient for velocity updates.
        cer_lambda: Regularization strength for Chrono-Entropic Regularizer.
        adaptive_lr: Scale learning rate inversely with paradox pressure.
    """

    def __init__(
        self,
        ctn: CausalTensorNetwork,
        lr: float = 0.01,
        momentum: float = 0.9,
        cer_lambda: float = 0.1,
        adaptive_lr: bool = True,
    ) -> None:
        self.ctn = ctn
        self.lr = lr
        self.momentum = momentum
        self.cer_lambda = cer_lambda
        self.adaptive_lr = adaptive_lr
        self.cer = ChronoEntropicRegularizer(ctn)

        self.velocities: dict[str, NDArray[np.floating]] = {
            nid: np.zeros(n.state_dim) for nid, n in ctn.nodes.items()
        }
        self.adaptive_lrs: dict[str, float] = {nid: lr for nid in ctn.nodes}
        self.loss_history: list[float] = []

    def compute_paradox_loss(
        self,
    ) -> tuple[float, dict[str, NDArray[np.floating]]]:
        """Compute total paradox pressure loss and per-node gradients.

        Returns:
            Tuple of (total_loss, {node_id: gradient_vector}).
        """
        total_loss = 0.0
        gradients: dict[str, NDArray[np.floating]] = {}

        for node_id, node in self.ctn.nodes.items():
            target_state = self.ctn.compute_node_state(node_id)
            residual = node.state - target_state
            node_loss = float(np.dot(residual, residual))

            cer_loss = self.cer.compute(node)
            total_loss += node_loss + self.cer_lambda * cer_loss

            cer_grad = self.cer.gradient(node)
            gradients[node_id] = 2 * residual + self.cer_lambda * cer_grad

        return total_loss, gradients

    def step(self) -> float:
        """Perform one PPGD optimization step.

        Returns:
            The total loss after this step.
        """
        loss, gradients = self.compute_paradox_loss()

        for node_id, node in self.ctn.nodes.items():
            grad = gradients[node_id]

            if self.adaptive_lr:
                pressure_factor = 1.0 / (1.0 + node.paradox_pressure)
                effective_lr = self.adaptive_lrs[node_id] * pressure_factor
            else:
                effective_lr = self.lr

            self.velocities[node_id] = (
                self.momentum * self.velocities[node_id] - effective_lr * grad
            )
            node.state = node.state + self.velocities[node_id]
            node.paradox_pressure = float(np.linalg.norm(grad))

        self.loss_history.append(loss)
        return loss

    def optimize(
        self,
        max_steps: int = 1000,
        target_loss: float = 1e-6,
        verbose: bool = False,
    ) -> dict[str, Any]:
        """Run full PPGD optimization.

        Args:
            max_steps: Maximum optimization steps.
            target_loss: Stop when loss drops below this value.
            verbose: Print progress to stdout.

        Returns:
            Dict with keys: final_loss, steps, loss_history, converged.
        """
        if verbose:
            print("\n+" + "-" * 44 + "+")
            print("|   PPGD Optimizer - Paradox Descent         |")
            print("+" + "-" * 44 + "+")

        loss = float("inf")
        step = 0

        for step in range(max_steps):
            loss = self.step()

            if verbose and step % 100 == 0:
                print(f"|  Step {step:4d} | Loss = {loss:.6f}              |")

            if loss < target_loss:
                if verbose:
                    print(f"|  Converged at step {step:<24}|")
                break

        if verbose:
            print("+" + "-" * 44 + "+")

        return {
            "final_loss": loss,
            "steps": step + 1,
            "loss_history": self.loss_history,
            "converged": loss < target_loss,
        }
