import torch
from torch.utils.data import DataLoader
from datasets import load_dataset
from torchvision import transforms
from huggingface_hub import HfApi, create_repo, upload_folder
import os
from tqdm.auto import tqdm
import json

class LouTrainer:
    def __init__(self, model, diffusion, device='cuda'):
        self.model = model
        self.diffusion = diffusion
        self.device = device
        self.optimizer = None
        self.current_epoch = 0
        
    def setup_optimizer(self, lr=1e-4):
        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr, weight_decay=0.01)
        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=1000)
    
    def prepare_dataset(self, dataset_name="diffusers/pokemon-gpt4-captions", batch_size=4, image_size=256):
        """Charge et prépare le dataset"""
        
        # Transformations
        transform = transforms.Compose([
            transforms.Resize(image_size),
            transforms.CenterCrop(image_size),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
        ])
        
        # Chargement dataset
        dataset = load_dataset(dataset_name, split="train")
        
        def transform_examples(examples):
            images = [transform(img.convert("RGB")) for img in examples["image"]]
            return {"pixel_values": images}
        
        dataset.set_transform(transform_examples)
        
        dataloader = DataLoader(
            dataset, 
            batch_size=batch_size, 
            shuffle=True, 
            num_workers=0,
            pin_memory=True
        )
        
        return dataloader
    
    def train_epoch(self, dataloader, progress_bar=None):
        """Entraîne une epoch"""
        total_loss = 0
        
        for batch_idx, batch in enumerate(dataloader):
            images = batch["pixel_values"].to(self.device)
            
            self.optimizer.zero_grad()
            loss = self.diffusion.training_step(images)
            loss.backward()
            
            # Gradient clipping
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
            
            self.optimizer.step()
            
            total_loss += loss.item()
            
            if progress_bar:
                progress_bar.update(1)
                progress_bar.set_postfix({"loss": f"{loss.item():.4f}"})
        
        self.scheduler.step()
        return total_loss / len(dataloader)
    
    def train(self, dataset_name, num_epochs, lr, batch_size=4, image_size=256, 
              save_every=5, repo_id=None, token=None, progress_callback=None):
        """Entraînement complet"""
        
        self.setup_optimizer(lr)
        dataloader = self.prepare_dataset(dataset_name, batch_size, image_size)
        
        total_steps = len(dataloader) * num_epochs
        
        with tqdm(total=total_steps, desc="Entraînement Lou") as pbar:
            for epoch in range(num_epochs):
                self.current_epoch = epoch
                avg_loss = self.train_epoch(dataloader, pbar)
                
                print(f"Epoch {epoch+1}/{num_epochs} - Loss: {avg_loss:.4f}")
                
                # Sauvegarde périodique
                if (epoch + 1) % save_every == 0 and repo_id and token:
                    self.save_checkpoint(repo_id, token, epoch + 1)
                
                if progress_callback:
                    progress_callback(epoch + 1, num_epochs, avg_loss)
        
        # Sauvegarde finale
        if repo_id and token:
            self.save_checkpoint(repo_id, token, "final")
        
        return True
    
    def save_checkpoint(self, repo_id, token, epoch):
        """Sauvegarde sur HF Hub"""
        os.makedirs("checkpoints", exist_ok=True)
        
        # Sauvegarde locale
        checkpoint_path = f"checkpoints/lou_epoch_{epoch}.pt"
        torch.save({
            'epoch': epoch,
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict() if self.optimizer else None,
        }, checkpoint_path)
        
        # Upload vers HF
        try:
            api = HfApi(token=token)
            try:
                create_repo(repo_id, private=False, token=token)
            except:
                pass
            
            api.upload_file(
                path_or_fileobj=checkpoint_path,
                path_in_repo=f"lou_epoch_{epoch}.pt",
                repo_id=repo_id,
                token=token
            )
            print(f"✓ Checkpoint epoch {epoch} uploadé sur {repo_id}")
        except Exception as e:
            print(f"✗ Erreur upload: {e}")
    
    def load_checkpoint(self, checkpoint_path):
        """Charge un checkpoint"""
        checkpoint = torch.load(checkpoint_path, map_location=self.device)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        if self.optimizer and checkpoint['optimizer_state_dict']:
            self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        self.current_epoch = checkpoint.get('epoch', 0)
        print(f"✓ Checkpoint chargé (epoch {self.current_epoch})")