from datetime import timedelta
import numpy as np
import pandas as pd

class ModelA(AlphaModel): 
    
    def __init__(self, resolution, insightsTimeDelta ):

        self.symbolDataBySymbol =   {}
        self.modelResolution    =   resolution
        self.insightsTimeDelta  =   insightsTimeDelta
        
    def OnSecuritiesChanged(self, algorithm, changes):
            for added in changes.AddedSecurities:
                symbolData = self.symbolDataBySymbol.get(added.Symbol)
                if symbolData is None:
                    symbolData = SymbolData(added.Symbol, algorithm, self.modelResolution)
                    self.symbolDataBySymbol[added.Symbol] = symbolData
 
    def Update(self, algorithm, data):
        
        insights=[]
        
        for symbol, symbolData in self.symbolDataBySymbol.items():

            symbolData.getInsight(algorithm.Securities[symbol].Price) # Latest known price; we are at 12:00 and the last trade at 10.57 
            
            algorithm.Log(f"{symbol}\tMOM\t[{symbolData.fmom}]\t{round(symbolData.mom.Current.Value,2)}\tKAMA\t[{symbolData.fkama}]\t{round(symbolData.kama.Current.Value,2)}\
                            \tPrice\t{symbolData.price}\tROC\t[{symbolData.froc}]\t{round(symbolData.roc.Current.Value,4)}\tEMA\t[{symbolData.fema}]\tEMA-13\t{round(symbolData.ema13.Current.Value,2)}\
                            \tEMA-63\t{round(symbolData.ema63.Current.Value,2)}\tEMA-150\t{round(symbolData.ema150.Current.Value,2)}\taction\t{symbolData.InsightDirection}")
                            
            insights.append(Insight(symbol, self.insightsTimeDelta, InsightType.Price, symbolData.InsightDirection, 0.0025,None, "ModelA",None))
        
        return insights

class FrameworkAlgorithm(QCAlgorithm):
    
    def Initialize(self):

        tickers     =   ["MSFT","MRNA","MELI","FSLY"]
        symbols     =   [Symbol.Create(x, SecurityType.Equity, Market.USA) for x in tickers]
        resolution  =   Resolution.Hour  #10-11, etc Daily data is midnight to mifnight, 12AM EST 
        warmup      =   28
        insightsTimeDelta=timedelta(days=1)
        
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.Every(TimeSpan.FromMinutes(60)), self.LogPortfolioDetail)

        self.SetStartDate(2020, 11, 29)   
        self.SetCash(100000)           
        self.SetBenchmark("SPY")
        self.UniverseSettings.Resolution = resolution
        self.SetWarmUp(timedelta(warmup)) 
        self.SetUniverseSelection(ManualUniverseSelectionModel(symbols))
        self.SetAlpha(ModelA(resolution,insightsTimeDelta))
        #self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetPortfolioConstruction(MeanVarianceOptimizationPortfolioConstructionModel(resolution,PortfolioBias.LongShort,1,63,resolution,0.02,MaximumSharpeRatioPortfolioOptimizer(0,1,0)))
        self.SetRiskManagement(MaximumDrawdownPercentPerSecurity(0.02)) # drop in profit from the max 
        self.SetExecution(ImmediateExecutionModel())

    def LogPortfolioDetail(self):
        
        dt=int(str(self.Time)[11:13]) # can be improved
        if dt>9 and dt<18:
            
            if self.Portfolio.Invested:
                self.Log("\n\nPortfolio at : {0}".format(self.Time))
                summary = {}

                invested = [ x.Symbol.Value for x in self.Portfolio.Values if x.Invested ]
                for symbol in invested:
                
                    hold_val    = round(self.Portfolio[symbol].HoldingsValue, 2)
                    abs_val     = round(self.Portfolio[symbol].AbsoluteHoldingsValue, 2)
                    pnl         = round(self.Portfolio[symbol].UnrealizedProfit, 2)
                    qty         = self.Portfolio[symbol].Quantity
                    price       = self.Portfolio[symbol].Price
                    
                    summary[symbol]=[hold_val,abs_val,pnl,qty,price]
                df=pd.DataFrame(summary)
                df.index = ['hold_val', 'abs_val', 'pnl', 'qty','price']
                df=df.T
                hold_val_total= df['hold_val'].sum()
                df = df.assign(weight=df['hold_val']/hold_val_total)
                self.Log(df)

            
class SymbolData:
    
    def __init__(self, symbol, algorithm, resolution):
        
        self.symbol             = symbol
        self.price              = 0.00
        self.InsightDirection   = InsightDirection.Flat
        self.kama               = algorithm.KAMA(symbol, 10,2,30, resolution)
        self.mom                = algorithm.MOM(symbol, 14, resolution)
        self.roc                = algorithm.ROC(symbol, 9, resolution) 
        self.ema13              = algorithm.EMA(symbol, 13, resolution)
        self.ema63              = algorithm.EMA(symbol, 63, resolution)
        self.ema150             = algorithm.EMA(symbol, 150, resolution)
        self.fkama              = False
        self.fmom               = False
        self.froc               = False
        self.fema               = False

    def getInsight(self, price):    
        
        self.price          = price
        self.fkama          = self.price>self.kama.Current.Value 
        self.fmom           = self.mom.Current.Value>0
        self.froc           = self.roc.Current.Value>0
        self.fema           = self.ema13.Current.Value>self.ema63.Current.Value>self.ema150.Current.Value
        
        if self.fmom and self.fkama and self.fema and self.froc:
            self.InsightDirection= InsightDirection.Up

        if not self.fmom or not self.fkama or not self.fema or not self.froc:
            self.InsightDirection = InsightDirection.Down
        
        self.InsightDirection   = InsightDirection.Flat