#!/usr/bin/env python3

import os
import requests
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 tqdm import tqdm
from opacus import PrivacyEngine
from opacus.validators import ModuleValidator
from safetensors.torch import save_file

# =============================================================================
# CONFIGURATION
# =============================================================================
BATCH_SIZE = 64
DP_EPOCHS = 20       # Increased to reach 55% accuracy target
DP_LR = 1e-3         # Adam LR
MAX_GRAD_NORM = 1.0
DELTA = 1e-5
TARGET_EPSILON = 6.0
DEVICE = torch.device('cpu')
MODELS_DIR = "models"
os.makedirs(MODELS_DIR, exist_ok=True)

# =============================================================================
# STABLE MODEL ARCHITECTURE
# =============================================================================
class ChallengeSVHN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.gn1 = nn.GroupNorm(4, 32)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.gn2 = nn.GroupNorm(8, 64)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 8 * 8, 64)
        self.fc2 = nn.Linear(64, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.gn1(self.conv1(x))))
        x = self.pool(F.relu(self.gn2(self.conv2(x))))
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        return self.fc2(x)

# =============================================================================
# MAIN PROCESS
# =============================================================================
def run_challenge():
    # 1. Data Setup
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4377, 0.4438, 0.4728), (0.1980, 0.2010, 0.1970))
    ])
    train_ds = datasets.SVHN("data", split='train', download=True, transform=transform)
    test_ds = datasets.SVHN("data", split='test', download=True, transform=transform)
    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True)
    test_loader = DataLoader(test_ds, batch_size=BATCH_SIZE, shuffle=False)

    # 2. Model & Optimizer
    model = ChallengeSVHN().to(DEVICE)
    model = ModuleValidator.fix(model)
    optimizer = optim.Adam(model.parameters(), lr=DP_LR)

    # 3. Privacy Engine (Stability Mode)
    privacy_engine = PrivacyEngine()
    model, optimizer, train_loader_dp = privacy_engine.make_private_with_epsilon(
        module=model,
        optimizer=optimizer,
        data_loader=train_loader,
        target_epsilon=TARGET_EPSILON,
        target_delta=DELTA,
        epochs=DP_EPOCHS,
        max_grad_norm=MAX_GRAD_NORM,
        grad_sample_mode="hooks", # CRITICAL: Prevents AMD CPU Segfault
    )

    criterion = nn.CrossEntropyLoss()

    # 4. Training Loop
    print(f"Starting Training for {DP_EPOCHS} epochs...")
    for epoch in range(DP_EPOCHS):
        model.train()
        pbar = tqdm(train_loader_dp, desc=f"Epoch {epoch+1}")
        for images, labels in pbar:
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            pbar.set_postfix({'eps': f'{privacy_engine.get_epsilon(DELTA):.2f}'})

    # 5. Final Evaluation
    model.eval()
    correct, total = 0, 0
    with torch.no_grad():
        for images, labels in test_loader:
            outputs = model(images.to(DEVICE))
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels.to(DEVICE)).sum().item()
    
    acc = 100.0 * correct / total
    eps = privacy_engine.get_epsilon(DELTA)
    print(f"\nTraining Complete. Final Accuracy: {acc:.2f}% | Final Epsilon: {eps:.2f}")

    # 6. Save Model
    # We save model._module because 'model' is wrapped by Opacus
    model_path = os.path.join(MODELS_DIR, "dp_model.safetensors")
    save_file(model._module.state_dict(), model_path)
    print(f"Model saved to {model_path}")

    # 7. Submit to Server
    if acc >= 55.0:
        print("Accuracy requirement met. Submitting for flag...")
        with open(model_path, "rb") as f:
            response = requests.post(
                "http://154.57.164.77:32346/validate",
                files={"model": ("dp_model.safetensors", f, "application/octet-stream")}
            )
        print("\n--- SERVER RESPONSE ---")
        print(response.text)
        print("-----------------------")
    else:
        print("Accuracy was below 55%. Try running again or adjusting DP_LR.")

if __name__ == "__main__":
    # Prevent memory collision on CPU
    torch.set_num_threads(1)
    run_challenge()