#!/usr/bin/env python3
"""
Hyper‑Gas CIFAR‑10 Classifier with Gas Relaxation and Compression

This script trains a Bayesian (Hyper‑Gas) CNN where each weight is a Gaussian distribution.
After training, we apply a gas‑relaxation step (diffusion + condensation) to rearrange the
weight distributions into a more compressible state. Then we quantize the weight means to 8‑bit
integers (int8) and optionally discard the variance parameters for high‑confidence weights.
The result is a compressed model that retains predictive accuracy while drastically reducing
memory footprint – exactly as a gas naturally arranges itself before being compressed.

Usage:
    python hypergas_cifar10.py
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from tqdm import tqdm
import numpy as np
import math

# -------------------------------
# Hyper‑Gas Layers: weights as distributions
# -------------------------------
class HyperGasLinear(nn.Module):
    def __init__(self, in_features, out_features, prior_sigma=0.1):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features

        self.w_mu = nn.Parameter(torch.Tensor(out_features, in_features).normal_(0, 0.02))
        self.w_log_sigma = nn.Parameter(torch.Tensor(out_features, in_features).normal_(-3, 0.1))
        self.b_mu = nn.Parameter(torch.Tensor(out_features).zero_())
        self.b_log_sigma = nn.Parameter(torch.Tensor(out_features).normal_(-3, 0.1))
        self.prior_sigma = prior_sigma

    def forward(self, x):
        w_eps = torch.randn_like(self.w_mu)
        b_eps = torch.randn_like(self.b_mu)
        w_sigma = torch.exp(self.w_log_sigma)
        b_sigma = torch.exp(self.b_log_sigma)
        w = self.w_mu + w_sigma * w_eps
        b = self.b_mu + b_sigma * b_eps

        kl_w = torch.sum(self.w_log_sigma - torch.log(torch.tensor(self.prior_sigma)) +
                         (w_sigma**2 + self.w_mu**2) / (2 * self.prior_sigma**2) - 0.5)
        kl_b = torch.sum(self.b_log_sigma - torch.log(torch.tensor(self.prior_sigma)) +
                         (b_sigma**2 + self.b_mu**2) / (2 * self.prior_sigma**2) - 0.5)
        self.kl = kl_w + kl_b
        return F.linear(x, w, b)


class HyperGasConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, prior_sigma=0.1):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size
        self.stride = stride
        self.padding = padding

        fan_in = in_channels * self.kernel_size[0] * self.kernel_size[1]
        self.w_mu = nn.Parameter(torch.Tensor(out_channels, in_channels, *self.kernel_size).normal_(0, 0.02))
        self.w_log_sigma = nn.Parameter(torch.Tensor(out_channels, in_channels, *self.kernel_size).normal_(-3, 0.1))
        self.b_mu = nn.Parameter(torch.Tensor(out_channels).zero_())
        self.b_log_sigma = nn.Parameter(torch.Tensor(out_channels).normal_(-3, 0.1))
        self.prior_sigma = prior_sigma

    def forward(self, x):
        w_eps = torch.randn_like(self.w_mu)
        b_eps = torch.randn_like(self.b_mu)
        w_sigma = torch.exp(self.w_log_sigma)
        b_sigma = torch.exp(self.b_log_sigma)
        w = self.w_mu + w_sigma * w_eps
        b = self.b_mu + b_sigma * b_eps

        kl_w = torch.sum(self.w_log_sigma - torch.log(torch.tensor(self.prior_sigma)) +
                         (w_sigma**2 + self.w_mu**2) / (2 * self.prior_sigma**2) - 0.5)
        kl_b = torch.sum(self.b_log_sigma - torch.log(torch.tensor(self.prior_sigma)) +
                         (b_sigma**2 + self.b_mu**2) / (2 * self.prior_sigma**2) - 0.5)
        self.kl = kl_w + kl_b
        return F.conv2d(x, w, b, stride=self.stride, padding=self.padding)


class HyperGasNet(nn.Module):
    def __init__(self, prior_sigma=0.1):
        super().__init__()
        self.conv1 = HyperGasConv2d(3, 32, 3, padding=1, prior_sigma=prior_sigma)
        self.conv2 = HyperGasConv2d(32, 64, 3, padding=1, prior_sigma=prior_sigma)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = HyperGasLinear(64 * 8 * 8, 256, prior_sigma=prior_sigma)
        self.fc2 = HyperGasLinear(256, 10, prior_sigma=prior_sigma)

    def forward(self, x, sample=True):
        if not sample:
            # deterministic mode using only means
            x = F.relu(F.conv2d(x, self.conv1.w_mu, self.conv1.b_mu, stride=1, padding=1))
            x = self.pool(x)
            x = F.relu(F.conv2d(x, self.conv2.w_mu, self.conv2.b_mu, stride=1, padding=1))
            x = self.pool(x)
            x = x.view(x.size(0), -1)
            x = F.relu(F.linear(x, self.fc1.w_mu, self.fc1.b_mu))
            x = F.linear(x, self.fc2.w_mu, self.fc2.b_mu)
            return x
        else:
            x = F.relu(self.conv1(x))
            x = self.pool(x)
            x = F.relu(self.conv2(x))
            x = self.pool(x)
            x = x.view(x.size(0), -1)
            x = F.relu(self.fc1(x))
            x = self.fc2(x)
            return x

    def kl_loss(self):
        return (self.conv1.kl + self.conv2.kl + self.fc1.kl + self.fc2.kl) / 60000.0   # normalize by dataset size


# -------------------------------
# Gas Relaxation (like gas molecules rearranging)
# -------------------------------
def relax_gas(model, steps=10, temperature=0.05, lr=0.01):
    """
    Rearrange the weight distributions like a gas:
    - Means diffuse toward local averages (pressure equalisation).
    - Variances shrink near clusters (condensation) and expand elsewhere.
    This prepares the numbers for compression.
    """
    with torch.no_grad():
        for step in range(steps):
            for module in model.modules():
                if hasattr(module, 'w_mu') and hasattr(module, 'w_log_sigma'):
                    mu = module.w_mu
                    log_sigma = module.w_log_sigma
                    sigma = torch.exp(log_sigma)

                    # Local average: use the mean of all weight means in this layer
                    local_avg = mu.mean()

                    # Gas diffusion: move means toward local average + thermal noise
                    noise = torch.randn_like(mu) * temperature
                    mu_new = mu + lr * (local_avg - mu) + noise
                    module.w_mu.data = mu_new

                    # Condensation/expansion: sigma decreases where mu is near local average
                    distance = torch.abs(mu_new - local_avg)
                    sigma_new = sigma * (0.9 + 0.1 * torch.tanh(distance))  # smooth scaling
                    module.w_log_sigma.data = torch.log(sigma_new + 1e-8)


# -------------------------------
# Compression: quantize means to 8‑bit integers
# -------------------------------
def compress_model(model, prune_sigma_threshold=0.01):
    """
    Compress the Hyper‑Gas model by quantizing weight means to int8.
    Optionally prune weights with very low sigma (deterministic).
    Returns a dictionary of quantized parameters and metadata.
    Also returns a "decompressed" model (using quantized means) for evaluation.
    """
    compressed = {}
    prior_sigma = model.conv1.prior_sigma if hasattr(model, 'conv1') else 0.1
    decompressed_model = HyperGasNet(prior_sigma=prior_sigma).to(next(model.parameters()).device)
    decompressed_model.load_state_dict(model.state_dict())
    total_original_params = 0
    total_compressed_bits = 0

    for name, module in decompressed_model.named_modules():
        if hasattr(module, 'w_mu'):
            # Original parameters count (floats, 32 bits each)
            original_params = module.w_mu.numel() + (module.b_mu.numel() if hasattr(module, 'b_mu') else 0)
            total_original_params += original_params

            # Quantize weight means
            mu = module.w_mu.detach().cpu().numpy()
            mu_min, mu_max = mu.min(), mu.max()
            # Avoid single value range
            if mu_max - mu_min < 1e-6:
                mu_min, mu_max = mu_min - 1.0, mu_max + 1.0
            mu_q = np.round((mu - mu_min) / (mu_max - mu_min) * 255).astype(np.uint8)
            # Store quantization metadata
            compressed[name + '.w_mu'] = {
                'data': mu_q,
                'min': mu_min,
                'max': mu_max,
                'shape': mu.shape,
                'dtype': 'uint8'
            }
            total_compressed_bits += mu_q.size * 8   # 8 bits per quantized value

            # Optionally prune weights with very small sigma (set to zero)
            if hasattr(module, 'w_log_sigma'):
                sigma = torch.exp(module.w_log_sigma).detach().cpu().numpy()
                prune_mask = sigma < prune_sigma_threshold
                # For decompressed model: replace pruned means with zero (or keep quantized but set to zero later)
                # We'll just keep quantized for now, but could set to zero.

            # Bias quantization (if exists)
            if hasattr(module, 'b_mu'):
                b = module.b_mu.detach().cpu().numpy()
                b_min, b_max = b.min(), b.max()
                if b_max - b_min < 1e-6:
                    b_min, b_max = b_min - 1.0, b_max + 1.0
                b_q = np.round((b - b_min) / (b_max - b_min) * 255).astype(np.uint8)
                compressed[name + '.b_mu'] = {
                    'data': b_q,
                    'min': b_min,
                    'max': b_max,
                    'shape': b.shape,
                    'dtype': 'uint8'
                }
                total_compressed_bits += b_q.size * 8

    # Update decompressed model with quantized (de‑quantized) weights for evaluation
    with torch.no_grad():
        for name, module in decompressed_model.named_modules():
            if name + '.w_mu' in compressed:
                q_info = compressed[name + '.w_mu']
                mu_deq = (q_info['data'].astype(np.float32) / 255.0) * (q_info['max'] - q_info['min']) + q_info['min']
                module.w_mu.data = torch.tensor(mu_deq.reshape(q_info['shape']), device=module.w_mu.device, dtype=module.w_mu.dtype)
            if name + '.b_mu' in compressed:
                q_info = compressed[name + '.b_mu']
                b_deq = (q_info['data'].astype(np.float32) / 255.0) * (q_info['max'] - q_info['min']) + q_info['min']
                module.b_mu.data = torch.tensor(b_deq.reshape(q_info['shape']), device=module.b_mu.device, dtype=module.b_mu.dtype)

    compression_ratio = total_original_params * 32 / total_compressed_bits if total_compressed_bits > 0 else 0
    print(f"Compression: original {total_original_params * 32} bits → compressed {total_compressed_bits} bits (ratio {compression_ratio:.2f}x)")
    return compressed, decompressed_model


# -------------------------------
# Training and evaluation utilities
# -------------------------------
def train_epoch(model, device, train_loader, optimizer, epoch, beta=0.1):
    model.train()
    total_loss = 0.0
    total_correct = 0
    for data, target in tqdm(train_loader, desc=f"Epoch {epoch}"):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        nll = F.cross_entropy(output, target, reduction='mean')
        kl = model.kl_loss()
        loss = nll + beta * kl
        loss.backward()
        optimizer.step()
        pred = output.argmax(dim=1)
        total_correct += pred.eq(target).sum().item()
        total_loss += loss.item() * len(data)
    avg_loss = total_loss / len(train_loader.dataset)
    acc = 100. * total_correct / len(train_loader.dataset)
    print(f"Train Epoch {epoch}: Loss={avg_loss:.4f}, Accuracy={acc:.2f}%")
    return avg_loss, acc


def evaluate(model, device, test_loader, mc_samples=10):
    """Evaluate accuracy using Monte Carlo sampling (gas phase)."""
    model.eval()
    correct = 0
    with torch.no_grad():
        for data, target in tqdm(test_loader, desc="Evaluating"):
            data, target = data.to(device), target.to(device)
            # Monte Carlo average
            outputs = torch.zeros(data.size(0), 10).to(device)
            for _ in range(mc_samples):
                logits = model(data)
                outputs += F.softmax(logits, dim=1)
            outputs /= mc_samples
            pred = outputs.argmax(dim=1)
            correct += pred.eq(target).sum().item()
    acc = 100. * correct / len(test_loader.dataset)
    return acc


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

    # Data preparation
    transform_train = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])
    transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])
    train_set = torchvision.datasets.CIFAR10(root='../data', train=True, download=True, transform=transform_train)
    test_set = torchvision.datasets.CIFAR10(root='../data', train=False, download=True, transform=transform_test)
    train_loader = torch.utils.data.DataLoader(train_set, batch_size=128, shuffle=True, num_workers=2)
    test_loader = torch.utils.data.DataLoader(test_set, batch_size=128, shuffle=False, num_workers=2)

    # Create Hyper‑Gas model
    model = HyperGasNet(prior_sigma=0.1).to(device)
    optimizer = optim.Adam(model.parameters(), lr=1e-3)

    # Training loop (fewer epochs for demo)
    print("\n===== TRAINING PHASE (20 epochs) =====")
    for epoch in range(1, 21):
        train_epoch(model, device, train_loader, optimizer, epoch, beta=0.1)
        if epoch % 10 == 0:
            acc = evaluate(model, device, test_loader, mc_samples=5)
            print(f"  → Test accuracy after epoch {epoch}: {acc:.2f}%")

    print("\n===== FINAL EVALUATION BEFORE RELAXATION =====")
    acc_before = evaluate(model, device, test_loader, mc_samples=20)
    print(f"Test accuracy (before relaxation): {acc_before:.2f}%")

    # Gas relaxation (batch iteration 10x as requested)
    print("\n===== GAS RELAXATION (10 steps) =====")
    relax_gas(model, steps=10, temperature=0.05, lr=0.01)

    print("\n===== EVALUATION AFTER GAS RELAXATION =====")
    acc_after_relax = evaluate(model, device, test_loader, mc_samples=20)
    print(f"Test accuracy (after gas relaxation): {acc_after_relax:.2f}%")

    # Compression (quantization)
    print("\n===== COMPRESSION =====")
    compressed_state, compressed_model = compress_model(model, prune_sigma_threshold=0.01)

    print("\n===== EVALUATION OF COMPRESSED MODEL =====")
    acc_compressed = evaluate(compressed_model, device, test_loader, mc_samples=1)  # deterministic after compression
    print(f"Test accuracy (compressed model, 1 sample): {acc_compressed:.2f}%")

    print("\n===== SUMMARY =====")
    print(f"Original model test acc: {acc_before:.2f}%")
    print(f"After gas relaxation: {acc_after_relax:.2f}%")
    print(f"After quantization (compressed): {acc_compressed:.2f}%")
    print("Compression done. Model weights quantized to 8‑bit ints.")


if __name__ == "__main__":
    main()