#!/usr/bin/python3.5

import copy
import datetime
import matplotlib.pyplot as plt
import numpy as np
import os
import pandas as pd
import quandl
import sklearn

from datetime import timedelta
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression

import util
import quandl_data

from util import *
from quandl_data import *

y_hat_frbny = None

def load_data (refresh=False):
    global y_hat_frbny
    if refresh:
        get_quandl_data ()
    else:
        for dc in QUANDL_DATA:
            file_name = Constant.DATA_DIR + "/" + dc + ".csv"
            if os.path.isfile (file_name):
                QUANDL_DATA [dc] ["data"] = pd.read_csv (file_name)
          
    # sn: needs manual updates - nowcasts not avaialable in feed format
    y_hat_frbny = pd.read_csv (Constant.DATA_DIR + "/ny_fed.csv")


# sn: this function returns X, y - the full data set since the beginning of y
# each day in the prior 2 years is a feature
# the function takes as argument the number of days prior to quarter end that we are
# predicting the GDP growth for
# x_df and y_df must have the column 'Date' and x_df must also have the column 'Total'
def featurize_series (days_to_qtr_end, x_df, y_df, num_days_per_period=1, num_years=2):
    X = []
    y = []
    quarters = []
    DAYS_PER_YEAR = 365

    for idx, row in y_df.iterrows ():
        #quarters.append (convert_yyyymmdd_str_to_qtr(row ['Date']))
        #qtr_end = datetime.datetime.strptime (row ['Date'], fmt)
        qtr_end = row ['Date']
        x = [0] * num_years * int (DAYS_PER_YEAR/num_days_per_period)
        for idx2, row2 in x_df.iterrows ():
            series_row_date = row2 ['Date']
            if (series_row_date <= qtr_end + timedelta (days=-days_to_qtr_end)) and \
                    (series_row_date > qtr_end + timedelta (days=-2 * num_years * DAYS_PER_YEAR)):
                feature_col_idx = int ((qtr_end - series_row_date).days / num_days_per_period)
                if feature_col_idx >= len (x):
                    feature_col_idx = len (x) - 1
                x [feature_col_idx] = row2 ['Total']
       
        X.append (x)
        quarters.append (qtr_end)
        y.append (row ['Gross domestic product'])

    return np.array (X), np.array (y)


def main ():
    load_data (refresh=False)
    # set up aliases
    frbny_nowcast_df = y_hat_frbny

    nonfarm_df = QUANDL_DATA ["EMPL_NONFARM_ADP"] ["data"]
    capex_df = QUANDL_DATA ['CAPEX_FC_FedNY']['data']
    act_gdp_df = QUANDL_DATA ["QTRLY_GDP_PCT_CHG_Bea"] ["data"]

    nonfarm_df ['Date'] = nonfarm_df ['Month'].apply (datetime.datetime.strptime, args=(Constant.DATE_STR_FMT_1,))
    nonfarm_df ['Total'] = nonfarm_df ['1-49'] + nonfarm_df ['50-499'] + nonfarm_df ['500+']

    capex_df ['Date'] = capex_df ['Date'].apply (datetime.datetime.strptime, args=(Constant.DATE_STR_FMT_1,))
    capex_df ['Total'] = capex_df ['Value']

    biz_acy

    act_gdp_df ['Date'] = act_gdp_df ['Date'].apply (datetime.datetime.strptime, args=(Constant.DATE_STR_FMT_1,))

    # featurize and normalize the data
    X = None
    y = None
    # first data point is 2005-01-31
    X1, y1 = featurize_series (0, nonfarm_df, \
        act_gdp_df [act_gdp_df ['Date'] >= datetime.datetime.strptime ('2005-01-01', \
        Constant.DATE_STR_FMT_1)], num_days_per_period=91)
    # first data point is 2001-07-31
    X2, y2 = featurize_series (0, capex_df, \
        act_gdp_df [act_gdp_df ['Date'] >= datetime.datetime.strptime ('2005-01-01', \
        Constant.DATE_STR_FMT_1)], num_days_per_period=91)

    scale_data = False
    
    if scale_data:
        min_y, ptp_y, scaled_y1 = normalize_series (y1)
        min_X1, ptp_X1, scaled_X1 = normalize_series (X1)
        min_X2, ptp_X2, scaled_X2 = normalize_series (X2)

        X = np.concatenate ((scaled_X1, scaled_X2), axis=1)
        y = scaled_y1
    else:
        X = np.concatenate ((X1, X2), axis=1)
        y = y1

    X_train, X_test, y_train, y_test = train_test_split (X, y, test_size=0.2, random_state=2)

    regr = LinearRegression ()
    regr.fit (X_train, y_train)
    y_preds = regr.predict (X_test)

    print ("Predictions vs actual values:\n")
    print (["{0:.1f}".format(y_pred) for y_pred in y_preds])
    print ("-----------------------------------------------------------")
    print (["{0:.1f}".format(y_act) for y_act in y_test])
    print ("-----------------------------------------------------------")
    print ("MSE: " + str (((y_preds - y_test) ** 2).mean (axis=None)))
    
if __name__ == "__main__":
    main()
