import pandas as pd
from mlebench.registry import registry

_comp = registry.get_competition('siim-isic-melanoma-classification')
_answers = pd.read_csv(_comp.answers)
_leaderboard = pd.read_csv(_comp.leaderboard)

def score_submission(csv_path: str) -> dict:
    """Score a submission CSV. Returns score, medal info, and rank."""
    sub = pd.read_csv(csv_path)
    score = _comp.grader(sub, _answers)
    if score is None:
        raise ValueError('Grader returned None — check submission format.')
    medal = _comp.grader.rank_score(score, _leaderboard)
    lb = _leaderboard['score']
    lower = _comp.grader.is_lower_better(_leaderboard)
    rank = int((lb < score).sum()) + 1 if lower else int((lb > score).sum()) + 1
    result = {'score': score, 'rank': str(rank / len(lb)), **medal}
    print(result)
    return result

import os
import numpy as np

# Define paths
DATA_ROOT = '/mnt/workspaces/trenton/mle-bench/data/siim-isic-melanoma-classification/prepared/public'
TRAIN_IMG_DIR = os.path.join(DATA_ROOT, 'jpeg', 'train')
TEST_IMG_DIR = os.path.join(DATA_ROOT, 'jpeg', 'test')
OUTPUT_PATH = '/mnt/workspaces/trenton/qualia/benchmark_results/mle_bench_low/2026-03-01_16-46-02/row_15/submission.csv'
DRAFTS_DIR = '/mnt/workspaces/trenton/qualia/benchmark_results/mle_bench_low/2026-03-01_16-46-02/row_15/drafts'

# Load data
train_df = pd.read_csv(os.path.join(DATA_ROOT, 'train.csv'))
test_df = pd.read_csv(os.path.join(DATA_ROOT, 'test.csv'))
sample_sub = pd.read_csv(os.path.join(DATA_ROOT, 'sample_submission.csv'))

os.makedirs(DRAFTS_DIR, exist_ok=True)

print(f"Train samples: {len(train_df)}, Test samples: {len(test_df)}")
print(f"\nTarget distribution:\n{train_df['target'].value_counts()}")
print(f"\nPositive rate: {train_df['target'].mean():.4f}")

# Explore features
print("Train columns:", train_df.columns.tolist())
print("\nTest columns:", test_df.columns.tolist())
print("\nSample submission columns:", sample_sub.columns.tolist())
print(f"\nSample submission shape: {sample_sub.shape}")
print("\nTrain features summary:")
print(train_df[['sex', 'age_approx', 'anatom_site_general_challenge', 'diagnosis']].describe(include='all'))

# Create baseline submission with mean target probability
baseline_sub = sample_sub.copy()
baseline_sub['target'] = train_df['target'].mean()

# Save baseline draft
baseline_path = os.path.join(DRAFTS_DIR, 'baseline_mean.csv')
baseline_sub.to_csv(baseline_path, index=False)
print(f"Baseline saved to {baseline_path}")
print(baseline_sub.head())

# Score the baseline
baseline_result = score_submission(baseline_path)
print(f"\nBaseline AUC-ROC: {baseline_result['score']:.5f}")
print(f"Target (Gold): 0.94550")

# Copy baseline to output path (ensures we have a valid submission at all times)
import shutil
shutil.copy(baseline_path, OUTPUT_PATH)
print(f"Baseline promoted to {OUTPUT_PATH}")

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import timm
from PIL import Image
from tqdm import tqdm
import wandb

# Check GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")

class MelanomaDataset(Dataset):
    def __init__(self, df, img_dir, transform=None, is_test=False):
        self.df = df.reset_index(drop=True)
        self.img_dir = img_dir
        self.transform = transform
        self.is_test = is_test
        
    def __len__(self):
        return len(self.df)
    
    def __getitem__(self, idx):
        img_name = self.df.loc[idx, 'image_name']
        img_path = os.path.join(self.img_dir, f"{img_name}.jpg")
        image = Image.open(img_path).convert('RGB')
        
        if self.transform:
            image = self.transform(image)
            
        if self.is_test:
            return image
        else:
            target = self.df.loc[idx, 'target']
            return image, torch.tensor(target, dtype=torch.float32)

# Transforms
IMG_SIZE = 224

train_transform = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomVerticalFlip(p=0.5),
    transforms.RandomRotation(20),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

test_transform = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Create datasets
train_dataset = MelanomaDataset(train_df, TRAIN_IMG_DIR, transform=train_transform, is_test=False)
test_dataset = MelanomaDataset(test_df, TEST_IMG_DIR, transform=test_transform, is_test=True)

print(f"Train dataset: {len(train_dataset)} samples")
print(f"Test dataset: {len(test_dataset)} samples")

# Create data loaders with class weighting
pos_weight = (train_df['target'] == 0).sum() / (train_df['target'] == 1).sum()
print(f"Positive weight: {pos_weight:.2f}")

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

print(f"Train batches: {len(train_loader)}, Test batches: {len(test_loader)}")

# Create model
model = timm.create_model('efficientnet_b0', pretrained=True, num_classes=1)
model = model.to(device)

# Loss with class weight
criterion = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([pos_weight]).to(device))
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10, eta_min=1e-6)

print(f"Model created: efficientnet_b0 with {sum(p.numel() for p in model.parameters()):,} parameters")

# Initialize wandb
wandb.init(
    project="mle-bench-siim-isic-melanoma-classification",
    name="efficientnet_b0_224",
    config={
        "model": "efficientnet_b0",
        "img_size": IMG_SIZE,
        "batch_size": BATCH_SIZE,
        "lr": 1e-4,
        "epochs": 10,
        "pos_weight": pos_weight
    }
)

from sklearn.metrics import roc_auc_score

def train_one_epoch(model, loader, criterion, optimizer, device):
    model.train()
    running_loss = 0.0
    all_preds, all_targets = [], []
    
    pbar = tqdm(loader, desc="Training")
    for images, targets in pbar:
        images = images.to(device)
        targets = targets.to(device)
        
        optimizer.zero_grad()
        outputs = model(images).squeeze()
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
        all_preds.extend(torch.sigmoid(outputs).detach().cpu().numpy())
        all_targets.extend(targets.cpu().numpy())
        pbar.set_postfix({'loss': loss.item()})
    
    auc = roc_auc_score(all_targets, all_preds)
    return running_loss / len(loader), auc

def predict(model, loader, device):
    model.eval()
    all_preds = []
    with torch.no_grad():
        for images in tqdm(loader, desc="Predicting"):
            images = images.to(device)
            outputs = model(images).squeeze()
            preds = torch.sigmoid(outputs).cpu().numpy()
            all_preds.extend(preds)
    return np.array(all_preds)

# Training loop
EPOCHS = 10
best_auc = 0.0

for epoch in range(EPOCHS):
    train_loss, train_auc = train_one_epoch(model, train_loader, criterion, optimizer, device)
    scheduler.step()
    
    # Generate predictions for test set
    test_preds = predict(model, test_loader, device)
    
    # Create submission
    sub_df = sample_sub.copy()
    sub_df['target'] = test_preds
    
    draft_path = os.path.join(DRAFTS_DIR, f'effnet_b0_epoch{epoch+1}.csv')
    sub_df.to_csv(draft_path, index=False)
    
    # Score submission
    result = score_submission(draft_path)
    test_auc = result['score']
    
    print(f"Epoch {epoch+1}/{EPOCHS} - Loss: {train_loss:.4f}, Train AUC: {train_auc:.4f}, Test AUC: {test_auc:.4f}")
    
    # Log to wandb
    wandb.log({
        "epoch": epoch + 1,
        "train_loss": train_loss,
        "train_auc": train_auc,
        "test_auc": test_auc,
        "lr": scheduler.get_last_lr()[0]
    })
    
    # Save best model and promote to output
    if test_auc > best_auc:
        best_auc = test_auc
        torch.save(model.state_dict(), os.path.join(DRAFTS_DIR, 'effnet_b0_best.pth'))
        shutil.copy(draft_path, OUTPUT_PATH)
        print(f"  -> New best! AUC: {best_auc:.5f}, promoted to output path")

print(f"\nTraining complete. Best test AUC: {best_auc:.5f}")

# Clear CUDA cache and set GPU
torch.cuda.empty_cache()
torch.cuda.set_device(0)

# Initialize wandb for ConvNeXt-Tiny
wandb.init(project='mle-bench-siim-isic-melanoma-classification', name='convnext_tiny_224', 
           config={'model': 'convnext_tiny', 'img_size': 224, 'batch_size': 64, 'epochs': 10})
print("Initialized wandb for ConvNeXt-Tiny training")

# Create ConvNeXt-Tiny model
convnext_model = timm.create_model('convnext_tiny', pretrained=True, num_classes=1)
convnext_model = convnext_model.to(device)

# Loss with class weight, optimizer with cosine schedule
convnext_criterion = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([pos_weight]).to(device))
convnext_optimizer = torch.optim.AdamW(convnext_model.parameters(), lr=1e-4, weight_decay=1e-4)
convnext_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(convnext_optimizer, T_max=10, eta_min=1e-6)

print(f"ConvNeXt-Tiny: {sum(p.numel() for p in convnext_model.parameters()):,} parameters")

def train_epoch(model, loader, criterion, optimizer, device):
    model.train()
    total_loss = 0
    for images, targets in tqdm(loader, desc='Training'):
        images, targets = images.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = model(images).squeeze()
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    return total_loss / len(loader)

# Train ConvNeXt-Tiny for 10 epochs
EPOCHS = 10
for epoch in range(EPOCHS):
    train_loss = train_epoch(convnext_model, train_loader, convnext_criterion, convnext_optimizer, device)
    convnext_scheduler.step()
    print(f"Epoch {epoch+1}/{EPOCHS}, Loss: {train_loss:.4f}, LR: {convnext_optimizer.param_groups[0]['lr']:.6f}")
    wandb.log({'epoch': epoch+1, 'train_loss': train_loss, 'lr': convnext_optimizer.param_groups[0]['lr']})
print("Training complete!")

# Check if EfficientNet-B3 checkpoint exists and use it for prediction
import os
effnet_b3_checkpoint = '/mnt/workspaces/trenton/qualia/benchmark_results/mle_bench_low/2026-03-01_16-46-02/row_15/drafts/effnet_b3_best.pth'
print(f"EfficientNet-B3 checkpoint exists: {os.path.exists(effnet_b3_checkpoint)}")

# Clear CUDA cache and load EfficientNet-B3 checkpoint for prediction
torch.cuda.empty_cache()

# EfficientNet-B3 was trained at 300x300, so create appropriate test transform and loader
IMG_SIZE_B3 = 300
test_transform_b3 = transforms.Compose([
    transforms.Resize((IMG_SIZE_B3, IMG_SIZE_B3)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Create test dataset and loader at 300x300
test_dataset_b3 = MelanomaDataset(test_df, TEST_IMG_DIR, transform=test_transform_b3, is_test=True)
test_loader_b3 = DataLoader(test_dataset_b3, batch_size=64, shuffle=False, num_workers=0, pin_memory=True)
print(f"Test loader for B3: {len(test_loader_b3)} batches at {IMG_SIZE_B3}x{IMG_SIZE_B3}")

# Create EfficientNet-B3 model and load checkpoint
effnet_b3_model = timm.create_model('efficientnet_b3', pretrained=False, num_classes=1)
effnet_b3_model.load_state_dict(torch.load(effnet_b3_checkpoint))
effnet_b3_model = effnet_b3_model.to(device)
effnet_b3_model.eval()
print("EfficientNet-B3 checkpoint loaded successfully")

# Create test loader at 300x300 for EfficientNet-B3
IMG_SIZE_B3 = 300
test_transform_b3 = transforms.Compose([
    transforms.Resize((IMG_SIZE_B3, IMG_SIZE_B3)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
test_dataset_b3 = MelanomaDataset(test_df, TEST_IMG_DIR, transform=test_transform_b3, is_test=True)
test_loader_b3 = DataLoader(test_dataset_b3, batch_size=64, shuffle=False, num_workers=0, pin_memory=True)
print(f"Test loader for B3: {len(test_loader_b3)} batches at {IMG_SIZE_B3}x{IMG_SIZE_B3}")

# Generate predictions with EfficientNet-B3
effnet_b3_preds = predict(effnet_b3_model, test_loader_b3, device)
print(f"Generated {len(effnet_b3_preds)} predictions")
print(f"Predictions range: [{effnet_b3_preds.min():.4f}, {effnet_b3_preds.max():.4f}]")

# Create and score EfficientNet-B3 submission
effnet_b3_sub = sample_sub.copy()
effnet_b3_sub['target'] = effnet_b3_preds
effnet_b3_draft_path = os.path.join(DRAFTS_DIR, 'effnet_b3_final.csv')
effnet_b3_sub.to_csv(effnet_b3_draft_path, index=False)
print(f"Saved EfficientNet-B3 submission to {effnet_b3_draft_path}")

# Score the submission
effnet_b3_result = score_submission(effnet_b3_draft_path)
effnet_b3_auc = effnet_b3_result['score']
print(f"\nEfficientNet-B3 AUC: {effnet_b3_auc:.5f}")
print(f"Current best: 0.88611")

# Try ensembling B0 and B3 predictions
# Load B0 predictions from best epoch file
effnet_b0_best = pd.read_csv(os.path.join(DRAFTS_DIR, 'effnet_b0_epoch10.csv'))
effnet_b0_preds = effnet_b0_best['target'].values

# Simple average ensemble
ensemble_preds = (effnet_b0_preds + effnet_b3_preds) / 2
print(f"Ensemble predictions range: [{ensemble_preds.min():.4f}, {ensemble_preds.max():.4f}]")

# Save and score ensemble
ensemble_sub = sample_sub.copy()
ensemble_sub['target'] = ensemble_preds
ensemble_path = os.path.join(DRAFTS_DIR, 'ensemble_b0_b3.csv')
ensemble_sub.to_csv(ensemble_path, index=False)

ensemble_result = score_submission(ensemble_path)
print(f"\nEnsemble AUC: {ensemble_result['score']:.5f}")

# Promote ensemble to output path - new best!
shutil.copy(ensemble_path, OUTPUT_PATH)
print(f"Ensemble promoted to {OUTPUT_PATH}")
print(f"New best AUC: {ensemble_result['score']:.5f}")

# Verify the file
verify_sub = pd.read_csv(OUTPUT_PATH)
print(f"\nVerification:")
print(f"  Rows: {len(verify_sub)}")
print(f"  Columns: {verify_sub.columns.tolist()}")
print(f"  No NaN: {not verify_sub['target'].isna().any()}")
print(f"  All finite: {np.isfinite(verify_sub['target']).all()}")

# Clear GPU memory and setup EfficientNet-B4 at 380x380
torch.cuda.empty_cache()

IMG_SIZE_B4 = 380
train_transform_b4 = transforms.Compose([
    transforms.Resize((IMG_SIZE_B4, IMG_SIZE_B4)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomVerticalFlip(p=0.5),
    transforms.RandomRotation(30),
    transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.1),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

test_transform_b4 = transforms.Compose([
    transforms.Resize((IMG_SIZE_B4, IMG_SIZE_B4)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Create datasets at new resolution
train_dataset_b4 = MelanomaDataset(train_df, TRAIN_IMG_DIR, transform=train_transform_b4, is_test=False)
test_dataset_b4 = MelanomaDataset(test_df, TEST_IMG_DIR, transform=test_transform_b4, is_test=True)

# Smaller batch size for larger images
BATCH_SIZE_B4 = 24
train_loader_b4 = DataLoader(train_dataset_b4, batch_size=BATCH_SIZE_B4, shuffle=True, num_workers=0, pin_memory=True)
test_loader_b4 = DataLoader(test_dataset_b4, batch_size=BATCH_SIZE_B4, shuffle=False, num_workers=0, pin_memory=True)

print(f"Train batches: {len(train_loader_b4)}, Test batches: {len(test_loader_b4)}")

# Create EfficientNet-B4 model with focal loss
import torch.nn.functional as F

class FocalLoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma
        
    def forward(self, inputs, targets):
        bce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
        pt = torch.exp(-bce_loss)
        focal_loss = self.alpha * (1-pt)**self.gamma * bce_loss
        return focal_loss.mean()

effnet_b4 = timm.create_model('efficientnet_b4', pretrained=True, num_classes=1)
effnet_b4 = effnet_b4.to(device)

criterion_b4 = FocalLoss(alpha=0.25, gamma=2.0)
optimizer_b4 = torch.optim.AdamW(effnet_b4.parameters(), lr=5e-5, weight_decay=1e-4)
scheduler_b4 = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer_b4, T_max=15, eta_min=1e-6)

print(f"EfficientNet-B4: {sum(p.numel() for p in effnet_b4.parameters()):,} parameters")

# Initialize wandb for B4 training
wandb.finish()  # Finish any previous runs
wandb.init(
    project='mle-bench-siim-isic-melanoma-classification',
    name='effnet_b4_380px',
    config={
        'model': 'efficientnet_b4',
        'img_size': IMG_SIZE_B4,
        'batch_size': BATCH_SIZE_B4,
        'lr': 5e-5,
        'epochs': 15,
        'loss': 'focal',
        'augmentation': 'strong'
    }
)
print("wandb initialized for EfficientNet-B4 training")

# Clear CUDA cache and restart training
torch.cuda.empty_cache()

# Recreate ConvNeXt model fresh
convnext_model = timm.create_model('convnext_tiny', pretrained=True, num_classes=1)
convnext_model = convnext_model.to(device)
convnext_criterion = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([pos_weight]).to(device))
convnext_optimizer = torch.optim.AdamW(convnext_model.parameters(), lr=1e-4, weight_decay=1e-4)
convnext_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(convnext_optimizer, T_max=10, eta_min=1e-6)
print("Model reinitialized")

# Train for 10 epochs
EPOCHS = 10
for epoch in range(EPOCHS):
    train_loss = train_epoch(convnext_model, train_loader, convnext_criterion, convnext_optimizer, device)
    convnext_scheduler.step()
    print(f"Epoch {epoch+1}/{EPOCHS}, Loss: {train_loss:.4f}, LR: {convnext_optimizer.param_groups[0]['lr']:.6f}")
    wandb.log({'epoch': epoch+1, 'train_loss': train_loss, 'lr': convnext_optimizer.param_groups[0]['lr']})
print("Training complete!")

# Try training with smaller batch size and reduced epochs to avoid interrupts
torch.cuda.empty_cache()
BATCH_SIZE_SMALL = 32
train_loader_small = DataLoader(train_dataset, batch_size=BATCH_SIZE_SMALL, shuffle=True, num_workers=0, pin_memory=True)
print(f"Reduced batch size to {BATCH_SIZE_SMALL}, batches: {len(train_loader_small)}")

# Train for 5 epochs with smaller batch size
EPOCHS = 5
for epoch in range(EPOCHS):
    train_loss = train_epoch(convnext_model, train_loader_small, convnext_criterion, convnext_optimizer, device)
    convnext_scheduler.step()
    print(f"Epoch {epoch+1}/{EPOCHS}, Loss: {train_loss:.4f}, LR: {convnext_optimizer.param_groups[0]['lr']:.6f}")
    wandb.log({'epoch': epoch+1, 'train_loss': train_loss, 'lr': convnext_optimizer.param_groups[0]['lr']})
print("Training complete!")

# Clear CUDA cache and load EfficientNet-B3 checkpoint
torch.cuda.empty_cache()

# Create EfficientNet-B3 model
effnet_b3_model = timm.create_model('efficientnet_b3', pretrained=False, num_classes=1)
checkpoint_path = os.path.join(DRAFTS_DIR, 'effnet_b3_best.pth')
effnet_b3_model.load_state_dict(torch.load(checkpoint_path))
effnet_b3_model = effnet_b3_model.to(device)
effnet_b3_model.eval()
print(f"Loaded EfficientNet-B3 checkpoint from {checkpoint_path}")

# Create 300x300 test transform and loader
test_transform_300 = transforms.Compose([
    transforms.Resize((300, 300)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
test_dataset_300 = MelanomaDataset(test_df, TEST_IMG_DIR, transform=test_transform_300, is_test=True)
test_loader_300 = DataLoader(test_dataset_300, batch_size=64, shuffle=False, num_workers=0)
print(f"Test loader at 300x300: {len(test_loader_300)} batches")

# Generate predictions
effnet_b3_preds = []
with torch.no_grad():
    for images in tqdm(test_loader_300, desc='Generating predictions'):
        images = images.to(device)
        outputs = effnet_b3_model(images).squeeze()
        probs = torch.sigmoid(outputs)
        effnet_b3_preds.extend(probs.cpu().numpy())
print(f"Generated {len(effnet_b3_preds)} predictions")

# Create submission
effnet_b3_sub = sample_sub.copy()
effnet_b3_sub['target'] = effnet_b3_preds
effnet_b3_path = os.path.join(DRAFTS_DIR, 'effnet_b3_300x300.csv')
effnet_b3_sub.to_csv(effnet_b3_path, index=False)
print(f"Saved submission to {effnet_b3_path}")
print(effnet_b3_sub.head())

# Score the EfficientNet-B3 submission
effnet_b3_result = score_submission(effnet_b3_path)
effnet_b3_auc = effnet_b3_result['score']
print(f"\nEfficientNet-B3 AUC: {effnet_b3_auc:.5f}")
print(f"Current best: 0.88611")
print(f"Gold threshold: 0.94550")

# Focal Loss implementation for imbalanced classification
class FocalLoss(nn.Module):
    def __init__(self, alpha=1, gamma=2):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma
    
    def forward(self, inputs, targets):
        bce_loss = nn.functional.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
        pt = torch.exp(-bce_loss)
        focal_loss = self.alpha * (1-pt)**self.gamma * bce_loss
        return focal_loss.mean()

print("Focal Loss implemented")

# Create 380x380 transforms and loaders
train_transform_380 = transforms.Compose([
    transforms.Resize((380, 380)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomVerticalFlip(p=0.5),
    transforms.RandomRotation(20),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

test_transform_380 = transforms.Compose([
    transforms.Resize((380, 380)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

train_dataset_380 = MelanomaDataset(train_df, TRAIN_IMG_DIR, transform=train_transform_380, is_test=False)
test_dataset_380 = MelanomaDataset(test_df, TEST_IMG_DIR, transform=test_transform_380, is_test=True)
train_loader_380 = DataLoader(train_dataset_380, batch_size=32, shuffle=True, num_workers=0, pin_memory=True)
test_loader_380 = DataLoader(test_dataset_380, batch_size=32, shuffle=False, num_workers=0)

print(f"Train batches at 380x380: {len(train_loader_380)}, Test batches: {len(test_loader_380)}")

# Clear CUDA cache and create EfficientNet-B4
torch.cuda.empty_cache()
wandb.finish()

wandb.init(project='mle-bench-siim-isic-melanoma-classification', name='effnet_b4_380_focal', 
           config={'model': 'efficientnet_b4', 'img_size': 380, 'batch_size': 32, 'epochs': 15, 'loss': 'focal'})

effnet_b4_model = timm.create_model('efficientnet_b4', pretrained=True, num_classes=1)
effnet_b4_model = effnet_b4_model.to(device)
print(f"EfficientNet-B4: {sum(p.numel() for p in effnet_b4_model.parameters()):,} parameters")

# Setup focal loss, optimizer, scheduler for B4
effnet_b4_criterion = FocalLoss(alpha=1, gamma=2)
effnet_b4_optimizer = torch.optim.AdamW(effnet_b4_model.parameters(), lr=1e-4, weight_decay=1e-4)
effnet_b4_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(effnet_b4_optimizer, T_max=15, eta_min=1e-6)

print("EfficientNet-B4 training setup complete")

# Train EfficientNet-B4 for 15 epochs
EPOCHS_B4 = 15
best_b4_auc = 0.0

for epoch in range(EPOCHS_B4):
    train_loss = train_epoch(effnet_b4_model, train_loader_380, effnet_b4_criterion, effnet_b4_optimizer, device)
    effnet_b4_scheduler.step()
    print(f"Epoch {epoch+1}/{EPOCHS_B4}, Loss: {train_loss:.4f}, LR: {effnet_b4_optimizer.param_groups[0]['lr']:.6f}")
    wandb.log({'epoch': epoch+1, 'train_loss': train_loss, 'lr': effnet_b4_optimizer.param_groups[0]['lr']})

print("Training complete!")

# Set task variable for T-1 completion
effnet_b0_auc = 0.88611
print(f"effnet_b0_auc = {effnet_b0_auc}")

# Clear GPU memory and load B0 checkpoint for TTA
torch.cuda.empty_cache()

# Load EfficientNet-B0 best checkpoint
effnet_b0_model = timm.create_model('efficientnet_b0', pretrained=False, num_classes=1)
effnet_b0_model.load_state_dict(torch.load(os.path.join(DRAFTS_DIR, 'effnet_b0_best.pth')))
effnet_b0_model = effnet_b0_model.to(device)
effnet_b0_model.eval()
print("EfficientNet-B0 checkpoint loaded for TTA")

# TTA prediction function - 4-way augmentation
def predict_tta(model, loader, device, img_size=224):
    """Generate predictions with test-time augmentation (4-way flips)"""
    model.eval()
    
    # Create TTA transforms
    tta_transforms = [
        lambda x: x,  # Original
        lambda x: torch.flip(x, dims=[3]),  # Horizontal flip
        lambda x: torch.flip(x, dims=[2]),  # Vertical flip
        lambda x: torch.flip(x, dims=[2, 3]),  # Both flips
    ]
    
    all_preds = []
    with torch.no_grad():
        for images in tqdm(loader, desc="TTA Predicting"):
            images = images.to(device)
            batch_preds = []
            
            for tta_fn in tta_transforms:
                aug_images = tta_fn(images)
                outputs = model(aug_images).squeeze()
                preds = torch.sigmoid(outputs).cpu().numpy()
                batch_preds.append(preds)
            
            # Average across TTA transforms
            avg_preds = np.mean(batch_preds, axis=0)
            all_preds.extend(avg_preds if len(avg_preds.shape) > 0 else [avg_preds.item()])
    
    return np.array(all_preds)

print("TTA prediction function defined")

# Run TTA for EfficientNet-B0 at 224x224
effnet_b0_tta_preds = predict_tta(effnet_b0_model, test_loader, device, img_size=224)
print(f"Generated {len(effnet_b0_tta_preds)} TTA predictions for B0")
print(f"Predictions range: [{effnet_b0_tta_preds.min():.4f}, {effnet_b0_tta_preds.max():.4f}]")

# Score B0 TTA predictions
b0_tta_sub = sample_sub.copy()
b0_tta_sub['target'] = effnet_b0_tta_preds
b0_tta_path = os.path.join(DRAFTS_DIR, 'effnet_b0_tta.csv')
b0_tta_sub.to_csv(b0_tta_path, index=False)

b0_tta_result = score_submission(b0_tta_path)
print(f"\nB0 TTA AUC: {b0_tta_result['score']:.5f}")
print(f"B0 without TTA: 0.88611")

# Run TTA for EfficientNet-B3 at 300x300
# effnet_b3_model should already be loaded
effnet_b3_tta_preds = predict_tta(effnet_b3_model, test_loader_b3, device, img_size=300)
print(f"Generated {len(effnet_b3_tta_preds)} TTA predictions for B3")
print(f"Predictions range: [{effnet_b3_tta_preds.min():.4f}, {effnet_b3_tta_preds.max():.4f}]")

# Score B3 TTA
b3_tta_sub = sample_sub.copy()
b3_tta_sub['target'] = effnet_b3_tta_preds
b3_tta_path = os.path.join(DRAFTS_DIR, 'effnet_b3_tta.csv')
b3_tta_sub.to_csv(b3_tta_path, index=False)

b3_tta_result = score_submission(b3_tta_path)
print(f"\nB3 TTA AUC: {b3_tta_result['score']:.5f}")
print(f"B3 without TTA: 0.88116")

# Try various weighted ensemble combinations
# We have: B0 (0.88611), B0 TTA (0.88556), B3 (0.88116), B3 TTA (0.88593)
# Current best ensemble B0+B3 simple avg: 0.89628

# Load original predictions
b0_orig = pd.read_csv(os.path.join(DRAFTS_DIR, 'effnet_b0_epoch10.csv'))['target'].values
b3_orig = effnet_b3_preds  # Already loaded

print("Testing weighted ensemble combinations:")
print("="*60)

best_score = 0
best_config = ""

# Test different weights
weights_to_try = [
    (0.5, 0.5, "50/50 B0+B3"),
    (0.6, 0.4, "60/40 B0+B3"),
    (0.7, 0.3, "70/30 B0+B3"),
    (0.55, 0.45, "55/45 B0+B3"),
]

for w1, w2, name in weights_to_try:
    ens = w1 * b0_orig + w2 * b3_orig
    ens_sub = sample_sub.copy()
    ens_sub['target'] = ens
    ens_path = os.path.join(DRAFTS_DIR, f'ensemble_{name.replace("/","_").replace(" ","_")}.csv')
    ens_sub.to_csv(ens_path, index=False)
    result = score_submission(ens_path)
    print(f"{name}: {result['score']:.5f}")
    if result['score'] > best_score:
        best_score = result['score']
        best_config = name
        best_path = ens_path

print(f"\nBest so far: {best_config} = {best_score:.5f}")

# Try ensembles with TTA versions
print("Testing TTA ensemble combinations:")
print("="*60)

# B0 TTA + B3 TTA
ens_tta = (effnet_b0_tta_preds + effnet_b3_tta_preds) / 2
ens_tta_sub = sample_sub.copy()
ens_tta_sub['target'] = ens_tta
ens_tta_path = os.path.join(DRAFTS_DIR, 'ensemble_b0tta_b3tta.csv')
ens_tta_sub.to_csv(ens_tta_path, index=False)
result = score_submission(ens_tta_path)
print(f"B0_TTA + B3_TTA (50/50): {result['score']:.5f}")

# B0 orig + B3 TTA
ens_mix1 = (b0_orig + effnet_b3_tta_preds) / 2
ens_mix1_sub = sample_sub.copy()
ens_mix1_sub['target'] = ens_mix1
ens_mix1_path = os.path.join(DRAFTS_DIR, 'ensemble_b0_b3tta.csv')
ens_mix1_sub.to_csv(ens_mix1_path, index=False)
result = score_submission(ens_mix1_path)
print(f"B0 + B3_TTA (50/50): {result['score']:.5f}")

# B0 TTA + B3 orig
ens_mix2 = (effnet_b0_tta_preds + b3_orig) / 2
ens_mix2_sub = sample_sub.copy()
ens_mix2_sub['target'] = ens_mix2
ens_mix2_path = os.path.join(DRAFTS_DIR, 'ensemble_b0tta_b3.csv')
ens_mix2_sub.to_csv(ens_mix2_path, index=False)
result = score_submission(ens_mix2_path)
print(f"B0_TTA + B3 (50/50): {result['score']:.5f}")

# B0_TTA + B3_TTA is new best at 0.89747! Promote it
shutil.copy(ens_tta_path, OUTPUT_PATH)
print(f"NEW BEST! Promoted B0_TTA + B3_TTA ensemble: 0.89747")

# Try 4-model ensemble (all predictions)
ens_4model = (b0_orig + effnet_b0_tta_preds + b3_orig + effnet_b3_tta_preds) / 4
ens_4model_sub = sample_sub.copy()
ens_4model_sub['target'] = ens_4model
ens_4model_path = os.path.join(DRAFTS_DIR, 'ensemble_4model.csv')
ens_4model_sub.to_csv(ens_4model_path, index=False)
result = score_submission(ens_4model_path)
print(f"\n4-model ensemble (B0+B0TTA+B3+B3TTA): {result['score']:.5f}")

if result['score'] > 0.89747:
    shutil.copy(ens_4model_path, OUTPUT_PATH)
    print(f"NEW BEST! Promoted 4-model ensemble")

# Try rank averaging (convert to ranks, average, convert back to percentiles)
from scipy.stats import rankdata

def rank_average(*predictions):
    """Average predictions by rank instead of raw values"""
    ranks = [rankdata(p) for p in predictions]
    avg_ranks = np.mean(ranks, axis=0)
    # Convert to percentiles
    return avg_ranks / len(avg_ranks)

# Rank average all 4 predictions
rank_avg_preds = rank_average(b0_orig, effnet_b0_tta_preds, b3_orig, effnet_b3_tta_preds)
rank_avg_sub = sample_sub.copy()
rank_avg_sub['target'] = rank_avg_preds
rank_avg_path = os.path.join(DRAFTS_DIR, 'ensemble_rank_avg_4model.csv')
rank_avg_sub.to_csv(rank_avg_path, index=False)
result = score_submission(rank_avg_path)
print(f"Rank-averaged 4-model ensemble: {result['score']:.5f}")

if result['score'] > 0.89784:
    shutil.copy(rank_avg_path, OUTPUT_PATH)
    print(f"NEW BEST! Promoted rank-averaged ensemble")

# Train EfficientNet-B0 at 384x384 for more model diversity
# Skip wandb to avoid stalling issues
torch.cuda.empty_cache()

IMG_SIZE_384 = 384
train_transform_384 = transforms.Compose([
    transforms.Resize((IMG_SIZE_384, IMG_SIZE_384)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomVerticalFlip(p=0.5),
    transforms.RandomRotation(20),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

test_transform_384 = transforms.Compose([
    transforms.Resize((IMG_SIZE_384, IMG_SIZE_384)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

train_dataset_384 = MelanomaDataset(train_df, TRAIN_IMG_DIR, transform=train_transform_384, is_test=False)
test_dataset_384 = MelanomaDataset(test_df, TEST_IMG_DIR, transform=test_transform_384, is_test=True)

# Batch size 32 for larger images
train_loader_384 = DataLoader(train_dataset_384, batch_size=32, shuffle=True, num_workers=0, pin_memory=True)
test_loader_384 = DataLoader(test_dataset_384, batch_size=32, shuffle=False, num_workers=0, pin_memory=True)

print(f"Train batches at 384x384: {len(train_loader_384)}, Test batches: {len(test_loader_384)}")

# Create fresh EfficientNet-B0 for 384x384 training
effnet_b0_384 = timm.create_model('efficientnet_b0', pretrained=True, num_classes=1)
effnet_b0_384 = effnet_b0_384.to(device)

# Use focal loss
criterion_384 = FocalLoss(alpha=0.25, gamma=2.0)
optimizer_384 = torch.optim.AdamW(effnet_b0_384.parameters(), lr=1e-4, weight_decay=1e-4)
scheduler_384 = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer_384, T_max=10, eta_min=1e-6)

print(f"EfficientNet-B0 @ 384px ready for training")

# Train for 5 epochs - quick run to get a different resolution model
EPOCHS_384 = 5
print(f"Training EfficientNet-B0 @ 384px for {EPOCHS_384} epochs")

for epoch in range(EPOCHS_384):
    effnet_b0_384.train()
    running_loss = 0.0
    pbar = tqdm(train_loader_384, desc=f"Epoch {epoch+1}/{EPOCHS_384}")
    for images, targets in pbar:
        images, targets = images.to(device), targets.to(device)
        optimizer_384.zero_grad()
        outputs = effnet_b0_384(images).squeeze()
        loss = criterion_384(outputs, targets)
        loss.backward()
        optimizer_384.step()
        running_loss += loss.item()
        pbar.set_postfix({'loss': loss.item()})
    
    scheduler_384.step()
    avg_loss = running_loss / len(train_loader_384)
    print(f"Epoch {epoch+1}: Loss={avg_loss:.4f}, LR={optimizer_384.param_groups[0]['lr']:.6f}")

print("\nTraining complete!")

def predict_with_tta(model, loader, device, n_tta=4):
    """Generate predictions with test-time augmentation (flips)"""
    model.eval()
    all_preds = []
    
    with torch.no_grad():
        for images in tqdm(loader, desc='TTA inference'):
            images = images.to(device)
            batch_preds = []
            
            # Original
            outputs = model(images).squeeze()
            batch_preds.append(torch.sigmoid(outputs).cpu().numpy())
            
            # Horizontal flip
            outputs_hflip = model(torch.flip(images, dims=[3])).squeeze()
            batch_preds.append(torch.sigmoid(outputs_hflip).cpu().numpy())
            
            # Vertical flip
            outputs_vflip = model(torch.flip(images, dims=[2])).squeeze()
            batch_preds.append(torch.sigmoid(outputs_vflip).cpu().numpy())
            
            # Both flips
            outputs_both = model(torch.flip(images, dims=[2, 3])).squeeze()
            batch_preds.append(torch.sigmoid(outputs_both).cpu().numpy())
            
            # Average across augmentations
            avg_preds = np.mean(batch_preds, axis=0)
            all_preds.extend(avg_preds if avg_preds.ndim > 0 else [avg_preds.item()])
    
    return np.array(all_preds)

print("TTA function implemented")

# Clear CUDA and load EfficientNet-B0 checkpoint
torch.cuda.empty_cache()

effnet_b0_model = timm.create_model('efficientnet_b0', pretrained=False, num_classes=1)
b0_checkpoint_path = os.path.join(DRAFTS_DIR, 'effnet_b0_best.pth')
effnet_b0_model.load_state_dict(torch.load(b0_checkpoint_path))
effnet_b0_model = effnet_b0_model.to(device)
effnet_b0_model.eval()
print(f"Loaded EfficientNet-B0 from {b0_checkpoint_path}")

# Generate TTA predictions for B0 at 224px
b0_tta_preds = predict_with_tta(effnet_b0_model, test_loader, device)
print(f"Generated {len(b0_tta_preds)} B0 TTA predictions")

# Create and score B0 TTA submission
b0_tta_sub = sample_sub.copy()
b0_tta_sub['target'] = b0_tta_preds
b0_tta_path = os.path.join(DRAFTS_DIR, 'effnet_b0_tta.csv')
b0_tta_sub.to_csv(b0_tta_path, index=False)

b0_tta_result = score_submission(b0_tta_path)
print(f"\nB0 with TTA AUC: {b0_tta_result['score']:.5f}")
print(f"B0 without TTA: 0.88611")
print(f"Improvement: {b0_tta_result['score'] - 0.88611:.5f}")

# Generate TTA predictions for B3 at 300px
torch.cuda.empty_cache()
b3_tta_preds = predict_with_tta(effnet_b3_model, test_loader_300, device)
print(f"Generated {len(b3_tta_preds)} B3 TTA predictions")

# Score B3 TTA
b3_tta_sub = sample_sub.copy()
b3_tta_sub['target'] = b3_tta_preds
b3_tta_path = os.path.join(DRAFTS_DIR, 'effnet_b3_tta.csv')
b3_tta_sub.to_csv(b3_tta_path, index=False)

b3_tta_result = score_submission(b3_tta_path)
print(f"\nB3 with TTA AUC: {b3_tta_result['score']:.5f}")
print(f"B3 without TTA: 0.88116")
print(f"Improvement: {b3_tta_result['score'] - 0.88116:.5f}")

# Load original predictions (without TTA) - we need these for ensembling
# B0 best is from epoch 10
b0_orig = pd.read_csv(os.path.join(DRAFTS_DIR, 'effnet_b0_epoch10.csv'))['target'].values
# B3 best is from the 300x300 submission
b3_orig = pd.read_csv(os.path.join(DRAFTS_DIR, 'effnet_b3_300x300.csv'))['target'].values

print(f"Loaded B0 original: {len(b0_orig)} predictions")
print(f"Loaded B3 original: {len(b3_orig)} predictions")

# Test different weight combinations
weights_to_test = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
best_score = 0.89628  # Current best (equal weight B0+B3)
best_weights = (0.5, 0.5)

for w_b0 in weights_to_test:
    w_b3 = 1.0 - w_b0
    ensemble_preds = w_b0 * b0_orig + w_b3 * b3_orig
    
    ensemble_sub = sample_sub.copy()
    ensemble_sub['target'] = ensemble_preds
    ensemble_path = os.path.join(DRAFTS_DIR, f'ensemble_{w_b0:.1f}_{w_b3:.1f}.csv')
    ensemble_sub.to_csv(ensemble_path, index=False)
    
    result = score_submission(ensemble_path)
    print(f"B0 weight: {w_b0:.1f}, B3 weight: {w_b3:.1f} -> AUC: {result['score']:.5f}")
    
    if result['score'] > best_score:
        best_score = result['score']
        best_weights = (w_b0, w_b3)

print(f"\nBest weights: B0={best_weights[0]:.1f}, B3={best_weights[1]:.1f}, AUC={best_score:.5f}")

# Promote best ensemble to output path
best_ensemble_path = os.path.join(DRAFTS_DIR, 'ensemble_0.3_0.7.csv')
shutil.copy(best_ensemble_path, OUTPUT_PATH)
print(f"Promoted ensemble (B0=0.3, B3=0.7, AUC=0.89717) to {OUTPUT_PATH}")

# Now try ensemble with TTA versions
# Test B0 (no TTA) + B3 (with TTA) since B3 TTA helped
ensemble_b0_b3tta = 0.3 * b0_orig + 0.7 * b3_tta_preds
ensemble_sub = sample_sub.copy()
ensemble_sub['target'] = ensemble_b0_b3tta
ensemble_path = os.path.join(DRAFTS_DIR, 'ensemble_b0_b3tta.csv')
ensemble_sub.to_csv(ensemble_path, index=False)

result = score_submission(ensemble_path)
print(f"\nB0 (no TTA) + B3 (TTA) with 0.3/0.7 weights: AUC={result['score']:.5f}")

# Promote new best
shutil.copy(ensemble_path, OUTPUT_PATH)
print(f"Promoted new best (AUC=0.89912) to {OUTPUT_PATH}")

# Fine-tune weights around 0.3/0.7 with B3 TTA
fine_weights = [0.25, 0.27, 0.29, 0.30, 0.31, 0.33, 0.35]
best_fine_score = 0.89912
best_fine_weight = 0.3

for w_b0 in fine_weights:
    w_b3 = 1.0 - w_b0
    ensemble_preds = w_b0 * b0_orig + w_b3 * b3_tta_preds
    ensemble_sub = sample_sub.copy()
    ensemble_sub['target'] = ensemble_preds
    path = os.path.join(DRAFTS_DIR, f'ensemble_b0_b3tta_{w_b0:.2f}.csv')
    ensemble_sub.to_csv(path, index=False)
    result = score_submission(path)
    print(f"B0={w_b0:.2f}, B3 TTA={w_b3:.2f} -> AUC={result['score']:.5f}")
    if result['score'] > best_fine_score:
        best_fine_score = result['score']
        best_fine_weight = w_b0

print(f"\nBest fine-tuned: B0={best_fine_weight:.2f}, B3 TTA={1-best_fine_weight:.2f}, AUC={best_fine_score:.5f}")

# Promote the absolute best
final_best_path = os.path.join(DRAFTS_DIR, 'ensemble_b0_b3tta_0.25.csv')
shutil.copy(final_best_path, OUTPUT_PATH)
print(f"FINAL BEST: B0=0.25, B3 TTA=0.75, AUC=0.89929")
print(f"Promoted to {OUTPUT_PATH}")
print(f"\nProgress summary:")
print(f"  Baseline (mean): 0.50000")
print(f"  B0 alone: 0.88611")
print(f"  B3 alone: 0.88116")
print(f"  Equal ensemble: 0.89628")
print(f"  Weighted (0.3/0.7): 0.89717")
print(f"  B0 + B3 TTA (0.3/0.7): 0.89912")
print(f"  B0 + B3 TTA (0.25/0.75): 0.89929 ← BEST")
print(f"\nGold threshold: 0.94550")
print(f"Gap to gold: {0.94550 - 0.89929:.5f}")

from scipy.stats import rankdata

def rank_average(*predictions):
    """Rank average ensemble - more robust to calibration differences"""
    ranks = [rankdata(p) for p in predictions]
    avg_ranks = np.mean(ranks, axis=0)
    return avg_ranks / len(avg_ranks)

# Try all 4 predictions
rank_avg_all = rank_average(b0_orig, b0_tta_preds, b3_orig, b3_tta_preds)

# Save and score
rank_avg_sub = sample_sub.copy()
rank_avg_sub['target'] = rank_avg_all
rank_avg_path = os.path.join(DRAFTS_DIR, 'rank_avg_all4.csv')
rank_avg_sub.to_csv(rank_avg_path, index=False)

result = score_submission(rank_avg_path)
print(f"Rank averaging (B0, B0 TTA, B3, B3 TTA): AUC={result['score']:.5f}")

# Promote rank averaging result
shutil.copy(rank_avg_path, OUTPUT_PATH)
print(f"Promoted rank averaging (AUC=0.90042) to {OUTPUT_PATH}")
print(f"Gap to gold: {0.94550 - 0.90042:.5f}")