#region imports
from AlgorithmImports import *
#endregion
from datetime import datetime, timedelta
from clr import AddReference
from pytz import timezone
import talib
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 import BaseData, SubscriptionDataSource, SubscriptionDataConfig
from QuantConnect.Python import PythonData
from QuantConnect.Data.Market import TradeBar

class HSI(PythonData):

    def GetSource(self, 
        config: SubscriptionDataConfig, 
        date: datetime, 
        isLive: bool) -> SubscriptionDataSource:

        source = "https://raw.githubusercontent.com/antonymawork/Data/main/HK50_1M_HKT.csv"
        return SubscriptionDataSource(source, SubscriptionTransportMedium.RemoteFile)

    def Reader(self, 
        config: SubscriptionDataConfig, 
        line: str, 
        date: datetime, 
        isLive: bool) -> BaseData:

        # If first character is not digit, pass
        if not (line.strip() and line[0].isdigit()): 
            return None

        data = line.split(',')

        if data[0].lower() == "date": 
            return None

        hsi = HSI() 
        hsi.Symbol = config.Symbol

        hsi.Time = datetime.strptime(data[0], '%Y-%m-%d %H:%M:%S')

        hsi.EndTime = hsi.Time + timedelta(minutes=1)
        hsi.Value = float(data[4])
        if hsi.Value == 0:
            self.Debug("No Value")
            return None
            
        hsi["Open"] = float(data[1])
        hsi["High"] = float(data[2])
        hsi["Low"] = float(data[3])
        hsi["Close"] = hsi.Value
        hsi["Volume"] = float(data[5])

        hsi.Currency = Currencies.HKD

        return hsi


class CustomAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2020, 5, 1)
        self.SetEndDate(2022, 6, 30)
        self.SetAccountCurrency("HKD", 1500000)
        self.initial_margin = 101000
        self.Debug("Range and Capital Set")

        self.symbol = self.AddData(HSI, 'HSI', Resolution.Minute).Symbol
        self.SetBenchmark(self.symbol)
        self.Debug(f"Symbol: {self.symbol}")

        self.ema50 = self.EMA(self.symbol, 55, Resolution.Minute)
        self.ema100 = self.EMA(self.symbol, 105, Resolution.Minute)
        self.stoch = self.STO(self.symbol, 5, 3, Resolution.Minute)
        self.rsi = self.RSI(self.symbol, 5, MovingAverageType.Simple, Resolution.Minute)
        self.atr = self.ATR(self.symbol, 3, MovingAverageType.Simple, Resolution.Minute)

        self.stochWindow = RollingWindow[IndicatorDataPoint](4)

        self.stopMarketTicket = None

        # Warmup for 100 minutes
        self.SetWarmUp(100, Resolution.Minute)

    def OnData(self, slice):
        if self.IsWarmingUp:
            self.Debug("Warmed Up")
            return

        if not slice.ContainsKey(self.symbol): 
            self.Debug(f"Data for {self.symbol} not found in slice")
            return

        if self.symbol not in slice: 
            self.Debug("symbol not in slice")
            self.Debug(f"Symbol: {slice[self.symbol]}")
            return

        bar = slice[self.symbol]
        #self.Debug(f"Open: {bar['Open']}, High: {bar['High']}, Low: {bar['Low']}, Close: {bar.Value}, Volume: {bar['Volume']}")
        
        close = bar.Close
        #self.Debug(f"Close: {close}")

        holdings = self.Portfolio[self.symbol].Quantity
        #self.Debug(f"Holdings: {holdings}")

        self.stochWindow.Add(self.stoch.Current)

        long_condition = (
            close > self.ema50.Current.Value and
            self.ema50.Current.Value > self.ema100.Current.Value and
            self.stoch.Current.Value > 20 and self.stochWindow[1].Value < 20 and
            self.rsi.Current.Value > 50
        )

        short_condition = (
            close < self.ema50.Current.Value and
            self.ema50.Current.Value < self.ema100.Current.Value and
            self.stoch.Current.Value < 80 and self.stochWindow[1].Value > 80 and
            self.rsi.Current.Value < 50     
        )

        quantity_long = 1
        quantity_short = 1

        m1 = self.Portfolio.TotalPortfolioValue > quantity_long * self.initial_margin

        if long_condition and m1:
            self.MarketOrder(self.symbol, quantity_long)
            self.stopMarketTicket = self.StopMarketOrder(self.symbol, -quantity_long, round(close - 2 * self.atr.Current.Value, 2))
            self.LimitOrder(self.symbol, -quantity_long, round(close + 5 * self.atr.Current.Value, 2))

        if short_condition and m1:
            self.MarketOrder(self.symbol, quantity_short)
            self.stopMarketTicket = self.StopMarketOrder(self.symbol, -quantity_short, round(close + 2 * self.atr.Current.Value, 2))
            self.LimitOrder(self.symbol, -quantity_short, round(close - 5 * self.atr.Current.Value, 2))
        
        if self.stopMarketTicket is not None and self.stopMarketTicket.Status == OrderStatus.Filled:
            if holdings > 0 and close - 2 * self.atr.Current.Value > self.stopMarketTicket.Get(OrderField.StopPrice):
                updateOrderFields = UpdateOrderFields()
                updateOrderFields.StopPrice = round(close - 2 * self.atr.Current.Value, 2)
                self.stopMarketTicket.Update(updateOrderFields)
            elif holdings < 0 and close + 2 * self.atr.Current.Value < self.stopMarketTicket.Get(OrderField.StopPrice):
                updateOrderFields = UpdateOrderFields()
                updateOrderFields.StopPrice = round(close + 2 * self.atr.Current.Value, 2)
                self.stopMarketTicket.Update(updateOrderFields)