"""
Campus Weather VAE
==================
VAE for learning weather embeddings from 40 NUS campus weather stations.
Compresses 6 meteorological variables into a 12-dimensional latent space.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict


class WeatherVAE(nn.Module):
    """
    Variational Autoencoder for multivariate weather station data.
    
    Input:  6 weather variables (WindSpeed, WindDir, AirTemp, RelHum, AtmPress, GlobalRad)
    Output: Compact latent embedding z ∈ R^d_latent
    
    The latent space is regularised (KL divergence) to be smooth and interpolable,
    enabling spatial interpolation, forecasting, clustering, and anomaly detection.
    """
    
    def __init__(self, n_vars: int = 6, d_hidden: int = 128, d_latent: int = 12,
                 n_layers: int = 3, beta: float = 0.001):
        super().__init__()
        self.n_vars = n_vars
        self.d_latent = d_latent
        self.beta = beta
        
        # Encoder
        enc = [nn.Linear(n_vars, d_hidden), nn.LayerNorm(d_hidden), nn.GELU()]
        for _ in range(n_layers - 1):
            enc += [nn.Linear(d_hidden, d_hidden), nn.LayerNorm(d_hidden), nn.GELU()]
        self.encoder = nn.Sequential(*enc)
        self.mu_head = nn.Linear(d_hidden, d_latent)
        self.logvar_head = nn.Linear(d_hidden, d_latent)
        
        # Decoder
        dec = [nn.Linear(d_latent, d_hidden), nn.GELU()]
        for _ in range(n_layers - 1):
            dec += [nn.Linear(d_hidden, d_hidden), nn.GELU()]
        dec.append(nn.Linear(d_hidden, n_vars))
        self.decoder = nn.Sequential(*dec)
        
        # Normalisation (set from data)
        self.register_buffer('x_mean', torch.zeros(n_vars))
        self.register_buffer('x_std', torch.ones(n_vars))
    
    def set_normalisation(self, mean: torch.Tensor, std: torch.Tensor):
        self.x_mean.copy_(mean)
        self.x_std.copy_(std)
    
    def normalise(self, x): 
        return (x - self.x_mean) / (self.x_std + 1e-8)
    
    def denormalise(self, x): 
        return x * (self.x_std + 1e-8) + self.x_mean
    
    def encode(self, x_raw: torch.Tensor) -> Dict[str, torch.Tensor]:
        h = self.encoder(self.normalise(x_raw))
        mu = self.mu_head(h)
        logvar = self.logvar_head(h)
        if self.training:
            z = mu + torch.exp(0.5 * logvar) * torch.randn_like(mu)
        else:
            z = mu
        return {'z': z, 'mu': mu, 'logvar': logvar}
    
    def decode(self, z: torch.Tensor) -> torch.Tensor:
        return self.denormalise(self.decoder(z))
    
    def forward(self, x_raw: torch.Tensor) -> Dict[str, torch.Tensor]:
        enc = self.encode(x_raw)
        x_hat = self.decode(enc['z'])
        
        loss_recon = F.mse_loss(self.normalise(x_hat), self.normalise(x_raw))
        loss_kl = -0.5 * torch.mean(1 + enc['logvar'] - enc['mu']**2 - enc['logvar'].exp())
        loss = loss_recon + self.beta * loss_kl
        
        return {'loss': loss, 'loss_recon': loss_recon.item(), 'loss_kl': loss_kl.item(),
                'x_hat': x_hat, 'z': enc['z'], 'mu': enc['mu']}
    
    def get_embedding(self, x_raw: torch.Tensor) -> torch.Tensor:
        """Deterministic embedding (mu)."""
        return self.encode(x_raw)['mu']


def get_config(size='base'):
    return {
        'small': dict(n_vars=6, d_hidden=64, d_latent=8, n_layers=2, beta=0.001),
        'base':  dict(n_vars=6, d_hidden=128, d_latent=12, n_layers=3, beta=0.001),
        'large': dict(n_vars=6, d_hidden=256, d_latent=24, n_layers=4, beta=0.001),
    }[size]