import torch
import numpy as np
import pandas as pd
import gradio as gr
import joblib
from model.ConvGRUTransformerHL import ConvGRUTransformerHL_Attn  # โมเดลของมึง

# ===== Config =====
LOOKBACK = 3
FOLDS = [1,2,3]
MODEL_DIR = "models"  # ปรับ path ตามจริง

# ===== Load scalers =====
scalers = {}
for fold in FOLDS:
    scalers[fold] = {
        "price": joblib.load(f"{MODEL_DIR}/scaler_fold{fold}/scaler_price.pkl"),
        "adx": joblib.load(f"{MODEL_DIR}/scaler_fold{fold}/scaler_adx.pkl"),
        "atr": joblib.load(f"{MODEL_DIR}/scaler_fold{fold}/scaler_atr.pkl"),
        "tick": joblib.load(f"{MODEL_DIR}/scaler_fold{fold}/scaler_tick_volume.pkl"),
        "target_high": joblib.load(f"{MODEL_DIR}/scaler_fold{fold}/scaler_target_high.pkl"),
        "target_low": joblib.load(f"{MODEL_DIR}/scaler_fold{fold}/scaler_target_low.pkl"),
    }

# ===== Load models =====
models = {}
for fold in FOLDS:
    checkpoint = torch.load(f"{MODEL_DIR}/EURUSDm_transformer_finetuned_fold_{fold}_01.pth", map_location="cpu")
    model = ConvGRUTransformerHL_Attn(input_dim=10, seq_len=LOOKBACK, kernel_size=LOOKBACK, output_steps=1)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()
    models[fold] = model

# ===== Prediction function =====
def predict_signal(open_, high_, low_, close_, tick_volume, adx, atr, hour, weekday, month):
    df = pd.DataFrame([{
        'open': open_, 'high': high_, 'low': low_, 'close': close_,
        'tick_volume': tick_volume, 'adx': adx, 'atr': atr,
        'hour': hour, 'weekday': weekday, 'month': month
    }])

    # เตรียม X_scaled
    pred_high_list = []
    pred_low_list = []

    for fold in FOLDS:
        s = scalers[fold]
        price_scaled = s["price"].transform(df[['open','high','low','close']].values)
        adx_scaled   = s["adx"].transform(df[['adx']].values)
        atr_scaled   = s["atr"].transform(df[['atr']].values)
        tick_scaled  = s["tick"].transform(df[['tick_volume']].values)
        time_data    = df[['hour','weekday','month']].values

        X_scaled = np.concatenate([price_scaled, adx_scaled, atr_scaled, tick_scaled, time_data], axis=1)
        X_tensor = torch.tensor(X_scaled, dtype=torch.float32).unsqueeze(0)  # (1, seq_len, features)

        # Predict
        with torch.no_grad():
            pred_high, pred_low = models[fold](X_tensor)
        
        # Inverse transform
        pred_high_np = s["target_high"].inverse_transform(pred_high.numpy())
        pred_low_np  = s["target_low"].inverse_transform(pred_low.numpy())

        pred_high_list.append(pred_high_np[0][0])
        pred_low_list.append(pred_low_np[0][0])

    # เอาค่าเฉลี่ย folds
    return float(np.mean(pred_high_list)), float(np.mean(pred_low_list))

# ===== Gradio UI =====
iface = gr.Interface(
    fn=predict_signal,
    inputs=["number","number","number","number","number","number","number","number","number","number"],  # 10 features
    outputs=["number","number"],
    title="AI Trading Signals",
    description="Enter latest OHLC + tick + adx + atr + hour + weekday + month to get predicted High/Low"
)

iface.launch()