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

# ------------------------------
# 1. Define a simple CNN for MNIST
# ------------------------------
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(64*7*7, 128)
        self.fc2 = nn.Linear(128, 10)
        self.pool = nn.MaxPool2d(2, 2)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# ------------------------------
# Training and testing utilities
# ------------------------------
def train(model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.cross_entropy(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 100 == 0:
            print(f'Train Epoch {epoch}: Loss: {loss.item():.4f}')

def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.cross_entropy(output, target, reduction='sum').item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
    test_loss /= len(test_loader.dataset)
    accuracy = 100. * correct / len(test_loader.dataset)
    print(f'Test set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)\n')
    return accuracy

# ------------------------------
# Compression #31: Unstructured Pruning + RLE Masks
# ------------------------------
def prune_weights_rle(model, sparsity_ratio=0.5):
    """
    Prune weights below magnitude threshold to achieve `sparsity_ratio` zeros.
    Returns a dictionary containing for each parameter name:
        - 'values': list of non‑zero weights (float16)
        - 'rle_mask': run‑length encoded binary mask (list of (run_length, bit_value))
        - 'shape': original shape
    """
    compressed = {}
    state_dict = model.state_dict()
    for name, param in state_dict.items():
        if 'weight' not in name:   # keep biases uncompressed for simplicity
            compressed[name] = {'type': 'bias', 'data': param.cpu().numpy()}
            continue

        w = param.data.cpu().numpy().flatten()
        # Determine threshold to reach desired sparsity
        k = int(len(w) * sparsity_ratio)
        threshold = np.sort(np.abs(w))[k] if k < len(w) else 0.0
        mask = np.abs(w) > threshold
        pruned_weights = w[mask]
        # Run‑length encode the binary mask
        rle = []
        prev = mask[0]
        run = 1
        for b in mask[1:]:
            if b == prev:
                run += 1
            else:
                rle.append((run, int(prev)))
                prev = b
                run = 1
        rle.append((run, int(prev)))

        compressed[name] = {
            'type': 'pruned',
            'values': pruned_weights.astype(np.float16),   # store as fp16
            'rle_mask': rle,
            'shape': param.shape,
            'original_dtype': param.dtype
        }
    return compressed

def decompress_pruned_rle(compressed):
    """Rebuild state dict from pruned+RLE representation."""
    state_dict = OrderedDict()
    for name, info in compressed.items():
        if info['type'] == 'bias':
            state_dict[name] = torch.tensor(info['data'])
            continue
        # Reconstruct mask from RLE
        mask_parts = []
        for length, bit in info['rle_mask']:
            mask_parts.extend([bit] * length)
        mask = np.array(mask_parts, dtype=bool)
        # Reconstruct full weight array
        full = np.zeros(len(mask), dtype=np.float32)
        values = info['values'].astype(np.float32)
        full[mask] = values
        state_dict[name] = torch.tensor(full.reshape(info['shape']), dtype=info.get('original_dtype', torch.float32))
    return state_dict

# ------------------------------
# Compression #16: Ternary Quantization {-1,0,1}
# ------------------------------
def ternary_quantize_weights(model, threshold_scale=0.5):
    """
    Quantize weights to {-scale, 0, +scale} where scale = mean(|w|) for non‑zero weights.
    For each weight tensor we store:
        - ternary codes packed into bytes (2 bits per weight: 0→0, 1→+1, 2→-1)
        - scale factor (float32)
    """
    compressed = {}
    state_dict = model.state_dict()
    for name, param in state_dict.items():
        if 'weight' not in name:
            compressed[name] = {'type': 'bias', 'data': param.cpu().numpy()}
            continue

        w = param.data.cpu().numpy().flatten()
        # Compute scale as mean absolute value of all weights (including zeros? only non‑zeros)
        nonzero = w[np.abs(w) > 1e-6]
        if len(nonzero) == 0:
            scale = 1.0
        else:
            scale = np.mean(np.abs(nonzero))
        # Ternary thresholds: if |w| > threshold_scale * scale  -> sign, else 0
        thr = threshold_scale * scale
        ternary = np.zeros_like(w, dtype=np.int8)
        ternary[w > thr] = 1
        ternary[w < -thr] = -1

        # Pack ternary codes (values -1,0,1) into 2‑bit fields
        # mapping: 0→0, 1→1, -1→2 (2 bits)
        packed_bytes = bytearray()
        for i in range(0, len(ternary), 4):
            byte = 0
            for j in range(4):
                if i+j < len(ternary):
                    val = ternary[i+j]
                    code = 0 if val == 0 else (1 if val == 1 else 2)
                    byte |= (code << (2*j))
            packed_bytes.append(byte)

        compressed[name] = {
            'type': 'ternary',
            'packed': packed_bytes,
            'scale': scale,
            'shape': param.shape,
            'num_weights': len(ternary),
            'original_dtype': param.dtype
        }
    return compressed

def decompress_ternary(compressed):
    """Rebuild state dict from ternary representation."""
    state_dict = OrderedDict()
    for name, info in compressed.items():
        if info['type'] == 'bias':
            state_dict[name] = torch.tensor(info['data'])
            continue
        # Unpack 2‑bit codes
        ternary = np.zeros(info['num_weights'], dtype=np.float32)
        for idx, byte in enumerate(info['packed']):
            for j in range(4):
                pos = idx*4 + j
                if pos >= info['num_weights']:
                    break
                code = (byte >> (2*j)) & 0x03
                if code == 0:
                    val = 0.0
                elif code == 1:
                    val = 1.0
                else:   # code == 2
                    val = -1.0
                ternary[pos] = val
        # Multiply by scale
        weights = ternary * info['scale']
        state_dict[name] = torch.tensor(weights.reshape(info['shape']), dtype=info['original_dtype'])
    return state_dict

# ------------------------------
# Combine pruning + ternary quantization
# ------------------------------
def compress_model_pipeline(model, sparsity_ratio=0.5, ternary_threshold=0.5):
    """
    Step 1: prune weights (unstructured) and store pruned sparse representation.
    Step 2: apply ternary quantization to the pruned weights.
    Returns compressed representation and the reconstructed model (for validation).
    """
    # Step 1: prune and store with RLE mask
    pruned_compressed = prune_weights_rle(model, sparsity_ratio)
    # Reconstruct pruned model
    pruned_state = decompress_pruned_rle(pruned_compressed)
    pruned_model = SimpleCNN()
    pruned_model.load_state_dict(pruned_state, strict=False)
    # Step 2: ternary quantize the pruned model weights
    ternary_compressed = ternary_quantize_weights(pruned_model, ternary_threshold)
    return ternary_compressed

def decompress_pipeline(compressed):
    """Decompress the combined representation (ternary after pruning)."""
    return decompress_ternary(compressed)

# ------------------------------
# Memory measurement helper
# ------------------------------
def model_size_mb(state_dict):
    """Compute size of state_dict in MB (approximate)."""
    total_bytes = 0
    for k, v in state_dict.items():
        total_bytes += v.numel() * v.element_size()
    return total_bytes / (1024 * 1024)

def compressed_size_mb(compressed):
    """Estimate size of compressed representation."""
    total_bytes = 0
    for info in compressed.values():
        if info['type'] == 'bias':
            total_bytes += info['data'].nbytes
        elif info['type'] == 'pruned':
            total_bytes += info['values'].nbytes  # fp16
            # RLE mask: each tuple (run, bit) stored as two ints
            total_bytes += len(info['rle_mask']) * 2 * 4  # assume 4 bytes each
        elif info['type'] == 'ternary':
            total_bytes += len(info['packed'])          # packed bytes
            total_bytes += 4                            # scale (float32)
    return total_bytes / (1024 * 1024)

# ------------------------------
# Main: train, compress, test
# ------------------------------
def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

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

    # Train original model
    model = SimpleCNN().to(device)
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    for epoch in range(1, 4):  # 3 epochs are enough for ~98% accuracy
        train(model, device, train_loader, optimizer, epoch)
        test(model, device, test_loader)

    original_state = model.state_dict()
    original_size = model_size_mb(original_state)
    print(f"Original model size: {original_size:.2f} MB")

    # Apply compression pipeline (pruning + ternary)
    sparsity = 0.6          # keep 40% of weights
    ternary_thr = 0.5
    compressed_repr = compress_model_pipeline(model, sparsity, ternary_thr)

    # Decompress and evaluate accuracy
    decompressed_state = decompress_pipeline(compressed_repr)
    model.load_state_dict(decompressed_state, strict=False)
    model.to(device)
    acc_compressed = test(model, device, test_loader)

    # Measure final compressed size
    final_size = compressed_size_mb(compressed_repr)
    print(f"Compressed model size (pruned + ternary): {final_size:.2f} MB")
    print(f"Compression ratio: {original_size / final_size:.2f}x")
    print(f"Accuracy after compression: {acc_compressed:.2f}%")

if __name__ == "__main__":
    main()