import os
from glob import glob
import re
import pandas as pd
import numpy as np
!pip install matplotlib
!pip install seaborn
!pip install sklearn
!pip install openpyxl
import matplotlib.pyplot as plt
import seaborn as sns

from IPython.display import Markdown
import plotly.express as px

# def _extract_year_from_dir(path):
#     '''
#     path: path to query extract with year in name
    
#     had to query CDC wonder GUI individually by year
#     due to time outs so need to add year variable
#     '''
#     return re.search('\d+',path).group()

def read_cdc_wonder_query(path):
    '''
    path: path to query extract with query info and footnotes
    
    reads text file and skips the last 83 lines which 
    is used for footnotes in CDC wonder extracts
    '''
    df = pd.read_table(path)
    #df['year'] = _extract_year_from_dir(path) #used in previous version
    
    #get rid of query details at end of text file
    no_notes = df.Notes.isna()
    is_total = df.Notes=='Total'
    #print(is_total.sum())
    return df.loc[(no_notes | is_total)]


def convert_float_to_str(series,fill_missing_to=0):
    '''
    for grouping/coding variables, want to be able to see unique 
    values for pandas fxns such as describe().
    
    '''
    
    return series.fillna(fill_missing_to).astype(int).astype(str)


def clean_column_names(df,replace_str='_'):
    #clean column names
    df.columns = [x.lower().replace(' ',replace_str) for x in df.columns]
    
    
def show_every_other_axis_label(fig,axis='x'):
    '''
    makes a plt figure only show every other label for given axis
    '''
    if axis=='x':
        tick_labels = fig.get_xticklabels()
    else:
        tick_labels = fig.get_yticklabels()
        
    for ind, label in enumerate(tick_labels):
        if ind % 2 == 0:  # show every odd every odd year
            label.set_visible(True)
        else:
            label.set_visible(False)
def format_wide_hedegaard_data_for_plot(df,variable_name_key):
    '''
    format cleaned data prepared from the CDC Hedegaard data brief
    for plotting trajectories based on different dimensions
    '''
    formatted_data = (
        df
        .assign(year=lambda x: x.year.astype(int).astype(str)) #make year factor for no decimals
        .set_index('year')
        [variable_name_key.keys()]
        .rename(columns=variable_name_key)
    )
    return formatted_data

def draw_hedegaard_lineplot(df,variable_name_key=None,data_structure='wide',
                            yvar=None,groupvar=None,lineplot_args={}):
    '''
    convert x variable (years) into a factor and format 
    for seaborn lineplot input using specificied 
    variables in df as keys and names rendered as values
    in the input dict (ie variable name key)
    '''
    
    if data_structure=='wide':
        fig = sns.lineplot(
           data= format_wide_hedegaard_data_for_plot(df,variable_name_key),
            **lineplot_args
        )
    else:
        fig = sns.lineplot(data=df, x="year", y=yvar, hue=groupvar,**lineplot_args)
    #customize aesthetics more like paper
    ##tick labels every other year
    show_every_other_axis_label(fig)
    return fig
    

#read in cleaned data from HEAL platform
opioid_gender = pd.read_excel("./deaths_gender.xlsx")
opioid_age = pd.read_excel("./deaths_age_cat.xlsx")
opioid_type = pd.read_excel("./deaths_type_opioid.xlsx")

#format dataframe for lineplot
#get rid of spaces in names
clean_column_names(opioid_gender,replace_str='')

#select and format categories and names
gender_column_labels = {
    'deaths_total_per1k':'Total',
    'deaths_male_per1k':'Male',
    'deaths_female_per1k':'Female'
}

total_df = read_cdc_wonder_query("./cdc_wonder_year_cause_hedegaard_et_al_2020.txt")
clean_column_names(total_df)
#mapping of ICD Codes to labels used in CDC data brief
type_column_labels = {
 'T40.1':'Heroin',
 'T40.2':'Natural and semisynthetic opioids',
 'T40.3': 'Methadone',
 'T40.4':'Synthetic opioids other than methadone'
}

#make a variable with name mapping
total_df['type'] = total_df['multiple_cause_of_death_code'].map(type_column_labels)
#convert year variable to string object for plotting purposes
total_df['year'] = convert_float_to_str(total_df['year'])

#create boolean vector for focus of subsequent analyses (for filtering)
#ie synthetic opioids
is_opioid = total_df['multiple_cause_of_death_code'].isin(type_column_labels.keys())
is_synthetic_type = total_df['multiple_cause_of_death_code']=='T40.4'

#filter only synthetic opioids other than methadone
state_df = read_cdc_wonder_query("./cdc_wonder_year_cause_state_hedegaard_et_al_2020.txt")
clean_column_names(state_df)
type_column_labels = {
 'T40.1':'Heroin',
 'T40.2':'Natural and semisynthetic opioids',
 'T40.3': 'Methadone',
 'T40.4':'Synthetic opioids other than methadone'
}

state_df['type'] = state_df['multiple_cause_of_death_code'].map(type_column_labels)
state_df['year'] = convert_float_to_str(state_df['year'])

#create boolean vector for focus of subsequent analyses (for filtering)
#ie synthetic opioids
is_opioid = state_df['multiple_cause_of_death_code'].isin(type_column_labels.keys())
is_synthetic_type = state_df['multiple_cause_of_death_code']=='T40.4'
#only want individual states and not totals (CDC queries give both)
is_total = np.where(state_df.notes.isna(),False,True) 

#make these objects given they are codes
state_df['state_code'] = convert_float_to_str(state_df['state_code'])
#state_df['county_code'] = convert_float_to_str(state_df['county_code'])

#convert metrics to numeric values and make a key for why missing
data = state_df
for metric_name in ['crude_rate','deaths','population','age_adjusted_rate']:
    #get missing val strings
    reasons_missing = ['Suppressed','Unreliable','Missing',None]
    is_missing = data[metric_name].isin(reasons_missing) 

    #make a key showing to retain info on why missing
    data[metric_name+'_missing_desc'] = data[metric_name]
    data.loc[~is_missing,metric_name+'_missing_desc'] = ''

    #missing for diff reasons but need to make float var from string -- if missing, make a NaN
    data.loc[is_missing,metric_name] = np.nan #change missing val strings to NaNs
    data[metric_name] = data[metric_name].astype(float) #change to var to float
del data

#2019,synthetic categorical bar plot organized from smallest to largest
is_2019 = state_df.year=='2019'
state_2019_synthetic_sorted_df = (
    state_df
    .loc[(is_2019 & is_synthetic_type & ~is_total)]
    .sort_values(by='age_adjusted_rate',ascending=False)
)

#median state death rate for comparison to individual
median_state_death_rate = state_2019_synthetic_sorted_df.age_adjusted_rate.median()

#plotly only accepts state abbreviations if not using GEOJSON
state_abbrevs = {'Alabama': 'AL',
        'Alaska': 'AK',
        'Arizona': 'AZ',
        'Arkansas': 'AR',
        'California': 'CA',
        'Colorado': 'CO',
        'Connecticut': 'CT',
        'Delaware': 'DE',
        'District of Columbia': 'DC',
        'Florida': 'FL',
        'Georgia': 'GA',
        'Hawaii': 'HI',
        'Idaho': 'ID',
        'Illinois': 'IL',
        'Indiana': 'IN',
        'Iowa': 'IA',
        'Kansas': 'KS',
        'Kentucky': 'KY',
        'Louisiana': 'LA',
        'Maine': 'ME',
        'Maryland': 'MD',
        'Massachusetts': 'MA',
        'Michigan': 'MI',
        'Minnesota': 'MN',
        'Mississippi': 'MS',
        'Missouri': 'MO',
        'Montana': 'MT',
        'Nebraska': 'NE',
        'Nevada': 'NV',
        'New Hampshire': 'NH',
        'New Jersey': 'NJ',
        'New Mexico': 'NM',
        'New York': 'NY',
        'North Carolina': 'NC',
        'North Dakota': 'ND',
        'Ohio': 'OH',
        'Oklahoma': 'OK',
        'Oregon': 'OR',
        'Pennsylvania': 'PA',
        'Rhode Island': 'RI',
        'South Carolina': 'SC',
        'South Dakota': 'SD',
        'Tennessee': 'TN',
        'Texas': 'TX',
        'Utah': 'UT',
        'Vermont': 'VT',
        'Virginia': 'VA',
        'Washington': 'WA',
        'West Virginia': 'WV',
        'Wisconsin': 'WI',
        'Wyoming': 'WY'}

state_2019_synthetic_sorted_df['State'] = (
    state_2019_synthetic_sorted_df['state']
    .map(state_abbrevs)
)

from sklearn.cluster import KMeans
kmeans = KMeans(3,random_state=100)

# make a Kmeans cluster model and assign the cluster to each state
kmeans_groups = kmeans.fit_predict(
    (state_2019_synthetic_sorted_df
         [['age_adjusted_rate']]
         .fillna(median_state_death_rate)),
    
    ) + 1
state_2019_synthetic_sorted_df['kmeans_cluster'] = kmeans_groups

cluster_name_map = {
    3:'High',
    1:'Medium',
    2:'Low',
    -1:'Missing data'
}

#needed to impute missing values. Now I am going to add back in the missing values for viz
is_na_rate = state_2019_synthetic_sorted_df.age_adjusted_rate.isna()
state_2019_synthetic_sorted_df.loc[is_na_rate,'kmeans_cluster'] = -1

state_2019_synthetic_sorted_df['Death Rate Cluster'] = (
    state_2019_synthetic_sorted_df
    ['kmeans_cluster']
    .map(cluster_name_map)
)