#region imports
from AlgorithmImports import *
#endregion

import numpy as np
from datetime import timedelta

from Risk.TrailingStopRiskManagementModel import TrailingStopRiskManagementModel
from Execution.StandardDeviationExecutionModel import StandardDeviationExecutionModel


class Two_2Algorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2015, 1, 1)
        self.SetCash(1_000_000)
        self.SetTimeZone("America/Chicago")

        self.benchmark_ticker = "SPY"
        self.reference_ticker_defensive = "DBP"
        self.reference_ticker_growth = "TQQQ"

        self.matrix_1_dates = 17          # defensive
        self.growth_lookback = 17         # growth 

        self.defensive_picks = 4
        self.growth_picks = 2

        self.base_position_weight = 0.1625
        self.max_total_exposure = 0.95

        self.SetBrokerageModel(AlphaStreamsBrokerageModel())
        self.SetRiskManagement(TrailingStopRiskManagementModel(0.054))
        self.SetExecution(StandardDeviationExecutionModel(40, 1.85, Resolution.Minute))

        self.defensive_tickers = [
            "GLD","SHY","JPST","SVOL","XLE","XME", "XLP", "UGL", "IAU", "XLU",
            "AAPL","MSFT","JNJ","HD","UNP","NVDA", "BOXX", "BIL", "VGSH", "SLV"]

        self.growth_candidates = [
            "SOXX","QQQ","XLK","SMH",
            "VGT","QLD","TQQQ","SOXL", "TECL"]

        self.defensive_symbols = []
        self.growth_symbols = []

        def safe_add(t):
            try:
                return self.AddEquity(t, Resolution.Minute).Symbol
            except:
                return None

        for t in self.defensive_tickers:
            s = safe_add(t)
            if s:
                self.defensive_symbols.append(s)

        for t in self.growth_candidates:
            s = safe_add(t)
            if s:
                self.growth_symbols.append(s)

        self._benchmark_symbol = safe_add(self.benchmark_ticker)
        self._reference_def_symbol = safe_add(self.reference_ticker_defensive)
        self._reference_growth_symbol = safe_add(self.reference_ticker_growth)

        self.SetBenchmark(self._benchmark_symbol)

        self.Schedule.On(
            self.DateRules.MonthStart(self._benchmark_symbol),
            self.TimeRules.AfterMarketOpen(self._benchmark_symbol, 12),
            self.rebalance_defensive
        )

        self.Schedule.On(
            self.DateRules.WeekStart(self._benchmark_symbol),
            self.TimeRules.AfterMarketOpen(self._benchmark_symbol, 10),
            self.rebalance_growth
        )

        self.shaken_defensive = []
        self.shaken_growth = []
        self.started = False

    def OnData(self, data: Slice):
        if not self.started:
            self.rebalance(run_defensive=True, run_growth=True)
            self.started = True

    def matrix_1(self, 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] 

    def rebalance_defensive(self):
        self.rebalance(run_defensive=True, run_growth=False)

    def rebalance_growth(self):
        self.rebalance(run_defensive=False, run_growth=True)

    def rebalance(self, run_defensive=True, run_growth=True):

        lookback = max(self.matrix_1_dates, self.growth_lookback)
        all_syms = list(set(
            self.defensive_symbols +
            self.growth_symbols +
            [self._reference_def_symbol, self._reference_growth_symbol]
        ))

        history = self.History(all_syms, lookback, Resolution.Daily)
        if history.empty:
            return

        if run_defensive:
            ranked = []
            for s in self.defensive_symbols:
                if s not in history.index.get_level_values(0):
                    continue

                a = history.loc[s]["close"][-self.matrix_1_dates:]
                b = history.loc[self._reference_def_symbol]["close"][-self.matrix_1_dates:]

                if len(a) != len(b) or len(a) < 3:
                    continue

                intercept, _ = self.matrix_1(list(a), list(b))
                ranked.append((s, intercept))

            ranked.sort(key=lambda x: x[1], reverse=True)
            self.shaken_defensive = [x[0] for x in ranked[:self.defensive_picks]]

        if run_growth:
            ranked = []
            for s in self.growth_symbols:
                if s == self._reference_growth_symbol:
                    continue

                if s not in history.index.get_level_values(0):
                    continue

                a = history.loc[s]["close"][-self.growth_lookback:]
                b = history.loc[self._reference_growth_symbol]["close"][-self.growth_lookback:]

                if len(a) != len(b) or len(a) < 3:
                    continue

                _, slope = self.matrix_1(list(a), list(b))
                ranked.append((s, slope))

            ranked.sort(key=lambda x: x[1], reverse=True)
            self.shaken_growth = [x[0] for x in ranked[:self.growth_picks]]

        targets = list(dict.fromkeys(self.shaken_defensive + self.shaken_growth))

        desired = self.base_position_weight * len(targets)
        scale = min(1.0, self.max_total_exposure / desired)
        w = self.base_position_weight * scale

        invested = [kv.Key for kv in self.Portfolio if kv.Value.Invested]

        for s in invested:
            if s not in targets and (run_defensive or s not in self.shaken_defensive):
                self.Liquidate(s)

        for s in targets:
            self.SetHoldings(s, w)