from QuantConnect import *
from QuantConnect.Parameters import *
from QuantConnect.Benchmarks import *
from QuantConnect.Brokerages import *
from QuantConnect.Util import *
from QuantConnect.Interfaces import *
from QuantConnect.Algorithm import *
from QuantConnect.Algorithm.Framework import *
from QuantConnect.Algorithm.Framework.Selection import *
from QuantConnect.Algorithm.Framework.Alphas import *
from QuantConnect.Algorithm.Framework.Portfolio import *
from QuantConnect.Algorithm.Framework.Execution import *
from QuantConnect.Algorithm.Framework.Risk import *
from QuantConnect.Indicators import *
from QuantConnect.Data import *
from QuantConnect.Data.Consolidators import *
from QuantConnect.Data.Custom import *
from QuantConnect.Data.Fundamental import *
from QuantConnect.Data.Market import *
from QuantConnect.Data.UniverseSelection import *
from QuantConnect.Notifications import *
from QuantConnect.Orders import *
from QuantConnect.Orders.Fees import *
from QuantConnect.Orders.Fills import *
from QuantConnect.Orders.Slippage import *
from QuantConnect.Scheduling import *
from QuantConnect.Securities import *
from QuantConnect.Securities.Equity import *
from QuantConnect.Securities.Forex import *
from QuantConnect.Securities.Interfaces import *
from datetime import date, datetime, timedelta
from QuantConnect.Python import *
from QuantConnect.Storage import *
QCAlgorithmFramework = QCAlgorithm
QCAlgorithmFrameworkBridge = QCAlgorithm
import math
import numpy as np
import pandas as pd
import scipy as sp










class RBreaker(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 3, 1)  # Set Start Date
        self.SetEndDate(2020, 3, 2)
        self.SetCash(100000)  # Set Strategy Cash
        self.spy = self.AddEquity("SPY", Resolution.Hour)
        self.spy.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        self.SetWarmup(2)
        self.R_breaker_dict = {}
        self.weight = 1.0
        
    def OnData(self, data):
        self.Debug('ENTER ONDATA')
        if not self.IsWarmingUp:
            self.Debug('IS READY')
            for symbol in self.R_breaker_dict:
                sym_str = symbol.Value
                self.Debug(f'ENTER {sym_str}')
                if data.ContainsKey(sym_str):
                    self.Debug(f'{sym_str} IS IN DICT')
                    resist_support = self.R_breaker_dict[symbol]
                    
                    r3 = resist_support.r3
                    r2 = resist_support.r2
                    r1 = resist_support.r1
                    s1 = resist_support.s1
                    s2 = resist_support.s2
                    s3 = resist_support.s3
                    
                    curr_high = data[sym_str].High
                    curr_close = data[sym_str].Close
                    curr_low = data[sym_str].Low
                    
                    if curr_close >= r3:
                        self.SetHoldings(symbol,self.weight)
                    elif curr_close <= s3:
                        self.SetHoldings(symbol,-self.weight)
                    elif self.Portfolio[sym_str].Quantity > 0 and curr_high >= r2 and curr_close <= r1 :
                        self.SetHoldings(symbol,-self.weight)
                    elif self.Portfolio[sym_str].Quantity < 0 and curr_low <= s2 and curr_close >= s1 :
                        self.SetHoldings(symbol,self.weight)
                    
                    self.R_breaker_dict[symbol] = Resist_Support([curr_high,curr_close,curr_low])
                else:
                    self.Debug(f'{sym_str} NOT IN DICT')
                    curr_high = data[sym_str].High
                    curr_close = data[sym_str].Close
                    curr_low = data[sym_str].Low
                    self.R_breaker_dict[symbol] = Resist_Support([curr_high,curr_close,curr_low])
                self.Debug(f'{sym_str} ENDED')
        else:
            self.Debug('IS WARMING UP')
            symbol = self.spy.Symbol
            sym_str = "SPY"
            curr_high = data[sym_str].High
            curr_close = data[sym_str].Close
            curr_low = data[sym_str].Low
            self.Debug(f'curr_high: {str(type(curr_high))}')
            self.Debug(f'curr_high: {str(type(curr_close))}')
            self.Debug(f'curr_high: {str(type(curr_low))}')
            self.R_breaker_dict[symbol] = Resist_Support([curr_high,curr_close,curr_low])
        self.Debug('===== END ===== \n\n\n')


class Resist_Support():
        
    def __init__(self, price_list):
        hi = price_list[0]
        cl = price_list[1]
        lo = price_list[2]
        self.pivot = np.mean(price_list)
        
        self.r3 = hi + 2*(self.pivot-lo)
        self.r2 = self.pivot + hi-lo
        self.r1 = 2*(self.pivot-lo)
        self.s1 = 2*self.pivot-hi
        self.s2 = self.pivot-(hi-lo)
        self.s3 = lo-2*(hi-self.pivot)