# model_analyzer.py
import torch
import torch.nn as nn
from torch.nn.utils import prune
from rich.console import Console
from rich.table import Table

console = Console()


def analyze_model(model: nn.Module, name: str = "Model"):
    """Analyze model size, FLOPs, and parameter distribution."""
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total_bytes = sum(
        p.numel() * p.element_size() for p in model.parameters()
    )
    buffer_bytes = sum(
        b.numel() * b.element_size() for b in model.buffers()
    )

    table = Table(title=f"{name} Analysis")
    table.add_column("Metric", style="cyan")
    table.add_column("Value", style="green")

    table.add_row("Total Parameters", f"{total_params:,}")
    table.add_row("Trainable Parameters", f"{trainable_params:,}")
    table.add_row("Model Size (MB)", f"{(total_bytes + buffer_bytes) / 1e6:.2f}")
    table.add_row("Weight Size (MB)", f"{total_bytes / 1e6:.2f}")
    table.add_row("Buffer Size (MB)", f"{buffer_bytes / 1e6:.2f}")
    table.add_row("FP32 Size (MB)", f"{total_params * 4 / 1e6:.2f}")
    table.add_row("INT8 Size (MB)", f"{total_params * 1 / 1e6:.2f}")
    table.add_row("INT4 Size (MB)", f"{total_params * 0.5 / 1e6:.2f}")

    console.print(table)

    by_type = {}
    for name_param, param in model.named_parameters():
        layer_type = type(param).__name__
        by_type.setdefault(layer_type, 0)
        by_type[layer_type] += param.numel()

    param_table = Table(title="Parameters by Layer Type")
    param_table.add_column("Layer Type", style="cyan")
    param_table.add_column("Parameters", style="green")
    param_table.add_column("Percentage", style="yellow")

    for layer_type, count in sorted(by_type.items(), key=lambda x: -x[1]):
        param_table.add_row(
            layer_type,
            f"{count:,}",
            f"{count / total_params * 100:.1f}%",
        )

    console.print(param_table)
    return total_params, total_bytes

# pruning_engine.py
import torch
import torch.nn as nn
from torch.nn.utils import prune
import copy
import math
from typing import Optional, Tuple


class PruningEngine:
    """Structured and unstructured pruning engine."""

    def __init__(self, model: nn.Module):
        self.original_model = copy.deepcopy(model)
        self.pruned_model = model
        self.pruning_history = []

    def magnitude_pruning(
        self,
        amount: float = 0.5,
        global_pruning: bool = True,
        exclude_bias: bool = True,
    ) -> nn.Module:
        """Apply magnitude-based unstructured pruning globally."""
        parameters_to_prune = []
        for name, module in self.pruned_model.named_modules():
            if isinstance(module, (nn.Linear, nn.Conv2d, nn.Conv1d)):
                if exclude_bias:
                    parameters_to_prune.append((module, "weight"))
                else:
                    parameters_to_prune.append((module, "weight"))
                    if module.bias is not None:
                        parameters_to_prune.append((module, "bias"))

        if global_pruning:
            prune.global_unstructured(
                parameters_to_prune,
                pruning_method=prune.L1Unstructured,
                amount=amount,
            )
        else:
            for module, param_name in parameters_to_prune:
                prune.l1_unstructured(module, param_name, amount=amount)

        sparsity = self._calculate_sparsity()
        self.pruning_history.append({
            "method": "magnitude_pruning",
            "amount": amount,
            "sparsity": sparsity,
        })
        return self.pruned_model

    def structured_pruning(
        self, amount: float = 0.3, dim: int = 0
    ) -> nn.Module:
        """Apply structured pruning (remove entire channels/neurons)."""
        for name, module in self.pruned_model.named_modules():
            if isinstance(module, nn.Linear):
                prune.ln_structured(
                    module, "weight", amount=amount, n=2, dim=dim
                )
            elif isinstance(module, nn.Conv2d):
                prune.ln_structured(
                    module, "weight", amount=amount, n=2, dim=0
                )

        self.pruning_history.append({
            "method": "structured_pruning",
            "amount": amount,
        })
        return self.pruned_model

    def gradual_pruning(
        self,
        target_sparsity: float = 0.9,
        initial_sparsity: float = 0.0,
        steps: int = 10,
        schedule: str = "cubic",
    ) -> nn.Module:
        """Gradually increase pruning over training steps."""
        for step in range(steps):
            t = step / (steps - 1)
            if schedule == "cubic":
                current_sparsity = target_sparsity - (
                    target_sparsity - initial_sparsity
                ) * (1 - t) ** 3
            elif schedule == "exponential":
                current_sparsity = target_sparsity * (1 - (1 - t) ** 2)
            else:
                current_sparsity = initial_sparsity + (
                    target_sparsity - initial_sparsity
                ) * t

            amount = current_sparsity
            self.magnitude_pruning(amount=amount)

        return self.pruned_model

    def unstructured_to_structured(self) -> nn.Module:
        """Convert unstructured sparsity to structured for hardware speedup."""
        for name, module in self.pruned_model.named_modules():
            if isinstance(module, (nn.Linear, nn.Conv2d)):
                weight = module.weight.data
                if weight.dim() > 1:
                    importance = weight.abs().sum(dim=tuple(range(1, weight.dim())))
                    threshold = torch.quantile(importance, 0.3)
                    mask = (importance > threshold).float()
                    mask = mask.view(
                        *mask.shape, *([1] * (weight.dim() - 1))
                    )
                    prune.custom_from_mask(module, "weight", mask)
        return self.pruned_model

    def remove_pruning_reparametrization(self):
        """Remove pruning reparametrization for clean export."""
        for name, module in self.pruned_model.named_modules():
            if isinstance(module, (nn.Linear, nn.Conv2d)):
                prune.remove(module, "weight")

    def _calculate_sparsity(self) -> float:
        """Calculate global sparsity ratio."""
        total_params = 0
        zero_params = 0
        for param in self.pruned_model.parameters():
            total_params += param.numel()
            zero_params += (param == 0).sum().item()
        return zero_params / total_params if total_params > 0 else 0

    def get_compression_stats(self) -> dict:
        """Get detailed compression statistics."""
        original_size = sum(
            p.numel() * 4 for p in self.original_model.parameters()
        )
        pruned_size = 0
        total_params = 0
        zero_params = 0

        for param in self.pruned_model.parameters():
            total_params += param.numel()
            zero_params += (param == 0).sum().item()
            if param.is_sparse:
                pruned_size += param._nnz() * param.element_size()
            else:
                non_zero = (param != 0).sum().item()
                pruned_size += non_zero * param.element_size()

        return {
            "original_parameters": sum(
                p.numel() for p in self.original_model.parameters()
            ),
            "total_parameters": total_params,
            "zero_parameters": zero_params,
            "sparsity": zero_params / total_params,
            "original_size_mb": original_size / 1e6,
            "compressed_size_mb": pruned_size / 1e6,
            "compression_ratio": original_size / pruned_size if pruned_size > 0 else 1,
        }

# quantization_engine.py
import torch
import torch.nn as nn
from torch.quantization import (
    quantize_dynamic,
    quantize_fx,
    get_default_qconfig,
    QConfigMapping,
)
from torch.ao.quantization.quantize_fx import prepare_fx, convert_fx
import copy
from typing import Optional


class QuantizationEngine:
    """Post-training and quantization-aware quantization."""

    def __init__(self, model: nn.Module):
        self.original_model = copy.deepcopy(model)
        self.quantized_model = None

    def dynamic_quantization(
        self,
        dtype: torch.dtype = torch.qint8,
        sparse: bool = False,
    ) -> nn.Module:
        """Apply dynamic quantization (weights quantized, activations at runtime)."""
        self.quantized_model = quantize_dynamic(
            self.original_model,
            {nn.Linear},
            dtype=dtype,
            sparse=sparse,
        )
        return self.quantized_model

    def static_quantization(
        self,
        calibration_data: Optional[torch.Tensor] = None,
        num_calibration_batches: int = 100,
    ) -> nn.Module:
        """Apply static quantization with calibration."""
        model_fp32 = copy.deepcopy(self.original_model)
        model_fp32.eval()

        model_fp32.qconfig = get_default_qconfig("qnnpack")

        prepared = prepare_fx(
            model_fp32,
            QConfigMapping().set_global(model_fp32.qconfig),
        )

        if calibration_data is not None:
            with torch.no_grad():
                for i in range(min(num_calibration_batches, len(calibration_data))):
                    prepared(calibration_data[i])

        self.quantized_model = convert_fx(prepared)
        return self.quantized_model

    def quantization_aware_training(
        self,
        train_loader,
        val_loader,
        epochs: int = 5,
        lr: float = 1e-4,
    ) -> nn.Module:
        """Apply QAT for better accuracy at low precision."""
        model_fp32 = copy.deepcopy(self.original_model)
        model_fp32.train()

        model_fp32.qconfig = get_default_qconfig("qnnpack")

        prepared = prepare_fx(
            model_fp32,
            QConfigMapping().set_global(model_fp32.qconfig),
        )

        optimizer = torch.optim.Adam(prepared.parameters(), lr=lr)
        criterion = nn.CrossEntropyLoss()

        for epoch in range(epochs):
            prepared.train()
            for batch_x, batch_y in train_loader:
                optimizer.zero_grad()
                output = prepared(batch_x)
                loss = criterion(output, batch_y)
                loss.backward()
                optimizer.step()

            prepared.eval()
            correct = 0
            total = 0
            with torch.no_grad():
                for batch_x, batch_y in val_loader:
                    output = prepared(batch_x)
                    _, predicted = output.max(1)
                    total += batch_y.size(0)
                    correct += predicted.eq(batch_y).sum().item()
            print(f"QAT Epoch {epoch+1}/{epochs}: Acc={100.*correct/total:.2f}%")

        self.quantized_model = convert_fx(prepared)
        return self.quantized_model

    def get_quantization_stats(self) -> dict:
        """Compare original vs quantized model."""
        orig_size = sum(
            p.numel() * p.element_size() for p in self.original_model.parameters()
        )

        quant_params = 0
        quant_size = 0
        if self.quantized_model:
            for name, param in self.quantized_model.named_parameters():
                quant_params += param.numel()
                quant_size += param.numel() * param.element_size()

        return {
            "original_size_mb": orig_size / 1e6,
            "quantized_size_mb": quant_size / 1e6,
            "compression_ratio": orig_size / quant_size if quant_size > 0 else 1,
            "bit_width": 8,
            "memory_savings_percent": (1 - quant_size / orig_size) * 100 if orig_size > 0 else 0,
        }

# distillation_engine.py
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Tuple


class DistillationLoss(nn.Module):
    """Combined loss for knowledge distillation."""

    def __init__(self, temperature: float = 4.0, alpha: float = 0.3):
        super().__init__()
        self.temperature = temperature
        self.alpha = alpha
        self.ce_loss = nn.CrossEntropyLoss()
        self.kl_loss = nn.KLDivLoss(reduction="batchmean")

    def forward(
        self,
        student_logits: torch.Tensor,
        teacher_logits: torch.Tensor,
        labels: torch.Tensor,
    ) -> Tuple[torch.Tensor, dict]:
        """Compute distillation loss."""
        T = self.temperature

        hard_loss = self.ce_loss(student_logits, labels)

        soft_student = F.log_softmax(student_logits / T, dim=1)
        soft_teacher = F.softmax(teacher_logits / T, dim=1)
        soft_loss = self.kl_loss(soft_student, soft_teacher) * (T ** 2)

        total_loss = self.alpha * soft_loss + (1 - self.alpha) * hard_loss

        loss_dict = {
            "total": total_loss.item(),
            "hard": hard_loss.item(),
            "soft": soft_loss.item(),
        }

        return total_loss, loss_dict


class DistillationEngine:
    """Full knowledge distillation pipeline."""

    def __init__(
        self,
        teacher: nn.Module,
        student: nn.Module,
        temperature: float = 4.0,
        alpha: float = 0.3,
    ):
        self.teacher = teacher
        self.student = student
        self.teacher.eval()
        self.criterion = DistillationLoss(temperature, alpha)

    def train(
        self,
        train_loader,
        val_loader,
        epochs: int = 10,
        lr: float = 3e-4,
        device: str = "cuda",
    ) -> nn.Module:
        """Train student with teacher distillation."""
        self.teacher.to(device)
        self.student.to(device)
        optimizer = torch.optim.AdamW(self.student.parameters(), lr=lr)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, epochs)

        best_val_acc = 0.0

        for epoch in range(epochs):
            self.student.train()
            total_loss = 0
            num_batches = 0

            for batch_x, batch_y in train_loader:
                batch_x, batch_y = batch_x.to(device), batch_y.to(device)

                with torch.no_grad():
                    teacher_logits = self.teacher(batch_x)

                student_logits = self.student(batch_x)
                loss, loss_dict = self.criterion(
                    student_logits, teacher_logits, batch_y
                )

                optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self.student.parameters(), 1.0)
                optimizer.step()

                total_loss += loss_dict["total"]
                num_batches += 1

            scheduler.step()

            val_acc = self._evaluate(val_loader, device)
            avg_loss = total_loss / num_batches

            print(
                f"Epoch {epoch+1}/{epochs} - "
                f"Loss: {avg_loss:.4f} - Val Acc: {val_acc:.2f}%"
            )

            if val_acc > best_val_acc:
                best_val_acc = val_acc
                torch.save(self.student.state_dict(), "best_student.pt")

        return self.student

    def _evaluate(self, val_loader, device: str) -> float:
        """Evaluate student model accuracy."""
        self.student.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for batch_x, batch_y in val_loader:
                batch_x, batch_y = batch_x.to(device), batch_y.to(device)
                output = self.student(batch_x)
                _, predicted = output.max(1)
                total += batch_y.size(0)
                correct += predicted.eq(batch_y).sum().item()
        return 100.0 * correct / total

    def create_student_architecture(
        self, teacher_params: int, reduction_factor: float = 0.1
    ) -> nn.Module:
        """Create a smaller student model from teacher."""
        student_params = int(teacher_params * reduction_factor)
        hidden_size = int((student_params / 2) ** 0.5)
        num_layers = max(2, int((student_params / (hidden_size * 2)) ** 0.5))

        layers = []
        layers.append(nn.Linear(784, hidden_size))
        layers.append(nn.ReLU())
        layers.append(nn.Dropout(0.2))

        for _ in range(num_layers - 1):
            layers.append(nn.Linear(hidden_size, hidden_size))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(0.2))

        layers.append(nn.Linear(hidden_size, 10))

        return nn.Sequential(*layers)

# compression_pipeline.py
import torch
import torch.nn as nn
from pruning_engine import PruningEngine
from quantization_engine import QuantizationEngine
from distillation_engine import DistillationEngine
from model_analyzer import analyze_model
from rich.console import Console

console = Console()


def run_compression_pipeline(model: nn.Module, train_loader, val_loader):
    """Execute the full compression pipeline."""
    console.print("[bold blue]Model Compression Pipeline[/bold blue]")
    console.print("=" * 60)

    console.print("\n[bold]Step 1: Baseline Analysis[/bold]")
    analyze_model(model, "Original Model")

    console.print("\n[bold]Step 2: Pruning[/bold]")
    pruning_engine = PruningEngine(model)
    pruned_model = pruning_engine.magnitude_pruning(amount=0.5)
    stats = pruning_engine.get_compression_stats()
    console.print(f"Sparsity: {stats['sparsity']:.2%}")
    console.print(f"Compression ratio: {stats['compression_ratio']:.2f}×")

    console.print("\n[bold]Step 3: Quantization[/bold]")
    quant_engine = QuantizationEngine(pruned_model)
    quantized_model = quant_engine.dynamic_quantization()
    q_stats = quant_engine.get_quantization_stats()
    console.print(f"Memory savings: {q_stats['memory_savings_percent']:.1f}%")

    console.print("\n[bold]Step 4: Knowledge Distillation[/bold]")
    student = nn.Sequential(
        nn.Linear(784, 128), nn.ReLU(), nn.Dropout(0.2),
        nn.Linear(128, 64), nn.ReLU(),
        nn.Linear(64, 10),
    )

    distill_engine = DistillationEngine(
        teacher=model, student=student, temperature=4.0, alpha=0.3
    )
    distilled_student = distill_engine.train(
        train_loader, val_loader, epochs=10
    )

    console.print("\n[bold]Step 5: Final Comparison[/bold]")
    analyze_model(distilled_student, "Distilled Student")

    return {
        "pruned": pruned_model,
        "quantized": quantized_model,
        "distilled": distilled_student,
    }


if __name__ == "__main__":
    model = nn.Sequential(
        nn.Linear(784, 512), nn.ReLU(), nn.Dropout(0.2),
        nn.Linear(512, 256), nn.ReLU(), nn.Dropout(0.2),
        nn.Linear(256, 128), nn.ReLU(),
        nn.Linear(128, 10),
    )

    from torch.utils.data import DataLoader, TensorDataset
    X = torch.randn(5000, 784)
    y = torch.randint(0, 10, (5000,))
    loader = DataLoader(TensorDataset(X, y), batch_size=64, shuffle=True)

    results = run_compression_pipeline(model, loader, loader)