"""
This code analyzes the data generated by the experiment analyze.py

"""
import argparse
import io
import matplotlib.pyplot as plt
import numpy as np
import os
import pickle
import sys
import torch
from tqdm import tqdm
import yaml
from datetime import datetime
from matplotlib import colors as mcolors
import matplotlib.pyplot as plt
from pathlib import Path

sys.path.append("../../")
from c3po.utils.general import generate_filename
from c3po.utils.visualization_utils import scatter_bar_plot
from c3po.datasets.shapenet import CLASS_NAME, CLASS_ID
from c3po.models.certifiability import certifiability

BASE_DIR = Path(__file__).parent
COLORS = dict(mcolors.BASE_COLORS, **mcolors.CSS4_COLORS)

plt.style.use('seaborn-whitegrid')

class LearningObjectsUnpickler(pickle.Unpickler):
    def find_class(self, module, name):
        if module == 'torch.storage' and name == '_load_from_bytes':
            return lambda b: torch.load(io.BytesIO(b), map_location='cpu')
        if module == 'learning_objects.models.certifiability':
            return super().find_class('c3po.models.certifiability', name)
        else:
            return super().find_class(module, name)

def masked_varul_mean(data, mask):
    """
    inputs:
    data    : torch.tensor of shape (B, n)
    mask    : torch.tensor of shape (B, n). dtype=torch.bool

    outputs:
    var     : torch.tensor of shape (B, 2)  :
        var[:, 0] = lower variance
        var[:, 1] = upper variance

    mean    : torch.tensor of shape (B,)

    """
    device_ = data.device
    batch_size = data.shape[0]

    var = torch.zeros(batch_size, 2).to(device_)
    mean = torch.zeros(batch_size).to(device_)

    for batch, (d, m) in enumerate(zip(data, mask)):
        dm = torch.masked_select(d, m)

        dm_mean = dm.mean()
        dm_centered = dm - dm_mean
        dm_centered_up = dm_centered*(dm_centered >= 0)
        dm_centered_lo = dm_centered*(dm_centered < 0)
        len = dm_centered.shape[0]

        dm_var_up = torch.sum(dm_centered_up**2)/(len + 0.001)
        dm_var_lo = torch.sum(dm_centered_lo**2)/(len + 0.001)

        mean[batch] = dm_mean
        var[batch, 0] = dm_var_lo
        var[batch, 1] = dm_var_up

    return var, mean


def varul_mean(data):
    """
    inputs:
    data    : torch.tensor of shape (B, n)

    outputs:
    var     : torch.tensor of shape (B, 2)  :
        var[:, 0] = lower variance
        var[:, 1] = upper variance

    mean    : torch.tensor of shape (B,)

    """

    mean = data.mean(dim=1).unsqueeze(-1)

    data_centered = data - mean
    data_pos = data_centered*(data_centered >= 0)
    data_neg = data_centered*(data_centered < 0)
    len = data_centered.shape[1]

    var_up = torch.sum(data_pos**2, dim=1)/(len + 0.001)
    var_low = torch.sum(data_neg**2, dim=1)/(len + 0.001)

    var_up = var_up.unsqueeze(-1)
    var_low = var_low.unsqueeze(-1)
    var = torch.cat([var_low, var_up], dim=1)

    return var, mean.squeeze(-1)


def certification(data, epsilon, delta, num_iterations=100, full_batch=False):
    device_ = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    certify=certifiability(epsilon=epsilon, delta=delta, radius=0.3)
    ###
    certi_naive = torch.zeros(size=(15, num_iterations), dtype=torch.bool).to(
        device=device_)
    certi_corrector = torch.zeros(size=(15, num_iterations), dtype=torch.bool).to(
        device=device_)

    ###
    sqdist_input_naiveest = data['sqdist_input_naiveest']
    sqdist_input_correctorest = data['sqdist_input_correctorest']
    sqdist_kp_naiveest = data['sqdist_kp_naiveest']
    sqdist_kp_correctorest = data['sqdist_kp_correctorest']
    pc_padding_masks = data['pc_padding_masks']

    for kp_noise_var_i in tqdm(range(len(sqdist_input_naiveest))):
        # print("device_", device_)
        # print("kp_noise_var_i", kp_noise_var_i)
        c_naive = torch.zeros((num_iterations, 1), dtype=torch.bool).to(device=device_)
        c_corrector = torch.zeros((num_iterations, 1), dtype=torch.bool).to(device=device_)
        # if experiments were full batch, just set output of certify.forward_with_distances
        # to c_naive and c_corrector
        for batch_i in range(len(sqdist_input_naiveest[kp_noise_var_i])):
            # print("batch_i", batch_i)

            #len 100 or batch size 100
            sqdist_input_naive = sqdist_input_naiveest[kp_noise_var_i][batch_i]
            sqdist_input_corrector = sqdist_input_correctorest[kp_noise_var_i][batch_i]
            sqdist_kp_naive = sqdist_kp_naiveest[kp_noise_var_i][batch_i]
            sqdist_kp_corrector = sqdist_kp_correctorest[kp_noise_var_i][batch_i]
            pc_padding = pc_padding_masks[kp_noise_var_i][batch_i]
            certi_naive_batch, _ = certify.forward_with_distances(
                sqdist_input_naive[0], sqdist_input_naive[1], sqdist_input_naive[2], sqdist_kp_naive, pc_padding)
            certi_corrector_batch, _ = certify.forward_with_distances(
                sqdist_input_corrector[0], sqdist_input_corrector[1], sqdist_input_naive[2], sqdist_kp_corrector, pc_padding)
            if full_batch: #full batch
                c_naive = certi_naive_batch
                c_corrector = certi_corrector_batch
            else:
                print("certi_naive_batch.shape", certi_naive_batch.shape)
                c_naive[batch_i] = certi_naive_batch
                c_corrector[batch_i] = certi_corrector_batch
        certi_naive[kp_noise_var_i, ...] = c_naive.squeeze(-1)
        certi_corrector[kp_noise_var_i, ...] = c_corrector.squeeze(-1)

    return certi_naive, certi_corrector


def expt_wrapper(class_id, model_id, show_plots=True):

    _dir_name =  str(BASE_DIR) + "/runs/" + class_id + '/' + model_id + '_wchamfer' + '/'
    for file in os.listdir(_dir_name):
        extension = os.path.splitext(file)[1]
        if extension == '.pickle':
            file_path = os.path.join(_dir_name, file)
            file_names = [file_path]

    for name in file_names:

        fp = open(name, 'rb')
        # print(fp)
        parameters, data = LearningObjectsUnpickler(fp).load()
        fp.close()

        # if use_adds_metric:
        for noise_idx in range(len(data['chamfer_pose_naive_to_gt_pose_list'])):
            data['chamfer_pose_naive_to_gt_pose_list'][noise_idx] = np.asarray(
                data['chamfer_pose_naive_to_gt_pose_list'][noise_idx][0].squeeze().to('cpu'))
            data['chamfer_pose_corrected_to_gt_pose_list'][noise_idx] = np.asarray(
                data['chamfer_pose_corrected_to_gt_pose_list'][noise_idx][0].squeeze().to('cpu'))

        # print("-" * 80)

        Rerr_naive = data['rotation_err_naive'].to('cpu')
        Rerr_corrector = data['rotation_err_corrector'].to('cpu')
        terr_naive = data['translation_err_naive'].to('cpu')
        terr_corrector = data['translation_err_corrector'].to('cpu')
        # certi_naive = data['certi_naive'].to('cpu')
        # certi_corrector = data['certi_corrector'].to('cpu')
        # CALCULATE CERTIFIABILITY METRICS DYNAMICALLY, NOT SAVED DURING EXPERIMENTS
        epsilon = .99
        delta = .7
        # fig_save_folder = '/'.join(name.split('/')[:-1] + ['eps' + str(epsilon)[2:]])
        # if not os.path.exists(fig_save_folder):
        #     os.makedirs(fig_save_folder)
        # fig_save_prefix = fig_save_folder + '/' + name.split('/')[-1]
        # print("figures will be saved at filepath:", fig_save_prefix)
        certi_naive, certi_corrector = certification(data, epsilon=epsilon, delta=delta, full_batch=True)
        certi_naive = certi_naive.to('cpu')
        certi_corrector = certi_corrector.to('cpu')

        sqdist_input_naiveest = data['sqdist_input_naiveest']
        sqdist_input_correctorest = data['sqdist_input_correctorest']
        sqdist_kp_naiveest = data['sqdist_kp_naiveest']
        sqdist_kp_correctorest = data['sqdist_kp_correctorest']

        # if use_adds_metric:
        chamfer_pose_naive_to_gt_pose_list = torch.from_numpy(
            np.asarray(data['chamfer_pose_naive_to_gt_pose_list']))
        chamfer_pose_corrected_to_gt_pose_list = torch.from_numpy(
            np.asarray(data['chamfer_pose_corrected_to_gt_pose_list']))

        kp_noise_var_range = parameters['kp_noise_var_range'].to('cpu')
        kp_noise_type = parameters['kp_noise_type']
        kp_noise_fra = parameters['kp_noise_fra']
        class_id = parameters['class_id']
        model_id = parameters['model_id']
        cad_model_name = CLASS_NAME[class_id]

        # Plotting rotation distribution

        # fig = plt.figure()
        #
        # plt = scatter_bar_plot(plt, x=kp_noise_var_range, y=Rerr_naive, label='naive', color='lightgray')
        # plt = scatter_bar_plot(plt, x=kp_noise_var_range, y=Rerr_naive * certi_naive, label='naive + certification',
        #                        color='royalblue')
        # plt.show()
        # plt.close(fig)
        #
        # fig = plt.figure()
        # plt = scatter_bar_plot(plt, x=kp_noise_var_range, y=Rerr_corrector, label='corrector', color='lightgray')
        # plt = scatter_bar_plot(plt, x=kp_noise_var_range, y=Rerr_corrector * certi_corrector,
        #                        label='corrector + certification',
        #                        color='orangered')
        # plt.show()
        # plt.close(fig)

        # if use_adds_metric:
        chamfer_metric_naive_var, chamfer_metric_naive_mean = varul_mean(chamfer_pose_naive_to_gt_pose_list)
        chamfer_metric_corrected_var, chamfer_metric_corrected_mean = varul_mean(
            chamfer_pose_corrected_to_gt_pose_list)
        chamfer_metric_naive_certi_var, chamfer_metric_naive_certi_mean = masked_varul_mean(
            chamfer_pose_naive_to_gt_pose_list, mask=certi_naive)
        chamfer_metric_corrected_certi_var, chamfer_metric_corrected_certi_mean = masked_varul_mean(
            chamfer_pose_corrected_to_gt_pose_list, mask=certi_corrector)

        Rerr_naive_var, Rerr_naive_mean = varul_mean(Rerr_naive)
        Rerr_corrector_var, Rerr_corrector_mean = varul_mean(Rerr_corrector)
        terr_naive_var, terr_naive_mean = varul_mean(terr_naive)
        terr_corrector_var, terr_corrector_mean = varul_mean(terr_corrector)

        Rerr_naive_certi_var, Rerr_naive_certi_mean = masked_varul_mean(Rerr_naive, mask=certi_naive)
        Rerr_corrector_certi_var, Rerr_corrector_certi_mean = masked_varul_mean(Rerr_corrector, mask=certi_corrector)
        terr_naive_certi_var, terr_naive_certi_mean = masked_varul_mean(terr_naive, mask=certi_naive)
        terr_corrector_certi_var, terr_corrector_certi_mean = masked_varul_mean(terr_corrector, mask=certi_corrector)

        fraction_not_certified_naive_var, fraction_not_certified_naive_mean = \
            varul_mean(1 - certi_naive.float())
        fraction_not_certified_corrector_var, fraction_not_certified_corrector_mean = \
            varul_mean(1 - certi_corrector.float())

        Rerr_naive_var = torch.sqrt(Rerr_naive_var).T
        Rerr_corrector_var = torch.sqrt(Rerr_corrector_var).T
        terr_naive_var = torch.sqrt(terr_naive_var).T
        terr_corrector_var = torch.sqrt(terr_corrector_var).T
        Rerr_naive_certi_var = torch.sqrt(Rerr_naive_certi_var).T
        Rerr_corrector_certi_var = torch.sqrt(Rerr_corrector_certi_var).T
        terr_naive_certi_var = torch.sqrt(terr_naive_certi_var).T
        terr_corrector_certi_var = torch.sqrt(terr_corrector_certi_var).T
        fraction_not_certified_corrector_var = torch.sqrt(fraction_not_certified_corrector_var).T
        fraction_not_certified_naive_var = torch.sqrt(fraction_not_certified_naive_var).T

        # if use_adds_metric:
        chamfer_metric_naive_var = torch.sqrt(chamfer_metric_naive_var).T
        chamfer_metric_corrected_var = torch.sqrt(chamfer_metric_corrected_var).T
        chamfer_metric_naive_certi_var = torch.sqrt(chamfer_metric_naive_certi_var).T
        chamfer_metric_corrected_certi_var = torch.sqrt(chamfer_metric_corrected_certi_var).T

        #
        if not os.path.exists('./runs'):
            os.makedirs('./runs')

        # plotting chamfer metric
        object_ = CLASS_NAME[class_id]
        # if use_adds_metric:
        fig = plt.figure()
        plt.errorbar(x=kp_noise_var_range, y=chamfer_metric_naive_mean, yerr=chamfer_metric_naive_var,
                     fmt='-x', color='black', ecolor='gray', elinewidth=1, capsize=3, label='naive')
        plt.errorbar(x=kp_noise_var_range, y=chamfer_metric_naive_certi_mean, yerr=chamfer_metric_naive_certi_var,
                     fmt='--o', color='grey', ecolor='lightgray', elinewidth=3, capsize=0,
                     label='naive + oc = 1')
        plt.errorbar(x=kp_noise_var_range, y=chamfer_metric_corrected_mean, yerr=chamfer_metric_corrected_var,
                     fmt='-x', color='red',
                     ecolor='salmon', elinewidth=1, capsize=3, label='corrector')
        plt.errorbar(x=kp_noise_var_range, y=chamfer_metric_corrected_certi_mean,
                     yerr=chamfer_metric_corrected_certi_var, fmt='--o',
                     color='orangered', ecolor='salmon', elinewidth=3, capsize=0, label='corrector + oc = 1')
        # alternate colors
        # plt.errorbar(x=kp_noise_var_range, y=chamfer_metric_naive_mean, yerr=chamfer_metric_naive_var,
        #              fmt='-', color='red', ecolor='red', elinewidth=1, capsize=3, label='naive')
        # plt.errorbar(x=kp_noise_var_range, y=chamfer_metric_naive_certi_mean, yerr=chamfer_metric_naive_certi_var,
        #              fmt='--o', color='red', ecolor=(1.0,0,0,0.3), elinewidth=3, capsize=3, label='naive + certification')
        # plt.errorbar(x=kp_noise_var_range, y=chamfer_metric_corrected_mean, yerr=chamfer_metric_corrected_var,
        #              fmt='-', color='green', ecolor='green', elinewidth=1, capsize=3, label='corrector')
        # plt.errorbar(x=kp_noise_var_range, y=chamfer_metric_corrected_certi_mean, yerr=chamfer_metric_corrected_certi_var,
        #              fmt='--o', color='green', ecolor=(0,.5,0,0.3), elinewidth=3, capsize=3, label='corrector + certification')
        plt.legend(loc='upper left')
        # title = 'CAD model: ' + cad_model_name + ', Noise: ' + kp_noise_type + f' with fra={kp_noise_fra:.2f}'
        # plt.title(title)
        plt.xlabel('Noise variance parameter $\sigma$')
        plt.ylabel('Normalized ADD-S')
        if show_plots:
            plt.show()
        # timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        # rand_string = generate_filename()
        filename = './runs/' + object_ + '-adds.jpg'
        fig.savefig(filename)
        plt.close(fig)

        # plotting rotation errors
        fig = plt.figure()
        plt.errorbar(x=kp_noise_var_range, y=Rerr_naive_mean, yerr=Rerr_naive_var, fmt='-x', color='black',
                     ecolor='gray',
                     elinewidth=1, capsize=3, label='naive')
        plt.errorbar(x=kp_noise_var_range, y=Rerr_naive_certi_mean, yerr=Rerr_naive_certi_var, fmt='--o', color='grey',
                     ecolor='lightgray', elinewidth=3, capsize=0, label='naive + oc = 1')
        plt.errorbar(x=kp_noise_var_range, y=Rerr_corrector_mean, yerr=Rerr_corrector_var, fmt='-x', color='red',
                     ecolor='salmon', elinewidth=1, capsize=3, label='corrector')
        plt.errorbar(x=kp_noise_var_range, y=Rerr_corrector_certi_mean, yerr=Rerr_corrector_certi_var, fmt='--o',
                     color='orangered', ecolor='salmon', elinewidth=3, capsize=0, label='corrector + oc = 1')

        plt.legend(loc='upper left')
        # title = 'CAD model: ' + cad_model_name + ', Noise: ' + kp_noise_type + f' with fra={kp_noise_fra:.2f}'
        # plt.title(title)
        plt.xlabel('Noise variance parameter $\sigma$')
        plt.ylabel('Rotation error')
        if show_plots:
            plt.show()
        # timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        # rand_string = generate_filename()
        filename = './runs/' + object_ + '-rerr.jpg'
        fig.savefig(filename)
        plt.close(fig)

        # Plotting translation errors
        fig = plt.figure()
        plt.errorbar(x=kp_noise_var_range, y=terr_naive_mean, yerr=terr_naive_var, fmt='-x', color='black',
                     ecolor='gray',
                     elinewidth=1, capsize=3, label='naive')
        plt.errorbar(x=kp_noise_var_range, y=terr_naive_certi_mean, yerr=terr_naive_certi_var, fmt='--o', color='grey',
                     ecolor='lightgray', elinewidth=3, capsize=0, label='naive + oc = 1')
        plt.errorbar(x=kp_noise_var_range, y=terr_corrector_mean, yerr=terr_corrector_var, fmt='-x', color='red',
                     ecolor='salmon', elinewidth=1, capsize=3, label='corrector')
        plt.errorbar(x=kp_noise_var_range, y=terr_corrector_certi_mean, yerr=terr_corrector_certi_var, fmt='--o',
                     color='salmon', ecolor='orangered', elinewidth=3, capsize=0, label='corrector + oc = 1')

        plt.legend(loc='upper left')
        # title = 'CAD model: ' + cad_model_name + ', Noise: ' + kp_noise_type + f' with fra={kp_noise_fra:.2f}'
        # plt.title(title)
        plt.xlabel('Noise variance parameter $\sigma$')
        plt.ylabel('Translation error')
        if show_plots:
            plt.show()
        # timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        # rand_string = generate_filename()
        filename = './runs/' + object_ + '-terr.jpg'
        fig.savefig(filename)
        plt.close(fig)

        # Plotting fraction not certified
        fig = plt.figure()
        plt.bar(x=kp_noise_var_range - 0.01, width=0.02, height=fraction_not_certified_naive_mean,
                color='grey', align='center', label='naive')
        # plt.errorbar(x=kp_noise_var_range-0.01, y=fraction_not_certified_naive_mean,
        #              yerr=fraction_not_certified_naive_var,
        #              fmt='o', color='black', ecolor='darkgray', elinewidth=1, capsize=3)
        plt.bar(x=kp_noise_var_range + 0.01, width=0.02, height=fraction_not_certified_corrector_mean,
                color='salmon', align='center', label='corrector')
        # plt.errorbar(x=kp_noise_var_range+0.01, y=fraction_not_certified_corrector_mean,
        #              yerr=fraction_not_certified_corrector_var,
        #              fmt='o', color='red', ecolor='orangered', elinewidth=1, capsize=3)
        plt.legend(loc='upper left')
        plt.xlabel('Noise variance parameter $\sigma$')
        plt.ylabel('Fraction not obs. correct (oc = 0)')
        # title = 'CAD model: ' + cad_model_name + ', Noise: ' + kp_noise_type + f' with fra={kp_noise_fra:.2f}'
        # plt.title(title)
        if show_plots:
            plt.show()
        # timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        # rand_string = generate_filename()
        filename = './runs/' + object_ + '-cert.jpg'
        fig.savefig(filename)
        plt.close(fig)


def plot(class_name, show_plots=True):

    class_id = CLASS_ID[class_name]

    stream = open("class_model_ids.yml", "r")
    model_class_ids = yaml.load(stream=stream, Loader=yaml.Loader)
    if class_name not in model_class_ids:
        raise Exception('Invalid class_name')
    else:
        model_id = model_class_ids[class_name]

    expt_wrapper(class_id=class_id, model_id=model_id, show_plots=show_plots)


if __name__ == '__main__':

    """
    usage: 
    >> python plot.py "chair"
    >> python plot.py "car"
    """

    parser = argparse.ArgumentParser()
    parser.add_argument("class_name", help="specify the ShapeNet class name.", type=str)

    args = parser.parse_args()
    class_name = args.class_name
    class_id = CLASS_ID[class_name]

    stream = open("class_model_ids.yml", "r")
    model_class_ids = yaml.load(stream=stream, Loader=yaml.Loader)
    if class_name not in model_class_ids:
        raise Exception('Invalid class_name')
    else:
        model_id = model_class_ids[class_name]

    expt_wrapper(class_id=class_id, model_id=model_id, use_adds_metric=True)



