import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
%matplotlib inline
from ch_api import ch_optimize, util, ch_portfolio

import importlib

USE_DB = True  # It DB connection is limited, let's use pickle files inluced in the repo not connecting DB
ID = ''
PW = ''

import yaml
from sqlalchemy import create_engine


engine = create_engine('mysql+pymysql://%s:%s@localhost/securities_master' % 
                       (ID, PW), 
                       echo=False)

if USE_DB:
    factors_ts = pd.read_sql("select * from symbol where instrument='factor' and ticker<>'RF'", engine)
    # factors data are available in the pickle file: %root%/input/factors.pkl
    #factors_ts.to_pickle("../data/factors.pkl")
else:
    factors_ts = pd.read_pickle('../data/factors.pkl')
    
tis = list(factors_ts['ticker'].values)

# Factors
tis

start_date = '2012-01-02'
end_date = '2018-08-24'

if USE_DB:
    stm = "select * from daily_price where ticker in %s and instrument_type='factor'" % str(tuple(tis))
    factors_ts = pd.read_sql(stm, engine)

    factors_ts = factors_ts[['price_date', 'ticker', 'adj_close_price']]
    idx = (factors_ts['price_date']>=start_date) & (factors_ts['price_date']<=end_date)

    df_sliced = factors_ts[idx].copy()
    df_sliced.set_index(['price_date', 'ticker'],inplace=True)
    factors_unstacked = df_sliced.unstack(level=-1)
    factors_unstacked_aligned = factors_unstacked.dropna()
    #factors_unstacked_aligned.to_pickle("../data/factors_return.py")
else:
    factors_unstacked_aligned = pd.read_pickle("../data/factors_return.py")

factors_unstacked_aligned.head()

factors_unstacked_aligned = factors_unstacked_aligned['adj_close_price'].copy()

factors_unstacked_aligned_cum = np.cumprod(factors_unstacked_aligned + 1.0) - 1.0
factors_unstacked_aligned_cum.to_csv("../output/factors_cum_returns.csv")

sns.heatmap(factors_unstacked_aligned.corr(), annot=True)

ew_portfolio_ts = pd.read_pickle('../output/ts_base_portfolio.pkl')

ew_portfolio_ts.rename(columns={'metric_date': 'date'}, inplace=True)
ew_portfolio_ts.set_index(['date', 'ticker'], inplace=True)

ew_portfolio_ts.head()

ew_portfolio_ts.tail()

# portfolio returns
port_ret = ew_portfolio_ts.groupby(level=0).apply(lambda x: (x['weight']*x['daily_return']).sum())
port_ret_df = pd.DataFrame(port_ret)
port_ret_df.columns = ['EW']
port_ret_df.head()

ew_portfolio_ts_reset = ew_portfolio_ts.reset_index()

exposure_ts_df = ew_portfolio_ts_reset.groupby('date').apply(lambda x: ch_portfolio.computing_exposures(x))

exposure_ts_df.head()

importlib.reload(ch_optimize)

# optimize each portfolio with daily rebalancing
# 1. portfolio return df : port_ret_df
# 2. portfolio exposure df: exposure_ts_df
from cfg import mv, value, value_mom_mix
lst_portfolio = [mv, value, value_mom_mix]
for pf in lst_portfolio:
    # 1. optimize ts data
    opt = ch_optimize.Optimizer().createOptimizer(pf, 
                                optimizer_type='MV', engine=engine, 
            lst_factor=[ u'CMA', u'HML', u'MOM', 'RMW',u'SMB',])
    optimized_df = opt.run_optimize(ew_portfolio_ts)
    optimized_df = optimized_df.rename(columns={
        'weight': 'old_weight', 
        'optimized_weight': 'weight'})
    optimized_df['pf_name'] = pf.pf_name
    
    # 2. shift optimized_weight
    opt_weights_sr = optimized_df.unstack()['weight'].shift().stack()
    opt_weights_df = pd.DataFrame(opt_weights_sr)
    opt_weights_df.columns = ['optimized_weight_shift']
    final_portfolio = optimized_df.merge(opt_weights_df, 
                                           left_index=True, 
                                           right_index=True, 
                                           how='left')
    # 3. performance computation
    final_portfolio_perf = final_portfolio.groupby(level=0).apply(
        lambda x: (x['daily_return'] * x['optimized_weight_shift']).sum())
    final_portfolio_perf_df = pd.DataFrame(final_portfolio_perf)
    final_portfolio_perf_df.columns = [pf.pf_name]
    port_ret_df= pd.concat([port_ret_df, final_portfolio_perf_df], axis=1)
    # 4. Exposure computation
    final_portfolio_exp_ready = final_portfolio.rename(columns={'weight': 'optimized_weight', 
                                                      'optimized_weight_shift': 'weight'})
    
    final_portfolio_exp_ready_reset = final_portfolio_exp_ready.reset_index()
    exposure_ts_tmp_df = final_portfolio_exp_ready_reset.groupby('date').apply(
                 lambda x: ch_portfolio.computing_exposures(x))
        
    
    exposure_ts_df = pd.concat([exposure_ts_df, exposure_ts_tmp_df])
        

exposure_ts_df['portfolio_name'].unique()

exposure_ts_df.to_csv('../output/exposure_ts_df_total.csv')

port_ret_df.head()

vol_df = pd.DataFrame(port_ret_df.std() * np.sqrt(250))
vol_df.columns = ['Vol']
vol_df

final_return_df = pd.DataFrame(port_cum_ret_df.iloc[-1])
final_return_df.columns = ['Return']
final_return_df

summary_df = pd.concat([final_return_df, vol_df], axis=1)
summary_df

summary_df['Risk Adj Ret'] = summary_df['Return']/summary_df['Vol']
summary_df

port_cum_ret_df = np.cumprod(port_ret_df + 1.0) - 1.0
port_cum_ret_df.to_csv('../output/port_cum_ret_df_total.csv')