# Default Imports START
from QuantConnect import *
from QuantConnect import Orders
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 *
from System.Drawing import Color
QCAlgorithmFramework = QCAlgorithm
QCAlgorithmFrameworkBridge = QCAlgorithm
# Default Imports ENDfrom typing import Dictfrom typing import Dict

from typing import Dict
from System import String

class MyAlgo(QCAlgorithm):
    
    # general
    secs = [ "SPY" ]
    maxPositions = len(secs)

    # indicators
    sma_duration = 200
    rsi_duraction = 14
    
    # buy
    rsi_buy_level = 30
    
    # sell
    rsi_sell_level = 40
    sell_after_days = 14

    def Initialize(self):
        self.SetStartDate(2005, 1, 1)
        self.SetEndDate(2020, 12, 31)
        
        self.SetCash(10000)

        self.smas: Dict[String, SimpleMovingAverage] = {}
        self.rsis: Dict[String, RelativeStrengthIndex] = {}
        self.bought_at: Dict[String, datetime] = {}

        for sec in self.secs:
            self.Debug("Add " + sec + " to view list!")
            self.AddEquity(sec, Resolution.Daily)
            self.smas[sec] = self.SMA(sec, self.sma_duration, Resolution.Daily, Field.Close)
            self.rsis[sec] = self.RSI(sec, self.rsi_duraction, MovingAverageType.Simple , Resolution.Daily, Field.Close)

            chart = Chart(sec)
            chart.AddSeries(Series("Price", SeriesType.Line, 0))
            chart.AddSeries(Series("Buy", SeriesType.Scatter, 0))
            chart.AddSeries(Series("Sell", SeriesType.Scatter, 0))
            chart.AddSeries(Series("SMA", SeriesType.Line, 0))
            chart.AddSeries(Series("RSI", SeriesType.Line, 1))
            chart.AddSeries(Series("RSI Buy Level", SeriesType.Line, 1))
            self.AddChart(chart)
        
        self.SetWarmUp(max(self.sma_duration, self.rsi_duraction))
        
    def OnData(self, data: Slice):
        for sec in self.secs:
            if not data.ContainsKey(sec):
                self.Debug("Data for " + sec + " not found. Skipping for this day!")
                break
            
            if self.IsWarmingUp:
                break
        
            self.handleSells(sec, data)
            self.handleBuys(sec, data)
            
            self.Plot(sec, "Price", data[sec].Close)
            self.Plot(sec, "SMA", self.smas[sec].Current.Value)
            self.Plot(sec, "RSI", self.rsis[sec].Current.Value)
            self.Plot(sec, "RSI Buy Level", self.rsi_buy_level)
        
    def handleSells(self, sec: String, data: Slice):
        if self.Portfolio[sec].Invested and self.isSellSignal(sec, data):
            self.Liquidate(sec)
            currentPrice = data[sec].Close
            self.Plot(sec, "Sell", currentPrice)
    
    def handleBuys(self, sec: String, data: Slice):
        if self.Portfolio[sec].Invested:
            return
        
        if self.isBuySignal(sec, data):
            
            currentPositions = self.countPositions()
            currentCash = self.Portfolio.Cash
            if currentPositions == self.maxPositions:
                return
            orderVolume = currentCash / (self.maxPositions - currentPositions)
            currentPrice = data[sec].Close
            quantity = orderVolume / currentPrice
            
            self.MarketOnOpenOrder(sec, quantity)
            self.bought_at[sec] = self.Time
            self.Plot(sec, "Buy", currentPrice)

    def isBuySignal(self, sec: String, data: Slice):
        if data[sec].Close <= self.smas[sec].Current.Value:
            return False
        return self.rsis[sec].Current.Value < self.rsi_sell_level

    def isSellSignal(self, sec: String, data: Slice):
        return self.rsis[sec].Current.Value >= self.rsi_sell_level or (self.Time - self.bought_at[sec]).days >= self.sell_after_days
        
    def countPositions(self):
        count = 0
        for sec in self.Portfolio:
            if sec.Value.Invested:
                count += 1
        return count