import numpy as np
import pandas as pd
import statsmodels.api as sm
from Selection.QC500UniverseSelectionModel import QC500UniverseSelectionModel
from QuantConnect.Python import PythonQuandl

class Oilsensibiltiy(QCAlgorithm):

    def Initialize(self):
        
        self.SetStartDate( 2012 , 1, 1)       # Set Start Date
        self.SetEndDate( 2020 , 10, 10)
        self.SetCash(100000)                # Set Strategy Cash

        self.lookback = 61                 # Length(days) of historical data
        self.weights_long,self.weights_short = pd.DataFrame(),pd.DataFrame()      # Pandas data frame (index: symbol) that stores the weight
        self.Portfolio.MarginModel = PatternDayTradingMarginModel()
        
        self.AGG = self.AddEquity("AGG", Resolution.Daily).Symbol
        self.nextLiquidate = self.Time   # Initialize last trade time
        self.rebalance_days = 30
        
        self.UniverseSettings.Resolution = Resolution.Daily   # Use hour resolution for speed
        self.oil = self.AddData(QuandlOil, 'FRED/DCOILBRENTEU', Resolution.Daily).Symbol
        self.AddUniverse(self.CoarseSelection, self.SelectFine)
        self.selectedequity = 1000
        self.numberOfSymbolsFine = 25
        self.Symbols_long = []
        self.Symbols_short = []
        self.zscore_keep_buy = []
        self.zscore_keep_short = []

    def CoarseSelection(self, coarse):
        
        if self.Time < self.nextLiquidate:
            return Universe.Unchanged
        
        selected = sorted([x for x in coarse if x.HasFundamentalData and x.Price > 5],
                          key=lambda x: x.DollarVolume, reverse=True)

        symbols = [x.Symbol for x in selected[:self.selectedequity ] ]
        return symbols
        
    def SelectFine(self, fine):
        
        filtered = [x.Symbol for x in fine if  x.AssetClassification.MorningstarSectorCode == 309]
        self.Symbols_long = filtered[:self.numberOfSymbolsFine]
        self.Symbols_short = self.Symbols_long
        return self.Symbols_long + self.Symbols_short

    def ZscoreGrade(self,sample, factors) :
        
        factors = sm.add_constant(factors)
        
        self.Log(f"Samples:\n{sample.index}")
        self.Log(f"Factors:\n{factors.index}")
        self.Quit()
        return
        
        # Train Ordinary Least Squares linear model for each stock
        OLSmodels = {ticker: sm.OLS(sample[ticker], factors).fit() for ticker in sample.columns}

    def GetWeights(self, history , crudeoil_history,Long):
        crudeoil_history = crudeoil_history.pct_change().dropna()
        sample = history.unstack(level=0).pct_change().dropna()
        zscore = self.ZscoreGrade(sample,crudeoil_history)

    def changeDataFrame(self,history_long):
        crudeoil_history = self.History(QuandlOil,self.oil , 300, Resolution.Daily).droplevel(level=0)
        
        idxb = history_long.unstack().columns[0]
        idxe = history_long.unstack().columns[-1]
        crudeoil_history = crudeoil_history[~crudeoil_history.index.duplicated(keep='last')].loc[idxb:idxe]
        return crudeoil_history
        
    def OnData(self, data):
        history_long = self.History(self.Symbols_long, self.lookback, Resolution.Daily).close
        crudeoil_history = self.changeDataFrame(history_long)
        
        self.GetWeights(history_long,crudeoil_history,Long=True)
        
    def OnSecuritiesChanged(self, changes):
        for security in changes.RemovedSecurities:
            if security.Invested:
                self.Liquidate(security.Symbol, 'Removed from Universe')
        
        
class QuandlOil(PythonQuandl):
    def __init__(self):
        self.ValueColumnName = 'Value'