"""
Core Unlearning Module for Verril-Learn

Implements Surgical Unlearning using:
1. Fisher Information Matrix (Diagonal Approximation)
   - Identifies critical weights for retain data
   
2. Gradient Ascent on Forget Data
   - Maximizes loss to "forget" the poisoned samples
   
3. Fisher-Weighted Regularization
   - Prevents catastrophic forgetting by penalizing
     changes to weights critical for retain data

4. Repair Phase (NEW)
   - After unlearning, briefly retrain on retain data
   - "Heals" the model by reinforcing good patterns

Mathematical Foundation:
    
    The unlearning objective is:
    
    L_total = -L_forget + λ * Σ F_ii * (θ_i - θ_i^original)²
    
    Where:
    - L_forget: Cross-entropy loss on forget data (we MAXIMIZE this)
    - F_ii: Diagonal Fisher Information for weight i
    - θ_i: Current weight value
    - θ_i^original: Original weight value before unlearning
    - λ: Fisher regularization weight (fisher_lambda)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from typing import Dict
import copy


def calculate_fisher(
    model: nn.Module,
    data_loader: DataLoader,
    device: str = 'cpu'
) -> Dict[str, torch.Tensor]:
    """
    Calculate the diagonal Fisher Information Matrix (FIM).
    
    The Fisher Information Matrix tells us how sensitive the model's
    predictions are to changes in each parameter. Parameters with high
    Fisher values are "critical" - changing them would significantly
    affect the model's performance on the data.
    
    Diagonal Approximation:
        F_ii = E[(∂L / ∂θ_i)²]
        
    This is computed by:
    1. Iterating through the data_loader
    2. Computing gradients for each batch
    3. Squaring the gradients
    4. Averaging across all samples
    
    Args:
        model: The neural network model
        data_loader: DataLoader for computing Fisher (typically retain data)
        device: Device to compute on ('cpu' or 'cuda')
    
    Returns:
        Dictionary mapping parameter names to Fisher diagonal values
        (importance weights for each parameter)
        
    Example:
        >>> fisher = calculate_fisher(model, retain_loader, device='cuda')
        >>> print(fisher['conv1.weight'].shape)  # Same as conv1.weight
    """
    model = model.to(device)
    model.eval()
    
    # Initialize Fisher dictionary with zeros
    fisher_dict = {}
    for name, param in model.named_parameters():
        fisher_dict[name] = torch.zeros_like(param.data)
    
    # Count total samples for averaging
    total_samples = 0
    
    print("[Verril-Learn] Calculating Fisher Information Matrix...")
    
    for batch_idx, (data, target) in enumerate(data_loader):
        data, target = data.to(device), target.to(device)
        batch_size = data.size(0)
        total_samples += batch_size
        
        # Zero gradients
        model.zero_grad()
        
        # Forward pass
        output = model(data)
        
        # Compute loss (cross-entropy)
        loss = F.cross_entropy(output, target)
        
        # Backward pass to get gradients
        loss.backward()
        
        # Accumulate squared gradients (Fisher diagonal)
        for name, param in model.named_parameters():
            if param.grad is not None:
                # Square the gradients and weight by batch size
                fisher_dict[name] += (param.grad.data ** 2) * batch_size
    
    # Average over all samples
    for name in fisher_dict:
        fisher_dict[name] /= total_samples
    
    print(f"[Verril-Learn] Fisher computed on {total_samples} samples")
    
    return fisher_dict


def repair_model(
    model: nn.Module,
    retain_loader: DataLoader,
    device: str = 'cpu',
    lr: float = 0.001,
    epochs: int = 1
) -> nn.Module:
    """
    Repair the model by training on retain data.
    
    This "heals" the model after aggressive unlearning by reminding
    it what the good data looks like. Uses standard gradient descent.
    
    Args:
        model: The model to repair
        retain_loader: DataLoader for clean retain data
        device: Device to train on
        lr: Learning rate for repair training
        epochs: Number of repair epochs
    
    Returns:
        The repaired model
    """
    model = model.to(device)
    model.train()
    
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()
    
    print(f"\n[Repair Phase] Training on retain data for {epochs} epoch(s)...")
    
    for epoch in range(epochs):
        running_loss = 0.0
        correct = 0
        total = 0
        
        for batch_idx, (data, target) in enumerate(retain_loader):
            data, target = data.to(device), target.to(device)
            
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item()
            _, predicted = output.max(1)
            total += target.size(0)
            correct += predicted.eq(target).sum().item()
        
        epoch_loss = running_loss / len(retain_loader)
        epoch_acc = 100. * correct / total
        print(f"  Repair Epoch {epoch+1}/{epochs} | "
              f"Loss: {epoch_loss:.4f} | Acc: {epoch_acc:.2f}%")
    
    return model


def surgical_unlearn(
    model: nn.Module,
    forget_loader: DataLoader,
    retain_loader: DataLoader,
    device: str = 'cpu',
    lr: float = 0.001,
    epochs: int = 5,
    fisher_lambda: float = 0.1,
    repair: bool = True,
    repair_lr: float = 0.0005,
    repair_epochs: int = 1
) -> nn.Module:
    """
    Perform Surgical Unlearning on the model.
    
    This is the main unlearning algorithm that:
    
    Step A: Calculate Fisher Information on retain_loader
            (identifies which weights are critical for good data)
            
    Step B: Run Gradient Ascent on forget_loader
            (maximize loss to make model "forget" bad data)
            
    Step C: Apply Fisher penalty to prevent catastrophic forgetting
            penalty = fisher_lambda * Σ F_ii * (θ_i - θ_original_i)²
    
    Step D: (NEW) Repair Phase - if repair=True, run 1 epoch of
            standard training on retain_loader to "heal" the model
    
    The combined objective we MINIMIZE is:
        L_total = -L_forget + fisher_lambda * Σ F_ii * (θ - θ_original)²
        
    Minimizing -L_forget = Maximizing L_forget (gradient ascent)
    Minimizing the penalty = staying close to original weights on critical params
    
    Args:
        model: The trained model to unlearn from
        forget_loader: DataLoader for data to forget (poisoned data)
        retain_loader: DataLoader for data to remember (clean data)
        device: Device to run on ('cpu' or 'cuda')
        lr: Learning rate for unlearning updates
        epochs: Number of unlearning epochs
        fisher_lambda: Weight for Fisher regularization penalty
        repair: If True, run repair phase after unlearning (NEW)
        repair_lr: Learning rate for repair phase
        repair_epochs: Number of repair epochs
    
    Returns:
        The unlearned (and optionally repaired) model
        
    Example:
        >>> unlearned_model = surgical_unlearn(
        ...     model=trained_model,
        ...     forget_loader=poison_loader,
        ...     retain_loader=clean_loader,
        ...     device='cuda',
        ...     lr=0.001,
        ...     epochs=2,
        ...     fisher_lambda=5000,
        ...     repair=True
        ... )
    """
    model = model.to(device)
    
    print("\n" + "=" * 60)
    print("  VERRIL-LEARN: Surgical Unlearning")
    print("=" * 60)
    print(f"  Learning Rate: {lr}")
    print(f"  Epochs: {epochs}")
    print(f"  Fisher Lambda: {fisher_lambda}")
    print(f"  Repair Phase: {'Enabled' if repair else 'Disabled'}")
    print("=" * 60)
    
    # =========================================
    # STEP A: Calculate Fisher Information
    # =========================================
    print("\n[Step A] Calculating Fisher Information on retain data...")
    fisher_dict = calculate_fisher(model, retain_loader, device)
    
    # =========================================
    # Store Original Weights
    # =========================================
    print("\n[Step B] Storing original model weights...")
    original_params = {}
    for name, param in model.named_parameters():
        original_params[name] = param.data.clone()
    print(f"  Stored {len(original_params)} parameter tensors")
    
    # =========================================
    # STEP C: Gradient Ascent with Fisher Penalty
    # =========================================
    print(f"\n[Step C] Running Gradient Ascent for {epochs} epochs...")
    
    model.train()
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
    
    for epoch in range(epochs):
        total_loss = 0.0
        total_penalty = 0.0
        num_batches = 0
        
        for batch_idx, (data, target) in enumerate(forget_loader):
            data, target = data.to(device), target.to(device)
            
            # Zero gradients
            optimizer.zero_grad()
            
            # Forward pass on forget data
            output = model(data)
            forget_loss = F.cross_entropy(output, target)
            
            # =========================================
            # Calculate Fisher Penalty
            # penalty = fisher_lambda * Σ F_ii * (θ - θ_original)²
            # =========================================
            fisher_penalty = torch.tensor(0.0, device=device)
            for name, param in model.named_parameters():
                if name in fisher_dict and name in original_params:
                    # Weight difference from original
                    weight_diff = param - original_params[name].to(device)
                    # Fisher-weighted squared difference
                    fisher_penalty += (fisher_dict[name].to(device) * (weight_diff ** 2)).sum()
            
            fisher_penalty = fisher_lambda * fisher_penalty
            
            # =========================================
            # Combined Loss
            # We want to MAXIMIZE forget_loss (gradient ascent)
            # We want to MINIMIZE fisher_penalty (stay close to original)
            # So total loss = -forget_loss + fisher_penalty
            # =========================================
            combined_loss = -forget_loss + fisher_penalty
            
            # Backward and optimize
            combined_loss.backward()
            optimizer.step()
            
            total_loss += forget_loss.item()
            total_penalty += fisher_penalty.item()
            num_batches += 1
        
        avg_loss = total_loss / num_batches
        avg_penalty = total_penalty / num_batches
        
        print(f"  Epoch {epoch + 1}/{epochs} | "
              f"Forget Loss: {avg_loss:.4f} | "
              f"Fisher Penalty: {avg_penalty:.4f}")
    
    print("\n[Unlearning Phase Complete]")
    
    # =========================================
    # STEP D: Repair Phase (NEW)
    # =========================================
    if repair:
        print("\n" + "-" * 60)
        print("  REPAIR PHASE: Healing the model")
        print("-" * 60)
        model = repair_model(
            model=model,
            retain_loader=retain_loader,
            device=device,
            lr=repair_lr,
            epochs=repair_epochs
        )
    
    print("\n" + "=" * 60)
    print("  Surgical Unlearning Complete!")
    print("=" * 60 + "\n")
    
    return model
