import argparse
from helper_funcs import *
import pandas as pd
import matplotlib.pyplot as plt
from pandas.plotting import register_matplotlib_converters
register_matplotlib_converters()
import pmdarima
from statsmodels.tsa.arima.model import ARIMA

# -------------------------------------------------- ARIMA --------------------------------------------------
def arima(l_args, s_ticker, s_interval, df_stock):
    parser = argparse.ArgumentParser(prog='arima',
                                     description="""In statistics and econometrics, and in particular in time
                                     series analysis, an autoregressive integrated moving average (ARIMA) model
                                     is a generalization of an autoregressive moving average (ARMA) model. Both
                                     of these models are fitted to time series data either to better understand
                                     the data or to predict future points in the series (forecasting).
                                     ARIMA(p,d,q) where parameters p, d, and q are non-negative integers, p is
                                     the order (number of time lags) of the autoregressive model, d is the degree
                                     of differencing (the number of times the data have had past values subtracted),
                                     and q is the order of the moving-average model.""")

    parser.add_argument('-d', "--days", action="store", dest="n_days", type=check_positive, default=5,
                        help='prediction days.')
    parser.add_argument('-i', "--ic", action="store", dest="s_ic", type=str, default='aic',
                        choices=['aic', 'aicc', 'bic', 'hqic', 'oob'], help='information criteria.')
    parser.add_argument('-s', "--seasonal", action="store_true", default=False, dest="b_seasonal",
                        help='Use weekly seasonal data.')
    parser.add_argument('-o', "--order", action="store", dest="s_order", type=str,
                        help='arima model order (p,d,q) in format: pdq.')
    parser.add_argument('-r', "--results", action="store_true", dest="b_results", default=False,
                        help='results about ARIMA summary flag.')

    try:
        (ns_parser, l_unknown_args) = parser.parse_known_args(l_args)

        if l_unknown_args:
            print(f"The following args couldn't be interpreted: {l_unknown_args}\n")
            return

        # Machine Learning model
        if ns_parser.s_order:
            t_order = tuple([int(ord) for ord in list(ns_parser.s_order)])
            model = ARIMA(df_stock['5. adjusted close'].values, order=t_order).fit()
            l_predictions = model.predict(start=len(df_stock['5. adjusted close'])+1, end=len(df_stock['5. adjusted close'])+ns_parser.n_days)
        else:
            if ns_parser.b_seasonal:
                model = pmdarima.auto_arima(df_stock['5. adjusted close'].values, error_action='ignore', seasonal=True, m=5, information_criteria=ns_parser.s_ic)
            else:
                model = pmdarima.auto_arima(df_stock['5. adjusted close'].values, error_action='ignore', seasonal=False, information_criteria=ns_parser.s_ic)
            l_predictions = model.predict(n_periods=ns_parser.n_days)

        # Prediction data
        l_pred_days = get_next_stock_market_days(last_stock_day=df_stock['5. adjusted close'].index[-1], n_next_days=ns_parser.n_days)
        df_pred = pd.Series(l_predictions, index=l_pred_days, name='Price')

        if ns_parser.b_results:
            print(model.summary())
            print("")

        # Plotting
        plt.plot(df_stock.index, df_stock['5. adjusted close'], lw=2)
        if ns_parser.s_order:
            plt.title(f"ARIMA {str(t_order)} on {s_ticker} - {ns_parser.n_days} days prediction")
        else:
            plt.title(f"ARIMA {model.order} on {s_ticker} - {ns_parser.n_days} days prediction")
        plt.xlim(df_stock.index[0], get_next_stock_market_days(df_pred.index[-1], 1)[-1] )
        plt.xlabel('Time')
        plt.ylabel('Share Price ($)')
        plt.grid(b=True, which='major', color='#666666', linestyle='-')
        plt.minorticks_on()
        plt.grid(b=True, which='minor', color='#999999', linestyle='-', alpha=0.2)
        plt.plot([df_stock.index[-1], df_pred.index[0]], [df_stock['5. adjusted close'].values[-1], df_pred.values[0]], lw=1, c='tab:green', linestyle='--')
        plt.plot(df_pred.index, df_pred, lw=2, c='tab:green')
        plt.axvspan(df_stock.index[-1], df_pred.index[-1], facecolor='tab:orange', alpha=0.2)
        xmin, xmax, ymin, ymax = plt.axis()
        plt.vlines(df_stock.index[-1], ymin, ymax, linewidth=1, linestyle='--', color='k')
        plt.show()

        # Print prediction data
        print("Predicted share price:")
        df_pred = df_pred.apply(lambda x: f"{x:.2f} $")
        print(df_pred.to_string())
        print("")

    except:
        print("")