import numpy as np              # linear algebra
import pandas as pd             # data processing, CSV file I/O (e.g. pd.read_csv)
import matplotlib.pyplot as plt # visualization
import seaborn as sns

from tqdm.auto import tqdm      # progress bar
from itertools import groupby

import random

from pulp import *
#import cython

#%load_ext Cython

from matplotlib import colors
from matplotlib import ticker
import warnings
import datetime

#plt.style.use('default')
warnings.filterwarnings("ignore", category=UserWarning)

def plot_ans(final_ans, labels=None, ticks=None, set_labels_axis=None, franjas_horarias=None, 
             ax_provided=None, colorbar=True, pad=0.15, enc_dict=None):
    """
    Funcion para visualizar por color los eventos del horario asignado a los trabajadores
    
    Parametros:
        final_ans: array que contiene el horario asignado a los trabajadores, codificado en numeros
        labels:    Nombre asignado a los eventos encontrados
        ticks:     Posicion de las labels seleccionadas en la barra de color
        set_labels_axis: diccionario que contiene el titulo, y los labels de los axis
        franjas_horaris: numero de las franjas horarias, sirve para asignar correctamente las horas en el eje x
        ax_provided: ax proviste de un subplots externo a la funcion
        colorbar:  bool, mostrar barra de color con la representacion de estados
    """
    
    c_map = colors.ListedColormap(['black', '#BBD0FF', '#0077B6', '#03045E'][: 4 if labels is None else len(labels)])
    
    if ax_provided is None:
        f,ax=plt.subplots(figsize=(10,8))
    else:
        ax = ax_provided
    
    h = ax.imshow(final_ans, cmap=c_map, extent=(-0.5, final_ans.shape[1]-0.45, final_ans.shape[0]-0.45, -0.5))
    
    ax.xaxis.set_major_locator(ticker.MultipleLocator(base=1))
    ax.yaxis.set_major_locator(ticker.MultipleLocator(base=1))

    ax.xaxis.set_minor_locator(ticker.IndexLocator(base=0.5, offset=1))
    ax.yaxis.set_minor_locator(ticker.IndexLocator(base=0.5, offset=1))
    #ax.grid(color='w', linestyle='-', linewidth=1)
    
    labels = list(enc_dict.keys()) if labels is None else labels
    ticks  = [(len(labels)-1)/(len(labels)*2) + ((len(labels)-1)/len(labels))*i for i in range(len(labels))]
    

    if colorbar:
        c_bar = plt.colorbar(h, orientation='horizontal',
                             format=ticker.FixedFormatter(labels),
                             ticks=ticks, pad=pad)# [0.75*i - 0.75/2 for i in range(1,5)] if ticks is None else ticks)
    
    
    time_init = datetime.datetime(year=1, month=1, day=1, hour=7, minute=30)
    horas = [(time_init + datetime.timedelta(minutes=15*i)).strftime('%H:%M') for i in range(franjas_horarias)] 
        
    ax.grid(which='minor', color='slategrey', linestyle='-', linewidth=1, alpha=0.7)
    ax.set_xticklabels(['']+horas)
    
    if set_labels_axis is None:
        ax.set(title="Horario de Trabajadores Ejemplo", ylabel="Trabajador ID", xlabel="Franja Horaria")
    else:
        ax.set(title=set_labels_axis['title'], ylabel = set_labels_axis['ylabel'], xlabel=set_labels_axis['xlabel'])
    plt.setp(ax.xaxis.get_majorticklabels(), rotation=90, ha='center')
    
    if ax_provided is None:
        plt.show()

#@memory.cache

def duration_block(block, t_zones):
    """
    Algoritmo que calcula la duracion en franjas horarios del bloque
    Parameters: 
    
    block:   {"1-2","1-3","1"} tipo de bloque
    t_zones: Duration del espacio de trabajo "1"
    
    return:   Franjas Horarias ocupadas por todo el bloque
    """
    l = block.split('-')
    
    if len(l)>1:
        t = 1 if l[1]=='2' else 6
        return t_zones + t
    else:
        return t_zones
    

def loop_search(sec, l=[], w=[], n=None, limit=32+6, verbose=True):
    """
    Algoritmo que itera usando recursion sobre todas las soluciones posibles para dar solucion 
    a la secuencia de acuerdo al limite de franjas horarias establecido
    
    Parametros:
    
    sec: Secuencia a buscar solucion: e.g  ['1-2','1-2','1-3']
    l :  Acumulador de las coeficientes que satisfacen la secuancia "sec"
    w :  Acumulador de franjas horarias en la posible solucion
    n :  Numero de coeficiente buscados: n = len(sec), este valor se va 
         reduciendo a medida que se acumulan los coeficientes
    limit: Franjas horarias de trabajo-pausa-almuerza necesarias
    verbose : Imprimir coeficientes de respuesta a medida que se encuentran 
    
    return
        final_r: variable local que referencia a una lista con las soluciones encontradas para la secuencia
    """
    
    if n is None:
        global final_r
        final_r = []
        
    n = len(sec) if n is None else n
    if n >= 1:
        for x in range(4,9):
            w1 = duration_block(block=sec[len(sec)-n], t_zones=x)
            if sum(w)+w1 > limit: break # excede jornada laboral
            
            loop_search(sec, l=l+[x], w=w+[w1], n=n-1, limit=limit, verbose=verbose)   
    else:
        if sum(w)== limit:
            if verbose: print(l)
            final_r = final_r + [l]
            
            
def generate_sol_per_sequence(limit=32+6, incluir_almuerzo=True, verbose=True, max_bloques=5):
    """
    Generar un espacio amplio de busqueda de secuencias simplificadas que sean viables, 
    respetando los limites impuestos
    
    Parametros:
        limit:            Franjas horarias de trabajo-pausa-almuerza necesarias
        incluir_almuerzo: Boleano que representa si considerar el bloque '1-3' que contiene el almuerzo
        verbose:          Imprimir por pantalla el numero de secuencias simplificadas encontradas
        max_bloques:      Numero maximo de bloques de las sequencias a buscar solucion
    Return:
        r_seq: Lista que contiene las secuencias simplificadas encontradas y el conjunto de coeficientes de solucion
    """
    if verbose: print("Forma general de secuencias con soluciones viables: ")
        
    r_seq = []
    for n_12_antes in range(max_bloques):
        for n_12_despues in range(max_bloques):
            t_sec  = ['1-2' for i in range(n_12_antes)]
            t_sec += ['1-3'] if incluir_almuerzo else []
            t_sec += ['1-2' for i in range(n_12_despues)] + ['1']
            
            # se genera una variable global 'final_r'
            loop_search(t_sec, limit=limit, verbose=False)
            
            if len(final_r)>0:
                
                if verbose: print(t_sec, len(final_r))
                    
                r_seq.append([t_sec, final_r])
    return r_seq

def retrieve_sequence(base_seq, config_block):
    """
    Funcion que estructura los coeficientes en la secuencia provista
    
    Parametros:
    
    base_seq : Secuencia provista         
    config_block:  Coeficientes a asignar
    
    Retorna la serie de eventos sin la parte temporal.  
    e.g.  
    
    base_seq     = ['1-3','1-2','1-2', '1']
    config_block = [    6,    8,    8,  8 ] 
    
    return  ->  11111133333311111111211111111211111111
    """
    seq = []
    for block, config in zip(base_seq, config_block):
        sp = block.split('-')
        if len(sp)>1:
            extra = [2] if sp[1]=='2' else [3 for i in range(6)]
        else:
            extra = []

        f = [int(sp[0]) for i in range(config)]
        seq += f + extra

    return seq


def return_complete_map(limit, incluir_almuerzo, verbose, max_bloques):    
    """
    Mapa de eventos de todas las serie de eventos encontradas sin parte temporal
    Parameters:
        list_secuencias: lista que contiene cada secuancia y sus posibles soluciones de acuerdo a las restricciones
    Return:
        lista de las secuencia de eventos encontrada
    """
    
    list_secuencias = generate_sol_per_sequence(limit, incluir_almuerzo, verbose, max_bloques)
    
    complete_map = []
    #########  TC, entre semana
    for (base_seq, configs) in list_secuencias:
        for config_block in configs:
            complete_map.append(retrieve_sequence(base_seq, config_block))
            
    return np.array(complete_map, dtype=np.int32)

enc_dict = {'Nada': 0, 'Trabaja': 1, 'Pausa Activa': 2, 'Almuerza': 3}

def generar_espacios_inicio(complete_map, posibilidades=9, key_prefix='TC', enc_dict=None, incluir_almuerzo=True, verbose=True): 
    """
    Prueba cada series de eventos con su parte temporal de acuerdo a los inicios posibles (Posibilidades), 
    y sus restricciones horarias (e.g. hora de almuerzo) dentro de la configuracion horaria y de acuerdo 
    al numero de "Pausas Activas".
    
    Guarda solamente los indices (idx) del mapa completo de posibles series de events (complete_map)
    Parameters:
    """
    
    espacio_posibilidad = {}
    for i in range(posibilidades):
        t_r = [] #{}
        for idx, t_seq in enumerate(complete_map):
            
            if incluir_almuerzo:
                almuerzo_init = np.where(t_seq == enc_dict['Almuerza'])[0][0] + i # posicion del almuerzo de acuerdo a la posibilidad
            else:
                almuerzo_init = None
            # 7:30 to 18:45, 46 franjas horarias
                # 11:30, position 16, 13:30,  position 24
            if (almuerzo_init in range(16,25)) or (almuerzo_init is None):
                #t_np = (t_seq==2).sum() # calcula el numero de descansos en la secuencia
                #if t_r.get(t_np) is None:
                #    t_r[t_np]  = [idx]
                #else:
                #    t_r[t_np] += [idx]
                t_r.append(idx)
                
        espacio_posibilidad[f"{key_prefix}_{i}"] = t_r
        if verbose:
            print(f"Posibilidad {key_prefix}_{i}, sequencias validas: {len(t_r):<4}")
    return espacio_posibilidad

def generar_espacios_almuerzos(complete_map, posibilidades=9, key_prefix='TC', enc_dict=None, verbose=True):
    """
    Prueba cada series de eventos con su parte temporal de acuerdo a los inicios posibles (Posibilidades), 
    y sus restricciones horarias (e.g. hora de almuerzo) dentro de la configuracion horaria y de acuerdo 
    al numero de "Pausas Activas".
    
    Guarda solamente los indices (idx) del mapa completo de posibles series de events (complete_map)
    """
    espacio_posibilidad = {}
    for i in range(posibilidades):
        t_r = {}
        for idx, t_seq in enumerate(complete_map):
            
            # posicion del almuerzo de acuerdo a la posibilidad
            almuerzo_init = np.where(t_seq == enc_dict['Almuerza'])[0][0] + i 
            
            # 7:30 to 18:45, 46 franjas horarias
            # 11:30, position 16, 13:30,  position 24
            
            if almuerzo_init in range(16,25):
                
                if t_r.get(almuerzo_init) is None:
                    t_r[almuerzo_init]  = [idx]
                else:
                    t_r[almuerzo_init] += [idx]
        
        espacio_posibilidad[f"{key_prefix}_{i}"] = t_r
        
        if verbose:
            print(f"Posibilidad {key_prefix}_{i}, sequencias validas: almuerzos {list(t_r.keys())} total: {sum([len(el[1]) for el in t_r.items()]):<4}")
    return espacio_posibilidad


def generate_complete(espacio_worker, complete_map, limit, limit_day):
    """
    Insertar espacios nada de acuerdo a limitaciones
    """
    espacio_compl = []
    for pos in espacio_worker.keys():
        base = np.zeros(shape=(len(espacio_worker[pos]), limit_day), dtype='int')
        
        pos_int = int(pos.split("_")[1])
        for i, idx in enumerate(espacio_worker[pos]):
            base[i, pos_int:pos_int+limit] = complete_map[idx]

        espacio_compl.append(base)
    return np.concatenate(espacio_compl, axis=0)

def cost_function(final_array, demanda, volume=False):
    """
    Funcion de costo a minimizar:
    Parametros:
        final_array: Horario de eventos de los empleados asignados
        demanda:     Demanda por franja horaria a optimizar
    Return
        Suma de las Diferencias con valor positivo entre Capacidad y Demanda por franja horaria
    """
    
    capacidad = (final_array==1).sum(axis=1 if volume else 0)
    ans = (demanda - capacidad)
    return ((ans>0) * ans).sum(), -((ans<0) * ans).sum() #unattended demand, overcapacity

def array_sol(x, y, espacios, dicts=True):
    """
    Retrieve Solution idx
    """
    if dicts:
        TC_idx_l = [k for k, v in x.items() for _ in range(int(v.varValue))] # considering that `v.varValue>0`
        MT_idx_l = [k for k, v in y.items() for _ in range(int(v.varValue))]
    else:
        TC_idx_l = [i for i, el in enumerate(x) for _ in range(int(el.varValue))] # considering that `el.varValue>0`
        MT_idx_l = [i for i, el in enumerate(y) for _ in range(int(el.varValue))]
    
    TC_arr = espacios['TC'][TC_idx_l]
    MT_arr = espacios['MT'][MT_idx_l]
    
    if len(x)==0:
        return MT_arr
    elif len(y)==0:
        return TC_arr
    else:
        return np.concatenate([MT_arr, TC_arr], axis=0)
    
def lp_method(demanda, franjas, espacios, n_workers, name_idx):
    """
    Solution for one day
    
    Input:
        demanda
        franjas
        espacios
        n_workers
    """
    
    # problems
    prob = LpProblem(f"Day for {name_idx}", LpMinimize)
    
    # variables
    x = LpVariable.dicts("x", range(espacios['TC'].shape[0]), lowBound=0, upBound=n_workers['TC'], cat=LpInteger)
    y = LpVariable.dicts("y", range(espacios['MT'].shape[0]), lowBound=0, upBound=n_workers['MT'], cat=LpInteger)
    
    ##### to save cost function result if > 0
    differences = LpVariable.dicts("diff", range(franjas['total']), lowBound=0, cat=LpInteger)
    
    # auxiliar function
    esp_TC = espacios['TC']==1 ; esp_MT = espacios['MT']==1
    cap_TC = lambda f: lpSum(el[f] * x[i] for i, el in enumerate(esp_TC) )
    cap_MT = lambda f: lpSum(el[f] * y[i] for i, el in enumerate(esp_MT) )
    ## calculate intermidiate states forms, capacity sum per day and franja
    Cap_TC = {f: cap_TC(f) for f in range(franjas['inter'])}
    Cap_MT = {f: cap_MT(f) for f in range(franjas['inter'])}
    
    # Add the objective function: minimize the sum of positive differences
    prob += lpSum(differences)

    for f in range(franjas['inter']):
        prob += differences[f] >= demanda[f] - Cap_TC[f] - Cap_MT[f]
    
    ###################### numero de turnos selecionados == trabajadores 
    prob += lpSum(v for k,v in x.items()) == n_workers['TC']
    prob += lpSum(v for k,v in y.items()) == n_workers['MT']
    #prob += lpSum(y[i] for i in range(pos_dict['MT'])) == n_workers['MT']
    
    ###################### al menos un trabajador por franja horaria
    for f in range(franjas['inter']): # solo donde existe demanda
        prob += Cap_TC[f] + Cap_MT[f] >= 1

    ######################
    status = prob.solve(PULP_CBC_CMD(msg=0))
    if status == LpStatusOptimal:
        return x, y, value(prob.objective)
    else:
        return 0, 0, None
    #print("#"*20)
    #if prob.status==1:
    #print("Status ", LpStatus[prob.status])

    #for el in prob.constraints:
    #    print(el, prob.constraints[el])
    #print("#"*20)
    #for var in prob.variables():
    #    print(var, value(var))
    #    print("Optimal solution", value(prob.objective))
    

def lp_method_etapa2(demanda, franjas, n_workers, espacios, espacios_turno, days, name_idx):
    # problems
    prob = LpProblem(f"Days for {name_idx}", LpMinimize)
    
    #variables, weekly
    x = LpVariable.matrix('x',
                          indices=(range(days), range(espacios['TC'].shape[0])), 
                          lowBound=0, upBound=n_workers['TC'], cat=LpInteger) #LpBinary)
    
    y = LpVariable.matrix('y',
                          indices=(range(days), range(espacios['MT'].shape[0])), 
                          lowBound=0, upBound=n_workers['MT'], cat=LpInteger)
    
    ##### to save cost function result if > 0
    differences = LpVariable.matrix("diff", 
                                    indices=(range(days), range(franjas['total'])), lowBound=0, cat=LpInteger)
    
    ########## auxiliar function, capacity per day
    ######
    esp_TC = espacios['TC']==1 ; esp_MT = espacios['MT']==1
    cap_TC = lambda f,d: lpSum(el[f] * x[d][i] for i, el in enumerate(esp_TC) )
    cap_MT = lambda f,d: lpSum(el[f] * y[d][i] for i, el in enumerate(esp_MT) )
    ## calculate intermidiate states forms, capacity sum per day and franja
    Cap_TC = {(f,d):cap_TC(f,d) for d in range(days) for f in range(franjas['inter'])}
    Cap_MT = {(f,d):cap_MT(f,d) for d in range(days) for f in range(franjas['inter'])}
    ######
    
    # Add the objective function: minimize the sum of positive differences
    prob += lpSum(differences)
    
    # define cost values as demanda - capacity (TC/MT), '>=' to store only if it is '>' lowBound=0  
    for d in range(days): 
        for f in range(franjas['inter']):
            prob += differences[d][f] >= demanda[d][f] - Cap_TC[(f,d)] - Cap_MT[(f,d)]
    
    ###################### numero de turnos selecionados == trabajadores
    for d in range(days):
        prob += lpSum(el for el in x[d]) == n_workers['TC']
        prob += lpSum(el for el in y[d]) == n_workers['MT']
    
    ###################### al menos un trabajador por franja horaria
    for d in range(days):
        for f in range(franjas['inter']): # solo donde existe demanda
            prob += Cap_TC[(f,d)] + Cap_MT[(f,d)] >= 1 #demanda[f]
            
    ###################### cada dia de la semana el mismo turno por trabajador
    ######## first day values variable to reuse result
    s_TC_d1 = LpVariable.dicts("TC_d1", range(len(espacios_turno['TC'].keys())), lowBound=0, cat=LpInteger)
    s_MT_d1 = LpVariable.dicts("MT_d1", range(len(espacios_turno['MT'].keys())), lowBound=0, cat=LpInteger)
    for cnt, (k,v) in enumerate(espacios_turno['TC'].items()):
        prob +=  s_TC_d1[cnt]==lpSum(x[0][i] for i in v)
    
    for cnt, (k,v) in enumerate(espacios_turno['MT'].items()):
        prob +=  s_MT_d1[cnt]==lpSum(y[0][i] for i in v)
    ########
    
    for d in range(1,days):
        for cnt, (k,v) in enumerate(espacios_turno['TC'].items()): # en cada grupo iniico_almuerzo, se seleciona el mismo numero cada dia
            #idx_i, idx_f = v
            prob += lpSum(x[d][i] for i in v) == s_TC_d1[cnt]
        
        for cnt, (k,v) in enumerate(espacios_turno['MT'].items()):
            #idx_i, idx_f = v
            prob += lpSum(y[d][i] for i in v) == s_MT_d1[cnt]
    ######################
    
    status = prob.solve(PULP_CBC_CMD(msg=0))
    if status == LpStatusOptimal:
            return x, y, value(prob.objective)
    else:
        return 0, 0, None
    #if prob.status==1:
    #print("Status ", LpStatus[prob.status])

    #for el in prob.constraints:
    #    print(el, prob.constraints[el])
    
    #for var in prob.variables():
    #    print(var, value(var))

import string

def return_pos(row):
    count = 0
    for el in row:
        if el!=0: break
        count +=1 
    return count

def group_dist(espacio):
    r = {}
    for i, el in enumerate(espacio):
        if r.get(el) is None:
            r[el] = [i]
        else:
            r[el] +=  [i]
    return r

### the mapped numbers cannot be decomposed into numbers inside selected indices   
#### considerando grupos de turnos y horario de almuerzo
espacio_TC_inter_pos = np.array([str(return_pos(row))+'_'+str(np.where(row==3)[0][0]) \
                                 for row in espacio_TC_inter_compl])
espacio_TC_inter_pos = group_dist(espacio_TC_inter_pos)

#espacio_TC_inter_pos = group_dist(espacio_TC_inter_pos)
#### considerando grupos solo de turnos
espacio_MT_inter_pos = np.array([return_pos(row) for row in espacio_MT_inter_compl])
espacio_MT_inter_pos = group_dist(espacio_MT_inter_pos)

#espacio_TC_alm = [abs(hash(str(el)))%10**6 for el in espacio_TC_alm]

def return_zeros(row):
    """
    Return the final no demand zones number
    """
    
    if row[-1]!=0:
        return 0
    else:
        count = 0
        for el in row[::-1]:
            if el!=0: break
            count +=1 
        return count

def return_alm_config(temp_array):
    """
    extraer almuerzo y franja de inicio de trabajo en array de configuracion diaria 
    """
    r = []; pos_l = []
    for row in temp_array:
        pos = 0
        for el in row:
            if el!=0: break
            pos+=1
        
        temp = np.where(row==3)[0]
        temp = -1 if len(temp)==0 else temp[0]
        
        r.append(temp)
        pos_l.append(f"{'MT_' if temp==-1 else 'TC_'}{pos}")
    
    return r, pos_l


def organize_data(configuration_group):
    """
    Retorna el array de configuraciones organizado por contrato y lugar de almuerzo
    """
    new_group = []
    for i, temp_config in enumerate(configuration_group): #final_config[0]):

        alm_config, pos_config = return_alm_config(temp_config)
        
        temp_list = []
        for  t_pos, t_alm, t_row_config in zip(pos_config, alm_config, temp_config):
            temp_list.append([t_pos+','+ str(t_alm), t_row_config])
        
        temp_list = sorted(temp_list, key=lambda t: t[0]) # sort by string encoded config
        #print([el[0] for el in temp_list])
        
        t_comb = [el[0].split(',')[0] for el in temp_list]
        t_final = np.concatenate([el[1][np.newaxis,...] for el in temp_list], axis=0)
        
        #new_group.append([t_comb, t_final])
        new_group.append(t_final)

    return new_group

for idx in range(5):
    final_config[idx]['inter_semana'] = organize_data(final_config[idx]['inter_semana'])
    final_config[idx]['sabado']       = organize_data([final_config[idx]['sabado']])[0]

import datetime

enc_dict_1 = {'Nothing': 0, 'Work': 1, 'Active Break': 2, 'Lunch': 3}
enc_dict_rev = {v:k for k,v in enc_dict_1.items()}
#print(enc_dict_rev)

def gen_format_output(best_result, workers_suc, franjas_horarias, fecha):
    """
    best_result, array por dia de respuesta
    """
    # horas del dia establecidas
    time_init = datetime.datetime(year=1, month=1, day=1, hour=7, minute=30)
    franjas = [(time_init + datetime.timedelta(minutes=15*i)).strftime('%H:%M') for i in range(franjas_horarias)]
    
    answer_df = None
    for i, row in enumerate(workers_suc.iterrows()):
        row = row[1]
        
        t_df = pd.DataFrame()
        t_df['interval'] = range(30,30+franjas_horarias)
        t_df['suc_cod']  = row['suc_cod']
        t_df['document'] = row['documento']

        t_df['date'] = fecha #fecha_enc[dia] #'22/04/2022'
        t_df['hour']  = franjas

        t_df['state'] = best_result[i].astype('int')
        t_df['state'] = t_df['state'].apply(lambda t: enc_dict_rev[t])

        if answer_df is None:
            answer_df = t_df
        else:
            answer_df = pd.concat([answer_df, t_df], axis=0, ignore_index=True)
            
    return answer_df