import backtrader as bt
from chronos import ChronosPipeline ①
import torch

class Agent(bt.Strategy):
    params = (
        ('lookback', 20),
        ('model_name', 'amazon/chronos-t5-small'), ②
    )

    def __init__(self):
        self.forecaster = ChronosPipeline.from_pretrained(
            self.params.model_name,
            device_map="auto",
            torch_dtype=torch.float32
        )

    def prenext(self):
        self.next()

    def next(self):
        for d in self.datas:
            if len(d) < self.params.lookback:
                continue

            prices = list(d.close.get(size=self.params.lookback)) ③
            context = torch.tensor(prices).unsqueeze(0)

            forecast = self.forecaster.predict(context, prediction_length=1)
            predicted = forecast[0].mean().item()

            pos = self.getposition(d).size ④
            if not pos and predicted > d.close[0]:
                size = int(self.broker.getcash() * 0.15 / d.close[0])
                self.buy(data=d, size=size)
            elif pos and predicted < d.close[0]:
                self.close(data=d)