# region imports
from AlgorithmImports import *

# endregion

class LongShortMLSentimentRankAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 1, 4)
        self.SetEndDate(2023, 3, 1)
        self.SetCash(1_000_000)

        self.AddUniverse(BrainSentimentIndicatorUniverse, "Universe", Resolution.Daily, self.select_universe_symbols)

        self.AddAlpha(LongShortMLSentimentRankAlphaModel())
    
    def select_universe_symbols(self, alt_coarse: List[BrainSentimentIndicatorUniverse]) -> List[Symbol]:
        # QC data starts in 1998. We need to drop Symbols that end in "2T", which have a start date of 1899.
        alt_coarse = [c for c in alt_coarse if c.Sentiment7Days is not None and c.Symbol.ID.Date.year >= 1998] 
        self.Log(f"In Universe at Time {self.Time} // UTC Time {self.UtcTime}. Selected {alt_coarse[0].Symbol} with sentiment {alt_coarse[0].Sentiment7Days}")
        return [alt_coarse[0].Symbol]

class LongShortMLSentimentRankAlphaModel(AlphaModel):

    def Update(self, algorithm: QCAlgorithm, data: Slice) -> List[Insight]:
        for dataset_symbol, data_point in data.Get(BrainSentimentIndicator7Day).items():
            algorithm.Quit(f"In Update at Time {algorithm.Time} // Utc Time {algorithm.UtcTime}. Sentiment {data_point.Sentiment}")

            for bar in algorithm.History[BrainSentimentIndicator7Day](dataset_symbol, 2, Resolution.Daily):
                algorithm.Quit(f"History at UtcTime {bar.Time} / UtcTime EndTime {bar.EndTime}. Sentiment: {bar.Sentiment}")
        return []

    def OnSecuritiesChanged(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        for security in changes.AddedSecurities:
            algorithm.AddData(BrainSentimentIndicator7Day, security.Symbol).Symbol