#!/usr/bin/env python3

import os
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from safetensors.torch import save_file
from tqdm import tqdm

# =============================================================================
# CONFIGURATION
# =============================================================================

RANDOM_SEED = 1337
BATCH_SIZE = 256
DP_EPOCHS = 10
DP_LR = 0.1
MAX_GRAD_NORM = 1.0
DELTA = 1e-5

# The key parameter: use ε=6 to balance accuracy and privacy
TARGET_EPSILON = 6.0

# Set reproducibility
torch.manual_seed(RANDOM_SEED)
np.random.seed(RANDOM_SEED)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(RANDOM_SEED)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Output directories
MODELS_DIR = "models"
os.makedirs(MODELS_DIR, exist_ok=True)

# SVHN normalization values
SVHN_MEAN = (0.4377, 0.4438, 0.4728)
SVHN_STD = (0.1980, 0.2010, 0.1970)

# =============================================================================
# MODEL ARCHITECTURE
# =============================================================================

class SVHNCNN(nn.Module):
    """CNN for SVHN classification (compatible with Opacus)."""

    def __init__(self):
        super(SVHNCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, padding=1)

        self.pool = nn.MaxPool2d(2, 2)

        self.fc1 = nn.Linear(64 * 4 * 4, 64)
        self.fc2 = nn.Linear(64, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = x.view(-1, 64 * 4 * 4)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# =============================================================================
# DATA LOADING
# =============================================================================

def get_svhn_loaders(batch_size=256, download=True):
    """Load SVHN dataset."""
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(SVHN_MEAN, SVHN_STD)
    ])

    train_dataset = datasets.SVHN(
        "data", split='train', download=download, transform=transform
    )
    test_dataset = datasets.SVHN(
        "data", split='test', download=download, transform=transform
    )

    train_loader = DataLoader(
        train_dataset, batch_size=batch_size, shuffle=True, num_workers=0
    )
    test_loader = DataLoader(
        test_dataset, batch_size=batch_size, shuffle=False, num_workers=0
    )

    return train_dataset, test_dataset, train_loader, test_loader


def evaluate_accuracy(model, data_loader, device):
    """Evaluate model accuracy."""
    model.eval()
    correct = 0
    total = 0

    with torch.no_grad():
        for images, labels in data_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()

    return 100.0 * correct / total


def compute_mia_advantage(model, train_loader, test_loader, device, num_samples=2000):
    """Compute MIA advantage using confidence threshold attack."""
    model.eval()

    def get_confidences(loader, max_samples):
        all_conf = []
        count = 0
        with torch.no_grad():
            for images, _ in loader:
                if count >= max_samples:
                    break
                images = images.to(device)
                outputs = model(images)
                probs = F.softmax(outputs, dim=1)
                conf = probs.max(dim=1).values.cpu().numpy()
                all_conf.extend(conf)
                count += len(conf)
        return np.array(all_conf[:max_samples])

    member_conf = get_confidences(train_loader, num_samples)
    nonmember_conf = get_confidences(test_loader, num_samples)

    n = min(len(member_conf), len(nonmember_conf))
    all_conf = np.concatenate([member_conf[:n], nonmember_conf[:n]])
    all_labels = np.concatenate([np.ones(n), np.zeros(n)])

    thresholds = np.percentile(all_conf, np.linspace(0, 100, 500))
    best_acc = 0.5

    for t in thresholds:
        acc = max(
            np.mean((all_conf >= t) == all_labels),
            np.mean((all_conf < t) == all_labels)
        )
        best_acc = max(best_acc, acc)

    return best_acc, best_acc - 0.5

# =============================================================================
# TRAINING
# =============================================================================

print("=" * 80)
print("  DP-SGD PRIVACY (SVHN)")
print("=" * 80)
print(f"\nDevice: {device}")
print(f"Target epsilon: {TARGET_EPSILON}")
print(f"Required: accuracy >= 55%, MIA advantage <= 5%")

print("\nLoading SVHN dataset...")
train_dataset, test_dataset, train_loader, test_loader = get_svhn_loaders(
    batch_size=BATCH_SIZE, download=True
)

print(f"Training samples: {len(train_dataset):,}")
print(f"Test samples: {len(test_dataset):,}")

# Import Opacus
from opacus import PrivacyEngine
from opacus.validators import ModuleValidator

print("\n" + "=" * 80)
print(f"  TRAINING: DP-SGD MODEL (Target ε={TARGET_EPSILON})")
print("=" * 80)

# Fresh data loaders for DP training
_, _, train_loader_dp, test_loader_dp = get_svhn_loaders(
    batch_size=BATCH_SIZE, download=False
)

# Initialize model
dp_model = SVHNCNN().to(device)
dp_model = ModuleValidator.fix(dp_model)
optimizer_dp = optim.Adam(dp_model.parameters(), lr=1e-3)
#optimizer_dp = optim.SGD(dp_model.parameters(), lr=DP_LR, momentum=0.9)

# Attach privacy engine
privacy_engine = PrivacyEngine(accountant="rdp")
dp_model, optimizer_dp, train_loader_dp = privacy_engine.make_private_with_epsilon(
    module=dp_model,
    optimizer=optimizer_dp,
    data_loader=train_loader_dp,
    target_epsilon=TARGET_EPSILON,
    target_delta=DELTA,
    epochs=DP_EPOCHS,
    max_grad_norm=MAX_GRAD_NORM,
)

print(f"\nConfiguration:")
print(f"  Target epsilon: {TARGET_EPSILON}")
print(f"  Delta: {DELTA}")
print(f"  Max gradient norm: {MAX_GRAD_NORM}")
print(f"  Epochs: {DP_EPOCHS}")

# Training loop
criterion = nn.CrossEntropyLoss()
dp_model.train()

for epoch in range(DP_EPOCHS):
    running_loss = 0.0
    correct = 0
    total = 0

    pbar = tqdm(train_loader_dp, desc=f"Epoch {epoch+1}/{DP_EPOCHS}")
    for images, labels in pbar:
        images, labels = images.to(device), labels.to(device)

        optimizer_dp.zero_grad()
        outputs = dp_model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer_dp.step()

        running_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

        pbar.set_postfix({
            'loss': f'{running_loss/total:.4f}',
            'acc': f'{100.*correct/total:.2f}%'
        })

    # Get current epsilon
    epsilon = privacy_engine.get_epsilon(DELTA)
    print(f"  Epoch {epoch+1}: ε = {epsilon:.2f}")

final_epsilon = privacy_engine.get_epsilon(DELTA)
print(f"\nFinal privacy guarantee: (ε={final_epsilon:.2f}, δ={DELTA})")

# Evaluate
train_acc = evaluate_accuracy(dp_model, train_loader, device)
#train_acc = evaluate_accuracy(dp_model, train_loader_dp, device)
test_acc = evaluate_accuracy(dp_model, test_loader_dp, device)

print(f"\nDP Model Performance:")
print(f"  Training accuracy: {train_acc:.2f}%")
print(f"  Test accuracy: {test_acc:.2f}%")
print(f"  Overfitting gap: {train_acc - test_acc:.2f}%")

# MIA evaluation
print("\n" + "=" * 80)
print("  MEMBERSHIP INFERENCE EVALUATION")
print("=" * 80)

mia_acc, mia_adv = compute_mia_advantage(
    dp_model, train_loader_dp, test_loader_dp, device
)

print(f"\nMIA Results:")
print(f"  Attack accuracy: {mia_acc:.4f}")
print(f"  Attack advantage: {mia_adv:.4f} ({mia_adv*100:.2f}%)")

# Save model
model_path = os.path.join(MODELS_DIR, "dp_model.safetensors")
save_file(dp_model._module.state_dict(), model_path)
print(f"\nModel saved to {model_path}")

# =============================================================================
# VALIDATION CHECK
# =============================================================================

print("\n" + "=" * 80)
print("  VALIDATION CHECK")
print("=" * 80)

accuracy_pass = test_acc >= 55.0
mia_pass = mia_adv <= 0.05

print(f"\nChallenge Requirements:")
print(f"  Accuracy >= 55%:       {test_acc:.2f}% {'✓ PASS' if accuracy_pass else '✗ FAIL'}")
print(f"  MIA advantage <= 5%: {mia_adv*100:.2f}% {'✓ PASS' if mia_pass else '✗ FAIL'}")

if accuracy_pass and mia_pass:
    print("\n" + "=" * 80)
    print("  SOLUTION READY FOR SUBMISSION")
    print("=" * 80)
    print(f"\nSubmit {model_path} to the validator to get the flag!")
else:
    print("\n" + "=" * 80)
    print("  SOLUTION NOT YET PASSING")
    print("=" * 80)
    print("\nConsider adjusting epsilon or other hyperparameters.")

import requests

with open("./models/dp_model.safetensors", "rb") as f:
    response = requests.post(
        "http://154.57.164.72:32303/validate",
        files={"model": ("dp_model.safetensors", f, "application/octet-stream")}
    )
print(response.text)