# 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.SetEndDate(2023, 2, 4)
        self.SetCash(50000)  # Set Strategy Cash
        self.UniverseSettings.Resolution = Resolution.Daily
        
        self.tickers = ["AAPL"]
        
        self.AddUniverse(self.CoarseSelection, self.FineSelection)
        
        self.symbolDataBySymbol = {}
        self.fineFundamentals = {}

        self.SAVE_EARNINGS = True
        self.loaded_earnings = {}

        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, 4)
            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)

    # 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 OnData(self, data: Slice):
        if self.SAVE_EARNINGS:
            for symbol, symbolData in self.symbolDataBySymbol.items():
                symbolData.update_earnings(self.fineFundamentals[symbol].EarningReports.FileDate)
        
    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            if security.Symbol.Value == "SPY": continue

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

    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):
        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


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