import torch
import torch.nn as nn
from fastapi import FastAPI
import numpy as np

# Định nghĩa lại mô hình
class TransformerModel(nn.Module):
    def __init__(self, input_dim, d_model=64, nhead=4, num_layers=2):
        super(TransformerModel, self).__init__()
        self.input_fc = nn.Linear(input_dim, d_model)
        self.transformer = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(d_model, nhead, batch_first=True), num_layers
        )
        self.fc_signal = nn.Linear(d_model, 2)
        self.fc_tp = nn.Linear(d_model, 1)

    def forward(self, x):
        x = self.input_fc(x)
        x = self.transformer(x)
        signal = torch.softmax(self.fc_signal(x[:, -1, :]), dim=-1)
        tp = self.fc_tp(x[:, -1, :])
        return signal, tp

# Khởi tạo mô hình
input_dim = 7
model = TransformerModel(input_dim)
model.load_state_dict(torch.load("tradingbot_model.pth"))
model.eval()

app = FastAPI()

@app.post("/predict")
async def predict(inputs: list):
    inputs = torch.FloatTensor(inputs).unsqueeze(1)  # Thêm chiều batch
    with torch.no_grad():
        signal_prob, tp_pred = model(inputs)
        signal = torch.argmax(signal_prob).item()
        tp = tp_pred.item()
    return {"signal": signal, "tp": tp}