'''
#########################################################
This file containts classes to handle data processing

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

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

from __future__ import print_function
import plasma.global_vars as g
import os
import time
import sys
import abc

import numpy as np
from scipy.signal import exponential, correlate
import pathos.multiprocessing as mp

from plasma.primitives.shots import ShotList, Shot

'''TODO
- incorporate stats, pass machine (perhaps save machine in stats object!)
- incorporate stats, have a dictionary of aggregate stats for every machine.
- check "is_previously_saved" by making sure there is a normalizer for every
machine

'''


#################
# NORMALIZATION #
#################
class Stats(object):
    pass


class Normalizer(object):
    def __init__(self, conf):
        self.num_processed = dict()
        self.num_disruptive = dict()
        self.conf = conf
        self.path = conf['paths']['normalizer_path']
        self.remapper = conf['data']['target'].remapper
        self.machines = set()
        self.inference_mode = False
        self.bound = np.Inf
        if 'norm_stat_range' in self.conf['data']:
            self.bound = self.conf['data']['norm_stat_range']

    @abc.abstractmethod
    def __str__(self):
        pass

    @abc.abstractmethod
    def extract_stats(self, shot):
        pass

    @abc.abstractmethod
    def incorporate_stats(self, stats):
        pass

    @abc.abstractmethod
    def apply(self, shot):
        pass

    @abc.abstractmethod
    def save_stats(self, verbose=False):
        pass

    @abc.abstractmethod
    def load_stats(self, verbose=False):
        pass

    def print_summary(self, action='loaded'):
        g.print_unique(
            '{} normalization data from {} shots ( {} disruptive )'.format(
                action, self.num_processed, self.num_disruptive))

    def set_inference_mode(self, val):
        self.inference_mode = val

    def ensure_machine(self, machine):
        if machine not in self.means:
            self.num_processed[machine] = 0
            self.num_disruptive[machine] = 0
    # Modify the above to change the specifics of the normalization scheme

    def train(self, verbose=False):
        conf = self.conf
        # only use training shots here!! "Don't touch testing shots"
        # + conf['paths']['shot_files_test']
        shot_files = conf['paths']['shot_files']
        shot_files_all = conf['paths']['shot_files_all']
        all_machines = set([file.machine for file in shot_files_all])
        train_machines = set([file.machine for file in shot_files])

        if train_machines >= all_machines:
            shot_files_use = shot_files
        else:
            print('Testing set contains new machine, using testing set ',
                  'to train normalizer for that machine.')
            shot_files_use = shot_files_all

        # shot_list_dir = conf['paths']['shot_list_dir']
        use_shots = max(400, conf['data']['use_shots'])
        return self.train_on_files(shot_files_use, use_shots, all_machines,
                                   verbose=verbose)

    def train_on_files(self, shot_files, use_shots, all_machines,
                       verbose=False):
        conf = self.conf
        all_signals = conf['paths']['all_signals']
        shot_list = ShotList()
        shot_list.load_from_shot_list_files_objects(shot_files, all_signals)
        shot_list_picked = shot_list.random_sublist(use_shots)

        previously_saved, machines_saved = self.previously_saved_stats()
        machines_to_compute = all_machines - machines_saved
        recompute = conf['data']['recompute_normalization']
        if recompute:
            machines_to_compute = all_machines
            previously_saved = False
        if not previously_saved or len(machines_to_compute) > 0:
            if previously_saved:
                self.load_stats(verbose=True)
            print('computing normalization for machines {}'.format(
                machines_to_compute))
            use_cores = max(1, mp.cpu_count()-2)
            pool = mp.Pool(use_cores)
            print('running in parallel on {} processes'.format(
                pool._processes))
            start_time = time.time()

            for (i, stats) in enumerate(pool.imap_unordered(
                    self.train_on_single_shot, shot_list_picked)):
                # for (i,stats) in
                # enumerate(map(self.train_on_single_shot,shot_list_picked)):
                if stats.machine in machines_to_compute:
                    self.incorporate_stats(stats)
                    self.machines.add(stats.machine)
                sys.stdout.write('\r'
                                 + '{}/{}'.format(i, len(shot_list_picked)))
            pool.close()
            pool.join()
            print('\nFinished Training Normalizer on ',
                  '{} files in {} seconds'.format(len(shot_list_picked),
                                                  time.time()-start_time))
            self.save_stats(verbose=True)
        else:
            self.load_stats(verbose=verbose)
        # print representation of trained Normalizer to stdout:
        # Machine, NormalizerName, per-signal normalization stats/params
        if verbose:
            g.print_unique(self)

    def cut_end_of_shot(self, shot):
        cut_shot_ends = self.conf['data']['cut_shot_ends']
        # only cut shots during training
        if not self.inference_mode and cut_shot_ends:
            T_min_warn = self.conf['data']['T_min_warn']
            if shot.ttd.shape[0] - T_min_warn <= max(
                    self.conf['model']['length'], 0):
                print("not cutting shot; length of shot after cutting by ",
                      "T_min_warn would be shorter than RNN length")
                return
            for key in shot.signals_dict:
                shot.signals_dict[key] = shot.signals_dict[key][:-T_min_warn,:]  # noqa
            shot.ttd = shot.ttd[:-T_min_warn]

    # def apply_mask(self,shot):
    #     use_signals = self.conf['paths']['use_signals']
    #     return shot.get_data_arrays(use_signals)

    # def apply_positivity_mask(self,shot):
    #     mask = self.conf['paths']['positivity_mask']
    #     mask = [np.array(subl) for subl in mask]
    #     indices = np.concatenate([indices_sublist[mask[i]] for
    #     i,indices_sublist in enumerate(self.get_indices_list())])
    #     shot.signals[:,indices] = np.clip(shot.signals[:,indices],0,np.Inf)

    def train_on_single_shot(self, shot):
        assert isinstance(shot, Shot), 'should be instance of shot'
        processed_prepath = self.conf['paths']['processed_prepath']
        shot.restore(processed_prepath)
        # print(shot)
        stats = self.extract_stats(shot)
        shot.make_light()
        return stats

    def ensure_save_directory(self):
        prepath = os.path.dirname(self.path)
        if not os.path.exists(prepath):
            os.makedirs(prepath)

    def previously_saved_stats(self):
        if not os.path.isfile(self.path):
            return False, set([])
        else:
            dat = np.load(self.path, encoding="latin1", allow_pickle=True)
            machines = dat['machines'][()]
            ret = all(
                [m in machines for m in self.conf['paths']['all_machines']])
            if not ret:
                print(machines)
                print(self.conf['paths']['all_machines'])
                print('Not all machines present. Recomputing normalizer.')
            return True, set(machines)

    # def get_indices_list(self):
    #     return get_signal_slices(self.conf['paths']['signals_dirs'])


class MeanVarNormalizer(Normalizer):
    def __init__(self, conf):
        Normalizer.__init__(self, conf)
        self.means = dict()
        self.stds = dict()
        self.bound = np.Inf
        if 'norm_stat_range' in self.conf['data']:
            self.bound = self.conf['data']['norm_stat_range']

    def __str__(self):
        s = ''
        for machine in self.means:
            means = np.median(self.means[machine], axis=0)
            stds = np.median(self.stds[machine], axis=0)
            s += 'Machine = {}:\nMean Var Normalizer.\n'.format(machine)
            s += 'means: {}\nstds: {}'.format(means, stds)
        return s

    def extract_stats(self, shot):
        stats = Stats()
        if shot.valid:
            list_of_signals = shot.get_individual_signal_arrays()
            num_signals = len(list_of_signals)
            stats.means = np.reshape(np.array([np.mean(sig) for
                                               sig in list_of_signals]),
                                     (1, num_signals))
            stats.stds = np.reshape(np.array([np.std(sig, dtype=np.float64) for
                                              sig in list_of_signals]),
                                    (1, num_signals))
            stats.is_disruptive = shot.is_disruptive
        else:
            print('Warning: shot {} not valid [omit]'.format(shot.number))
        stats.valid = shot.valid
        stats.machine = shot.machine
        return stats

    def incorporate_stats(self, stats):
        machine = stats.machine
        self.ensure_machine(stats.machine)
        if stats.valid:
            means = stats.means
            stds = stats.stds
            if self.num_processed[machine] == 0:
                self.means[machine] = means
                self.stds[machine] = stds
            else:
                self.means[machine] = np.concatenate(
                    (self.means[machine], means), axis=0)
                self.stds[machine] = np.concatenate(
                    (self.stds[machine], stds), axis=0)
            self.num_processed[machine] = self.num_processed[machine] + 1
            self.num_disruptive[machine] = (
                self.num_disruptive[machine]
                + (1 if stats.is_disruptive else 0))

    def apply(self, shot):
        apply_positivity(shot)
        m = shot.machine
        assert self.means[m] is not None and self.stds[m] is not None, (
            "self.means or self.stds not initialized")
        means = np.median(self.means[m], axis=0)
        stds = np.median(self.stds[m], axis=0)
        for (i, sig) in enumerate(shot.signals):
            if sig.normalize:
                stds_curr = stds[i]
                if stds_curr == 0.0:
                    stds_curr = 1.0
                shot.signals_dict[sig] = (
                    shot.signals_dict[sig] - means[i])/stds_curr
                shot.signals_dict[sig] = np.clip(
                    shot.signals_dict[sig], -self.bound, self.bound)

        shot.ttd = self.remapper(shot.ttd, self.conf['data']['T_warning'])
        self.cut_end_of_shot(shot)
        # self.apply_positivity_mask(shot)
        # self.apply_mask(shot)

    def save_stats(self, verbose=False):
        # standard_deviations = dat['standard_deviations']
        # num_processed = dat['num_processed']
        # num_disruptive = dat['num_disruptive']
        self.ensure_save_directory()
        np.savez(self.path, means=self.means, stds=self.stds,
                 num_processed=self.num_processed,
                 num_disruptive=self.num_disruptive, machines=self.machines)
        if verbose:
            self.print_summary(action='saved')

    def load_stats(self, verbose=False):
        assert self.previously_saved_stats()[0], "stats not saved before"
        dat = np.load(self.path, encoding="latin1", allow_pickle=True)
        self.means = dat['means'][()]
        self.stds = dat['stds'][()]
        self.num_processed = dat['num_processed'][()]
        self.num_disruptive = dat['num_disruptive'][()]
        self.machines = dat['machines'][()]
        # for machine in self.means:
        #     g.print_unique('Machine = {}:'.format(machine))
        if verbose:
            self.print_summary()


class VarNormalizer(MeanVarNormalizer):
    def apply(self, shot):
        apply_positivity(shot)
        assert self.means is not None and self.stds is not None, (
            "self.means or self.stds not initialized")
        m = shot.machine
        stds = np.median(self.stds[m], axis=0)
        for (i, sig) in enumerate(shot.signals):
            if sig.normalize:
                stds_curr = stds[i]
                if stds_curr == 0.0:
                    stds_curr = 1.0
                shot.signals_dict[sig] = (shot.signals_dict[sig])/stds_curr
                shot.signals_dict[sig] = np.clip(
                    shot.signals_dict[sig], -self.bound, self.bound)
        shot.ttd = self.remapper(shot.ttd, self.conf['data']['T_warning'])
        self.cut_end_of_shot(shot)

    def __str__(self):
        s = ''
        for m in self.stds:
            stds = np.median(self.stds[m], axis=0)
            s += 'Machine: {}:\n'.format(m)
            s += 'Var Normalizer.\nstds: {}\n'.format(stds)
        return s


class AveragingVarNormalizer(VarNormalizer):
    def apply(self, shot):
        apply_positivity(shot)
        super(AveragingVarNormalizer, self).apply(shot)
        window_decay = self.conf['data']['window_decay']
        window_size = self.conf['data']['window_size']
        window = exponential(window_size, 0, window_decay, False)
        window /= np.sum(window)
        for (i, sig) in enumerate(shot.signals):
            if sig.normalize:
                shot.signals_dict[sig] = np.apply_along_axis(
                    lambda m: correlate(m, window, 'valid'),
                    axis=0, arr=shot.signals_dict[sig])
                shot.signals_dict[sig] = np.clip(
                    shot.signals_dict[sig], -self.bound, self.bound)
        shot.ttd = shot.ttd[-shot.signals.shape[0]:]

    def __str__(self):
        window_decay = self.conf['data']['window_decay']
        window_size = self.conf['data']['window_size']
        s = ''
        for m in self.stds:
            stds = np.median(self.stds[m], axis=0)
            s += 'Machine: {}:\n'.format(m)
            s += 'Averaging Var Normalizer.\nstds: '
            s += ' {}\nWindow size: {}, Window decay: {}'.format(
                stds, window_size, window_decay)
        return s


class MinMaxNormalizer(Normalizer):
    def __init__(self, conf):
        Normalizer.__init__(self, conf)
        self.minimums = None
        self.maximums = None
        self.bound = np.Inf
        if 'norm_stat_range' in self.conf['data']:
            self.bound = self.conf['data']['norm_stat_range']

    def __str__(self):
        s = ''
        for m in self.minimums:
            s += 'Machine {}:\n.Min Max Normalizer.\n'.format(m,
                                                              self.minimums[m])
            s += 'minimums: {}\nmaximums: {}'.format(self.maximums[m])
        return s

    def extract_stats(self, shot):
        stats = Stats()
        if shot.valid:
            list_of_signals = shot.get_individual_signal_arrays()
            stats.minimums = np.array([np.min(sig) for sig in list_of_signals])
            stats.maximums = np.array([np.max(sig) for sig in list_of_signals])
            stats.is_disruptive = shot.is_disruptive
        else:
            print('Warning: shot {} not valid [omit]'.format(shot.number))
        stats.valid = shot.valid
        stats.machine = shot.machine
        return stats

    def incorporate_stats(self, stats):
        self.ensure_machine(stats.machine)
        if stats.valid:
            m = stats.machine
            minimums = stats.minimums
            maximums = stats.maximums
            if self.num_processed == 0:
                self.minimums[m] = minimums
                self.maximums[m] = maximums
            else:
                self.minimums[m] = (self.num_processed[m]*self.minimums
                                    + minimums)/(self.num_processed[m] + 1.0)
                self.maximums[m] = (self.num_processed[m]*self.maximums
                                    + maximums)/(self.num_processed[m] + 1.0)
            self.num_processed[m] = self.num_processed[m] + 1
            self.num_disruptive[m] = (self.num_disruptive[m]
                                      + (1 if stats.is_disruptive else 0))

    def apply(self, shot):
        apply_positivity(shot)
        assert self.minimums is not None and self.maximums is not None
        m = shot.machine
        curr_range = (self.maximums[m] - self.minimums[m])
        if curr_range == 0.0:
            curr_range = 1.0
        shot.signals = (shot.signals - self.minimums[m])/curr_range
        for (i, sig) in enumerate(shot.signals):
            if sig.normalize:
                shot.signals_dict[sig] = (
                    shot.signals_dict[sig] - self.minimums[m])/(
                        self.maximums[m] - self.minimums[m])
                shot.signals_dict[sig] = np.clip(
                    shot.signals_dict[sig], -self.bound, self.bound)
        shot.ttd = self.remapper(shot.ttd, self.conf['data']['T_warning'])
        self.cut_end_of_shot(shot)
        # self.apply_positivity_mask(shot)
        # self.apply_mask(shot)

    def save_stats(self, verbose=False):
        # standard_deviations = dat['standard_deviations']
        # num_processed = dat['num_processed']
        # num_disruptive = dat['num_disruptive']
        self.ensure_save_directory()
        np.savez(self.path, minimums=self.minimums, maximums=self.maximums,
                 num_processed=self.num_processed,
                 num_disruptive=self.num_disruptive, machines=self.machines)
        if verbose:
            self.print_summary(action='saved')

    def load_stats(self, verbose=False):
        assert self.previously_saved_stats()[0]
        dat = np.load(self.path, encoding="latin1", allow_pickle=True)
        self.minimums = dat['minimums'][()]
        self.maximums = dat['maximums'][()]
        self.num_processed = dat['num_processed'][()]
        self.num_disruptive = dat['num_disruptive'][()]
        self.machines = dat['machines'][()]
        # for machine in self.means:
        #     g.print_unique('Machine {}:'.format(machine))
        if verbose:
            self.print_summary()


def apply_positivity(shot):
    # if shot.signals_dict is None:
    #     print(shot)

    # shot.valid is <class 'numpy.bool_'>; next comparison will always fail

    # if shot.valid is False:
    #     print(shot)
    for (i, sig) in enumerate(shot.signals):
        if hasattr(sig, "is_strictly_positive"):
            # backwards compatibility when this attribute didn't exist
            if sig.is_strictly_positive:
                # print ('Applying positivity constraint to {}
                # signal'.format(sig.description))
                shot.signals_dict[sig] = np.clip(
                    shot.signals_dict[sig], 0, np.inf)
    # KGF: if "TypeError: 'NoneType' object is not subscriptable" occurs at
    # this line during inter-epoch inference, it typically indicates a mismatch
    # of preprocesed_shots/signal_group_X/*.npz, causing a shot to appear in
    # the validation and/or test set with shot.valid==False and/or missing
    # signals_dict even after restore(). Need to trap this earlier.

    # TODO(KGF): why would d3d_all preprocessed shots result in this behavior
    # for d3d_0D training?
