import pandas as pd
import numpy as np
from quantopian.pipeline import Pipeline, CustomFactor

from quantopian.pipeline.factors import Latest
from quantopian.pipeline.filters import Q500US
from quantopian.research import run_pipeline
from time import strftime
from datetime import timedelta, datetime
from quantopian.pipeline.data.builtin import USEquityPricing
import seaborn as sns
import alphalens
import matplotlib.pyplot as plt
import talib as ta
from quantopian.pipeline.classifiers.morningstar import Sector
from quantopian.pipeline.data.user_5aba7037bad61a0013e6e687 import test
import math
from quantopian.pipeline.data import Fundamentals

MORNINGSTAR_SECTOR_CODES = {
     -1: 'Misc',
    101: 'Basic Materials',
    102: 'Consumer Cyclical',
    103: 'Financial Services',
    104: 'Real Estate',
    205: 'Consumer Defensive',
    206: 'Healthcare',
    207: 'Utilities',
    308: 'Communication Services',
    309: 'Energy',
    310: 'Industrials',
    311: 'Technology' ,    
}

def auto_heatmap(factor1, factor2, factor_1_name, factor_2_name, quantile_groups = 4, start_date = '2018-11-06', end_date = '2018-11-08'):
    class PreviousClose(CustomFactor):  
        inputs = [USEquityPricing.close]  
        window_length = 1

        def compute(self, today, assets, out, close):  
            out[:] = close[0]

    class PreviousOpen(CustomFactor):  
        inputs = [USEquityPricing.open]  
        window_length = 1

        def compute(self, today, assets, out, open):  
            out[:] = open[0]

    def make_pipeline():
        # define our fundamental factor pipeline
        pipe = Pipeline()
        
        universe = Q500US()
        
        factor_1 = Latest([factor1])
        factor_2 = Latest([factor2])

        factor_1_dec = factor_1.quantiles(quantile_groups, mask=universe)
        factor_2_dec = factor_2.quantiles(quantile_groups, mask=universe)


    #     PreviousClose(window_length = 1) is the same as USEquityPricing.close.latest
        previous_close_price = PreviousClose(window_length = 2)
        close_price = USEquityPricing.close.latest
        open_price = USEquityPricing.open.latest

        Overnight_return = open_price / previous_close_price - 1
        Intraday_return = close_price / open_price - 1

        pipe = Pipeline(
            columns = {
                'open':open_price,
                'close':close_price,
                'close_previous':previous_close_price,
                'Overnight_return':Overnight_return,
                'Intraday_return':Intraday_return,
                factor_1_name:factor_1,
                factor_2_name:factor_2,
                factor_1_name+" quantile":factor_1_dec,
                factor_2_name+" quantile":factor_2_dec,
                'Sector': Sector(mask=universe),
                'first_30': test.first_30.latest,
                'last_30': test.last_30.latest
            },
            screen=universe
        )

        return pipe

    pipe = make_pipeline()

    results = run_pipeline(pipe, start_date, end_date).dropna()
    
    decile_means = results.groupby([factor_1_name+" quantile", factor_2_name+" quantile"])['Intraday_return', 'Overnight_return', 'first_30', 'last_30'].mean()
    decile_count = results.groupby([factor_1_name+" quantile", factor_2_name+" quantile"])['Overnight_return'].count()
    
    for ret_type in ['Intraday_return', 'Overnight_return', 'first_30', 'last_30']:
        sns.heatmap(decile_means[ret_type].unstack(), annot=True, linewidths=.5, fmt=".2%", center=0, cmap=sns.diverging_palette(20, 150, as_cmap=True), vmin=-0.005, vmax=0.005,).set_title(ret_type)
        plt.figure()
    
    sns.heatmap(decile_count.unstack(), annot=True, linewidths=.5, fmt='g', cmap="Blues").set_title('Number of Observations')
    plt.figure()
    
    decile_means_sector = results.groupby([factor_1_name+" quantile", factor_2_name+" quantile", "Sector"])['Intraday_return', 'Overnight_return', 'first_30', 'last_30'].mean()
    
    
    for ret_type in ['Intraday_return', 'Overnight_return', 'first_30', 'last_30']:
        def facet_heatmap(data, color, value_col, **kws):
            data = data.pivot(index=factor_1_name+" quantile", columns=factor_2_name+" quantile", values = value_col)
            sns.heatmap(data, annot=True, linewidths=.5, fmt=".2%", center=0, cmap=sns.diverging_palette(20, 150, as_cmap=True), vmin=-0.005, vmax=0.005, cbar = False, **kws)

        with sns.plotting_context(font_scale=5):
            g = sns.FacetGrid(decile_means_sector.reset_index(), col="Sector", col_wrap=3, size=4.5, aspect=1)

        g = g.map_dataframe(facet_heatmap, value_col = ret_type)

        g.set_titles(col_template="{col_name}", fontweight='bold', fontsize=18)
        g.fig.suptitle(ret_type + ' by sector')
        plt.figure()
        
    decile_count_sector = results.groupby([factor_1_name+" quantile", factor_2_name+" quantile", "Sector"])['Overnight_return'].count()
    
    def facet_heatmap_count(data, color, value_col, **kws):
            data = data.pivot(index=factor_1_name+" quantile", columns=factor_2_name+" quantile", values = value_col)
            sns.heatmap(data, annot=True, linewidths=.5, fmt='g', cmap="Blues", **kws)


    with sns.plotting_context(font_scale=5):
        g = sns.FacetGrid(decile_count_sector.reset_index(), col="Sector", col_wrap=3, size=4.5, aspect=1)

        g = g.map_dataframe(facet_heatmap_count, value_col = "Overnight_return")

        g.set_titles(col_template="{col_name}", fontweight='bold', fontsize=18)
        g.fig.suptitle('Number of observations by sector')
        plt.figure()
    
#     results['cum_intraday']=results.groupby(level = 1)['Intraday_return'].apply(lambda x : x.shift().fillna(0).add(1).cumprod()).values
#     results['cum_overnight']=results.groupby(level = 1)['Overnight_return'].apply(lambda x : x.shift().fillna(0).add(1).cumprod()).values
        
    return results

    

class Volatility(CustomFactor):  
    inputs = [USEquityPricing.close]  
    window_length = 20  
    def compute(self, today, assets, out, close):  
        # [0:-1] is needed to remove last close since diff is one element shorter  
        daily_returns = np.diff(close, axis = 0) / close[0:-1]  
        out[:] = daily_returns.std(axis = 0) * math.sqrt(252)

a