import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import timm
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Subset
from sklearn.model_selection import train_test_split
from torch.optim.lr_scheduler import CosineAnnealingLR

# =========================
# CONFIG
# =========================
DATA_DIR = r"D:\BPA PROJECT\train"
BATCH_SIZE = 32  # Tiny model allows for larger batches = faster training
EPOCHS = 30
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print(f"Checking hardware... Using: {DEVICE}")

# =========================
# MODEL
# =========================
class BreedClassifier(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        print("📥 Downloading/Loading ConvNeXt-Tiny weights (this may take a minute)...")
        # Tiny is much faster than Base but still extremely accurate
        self.model = timm.create_model(
            "convnext_tiny.fb_in22k_ft_in1k", 
            pretrained=True, 
            num_classes=num_classes
        )

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

# =========================
# TRANSFORMS
# =========================
norm_stats = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])

train_tfms = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(*norm_stats)
])

val_tfms = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(*norm_stats)
])

# =========================
# DATA LOADING
# =========================
def load_data():
    print("📂 Scanning dataset folders...")
    full_dataset = datasets.ImageFolder(DATA_DIR)
    class_names = full_dataset.classes
    
    indices = list(range(len(full_dataset)))
    train_idx, val_idx = train_test_split(
        indices, 
        test_size=0.2, 
        stratify=full_dataset.targets,
        random_state=42
    )

    train_ds = Subset(datasets.ImageFolder(DATA_DIR, transform=train_tfms), train_idx)
    val_ds = Subset(datasets.ImageFolder(DATA_DIR, transform=val_tfms), val_idx)

    # Set num_workers=0 if you are on Windows to avoid multi-processing hangs
    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
    val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, num_workers=0)

    return train_loader, val_loader, len(class_names), class_names

# =========================
# TRAINING
# =========================
def train():
    train_loader, val_loader, num_classes, class_names = load_data()
    
    model = BreedClassifier(num_classes).to(DEVICE)

    # Higher learning rate for the smaller 'tiny' model
    optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.05)
    criterion = nn.CrossEntropyLoss() 
    scheduler = CosineAnnealingLR(optimizer, T_max=EPOCHS)

    print(f"✅ Setup complete. Starting {EPOCHS} epochs of training...")

    best_acc = 0
    for epoch in range(EPOCHS):
        model.train()
        total_loss = 0
        correct = total = 0

        for batch_idx, (x, y) in enumerate(train_loader):
            x, y = x.to(DEVICE), y.to(DEVICE)

            optimizer.zero_grad()
            out = model(x)
            loss = criterion(out, y)
            loss.backward()
            optimizer.step()

            pred = out.argmax(1)
            correct += (pred == y).sum().item()
            total += y.size(0)
            total_loss += loss.item()
            
            if batch_idx % 5 == 0:
                print(f"Epoch {epoch+1} | Batch {batch_idx}/{len(train_loader)} | Loss: {loss.item():.4f}", end='\r')

        # VALIDATION
        model.eval()
        val_correct = val_total = val_conf = 0
        with torch.no_grad():
            for x, y in val_loader:
                x, y = x.to(DEVICE), y.to(DEVICE)
                out = model(x)
                probs = F.softmax(out, dim=1)
                conf, pred = torch.max(probs, dim=1)

                val_correct += (pred == y).sum().item()
                val_total += y.size(0)
                val_conf += conf.sum().item()

        val_acc = 100 * val_correct / val_total
        avg_conf = 100 * val_conf / val_total
        
        print(f"\n✨ Epoch [{epoch+1}/{EPOCHS}] - Val Acc: {val_acc:.2f}% | Confidence: {avg_conf:.2f}%")
        
        scheduler.step()

        if val_acc > best_acc:
            best_acc = val_acc
            torch.save(model.state_dict(), "best_breed_model.pth")
            print("💾 Model Saved!")

if __name__ == "__main__":
    train()