import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from tqdm import tqdm

# ---------------- Structured Linear Layer ----------------
class BlockDiagonalLowRank(nn.Module):
    """
    A linear layer with structured weight matrix:
    - empty space: zero blocks everywhere except on the block‑diagonal
    - encapsulated internal function: each block is a low‑rank product U @ V^T
    """
    def __init__(self, in_features, out_features, num_blocks, rank, bias=True):
        super().__init__()
        assert in_features % num_blocks == 0, "in_features must be divisible by num_blocks"
        assert out_features % num_blocks == 0, "out_features must be divisible by num_blocks"
        self.num_blocks = num_blocks
        self.rank = rank
        self.in_block = in_features // num_blocks
        self.out_block = out_features // num_blocks

        # For each block, store U (out_block x rank) and V (in_block x rank)
        self.U = nn.Parameter(torch.empty(num_blocks, self.out_block, rank))
        self.V = nn.Parameter(torch.empty(num_blocks, self.in_block, rank))
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.U, a=5**0.5)
        nn.init.kaiming_uniform_(self.V, a=5**0.5)

    def forward(self, x):
        # x: (batch, in_features)
        x_blocks = x.view(-1, self.num_blocks, self.in_block)   # (b, B, inB)

        # V: (B, inB, r) → need to multiply each block's V with the corresponding x_blocks slice
        # We want: for each b, for each B: x_blocks[b, B, :] @ V[B]  → (b, B, r)
        # Using broadcasted matmul: add dims to align the block dimension
        # x_blocks.unsqueeze(2)  : (b, B, 1, inB)
        # V.unsqueeze(0)         : (1, B, inB, r)
        # matmul -> (b, B, 1, r) then squeeze
        x_proj = torch.matmul(x_blocks.unsqueeze(2), self.V.unsqueeze(0)).squeeze(2)  # (b, B, r)

        # U: (B, outB, r) → project to output blocks
        # x_proj: (b, B, r) unsqueeze(2) -> (b, B, 1, r)
        # U.unsqueeze(0).transpose(-1, -2) -> (1, B, r, outB)
        out_blocks = torch.matmul(x_proj.unsqueeze(2),
                                  self.U.unsqueeze(0).transpose(-1, -2)).squeeze(2)  # (b, B, outB)

        out = out_blocks.reshape(-1, self.num_blocks * self.out_block)
        if self.bias is not None:
            out += self.bias
        return out
    
# ---------------- MNIST Classifier ----------------
class StructuredMNISTNet(nn.Module):
    def __init__(self, hidden_size=256, num_blocks=8, rank=8):
        super().__init__()
        # Flatten MNIST images from 28x28 = 784 to hidden_size
        self.linear1 = BlockDiagonalLowRank(784, hidden_size, num_blocks, rank)
        self.linear2 = nn.Linear(hidden_size, 10)  # standard dense for final class.
        self.relu = nn.ReLU()

    def forward(self, x):
        x = x.view(x.size(0), -1)  # flatten
        x = self.relu(self.linear1(x))
        x = self.linear2(x)
        return x

# ---------------- Training & Evaluation ----------------
def train_epoch(model, loader, optimizer, criterion, device):
    model.train()
    total_loss, correct = 0, 0
    for data, target in tqdm(loader, desc='Train', leave=False):
        data, target = data.to(device), target.to(device)
        for i in range(1):
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            if i==0:
                total_loss += loss.item() * data.size(0)
                pred = output.argmax(dim=1)
                correct += pred.eq(target).sum().item()
    return total_loss / len(loader.dataset), correct / len(loader.dataset)

def test(model, loader, criterion, device):
    model.eval()
    total_loss, correct = 0, 0
    with torch.no_grad():
        for data, target in loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            loss = criterion(output, target)
            total_loss += loss.item() * data.size(0)
            pred = output.argmax(dim=1)
            correct += pred.eq(target).sum().item()
    return total_loss / len(loader.dataset), correct / len(loader.dataset)

def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # Hyperparams
    batch_size = 512
    lr = 0.001
    epochs = 5
    hidden_size = 128
    num_blocks = 8
    rank = 8

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

    # Model, optimizer, loss
    model = StructuredMNISTNet(hidden_size, num_blocks, rank).to(device)
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    # Print parameter stats
    total_params = sum(p.numel() for p in model.parameters())
    print(f"Total parameters: {total_params}")

    for epoch in range(1, epochs+1):
        train_loss, train_acc = train_epoch(model, train_loader, optimizer, criterion, device)
        test_loss, test_acc = test(model, test_loader, criterion, device)
        print(f"Epoch {epoch}: Train Loss={train_loss:.4f}, Train Acc={train_acc:.4f} | "
              f"Test Loss={test_loss:.4f}, Test Acc={test_acc:.4f}")

if __name__ == '__main__':
    main()