# region imports
from AlgorithmImports import *
import datetime as dt
import ast
import json
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel 
from io import StringIO
import pandas as pd

# endregion

class AdaptableVioletJaguar(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2022, 1, 1)  # Set Start Date
        self.SetCash(50000)  # Set Strategy Cash
        self.UniverseSettings.Resolution = Resolution.Minute
        
        self.tickers = ["AAPL"]
        
        self.AddUniverse(self.CoarseSelection, self.FineSelection)
        
        self.symbolDataBySymbol = {}
        self.fineFundamentals = {}

        self.SAVE_EARNINGS = False
        self.loaded_earnings = {}

        self.slice = None

        if not self.SAVE_EARNINGS:
            self.SetEndDate(2023, 2, 1)
            earnings = json.loads(self.ObjectStore.Read("Earnings_dates"))
            
            temp_earnings = dict(earnings)
            for ticker, str_dates in earnings.items():
                temp_earnings[ticker] = [dt.datetime.strptime(x, '%m/%d/%Y').date() for x in ast.literal_eval(str_dates)]

            self.loaded_earnings = dict(temp_earnings)
            self.Debug(temp_earnings)
            del temp_earnings
            del earnings
            
        elif self.SAVE_EARNINGS and self.ObjectStore.ContainsKey("Earnings_dates"):
            self.SetEndDate(2023, 2, 10)
            self.ObjectStore.Delete("Earnings_dates")     

        self.spy = self.AddEquity("SPY", Resolution.Daily)

        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.AfterMarketOpen("SPY", -10),
                 self.CheckEarningsRules)

        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.BeforeMarketClose("SPY", 60),
                 self.UpdateEarningsObject)

        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.BeforeMarketClose("SPY", 30),
                 self.TradeOptions)

    def UpdateEarningsObject(self):
        if self.SAVE_EARNINGS:
            for symbol, symbolData in self.symbolDataBySymbol.items():
                symbolData.update_earnings(self.fineFundamentals[symbol].EarningReports.FileDate)

    # Long Entry Criteria
    #    at least 30 days since the last earnings report
    #    over 30 days until the next earnings report
    #
    # Short Entry Criteria
    #    less than 10 days until the next earnings report
    #
    # Short Exit Criteria
    #    over 30 days since last earnings report
    #
    def CheckEarningsRules(self):
        if self.SAVE_EARNINGS: return

        cur_date = self.Time.date()

        for symbol, symbolData in self.symbolDataBySymbol.items():
            symbol_earnings = symbolData.earnings_dates
            symbol_earnings_series = pd.Series(symbol_earnings)
            insertion_index = symbol_earnings_series.searchsorted(cur_date) ## Result is the index at which it will exit in the next list

            prev_earnings_date = symbol_earnings[insertion_index-1]
            next_earnings_date = symbol_earnings[insertion_index]  ## WHEN RUNNING ALGORITHM, MAKE SURE self.EndDate IS LESS THAN THE LAST EARNINGS VALUE

            delta_prev = cur_date - prev_earnings_date
            delta_next = next_earnings_date - cur_date

            if delta_prev.days >= 30 and delta_next.days > 30:
                symbolData.long_entry_earnings = True
            else:
                symbolData.long_entry_earnings = False

            if delta_next.days < 10:
                symbolData.short_entry_earnings = True
            else:
                symbolData.short_entry_earnings = False

            if delta_prev.days >= 30:
                symbolData.short_exit_earnings = True
            else:
                symbolData.short_exit_earnings = False

            self.Debug(f"{insertion_index} {self.Time.date()} {symbol.Value} : long etr - {symbolData.long_entry_earnings}, short etr - {symbolData.short_entry_earnings}, short ext - {symbolData.short_exit_earnings}")

    def ContractFilter(self, symbol, min_strike, max_strike, min_expiry_days, max_expiry_days):
        contracts = self.OptionChainProvider.GetOptionContractList(symbol, self.Time.date())
        if len(contracts) == 0 : return []
        
        contract_list = [i for i in contracts if min_expiry_days < (i.ID.Date.date() - self.Time.date()).days < max_expiry_days]
        
        if len(contract_list) == 0: return []
        
        min_strike_price = sorted(contract_list, key = lambda x: abs(x.ID.StrikePrice - min_strike))[0].ID.StrikePrice
        max_strike_price = sorted(contract_list, key = lambda x: abs(x.ID.StrikePrice - max_strike))[0].ID.StrikePrice
        self.Debug(f"Min: {min_strike}")
        self.Debug(f"Max: {max_strike}")
        
        strike_list = sorted(set([i.ID.StrikePrice for i in contract_list]))
                
        min_strike_rank = strike_list.index(min_strike_price)
        max_strike_rank = strike_list.index(max_strike_price)
        
        try: 
            strikes = strike_list[min_strike_rank:max_strike_rank]
        except:
            strikes = strike_list
        
        filtered_contracts = [i for i in contract_list if i.ID.StrikePrice in strikes]

        return filtered_contracts 

    def AddContract(self, slice, symbol, symbolData, position):
        security_price = self.Securities[symbol].Price
        lower_atr_strike = (symbolData.atr.Current.Value * 8) + security_price
        upper_atr_strike = (symbolData.atr.Current.Value * 12) + security_price
        cur_atr_strike = (symbolData.atr.Current.Value * 10) + security_price
        cur_atr = symbolData.atr.Current.Value
        if position == 'long':
            self.Debug('long')
            self.Debug(f"p: {security_price}, atr_l: {lower_atr_strike}, atr_u: {upper_atr_strike}, atr_c: {cur_atr_strike}, atr: {cur_atr}")
            filtered_contracts = self.ContractFilter(symbol, lower_atr_strike, upper_atr_strike, 300, 400)
        elif position == 'short':
            self.Debug('short')
            self.Debug(f"atr_l: {lower_atr_strike}, atr_u: {upper_atr_strike}, atr_c: {cur_atr_strike}, atr: {cur_atr}")
            filtered_contracts = self.ContractFilter(symbol, lower_atr_strike, upper_atr_strike, 20, 30)

        if len(filtered_contracts) == 0: 
            return []
        else:
            calls = [x for x in filtered_contracts if x.ID.OptionRight == OptionRight.Call] 
            test = [x.ID.StrikePrice for x in calls]
            test2 = [x.ID.Date.date() for x in calls]
            self.Debug(f"Strikes: {test}")
            self.Debug(f"Expirys: {test2}")
            # contracts = sorted(sorted(calls, key = lambda x: abs(self.Securities[self.underlying].Price- x.ID.StrikePrice)), 
            #                                 key = lambda x: x.ID.Date, reverse=True)
            # if len(contracts) > 0:
            #     self.AddOptionContract(contracts[0], Resolution.Minute)
            #     return contracts[0]
            # else:
            #     return str()

    def TradeOptions(self):
        if self.SAVE_EARNINGS: return
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if symbolData.long_call == None:
                symbolData.long_call = self.AddContract(self.slice, symbol, symbolData, 'long')
            if symbolData.short_call == None:
                symbolData.short_call = self.AddContract(self.slice, symbol, symbolData, 'short')

    def OnData(self, data: Slice):
        self.slice = Slice
        
    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            if security.Symbol.Value == "SPY" or security.Symbol.SecurityType != SecurityType.Equity: continue

            atr = self.ATR(security.Symbol, 20, MovingAverageType.Simple, Resolution.Daily)
            history = self.History(security.Symbol, 20, Resolution.Daily)
            for bar in history.itertuples():
                tradebar = TradeBar(bar.Index[1], security.Symbol, bar.open, bar.high, bar.low, bar.close, bar.volume)
                atr.Update(tradebar)

            symbolData = SymbolData(self, security, atr)
            if not self.SAVE_EARNINGS:
                symbolData.earnings_dates = self.loaded_earnings[security.Symbol.Value]
            self.symbolDataBySymbol[security.Symbol] = symbolData

        for security in changes.RemovedSecurities:
            if security.Symbol.SecurityType == SecurityType.Equity:
                self.symbolDataBySymbol.pop(security.Symbol)

            else:
            # self.RemoveSecurity(x.Symbol)
                for symbol in self.Securities.Keys:
                    if symbol.SecurityType == SecurityType.Option and symbol.Underlying == security.Symbol:
                        self.RemoveSecurity(symbol)

    def CoarseSelection(self, coarse):
        filteredCoarse = [x.Symbol for x in coarse if x.Symbol.Value in self.tickers and x.HasFundamentalData]
        return filteredCoarse
    
    def FineSelection(self, fine):
        for x in fine: 
            self.fineFundamentals[x.Symbol] = x
        
        return [x.Symbol for x in fine]

    def OnEndOfAlgorithm(self):
        if self.SAVE_EARNINGS:
            earnings_dict = {}
            for symbol, symbolData in self.symbolDataBySymbol.items():
                self.Debug(f"{symbol.Value} : {symbolData.earnings_dates}")
                string_earnings = [t.strftime('%m/%d/%Y') for t in symbolData.earnings_dates]
                earnings_dict[str(symbol.Value)] = str(string_earnings)

            dump = json.dumps(earnings_dict)
            self.ObjectStore.Save("Earnings_dates", dump)
            self.Debug(earnings_dict)

        else:
            for symbol, symbolData in self.symbolDataBySymbol.items():
                self.Debug(symbolData.earnings_dates)

    
class SymbolData:
    def __init__(self, algo, security, atr):
        self.algo = algo
        self.security = security
        self.symbol = security.Symbol
        self.earnings_dates = []
        self.long_entry_earnings = False
        self.short_entry_earnings = False
        self.short_exit_earnings = False
        self.long_call = None
        self.short_call = None
        self.atr = atr

    def update_earnings(self, date):
        if date.date() not in self.earnings_dates:
            self.earnings_dates.append(date.date())