# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

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

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Indicators import *
from QuantConnect.Data.Market import TradeBar
import numpy as np
import math

### <summary>
### Basic template algorithm simply initializes the date range and cash. This is a skeleton
### framework you can use for designing an algorithm.
### </summary>
### <meta name="tag" content="using data" />
### <meta name="tag" content="using quantconnect" />
### <meta name="tag" content="trading and orders" />
class TosAlgorithm(QCAlgorithm):
    '''Basic template algorithm simply initializes the date range and cash'''

    def Initialize(self):
        '''Initialise the data and resolution required, as well as the cash and start-end dates for your algorithm. All algorithms must initialized.'''

        self.SetStartDate(2019, 3, 31)  # Set Start Date
        self.SetEndDate(2020, 8, 30)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        self.AddEquity("IBM", Resolution.Daily) ## Subscribe to hourly TradeBars
        self.ibm = RollingWindow[TradeBar](2) # spy trade bar window
        self.ST_Period = 5
        self.ST_Coeff = 1
        self.SetWarmUp(self.ST_Period)
        self.superTrend = SuperTrend('custom', self.ST_Period,self.ST_Coeff) # suepr trend
        self.lastSuperTrendTrend = 0 # the last self.superTrend.Trend
        #self.PlotIndicator("ST",self.superTrend)


    def OnData(self, data):
        self.ibm.Add(data["IBM"])
        if not self.ibm.IsReady:  #  not self.superTrend.IsReady
            return 
        self.superTrend.Update(data["IBM"], self.ibm[1])
        self.Debug(" time:" + str(self.Time) +
                   " open:" + str(self.superTrend.o) +
                   " high:" + str(self.superTrend.h) +
                   " low:" +  str(self.superTrend.l) +
                   " close:" + str(self.superTrend.c) +
                   " ha_close:" + str(round(self.superTrend.ha_close,2)) +
                   " ha_open:"  +  str(round(self.superTrend.ha_open,2)) +
                   " ha_high:"  +  str(round(self.superTrend.ha_high,2)) +
                   " ha_low:"   +  str(round(self.superTrend.ha_low,2)) +
                   " hahl2:"    +  str(round(self.superTrend.hahl2_1,2)) +
                   " tmpUp1:"   +  str(round(self.superTrend.tmpUp1,2)) +
                   " tmpDn1:"   +  str(round(self.superTrend.tmpDn1,2)) +
                   " finalUp1:"   +  str(round(self.superTrend.finalUp1,2)) +
                   " finalDn1:"   +  str(round(self.superTrend.finalDn1,2)) +
                   " atr:"      +  str(self.superTrend.atr) +
                   " super trend:" + str(self.superTrend.trendDir1) 
                   )
        if self.lastSuperTrendTrend == -1 and self.superTrend.trendDir1 == 1:
            if self.Portfolio.Invested: 
                self.Liquidate()
            self.SetHoldings("IBM", 1)
            self.Debug("Long")
        if self.lastSuperTrendTrend == 1 and self.superTrend.trendDir1 == -1:
            if self.Portfolio.Invested: 
                self.Liquidate()
            self.SetHoldings("IBM", -1)
            self.Debug("Short")
        self.lastSuperTrendTrend = self.superTrend.trendDir1

        


class SuperTrend:
    def __init__(self, name, period,coeff):
        self.Name = name
        self.IsReady = False
        self.o = 0
        self.h = 0
        self.l = 0
        self.c = 0
        self.ha_close = 0
        self.ha_open = 0
        self.ha_high = 0
        self.ha_low = 0
        self.hahl2_1 = 0
        self.trendDir1 = 0
        self.tmpUp1 = 0
        self.tmpDn1 = 0
        self.atr = AverageTrueRange(period, MovingAverageType.Wilders) 
        self.test= 0
        self.finalUp1 = 0
        self.finalDn1 = 0
        self.coeff = coeff
        self.qca = QCAlgorithm()
        self.result = 0
       
        
    def Update(self, input, lastBar):
        self.atr.Update(input)
        if self.atr.IsReady: 
            self.test = 0
            self.finalUp1 = self.tmpUp1
            self.finalDn1 = self.tmpDn1

            self.o  =  input.Open         ## Open price
            self.h  =  input.High         ## High price
            self.l  =  input.Low          ## Low price
            self.c  =  input.Close        ## Close price

            self.ha_close = (self.o + self.h + self.l + self.c) / 4.0
            self.ha_open = (lastBar.Open + self.ha_close) / 2.0
            self.ha_high = Math.Max(self.o, Math.Max(self.ha_open, self.ha_close))
            self.ha_low = Math.Min(self.l, Math.Min(self.ha_open, self.ha_close))
            self.hahl2_1 = (self.ha_high  + self.ha_low) / 2
            #hl2 = ((input.High - input.Low) / 2 ) + input.Low
            self.tmpUp1 = self.hahl2_1 + (self.atr.Current.Value * self.coeff)
            self.tmpDn1 = self.hahl2_1 - (self.atr.Current.Value * self.coeff)
             
            if lastBar.Close > self.finalUp1:
                self.finalUp1 = max(self.tmpUp1, self.finalUp1)
            else:
                self.finalUp1 = self.tmpUp1
               
            if lastBar.Close < self.finalDn1:
                self.finalDn1 = min(self.tmpDn1, self.finalDn1)
            else:
                self.finalDn1 = self.tmpDn1

            # calculate trendDir1

            if input.Close > self.finalDn1:
                self.trendDir1 = 1
            elif input.Close < self.finalUp1:
                self.trendDir1 = -1

            if not math.isnan(self.trendDir1):
                self.trendDir = self.trendDir1
            else:
                self.trendDir1 = 1



            #if self.finalUp1 != 0 and input.Close > self.finalUp1:
            #    self.trendDir1 = 1
            #if self.finalDn1 != 0 and input.Close < self.finalDn1:
             #   self.trendDir1 = -1
            self.IsReady = self.trendDir1 != 0