import os
import argparse
import math
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader

from mpas import *
from generator import Generator
from discriminator import Discriminator
from graph import load_graph
from utils import *

import pdb

# parse arguments
def parse_args():
    parser = argparse.ArgumentParser(description="Deep Learning Model")

    parser.add_argument("--seed", type=int, default=1,
                        help="random seed (default: 1)")

    parser.add_argument("--root", required=True, type=str,
                        help="root of the dataset")
    parser.add_argument("--resume", type=str, default="",
                        help="path to the latest checkpoint (default: none)")

    parser.add_argument("--dsp", type=int, default=4,
                        help="dimensions of the simulation parameters (default: 4)")
    parser.add_argument("--dspe", type=int, default=512,
                        help="dimensions of the simulation parameters' encode (default: 512)")
    parser.add_argument("--ch", type=int, default=64,
                        help="channel multiplier (default: 64)")

    parser.add_argument("--sn", action="store_true", default=False,
                        help="enable spectral normalization")

    parser.add_argument("--gan-loss", type=str, default="none",
                        help="gan loss (default: none): none, vanilla, wgan")
    parser.add_argument("--gan-loss-weight", type=float, default=1e-2,
                        help="weight of the gan loss (default: 1e-2)")

    parser.add_argument("--batch-size", type=int, default=1,
                        help="batch size for training (default: 1)")

    parser.add_argument("--bwsa", type=float, default=2.5,
                        help="config_bulk_wind_stress_amp (BwsA) (default 2.5)")
    parser.add_argument("--kappa", type=float, default=900.0,
                        help="config_gm_constant_kappa (default 900.0)")
    parser.add_argument("--cvmix", type=float, default=0.625,
                        help="config_cvmix_kpp_criticalbulkrichardsonnumber (default 0.625)")
    parser.add_argument("--mom", type=float, default=200.0,
                        help="config_mom_del2 (default: 200.0)")
    parser.add_argument("--mode", required=True, type=str,
                        help="global, equator, north_atlantic")

    return parser.parse_args()

# the main function
def main(args):
    # log hyperparameters
    print(args)

    # set random seed
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)

    # model
    def add_sn(m):
        for name, c in m.named_children():
            m.add_module(name, add_sn(c))
        if isinstance(m, (nn.Linear, nn.Conv1d)):
            return nn.utils.spectral_norm(m, eps=1e-4)
        else:
            return m

    graphSizes, adjValues, edgeOnes, E_starts, E_ends, \
    avgPoolAsgnIndices, avgPoolAsgnValues, upAsgnIndices, upAsgnValues = \
        load_graph(os.path.join(args.root, "graphM"))
    upAsgnIdx = torch.from_numpy(np.load(os.path.join(args.root, "graphM", "ghtUpAsgnIdx1.npy"))).type(torch.LongTensor).to('cuda:0')
    upAsgnValue = torch.from_numpy(np.load(os.path.join(args.root, "graphM", "ghtUpAsgnValue1.npy")).astype(np.float32)).to('cuda:0')

    g_model = Generator(graphSizes, adjValues, edgeOnes, E_starts, E_ends,
                        upAsgnIndices, upAsgnValues,
                        args.dsp, args.dspe, args.ch)
    if args.sn:
        g_model = add_sn(g_model)

    if args.gan_loss != "none":
        d_model = Discriminator(graphSizes, adjValues, edgeOnes, E_starts, E_ends,
                                avgPoolAsgnIndices, avgPoolAsgnValues,
                                args.dsp, args.dspe, args.ch)

    # load checkpoint
    if args.resume:
        if os.path.isfile(args.resume):
            print("=> loading checkpoint {}".format(args.resume))
        checkpoint = torch.load(args.resume, map_location=torch.device("cuda:0"))
        args.start_epoch = checkpoint["epoch"]
        g_model.load_state_dict(checkpoint["g_model_state_dict"])
        # g_optimizer.load_state_dict(checkpoint["g_optimizer_state_dict"])
        if args.gan_loss != "none":
            # d_model.load_state_dict(checkpoint["d_model_state_dict"])
            # d_optimizer.load_state_dict(checkpoint["d_optimizer_state_dict"])
            d_losses = checkpoint["d_losses"]
            g_losses = checkpoint["g_losses"]
        train_losses = checkpoint["train_losses"]
        test_losses = checkpoint["test_losses"]
        print("=> loaded checkpoint {} (epoch {})"
                .format(args.resume, checkpoint["epoch"]))

    # equator = torch.from_numpy(np.load(os.path.join(args.root, "equator_patch.npy"))).cuda()
    # north_atlantic = np.load(os.path.join(args.root, "north_atlantic.npy"))
    # north_atlantic = torch.from_numpy(north_atlantic.reshape((1, north_atlantic.shape[0], 1))).cuda()
    BwsAMin, BwsAMax = 0.0, 5.0
    kappaMin, kappaMax = 300.0, 1500.0
    cvmixMin, cvmixMax = 0.25, 1.0
    momMin, momMax = 100.0, 300.0

    g_model.train()  # In BatchNorm, we still want the mean and var calculated from the current instance

    step = 2. / 99.
    sparams_sweep = Variable(torch.zeros(1, 4), requires_grad=True)

    grads = np.zeros((4, 100))
    for i in tqdm(range(100)):
        sparams_sweep[0, 0] = -1. + i * step
        sparams_sweep[0, 1] = (args.kappa - kappaMin) / (kappaMax - kappaMin)
        sparams_sweep[0, 2] = (args.cvmix - cvmixMin) / (cvmixMax - cvmixMin)
        sparams_sweep[0, 3] = (args.mom - momMin) / (momMax - momMin)

        fake_data = g_model(sparams_sweep)
        fake_data = batch_spmm(upAsgnIdx, upAsgnValue, upAsgnValue.shape[0], graphSizes[0], fake_data)
        if args.mode == "north_atlantic":
            fake_data = fake_data * north_atlantic
        grad = torch.autograd.grad(fake_data.norm(p=1), sparams_sweep, retain_graph=True)
        grads[0, i] = grad[0][0, 0].item()
        del fake_data
        del grad

    for i in tqdm(range(100)):
        sparams_sweep[0, 0] = (args.bwsa - BwsAMin) / (BwsAMax - BwsAMin)
        sparams_sweep[0, 1] = -1. + i * step
        sparams_sweep[0, 2] = (args.cvmix - cvmixMin) / (cvmixMax - cvmixMin)
        sparams_sweep[0, 3] = (args.mom - momMin) / (momMax - momMin)

        fake_data = g_model(sparams_sweep)
        fake_data = batch_spmm(upAsgnIdx, upAsgnValue, upAsgnValue.shape[0], graphSizes[0], fake_data)
        if args.mode == "north_atlantic":
            fake_data = fake_data * north_atlantic
        grad = torch.autograd.grad(fake_data.norm(p=1), sparams_sweep, retain_graph=True)
        grads[1, i] = grad[0][0, 1].item()
        del fake_data
        del grad

    for i in tqdm(range(100)):
        sparams_sweep[0, 0] = (args.bwsa - BwsAMin) / (BwsAMax - BwsAMin)
        sparams_sweep[0, 1] = (args.kappa - kappaMin) / (kappaMax - kappaMin)
        sparams_sweep[0, 2] = -1. + i * step
        sparams_sweep[0, 3] = (args.mom - momMin) / (momMax - momMin)

        fake_data = g_model(sparams_sweep)
        fake_data = batch_spmm(upAsgnIdx, upAsgnValue, upAsgnValue.shape[0], graphSizes[0], fake_data)
        if args.mode == "north_atlantic":
            fake_data = fake_data * north_atlantic
        grad = torch.autograd.grad(fake_data.norm(p=1), sparams_sweep, retain_graph=True)
        grads[2, i] = grad[0][0, 2].item()
        del fake_data
        del grad

    for i in tqdm(range(100)):
        sparams_sweep[0, 0] = (args.bwsa - BwsAMin) / (BwsAMax - BwsAMin)
        sparams_sweep[0, 1] = (args.kappa - kappaMin) / (kappaMax - kappaMin)
        sparams_sweep[0, 2] = (args.cvmix - cvmixMin) / (cvmixMax - cvmixMin)
        sparams_sweep[0, 3] = -1. + i * step

        fake_data = g_model(sparams_sweep)
        fake_data = batch_spmm(upAsgnIdx, upAsgnValue, upAsgnValue.shape[0], graphSizes[0], fake_data)
        if args.mode == "north_atlantic":
            fake_data = fake_data * north_atlantic
        grad = torch.autograd.grad(fake_data.norm(p=1), sparams_sweep, retain_graph=True)
        grads[3, i] = grad[0][0, 3].item()
        del fake_data
        del grad

    np.save(os.path.join(args.root, "grads.npy"), grads)

if __name__ == "__main__":
  main(parse_args())
