import backtrader as bt
import torch
import numpy as np
import pandas as pd
from ml.models.forex_mlp import ForexMLP

class PyTorchAIModel(bt.Strategy):
    def __init__(self):
        self.model = ForexMLP()
        self.model.load_state_dict(torch.load("ml/models/forex_mlp.pt", map_location=torch.device("cpu")))
        self.model.eval()

        self.buy_threshold = 0.7
        self.sell_threshold = 0.6

    def next(self):
        # Skip early bars (for indicators if you add them)
        if len(self.datas[0]) < 30:
            return

        # Create feature vector for the current candle
        features = np.array([[
            self.data.open[0],
            self.data.high[0],
            self.data.low[0],
            self.data.close[0],
            self.data.volume[0]
        ]], dtype=np.float32)

        inputs = torch.tensor(features)
        with torch.no_grad():
            output = self.model(inputs)
            buy_score, sell_score = output[0].numpy()

        print(f"[AI] Buy: {buy_score:.3f}, Sell: {sell_score:.3f}")

        # Trade logic
        if buy_score > self.buy_threshold and not self.position:
            self.buy()
        elif sell_score > self.sell_threshold and self.position:
            self.close()