"""
Flux Algebra for PyTorch - Matrix/Tensor Implementation
========================================================
Extends the FluxNumber system to work with PyTorch tensors for ML iterations.

Key Features:
- Batched Flux operations on tensors
- GPU-compatible flux algebra operators
- Integration with PyTorch autograd for gradient-based learning
- Support for neural network weight updates with uncertainty tracking

Based on: Flux Algebra Book (Conditional Collapse Theory)
"""

import torch
import torch.nn as nn
from typing import Tuple, Optional, Union


class FluxTensor:
    """
    Flux Tensor: Z = ⟨V, Σ, Τ⟩
    
    Extends FluxNumber to work with PyTorch tensors/matrices.
    
    Attributes:
        v (torch.Tensor): Value tensor (stationary component)
        s (torch.Tensor): Entropy tensor (uncertainty, >= 0)
        t (torch.Tensor): Flux tensor (rate of change)
    """
    
    def __init__(
        self, 
        v: torch.Tensor, 
        s: Optional[torch.Tensor] = None, 
        t: Optional[torch.Tensor] = None
    ):
        """
        Initialize a FluxTensor.
        
        Args:
            v: Value tensor (any shape)
            s: Entropy tensor (same shape as v, or scalar, defaults to zeros)
            t: Flux tensor (same shape as v, or scalar, defaults to zeros)
        """
        self.v = v.clone().detach().float()
        
        if s is None:
            self.s = torch.zeros_like(self.v)
        else:
            self.s = s.clone().detach().float()
            # Ensure entropy is non-negative
            self.s = torch.clamp(self.s, min=0.0)
            
        if t is None:
            self.t = torch.zeros_like(self.v)
        else:
            self.t = t.clone().detach().float()
        
        # Validate shapes
        if self.s.shape != self.v.shape:
            self.s = self.s.expand_as(self.v).clone()
        if self.t.shape != self.v.shape:
            self.t = self.t.expand_as(self.v).clone()
    
    def to(self, device: torch.device) -> 'FluxTensor':
        """Move all components to specified device."""
        self.v = self.v.to(device)
        self.s = self.s.to(device)
        self.t = self.t.to(device)
        return self
    
    def cuda(self) -> 'FluxTensor':
        """Move to GPU."""
        return self.to(torch.device('cuda'))
    
    def cpu(self) -> 'FluxTensor':
        """Move to CPU."""
        return self.to(torch.device('cpu'))
    
    def detach(self) -> 'FluxTensor':
        """Detach all components from computation graph."""
        return FluxTensor(self.v.detach(), self.s.detach(), self.t.detach())
    
    def clone(self) -> 'FluxTensor':
        """Create a deep copy."""
        return FluxTensor(self.v.clone(), self.s.clone(), self.t.clone())
    
    @property
    def shape(self):
        """Return the shape of the value tensor."""
        return self.v.shape
    
    @property
    def device(self):
        """Return the device."""
        return self.v.device
    
    def __repr__(self) -> str:
        return (f"FluxTensor(shape={self.shape}, "
                f"v_mean={self.v.mean().item():.4f}, "
                f"s_mean={self.s.mean().item():.4f}, "
                f"t_mean={self.t.mean().item():.4f})")
    
    # ==================== FLUX OPERATORS ====================
    
    def __add__(self, other: Union['FluxTensor', torch.Tensor, float]) -> 'FluxTensor':
        """
        Flux Addition (⊕): Merging two trajectories.
        
        z₁ ⊕ z₂ = ⟨v₁ + v₂, √(σ₁² + σ₂²), τ₁ + τ₂⟩
        """
        if isinstance(other, (torch.Tensor, float, int)):
            if isinstance(other, (float, int)):
                other = torch.tensor(other, device=self.device)
            other = FluxTensor(other, torch.zeros_like(other), torch.zeros_like(other))
        
        new_v = self.v + other.v
        new_s = torch.sqrt(self.s**2 + other.s**2)
        new_t = self.t + other.t
        return FluxTensor(new_v, new_s, new_t)
    
    def __radd__(self, other: Union[torch.Tensor, float]) -> 'FluxTensor':
        """Right addition for scalar/tensor + FluxTensor."""
        return self.__add__(other)
    
    def __mul__(self, other: Union['FluxTensor', torch.Tensor, float]) -> 'FluxTensor':
        """
        Flux Multiplication (⊗): Scaling of complexity.
        
        z₁ ⊗ z₂ = ⟨v₁·v₂, √((v₁·σ₂)² + (v₂·σ₁)²), v₁·τ₂ + v₂·τ₁⟩
        """
        if isinstance(other, (torch.Tensor, float, int)):
            if isinstance(other, (float, int)):
                other = torch.tensor(other, device=self.device)
            other = FluxTensor(other, torch.zeros_like(other), torch.zeros_like(other))
        
        new_v = self.v * other.v
        new_s = torch.sqrt((self.v * other.s)**2 + (other.v * self.s)**2)
        new_t = self.v * other.t + other.v * self.t
        return FluxTensor(new_v, new_s, new_t)
    
    def __rmul__(self, other: Union[torch.Tensor, float]) -> 'FluxTensor':
        """Right multiplication for scalar/tensor * FluxTensor."""
        return self.__mul__(other)
    
    def __matmul__(self, other: Union['FluxTensor', torch.Tensor]) -> 'FluxTensor':
        """
        Flux Matrix Multiplication.
        
        Extends flux algebra to matrix operations with uncertainty propagation.
        """
        if isinstance(other, torch.Tensor):
            other = FluxTensor(other, torch.zeros_like(other), torch.zeros_like(other))
        
        # Value: standard matrix multiplication
        new_v = self.v @ other.v
        
        # Entropy: propagate through matrix multiplication
        # σ_new = √(V₁·Σ₂²·V₁ᵀ + V₂·Σ₁²·V₂ᵀ)
        # Simplified: use Frobenius norm approximation
        s_1_sq = self.s**2
        s_2_sq = other.s**2
        
        # Conservative entropy propagation for matrix multiply
        new_s = (torch.sqrt(s_1_sq) @ other.v.abs() + 
                 self.v.abs() @ torch.sqrt(s_2_sq))
        
        # Flux: product rule for matrices
        new_t = self.v @ other.t + other.v @ self.t
        
        return FluxTensor(new_v, new_s, new_t)
    
    # ==================== COLLAPSE & EVOLVE ====================
    
    def collapse(self, work: Union[float, torch.Tensor]) -> 'FluxTensor':
        """
        Collapse Operator (𝒞): Reduce entropy by paying work.
        
        𝒞(z, W) = ⟨v, max(0, σ - W), τ⟩
        
        Args:
            work: Work budget (scalar or tensor matching shape)
        """
        if isinstance(work, (float, int)):
            work = torch.tensor(work, device=self.device)
        
        new_s = torch.clamp(self.s - work, min=0.0)
        return FluxTensor(self.v, new_s, self.t)
    
    def evolve(self, dt: float = 1.0, entropy_decay: float = 0.0) -> 'FluxTensor':
        """
        Time evolution / derivative step.
        
        Updates value by flux * dt, optionally decays entropy.
        
        Args:
            dt: Time step
            entropy_decay: Decay factor (0 = no decay, 1 = full decay)
        """
        new_v = self.v + self.t * dt
        new_s = self.s * (1.0 - entropy_decay)
        new_t = self.t  # flux unchanged
        return FluxTensor(new_v, new_s, new_t)
    
    # ==================== LEARNING INTEGRATION ====================
    
    def gradient_step(
        self, 
        loss: torch.Tensor, 
        lr: float = 0.01,
        work: float = 0.0,
        create_graph: bool = False
    ) -> 'FluxTensor':
        """
        Perform a gradient descent step with flux algebra.
        
        Implements learning as entropy collapse:
        θ_{t+1} = θ_t ⊕ (-lr · ∇L) then collapse with work
        
        Args:
            loss: Loss tensor (scalar)
            lr: Learning rate (becomes flux component)
            work: Work budget for collapse
            create_graph: Whether to maintain computation graph
        
        Returns:
            Updated FluxTensor
        """
        # Ensure value tensor requires grad for this operation
        v_with_grad = self.v.clone().requires_grad_(True)
        
        # Recompute loss dependency on this specific tensor
        # Note: This assumes loss was computed using self.v
        # For proper integration, use with FluxLinear or track gradients manually
        
        # Try to get existing gradients if available
        if self.v.grad is not None:
            grads = self.v.grad
        else:
            # If no pre-computed gradients, compute them from loss
            # This requires loss to depend on self.v
            grads = torch.autograd.grad(loss, self.v, retain_graph=True)[0]
        
        # Flux update: value moves along gradient
        delta_v = -lr * grads
        
        # Treat gradient as flux (learning trajectory)
        new_t = self.t + delta_v
        
        # Evolve
        new_v = self.v + delta_v
        
        # Entropy grows with gradient magnitude (uncertainty in direction)
        new_s = self.s + torch.abs(grads) * lr
        
        # Collapse if work is provided
        if work > 0:
            new_s = torch.clamp(new_s - work, min=0.0)
        
        return FluxTensor(new_v, new_s, new_t)
    
    def adam_step(
        self,
        loss: torch.Tensor,
        lr: float = 0.001,
        beta1: float = 0.9,
        beta2: float = 0.999,
        eps: float = 1e-8,
        step: int = 1
    ) -> 'FluxTensor':
        """
        Adam optimizer step with flux tracking.
        
        Integrates Adam's adaptive learning with flux algebra's uncertainty.
        """
        grad = torch.autograd.grad(loss, self.v, retain_graph=True)[0]
        
        # Adam momentum updates
        new_t = beta1 * self.t + (1 - beta1) * grad
        new_s = beta2 * self.s + (1 - beta2) * grad**2
        
        # Bias correction
        t_corrected = new_t / (1 - beta1**step)
        s_corrected = new_s / (1 - beta2**step)
        
        # Value update
        new_v = self.v - lr * t_corrected / (torch.sqrt(s_corrected) + eps)
        
        return FluxTensor(new_v, torch.sqrt(torch.clamp(s_corrected, min=0)), new_t)
    
    # ==================== PROPERTIES & METRICS ====================
    
    def entropy_velocity(self, prev_s: torch.Tensor, dt: float = 1.0) -> torch.Tensor:
        """
        Compute entropy velocity (σ̇).
        
        σ̇ < 0: System is collapsing (learning/stabilizing)
        σ̇ > 0: System is expanding (hallucinating/diverging)
        """
        return (self.s - prev_s) / dt
    
    def is_collapsing(self, prev_s: torch.Tensor, threshold: float = 0.0) -> torch.Tensor:
        """Check if entropy is collapsing (learning)."""
        return self.entropy_velocity(prev_s) < threshold
    
    def total_entropy(self) -> torch.Tensor:
        """Get total entropy (sum over all elements)."""
        return self.s.sum()
    
    def entropy_norm(self) -> torch.Tensor:
        """Get L2 norm of entropy."""
        return torch.norm(self.s)
    
    def to_tensor(self) -> torch.Tensor:
        """Extract value tensor."""
        return self.v
    
    def item(self):
        """Get single value (only for scalar tensors)."""
        if self.v.numel() != 1:
            raise ValueError("Can only call .item() on scalar FluxTensor")
        return self.v.item()


# ==================== FLUX NEURAL NETWORK LAYER ====================

class FluxLinear(nn.Module):
    """
    Linear layer with Flux Algebra tracking.
    
    Wraps nn.Linear to track weight uncertainty and evolution.
    """
    
    def __init__(self, in_features: int, out_features: int, bias: bool = True):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features, bias=bias)
        
        # Initialize flux state for weights
        self.weight_flux = FluxTensor(
            self.linear.weight.data,
            s=torch.ones_like(self.linear.weight.data) * 0.1,  # Initial uncertainty
            t=torch.zeros_like(self.linear.weight.data)
        )
        
        if bias:
            self.bias_flux = FluxTensor(
                self.linear.bias.data,
                s=torch.ones_like(self.linear.bias.data) * 0.1,
                t=torch.zeros_like(self.linear.bias.data)
            )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass using current weight values."""
        # Use F.linear with flux tensor values (maintains computation graph)
        import torch.nn.functional as F
        result = F.linear(x, self.weight_flux.v, self.bias_flux.v if self.linear.bias is not None else None)
        return result
    
    def flux_update(self, loss: torch.Tensor, lr: float = 0.01, work: float = 0.1):
        """
        Update flux state after computing loss.
        
        Args:
            loss: Loss tensor
            lr: Learning rate
            work: Work budget for collapse
        """
        # Update weights using standard optimizer step
        with torch.no_grad():
            # Compute gradient step
            grad = self.weight_flux.v.grad
            if grad is None:
                # Manually compute if needed
                grad = torch.autograd.grad(loss, self.weight_flux.v, retain_graph=True)[0]
            
            delta = -lr * grad
            new_v = self.weight_flux.v + delta
            new_t = self.weight_flux.t + delta
            new_s = self.weight_flux.s + torch.abs(delta)
            
            # Apply collapse
            if work > 0:
                new_s = torch.clamp(new_s - work, min=0.0)
            
            # Update the flux state
            self.weight_flux = FluxTensor(new_v, new_s, new_t)
            self.linear.weight.data = new_v
            
            # Same for bias
            if self.linear.bias is not None:
                grad_bias = self.bias_flux.v.grad
                if grad_bias is None:
                    grad_bias = torch.autograd.grad(loss, self.bias_flux.v, retain_graph=True)[0]
                
                delta_bias = -lr * grad_bias
                new_v_bias = self.bias_flux.v + delta_bias
                new_t_bias = self.bias_flux.t + delta_bias
                new_s_bias = self.bias_flux.s + torch.abs(delta_bias)
                
                if work > 0:
                    new_s_bias = torch.clamp(new_s_bias - work, min=0.0)
                
                self.bias_flux = FluxTensor(new_v_bias, new_s_bias, new_t_bias)
                self.linear.bias.data = new_v_bias
    
    def collapse(self, work: float):
        """Collapse entropy for weights and biases."""
        self.weight_flux = self.weight_flux.collapse(work)
        if self.linear.bias is not None:
            self.bias_flux = self.bias_flux.collapse(work)
    
    def entropy_report(self) -> dict:
        """Report current entropy levels."""
        report = {
            'weight_entropy_mean': self.weight_flux.s.mean().item(),
            'weight_entropy_total': self.weight_flux.total_entropy().item(),
        }
        if self.linear.bias is not None:
            report['bias_entropy_mean'] = self.bias_flux.s.mean().item()
            report['bias_entropy_total'] = self.bias_flux.total_entropy().item()
        return report


# ==================== FLUX OPTIMIZER ====================

class FluxOptimizer:
    """
    Optimizer wrapper that applies flux algebra to standard PyTorch optimizers.
    
    Tracks uncertainty and applies collapse operators during training.
    """
    
    def __init__(self, optimizer: torch.optim.Optimizer, work_schedule: Optional[callable] = None):
        """
        Initialize Flux Optimizer.
        
        Args:
            optimizer: PyTorch optimizer
            work_schedule: Function(step) -> work_budget (optional)
        """
        self.optimizer = optimizer
        self.work_schedule = work_schedule or (lambda step: 0.01)
        self.step_count = 0
        
        # Initialize flux tracking for all parameters
        self.flux_states = {}
        for group in optimizer.param_groups:
            for param in group['params']:
                self.flux_states[id(param)] = FluxTensor(
                    param.data,
                    s=torch.ones_like(param.data) * 0.1,
                    t=torch.zeros_like(param.data)
                )
    
    def step(self, loss: Optional[torch.Tensor] = None):
        """
        Perform optimization step with flux tracking.
        
        Args:
            loss: Loss tensor (if None, uses standard optimizer step)
        """
        self.step_count += 1
        work = self.work_schedule(self.step_count)
        
        # Standard optimizer step
        if loss is not None:
            loss.backward()
        
        self.optimizer.step()
        self.optimizer.zero_grad()
        
        # Update flux states
        for group in self.optimizer.param_groups:
            for param in group['params']:
                param_id = id(param)
                old_flux = self.flux_states[param_id]
                
                # Update flux with new parameter values
                delta = param.data - old_flux.v
                new_t = old_flux.t + delta
                new_s = old_flux.s + torch.abs(delta)
                
                # Apply collapse
                new_s = torch.clamp(new_s - work, min=0.0)
                
                self.flux_states[param_id] = FluxTensor(
                    param.data.clone(),
                    new_s,
                    new_t
                )
    
    def get_entropy_stats(self) -> dict:
        """Get aggregate entropy statistics."""
        total_entropy = 0
        max_entropy = 0
        count = 0
        
        for param_id, flux in self.flux_states.items():
            total_entropy += flux.total_entropy().item()
            max_entropy = max(max_entropy, flux.s.max().item())
            count += 1
        
        return {
            'total_entropy': total_entropy,
            'max_entropy': max_entropy,
            'avg_entropy': total_entropy / count if count > 0 else 0,
            'step': self.step_count
        }
    
    def collapse_all(self, work: float):
        """Apply collapse to all tracked parameters."""
        for param_id, flux in self.flux_states.items():
            self.flux_states[param_id] = flux.collapse(work)


# ==================== EXAMPLE USAGE ====================

def example_basic_flux_tensor():
    """Basic example of FluxTensor operations."""
    print("\n" + "="*60)
    print("Example 1: Basic FluxTensor Operations")
    print("="*60)
    
    # Create flux tensors
    z1 = FluxTensor(
        v=torch.tensor([2.0, 3.0, 4.0]),
        s=torch.tensor([0.5, 0.3, 0.2]),
        t=torch.tensor([0.1, 0.2, 0.15])
    )
    
    z2 = FluxTensor(
        v=torch.tensor([1.0, 2.0, 3.0]),
        s=torch.tensor([0.3, 0.4, 0.1]),
        t=torch.tensor([0.05, 0.1, 0.2])
    )
    
    print(f"z1 = {z1}")
    print(f"z2 = {z2}")
    
    # Addition
    z_add = z1 + z2
    print(f"z1 ⊕ z2 = {z_add}")
    
    # Multiplication
    z_mul = z1 * z2
    print(f"z1 ⊗ z2 = {z_mul}")
    
    # Collapse
    z_collapsed = z1.collapse(work=0.2)
    print(f"𝒞(z1, W=0.2) = {z_collapsed}")
    
    # Evolution
    z_evolved = z1.evolve(dt=0.5, entropy_decay=0.01)
    print(f"Evolve(z1, dt=0.5) = {z_evolved}")


def example_matrix_operations():
    """Example with matrix operations."""
    print("\n" + "="*60)
    print("Example 2: Matrix Operations")
    print("="*60)
    
    # Create flux matrices
    A = FluxTensor(
        v=torch.tensor([[1.0, 2.0], [3.0, 4.0]]),
        s=torch.tensor([[0.1, 0.2], [0.15, 0.1]]),
        t=torch.tensor([[0.05, 0.1], [0.08, 0.06]])
    )
    
    B = FluxTensor(
        v=torch.tensor([[2.0, 0.0], [1.0, 3.0]]),
        s=torch.tensor([[0.05, 0.1], [0.2, 0.15]]),
        t=torch.tensor([[0.1, 0.05], [0.07, 0.09]])
    )
    
    print(f"A = {A}")
    print(f"B = {B}")
    
    # Matrix multiplication
    C = A @ B
    print(f"A @ B = {C}")


def example_ml_training():
    """Example of using flux algebra in ML training."""
    print("\n" + "="*60)
    print("Example 3: ML Training with Flux Algebra")
    print("="*60)
    
    # Simple regression problem
    torch.manual_seed(42)
    
    # Create a standard linear layer
    layer = nn.Linear(10, 5)
    
    # Initialize flux tracking for weights
    weight_flux = FluxTensor(
        layer.weight.data.clone(),
        s=torch.ones_like(layer.weight.data) * 0.1,
        t=torch.zeros_like(layer.weight.data)
    )
    bias_flux = FluxTensor(
        layer.bias.data.clone(),
        s=torch.ones_like(layer.bias.data) * 0.1,
        t=torch.zeros_like(layer.bias.data)
    )
    
    # Dummy data
    x = torch.randn(32, 10)
    y = torch.randn(32, 5)
    
    # Training loop
    print("\nTraining with flux tracking...")
    for epoch in range(5):
        # Forward pass (using actual layer weights)
        pred = layer(x)
        loss = ((pred - y)**2).mean()
        
        # Backward pass
        loss.backward()
        
        # Standard gradient step
        with torch.no_grad():
            lr = 0.01
            work = 0.05
            
            # Update weights
            grad_w = layer.weight.grad
            delta_w = -lr * grad_w
            weight_flux = FluxTensor(
                layer.weight + delta_w,
                weight_flux.s + torch.abs(delta_w),
                weight_flux.t + delta_w
            ).collapse(work)
            layer.weight.data = weight_flux.v
            
            # Update bias
            grad_b = layer.bias.grad
            delta_b = -lr * grad_b
            bias_flux = FluxTensor(
                layer.bias + delta_b,
                bias_flux.s + torch.abs(delta_b),
                bias_flux.t + delta_b
            ).collapse(work)
            layer.bias.data = bias_flux.v
        
        # Zero gradients
        layer.zero_grad()
        
        # Report
        print(f"Epoch {epoch+1}: Loss={loss.item():.4f}, "
              f"Weight Entropy={weight_flux.s.mean().item():.4f}")
    
    # Final collapse
    weight_flux = weight_flux.collapse(0.2)
    print(f"\nAfter collapse: Weight Entropy={weight_flux.s.mean().item():.4f}")


def example_flux_optimizer():
    """Example using FluxOptimizer with PyTorch."""
    print("\n" + "="*60)
    print("Example 4: FluxOptimizer")
    print("="*60)
    
    # Simple model
    model = nn.Sequential(
        nn.Linear(10, 20),
        nn.ReLU(),
        nn.Linear(20, 5)
    )
    
    # Wrap with flux optimizer
    base_optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    optimizer = FluxOptimizer(
        base_optimizer,
        work_schedule=lambda step: 0.01 * (1 + step * 0.1)
    )
    
    # Dummy training
    x = torch.randn(64, 10)
    y = torch.randn(64, 5)
    
    print("\nTraining with FluxOptimizer...")
    for epoch in range(5):
        pred = model(x)
        loss = ((pred - y)**2).mean()
        
        optimizer.step(loss)
        
        stats = optimizer.get_entropy_stats()
        print(f"Epoch {epoch+1}: Loss={loss.item():.4f}, "
              f"Total Entropy={stats['total_entropy']:.4f}")


if __name__ == "__main__":
    print("Flux Algebra for PyTorch - Matrix/Tensor Implementation")
    print("Based on: The Algebra of Flux (Conditional Collapse Theory)")
    
    # Run examples
    example_basic_flux_tensor()
    example_matrix_operations()
    example_ml_training()
    example_flux_optimizer()
    
    print("\n" + "="*60)
    print("All examples completed successfully!")
    print("="*60)