import numpy as np
from sklearn import linear_model

reg = linear_model.LinearRegression(fit_intercept=True)
import pandas as pd

from scipy import signal


def get_doubling_time(array):
    # Use a linear regression to approximate the doubling rate

    y = np.array(array)
    X = np.arange(-1, 2).reshape(-1, 1)

    assert len(array) == 3
    reg.fit(X, y)
    intercept = reg.intercept_
    slope = reg.coef_

    return intercept / slope


def savgol_filter(df_input, column='confirmed', window=5):
    #  Savgol Filter which can be used in groupby apply function (data structure kept)

    df_result = df_input

    filter_in = df_input[column].fillna(0)  # attention with the neutral element here

    result = signal.savgol_filter(np.array(filter_in),
                                  window,  # window size used for filtering
                                  1)
    df_result[str(column + '_filtered')] = result
    return df_result


def rolling_regression(df_input, col='confirmed'):
    # Rolling Regression to approximate the doubling time'

    days_back = 3
    result = df_input[col].rolling(
        window=days_back,
        min_periods=days_back).apply(get_doubling_time, raw=False)

    return result


def calculate_doubling_rate(df_input, filter_on='confirmed'):
    # Calculate approximated doubling rate and return merged data frame

    must_contain = set(['state', 'country', filter_on])
    assert must_contain.issubset(set(df_input.columns)), ' Erro in calc_filtered_data not all columns in data frame'

    pd_DR_result = df_input.groupby(['state', 'country']).apply(rolling_regression, filter_on).reset_index()

    pd_DR_result = pd_DR_result.rename(columns={filter_on: filter_on + '_DR',
                                                'level_2': 'index'})

    df_output = pd.merge(df_input, pd_DR_result[['index', str(filter_on + '_DR')]], left_index=True, right_on=['index'],
                         how='left')
    df_output = df_output.drop(columns=['index'])

    return df_output


def calculate_filtered_data(df_input, filter_on='confirmed'):
    # Calculate savgol filter and return merged data frame

    must_contain = set(['state', 'country', filter_on])
    assert must_contain.issubset(set(df_input.columns)), ' Erro in calc_filtered_data not all columns in data frame'

    df_output = df_input.copy()  # we need a copy here otherwise the filter_on column will be overwritten

    pd_filtered_result = df_output[['state', 'country', filter_on]].groupby(['state', 'country']).apply(
        savgol_filter)  # .reset_index()

    df_output = pd.merge(df_output, pd_filtered_result[[str(filter_on + '_filtered')]], left_index=True,
                         right_index=True, how='left')
    return df_output.copy()


def pd_result_large():
    pd_data = pd.read_csv('../../data/processed/COVID_relational_confirmed.csv', sep=';', parse_dates=[0])
    pd_data = pd_data.sort_values('date', ascending=True).copy()

    pd_result_large = calculate_filtered_data(pd_data)
    pd_result_large = calculate_doubling_rate(pd_result_large)
    pd_result_large = calculate_doubling_rate(pd_result_large, 'confirmed_filtered')

    mask = pd_result_large['confirmed'] > 100
    pd_result_large['confirmed_filtered_DR'] = pd_result_large['confirmed_filtered_DR'].where(mask, other=np.NaN)
    pd_result_large.to_csv('../../data/processed/COVID_final_set.csv', sep=';', index=False)
