# AUTOGENERATED! DO NOT EDIT! File to edit: 01_vulnerabilidad.ipynb (unless otherwise specified).

__all__ = ['show_feature_importances', 'mostrar_coeficientes_PLS', 'agregar_conteo_pruebas',
           'agregar_tasas_municipales', 'caracteristicas_modelos_municipios', 'ajustar_pls_letalidad',
           'ajustar_pls_columna', 'ajustar_rf_letalidad', 'calificar_municipios_letalidad',
           'calificar_municipios_letalidad_formato_largo', 'seleccionar_caracteristicas',
           'calcular_indices_vulnerabilidad', 'calcular_indices_vulnerabilidad_formato_largo',
           'calcular_periodo_vulnerabilidad', 'periodo_vulnerabilidad_con_dataframe',
           'calcular_periodo_vulnerabilidad_2', 'agregar_periodo_vulnerabilidad', 'calcular_vulnerabilidad_urbana',
           'agregar_vulnerabilidad_entidades', 'guardar_resultados_csv', 'mapas_serie_letalidad',
           'mapas_serie_vulnerabilidad', 'guardar_shape_vulnerabilidad', 'checkpoint_vulnerabilidad',
           'cargar_checkpoint_vulnerabilidad']

# Cell
from ipywidgets import IntProgress
from IPython.display import display
import time

from .datos import *

import pandas as pd
import geopandas as gpd

import glob
import os
import matplotlib.pyplot as plt
import numpy as np

from datetime import timedelta
import datetime

from sklearn.feature_selection import SelectKBest, chi2
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.feature_selection import SelectFromModel
from sklearn.cross_decomposition import PLSRegression

import imageio
from pygifsicle import optimize

# Cell

def show_feature_importances(rf):
    importances = rf.feature_importances_

    std = np.std([tree.feature_importances_ for tree in rf.estimators_],
                 axis=0)
    indices = np.argsort(importances)[::-1]

    # Print the feature ranking
    print("Feature ranking:")

    for f in range(X.shape[1]):
        print("%d. feature %s (%f)" % (f + 1, X.columns[indices[f]], importances[indices[f]]))

    # Plot the feature importances of the forest
    f, ax = plt.subplots(figsize=(15, 10))
    plt.title("Feature importances")
    plt.bar(range(X.shape[1]), importances[indices],
           color="r", yerr=std[indices], align="center")
    plt.xticks(range(X.shape[1]), X.columns[indices])
    plt.xlim([-1, X.shape[1]])
    plt.show()

# Cell

def mostrar_coeficientes_PLS(pls):
    f, ax = plt.subplots(figsize=(15, 10))
    sns.barplot(y='nombre', x='coef', data=feats_df, color='c')

# Cell

def agregar_conteo_pruebas(covid_municipal, solo_covid=True):
    cols_localidad = ['ENTIDAD_RES', 'MUNICIPIO_RES', 'CLAVE_MUNICIPIO_RES']
    count_tested = covid_municipal[cols_localidad + ['conteo']].groupby(cols_localidad).sum()
    count_tested.reset_index(inplace=True)
    count_tested.rename(columns={'conteo': 'total_pruebas'}, inplace=True)

    covid_municipal = covid_municipal.merge(count_tested, on=cols_localidad, how='left')


    covid_municipal['casos_frac'] = 100 * covid_municipal['conteo'] / covid_municipal['total_pruebas']
    covid_municipal['tasa_covid_letal'] = 100 * covid_municipal['defunciones'] / covid_municipal['conteo']

    if solo_covid:
        covid_municipal = covid_municipal.query('RESULTADO == "Positivo SARS-CoV-2"')
    return covid_municipal

# Cell

def agregar_tasas_municipales(casos_mun_df):
    casos_mun_df['covid_confirmados_100k'] = 100000 * casos_mun_df['conteo'] / casos_mun_df['pob2020']
    casos_mun_df['covid_defun_100k'] = 100000 * casos_mun_df['defunciones'] / casos_mun_df['pob2020']
    casos_mun_df['tasa_covid_letal'] = 100 * casos_mun_df['defunciones'] / casos_mun_df['conteo']

    # covid_municipal = covid_municipal.query('RESULTADO == "Positivo SARS-CoV-2"').copy()

    # casos_mun_df = gpd.GeoDataFrame(casos_mun_df, geometry='geometry')

    return casos_mun_df

# Cell

def caracteristicas_modelos_municipios(mun_df, poblaciones=False, i_vuln=False):
    pob_vars = []
    if not poblaciones:
        pob_vars = list(mun_df.columns[mun_df.columns.str.contains('_pob')])
        pob_vars = pob_vars + ['mayores_65', 'pob2020', 'pt_2015', 'pob_menore', 'poblacion', 'pob_total']
        pob_vars = pob_vars + ['sin_dere_1', 'sin_derech', 'carencias_']
    i_vuln_vars = []
    if not i_vuln:
        i_vuln_vars = ['i_vuln_salud', 'i_vuln_cobertura', 'i_vuln_econo', 'i_vuln_social',
                       'i_vuln_gen', 'i_vuln_infraestructura']

    columnas_numericas = mun_df.select_dtypes(include=np.number).columns

    otras_vars = ['covid_defun_100k', 'tasa_covid_letal', 'defunciones',
                  'INTUBADO_BIN', 'UCI_BIN', 'UCIs', 'tasa_uci', 'total_pruebas',
                  'area_cart', 'area', 'densi', 'casos_frac', 'area_km2', 'conteo',
                  'oid', 'covid_confrimados_100k', 'id', 'vulnerabilidad_ambiental_num',
                  'covid_confirmados_100k', 'index']

    caracteristicas = list(set(columnas_numericas).difference(pob_vars + i_vuln_vars + otras_vars))

    return caracteristicas


# Cell

def ajustar_pls_letalidad(municipios_df, caracteristicas, min_casos=20, min_defunciones=0):
    data_train = municipios_df.loc[municipios_df[caracteristicas].notna().all(axis=1)]

    X = data_train.query(f'(conteo > {min_casos}) & (defunciones > {min_defunciones})')[caracteristicas]
    Y = data_train.query(f'(conteo > {min_casos}) & (defunciones > {min_defunciones})')['tasa_covid_letal']

    # X['i_vuln_econo'] = -X['i_vuln_econo']

    pls2 = PLSRegression(n_components=1)
    pls2.fit(X, Y)
    pls2.coef_ = pls2.coef_.flatten()

    return pls2

# Cell

def ajustar_pls_columna(municipios_df, caracteristicas, columna, min_casos=20, min_defunciones=0):
    data_train = municipios_df.loc[municipios_df[caracteristicas].notna().all(axis=1)]

    X = data_train.query(f'(conteo > {min_casos}) & (defunciones > {min_defunciones})')[caracteristicas]
    try:
        Y = data_train.query(f'(conteo > {min_casos}) & (defunciones > {min_defunciones})')[columna]
    except KeyError:
        print(f"No existe la columna {columna}")
    # X['i_vuln_econo'] = -X['i_vuln_econo']

    pls2 = PLSRegression(n_components=1)
    pls2.fit(X, Y)
    pls2.coef_ = pls2.coef_.flatten()

    return pls2

# Cell

def ajustar_rf_letalidad(municipios_df, caracteristicas, min_casos=20, min_defunciones=0, random_seed=0):
    data_train = municipios_df.loc[municipios_df[caracteristicas].notna().all(axis=1)]

    X = data_train.query(f'(conteo > {min_casos}) & (defunciones > {min_defunciones})')[caracteristicas]
    Y = data_train.query(f'(conteo > {min_casos}) & (defunciones > {min_defunciones})')['tasa_covid_letal']

    # X['i_vuln_econo'] = -X['i_vuln_econo']
    rf = RandomForestRegressor(random_state=random_seed)
    rf.fit(X, Y)

    return rf

# Cell

def calificar_municipios_letalidad(mun_df, regr, caracteristicas, etiqueta):
    not_na_row = mun_df[caracteristicas].notnull().all(axis=1)
    X = mun_df.loc[not_na_row, caracteristicas]
    Y_pred = regr.predict(X)
    mun_df.loc[not_na_row, etiqueta] = Y_pred

    return mun_df


# Cell

def calificar_municipios_letalidad_formato_largo(mun_df, regr, caracteristicas, modelo, dia_ajuste):
    mun_df = mun_df.copy()
    not_na_row = mun_df[caracteristicas].notnull().all(axis=1)
    X = mun_df.loc[not_na_row, caracteristicas]
    Y_pred = regr.predict(X)
    mun_df['modelo'] = modelo
    mun_df['dia_ajuste'] = dia_ajuste
    mun_df.loc[not_na_row, 'valor'] = Y_pred

    return mun_df


# Cell

def seleccionar_caracteristicas(regr, X, caracteristicas):
    sel = SelectFromModel(regr, prefit=True)

    caracteristicas_selec = list(X[caracteristicas].columns[sel.get_support()])
    return caracteristicas_selec

# Cell

def calcular_indices_vulnerabilidad(fecha_0, fecha_1, vulnerabilidad='fecha_0'):
    covid_municipal_0 = tabla_covid_indicadores_municipales(fecha_0)
    covid_municipal_0 = agregar_tasas_municipales(covid_municipal_0)

    covid_municipal_1 = tabla_covid_indicadores_municipales(fecha_1)
    covid_municipal_1 = agregar_tasas_municipales(covid_municipal_1)

    caracteristicas = caracteristicas_modelos_municipios(covid_municipal_0)

    resultados = covid_municipal_1.copy()

    # rf = ajustar_rf_municipios(covid_municipal_0, fecha_0, caracteristicas)
    # caracteristicas_rf = seleccionar_caracteristicas(rf, covid_municipal_0, caracteristicas)
    # resultados = calificar_municipios_letalidad(resultados, rf, caracteristicas, etiqueta='i_RF_0')

    pls = ajustar_pls_letalidad(covid_municipal_0, caracteristicas)
    resultados = calificar_municipios_letalidad(resultados, pls, caracteristicas, etiqueta='i_PLS_0')

    pls = ajustar_pls_letalidad(covid_municipal_1, caracteristicas)
    resultados = calificar_municipios_letalidad(resultados, pls, caracteristicas, etiqueta='i_PLS_1')

    # pls = ajustar_pls_letalidad(covid_municipal_0, caracteristicas_rf)
    # resultados = calificar_municipios_letalidad(resultados, pls, caracteristicas, etiqueta='i_PLS_RF_1')

    return resultados


# Cell

def calcular_indices_vulnerabilidad_formato_largo(fecha_0, fecha_1):
    covid_municipal_0 = tabla_covid_indicadores_municipales(fecha_0)
    covid_municipal_0 = agregar_tasas_municipales(covid_municipal_0)

    covid_municipal_1 = tabla_covid_indicadores_municipales(fecha_1)
    covid_municipal_1 = agregar_tasas_municipales(covid_municipal_1)

    caracteristicas = caracteristicas_modelos_municipios(covid_municipal_0)

    # rf = ajustar_rf_municipios(covid_municipal_0, fecha_0, caracteristicas)
    # caracteristicas_rf = seleccionar_caracteristicas(rf, covid_municipal_0, caracteristicas)
    # resultados = calificar_municipios_letalidad(resultados, rf, caracteristicas, etiqueta='i_RF_0')

    pls = ajustar_pls_letalidad(covid_municipal_0, caracteristicas)
    resultados = calificar_municipios_letalidad_formato_largo(covid_municipal_0, pls, caracteristicas,
                                                modelo='PLS', dia_ajuste=fecha_0)

    pls = ajustar_pls_letalidad(covid_municipal_1, caracteristicas)
    resultados_temp = calificar_municipios_letalidad_formato_largo(covid_municipal_1, pls, caracteristicas,
                                                     modelo='PLS', dia_ajuste=fecha_1)

    resultados = pd.concat([resultados, resultados_temp], ignore_index=True)
    # pls = ajustar_pls_letalidad(covid_municipal_0, caracteristicas_rf)
    # resultados = calificar_municipios_letalidad(resultados, pls, caracteristicas, etiqueta='i_PLS_RF_1')

    return resultados


# Cell

def calcular_periodo_vulnerabilidad(inicio, fin, min_defunciones=-1):
    inicio = pd.to_datetime(inicio, yearfirst=True)
    fin = pd.to_datetime(fin, yearfirst=True)

    fechas = pd.date_range(inicio, fin)
    resultados = []
    modelos = []

    asegura_archivos_covid_disponibles(fechas)

    f = IntProgress(min=0, max=len(fechas) - 1) # instantiate the bar
    display(f) # display the bar

    for count, fecha in enumerate(fechas):
        covid_municipal = tabla_covid_indicadores_municipales(fecha.strftime("%y%m%d"))
        covid_municipal = agregar_tasas_municipales(covid_municipal)

        caracteristicas = caracteristicas_modelos_municipios(covid_municipal)
        pls = ajustar_pls_letalidad(covid_municipal, caracteristicas, min_defunciones=min_defunciones)
        df = calificar_municipios_letalidad_formato_largo(covid_municipal, pls, caracteristicas,
                                                    modelo='PLS', dia_ajuste=fecha)
        resultados.append(df)
        modelo = pd.DataFrame({'caracteristica': caracteristicas, 'coef': pls.coef_})
        modelo['dia_ajuste'] = fecha
        modelo['modelo'] = 'PLS'
        modelos.append(modelo)

        rf = ajustar_rf_letalidad(covid_municipal, caracteristicas, min_defunciones=min_defunciones)
        df = calificar_municipios_letalidad_formato_largo(covid_municipal, rf, caracteristicas,
                                            modelo='RF', dia_ajuste=fecha)
        resultados.append(df)
        modelo = pd.DataFrame({'caracteristica': caracteristicas, 'coef': rf.feature_importances_})
        modelo['dia_ajuste'] = fecha
        modelo['modelo'] = 'RF'
        modelos.append(modelo)

        f.value = count

    resultados_df = pd.concat(resultados, ignore_index=True)
    modelos_df = pd.concat(modelos, ignore_index=True)
    return modelos_df, resultados_df

# Cell


def periodo_vulnerabilidad_con_dataframe(covid_municipal, inicio, fin, columna='tasa_covid_letal',
                                         min_casos=20, min_defunciones=-1, rf=True):
    """Calcula la vulnerabilidad (PLS) para todo el periodo usando como objetivo
       la columna que se le pase.

    :param df: el dataframe con los datos para ajustar el modelo. Debe traer ya las tasas municipales
    :type df: pd.DataFrame
    :param inicio: fecha inicial (Y-m-d)
    :type inicio: str
    :param fin: fecha final (Y-m-d)
    :type fin: str
    :param columna: la columna para usar como objetivo, el default es 'tasa_covid_letal'
    :type columna: str
    :param min_casos: Número mínimo de casos para considerar a un municipio
    :type min_casos: int
    :param min_defunciones: Número mínimo de defunciones para considerar a un municipio
    :type min_defunciones: int
    :param rf: True/False ajustar también un nmodelo de Random Forest a los dato
    :type rf: bool

    :returns: Un DataFrame igual que el de entrada pero cun una columna extra con el resultado
              del modelo. La columna se llama 'valor_{columna}'
    :rtype: gpd.GeoDataFrame

    """
    inicio = pd.to_datetime(inicio, yearfirst=True)
    fin = pd.to_datetime(fin, yearfirst=True)
    fin = min(covid_municipal.FECHA_INGRESO.max(), fin)
    fechas = pd.date_range(inicio, fin)
    resultados = []
    modelos = []
    f = IntProgress(min=0, max=len(fechas) - 1) # instantiate the bar
    display(f) # display the bar
    # covid_municipal = agregar_tasas_municipales(df)
    caracteristicas = caracteristicas_modelos_municipios(covid_municipal)
    for count, fecha in enumerate(fechas):
        covid_municipal_fecha = covid_municipal.query(
            f'FECHA_INGRESO == "{fecha.strftime("%Y-%m-%d")}"')
        pls = ajustar_pls_columna(covid_municipal_fecha,
                                    caracteristicas,
                                    columna=columna,
                                    min_casos=min_casos,
                                    min_defunciones=min_defunciones)
        df = calificar_municipios_letalidad_formato_largo(covid_municipal_fecha,
                                                          pls,
                                                          caracteristicas,
                                                          modelo='PLS', dia_ajuste=fecha)
        resultados.append(df)
        modelo = pd.DataFrame({'caracteristica': caracteristicas, 'coef': pls.coef_})
        modelo['dia_ajuste'] = fecha
        modelo['modelo'] = 'PLS'
        modelos.append(modelo)
        if rf:
            rf = ajustar_rf_letalidad(covid_municipal_fecha,
                                    caracteristicas,
                                    min_casos=min_casos,
                                    min_defunciones=min_defunciones)
            df = calificar_municipios_letalidad_formato_largo(covid_municipal_fecha,
                                                            rf, caracteristicas,
                                                            modelo='RF',
                                                            dia_ajuste=fecha)
            resultados.append(df)
            modelo = pd.DataFrame({'caracteristica': caracteristicas, 'coef': rf.feature_importances_})
            modelo['dia_ajuste'] = fecha
            modelo['modelo'] = 'RF'
            modelos.append(modelo)
        f.value = count

    resultados_df = pd.concat(resultados, ignore_index=True)
    modelos_df = pd.concat(modelos, ignore_index=True)
    resultados_df = gpd.GeoDataFrame(resultados_df, geometry='geometry')
    resultados_df.rename({'valor':f'valor_{columna}'}, axis=1, inplace=True)
    return modelos_df, resultados_df

def calcular_periodo_vulnerabilidad_2(fecha_archivo, inicio, fin, min_casos=10, min_defunciones=-1):
    asegura_archivos_covid_disponibles([fecha_archivo])
    covid_municipal = serie_covid_indicadores_municipales(fecha_archivo.strftime("%y%m%d"))

    modelos_df, resultados_df = periodo_vulnerabilidad_con_dataframe(covid_municipal,
                                                                     inicio,
                                                                     fin,
                                                                     min_casos=min_casos,
                                                                     min_defunciones=min_defunciones)

    return modelos_df, resultados_df


# Cell

def agregar_periodo_vulnerabilidad(modelos_df, serie_vulnerabilidad, inicio, fin):
    inicio = pd.to_datetime(inicio, yearfirst=True)
    fin = pd.to_datetime(fin, yearfirst=True)
    fechas = pd.date_range(inicio, fin)
    resultados = []
    modelos = []

    asegura_archivos_covid_disponibles(fechas)
    f = IntProgress(min=0, max=len(fechas) - 1) # instantiate the bar
    display(f) # display the bar

    for count, fecha in enumerate(fechas):
        covid_municipal = tabla_covid_indicadores_municipales(fecha.strftime("%y%m%d"))
        covid_municipal = agregar_tasas_municipales(covid_municipal)

        caracteristicas = caracteristicas_modelos_municipios(covid_municipal)
        pls = ajustar_pls_letalidad(covid_municipal, caracteristicas)
        df = calificar_municipios_letalidad_formato_largo(covid_municipal, pls, caracteristicas,
                                                    modelo='PLS', dia_ajuste=fecha)
        resultados.append(df)
        modelo = pd.DataFrame({'caracteristica': caracteristicas, 'coef': pls.coef_})
        modelo['dia_ajuste'] = fecha
        modelos.append(modelo)
        f.value = count

    resultados_df = pd.concat(resultados, ignore_index=True)
    nuevos_modelos_df = pd.concat(modelos, ignore_index=True)

    serie_vulnerabilidad = serie_vulnerabilidad.append(resultados_df)
    modelos_df = modelos_df.append(nuevos_modelos_df)

    return modelos_df, serie_vulnerabilidad


# Cell
def calcular_vulnerabilidad_urbana(inicio, fin, min_casos=20, min_defunciones=0):
    ''' Ajustar un modelo para los municipios urbanos y otro para los no urbanos '''
    inicio = pd.to_datetime(inicio, yearfirst=True)
    fin = pd.to_datetime(fin, yearfirst=True)
    fechas = pd.date_range(inicio, fin)
    resultados = []
    modelos = []

    asegura_archivos_covid_disponibles(fechas)
    f = IntProgress(min=0, max=len(fechas) - 1) # instantiate the bar
    display(f) # display the bar

    urbanos = municipios_urbanos()

    for count, fecha in enumerate(fechas):
        covid_municipal = tabla_covid_indicadores_municipales(fecha.strftime("%y%m%d"))
        covid_municipal = agregar_tasas_municipales(covid_municipal)

        caracteristicas = caracteristicas_modelos_municipios(covid_municipal)

        serie_urbanos = covid_municipal[covid_municipal.CLAVE_MUNICIPIO_RES.isin(urbanos.CLAVE_MUNICIPIO)]

        pls_urbanos = ajustar_pls_letalidad(serie_urbanos, caracteristicas,
                                            min_defunciones=min_defunciones, min_casos=min_casos)
        df_urbanos = calificar_municipios_letalidad_formato_largo(serie_urbanos,
                                                                  pls_urbanos,
                                                                  caracteristicas,
                                                                  modelo='PLS_URBANO',
                                                                  dia_ajuste=fecha)
        resultados.append(df_urbanos)
        modelo = pd.DataFrame({'caracteristica': caracteristicas,
                               'coef': pls_urbanos.coef_})
        modelo['dia_ajuste'] = fecha
        modelo['modelo'] = 'PLS_URBANO'

        modelos.append(modelo)

        serie_no_urbanos = covid_municipal[~covid_municipal.CLAVE_MUNICIPIO_RES.isin(urbanos.CLAVE_MUNICIPIO)]
        pls_no_urbanos = ajustar_pls_letalidad(serie_no_urbanos, caracteristicas,
                                               min_defunciones=min_defunciones, min_casos=min_casos)
        df_no_urbanos = calificar_municipios_letalidad_formato_largo(serie_no_urbanos,
                                                                  pls_no_urbanos,
                                                                  caracteristicas,
                                                                  modelo='PLS_NO_URBANO',
                                                                  dia_ajuste=fecha)
        resultados.append(df_no_urbanos)
        modelo = pd.DataFrame({'caracteristica': caracteristicas, 'coef': pls_no_urbanos.coef_})
        modelo['dia_ajuste'] = fecha
        modelo['modelo'] = 'PLS_NO_URBANO'

        modelos.append(modelo)

        f.value = count

    resultados_df = pd.concat(resultados, ignore_index=True)
    modelos_df = pd.concat(modelos, ignore_index=True)
    return modelos_df, resultados_df

# Cell

def agregar_vulnerabilidad_entidades(serie_vulnerabilidad, campo_indice='valor'):
    df = serie_vulnerabilidad.copy()

    df['poblacion_entidad'] = df.groupby(['CLAVE_ENTIDAD_RES', 'dia_ajuste', 'modelo'])['pob2020'].transform('sum')
    df['proporcion_pob_entidad'] = df['pob2020'] / df['poblacion_entidad']

    df[f'{campo_indice}_entidad'] = df[campo_indice] * df['proporcion_pob_entidad']
    df[f'{campo_indice}_entidad'] = df.groupby(['CLAVE_ENTIDAD_RES', 'dia_ajuste', 'modelo'])[f'{campo_indice}_entidad'].transform('sum')
    df['letalidad_entidad'] = df['tasa_covid_letal'] * df['proporcion_pob_entidad']
    df['letalidad_entidad'] = df.groupby(['CLAVE_ENTIDAD_RES', 'dia_ajuste', 'modelo'])[f'letalidad_entidad'].transform('sum')


    return df

# Cell

def guardar_resultados_csv(resultados):
    reporte_csv = resultados.rename(columns={'tasa_covid_conocidos': 'casos_100k',
                             'tasa_covid_defun': 'defunciones_100k',
                             'tasa_covid_letal': 'letalidad',
                             'i_PLS_0': 'i_vulnerabilidad'})

    resultados_dir = 'resultados'
    sem_ent_cols = ['CLAVE_ENTIDAD_RES', 'nom_ent'] + ['i_vuln_estatal', 'semaforo_estado', 'semaforo_tend_estado']
    sem_mun_cols = ['CLAVE_MUNICIPIO_RES', 'nom_ent', 'nom_mun'] + ['i_vulnerabilidad', 'semaforo_municipio']

    nombre_csv_mun = os.path.join(resultados_dir, 'i_semaforo_municipios.csv')
    nombre_csv_ent = os.path.join(resultados_dir, 'i_semaforo_entidades.csv')
    nombre_csv_full = os.path.join(resultados_dir, 'semaforo_municipios.csv')

    reporte[sem_mun_cols].to_csv(nombre_csv_mun, index=False)
    reporte[sem_ent_cols].drop_duplicates().to_csv(nombre_csv_ent, index=False)
    reporte.drop(columns='geometry').to_csv(nombre_csv_full, index=False)


# Cell

def mapas_serie_letalidad(serie_vulnerabilidad, gif=True, lapso_dias=1):
    '''
        Recibe el DataFrame con la vulnerabilidad calculada para un rango de fechas y genera los mapas de cada fecha.
    '''
    dir_salida = 'resultados/mapas/letalidad'
    if not os.path.exists(dir_salida):
        os.mkdir(dir_salida)

    serie_vulnerabilidad = serie_vulnerabilidad[['dia_ajuste', 'geometry', 'tasa_covid_letal', 'CLAVE_MUNICIPIO_RES']].copy()
    serie_vulnerabilidad = serie_vulnerabilidad.drop_duplicates()

    inicio = serie_vulnerabilidad.dia_ajuste.min()
    fin = serie_vulnerabilidad.dia_ajuste.max()

    imagenes_mapas = []
    lista_mapas = []
    for fecha in pd.date_range(inicio, fin, freq=f'{lapso_dias}D'):
        datos_mapa = serie_vulnerabilidad[(serie_vulnerabilidad.dia_ajuste == fecha)]
        datos_mapa = datos_mapa[datos_mapa.tasa_covid_letal.notna()]
        fig = datos_mapa.plot(column='tasa_covid_letal', cmap='YlOrRd', scheme="equalinterval",
                        figsize=(10, 10), legend=True, k=10, edgecolor="face")
        # remove axis of chart
        fig.axis('off')

        # add a title
        fig.set_title('Letalidad COVID-19 Municipios de México', \
                  fontdict={'fontsize': '25', 'fontweight' : '3'})
        nombre_mapa = fecha.strftime("%y%m%d") + '_letalidad.jpg'
        ruta_archivo = os.path.join(dir_salida, nombre_mapa)
        lista_mapas.append(ruta_archivo)
        # position the annotation to the bottom left
        fig.annotate(fecha.strftime("%d/%m"),
                xy=(0.1, .225), xycoords='figure fraction',
                horizontalalignment='left', verticalalignment='top',
                fontsize=35)

        chart = fig.get_figure()
        chart.savefig(ruta_archivo, dpi=300)
        plt.close(chart)
        imagenes_mapas.append(imageio.imread(ruta_archivo))

    inicio_str = inicio.strftime('%m%d')
    fin_str = fin.strftime('%m%d')
    imageio.mimsave(f'letalidad_covid_{inicio_str}_{fin_str}_{f"{lapso_dias}D"}.gif', imagenes_mapas, duration=1)
    optimize(f'letalidad_covid_{inicio_str}_{fin_str}_{f"{lapso_dias}D"}.gif')
    return lista_mapas



# Cell

def mapas_serie_vulnerabilidad(serie_vulnerabilidad, modelo='PLS', gif=True, lapso_dias=1):
    '''
        Recibe el DataFrame con la vulnerabilidad calculada para un rango de fechas y genera los mapas de cada fecha.
    '''
    dir_salida = 'resultados/mapas/vulnerabilidad'
    if not os.path.exists(dir_salida):
        os.mkdir(dir_salida)

    serie_vulnerabilidad = serie_vulnerabilidad.copy()
    inicio = serie_vulnerabilidad.dia_ajuste.min()
    fin = serie_vulnerabilidad.dia_ajuste.max()
    lista_mapas = []
    imagenes_mapas = []

    for fecha in pd.date_range(inicio, fin, freq=f'{lapso_dias}D'):
        datos_mapa = serie_vulnerabilidad[(serie_vulnerabilidad.dia_ajuste == fecha) &
                                         (serie_vulnerabilidad.modelo == modelo)]
        datos_mapa = datos_mapa[datos_mapa.valor.notna()]
        fig = datos_mapa.plot(column='valor', cmap='RdYlBu_r', scheme="Quantiles",
                        figsize=(10, 10), legend=True, k=5, edgecolor="face")
        # remove axis of chart
        fig.axis('off')

        # add a title
        fig.set_title('Vulnerabilidad COVID-19 Municipios de México', \
                  fontdict={'fontsize': '25', 'fontweight' : '3'})
        nombre_mapa = fecha.strftime("%y_%m_%d") + '_vulnerabilidad.jpg'
        ruta_archivo = os.path.join(dir_salida, nombre_mapa)
        lista_mapas.append(ruta_archivo)

        # position the annotation to the bottom left
        fig.annotate(fecha.strftime("%d/%m"),
                xy=(0.1, .225), xycoords='figure fraction',
                horizontalalignment='left', verticalalignment='top',
                fontsize=35)

        chart = fig.get_figure()
        chart.savefig(ruta_archivo, dpi=300)
        imagenes_mapas.append(imageio.imread(ruta_archivo))

        plt.close(chart)
    inicio_str = inicio.strftime('%m%d')
    fin_str = fin.strftime('%m%d')
    imageio.mimsave(f'vulnerabilidad_covid_{inicio_str}_{fin_str}_{f"{lapso_dias}D"}.gif', imagenes_mapas, duration=1)
    optimize(f'vulnerabilidad_covid_{inicio_str}_{fin_str}_{f"{lapso_dias}D"}.gif')

    return lista_mapas




# Cell
def guardar_shape_vulnerabilidad(vulnerabilidad_gdf, nombre_archivo, fecha_ajuste):
    fecha_ajuste = pd.to_datetime(fecha_ajuste, yearfirst=True)

    vulnerabilidad_dia = vulnerabilidad_gdf[vulnerabilidad_gdf.dia_ajuste == fecha_ajuste].copy()
    vulnerabilidad_dia.dia_ajuste = vulnerabilidad_dia.dia_ajuste.dt.strftime("%d-%m-%y")
    vulnerabilidad_dia.to_file(nombre_archivo)

# Cell

def checkpoint_vulnerabilidad(serie_vulnerabilidad, nombre_archivo):
    serie_vulnerabilidad.drop(columns='geometry').to_feather(nombre_archivo)

# Cell

def cargar_checkpoint_vulnerabilidad(nombre_archivo):
    vulnerabilidad_df = pd.read_feather(nombre_archivo)
    geometria_municipal = gpd.read_file('datos/municipios/marco_2019.json')

    vulnerabilidad_df= vulnerabilidad_df.merge(geometria_municipal[['geometry', 'municipio_cvegeo']],
                            left_on=['CLAVE_MUNICIPIO_RES'], right_on='municipio_cvegeo')

    return vulnerabilidad_df

# AUTOGENERATED! DO NOT EDIT! File to edit: 01_vulnerabilidad.ipynb (unless otherwise specified).

__all__ = ['show_feature_importances', 'mostrar_coeficientes_PLS', 'agregar_conteo_pruebas',
           'agregar_tasas_municipales', 'caracteristicas_modelos_municipios', 'ajustar_pls_letalidad',
           'ajustar_pls_columna', 'ajustar_rf_letalidad', 'calificar_municipios_letalidad',
           'calificar_municipios_letalidad_formato_largo', 'seleccionar_caracteristicas',
           'calcular_indices_vulnerabilidad', 'calcular_indices_vulnerabilidad_formato_largo',
           'calcular_periodo_vulnerabilidad', 'periodo_vulnerabilidad_con_dataframe',
           'calcular_periodo_vulnerabilidad_2', 'agregar_periodo_vulnerabilidad', 'calcular_vulnerabilidad_urbana',
           'agregar_vulnerabilidad_entidades', 'guardar_resultados_csv', 'mapas_serie_letalidad',
           'mapas_serie_vulnerabilidad', 'guardar_shape_vulnerabilidad', 'checkpoint_vulnerabilidad',
           'cargar_checkpoint_vulnerabilidad']

# Cell
from ipywidgets import IntProgress
from IPython.display import display
import time

from .datos import *

import pandas as pd
import geopandas as gpd

import glob
import os
import matplotlib.pyplot as plt
import numpy as np

from datetime import timedelta
import datetime

from sklearn.feature_selection import SelectKBest, chi2
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.feature_selection import SelectFromModel
from sklearn.cross_decomposition import PLSRegression

import imageio
from pygifsicle import optimize

# Cell

def show_feature_importances(rf):
    importances = rf.feature_importances_

    std = np.std([tree.feature_importances_ for tree in rf.estimators_],
                 axis=0)
    indices = np.argsort(importances)[::-1]

    # Print the feature ranking
    print("Feature ranking:")

    for f in range(X.shape[1]):
        print("%d. feature %s (%f)" % (f + 1, X.columns[indices[f]], importances[indices[f]]))

    # Plot the feature importances of the forest
    f, ax = plt.subplots(figsize=(15, 10))
    plt.title("Feature importances")
    plt.bar(range(X.shape[1]), importances[indices],
           color="r", yerr=std[indices], align="center")
    plt.xticks(range(X.shape[1]), X.columns[indices])
    plt.xlim([-1, X.shape[1]])
    plt.show()

# Cell

def mostrar_coeficientes_PLS(pls):
    f, ax = plt.subplots(figsize=(15, 10))
    sns.barplot(y='nombre', x='coef', data=feats_df, color='c')

# Cell

def agregar_conteo_pruebas(covid_municipal, solo_covid=True):
    cols_localidad = ['ENTIDAD_RES', 'MUNICIPIO_RES', 'CLAVE_MUNICIPIO_RES']
    count_tested = covid_municipal[cols_localidad + ['conteo']].groupby(cols_localidad).sum()
    count_tested.reset_index(inplace=True)
    count_tested.rename(columns={'conteo': 'total_pruebas'}, inplace=True)

    covid_municipal = covid_municipal.merge(count_tested, on=cols_localidad, how='left')


    covid_municipal['casos_frac'] = 100 * covid_municipal['conteo'] / covid_municipal['total_pruebas']
    covid_municipal['tasa_covid_letal'] = 100 * covid_municipal['defunciones'] / covid_municipal['conteo']

    if solo_covid:
        covid_municipal = covid_municipal.query('RESULTADO == "Positivo SARS-CoV-2"')
    return covid_municipal

# Cell

def agregar_tasas_municipales(casos_mun_df):
    casos_mun_df['covid_confirmados_100k'] = 100000 * casos_mun_df['conteo'] / casos_mun_df['pob2020']
    casos_mun_df['covid_defun_100k'] = 100000 * casos_mun_df['defunciones'] / casos_mun_df['pob2020']
    casos_mun_df['tasa_covid_letal'] = 100 * casos_mun_df['defunciones'] / casos_mun_df['conteo']

    # covid_municipal = covid_municipal.query('RESULTADO == "Positivo SARS-CoV-2"').copy()

    # casos_mun_df = gpd.GeoDataFrame(casos_mun_df, geometry='geometry')

    return casos_mun_df

# Cell

def caracteristicas_modelos_municipios(mun_df, poblaciones=False, i_vuln=False):
    pob_vars = []
    if not poblaciones:
        pob_vars = list(mun_df.columns[mun_df.columns.str.contains('_pob')])
        pob_vars = pob_vars + ['mayores_65', 'pob2020', 'pt_2015', 'pob_menore', 'poblacion', 'pob_total']
        pob_vars = pob_vars + ['sin_dere_1', 'sin_derech', 'carencias_']
    i_vuln_vars = []
    if not i_vuln:
        i_vuln_vars = ['i_vuln_salud', 'i_vuln_cobertura', 'i_vuln_econo', 'i_vuln_social',
                       'i_vuln_gen', 'i_vuln_infraestructura']

    columnas_numericas = mun_df.select_dtypes(include=np.number).columns

    otras_vars = ['covid_defun_100k', 'tasa_covid_letal', 'defunciones',
                  'INTUBADO_BIN', 'UCI_BIN', 'UCIs', 'tasa_uci', 'total_pruebas',
                  'area_cart', 'area', 'densi', 'casos_frac', 'area_km2', 'conteo',
                  'oid', 'covid_confrimados_100k', 'id', 'vulnerabilidad_ambiental_num',
                  'covid_confirmados_100k', 'index']

    caracteristicas = list(set(columnas_numericas).difference(pob_vars + i_vuln_vars + otras_vars))

    return caracteristicas


# Cell

def ajustar_pls_letalidad(municipios_df, caracteristicas, min_casos=20, min_defunciones=0):
    data_train = municipios_df.loc[municipios_df[caracteristicas].notna().all(axis=1)]

    X = data_train.query(f'(conteo > {min_casos}) & (defunciones > {min_defunciones})')[caracteristicas]
    Y = data_train.query(f'(conteo > {min_casos}) & (defunciones > {min_defunciones})')['tasa_covid_letal']

    # X['i_vuln_econo'] = -X['i_vuln_econo']

    pls2 = PLSRegression(n_components=1)
    pls2.fit(X, Y)
    pls2.coef_ = pls2.coef_.flatten()

    return pls2

# Cell

def ajustar_pls_columna(municipios_df, caracteristicas, columna, min_casos=20, min_defunciones=0):
    data_train = municipios_df.loc[municipios_df[caracteristicas].notna().all(axis=1)]

    X = data_train.query(f'(conteo > {min_casos}) & (defunciones > {min_defunciones})')[caracteristicas]
    try:
        Y = data_train.query(f'(conteo > {min_casos}) & (defunciones > {min_defunciones})')[columna]
    except KeyError:
        print(f"No existe la columna {columna}")
    # X['i_vuln_econo'] = -X['i_vuln_econo']

    pls2 = PLSRegression(n_components=1)
    pls2.fit(X, Y)
    pls2.coef_ = pls2.coef_.flatten()

    return pls2

# Cell

def ajustar_rf_letalidad(municipios_df, caracteristicas, min_casos=20, min_defunciones=0, random_seed=0):
    data_train = municipios_df.loc[municipios_df[caracteristicas].notna().all(axis=1)]

    X = data_train.query(f'(conteo > {min_casos}) & (defunciones > {min_defunciones})')[caracteristicas]
    Y = data_train.query(f'(conteo > {min_casos}) & (defunciones > {min_defunciones})')['tasa_covid_letal']

    # X['i_vuln_econo'] = -X['i_vuln_econo']
    rf = RandomForestRegressor(random_state=random_seed)
    rf.fit(X, Y)

    return rf

# Cell

def calificar_municipios_letalidad(mun_df, regr, caracteristicas, etiqueta):
    not_na_row = mun_df[caracteristicas].notnull().all(axis=1)
    X = mun_df.loc[not_na_row, caracteristicas]
    Y_pred = regr.predict(X)
    mun_df.loc[not_na_row, etiqueta] = Y_pred

    return mun_df


# Cell

def calificar_municipios_letalidad_formato_largo(mun_df, regr, caracteristicas, modelo, dia_ajuste):
    mun_df = mun_df.copy()
    not_na_row = mun_df[caracteristicas].notnull().all(axis=1)
    X = mun_df.loc[not_na_row, caracteristicas]
    Y_pred = regr.predict(X)
    mun_df['modelo'] = modelo
    mun_df['dia_ajuste'] = dia_ajuste
    mun_df.loc[not_na_row, 'valor'] = Y_pred

    return mun_df


# Cell

def seleccionar_caracteristicas(regr, X, caracteristicas):
    sel = SelectFromModel(regr, prefit=True)

    caracteristicas_selec = list(X[caracteristicas].columns[sel.get_support()])
    return caracteristicas_selec

# Cell

def calcular_indices_vulnerabilidad(fecha_0, fecha_1, vulnerabilidad='fecha_0'):
    covid_municipal_0 = tabla_covid_indicadores_municipales(fecha_0)
    covid_municipal_0 = agregar_tasas_municipales(covid_municipal_0)

    covid_municipal_1 = tabla_covid_indicadores_municipales(fecha_1)
    covid_municipal_1 = agregar_tasas_municipales(covid_municipal_1)

    caracteristicas = caracteristicas_modelos_municipios(covid_municipal_0)

    resultados = covid_municipal_1.copy()

    # rf = ajustar_rf_municipios(covid_municipal_0, fecha_0, caracteristicas)
    # caracteristicas_rf = seleccionar_caracteristicas(rf, covid_municipal_0, caracteristicas)
    # resultados = calificar_municipios_letalidad(resultados, rf, caracteristicas, etiqueta='i_RF_0')

    pls = ajustar_pls_letalidad(covid_municipal_0, caracteristicas)
    resultados = calificar_municipios_letalidad(resultados, pls, caracteristicas, etiqueta='i_PLS_0')

    pls = ajustar_pls_letalidad(covid_municipal_1, caracteristicas)
    resultados = calificar_municipios_letalidad(resultados, pls, caracteristicas, etiqueta='i_PLS_1')

    # pls = ajustar_pls_letalidad(covid_municipal_0, caracteristicas_rf)
    # resultados = calificar_municipios_letalidad(resultados, pls, caracteristicas, etiqueta='i_PLS_RF_1')

    return resultados


# Cell

def calcular_indices_vulnerabilidad_formato_largo(fecha_0, fecha_1):
    covid_municipal_0 = tabla_covid_indicadores_municipales(fecha_0)
    covid_municipal_0 = agregar_tasas_municipales(covid_municipal_0)

    covid_municipal_1 = tabla_covid_indicadores_municipales(fecha_1)
    covid_municipal_1 = agregar_tasas_municipales(covid_municipal_1)

    caracteristicas = caracteristicas_modelos_municipios(covid_municipal_0)

    # rf = ajustar_rf_municipios(covid_municipal_0, fecha_0, caracteristicas)
    # caracteristicas_rf = seleccionar_caracteristicas(rf, covid_municipal_0, caracteristicas)
    # resultados = calificar_municipios_letalidad(resultados, rf, caracteristicas, etiqueta='i_RF_0')

    pls = ajustar_pls_letalidad(covid_municipal_0, caracteristicas)
    resultados = calificar_municipios_letalidad_formato_largo(covid_municipal_0, pls, caracteristicas,
                                                modelo='PLS', dia_ajuste=fecha_0)

    pls = ajustar_pls_letalidad(covid_municipal_1, caracteristicas)
    resultados_temp = calificar_municipios_letalidad_formato_largo(covid_municipal_1, pls, caracteristicas,
                                                     modelo='PLS', dia_ajuste=fecha_1)

    resultados = pd.concat([resultados, resultados_temp], ignore_index=True)
    # pls = ajustar_pls_letalidad(covid_municipal_0, caracteristicas_rf)
    # resultados = calificar_municipios_letalidad(resultados, pls, caracteristicas, etiqueta='i_PLS_RF_1')

    return resultados


# Cell

def calcular_periodo_vulnerabilidad(inicio, fin, min_defunciones=-1):
    inicio = pd.to_datetime(inicio, yearfirst=True)
    fin = pd.to_datetime(fin, yearfirst=True)

    fechas = pd.date_range(inicio, fin)
    resultados = []
    modelos = []

    asegura_archivos_covid_disponibles(fechas)

    f = IntProgress(min=0, max=len(fechas) - 1) # instantiate the bar
    display(f) # display the bar

    for count, fecha in enumerate(fechas):
        covid_municipal = tabla_covid_indicadores_municipales(fecha.strftime("%y%m%d"))
        covid_municipal = agregar_tasas_municipales(covid_municipal)

        caracteristicas = caracteristicas_modelos_municipios(covid_municipal)
        pls = ajustar_pls_letalidad(covid_municipal, caracteristicas, min_defunciones=min_defunciones)
        df = calificar_municipios_letalidad_formato_largo(covid_municipal, pls, caracteristicas,
                                                    modelo='PLS', dia_ajuste=fecha)
        resultados.append(df)
        modelo = pd.DataFrame({'caracteristica': caracteristicas, 'coef': pls.coef_})
        modelo['dia_ajuste'] = fecha
        modelo['modelo'] = 'PLS'
        modelos.append(modelo)

        rf = ajustar_rf_letalidad(covid_municipal, caracteristicas, min_defunciones=min_defunciones)
        df = calificar_municipios_letalidad_formato_largo(covid_municipal, rf, caracteristicas,
                                            modelo='RF', dia_ajuste=fecha)
        resultados.append(df)
        modelo = pd.DataFrame({'caracteristica': caracteristicas, 'coef': rf.feature_importances_})
        modelo['dia_ajuste'] = fecha
        modelo['modelo'] = 'RF'
        modelos.append(modelo)

        f.value = count

    resultados_df = pd.concat(resultados, ignore_index=True)
    modelos_df = pd.concat(modelos, ignore_index=True)
    return modelos_df, resultados_df

# Cell


def periodo_vulnerabilidad_con_dataframe(covid_municipal, inicio, fin, columna='tasa_covid_letal',
                                         min_casos=20, min_defunciones=-1, rf=True):
    """Calcula la vulnerabilidad (PLS) para todo el periodo usando como objetivo
       la columna que se le pase.

    :param df: el dataframe con los datos para ajustar el modelo. Debe traer ya las tasas municipales
    :type df: pd.DataFrame
    :param inicio: fecha inicial (Y-m-d)
    :type inicio: str
    :param fin: fecha final (Y-m-d)
    :type fin: str
    :param columna: la columna para usar como objetivo, el default es 'tasa_covid_letal'
    :type columna: str
    :param min_casos: Número mínimo de casos para considerar a un municipio
    :type min_casos: int
    :param min_defunciones: Número mínimo de defunciones para considerar a un municipio
    :type min_defunciones: int
    :param rf: True/False ajustar también un nmodelo de Random Forest a los dato
    :type rf: bool

    :returns: Un DataFrame igual que el de entrada pero cun una columna extra con el resultado
              del modelo. La columna se llama 'valor_{columna}'
    :rtype: gpd.GeoDataFrame

    """
    inicio = pd.to_datetime(inicio, yearfirst=True)
    fin = pd.to_datetime(fin, yearfirst=True)
    fin = min(covid_municipal.FECHA_INGRESO.max(), fin)
    fechas = pd.date_range(inicio, fin)
    resultados = []
    modelos = []
    f = IntProgress(min=0, max=len(fechas) - 1) # instantiate the bar
    display(f) # display the bar
    # covid_municipal = agregar_tasas_municipales(df)
    caracteristicas = caracteristicas_modelos_municipios(covid_municipal)
    for count, fecha in enumerate(fechas):
        covid_municipal_fecha = covid_municipal.query(
            f'FECHA_INGRESO == "{fecha.strftime("%Y-%m-%d")}"')
        pls = ajustar_pls_columna(covid_municipal_fecha,
                                    caracteristicas,
                                    columna=columna,
                                    min_casos=min_casos,
                                    min_defunciones=min_defunciones)
        df = calificar_municipios_letalidad_formato_largo(covid_municipal_fecha,
                                                          pls,
                                                          caracteristicas,
                                                          modelo='PLS', dia_ajuste=fecha)
        resultados.append(df)
        modelo = pd.DataFrame({'caracteristica': caracteristicas, 'coef': pls.coef_})
        modelo['dia_ajuste'] = fecha
        modelo['modelo'] = 'PLS'
        modelos.append(modelo)
        if rf:
            rf = ajustar_rf_letalidad(covid_municipal_fecha,
                                    caracteristicas,
                                    min_casos=min_casos,
                                    min_defunciones=min_defunciones)
            df = calificar_municipios_letalidad_formato_largo(covid_municipal_fecha,
                                                            rf, caracteristicas,
                                                            modelo='RF',
                                                            dia_ajuste=fecha)
            resultados.append(df)
            modelo = pd.DataFrame({'caracteristica': caracteristicas, 'coef': rf.feature_importances_})
            modelo['dia_ajuste'] = fecha
            modelo['modelo'] = 'RF'
            modelos.append(modelo)
        f.value = count

    resultados_df = pd.concat(resultados, ignore_index=True)
    modelos_df = pd.concat(modelos, ignore_index=True)
    resultados_df = gpd.GeoDataFrame(resultados_df, geometry='geometry')
    resultados_df.rename({'valor':f'valor_{columna}'}, axis=1, inplace=True)
    return modelos_df, resultados_df

def calcular_periodo_vulnerabilidad_2(fecha_archivo, inicio, fin, min_casos=10, min_defunciones=-1):
    asegura_archivos_covid_disponibles([fecha_archivo])
    covid_municipal = serie_covid_indicadores_municipales(fecha_archivo.strftime("%y%m%d"))

    modelos_df, resultados_df = periodo_vulnerabilidad_con_dataframe(covid_municipal,
                                                                     inicio,
                                                                     fin,
                                                                     min_casos=min_casos,
                                                                     min_defunciones=min_defunciones)

    return modelos_df, resultados_df


# Cell

def agregar_periodo_vulnerabilidad(modelos_df, serie_vulnerabilidad, inicio, fin):
    inicio = pd.to_datetime(inicio, yearfirst=True)
    fin = pd.to_datetime(fin, yearfirst=True)
    fechas = pd.date_range(inicio, fin)
    resultados = []
    modelos = []

    asegura_archivos_covid_disponibles(fechas)
    f = IntProgress(min=0, max=len(fechas) - 1) # instantiate the bar
    display(f) # display the bar

    for count, fecha in enumerate(fechas):
        covid_municipal = tabla_covid_indicadores_municipales(fecha.strftime("%y%m%d"))
        covid_municipal = agregar_tasas_municipales(covid_municipal)

        caracteristicas = caracteristicas_modelos_municipios(covid_municipal)
        pls = ajustar_pls_letalidad(covid_municipal, caracteristicas)
        df = calificar_municipios_letalidad_formato_largo(covid_municipal, pls, caracteristicas,
                                                    modelo='PLS', dia_ajuste=fecha)
        resultados.append(df)
        modelo = pd.DataFrame({'caracteristica': caracteristicas, 'coef': pls.coef_})
        modelo['dia_ajuste'] = fecha
        modelos.append(modelo)
        f.value = count

    resultados_df = pd.concat(resultados, ignore_index=True)
    nuevos_modelos_df = pd.concat(modelos, ignore_index=True)

    serie_vulnerabilidad = serie_vulnerabilidad.append(resultados_df)
    modelos_df = modelos_df.append(nuevos_modelos_df)

    return modelos_df, serie_vulnerabilidad


# Cell
def calcular_vulnerabilidad_urbana(inicio, fin, min_casos=20, min_defunciones=0):
    ''' Ajustar un modelo para los municipios urbanos y otro para los no urbanos '''
    inicio = pd.to_datetime(inicio, yearfirst=True)
    fin = pd.to_datetime(fin, yearfirst=True)
    fechas = pd.date_range(inicio, fin)
    resultados = []
    modelos = []

    asegura_archivos_covid_disponibles(fechas)
    f = IntProgress(min=0, max=len(fechas) - 1) # instantiate the bar
    display(f) # display the bar

    urbanos = municipios_urbanos()

    for count, fecha in enumerate(fechas):
        covid_municipal = tabla_covid_indicadores_municipales(fecha.strftime("%y%m%d"))
        covid_municipal = agregar_tasas_municipales(covid_municipal)

        caracteristicas = caracteristicas_modelos_municipios(covid_municipal)

        serie_urbanos = covid_municipal[covid_municipal.CLAVE_MUNICIPIO_RES.isin(urbanos.CLAVE_MUNICIPIO)]

        pls_urbanos = ajustar_pls_letalidad(serie_urbanos, caracteristicas,
                                            min_defunciones=min_defunciones, min_casos=min_casos)
        df_urbanos = calificar_municipios_letalidad_formato_largo(serie_urbanos,
                                                                  pls_urbanos,
                                                                  caracteristicas,
                                                                  modelo='PLS_URBANO',
                                                                  dia_ajuste=fecha)
        resultados.append(df_urbanos)
        modelo = pd.DataFrame({'caracteristica': caracteristicas,
                               'coef': pls_urbanos.coef_})
        modelo['dia_ajuste'] = fecha
        modelo['modelo'] = 'PLS_URBANO'

        modelos.append(modelo)

        serie_no_urbanos = covid_municipal[~covid_municipal.CLAVE_MUNICIPIO_RES.isin(urbanos.CLAVE_MUNICIPIO)]
        pls_no_urbanos = ajustar_pls_letalidad(serie_no_urbanos, caracteristicas,
                                               min_defunciones=min_defunciones, min_casos=min_casos)
        df_no_urbanos = calificar_municipios_letalidad_formato_largo(serie_no_urbanos,
                                                                  pls_no_urbanos,
                                                                  caracteristicas,
                                                                  modelo='PLS_NO_URBANO',
                                                                  dia_ajuste=fecha)
        resultados.append(df_no_urbanos)
        modelo = pd.DataFrame({'caracteristica': caracteristicas, 'coef': pls_no_urbanos.coef_})
        modelo['dia_ajuste'] = fecha
        modelo['modelo'] = 'PLS_NO_URBANO'

        modelos.append(modelo)

        f.value = count

    resultados_df = pd.concat(resultados, ignore_index=True)
    modelos_df = pd.concat(modelos, ignore_index=True)
    return modelos_df, resultados_df

# Cell

def agregar_vulnerabilidad_entidades(serie_vulnerabilidad, campo_indice='valor'):
    df = serie_vulnerabilidad.copy()

    df['poblacion_entidad'] = df.groupby(['CLAVE_ENTIDAD_RES', 'dia_ajuste', 'modelo'])['pob2020'].transform('sum')
    df['proporcion_pob_entidad'] = df['pob2020'] / df['poblacion_entidad']

    df[f'{campo_indice}_entidad'] = df[campo_indice] * df['proporcion_pob_entidad']
    df[f'{campo_indice}_entidad'] = df.groupby(['CLAVE_ENTIDAD_RES', 'dia_ajuste', 'modelo'])[f'{campo_indice}_entidad'].transform('sum')
    df['letalidad_entidad'] = df['tasa_covid_letal'] * df['proporcion_pob_entidad']
    df['letalidad_entidad'] = df.groupby(['CLAVE_ENTIDAD_RES', 'dia_ajuste', 'modelo'])[f'letalidad_entidad'].transform('sum')


    return df

# Cell

def guardar_resultados_csv(resultados):
    reporte_csv = resultados.rename(columns={'tasa_covid_conocidos': 'casos_100k',
                             'tasa_covid_defun': 'defunciones_100k',
                             'tasa_covid_letal': 'letalidad',
                             'i_PLS_0': 'i_vulnerabilidad'})

    resultados_dir = 'resultados'
    sem_ent_cols = ['CLAVE_ENTIDAD_RES', 'nom_ent'] + ['i_vuln_estatal', 'semaforo_estado', 'semaforo_tend_estado']
    sem_mun_cols = ['CLAVE_MUNICIPIO_RES', 'nom_ent', 'nom_mun'] + ['i_vulnerabilidad', 'semaforo_municipio']

    nombre_csv_mun = os.path.join(resultados_dir, 'i_semaforo_municipios.csv')
    nombre_csv_ent = os.path.join(resultados_dir, 'i_semaforo_entidades.csv')
    nombre_csv_full = os.path.join(resultados_dir, 'semaforo_municipios.csv')

    reporte[sem_mun_cols].to_csv(nombre_csv_mun, index=False)
    reporte[sem_ent_cols].drop_duplicates().to_csv(nombre_csv_ent, index=False)
    reporte.drop(columns='geometry').to_csv(nombre_csv_full, index=False)


# Cell

def mapas_serie_letalidad(serie_vulnerabilidad, gif=True, lapso_dias=1):
    '''
        Recibe el DataFrame con la vulnerabilidad calculada para un rango de fechas y genera los mapas de cada fecha.
    '''
    dir_salida = 'resultados/mapas/letalidad'
    if not os.path.exists(dir_salida):
        os.mkdir(dir_salida)

    serie_vulnerabilidad = serie_vulnerabilidad[['dia_ajuste', 'geometry', 'tasa_covid_letal', 'CLAVE_MUNICIPIO_RES']].copy()
    serie_vulnerabilidad = serie_vulnerabilidad.drop_duplicates()

    inicio = serie_vulnerabilidad.dia_ajuste.min()
    fin = serie_vulnerabilidad.dia_ajuste.max()

    imagenes_mapas = []
    lista_mapas = []
    for fecha in pd.date_range(inicio, fin, freq=f'{lapso_dias}D'):
        datos_mapa = serie_vulnerabilidad[(serie_vulnerabilidad.dia_ajuste == fecha)]
        datos_mapa = datos_mapa[datos_mapa.tasa_covid_letal.notna()]
        fig = datos_mapa.plot(column='tasa_covid_letal', cmap='YlOrRd', scheme="equalinterval",
                        figsize=(10, 10), legend=True, k=10, edgecolor="face")
        # remove axis of chart
        fig.axis('off')

        # add a title
        fig.set_title('Letalidad COVID-19 Municipios de México', \
                  fontdict={'fontsize': '25', 'fontweight' : '3'})
        nombre_mapa = fecha.strftime("%y%m%d") + '_letalidad.jpg'
        ruta_archivo = os.path.join(dir_salida, nombre_mapa)
        lista_mapas.append(ruta_archivo)
        # position the annotation to the bottom left
        fig.annotate(fecha.strftime("%d/%m"),
                xy=(0.1, .225), xycoords='figure fraction',
                horizontalalignment='left', verticalalignment='top',
                fontsize=35)

        chart = fig.get_figure()
        chart.savefig(ruta_archivo, dpi=300)
        plt.close(chart)
        imagenes_mapas.append(imageio.imread(ruta_archivo))

    inicio_str = inicio.strftime('%m%d')
    fin_str = fin.strftime('%m%d')
    imageio.mimsave(f'letalidad_covid_{inicio_str}_{fin_str}_{f"{lapso_dias}D"}.gif', imagenes_mapas, duration=1)
    optimize(f'letalidad_covid_{inicio_str}_{fin_str}_{f"{lapso_dias}D"}.gif')
    return lista_mapas



# Cell

def mapas_serie_vulnerabilidad(serie_vulnerabilidad, modelo='PLS', gif=True, lapso_dias=1):
    '''
        Recibe el DataFrame con la vulnerabilidad calculada para un rango de fechas y genera los mapas de cada fecha.
    '''
    dir_salida = 'resultados/mapas/vulnerabilidad'
    if not os.path.exists(dir_salida):
        os.mkdir(dir_salida)

    serie_vulnerabilidad = serie_vulnerabilidad.copy()
    inicio = serie_vulnerabilidad.dia_ajuste.min()
    fin = serie_vulnerabilidad.dia_ajuste.max()
    lista_mapas = []
    imagenes_mapas = []

    for fecha in pd.date_range(inicio, fin, freq=f'{lapso_dias}D'):
        datos_mapa = serie_vulnerabilidad[(serie_vulnerabilidad.dia_ajuste == fecha) &
                                         (serie_vulnerabilidad.modelo == modelo)]
        datos_mapa = datos_mapa[datos_mapa.valor.notna()]
        fig = datos_mapa.plot(column='valor', cmap='RdYlBu_r', scheme="Quantiles",
                        figsize=(10, 10), legend=True, k=5, edgecolor="face")
        # remove axis of chart
        fig.axis('off')

        # add a title
        fig.set_title('Vulnerabilidad COVID-19 Municipios de México', \
                  fontdict={'fontsize': '25', 'fontweight' : '3'})
        nombre_mapa = fecha.strftime("%y_%m_%d") + '_vulnerabilidad.jpg'
        ruta_archivo = os.path.join(dir_salida, nombre_mapa)
        lista_mapas.append(ruta_archivo)

        # position the annotation to the bottom left
        fig.annotate(fecha.strftime("%d/%m"),
                xy=(0.1, .225), xycoords='figure fraction',
                horizontalalignment='left', verticalalignment='top',
                fontsize=35)

        chart = fig.get_figure()
        chart.savefig(ruta_archivo, dpi=300)
        imagenes_mapas.append(imageio.imread(ruta_archivo))

        plt.close(chart)
    inicio_str = inicio.strftime('%m%d')
    fin_str = fin.strftime('%m%d')
    imageio.mimsave(f'vulnerabilidad_covid_{inicio_str}_{fin_str}_{f"{lapso_dias}D"}.gif', imagenes_mapas, duration=1)
    optimize(f'vulnerabilidad_covid_{inicio_str}_{fin_str}_{f"{lapso_dias}D"}.gif')

    return lista_mapas




# Cell
def guardar_shape_vulnerabilidad(vulnerabilidad_gdf, nombre_archivo, fecha_ajuste):
    fecha_ajuste = pd.to_datetime(fecha_ajuste, yearfirst=True)

    vulnerabilidad_dia = vulnerabilidad_gdf[vulnerabilidad_gdf.dia_ajuste == fecha_ajuste].copy()
    vulnerabilidad_dia.dia_ajuste = vulnerabilidad_dia.dia_ajuste.dt.strftime("%d-%m-%y")
    vulnerabilidad_dia.to_file(nombre_archivo)

# Cell

def checkpoint_vulnerabilidad(serie_vulnerabilidad, nombre_archivo):
    serie_vulnerabilidad.drop(columns='geometry').to_feather(nombre_archivo)

# Cell

def cargar_checkpoint_vulnerabilidad(nombre_archivo):
    vulnerabilidad_df = pd.read_feather(nombre_archivo)
    geometria_municipal = gpd.read_file('datos/municipios/marco_2019.json')

    vulnerabilidad_df= vulnerabilidad_df.merge(geometria_municipal[['geometry', 'municipio_cvegeo']],
                            left_on=['CLAVE_MUNICIPIO_RES'], right_on='municipio_cvegeo')

    return vulnerabilidad_df

# AUTOGENERATED! DO NOT EDIT! File to edit: 01_vulnerabilidad.ipynb (unless otherwise specified). # Cell # Cell # Print the feature ranking # Plot the feature importances of the forest # Cell # Cell # Cell # covid_municipal = covid_municipal.query('RESULTADO == "Positivo SARS-CoV-2"').copy() # casos_mun_df = gpd.GeoDataFrame(casos_mun_df, geometry='geometry') # Cell # Cell # X['i_vuln_econo'] = -X['i_vuln_econo'] # Cell # X['i_vuln_econo'] = -X['i_vuln_econo'] # Cell # X['i_vuln_econo'] = -X['i_vuln_econo'] # Cell # Cell # Cell # Cell # rf = ajustar_rf_municipios(covid_municipal_0, fecha_0, caracteristicas) # caracteristicas_rf = seleccionar_caracteristicas(rf, covid_municipal_0, caracteristicas) # resultados = calificar_municipios_letalidad(resultados, rf, caracteristicas, etiqueta='i_RF_0') # pls = ajustar_pls_letalidad(covid_municipal_0, caracteristicas_rf) # resultados = calificar_municipios_letalidad(resultados, pls, caracteristicas, etiqueta='i_PLS_RF_1') # Cell # rf = ajustar_rf_municipios(covid_municipal_0, fecha_0, caracteristicas) # caracteristicas_rf = seleccionar_caracteristicas(rf, covid_municipal_0, caracteristicas) # resultados = calificar_municipios_letalidad(resultados, rf, caracteristicas, etiqueta='i_RF_0') # pls = ajustar_pls_letalidad(covid_municipal_0, caracteristicas_rf) # resultados = calificar_municipios_letalidad(resultados, pls, caracteristicas, etiqueta='i_PLS_RF_1') # Cell # instantiate the bar # display the bar # Cell Calcula la vulnerabilidad (PLS) para todo el periodo usando como objetivo        la columna que se le pase.      :param df: el dataframe con los datos para ajustar el modelo. Debe traer ya las tasas municipales     :type df: pd.DataFrame     :param inicio: fecha inicial (Y-m-d)     :type inicio: str     :param fin: fecha final (Y-m-d)     :type fin: str     :param columna: la columna para usar como objetivo, el default es 'tasa_covid_letal'     :type columna: str     :param min_casos: Número mínimo de casos para considerar a un municipio     :type min_casos: int     :param min_defunciones: Número mínimo de defunciones para considerar a un municipio     :type min_defunciones: int     :param rf: True/False ajustar también un nmodelo de Random Forest a los dato     :type rf: bool      :returns: Un DataFrame igual que el de entrada pero cun una columna extra con el resultado               del modelo. La columna se llama 'valor_{columna}'     :rtype: gpd.GeoDataFrame # instantiate the bar # display the bar # covid_municipal = agregar_tasas_municipales(df) # Cell # instantiate the bar # display the bar # Cell Ajustar un modelo para los municipios urbanos y otro para los no urbanos # instantiate the bar # display the bar # Cell # Cell # Cell Recibe el DataFrame con la vulnerabilidad calculada para un rango de fechas y genera los mapas de cada fecha. # remove axis of chart # add a title # position the annotation to the bottom left # Cell Recibe el DataFrame con la vulnerabilidad calculada para un rango de fechas y genera los mapas de cada fecha. # remove axis of chart # add a title # position the annotation to the bottom left # Cell # Cell # Cell