"""
Sparse Tensor — Lightweight ndarray wrapper with brain-inspired sparsity.

The brain is lazy and sparse: ~1-5% of neurons fire at any moment.
This module provides sparse operations that model this biological constraint.
Sparsity enables "common sense" — knowing ~60% is enough, then filling gaps.

Author: Algorembrant, Rembrant Oyangoren Albeos (2026)
"""

import numpy as np
from typing import Optional, Tuple, Union


class SparseTensor:
    """
    A sparse tensor wrapper over NumPy arrays that enforces biological sparsity.
    
    Key biological properties:
    - Top-k sparsification (winner-take-all inhibition)
    - Threshold activation (firing threshold)
    - Lazy computation (only compute when needed)
    - Efficient sparse dot products
    """

    def __init__(self, data: np.ndarray, sparsity_mask: Optional[np.ndarray] = None):
        """
        Args:
            data: Dense NumPy array (the raw signal)
            sparsity_mask: Boolean mask of active units (True = active/firing)
        """
        self._data = np.asarray(data, dtype=np.float64)
        if sparsity_mask is not None:
            self._mask = np.asarray(sparsity_mask, dtype=bool)
            assert self._mask.shape == self._data.shape, \
                f"Mask shape {self._mask.shape} != data shape {self._data.shape}"
        else:
            # Fully dense by default (all active)
            self._mask = np.ones(self._data.shape, dtype=bool)

    @property
    def data(self) -> np.ndarray:
        """Raw underlying data (masked values are zeroed)."""
        return self._data * self._mask

    @property
    def dense(self) -> np.ndarray:
        """Full dense representation (unmasked)."""
        return self._data

    @property
    def mask(self) -> np.ndarray:
        """Boolean sparsity mask: True where units are active."""
        return self._mask

    @property
    def shape(self) -> Tuple[int, ...]:
        return self._data.shape

    @property
    def sparsity(self) -> float:
        """Fraction of zeros (inactive units). 0.0 = fully dense, 1.0 = fully sparse."""
        return 1.0 - (np.sum(self._mask) / self._mask.size)

    @property
    def num_active(self) -> int:
        """Number of active (non-zero) units."""
        return int(np.sum(self._mask))

    # ----- Sparsification Operations (biological inhibition) -----

    def threshold(self, theta: float) -> 'SparseTensor':
        """
        Threshold activation — only units above theta fire.
        Models neuronal firing threshold / activation threshold.
        
        Args:
            theta: Firing threshold
        Returns:
            New SparseTensor with sub-threshold units masked out
        """
        new_mask = self._mask & (np.abs(self._data) >= theta)
        return SparseTensor(self._data, new_mask)

    def top_k(self, k: int, axis: Optional[int] = None) -> 'SparseTensor':
        """
        Top-k sparsification — winner-take-all competitive inhibition.
        Only the k strongest activations survive. This is how lateral inhibition
        in cortex creates sparse population codes.
        
        Args:
            k: Number of top activations to keep
            axis: Axis along which to apply top-k (None = global)
        Returns:
            New SparseTensor with only top-k values active
        """
        if axis is None:
            flat = np.abs(self._data).ravel()
            if k >= flat.size:
                return SparseTensor(self._data.copy(), self._mask.copy())
            # Find the k-th largest value
            threshold_val = np.partition(flat, -k)[-k]
            new_mask = self._mask & (np.abs(self._data) >= threshold_val)
            # If too many elements equal the threshold, keep only k total
            active_count = np.sum(new_mask)
            if active_count > k:
                active_indices = np.argwhere(new_mask.ravel()).ravel()
                active_vals = np.abs(self._data.ravel()[active_indices])
                # Sort by value descending, keep only k
                sorted_order = np.argsort(-active_vals)
                kill = active_indices[sorted_order[k:]]
                flat_mask = new_mask.ravel().copy()
                flat_mask[kill] = False
                new_mask = flat_mask.reshape(self._data.shape)
            return SparseTensor(self._data, new_mask)
        else:
            # Per-slice top-k along axis
            new_mask = np.zeros_like(self._mask)
            nd = self._data.ndim
            slices = [slice(None)] * nd
            for i in range(self._data.shape[axis]):
                slices[axis] = i
                sl = tuple(slices)
                vals = np.abs(self._data[sl])
                flat = vals.ravel()
                actual_k = min(k, flat.size)
                if actual_k == flat.size:
                    new_mask[sl] = self._mask[sl]
                else:
                    thresh = np.partition(flat, -actual_k)[-actual_k]
                    new_mask[sl] = self._mask[sl] & (vals >= thresh)
            return SparseTensor(self._data, new_mask)

    def sparsify(self, target_sparsity: float) -> 'SparseTensor':
        """
        Achieve a target sparsity level (fraction of zeros).
        Brain typically has 95-99% sparsity in any population code.
        
        Args:
            target_sparsity: Desired fraction of inactive units (0.0 to 1.0)
        Returns:
            New SparseTensor with approximately target_sparsity inactive
        """
        k = max(1, round((1.0 - target_sparsity) * self._data.size))
        return self.top_k(k)

    # ----- Activation Functions (neuronal nonlinearities) -----

    def relu(self) -> 'SparseTensor':
        """Half-wave rectification — models neuronal firing rate (no negative rates)."""
        new_data = np.maximum(self._data, 0.0)
        new_mask = self._mask & (new_data > 0.0)
        return SparseTensor(new_data, new_mask)

    def sigmoid(self, gain: float = 1.0) -> 'SparseTensor':
        """Sigmoidal activation — saturating firing rate."""
        new_data = 1.0 / (1.0 + np.exp(-gain * self._data))
        return SparseTensor(new_data, self._mask)

    def softmax(self, axis: int = -1) -> 'SparseTensor':
        """Softmax normalization — competitive normalization across a population."""
        shifted = self._data - np.max(self._data, axis=axis, keepdims=True)
        exp_vals = np.exp(shifted) * self._mask
        sums = np.sum(exp_vals, axis=axis, keepdims=True)
        sums = np.where(sums == 0, 1.0, sums)
        new_data = exp_vals / sums
        return SparseTensor(new_data, self._mask)

    def divisive_normalization(self, sigma: float = 1.0, axis: int = -1) -> 'SparseTensor':
        """
        Divisive normalization — the canonical neural computation.
        r_i = r_i^n / (sigma^n + sum(r_j^n))
        Models gain control in visual cortex.
        """
        n = 2.0  # Exponent (typically 2)
        powered = np.abs(self.data) ** n
        pool = np.sum(powered, axis=axis, keepdims=True)
        normalized = powered / (sigma ** n + pool)
        # Restore sign
        signs = np.sign(self._data)
        new_data = signs * (normalized ** (1.0 / n))
        return SparseTensor(new_data, self._mask)

    # ----- Linear Algebra -----

    def dot(self, other: Union['SparseTensor', np.ndarray]) -> 'SparseTensor':
        """
        Sparse dot product — only active units contribute.
        Efficient because inactive synapses don't transmit.
        """
        if isinstance(other, SparseTensor):
            result = np.dot(self.data, other.data)
        else:
            result = np.dot(self.data, np.asarray(other, dtype=np.float64))
        return SparseTensor(result)

    def outer(self, other: 'SparseTensor') -> 'SparseTensor':
        """Outer product — used for Hebbian learning (pre × post)."""
        result = np.outer(self.data.ravel(), other.data.ravel())
        return SparseTensor(result)

    # ----- Element-wise operations -----

    def __add__(self, other: Union['SparseTensor', np.ndarray, float]) -> 'SparseTensor':
        if isinstance(other, SparseTensor):
            return SparseTensor(self._data + other._data, self._mask | other._mask)
        return SparseTensor(self._data + np.float64(other), self._mask)

    def __sub__(self, other: Union['SparseTensor', np.ndarray, float]) -> 'SparseTensor':
        if isinstance(other, SparseTensor):
            return SparseTensor(self._data - other._data, self._mask | other._mask)
        return SparseTensor(self._data - np.float64(other), self._mask)

    def __mul__(self, other: Union['SparseTensor', np.ndarray, float]) -> 'SparseTensor':
        if isinstance(other, SparseTensor):
            return SparseTensor(self._data * other._data, self._mask & other._mask)
        return SparseTensor(self._data * np.float64(other), self._mask)

    def __neg__(self) -> 'SparseTensor':
        return SparseTensor(-self._data, self._mask.copy())

    def __repr__(self) -> str:
        return (f"SparseTensor(shape={self.shape}, "
                f"active={self.num_active}/{self._data.size}, "
                f"sparsity={self.sparsity:.1%})")

    # ----- Utility -----

    def copy(self) -> 'SparseTensor':
        return SparseTensor(self._data.copy(), self._mask.copy())

    def reshape(self, *shape) -> 'SparseTensor':
        return SparseTensor(self._data.reshape(*shape), self._mask.reshape(*shape))

    def flatten(self) -> 'SparseTensor':
        return SparseTensor(self._data.ravel(), self._mask.ravel())

    @staticmethod
    def from_dense(data: np.ndarray, threshold: float = 0.0) -> 'SparseTensor':
        """Create from dense array, automatically masking near-zero values."""
        mask = np.abs(data) > threshold
        return SparseTensor(data, mask)

    @staticmethod
    def zeros(shape: Tuple[int, ...]) -> 'SparseTensor':
        return SparseTensor(np.zeros(shape), np.zeros(shape, dtype=bool))

    @staticmethod
    def ones(shape: Tuple[int, ...]) -> 'SparseTensor':
        return SparseTensor(np.ones(shape))

    @staticmethod
    def random(shape: Tuple[int, ...], sparsity: float = 0.95) -> 'SparseTensor':
        """Random sparse tensor — models spontaneous neural activity."""
        data = np.random.randn(*shape)
        mask = np.random.random(shape) > sparsity
        return SparseTensor(data, mask)