import torch
from ms_training import load_trajs # assuming the model is defined in ms_training.py
import numpy as np
import json
from neuralop.models import FNO, TFNO
import seaborn
from neuralop.utils import count_params
import os
import cv2

MODELS = {
    "FNO": FNO,
    "TFNO": TFNO,
}

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

def vorticity2rgb(
    w,
    vmin = -1.25,
    vmax = 1.25,
):
    w = (w - vmin) / (vmax - vmin)
    w = 2 * w - 1
    w = np.sign(w) * np.abs(w) ** 0.8
    w = (w + 1) / 2
    w = seaborn.cm.icefire(w)
    w = 256 * w[..., :3]
    w = w.astype(np.uint8)

    return w

def load_model(dir):
    with open(dir + "/model_config.json", "r") as f:
        model_config = json.load(f)

    print(model_config)

    model_config['model'] = MODELS[model_config['model']]
    model = model_config['model'](**model_config['params'])
    model.load_state_dict(torch.load(dir + "/model.pt", map_location=torch.device('cpu')))
    model.eval()

    return model


def sim_traj(model, gt_traj):
    """
    Traj: [T, H, W]

    Output: pred_traj [T, H, W]
    """

    pred_traj = torch.zeros_like(gt_traj).to(gt_traj.device)
    pred_traj[0] = gt_traj[0]
    with torch.no_grad():
        for i in range(1, gt_traj.shape[0]):
            w = pred_traj[i-1].unsqueeze(0).unsqueeze(0)
            pred_traj[i] = model(w)

    return pred_traj


def sim_and_render(folder_name, model, gt_traj):
    pred_traj = sim_traj(model, gt_traj)

    os.makedirs(f"renders/{folder_name}", exist_ok=True)

    def write_traj(traj, name, fps=10):
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(f'renders/{folder_name}/{name}.mp4', fourcc, 10, (traj.shape[2], traj.shape[1]))
        for frame in traj:
            if traj.shape[1] != traj.shape[2]:
                frame = cv2.putText(frame, 'G', (1, 8), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (255, 255, 255), 1, cv2.LINE_AA)
                frame = cv2.putText(frame, 'P', (traj.shape[2] // 2 + 1, 8), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (255, 255, 255), 1, cv2.LINE_AA)
            out.write(frame)

        out.release()

    pred_rgb, gt_rgb = vorticity2rgb(pred_traj), vorticity2rgb(gt_traj)
    cat_rgb = np.concatenate([gt_rgb, pred_rgb], axis=2)
    diff_traj = np.abs(gt_traj - pred_traj)
    diff_rgb = (seaborn.cm.rocket(diff_traj)[..., :3] * 256).astype(np.uint8)

    write_traj(pred_rgb, "pred")
    write_traj(gt_rgb, "gt")
    write_traj(cat_rgb, "both")
    write_traj(diff_rgb, "diff")


data_name = "navier_stokes"
data = np.load("../datasets/data/ns_vorticity_1e-3.npy")  # [50, 64, 64, B=500]
data = data[:41, ..., 100:]  # [41, 64, 64, 100]
data = np.transpose(data, (3, 0, 1, 2))  # [100, 41, 64, 64]
data = torch.tensor(data, dtype=torch.float32)

print(data.shape)

model = load_model("models/nmodes_(6, 6)")

print(count_params(model))

sim_and_render("6modes", model, data[0])