import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from datasets import load_dataset
import open_clip
from tqdm.auto import tqdm
import os
import random
from PIL import Image
from huggingface_hub import login
from kaggle_secrets import UserSecretsClient

try:
    user_secrets = UserSecretsClient()
    hf_token = user_secrets.get_secret("HF_TOKEN")
    login(token=hf_token)
    print("Successfully logged into Hugging Face")
except Exception as e:
    print("Warning: Could not find HF_TOKEN in Kaggle Secrets. Proceeding anonymously")

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

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

# data loading with streaming
print(f"Streaming {len(TARGET_SHARDS)} shards from Hugging Face")
final_data = []

for shard in TARGET_SHARDS:
    print(f"Opening stream for {shard}")
    shard_stream = load_dataset(HF_AI_DATASET, split=shard, streaming=True)

    for item in tqdm(shard_stream, total=1000, desc=f"Streaming {shard}"):
        # Resize immediately to keep RAM usage low
        img = item["image1"].convert("RGB").resize(IMG_SIZE)
        final_data.append({
            "image": img,
            "label": 1
        })

num_ai_images = len(final_data)
print(f"Total AI images collected: {num_ai_images}")

print("Loading Real Images from Kaggle")
real_images_list = [os.path.join(KAGLE_REAL_PATH, f) for f in os.listdir(KAGLE_REAL_PATH) if
                    f.endswith(('.jpg', '.jpeg', '.png'))]
random.shuffle(real_images_list)

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

# shuffle split
random.seed(42)
random.shuffle(final_data)

split_idx = int(len(final_data) * 0.85)
train_list = final_data[:split_idx]
val_list = final_data[split_idx:]

print(f"Dataset prepared: Train size = {len(train_list)}, Val size = {len(val_list)}")

# model init
print(f"Initializing ViT-L-14 on {DEVICE}")
model, _, preprocess_val = open_clip.create_model_and_transforms(
    'ViT-L-14',
    pretrained='datacomp_xl_s13b_b90k'
)
model = model.to(DEVICE)

# freeze backbone
for param in model.parameters():
    param.requires_grad = False

print("Detecting feature dimensions...")
with torch.no_grad():
    dummy_input = torch.randn(1, 3, 224, 224).to(DEVICE)
    dummy_feature = model.encode_image(dummy_input)
    detected_dim = dummy_feature.shape[1]
    print(f"Backbone output dimension: {detected_dim}")


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),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.net(x)


# Initialize head with detected dimension (768 for DataComp ViT-L-14)
head = ForensicHead(input_dim=detected_dim).to(DEVICE)


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


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.BCELoss()
best_acc = 0.0

print(f"Starting training on {len(train_list)} images")

for epoch in range(EPOCHS):
    head.train()
    train_pbar = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{EPOCHS} [Train]")

    epoch_loss = 0
    for imgs, lbls in train_pbar:
        imgs, lbls = imgs.to(DEVICE), lbls.to(DEVICE)

        with torch.no_grad():
            features = model.encode_image(imgs)
            features /= features.norm(dim=-1, keepdim=True)

        optimizer.zero_grad()
        outputs = head(features)
        loss = criterion(outputs, lbls)
        loss.backward()
        optimizer.step()

        epoch_loss += loss.item()
        train_pbar.set_postfix(loss=f"{loss.item():.4f}")

    # validation
    head.eval()
    val_correct = 0
    val_pbar = tqdm(val_loader, desc=f"Epoch {epoch + 1}/{EPOCHS} [Val]")

    with torch.no_grad():
        for imgs, lbls in val_pbar:
            imgs, lbls = imgs.to(DEVICE), lbls.to(DEVICE)
            feat = model.encode_image(imgs)
            feat /= feat.norm(dim=-1, keepdim=True)
            preds = (head(feat) > 0.5).float()
            val_correct += (preds == lbls).sum().item()

    val_acc = val_correct / len(val_list)
    print(f"Epoch {epoch + 1} Results | Loss: {epoch_loss / len(train_loader):.4f} | Val Acc: {val_acc:.4f}")

    if val_acc > best_acc:
        best_acc = val_acc
        torch.save(head.state_dict(), SAVE_PATH)
        print(f"New best model saved with {val_acc:.4f} accuracy")

print(f"Training complete. Model saved in: {SAVE_PATH}")