import tables
import blechpy
import os
import numpy as np
import pandas as pd
import itertools as it
from blechpy.analysis import poissonHMM as phmm
from blechpy.utils import write_tools as wt
from blechpy.plotting import hmm_plot as hmmplt
import pickle
import pylab as plt
import multiprocessing as mp

HMM_PARAMS = {'time_window': [-1000, 2000], 'convergence_thresh': 1e-4,
              'max_iter': 2000, 'n_repeats': 3, 'unit_type': 'single',
              'bin_size': 0.001}

class InfoParticle(tables.IsDescription):
    taste = tables.StringCol(20)
    channel = tables.Int16Col()
    n_cells = tables.Int32Col()
    n_trials = tables.Int32Col()
    n_states = tables.Int32Col()
    dt = tables.Float64Col()
    BIC = tables.Float64Col()
    converged = tables.BoolCol()
    threshold = tables.Float64Col()
    cost = tables.Float64Col()

class MultiInfoParticle(tables.IsDescription):
    #   HMM_ID, taste, din_channel, n_cells, time_start, time_end, thresh,
    #   unit_type, n_repeats, dt, n_states, n_iters, BIC, cost, converged
    hmm_id = tables.Int16Col()
    taste = tables.StringCol(20)
    channel = tables.Int16Col()
    n_cells = tables.Int32Col()
    unit_type = tables.StringCol(15)
    n_trials = tables.Int32Col()
    dt = tables.Float64Col()
    max_iter = tables.Int32Col()
    threshold = tables.Float64Col()
    time_start = tables.Int32Col()
    time_end = tables.Int32Col()
    n_repeats = tables.Int16Col()
    n_states = tables.Int32Col()
    n_iterations = tables.Int32Col()
    BIC = tables.Float64Col()
    cost = tables.Float64Col()
    converged = tables.BoolCol()
    fitted = tables.BoolCol()

class HmmAnalysis(object):
    def __init__(self, dat, n_states, save_dir=None, params=None):
        if isinstance(dat, str):
            dat = blechpy.load_dataset(dat)
            if dat is None:
                raise FileNotFoundError('No dataset.p file found given directory')

        if save_dir is None:
            save_dir = os.path.join(dat.root_dir,
                                    '%s_analysis' % dat.data_name)

        self._dataset = dat
        self.root_dir = dat.root_dir
        self.save_dir = save_dir
        self.n_states = n_states

        if not os.path.isdir(save_dir):
            os.makedirs(save_dir)

        plot_dir = os.path.join(save_dir, '%i_states' % n_states, 'plots')
        if not os.path.isdir(plot_dir):
            os.makedirs(plot_dir)

        self._plot_dir = plot_dir

        self._files = {'hmm_data': os.path.join(save_dir, 'hmm_data.hdf5'),
                       'params' : os.path.join(save_dir, 'hmm_params.json')}
        file_check = self._file_check()
        self.update_params(params)

        dig_in_map = dat.dig_in_mapping.query('spike_array == True and exclude == False')
        self._dig_ins = dig_in_map.set_index('name')
        self._fitted_models = {}
        self._setup_hdf5()

    def _file_check(self):
        out = {}
        for k, v in self._files.items():
            if os.path.isfile(v):
                out[k] = True
            else:
                out[k] = False

        return out

    def update_params(self, params=None):
        file_check = self._file_check()
        if file_check['params']:
            old_params = wt.read_dict_from_json(self._files['params'])
        else:
            old_params = HMM_PARAMS.copy()

        if params:
            for k, v in params.items():
                if isinstance(v, dict):
                    old_params.update(v)
                else:
                    old_params[k] = v

        self.params = old_params
        wt.write_dict_to_json(self.params, self._files['params'])

    def _setup_hdf5(self):
        dig_ins = self._dig_ins
        h5_file = self._files['hmm_data']
        n_states = self.n_states
        state_str = '%i_states' % n_states

        with tables.open_file(h5_file, 'a') as hf5:
            # Taste -> PI, A, B, BIC, state_sequences, nStates, nCells, dt
            if not 'data_overview' in hf5.root:
                # Contains taste, channel, n_cells, n_trials, n_states, dt, BIC
                table = hf5.create_table('/', 'data_overview', InfoParticle,
                                         'Basic info for each digital_input')
                table.flush()

            if not state_str in hf5.root:
                hf5.create_group('/', state_str, '%i State HMM Solutions' % n_states)

            for taste in dig_ins.index:
                t = taste.replace(' ', '_')
                if t not in hf5.root[state_str]:
                    # Contains initial_distribution, transition, emission, state_sequences
                    # time vector
                    hf5.create_group('/%s' % state_str, t, '%s HMM Data' % taste)

            hf5.flush()

    def get_unit_names(self):
        dat = self._dataset
        unit_type = self.params['unit_type']
        unit_table = dat.get_unit_table()

        if unit_type == 'single':
            ut = unit_table.query('single_unit == True')
            units = ut['unit_name'].tolist()
        elif unit_type == 'all':
            units = unit_table['unit_name'].tolist()
        elif unit_type == 'pyramidal':
            ut = unit_table.query('single_unit == True and regular_spiking == True')
            units = ut['unit_name'].tolist()
        elif unit_type == 'interneuron':
            ut = unit_table.query('single_unit == True and fast_spiking == True')
            units = ut['unit_name'].tolist()
        else:
            raise ValueError('unit_type must be single, all, pyramidal or interneuron')

        return units

    def get_spike_data(self, taste):
        bin_size = self.params['bin_size']
        dig_ins = self._dig_ins
        channel = dig_ins.loc[taste]['channel']
        rec_dir = self.root_dir
        window = self.params['time_window']
        units = self.get_unit_names()
        time, spike_array = blechpy.dio.h5io.get_spike_data(rec_dir, units, channel)
        idx = np.where((time >= window[0]) & (time < window[1]))[0]
        time = time[idx]
        spike_array = spike_array[:, :, idx].astype('int32')
        dt = np.unique(np.diff(time))[0] / 1000
        if bin_size > dt:
            spike_array, dt, time = rebin_spike_array(spike_array, dt, time, bin_size)

        return spike_array, dt, time

    def run(self, overwrite=False, parallel=True):
        params = self.params
        rec_dir = self.root_dir
        n_states = self.n_states
        window = params['time_window']
        unit_type = params['unit_type']
        n_repeats = params['n_repeats']
        max_iter = params['max_iter']
        thresh = params['convergence_thresh']
        dig_ins = self._dig_ins
        h5_file = self._files['hmm_data']
        plot_dir = self._plot_dir

        for taste, row in dig_ins.iterrows():
            # Check if trained HMM already exists
            # Add overwrite flag
            if HMM_exists(h5_file, taste, n_states) and not overwrite:
                print('Detected existing fitted model for %i states - %s.'
                      '\nLoading...' % (n_states, taste))
                self._fitted_models[taste] = self.load_hmm_from_h5(taste)
                continue

            channel = row['channel']
            spike_array, dt, time = self.get_spike_data(taste)
            n_trials, n_cells, n_steps = spike_array.shape

            print('='*80 + '\nTraining HMMs for %s\n' % taste + '='*80)
            convergence = []
            BIC = []
            hmms = []

            def update(ans):
                hmms.append(ans)

            def error_call(e):
                print(e)
                raise ValueError(e)

            if parallel:
                pool = mp.get_context('spawn').Pool(mp.cpu_count()-1)
                for i in range(n_repeats):
                    pool.apply_async(phmm.fit_hmm_mp,
                                     (n_states, spike_array, dt, max_iter, thresh),
                                     callback=update, error_callback=error_call)

                pool.close()
                pool.join()
            else:
                tmp_hmm = phmm.fit_hmm_mp(n_states, spike_array, dt,
                                          max_iter, thresh)
                update(tmp_hmm)

            print('Fitting Complete!')
            print('Picking Best HMM and Saving...')
            for hmm in hmms:
                convergence.append(hmm.isConverged(thresh))
                tmp_bic, _ = hmm.get_BIC()
                BIC.append(tmp_bic)

            if all(convergence):
                print('All HMMs converged...picking best...')
                idx = np.argmin(BIC)
                bestHMM = hmms[idx]
                bestBIC = BIC[idx]
            elif any(convergence):
                print('Some HMMs converged...picking best...')
                bestHMM = None
                while bestHMM is None:
                    idx = np.argmin(BIC)
                    if hmms[idx].isConverged(thresh):
                        bestHMM = hmms[idx]
                        bestBIC = BIC[idx]
                    else:
                        hmms.pop(idx)
                        BIC.pop(idx)

            else:
                print('No HMMs converged...picking best...')
                idx = np.argmin(BIC)
                minHMM = None
                minBIC = None
                for h in hmms:
                    mats, iteration, BICs = h.find_best_in_history()
                    if minBIC is None:
                        minBIC = np.min(BICs)
                        h.set_matrices(mats)
                        minHMM = h
                    elif np.min(BICs) < minBIC:
                        minBIC = np.min(BICs)
                        h.set_matrices(mats)
                        minHMM = h

                bestHMM = minHMM
                bestBIC = minBIC

            # Save HMM to h5
            self._fitted_models[taste] = bestHMM
            self.add_hmm_to_h5(taste, bestHMM, overwrite=overwrite)
            # make plots
            save_dir = os.path.join(plot_dir, taste)
            if not os.path.isdir(save_dir):
                os.makedirs(save_dir)

            hmmplt.plot_hmm_figures(hmm, self.params['time_window'], save_dir=save_dir)

    def add_hmm_to_h5(self, taste, hmm, overwrite=False):
        print('Writing best HMM for %s to hdf5...' % taste)
        plot_dir = self._plot_dir
        thresh = self.params['convergence_thresh']
        converged = hmm.isConverged(thresh)
        n_states = hmm.n_states
        n_trials, n_cells, n_steps = hmm.data.shape
        channel = self._dig_ins.loc[taste]['channel']
        state_str = '%i_states' % n_states

        dt = hmm.dt
        PI = hmm.initial_distribution
        A = hmm.transition
        B = hmm.emission
        BIC, bestPaths = hmm.get_BIC()
        cost = hmm.cost

        window = self.params['time_window']
        time = np.arange(window[0], window[1], dt*1000)

        h5_file = self._files['hmm_data']
        exists = HMM_exists(h5_file, taste, n_states)
        if exists and not overwrite:
            return

        with tables.open_file(h5_file, 'a') as hf5:
            tmp_taste = taste.replace(' ', '_')

            # Add row to data overview
            dat = hf5.root.data_overview
            if exists:
                hf5.remove_node('/'+state_str, tmp_taste, recursive=True)
                idx = np.where((dat['taste'] == taste) &
                               (dat['n_states'] == n_states))[0]
                if len(idx) == 1:
                    dat.remove_rows(idx[0], idx[0]+1)
                elif len(idx) > 1:
                    raise ValueError('Multiple entries for %i states: %s' %
                                     (n_states, taste))

                dat.flush()
                hf5.flush()

            if state_str not in hf5.root:
                hf5.create_group('/', state_str, '%i State HMM Solutions' % n_states)

            if tmp_taste not in hf5.root[state_str]:
                hf5.create_group('/' + state_str, tmp_taste, '%s HMM Data' % taste)

            row = dat.row
            row['taste'] = taste
            row['channel'] = channel
            row['n_cells'] = n_cells
            row['n_trials'] = n_trials
            row['n_states'] = n_states
            row['dt'] = dt
            row['BIC'] = BIC
            row['converged'] = converged
            row['threshold'] = thresh
            row['cost'] = cost
            row.append()
            dat.flush()

            # Add data arrays
            array_loc = '/%s/%s' % (state_str, tmp_taste)
            hf5.create_array(array_loc, 'initial_distribution', PI)
            hf5.create_array(array_loc, 'transition', A)
            hf5.create_array(array_loc, 'emission', B)
            hf5.create_array(array_loc, 'state_sequences', bestPaths)
            hf5.flush()

    def load_hmm_from_h5(self, taste):
        n_states = self.n_states
        h5_file = self._files['hmm_data']
        spike_array, dt, _ = self.get_spike_data(taste)
        hmm = load_hmm_from_h5(h5_file, taste, n_states, spike_array, dt)
        return hmm

    def get_data_overview(self):
        h5_file = self._files['hmm_data']
        return get_data_overview(h5_file)

    def HMM_exists(taste):
        return HMM_exists(self._files['hmm_data'], taste, self.n_states)

def load_hmm_from_h5(h5_file, taste, n_states, spike_array, dt):
    state_str = '%i_states' % n_states
    taste_str = taste.replace(' ', '_')
    if not HMM_exists(h5_file, taste, n_states):
        return None

    hmm = phmm.PoissonHMM(n_states, spike_array, dt)
    with tables.open_file(h5_file, 'r') as hf5:
        taste_node = hf5.root[state_str][taste_str]
        PI = taste_node['initial_distribution'][:]
        A  = taste_node['transition'][:]
        B  = taste_node['emission'][:]
        hmm.set_matrices({'PI': PI, 'A': A, 'B': B})
        hmm.fitted = True

    return hmm

def get_data_overview(h5_file):
    with tables.open_file(h5_file, 'r') as hf5:
        return hf5.root.data_overview[:]

def HMM_exists(h5_file, taste, n_states):
    state_str = '%i_states' % n_states
    taste_str = taste.replace(' ', '_')
    dat = get_data_overview(h5_file)
    idx = np.where((dat['taste'].astype('str') == taste) &
                   (dat['n_states'] == n_states))[0]
    if len(idx) == 0:
        return False

    # Check if actual data exists
    with tables.open_file(h5_file, 'r') as hf5:
        if not state_str in hf5.root:
            return False

        if not taste_str in hf5.root[state_str]:
            return False

        if not any([x in hf5.root[state_str][taste_str]
                    for x in ['initial_distribution', 'transition',
                              'emission', 'state_sequences']]):
            return False

    return True

def rebin_spike_array(spikes, dt, time, new_dt):
    if spikes.ndim == 2:
        spikes = np.expand_dims(spikes,0)

    n_trials, n_cells, n_steps = spikes.shape
    n_bins = int(new_dt/dt)
    new_time = np.arange(time[0], time[-1], new_dt)
    new_spikes = np.zeros((n_trials, n_cells, len(new_time)))
    for i, w in enumerate(new_time):
        idx = np.where((time >= w) & (time < w+new_dt))[0]
        new_spikes[:,:,i] = np.sum(spikes[:,:,idx], axis=-1)

    return new_spikes, new_time

# Cut data to 0-2000 sec, fit hmm with 2 states and 3 states and try a different bin size
# Bin sizes: 0.001 and 0.10
# Get program to make all hmms and put into an array first and then fit them all
HMM_PARAMS = {'unit_type': 'single', 'dt': 0.001, 'threshold': 1e-4, 'max_iter': 1000,
              'time_start': 0, 'time_end': 2000, 'n_repeats': 3, 'n_states': 3}

class HmmHandler(object):
    def __init__(self, dat, params=None, save_dir=None):
        '''Takes a blechpy dataset object and fits HMMs for each tastant

        Parameters
        ----------
        dat: blechpy.dataset
        params: dict or list of dicts
            each dict must have fields:
                time_window: list of int, time window to cut around stimuli in ms
                convergence_thresh: float
                max_iter: int
                n_repeats: int
                unit_type: str, {'single', 'pyramidal', 'interneuron', 'all'}
                bin_size: time bin for spike array when fitting in seconds
                n_states: predicted number of states to fit
        '''
        if isinstance(params, dict):
            params = [params]

        if isinstance(dat, str):
            dat = blechpy.load_dataset(dat)
            if dat is None:
                raise FileNotFoundError('No dataset.p file found given directory')

        if save_dir is None:
            save_dir = os.path.join(dat.root_dir,
                                    '%s_analysis' % dat.data_name)

        self._dataset = dat
        self.root_dir = dat.root_dir
        self.save_dir = save_dir
        self.h5_file = os.path.join(save_dir, '%s_HMM_Analysis.hdf5' % dat.data_name)
        dim = dat.dig_in_mapping.query('exclude==False')
        tastes = dim['name'].tolist()
        if params is None:
            # Load params and fitted models
            self.load_data()
        else:
            self.init_params(params)

        self.params = params

        if not os.path.isdir(save_dir):
            os.makedirs(save_dir)

        self.plot_dir = os.path.join(save_dir, 'HMM_Plots')
        if not os.path.isdir(self.plot_dir):
            os.makedirs(self.plot_dir)

        self._setup_hdf5()

    def init_params(self, params):
        dat = self._dataset
        dim = dat.dig_in_mapping.query('exclude == False')
        tastes = dim['name'].tolist()
        dim = dim.set_index('name')
        if not hasattr(dat, 'dig_in_trials'):
            dat.create_trial_list()

        trials = dat.dig_in_trials
        data_params = []
        fit_objs = []
        for i, X in enumerate(it.product(params,tastes)):
            p = X[0].copy()
            t = X[1]
            p['hmm_id'] = i
            p['taste'] = t
            p['channel'] = dim.loc[t, 'channel']
            unit_names = query_units(dat, p['unit_type'])
            p['n_cells'] = len(unit_names)
            p['n_trials'] = len(trials.query('name == @t'))

            data_params.append(p)
            # Make fit object for each repeat
            # During fitting compare HMM as ones with the same ID are returned
            for i in range(p['n_repeats']):
                hmmFit = HMMFit(dat.root_dir, p)
                fit_objs.append(hmmFit)

        self._fit_objects = fit_objs
        self._data_params = data_params
        self._fitted_models = dict.fromkeys([x['hmm_id'] for x in data_params])

    def load_params(self):
        h5_file = self.h5_file
        if not os.path.isfile(h5_file):
            raise ValueError('No params to load')

        rec_dir = self._dataset.root_dir
        params = []
        fit_objs = []
        fitted_models = {}
        with tables.open_file(h5_file, 'r') as hf5:
            table = hf5.root.data_overview
            for row in table[:]:
                p = {}
                for k in NEW_HMM_PARAMS.keys():
                    p[k] = row[k]

                params.append(p)
                for i in range(p['n_repeats']):
                    hmmFit = HMMFit(rec_dir, p)
                    fit_objs.append(hmmFit)

        for p in params:
            hmm_id = p['hmm_id']
            fitted_models[hmm_id] = read_hmm_from_hdf5(h5_file, hmm_id)

        self._data_params = params
        self._fit_objects = fit_objs
        self._fitted_models = fitted_models

    def write_overview_to_hdf5(self):
        params = self._data_params
        h5_file = self.h5_file
        if hasattr(self, '_fitted_models'):
            models = self._fitted_models
        else:
            models = dict.fromkeys([x['hmm_id']
                                    for x in data_params])
            self._fitted_models = models

        if not os.path.isfile(h5_file):
            self._setup_hdf5()

        print('Writing data overview table to hdf5...')
        with tables.open_file(h5_file, 'a') as hf5:
            table = hf5.root.data_overview
            # Clear old table
            table.remove_rows(start=0)

            # Add new rows
            for p in params:
                row = table.row
                for k, v in p.items():
                    row[k] = v

                if models[p['hmm_id']] is not None:
                   hmm = models[p['hmm_id']]
                   row['n_iterations'] =  hmm.iterations
                   row['BIC'] = hmm.BIC
                   row['cost'] = hmm.cost
                   row['converged'] = hmm.isConverged(p['threshold'])
                   row['fitted'] = hmm.fitted

                row.append()

            table.flush()
            hf5.flush()

        print('Done!')

    def _setup_hdf5(self):
        h5_file = self.h5_file

        with tables.open_file(h5_file, 'a') as hf5:
            # Taste -> PI, A, B, BIC, state_sequences, nStates, nCells, dt
            if not 'data_overview' in hf5.root:
                # Contains taste, channel, n_cells, n_trials, n_states, dt, BIC
                table = hf5.create_table('/', 'data_overview', MultiInfoParticle,
                                         'Basic info for each digital_input')
                table.flush()

            if hasattr(self, '_data_params') and self._data_params is not None:
                for p in self._data_params:
                    hmm_str = 'hmm_%i' % p['hmm_id']
                    if hmm_str not in hf5.root:
                        hf5.create_group('/', hmm_str, 'Data for HMM #%i' % p['hmm_id'])

            hf5.flush()

    def run(self, parallel=True):
        h5_file = self.h5_file
        fit_objs = self._fit_objects
        HMMs = dict.fromkeys([x['hmm_id'] for x in self._data_params])
        errors = []

        def update(ans):
            hmm_id = ans[0]
            hmm = ans
            if hmm_id in HMMs:
                new_hmm = pick_best_hmm([HMMs[hmm_id], hmm])
                HMMs[hmm_id] = new_hmm
            else:
                # Check history for lowest BIC
                HMMs[hmm_id] = hmm.set_to_lowest_BIC()

        def error_call(e):
            errors.append(e)

        if parallel:
            n_cpu = np.min((mp.cpu_count(), len(fit_objs)))
            pool = mp.get_context('spawn').Pool(mp.cpu_count())
            for f in fit_objs:
                pool.apply_async(f.run, callback=update, error_callback=error_call)

            pool.close()
            pool.join()
        else:
            for f in fit_objs:
                try:
                    ans = f.run()
                    update(ans)
                except Exception as e:
                    raise Exception(e)
                    error_call(e)

        self._fitted_models = HMMs
        self.write_overview_to_hdf5()
        self.save_fitted_models()
        if len(errors) > 0:
            print('Encountered errors: ')
            for e in errors:
                print(e)

    def save_fitted_models(self):
        models = self._fitted_models
        for k, v in models:
            write_hmm_to_hdf5(self.h5_file, k, v)
            plot_dir = os.path.join(self.plot_dir, 'HMM_%i' % k)
            if not os.path.isdir(plot_dir):
                os.makedirs(plot_dir)

            ids = [x['hmm_id'] for x in self._data_params]
            idx = np.where(ids == k)[0]
            params = self._data_params[idx]
            time_window = [params['time_start'], params['time_end']]
            hmmplt.plot_hmm_figures(v, time_window, save_dir=plot_dir)

def read_hmm_from_hdf5(h5_file, hmm_id, rec_dir):
    with tables.open_file(h5_file, 'r') as hf5:
        h_str = 'hmm_%i' % hmm_id
        if h_str not in hf5.root:
            return None

        table = hf5.root.data_overview
        row = table.where('hmm_id == id', condvars={'id':hmm_id})
        if len(row) == 0:
            raise ValueError('Parameters not found for hmm %i' % hmm_id)
        elif len(row) > 1:
            raise ValueError('Multiple parameters found for hmm %i' % hmm_id)

        units = query_units(rec_dir, row['unit_type'])
        spikes, dt, time = get_spike_data(rec_dir, units, row['channel'],
                                          dt=row['dt'],
                                          time_start=row['time_start'],
                                          time_end=row['time_end'])
        tmp = hf5.root[h_str]
        mats = {'initial_distribution': tmp['initial_distribution'][:],
                'transition': tmp['transition'][:],
                'emission': tmp['emission'][:],
                'fitted': row['fitted']}
        hmm = phmm.PoissonHMM(row['n_states'], spikes, dt, set_data=mats)

    return hmm

def write_hmm_to_hdf5(h5_file, hmm_id, hmm):
    h_str = 'hmm_%i' % hmm_id
    with tables.open_file(h5_file, 'a') as hf5:
        if hmm_id in hf5.root:
            hf5.remove_node('/', h_str, recursive=True)

        hf5.create_group('/', h_str, 'Data for HMM #%i' % i)
        hf5.create_array('/'+h_str, 'initial_distribution',
                         hmm.initial_distribution)
        hf5.create_array('/'+h_str, 'transition', hmm.transition)
        hf5.create_array('/'+h_str, 'emission', hmm.emission)

        best_paths, _ = hmm.get_best_paths()
        hf5.create_array('/'+h_str, 'state_sequences', best_paths)

def query_units(dat, unit_type):
    '''Returns the units names of all units in the dataset that match unit_type

    Parameters
    ----------
    dat : blechpy.dataset or str
        Can either be a dataset object or the str path to the recording
        directory containing that data .h5 object
    unit_type : str, {'single', 'pyramidal', 'interneuron', 'all'}
        determines whether to return 'single' units, 'pyramidal' (regular
        spiking single) units, 'interneuron' (fast spiking single) units, or
        'all' units

    Returns
    -------
        list of str : unit_names
    '''
    if isinstance(dat, str):
        units = blechpy.dio.h5io.get_unit_table(dat)
    else:
        units = dat.get_unit_table()

    u_str = unit_type.lower()
    q_str = ''
    if u_str == 'single':
        q_str = 'single_unit == True'
    elif u_str == 'pyramidal':
        q_str = 'single_unit == True and regular_spiking == True'
    elif u_str == 'interneuron':
        q_str = 'single_unit == True and fast_spiking == True'
    elif u_str == 'all':
        return units['unit_name'].tolist()
    else:
        raise ValueError('Invalid unit_type %s. Must be '
                         'single, pyramidal, interneuron or all' % u_str)

    return units.query(q_str)['unit_name'].tolist()

    # Parameters
    # hmm_id
    # taste
    # channel
    # n_cells
    # unit_type
    # n_trials
    # dt
    # threshold
    # time_start
    # time_end
    # n_repeats
    # n_states
    # n_iterations
    # BIC
    # cost
    # converged
    # fitted
    #
    # Extras: unit_names, rec_dir

class HMMFit(object):
    def __init__(self, rec_dir, params):
        self._rec_dir = rec_dir
        self._params = params

    def run(self):
        params = self._params
        spikes, dt, time = self.get_spike_data()
        hmm = phmm.fit_hmm_mp(params['n_states'], spikes, dt,
                              params['max_iter'], params['threshold'])
        return params['hmm_id'], hmm

    def get_spike_data(self):
        p = self._params
        units = query_units(self._rec_dir, p['unit_type'])
        # Get stored spike array, time is in ms, dt is usually 1 ms
        spike_array, dt, time = get_spike_data(self._rec_dir, units,
                                               p['channel'], dt=p['dt'],
                                               time_start=p['time_start'],
                                               time_end=p['time_end'])
        return spike_array, dt, time

def get_spike_data(rec_dir, units, channel, dt=None, time_start=None, time_end=None):
    time, spike_array = blechpy.dio.h5io.get_spike_data(rec_dir, units, channel)
    curr_dt = np.unique(np.diff(time))[0] / 1000
    if dt is not None and curr_dt < dt:
        spike_array, time = rebin_spike_array(spike_array, dt, time, p['dt'])
    elif dt is not None and curr_dt > dt:
        raise ValueError('Cannot upsample spike array from %f ms '
                         'bins to %f ms bins' % (dt, curr_dt))
    else:
        dt = curr_dt

    if time_start and time_end:
        idx = np.where((time >= time_start) & (time < time_end))[0]
        time = time[idx]
        spike_array = spike_array[:, :, idx]

    return spike_array.astype('int32'), dt, time

def pick_best_hmm(HMMs):
    '''For each HMM it searches the history for the HMM with lowest BIC Then it
    compares HMMs. Those with same # of free parameters are compared by BIC
    Those with different # of free parameters (namely # of states) are compared
    by cost Best HMM is returned

    Parameters
    ----------
    HMMs : list of phmm.PoissonHmm objects

    Returns
    -------
    phmm.PoissonHmm
    '''
    # First optimize each HMMs and sort into groups based on # of states
    groups = {}
    for hmm in HMMs:
        hmm.set_to_lowest_BIC()
        if hmm.n_states in groups:
            groups[hmm.n_states].append(hmm)
        else:
            groups[hmm.n_states] = [hmm]

    best_per_state = {}
    for k, v in groups:
        BICs = np.array([x.get_BIC() for x in v])
        idx = np.argmin(BICs)
        best_per_state[k] = v[idx]

    hmm_list = best_per_state.values()
    costs = np.array([x.get_cost() for x in hmm_list])
    idx = np.argmin(costs)
    return hmm_list[idx]

    # Compare HMMs with same number of states by BIC


"""
Sprite Collect Coins

Artwork from http://kenney.nl

http://arcade.academy/examples/index.html

"""
import arcade
import random
import pymunk
import os

SCREEN_WIDTH = 800
SCREEN_HEIGHT = 600

COIN_SCALE = 0.2
COIN_COUNT = 30

VIEWPORT_MARGIN = 40
RIGHT_MARGIN = 150

MOVEMENT_SPEED = 2

"""
 MAP=[[1,1,1,1,1,1,1,1,1,1,1,1,1],
     [0,0,0,0,0,0,0,0,1,1,0,0,0],
     [0,0,0,0,0,0,0,0,0,0,0,0,0],
     [1,1,1,1,1,1,1,1,1,1,1,1,1],
     ]
"""

def detect_collision(player, items_list):
    #brute force method
    ret_index = []
    p_width, p_height = (player.width, player.height)
    p_x, p_y = player.get_position()
    i = 0
    for item in items_list:
        i_width, i_height = (item.width, item.height)
        i_x, i_y = item.get_position()
        if (p_x < (i_x + i_width) and (p_x + p_width) > i_x and \
           p_y < (i_y + i_height) and p_height + p_y > i_y):
            ret_index.append(item)
        i += 1
    return ret_index

class Bullet(arcade.Sprite):
    def __init__(self):
        super().__init__()

class Coin(arcade.AnimatedTimeSprite):
    """ Coin class """

    def __init__(self): #  image, scale):
        """ Set up the player """
        super().__init__()
        self.textures = []
        self.center_x = random.randrange(SCREEN_WIDTH)
        self.center_y = random.randrange(SCREEN_HEIGHT)
        coin_img_path = "img/silver_coins/Silver_%i.png"
        for y in range(1, 11):
            self.textures.append(arcade.load_texture(coin_img_path % y, scale=0.1))
        self.cur_texture_index = random.randrange(len(self.textures))

class Zombie(arcade.AnimatedTimeSprite):
    """ Player class """

    def __init__(self, gender="male"):#, image, scale):
        """ Set up the player """

        # Call the parent init
        super().__init__() #image, scale)
        self.walk_direction = 1
        self.walk_from_x = 0
        self.walk_to_x = 0
        self.textures_walk_left = []
        self.textures_walk_right = []
        self.textures_idle_left = []
        self.textures_idle_right = []
        self.textures_attack_right = []
        self.textures_attack_left = []
        self.textures_dead_right = []
        self.textures_dead_left = []
        self.textures = []
        self.center_x = random.randrange(SCREEN_WIDTH)
        self.center_y = random.randrange(SCREEN_HEIGHT)
        attack_img = "img/" + gender + "/Attack (%i).png"
        walk_img = "img/" + gender + "/Walk (%i).png"
        idle_img = "img/" + gender + "/Idle (%i).png"
        dead_img = "img/" + gender + "/Dead (%i).png"
        self.direction = "left"
        for y in range(1, 9):
            self.textures_attack_right.append(arcade.load_texture(attack_img % y, scale=COIN_SCALE))
            self.textures_attack_left.append(arcade.load_texture(attack_img % y, mirrored=True, scale=COIN_SCALE))
        for y in range(1, 11):
            self.textures_walk_right.append(arcade.load_texture(walk_img % y, scale=COIN_SCALE))
            self.textures_walk_left.append(arcade.load_texture(walk_img % y, mirrored=True, scale=COIN_SCALE))
        for y in range(1, 16):
            self.textures_idle_left.append(arcade.load_texture(idle_img % y, mirrored=True, scale=COIN_SCALE))
            self.textures_idle_right.append(arcade.load_texture(idle_img % y, scale=COIN_SCALE))
        for y in range(1, 13):
            self.textures_dead_left.append(arcade.load_texture(dead_img % y, mirrored=True, scale=COIN_SCALE))
            self.textures_dead_right.append(arcade.load_texture(dead_img % y, scale=COIN_SCALE))

        self.textures = self.textures_walk_left
        self.cur_texture_index = random.randrange(len(self.textures))

        # Create a variable to hold our speed. 'angle' is created by the parent
        self.speed = 0

    def get_last_side(self):
        return self.action.split("_")[0]

    def get_action(self):
        return self.action

    def set_action(self, action):
        self.action = action
        if action == "left_walk":
            self.textures = self.textures_walk_left
        elif action == "right_walk":
            self.textures = self.textures_walk_right
        elif action == "left_idle":
            self.textures = self.textures_idle_left
        elif action == "right_idle":
            self.textures = self.textures_idle_right
        elif action == "left_attack":
            self.textures = self.textures_attack_left
        elif action == "right_attack":
            self.textures = self.textures_attack_right

    def update(self):
        if self.center_x <= self.walk_from_x:
            self.walk_direction = +1.0
            self.set_action("right_walk")
            print("right_walk")
        if self.center_x >= self.walk_to_x:
            print("left_walk")
            self.walk_direction = -1.0
            self.set_action("left_walk")
        self.center_x += self.walk_direction*MOVEMENT_SPEED

class Player(arcade.AnimatedTimeSprite):
    """ Player class """

    def __init__(self, scale=COIN_SCALE):#, image, scale):
        """ Set up the player """

        # Call the parent init
        super().__init__()#image, scale)
        self.textures_walk_left = []
        self.textures_walk_right = []
        self.textures_idle_left = []
        self.textures_idle_right = []
        self.textures_attack_right = []
        self.textures_attack_left = []
        self.textures_dead_right = []
        self.textures_dead_left = []
        self.textures_jump_right = []
        self.textures_jump_left = []
        self.textures = []
        self.center_x = random.randrange(SCREEN_WIDTH)
        self.center_y = random.randrange(SCREEN_HEIGHT)
        attack_img = "img/sheriff_girl/Shoot (%i).png"
        walk_img = "img/sheriff_girl/Run (%i).png"
        idle_img = "img/sheriff_girl/Idle (%i).png"
        dead_img = "img/sheriff_girl/Dead (%i).png"
        jump_img = "img/sheriff_girl/Jump (%i).png"
        self.action = "left_idle"

        for y in range(1, 4):
            self.textures_attack_right.append(arcade.load_texture(attack_img % y, scale=COIN_SCALE))
            self.textures_attack_left.append(arcade.load_texture(attack_img % y, mirrored=True, scale=COIN_SCALE))
        for y in range(1, 9):
            self.textures_walk_right.append(arcade.load_texture(walk_img % y, scale=COIN_SCALE))
            self.textures_walk_left.append(arcade.load_texture(walk_img % y, mirrored=True, scale=COIN_SCALE))
        for y in range(1, 11):
            self.textures_idle_left.append(arcade.load_texture(idle_img % y, mirrored=True, scale=COIN_SCALE))
            self.textures_idle_right.append(arcade.load_texture(idle_img % y, scale=COIN_SCALE))
        for y in range(1, 11):
            self.textures_dead_left.append(arcade.load_texture(dead_img % y, mirrored=True, scale=COIN_SCALE))
            self.textures_dead_right.append(arcade.load_texture(dead_img % y, scale=COIN_SCALE))
        for y in range(1, 11):
            self.textures_jump_left.append(arcade.load_texture(jump_img % y, mirrored=True, scale=COIN_SCALE))
            self.textures_jump_right.append(arcade.load_texture(jump_img % y, scale=COIN_SCALE))

        self.textures = self.textures_idle_left
        self.cur_texture_index = random.randrange(len(self.textures))

        # Create a variable to hold our speed. 'angle' is created by the parent
        self.speed = 0

    def get_last_side(self):
        return self.action.split("_")[0]

    def get_action(self):
        return self.action

    def set_action(self, action):
        self.action = action
        if action == "left_walk":
            self.textures = self.textures_walk_left
        elif action == "right_walk":
            self.textures = self.textures_walk_right
        elif action == "left_idle":
            self.textures = self.textures_idle_left
        elif action == "right_idle":
            self.textures = self.textures_idle_right
        elif action == "left_attack":
            self.textures = self.textures_attack_left
        elif action == "right_attack":
            self.textures = self.textures_attack_right
        elif action == "left_jump":
            self.textures = self.textures_jump_left
        elif action == "right_jump":
            self.textures = self.textures_jump_right
        elif action == "left_dead":
            self.textures = self.textures_dead_left
        elif action == "right_dead":
            self.textures = self.textures_dead_right

class MyGame(arcade.Window):
    """ Main application class. """

    def __init__(self, width, height):
        """
        Initializer
        """
        super().__init__(width, height)

        # Set the working directory (where we expect to find files) to the same
        # directory this .py file is in. You can leave this out of your own
        # code, but it is needed to easily run the examples using "python -m"
        # as mentioned at the top of this program.
        file_path = os.path.dirname(os.path.abspath(__file__))
        os.chdir(file_path)

        """ Set up the game and initialize the variables. """

        # Sprite lists
        self.all_sprites_list = None
        self.coin_list = None
        self.wall_list = None
        self.bullet_list = None
        self.zombie_list = None
        self.physics_engine = None

        # Set up the player
        self.score = 0
        self.player = None
        self.key_pressed = False
        self.view_left = 0
        self.view_bottom = 0

    def setup(self):
        # Set Spritelists
        self.all_sprites_list = arcade.SpriteList()
        self.coin_list = arcade.SpriteList()
        self.zombie_list = arcade.SpriteList()

        # Set up the player
        self.score = 0
        self.wall_list = arcade.SpriteList()
        self.bullet_list = arcade.SpriteList()
        character_scale = 0.3
        SPRITE_SCALING = 0.5
        # crete enemy
        zombie = Zombie("female")
        zombie.set_action("right_walk")
        zombie.walk_from_x = 100
        zombie.walk_to_x = 300
        zombie.set_position(200, 120)
        self.zombie_list.append(zombie)
        self.all_sprites_list.append(zombie)

        zombie = Zombie("male")
        zombie.set_action("right_walk")
        zombie.walk_from_x = 400
        zombie.walk_to_x = 700
        zombie.set_position(200, 120)
        self.zombie_list.append(zombie)
        self.all_sprites_list.append(zombie)

        # -- Set up the walls
        # Create a row of boxes
        # x_index = 0
        # y_index = 0
        # for y_arr in MAP:
        #    # x_index=MAP.index(x_arr)
        #     for item in y_arr:
        #        # y_index=x_arr.index(item)
        #         if item == 1:
        #             pass
        #            # wall = arcade.Sprite("boxCrate_double.png", SPRITE_SCALING)
        #            # wall.top = (7-x_index)*64
        #            # wall.right = y_index*64
        #            # self.wall_list.append(wall)
        #         x_index += 1
        #         print(x_index, y_index)
        #     y_index+=1

        # Create box wall
        for x in range(0, 1500, 64):
            wall = arcade.Sprite("img/land/boxCrate_double.png", SPRITE_SCALING)
            #print(wall.width)
            wall.center_x = x
            wall.center_y = 64/2
            self.wall_list.append(wall)

        self.player = Player()
        self.player.texture_change_distance = 20

        self.player.center_x = SCREEN_WIDTH // 2
        self.player.center_y = SCREEN_HEIGHT // 2
        self.player.scale = 0.8

        self.all_sprites_list.append(self.player)

        counter = 0
        for i in range(COIN_COUNT):
            while True:
                coin = Coin()
                collision = detect_collision(coin, self.coin_list)
                counter += 1
                print("break", counter)
                if len(collision) == 0:
                    self.coin_list.append(coin)
                    break
            # coin = Coin()#arcade.AnimatedTimeSprite(scale=0.3)
            self.all_sprites_list.append(coin)

        # Set the background color
        self.physics_engine = arcade.PhysicsEnginePlatformer(self.player, self.wall_list, gravity_constant=0.5)
        arcade.set_background_color(arcade.color.AMAZON)

    def on_draw(self):
        """
        Render the screen.
        """

        # This command has to happen before we start drawing
        arcade.start_render()

        # Draw all the sprites.
        self.all_sprites_list.draw()
        self.wall_list.draw()
        self.bullet_list.draw()

        # Put the text on the screen.
        output = "Score: {}".format(self.score)
        arcade.draw_text(output, 10 + self.view_left, 20, arcade.color.WHITE, 14)

    def on_key_press(self, key, modifiers):
        """
        Called whenever the mouse moves.
        """
        self.key_pressed = True
        if key == arcade.key.UP:
            if self.physics_engine.can_jump():
                print("can_jump")
                self.player.change_y = 10
                if self.player.get_last_side() == "left":
                    self.player.set_action("left_jump")
                else:
                    self.player.set_action("right_jump")
        elif key == arcade.key.DOWN:
            pass
        elif key == arcade.key.A:
            # dont move while shooting
            self.player.change_x = 0
            # Create a bullet
            bullet = arcade.Sprite("img/laserBlue01.png", 0.2)

            # Position the bullet at the player's current location
            start_x, start_y = (self.player.center_x, self.player.center_y)
            bullet.center_x, bullet.center_y = (start_x, start_y)
            if self.player.get_last_side() == "right":
                bullet.change_x = 10
            else:
                bullet.change_x = -10

            self.bullet_list.append(bullet)
            if self.player.get_last_side() == "left":
                self.player.set_action("left_attack")
            else:
                self.player.set_action("right_attack")
        elif key == arcade.key.S:
            # dont move while shooting
            self.player.change_x = 0
            if self.player.get_last_side() == "left":
                self.player.set_action("left_dead")
            else:
                self.player.set_action("right_dead")
        elif key == arcade.key.LEFT:
            self.player.change_x = -MOVEMENT_SPEED
            self.player.set_action("left_walk")
        elif key == arcade.key.RIGHT:
            self.player.change_x = MOVEMENT_SPEED
            self.player.set_action("right_walk")
        elif key == arcade.key.SPACE:
            # print("space ")
            self.player.set_action("right_walk")

    def on_key_release(self, key, modifiers):
        """
        Called when the user presses a mouse button.
        """
        self.key_pressed = False
        if key == arcade.key.UP or key == arcade.key.DOWN:
            self.player.change_y = 0
        elif key == arcade.key.LEFT or key == arcade.key.RIGHT:
            self.player.change_x = 0
        if self.player.get_last_side() == "left":
            self.player.set_action("left_idle")
        else:
            self.player.set_action("right_idle")

    def update(self, delta_time):
        """ Movement and game logic """
        self.all_sprites_list.update()
        self.all_sprites_list.update_animation()
        self.physics_engine.update()
        self.bullet_list.update()

        # Generate a list of all sprites that collided with the player.
        hit_list = detect_collision(self.player, self.coin_list)
        for coin_index in hit_list:
            # print(coin_index, len(self.coin_list))
            self.coin_list.remove(coin_index)
            # print(coin_index.width)
            coin_index.kill()
            self.score += 1

        zombie_hit = detect_collision(self.player, self.zombie_list)
        for zombie in zombie_hit:
            # print("zombie bite",zombie)
            if zombie.get_last_side() == "left":
                zombie.set_action("left_attack")
            else:
                zombie.set_action("right_attack")
        changed = False
        # Scroll left
        left_bndry = self.view_left + VIEWPORT_MARGIN
        if self.player.left < left_bndry:
            self.view_left -= int(left_bndry - self.player.left)
            changed = True

        # Scroll right
        right_bndry = self.view_left + SCREEN_WIDTH - RIGHT_MARGIN
        if self.player.right > right_bndry:
            self.view_left += int(self.player.right - right_bndry)
            changed = True
        # If we need to scroll, go ahead and do it.
        if changed:
            arcade.set_viewport(self.view_left,
                                SCREEN_WIDTH + self.view_left,
                                self.view_bottom,
                                SCREEN_HEIGHT + self.view_bottom)

def main():
    """ Main method """
    window = MyGame(SCREEN_WIDTH, SCREEN_HEIGHT)
    window.setup()
    arcade.run()

if __name__ == "__main__":
    main()


from django.conf import settings
from django.core.exceptions import ValidationError
from django.core.validators import MinValueValidator
from django.db import models

from station.upload_to_path import UploadToPath

class TrainType(models.Model):
    name = models.CharField(max_length=255)

    class Meta:
        ordering = ["name"]

    def __str__(self):
        return self.name

class Train(models.Model):
    name = models.CharField(max_length=255)
    number = models.CharField(max_length=8, unique=True)
    train_type = models.ForeignKey(
        TrainType, on_delete=models.CASCADE, related_name="trains"
    )
    image = models.ImageField(
        blank=True, null=True, upload_to=UploadToPath("train-images/")
    )

    @property
    def capacity(self):
        return sum(carriage.seats for carriage in self.carriages.all())

    class Meta:
        ordering = ["number"]

    def __str__(self):
        return self.name

class Carriage(models.Model):
    class CarriageType(models.TextChoices):
        ECONOMY = ("Economy", "Economy Class")
        BUSINESS = ("Business", "Business Class")
        PREMIUM = ("Premium", "Premium Class")

    CARRIAGE_TYPE_SEAT_PRICES = {
        CarriageType.ECONOMY: 50,
        CarriageType.BUSINESS: 100,
        CarriageType.PREMIUM: 150,
    }

    carriage_type = models.CharField(
        max_length=10,
        choices=CarriageType.choices,
        default=CarriageType.ECONOMY,
    )

    number = models.IntegerField(validators=[MinValueValidator(1)])
    seats = models.IntegerField()
    train = models.ForeignKey(Train, on_delete=models.CASCADE, related_name="carriages")

    @property
    def seat_price(self):
        return self.CARRIAGE_TYPE_SEAT_PRICES[Carriage.CarriageType(self.carriage_type)]

    @staticmethod
    def validate_carriage_number(number, train, error_to_raise):
        if Carriage.objects.filter(number=number, train=train).exists():
            raise error_to_raise(
                {
                    "number": "Carriage with this number "
                    "already exists for this train."
                }
            )

    def is_seat_number_valid(self, seat_number):
        return 1 <= seat_number <= self.seats

    class Meta:
        ordering = ["number"]
        constraints = [
            models.UniqueConstraint(
                fields=["number", "train"], name="unique_carriage_number"
            )
        ]

    def __str__(self):
        return f"Carriage {self.number} of {self.train}"

class Station(models.Model):
    name = models.CharField(max_length=255)
    latitude = models.FloatField()
    longitude = models.FloatField()

    class Meta:
        ordering = ["name"]

    def __str__(self):
        return self.name

class Route(models.Model):
    name = models.CharField(max_length=255)
    distance = models.IntegerField()
    from_station = models.ForeignKey(
        Station, on_delete=models.CASCADE, related_name="from_station"
    )
    to_station = models.ForeignKey(
        Station, on_delete=models.CASCADE, related_name="to_station"
    )

    class Meta:
        ordering = ["name"]

    def __str__(self):
        return self.name

class Crew(models.Model):
    first_name = models.CharField(max_length=255)
    last_name = models.CharField(max_length=255)
    image = models.ImageField(
        blank=True, null=True, upload_to=UploadToPath("crew-images/")
    )

    class Meta:
        ordering = ["first_name", "last_name"]
        verbose_name_plural = "crew"

    def __str__(self):
        return f"{self.first_name} {self.last_name}"

    @property
    def full_name(self):
        return f"{self.first_name} {self.last_name}"

class Journey(models.Model):
    crew = models.ManyToManyField(Crew)
    departure_time = models.DateTimeField()
    arrival_time = models.DateTimeField()
    route = models.ForeignKey(Route, on_delete=models.CASCADE, related_name="journeys")
    train = models.ForeignKey(Train, on_delete=models.CASCADE, related_name="journeys")
    image = models.ImageField(
        blank=True, null=True, upload_to=UploadToPath("journey-images/")
    )

    def clean(self):
        if self.departure_time >= self.arrival_time:
            raise ValidationError(
                {"arrival_time": "Arrival time must be greater than departure time"}
            )

    def save(self, *args, **kwargs):
        self.full_clean()
        return super(Journey, self).save(*args, **kwargs)

    class Meta:
        ordering = ["-departure_time"]

    def __str__(self):
        return (
            f"Journey {self.route.name}, train number — {self.train.number} "
            f"[{self.departure_time} - {self.arrival_time}]"
        )

class Order(models.Model):
    created_at = models.DateTimeField(auto_now_add=True)
    user = models.ForeignKey(
        settings.AUTH_USER_MODEL, on_delete=models.CASCADE, related_name="orders"
    )

    class Meta:
        ordering = ["-created_at"]

    def __str__(self):
        return f"Order {self.id} by {self.user} at {self.created_at}"

class Ticket(models.Model):
    seat = models.IntegerField()
    carriage = models.ForeignKey(
        Carriage, on_delete=models.CASCADE, related_name="tickets"
    )
    journey = models.ForeignKey(
        Journey, on_delete=models.CASCADE, related_name="tickets"
    )
    order = models.ForeignKey(Order, on_delete=models.CASCADE, related_name="tickets")

    @property
    def price(self):
        return self.carriage.seat_price

    @staticmethod
    def validate_ticket(seat, carriage, journey, error_to_raise):
        if Ticket.objects.filter(
            carriage=carriage, seat=seat, journey=journey
        ).exists():
            raise error_to_raise(
                {
                    "seat": f"A ticket for seat: {seat}, "
                    f"in carriage: {carriage.number}, "
                    f"on train: {journey.train}, "
                    f"on journey route: {journey.route.name}, "
                    f"already exists."
                }
            )

        if not carriage.is_seat_number_valid(seat):
            raise error_to_raise(
                {
                    "seat": f"Seat number must be in available range: "
                    f"(1, {carriage.seats}), "
                    f"but got: {seat}"
                }
            )

    def clean(self):
        self.validate_ticket(
            self.seat,
            self.carriage,
            self.journey,
            ValidationError,
        )

    def save(
        self, force_insert=False, force_update=False, using=None, update_fields=None
    ):
        self.full_clean()
        return super(Ticket, self).save(
            force_insert, force_update, using, update_fields
        )

    class Meta:
        constraints = [
            models.UniqueConstraint(
                fields=["carriage", "seat", "journey"], name="unique_ticket"
            )
        ]
        ordering = ["carriage", "seat"]

    def __str__(self):
        return (
            f"Ticket for {self.journey} "
            f"in carriage {self.carriage.number} seat {self.seat}"
        )


from django.conf.urls import url

from . import views

urlpatterns = [
    url(r'^$', views.index, name='index'),
    url(r'^(?P<id>[0-9]+)/vendor_list/$', views.vendor_list, name='vendor_list'),
    url(r'^(?P<pk>[0-9]+)/$', views.DetailView.as_view(), name='detail'),
    url(r'^(?P<pk>[0-9]+)/results/$', views.ResultsView.as_view(), name='results'),
]


from datetime import timezone
from django.db.models import Q
from rest_framework import permissions
from rest_framework import status
from rest_framework import viewsets
from rest_framework.decorators import action
from rest_framework.response import Response
from rest_framework.views import APIView
from procurement.serializers import *
from .serializers import *
from master.serializers import ModuleAccessSerializer, ModuleMasterSerializer

from approval.SendMail import alertMail

# from approval.sendActionableMessage import alertMail

User = get_user_model()

# Create your views here.

class GetModuleAccessViewSet(viewsets.ModelViewSet):
    serializer_class = ModuleAccessSerializer
    queryset = ModuleAccess.objects.all()

    def get_queryset(self):
        queryset = ModuleAccess.objects.all()
        return queryset

    @action(methods=['post'], detail=False, url_path='getModuleAccess')
    def getModuleAccess(self, request):
        data = request.data
        is_admin = data['is_admin']
        is_approver = data['is_approver']
        is_dsin = data['is_dsin']
        is_requester = data['is_requester']
        arr = []
        module_ids = []

        if is_admin == True:
            queryset = ModuleAccess.objects.filter(is_admin=True).values_list('moduleId_id', flat=True)
            module_ids = list(queryset)

        if is_approver == True:
            queryset = ModuleAccess.objects.filter(is_approver=True).values_list('moduleId_id', flat=True)
            module_ids.extend(list(queryset))

        if is_dsin == True:
            queryset = ModuleAccess.objects.filter(is_dsin=True).values_list('moduleId_id', flat=True)
            module_ids.extend(list(queryset))

        if is_requester == True:
            queryset = ModuleAccess.objects.filter(is_requester=True).values_list('moduleId_id', flat=True)
            module_ids.extend(list(queryset))

        module_ids = list(set(module_ids))

        for i in module_ids:
            root = ModuleMaster.objects.filter(moduleId=i).values('root')[0]['root']
            if root == 'ROOT':
                filter_module_root = ModuleMaster.objects.filter(moduleId=i)
                filter_module_root = ModuleMasterSerializer(filter_module_root, many=True, context={'request': request})
                filter_module_root = filter_module_root.data
                filter_module_submenu = ModuleMaster.objects.filter(root=i, moduleId__in=module_ids)
                filter_module_submenu = ModuleMasterSerializer(filter_module_submenu, many=True,
                                                               context={'request': request})
                filter_module_submenu = filter_module_submenu.data

                arr.append({"module_id": filter_module_root[0]['moduleId'],
                            "module_name": filter_module_root[0]['module_name'],
                            "module_slug": filter_module_root[0]['module_slug'],
                            "root": filter_module_root[0]['root'],
                            "m_color": filter_module_root[0]['m_color'],
                            "m_icon_name": filter_module_root[0]['m_icon_name'],
                            "m_link": filter_module_root[0]['m_link'],
                            "root_module": filter_module_submenu})

        return Response({
            'message': 'success',
            'module_ids': arr
        })

class ApprovalProcurementPendingList(APIView):
    permission_classes = [permissions.IsAuthenticated]

    def get(self, request, *args, **kwargs):
        user = request.user
        ApprovalUser = ApprovalTransaction.objects.filter(
            Q(approvalUserName=user) & Q(status='Pending') & Q(is_active=True))
        if ApprovalUser.exists():
            serializer = ApprovalTransactionSerializer(ApprovalUser, many=True)
            data_list = []
            for data in serializer.data:
                procurement = MasterProcurement.objects.get(id=data['procurementId'])
                procurement_serializer = MasterProcurementSerializer(procurement)
                prevoius_approver_val = None
                if data['sequence'] > 1:
                    prevoius_approver_val = ApprovalTransaction.objects.filter(sequence__lt=data['sequence'],
                                                                               procurementId=data['procurementId'],
                                                                               is_active=True).order_by('-sequence').first()
                    if prevoius_approver_val:
                        prevoius_approver_val = ApprovalTransactionSerializer(prevoius_approver_val).data
                        prevoius_approver_val = prevoius_approver_val['approvalUserName']
                        username = User.objects.filter(username=prevoius_approver_val).values('name')[0]['name']
                    else:
                        prevoius_approver_val = 'Request Submitted'

                updated_data = {
                    'id': data['id'],
                    'approvalUserName': data['approvalUserName'],
                    'approverEmail': data['approverEmail'],
                    'sequence': data['sequence'],
                    'approverType': data['approverType'],
                    'status': data['status'],
                    'prevoius_approver_data': prevoius_approver_val,
                    'procurementId': {
                        'id': data['procurementId'],
                        'sequence': data['sequence'],  # Include the sequence here
                        'approverType': data['approverType'],
                        'sequence_status': data['status'],  # Include the status here
                        'RequestNumber': procurement_serializer.data['RequestNumber'],
                        'RequestBy': procurement_serializer.data['RequestBy'],
                        'RequestType': procurement_serializer.data['RequestType'],
                        'Status': procurement_serializer.data['Status'],
                        'TotalBudget': procurement_serializer.data['TotalBudget'],
                        'UtilizedBudget': procurement_serializer.data['UtilizedBudget'],
                        'Department': procurement_serializer.data['Department'],
                        'TotalAmount': procurement_serializer.data['TotalAmount'],
                        'DeviceType': procurement_serializer.data['DeviceType'],
                        # 'inlineitem': procurement_serializer.data['inlineitem'],
                        'previous_approver_details': prevoius_approver_val,
                    }
                }
                data_list.append(updated_data)
            return Response(data_list)
        else:
            return Response({'error': 'No pending approval'})

class GetProcurementApprovalTransactionDetails(APIView):
    serializer_class = ApprovalTransactionSerializer

    def get(self, request, *args, **kwargs):
        if request.query_params.get('procurementId') and request.query_params.get('sequenceId'):
            # get the procurement id and sequence id
            procurement_id = int(request.query_params.get('procurementId'))
            sequence = int(request.query_params.get('sequenceId'))
            try:
                instance = ApprovalTransaction.objects.get(procurementId=procurement_id, sequence=sequence,
                                                           is_active=True)
            except ApprovalTransaction.DoesNotExist:
                return Response({'error': 'Approval transaction does not exist'}, status=status.HTTP_400_BAD_REQUEST)
            #
            serializer = ApprovalTransactionSerializer(instance)
            procurement_instance = MasterProcurement.objects.get(id=serializer.data['procurementId'])
            procurement_serializer = MasterProcurementSerializer(procurement_instance)

            files_val = []
            files = MoreAttachments.objects.filter(procurement_id=procurement_id, is_active=True)
            files_serializer = MoreAttachmentsSerializer(files, many=True)
            files_val = files_serializer.data

            final_previous_approver = []
            prevoius_approver_data = []
            if sequence > 1:
                # prevoius_approver_val = ApprovalTransaction.objects.filter(sequence__lt=sequence, need to discuss this
                #                                                             procurementId=procurement_id, is_active=True)
                prevoius_approver_val = ApprovalTransaction.objects.filter(sequence__lt=sequence,
                                                                           procurementId=procurement_id,
                                                                           is_active=True)
                prevoius_approver_val = ApprovalTransactionSerializer(prevoius_approver_val, many=True)
                prevoius_approver_data.append(prevoius_approver_val.data)

                for val in prevoius_approver_data[0]:
                    # if val['approverType'] == 'BuHead':
                    #     val['approverTitle'] == 'BU Head'
                    if User.objects.filter(username=val['approvalUserName']).exists():
                        val['approverName'] = User.objects.filter(username=val['approvalUserName']).values('name')[0][
                            'name']

                if prevoius_approver_data is not None:
                    final_previous_approver = prevoius_approver_data[0]
                else:
                    final_previous_approver = prevoius_approver_data
            return Response({
                'procurementId': serializer.data['procurementId'],
                'sequence': serializer.data['sequence'],
                'Approver': serializer.data['approvalUserName'],
                'Email': serializer.data['approverEmail'],
                'RequestNumber': procurement_serializer.data['RequestNumber'],
                'RequestBy': procurement_serializer.data['RequestBy'],
                'Remarks': procurement_serializer.data['Remarks'],
                'RequestType': procurement_serializer.data['RequestType'],
                'Status': procurement_serializer.data['Status'],
                'IsExpenditure': procurement_serializer.data['IsExpenditure'],
                'TotalBudget': procurement_serializer.data['TotalBudget'],
                'UtilizedBudget': procurement_serializer.data['UtilizedBudget'],
                'Department': procurement_serializer.data['Department'],
                'TotalAmount': procurement_serializer.data['TotalAmount'],
                'DeviceType': procurement_serializer.data['DeviceType'],
                'inlineitem': procurement_serializer.data['inlineitem'],
                'previous_approver_details': final_previous_approver,
                'attachments': files_val
                # 'previous_approver_details': prevoius_approver_data
            })
        else:
            return Response({'error': 'Invalid request'}, status=status.HTTP_400_BAD_REQUEST)

class ApprovalUpdateStatusAPIViewCorrect(APIView):
    permission_classes = [permissions.IsAuthenticated]

    def put(self, request, *args, **kwargs):
        if request.query_params.get('procurementId') and request.query_params.get('sequenceId'):
            procurement_id = int(request.query_params.get('procurementId'))
            sequence = int(request.query_params.get('sequenceId'))
            try:
                instance = ApprovalTransaction.objects.get(procurementId=procurement_id, sequence=sequence)
            except ApprovalTransaction.DoesNotExist:
                return Response({
                    'message': 'Approval transaction does not exist',
                    'status': 'failed'
                }, status=status.HTTP_400_BAD_REQUEST)

            serializer = ApprovalTransactionSerializer(instance, data=request.data, partial=True)
            serializer.is_valid(raise_exception=True)

            if sequence == 1:
                previous_approver_email = instance.approverEmail
                previous_approver_sequence = instance.sequence
                instance.status = request.data['status']
                instance.approvaldatetime = datetime.datetime.now(tz=timezone.utc)
                instance.create_by = request.user
                instance.update_by = request.user
                instance.save()
                # Update the next approver status
                new_sequence_number = previous_approver_sequence + 1
                if new_sequence_number == 2:
                    next_approver = ApprovalTransaction.objects.filter(procurementId=procurement_id,
                                                                       sequence=new_sequence_number).values(
                        'approverEmail')[0]['approverEmail']
                    if next_approver == previous_approver_email:
                        ApprovalTransaction.objects.filter(procurementId=procurement_id,
                                                           sequence=new_sequence_number).update(
                            status=request.data['status'], approvaldatetime=datetime.datetime.now(tz=timezone.utc))
                # Update the MasterProcurement instance
                procurement_instance = MasterProcurement.objects.get(id=instance.procurementId.id)
                procurement_instance.Status = request.data['status']
                procurement_instance.save()
                serializer.save()
                return Response(serializer.data)
            elif sequence != 1:
                instance.status = request.data['status']
                instance.approvaldatetime = datetime.datetime.now(tz=timezone.utc)
                instance.create_by = request.user
                instance.update_by = request.user
                instance.save()
                # Update the MasterProcurement instance
                procurement_instance = MasterProcurement.objects.get(id=instance.procurementId.id)
                procurement_instance.Status = request.data['status']
                procurement_instance.save()
                serializer.save()
                return Response(serializer.data)
        else:
            return Response({
                'message': 'Please provide the procurementId and sequenceId',
                'status': 'failed'
            }, status=status.HTTP_400_BAD_REQUEST)

class ApprovalUpdateStatusAPIView(APIView):
    permission_classes = [permissions.IsAuthenticated]

    def put(self, request, *args, **kwargs):
        previous_approver_email = None
        previous_approver_sequence = None
        next_approver_email = None
        if request.query_params.get('procurementId') and request.query_params.get('sequenceId'):
            procurement_id = int(request.query_params.get('procurementId'))
            sequence = int(request.query_params.get('sequenceId'))
            try:
                instance = ApprovalTransaction.objects.get(procurementId=procurement_id, sequence=sequence,
                                                           is_active=True)
            except ApprovalTransaction.DoesNotExist:
                return Response({
                    'message': 'Approval transaction does not exist',
                    'status': 'failed'
                }, status=status.HTTP_400_BAD_REQUEST)

            serializer = ApprovalTransactionSerializer(instance, data=request.data, partial=True)
            serializer.is_valid(raise_exception=True)

            # procurementMasterData For all mail Data
            procurement_Details = MasterProcurement.objects.get(id=instance.procurementId.id)

            requesterEmail = None
            if User.objects.filter(username=procurement_Details.Created_by_id).exists():
                requesterEmail = User.objects.filter(username=procurement_Details.Created_by_id).values('email')[0][
                    'email']

            # current seq initail values
            present_approver_email = instance.approverEmail
            previous_approver_email = instance.approverEmail
            previous_approver_sequence = instance.sequence
            total_transaction_count = ApprovalTransaction.objects.filter(procurementId=procurement_id,
                                                                         is_active=True).count()
            current_sequence_no = instance.sequence

            # updating the first sequence and procurement approval status
            instance.status = request.data['status']
            instance.approvaldatetime = datetime.datetime.now(tz=timezone.utc)
            instance.create_by = request.user
            instance.update_by = request.user
            instance.remarks = request.data['remarks']
            instance.save()

            if request.data['status'] == 'Modification':
                # setting Records in Approval Transaction to False
                instance.is_active = False
                if ApprovalTransaction.objects.filter(procurementId=procurement_id, is_active=True).exists():
                    ApprovalTransaction.objects.filter(procurementId=procurement_id, is_active=True).update(
                        is_active=False)

                procurement_instance = MasterProcurement.objects.get(id=instance.procurementId.id)
                procurement_instance.Status = request.data['status']
                procurement_instance.save()
                serializer.save()

                requester_mail_data = []
                cc = None
                cc = present_approver_email
                mail_data_requester = {
                    'RequestNumber': procurement_Details.RequestNumber,
                    'RequestType': procurement_Details.RequestType,
                    'department': procurement_Details.Department,
                    'expenditureBudgeted': procurement_Details.IsExpenditure,
                    'cc': present_approver_email,
                    'requested_person_email': requesterEmail,
                    'module': 'Requester',
                    'subject': procurement_Details.RequestType + " " + procurement_Details.RequestNumber + " " + 'ITOPEX' + " " + 'Request',
                    'Heading': 'Your Procurement Request Has Been Returned For Modification.'
                }
                requester_mail_data.append(mail_data_requester)
                alertMail(requesterEmail, requester_mail_data[0], cc)

            if request.data['status'] == 'Reject':
                procurement_instance = MasterProcurement.objects.get(id=instance.procurementId.id)
                procurement_instance.Status = request.data['status']
                procurement_instance.save()
                serializer.save()

                requester_mail_data = []
                cc = None
                cc = present_approver_email
                mail_data_requester = {
                    'RequestNumber': procurement_Details.RequestNumber,
                    'RequestType': procurement_Details.RequestType,
                    'Name': procurement_Details.Name,
                    'department': procurement_Details.Department,
                    'expenditureBudgeted': procurement_Details.IsExpenditure,
                    'cc': present_approver_email,
                    'requested_person_email': requesterEmail,
                    'module': 'Requester',
                    'Heading': 'Your Procurement Request Has Been Rejected.',
                    'subject': procurement_Details.RequestType + " " + procurement_Details.RequestNumber + " " + 'ITOPEX' + ' Request'
                }
                requester_mail_data.append(mail_data_requester)
                alertMail(requesterEmail, requester_mail_data[0], cc)

            # mail += User.objects.filter(id=obj['emp_id']).values('email')[0]['email'] + ',';
            # Update the next approver status in case if they have the same email as previous then
            # status = ui status or staus = 'Pending' for next approver and they have different email id

            if request.data['status'] == 'Approved':
                while current_sequence_no <= total_transaction_count:
                    new_sequence_number = previous_approver_sequence + 1
                    if ApprovalTransaction.objects.filter(procurementId=procurement_id, sequence=new_sequence_number,
                                                          is_active=True).exists():
                        next_approver_email = ApprovalTransaction.objects.filter(procurementId=procurement_id,
                                                                                sequence=new_sequence_number,
                                                                                is_active=True).values(
                            'approverEmail')[0]['approverEmail']

                        if next_approver_email == previous_approver_email:
                            ApprovalTransaction.objects.filter(procurementId=procurement_id,
                                                               sequence=new_sequence_number, is_active=True).update(
                                status=request.data['status'], remarks=request.data['remarks'],
                                approvaldatetime=datetime.datetime.now(tz=timezone.utc))
                            current_sequence_no = current_sequence_no + 1
                            previous_approver_sequence = previous_approver_sequence + 1

                        elif next_approver_email != previous_approver_email:
                            ApprovalTransaction.objects.filter(procurementId=procurement_id,
                                                               sequence=new_sequence_number, is_active=True).update(
                                status='Pending')

                            # next_approver_mail_id = ApprovalTransaction.objects.filter(procurementId=procurement_id,
                            #                                    sequence=new_sequence_number, is_active=True).values('approverEmail')[0]['approverEmail']

                            next_approver_details = ApprovalTransaction.objects.filter(procurementId=procurement_id,
                                                                                sequence=new_sequence_number,
                                                                                is_active=True).values_list(
                                'approverEmail', 'approverType')

                            next_approver_mail_id, next_approver_type, mailTitle = None, None, None
                            if len(next_approver_details) != 0:
                                next_approver_mail_id = next_approver_details[0][0]
                                next_approver_type = next_approver_details[0][1]
                                if next_approver_type is not None:
                                    if next_approver_type == 'BuHead':
                                        mailTitle = 'Department Approval'
                                    elif next_approver_type == 'DSINHead':
                                        mailTitle = 'IT Approval'
                                    elif next_approver_type == 'FinanceHead':
                                        mailTitle = 'Finance Approval'
                                    elif next_approver_type == 'MD':
                                        mailTitle = 'Final Approval'

                            approver_mail_data = []
                            cc = None
                            if (requesterEmail is not None) and (present_approver_email is None):
                                cc = requesterEmail
                            elif (requesterEmail is None) and (present_approver_email is not None):
                                cc = present_approver_email
                            elif (requesterEmail is not None) and (present_approver_email is not None):
                                cc = requesterEmail + ',' + present_approver_email
                            mail_data_approver = {
                                'RequestType': procurement_Details.RequestType,
                                'RequestNumber': procurement_Details.RequestNumber,
                                'Name': procurement_Details.Name,
                                'department': procurement_Details.Department,
                                'expenditureBudgeted': procurement_Details.IsExpenditure,
                                'cc': cc,
                                'requested_person_email': next_approver_mail_id,
                                'module': 'Approver',
                                'Heading': 'Your Approval Is Required.',
                                'subject': procurement_Details.RequestType + " " + procurement_Details.RequestNumber + " " + 'ITOPEX' + " " + mailTitle
                            }
                            approver_mail_data.append(mail_data_approver)
                            # alertMail(next_approver_mail_id, approver_mail_data[0], cc)
                            return Response(True)
                        else:
                            ApprovalTransaction.objects.filter(procurementId=procurement_id,
                                                               sequence=new_sequence_number, is_active=True).update(
                                status='Pending')

                            # 'cc': requesterEmail + ',' + present_approver_email,
                            # 'cc': 'punya.hc@yokogawa.com' + ',' + present_approver_email,
                            next_approver_details = ApprovalTransaction.objects.filter(procurementId=procurement_id,
                                                                                sequence=new_sequence_number,
                                                                                is_active=True).values_list(
                                'approverEmail', 'approverType')

                            next_approver_mail_id, next_approver_type, mailTitle = None, None, None
                            if len(next_approver_details) != 0:
                                next_approver_mail_id = next_approver_details[0][0]
                                next_approver_type = next_approver_details[0][1]
                                if next_approver_type is not None:
                                    if next_approver_type == 'BuHead':
                                        mailTitle = 'Department Approval'
                                    elif next_approver_type == 'DSINHead':
                                        mailTitle = 'IT Approval'
                                    elif next_approver_type == 'FinanceHead':
                                        mailTitle = 'Finance Approval'
                                    elif next_approver_type == 'MD':
                                        mailTitle = 'Final Approval'

                            cc = None
                            if (requesterEmail is not None) and (present_approver_email is None):
                                cc = requesterEmail
                            elif (requesterEmail is None) and (present_approver_email is not None):
                                cc = present_approver_email
                            elif (requesterEmail is not None) and (present_approver_email is not None):
                                cc = requesterEmail + ',' + present_approver_email
                            approver_mail_data = []
                            mail_data_approver = {
                                'RequestType': procurement_Details.RequestType,
                                'RequestNumber': procurement_Details.RequestNumber,
                                'Name': procurement_Details.Name,
                                'department': procurement_Details.Department,
                                'expenditureBudgeted': procurement_Details.IsExpenditure,
                                'cc': cc,
                                'requested_person_email': requesterEmail,
                                'module': 'Approver',
                                'Heading': 'Your Approval Is Required.',
                                'subject': procurement_Details.RequestType + " " + procurement_Details.RequestNumber + " " + 'ITOPEX' + " " + mailTitle
                            }
                            # approver_mail_data.append(mail_data_approver)
                            alertMail(next_approver_mail_id, approver_mail_data[0], cc)

                            return Response(True)

                    # this is for last Approval Transaction status (eg.DSINMPR person approves)
                    elif not ApprovalTransaction.objects.filter(procurementId=procurement_id,
                                                                sequence=new_sequence_number, is_active=True).exists():
                        if new_sequence_number > total_transaction_count:
                            procurement_instance = MasterProcurement.objects.get(id=instance.procurementId.id)
                            procurement_instance.Status = request.data['status']
                            procurement_instance.save()
                            serializer.save()

                            requester_mail_data = []
                            cc = None
                            cc = present_approver_email
                            mail_data_requester = {
                                'RequestNumber': procurement_Details.RequestNumber,
                                'RequestType': procurement_Details.RequestType,
                                'Name': procurement_Details.Name,
                                'department': procurement_Details.Department,
                                'expenditureBudgeted': procurement_Details.IsExpenditure,
                                'cc': present_approver_email,
                                'requested_person_email': requesterEmail,
                                'module': 'Requester',
                                'subject': procurement_Details.RequestType + " " + procurement_Details.RequestNumber + " " + 'ITOPEX ' + ' MPR Process',
                                'Heading': 'Your Procurment Request Has Been Sent For MPR Creation!!!.'
                            }
                            requester_mail_data.append(mail_data_requester)
                            # alertMail(requesterEmail, requester_mail_data[0], cc)
                            return Response(True)
            return Response(serializer.data)
        else:
            return Response({
                'message': 'Please provide the procurementId and sequenceId',
                'status': 'failed'
            }, status=status.HTTP_400_BAD_REQUEST)

class ApprovalTransactionViewSet(viewsets.ModelViewSet):
    queryset = ApprovalTransaction.objects.all()
    serializer_class = ApprovalTransactionSerializer
    permission_classes = [permissions.IsAuthenticated]

    def list(self, request, *args, **kwargs):
        user = request.user
        queryset = ApprovalTransaction.objects.filter(id=id).order_by('-id')
        serializer = ApprovalTransactionSerializer(queryset, many=True)
        return Response(serializer.data)

    @action(methods=['post'], detail=False, url_path='alert_mail')
    def alertMailFunc(self, request):
        data = request.data

        if data['module'] == "Requester":
            alertMail(data['requested_person'], data, data['cc'])

        elif data['module'] == "BUHead":
            alertMail(data['requested_person'], data, data['cc'])

        elif data['module'] == "DSINHead":
            alertMail(data['requested_person'], data, data['cc'])

        elif data['module'] == "FinanceHead":
            alertMail(data['requested_person'], data, data['cc'])

        elif data['module'] == "MD":
            alertMail(data['requested_person'], data, data['cc'])

        elif data['module'] == "MPRCreation":
            alertMail(data['requested_person'], data, data['cc'])

        return Response(True)


#%%
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
from tabulate import tabulate

import pandas as pd
import math
import keras as k

import datetime
import matplotlib

from keras.models import load_model
from keras.models import Sequential
from keras.layers import Dense
from keras.layers import Activation
from keras.layers import Dropout
from keras.optimizers import Adam
from keras.constraints import maxnorm
from keras.callbacks import EarlyStopping
from keras import backend as B

from sklearn.externals import joblib
from sklearn.metrics import mean_squared_error
from sklearn.metrics import mean_absolute_error
from sklearn.metrics import r2_score
from sklearn.model_selection import TimeSeriesSplit
from sklearn.preprocessing import MinMaxScaler
from sklearn.decomposition import PCA

# Seed for reproducibility
from numpy.random import seed
seed(9)
from tensorflow import set_random_seed
set_random_seed(9)

early_stopper = EarlyStopping(monitor = 'val_loss', patience=10)

# Sensitivity Analysis

def split_dataset(fgroup, model, dataset: np.ndarray, test: np.ndarray):
    """Retrieve the dengue dataset and process the data."""
    # Set defaults.
    batch_size = 32 # not used

    train = dataset

    # Split training and test sets to x and y.
    train_x, train_y = train[:,:-1], train[:, -1]
    test_x, test_y = test[:,:-1], test[:, -1]

    fgroup = str(fgroup)
    model = str(model)
    # Scale test set.
    filenamex = r"C:\Users\Eric\Documents\VSCode\neural-network-genetic-algorithm-master\fgroup" + fgroup + r"\fgroup" + fgroup + "-scalerx" + model + ".save"
    scalerx = joblib.load(filenamex)
    test_xsc = scalerx.transform(test_x)

    # reshape outputs
    test_y = test_y.reshape(-1, 1)

    # Scaled test set.
    filenamey = r"C:\Users\Eric\Documents\VSCode\neural-network-genetic-algorithm-master\fgroup" + fgroup + r"\fgroup" + fgroup + "-scalery" + model + ".save"
    scalery = joblib.load(filenamey)
    test_ysc = scalery.transform(test_y)

    return (test_x, test_xsc, test_y, test_ysc, scalerx, scalery)

def load_dataset(datasource: str, fgroup: int, week: int) -> (np.ndarray):
    """
    The function loads dataset from given file name
    :param datasource: file name of data source
    :return: tuple of dataset and the used MinMaxScaler
    """
    # load the dataset
    dataframe = pd.read_csv(datasource)
    i = week
    if fgroup == 1: # model id 1
        dataset = dataframe[['cases'+str(i),'cases'+str(i+1),'cases'+str(i+2),'cases'+str(i+3),'rainfall11','rainfall13',
        'tmax11','tmax12','tmin9','tmin11','tmean11','tmean12', 'rh5' if i==4  else 'rh3','rh4','wind_speed4','wind_speed3','target']]
    elif fgroup == 2: # model id 4
        dataset = dataframe[['cases'+str(i),'cases'+str(i+1),'cases'+str(i+2),'rainfall11','tmax12','tmin9','tmean11','rh4','wind_speed4','target']]
    elif fgroup == 3: # model id 5
        dataset = dataframe[['cases'+str(i),'cases'+str(i+1),'rainfall11','tmax12','tmin9','tmean11','rh4','wind_speed4','target']]
    elif fgroup == 4: # model id 6
        dataset = dataframe[['cases'+str(i+1),'rainfall11','tmax12','tmin9','tmean11','rh4','wind_speed4','target']]
    dataset = dataset.values
    dataset = dataset.astype('float32')

    return dataset

def main():
    fgroup = 1
    print("Fgroup " + str(fgroup))
    fgrouptbl = []

    model = 3
    week = 1

    print("Model " + str(model))

    datasource = r'C:\Users\Eric\Documents\VSCode\neural-network-genetic-algorithm-master\csvs\full_dataset_extended.csv'
    dataset = load_dataset(datasource, fgroup, week)
    testdatasource = r'C:\Users\Eric\Documents\VSCode\neural-network-genetic-algorithm-master\csvs\test_dataset_extended.csv'
    testdataset = load_dataset(testdatasource, fgroup, week)

    # split dataset
    x_test, x_scaled, y_test, y_scaled, scalerx, scalery = split_dataset(fgroup, model, dataset, testdataset)

    # sens analysis
    values = [0.80  , 0.85, 0.90, 0.95, 1.05, 1.10, 1.15, 1.20]
    fgroup_name = r'C:\Users\Eric\Documents\VSCode\neural-network-genetic-algorithm-master\fgroup' + str(fgroup) + r"\fgroup " + str(fgroup)
    model_name = fgroup_name + " - model" + str(model) + ".h5"
    model = load_model(model_name)

    table = []
    # print(x_test[:, 15])
    para = ""

    # iterate through values
    # og_in = x values, og_out = y values
    for i in range(len(values)):
        og_in = x_test.copy()
        og_out = scalery.inverse_transform(model.predict(x_scaled, verbose=0))
        #  [:, x] replace x with data column to be tested (e.g. 0 = cases1, 1 = cases2, etc.)
        og_in[:, 15] = og_in[:, 15] * values[i]
        new_in = scalerx.transform(og_in)
        new_out = scalery.inverse_transform(model.predict(new_in, verbose=0))
        table.append(percent_change(og_out,new_out))
    send = []
    send.append(table)
    drawtable("para", send)

    # Print basic evaluative data.
    # fgroup_name = r'C:\Users\Eric\Documents\VSCode\neural-network-genetic-algorithm-master\fgroup' + str(fgroup) + r"\fgroup " + str(fgroup)
    # model_name = fgroup_name + " - model" + str(model) + ".h5"
    # model = load_model(model_name)
    # y = scalery.inverse_transform(model.predict(x_scaled, verbose=0))
    # print(str(np.mean(y)) + " " + str(np.std(y)) + " " + str(np.median(y)))

def percent_change(old, new):
    changes = []
    for i in range(len(old)):
        if new[i] == -1:
            continue
        if old[i] != 0:
            changes.append(float(new[i] - old[i]) / old[i] * 100)
    return np.mean(changes)

def drawtable(parameter: str, table: list):

    rows = [parameter]
    # columns = ("week", "rmse", "mae", "r2s")

    df = pd.DataFrame(table, index = rows)
    df.columns = ["0.80", "0.85", "0.90", "0.95", "1.05", "1.10", "1.15", "1.20"]
    df.index.names = ['para']
    df.columns.names = ['scale']
    h = [df.index.names[0] +'/'+ df.columns.names[0]] + list(df.columns)
    print(tabulate(df, headers= h, tablefmt= 'grid'))

main()


class Mensaje:
    def __init__(self,lugar,fecha,hora,user,red, contenido):
        self.lugar=lugar
        self.fecha=fecha
        self.hora = hora
        self.user=user
        self.red =red
        self.contenido = contenido
        #Las primeras son variables de declaracion normal, las otras se declaran con funciones especiales
        self.texto = ""
        self.empresa = 'None'
        self.servicio = 'Desconocido'
        self.positivos = 0
        self.negativos = 0
        self.tipo = 'None'
        self.balancePositivo = 0.0
        self.balanceNegativo = 0.0
        self.Contener()

    def obtenerMensaje(self):
        return self
    #CLASIFICA EN FUNCION SI SON MAYORES, MENORES IGUALES LOS MENSAJES POSITIVOS O NEGATIVOS
    def clasificarMensaje (self):
        if self.positivos > self.negativos:
            self.tipo = 'POSITIVO'
        elif self.negativos > self.positivos:
            self.tipo = 'NEGATIVO'
        elif self.positivos == self.negativos:
            self.tipo = 'NEUTRAL'

    def Contener (self):
        for text in self.contenido:
            self.texto += text
            self.texto += " "

    def balancear(self):
        #Saca el balance de positivos y el balance de negativos
        self.balancePositivo = (float(self.positivos)//(self.positivos+self.negativos))*100
        self.balanceNegativo = (float(self.negativos)//(self.positivos+self.negativos))*100

class servicio:
    def __init__(self, nombre):
        self.nombre = nombre
        self.total = 0
        self.positivo = 0
        self.negativo = 0
        self.neutral = 0

class empresa:
    def __init__(self , nombre):
        self.nombre = nombre
        self.total = 0
        self.positivo = 0
        self.negativo = 0
        self.neutral = 0
        self.servicios = []


#Song XML info getter

#import relevant modules
import csv
import time
import xml.etree.ElementTree as etree
import urllib.request
import time

#open artist CSV
f = open("artists.csv")

#create CSV reader object
reader = csv.reader(f)

#loop through reader object and grab artist names
allArtists = []
for artist in reader:
    allArtists.append(artist[0])

f.close()

xmlCount = 0
failedArtists = []
#loop through all artist names
for artist in allArtists:
    artist = artist.replace(" ","%20")
    #build relevant query for each artist
    query = "http://developer.echonest.com/api/v4/song/search?api_key=H0QIAQ1LC7ELW8V75&format=xml&artist="+artist+"&bucket=audio_summary&results=100"

    #write XML output
    try:
        response = urllib.request.urlopen(query)
        tree = etree.parse(response)
        artist = artist.replace("%20"," ")
        tree.write(artist+".xml", "UTF-8")
        xmlCount += 1
    #if xml generation fails, add artist to failed artist list
    except:
        artist = artist.replace("%20"," ")
        failedArtists.append([artist])

    #wait 3 seconds between each query to avoid hitting EchoNest request limit
    time.sleep(3)

#write out failed artists csv
f = open("failedArtists.csv",'w')
writer = csv.writer(f, lineterminator = '\n')
writer.writerows(failedArtists)
f.close()