import clr
clr.AddReference("System")
clr.AddReference("QuantConnect.Algorithm")
clr.AddReference("QuantConnect.Indicators")
clr.AddReference("QuantConnect.Common")

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Indicators import *


class QuantumCalibratedComputer(QCAlgorithm):

    def Initialize(self):
        self.SetBenchmark('SPY')
        self.SetStartDate(2020, 1, 1)    #Set Start Date
        self.SetEndDate(2021, 1, 25)      #Set End Date
        self.SetCash(2400)             #Set Strategy Cash
        self.symbols = {
            'a': ['PSQ', 34, 100, 2],
            }
        self.fast, self.slow, self.atr = {}, {}, {}
        for key, symbol in self.symbols.items():
            code = symbol[0]
            self.AddEquity(code)
            # self.AddForex("EURUSD", Resolution.Hour, Market.Oanda)
            slow, fast = 120, 60
            self.fast[code] = self.SMA(code, fast, Resolution.Daily)
            self.slow[code] = self.SMA(code, slow, Resolution.Daily)
            self.atr[code] = self.ATR(code, 15, MovingAverageType.Simple, Resolution.Daily)
            self.PlotIndicator('atr' + code, self.atr[code])
            self.PlotIndicator('fast' + code, self.fast[code])
            self.PlotIndicator('slow' + code, self.slow[code])
            self.SetWarmup(150)
            
        self.Schedule.On(self.DateRules.EveryDay(self.symbols['a'][0]), 
                self.TimeRules.AfterMarketOpen(self.symbols['a'][0], 10), self.SetSignal_a)
            


    def SetSignal_a(self):
        self.SetSignal(self.symbols['a'])
        
    def SetSignal_b(self):
        self.SetSignal(self.symbols['b'])

    def SetSignal(self, symbol):
        code = symbol[0]
        pos_per_time = symbol[1]
        max_pos = symbol[2]
        atr_number = symbol[3]
        breakout = 5
        opens = self.History(self.Symbol(code), breakout, Resolution.Daily)['open']
        
        if not self.slow[code].IsReady or not self.fast[code].IsReady or not self.atr[code].IsReady:
            # self.Log('indicators is not ready')
            return

        p = self.Securities[code].Price
        ma_slow = self.slow[code].Current.Value
        ma_fast = self.fast[code].Current.Value
        atr = self.atr[code].Current.Value
        highest = max(opens)
        lowest = min(opens)
        
        pos = self.Portfolio[code].Quantity
        avg_open = self.Portfolio[code].AveragePrice
        pl = 0
        if avg_open > 0:
            if pos > 0:
                pl = (p - avg_open) / avg_open
            if pos < 0:
                pl = (avg_open - p) / avg_open
        
        self.Log(code 
            + ' price ' + self.f(p) 
            + ' position ' + self.f(pos)
            + ' avg_open ' + self.f(avg_open)
            + ' atr ' + self.f(atr) 
            + ' ma_slow ' + self.f(ma_slow)
            + ' ma_fast ' + self.f(ma_fast)
            + ' highest ' + self.f(highest)
            + ' lowest ' + self.f(lowest)
            )
        
        is_long, is_short = pos > 0, pos < 0
        
        if is_long:
            stop_loss_p = avg_open - atr * atr_number
            if p <= stop_loss_p:
                self.Liquidate(code)
                is_long = False
        if is_short:
            stop_loss_p = avg_open + atr * atr_number
            if p >= stop_loss_p:
                self.Liquidate(code)
                is_short = False
                pl = -1
        
        if is_long and (p < ma_fast or p < ma_slow):
            self.Liquidate(code)
            is_long = False
            pl = -1
        if is_short and (p > ma_fast or p > ma_slow):
            self.Liquidate(code)
            is_short = False
            pl = -1
        
        if (p >= ma_slow and p >= ma_fast and p >= highest 
            and abs(pos) < max_pos and pl >= 0):
            self.MarketOrder(code, pos_per_time)
            stop_loss_p = p - atr * atr_number
            self.StopMarketOrder(code, pos_per_time * -1, stop_loss_p)
        if (p <= ma_slow and p <= ma_fast and p <= lowest 
            and abs(pos) < max_pos and pl >= 0):
            self.MarketOrder(code, pos_per_time * -1)
            stop_loss_p = p + atr * atr_number
            self.StopMarketOrder(code, pos_per_time, stop_loss_p)

        
    
    def OnEndOfDay(self, data):
        pass
        
    def f(self, f):
        if f is None:
            return 'None'
        if type(f) != float:
            f = float(f)
        s = str('{:.6f}'.format(f))
        if s.endswith('.000000'):
            s = s[:-7]
        return s