from AlgorithmImports import *
from symbol_calcs import *
from datetime import timedelta
import numpy as np

class EldersTripleScreenAlpha(QCAlgorithm):

    def Initialize(self):
        ''' Initial Algo parameters and QuantConnect methods'''

        ########### Strategy Params ###########
        self.SetStartDate(2017, 1, 1)
        self.SetEndDate(2024, 3, 31)
        self.SetCash(100000000)

        # Warmup the algorithm with prior data for backtesting
        self.SetWarmup(timedelta(30))   
        
        # Additional variables for Sharpe Ratio calculation
        self.dailyReturns = []
        self.sharpeRatios = []
        self.previousPortfolioValue = self.Portfolio.TotalPortfolioValue

        ########### Universe Selection ###########
        self.UniverseSettings.Resolution = Resolution.Hour
        self.UniverseSettings.Leverage = 2
        self.AddUniverse(self._fundamental_selection_function)

        # Variables for universe selection model
        self.coarse_count = 10
        
        self.symbols = []
        self.data = {}

        ########### Elders Triple Screen Signals ###########
        self.ema_period = 20
        self.rsi_period = 14
        self.macd_fast_period = 12
        self.macd_slow_period = 26
        self.macd_signal_period = 6

        ########### Risk Management Model ###########
        self.AddRiskManagement(MaximumDrawdownPercentPerSecurity(0.05))
        
        ########### Order Execution Model ###########
        self.SetExecution(ImmediateExecutionModel())

        ########### Portfolio Construction Model ###########
        self.SetPortfolioConstruction(RiskParityPortfolioConstructionModel(portfolioBias=PortfolioBias.Long))

        # Required portfolio free cash value
        self.Settings.FreePortfolioValuePercentage=0.05
        
        ########### Reality Modeling Parameters ###########
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        self.SetRiskFreeInterestRateModel(InterestRateProvider())
        self.Portfolio.MarginCallModel = DefaultMarginCallModel(self.Portfolio, self.DefaultOrderProperties)
        
        ########### Scheduler ###########
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.At(11,30,0, TimeZones.NewYork), self.EvaluateIndicators)
        
        ########### Object Store ##########
        self.universe_store = ''
        self.EMA_store = ''
        self.RSI_store = ''
        self.MACD_store = ''
        
        ########### Load Existing Universe ###########
        if self.IsLiveEnvironment():
            self.LoadExistingUniverse()

        # Initialize plotting
        self.Plot("Sharpe Ratio", "Daily Annualized Sharpe Ratio", 0)

    def IsLiveEnvironment(self):
        ''' Check if the algorithm is running in a live environment '''
        return self.LiveMode

    def LoadExistingUniverse(self):
        ''' Load the existing universe from the live brokerage account holdings '''
        self.Debug("Loading existing universe from brokerage account...")

        # Fetch existing holdings
        holdings = [x.Symbol for x in self.Portfolio.Values if x.Invested]
        self.Debug(f"Found {len(holdings)} existing holdings.")

        # Register indicators for the existing holdings
        for symbol in holdings:
            self.data[symbol] = {
                'ema': ExponentialMovingAverage(self.ema_period),
                'rsi': RelativeStrengthIndex(self.rsi_period, MovingAverageType.Wilders),
                'macd': MovingAverageConvergenceDivergence(self.macd_fast_period, self.macd_slow_period, self.macd_signal_period, MovingAverageType.Wilders)
            }
            self.RegisterIndicator(symbol, self.data[symbol]['ema'], Resolution.Hour)
            self.RegisterIndicator(symbol, self.data[symbol]['rsi'], Resolution.Hour)
            self.RegisterIndicator(symbol, self.data[symbol]['macd'], Resolution.Daily)

        return holdings
        

    def _fundamental_selection_function(self, fundamental: List[Fundamental]) -> List[Symbol]:
        ''' Fundamental Filter to produce a list of securities within the universe'''
        filtered = [f for f in fundamental if f.Price > 10 and \
                                            f.HasFundamentalData and \
                                            not np.isnan(f.ValuationRatios.PERatio) and \
                                            f.ValuationRatios.ForwardPERatio > 5 and \
                                            f.MarketCap > 500000]
        
        self.Debug(f"Filtered securities: {', '.join([f.Symbol.Value for f in filtered])}")
        
        sorted_by_dollar_volume = sorted(filtered, key=lambda f: f.DollarVolume, reverse=True)[:100]
        self.Debug(f"Top 100 by dollar volume: {', '.join([f.Symbol.Value for f in sorted_by_dollar_volume])}")
        
        sorted_by_pe_ratio = sorted(sorted_by_dollar_volume, key=lambda f: f.ValuationRatios.PERatio, reverse=False)[:20]
        self.Debug(f"Final selected securities: {', '.join([f.Symbol.Value for f in sorted_by_pe_ratio])}")
        
        self.Debug(f"Final selected securities: {sorted_by_pe_ratio[:20]}")
        
        return [f.Symbol for f in sorted_by_pe_ratio[:20]]
    
    def FineFilter(self, fine):
        ''' Placeholder for further filtering of the universe '''
        pass

    ''' Simpler changes to universe
    # this event fires whenever we have changes to our universe
    def OnSecuritiesChanged(self, changes):
        # liquidate removed securities
        for security in changes.RemovedSecurities:
            if security.Invested:
                self.Liquidate(security.Symbol)
                self.Debug(f"Liquidated {security.Symbol}")

        # we want 20% allocation in each security in our universe - This gives 10% allocation
        for security in changes.AddedSecurities:
            self.SetHoldings(security.Symbol, 0.1)
            self.Debug(f"Bought {security.Symbol}")
    '''

    ############### Evaluate Indicators #########################
    ''' Elder's Triple Screen Alpha Logic'''

    def OnSecuritiesChanged(self, changes):
        ''' Add indicators whenever securities are added to or removed from the universe'''
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            self.data[symbol] = {
                'ema': ExponentialMovingAverage(self.ema_period),
                'rsi': RelativeStrengthIndex(self.rsi_period, MovingAverageType.Wilders),
                'macd': MovingAverageConvergenceDivergence(self.macd_fast_period, self.macd_slow_period, self.macd_signal_period, MovingAverageType.Wilders)
            }
            self.RegisterIndicator(symbol, self.data[symbol]['ema'], Resolution.Hour)
            self.RegisterIndicator(symbol, self.data[symbol]['rsi'], Resolution.Hour)
            self.RegisterIndicator(symbol, self.data[symbol]['macd'], Resolution.Daily)
            
            # Reality Modeling for slippage
            #security.SetSlippageModel(VolumeShareSlippageModel(0.025, 0.1))
            #security.SetSlippageModel(MarketImpactSlippageModel(self))
            

        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if symbol in self.data:
                del self.data[symbol]
    
    def EvaluateIndicators(self):
        ''' Evaluate the indicators for the securities '''
        # Maybe change this self.data.items() to something else. Rn its not iterating through all the securities like it should
        for symbol, indicators in self.data.items():
            ema = indicators['ema']
            rsi = indicators['rsi']
            macd = indicators['macd']
            
            if not (ema.IsReady and rsi.IsReady and macd.IsReady):
                continue
            
            ema_value = ema.Current.Value
            rsi_value = rsi.Current.Value
            macd_value = macd.Current.Value
            macd_signal_value = macd.Signal.Current.Value
            price = self.Securities[symbol].Price
            
            direction = InsightDirection.Flat
            confidence = 0
            magnitude = 0
            
            if rsi_value < 30 and macd_value > macd_signal_value: #and price > ema_value:
                direction = InsightDirection.Up
                confidence = (70 - rsi_value) / 100
                magnitude = macd_value - macd_signal_value
                
            # Consider changing this to an AND instead of an or, since both need to be satisfied
            elif rsi_value > 70 or macd_value < macd_signal_value:# or price < ema_value:
                direction = InsightDirection.Down
                confidence = (rsi_value - 30) / 100
                magnitude = macd_signal_value - macd_value
            
            insight = Insight.Price(symbol, timedelta(days=1), direction, magnitude, confidence)
            #self.EmitInsights(insight)
            #self.Debug(f"Generated Insight: {insight}")
            
            if confidence > 0.2:
                self.EmitInsights(insight)
                self.Debug(f"Generated Insight: {insight}")
            
            return insight
    
    def OnEndOfDay(self):
        ''' Calculate daily return and Sharpe Ratio '''
        currentPortfolioValue = self.Portfolio.TotalPortfolioValue
        dailyReturn = (currentPortfolioValue - self.previousPortfolioValue) / self.previousPortfolioValue
        self.dailyReturns.append(dailyReturn)
        self.previousPortfolioValue = currentPortfolioValue

        if len(self.dailyReturns) > 1:
            meanReturn = np.mean(self.dailyReturns)
            stdDeviation = np.std(self.dailyReturns)
            if stdDeviation != 0:
                sharpeRatio = (meanReturn / stdDeviation) * np.sqrt(252)  # Annualizing the Sharpe Ratio
                self.sharpeRatios.append(sharpeRatio)
            else:
                self.sharpeRatios.append(0)

            # Update plot
            self.Plot("Sharpe Ratio", "Daily Annualized Sharpe Ratio", self.sharpeRatios[-1])

    def OnEndOfAlgorithm(self):
        ''' Calculate the Sharpe Ratio at the end of the algorithm and plot it '''
        if len(self.sharpeRatios) > 1:
            self.Plot("Sharpe Ratio", "Final Sharpe Ratio", self.sharpeRatios[-1])
        else:
            self.Log("Not enough data to plot Sharpe Ratio.")

#region imports
from AlgorithmImports import *
#endregion

# Note this class was copied from bottom of main.py, it can technically be run seperately as long as the import is there

class SymbolData(object):
    def __init__(self, symbol):
        self._symbol = symbol
        self.tolerance = 1.01
        self.fast = ExponentialMovingAverage(100)
        self.slow = ExponentialMovingAverage(300)
        self.is_uptrend = False
        self.scale = 1
        
        print("SymbolData Init Complete")

    def update(self, time, value):
        if self.fast.update(time, value) and self.slow.update(time, value):
            fast = self.fast.current.value
            slow = self.slow.current.value
            self.is_uptrend = fast > slow * self.tolerance

        if self.is_uptrend:
            self.scale = (fast - slow) / ((fast + slow) / 2.0)
            return
        print("Update method ran successfully")