import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np

# ------------------------------------------------------------
# CCT Linear layer with conditional collapse of dot products
# ------------------------------------------------------------
class CCTLinear(nn.Module):
    def __init__(self, in_features, out_features, rank_threshold=20, sparsity_threshold=0.9):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        
        # Full weight matrix (trainable)
        self.weight = nn.Parameter(torch.Tensor(out_features, in_features))
        self.bias = nn.Parameter(torch.Tensor(out_features))
        nn.init.kaiming_uniform_(self.weight, a=np.sqrt(5))
        fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
        bound = 1 / np.sqrt(fan_in)
        nn.init.uniform_(self.bias, -bound, bound)
        
        # Precomputed low-rank factors (stationary) – updated after training
        self.register_buffer('U', None)
        self.register_buffer('S', None)
        self.register_buffer('Vt', None)
        self.rank_threshold = rank_threshold
        self.sparsity_threshold = sparsity_threshold
        
    def _compute_low_rank(self):
        """Compute SVD of weight for low-rank collapse."""
        U, S, Vt = torch.linalg.svd(self.weight, full_matrices=False)
        # Keep only significant singular values
        k = min(self.rank_threshold, torch.sum(S > 1e-4).item())
        self.U = U[:, :k]
        self.S = S[:k]
        self.Vt = Vt[:k, :]
        return k
    
    def forward(self, x, exact_required=True):
        """
        x: input tensor (batch, in_features)
        exact_required: if True, must return exact result (training or high-stakes)
        """
        batch_size = x.shape[0]
        
        # ---------- CCT Question Sequence ----------
        # Q1: Can we use low-rank approximation?
        if self.U is not None and not exact_required:
            k = self.U.shape[1]
            # Work cost: O(batch * (in*k + k + k*out))
            # Collapse potential: high if effective rank is low
            # Compute y = U @ (S * (Vt @ x.T)).T
            Vt_x = torch.mm(self.Vt, x.T)           # (k, batch)
            S_Vt_x = self.S.unsqueeze(1) * Vt_x     # (k, batch)
            y_low = torch.mm(self.U, S_Vt_x).T      # (batch, out)
            return y_low + self.bias
        
        # Q2: Is the input sparse (many zeros)?
        sparsity = (x == 0).float().mean().item()
        if sparsity > self.sparsity_threshold and not exact_required:
            # Work cost: O(nnz(x) * out)
            # Use torch.sparse but here a simple mask
            mask = x != 0
            # Weighted sum only over non-zero entries
            y_sparse = torch.mm(x * mask, self.weight.T) + self.bias
            return y_sparse
        
        # Q3: Fallback to exact full dot product (worst case)
        return F.linear(x, self.weight, self.bias)
    
    def prepare_for_inference(self):
        """After training, precompute low-rank factors for fast collapse."""
        self._compute_low_rank()
        # Optionally freeze weights
        self.weight.requires_grad = False
        self.bias.requires_grad = False

class MNIST_CCT_Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = CCTLinear(784, 256)
        self.fc2 = CCTLinear(256, 64)
        self.fc3 = CCTLinear(64, 10)
        
    def forward(self, x, exact_required=True):
        x = x.view(-1, 784)
        x = F.relu(self.fc1(x, exact_required))
        x = F.relu(self.fc2(x, exact_required))
        return self.fc3(x, exact_required)

# Data
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_loader = DataLoader(datasets.MNIST('./data', train=True, download=True, transform=transform), batch_size=64, shuffle=True)
test_loader = DataLoader(datasets.MNIST('./data', train=False, transform=transform), batch_size=64)

# Training
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = MNIST_CCT_Net().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

model.train()
for epoch in range(5):
    for data, target in train_loader:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data, exact_required=True)   # exact during training
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
    print(f"Epoch {epoch} done")