# This is code I used to fix the codebook collapse and train atom_vqvae_2_full model.

import os
import glob
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
from safetensors.torch import load_model, save_model

# Import the VQ-VAE model implementation
from AtomVQVAE import ATOMVQVAE

# ---------------------------------------------------------
# 1. Dataset for Loading Images from a Single Directory
# ---------------------------------------------------------
class FlatFolderDataset(Dataset):
    def __init__(self, folder_path, transform=None):
        # Collect all supported image files from the target directory
        self.image_paths = []
        for ext in ('*.png', '*.jpg', '*.jpeg', '*.PNG', '*.JPG', '*.JPEG'):
            self.image_paths.extend(glob.glob(os.path.join(folder_path, ext)))
        
        if len(self.image_paths) == 0:
            raise ValueError(f"No images found in {folder_path}!")
            
        self.transform = transform

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        try:
            image = Image.open(img_path).convert('RGB')
        except Exception as e:
            # Fallback to the first valid image if loading fails
            image = Image.open(self.image_paths[0]).convert('RGB')
            
        if self.transform:
            image = self.transform(image)
        return image

# ---------------------------------------------------------
# 2. Dead Token Recovery Utility
# ---------------------------------------------------------
def revive_dead_tokens(model, dead_tokens, dataloader, device):
    """Replace unused codebook entries with encoder feature vectors."""
    print(f"--- Reviving {len(dead_tokens)} dead tokens ---")
    model.eval()
    
    # Use a batch of real images to generate replacement feature vectors
    images = next(iter(dataloader)).to(device)
    
    with torch.no_grad():
        # Update this reference if your encoder uses a different attribute name
        try:
            continuous_features = model.encoder(images)
            
            # Apply the pre-quantization convolution if present
            if hasattr(model, 'quant_conv'):
                continuous_features = model.quant_conv(continuous_features)
                
        except AttributeError:
            print("ERROR: Could not find 'model.encoder'. Please update the variable name in the script!")
            return

        # Flatten features from [B, C, H, W] to [N, C]
        # Expected channel size is 128 for this model
        flat_features = continuous_features.permute(0, 2, 3, 1).reshape(-1, 128)
        
        # Reinitialize each inactive token with a sampled feature vector
        for token_id in dead_tokens:
            random_idx = torch.randint(0, flat_features.size(0), (1,)).item()
            
            # Update this reference if the VQ embedding layer uses a different name
            model.vq_layer.embedding.weight.data[token_id] = flat_features[random_idx]

    model.train()
    print("--- Revival Complete! ---")

# ---------------------------------------------------------
# 3. Training Entry Point
# ---------------------------------------------------------
def train():
    # --- Training Configuration ---
    DATA_FOLDER = "/workspace/10kimages"  # Update this path to your image dataset
    WEIGHTS_PATH = "atom_vqvae.safetensors"
    EPOCHS = 10
    BATCH_SIZE = 64
    LEARNING_RATE = 5e-5  # Reduced learning rate for fine-tuning
    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    VOCAB_SIZE = 2048
    
    print(f"Using device: {DEVICE}")

    # --- Model Initialization ---
    model = ATOMVQVAE().to(DEVICE)
    if os.path.exists(WEIGHTS_PATH):
        print(f"Loading existing weights from {WEIGHTS_PATH}...")
        load_model(model, WEIGHTS_PATH)
    else:
        print("No existing weights found. Starting from scratch!")

    # --- Dataset and DataLoader Setup ---
    transform = transforms.Compose([
        transforms.Resize((128, 128)),
        transforms.ToTensor()  # Converts image values to the [0, 1] range
    ])
    
    dataset = FlatFolderDataset(DATA_FOLDER, transform=transform)
    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=True)
    print(f"Loaded {len(dataset)} images.")

    # --- Optimizer and Loss Functions ---
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
    mse_loss_fn = nn.MSELoss()

    # --- Main Training Loop ---
    for epoch in range(EPOCHS):
        model.train()
        total_loss = 0
        active_tokens_this_epoch = set()

        for batch_idx, images in enumerate(dataloader):
            images = images.to(DEVICE)
            optimizer.zero_grad()

            # Forward pass
            # Adjust unpacking if your model returns different outputs
            reconstructed, vq_loss, _ = model(images)
            
            # Reconstruction loss measures image quality
            reconstruction_loss = mse_loss_fn(reconstructed, images)
            
            # Combined training objective
            loss = reconstruction_loss + vq_loss
            
            # Backpropagation and optimizer step
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

            # Track which tokens were used during this epoch
            with torch.no_grad():
                tokens = model.encode_to_tokens(images)
                active_tokens_this_epoch.update(tokens.unique().tolist())

            if batch_idx % 10 == 0:
                print(f"Epoch [{epoch+1}/{EPOCHS}] Batch [{batch_idx}/{len(dataloader)}] Loss: {loss.item():.4f}")

        # --- Epoch Summary and Token Usage Analysis ---
        avg_loss = total_loss / len(dataloader)
        print(f"\n=== Epoch {epoch+1} Summary ===")
        print(f"Average Loss: {avg_loss:.4f}")
        print(f"Active Tokens this Epoch: {len(active_tokens_this_epoch)} / {VOCAB_SIZE}")
        
        # Identify tokens that were never activated
        dead_tokens = [i for i in range(VOCAB_SIZE) if i not in active_tokens_this_epoch]
        
        # Reinitialize inactive tokens if any are found
        if len(dead_tokens) > 0:
            revive_dead_tokens(model, dead_tokens, dataloader, DEVICE)
            
        # Save model weights after each epoch
        save_model(model, WEIGHTS_PATH)
        print(f"Model saved to {WEIGHTS_PATH}\n")

if __name__ == "__main__":
    train()