#!/usr/bin/env python3

import os
import sys
import torch
import logging
import argparse
import fnmatch
import numpy as np
import pandas as pd
import seaborn as sns
import os.path as osp
import torch.nn as nn
import matplotlib.pyplot as plt
from torchvision import transforms
from torch.utils.data import DataLoader
import torchvision.datasets as datasets
from torchvision import datasets

sys.path.append(osp.dirname(os.getcwd()))
from models import Network, model_urls, download_model
from hessian import LayerHessian
from utils import get_mean_std, get_subset_dataset, make_deterministic, Config


def parse_args():

    parser = argparse.ArgumentParser(description='layerwise hessian dynamic decomposition')

    model_choices = ['VGG11_bn', 'ResNet18', 'DenseNet3_40', 'MobileNet', 'LeNet']
    dataset_choices = ['CIFAR10', 'CIFAR100', 'MNIST', 'FashionMNIST']

    default_model = model_choices[0]
    default_dataset = dataset_choices[0]
    default_dataset_root = osp.join(osp.dirname(os.getcwd()), 'datasets')
    default_examples_per_class = 10
    default_batch_size = 1024
    default_dpi = 600

    parser.add_argument('--cuda', type=int, help='number of gpu to be used, default=None')
    parser.add_argument('-d', '--dataset', type=str, default=default_dataset, choices=dataset_choices, help='dataset to analyze, default={}'.format(default_dataset))
    parser.add_argument('-m', '--model', type=str, default=default_model, choices=model_choices, help='model for analysis, default={}'.format(default_model))
    # parser.add_argument('-l', '--layer', type=str, required=True, help='layer to be analyzed')
    parser.add_argument('-r', '--run', type=str, required=True, help='run directory to analyze')
    parser.add_argument('-e', '--epoch', type=int, help='epoch number of analyze. for last epoch, this option is not required')
    parser.add_argument('--dataset_root', type=str, default=default_dataset_root, help='dataset root, irrespective of dataset type, default={}'.format(default_dataset_root))
    parser.add_argument('--examples_per_class', type=int, default=default_examples_per_class, help='examples per class, default={}'.format(default_examples_per_class))
    parser.add_argument('-b', '--batch_size', type=int, default=default_batch_size, help='batch size, default={}'.format(default_batch_size))
    parser.add_argument('--dpi', type=int, default=default_dpi, help='dpi for saving images, default={}'.format(default_dpi))
    parser.add_argument('--pdb', action='store_true', help='run with debuggger')

    return parser.parse_args()


if __name__ == '__main__':

    args = parse_args()
    if args.pdb:
        import pdb
        pdb.set_trace()
    make_deterministic(args.cuda)
    sns.set_style('darkgrid')
    device = torch.device('cpu' if args.cuda is None else 'cuda:{}'.format(args.cuda))

    if not osp.exists(args.run):
        os.makedirs(args.run)

    ckpt_dir = osp.join(args.run, 'ckpt')
    images_dir = osp.join(args.run, 'images')
    log_dir = osp.join(args.run, 'logs')

    if not osp.exists(ckpt_dir):
        os.makedirs(ckpt_dir)
    if not osp.exists(images_dir):
        os.makedirs(images_dir)
    if not osp.exists(log_dir):
        os.makedirs(log_dir)

    logging.basicConfig(level=logging.INFO,  format='%(message)s')
    logger = logging.getLogger('layerwise_hessian_dyn_decomp')
    logging_file = osp.join(log_dir, 'layerwise_hessian_dyn_decomp.log')
    with open(logging_file, 'w+') as f:
        pass

    logging_file_handler = logging.FileHandler(logging_file)
    logger.addHandler(logging_file_handler)
    logger.info('Arguments: {}'.format(args))
    if args.epoch is None:
        model_weights_path = osp.join(ckpt_dir, 'model_weights.pth')
    else:
        model_weights_path = osp.join(ckpt_dir, 'model_weights_epochs_{}.pth'.format(args.epoch))
    assert osp.exists(model_weights_path), '{} was not found'.format(model_weights_path)
    assert osp.isfile(model_weights_path), '{} is not a file'.format(model_weights_path)
    if args.dataset in ['MNIST', 'FashionMNIST']:
        input_ch = 1
        padded_im_size = 28
        num_classes = 10
        im_size = 28
        epc_seed = 0
        row = Config(input_ch=input_ch, 
                    padded_im_size=padded_im_size, 
                    num_classes=num_classes,
                    im_size=im_size,
                    epc_seed=epc_seed
                    )
    elif args.dataset in ['CIFAR10', 'CIFAR100']:
        input_ch = 3
        padded_im_size = 32
        if args.dataset == 'CIFAR10':
            num_classes = 10
        elif args.dataset == 'CIFAR100':
            num_classes = 100
        else:
            raise Exception('Should not have reached here')
        im_size = 32
        epc_seed = 0
        row = Config(input_ch=input_ch, 
                    padded_im_size=padded_im_size,
                    num_classes=num_classes,
                    im_size=im_size,
                    epc_seed=epc_seed
                    )
    else:
        raise Exception('this was expected to be an unreachable line')        

    if args.model in ['VGG11_bn', 'ResNet18', 'DenseNet3_40', 'MobileNet', 'LeNet']:
        model = Network().construct(args.model, row)
    else:
        raise Exception('Unknown model argument: {}'.format(args.model))

    state_dict = torch.load(model_weights_path, map_location=lambda storage, loc: storage)
    if 'model' in state_dict.keys():
        state_dict = state_dict['model']
    model.load_state_dict(state_dict, strict=True)
    model = model.to(device)

    model = model.eval()

    mean, std = get_mean_std(args.dataset)
    pad = int((row.padded_im_size-row.im_size)/2)
    transform = transforms.Compose([transforms.Pad(pad),
                                    transforms.ToTensor(),
                                    transforms.Normalize(mean,std)])

    full_dataset = getattr(datasets, args.dataset)
    subset_dataset = get_subset_dataset(full_dataset=full_dataset,
                                        examples_per_class=args.examples_per_class,
                                        epc_seed=row.epc_seed,
                                        root=osp.join(args.dataset_root, args.dataset),
                                        train=True,
                                        transform=transform,
                                        download=True
                                        )

    loader = DataLoader(dataset=subset_dataset,
                        drop_last=False,
                        batch_size=args.batch_size)

    C = row.num_classes
    logger.info('Starting layerwise hessian decomposition...')
    for layer_index, (layer_name, _) in enumerate(model.named_parameters()):
        logger.info('starting decomposition for layer_index: {} layer_name: {}'.format(layer_index, layer_name))
        logger.info('starting hessian eigval density')
        Hess = LayerHessian(crit='CrossEntropyLoss',
                            loader=loader,
                            device=device,
                            model=model,
                            num_classes=C,
                            layer_name=layer_name,
                            hessian_type='Hessian',
                            init_poly_deg=64,
                            poly_deg=128,
                            spectrum_margin=0.05,
                            poly_points=1024,
                            SSI_iters=128
                            )

        Hess_eigval, \
        Hess_eigval_density = Hess.LanczosLoop(denormalize=True)     # the higher the parameter the better the approximation
        logger.info('hessian eigval density computation finished')
        # H = FullHessian(crit='CrossEntropyLoss',
        #                     loader=loader,
        #                     device=device,
        #                     model=model,
        #                     num_classes=C,
        #                     hessian_type='H',
        #                     init_poly_deg=64,
        #                     poly_deg=128,
        #                     spectrum_margin=0.05,
        #                     poly_points=1024,
        #                     SSI_iters=128
        #                     )

        # H_eigval, \
        # H_eigval_density = H.LanczosLoop(denormalize=True)     # the higher the parameter the better the approximation

        logger.info('starting g deigval density computation')
        G = LayerHessian(crit='CrossEntropyLoss',
                            loader=loader,
                            device=device,
                            model=model,
                            num_classes=C,
                            layer_name=layer_name,
                            hessian_type='G',
                            init_poly_deg=64,
                            poly_deg=128,
                            spectrum_margin=0.05,
                            poly_points=1024,
                            SSI_iters=128
                            )

        G_eigval, \
        G_eigval_density = G.LanczosLoop(denormalize=True)     # the higher the parameter the better the approximation
        logger.info('g eigval density computation finished')
        logger.info('creating plots')
        plt.figure()

        plt.semilogy(Hess_eigval, Hess_eigval_density, label='Hessian')
        # plt.semilogy(H_eigval, H_eigval_density, label='H')
        plt.semilogy(G_eigval, G_eigval_density, label='G')
        plt.legend()
        plt.xlabel('Eigenvalue')
        plt.ylabel('Density of Spectrum')
        if args.epoch is not None:
            image_path = osp.join(images_dir, '{}_{}_hessian_dyn_decomp.png'.format(layer_name, args.epoch))
        else:
            image_path = osp.join(images_dir, '{}_hessian_dyn_decomp.png'.format(layer_name))
        plt.savefig(image_path, cpi=args.dpi)
        logger.info('Saved plots for layer: {} ...'.format(layer_name))

    logger.info('layerwise hessian decomposition finished for model:{}, dataset: {}'.format(args.model, args.dataset))