import backtrader as bt

class SizeInvStrategy(bt.Strategy):
    def __init__(self):
        self.market_caps = {}
        self.total_assets = {}
        self.investments = {}

    def next(self):
        # Check if June
        if self.datetime.date().month != 6:
            return

        # Allocate stocks to Size and Investment groups
        self.market_caps.clear()
        self.total_assets.clear()
        self.investments.clear()

        for data in self.datas:
            symbol = data._name
            self.market_caps[symbol] = data.market_cap[0]
            self.total_assets[symbol] = data.total_assets[0]
            self.investments[symbol] = (data.total_assets[0] - data.total_assets[-252]) / data.total_assets[-252]

        size_sorted = sorted(self.market_caps.items(), key=lambda x: x[1])
        inv_sorted = sorted(self.investments.items(), key=lambda x: x[1])

        size_groups = [size_sorted[i::5] for i in range(5)]
        inv_groups = [inv_sorted[i::5] for i in range(5)]

        # Find long and short positions
        long_position = size_groups[-1][0]
        short_position = size_groups[-1][-1]

        # Close existing positions
        for data in self.datas:
            if self.getposition(data).size > 0:
                self.sell(data)
            elif self.getposition(data).size < 0:
                self.buy(data)

        # Open new long and short positions
        self.buy(data=self.getdatabyname(long_position), target=0.5)
        self.sell(data=self.getdatabyname(short_position), target=0.5)


# Load data, set up cerebro, and run the strategy
cerebro = bt.Cerebro()

# Add your data feed here using the 'cerebro.adddata()' method

cerebro.addstrategy(SizeInvStrategy)

results = cerebro.run()