from AlgorithmImports import *
import numpy as np

class SymbolIndicatorData():
    def __init__(self, algorithm, symbol) -> None:
        self.symbol = symbol
        self.algorithm = algorithm

        self.spyRollingWindow = RollingWindow[float](21*3)
        self.ema_21_vol = None
        self.sma_63_vol = None
        self.spyEma21VolRollingWindow = RollingWindow[float](21*6)
        self.spySma63VolRollingWindow = RollingWindow[float](21*6)

        warmUpData = self.algorithm.History[TradeBar](
            symbol,
            21*6 + 21*3 +10,
            Resolution.Daily
        )

        for bar in warmUpData:
            self.spyRollingWindow.Add(bar.Close)
            if self.spyRollingWindow.IsReady:
                l = list(self.spyRollingWindow)
                self.ema_21_vol = np.std(l[:21])
                self.sma_63_vol = np.std(l[:21*3])
                self.spyEma21VolRollingWindow.Add(self.ema_21_vol)
                self.spySma63VolRollingWindow.Add(self.sma_63_vol)

    def Update(self, close):
        self.spyRollingWindow.Add(close)
        if self.spyRollingWindow.IsReady:
            l = list(self.spyRollingWindow)
            self.ema_21_vol = np.std(l[:21])
            self.sma_63_vol = np.std(l[:21*3])
            self.spyEma21VolRollingWindow.Add(self.ema_21_vol)
            self.spySma63VolRollingWindow.Add(self.sma_63_vol)

    def Reset(self):
        self.spyRollingWindow.Reset()
        self.ema_21_vol = None
        self.sma_63_vol = None
        self.spyEma21VolRollingWindow.Reset()
        self.spyEma21VolRollingWindow.Reset()

    @property
    def IsReady(self):
        if self.spyEma21VolRollingWindow.IsReady and self.spySma63VolRollingWindow.IsReady:
            return True
        return False

    @property
    def IsHighVol(self):
        if not self.IsReady:
            return 0

        if self.ema_21_vol > np.mean(list(self.spyEma21VolRollingWindow)) and self.sma_63_vol > np.mean(list(self.spySma63VolRollingWindow)):
            return 1
        elif self.ema_21_vol < np.mean(list(self.spyEma21VolRollingWindow)) and self.sma_63_vol < np.mean(list(self.spySma63VolRollingWindow)):
            return -1
        return 0

#region imports
from AlgorithmImports import *
from enum import Enum
from SymbolData import SymbolIndicatorData
'''
1. At every point in time, you can allocate multiplier M of the cushion to the risky assets. Thus, as the cushion decreases, you are reducing the risky asset allocation. If the cushion goes to zero (the value of the portfolio hits the floor), your allocation to the risky assets is the multiplier M times zero. This means you are allocating nothing to the portfolio's risky portion, and 100% is invested in the safe component of your portfolio.
2. Because of this, it is recommended to set the multiplier as a function of the maximum potential loss within a given trading interval.
'''

# Benchmark
BENCHMARK = 'SPY'

# Risky asset
# RISKY_ASSET = 'SPY'
RISKY_ASSET = 'SSO' # 2x SPY ETF

# Risk free asset
RISKFREE_ASSET = 'SHV'  # short term Treasury ETF
# RISKFREE_ASSET = 'SHY' # short term Treasury ETF'
# RISKFREE_ASSET = 'TLT' # long term Treasury ETF'
# RISKFREE_ASSET = 'IEF' # long term Treasury ETF'

class CppiMode(Enum):
    BEHCNMARK = 0
    BASIC = 1
    NEW_HIGH = 2  # TIPP
    LEARNING = 3

class MultiplierMode(Enum):
    FIXED_MULTIPLIER = 0
    DYNAMIC_MULTIPLIER = 1

class BacktestTimePeriod(Enum):
    ALL = 1
    LOW = 2
    HIGH = 3

class OvernightTradeAlgorithm(QCAlgorithm):

    def Initialize(self):
        self._init_cash = 100000
        self.SetCash(self._init_cash)            #Set Strategy Cash
        # self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage) # T+2

        ##########################################
        # Symbol management
        self.benchmark = self.AddEquity(BENCHMARK, Resolution.Daily).Symbol
        self.spy = self.AddEquity(RISKY_ASSET, Resolution.Daily).Symbol
        self.riskfree = self.AddEquity(RISKFREE_ASSET, Resolution.Daily).Symbol

        ##########################################
        # CPPI setup
        self.cppi_mode = CppiMode.LEARNING
        self.mutiplier_mdoe = MultiplierMode.DYNAMIC_MULTIPLIER
        self.scenario = BacktestTimePeriod.LOW

        self.max_asset = self._init_cash
        self.FloorPercentage = 0.8
        self.Floor = self.FloorPercentage * self._init_cash
        self.Cushion = self._init_cash - self.Floor
        self.Multiplier = 3
        self.UpdateUpperBound = 1.1
        self.learning_alpha = 0.2
        # By percentage % of total asset
        self.E = (self.Cushion * self.Multiplier) / self._init_cash
        self.B = (self._init_cash - self.Cushion * self.Multiplier) / self._init_cash

        if self.scenario == BacktestTimePeriod.ALL:
            self.SetStartDate(2010, 1, 1)   #Set Start Date
            self.SetEndDate(2022, 11, 1)     #Set End Date
        elif self.scenario == BacktestTimePeriod.LOW:
            self.SetStartDate(2007, 1, 1)   #Set Start Date
            self.SetEndDate(2012, 1, 1)     #Set End Date
        elif self.scenario == BacktestTimePeriod.HIGH:
            self.SetStartDate(2015, 1, 1)   #Set Start Date
            self.SetEndDate(2019, 1, 1)     #Set End Date

        # Construct indicators
        if self.mutiplier_mdoe == MultiplierMode.DYNAMIC_MULTIPLIER:
            self.dynamicMultiplierIndicator = SymbolIndicatorData(self, self.spy)

        ##########################################
        # Scheduling
        self.Schedule.On(
            self.DateRules.WeekStart(self.benchmark),
            self.TimeRules.At(6, 30),
            self.UpdateParameters
        )

        # Place the order/SetHolding before market open, then order price will be decided by the market open price
        self.Schedule.On(
            self.DateRules.WeekStart(self.benchmark),
            self.TimeRules.At(8, 30),
            self.EveryDayAfterMarketOpen
        )

        ##########################################
        # Charting
        self._benchmark_queue = []
        self._perf_start = None

        self._my_chart = Chart('AvailCash')
        self._my_chart.AddSeries(Series("Cash", SeriesType.Line, 0))
        self.AddChart(self._my_chart)

        self._my_chart2 = Chart('BnE')
        self._my_chart2.AddSeries(Series("B", SeriesType.Line, 0))
        self._my_chart2.AddSeries(Series("E", SeriesType.Line, 0))
        self.AddChart(self._my_chart2)

        self._my_chart3 = Chart('Multiplier')
        self._my_chart3.AddSeries(Series("Multi", SeriesType.Line, 0))
        self.AddChart(self._my_chart3)

    def UpdateParameters(self):
        if self.mutiplier_mdoe == MultiplierMode.DYNAMIC_MULTIPLIER:
            if self.dynamicMultiplierIndicator.IsReady:
                if self.dynamicMultiplierIndicator.IsHighVol > 0:
                    self.Multiplier = 2
                elif self.dynamicMultiplierIndicator.IsHighVol == 0:
                    self.Multiplier = 3
                elif self.dynamicMultiplierIndicator.IsHighVol < 0:
                    self.Multiplier = 4
            else:
                self.Debug(f'The indicator is not ready')
                self.Multiplier = 3
        elif self.mutiplier_mdoe == MultiplierMode.FIXED_MULTIPLIER:
            pass

        cur_total_asset = self.Portfolio.TotalPortfolioValue

        if self.cppi_mode == CppiMode.BASIC:
            pass
        elif self.cppi_mode == CppiMode.NEW_HIGH:
            self.max_asset = cur_total_asset if (cur_total_asset > self.max_asset * self.UpdateUpperBound) else self.max_asset
            self.Floor = self.max_asset * self.FloorPercentage
        elif self.cppi_mode == CppiMode.LEARNING:
            diff = 0
            # if abs(cur_total_asset - self.max_asset * self.UpdateUpperBound) > (self.max_asset * (self.UpdateUpperBound - 1)):
            if abs(cur_total_asset - self.max_asset) > (self.max_asset * (self.UpdateUpperBound - 1)):
                diff = (cur_total_asset - self.max_asset)
            self.max_asset = self.max_asset + diff * self.learning_alpha
            self.Floor = self.max_asset * self.FloorPercentage

        self.Cushion = cur_total_asset - self.Floor
        self.E = (self.Cushion * self.Multiplier) / self.max_asset
        self.E = min(1, self.E)
        self.E = max(0, self.E)
        self.B = 1 - self.E

    def EveryDayAfterMarketOpen(self):
        if self.cppi_mode == CppiMode.BEHCNMARK:
            self.SetHoldings(self.spy, 1)
            self.SetHoldings(self.riskfree, 0)
        else:
            self.SetHoldings(self.spy, self.E)
            self.SetHoldings(self.riskfree, self.B)

    def OnData(self, data):
        if len(data.Bars) <= 0:
            self.Debug(f'{self.Time}: Data is empty')
            return

        if self.mutiplier_mdoe == MultiplierMode.DYNAMIC_MULTIPLIER:
            if data.ContainsKey(self.spy):
                self.dynamicMultiplierIndicator.Update(
                    data[self.spy].Close
                )
            else:
                self.Debug(f'[{self.Time}]: {self.spy} is not found in the slice data!!!!')

    def OnEndOfDay(self, symbol):
        self.Plot('AvailCash', 'Cash', self.Portfolio.Cash)
        self.Plot('AvailCash', 'Portfolio', self.Portfolio.TotalPortfolioValue)

        self.Plot('BnE', 'B', self.B)
        self.Plot('BnE', 'E', self.E)
        self.Plot('Multiplier', 'Multi', self.Multiplier)

        # Plot the benchmark return
        hist = self.History(self.benchmark, 2, Resolution.Daily)['close'].unstack(level= 0).dropna()

        self._benchmark_queue.append(hist[self.benchmark].iloc[-1])
        if self._perf_start is None:
            self._perf_start = hist[self.benchmark].iloc[0]
        spy_perf = self._benchmark_queue[-1] / self._perf_start * self._init_cash
        self.Plot('Strategy Equity', self.benchmark, spy_perf)