from plasma.models.mpi_runner import (
    mpi_make_predictions, mpi_make_predictions_and_evaluate
)
from mpi4py import MPI
from plasma.preprocessor.preprocess import guarantee_preprocessed
from plasma.preprocessor.augment import ByShotAugmentator
from plasma.primitives.shots import ShotList
from plasma.models.loader import Loader
from plasma.conf import conf
from pprint import pprint
'''
#########################################################
This file trains a deep learning model to predict
disruptions on time series data from plasma discharges.

Must run guarantee_preprocessed.py in order for this to work.

Dependencies:
conf.py: configuration of model,training,paths, and data
model_builder.py: logic to construct the ML architecture
data_processing.py: classes to handle data processing

Author: Julian Kates-Harbeck, jkatesharbeck@g.harvard.edu

This work was supported by the DOE CSGF program.
#########################################################
'''

import sys
import random
import numpy as np
import copy
from functools import partial

import matplotlib
matplotlib.use('Agg')

sys.setrecursionlimit(10000)


if conf['model']['shallow']:
    print(
        "Shallow learning using MPI is not supported yet. ",
        "Set conf['model']['shallow'] to False.")
    exit(1)
if conf['data']['normalizer'] == 'minmax':
    from plasma.preprocessor.normalize import MinMaxNormalizer as Normalizer
elif conf['data']['normalizer'] == 'meanvar':
    from plasma.preprocessor.normalize import MeanVarNormalizer as Normalizer
elif conf['data']['normalizer'] == 'var':
    # performs !much better than minmaxnormalizer
    from plasma.preprocessor.normalize import VarNormalizer as Normalizer
elif conf['data']['normalizer'] == 'averagevar':
    # performs !much better than minmaxnormalizer
    from plasma.preprocessor.normalize import (
        AveragingVarNormalizer as Normalizer
    )
else:
    print('unkown normalizer. exiting')
    exit(1)

comm = MPI.COMM_WORLD
task_index = comm.Get_rank()
num_workers = comm.Get_size()
NUM_GPUS = conf['num_gpus']
MY_GPU = task_index % NUM_GPUS


np.random.seed(task_index)
random.seed(task_index)
if task_index == 0:
    pprint(conf)

only_predict = len(sys.argv) > 1
custom_path = None
if only_predict:
    custom_path = sys.argv[1]
print("predicting using path {}".format(custom_path))

assert(only_predict)


#####################################################
#                 NORMALIZATION                     #
#####################################################
# TODO(KGF): identical in at least 3x files in examples/
# make sure preprocessing has been run, and is saved as a file
if task_index == 0:
    # TODO(KGF): check tuple unpack
    (shot_list_train, shot_list_validate,
     shot_list_test) = guarantee_preprocessed(conf)
comm.Barrier()
(shot_list_train, shot_list_validate,
 shot_list_test) = guarantee_preprocessed(conf)


def chunks(l, n):
    """Yield successive n-sized chunks from l."""
    return[l[i:i + n] for i in range(0, len(l), n)]


def hide_signal_data(shot, t=0, sigs_to_hide=None):
    for sig in shot.signals:
        if sigs_to_hide is None or (
                sigs_to_hide is not None and sig in sigs_to_hide):
            shot.signals_dict[sig][t:, :] = shot.signals_dict[sig][t, :]


def create_shot_list_tmp(original_shot, time_points, sigs=None):
    shot_list_tmp = ShotList()
    T = len(original_shot.ttd)
    t_range = np.linspace(0, T-1, time_points, dtype=np.int)
    for t in t_range:
        new_shot = copy.copy(original_shot)
        assert(new_shot.augmentation_fn is None)
        new_shot.augmentation_fn = partial(
            hide_signal_data, t=t, sigs_to_hide=sigs)
        # new_shot.number = original_shot.number
        shot_list_tmp.append(new_shot)
    return shot_list_tmp, t_range


def get_importance_measure(
        original_shot,
        loader,
        custom_path,
        metric,
        time_points=10,
        sigs=None):
    shot_list_tmp, t_range = create_shot_list_tmp(
        original_shot, time_points, sigs)
    y_prime, y_gold, disruptive = mpi_make_predictions(
        conf, shot_list_tmp, loader, custom_path)
    shot_list_tmp.make_light()
    return t_range, get_importance_measure_given_y_prime(
        y_prime, metric), y_prime[-1]


def difference_metric(y_prime, y_prime_orig):
    idx = np.argmax(y_prime_orig)
    return ((np.max(y_prime_orig) - y_prime[idx])
            / (np.max(y_prime_orig) - np.min(y_prime_orig)))


def get_importance_measure_given_y_prime(y_prime, metric):
    differences = [metric(y_prime[i], y_prime[-1])
                   for i in range(len(y_prime))]
    return 1.0-np.array(differences)  # /np.max(differences)


print("normalization", end='')
normalizer = Normalizer(conf)
normalizer.train()
normalizer = ByShotAugmentator(normalizer)
loader = Loader(conf, normalizer)
print("...done")

# if not only_predict:
#     mpi_train(conf,shot_list_train,shot_list_validate,loader)

# load last model for testing
loader.set_inference_mode(True)
use_signals = copy.copy(conf['paths']['use_signals'])
use_signals.append(None)


for shot in shot_list_test:
    # partial(hide_signal_data,t = 0,sigs_to_hide = sigs_to_hide)
    shot.augmentation_fn = None

print("All signals:")
y_prime, y_gold, disruptive, roc, loss = mpi_make_predictions_and_evaluate(
    conf, shot_list_test, loader, custom_path)
print(roc)
print(loss)

# for sigs_to_hide in [[s] for s in use_signals[:-3]] +
# [use_signals[-3:-1]] + [use_signals[-1]]:
for sigs_to_hide in ([[s] for s in use_signals[:-3]]
                     + [[s] for s in use_signals[-3:-1]]
                     + [use_signals[-3:-1]]):
    for shot in shot_list_test:
        shot.augmentation_fn = partial(
            hide_signal_data, t=0, sigs_to_hide=sigs_to_hide)
    print("Hiding: {}".format(sigs_to_hide))
    y_prime, y_gold, disruptive, roc, loss = mpi_make_predictions_and_evaluate(
        conf, shot_list_test, loader, custom_path)
    print(roc)
    print(loss)


if task_index == 0:
    print('finished.')
