import numpy as np


class FluxMatrix:
    """
    A matrix in Flux Algebra where each element is a Flux Number.
    
    The matrix is represented as three aligned matrices:
    - V: Value matrix (the current "mean" state)
    - S: Entropy matrix (uncertainty/spread of each element)
    - T: Flux matrix (rate of change/velocity of each element)
    
    This enables tracking uncertainty and evolution in ML weights,
    gradients, and activations.
    """
    
    def __init__(self, V, S=None, T=None):
        """
        Initialize a FluxMatrix.
        
        Args:
            V: Value matrix (numpy array or scalar)
            S: Entropy matrix (same shape as V, default: zeros)
            T: Flux matrix (same shape as V, default: zeros)
        """
        self.V = np.array(V, dtype=np.float64)
        if S is None:
            self.S = np.zeros_like(self.V)
        else:
            self.S = np.array(S, dtype=np.float64)
        if T is None:
            self.T = np.zeros_like(self.V)
        else:
            self.T = np.array(T, dtype=np.float64)
        
        assert self.V.shape == self.S.shape == self.T.shape, \
            "V, S, T must have the same shape"
    
    @property
    def shape(self):
        return self.V.shape
    
    def __repr__(self):
        return (f"FluxMatrix(shape={self.shape}, "
                f"V_range=[{self.V.min():.3f}, {self.V.max():.3f}], "
                f"S_range=[{self.S.min():.3f}, {self.S.max():.3f}], "
                f"T_range=[{self.T.min():.3f}, {self.T.max():.3f}])")
    
    # ──────────────────────────────────────────────────────────
    # FLUX ALGEBRA OPERATORS
    # ──────────────────────────────────────────────────────────
    
    def __add__(self, other):
        """
        Flux Addition (⊕): Merging of two trajectories.
        
        z1 ⊕ z2 = ⟨v1 + v2, √(σ1² + σ2²), τ1 + τ2⟩
        """
        if isinstance(other, (int, float)):
            return FluxMatrix(
                self.V + other,
                self.S.copy(),
                self.T.copy()
            )
        elif isinstance(other, FluxMatrix):
            new_V = self.V + other.V
            new_S = np.sqrt(self.S**2 + other.S**2)
            new_T = self.T + other.T
            return FluxMatrix(new_V, new_S, new_T)
        else:
            raise TypeError(f"Cannot add FluxMatrix and {type(other)}")
    
    def __radd__(self, other):
        return self.__add__(other)
    
    def __sub__(self, other):
        """Flux Subtraction."""
        if isinstance(other, (int, float)):
            return FluxMatrix(
                self.V - other,
                self.S.copy(),
                self.T.copy()
            )
        elif isinstance(other, FluxMatrix):
            new_V = self.V - other.V
            new_S = np.sqrt(self.S**2 + other.S**2)
            new_T = self.T - other.T
            return FluxMatrix(new_V, new_S, new_T)
        else:
            raise TypeError(f"Cannot subtract {type(other)} from FluxMatrix")
    
    def __mul__(self, scalar):
        """
        Scalar Multiplication.
        
        α ⊗ z = ⟨αv, |α|σ, ατ⟩
        """
        if isinstance(scalar, (int, float)):
            new_V = self.V * scalar
            new_S = self.S * abs(scalar)
            new_T = self.T * scalar
            return FluxMatrix(new_V, new_S, new_T)
        else:
            raise TypeError(f"Cannot multiply FluxMatrix by {type(scalar)}")
    
    def __rmul__(self, scalar):
        return self.__mul__(scalar)
    
    def hadamard(self, other):
        """
        Element-wise Flux Multiplication (⊗).
        
        z1 ⊗ z2 = ⟨v1*v2, √((v1*σ2)² + (v2*σ1)²), v1*τ2 + v2*τ1⟩
        """
        if isinstance(other, FluxMatrix):
            new_V = self.V * other.V
            new_S = np.sqrt((self.V * other.S)**2 + (other.V * self.S)**2)
            new_T = self.V * other.T + other.V * self.T
            return FluxMatrix(new_V, new_S, new_T)
        elif isinstance(other, (int, float)):
            return self.__mul__(other)
        else:
            raise TypeError(f"Cannot hadamard FluxMatrix and {type(other)}")
    
    def matmul(self, other):
        """
        Flux Matrix Multiplication.
        
        For C = A @ B:
        C[i,j] = Σ_k A[i,k] ⊗ B[k,j]
        
        This propagates entropy through the matrix multiplication chain.
        """
        if not isinstance(other, FluxMatrix):
            raise TypeError(f"Cannot matmul FluxMatrix and {type(other)}")
        
        # Standard matrix multiplication for values
        new_V = self.V @ other.V
        
        # Entropy propagation through matrix multiplication
        # For each element C[i,j] = Σ_k A[i,k] * B[k,j]
        # Entropy: σ_C[i,j] = √(Σ_k [(A[i,k]*σ_B[k,j])² + (B[k,j]*σ_A[i,k])²])
        new_S = np.zeros(new_V.shape)
        new_T = np.zeros(new_V.shape)
        
        for i in range(self.V.shape[0]):
            for j in range(other.V.shape[1]):
                s_sq_sum = 0.0
                t_sum = 0.0
                for k in range(self.V.shape[1]):
                    # Entropy accumulation
                    s_sq_sum += (self.V[i, k] * other.S[k, j])**2 + \
                                (other.V[k, j] * self.S[i, k])**2
                    # Flux accumulation (product rule)
                    t_sum += self.V[i, k] * other.T[k, j] + \
                             other.V[k, j] * self.T[i, k]
                new_S[i, j] = np.sqrt(s_sq_sum)
                new_T[i, j] = t_sum
        
        return FluxMatrix(new_V, new_S, new_T)
    
    def __matmul__(self, other):
        """Matrix multiplication operator (@)."""
        return self.matmul(other)
    
    # ──────────────────────────────────────────────────────────
    # FLUX OPERATORS
    # ──────────────────────────────────────────────────────────
    
    def collapse(self, W):
        """
        Collapse Operator (C): Reduce entropy by paying Work.
        
        C(z, W) = ⟨v, max(0, σ - W), τ⟩
        
        Args:
            W: Work budget (scalar or matrix of same shape)
        
        Returns:
            New FluxMatrix with reduced entropy
        """
        if isinstance(W, (int, float)):
            new_S = np.maximum(0, self.S - W)
        else:
            new_S = np.maximum(0, self.S - W)
        return FluxMatrix(self.V.copy(), new_S, self.T.copy())
    
    def evolve(self, dt):
        """
        Flux Derivative / ODE step.
        
        d/dt ⟨v, σ, τ⟩ = ⟨τ, σ̇, τ̇⟩
        
        For discrete evolution:
        - v_new = v + τ * dt
        - σ_new = σ * decay (natural entropy decay)
        - τ stays constant (unless specified otherwise)
        
        Args:
            dt: Time step
        
        Returns:
            New evolved FluxMatrix
        """
        new_V = self.V + self.T * dt
        new_S = self.S * 0.99  # Natural entropy decay
        return FluxMatrix(new_V, new_S, self.T.copy())
    
    def entropy_velocity(self, dt=1.0):
        """
        Compute entropy velocity (σ̇).
        
        σ̇ = (σ_t - σ_{t-1}) / dt
        
        Returns:
            Matrix of entropy velocities
        """
        return self.S / dt
    
    def is_collapsing(self, threshold=0.0):
        """
        Check if the system is collapsing (learning/stabilizing).
        
        Returns:
            Boolean matrix where True indicates collapsing elements
        """
        return self.entropy_velocity() < threshold
    
    def is_expanding(self, threshold=0.0):
        """
        Check if the system is expanding (hallucinating/diverging).
        
        Returns:
            Boolean matrix where True indicates expanding elements
        """
        return self.entropy_velocity() > threshold
    
    # ──────────────────────────────────────────────────────────
    # ML-SPECIFIC OPERATIONS
    # ──────────────────────────────────────────────────────────
    
    @classmethod
    def from_random(cls, shape, v_scale=0.1, s_init=0.01, t_scale=0.001):
        """
        Create a random FluxMatrix for initializing ML weights.
        
        Args:
            shape: Tuple specifying matrix shape
            v_scale: Scale for value initialization (Xavier-like)
            s_init: Initial entropy
            t_scale: Scale for flux initialization
        
        Returns:
            Random FluxMatrix
        """
        V = np.random.randn(*shape) * v_scale
        S = np.ones(shape) * s_init
        T = np.random.randn(*shape) * t_scale
        return cls(V, S, T)
    
    @classmethod
    def from_array(cls, array, entropy=0.01, flux=0.0):
        """
        Create a FluxMatrix from a standard numpy array.
        
        Args:
            array: numpy array of values
            entropy: Initial entropy for all elements
            flux: Initial flux for all elements
        
        Returns:
            FluxMatrix with given values and uniform entropy/flux
        """
        V = np.array(array, dtype=np.float64)
        S = np.ones_like(V) * entropy
        T = np.ones_like(V) * flux
        return cls(V, S, T)
    
    def gradient_step(self, gradient, learning_rate=0.01, work_budget=0.001):
        """
        Perform a gradient descent step in Flux Algebra.
        
        This is the Flux Algebra equivalent of:
        θ_new = θ_old - lr * gradient
        
        Implemented as:
        θ_{t+1} = θ_t ⊕ (-lr * gradient)
        Then collapse to reduce uncertainty.
        
        Args:
            gradient: FluxMatrix or standard matrix of gradients
            learning_rate: Learning rate (scalar)
            work_budget: Work to pay for collapse
        
        Returns:
            Updated FluxMatrix
        """
        if isinstance(gradient, FluxMatrix):
            delta = gradient * (-learning_rate)
        else:
            # Convert standard gradient to FluxMatrix
            delta_V = -learning_rate * gradient
            delta_S = np.zeros_like(delta_V)
            delta_T = gradient.copy() if hasattr(gradient, 'copy') else np.array(gradient)
            delta = FluxMatrix(delta_V, delta_S, delta_T)
        
        # Flux addition: θ ⊕ Δθ
        new_self = self + delta
        
        # Collapse to reduce entropy (learning = entropy collapse)
        return new_self.collapse(work_budget)
    
    def predict(self, X):
        """
        Make predictions with uncertainty propagation.
        
        For y = X @ W:
        Returns both prediction and uncertainty.
        
        Args:
            X: Input FluxMatrix or standard matrix
        
        Returns:
            Tuple of (prediction, uncertainty) matrices
        """
        if isinstance(X, FluxMatrix):
            result = X @ self
            return result.V, result.S
        else:
            # Standard matrix multiplication, propagate entropy
            X_arr = np.array(X)
            result_V = X_arr @ self.V
            result_S = np.sqrt((X_arr**2 @ self.S**2))
            return result_V, result_S
    
    def total_entropy(self):
        """Get total entropy of the matrix (Frobenius norm of S)."""
        return np.linalg.norm(self.S, 'fro')
    
    def total_flux(self):
        """Get total flux magnitude (Frobenius norm of T)."""
        return np.linalg.norm(self.T, 'fro')
    
    def stability_score(self):
        """
        Compute stability score: ratio of collapsed entropy to total.
        
        Higher score = more stable/certain system.
        """
        total = self.total_entropy()
        if total == 0:
            return 1.0
        collapsed = np.sum(np.maximum(0, -self.entropy_velocity()))
        return 1.0 - (collapsed / total)
    
    def to_array(self):
        """Extract just the value matrix as a numpy array."""
        return self.V.copy()
    
    def copy(self):
        """Create a deep copy of this FluxMatrix."""
        return FluxMatrix(self.V.copy(), self.S.copy(), self.T.copy())


# ──────────────────────────────────────────────────────────────
# CONVENIENCE FUNCTIONS
# ──────────────────────────────────────────────────────────────

def flux_eye(n, s=0.0, t=0.0):
    """Create a FluxMatrix identity matrix."""
    return FluxMatrix(np.eye(n), np.zeros((n, n)) + s, np.zeros((n, n)) + t)


def flux_zeros(shape):
    """Create a FluxMatrix of zeros."""
    return FluxMatrix(np.zeros(shape))


def flux_ones(shape):
    """Create a FluxMatrix of ones."""
    return FluxMatrix(np.ones(shape))


# ──────────────────────────────────────────────────────────────
# EXAMPLE: FLUX MATRIX IN A NEURAL NETWORK LAYER
# ──────────────────────────────────────────────────────────────

class FluxLinearLayer:
    """
    A linear layer using Flux Algebra for uncertainty-aware learning.
    
    This tracks not just the weights, but their uncertainty and
    evolution over time—enabling automatic confidence estimation.
    """
    
    def __init__(self, in_features, out_features, init_entropy=0.01):
        # Initialize weights as FluxMatrix
        scale = np.sqrt(2.0 / in_features)  # He initialization
        self.weights = FluxMatrix.from_random(
            (in_features, out_features),
            v_scale=scale,
            s_init=init_entropy,
            t_scale=scale * 0.01
        )
        self.bias = FluxMatrix.from_random(
            (1, out_features),
            v_scale=0.0,
            s_init=init_entropy,
            t_scale=0.001
        )
    
    def forward(self, X):
        """
        Forward pass with uncertainty propagation.
        
        Args:
            X: Input matrix (standard numpy array)
        
        Returns:
            Tuple of (output_value, output_entropy)
        """
        # X @ W + b
        out_V = X @ self.weights.V + self.bias.V
        out_S = np.sqrt((X**2 @ self.weights.S**2) + self.bias.S**2)
        return out_V, out_S
    
    def update(self, grad_weights, grad_bias, lr=0.01, work=0.001):
        """
        Update weights using gradient descent with entropy collapse.
        
        Args:
            grad_weights: Gradient for weights
            grad_bias: Gradient for bias
            lr: Learning rate
            work: Work budget for collapse
        """
        self.weights = self.weights.gradient_step(grad_weights, lr, work)
        self.bias = self.bias.gradient_step(grad_bias, lr, work)
    
    def confidence(self):
        """
        Get confidence in this layer's weights.
        1.0 = fully confident, 0.0 = maximum uncertainty.
        """
        total_entropy = self.weights.total_entropy()
        # Normalize to [0, 1] range
        return np.exp(-total_entropy)


# ──────────────────────────────────────────────────────────────
# DEMONSTRATION
# ──────────────────────────────────────────────────────────────

if __name__ == "__main__":
    print("=" * 60)
    print("FLUX ALGEBRA FOR MATRICES - ML Iteration Demo")
    print("=" * 60)
    
    # Example 1: Basic Flux Matrix Operations
    print("\n--- Example 1: Basic Operations ---")
    A = FluxMatrix.from_random((3, 3), v_scale=1.0, s_init=0.5, t_scale=0.1)
    B = FluxMatrix.from_random((3, 3), v_scale=1.0, s_init=0.3, t_scale=0.05)
    
    print(f"A = {A}")
    print(f"B = {B}")
    
    C = A + B
    print(f"A ⊕ B = {C}")
    
    D = A.matmul(B)
    print(f"A ⊗ B (matmul) = {D}")
    
    # Example 2: Learning as Entropy Collapse
    print("\n--- Example 2: Learning as Entropy Collapse ---")
    model = FluxMatrix.from_random((4, 2), v_scale=0.5, s_init=1.0, t_scale=0.1)
    print(f"Initial total entropy: {model.total_entropy():.4f}")
    
    work_budget = 0.1
    for epoch in range(5):
        model = model.collapse(work_budget)
        print(f"Epoch {epoch+1}: entropy = {model.total_entropy():.4f}")
    
    # Example 3: Flux Neural Layer
    print("\n--- Example 3: Flux Neural Network Layer ---")
    layer = FluxLinearLayer(in_features=10, out_features=5)
    X = np.random.randn(32, 10)  # Batch of 32 samples
    
    output, uncertainty = layer.forward(X)
    print(f"Input shape: {X.shape}")
    print(f"Output shape: {output.shape}")
    print(f"Output uncertainty (mean): {uncertainty.mean():.4f}")
    print(f"Layer confidence: {layer.confidence():.4f}")
    
    # Simulate training
    print("\nTraining simulation:")
    for epoch in range(10):
        grad_W = np.random.randn(10, 5) * 0.1
        grad_b = np.random.randn(1, 5) * 0.01
        layer.update(grad_W, grad_b, lr=0.01, work=0.005)
        
        if epoch % 3 == 0:
            print(f"  Epoch {epoch}: weight_entropy={layer.weights.total_entropy():.4f}, "
                  f"confidence={layer.confidence():.4f}")
    
    print("\n" + "=" * 60)
    print("Flux Algebra enables uncertainty-aware ML iterations!")
    print("=" * 60)