#region imports
from AlgorithmImports import *
from collections import deque
import numpy as np
from typing import List, Dict
from pandas.core.frame import DataFrame
#endregion

class Algo(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2008, 1, 1)
        self.capital = 100000
        self.SetCash(self.capital)
        
        self.market = self.AddEquity("SPY", Resolution.Daily).Symbol
        
        self.period:int = 5
        self.leverage:float = 0.99
        self.threshold:float = 0.0
        self.selection_flag:bool = False

        self.data:Dict[Symbol, RollingWindow] = {}
        self.weight:Dict[Symbol, float] = {}

        # self.benchmark_symbol = self.AddEquity("TQQQ", Resolution.Daily).Symbol
        # self.benchmark_values = deque(maxlen=252)
        
        self.coarse_count:int = 250
        self.stocks_to_hold:int = 10
        self.UniverseSettings.Resolution = Resolution.Daily
        self.SetSecurityInitializer(lambda security: security.SetMarketPrice(self.GetLastKnownPrice(security)))
        self.AddUniverseSelection(FineFundamentalUniverseSelectionModel(self.CoarseSelectionFunction, self.FineSelectionFunction))
        self.Schedule.On(self.DateRules.MonthStart(self.market), self.TimeRules.AfterMarketOpen(self.market), self.Selection)
        # self.Schedule.On(self.DateRules.EveryDay(self._benchmark_symbol), self.TimeRules.BeforeMarketClose(self._benchmark_symbol), self.UpdateChart)

    # def UpdateChart(self):
    #     hist = self.History(self.benchmark_symbol, 2, Resolution.Daily)
    #     if len(hist) == 2:
    #         bar = hist.iloc[0]
    #         self.benchmark_values.append(bar["close"])
    #         benchmark_perf = (self.benchmark_values[-1] / self.benchmark_values[0]) * self._capital
    #         self.Plot("Strategy Equity", self.benchmark_symbol.Value, benchmark_perf)
    #     else:
    #         self.Plot("Strategy Equity", self.benchmark_symbol.Value, 0)

    def OnSecuritiesChanged(self, changes:SecurityChanges) -> None:
        for security in changes.AddedSecurities:
            security.SetMarketPrice(self.GetLastKnownPrice(security))

    def CoarseSelectionFunction(self, coarse:List[CoarseFundamental]) -> List[Symbol]:
        # update the rolling window every day
        for stock in coarse:
            symbol = stock.Symbol

            # store daily price
            if symbol in self.data:
                self.data[symbol].Add(stock.AdjustedPrice)

        selected:List[Symbol] = [x.Symbol
            for x in sorted([x for x in coarse if x.Symbol.Value not in ["AMC", "GME", "UVXY"]],
                key = lambda x: x.DollarVolume, reverse = True)[:self.coarse_count]]

        # warmup daily prices 
        for symbol in selected:
            if symbol in self.data:
                continue

            self.data[symbol] = RollingWindow[float](self.period)
            history:DataFrame = self.History(symbol, self.period, Resolution.Daily)
            if history.empty:
                self.Log(f"Not enough data for {symbol} yet")
                continue
            for bar in history.itertuples():
                self.data[symbol].Add(bar.close)

        # return the symbols of the selected stocks
        return [x for x in selected if self.data[x].IsReady]

    def FineSelectionFunction(self, fine:List[FineFundamental]) -> List[Symbol]:
        fine:List[Symbol] = list(map(lambda x: x.Symbol, fine))
        avg_ratio:Dict[Symbol, float] = {}

        # calculate avg_ratio
        for symbol1 in fine:
            prices1:np.ndarray = np.array(list(self.data[symbol1]))
            average_price1:float = np.mean(prices1)
            normalized_prices1:np.ndarray = prices1 / average_price1
            ratio_sum:float = 0.

            for symbol2 in fine:
                if symbol1 != symbol2:
                    prices2:np.ndarray = np.array(list(self.data[symbol2]))
                    average_price2:float = np.mean(prices2)
                    normalized_prices2:np.ndarray = prices2 / average_price2
                    differences:np.ndarray = normalized_prices1 - normalized_prices2

                    max_diff:float = max(differences)
                    min_diff:float = min(differences)
                    diff_range:float = max_diff - min_diff
                    curr_diff:float = differences[0]
                    if diff_range != 0:
                        ratio:float = curr_diff / diff_range
                        ratio_sum += ratio

            avg_ratio_value:float = ratio_sum / (len(fine) - 1)
            if avg_ratio_value != 0:
                avg_ratio[symbol1] = avg_ratio_value

        if len(avg_ratio) >= self.stocks_to_hold:
            long:List[Tuple] = sorted(avg_ratio.items(), key=lambda x: abs(x[1]), reverse=True)[:self.stocks_to_hold]
            for symbol, _ in long:
                if avg_ratio[symbol] < 0:
                    self.weight[symbol] = 1.0 / len(long)
                elif avg_ratio[symbol] > 0:
                    # self.weight[symbol] = 0
                    self.weight[symbol] = -1.0 / len(long)

        # return the symbols of the selected stocks
        return list(self.weight.keys())
    
    def OnData(self, data:Slice) -> None:
        # liquidate
        invested:List[Symbol] = [x.Key for x in self.Portfolio if x.Value.Invested]
        for symbol in invested:
            if symbol not in self.weight:
                self.Liquidate()

        # submit MOO orders
        for symbol, w in self.weight.items():
            if symbol in data and data[symbol]:
                self.SetHoldings(symbol, w)

        self.weight.clear()

    def Selection(self) -> None:
        self.selection_flag = True

class CustomFeeModel(FeeModel):
    def GetOrderFee(self, parameters):
        fee = parameters.Security.Price * parameters.Order.AbsoluteQuantity * 0.00005
        return OrderFee(CashAmount(fee, "USD"))