import os

import joblib
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset


class MLP(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, d),
        )

    def forward(self, x):
        return self.net(x)


class RealNVP(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.d = d
        self.s1 = MLP(d // 2)
        self.t1 = MLP(d // 2)
        self.s2 = MLP(d // 2)
        self.t2 = MLP(d // 2)

    def forward(self, x):
        x1, x2 = x[:, : self.d // 2], x[:, self.d // 2 :]
        s = self.s1(x1)
        t = self.t1(x1)
        y2 = x2 * torch.exp(s) + t
        y1 = x1

        s2 = self.s2(y2)
        t2 = self.t2(y2)
        z1 = y1 * torch.exp(s2) + t2
        z2 = y2

        z = torch.cat([z1, z2], dim=1)
        log_det = s.sum(1) + s2.sum(1)
        return z, log_det


def train_flow(
    x_profile_joblib: str = "outputs/artifacts/X_profile.joblib",
    artifacts_dir: str = "outputs/artifacts",
    epochs: int = 5,
    batch_size: int = 256,
    lr: float = 1e-5,
):
    os.makedirs(artifacts_dir, exist_ok=True)

    X = joblib.load(x_profile_joblib)
    if hasattr(X, "toarray"):
        X = X.toarray()  # WARNING: can be huge

    X = torch.tensor(X, dtype=torch.float32)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    loader = DataLoader(TensorDataset(X), batch_size=batch_size, shuffle=True)

    D = X.shape[1]
    flow = RealNVP(D).to(device)
    opt = torch.optim.Adam(flow.parameters(), lr=lr)

    def log_prob(z):
        return -0.5 * (z**2).sum(1)

    for epoch in range(epochs):
        flow.train()
        losses = []
        for (x,) in loader:
            x = x.to(device)
            z, log_det = flow(x)
            loss = -(log_prob(z) + log_det).mean()
            opt.zero_grad()
            loss.backward()
            opt.step()
            losses.append(loss.item())
        print(f"Epoch {epoch+1}/{epochs} loss={np.mean(losses):.4f}")

    torch.save(flow.state_dict(), os.path.join(artifacts_dir, "realnvp.pt"))
    print("✅ Saved:", os.path.join(artifacts_dir, "realnvp.pt"))