# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import pandas as pd
import numpy as np
from IPython.display import display, HTML
from collections import defaultdict

cities = ['bogota','boston', 'LA', 'chicago']
models = [{'core': True}, {'sd': True}, {'uf': True}, {'m': True}, {'sd': True, 'uf':True}, {'sd': True, 'm': True}, {'uf': True, 'm': True}, {'full': True}] #, 
neigh = 'neigh'


# +
def create_model_string(full=True, sd=False, uf=False, m=False, core=False):
    if sd or uf or m:
        full=False
    name = ['full']
    if core or sd or uf or m:
        name = []
        if core:
            name.append('core')
        else:
            if sd:
                name.append('sd')
            if uf:
                name.append('uf')
            if m:
                name.append('m')
    return '_'.join(name)

def load_results(cities, models, neigh, model_type, target='', od='', load_loo=False, niterations=15000):
    
    frame = pd.DataFrame()
    list_ = []
    for c in cities:
        for m in models:
            if ('m' in m or 'full' in m) and c == 'chicago':
                continue
            model_string = create_model_string(**m)
            filenameformat = '{model_type}_{city}_ego_{model_name}_{od}{niterations}_ncrimes.csv'.format(model_type=model_type, od=od, city=c, model_name=model_string, niterations=niterations)
            df_t = pd.read_csv('../data/generated_files/model_results/pystan_{}'.format(filenameformat), index_col=None, header=0)
            if load_loo:
                filenameformat = '{model_type}_CV_{city}_ego_{model_name}_{od}{niterations}_ncrimes.csv'.format(model_type=model_type, od=od, city=c, model_name=model_string, niterations=niterations)
                df_cv = pd.read_csv('../data/generated_files/model_results/pystanCV_{}'.format(filenameformat),index_col=None, header=0)
                df_cv = df_cv[df_cv.Metric == 'LOO']
                df_t = pd.concat((df_t[df_t.Metric != 'LOO'], df_cv))
            
            df_t['city'] = c
            df_t['model'] = model_string
            df_t.loc[df_t.Metric == 'LOO', 'R1'] = np.around(df_t.loc[df_t.Metric == 'LOO', 'R1'])
            list_.append(df_t)
    frame = pd.concat(list_)
    return frame


# +
model_strings = [create_model_string(**m) for m in models]

model_nice_names = {
    'core': 'Core',
    'sd': 'Social-disorganization (SD)',
    'uf': 'Built environment (BE)',
    'm': 'Mobility (M)',
    'sd_uf': 'SD+BE',
    'sd_m': 'SD+M',
    'uf_m': 'BE+M',
    'full': 'SD+BE+M (Full)'
}

model_strings

# +
model_type = 'BSF'
cities_to_compare = ['bogota', 'boston', 'LA', 'chicago']
metrics_to_show = ['R2 marginal', 'R2 conditional', 'LOO']
default_for_chicago = [-1.0, -1.0, -100000]

model_results = defaultdict(list)
for i, m in enumerate(model_strings):
    fetch_cities = cities[:]
    if ('m' in m or 'full' in m):
        fetch_cities.remove('chicago')
        
    m1_df = load_results(fetch_cities, models, neigh, model_type=model_type)
    m1_df['model'] = pd.Categorical(m1_df['model'].values, [create_model_string(**m) for m in models])
    m1_df['city'] = pd.Categorical(m1_df['city'], cities)
            
    m1_df['R1'] = m1_df['R1'].astype('float32')
    
    for c in cities_to_compare:
        if c == 'chicago' and ('m' in m or 'full' in m):
            model_results[c].append(default_for_chicago)
            continue
        
        metrics = []
        for k in metrics_to_show:
            metrics.append(m1_df[(m1_df.Metric == k) & (m1_df.city == c) & (m1_df.model == m)]['R1'].values[0])
        model_results[c].append(metrics)


results_string = model_results.copy()
for c in cities_to_compare:
    results_string[c] = np.array(results_string[c]).astype('str')
    model_results[c] = np.array(model_results[c])
    
    for i, k in enumerate(metrics_to_show):
        bestjs = np.argwhere(model_results[c][:, i] == np.amax(model_results[c][:, i])).flatten()
        
        for j in range(results_string[c].shape[0]):
            if k in {'R2 conditional', 'R2 marginal'}:
                results_string[c][j, i] = '{:0.2f}'.format(model_results[c][j, i])
            else:
                results_string[c][j, i] = '{:0.0f}'.format(model_results[c][j, i])
        if k != 'R2 conditional':
            for b in bestjs:
                results_string[c][b, i] = '$\\bst{{{}}}$'.format(results_string[c][b, i])
        

for i, m in enumerate(model_strings):
    final_string = []
    for c in cities_to_compare:
        chicago_filter = [x.replace('-100000.00', '-').replace('-1.00', '-').replace('-100000', '-') for x in results_string[c][i].tolist()]
        final_string.extend(chicago_filter)
    
    print('{} & {} ({}) & {} && {} ({}) & {} && {} ({}) & {} && {} ({}) & {} \\\\'.format(*([model_nice_names[m]] + final_string)))

# -

# ## The need for spatial models

# +
from collections import defaultdict
models_to_compare = ['BSF', 'nb']
metrics_to_show = ['LOO', 'R2 conditional', 'MC_p']

for c in cities:
    for i, m in enumerate(model_strings):
        if ('m' in m or 'full' in m) and c == 'chicago':
                continue
        
        model_results = defaultdict(list)
        for model_type in models_to_compare:
            m1_df = load_results(cities, models, neigh, model_type=model_type)
            m1_df['model'] = pd.Categorical(m1_df['model'].values, [create_model_string(**m) for m in models])
            m1_df['city'] = pd.Categorical(m1_df['city'], cities)
            
            row = m1_df[(m1_df.city == c) & (m1_df.model == m)].copy()
            row['R1'] = row['R1'].astype('float32')
            
            for k in metrics_to_show:
                model_results[k].append(row[row.Metric == k]['R1'].values[0])
        
        city_column = ''
        if i == 0:
            city_column = '\multirow{{1}}{{*}}{{{cityname}}} '.format(cityname=c.title() if c != 'LA' else c)
        
        # A string for each metric
        results_string = model_results.copy()
        for k in metrics_to_show:
            bestjs = np.argwhere(model_results[k] == np.amin(model_results[k])).flatten()
            if k == 'MC_p':
                results_string[k] = ['{:0.3f}'.format(x) for x in results_string[k]]
            elif k == 'R2 conditional':
                bestjs = np.argwhere(model_results[k] == np.amax(model_results[k])).flatten()
                results_string[k] = ['{:0.2f}'.format(x) for x in results_string[k]]
            elif k == 'LOO':
                bestjs = np.argwhere(model_results[k] == np.amax(model_results[k])).flatten()
                results_string[k] = ['{:0.0f}'.format(x) for x in results_string[k]]
            for b in bestjs:
                results_string[k][b] = '\\textbf{{{}}}'.format(results_string[k][b])
        
        final_string = [city_column, model_nice_names[m]]
        for i in range(len(models_to_compare)):
            for k in metrics_to_show:
                final_string.append(results_string[k][i])
        
        print('{}& {} & {} & {} & {} && {} & {} & {} \\\\'.format(*final_string))
    print('\midrule')
# -

# ## BSF vs MSF vs RSR

# +
from collections import defaultdict
models_to_compare = ['BSF', 'REESF-a', 'ESF']
metrics_to_show = ['LOO', 'MC_p']

for c in cities:
    for i, m in enumerate(model_strings):
        if ('m' in m or 'full' in m) and c == 'chicago':
                continue
        
        model_results = defaultdict(list)
        for model_type in models_to_compare:
            m1_df = load_results(cities, models, neigh, model_type=model_type)
            m1_df['model'] = pd.Categorical(m1_df['model'].values, [create_model_string(**m) for m in models])
            m1_df['city'] = pd.Categorical(m1_df['city'], cities)
            
            row = m1_df[(m1_df.city == c) & (m1_df.model == m)].copy()
            row['R1'] = row['R1'].astype('float32')
            row.loc[row.Metric == 'LOO', 'R1'] = row[row.Metric == 'LOO']['R1'].astype('int')
            
            for k in metrics_to_show:
                model_results[k].append(row[row.Metric == k]['R1'].values[0])
        
        city_column = ''
        if i == 0:
            city_column = '\multirow{{1}}{{*}}{{{cityname}}} '.format(cityname=c.title() if c != 'LA' else c)
        
        # A string for each metric
        results_string = model_results.copy()
        for k in metrics_to_show:
            bestjs = np.argwhere(model_results[k] == np.amax(model_results[k])).flatten()
            if k == 'MC_p':
                results_string[k] = ['{:0.3f}'.format(x) for x in results_string[k]]
            else:
                results_string[k] = ['{:0.0f}'.format(x) for x in results_string[k]]
                for b in bestjs:
                    results_string[k][b] = '\\textbf{{{}}}'.format(results_string[k][b])
        
        final_string = [city_column, model_nice_names[m]]
        for i in range(len(models_to_compare)):
            for k in metrics_to_show:
                final_string.append(results_string[k][i])
        
        print('{}& {} & {} & {} && {} & {} && {} & {}\\\\'.format(*final_string))
    print('\midrule')
# -

# ## Od-d vs Od-m vs binary

# +
from collections import defaultdict
models_to_compare = ['', 'ODd_', 'ODm_']
metrics_to_show = ['LOO', 'MC_p']

for c in cities:
    for i, m in enumerate(model_strings):
        if ('m' in m or 'full' in m) and c == 'chicago':
                continue
        
        model_results = defaultdict(list)
        for model_type in models_to_compare:
            
            # Fix for chicago
            if model_type == 'ODm_' and c == 'chicago':
                for k in metrics_to_show:
                    model_results[k].append(-1000000)
            else:
                m1_df = load_results([c], models, neigh, model_type='BSF', od=model_type, niterations=20000 if model_type else 15000)
                m1_df['model'] = pd.Categorical(m1_df['model'].values, [create_model_string(**m) for m in models])
                m1_df['city'] = pd.Categorical(m1_df['city'], cities)

                row = m1_df[(m1_df.city == c) & (m1_df.model == m)].copy()
                row['R1'] = row['R1'].astype('float32')

                for k in metrics_to_show:
                    model_results[k].append(row[row.Metric == k]['R1'].values[0])
        
        city_column = ''
        if i == 0:
            city_column = '\multirow{{1}}{{*}}{{{cityname}}} '.format(cityname=c.title() if c != 'LA' else c)
        
        # A string for each metric
        results_string = model_results.copy()
        for k in metrics_to_show:
            bestjs = np.argwhere(model_results[k] == np.amax(model_results[k])).flatten()
            if k == 'MC_p':
                results_string[k] = ['{:0.3f}'.format(x) for x in results_string[k]]
            else:
                results_string[k] = ['{:0.0f}'.format(x) for x in results_string[k]]
                for b in bestjs:
                    results_string[k][b] = '\\textbf{{{}}}'.format(results_string[k][b])
        
        final_string = [city_column, model_nice_names[m]]
        for i in range(len(models_to_compare)):
            for k in metrics_to_show:
                final_string.append(results_string[k][i])
        
        print('{}& {} & {} & {} && {} & {} && {} & {}\\\\'.format(*final_string))
    print('\midrule')
# -



StarcoderdataJupyterScriptsDedupFiltered