import datetime
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
# import pandas_datareader as web
import pymc3 as pm
from pymc3.distributions.timeseries import GaussianRandomWalk
import seaborn as sns

import pickle
import warnings
warnings.filterwarnings('ignore')


def obtain_plot_amz_dataset(start_date, end_date):
    amzn = pd.read_csv('data/amzn.csv')
    print("Loading and plotting AMZN log returns...")
    amzn['returns'] = amzn['Adj Close']/amzn['Adj Close'].shift(1)
    amzn.dropna(inplace=True)
    amzn['log_returns'] = np.log(amzn['returns'])
    amzn['log_returns'].plot(linewidth=0.5)
    plt.ylabel('AMZN daily percentage returns')
    plt.show()
    return amzn
# Model Specification in PyMC3 ------------------------


def configure_sample_stoch_vol_model(log_returns, samples):
    print("Configuring stochastic volatility with PyMC3...")
    model = pm.Model()

    with model:
        sigma = pm.Exponential('sigma', 50.0, testval=0.1)
        nu = pm.Exponential('nu', 0.1)
        s = GaussianRandomWalk('s', sigma=(sigma**-2), shape=len(log_returns))
        logrets = pm.StudentT(
            'logrets', nu, lam=pm.math.exp(-2.0*s), observed=log_returns)

    # Fitting the Model with NUTS
    print('Fitting the stochastic volatility model')
    with model:
        trace = pm.sample(samples)

        with open('data\\trace.pickle', 'wb') as f:
            pickle.dump(trace, f, protocol=pickle.HIGHEST_PROTOCOL)

        # with open('trace.pickle', 'rb') as f:
        #     trace = pickle.load(f)

        # with open('trace2.pickle', 'rb') as f:
        #     trace = pickle.load(f)
    '''
    # print(model.vars[:-1])
    # error when call traceplot
    pm.traceplot(trace, model.vars)
    plt.show()
    '''

    print("Plotting the log volatility...")
    k = 10      # step-size
    opacity = 0.03

    plt.plot(trace[s][::k].T, 'b', alpha=opacity)       # 'b' mean blue
    plt.xlabel('Time')
    plt.ylabel('Log Volatility')
    plt.show()

    print("Plotting the absolute returns overlaid with vol...")
    plt.plot(np.abs(np.exp(log_returns))-1.0, linewidth=0.5)
    plt.plot(np.exp(trace[s][::k].T), 'r', alpha=opacity)

    plt.xlabel("Trading Days")
    plt.ylabel("Absolute Returns/Volatility")
    plt.show()


if __name__ == "__main__":
    start_date = datetime.datetime(2006, 1, 1)
    end_date = datetime.datetime(2015, 12, 31)

    # Obtain and plot the logarithmic returns of amz prices
    df = obtain_plot_amz_dataset(start_date, end_date)
    log_returns = np.array(df['log_returns'])

    samples = 500   # 2000
    configure_sample_stoch_vol_model(log_returns, samples)
    pass
