from AlgorithmImports import *
import numpy as np


class CryingVioletZebra(QCAlgorithm):

    def Initialize(self):
        
        # set up
        self.SetStartDate(2003, 2, 1)
        self.SetEndDate(2022, 1, 1)
        self.SetCash(1000000)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        # parameters
        self.sd_threshold = 0.004       # standard deviation threshold after which we sell
        self.predictions_sd_window = 3  # window size when calculating standard deviation of predictions
        self.safe_asset = "TLT"          # safe asset we buy, when we sell SPY. Chouces: (None, ticker: Inverse SPY: SH, SDS, SPXU, Bonds: TLT, VFSTX , money market FX: FXE)
        # self.thresold = 0.00          # ne koristimo, prvotno zamisljeno za predikcije
        #  self.sma_width = 1           # ne koristimo
        # self.close_window_length = 7  # ne koristimo

        # optimization parameters
        # self.sd_threshold = float(self.GetParameter("threshold-sd"))
        # self.predictions_sd_window = int(self.GetParameter("preds-window"))

        # data feeds
        self.UniverseSettings.Leverage = 4
        PRAData.set_algo(self)
        if self.LiveMode:
            self.spy = self.AddEquity("SPY", Resolution.Minute).Symbol
            self.symbol = self.AddData(PRAData, "PR", Resolution.Minute, TimeZones.Utc).Symbol
        else:
            equity = self.AddEquity("SPY", Resolution.Daily)
            equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
            self.spy = equity.Symbol
            self.SetBenchmark(lambda dt: self.Securities[self.spy].Price)
            self.symbol = self.AddData(PRAData, "PR", Resolution.Daily).Symbol
            if self.safe_asset is not None: 
                self.AddEquity(self.safe_asset, Resolution.Daily)
            # self.pra_sma = self.SMA("PR", self.sma_width, Resolution.Daily)
        
        # init
        self.pred_window = RollingWindow[float](self.predictions_sd_window)
        self.std_window = RollingWindow[float](2)
        self.high_month = 0 
        
        # consolidator
        self.Consolidate("SPY", Calendar.Monthly, self.MonthBarHandler)
        
        # set and warmup rolling close
        # self.close_radf = RollingWindow[float](self.close_window_length)
        # self.volume = RollingWindow[float](8)
        # history = self.History(["SPY"], self.close_window_length, Resolution.Hour)
        # for index, ohlcv in history.loc["SPY"].iterrows():
        #     self.close_radf.Add(ohlcv["close"])
        #     self.volume.Add(ohlcv["volume"])
        
        # define margins
        # self.Securities["SPY"].SetLeverage(4)
        # self.Securities[self.safe_asset].SetLeverage(4)
 
    def MonthBarHandler(self, bar):
        self.high_month = bar.High
 
    def OnData(self, data):

        # extract indicator
        if self.LiveMode:
            # get SPY data
            spy = self.Securities[self.spy]
            o, h, l, c = spy.Open, spy.High, spy.Low, spy.Close
            if c == 0: 
                self.Log("SPY close price is 0.")    
                return
            # check if custom data is ready
            if not data.ContainsKey(self.symbol):
                # self.Log(f'There is no indicator data for the bar.')
                return
            # define radf value and spy close
            pr = data[self.symbol].Value
        else:
            # check if SPY and safe assets are ready
            if not data.Bars.ContainsKey(self.spy) and not data.Bars.ContainsKey(self.spy) or data[self.spy] is None:
                self.Log(f'SPY is not ready')
                return
            if self.safe_asset is not None:
                if not data.Bars.ContainsKey(self.safe_asset)  and not data.Bars.ContainsKey(self.safe_asset) or data[self.safe_asset] is None:
                    self.Log(f'TLT is not ready')
                    return
                
            # check if custom data is ready
            if not data.ContainsKey(self.symbol):
                self.Log(f'There is no indicator data for the bar.')
                return
            pr = data[self.symbol].Value
            ############ SMA ############
            # if not self.pra_sma.IsReady:
            #     self.Log(f'Radf agg SMA is not ready')
            #     return
            # pr = self.pra_sma.Current.Value
            ############ SMA ############
            
        # populate predictins window
        self.pred_window.Add(pr)
        if not self.pred_window.IsReady: return

        # predictions standard deviation
        preds_std = np.std(list(self.pred_window))
        self.std_window.Add(preds_std)
        if not self.pred_window.IsReady: return
        if not self.std_window.IsReady: return

        # growth rate of indicator
        indicator_gw = list(self.std_window)
        indicator_gw = indicator_gw[0] / indicator_gw[1] - 1
        self.Debug(f"Indicator growth rate: {indicator_gw}")

        # current return
        # close_diff = self.close_radf[0] - self.close_radf[self.close_window_length - 1]
        
        # plot indicator
        # self.Plot("Indicators", "PR", pr)
        # self.Plot("Indicators", "Threshold", self.thresold)
        self.Plot("Indicators", "Predictoins sd", preds_std)
        self.Plot("Indicators", "Threshold sd", self.sd_threshold)
        self.Plot("Indicators GR", "SD", indicator_gw)

        ## buy or sell all
        if not self.Portfolio["SPY"].Invested and preds_std <= self.sd_threshold: # pr > self.thresold
            if self.safe_asset is not None:
                if self.Portfolio[self.safe_asset].Invested:
                    self.Liquidate(self.safe_asset)
            self.SetHoldings("SPY", 1)
        elif self.Portfolio["SPY"].Invested and preds_std > self.sd_threshold:
            self.Liquidate("SPY")
            if self.safe_asset is not None:
                self.SetHoldings(self.safe_asset, 1)
        

class PRAData(PythonData):
    
    algo = None
    
    @staticmethod
    def set_algo(algo):
        PRAData.algo = algo
        
    def GetSource(self, config, date, isLive):
        if isLive:
            url = "https://contentiobatch.blob.core.windows.net/qc-live/pra.csv?sp=r&st=2021-09-29T14:42:22Z&se=2025-04-28T22:42:22Z&sv=2020-08-04&sr=b&sig=gJCmQj%2FU6%2FwL2pPEytx6nwtfbWAoKXDLYQ%2FmgeieGGM%3D"
            return SubscriptionDataSource(url, SubscriptionTransportMedium.RemoteFile) 
        else:
            url = "https://contentiobatch.blob.core.windows.net/qc-backtest/systemic_risk.csv"
            return SubscriptionDataSource(url, SubscriptionTransportMedium.RemoteFile) 
        
    def Reader(self, config, line, date, isLive):
        # If first character is not digit, pass
        if not (line.strip() and line[0].isdigit()): return None

        if isLive:
            data = line.split(';') 
            index = PRAData()
            index.Symbol = config.Symbol
            index.EndTime = datetime.utcnow()
            index.Time = datetime.utcnow()
            index.Value = int(data[1])

        else:
            data = line.split(',')
            index = PRAData()
            index.Symbol = config.Symbol
            index.Time = datetime.strptime(data[0], '%Y-%m-%d') #+ timedelta(hours=1) # Make sure we only get this data AFTER trading day - don't want forward bias.
            index.Value = float(data[1])      # pr_below_dummy_176

        return index