import numpy as np
import math

# T is short for either Type or Mr. T (your choice)
class PositionT(Enum):
    I = 1 # just so we can set our inital enum state to a value
    T1 = 2 # long ZNH short CEA
    T2 = 2 # short ZNH long CEA
    T3 = 3 # liquidated (no holdings)

stock1 = 'ZNH'
stock2 = 'CEA'

class TachyonCalibratedThrustAssembly(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2020, 4, 1)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        
        
        self.UniverseSettings.Resolution = Resolution.Hour
        # fun fact: both of these companies are Chinese Airlines listed on NYSE!
        self.znh = self.AddEquity(stock1)
        self.cea = self.AddEquity(stock2)
        self.LogRatioMovingAverage = SimpleMovingAverage(50)
        self.LogRatioMovingStd = StandardDeviation(50)
        self.PreviousHour = -1
        self.T = PositionT.I
        history = self.History(self.Securities.Keys, 100)
        history = history['open'].unstack(level=0).dropna()
        history = history[(history!=0).all(1)]
    
        for tuple in history.itertuples():
            temp = math.log(tuple[2] / tuple[1])
            self.LogRatioMovingAverage.Update(tuple[0], math.log(tuple[1] / tuple[2]))
            self.LogRatioMovingStd.Update(tuple[0], math.log(tuple[1] / tuple[2]))

    def OnData(self, data):
        
        if data.ContainsKey(stock1) and data.ContainsKey(stock2) and self.PreviousHour != self.Time.hour:
            self.PreviousHour = self.Time.hour
            log_ratio = math.log(data[stock1].Open / data[stock2].Open)
            self.LogRatioMovingStd.Update(self.Time, log_ratio)
            self.LogRatioMovingAverage.Update(self.Time, log_ratio)
            if not self.LogRatioMovingStd.IsReady or not self.LogRatioMovingAverage.IsReady:
                return
            std = self.LogRatioMovingStd.Current.Value
            ma = self.LogRatioMovingAverage.Current.Value
        
            if log_ratio > ma + std and self.T != PositionT.T1:
                self.T = PositionT.T1
                self.SetHoldings(stock1, .1)
                self.SetHoldings(stock2, -.1)
            elif log_ratio < ma + std and self.T != PositionT.T2:
                self.T = PositionT.T2
                self.SetHoldings(stock1, -.1)
                self.SetHoldings(stock2, .1)
            elif self.T == PositionT.T1 and log_ratio < ma:
                self.Liquidate()
                self.T = PositionT.T3
            elif self.T == PositionT.T2 and log_ratio > ma:
                self.Liquidate()
                self.T = PositionT.T3