#region imports
from AlgorithmImports import *
import numpy as np
import pandas as pd
#endregion


class SquareBrownBarracuda(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 1)
        #self.SetEndDate(2023, 4, 30)
        self.SetCash(100000)

        symbols = ['MMM', 'ABT', 'ABBV', 'ABMD', 'ACN', 'ATVI', 'ADBE', 'AAP', 'AMD', 'AES', 'AFL', 'A', 'APD', 'AKAM', 'ALK', 'ALB', 'ARE', 'ALGN', 'ALLE', 'LNT', 'ALL', 'GOOGL', 'GOOG', 'MO', 'AMZN', 'AMCR', 'AEE', 'AAL', 'AEP', 'AXP', 'AIG', 'AMT', 'AWK', 'AMP', 'ABC', 'AME', 'AMGN', 'APH', 'ADI', 'ANSS', 'ANTM', 'AON', 'AOS', 'APA', 'AAPL', 'AMAT', 'APTV', 'ADM', 'ANET', 'AJG', 'AIZ', 'ATO', 'T', 'ADSK']
        self.equities = []
        self.rankings = []

        for symbol in symbols:
            equity = self.AddEquity(symbol, Resolution.Daily)
            self.equities.append(equity)

        self.CustomIndicator = {}

        for equity in self.equities:
            self.CustomIndicator[equity.Symbol] = self.SMA(equity.Symbol, 50)

    def OnData(self, data):
        if self.IsWarmingUp:
            return

        for equity in self.equities:
            if self.CustomIndicator[equity.Symbol].Current.Value is None:
                return

        rankings = []

        for equity in self.equities:
            history = self.History(equity.Symbol, 90, Resolution.Daily)

            if history.empty:
                return

            z_score = (history['close'] - history['close'].rolling(window=30).mean()) / history['close'].rolling(window=30).std()

            history['z_score'] = z_score
            history = history.tail(90)

            avg_z_score = history['z_score'].mean()

            rankings.append(pd.DataFrame({'symbol': equity.Symbol, 'avg_z_score': avg_z_score}, index=[0]))

        rankings_df = pd.concat(rankings, ignore_index=True)

        rankings_df['rank'] = rankings_df['avg_z_score'].rank(method='dense')
        rankings_df['avg_rank'] = rankings_df.groupby('symbol')['rank'].transform(lambda x: x.mean())

        rankings_df = rankings_df.sort_values('avg_rank')
        top_stocks = rankings_df.head(3)['symbol'].tolist()
        print('Top 3 stocks:', top_stocks)

        for equity in self.equities:
            if equity.Symbol in top_stocks:
                self.SetHoldings(equity.Symbol, 1/3)
            else:
                self.SetHoldings(equity.Symbol, 0)