#!/usr/bin/env python
# -*- coding: utf-8 -*-

from matplotlib import pyplot as plt
import pandas as pd
import numpy as np
from sklearn.neighbors import KernelDensity
from sklearn.model_selection import GridSearchCV
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras import callbacks

import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow.keras.layers import Dense
from tensorflow.keras.regularizers import l2
from tensorflow.keras import backend as k

from ortools.sat.python import cp_model

import collections
Task = collections.namedtuple('task', 'start end interval ttype')

import copy
from ortools.linear_solver import pywraplp
from scipy import sparse

status_string = {
    cp_model.OPTIMAL: 'optimal',
    cp_model.FEASIBLE: 'feasible',
    cp_model.INFEASIBLE: 'infeasible',
    cp_model.MODEL_INVALID: 'invalid',
    cp_model.UNKNOWN: 'unknown'
}

# Configuration
figsize = (9, 3)

def load_raw_data(data_folder):
    # Read the CSV files
    fnames = ['2018', '2019']
    datalist = []
    for fstem in fnames:
        # Read data and append
        data = pd.read_csv(f'{data_folder}/{fstem}.csv', sep='|')
        data['Anno'] = int(fstem)
        data = data[['Anno'] + data.columns[:-1].tolist()]
        datalist.append(data)
    # Concatenate
    data = pd.concat(datalist)
    return data

def translate_data(data_folder):
    # Read the CSV files
    fnames = ['2018', '2019']
    datalist = []
    for fstem in fnames:
        # Read data and append
        data = pd.read_csv(f'{data_folder}/{fstem}.csv', sep='|')
        data['Anno'] = int(fstem)
        data = data[['Anno'] + data.columns[:-1].tolist()]
        datalist.append(data)
    # Concatenate
    data = pd.concat(datalist)
    # Translate the "CODICE" entry
    data['CODICE'].replace('VERDE', 'green', inplace=True)
    data['CODICE'].replace('BIANCO', 'white', inplace=True)
    data['CODICE'].replace('ROSSO', 'red', inplace=True)
    data['CODICE'].replace('GIALLO', 'yellow', inplace=True)
    # Translate the "CODICE" entry
    data['Esito'].replace('Ammessi', 'admitted', inplace=True)
    data['Esito'].replace('Abbandoni', 'abandoned', inplace=True)
    # Translate the "path" entry
    data['Percorso'] = data['Percorso'].str.replace('"Visita"', 'visit')
    data['Percorso'] = data['Percorso'].str.replace('"LABORATORIO"', 'lab')
    data['Percorso'] = data['Percorso'].str.replace('"VISITA ORTOPEDICA"', 'orthopaedic visit')
    data['Percorso'] = data['Percorso'].str.replace('"Prescrizione Terapia"', 'prescription')
    data['Percorso'] = data['Percorso'].str.replace('"VISITA OTORINOLARINGOIATRICA"', 'otolaryngological visit')
    data['Percorso'] = data['Percorso'].str.replace('"ECOGRAFIA"', 'ultrasound')
    data['Percorso'] = data['Percorso'].str.replace('"TAC"', 'CT scan')
    data['Percorso'] = data['Percorso'].str.replace('"VISITA NEUROLOGICA"', 'neurological visit')
    data['Percorso'] = data['Percorso'].str.replace('"CONSULENZA CHIRURGICA VASCOLARE"', 'vascular surgery consultation')
    data['Percorso'] = data['Percorso'].str.replace('"CORONAROGRAFIA"', 'coronarography')
    data['Percorso'] = data['Percorso'].str.replace('"ESAME COMPLESSIVO DELL\'OCCHIO"', 'eye examination')
    data['Percorso'] = data['Percorso'].str.replace('"VISITA CHIRURGICA"', 'surgical visit')
    data['Percorso'] = data['Percorso'].str.replace('"VISITA NEUROPSICHIATRICA INFANTILE"', 'child psychiatrist visit')
    data['Percorso'] = data['Percorso'].str.replace('"ORTOPANTOMOGRAFIA DIGITALE"', 'digital orthopantomography')
    data['Percorso'] = data['Percorso'].str.replace('"VISITA UROLOGICA"', 'urological visit')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIO CAROTIDE COMUNE SN."', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIO CAROTIDE INTERNA SN"', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIO CEREBRALE PRIMO  VASO"', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"TROMBOLISI"', 'thrombolysis')
    data['Percorso'] = data['Percorso'].str.replace('"ARTERIOGRAFIA SELETTIVA"', 'arteriography')
    data['Percorso'] = data['Percorso'].str.replace('"VISITA DIABETOLOGICA"', 'diabetological visit')
    data['Percorso'] = data['Percorso'].str.replace('"ESECUZIONE CD / COPIE CD"', 'CD execution and copy')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIO CAROTIDE COMUNE DS."', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIO CAROTIDE INTERNA DS."', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"PROVA DA SFORZO"', 'stress test')
    data['Percorso'] = data['Percorso'].str.replace('"VISITA DERMATOLOGICA"', 'dermatological visit')
    data['Percorso'] = data['Percorso'].str.replace('"VISITA CARDIOLOGICA"', 'cardiological visit')
    data['Percorso'] = data['Percorso'].str.replace('"INVIO/SPEDIZIONE ESAME AD ALTRO OSPEDALE"', 'sending results to a second hospital')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIO \(VASI NON CODIFICATI\)"', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"POSIZIONAMENTO DI CVC IN SAL"', 'CVC positioning')
    data['Percorso'] = data['Percorso'].str.replace('"ELETTROENCEFALOGRAMMA"', 'EEG')
    data['Percorso'] = data['Percorso'].str.replace('"SIGMOIDOSCOPIA CON ENDOSCOPIO FLESSIBILE"', 'sigmodoscopy')
    data['Percorso'] = data['Percorso'].str.replace('"VISITA INFETTIVOLOGICA"', 'infectious disease visit')
    data['Percorso'] = data['Percorso'].str.replace('"EMBOLIZZAZIONE"', 'embolization')
    data['Percorso'] = data['Percorso'].str.replace('"EGDS CON BIOPSIA"', 'EGDS')
    data['Percorso'] = data['Percorso'].str.replace('"PEG - POSIZIONAMENTO"', 'PEG')
    data['Percorso'] = data['Percorso'].str.replace('"PEG - RIMOZIOE"', 'PEG')
    data['Percorso'] = data['Percorso'].str.replace('"2\^ VISITA NEUROLOGICA"', 'neurological visit')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIO  ARCO E VASI EPIAORTICI"', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIO CAROTIDE ESTERNA DS."', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIO CAROTIDE ESTERNA SN."', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIO PERCUTANEA \(PER VASO\)"', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIO VERTEBRALE DS."', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIO VERTEBRALE SN."', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIO-RM CEREBRALE ARTERIOSA"', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIO-RM CIRCOLO ARTERIOSO INTRA-CRANICO"', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIOGRAFIA RICOSTRUZIONI"', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIOGRAFIA ROTATORIA"', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIOPLASTICA PERCUTANEA CAROTIDI"', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIOPLASTICA PERCUTANEA CORONARICA MULTIVASALE"', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIOPLASTICA PERCUTANEA ILIACA"', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIOPLASTICA PERCUTANEA"', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"ANGIOPLASTICA PERCUTANEA"', 'angiography')
    data['Percorso'] = data['Percorso'].str.replace('"ARCHIVIAZIONE DIGITALE PRECEDENTI"', 'medical record digitalization')
    data['Percorso'] = data['Percorso'].str.replace('"ARTERIOGRAFIA ARTO INF. DX"', 'arteriography')
    data['Percorso'] = data['Percorso'].str.replace('"ARTERIOGRAFIA ARTO INF. SX"', 'arteriography')
    data['Percorso'] = data['Percorso'].str.replace('"ARTERIOGRAFIA SELETTIVA ARTERIE BRONCHIALI"', 'arteriography')
    data['Percorso'] = data['Percorso'].str.replace('"ARTERIOGRAFIA SELETTIVA RENALE"', 'arteriography')
    data['Percorso'] = data['Percorso'].str.replace('"CISTOGRAFIA"', 'cystography')
    data['Percorso'] = data['Percorso'].str.replace('"COLLOQUIO PSICHIATRICO"', 'psychiatrict counseling')
    data['Percorso'] = data['Percorso'].str.replace('"COLONSCOPIA"', 'colonscopy')
    data['Percorso'] = data['Percorso'].str.replace('"CONSULENZA NEURORADIOLOGICA"', 'neuroradiological consultation')
    data['Percorso'] = data['Percorso'].str.replace('"CONSULENZA RADIOLOGICA"', 'radiological consultation')
    data['Percorso'] = data['Percorso'].str.replace('"CONSULENZA RIANIMATORIA"', 'IC consultation')
    data['Percorso'] = data['Percorso'].str.replace('"CURA ODONTOIATRICA \(BERETTA\)"', 'dental care')
    data['Percorso'] = data['Percorso'].str.replace('"DRENAGGIO BILIARE ESTERNO"', 'biliary drainage')
    data['Percorso'] = data['Percorso'].str.replace('"ELETTROENCEFALOGRAMMA CON VIDEOREGISTRAZIONE"', 'EEG')
    data['Percorso'] = data['Percorso'].str.replace('"EMBOLIZZAZIONE CEREB. \(SPIRALI MAX.N.3\)"', 'embolization')
    data['Percorso'] = data['Percorso'].str.replace('"EMBOLIZZAZIONE COLLO-MASS. FACC. \(COLLE O PART.\)"', 'embolization')
    data['Percorso'] = data['Percorso'].str.replace('"EMBOLIZZAZIONE COLLO-MASS. FACC. \(SPIRALI\)"', 'embolization')
    data['Percorso'] = data['Percorso'].str.replace('"ESAME DEL FUNDUS OCULI"', 'eye examination')
    data['Percorso'] = data['Percorso'].str.replace('"GATED SPECT"', 'gated spect')
    data['Percorso'] = data['Percorso'].str.replace('"NEFROPIELOSTOMIA PERCUTANEA"', 'nephrolytotomy')
    data['Percorso'] = data['Percorso'].str.replace('"PEG - RIMOZIONE"', 'PEG')
    data['Percorso'] = data['Percorso'].str.replace('"PIELOGRAFIA TRANSPIELOSTOMICA DX"', 'pyelography')
    data['Percorso'] = data['Percorso'].str.replace('"PIELOGRAFIA TRANSPIELOSTOMICA SX"', 'pyelography')
    data['Percorso'] = data['Percorso'].str.replace('"POSIZIONAMENTO DI CVC IN SALA"', 'CVC positioning')
    data['Percorso'] = data['Percorso'].str.replace('"POSIZIONAMENTO DI CVC"', 'CVC positioning')
    data['Percorso'] = data['Percorso'].str.replace('"RM ADDOME INFERIORE"', 'MRI')
    data['Percorso'] = data['Percorso'].str.replace('"RM ADDOME SUPERIORE"', 'MRI')
    data['Percorso'] = data['Percorso'].str.replace('"RM DEL BACINO SMDC"', 'MRI')
    data['Percorso'] = data['Percorso'].str.replace('"RM DIFFUSIONE CEREBRALE"', 'MRI')
    data['Percorso'] = data['Percorso'].str.replace('"RM ENCEFALO CON MDC"', 'MRI')
    data['Percorso'] = data['Percorso'].str.replace('"RM ENCEFALO SMDC"', 'MRI')
    data['Percorso'] = data['Percorso'].str.replace('"RM RACHIDE CERVICALE CON MDC"', 'MRI')
    data['Percorso'] = data['Percorso'].str.replace('"RM RACHIDE CERVICALE SMDC"', 'MRI')
    data['Percorso'] = data['Percorso'].str.replace('"RM RACHIDE DORSALE CON MDC"', 'MRI')
    data['Percorso'] = data['Percorso'].str.replace('"RM RACHIDE DORSALE SMDC"', 'MRI')
    data['Percorso'] = data['Percorso'].str.replace('"RM RACHIDE LOMBARE CON MDC"', 'MRI')
    data['Percorso'] = data['Percorso'].str.replace('"RM RACHIDE LOMBARE SMDC"', 'MRI')
    data['Percorso'] = data['Percorso'].str.replace('"SCOPIA INTRAOPERATORIA  60 MIN."', 'intraoperative scopy')
    data['Percorso'] = data['Percorso'].str.replace('"SPECT MIOCARDICA RIPOSO"', 'myocardial SPECT')
    data['Percorso'] = data['Percorso'].str.replace('"STENT BILIARE \(X PEZZO\)"', 'biliary stent')
    data['Percorso'] = data['Percorso'].str.replace('"STENT RENALI.F114"', 'renal stent')
    data['Percorso'] = data['Percorso'].str.replace('"STUDIO ETA\' OSSEA"', 'bone age estimation')
    data['Percorso'] = data['Percorso'].str.replace('"TROMBECTOMIA"', 'thrombectomy')
    data['Percorso'] = data['Percorso'].str.replace('"VISITA ANESTESIOLOGICA"', 'anesthesiological visit')
    data['Percorso'] = data['Percorso'].str.replace('"VISITA ANGIOLOGICA"', 'angiological visit')
    data['Percorso'] = data['Percorso'].str.replace('"VISITA DIETOLOGICA"', 'diet visit')
    data['Percorso'] = data['Percorso'].str.replace('"VISITA GASTROENTEROLOGICA"', 'gastroenterological visit')
    data['Percorso'] = data['Percorso'].str.replace('"VISITA MAXILLO FACCIALE"', 'maxillofacial visit')
    data['Percorso'] = data['Percorso'].str.replace('"VISITA ODONTOIATRICA"', 'dental visit')
    data['Percorso'] = data['Percorso'].str.replace('"VISITA PNEUMOLOGICA"', 'pneumological visit')
    data['Percorso'] = data['Percorso'].str.replace('"Triage"', 'triage')
    data['Percorso'] = data['Percorso'].str.replace('"RX"', 'RX')
    data['Percorso'] = data['Percorso'].str.replace('"EGDS"', 'EGDS')
    data['Percorso'] = data['Percorso'].str.replace(', ', ',')
    # Fix NaN flows
    data['Percorso'].fillna('[]', inplace=True)
    # # Split flow strings
    # data['Percorso'] = data['Percorso'].apply(lambda s: s[1:-1].split(',') if s != np.NaN else np.NaN)
    # Change the column names
    data.columns = ['year', 'ID', 'Triage', 'TkCharge', 'Code', 'Outcome', 'Flow']
    # Write back
    data.to_csv(f'{data_folder}/er.csv', sep=';', index=False)

def load_data(data_folder):
    # Read the CSV file
    data = pd.read_csv(f'{data_folder}/er.csv', sep=';', parse_dates=[2, 3])
    # Convert a few fields to categorical format
    data['Code'] = data['Code'].astype('category')
    data['Outcome'] = data['Outcome'].astype('category')
    # Sort by triage time
    data.sort_values(by='Triage', inplace=True)
    # Discard the firl
    return data

def plot_bars(data, figsize=figsize, autoclose=True, tick_gap=1,
        series=None):
    if autoclose: plt.close('all')
    plt.figure(figsize=figsize)
    # x = np.arange(len(data))
    # x = 0.5 + np.arange(len(data))
    # plt.bar(x, data, width=0.7)
    # x = data.index-0.5
    x = data.index
    plt.bar(x, data, width=0.7)
    # plt.bar(x, data, width=0.7)
    if series is not None:
        # plt.plot(series.index-0.5, series, color='tab:orange')
        plt.plot(series.index, series, color='tab:orange')
    if tick_gap > 0:
        plt.xticks(x[::tick_gap], data.index[::tick_gap], rotation=45)
    plt.tight_layout()

def plot_series(series, std=None,
        figsize=figsize,
        autoclose=True, s=0):
    if autoclose:
        plt.close('all')
    plt.figure(figsize=figsize)
    if std is not None:
        lb = series.values.ravel() - std.values.ravel()
        ub = series.values.ravel() + std.values.ravel()
        plt.fill_between(std.index, lb, ub, alpha=0.3, label='+/- std')
    if s > 0:
        plt.plot(series.index, series, label='data')
    else:
        plt.plot(series.index, series, label='data',
                marker='.', markersize=s)
    if std is not None:
        plt.legend()
    plt.tight_layout()

def plot_pred_scatter(y_pred, y_true, figsize=figsize, autoclose=True):
    if autoclose:
        plt.close('all')
    plt.figure(figsize=figsize)
    plt.scatter(y_pred, y_true, marker='.', alpha=0.1)
    xl, xu = plt.xlim()
    yl, yu = plt.ylim()
    l, u = min(xl, yl), max(xu, yu)
    plt.plot([l, u], [l, u], ':', c='0.3')
    plt.xlim(l, u)
    plt.ylim(l, u)
    plt.xlabel('prediction')
    plt.ylabel('target')
    plt.tight_layout()

def plot_training_history(history,
        figsize=figsize, autoclose=True):
    if autoclose:
        plt.close('all')
    plt.figure(figsize=figsize)
    plt.plot(history.history['loss'], label='loss')
    if 'val_loss' in history.history.keys():
        plt.plot(history.history['val_loss'], label='val. loss')
        plt.legend()
    plt.tight_layout()

def sliding_window_2D(data, wlen, stride=1):
    # Get shifted tables
    m = len(data)
    lt = [data.iloc[i:m-wlen+i+1:stride, :].values for i in range(wlen)]
    # Reshape to add a new axis
    s = lt[0].shape
    for i in range(wlen):
        lt[i] = lt[i].reshape(s[0], 1, s[1])
    # Concatenate
    wdata = np.concatenate(lt, axis=1)
    return wdata

def plot_autocorrelation(data, max_lag=100, figsize=figsize):
    # Open a new figure
    plt.close('all')
    plt.figure(figsize=figsize)
    # Autocorrelation plot
    pd.plotting.autocorrelation_plot(data)
    # Customized x limits
    plt.xlim(0, max_lag)
    # Rotated x ticks
    plt.xticks(rotation=45)
    plt.tight_layout()

def flow_to_levels(fstring):
    # Parse all the activity names
    activities = fstring[1:-1].split(',')
    # Partition into levels
    levels = [[]]
    for a in activities:
        if a == 'triage':
            pass
        elif a == 'visit':
            if len(levels[-1]) == 0:
                levels[-1].append(a)
            else:
                levels.append([a])
            levels.append([])
        else:
            levels[-1].append(a)
    return levels[:-1]

def get_dur(ttype):
    durs = {'visit': 1, 'ultrasound': 2, 'RX': 2}
    return durs[ttype] if ttype in durs else 4

def get_level_dur(levels):
    res = 0
    for lvl in levels:
        res += sum(get_dur(a) for a in lvl)
    return res

def get_horizon(levels, setups=None):
    eoh = sum(get_level_dur(l) for l in levels.values())
    return eoh

def levels_to_vars(mdl, levels, idx, tasks, last, rl, dl):
    # Build a variable for each activity
    for k, lvl in enumerate(levels):
        for i, a in enumerate(lvl):
            # Obtain duration and build variables
            dur = get_dur(a)
            start = mdl.NewIntVar(rl, dl, f's_{idx}_{k}_{i}')
            end = mdl.NewIntVar(rl, dl, f'e_{idx}_{k}_{i}')
            interval = mdl.NewIntervalVar(start, dur, end, f'i_{idx}_{k}_{i}')
            # Store
            tasks[(idx,k,i)] = Task(start, end, interval, a)
    last[idx] = Task(start, end, interval, a)

def add_all_vars(mdl, levels, rl, dl):
    tasks, last = {}, {}
    for idx, lvl in levels.items():
        levels_to_vars(mdl, lvl, idx, tasks, last, rl[idx], dl=dl[idx])
    return tasks, last

def build_levels(fdata):
    levels, codes_by_idx = {}, {}
    for idx in fdata.index:
        levels[idx] = flow_to_levels(fdata.loc[idx]['Flow'])
        codes_by_idx[idx] = fdata.loc[idx]['Code']
    return levels, codes_by_idx

def print_solution(slv, levels, tasks, codes_by_idx):
    # Collect tasks by idx
    tasks_by_idx = {idx: [] for idx in levels.keys()}
    for (idx, k, i), t in tasks.items():
        tasks_by_idx[idx].append(t)
    # Sort by increasing start time
    strings_by_idx = {}
    starts_by_idx = {}
    for idx in tasks_by_idx.keys():
        ttasks = sorted(tasks_by_idx[idx], key=lambda t: slv.Value(t.start))
        starts_by_idx[idx] = slv.Value(ttasks[0].start)
        stasks = [f'{t.ttype}({slv.Value(t.start)}-{slv.Value(t.end)})' for t in ttasks]
        stasks = ', '.join(stasks)
        strings_by_idx[idx] = f'{idx}({codes_by_idx[idx]}): {stasks}'
    # Sort idx by starting time
    sorted_idx = sorted([k for k in starts_by_idx], key=lambda k: starts_by_idx[k])
    for idx in sorted_idx:
        print(strings_by_idx[idx])

def print_outcome(slv, levels, tasks, codes_by_idx, status, verbose=1):
    print(f'Solver status: {status_string[status]}', end='')
    print(f', time(CPU sec): {slv.UserTime():.2f}', end='')
    if status in (cp_model.OPTIMAL, cp_model.FEASIBLE):
        print(f', objective: {slv.ObjectiveValue()}')
        if verbose > 0:
            print()
            print_solution(slv, levels, tasks, codes_by_idx)
    else:
        print()

def add_precedences(mdl, levels, idx, tasks):
    for k, lvl in enumerate(levels[:-1]):
        for i, _ in enumerate(lvl):
            for j, _ in enumerate(levels[k+1]):
                mdl.Add(tasks[idx,k,i].end <= tasks[idx,k+1,j].start)

def add_all_precedences(mdl, levels, tasks):
    for idx, lvl in levels.items():
        add_precedences(mdl, lvl, idx, tasks)

def add_makespan_variables(mdl, levels, codes_by_idx, last, eoh):
    codes = ['red', 'yellow', 'green', 'white']
    last_by_code = {c: [] for c in codes}
    for idx in levels.keys():
        last_by_code[codes_by_idx[idx]].append(last[idx])

    obj_by_code = {}
    for code in codes:
        obj_by_code[code] = mdl.NewIntVar(0, eoh, f'mk_{code}')
        if len(last_by_code[code]) > 0:
            mdl.AddMaxEquality(obj_by_code[code], [t.end for t in last_by_code[code]])
    obj_by_code['all'] = mdl.NewIntVar(0, eoh, f'mk_global')
    mdl.AddMaxEquality(obj_by_code['all'], [obj_by_code[c] for c in codes])
    return obj_by_code

def add_waittime_variables(mdl, levels, codes_by_idx, tasks, eoh):
    codes = ['red', 'yellow', 'green', 'white']
    fbc = {c: [] for c in codes}
    for idx in levels.keys():
        fbc[codes_by_idx[idx]].append(tasks[idx, 0, 0])

    obj_by_code = {}
    for code in codes:
        obj_by_code[code] = mdl.NewIntVar(0, len(codes_by_idx)*eoh, f'wt_{code}')
        if len(fbc[code]) > 0:
            mdl.Add(obj_by_code[code] ==
                    sum(t.start for t in fbc[code]))
    obj_by_code['all'] = mdl.NewIntVar(0, len(codes_by_idx) * eoh, f'wt_globai')
    mdl.Add(obj_by_code['all'] ==  sum(obj_by_code[c] for c in codes))
    return obj_by_code

def add_cumulatives(mdl, tasks, capacities):
    res_intervals = {r: [] for r in capacities.keys()}
    for task in tasks.values():
        if task.ttype in capacities.keys():
            res_intervals[task.ttype].append(task.interval)

    for rtype, cap in capacities.items():
        intervals = res_intervals[rtype]
        mdl.AddCumulative(intervals, [1]*len(intervals), cap)

def add_no_overlap(mdl, levels, idx, tasks):
    intervals = []
    for k, lvl in enumerate(levels):
        for i, _ in enumerate(lvl):
            intervals.append(tasks[idx, k, i].interval)
    mdl.AddNoOverlap(intervals)

def add_all_no_overlap(mdl, levels, tasks):
    for idx, lvl in levels.items():
        add_no_overlap(mdl, lvl, idx, tasks)

def plot_scalability_evalulation(plist, tlist, plist_cmp=None, tlist_cmp=None,
                                 lbl=None, lbl_cmp=None,
                                 figsize=figsize, autoclose=True):
    if autoclose:
        plt.close('all')
    plt.figure(figsize=figsize)
    plt.plot(plist, tlist, label=lbl)
    if plist_cmp is not None:
        plt.plot(plist_cmp, tlist_cmp, label=lbl_cmp)
    if lbl is not None and lbl_cmp is not None:
        plt.legend()
    # sl2 = [f'{p}({status_string[s]})' for p, s in zip(plist, slist)]
    # plt.xticks(plist, sl2, rotation=45)
    plt.tight_layout()

def get_seq_hints(levels, codes_by_idx):
    codes = ['red', 'yellow', 'green', 'white']
    hints = {}
    totdur = 0
    for idx, lvls in sorted(levels.items(),
            key=lambda t: codes.index(codes_by_idx[t[0]])):
        for k, lvl in enumerate(lvls):
            for i, a in enumerate(lvl):
                hints[idx,k,i] = totdur
                totdur += get_dur(a)
    return hints

def solve_bounded_waittime_problem(levels, codes_by_idx, codes, capacities,
        ub_by_code={}, tlim=None, hints=None, aplus=None):
    mdl = cp_model.CpModel()
    eoh = get_horizon(levels)
    rl, dl = {idx:0 for idx in levels}, {idx:eoh for idx in levels}
    tasks, last = add_all_vars(mdl, levels, rl, dl)
    add_all_precedences(mdl, levels, tasks)

    # Add extra precedences
    if aplus is not None:
        for (idx, k, i), out_arcs in aplus.items():
            for idx2, k2, i2 in out_arcs:
                mdl.Add(tasks[idx,k,i].end <= tasks[idx2,k2,i2].start)

    add_cumulatives(mdl, tasks, {'visit': 3, 'ultrasound': 2, 'RX': 2})
    add_all_no_overlap(mdl, levels, tasks)
    obj_by_code = add_waittime_variables(mdl, levels, codes_by_idx, tasks, eoh)

    for code in ub_by_code: # <-- Enforce upper bounds
        mdl.Add(obj_by_code[code] <= ub_by_code[code])

    # Add hints
    if hints is not None:
        for (idx, k, i), stval in hints.items():
            mdl.AddHint(tasks[idx,k,i].start, stval)

    mdl.Minimize(obj_by_code[codes[-1]]) # <-- Focus by default on the last code
    slv = cp_model.CpSolver()
    if tlim is not None: slv.parameters.max_time_in_seconds = tlim
    status = slv.Solve(mdl)
    return status, slv, tasks

def goal_programming(levels, codes_by_idx, codes, capacities,
        tlim=None, verbose=1, hints=None, ub_by_code=None, aplus=None):
    # Handle the upper bounds
    if ub_by_code is None:
        wt_by_code = {}
    else:
        wt_by_code = copy.deepcopy(ub_by_code)
    ttime = 0
    for i, code in enumerate(codes):
        l_tlim = None if tlim is None else (tlim-ttime) / (len(codes)-i)
        status, slv, tasks = solve_bounded_waittime_problem(levels, codes_by_idx,
                codes[:i+1], capacities, ub_by_code=wt_by_code, tlim=l_tlim,
                hints=hints, aplus=aplus)
        ttime += slv.UserTime()
        if verbose:
            print_outcome(slv, levels, tasks, codes_by_idx, status, verbose=(i == len(codes)-1))
        if status in (cp_model.OPTIMAL, cp_model.FEASIBLE):
            wt_by_code[code] = int(slv.ObjectiveValue())
            starts = {tidx:slv.Value(tasks[tidx].start) for tidx in tasks}
            hints = starts
        else:
            wt_by_code, starts = None, None
            break
    if wt_by_code is not None:
        wt_by_code['all'] = sum(wt_by_code[code] for code in codes)
    return ttime, wt_by_code, starts

def scheduling_lns(levels, codes_by_idx, codes, capacities,
        tlim, init_time, it_time, nb_size, verbose=1):
    # Build initial solution
    ttime, wt_by_code, starts = goal_programming(levels, codes_by_idx, codes,
            capacities=capacities,
                               tlim=init_time, verbose=0)
    if verbose:
        print(f'Initial solution in {ttime:.2f} sec, {wt_by_code}')

    # Loop until the time limit is not over
    patients = [idx for idx in codes_by_idx.keys()]
    np.random.seed(42)
    while tlim - ttime > 0:
        aplus, aminus = sol_to_pos(levels, starts, capacities) # Obtain a POS
        relaxed = np.random.choice(patients, nb_size, replace=False) # Relax patients
        for idx in relaxed:
            aplus, aminus = remove_patient_from_pos(aplus, aminus, idx)
        ub_by_code = copy.deepcopy(wt_by_code) # Require an improvement
        ub_by_code['all'] -= 1
        # Re-solve
        itlim = min(it_time, tlim - ttime)
        itime, iwt_by_code, istarts = goal_programming(levels, codes_by_idx, codes,
                            capacities=capacities, tlim=itlim, verbose=0, aplus=aplus,
                            ub_by_code=ub_by_code)
        ttime += max(itime, 0.1) # Update the time limit
        # Update the best solution
        if iwt_by_code is not None:
            if verbose:
                objs = ','.join(f'{iwt_by_code[c]}' for c in codes + ['all'])
                print(f'Total time: {ttime:.2f}, neighborhood explored in {itime:.2f} sec, waiting times {objs}')
            wt_by_code = iwt_by_code
            starts = istarts
    return ttime, wt_by_code, starts

# def goal_programming(levels, codes_by_idx, codes, capacities,
#                      tlim=None, verbose=2, aplus={}, ub_by_code=None, hints=None):
#     if ub_by_code is None:
#         wt_by_code = {}
#     else:
#         wt_by_code = copy.deepcopy(ub_by_code)
#     ttime = 0
#     # Start the main loop
#     for i, code in enumerate(codes):
#         # Determine what to print
#         if verbose > 0:
#             if code == codes[-1]: vrb = 2
#             else: vrb = 1
#         else:
#             vrb = 0
#         # Determine the time limit
#         ctlim = None
#         if tlim is not None:
#             if ub_by_code is not None:
#                 ctlim = tlim-ttime
#             else:
#                 ctlim = (tlim-ttime) / (len(codes)-i)
#         # Solve
#         status, slv, tasks = solve_bounded_waittime_problem(levels, codes_by_idx,
#                 codes[:i+1], capacities, ub_by_code=wt_by_code, tlim=ctlim,
#                 hints=hints, aplus=aplus)
#         ttime += slv.UserTime()
#         if vrb:
#             print_outcome(slv, levels, tasks, codes_by_idx, status, verbose=vrb)
#         if status in (cp_model.FEASIBLE, cp_model.OPTIMAL):
#             starts = {tidx:slv.Value(tasks[tidx].start) for tidx in tasks}
#             hints = starts
#             wt_by_code[code] = int(slv.ObjectiveValue())
#         else:
#             starts = None
#             wt_by_code = None
#             break
#     if wt_by_code is not None:
#         wt_by_code['all'] = sum(wt_by_code[code] for code in codes)
#     return ttime, wt_by_code, starts

def unit_flow(levels, stasks, starts, aplus, aminus):
    nonprocessed = []
    src, dur = None, 0
    for dst in stasks:
        if src is not None and starts[src] + dur <= starts[dst]:
            # Build an arc
            aplus[src].append(dst)
            aminus[dst].append(src)
            # Reset the source
            src = None
        if src is None:
            # Set a new source
            src = dst
            dur = get_dur(levels[src[0]][src[1]][src[2]])
        else:
            # Store as non-processed
            nonprocessed.append(dst)
    return nonprocessed

def sol_to_pos(levels, starts, capacities):
    # Prepare data structures to store the graph
    aplus, aminus = {}, {}
    # Split tasks by resource
    stasks = {r:[] for r in capacities}
    for idx, k, i in starts:
        ttype = levels[idx][k][i]
        if ttype in capacities:
            stasks[ttype].append((idx, k, i))
            aplus[idx, k, i] = []
            aminus[idx, k, i] = []
    # Sort all collections by increasing start
    for res in capacities:
        stasks[res] = sorted(stasks[res], key=lambda t: starts[t])
    # Loop over all resources
    for res, cap in capacities.items():
        for _ in range(cap):
            # Route one flow unit
            stasks[res] = unit_flow(levels, stasks[res], starts, aplus, aminus)
    return aplus, aminus

def validate_pos(aplus, aminus):
    for src in aplus:
        for dst in aplus[src]:
            assert(src in aminus[dst])
    for dst in aminus:
        for src in aminus[dst]:
            assert(dst in aplus[src])

def remove_task_from_pos(aplus, aminus, task_key):
    # Process all predecessors
    for src in aminus[task_key]:
        # Process all successors
        for dst in aplus[task_key]:
            # Transfer arcs
            aminus[dst].append(src)
            aplus[src].append(dst)
    # Remove ingoing arcs
    for src in aminus[task_key]:
        aplus[src].remove(task_key)
    # Remove outgoing arcs
    for dst in aplus[task_key]:
        aminus[dst].remove(task_key)
    # Remove the node from the arc
    del aplus[task_key]
    del aminus[task_key]

def remove_patient_from_pos(aplus, aminus, target_idx):
    aplus_res = copy.deepcopy(aplus)
    aminus_res = copy.deepcopy(aminus)
    for idx, k, i in aplus:
        if idx == target_idx:
            remove_task_from_pos(aplus_res, aminus_res, (idx, k, i))
    return aplus_res, aminus_res

def generate_compatibility_graph(size, seed=42, icp_chance=0.05):
    # Blood type prevalence
    brp = {'O+':.385, 'A+':.355, 'B+':.075, 'AB+':.025,
           'O-':.070, 'A-':.060, 'B-':.015, 'AB-':.015}
    # Compatibility matrix (receipient, donor)
    cpm = {('O-', 'O-'),
            ('O+', 'O-'), ('O+', 'O+'),
            ('A-', 'O-'), ('A-', 'A-'),
            ('A+', 'O-'), ('A+', 'O+'), ('A+', 'A-'), ('A+', 'A+'),
            ('B-', 'O-'), ('B-', 'B-'),
            ('B+', 'O-'), ('B+', 'O+'), ('B+', 'B-'), ('B+', 'B+'),
            ('AB-', 'O-'), ('AB-', 'A-'), ('AB-', 'B-'), ('AB-', 'AB-'),
            ('AB+', 'O-'), ('AB+', 'O+'), ('AB+', 'A-'), ('AB+', 'A+'),
                ('AB+', 'B-'), ('AB+', 'B+'), ('AB+', 'AB-'), ('AB+', 'AB+'),
            }
    # Build an incompatibility matrix
    icpm = {(b1, b2) for b1 in brp.keys() for b2 in brp.keys()}
    icpm -= cpm
    icpm = sorted(list(icpm))
    # Define the probability of each pair
    icpm_p = np.array([brp[t[0]]*brp[t[1]] for t in icpm])
    icpm_p /= icpm_p.sum()
    # Generate incompatible pairs
    rng = np.random.default_rng(seed)
    pair_idxs = rng.choice(range(len(icpm)), size, p=icpm_p)
    Pair = collections.namedtuple('pair', 'recipient donor')
    pairs = {i:Pair(*icpm[j]) for i, j in enumerate(pair_idxs)}
    # Generate basic compatibility arcs
    arcs = [(i, j) for i, pd in pairs.items()
                   for j, pr in pairs.items()
                   if (pr.recipient, pd.donor) in cpm]
    # Remove some arcs to simulate other sources of incompatibility
    aidx = np.arange(len(arcs))
    np.random.seed(seed)
    mask = rng.choice([False, True], len(aidx),
            p=[icp_chance, 1-icp_chance])
    aidx = aidx[mask]
    arcs = [arcs[i] for i in aidx]
    # Convert the graph in forward star and sparse matrix format
    aplus = {i:[] for i in pairs.keys()}
    for src, dst in arcs:
        aplus[src].append(dst)
    # Convert the graph in to a sparse matrix
    return pairs, arcs, aplus

# def shortest_cycles2(aplus, weights, max_len, cap=None):
#     # Build an adjacency matrix
#     nnodes = len(weights)
#     adj = np.full((nnodes, nnodes), np.inf)
#     for i, alist in aplus.items():
#         for j in alist:
#             adj[i,j] = 1
#     # Prepare a list with candidate roots
#     roots = set(range(nnodes))
#     cycles = []
#     ccosts = []
#     processed = set()
#     while len(roots) > 0:
#         root = roots.pop()
#         tcl, tct = shortest_cycles_from_root2(root, adj, weights, max_len)
#         cycles += tcl
#         ccosts += tct
#         # Stop if at full capcity
#         if cap is not None and len(cycles) >= cap:
#             return cycles[:cap]
#         # Mark the nodes in the cycles as processed
#         for cl in tcl:
#             for i in cl:
#                 if i in roots:
#                     roots.remove(i)
#     return cycles, ccosts

# def shortest_cycles_from_root2(root, adj, weights, max_len):
#     nnodes = len(weights)
#     # Initial distance
#     dst = np.full(nnodes, np.inf)
#     dst[root] = 0
#     # Initial predecessor along the shortest path
#     prd = np.full((nnodes, max_len), -1)
#     # Time-unfolded graph traversal
#     idx = np.arange(nnodes)
#     cycles = []
#     ccosts = []
#     for k in range(max_len):
#         # Compute distances from all predecessors
#         # print(root, k)
#         # print(dst)
#         prd_dst = (adj * weights).T + dst
#         prd_dst = np.nan_to_num(prd_dst, nan=np.inf,
#                 posinf=np.inf, neginf=np.inf)
#         # print(prd_dst)
#         # Compute the best predecessors
#         prd[:, k] = np.argmin(prd_dst, axis=1).ravel()
#         # print(prd[:, k])
#         # Recomput the mimum distances
#         dst = prd_dst[idx, prd[:, k]]
#         # print(dst)
#         # Detect cycles
#         if dst[root] < np.inf:
#             cnode = root
#             cycle = []
#             for h in range(k, -1, -1):
#                 cpred = prd[cnode, h]
#                 cycle.append(cpred)
#                 cnode = cpred
#             # print('cycle detected:', cycle)
#             cycles.append(tuple(cycle))
#             ccosts.append(dst[root])
#     return cycles, ccosts

def shortest_cycles(aplus, weights, max_len):
    # Store the graph in backward star format
    aminus = {i:[] for i in aplus}
    for i, alist in aplus.items():
        for j in alist:
            aminus[j].append(i)
    aplus = copy.deepcopy(aplus)
    # Prepare a list with candidate roots
    cycles, ccosts = [], []
    for root in range(len(weights)):
        tcl, tct = shortest_cycles_from_root(root, aplus, weights, max_len)
        cycles += tcl
        ccosts += tct
        # Remove all forward arcs to the processed root
        for i in aminus[root]:
            aplus[i].remove(root)
    return cycles, ccosts

def shortest_cycles_from_root(root, aplus, weights, max_len):
    spt = {root: {root}} # initial shortest paths
    dst = {root: weights[root]} # shortest path distances
    cycles, ccosts = [], []
    for k in range(max_len): # loop over the possible cycle lengths
        ndst, nspt = {}, {}
        for i in dst: # process all visited nodes
            for j in aplus[i]: # loop over outgoing arcs
                if j == root: # detect cycles
                    cycles.append(spt[i])
                    ccosts.append(dst[i])
                elif j in spt[i]: # skip subcycles
                    continue
                elif j not in ndst or dst[i] + weights[j] < ndst[j]:
                    ndst[j] = dst[i] + weights[j]
                    nspt[j] = spt[i] | {j}
        dst, spt = ndst, nspt
    return cycles, ccosts

def cycle_next(seq, nsteps, aplus, cycles, cap=None):
    node = seq[-1]
    # Consider all possible extensions
    successors = np.array(aplus[node])
    np.random.shuffle(successors)
    for dst in successors:
        # Early exit
        if cap is not None and len(cycles) >= cap: return
        # Try to close the cycle
        if dst == seq[0] and dst == min(seq):
            cycles.add(tuple(seq))
        elif nsteps > 0 and dst not in seq:
            cycle_next(seq+[dst], nsteps-1, aplus, cycles, cap)

def find_all_cycles(aplus, max_length, cap=None, seed=42):
    cycles = set()
    roots = np.array(list(aplus.keys()))
    np.random.seed(seed)
    np.random.shuffle(roots)
    for node in roots:
        if cap is None or len(cycles) < cap:
            cycle_next([node], max_length-1, aplus, cycles, cap)
    return list(cycles)

def cycle_formulation(pairs, cycles, tlim=None, relaxation=False, verbose=1):
    # Build the solver
    if relaxation:
        slv = pywraplp.Solver.CreateSolver('CLP')
    else:
        slv = pywraplp.Solver.CreateSolver('CBC')
    # Some convenience values
    infinity = slv.infinity()
    ncycles = len(cycles)
    npairs = len(pairs)
    # Group cycles by pair (saves time later on)
    cpp = {i:[] for i in range(npairs)}
    for j, cycle in enumerate(cycles):
        for i in cycle: cpp[i].append(j)
    # Build all variables
    if relaxation:
        x = [slv.NumVar(0, infinity, f'x_{j}') for j in range(ncycles)]
    else:
        x = [slv.IntVar(0, 1, f'x_{j}') for j in range(ncycles)]
    # Build all the constraints
    for i in range(npairs):
        if relaxation:
            slv.Add(-sum(x[j] for j in cpp[i]) >= -1)
        else:
            slv.Add(sum(x[j] for j in cpp[i]) <= 1)
    # Define the objective
    obj = sum(len(c) * x[j] for j, c in enumerate(cycles))
    if relaxation:
        slv.Minimize(-obj)
    else:
        slv.Maximize(obj)
    # Set a time limit
    if tlim is not None:
        slv.SetTimeLimit(1000*tlim)
    # Solve
    status = slv.Solve()
    # Extract results
    duals = None
    sol = None
    if status in (slv.OPTIMAL, slv.FEASIBLE):
        sol = {f'x_{j}':x[j].solution_value() for j in range(len(x))}
        sol['objective'] = slv.Objective().Value()
        if relaxation:
            duals = np.array([c.dual_value() for c in slv.constraints()])
    # Print stuff
    if verbose > 0:
        print_cycle_formulation_outcome(status, sol, slv, verbose, relaxation)
    # Return results
    return sol, slv.WallTime()/1000, duals

def print_cycle_formulation_outcome(status, sol, slv, verbose, relaxation):
    print(f'Solution time: {slv.WallTime()/1000:.3f}', end='')
    if status not in (slv.OPTIMAL, slv.FEASIBLE):
        print(', no feasible solution found')
    else:
        print(f', objective value: {slv.Objective().Value()}', end='')
        if status == slv.OPTIMAL: print(' (optimal)', end='')
        else: print(' (optimal)', end='')
        if verbose > 1:
            if not relaxation:
                sol = {k:int(v) for k, v in sol.items()}
            print(f', solution: {sol}')
        else:
            print()

def cycle_formulation_cg(pairs, aplus, max_len, itcap=10, tol=1e-3, verbose=1):
    # Initial pool of arcs
    weights = -np.ones(len(pairs))
    cycles, _ = shortest_cycles(aplus, weights, max_len=max_len)
    # Start the main loop
    converged = False
    for itn in range(itcap):
        # Solve the LP relaxation
        sol, stime, duals = cycle_formulation(pairs, cycles, verbose=0, relaxation=True)
        if verbose > 0:
            print(f'(CG, it. {itn}), #cycles: {len(cycles)}, time: {stime}, relaxation objective: {sol["objective"]:.2f}')
        # Find shortest paths
        weights = -np.ones(len(pairs)) + duals
        scl, sct = shortest_cycles(aplus, weights, max_len=max_len)
        # Cycles with negative reduced cost
        nrc_cycles = [scl[i] for i, c in enumerate(sct) if c < -tol]
        if verbose > 0:
            print(f'(CG, it. {itn}), #cycles with negative reduced cost: {len(nrc_cycles)}')
        if len(nrc_cycles) == 0:
            converged = True
            break
        else: cycles += nrc_cycles
    return cycles, converged

# ==============================================================================
# OLD STUFF
# ==============================================================================

# def load_data(data_folder):
#     # Read the CSV files
#     fnames = ['train_FD001', 'train_FD002', 'train_FD003', 'train_FD004']
#     cols = ['machine', 'cycle', 'p1', 'p2', 'p3'] + [f's{i}' for i in range(1, 22)]
#     datalist = []
#     nmcn = 0
#     for fstem in fnames:
#         # Read data
#         data = pd.read_csv(f'{data_folder}/{fstem}.txt', sep=' ', header=None)
#         # Drop the last two columns (parsing errors)
#         data.drop(columns=[26, 27], inplace=True)
#         # Replace column names
#         data.columns = cols
#         # Add the data source
#         data['src'] = fstem
#         # Shift the machine numbers
#         data['machine'] += nmcn
#         nmcn += len(data['machine'].unique())
#         # Generate RUL data
#         cnts = data.groupby('machine')[['cycle']].count()
#         cnts.columns = ['ftime']
#         data = data.join(cnts, on='machine')
#         data['rul'] = data['ftime'] - data['cycle']
#         data.drop(columns=['ftime'], inplace=True)
#         # Store in the list
#         datalist.append(data)
#     # Concatenate
#     data = pd.concat(datalist)
#     # Put the 'src' field at the beginning and keep 'rul' at the end
#     data = data[['src'] + cols + ['rul']]
#     # data.columns = cols
#     return data

# def plot_dataframe(data, labels=None, vmin=-1.96, vmax=1.96,
#         figsize=figsize, autoclose=True, s=4):
#     if autoclose: plt.close('all')
#     plt.figure(figsize=figsize)
#     plt.imshow(data.T.iloc[:, :], aspect='auto',
#             cmap='RdBu', vmin=vmin, vmax=vmax)
#     if labels is not None:
#         # nonzero = data.index[labels != 0]
#         ncol = len(data.columns)
#         lvl = - 0.05 * ncol
#         # plt.scatter(nonzero, lvl*np.ones(len(nonzero)),
#         #         s=s, color='tab:orange')
#         plt.scatter(labels.index, np.ones(len(labels)) * lvl,
#                 s=s,
#                 color=plt.get_cmap('tab10')(np.mod(labels, 10)))
#     plt.tight_layout()

# def split_by_field(data, field):
#     res = {}
#     for fval, gdata in data.groupby(field):
#         res[fval] = gdata
#     return res

# def partition_by_machine(data, tr_machines):
#     # Separate
#     tr_machines = set(tr_machines)
#     tr_list, ts_list = [], []
#     for mcn, gdata in data.groupby('machine'):
#         if mcn in tr_machines:
#             tr_list.append(gdata)
#         else:
#             ts_list.append(gdata)
#     # Collate again
#     tr_data = pd.concat(tr_list)
#     ts_data = pd.concat(ts_list)
#     return tr_data, ts_data

# def sliding_window_by_machine(data, wlen, cols, stride=1):
#     l_w, l_m, l_r = [], [], []
#     for mcn, gdata in data.groupby('machine'):
#         # Apply a sliding window
#         tmp_w = sliding_window_2D(gdata[cols], wlen, stride)
#         # Build the machine vector
#         tmp_m = gdata['machine'].iloc[wlen-1::stride]
#         # Build the RUL vector
#         tmp_r = gdata['rul'].iloc[wlen-1::stride]
#         # Store everything
#         l_w.append(tmp_w)
#         l_m.append(tmp_m)
#         l_r.append(tmp_r)
#     res_w = np.concatenate(l_w)
#     res_m = np.concatenate(l_m)
#     res_r = np.concatenate(l_r)
#     return res_w, res_m, res_r

# def plot_rul(pred=None, target=None,
#         stddev=None,
#         q1_3=None,
#         same_scale=True,
#         figsize=figsize, autoclose=True):
#     if autoclose:
#         plt.close('all')
#     plt.figure(figsize=figsize)
#     if target is not None:
#         plt.plot(range(len(target)), target, label='target',
#                 color='tab:orange')
#     if pred is not None:
#         if same_scale or target is None:
#             ax = plt.gca()
#         else:
#             ax = plt.gca().twinx()
#         ax.plot(range(len(pred)), pred, label='pred',
#                 color='tab:blue')
#         if stddev is not None:
#             ax.fill_between(range(len(pred)),
#                     pred-stddev, pred+stddev,
#                     alpha=0.3, color='tab:blue', label='+/- std')
#         if q1_3 is not None:
#             ax.fill_between(range(len(pred)),
#                     q1_3[0], q1_3[1],
#                     alpha=0.3, color='tab:blue', label='1st/3rd quartile')
#     plt.legend()
#     plt.tight_layout()

# class RULCostModel:
#     def __init__(self, maintenance_cost, safe_interval=0):
#         self.maintenance_cost = maintenance_cost
#         self.safe_interval = safe_interval

#     def cost(self, machine, pred, thr, return_margin=False):
#         # Merge machine and prediction data
#         tmp = np.array([machine, pred]).T
#         tmp = pd.DataFrame(data=tmp,
#                            columns=['machine', 'pred'])
#         # Cost computation
#         cost = 0
#         nfails = 0
#         slack = 0
#         for mcn, gtmp in tmp.groupby('machine'):
#             idx = np.nonzero(gtmp['pred'].values < thr)[0]
#             if len(idx) == 0:
#                 cost += self.maintenance_cost
#                 nfails += 1
#             else:
#                 cost -= max(0, idx[0] - self.safe_interval)
#                 slack += len(gtmp) - idx[0]
#         if not return_margin:
#             return cost
#         else:
#             return cost, nfails, slack

# def opt_threshold_and_plot(machine, pred, th_range, cmodel,
#         plot=True, figsize=figsize, autoclose=True):
#     # Compute the optimal threshold
#     costs = [cmodel.cost(machine, pred, thr) for thr in th_range]
#     opt_th = th_range[np.argmin(costs)]
#     # Plot
#     if plot:
#         if autoclose:
#             plt.close('all')
#         plt.figure(figsize=figsize)
#         plt.plot(th_range, costs)
#         plt.tight_layout()
#     # Return the threshold
#     return opt_th

# class MLPplusNormal(keras.Model):
#     def __init__(self, input_shape, hidden):
#         super(MLPplusNormal, self).__init__()
#         # Build the estimation model for mu and sigma
#         model_in = keras.Input(shape=input_shape, dtype='float32')
#         x = model_in
#         for h in hidden:
#             x = layers.Dense(h, activation='relu')(x)
#         mu = layers.Dense(1, activation='linear')(x)
#         logsigma = layers.Dense(1, activation='linear')(x)
#         self.model = keras.Model(model_in, [mu, logsigma])
#         # Choose what to track at training time
#         self.loss_tracker = keras.metrics.Mean(name="loss")

#     @property
#     def metrics(self):
#         return [self.loss_tracker]

#     def call(self, data, training=True):
#         if training:
#             x, y = data
#         else:
#             x = data
#         mu, logsigma = self.model(x)
#         dist = tfp.distributions.Normal(mu, k.exp(logsigma))
#         print(dist)
#         raise Exception('HAHAHA')
#         if training:
#             return dist.log_prob(y)
#         else:
#             return dist.mean(), dist.stddev()

#     # def log_loss(self, data):
#     #     log_densities = self(data)
#     #     return -tf.reduce_mean(log_densities)

#     def train_step(self, data):
#         with tf.GradientTape() as tape:
#             log_densities = self(data)
#             loss = -tf.reduce_mean(log_densities)
#             # loss = self.log_loss(data)
#         g = tape.gradient(loss, self.trainable_variables)
#         self.optimizer.apply_gradients(zip(g, self.trainable_variables))
#         self.loss_tracker.update_state(loss)
#         return {"loss": self.loss_tracker.result()}

# ==============================================================================
# OLD FUNCTIONS
# ==============================================================================

# def plot_distribution_2D(estimator=None, samples=None,
#         xr=None, yr=None,
#         figsize=figsize, autoclose=True):
#     if autoclose:
#         plt.close('all')
#     plt.figure(figsize=figsize)
#     if samples is not None:
#         plt.scatter(samples[:, 0], samples[:, 1], alpha=0.3)
#     if estimator is not None:
#         if xr is None:
#             xr = np.linspace(-1, 1, 100)
#         if yr is None:
#             yr = np.linspace(-1, 1, 100)
#         nx = len(xr)
#         ny = len(yr)
#         xc = np.repeat(xr, ny)
#         yc = np.tile(yr, nx)
#         data = np.vstack((xc, yc)).T
#         dvals = np.exp(estimator.score_samples(data))
#         dvals = dvals.reshape((nx, ny))
#         plt.imshow(dvals.T[::-1, :], aspect='auto')
#         plt.xticks(np.linspace(0, len(xr), 5), np.linspace(xr[0], xr[-1], 5))
#         plt.xticks(np.linspace(0, len(yr), 5), np.linspace(yr[0], yr[-1], 5))
#     plt.tight_layout()

# def get_errors(signal, labels, thr, tolerance=1):
#     pred = signal[signal > thr].index
#     anomalies = labels[labels != 0].index

#     fp = set(pred)
#     fn = set(anomalies)
#     for lag in range(-tolerance, tolerance+1):
#         fp = fp - set(anomalies+lag)
#         fn = fn - set(pred+lag)
#     return fp, fn

# class HPCMetrics:
#     def __init__(self, c_alarm, c_missed, tolerance):
#         self.c_alarm = c_alarm
#         self.c_missed = c_missed
#         self.tolerance = tolerance

#     def cost(self, signal, labels, thr):
#         # Obtain errors
#         fp, fn = get_errors(signal, labels, thr, self.tolerance)

#         # Compute the cost
#         return self.c_alarm * len(fp) + self.c_missed * len(fn)

# def opt_threshold(signal, labels, th_range, cmodel):
#     costs = [cmodel.cost(signal, labels, th) for th in th_range]
#     best_th = th_range[np.argmin(costs)]
#     best_cost = np.min(costs)
#     return best_th, best_cost

# def collect_training(data, tr_ts_ratio):
#     if isinstance(data, dict):
#         tr_list = []
#         for key, kdata in data.items():
#             sep = int(np.round(tr_ts_ratio * len(kdata)))
#             tr_list.append(kdata.iloc[:sep])
#         tr_data = pd.concat(tr_list)
#     else:
#         sep = int(np.round(tr_ts_ratio * len(data)))
#         tr_data = data.iloc[:sep]
#     return tr_data

# def kde_ad(x_tr, x_vs, y_vs, x_ts, y_ts, th_range, h_range, cmodel):
#     # Optimize the bandwidth, if more than one value is given
#     if len(h_range) > 1:
#         params = {'bandwidth': h_range}
#         opt = GridSearchCV(KernelDensity(kernel='gaussian'), params, cv=5)
#         opt.fit(x_tr)
#         h = opt.best_params_['bandwidth']
#     else:
#         h = h_range[0]
#     # Traing a KDE estimator
#     kde = KernelDensity(bandwidth=h)
#     kde.fit(x_tr)
#     # Generate a signal for the validation set
#     ldens_vs = kde.score_samples(x_vs)
#     signal_vs = pd.Series(index=x_vs.index, data=-ldens_vs)
#     # Threshold optimization
#     th, val_cost = opt_threshold(signal_vs, y_vs, th_range, cmodel)
#     # Compute the cost over the test set
#     ldens_ts = kde.score_samples(x_ts)
#     signal_ts = pd.Series(index=x_ts.index, data=-ldens_ts)
#     ts_cost = cmodel.cost(signal_ts, y_ts, th)
#     # Return results
#     return kde, ts_cost

# def ae_ad(ae, x_tr, x_vs, y_vs, x_ts, y_ts, th_range, cmodel):
#     # Traing the autoencoder
#     ae.compile(optimizer='RMSProp', loss='mse')
#     cb = [callbacks.EarlyStopping(patience=3, restore_best_weights=True)]
#     ae.fit(x_tr, x_tr, validation_split=0.1, callbacks=cb,
#            batch_size=32, epochs=20, verbose=0)
#     # Generate a signal for the validation set
#     preds_vs = pd.DataFrame(index=x_vs.index,
#             columns=x_vs.columns, data=ae.predict(x_vs))
#     sse_vs = np.sum(np.square(preds_vs - x_vs), axis=1)
#     signal_vs = pd.Series(index=x_vs.index, data=sse_vs)
#     # Threshold optimization
#     th, val_cost = opt_threshold(signal_vs, y_vs, th_range, cmodel)
#     # Compute the cost over the test set
#     preds_ts = pd.DataFrame(index=x_ts.index,
#             columns=x_ts.columns, data=ae.predict(x_ts))
#     sse_ts = np.sum(np.square(preds_ts - x_ts), axis=1)
#     signal_ts = pd.Series(index=x_ts.index, data=sse_ts)
#     ts_cost = cmodel.cost(signal_ts, y_ts, th)
#     # Return results
#     return ts_cost

# def coupling(input_shape, nunits=64, nhidden=2, reg=0.01):
#     assert(nhidden >= 0)
#     x = keras.layers.Input(shape=input_shape)
#     # Build the layers for the t transformation (translation)
#     t = x
#     for i in range(nhidden):
#         t = Dense(nunits, activation="relu", kernel_regularizer=l2(reg))(t)
#     t = Dense(input_shape, activation="linear", kernel_regularizer=l2(reg))(t)
#     # Build the layers for the s transformation (scale)
#     s = x
#     for i in range(nhidden):
#         s = Dense(nunits, activation="relu", kernel_regularizer=l2(reg))(s)
#     s = Dense(input_shape, activation="tanh", kernel_regularizer=l2(reg))(s)
#     # Return the layers, wrapped in a keras Model object
#     return keras.Model(inputs=x, outputs=[s, t])

# class RealNVP(keras.Model):
#     def __init__(self, input_shape, num_coupling, units_coupling=32, depth_coupling=0,
#             reg_coupling=0.01):
#         super(RealNVP, self).__init__()
#         self.num_coupling = num_coupling
#         # Distribution of the latent space
#         self.distribution = tfp.distributions.MultivariateNormalDiag(
#             loc=np.zeros(input_shape, dtype=np.float32),
#             scale_diag=np.ones(input_shape, dtype=np.float32)
#         )
#         # Build a mask
#         half_n = int(np.ceil(input_shape/2))
#         m1 = ([0, 1] * half_n)[:input_shape]
#         m2 = ([1, 0] * half_n)[:input_shape]
#         self.masks = np.array([m1, m2] * (num_coupling // 2), dtype=np.float32)
#         # Choose what to track at training time
#         self.loss_tracker = keras.metrics.Mean(name="loss")
#         #  Build layers
#         self.layers_list = [coupling(input_shape, units_coupling, depth_coupling, reg_coupling)
#                             for i in range(num_coupling)]

#     @property
#     def metrics(self):
#         """List of the model's metrics.
#         We make sure the loss tracker is listed as part of `model.metrics`
#         so that `fit()` and `evaluate()` are able to `reset()` the loss tracker
#         at the start of each epoch and at the start of an `evaluate()` call.
#         """
#         return [self.loss_tracker]

#     def call(self, x, training=True):
#         log_det_inv, direction = 0, 1
#         if training: direction = -1
#         for i in range(self.num_coupling)[::direction]:
#             x_masked = x * self.masks[i]
#             reversed_mask = 1 - self.masks[i]
#             s, t = self.layers_list[i](x_masked)
#             s, t = s*reversed_mask, t*reversed_mask
#             gate = (direction - 1) / 2
#             x = reversed_mask * (x * tf.exp(direction * s) + direction * t * tf.exp(gate * s)) \
#                 + x_masked
#             log_det_inv += gate * tf.reduce_sum(s, axis=1)
#         return x, log_det_inv

#     def log_loss(self, x):
#         log_densities = self.score_samples(x)
#         return -tf.reduce_mean(log_densities)

#     def score_samples(self, x):
#         y, logdet = self(x)
#         log_probs = self.distribution.log_prob(y) + logdet
#         return log_probs

#     def train_step(self, data):
#         with tf.GradientTape() as tape:
#             loss = self.log_loss(data)
#         g = tape.gradient(loss, self.trainable_variables)
#         self.optimizer.apply_gradients(zip(g, self.trainable_variables))
#         self.loss_tracker.update_state(loss)
#         return {"loss": self.loss_tracker.result()}

#     def test_step(self, data):
#         loss = self.log_loss(data)
#         self.loss_tracker.update_state(loss)
#         return {"loss": self.loss_tracker.result()}

# def plot_rnvp_transformation(rnvp, xr=None, yr=None,
#         figsize=figsize, autoclose=True):
#     if autoclose:
#         plt.close('all')
#     plt.figure(figsize=figsize)
#     # Define ranges
#     if xr is None:
#         xr = np.linspace(-1, 1, 7, dtype=np.float32)
#     if yr is None:
#         yr = np.linspace(-1, 1, 7, dtype=np.float32)
#     # Build the input set
#     nx = len(xr)
#     ny = len(yr)
#     xc = np.repeat(xr, ny)
#     yc = np.tile(yr, nx)
#     data = np.vstack((xc, yc)).T
#     # Transform the input step
#     z, _ = rnvp(data, training=False)
#     # Obtain traces
#     traces = np.concatenate((
#         data.reshape(1, -1, 2),
#         z.numpy().reshape(1, -1, 2),
#         ))
#     # Plot traces
#     for i in range(traces.shape[1]):
#         plt.plot(traces[:, i, 0], traces[:, i, 1], ':',
#                 color='0.8', zorder=0)
#         xh = plt.scatter(data[:, 0], data[:, 1],
#                 color='tab:blue', s=4, zorder=1)
#         zh = plt.scatter(z.numpy()[:, 0], z.numpy()[:, 1],
#                 color='tab:orange', s=4, zorder=1)
#     plt.legend([xh, zh], ['x', 'z'])
#     plt.tight_layout()


import torch
import numpy as np
import sys
from sklearn.metrics import confusion_matrix

def compute_confusion_matrix(targets, predictions, num_classes):
    confusion_mat = confusion_matrix(y_true=targets.flatten(), y_pred=predictions.flatten(), labels=range(num_classes))
    return confusion_mat

def measure_performance(confusion_mat, num_classes, ignore_label=None, smooth=sys.float_info.epsilon):
    ious = []
    classes_with_missing_representatives = []

    for c in range(num_classes):
        if c != ignore_label:
            intersection = confusion_mat[c, c]
            union = confusion_mat[c, :].sum() + confusion_mat[:, c].sum() - intersection

            if confusion_mat[c, :].sum() == 0:
                classes_with_missing_representatives += [c]
                ious.append(0)
            else:
                ious.append((intersection + smooth) / (union + smooth) if union > 0 else 0)
        else:
            ious.append(0)

    miou = np.array(ious).sum()
    ious = np.array(ious)

    if ignore_label != None:
        miou -= ious[ignore_label]
        miou /= (num_classes - 1 - len(classes_with_missing_representatives))
    else:
        miou /= (num_classes - len(classes_with_missing_representatives))

    for c in classes_with_missing_representatives:
        ious[c] = -1

    return miou, ious

def add_weight_decay(net, l2_value, skip_list=()):
    decay, no_decay = [], []
    for name, param in net.named_parameters():
        if not param.requires_grad: continue # frozen weights
        if len(param.shape) == 1 or name.endswith(".bias") or name in skip_list: no_decay.append(param)
        else: decay.append(param)
    return [{'params': no_decay, 'weight_decay': 0.}, {'params': decay, 'weight_decay': l2_value}]


import sys
import unittest
try:
    from unittest.mock import MagicMock
except ImportError:
    from mock import MagicMock
import botocore

##############
# Parameters #
##############

# Define the default resource to report to Config Rules
DEFAULT_RESOURCE_TYPE = 'AWS::ECR::Repository'

#############
# Main Code #
#############

CONFIG_CLIENT_MOCK = MagicMock()
STS_CLIENT_MOCK = MagicMock()
ECR_CLIENT_MOCK = MagicMock()

class Boto3Mock():
    @staticmethod
    def client(client_name, *args, **kwargs):
        if client_name == 'config':
            return CONFIG_CLIENT_MOCK
        if client_name == 'sts':
            return STS_CLIENT_MOCK
        if client_name == 'ecr':
            return ECR_CLIENT_MOCK
        raise Exception("Attempting to create an unknown client")

sys.modules['boto3'] = Boto3Mock()

RULE = __import__('ECR_REPOSITORY_SCAN_ON_PUSH_CHECK')

class NotApplicable(unittest.TestCase):
    def test_scenario_1_not_applicable(self):
        list_repositories_result = {"repositories": []}

        ECR_CLIENT_MOCK.describe_repositories = MagicMock(return_value=list_repositories_result)
        rule_parameters = '{}'
        lambda_result = RULE.lambda_handler(build_lambda_scheduled_event(rule_parameters), {})
        expected_response = [build_expected_response("NOT_APPLICABLE", '123456789012', 'AWS::::Account')]
        assert_successful_evaluation(self, lambda_result, expected_response, len(lambda_result))

class InvalidParameter(unittest.TestCase):
    def test_scenario_2_invalid_parameter(self):
        list_repositories_result = {"repositories": []}

        ECR_CLIENT_MOCK.describe_repositories = MagicMock(return_value=list_repositories_result)
        rule_parameters = '{\"Parameter\":\"Value\"}'
        lambda_result = RULE.lambda_handler(build_lambda_scheduled_event(rule_parameters), {})
        expected_error_code = 'InvalidParameterValueException'
        expected_error_message = 'This rule is not configured to take any input parameters.'
        assert_customer_error_response(self, lambda_result, expected_error_code, expected_error_message)

class NonComplianceTest(unittest.TestCase):
    def test_scenario_3_compliant_resources_with_key(self):
        list_repositories_result = {"repositories":
            [
                {
                    "repositoryArn": "arn:aws:sns:us-east-1:123456789012:testRepository",
                    "imageScanningConfiguration": {"scanOnPush": False}
                }
            ]
        }

        ECR_CLIENT_MOCK.describe_repositories = MagicMock(return_value=list_repositories_result)
        rule_parameters = '{}'
        lambda_result = RULE.lambda_handler(build_lambda_scheduled_event(rule_parameters), {})
        expected_response = [build_expected_response(
            'NON_COMPLIANT',
            'arn:aws:sns:us-east-1:123456789012:testRepository',
            annotation="The Amazon Elastic Container Registry repository is not configured to scan images on push."
        )]
        assert_successful_evaluation(self, lambda_result, expected_response, len(lambda_result))

class ComplianceTest(unittest.TestCase):
    def test_scenario_4_compliant_resources_with_key(self):
        list_repositories_result = {"repositories":
            [
                {
                    "repositoryArn": "arn:aws:sns:us-east-1:123456789012:testRepository",
                    "imageScanningConfiguration": {"scanOnPush": True}
                }
            ]
        }

        ECR_CLIENT_MOCK.describe_repositories = MagicMock(return_value=list_repositories_result)
        rule_parameters = '{}'
        lambda_result = RULE.lambda_handler(build_lambda_scheduled_event(rule_parameters), {})
        expected_response = [build_expected_response(
            'COMPLIANT',
            'arn:aws:sns:us-east-1:123456789012:testRepository'
        )]
        assert_successful_evaluation(self, lambda_result, expected_response, len(lambda_result))

####################
# Helper Functions #
####################

def build_lambda_configurationchange_event(invoking_event, rule_parameters=None):
    event_to_return = {
        'configRuleName':'myrule',
        'executionRoleArn':'roleArn',
        'eventLeftScope': False,
        'invokingEvent': invoking_event,
        'accountId': '123456789012',
        'configRuleArn': 'arn:aws:config:us-east-1:123456789012:config-rule/config-rule-8fngan',
        'resultToken':'token'
    }
    if rule_parameters:
        event_to_return['ruleParameters'] = rule_parameters
    return event_to_return

def build_lambda_scheduled_event(rule_parameters=None):
    invoking_event = '{"messageType":"ScheduledNotification","notificationCreationTime":"2017-12-23T22:11:18.158Z"}'
    event_to_return = {
        'configRuleName':'myrule',
        'executionRoleArn':'roleArn',
        'eventLeftScope': False,
        'invokingEvent': invoking_event,
        'accountId': '123456789012',
        'configRuleArn': 'arn:aws:config:us-east-1:123456789012:config-rule/config-rule-8fngan',
        'resultToken':'token'
    }
    if rule_parameters:
        event_to_return['ruleParameters'] = rule_parameters
    return event_to_return

def build_expected_response(compliance_type, compliance_resource_id, compliance_resource_type=DEFAULT_RESOURCE_TYPE, annotation=None):
    if not annotation:
        return {
            'ComplianceType': compliance_type,
            'ComplianceResourceId': compliance_resource_id,
            'ComplianceResourceType': compliance_resource_type
            }
    return {
        'ComplianceType': compliance_type,
        'ComplianceResourceId': compliance_resource_id,
        'ComplianceResourceType': compliance_resource_type,
        'Annotation': annotation
        }

def assert_successful_evaluation(test_class, response, resp_expected, evaluations_count=1):
    if isinstance(response, dict):
        test_class.assertEquals(resp_expected['ComplianceResourceType'], response['ComplianceResourceType'])
        test_class.assertEquals(resp_expected['ComplianceResourceId'], response['ComplianceResourceId'])
        test_class.assertEquals(resp_expected['ComplianceType'], response['ComplianceType'])
        test_class.assertTrue(response['OrderingTimestamp'])
        if 'Annotation' in resp_expected or 'Annotation' in response:
            test_class.assertEquals(resp_expected['Annotation'], response['Annotation'])
    elif isinstance(response, list):
        test_class.assertEquals(evaluations_count, len(response))
        for i, response_expected in enumerate(resp_expected):
            test_class.assertEquals(response_expected['ComplianceResourceType'], response[i]['ComplianceResourceType'])
            test_class.assertEquals(response_expected['ComplianceResourceId'], response[i]['ComplianceResourceId'])
            test_class.assertEquals(response_expected['ComplianceType'], response[i]['ComplianceType'])
            test_class.assertTrue(response[i]['OrderingTimestamp'])
            if 'Annotation' in response_expected or 'Annotation' in response[i]:
                test_class.assertEquals(response_expected['Annotation'], response[i]['Annotation'])

def assert_customer_error_response(test_class, response, customer_error_code=None, customer_error_message=None):
    if customer_error_code:
        test_class.assertEqual(customer_error_code, response['customerErrorCode'])
    if customer_error_message:
        test_class.assertEqual(customer_error_message, response['customerErrorMessage'])
    test_class.assertTrue(response['customerErrorCode'])
    test_class.assertTrue(response['customerErrorMessage'])
    if "internalErrorMessage" in response:
        test_class.assertTrue(response['internalErrorMessage'])
    if "internalErrorDetails" in response:
        test_class.assertTrue(response['internalErrorDetails'])

def sts_mock():
    assume_role_response = {
        "Credentials": {
            "AccessKeyId": "string",
            "SecretAccessKey": "string",
            "SessionToken": "string"}}
    STS_CLIENT_MOCK.reset_mock(return_value=True)
    STS_CLIENT_MOCK.assume_role = MagicMock(return_value=assume_role_response)

##################
# Common Testing #
##################

class TestStsErrors(unittest.TestCase):

    def test_sts_unknown_error(self):
        RULE.ASSUME_ROLE_MODE = True
        RULE.evaluate_parameters = MagicMock(return_value=True)
        STS_CLIENT_MOCK.assume_role = MagicMock(side_effect=botocore.exceptions.ClientError(
            {'Error': {'Code': 'unknown-code', 'Message': 'unknown-message'}}, 'operation'))
        response = RULE.lambda_handler(build_lambda_configurationchange_event('{}'), {})
        assert_customer_error_response(
            self, response, 'InternalError', 'InternalError')

    def test_sts_access_denied(self):
        RULE.ASSUME_ROLE_MODE = True
        RULE.evaluate_parameters = MagicMock(return_value=True)
        STS_CLIENT_MOCK.assume_role = MagicMock(side_effect=botocore.exceptions.ClientError(
            {'Error': {'Code': 'AccessDenied', 'Message': 'access-denied'}}, 'operation'))
        response = RULE.lambda_handler(build_lambda_configurationchange_event('{}'), {})
        assert_customer_error_response(
            self, response, 'AccessDenied', 'AWS Config does not have permission to assume the IAM role.')


import webbrowser
from os import environ
from typing import Tuple

from flask import abort, request

from kanmail.log import logger
from kanmail.server.app import add_route
from kanmail.window import create_window, destroy_window, resize_window

@add_route("/window/open-link", methods=("GET",))
def open_link() -> Tuple[str, int]:
    link = request.args["url"]

    # https://github.com/pyinstaller/pyinstaller/issues/3668
    xdg_data_dirs = environ.pop("XDG_DATA_DIRS", None)
    try:
        if webbrowser.open(link):
            return "", 204
    finally:
        if xdg_data_dirs:
            environ["XDG_DATA_DIRS"] = xdg_data_dirs

    logger.critical(f"Failed to open browser link: {link}!")
    abort(500, "Could not open link!")

@add_route("/window/open", methods=("GET",))
def open_window() -> Tuple[str, int]:
    link = request.args["url"]

    if not create_window(
        link,
        width=int(request.args["width"]),
        height=int(request.args["height"]),
        unique_key=request.args.get("unique_key"),
        confirm_close=request.args.get("confirm_close"),
        resizable=not (request.args.get("resizable") == "false"),
    ):
        abort(500, f"Could not open {link} window")
    return "", 204

@add_route("/window/close", methods=("GET",))
def close_window() -> Tuple[str, int]:
    destroy_window(request.args["window_id"])
    return "", 204

@add_route("/window/resize", methods=("GET",))
def window_resize() -> Tuple[str, int]:
    resize_window(
        request.args["window_id"],
        width=int(request.args["width"]),
        height=int(request.args["height"]),
    )
    return "", 204


"""
This module contains code to find starting point on each main stream.
"""

from pathlib import Path
from shutil import copy2
from tempfile import TemporaryDirectory
from typing import Callable, Generator, List, Optional, Tuple, Union

import fiona
import geosardine as dine
import numpy as np
import rasterio
import whitebox
from rasterio.features import geometry_mask
from shapely import geometry, ops, speedups
from shapely.geometry import Point
from shapely.geometry.base import BaseGeometry
from tqdm.autonotebook import tqdm

from .custom_types import GeoJsonDict

speedups.disable()

wbt = whitebox.WhiteboxTools()

class StemTooShort(Exception):
    """Stream/Stem is too short to calculate volume. It is near the start vertex"""

    pass

class DifferentCRS(Exception):
    """two data have different CRS"""

    pass

class ProcessingData:
    """
    Contains processing data in temporary directory
    """

    def __init__(self, flow_direction: Path, flow_stream: Path) -> None:
        self.temp_folder = TemporaryDirectory()
        self.temp_path = Path(self.temp_folder.name)
        self.flow_direction = flow_direction
        self.flow_stream = flow_stream

        self.main_stem_rasterfile = self.temp_path / "main_stem.tif"
        self.main_stem_vectorfile = self.temp_path / "main_stem.shp"
        self.link_class_file = self.temp_path / "link_class.tif"

        self.link_class: Optional["dine.Raster"] = None

        self.generate()

    def generate(self) -> None:
        """generate processing data"""
        wbt.find_main_stem(
            self.flow_direction,
            self.flow_stream,
            self.main_stem_rasterfile,
            esri_pntr=True,
        )
        wbt.stream_link_class(
            self.flow_direction, self.flow_stream, self.link_class_file, esri_pntr=True
        )
        wbt.raster_streams_to_vector(
            self.main_stem_rasterfile,
            self.flow_direction,
            self.main_stem_vectorfile,
            esri_pntr=True,
        )

        self.link_class = dine.Raster.from_rasterfile(str(self.link_class_file))
        print("processing data has been generated")

    def save(self, output_directory: Path) -> None:
        """Copy processing data to outside temporary directory

        Parameters
        ----------
        output_directory : Path
            directory to copy processing data
        """
        for f in self.temp_path.glob("*"):
            if f.is_file():
                copy2(f.absolute(), output_directory)

    def cleanup(self) -> None:
        """remove  temporary directory"""
        self.temp_folder.cleanup()

def interpolate_points(
    polyline: geometry.LineString, distance: float = 0.25, upstream: int = 0
) -> Generator[geometry.Point, None, None]:
    """interpolate points along line

    Parameters
    ----------
    polyline : geometry.LineString
        stream
    distance : float, optional
        interpolation distance, by default 0.25
    upstream : int, optional
        [description], by default 0

    Yields
    -------
    Generator[geometry.Point, None, None]
        interpolated points
    """
    interpolate_distances = np.arange((polyline.length // distance) + 1)
    for interpolate_distance in interpolate_distances:
        yield polyline.interpolate(interpolate_distance)

def estimate_volume(array: np.ndarray, dsm_diff: dine.Raster) -> float:
    """estimate volume

    Parameters
    ----------
    array : np.ndarray
        clipped dsm difference of 2 epoch by buffered starting point
    dsm_diff : dine.Raster
        dsm difference of 2 epoch

    Returns
    -------
    float
        volume
    """
    return float(
        np.sum(array[(array > dsm_diff.no_data) & (array > 0)]) * dsm_diff.resolution[0]
    )

def calculate_volume_stem(
    stem: geometry.LineString,
    point: geometry.Point,
    dsm_diff: "dine.Raster",
    stream_buffer_size: float,
) -> float:
    """Calculate stem by buffering the stem

    Parameters
    ----------
    stem : geometry.LineString
        stream
    point : geometry.Point
        starting point
    dsm_diff : dine.Raster
        dsm difference of 2 epoch
    stream_buffer_size : float
        buffer length for stream

    Returns
    -------
    float
        volume

    Raises
    ------
    StemTooShort
        stem is below 5 times spatial resolution which will caused too small volume
    """
    splitted = ops.split(stem, point.buffer(0.1))
    if splitted[0].length < dsm_diff.resolution[0] * 5:
        raise StemTooShort
    mask = geometry_mask(
        [geometry.mapping(splitted[0].buffer(stream_buffer_size))],
        transform=dsm_diff.transform,
        out_shape=dsm_diff.shape[:2],
    )
    return float(
        (
            dsm_diff.array[
                (~mask.reshape(mask.shape[0], mask.shape[1], 1))
                & (dsm_diff > 0)
                & (dsm_diff != dsm_diff.no_data)
            ]
            * dsm_diff.resolution[0]
        ).sum()
    )

def find_starting_point(
    stem: geometry.LineString,
    distance: float,
    dsm_diff: "dine.Raster",
    link_class: "dine.Raster",
    stream_buffer_size: float,
) -> Tuple[Optional[geometry.Point], Optional[float]]:
    """find starting point over main stem

    Parameters
    ----------
    stem : geometry.LineString
        main stem
    distance : float
        interpolation distance
    dsm_diff : dine.Raster
        dsm difference of 2 epoch
    link_class : dine.Raster
        stream link class
    stream_buffer_size : float
        buffer length for stream

    Returns
    -------
    Tuple[Optional[geometry.Point], Optional[float]]
        Starting point and volume
        if geometry.Point & float, there is starting point
        if None & None, there isn't any starting point

    Raises
    ------
    ValueError
        Wrong type
    """

    if stem.length < 100:
        return None, None

    buffer_pixel = 5
    buffer_area = (buffer_pixel * 2) ** 2
    minimum_pixel = buffer_area * 0.5

    point_in_stem = interpolate_points(stem, distance)

    sp_candidates: List[geometry.Point] = []
    sp_candidates_junction: List[geometry.Point] = []
    for point in point_in_stem:
        point_diff = dsm_diff.xy_value(point.x, point.y)
        if point_diff != dsm_diff.no_data and point_diff >= 0:
            row, col = dsm_diff.xy2rowcol(point.x, point.y)

            bounding_box = np.array(
                [
                    row - buffer_pixel,
                    col - buffer_pixel,
                    row + buffer_pixel,
                    col + buffer_pixel,
                ]
            )

            deposition_pixel_count: int = np.sum(
                dsm_diff.array[
                    bounding_box[0] : bounding_box[2], bounding_box[1] : bounding_box[3]
                ]
                > 0
            )

            if (
                np.all(bounding_box > 0)
                and dsm_diff.array[row, col] > 0
                and deposition_pixel_count > minimum_pixel
            ):
                link_type = link_class.xy_value(point.x, point.y)

                if not isinstance(link_type, np.ndarray):
                    raise ValueError("Unexpected")

                if link_type[0] == 4:
                    sp_candidates_junction.append(point)
                else:
                    sp_candidates.append(point)
    starting_point = None
    if sp_candidates_junction:
        starting_point = sp_candidates_junction[-1]
    elif sp_candidates:
        starting_point = sp_candidates[-1]

    if starting_point is not None:
        try:
            volume: Optional[float] = calculate_volume_stem(
                stem, starting_point, dsm_diff, stream_buffer_size
            )
        except StemTooShort:
            row, col = dsm_diff.xy2rowcol(starting_point.x, starting_point.y)

            bounding_box = np.array(
                [
                    row - buffer_pixel,
                    col - buffer_pixel,
                    row + buffer_pixel,
                    col + buffer_pixel,
                ]
            )

            buffered_point = dsm_diff.array[
                bounding_box[0] : bounding_box[2], bounding_box[1] : bounding_box[3]
            ]
            volume = estimate_volume(
                buffered_point[buffered_point != dsm_diff.no_data],
                dsm_diff,
            )
    else:
        volume = None

    return starting_point, volume

def find_starting_points(
    input_earlier_dsm: str,
    input_later_dsm: str,
    input_flow_direction: str,
    input_flow_stream: str,
    max_percent_length: float,
    stream_buffer_size: float,
    return_processing_data: bool = False,
    progress_callback: Optional[Callable[[int, int], None]] = None,
) -> Tuple[List[Tuple[Point, float]], Optional[ProcessingData]]:
    """Batch find starting point for streams

    Parameters
    ----------
    input_earlier_dsm : str
        earlier dsm (first epoch) location
    input_later_dsm : str
        later dsm (second epoch) location
    input_flow_direction : str
        flow direction location d8-esri-style
    input_flow_stream : str
        flow stream location
    max_percent_length : float
        maximum percentage of stream to be included
    stream_buffer_size : float
        buffer length for stream
    return_processing_data : bool, optional
        return processing data as variable, by default False
    progress_callback : Optional[Callable[[int, int], None]], optional
        callback to be called after each loop, by default None

    Returns
    -------
    Tuple[List[Tuple[Point, float]], Optional[ProcessingData]]
        starting point, volume & processing data
        if Tuple[List[Tuple[Point, float]], ProcessingData], processing data is returned
        if Tuple[List[Tuple[Point, float]], None], processing data is not returned

    Raises
    ------
    ValueError
        Different CRS
    ValueError
        Wrong type
    e
        There is any exception raised
    """
    if progress_callback is not None:
        progress_callback(0, 0)

    processing_data = ProcessingData(
        Path(input_flow_direction), Path(input_flow_stream)
    )

    try:
        print(processing_data.main_stem_vectorfile)
        print("r", input_earlier_dsm, input_later_dsm)
        with fiona.open(processing_data.main_stem_vectorfile) as lines, rasterio.open(
            input_earlier_dsm
        ) as earlier_dsm, rasterio.open(input_later_dsm) as later_dsm:
            earlier_dsm = dine.Raster.from_rasterfile(input_earlier_dsm)
            later_dsm = dine.Raster.from_rasterfile(input_later_dsm)

            if later_dsm.epsg != earlier_dsm.epsg:
                raise ValueError("Both dsm should be in the same reference system")

            dsm_diff = later_dsm - earlier_dsm

            starting_points: List[Tuple[Point, float]] = []
            sp_coordinates: np.ndarray = np.empty((1, 2), dtype=np.float32)

            features: Tuple[Tuple[int, GeoJsonDict], ...] = tuple(
                lines.items(
                    bbox=(
                        later_dsm.x_min,
                        later_dsm.y_min,
                        later_dsm.x_max,
                        later_dsm.y_max,
                    )
                )
            )

            progress_total = len(features)

            for i, feature in tqdm(features, total=progress_total):

                line = feature["geometry"]["coordinates"]
                if later_dsm.xy_value(
                    *feature["geometry"]["coordinates"][0]
                ) < later_dsm.xy_value(*feature["geometry"]["coordinates"][-1]):
                    line = feature["geometry"]["coordinates"][::-1]

                stem = geometry.LineString(line)
                stem = ops.substring(stem, 0, stem.length * max_percent_length / 100)

                # if isinstance(stem, BaseGeometry) or processing_data.link_class is None:
                #     raise ValueError("Unexpected")

                starting_point, volume = find_starting_point(
                    stem, 0.25, dsm_diff, processing_data.link_class, stream_buffer_size
                )
                if starting_point is not None and volume is not None:
                    near_sp_exist = any(
                        (
                            np.abs(sp_coordinates[:, 0] - starting_point.x)
                            + np.abs(sp_coordinates[:, 1] - starting_point.y)
                        )
                        <= 3
                    )

                    if not near_sp_exist:
                        starting_points.append((starting_point, volume))

                if progress_callback is not None:
                    progress_callback(progress_total, i + 1)
            if return_processing_data:
                return starting_points, processing_data
            else:
                # processing_data.cleanup()
                return starting_points, None
    except Exception as e:
        processing_data.cleanup()
        raise e

def save2txt(
    starting_points: List[Tuple[Point, float]], output_location: Union[Path, str]
) -> None:
    """Save starting points and volume as csv

    Parameters
    ----------
    starting_points : List[Tuple[Point, float]]
        list of starting point as Point and volume
    output_location : Union[Path, str]
        location to save starting point
    """
    with open(output_location, "w") as out:
        for starting_point, volume in starting_points:
            out.writelines(f"{starting_point.x},{starting_point.y},{volume}\n")


import sys
from collections import defaultdict
input = sys.stdin.readline
sys.setrecursionlimit(10**9)

def dfs(cur, depth):
    node[depth].append(cur)
    for c in down[cur]:
        dfs(c, depth+1)

n = int(input())
up = defaultdict(dict)
down = defaultdict(list)
for _ in range(n-1):
    u, v, w = map(int, input().split())
    up[v][u] = w
    down[u].append(v)

node = defaultdict(list)
ans = [[0] for _ in range(n+1)]
dfs(1, 0)
depth = max(node.keys())

for d in range(depth, 0, -1):
    for v in node[d]:
        for p, w in up[v].items():
            ans[p].append(max(ans[v])+w)
answer = 0
for a in ans:
    if len(a) > 2:
        a.sort()
        answer = max(answer, a[-1]+a[-2])

b = sorted(ans[1])[-1]
print(max(answer, b))

알고리즘 스터디/1967번 트리의 지름.py


from transformers import BertPreTrainedModel
import torch.nn as nn
from kobert_transformers import get_kobert_model

import logging
_logger = logging.getLogger('train')

class BERT(BertPreTrainedModel):
    def __init__(self, config: dict, num_classes: int):
        super().__init__(config)

        self.bert = get_kobert_model() #monologg/kobert
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, num_classes)

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        output_attentions=None,
    ):

        outputs = self.bert(
            input_ids,
            attention_mask       = attention_mask,
            token_type_ids       = token_type_ids,
            position_ids         = None,
            head_mask            = None,
            inputs_embeds        = None,
            output_attentions    = output_attentions,
            output_hidden_states = None,
        )

        pooled_output = outputs[1]

        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)

        if output_attentions:
            return logits, outputs[-1]
        else:
            return logits


from django.test import TestCase, Client, RequestFactory
from django.core.urlresolvers import reverse
from food_core.models import Wishlist, User
import logging
from django.forms.models import model_to_dict

class wishlistTest(TestCase):
    #setUp method is called before each test in this class
    def setUp(self):
        #s = Wishlist.objects.create(pk=1, name='Chickpeas', price=15.25, description='Dired, salted and spiced chickpeas', nutrition_info='Great for you!')
        pass
    def test_create_wishlist(self):
        u2= User.objects.create(first_name='Test2', last_name='User', email='test@user.com', phone_number='123-456-7890')
        u2.save()
        data = {
            'pk': 2,
            'total_price': 12.0,
            'num_items': 5,
            'User': u2.id,
        }
        response = self.client.post(reverse('create_wishlist'),data)
        self.assertEquals(response.json()['status_code'], str(200), msg="Return status code is not 200! Error with create_wishlist API!")

    def test_get_wishlist(self):
        u = User.objects.create(first_name='Test', last_name='User', email='test@user.com', phone_number='123-456-7890')
        u.save()
        s = Wishlist.objects.create(user = u, total_price = 500.26, num_items=5)
        s.save()
        response = self.client.get(reverse('get_wishlists', kwargs={'pk': s.id}))
        self.assertContains(response, 'total_price', msg_prefix= str(response.json()))

    def test_update_wishlist(self):
        u = User.objects.create(first_name='Test', last_name='User', email='test@user.com', phone_number='123-456-7890')
        u.save()
        s = Wishlist.objects.create(user = u, total_price = 500.26, num_items=5)
        s.save()
        data = {
            'total_price': 12.0,
         }
        response = self.client.post(reverse('update_wishlist', kwargs={'pk':s.id}),data)

        self.assertTrue(Wishlist.objects.all().filter(total_price=12.0).exists(), msg="User did not actually get updated with new name!")
        self.assertEquals(response.json()['total_price'], str(12.0), msg="Change did not happen! Error")

    def test_delete_wishlist(self):
        u = User.objects.create(first_name='Test', last_name='User', email='test@user.com', phone_number='123-456-7890')
        u.save()
        s = Wishlist.objects.create(user = u, total_price = 500.26, num_items=5)
        s.save()

        response = self.client.post(reverse('destroy_wishlist', kwargs={'pk': s.id}))
        self.assertFalse(Wishlist.objects.all().filter(pk=1).exists(), msg="Wishlist did not actually get deleted in db with API!")
        self.assertEquals(response.json()['status_code'], str(200), msg="Wishlist not deleted!")

    def test_get_invalid_wishlist(self):
        response = self.client.get(reverse('get_wishlists', kwargs={'pk':20}))
        self.assertNotContains(response, 'total_price', msg_prefix= str(response))

    def test_delete_invalid_wishlist(self):
        response = self.client.post(reverse('destroy_wishlist', kwargs={'pk': 12}))
        self.assertEquals(response.json()['status_code'], str(500), msg="Wishlist that shouldn't exist, got deleted")

    def test_update_invalid_wishlist(self):
        data = {
            'total_price' : 6.99,
        }
        response = self.client.post(reverse('update_snack', kwargs={'pk': 12}),data)
        self.assertFalse(Wishlist.objects.all().filter(total_price='6.99').exists(), msg="Wishlist that does not exist, now exists, bad!")
        self.assertEquals(response.json()['status_code'], '500', msg="Object shouln't exist yet was updated")

    #tearDown method is called after each test
    def tearDown(self):

        pass #nothing to tear down


"""
Provides testing utilities.
"""

import numpy as np
import matplotlib.pyplot as plt
import bma.linear_averaging as linear_averaging
from sklearn.metrics import mean_squared_error, r2_score
import time

plt.style.use("ggplot")
plt.rcParams.update({'font.size': 16})

def simulate_data(nobs, weights, cor=0.5):
    """Simulate a simple dataset where 10 predictors are irrelevant and 5 increasingly relevant.

    Predictors are jointly gaussian with 0 mean, given correlation and variance 1.

    Residuals are gaussian with 0 mean and variance 1.

    Parameters
    ----------
    nobs : int {1, .., inf}
        number of samples to be drawn
    weights : np.ndarray
        model weights
    cor : float (-1, 1)
        correlation of the predictors

    Returns
    -------
    tup
        with [0] feature matrix and [1] response
    """

    ndim = len(weights)
    cov = (1 - cor) * np.identity(ndim) + cor * np.ones((ndim, ndim))
    mean = np.zeros(ndim)

    X = np.random.multivariate_normal(mean, cov, nobs)
    e = np.random.normal(0, 1, nobs)
    y = np.dot(X, weights) + e

    return (X, y)

def replicate_trial(trial, n):
    """Repeat a random trial n times.

    Parameters
    ----------
    trial : func
        call to the generating function
    n : int {1, .., inf}
        number of trials

    Returns
    -------
    np.ndarray
        where rows are trials and columns are variables
    """

    return np.array([trial() for i in range(n)])

start = time.time()

# set coefficients
weights = np.hstack((np.zeros(10), np.arange(0.2, 1.2, 0.2)))

# simulate data
np.random.seed(2015)
X, y = simulate_data(100, weights, 0.5)

# mcmc approximation
mc3 = linear_averaging.LinearMC3(X, y, 15**2, 1/3)
mc3.select(niter=10000, method="random")
mc3.estimate()

assert 1==0
# full enumeration
enumerator = linear_averaging.LinearEnumerator(X, y, 15**2, 1/3)
enumerator.select()
enumerator.estimate()

computed_weights = enumerator.estimates['coefficients']
y_hat =  np.dot(X, computed_weights[1:16]) # ή 0:15, δες το (mallon 1:16 thelei γτ βγαζει καλυτερο R2 score
end = time.time()
print(r2_score(y,y_hat))  #bale ayto sto base.py
print(end-start)


import matplotlib.pyplot as plt

from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

from data_access import load_iris_data
from plots.learning_results import plot_decision_regions

iris_df = load_iris_data()

predictor_columns = ["petal length (cm)", "petal width (cm)"]
target_column = "int class"

X = iris_df[predictor_columns].values
y = iris_df[target_column].values

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0)

sc = StandardScaler()
sc.fit(X_train)
X_train_std = sc.transform(X_train)
X_test_std = sc.transform(X_test)

lr = LogisticRegression(C=1000.0, random_state=0)
lr.fit(X_train_std, y_train)
plot_decision_regions(X=X_train_std, y=y_train, X_test=X_test_std, y_test=y_test, classifier=lr)

plt.xlabel('petal length [standardized]')
plt.ylabel('petal width [standardized]')
plt.legend(loc='upper left')
plt.show()


#!/usr/bin/python3.4
# -*-coding:Utf-8 -*

from argparse import ArgumentParser
from math import inf
from time import time

from lib.eventdispatcher import EventDispatcher
from mario.bridge.frame_reader import FrameReader
from mario.bridge.config import Config
from .EvolutiveGenerator.Generator import Generator
from .factories.IAFactory import IAFactory
from .graduators.IAGraduator import IAGraduator
from .graduators.GameOptimizer import GameOptimizer
from .Writer.Writer import Writer
from .Logger.FileLogger import FileLogger
from .Logger.ConsoleLogger import ConsoleLogger
from .Writer.PathManager import PathManager
from .Writer.Reader import Reader

def instanciateGenerator(show):
	event_dispatcher = EventDispatcher()
	FrameReader(event_dispatcher)
	GameOptimizer(event_dispatcher)
	return Generator(IAFactory, IAGraduator(event_dispatcher, show), [Writer(), FileLogger(), ConsoleLogger()],
		lambda state: True in [score.percent >= 1. for score, individual in state.grading]
	)

def checkProcessusExists(processus_id):
	if not Reader.processusExists(processus_id):
		raise ValueError("Processus with id={} doesn't exist.".format(processus_id))

def new(args):
	"""New processus"""
	population = instanciateGenerator(args.show).process(
		PathManager.newProcessusId(), args.generations, args.pop_length, args.proportion, args.chance
	)

def resume(args):
	"""Resume a processus"""
	checkProcessusExists(args.processus_id)
	population = instanciateGenerator(args.show).resume(Reader.getProcessusState(args.processus_id))

def play(args):
	"""Play the best individual of a processus' last generation"""
	checkProcessusExists(args.processus_id)
	# Get IA
	if args.ia_id is None:
		ia, generation_id = Reader.getBestIa(args.processus_id, args.generation_id)
		print('The best AI is {}.'.format(ia.id), flush=True)
	else:
		ia, generation_id = Reader.getIa(args.processus_id, args.ia_id)
	# Play IA
	event_dispatcher = EventDispatcher()
	FrameReader(event_dispatcher)
	graduator = IAGraduator(event_dispatcher, show=True)
	if args.as_grading:
		print(
			"Attention : Malgré que le visionnage présenté soit le plus proche possible des conditions d'évaluation, des aléas subsistent. "
			"Si vous cherchez à visionner une performance difficile à reproduire, n'hésitez pas à rééssayer plusieurs fois."
		, flush=True)
		GameOptimizer(event_dispatcher)
		graduator.grade(ia, generation_id)
	else:
		graduator.gradeIAWithConfig(ia, Config(True, event_dispatcher))

def print_data(args):
	checkProcessusExists(args.processus_id)

	data = Reader.getData(args.processus_id)
	txt1 = 'Générations,Scores des intelligences'
	for generation_id, grading in data:
		txt1 += '\n' + str(generation_id)
		for result, ia_id in grading:
			txt1 += ',' + str(result['score'])
	txt2 = 'Générations,Scores des intelligences'
	for generation_id, grading in data:
		txt2 += '\n' + str(generation_id)
		for result, ia_id in grading:
			txt2 += ',' + str(result['max_x'])

	path1 = PathManager.getPath(args.processus_id, read_only=True).parent / 'data' / (str(time()) + '.score.csv')
	path2 = PathManager.getPath(args.processus_id, read_only=True).parent / 'data' / (str(time()) + '.distance.csv')
	PathManager.makeDir(path1.parent)
	path1.write_text(txt1)
	path2.write_text(txt2)

# Build parser
parser = ArgumentParser()
subparsers = parser.add_subparsers()

new_parser = subparsers.add_parser('new')
new_parser.add_argument('pop_length', type=int)
new_parser.add_argument('--generations', default=inf, type=int)
new_parser.add_argument('--proportion', default=0.5, type=float)
new_parser.add_argument('--chance', default=0, type=float)
new_parser.add_argument('--show', dest='show', action='store_true')
new_parser.set_defaults(command=new, show=False)

resume_parser = subparsers.add_parser('resume')
resume_parser.add_argument('processus_id', type=int)
resume_parser.add_argument('--show', dest='show', action='store_true')
resume_parser.set_defaults(command=resume, show=False)

play_parser = subparsers.add_parser('play')
play_parser.add_argument('processus_id', type=int)
play_parser.add_argument('--generation_id', type=int)
play_parser.add_argument('--ia_id', type=int)
play_parser.add_argument('--as_grading', dest='as_grading', action='store_true')
play_parser.set_defaults(command=play, as_grading=False)

print_parser = subparsers.add_parser('print')
print_parser.add_argument('processus_id', type=int)
print_parser.set_defaults(command=print_data)

# Parse arguments
args = parser.parse_args()
if hasattr(args, 'command'):
	args.command(args)
else:
	print('No command given, use --help')


from flask import Flask, render_template, request
import pandas as pd
from bokeh.plotting import figure
from bokeh.transform import factor_cmap, factor_mark
from bokeh.models import HoverTool
from bokeh.models import ColumnDataSource

from bokeh.embed import components

app = Flask(__name__)

def load_data():
    # Load the Iris Data Set
    df_cars = pd.read_csv('data/cars_publish.csv', index_col=0)
    return df_cars

df_cars = load_data()
numerics = ['int16', 'int32', 'int64', 'float16', 'float32', 'float64']
y_columns = list(df_cars.select_dtypes(include=numerics))

# Create the main plot
def create_figure(y_column,x_column):

    df_cars = load_data()
    source = ColumnDataSource(data=df_cars)

    hover = HoverTool(
        tooltips=[
            ("Modell", "@Modell"),
            ("price", "@price"),

        ]
    )

    p = figure(title='Cars', tools=[hover])

    #x_column = 'Tillverkningsår'
    #y_column = 'price'

    markers = ['circle']

    p.scatter(x=x_column, y=y_column, source=source,
              legend="car type",
              marker=factor_mark('car type', markers=markers, factors=df_cars['car type'].unique()),
              color=factor_cmap('car type', palette='Category10_3', factors=df_cars['car type'].unique()),
              size=8, alpha=0.25)

    p.xaxis.axis_label = x_column
    p.yaxis.axis_label = y_column
    return p

# Index page
@app.route('/')
def index():
    # Determine the selected feature
    y_column = request.args.get("y_column")
    if y_column == None:
        y_column = "price"

    x_column = request.args.get("x_column")
    if x_column == None:
        x_column = "Tillverkningsår"

    # Create the plot
    plot = create_figure(y_column, x_column)

    # Embed plot into HTML via Flask Render
    script, div = components(plot)
    return render_template("plot.html", script=script, div=div,y_column = y_column,x_column = x_column,y_columns = y_columns)

# With debug=True, Flask server will auto-reload
# when there are code changes
if __name__ == '__main__':
    app.run(port=5000, debug=True)