"""
Use ratings data to look for initiations that have not previously been
covered. If these are found, hold the equity for 3 business days then
sell. 
"""

from AlgorithmImports import *

def add_business_days(from_date, add_days):
    business_days_to_add = add_days
    current_date = from_date
    while business_days_to_add > 0:
        current_date =current_date+ timedelta(days=1)
        weekday = current_date.weekday()
        if weekday >= 5: # sunday = 6
            continue
        business_days_to_add -= 1
    return current_date

class InitiationsWithoutCoverage(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2023, 1, 1)
        self.SetEndDate(2023, 3, 1)
        self.SetCash(10000)
        self.ratings=  self.AddData(Ratings,'RATINGS').Symbol
        self.buy_time = datetime.max
        self.to_sell_date= datetime.max.date()
        self.symbol = None
        self.invested = False
        self.good_rating_current = ['Buy','Outperform','Market Outperform','Strong Buy','hold']
    
    def OnData(self, data: Slice):

        if not self.invested and self.symbol:
            # if not self.symbol in data: return 
            if data.ContainsKey(self.symbol):
                self.Log(f"Setting symbol holdings: {self.symbol}")
                self.SetHoldings(self.symbol, 0.9)
                self.buy_time = self.Time
                self.symbol = None # reset the symbol, so we can look for another
                self.Log(f"Removing symbol to find another: {self.symbol}")
                self.invested = True

        if ((self.Time - self.buy_time).days > 3 ) and self.invested:
            self.Liquidate()
            self.invested = False
            self.Log('liquidated')
        
        if not self.ratings in data: return 

        ratings = data[self.ratings]

        if not ratings.Seen_Before and ratings.Initiation and (ratings.Rating_Current in self.good_rating_current):

            self.Log(f'Ticker:{ratings.Ticker} Initialisation with no previous coverage')
            if not self.invested:
                try:
                    # self.to_sell_date =add_business_days(self.Time.date() ,3)
                    self.symbol= self.AddEquity(ratings.Ticker).Symbol
                    self.Log(f"Chosen symbol: {self.symbol}")
                    
                    # Sell in 3 business days
                except Exception as e:
                    self.Log(e)
                    self.Log(f'didnt work')

class Ratings(PythonData):
    '''Ratings Custom Data Class'''
    def GetSource(self, config: SubscriptionDataConfig, date: datetime, isLiveMode: bool) -> SubscriptionDataSource:
        url = "https://www.dropbox.com/s/qaxr9t3ppxp07ry/ratings.csv?dl=1"
        return SubscriptionDataSource(url, SubscriptionTransportMedium.RemoteFile)

    def Reader(self, config: SubscriptionDataConfig, line: str, date: datetime, isLiveMode: bool) -> BaseData:
        if not (line.strip() and line[0].isdigit()):
            print('No data yet')
            return None
        # New Ratings object
        index = Ratings()
        index.Symbol = config.Symbol
        try:
            # File Format:
            # datetime,             action_company, exchange,   ticker,     analyst,    initiation, seen_before
            # 2022-05-06 04:31:18,  Maintains,      NYSE,       WCC,        Keybanc,    False,      False
            data = line.split(',')
            index.Time = datetime.strptime(data[0], "%Y-%m-%d %H:%M:%S")
            index.EndTime = index.Time
            index.Value = 1 if data[5]=='True' else 0
            index["action_company"] = str(data[1])
            index["ticker"] = str(data[3])
            index["analyst"] = str(data[4])
            index["initiation"] = True if data[5]=='True' else False
            index["seen_before"] = True if data[6]=='True' else False
            index['rating_current'] = data[7]
        except:
            print('e')
        return index