# region imports
from AlgorithmImports import *
from datetime import datetime, timedelta
from dataclasses import dataclass
import csv
from io import StringIO
# endregion

@dataclass
class Signal:
    Number: int
    Ticker: str
    startDate: str
    endDate: str
    startPrice: float
    endPrice: float

class StockTradingAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetCash(100000)
        self.current_index = 0
        self.active_stocks = {}
        
        self.number_of_stocks = int(
            self.GetParameter("number_of_stocks", "7")
        )
        
        self.holding_period = int(
            self.GetParameter("holding_period", "7")
        )
        
        self.frequency = int(
            self.GetParameter("frequency", "4")
        )
        
        self.dropbox_url =str(
            # 5 stocks - FZR5_NOOTC_1W_1YR
            # self.GetParameter("dropbox_url", "https://www.dropbox.com/scl/fi/07loovahjlaxmuujg41i7/FZR5_NOOTC_1W_1YR.csv?rlkey=cl7mm9h7usfashr7xjdd1iim3&st=n012ed20&dl=1")
            # 7 stocks - VM1_NOOTC_1W_1YR
            self.GetParameter("dropbox_url", "https://www.dropbox.com/scl/fi/323i4t3lq5u6524flyb76/VM1_NOOTC_1W_1YR.csv?rlkey=2jimy0z492xq6przr4540arr3&st=w0v3ky1s&dl=1")
            # 3 stocks - BMZ_NOOTC_1W_1YR
            # self.GetParameter("dropbox_url", "https://www.dropbox.com/scl/fi/ax36xcsx8bcf1zon6rnzz/BMZ_NOOTC_1W_1YR.csv?rlkey=4xhilujlygqm8dkpg3t9q25cl&st=yqlbitlx&dl=1")
        )
        
        self.signals = self.GetAllSignals(self.dropbox_url)

        self.SetStartDate(self.signals[0].startDate - timedelta(days = 10))
        self.SetEndDate(self.signals[len(self.signals) -1].endDate + timedelta(days = 10))
        self.next_rotation_date = self.signals[0].startDate - timedelta(days = 1)
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(9, 45), self.CheckForRotation)
        
    def GetAllSignals(self, dropbox_url) -> List[Signal]:
        data = self.Download(dropbox_url)
        signals= []
        if data:
            csv_data = StringIO(data)
            reader = csv.DictReader(csv_data)
            signals = [Signal(
                                Number=int(row['Number']),
                                Ticker=row['Ticker'],
                                startDate=datetime.strptime(row['Start Date'], "%m/%d/%Y"),
                                endDate=datetime.strptime(row['End Date'], "%m/%d/%Y"),
                                startPrice=float(row['Start Price']),
                                endPrice = float(row['End Price'])
                            ) for row in reader]
            return signals
    
    def CheckForRotation(self):
        if self.Time >= self.next_rotation_date:
            self.RotateStocks()
            self.next_rotation_date+=timedelta(days=self.holding_period)
    
    def RotateStocks(self):
        if self.current_index >= len(self.signals):
            for ticker in list(self.active_stocks):
                self.Schedule.On(self.DateRules.On(self.active_stocks[ticker]['endDate']), self.TimeRules.BeforeMarketClose(ticker, 1), self.ActionOnEnd(ticker, self.active_stocks[ticker]['endPrice']))
                del self.active_stocks[ticker]
            self.Debug("All signals processed")
            return

        # Extract the next batch of 3/7 stocks
        next_stocks_data = self.signals[self.current_index:self.current_index + self.number_of_stocks]
        self.current_index += self.number_of_stocks
        next_tickers = {stock.Ticker for stock in next_stocks_data}

        # Liquidate stocks not in the new list
        for ticker in list(self.active_stocks):
            if ticker not in next_tickers:
                self.Schedule.On(self.DateRules.On(self.active_stocks[ticker]['endDate']), self.TimeRules.BeforeMarketClose(ticker, 1), self.ActionOnEnd(ticker,self.active_stocks[ticker]['endPrice']))
                del self.active_stocks[ticker]

        # Initialize or continue trading the new/continuing stocks
        for signal in next_stocks_data:
            if signal.Ticker not in self.active_stocks:
                self.AddEquity(signal.Ticker, self.frequency)
                
                # we buy a minute later than we sell to make sure there is a balance
                self.Schedule.On(self.DateRules.On(signal.startDate), self.TimeRules.BeforeMarketClose(signal.Ticker, 0), self.ActionOnStart(signal))
            self.active_stocks[signal.Ticker] = {
                'endDate': signal.endDate,
                'endPrice': signal.endPrice
            }

    # Buy function - execute open position
    def ActionOnStart(self, signal):
        def Action():
            # if (self.active_stocks):
            self.SetHoldings(signal.Ticker, 1.0 / len(self.active_stocks), False, f"Open ${signal.startPrice}")
            self.Debug(f"Bought {signal.Ticker} on {signal.startDate}")
        return Action
    
    # Sell function
    def ActionOnEnd(self, ticker, endPrice):
        def Action():
            self.Liquidate(ticker, f"Close ${endPrice}")
            self.Debug(f"Sold {ticker}")
        return Action
    
    def OnData(self, data):
        pass