# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


'''
Naked Momentum -- V1.0.3
    8.23.20 -- Added Qty (Estimation), not sure if correct + tested Bull Call logic (+)
    8.24.20 -- Added Insights, + RiskManagement, etc.

Potential Improvements -- 
    Filter for TRADEABLE option symbols IN UNIVERSE ? (Full universe is filled, then QTY works)
        #Or just an IsTradeable list comp prior to entry loop?
    Find LEAST CORRELATED symbols within top x Momentum + Mkt Cap 
        https://www.quantconnect.com/forum/discussion/6780/from-research-to-production-uncorrelated-assets/p1
        DONE
        
    Add a dynamic hedge w realized vol > 12.5% -- Deep OTM Put ?
    
        
'''

from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel
from Risk.MaximumDrawdownPercentPortfolio import MaximumDrawdownPercentPortfolio
from Risk.TrailingStopRiskManagementModel import TrailingStopRiskManagementModel

from datetime import timedelta
from GetUncorrelatedAssets import GetUncorrelatedAssets



class CoveredCallAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 1, 1)
        self.SetEndDate(2019, 1, 31)
        self.SetCash(100000)
        
        '''Dynamic Universe'''
        self.AddUniverse(self.SelectCoarse,self.SelectFine)
        self.UniverseSettings.Resolution = Resolution.Daily
        
        self.SetExecution(ImmediateExecutionModel())

        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        #self.SetRiskManagement(MaximumDrawdownPercentPortfolio(maximumDrawdownPercent = .05, isTrailing=True))
        #self.SetRiskManagement(TrailingStopRiskManagementModel(.05)
        
        # ---------- Universe Params ----------- #
        

        
        self.mkt_cap_sort = True                                                #Switched On
        
        if self.mkt_cap_sort:
            self.mom_x = 100
            self.top_x = 10
        else:
            self.mom_x = 10
        
        # -- Momentum Params
        self.momentum_type = 0                                                  #0 = OFF, 1 = year, 2 = month
        
        self.momBySym = {}
        self.momValues = None
        
        # -- End Momentum 
        
        #Uncorrelated Pairs Switch
        self.uncorr = False

        
        #Option Selection + Execution Params
        self.symbol_list = []
        self.min_dte = 0 #0            #DONT think these are actually plugged in ?
        self.max_dte = 5 #5
        
        self.hold_days = 5 if self.min_dte <= 1 else self.min_dte
        
        self.max_positions = 10
        
        self.spread_type = 1 #0 #1                                                 #0 = Long Call, 1 = Bull Call
        
        
        #self.SetWarmUp()
        

        
    def SelectCoarse(self, coarse):
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)[:200]
        
        filtCoarse = [c for c in sortedByDollarVolume if c.Price > 10]
        
        if self.momentum_type == 0:
            self.symbol_list = [f.Symbol for f in filtCoarse]
            return self.symbol_list
        
        # ---------- Begin Momentum ---------- #
        selected = []
        for c in filtCoarse:
            pass
            symbol = c.Symbol
            if symbol not in self.momBySym:
                hist = self.History(symbol, 253, Resolution.Daily)
                self.momBySym[symbol] = Momentum(symbol, self, hist)            #pass in symbol, ALGO instance, and hist
            self.momBySym[symbol].Update(c.AdjustedPrice)
            
        if self.momentum_type == 1:
            sorted_by_momentum = {key:value for key, value in sorted(self.momentumBySymbol.items(),\
                    key=lambda kv: kv[1].mom_yr, reverse=True)[:self.mom_x]}            
        else:
            sorted_by_momentum = {key:value for key, value in sorted(self.momentumBySymbol.items(),\
                    key=lambda kv: kv[1].mom_mo, reverse=True)[:self.mom_x]}
                    
        self.momValues = sorted_by_momentum
        
        selected = list(sorted_by_momentum.keys())
        self.symbol_list = selected
        # ---------- End Momentum ---------- #

        return self.symbol_list
        
    def SelectFine(self, fine):
        #Mkt Cap Filter
        if self.mkt_cap_sort:
            #filteredByMktCap = [x for x in fine if 1e10 < x.MarketCap] #< 1e9]
            sortedByMktCap = sorted(fine, key = lambda f: f.MarketCap, reverse=True)[:self.top_x]     #reverse = Descending
            self.symbol_list = [f.Symbol for f in sortedByMktCap]
        else:
            self.symbol_list = [f.Symbol for f in fine]
        return self.symbol_list
        
        

        
    def OnSecuritiesChanged (self, changes):
        symbols = [x.Symbol for x in changes.AddedSecurities]
        #Returns <GOOG SYMBOLID>
        #init = symbols
        
        if self.uncorr and len(symbols) > 0:                                    #Should eb only check needed
            top = len(symbols) #self.top_x #int(self.top_x / 2) if self.mkt_cap_sort else self.top_x
            
            history = self.History(symbols, 150, Resolution.Hour)
            if history.shape[1] > 1:                                                #NEED better way to do this...
                hist = history.unstack(level = 1).close.transpose().pct_change().dropna()
                #WHY does this get out of index error ^^ 
                
                symbols_rank = GetUncorrelatedAssets(hist, top)
                s2 = [symbol for symbol, corr_rank in symbols_rank]            #Why not working right?
                
                #for s, s2 in zip(symbols, symbols_new):
                #    self.Debug(f'{s} - {s2}')                                  #Identical? WHY 168 not working?
                #s3 = [s for s in init if s in s2]
                s4 = [x.Symbol for x in changes.AddedSecurities if x.Symbol in symbols_rank]
                symbols = [s for s in symbols if s in s2]

        for x in changes.AddedSecurities:
            if x.Symbol not in symbols: continue                                #Matches w Uncorr or Regular -- DOES NOT WORK

            if x.Symbol.SecurityType != SecurityType.Equity: continue
            option = self.AddOption(x.Symbol.Value, Resolution.Minute)
            option.SetFilter(-1, +1, timedelta(self.min_dte), timedelta(self.max_dte))
            '''IF buying CALLS OUTRIGHT -- need more time -- 15 - 30 probably'''
    
        for x in changes.RemovedSecurities:
            if x.Symbol.SecurityType != SecurityType.Equity: continue
            # Loop through the securities
            # If it is an option and the underlying matches the removed security, remove it
            for symbol in self.Securities.Keys:
                if symbol.SecurityType == SecurityType.Option and symbol.Underlying == x.Symbol:
                    self.RemoveSecurity(symbol)
     
        
    def OnData(self,slice):
        #if not self.Portfolio["AAPL"].Invested:
        #    self.MarketOrder("AAPL",100)     # buy 100 shares of underlying stocks
        #    self.Log(str(self.Time) + " bought SPY " + "@" + str(self.Securities["SPY"].Price) 
        #            + " Cash balance: " + str(self.Portfolio.Cash)
        #            + " Equity: " + str(self.Portfolio.HoldStock))
        if len(self.symbol_list) == 0:
            return
        
        option_invested = [x.Key for x in self.Portfolio if x.Value.Invested and x.Value.Type==SecurityType.Option]
        if len(option_invested) < self.max_positions:
            self.TradeOptions(slice) 
 
    def TradeOptions(self,slice):
        #Filter for only ones NOT invested?
        invested = [x.Key for x in self.Portfolio if x.Value.Invested] #and x.Value.Type == SecurityType.Option]
        margin_remaining = self.Portfolio.MarginRemaining
        margin_per_position = margin_remaining * .99  * ( 1 / len(self.symbol_list) )   #Margin available for EACH symbol / option position
        insights = []
        for i in slice.OptionChains:
            #if i.Key in invested: continue                             #Addit -- to skip open symbols
            
            #if i.Key != self.symbol: continue
            #if i.Key not in self.FixedUniv: continue                           #FixedUniverse
            
            #THIS is really CONTRACTS
            chain = i.Value
            # filter the call options contracts -- CALLS
            call = [x for x in chain if x.Right == OptionRight.Call] 
            # sorted the contracts according to their expiration dates and choose the ATM options
            contracts = sorted(sorted(call, key = lambda x: abs(chain.Underlying.Price - x.Strike)), 
                                            key = lambda x: x.Expiry, reverse=True)
                                            
            #NEED to CHECK remaining margin available! -- SET quantity!
            #quantity = self.CalculateOrderQuantity(chain.LastPrice, 0.1) #self.LastPrice
            #Takes SHARES available in 10% of
            
            if len(contracts) != 0:     
                #opt_price = contracts[0].LastPrice * 100
                opt_price = contracts[0].TheoreticalPrice * 100                 #Per contract
                if opt_price != 0:
                    qty = margin_per_position / opt_price
                    
                #qty = margin_per_position / contracts[0].Underlying.Price      #Approximation -- ATM, low DTE calls will have limited extrinsic
                #qty = margin_per_position / contracts[0].LastPrice * 100        #take margin per position, divide by cost of option contract
                self.Debug(f'option_price -- {opt_price} -- qty: {qty} -- basis: {qty * opt_price}')
                self.long_call = contracts[0].Symbol
                
                #Need a check if tradeable?
                #if self.long_call.IsTradeable:
                
                #self.MarketOrder(self.long_call, 1)  
                insights += [Insight.Price(self.long_call, timedelta(days=5),InsightDirection.Up)]

                
                self.Debug(f'LE -- {self.Time.date()}: {self.long_call}')         #Looks right ? 
                
                #If turned on, enter BULL CALL leg of trade.
                if self.spread_type == 1 and len(contracts)  > 1:
                    self.short_call = contracts[1].Symbol
                    #self.MarketOrder(self.short_call, -1)
                    insights += [Insight.Price(self.short_call, timedelta(days=5), InsightDirection.Down)]
            
        self.EmitInsights(insights)
                    
                    
            
            

    
    def OnOrderEvent(self, orderEvent):
        self.Log(str(orderEvent))
        self.Log("Cash balance: " + str(self.Portfolio.Cash))
        
        
class Momentum():
    
    def __init__(self,symbol, algorithm, history):
        self.algo = algorithm
        self.sym = symbol
        self.window = RollingWindow[float](252)
        self.dailyBars = RollingWindow[TradeBar](10)
        self.mom_yr = 0
        self.mom_mo = 0
        
        #Consolidator ...
        self.dailyCons = TradeBarConsolidator(timedelta(days=1))
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.dailyCons)
        
        self.dailyCons.DataConsolidated += self.onDailyBar                      #Each daily bar, call handler
        
        for bar in history.itertuples():
            tb = TradeBar(bar.Index[1], bar.open, bar.high, bar.low, bar.close, bar.volume)
            self.Update(bar.close)
            self.dailyBars.Add(tb)
            
            
    def Update(self, price):
        self.window.Add(price)
        if self.window.IsReady:
            self.mom_yr = (self.window[0] - self.window[251]) / self.window[251] #Was [-252] and [-25]
            self.mom_mo = (self.window[0] - self.window[25]) / self.window[25]
    
    #Event handler for NEW DAILY BAR
    def onDailyBar(self, sender, bar):
        self.dailyBars.Add(bar)
        
    @property
    def IsReady(self):
        return self.window.IsReady #and self.dailyBars.IsReady
        
    @property
    def pivlo(self):
        for i in range(1,8):
            if self.dailyBars[i].Close > self.dailyBars[i + 1].Close:
                return False
        if self.dailyBars[0].Close < self.dailyBars[1].Close: 
            return False
        return True
        
        
    @property
    def pivhi(self):
        for i in range(1,8):
            if self.dailyBars[i].Close < self.dailyBars[i + 1].Close:
                return False
        if self.dailyBars[0].Close > self.dailyBars[1].Close:
            return False
        return True
        
    #def onIndicUpdate(self, sender, updated):
    #    if self.BBD.IsReady:
    #        self.uppers.Add(self.BBD.UpperBand.Current.Value)
    #        self.lowers.Add(self.BBD.LowerBand.Current.Value)
    
                
    #Manual way to do this... better to use event handler
    #def WindowUpdate(self, bar):
    #    if self.dailyBars.IsReady:
    #       self.dailyBars.Add(bar)

import numpy as np
import pandas as pd

def GetUncorrelatedAssets(returns, num_assets):
    '''
    Passed in HIST dataframe -- transformed slightly + unstacked
    history = qb.History(symbols, 150, Resolution.Hour)
        
    # Get hourly returns
    returns = history.unstack(level = 1).close.transpose().pct_change().dropna()
    https://www.quantconnect.com/forum/discussion/6780/from-research-to-production-uncorrelated-assets/p1

    '''
    # Get correlation
    correlation = returns.corr()
    
    # Find assets with lowest mean correlation, scaled by STD
    selected = []
    for index, row in correlation.iteritems():
        corr_rank = row.abs().mean()/row.abs().std()
        selected.append((index, corr_rank))

    # Sort and take the top num_assets
    selected = sorted(selected, key = lambda x: x[1])[:num_assets]
    
    return selected
    
    
'''

    #In self.initialize 
        
        self.FixedUniv = ['AAPL','AMD','TSLA','BABA','ROKU']                    #Filter by IV Eventually?
        
        #equity = self.AddEquity("AAPL", Resolution.Minute)
        equities = [self.AddEquity(sym, Resolution.Minute) for sym in self.FixedUniv]
        option = [self.AddOption(symbol) for symbol in self.FixedUniv]
        #option = self.AddOption("AAPL", Resolution.Minute)                     #Original
        self.symbols = [option.Symbol for option in option]
        
        # set strike/expiry filter for this option chain
        for opt in option:
            #opt.SetFilter(-3, +3, timedelta(0), timedelta(30))
            # use the underlying equity as the benchmark
        
            #ONLY for naked calls -- maybe bull calls?
            opt.SetFilter(-1,+1, timedelta(0), timedelta(5))
            #self.SetBenchmark()
        
-------------------------------------------------------------------------------
        
        #Uncorr Assets Filter in FINE
        #if self.uncorr:

        #    top = int(self.top_x / 2) if self.mkt_cap_sort else self.top_x
            
        #    history = self.History(self.symbol_list, 150, Resolution.Hour)
        #    history.unstack(level = 1).close.transpose().pct_change().dropna()
            
        #    symbols_rank = GetUncorrelatedAssets(history, top)
        #    symbols = [symbol for symbol, corr_rank in symbols_rank]

        #    self.symbol_list = [s for s in self.symbol_list if s in symbols]
            #self.Debug(f'Post Uncorr - {self.symbol_list}')

'''