# region imports
from AlgorithmImports import *
import datetime
from datetime import timedelta

# endregion


class Yahoodata(QCAlgorithm):
    def Initialize(self):
        # Locally Lean installs free sample data, to download more data please visit https://www.quantconnect.com/docs/v2/lean-cli/datasets/downloading-data

        self.SetStartDate(2023, 7, 1)  # Set Start Date
        self.SetEndDate(2023, 7, 19)  # Set End Date
        self.SetCash(100000)  # Set Strategy Cash
        self.SetBenchmark("QQQ")

        self.symbol = self.AddEquity(
            "QQQ",
            Resolution.Minute,
            extendedMarketHours=True,
            dataNormalizationMode=DataNormalizationMode.Raw,
        ).Symbol
        security = self.AddEquity("QQQ", Resolution.Daily, dataNormalizationMode=DataNormalizationMode.Raw)
        security.SetFillModel(EquityFillModel())

        self.SetWarmUp(timedelta(weeks=5))

        self.SetBrokerageModel(BrokerageName.TDAmeritrade, AccountType.Cash)

        self.sma5_week = SimpleMovingAverage(5)
        self.sma5_day = SimpleMovingAverage(5)
        self.sma5_15m = SimpleMovingAverage(5)

        self.consolidator = TradeBarConsolidator(Calendar.Weekly)
        self.consolidator.DataConsolidated += self.OnWeeklyData

        self.Consolidate(self.symbol, timedelta(minutes=15), self.OnFifteenMinuteData)

        self.quantity = 0
        self.max_shot_num = 1
        self.today_shot_num = 0

        self.today_open_order = False
        self.today_stop_order = None

        self.dailyBars = {}

        self.last_week_close = None
        self.last_week_sma5 = None

        self.last_day_close = None
        self.last_day_sma5 = None

        self.window = RollingWindow[TradeBar](4)

        self.last_15m_sma5 = 0

        self.this_week_long = False
        self.this_week_short = False

        self.today_long = False
        self.fourhour_long = False

        self.fourhour_color = None

    def OnWeeklyData(self, sender, tradeBar):
        self.last_week_close = tradeBar.Close

        self.sma5_week.Update(tradeBar.EndTime, tradeBar.Close)
        if self.sma5_week.IsReady:
            self.last_week_sma5 = self.sma5_week.Current.Value
            if self.last_week_close >= self.last_week_sma5:
                self.this_week_long = True
            else:
                self.this_week_long = False

    def OnDailyData(self, tradeBar):
        self.last_day_close = tradeBar.Close
        self.sma5_day.Update(tradeBar.EndTime, tradeBar.Close)
        if self.sma5_day.IsReady:
            self.last_day_sma5 = self.sma5_day.Current.Value
            if tradeBar.Close > tradeBar.Open and self.last_day_close >= self.last_day_sma5:
                self.today_long = True
            else:
                self.today_long = False

    def OnFifteenMinuteData(self, tradeBar):
        self.last_15m_close = tradeBar.Close
        self.last_15m_low = tradeBar.Low

        self.sma5_15m.Update(tradeBar.EndTime, tradeBar.Close)
        if self.sma5_15m.IsReady:
            self.last_15m_sma5 = self.sma5_15m.Current.Value

    def OnWarmupFinished(self) -> None:
        self.quantity = self.CalculateOrderQuantity("QQQ", 0.1)
        self.Log(f"*********  Warm up Finished and one shot quantity: {self.quantity} **********")

    def OnOrderEvent(self, orderEvent: OrderEvent) -> None:
        order_tickets = self.Transactions.GetOrderTickets().ToList()
        open_order_tickets = self.Transactions.GetOpenOrderTickets().ToList()
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        if orderEvent.Status == OrderStatus.Filled:
            if order.Type == 2:
                self.Log(f"OnOrderEvent {self.Time}: {order.Type}: StopPrice:{orderEvent.StopPrice}")
            else:
                self.Log(f"OnOrderEvent {self.Time}: {order.Type}: FillPrice:{orderEvent.FillPrice}")

    def OnEndOfDay(self, symbol):
        # reset
        self.today_shot_num = 0
        self.today_open_order = False
        self.today_stop_order = None

        self.today_long = False

    def OnData(self, data: Slice):
        """OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
        Arguments:
            data: Slice object keyed by symbol containing the stock data
        """

        # last trading day 00:00:00 receive daily data.
        # this week first trading day pre market opened, emit a weekly bar.
        if self.Time.time() == datetime.time(4, 1, 0):
            self.consolidator.Scan(self.Time)

        # 00:00:00 receive daily data and update weekly consolidator
        if self.Time.time() == datetime.time(0, 0, 0) and data.Bars.ContainsKey(self.symbol):
            dailybar = data.Bars[self.symbol]
            self.OnDailyData(dailybar)
            self.consolidator.Update(dailybar)

        # In OnData: Don't run if the indicators aren't ready
        if self.IsWarmingUp:
            return

        if self.IsMarketOpen(self.symbol):
            assert self.sma5_15m.IsReady

            if self.today_stop_order != None:
                # if today stop order filled
                # not filled then return
                if self.today_stop_order.Status == 3:
                    self.today_stop_order = None
                else:
                    return
            # reach max shot num, stop make order
            if self.today_shot_num >= self.max_shot_num:
                return

            if self.this_week_long:
                # check is able to buy
                quantity = self.CalculateOrderQuantity("QQQ", 1)
                if quantity > self.quantity:
                    # if self.today_long and self.fourhour_long and self.Time.time() >= datetime.time(9, 45, 0):
                    if self.today_long and self.Time.time() >= datetime.time(9, 45, 0):
                        # check is a buy point
                        if self.last_15m_close > self.last_15m_sma5:
                            orderticket = self.MarketOrder("QQQ", self.quantity)
                            self.today_stop_order = self.StopMarketOrder("QQQ", -self.quantity, self.last_15m_low)
                            self.today_shot_num += 1
            else:
                order_id = self.Transactions.LastOrderId
                ticket = self.Transactions.GetOrderTicket(order_id)
                # ticket_time = ticket.get_Time()
                # order_events = ticket.OrderEvents()
                order_tickets = self.Transactions.GetOrderTickets()

                # sell all
                if self.Portfolio["QQQ"].Quantity > 0:
                    orderticket = self.MarketOrder("QQQ", -self.Portfolio["QQQ"].Quantity)

import pandas as pd
from pathlib import Path
from datetime import timedelta
from yahoo_fin.stock_info import get_data

data_folder = '../data'
yahoo_dir = './yahoo'

def get_yahoo_ticker(ticker, folder, start_date, end_date):

    # check if the ticker file exists
    fname = ticker.lower() + '.csv'
    path = Path(folder)/fname

    # get the dates if the file aready exists
    if path.exists():
        # open the file and get the dates
        df = pd.read_csv(path, index_col=0)
        dates = pd.DatetimeIndex(df.index.sort_values(ascending=True))

    else:
        dates = None

    start_date = pd.to_datetime(start_date)
    end_date = pd.to_datetime(end_date)

    # if the range is not included in the file or if there is no file at all
    if dates is None or start_date < dates[0] or end_date > dates[-1]:

        # try to retrieve the data from Yahoo_fin
        try:
            delta = timedelta(days=3)
            df = get_data(ticker, start_date=start_date-delta, end_date=end_date+delta)
            df.to_csv(path)
            print(f'Retrieving ticker: {ticker}')

        except BaseException as e:
            print(f'Problem getting ticker {ticker}')
            return None

    else:
        print(f'Ticker {ticker} already loaded')

    return ticker

def get_yahoo_data(tickers: list, start_date, end_date):
    """Get a list of tickers from yahoo and save them in the Default LEAN data directory"""
    # transform tickers into a list (if it is not)
    tickers = tickers if isinstance(tickers, list) else [tickers]

    # check the directory
    folder = Path(data_folder)/yahoo_dir

    if not folder.exists():
        folder.mkdir()
        print(f'Folder {str(folder)} - Created')

    else:
        print(f'Folder {str(folder)} - Ok')

    # create a list to store all loaded tickers
    loaded_tickers = []
    for ticker in tickers:
        loaded_tickers.append(get_yahoo_ticker(ticker, folder, start_date, end_date))

    return [ticker for ticker in loaded_tickers if ticker is not None]
  

from AlgorithmImports import *
from pathlib import Path
from datetime import datetime, timedelta

class YahooData(PythonData):
    def GetSource(self, config, date, isLiveMode):
        # print(f'GetSource YAHOO for date {date}')

        # The name of the asset is the symbol in lowercase .csv (ex. spy.csv)
        fname = config.Symbol.Value.lower() + '.csv'

        # The source folder depends on the directory initialized in lean-cli
        # https://www.quantconnect.com/docs/v2/lean-cli/tutorials/local-data/importing-custom-data
        source = Path(Globals.DataFolder)/'yahoo'/fname

        # The subscription method is LocalFile in this case
        return SubscriptionDataSource(source.as_posix(), SubscriptionTransportMedium.LocalFile)

    def Reader(self, config, line, date, isLiveMode):

        # print(f'Reading date {date}')
        # print(f'line ==> {line}')

        equity = YahooData()
        equity.Symbol = config.Symbol

        # Parse the Line from the Yahoo CSV
        try:
            data = line.split(',')

            # If value is zero, return None
            value = data[4]
            if value == 0: return None

            equity.Time = datetime.strptime(data[0], "%Y-%m-%d")
            equity.EndTime = equity.Time + timedelta(days=1)
            equity.Value = value
            equity["Open"] = float(data[1])
            equity["High"] = float(data[2])
            equity["Low"] = float(data[3])
            equity["Close"] = float(data[4])
            equity["AdjClose"] = float(data[5])
            equity["VolumeUSD"] = float(data[6])

            # print(f'Returning --> {coin.EndTime} - {coin}')
            return equity

        except ValueError:
            # Do nothing, possible error in csv decoding
            return None