import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import transforms, models
from datasets import load_dataset
from huggingface_hub import login
from kaggle_secrets import UserSecretsClient
from PIL import Image
import os
import random
import pandas as pd
from tqdm.auto import tqdm

# config
user_secrets = UserSecretsClient()
try:
    hf_token = user_secrets.get_secret("HF_TOKEN")
    login(token=hf_token)
except:
    print("HF_TOKEN not found in Secrets. Ensure you added it!")

KAGLE_REAL_PATH = "/kaggle/input/datasets/matthewjansen/unsplash-lite-5k-colorization/train/color"
HF_AI_DATASET = "Rapidata/Flux_SD3_MJ_Dalle_Human_Alignment_Dataset"
TARGET_SHARDS = ["train_0001", "train_0002", "train_0003", "train_0004"]
SAVE_PATH = "/kaggle/working/convnext_forensic_head.pth"
LOG_PATH = "/kaggle/working/convnext_training_log.csv"

# training params
BATCH_SIZE = 32
EPOCHS = 5
LR = 1e-4
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
IMG_SIZE = (224, 224)

final_data = []
print(f"Streaming AI shards: {TARGET_SHARDS}")
for shard in TARGET_SHARDS:
    shard_stream = load_dataset(HF_AI_DATASET, split=shard, streaming=True)
    for item in tqdm(shard_stream, total=1000, desc=f"AI Shard {shard}"):
        img = item["image1"].convert("RGB").resize(IMG_SIZE)
        final_data.append({"image": img, "label": 1})

real_files = [os.path.join(KAGLE_REAL_PATH, f) for f in os.listdir(KAGLE_REAL_PATH) 
              if f.lower().endswith(('.jpg', '.png', '.jpeg'))]
random.shuffle(real_files)

print(f"Balancing with {len(final_data)} Real images")
for i in tqdm(range(min(len(final_data), len(real_files))), desc="Processing Real"):
    try:
        img = Image.open(real_files[i]).convert("RGB").resize(IMG_SIZE)
        final_data.append({"image": img, "label": 0})
    except: continue

random.shuffle(final_data)
split_idx = int(len(final_data) * 0.85)
train_list, val_list = final_data[:split_idx], final_data[split_idx:]

# model details
backbone = models.convnext_base(weights='IMAGENET1K_V1')
backbone = backbone.to(DEVICE)
for param in backbone.parameters():
    param.requires_grad = False
backbone.eval() 

class ForensicHead(nn.Module):
    def __init__(self, input_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 1)
        )
    def forward(self, x): return self.net(x)

feature_dim = backbone.classifier[2].in_features 
head = ForensicHead(input_dim=feature_dim).to(DEVICE)

if torch.cuda.device_count() > 1:
    print(f"Activating Dual-GPU Mode with {torch.cuda.device_count()} T4s")
    head = nn.DataParallel(head)
    backbone = nn.DataParallel(backbone)

# preprocessing
preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

def collate_fn(batch):
    imgs = torch.stack([preprocess(item['image']) for item in batch])
    lbls = torch.tensor([item['label'] for item in batch]).float().view(-1, 1)
    return imgs, lbls

train_loader = DataLoader(train_list, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)
val_loader = DataLoader(val_list, batch_size=BATCH_SIZE, collate_fn=collate_fn)

# training loop
optimizer = optim.Adam(head.parameters(), lr=LR)
criterion = nn.BCEWithLogitsLoss() 
scaler = torch.amp.GradScaler('cuda') 

best_acc, history = 0.0, []

for epoch in range(EPOCHS):
    head.train()
    train_loss = 0
    pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{EPOCHS}")
    
    for imgs, lbls in pbar:
        imgs, lbls = imgs.to(DEVICE), lbls.to(DEVICE)
        optimizer.zero_grad()

        with torch.amp.autocast('cuda'):
            with torch.no_grad():
                # Extract features handling DataParallel wrapper
                if isinstance(backbone, nn.DataParallel):
                    feat = backbone.module.features(imgs)
                    feat = backbone.module.avgpool(feat)
                else:
                    feat = backbone.features(imgs)
                    feat = backbone.avgpool(feat)
                feat = torch.flatten(feat, 1)
            
            logits = head(feat)
            loss = criterion(logits, lbls)

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        train_loss += loss.item()
        pbar.set_postfix(loss=f"{loss.item():.4f}")

    # validation
    head.eval()
    val_correct = 0
    with torch.no_grad():
        for imgs, lbls in val_loader:
            imgs, lbls = imgs.to(DEVICE), lbls.to(DEVICE)
            with torch.amp.autocast('cuda'):
                if isinstance(backbone, nn.DataParallel):
                    feat = backbone.module.features(imgs)
                    feat = backbone.module.avgpool(feat)
                else:
                    feat = backbone.features(imgs)
                    feat = backbone.avgpool(feat)
                feat = torch.flatten(feat, 1)
                
                logits = head(feat)
                preds = (torch.sigmoid(logits) > 0.5).float()
                
            val_correct += (preds == lbls).sum().item()

    val_acc = val_correct / len(val_list)
    avg_loss = train_loss / len(train_loader)
    print(f"Epoch {epoch+1} | Loss: {avg_loss:.4f} | Val Acc: {val_acc:.4f}")
    
    history.append({'epoch': epoch+1, 'val_acc': val_acc, 'train_loss': avg_loss})
    pd.DataFrame(history).to_csv(LOG_PATH, index=False)

    if val_acc > best_acc:
        best_acc = val_acc
        save_state = head.module.state_dict() if isinstance(head, nn.DataParallel) else head.state_dict()
        torch.save(save_state, SAVE_PATH)
        print("--> Best Model Saved!")

print(f"Training Complete. File saved: {SAVE_PATH}")