import argparse
import os

import torch
import h5py
from magic.mixture.dataset import MixtureDataset
from magic.mixture.metrics import difference_between_quaternions_tensors
from magic.mixture.models import KinematicMDNv3
from magic.mixture.utils import *
from magic.mixture.mdn import sample as mdn_sample


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Test model for articulated object dataset."
    )
    parser.add_argument("--model-dir", type=str, default="models/")
    parser.add_argument("--model-name", type=str, default="test_lstm")
    parser.add_argument("--test-dir", type=str, default="../data/test/microwave/")
    parser.add_argument("--output-dir", type=str, default="./plots/")
    """
    parser.add_argument(
        "--ntest",
        type=int,
        default=100,
        help="number of test samples (n_object_instants)",
    )
    """
    parser.add_argument(
        "--ndof",
        type=int,
        default=1,
        help="how many degrees of freedom in the object class?",
    )
    parser.add_argument("--batch", type=int, default=40, help="batch size")
    parser.add_argument("--nwork", type=int, default=8, help="num_workers")
    parser.add_argument("--device", type=int, default=0, help="cuda device")

    parser.add_argument("--obj", type=str, default="microwave")
    args = parser.parse_args()

    output_dir = os.path.join(
        os.path.abspath(args.output_dir), args.model_name
    )
    os.makedirs(output_dir, exist_ok=True)

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

    #test_file = os.path.join(args.test_dir, "complete_data_compressed.hdf5")
    test_file = os.path.join(args.test_dir, "complete_data.hdf5")
    with h5py.File(test_file, "r") as f:
        args.ntest = len(f)*16

    print("Testing Model: Abbattematteo et al.")

    bounds = np.load(os.path.join(args.test_dir, "bounds.npy"))
    keep_columns = np.load(os.path.abspath("keep_columns_" + args.obj + ".npy"))
    one_columns = np.load(os.path.abspath("one_columns_" + args.obj + ".npy"))
    testset = MixtureDataset(
        args.ntest,
        args.test_dir,
        n_dof=args.ndof,
        normalize=True,
        bounds=bounds,
        keep_columns=keep_columns,
        one_columns=one_columns,
    )

    best_model = KinematicMDNv3(n_gaussians=20, out_features=testset.labels.shape[1])
    best_model.load_state_dict(
        torch.load(os.path.join(args.model_dir, args.model_name + ".net"))
    )
    best_model.float().to(device)
    best_model.eval()

    testloader = torch.utils.data.DataLoader(
        testset,
        batch_size=args.batch,
        shuffle=False,
        num_workers=args.nwork,
        pin_memory=True,
    )

    all_labels = torch.zeros(args.ntest, testloader.dataset.full_labels.shape[1])
    all_output = torch.zeros(args.ntest, testloader.dataset.full_labels.shape[1])

    with torch.no_grad():
        for i, X in enumerate(testloader):
            depth, labels = X["depth"].to(device), X["label"].to(device)
            pi, sigma, mu = best_model(depth)
            output = mdn_sample(pi, sigma, mu)

            # expand labels and output back to full dataset size
            output = expand_labels(
                testloader.dataset.full_labels,
                output,
                testloader.dataset.keep_columns,
                testloader.dataset.one_columns,
            )

            labels = expand_labels(
                testloader.dataset.full_labels,
                labels,
                testloader.dataset.keep_columns,
                testloader.dataset.one_columns,
            )

            all_labels[i * labels.shape[0] : (i + 1) * labels.shape[0], :] = labels
            all_output[i * labels.shape[0] : (i + 1) * labels.shape[0], :] = output

    # interpret and convert to real.
    ground_truth_params = convert_dict_to_real(
        interpret_labels(all_labels, args.ndof), bounds, args.ndof
    )
    param_dict = convert_dict_to_real(
        interpret_labels(all_output, args.ndof), bounds, args.ndof
    )

    n_imgs_per_obj = 16
    axis_len = 7
    config_len = 1

    real_axis = ground_truth_params["axis"].view(
        -1, args.ndof, n_imgs_per_obj, axis_len
    )  # Axis len 7
    real_net_axis = param_dict["axis"].view(-1, args.ndof, n_imgs_per_obj, axis_len)

    # Calculate Error statistics
    all_ori_err_std, all_ori_err_mean = torch.std_mean(
        difference_between_quaternions_tensors(
            real_axis[:, :, :, 3:7], real_net_axis[:, :, :, 3:7]
        ),
        dim=-1,
    )
    all_ori_err_mean = all_ori_err_mean.squeeze().mean().cpu()
    all_ori_err_std = all_ori_err_std.squeeze().cpu()

    all_dist_err_std, all_dist_err_mean = torch.std_mean(
        torch.norm(real_axis[:, :, :, :3] - real_net_axis[:, :, :, :3], dim=-1), dim=-1
    )
    all_dist_err_mean = all_dist_err_mean.squeeze().mean().cpu()
    all_dist_err_std = all_dist_err_std.squeeze().cpu()

    # Configuration error
    real_configs = ground_truth_params["config"].view(
        -1, args.ndof, n_imgs_per_obj, config_len
    )  # Config len 1
    net_configs = param_dict["config"].view(-1, args.ndof, n_imgs_per_obj, config_len)

    all_q_err_std, all_q_err_mean = torch.std_mean(
        torch.norm(real_configs[:, :, :, :] - net_configs[:, :, :, :], dim=-1), dim=-1
    )
    all_q_err_mean = all_q_err_mean.squeeze().mean().cpu()
    all_q_err_std = all_q_err_std.squeeze().cpu()

    print(
        "For object: {} MAD Losses:\nl_ori: {:.4f}, dist: {:.4f}, config: {:.4f}".format(
            args.obj, all_ori_err_mean.item(), all_dist_err_mean.item(), all_q_err_mean.item()
        )
    )
