import torch
import numpy as np
from .model import VAE
import cv2
import base64
import os

class Dreamer:
    def __init__(self, model_path=None, latent_dim=128):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.latent_dim = latent_dim
        self.model = VAE(latent_dim).to(self.device)
        
        default_weights = os.path.join(os.path.dirname(__file__), '../../data/weights/vae_circuit.pth')
        if model_path:
            self.model.load_state_dict(torch.load(model_path, map_location=self.device))
        elif os.path.exists(default_weights):
            print(f"Loading weights from {default_weights}")
            self.model.load_state_dict(torch.load(default_weights, map_location=self.device))
        else:
            print("No weights found, using random initialization.")
        self.model.eval()
        
        self.z1 = torch.randn(1, latent_dim).to(self.device)
        self.z2 = torch.randn(1, latent_dim).to(self.device)
        self.t = 0.0
        self.speed = 0.03
        
    def step(self):
        self.t += self.speed
        if self.t >= 1.0:
            self.t = 0.0
            self.z1 = self.z2
            self.z2 = torch.randn(1, self.latent_dim).to(self.device)
            
        z = self.z1 * (1.0 - self.t) + self.z2 * self.t
        with torch.no_grad():
            img_tensor = self.model.decode(z)
            
        img = img_tensor.squeeze().cpu().numpy()
        
        img_min, img_max = img.min(), img.max()
        if img_max > img_min + 1e-5:
            img = (img - img_min) / (img_max - img_min)
        else:
            img = np.random.uniform(0, 0.05, size=img.shape)
        
        img = (img * 255).astype(np.uint8)
        
        _, buffer = cv2.imencode('.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), 80])
        jpg_as_text = base64.b64encode(buffer).decode('utf-8')
        
        return jpg_as_text, z.squeeze().cpu().numpy().tolist()

if __name__ == "__main__":
    dreamer = Dreamer()
    frame, z = dreamer.step()
    print(f"Frame length: {len(frame)}")
    print(f"Latent vector: {z[:5]}...")
