# Basic imports
import os
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import datetime as dt
import scipy.optimize as spo
import sys
from time import time
from sklearn.metrics import r2_score, median_absolute_error
from multiprocessing import Pool
import pickle

%matplotlib inline

%pylab inline
pylab.rcParams['figure.figsize'] = (20.0, 10.0)

%load_ext autoreload
%autoreload 2

sys.path.append('../../')

import recommender.simulator as sim
from utils.analysis import value_eval
from recommender.agent import Agent
from functools import partial

def show_results(results_list, data_in_df, symbol, graph=False):
    for values in results_list:
        total_value = values.sum(axis=1)
        print('Sharpe ratio: {}\nCum. Ret.: {}\nAVG_DRET: {}\nSTD_DRET: {}\nFinal value: {}'.format(*value_eval(pd.DataFrame(total_value))))
        print('-'*100)
        initial_date = total_value.index[0]
        compare_results = data_in_df.loc[initial_date:, 'Close'].copy()
        compare_results.name = symbol
        compare_results_df = pd.DataFrame(compare_results)
        compare_results_df['portfolio'] = total_value
        std_comp_df = compare_results_df / compare_results_df.iloc[0]
        if graph:
            plt.figure()
            std_comp_df.plot()

"""
Explanation: In this notebook a sensitivity analysis will be performed. The best Automatic Trader will be trained and tested many times, with the SPY symbol, to assess the sensitivity to changes in the period chosen for training and testing.
End of explanation
"""


n_train = 20

data_df = pd.read_pickle('../../data/data_df.pkl')
u_data_df = data_df.unstack()
n_samples = u_data_df.shape[0]
n_test = 504
first_test_start = n_samples - 4 * n_test
test_starts = np.linspace(first_test_start, n_samples - n_test, n_train).astype(np.int)
train_test_pairs = [(u_data_df.iloc[:test_start], u_data_df.iloc[test_start:test_start + n_test]) 
                    for test_start in test_starts]

"""
Explanation: Let's create some training / testing pairs
End of explanation
"""


for data_train_df, data_test_df in train_test_pairs:
    print(data_train_df.shape)
    print(data_test_df.shape)
    print('-'*100)

"""
Explanation: Just to check, let's print the shapes of the different train/test pairs
End of explanation
"""


n_train = 20

# Set the train/test pairs
n_test = 504
data_df = pd.read_pickle('../../data/data_df.pkl')
u_data_df = data_df.unstack()
n_samples = u_data_df.shape[0]
first_test_start = n_samples - 4 * n_test
test_starts = np.linspace(first_test_start, n_samples - n_test, n_train).astype(np.int)
train_test_pairs = [(u_data_df.iloc[:test_start], u_data_df.iloc[test_start:test_start + n_test]) 
                    for test_start in test_starts]

STARTING_DAYS_AHEAD = 252
POSSIBLE_FRACTIONS = [0.0, 1.0]

symbol = 'SPY'
print('Testing with:\n{}'.format(symbol))

# Results holder
results_df = pd.DataFrame(columns=['sharpe', 
                                   'sharpe_bench', 
                                   'cum_ret', 
                                   'cum_ret_bench', 
                                   'sharpe_quotient', 
                                   'cum_ret_quotient',])

for data_train_all_df, data_test_all_df in train_test_pairs:
    # Set the datasets to use
    total_data_train_df = data_train_all_df.stack()
    total_data_test_df = data_test_all_df.stack()
    data_train_df = data_train_all_df[symbol]
    data_test_df = data_test_all_df[symbol]
    data_in_df = data_train_df
    
    # Set the environment
    env, num_states, num_actions = sim.initialize_env(total_data_train_df, 
                                                  symbol, 
                                                  starting_days_ahead=STARTING_DAYS_AHEAD,
                                                  possible_fractions=POSSIBLE_FRACTIONS,
                                                  n_levels=10)
    # Create the Agent
    agent = Agent(num_states=num_states, 
                  num_actions=num_actions, 
                  random_actions_rate=0.98, 
                  random_actions_decrease=0.9999,
                  dyna_iterations=0,
                  name='Agent_0')
    
    # Train
    print('Training: {}'.format(symbol))
    n_epochs = 7

    for i in range(n_epochs):
        tic = time()
        env.reset(STARTING_DAYS_AHEAD)
        results_list = sim.simulate_period(total_data_train_df, 
                                           symbol,
                                           agent,
                                           starting_days_ahead=STARTING_DAYS_AHEAD,
                                           possible_fractions=POSSIBLE_FRACTIONS,
                                           verbose=False,
                                           other_env=env)
        toc = time()
        print('Epoch: {}'.format(i))
        print('Elapsed time: {} seconds.'.format((toc-tic)))
        print('Random Actions Rate: {}'.format(agent.random_actions_rate))
        show_results([results_list], data_in_df, symbol)
    
    
    # Test
    print('-'*100)
    print('Testing: {}'.format(symbol))
    TEST_DAYS_AHEAD = 20
    env.set_test_data(total_data_test_df, TEST_DAYS_AHEAD)
    tic = time()
    results_list = sim.simulate_period(total_data_test_df, 
                                    symbol,
                                    agent,
                                    learn=False,
                                    starting_days_ahead=TEST_DAYS_AHEAD,
                                    possible_fractions=POSSIBLE_FRACTIONS,
                                    verbose=False,
                                    other_env=env)
    toc = time()
    
    # Show results
    print('Elapsed time: {} seconds.'.format((toc-tic)))
    print('Random Actions Rate: {}'.format(agent.random_actions_rate))
    show_results([results_list], data_test_df, symbol, graph=False)
    print('-'*100)
    
    # Save some results
    total_value = results_list.sum(axis=1)
    sharpe, cumret, _, _, _ = value_eval(pd.DataFrame(total_value))
    sharpe_bench, cumret_bench, _, _, _ = value_eval(pd.DataFrame(data_test_df['Close'].iloc[TEST_DAYS_AHEAD:]))
    sharpe_q = sharpe / sharpe_bench
    cumret_q = cumret / cumret_bench
    results_df = results_df.append(pd.DataFrame({symbol: {'sharpe': sharpe,
                                         'sharpe_bench': sharpe_bench,
                                         'cum_ret': cumret,
                                         'cum_ret_bench': cumret_bench,
                                         'sharpe_quotient': sharpe_q,
                                         'cum_ret_quotient': cumret_q}}).T)

results_df

results_df['sharpe_delta'] = results_df['sharpe'] - results_df['sharpe_bench']
results_df

delta = results_df['sharpe_delta']
pd.DataFrame(delta.describe())

delta.hist(bins=20)

results_df['sharpe_delta'].plot()

"""
Explanation: Let's test training with SPY n_train times
End of explanation
"""


