# /// script
# dependencies = ["torch", "transformers", "accelerate", "pandas", "numpy", "datasets", "tqdm"]
# ///

import os
import json
import pandas as pd
import numpy as np
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from datasets import load_dataset
from tqdm import tqdm

# --- CONFIG ---
MODEL_ID = "AbdelrehmanFouad/t5-efficient-base-usdjpy-forecaster"
SPREAD = 0.015  # 1.5 pips average spread for USDJPY
HORIZON = 24    # hours

def load_backtest_data():
    print("Loading data for backtest...")
    ds_jpy = load_dataset("huggingXG/forex_USDJPY", split="train", streaming=True)
    data = []
    # Collect enough to have history + test slice
    for i, row in enumerate(ds_jpy):
        data.append(row)
        if i >= 3000000: break
        
    df = pd.DataFrame(data)
    df['timestamp'] = pd.to_datetime(df['timestamp'], format='ISO8601')
    df.set_index('timestamp', inplace=True)
    
    df['mid'] = (df['ask'] + df['bid']) / 2
    resampled = df['mid'].resample('1h').ohlc()
    resampled.dropna(inplace=True)
    
    ds_cal = load_dataset("Ehsanrs2/Forex_Factory_Calendar", split="train")
    df_cal = ds_cal.to_pandas()
    df_cal['DateTime'] = pd.to_datetime(df_cal['DateTime'], utc=True)
    
    split_idx = int(len(resampled) * 0.8)
    test_resampled = resampled.iloc[split_idx - 168:]
    
    return test_resampled, df_cal

def run_backtest():
    test_data, df_cal = load_backtest_data()
    
    print("Loading model...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID, device_map="auto")
    model.eval()
    
    context_size = 168
    results = []
    
    if test_data.index.tz is None:
        test_data.index = test_data.index.tz_localize('UTC')
    
    indices = range(context_size, len(test_data) - HORIZON, 12) # Step 12 for speed
    
    print(f"Running backtest over {len(indices)} trades...")
    
    for i in tqdm(indices):
        current_time = test_data.index[i]
        context_window = test_data.iloc[i-context_size:i]['close'].tolist()
        entry_price = test_data.iloc[i]['close']
        actual_exit_price = test_data.iloc[i+HORIZON]['close']
        
        # Calendar events
        end_time = test_data.index[i+HORIZON-1]
        events = df_cal[(df_cal['DateTime'] >= current_time) & (df_cal['DateTime'] <= end_time)]
        macro_events = []
        for _, event in events.iterrows():
            macro_events.append({
                "time": event['DateTime'].strftime("%Y-%m-%dT%H:%M:%SZ"),
                "event": event['Event'],
                "impact": event['Impact'].lower().split()[0]
            })
            
        input_json = {"context_window": [round(float(c), 4) for c in context_window], "macro_events": macro_events}
        prompt = "forecast usdjpy: " + json.dumps(input_json)
        
        inputs = tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True).to(model.device)
        with torch.no_grad():
            outputs = model.generate(**inputs, max_new_tokens=512)
        
        raw_pred = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Robust parsing
        direction = "sideways"
        try:
            parse_target = raw_pred.strip()
            if not parse_target.startswith("{"): parse_target = "{" + parse_target + "}"
            pred_data = json.loads(parse_target)
            direction = pred_data.get("direction", "sideways")
        except:
            if '"direction": "up"' in raw_pred: direction = "up"
            elif '"direction": "down"' in raw_pred: direction = "down"
            
        pnl = 0
        if direction == "up":
            pnl = (actual_exit_price - (entry_price + SPREAD))
        elif direction == "down":
            pnl = ((entry_price - SPREAD) - actual_exit_price)
            
        results.append({
            "time": str(current_time),
            "direction": direction,
            "pnl": float(pnl),
            "win": (pnl > 0) if direction != "sideways" else None
        })
        
    df_results = pd.DataFrame(results)
    trades = df_results[df_results['direction'] != 'sideways']
    
    total_pips = df_results['pnl'].sum()
    win_rate = trades['win'].mean() if len(trades) > 0 else 0
    
    stats = {
        "total_periods": len(df_results),
        "total_trades": len(trades),
        "win_rate": f"{win_rate:.2%}",
        "net_pips": f"{total_pips:.4f}",
        "avg_pnl": f"{(total_pips/len(trades)):.4f}" if len(trades) > 0 else "0"
    }
    print("\n--- RESULTS ---")
    print(json.dumps(stats, indent=2))
    print("----------------")

if __name__ == "__main__":
    run_backtest()