from AlgorithmImports import *
from QuantConnect.DataSource import *

import statistics

# TODO:
# - combine ranking values from various timeframes
# - combine with Brain Language Metrics and Sentiment Indicator

class BrainMLRankingDataAlgorithm(QCAlgorithm):
    def initialize(self) -> None:
        self.set_start_date(2020, 1, 1)
        self.set_end_date(2024, 4, 1)
        self.set_cash(100000)

        self.position_keep_days = int(self.GetParameter("position_keep_days")) if self.GetParameter("position_keep_days") else 3
        self.sentiment_days = int(self.GetParameter("sentiment_days")) if self.GetParameter("sentiment_days") else 7
        self.no_positions = int(self.GetParameter("no_positions")) if self.GetParameter("no_positions") else 10
        self.universe_size = int(self.GetParameter("universe_size")) if self.GetParameter("universe_size") else 100
        self.do_short = int(self.GetParameter("do_short")) if self.GetParameter("do_short") else 1
        self.ranking_weight = float(self.GetParameter("ranking_weight")) if self.GetParameter("ranking_weight") else 0.33
        self.sentiment_weight = float(self.GetParameter("sentiment_weight")) if self.GetParameter("sentiment_weight") else 0.33
        self.company_filing_weight = float(self.GetParameter("company_filing_weight")) if self.GetParameter("company_filing_weight") else 0.33

        self.Debug(f"position_keep_days: {self.position_keep_days}")
        self.Debug(f"sentiment_days: {self.sentiment_days}")
        self.Debug(f"no_positions: {self.no_positions}")
        self.Debug(f"universe_size: {self.universe_size}")
        self.Debug(f"do_short: {self.do_short}")
        self.Debug(f"ranking_weight: {self.ranking_weight}")
        self.Debug(f"sentiment_weight: {self.sentiment_weight}")
        self.Debug(f"company_filing_weight: {self.company_filing_weight}")

        if self.do_short == 0:
            self.no_positions *= 2

        self.brain_stock_ranking_class = self.get_brain_stock_ranking_class()
        self.brain_sentiment_class = self.get_brain_sentiment_class()
        self.brain_filing_class = BrainCompanyFilingLanguageMetrics10K

        self.UniverseSettings.Resolution = Resolution.Daily
        self.SetUniverseSelection(FineFundamentalUniverseSelectionModel(self.SelectCoarse, self.SelectFine))

        self.symbol_str_to_symbol = {}
        self.symbol_by_ranking_symbol = {}
        self.symbol_by_sentiment_symbol = {}
        self.symbol_by_filing_symbol = {}
        self.history_14_days = {}
        self.atr_14_days = {}
        self.open_dates = {}

        benchmark_chart = Chart("Benchmark SPY")
        benchmark_chart.AddSeries(Series("SPY", SeriesType.Line))
        benchmark_chart.AddSeries(Series("Portfolio", SeriesType.Line))
        self.add_chart(benchmark_chart)
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol    # for benchmarking
        self.starting_spy_price = None
        self.starting_portfolio_value = None

        rank_chart = Chart("Brain Stock Ranking")
        rank_chart.AddSeries(Series("Max Rank", SeriesType.Line))
        rank_chart.AddSeries(Series("Min Rank", SeriesType.Line))
        self.add_chart(rank_chart)

        sentiment_chart = Chart("Brain Sentiment Indicator")
        sentiment_chart.AddSeries(Series("Max Sentiment", SeriesType.Line))
        sentiment_chart.AddSeries(Series("Min Sentiment", SeriesType.Line))
        self.add_chart(sentiment_chart)

        filing_chart = Chart("Brain Company Filing Language Metrics")
        filing_chart.AddSeries(Series("Max Sentiment", SeriesType.Line))
        filing_chart.AddSeries(Series("Min Sentiment", SeriesType.Line))
        filing_chart.AddSeries(Series("Max Uncertainty", SeriesType.Line))
        filing_chart.AddSeries(Series("Min Uncertainty", SeriesType.Line))
        filing_chart.AddSeries(Series("Max Litigious", SeriesType.Line))
        filing_chart.AddSeries(Series("Min Litigious", SeriesType.Line))
        self.add_chart(filing_chart)

        self.set_warmup(14)

    def get_brain_stock_ranking_class(self):
        if self.position_keep_days <= 2:
            return BrainStockRanking2Day
        elif self.position_keep_days <= 3:
            return BrainStockRanking3Day
        elif self.position_keep_days <= 5:
            return BrainStockRanking5Day
        elif self.position_keep_days <= 10:
            return BrainStockRanking10Day
        elif self.position_keep_days <= 21:
            return BrainStockRanking21Day
        else:
            raise Exception("Invalid position_keep_days")

    def get_brain_sentiment_class(self):
        if self.sentiment_days <= 7:
            return BrainSentimentIndicator30Day
        elif self.position_keep_days <= 30:
            return BrainSentimentIndicator30Day
        else:
            raise Exception("Invalid sentiment_days")

    def SelectCoarse(self, coarse):
        filtered_coarse = [x for x in coarse if x.price > 5]
        filtered_coarse = [x for x in filtered_coarse if x.volume > 1_000_000]

        sortedByDollarVolume = sorted(filtered_coarse, key=lambda x: x.dollar_volume, reverse=True) 
        symbols = [ x.Symbol for x in sortedByDollarVolume[:self.universe_size] ]

        return symbols

    def SelectFine(self, fine):
        # we already have our self.universe_size symbols from coarse selection, no further selection needed
        return [ x.Symbol for x in fine ]

    def OnSecuritiesChanged(self, changes):
        for added in changes.AddedSecurities:
            symbol = added.Symbol
            self.symbol_str_to_symbol[symbol.value] = symbol
            ranking_dataset_symbol = self.add_data(self.brain_stock_ranking_class, symbol).symbol
            sentiment_dataset_symbol = self.add_data(self.brain_sentiment_class, symbol).symbol
            filing_dataset_symbol = self.add_data(self.brain_filing_class, symbol).symbol
            self.symbol_by_ranking_symbol[ranking_dataset_symbol] = symbol
            self.symbol_by_sentiment_symbol[sentiment_dataset_symbol] = symbol
            self.symbol_by_filing_symbol[filing_dataset_symbol] = symbol

            self.history_14_days[symbol.value] = RollingWindow[TradeBar](14)
            self.atr_14_days[symbol.value] = AverageTrueRange(14)

            # warm up the ATR indicator and the history
            history = self.history[TradeBar](symbol, 14, Resolution.DAILY)
            for bar in history:
                self.history_14_days[symbol.value].Add(bar)
                self.atr_14_days[symbol.value].Update(bar)

            #self.Debug(f"added {symbol} with dataset symbol {dataset_symbol}")
            
        for removed in changes.RemovedSecurities:
            symbol = removed.Symbol
            if symbol.value in self.symbol_str_to_symbol:
                del self.symbol_str_to_symbol[symbol.value]
            if symbol in self.symbol_by_ranking_symbol:
                del self.symbol_by_ranking_symbol[symbol]
            if symbol in self.symbol_by_sentiment_symbol:
                del self.symbol_by_sentiment_symbol[symbol]
            if symbol in self.symbol_by_filing_symbol:
                del self.symbol_by_filing_symbol[symbol]
            if symbol.value in self.history_14_days:
                del self.history_14_days[symbol.value]
            if symbol.value in self.atr_14_days:
                del self.atr_14_days[symbol.value]

    def on_data(self, slice: Slice) -> None:
        
        for symbol, trade_bar in slice.Bars.items():
            #self.Debug(f"adding {symbol} to history {symbol in self.history_14_days} {symbol in self.atr_14_days}")
            if symbol in self.history_14_days:
                self.history_14_days[symbol.value].Add(trade_bar)
            if symbol in self.atr_14_days:
                self.atr_14_days[symbol.value].Update(trade_bar)

        if self.is_warming_up:
            return

        if self.spy in slice.Bars:
            spy_price = slice[self.spy].Close
            if self.starting_spy_price is None:
                self.starting_spy_price = spy_price
                self.starting_portfolio_value = self.Portfolio.TotalPortfolioValue
            normalized_spy = spy_price / self.starting_spy_price
            normalized_portfolio = self.Portfolio.TotalPortfolioValue / self.starting_portfolio_value
            self.Plot("Benchmark SPY", "SPY", normalized_spy)
            self.Plot("Benchmark SPY", "Portfolio", normalized_portfolio)

        # Collect rankings for all symbols
        ranking_points = slice.Get(self.brain_stock_ranking_class)
        sentiment_points = slice.Get(self.brain_sentiment_class)
        filing_points = slice.Get(self.brain_filing_class)
        if ranking_points is None or len(ranking_points) == 0:
            self.Debug("No ranking points")
            return

        # Create new dictionaries
        ranking_points_dict = {point.symbol.Value: point for point in ranking_points.Values if point.symbol in self.symbol_by_ranking_symbol}
        rank_values = [point.rank for point in ranking_points_dict.values()]
        min_rank, max_rank = min(rank_values), max(rank_values)
        self.Plot("Brain Stock Ranking", "Max Rank", max_rank)
        self.Plot("Brain Stock Ranking", "Min Rank", min_rank)

        if sentiment_points is None or len(sentiment_points) == 0:
            sentiment_points_dict = {}
            sentiment_values = []
            min_sentiment, max_sentiment = 0, 0
        else:
            sentiment_points_dict = {point.symbol.Value: point for point in sentiment_points.Values if point.symbol in self.symbol_by_sentiment_symbol}
            sentiment_values = [point.sentiment for point in sentiment_points_dict.values()]
            min_sentiment, max_sentiment = min(sentiment_values), max(sentiment_values)
            self.Plot("Brain Sentiment Indicator", "Max Sentiment", max_sentiment)
            self.Plot("Brain Sentiment Indicator", "Min Sentiment", min_sentiment)

        if filing_points is None or len(filing_points) == 0:
            filing_points_dict = {}
            filing_values = []
            min_filing, max_filing = 0, 0
        else:
            filing_points_dict = {point.symbol.Value: point for point in filing_points.Values if point.symbol in self.symbol_by_filing_symbol}

            filing_sentiment_values = [point.report_sentiment.sentiment for point in filing_points_dict.values()]
            filing_uncertainity_values = [point.report_sentiment.uncertainty for point in filing_points_dict.values()]
            filing_litigious_values = [point.report_sentiment.litigious for point in filing_points_dict.values()]

            min_filing_sentiment, max_filing_sentiment = min(filing_sentiment_values), max(filing_sentiment_values)
            min_filing_uncertainty, max_filing_uncertainty = min(filing_uncertainity_values), max(filing_uncertainity_values)
            min_filing_litigious, max_filing_litigious = min(filing_litigious_values), max(filing_litigious_values)

            self.Plot("Brain Company Filing Language Metrics", "Max Sentiment", max_filing_sentiment)
            self.Plot("Brain Company Filing Language Metrics", "Min Sentiment", min_filing_sentiment)
            self.Plot("Brain Company Filing Language Metrics", "Max Uncertainty", max_filing_uncertainty)
            self.Plot("Brain Company Filing Language Metrics", "Min Uncertainty", min_filing_uncertainty)
            self.Plot("Brain Company Filing Language Metrics", "Max Litigious", max_filing_litigious)
            self.Plot("Brain Company Filing Language Metrics", "Min Litigious", min_filing_litigious)

            filing_values = filing_sentiment_values + filing_uncertainity_values + filing_litigious_values
            min_filing, max_filing = min(filing_values), max(filing_values)


        symbol_rank_pairs = []
        for symbol in ranking_points_dict.keys():
            rank = ranking_points_dict[symbol]

            sentiment = sentiment_points_dict[symbol] if symbol in sentiment_points_dict else None
            filing = filing_points_dict[symbol] if symbol in filing_points_dict else None

            ranking_weight = self.ranking_weight

            normalized_rank = (rank.rank - min_rank) / (max_rank - min_rank) if max_rank != min_rank else 0

            if sentiment is None or max_sentiment == min_sentiment:
                normalized_sentiment = 0
                ranking_weight += self.sentiment_weight
            else:
                normalized_sentiment = (sentiment.sentiment - min_sentiment) / (max_sentiment - min_sentiment)

            if filing is None or max_filing == min_filing:
                normalized_filing = 0
                ranking_weight += self.company_filing_weight
            else:
                normalized_filing = (filing.report_sentiment.sentiment - min_filing) / (max_filing - min_filing)

            combined_rank = ranking_weight * normalized_rank + self.sentiment_weight * normalized_sentiment + self.company_filing_weight * normalized_filing
            symbol_rank_pairs.append((symbol, combined_rank))               

        if len(symbol_rank_pairs) == 0:
            self.Debug("No symbols with rank after step 1")
            return []

        symbol_rank_pairs = [p for p in symbol_rank_pairs if p[0] in self.history_14_days and self.history_14_days[p[0]].Count >= 14]
        if len(symbol_rank_pairs) == 0:
            self.Debug("No symbols with rank after step 2")
            return []

        symbol_rank_pairs = [p for p in symbol_rank_pairs if self.history_14_days[p[0]][0].close >= 5]
        if len(symbol_rank_pairs) == 0:
            self.Debug("No symbols with rank after step 3")
            return []

        symbol_rank_pairs = [p for p in symbol_rank_pairs if statistics.mean(bar.volume for bar in self.history_14_days[p[0]]) >= 1_000_000]
        if len(symbol_rank_pairs) == 0:
            self.Debug("No symbols with rank after step 4")
            return []

        symbol_rank_pairs = [p for p in symbol_rank_pairs if self.atr_14_days[p[0]][0].value >= 0.5]
        if len(symbol_rank_pairs) == 0:
            self.Debug("No symbols with rank after step 5")
            return []

        sorted_symbols = sorted(symbol_rank_pairs, key=lambda x: x[1], reverse=True)

        top_symbols = sorted_symbols[:int(self.no_positions/2)]
        bottom_symbols = sorted_symbols[-int(self.no_positions/2):]

        still_holding = 0
        for holding in self.Portfolio.Values:
            if holding.Invested and holding.Symbol.value not in [symbol for symbol, rank in top_symbols] and holding.Symbol.value not in [symbol for symbol, rank in bottom_symbols]:
                open_date = self.open_dates.get(holding.Symbol.value)
                if open_date is None:
                    # if we don't have an open date, this position should be here
                    self.set_holdings(holding.Symbol, 0)                    
                elif (self.Time - open_date).days >= self.position_keep_days:
                    self.set_holdings(holding.Symbol, 0)
                    del self.open_dates[holding.Symbol.value]
                    #self.Debug(f"liquidate {holding.Symbol.value}")
                else:
                    still_holding += 1
                #if open_date is None:
                #    self.Debug(f"open_date not none {holding.Symbol.value}")

        total_invested = sum([abs(holding.HoldingsValue) for holding in self.Portfolio.Values if holding.Invested])
        portfolio_percentage_used = total_invested / self.Portfolio.TotalPortfolioValue

        if portfolio_percentage_used > 1.0:
            return

        new_positions = self.no_positions - still_holding
        top_symbols = top_symbols[:int(new_positions/2)]
        bottom_symbols = bottom_symbols[-int(new_positions/2):]
        #self.Debug(f"bottom_symbols: {len(bottom_symbols)}")

        # Place orders
        sum_ranks = sum(abs(rank) for symbol, rank in top_symbols)
        if self.do_short:
            sum_ranks += sum(abs(rank) for symbol, rank in bottom_symbols)

        for symbol, rank in top_symbols:
            weight = (abs(rank) / sum_ranks) * (1 - portfolio_percentage_used)
            self.set_holdings(self.symbol_str_to_symbol[symbol], weight)
            self.open_dates[symbol] = self.Time
            #self.Debug(f"long {symbol}")
        if self.do_short > 0:
            for symbol, rank in bottom_symbols:
                weight = (-(abs(rank) / sum_ranks)) * (1 - portfolio_percentage_used)
                self.set_holdings(self.symbol_str_to_symbol[symbol], weight)
                self.open_dates[symbol] = self.Time
                #self.Debug(f"short {symbol}")