import sys
import pandas

file_loc = 'data/influenza/US/US.6142004.csv'
df = pandas.read_csv(
    file_loc, 
    low_memory=False, 
    parse_dates=['PeriodStartDate', 'PeriodEndDate'], 
    date_parser=lambda x: pandas.datetime.strptime(x, '%Y-%m-%d')
)

pandas.unique(df[['PathogenName', 'AgeRange', 'DiagnosisCertainty']].values.ravel('K'))

from src.lib.parse import overlaps, combine_rows

# these are the columns used to determine 
#  if two rows should be checked for overlapping date ranges
# *in order* (least to most specific)
match_on = [
    'Admin1ISO',
    'Admin2Name',
    'CityName',
    'Fatalities'
]

sort_by = ['PeriodStartDate']
found = overlaps(df, match_on, sort_by)
dupes = pandas.DataFrame(f[1][1] for f in found)
dupes.sort_values(by=['CountValue','CityName'], ascending=False)[[
    'CityName', 
    'PeriodStartDate', 
    'PeriodEndDate', 
    'CountValue'
]]

ex_1 = (df.CityName == 'NEWARK') &(df.PeriodStartDate >= '1923-11-01') &(df.PeriodStartDate <= '1923-11-15')
ex_2 = (df.CityName == 'LOS ANGELES') &(df.PeriodStartDate >= '1937-05-15') &(df.PeriodStartDate <= '1937-05-31')
ex_3 = (df.CityName == 'CLEVELAND') &(df.PeriodStartDate >= '1923-11-01') &(df.PeriodStartDate <= '1923-11-30')
df.loc[ex_1 | ex_2 | ex_3].sort_values(by=[
    'CityName',
    'Fatalities', 
    'PeriodStartDate'
])[['CityName', 'PeriodStartDate', 'PeriodEndDate', 'CountValue']]

to_drop = [combine_rows(prev, cur, df) for prev, cur in overlaps(df, match_on, sort_by)]
df.drop(to_drop, inplace=True)

found = overlaps(df, match_on, sort_by)
try:
    print(next(found))
except StopIteration:
    print("No overlaps found!")

import seaborn as sns
import matplotlib.pyplot as plt

%matplotlib inline

sns.set(rc={'figure.figsize':(20,15)})

ax = sns.scatterplot(
    x='PeriodStartDate', 
    y='CountValue', 
    hue='Admin1ISO', 
    data=df[df.Fatalities == 1]
)

ax.set_yscale('symlog')
ax.set(xlim=(min(df.PeriodStartDate), max(df.PeriodStartDate) + pandas.DateOffset(years=1)))

plt.show()

%matplotlib inline

sns.set(rc={'figure.figsize':(25,15)})

ax = sns.scatterplot(
    x='PeriodStartDate', 
    y='CountValue', 
    hue='Admin1ISO', 
    data=df[df.Fatalities == 0]
)

ax.set_yscale('symlog')
ax.set(xlim=(min(df.PeriodStartDate), max(df.PeriodStartDate) + pandas.DateOffset(years=1)))

plt.show()

%matplotlib inline

sns.set(rc={'figure.figsize':(20,15)})

ax = sns.scatterplot(
    x='PeriodStartDate', 
    y='CountValue', 
    hue='Admin1ISO', 
    data=df
)

ax.set_yscale('symlog')
ax.set(xlim=(min(df.PeriodStartDate), max(df.PeriodStartDate) + pandas.DateOffset(years=1)))

plt.show()

states_east_to_west = [
    'US-ME', 'US-NH', 'US-VT', 'US-MA', 'US-RI', 'US-CT', 'US-NY', 'US-NJ', 'US-PA', 'US-DE', 'US-MD', 'US-DC',
    'US-WV', 'US-VA', 'US-NC', 'US-SC', 'US-GA', 'US-FL',
    'US-MI', 'US-IN', 'US-OH', 'US-KY', 'US-TN', 'US-AL', 'US-MS',
    'US-WI', 'US-IL', 
    'US-MN', 'US-IA', 'US-MO', 'US-AR', 'US-LA',
    'US-ND', 'US-SD', 'US-NE', 'US-KS', 'US-OK', 'US-TX',
    'US-MT', 'US-WY', 'US-CO', 'US-NM',
    'US-ID', 'US-NV', 'US-UT', 'US-AZ',
    'US-WA', 'US-OR', 'US-CA', 
    'US-HI', 'US-AK'


]

df_pv = df[df.Fatalities == 1].groupby(['Admin1ISO', df['PeriodStartDate'].dt.year])['CountValue'].sum()
df_pv = df_pv.reset_index().pivot('Admin1ISO','PeriodStartDate', 'CountValue').reindex(states_east_to_west, axis=0).fillna(0)

sns.heatmap(
    df_pv, 
    cmap=sns.diverging_palette(200, 275, s=99, l=70, as_cmap=True),
    center=800
)

df_pv = df[df.Fatalities == 0].groupby(['Admin1ISO', df['PeriodStartDate'].dt.year])['CountValue'].sum()
df_pv = df_pv.reset_index().pivot('Admin1ISO','PeriodStartDate', 'CountValue').reindex(states_east_to_west, axis=0).fillna(0)

sns.heatmap(
    df_pv, 
    cmap=sns.diverging_palette(200, 275, s=99, l=70, as_cmap=True),
    center=60000
)

df_pv = df[df.Fatalities == 0].groupby(['Admin1ISO', df['PeriodStartDate'].dt.year])['CountValue'].sum()
df_pv = df_pv.reset_index().pivot('Admin1ISO','PeriodStartDate', 'CountValue').reindex(states_east_to_west, axis=0).fillna(0)

sns.heatmap(
    df_pv, 
    cmap=sns.diverging_palette(200, 275, s=99, l=70, as_cmap=True),
    center=60000
)

# %load src/lib/parse.py
def between(s, e, d):
    return s <= d <= e


def is_row_between(prev, cur, match_on):
    """ 
        returns if the date range of prev 
        has any overlap 
        with the date range cur
        and if they have matching match_on values
    """
    (_, r1), (_, r2) = prev, cur
    d1 = (r1.PeriodStartDate, r1.PeriodEndDate, r2.PeriodStartDate)
    d2 = (r1.PeriodStartDate, r1.PeriodEndDate, r2.PeriodEndDate)
    return all(r1[m] == r2[m] for m in match_on) and (between(*d1) or between(*d2))


def iter_with_prev(df):
    """ 
        return an iterator that gives back a tuple on next
        ((i, previous_row), (i+1, previous_row))
        first value is ((0, df[0]), (1, df[1]))
    """
    try:
        rows = df.iterrows()
        prev, curr = next(rows), next(rows)
        yield prev, curr
        prev = curr
        for row in rows:
            yield prev, row
            prev = row
    except StopIteration: # could yield back ((1, df[1]), None) ?
        pass


def combine_rows(prev, cur, df): #, r2_i, df
    """
        combines the two rows into r2, drops r_1
            date range is smallest date range that includes r1 & r2
            count is sum(counts) (the worst caste)
                this could be avg, min, max etc.
        returns the index of the previous row
    """
    (r1_i, r1), (r2_i, r2) = prev, cur
    dates = (r1.PeriodStartDate, r1.PeriodEndDate, r2.PeriodStartDate, r2.PeriodEndDate)
    counts = (r1.CountValue, r2.CountValue)
    df.at[r2_i, 'PeriodStartDate'] = min(*dates)
    df.at[r2_i, 'PeriodEndDate'] = max(*dates)
    df.at[r2_i, 'CountValue'] = sum(counts)
    return r1_i


def overlaps(df, match_on, sort_by):
    """
        yields the previous and current row
        given rows to match on, sort by and a data frame
    """
    for prev, curr in iter_with_prev(df.sort_values(by=match_on + sort_by)):
        if is_row_between(prev, curr, match_on):
            yield prev, curr

df_pv2 = df_pv.div(df_pv.max(axis=1), axis=0)

sns.heatmap(
    df_pv2,
    cmap = 'Spectral_r', 
    linewidths=1,
    square=True,
    cbar=False
)

df_pv2 = df_pv.div(df_pv.max(axis=0), axis=1)

sns.heatmap(
    df_pv2,
    cmap = 'Spectral_r', 
    linewidths=1,
    square=True,
    cbar=False
)