# Import packages

import numpy as np

import matplotlib

import pd

import pandas as pd

from datetime import datetime

import us

import uszipcode

matplotlib.rcParams['figure.figsize'] = (10, 10)

# first step: vectorize the columns to convert strings into floats

# grab the columns

clin_addresses = pd.read_csv("/Users/shania/PycharmProjects/ClinicalAttritionRateMap/updated_clinical_attrition.csv")

# define a function that converts the dates into just the month and year

def parse_date(date_str):

    if pd.isna(date_str):

        return np.nan  # Return null

    try: # if it's in a month day year format, convert to month year

        return datetime.strptime(date_str, '%B %d, %Y').strftime('%B %Y')

    except ValueError:

        try:

            # already in the month year format, keep it/no change

            return datetime.strptime(date_str, '%B %Y').strftime('%B %Y')

        except ValueError:

            return np.nan  # Return null if both formats fail

# Apply the parse_date function the dates columns

clin_addresses['New Completion Date'] = clin_addresses['Completion Date'].apply(parse_date)

clin_addresses['New Start Date'] = clin_addresses['Start Date'].apply(parse_date)

# convert from strings to integers

clin_addresses['New Completion Date'] = pd.to_datetime(clin_addresses['New Completion Date'], errors='coerce', format='%B %Y')

clin_addresses['New Start Date'] = pd.to_datetime(clin_addresses['New Start Date'], errors='coerce', format='%B %Y')

#  difference in days

clin_addresses['length_of_trial'] = (clin_addresses['New Completion Date'] - clin_addresses['New Start Date']).dt.days

# Drop the original start_date and completion_date columns

# clin_addresses = clin_addresses.drop(columns=['Start Date', 'Completion Date'])

# Save the updated DataFrame to a new CSV file

clin_addresses.to_csv("/Users/shania/PycharmProjects/ClinicalAttritionRateMap/updated_clinical_attrition.csv", index=False)

# add the ruca codes

ruca_codes = pd.read_excel("2006 Complete Excel RUCA file 3.xls")

clin_addresses['Zipcode'] = clin_addresses['Zipcode'].astype(str)

ruca_codes['ZIPA'] = ruca_codes['ZIPA'].astype(str)

final_table = pd.merge(clin_addresses, ruca_codes, left_on='Zipcode', right_on='ZIPA', how='left')

drop_columns = ['ZIPN', 'ZIPA', 'STNAME','COMMFLAG']

final_table = final_table.drop(columns=drop_columns)

final_table.to_csv("/Users/shania/PycharmProjects/ClinicalAttritionRateMap/final_table.csv", index=False)

# define a function for the zipcodes

## if the zipcodes do not have exactly 5 digits

## if the zipcodes contain a letter

## if the zipcodes contain a "dash" take the first part

def cleaned_zipcodes(final_table):

    search = uszipcode.SearchEngine() # search engine comes with package

    clean_zipcode = []

    for i, city in zip(final_table['Zipcode'], final_table['City']): # zip zipcodes and city together

        if "-" in str(i):

            clean_zipcode.append(str(i).split('-')[0])

        elif pd.isna(i):

            result = search.by_city(city)

            clean_zipcode.append(result[0].zipcode if result else None)

        else:

            clean_zipcode.append(i)

    return clean_zipcode

final_table['Cleaned_Zipcodes'] = cleaned_zipcodes(final_table)

final_table.to_csv("/Users/shania/PycharmProjects/ClinicalAttritionRateMap/final_table.csv", index=False)

def cleaned_states(final_table):

    clean_state = []

    for state in final_table['State']:

        if state in [str(s) for s in us.states.STATES]:

            clean_state.append(state)

        else:

            clean_state.append(None)

    return clean_state

final_table['Cleaned States'] = cleaned_states(final_table)

final_table.dropna(subset=['Cleaned States'], inplace=True)

final_table.drop(columns='Cleaned States', inplace=True)

final_table.to_csv("/Users/shania/PycharmProjects/ClinicalAttritionRateMap/final_table.csv", index=False)

def convert_minimum_age(age):

    age = str(age)

    if "Years" in age:

        return int(''.join(filter(str.isdigit, age)))

    elif "Months" in age:

        return float(''.join(filter(str.isdigit, age))) / 12

    else:

        return None

final_table['Minimum Age'] = final_table['Minimum Age'].apply(convert_minimum_age)

final_table.to_csv("/Users/shania/PycharmProjects/ClinicalAttritionRateMap/final_table.csv", index=False)