#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(2013, 1, 1)
        self.SetCash(1_000_000)

        self.SetBrokerageModel(AlphaStreamsBrokerageModel())
        self.SetTimeZone(TimeZones.Chicago)

        self.tickers = ["AAPL", "MSFT", "AMZN", "FB", "TSLA", "GOOGL", "GOOG", "NVDA", "JNJ", "JPM"]
        self.symbols = []

        self._benchmark_ticker = "SPY"
        self._reference_ticker = "TECS"
        self.matrix_1_dates = 15

        self.short_target_weight = -0.09
        self.long_target_weight = 0.165
        self.short_update_weight = -0.08

        self.SetRiskManagement(TrailingStopRiskManagementModel(0.047))
        self.SetExecution(StandardDeviationExecutionModel(30, 1.25, Resolution.Minute))

        for t in self.tickers:
            self.symbols.append(self.AddEquity(t, Resolution.Minute).Symbol)

        self._benchmark = self.AddEquity(self._benchmark_ticker, Resolution.Minute).Symbol
        self._reference = self.AddEquity(self._reference_ticker, Resolution.Minute).Symbol
        self.SetBenchmark(self._benchmark_ticker)

        daterule_week = self.DateRules.WeekStart(self._benchmark)
        daterule_day = self.DateRules.EveryDay(self._benchmark)

        self.Schedule.On(daterule_week, self.TimeRules.AfterMarketOpen(self._benchmark, 8), self.rebalance)
        self.Schedule.On(daterule_day, self.TimeRules.AfterMarketOpen(self._benchmark, 18), self.update)
        self.Schedule.On(daterule_day, self.TimeRules.AfterMarketOpen(self._benchmark, 54), self.update)
        self.Schedule.On(daterule_day, self.TimeRules.AfterMarketOpen(self._benchmark, 164), self.update)
        self.Schedule.On(daterule_day, self.TimeRules.AfterMarketOpen(self._benchmark, 264), self.update)
        self.Schedule.On(daterule_day, self.TimeRules.AfterMarketOpen(self._benchmark, 324), self.update)
        self.Schedule.On(daterule_day, self.TimeRules.BeforeMarketClose(self._benchmark, 1), self.update)

        self.short_candidates = []
        self.long_candidates = []

        self.started = False

    def OnData(self, data: Slice):
        if self.started:
            return
        self.rebalance()
        self.started = True

    def _regression_intercept_slope(self, first, second):
        first = np.asarray(first, dtype=float)
        second = np.asarray(second, dtype=float)

        if len(first) < 2 or len(second) < 2:
            return 0.0, 0.0

        first = np.diff(first) / first[:-1]
        second = np.diff(second) / second[:-1]

        n = min(len(first), len(second))
        first = first[-n:]
        second = second[-n:]

        A = np.vstack([first, np.ones(len(first))]).T
        slope, intercept = np.linalg.lstsq(A, second, rcond=None)[0]
        return float(intercept), float(slope)

    def rebalance(self):
        hist = self.History(self.symbols + [self._reference], self.matrix_1_dates, Resolution.Daily)
        if hist.empty:
            return

        idx0 = hist.index.get_level_values(0)
        if self._reference not in idx0:
            return

        ranked = []
        ref_close = hist.loc[self._reference]["close"].tolist()
        if len(ref_close) < 2:
            return

        for sym in self.symbols:
            if sym not in idx0:
                continue

            sym_close = hist.loc[sym]["close"].tolist()
            if len(sym_close) != len(ref_close) or len(sym_close) < 2:
                continue

            intercept, _ = self._regression_intercept_slope(sym_close, ref_close)
            ranked.append((sym, intercept))

        ranked.sort(key=lambda x: x[1], reverse=True)
        if len(ranked) < 3:
            return

        self.short_candidates = [ranked[0][0]]
        self.long_candidates = [x[0] for x in ranked[-5:]]

        invested_symbols = [kvp.Key for kvp in self.Portfolio if kvp.Value.Invested]
        for sym in invested_symbols:
            if sym not in self.long_candidates:
                self.EmitInsights(Insight.Price(sym, timedelta(seconds=1), InsightDirection.Flat))
                self.Liquidate(sym)

        for sym in self.short_candidates:
            self.EmitInsights(Insight.Price(sym, Expiry.EndOfWeek, InsightDirection.Down))
            self.SetHoldings(sym, self.short_target_weight)

        for sym in self.long_candidates:
            self.EmitInsights(Insight.Price(sym, Expiry.EndOfWeek, InsightDirection.Up))
            self.SetHoldings(sym, self.long_target_weight)

    def update(self):
        for sym in self.short_candidates:
            self.EmitInsights(Insight.Price(sym, Expiry.EndOfWeek, InsightDirection.Down))
            self.SetHoldings(sym, self.short_update_weight)