import os
import argparse
import torch
import numpy as np
from model import LitDiffusionModel
from finetune_model import FineTuneLitDiffusionModel
from eval_utils import *
from chamferdist import ChamferDistance

parser = argparse.ArgumentParser()

model_args = parser.add_argument_group('model')
model_args.add_argument('--ckpt_path', type=str, help='Path to the model checkpoint', required=True)
model_args.add_argument('--hparams_path', type=str, help='Path to model hyperparameters', required=True)

data_args = parser.add_argument_group('data')
data_args.add_argument('--train_data_path', type=str, default='./data/gaussian_3D_test.npy', help='Path to training data numpy file')
data_args.add_argument('--test_data_path', type=str, default='./data/gaussian_3D_test.npy', help='Path to test data numpy file')

eval_args = parser.add_argument_group('evaluation')
eval_args.add_argument('--savedir', type=str, default='./results/', help='Path to directory for saving evaluation results')
eval_args.add_argument('--n_runs', type=int, default=3, help='Number of runs of evaluation')
eval_args.add_argument('--eval_emd', action='store_true', help='Calculate Earth Mover\'s Distance')
eval_args.add_argument('--eval_emd_samples', type=int, default=128, help='Number of random samples to be sampled for calculating EMD')
eval_args.add_argument('--eval_nll', action='store_true', help='Calculate negative log likelihood')
eval_args.add_argument('--eval_chamfer', action='store_true', help='Calculate Chamfer Distance (using `chamferdist`)')
eval_args.add_argument('--vis_overlay', action='store_true', help='Overlays predicted distribution on top of ground truth')
eval_args.add_argument('--vis_diffusion', action='store_true', help='Shows the evolution of samples through the diffusion process via an animation')
eval_args.add_argument('--vis_track_max', action='store_true', help='Track the point with highest Z in the predicted distribution in diffusion animation')
eval_args.add_argument('--vis_track_min', action='store_true', help='Track the point with lowest Z in the predicted distribution in diffusion animation')
eval_args.add_argument('--vis_smoothed_end', action='store_true', help='Smooths the end of animation by repeating the last frame of animation')
args = parser.parse_args()

# litmodel = LitDiffusionModel.load_from_checkpoint(
#    args.ckpt_path, 
#    hparams_file=args.hparams_path
# )
litmodel = FineTuneLitDiffusionModel.load_from_checkpoint(
    args.ckpt_path,
    hparams_file=args.hparams_path
)
litmodel.eval()

traindata = np.load(args.train_data_path)
testdata = np.load(args.test_data_path)

# mean = np.mean(traindata, axis=0)
# std = np.std(traindata, axis=0)
# mean[0] = 0
# mean[1] = 0
# std[0] = 1
# std[1] = 1
# print(f'Mean = {mean}')
# print(f'Std dev = {std}')

# traindata = (traindata-mean)/std
# testdata = (testdata-mean)/std

traindata = torch.from_numpy(traindata)
testdata = torch.from_numpy(testdata)

os.makedirs(args.savedir, exist_ok=True)

for i_run in range(args.n_runs):
    print(64*'-')
    print(f'Evaluation run {i_run+1}/{args.n_runs}')
    print(64*'-')
    with torch.no_grad():
        gendata, intermediate = litmodel.sample(testdata.size(0), progress=True, return_intermediate=True)
    

    #temp_data = np.concatenate((testdata.numpy(), gendata.numpy()), axis=0)
    #print(f'gendata.shape = {gendata.shape}')
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(gendata[:, 0], gendata[:, 1], gendata[:, 2], c=gendata[:, 2], cmap=cm.spring, alpha=0.1)
    #ax.scatter(testdata[:, 0], testdata[:, 1], testdata[:, 2], c=testdata[:, 2], alpha=0.1)
    ax.set_xlim(-3, 3)
    ax.set_ylim(-3, 3)
    #ax.set_zlim(0, 1)
    plt.show()

    with open(f'{args.savedir}/{i_run:02d}_log.txt', 'w') as f:
        f.write('Results\n')
    # EMD
    if args.eval_emd:
        idx = np.random.choice(np.arange(gendata.size(0)), size=args.eval_emd_samples, replace=False)
        test_emd = get_emd(testdata[idx].numpy(), gendata[idx].numpy())
        train_emd = get_emd(traindata[idx].numpy(), gendata[idx].numpy())
        print(f'test_emd: {test_emd}')
        print(f'train_emd: {train_emd}')
        with open(f'{args.savedir}/{i_run:02d}_log.txt', 'a') as f:
            f.write(f'test_emd: {test_emd}\n')
            f.write(f'train_emd: {train_emd}\n')

    # NLL
    if args.eval_nll:
        test_nll = get_nll(testdata, gendata).item()
        train_nll = get_nll(traindata, gendata).item()
        print(f'test_nll: {test_nll}')
        print(f'train_nll: {train_nll}')
        with open(f'{args.savedir}/{i_run:02d}_log.txt', 'a') as f:
            f.write(f'test_nll: {test_nll}\n')
            f.write(f'train_nll: {train_nll}\n')

    # Chamfer
    if args.eval_chamfer:
        cd = ChamferDistance()
        test_chamfer = cd(
            testdata.unsqueeze(0).float(), 
            gendata.unsqueeze(0).float()
        ).item()
        train_chamfer = cd(
            traindata.unsqueeze(0).float(),
            gendata.unsqueeze(0).float()
        ).item()
        print(f'test_chamfer: {test_chamfer}')
        print(f'train_chamfer: {train_chamfer}')
        with open(f'{args.savedir}/{i_run:02d}_log.txt', 'a') as f:
            f.write(f'test_chamfer: {test_chamfer}\n')
            f.write(f'train_chamfer: {train_chamfer}\n')
    
    # Visualize overlay
    if args.vis_overlay:
        # Only performed with test since it allows for sparser and better visualizations
        print('Visualizing predicted distribution by overlaying it on top of ground truth distribution')
        plot_final_distributions(
                f'{args.savedir}/{i_run:02d}_overlayvis.pdf', 
            testdata, gendata
        )
        print(f'Output: {args.savedir}/{i_run:02d}_overlayvis.pdf')
    
    # Visualize diffusion
    if args.vis_diffusion:
        print('Visualizing evolution of samples through the diffusion process')
        print('WARN: this will take a long time depending on the number of diffusion steps')
        fname = f'{i_run:02d}.diffusionvis.track_max={args.vis_track_max}.track_min={args.vis_track_min}.smoothed_end={args.vis_smoothed_end}.gif'
        animate_intermediate_distributions(
            f'{args.savedir}/{fname}', 
            testdata, intermediate, 
            track_max=args.vis_track_max, 
            track_min=args.vis_track_min, 
            smoothed_end=args.vis_smoothed_end
        )
        print(f'Output: {args.savedir}/{fname}')
    print(64*'-')
