import torch
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import cv2
import os
import glob
import numpy as np
from core.model import VAE
from flask_socketio import SocketIO
import gc
import base64

class CircuitDataset(Dataset):
    def __init__(self, data_dir):
        all_files = glob.glob(os.path.join(data_dir, "*.png"))
        self.file_list = []
        for f in all_files:
            if os.path.getsize(f) > 0:
                self.file_list.append(f)
        
        if len(all_files) != len(self.file_list):
            print(f"[DATA] Filtered {len(all_files) - len(self.file_list)} empty/corrupt files.", flush=True)
        
    def __len__(self):
        return len(self.file_list)
        
    def __getitem__(self, idx):
        img_path = self.file_list[idx]
        img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
        
        if img is None:
            print(f"[WARN] Failed to load image: {img_path}. Skipping.", flush=True)
            img = np.zeros((512, 512), dtype=np.uint8)
            
        img = img.astype(np.float32) / 255.0
        img = torch.from_numpy(img).unsqueeze(0)
        return img

def loss_function(recon_logits, x, mu, logvar):
    BCE = torch.nn.functional.binary_cross_entropy_with_logits(recon_logits, x, reduction='sum')
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return BCE + KLD

def train_with_socket(socketio=None, epochs=10, batch_size=8, lr=5e-4):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[TRAIN] Device: {device} (AMP_ENABLED, CLIP_GRAD)", flush=True)
    
    dataset = CircuitDataset("data/training")
    if len(dataset) == 0:
        if socketio: socketio.emit('train_error', {'msg': 'No training data found.'})
        return
        
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, pin_memory=True)
    model = VAE().to(device)
    optimizer = optim.Adam(model.parameters(), lr=lr)
    
    scaler = torch.amp.GradScaler('cuda') if device.type == 'cuda' else None
    
    fixed_z = torch.randn(1, 128).to(device)
    
    model.train()
    for epoch in range(epochs):
        train_loss = 0
        for batch_idx, data in enumerate(dataloader):
            data = data.to(device)
            optimizer.zero_grad()
            
            if scaler:
                with torch.amp.autocast('cuda'):
                    recon_logits, mu, logvar = model(data)
                    loss = loss_function(recon_logits, data, mu, logvar)
                
                scaler.scale(loss).backward()
                
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                
                scaler.step(optimizer)
                scaler.update()
            else:
                recon_logits, mu, logvar = model(data)
                loss = loss_function(recon_logits, data, mu, logvar)
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                optimizer.step()
            
            train_loss += loss.item()
            
            if batch_idx % 10 == 0:
                del recon_logits, mu, logvar
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                gc.collect()

            if socketio and batch_idx % 10 == 0:
                progress = (batch_idx + 1) / len(dataloader) * 100
                
                with torch.no_grad():
                    img_tensor = model.decode(fixed_z)
                    img = img_tensor.squeeze().cpu().numpy()
                    
                    img_min, img_max = img.min(), img.max()
                    if img_max > img_min:
                        img = (img - img_min) / (img_max - img_min)
                    
                    img = (img * 255).astype(np.uint8)
                    _, buffer = cv2.imencode('.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), 50])
                    img_b64 = base64.b64encode(buffer).decode('utf-8')

                socketio.emit('train_progress', {
                    'epoch': epoch,
                    'progress': progress,
                    'loss': loss.item() / len(data),
                    'imagination': img_b64
                })
                socketio.sleep(0.01)

        avg_loss = train_loss / len(dataset)
        print(f'====> Epoch: {epoch} Average loss: {avg_loss:.4f}', flush=True)
        if socketio:
            socketio.emit('epoch_complete', {'epoch': epoch, 'avg_loss': avg_loss})
            socketio.sleep(0.01)
            
    torch.save(model.state_dict(), "data/weights/vae_circuit.pth")
    if socketio: socketio.emit('train_complete', {'path': 'data/weights/vae_circuit.pth'})

if __name__ == "__main__":
    train_with_socket(epochs=10)
