import os
import random
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
from model import Generator, Discriminator, weights_init

# Configuration
DATAROOT = "data/textures"
WORKERS = 0  # 0 for Windows compatibility
BATCH_SIZE = 32
IMAGE_SIZE = 64
NC = 3
NZ = 100
NGF = 64
NDF = 64
NUM_EPOCHS = 50  # Quick training for demo
LR = 0.0002
BETA1 = 0.5
TEXTURE_TYPES = ["wood", "marble", "fabric", "brick", "noise"]

def train_model(texture_type):
    print(f"\n🚀 Starting training for: {texture_type}")
    
    # Path to specific texture data
    # We need a structure like root/class/img.png for ImageFolder
    # But our generate_dataset.py creates data/textures/wood/*.png
    # So to train 'wood', we need data/textures_wood/wood/*.png? 
    # Or just use the 'wood' folder but ImageFolder expects a root with subfolders.
    # Hack: We can use a custom Dataset or just point ImageFolder to 'data/textures' and filter?
    # Better: Point ImageFolder to 'data/textures' and use 'classes' argument?
    # Actually ImageFolder loads ALL classes.
    # To train specific models, we need to filter the dataset.
    
    dataset = dset.ImageFolder(root=DATAROOT,
                               transform=transforms.Compose([
                                   transforms.Resize(IMAGE_SIZE),
                                   transforms.CenterCrop(IMAGE_SIZE),
                                   transforms.ToTensor(),
                                   transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                               ]))
    
    # Filter dataset indices for the current texture type
    # ImageFolder.class_to_idx gives {'brick': 0, 'fabric': 1, ...}
    target_idx = dataset.class_to_idx[texture_type]
    indices = [i for i, (path, label) in enumerate(dataset.samples) if label == target_idx]
    
    subset = torch.utils.data.Subset(dataset, indices)
    dataloader = torch.utils.data.DataLoader(subset, batch_size=BATCH_SIZE,
                                             shuffle=True, num_workers=WORKERS)

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(f"   Using device: {device}")

    # Initialize models
    netG = Generator(NZ, NGF, NC).to(device)
    netG.apply(weights_init)
    
    netD = Discriminator(NDF, NC).to(device)
    netD.apply(weights_init)

    criterion = nn.BCELoss()

    fixed_noise = torch.randn(64, NZ, 1, 1, device=device)
    real_label = 1.
    fake_label = 0.

    optimizerD = optim.Adam(netD.parameters(), lr=LR, betas=(BETA1, 0.999))
    optimizerG = optim.Adam(netG.parameters(), lr=LR, betas=(BETA1, 0.999))

    print(f"   Training {len(subset)} images...")

    for epoch in range(NUM_EPOCHS):
        errorG = 0.0
        errorD = 0.0
        
        for i, data in enumerate(dataloader, 0):
            ############################
            # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
            ###########################
            ## Train with all-real batch
            netD.zero_grad()
            real_cpu = data[0].to(device)
            b_size = real_cpu.size(0)
            label = torch.full((b_size,), real_label, dtype=torch.float, device=device)
            output = netD(real_cpu).view(-1)
            errD_real = criterion(output, label)
            errD_real.backward()
            D_x = output.mean().item()

            ## Train with all-fake batch
            noise = torch.randn(b_size, NZ, 1, 1, device=device)
            fake = netG(noise)
            label.fill_(fake_label)
            output = netD(fake.detach()).view(-1)
            errD_fake = criterion(output, label)
            errD_fake.backward()
            D_G_z1 = output.mean().item()
            errD = errD_real + errD_fake
            optimizerD.step()

            ############################
            # (2) Update G network: maximize log(D(G(z)))
            ###########################
            netG.zero_grad()
            label.fill_(real_label)  # fake labels are real for generator cost
            output = netD(fake).view(-1)
            errG = criterion(output, label)
            errG.backward()
            D_G_z2 = output.mean().item()
            optimizerG.step()
            
            errorG += errG.item()
            errorD += errD.item()

        if epoch % 10 == 0 or epoch == NUM_EPOCHS - 1:
            print(f"   [{epoch}/{NUM_EPOCHS}] Loss_D: {errorD:.4f} Loss_G: {errorG:.4f}")

    # Save generator
    os.makedirs("checkpoints", exist_ok=True)
    save_path = f"checkpoints/generator_{texture_type}.pth"
    torch.save(netG.state_dict(), save_path)
    print(f"✅ Saved model to {save_path}")

if __name__ == "__main__":
    for t_type in TEXTURE_TYPES:
        train_model(t_type)