from typing import Dict
from datetime import timedelta
from AlgorithmImports import *

class Centurion(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 2, 20)  # Set Start Date
        self.SetEndDate(2021, 3, 25)
        # Set Start Date
        
        self.SetCash(100000)  # Set Strategy Cash
        # self.AddEquity("SPY", Resolution.Minute)
        '''
        Would it be easier to have a symboldata class?
        What would it contain?
            - initialization of Data
            - dictionaries for stoploss, option types, rebalance weights 
        '''
        
        # Assets=====================================
        self.bndx = self.AddEquity("BNDX", Resolution.Minute).Symbol
        
        self.bnd = self.AddEquity("BND", Resolution.Minute).Symbol
        
        spy = self.AddEquity("SPY", Resolution.Minute)
        spy.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.spy = spy.Symbol
        self.spycontract: Symbol = None

        tqqq = self.AddEquity("TQQQ", Resolution.Minute)
        tqqq.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.tqqq = tqqq.Symbol
        self.tqqqcontract: Symbol = None

        soxl = self.AddEquity("SOXL", Resolution.Minute)
        soxl.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.soxl = soxl.Symbol
        self.soxlcontract: Symbol = None

        
        tecl = self.AddEquity("TECL", Resolution.Minute)
        tecl.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.tecl = tecl.Symbol
        self.teclcontract: Symbol = None

        
        #Scheduled Events============================================
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.AfterMarketOpen("SPY", 1), self.monthlyRebalance)
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.AfterMarketOpen("SPY", 1), self.captureOpening)
        #Variables needed for stoploss
        self.stoplosshold = 0
        
        self.dailyOpen: Dict[Symbol,float] = {
            self.spy: 0,
            self.bndx:0,
            self.bnd:0
        }
        
        # Rebalancing Weights
        self.weights = {
            self.spy: .2,
            self.bndx: .3,
            self.bnd: .4
        }

        # Underlying & Options (UNOs)
        self.uNos = {
            self.spy : self.spycontract,
            self.tqqq : self.tqqqcontract,
            self.tecl : self.teclcontract
            #self.soxl : self.soxlcontract
        }

    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''
        
        if self.IsWarmingUp:
            return
        
        if self.Portfolio.Invested:
            self.stoploss(data)
        
        if not self.Portfolio.Invested:
            self.monthlyRebalance()
            
        
        self.setPuts()

        self.timeDecayCheck()
        
        self.buyPuts()
        
        self.exerciseLoki()    



    def setPuts(self):
        '''
        I want to purchase a long puts for every key in dictionary uNos ((u)underlying (N)and (os)Options)
        '''
        
        for underlying in self.uNos.keys():
            targetStrike = (self.Securities[underlying].Price * 0.60) - (self.Securities[underlying].Price * 0.60)%5
            contracts = self.OptionChainProvider.GetOptionContractList(underlying, self.Time)
            puts = [x for x in contracts if x.ID.OptionRight == OptionRight.Put]
            puts = sorted( sorted(puts, key = lambda x: x.ID.Date, reverse = True),
                       key = lambda x: x.ID.StrikePrice)
            puts = [x for x in puts if x.ID.StrikePrice == targetStrike]
            puts = [x for x in puts if 270 < (x.ID.Date - self.Time).days <= 420]
            if len(puts) == 0:
                self.Log("No Puts")
                continue
            self.AddOptionContract(puts[0], Resolution.Minute)
            self.uNos[underlying] = puts[0]


    def timeDecayCheck(self):
        #Don't let an options time to expiration be less than 6 months
        
        for underlying, contract in self.uNos.items():
            if contract is None:
                continue
            if (contract.ID.Date - self.Time).days < 180:
                self.Liquidate(contract)
                self.RemoveSecurity(contract)
                self.uNos[underlying] = None
        
    
    def buyPuts(self):
        for _,contract in self.uNos.items():
            if contract is None:
                continue
            if not self.Portfolio[contract].Invested:    
                self.SetHoldings(contract,  0.02)
        
    

    def exerciseLoki(self):
        for underlying,contract in self.uNos.items():
            if contract is None:
                continue
            #Liquidate a contract if its increases to 30% OTM
            if self.Securities[underlying].Price < contract.ID.StrikePrice * 1:
                # self.ExerciseOption(contract, self.Portfolio[contract].Quantity)
                self.Liquidate(contract)
                self.RemoveSecurity(contract)
                self.uNos[underlying] = None
                self.Log(f'{contract} has been sold!')

    def captureOpening(self):
        #Grabs the daily opening price of spy for our stoploss method
        if self.IsWarmingUp:
            return
        for key, value in self.dailyOpen.items():
            if self.CurrentSlice.Bars.ContainsKey(key):
                self.dailyOpen[key] = self.CurrentSlice[key].Open
            self.stoplosshold = 0
            
          
            
    def monthlyRebalance(self):
        # Rebalance portfolio monthly 
        if self.IsWarmingUp:
            return
        for key,value in self.weights.items():
            self.SetHoldings(key,value)
            
        

    def stoploss(self, data):
        '''
        Stoploss logic:
            - If spy drops more than 5% liquidate entire equity portfolio
            - Change stoplosshold value to 1, this indicates that the portfolios SL has been hit 
            and were going to hold until the next trading day
        '''
        if self.IsWarmingUp:
            return
        
        
        for symbol, weight in self.weights.items():
            if self.Securities[symbol].Price == 0:
                continue
            open = self.dailyOpen[symbol]
            #See if any symbol has decreased more than 5% in a given day, liquidate if true and check next one... 
            if ((self.Securities[symbol].Price-open)/self.Securities[symbol].Price) < -.05:
                self.SetHoldings(symbol, 0)
                
                self.stoplosshold = 1
                #self.Log('HIT')
        
        self.buyPuts()
        
        self.exerciseLoki()