"""Lecture 3: deterministic AE and Gaussian VAE on the shared observations."""
import argparse
import json
from pathlib import Path
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch import nn
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler
from common import make_dataset


class Autoencoder(nn.Module):
    def __init__(self, n_features):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(n_features, 32), nn.Tanh(), nn.Linear(32, 2))
        self.decoder = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, n_features))

    def forward(self, x):
        z = self.encoder(x)
        return self.decoder(z), z


class VAE(nn.Module):
    def __init__(self, n_features):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(n_features, 32), nn.Tanh())
        self.mu = nn.Linear(32, 2)
        self.logvar = nn.Linear(32, 2)
        self.decoder = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, n_features))

    def forward(self, x, sample=True):
        h = self.encoder(x)
        mu = self.mu(h)
        logvar = self.logvar(h).clamp(-10, 6)
        z = mu + torch.exp(0.5 * logvar) * torch.randn_like(mu) if sample else mu
        return self.decoder(z), mu, logvar


def vae_loss(x, reconstruction, mu, logvar, observation_sd=0.3):
    # Negative Gaussian log likelihood, omitting a constant independent of parameters.
    nll = ((x - reconstruction) ** 2).sum(1) / (2 * observation_sd ** 2)
    kl = 0.5 * (mu.square() + logvar.exp() - 1 - logvar).sum(1)
    return (nll + kl).mean(), kl.mean()


def train_models(data, epochs=300):
    torch.manual_seed(7)
    torch.set_num_threads(1)
    train = data['train']
    scaler = StandardScaler().fit(data['X'][train])
    X = scaler.transform(data['X']).astype(np.float32)
    tensor = torch.from_numpy(X)
    training = tensor[train]
    ae, vae = Autoencoder(X.shape[1]), VAE(X.shape[1])
    ae_optimizer = torch.optim.Adam(ae.parameters(), lr=0.01)
    vae_optimizer = torch.optim.Adam(vae.parameters(), lr=0.01)
    for _ in range(epochs):
        ae_optimizer.zero_grad()
        reconstruction, _ = ae(training)
        ((reconstruction-training) ** 2).mean().backward()
        ae_optimizer.step()
        vae_optimizer.zero_grad()
        reconstruction, mu, logvar = vae(training)
        loss, _ = vae_loss(training, reconstruction, mu, logvar)
        loss.backward()
        vae_optimizer.step()
    ae.eval()
    vae.eval()
    with torch.no_grad():
        ae_x, ae_z = ae(tensor)
        vae_x, vae_z, logvar = vae(tensor, sample=False)
        kl = 0.5 * (vae_z.square() + logvar.exp() - 1 - logvar).sum(1)
    pca = PCA(2, svd_solver='full').fit(X[train])
    pca_z = pca.transform(X)
    pca_x = pca.inverse_transform(pca_z)
    metrics = {'pca_test_mse': float(np.mean((pca_x[~train] - X[~train]) ** 2)),
               'ae_test_mse': float(np.mean((ae_x.numpy()[~train] - X[~train]) ** 2)),
               'vae_mean_decode_test_mse': float(np.mean((vae_x.numpy()[~train] - X[~train]) ** 2)),
               'vae_test_kl': float(kl.numpy()[~train].mean()), 'epochs': int(epochs)}
    return {'PCA': pca_z, 'AE': ae_z.numpy(), 'VAE mean': vae_z.numpy()}, metrics


def run(output_dir: Path) -> dict:
    output_dir.mkdir(parents=True, exist_ok=True)
    data = make_dataset()
    embeddings, metrics = train_models(data)
    test = ~data['train']
    fig, axes = plt.subplots(1, 3, figsize=(13, 4), constrained_layout=True)
    for ax, (name, z) in zip(axes, embeddings.items()):
        ax.scatter(z[test, 0], z[test, 1], c=data['time'][test], s=8)
        ax.set(title=name, xlabel='latent coordinate 1', ylabel='latent coordinate 2')
    fig.savefig(output_dir / 'lesson03.png', dpi=160)
    plt.close(fig)
    return metrics


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--output', type=Path, default=Path('outputs'))
    args = parser.parse_args()
    print(json.dumps(run(args.output), indent=2))