from AlgorithmImports import *
from datetime import datetime
from QuantConnect.Algorithm.Framework import *
from QuantConnect.Algorithm.Framework.Alphas import *
import numpy as np
import pandas as pd
from Risk.TrailingStopRiskManagementModel import TrailingStopRiskManagementModel
from Execution.StandardDeviationExecutionModel import StandardDeviationExecutionModel

class Two_2Algorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2022, 5, 1)  
        self.SetCash(80000)

        self.SetRiskManagement(TrailingStopRiskManagementModel(0.046))
        self.SetExecution(StandardDeviationExecutionModel(20, 4.0, Resolution.Minute))

        # Define the universes of securities
        self.universe_A = [self.AddEquity(ticker).Symbol
                           for ticker in ["SCHG", "TQQQ", "TNA", "UWM", "JPST", "VUG", "QLD", "QQQ", "SOXL", "FNGU", "TECL", "URTY", "USD", "IBIT", "ETHE", "BTF", "BETH", "XSD", "QQQU", "YBTC", "VGT", "DXD", "UVXY", "UVIX", "BITI", "SETH", "FNGO", "IWM", "SMH"]]

        self.universe_C = [self.AddEquity(ticker).Symbol
                           for ticker in ["SVOL", "SPXU", "TZA", "UGL", "GLD", "GLTR", "AGQ"]]

        self.universe_D = [self.AddEquity(ticker).Symbol
                           for ticker in ["JPST", "SPY", "TQQQ", "USMV", "SSO", "TNA", "UPRO", "BOXX", "QID", "SDS", "SSG"]]

        self.universe_F = [self.AddEquity(ticker).Symbol
                           for ticker in ["SCHG", "TQQQ", "TNA", "UWM", "SVOL", "JPST", "SSO", "QLD", "SOXL", "MIDU", "FNGU", "URTY", "IBIT", "ETHE", "BTF", "BETH", "SDS", "QID", "UVXY", "UVIX"]]

        self.matrix_1_dates = 18
        self.reference = self.AddEquity("SVIX").Symbol

        self.Schedule.On(self.DateRules.EveryDay(self.reference), self.TimeRules.AfterMarketOpen(self.reference, 20), self.rebalance)
        #self.Schedule.On(self.DateRules.EveryDay(self.reference), self.TimeRules.BeforeMarketClose(self.reference, 10), self.rebalance)
        # Set other initial states
        self.shaken = []
        self.toShort = []
        self.toShort_A = []
        self.toShort_C = []
        self.toShort_D = []
        self.toShort_F = []

        self.SetWarmup(timedelta(hours=4))
        self.started = False

    def OnData(self, data):
        return

    def rebalance(self):
        universes = {
            "universe_A": self.universe_A,
            "universe_C": self.universe_C,
            "universe_D": self.universe_D,
            "universe_F": self.universe_F
        }

        # Get the selected symbols for each universe
        to_short_lists = {}
        for name, universe in universes.items():
            to_short_lists[name] = self.GetSymbols(universe, self.reference, self.matrix_1_dates, name)

        # Liquidate symbols not in the selected lists
        all_symbols = set(self.universe_A + self.universe_C + self.universe_D + self.universe_F)
        selected_symbols = set(symbol for sublist in to_short_lists.values() for symbol in sublist)

        for symbol in all_symbols - selected_symbols:
            if self.Portfolio[symbol].Invested:
                self.Liquidate(symbol)

        # To ensure the least volatile security gets more than 24.8% allocation
        total_allocation = 0.99  # Total allocation across selected securities
        least_volatile_allocation = 0.00  # Allocate more than 24.8%

        # First, find the least volatile security (the one with the lowest volatility)
        all_selected = [symbol for sublist in to_short_lists.values() for symbol in sublist]

        # Check if there are any selected securities
        if not all_selected:
            self.Debug("No selected securities. Skipping rebalance.")
            return  # Skip the rebalance if no symbols were selected
        
        volatilities = {symbol: self.GetVolatility(symbol) for symbol in all_selected}

        # Check if the volatilities dictionary is empty
        if not volatilities:
            self.Debug("No valid volatilities data. Skipping rebalance.")
            return  # Skip rebalance if there are no volatilities data

        # Find the least volatile security
        least_volatile_symbol = min(volatilities, key=volatilities.get)

        # Allocate the least volatile security more than 24.8%
        allocations = {least_volatile_symbol: least_volatile_allocation}
        remaining_allocation = total_allocation - least_volatile_allocation

        # Now, allocate the remaining securities with the rest of the percentage
        other_selected = [symbol for symbol in all_selected if symbol != least_volatile_symbol]

        # Exclude SVOL and JPST from allocations
        other_selected = [symbol for symbol in other_selected if symbol not in ["SVOL", "JPST", "BOXX", "AGQ", "GLD", "UVXY", "USMV", "UVIX", "SSO", "GLTR", "XSD", "SPY", "YBTC", "VUG", "VGT", "SDS", "SPYI", "BITI", "SETH", "QQQ", "IWM", "SVXY", "SMH", "SSG", "LVHI"]]

        num_other_selected = len(other_selected)

        if num_other_selected > 0:
            remaining_per_symbol_allocation = remaining_allocation / num_other_selected
            for symbol in other_selected:
                allocations[symbol] = remaining_per_symbol_allocation

        # Set holdings based on the computed allocations
        for symbol, allocation in allocations.items():
            self.EmitInsights(Insight.Price(symbol, Expiry.EndOfWeek, InsightDirection.Up, None, None, None, allocation))
            self.SetHoldings(symbol, allocation)

    def GetSymbols(self, universe, reference, lookback, universe_name):
        def matrix_1(first, second):
            first = np.array(first)
            first = np.diff(first) / first[:-1]
            second = np.array(second)
            second = np.diff(second) / second[:-1]
            A = np.vstack([first, np.ones(len(first))]).T
            basis = np.linalg.lstsq(A, second, rcond=None)[0]
            return basis[1], basis[0]

        history = self.History(universe + [reference], lookback, Resolution.Daily)
        filter = []

        for symbol in universe:
            if str(symbol.ID) not in history.index.get_level_values(0):
                continue

            symbol_price = list(history.loc[symbol]['close'])
            reference_price = list(history.loc[reference]['open'])

            if len(symbol_price) == len(reference_price):
                filter.append((symbol, matrix_1(symbol_price, reference_price)[0]))

        # If filter list is empty, return an empty list
        if not filter:
            return []

        filter.sort(key=lambda first: first[1], reverse=True)

        # Log the ranked list of securities by ticker symbol
        ranked_symbols = [(symbol.Value, score) for symbol, score in filter]
        self.Debug(f"Ranked securities for {universe_name} on {self.Time}: {ranked_symbols}")

        # Return the selected securities
        return [x[0] for x in filter[-1:]]

    def GetVolatility(self, symbol):
        """ 
        A helper function to calculate the volatility of a security.
        For simplicity, assuming volatility is the standard deviation of returns over the lookback period.
        """
        history = self.History(symbol, self.matrix_1_dates, Resolution.Daily)
        if history.empty:
            return float('inf')  # If no data, assume infinite volatility (avoid selection)
        
        # Calculate daily returns
        daily_returns = history['close'].pct_change().dropna()
        
        # Return the standard deviation (volatility) of daily returns
        volatility = daily_returns.std()
        return volatility