import os
import pygsheets
import sys
import inspect
import datetime
import requests

currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
parentdir = os.path.dirname(currentdir)
sys.path.append(parentdir)
sys.path.append(parentdir + '/modeling')
import load_data
from fit_and_predict import add_preds
from functions import merge_data
from viz import viz_interactive

import numpy as np
from os.path import join as oj
import os
import pandas as pd
from scipy.stats import percentileofscore


def apply_manual_thresholds(vals, manual_thresholds={3: 6,
                                                     2: 2,
                                                     1: 0}):
    new_col = vals * 0
    for key in sorted(manual_thresholds.keys()):
        new_col[vals >= manual_thresholds[key]] = key
    return new_col.astype(int)


def cut_with_manual_low(vals, LOW_THRESH=1):
    '''Everything below LOW_THRESH gets severity 1
    All other things are split evenly by percentile
    '''
    new_col = vals * 0
    new_col[vals < LOW_THRESH] = 1
    new_col[vals >= LOW_THRESH] = pd.qcut(vals[vals >= LOW_THRESH], 2, labels=False) + 2
    return new_col.astype(int)


def percentiles(vals):
    '''Map each value to its percentile
    '''
    new_col = vals * 0
    new_col = [percentileofscore(vals, v) for v in vals]
    return np.array(new_col).astype(int)


def cut_into_categories(vals, NUM_CATEGORIES=3):
    '''Evenly divide values into NUM_CATEGORIES categories
    (higher values get higher categories)
    '''
    return pd.qcut(vals, NUM_CATEGORIES, labels=False, duplicates='drop') + 1


def add_severity_county(df, NUM_DAYS_LIST):
    '''Add county-level stuff
    '''
    for num_days in NUM_DAYS_LIST:
        df[f'Predicted New Deaths {num_days}-day'] = df[f'Predicted Deaths {num_days}-day'] - df['tot_deaths']

    index_county_keys = ['Predicted New Deaths 5-day', 'tot_deaths']
    for k in index_county_keys:
        df[k + ' Percentile'] = percentiles(df[k])
    df_perc = df[[k + ' Percentile' for k in index_county_keys]]
    df['Severity County 5-day'] = cut_into_categories(df_perc.mean(axis=1))
    df['Surge County 3-day'] = (2 * df['Predicted Deaths 3-day'] - df['#ICU_beds']) / df['#ICU_beds']

    return df


def add_severity_index(df, NUM_DAYS_LIST=[1, 2, 3, 4, 5, 6, ]):
    '''Add county + hospital indexes
    '''
    df = add_severity_county(df, NUM_DAYS_LIST)

    # loop over num day    
    df['Total Deaths Hospital'] = (df['tot_deaths'] * df['Frac Hospital Employees of County']).fillna(0)
    for num_days in NUM_DAYS_LIST:

        # hospital-level deaths
        df[f'Predicted Deaths Hospital {num_days}-day'] = (
                (df[f'Predicted Deaths {num_days}-day']) * df['Frac Hospital Employees of County']).fillna(0)
        df[f'Predicted New Deaths Hospital {num_days}-day'] = (
                df[f'Predicted New Deaths {num_days}-day'] * df['Frac Hospital Employees of County']).fillna(0)

        # severity
        index_hosp_keys = [f'Predicted New Deaths {num_days}-day', 'Total Deaths Hospital']
        for k in index_hosp_keys:
            df[k + ' Percentile'] = percentiles(df[k])
        df_perc = df[[k + ' Percentile' for k in index_hosp_keys]]
        df[f'Severity {num_days}-day'] = cut_into_categories(df_perc.mean(axis=1))
        # df[f'Severity {num_days}-day'] = cut_with_manual_low(df[f'Predicted Deaths Hospital {num_days}-day']) 
        df[f'Severity Emerging {num_days}-day'] = cut_with_manual_low(
            df[f'Predicted New Deaths Hospital {num_days}-day'])
        df[f'Rural Severity {num_days}-day'] = [-1] * df.shape[0]
        idxs_rural = df['Urban or Rural Designation'] == 'Rural'
        df[f'Rural Severity {num_days}-day'][idxs_rural] = cut_with_manual_low(
            df[f'Predicted Deaths Hospital {num_days}-day'][idxs_rural])

        # surge
        df[f'Surge {num_days}-day'] = (2 * df[f'Predicted Deaths Hospital {num_days}-day'] - df['ICU Beds'])

    s_hosp = f'Predicted Deaths Hospital 3-day'
    return df.sort_values(s_hosp, ascending=False).round(2)


def write_to_gsheet(d, sheet_name='COVID Severity Index', service_file='../creds.json'):
    '''Write data to google sheet
    '''
    print('writing to gsheets...')
    gc = pygsheets.authorize(service_file=service_file)
    sh = gc.open(sheet_name)  # name of the hospital
    wks = sh[0]  # select a sheet
    wks.update_value('A1', "Note: this sheet is read-only (automatically generated by the data and model)")
    wks.set_dataframe(d, (3, 1))  # update the first sheet with df, starting at cell B2.


def unpack_preds_list(r, NUM_DAYS_LIST, key="Predicted Deaths Intervals"):
    '''Takes a row r a df and a key for a prediction list col and splits into
    multiple columns with format f'{key} {i}-day'.
    '''
    preds = r[key]
    for i, num_days in enumerate(NUM_DAYS_LIST):
        r[f'{key} {num_days}-day'] = preds[i]
    return r


def prep_county_df(df, NUM_DAYS_LIST):
    '''Choose and rename appropriate keys
    '''
    df = add_severity_county(df, NUM_DAYS_LIST)
    k_severity = 'Severity County 5-day'
    latest_date_str = df.filter(regex='#Deaths_').columns[-1].replace('#Deaths_', '')
    latest_date = datetime.datetime.strptime(latest_date_str, '%m-%d-%Y').date()
    days = [
        prefix + (latest_date + datetime.timedelta(days=i)).strftime("%B %d")
        for i in NUM_DAYS_LIST for prefix in ('Predicted Deaths by ', 'Predicted Deaths Intervals by ')
    ]
    remap = {
        k: v for i in NUM_DAYS_LIST
        for k, v in zip(
            # keys
            (f'Predicted Deaths {i}-day', f'Predicted Deaths Intervals {i}-day'),
            # values
            ('Predicted Deaths by ' + \
             (latest_date + datetime.timedelta(days=i)).strftime("%B %d"),
             'Predicted Deaths Intervals by ' + \
             (latest_date + datetime.timedelta(days=i)).strftime("%B %d"))
        )
    }
    dfc = df.apply(lambda r: unpack_preds_list(r, NUM_DAYS_LIST, "Predicted Deaths Intervals"), axis=1)
    dfc = dfc.rename(columns=remap).sort_values(by=days[0], ascending=False).round(decimals=1)
    ks = ['countyFIPS', 'CountyName', 'StateName'] + days + [k_severity]
    return dfc[ks]


def write_to_api(d, api_file='../ian_key.env', csv_file='_hidden_hosp.csv'):
    '''Write to api for R4L
    '''
    print('writing to api')
    with open(api_file, 'r') as f:
        auth_token = f.read()
    hed = {'Authorization': 'Bearer ' + auth_token}
    url = 'https://api-r4l-ventilator-prediction.herokuapp.com/berkeley/severity/auto-upload'
    d.to_csv(csv_file)
    with open(csv_file, 'rb') as f:
        r = requests.post(url, files={'file': (csv_file, f, 'text/csv', {'Expires': '0'})},
                          headers=hed)
    print('api post succeeded?', r.text)


def df_to_plot(df, NUM_DAYS_LIST):
    ks = ['Total Deaths Hospital', 'Hospital Employees', 'Hospital Name', 'CountyName', 'StateName', 'ICU Beds',
          'CMS Certification Number', 'countyFIPS']
    remap = {1: 'Low', 2: 'Medium', 3: 'High'}
    for i in NUM_DAYS_LIST:
        ks.append(f'Severity {i}-day')
        ks.append(f'Rural Severity {i}-day')
        ks.append(f'Surge {i}-day')
        ks.append(f'Predicted New Deaths Hospital {i}-day')
        ks.append(f'Predicted Deaths Hospital {i}-day')
        df[f'Severity Index {i}-day'] = [remap[x] for x in df[f'Severity {i}-day']]
        ks.append(f'Severity Index {i}-day')
    ks += ['Surge County 3-day', 'tot_deaths', 'SVIPercentile']  # county keys
    return df[ks]


if __name__ == '__main__':
    # load and merge data
    print('severity index loading data...')
    NUM_DAYS_LIST = [1, 2, 3, 4, 5, 6, 7]
    df_county = load_data.load_county_level(data_dir=oj(parentdir, 'data'))
    print('loaded county level!')
    import time
    begin = time.time()
    print('severity index adding predictions...')

    # TODO: remove
    # df_county = add_preds(df_county, NUM_DAYS_LIST=[21], # should save the cached pkl
    #                       cached_dir=None,
    #                       add_predict_interval=True,
    #                       outcomes=['Cases'],
    #                       interval_target_days=[21],
    #                       force_predict=True,
    #                       # truncate to use only the most recent 25% of data and at most 365 days
    #                       expanded_shared_time_truncation=0.9,
    #                       expanded_shared_max_days=90)  # adds keys like "Predicted Deaths 1-day"


    df_county = add_preds(df_county, NUM_DAYS_LIST=NUM_DAYS_LIST + [14, 21, 28], # should save the cached pkl
                          cached_dir=oj(parentdir, 'data'),
                          add_predict_interval=True,
                          interval_target_days=NUM_DAYS_LIST,
                          force_predict=True,
                          # truncate to use only the most recent 25% of data and at most 365 days
                          # TODO: tweak expanded_shared_time_truncation and expanded_shared_max_days for AWS
                          expanded_shared_time_truncation=0.9,
                          expanded_shared_max_days=90)  # adds keys like "Predicted Deaths 1-day"
    end = time.time()
    print('severity index adding predictions took {end - begin:.3f} s')
    print('loading hosp data...')
    # import pdb; pdb.set_trace()
    df_hospital = load_data.load_hospital_level(data_dir=oj(os.path.dirname(parentdir),
                                                            'covid-19-private-data'))
    df = merge_data.merge_county_and_hosp(df_county, df_hospital)
    print('adding severity index...')
    df = add_severity_index(df, NUM_DAYS_LIST)
    df = df.sort_values('Total Deaths Hospital', ascending=False)

    # write to gsheets
    print('\tprep for gsheets')
    dfc = prep_county_df(df_county, NUM_DAYS_LIST)  # data for chicago mapping team
    print('dfc.shape', dfc.shape)
    write_to_gsheet(dfc, sheet_name='County-level Predictions',
                    service_file=oj(parentdir, 'creds.json'))
    ks_output = ['Severity 1-day', 'Severity 2-day', 'Severity 3-day', 'Severity 4-day',
                 'Severity 5-day', 'Severity 6-day', 'Severity 7-day',
                 'Total Deaths Hospital', 'Hospital Name', 'CMS Certification Number',
                 'countyFIPS', 'CountyName', 'StateName',
                 'System Affiliation', 'Latitude', 'Longitude']
    write_to_gsheet(df[ks_output], sheet_name='COVID Severity Index',
                    service_file=oj(parentdir, 'creds.json'))
    print('succesfully wrote to gsheets')

    # write to api
    print('writing to api...')
    da = df.rename(columns={'tot_deaths': 'Total Deaths County',
                            'Predicted New Deaths 3-day': 'Predicted New Deaths County 3-day'})
    extra_ks = ['Surge 3-day', 'Rural Severity 3-day', 'Predicted New Deaths Hospital 3-day',
                'Total Deaths County', 'Predicted New Deaths County 3-day',
                'ICU Beds', 'Hospital Employees']
    da = da[ks_output + extra_ks]
    write_to_api(da, api_file=oj(parentdir, 'ian_key.env'),
                 csv_file=oj(currentdir, '_hidden_hosp.csv'))
    print('succesfully wrote to api')

    # make animated index plot    
    d = df_to_plot(df, NUM_DAYS_LIST)
    print('writing viz index animated...')
    viz_interactive.viz_index_animated(d, [2, 5], out_name=oj(parentdir, 'results', 'hospital_index_animated.html'))
    print('succesfully wrote viz index animated')

    # print
    d = datetime.datetime.today()
    print(f'success! {d.month}_{d.day}_{d.hour}')
