import os
import sys
import subprocess

KRONOS_DIR = "Kronos"
if not os.path.exists(KRONOS_DIR):
    subprocess.run(["git", "clone", "https://github.com/shiyu-coder/Kronos.git", KRONOS_DIR], check=True)

sys.path.insert(0, os.path.abspath(KRONOS_DIR))

import gradio as gr
import pandas as pd
import torch
from transformers import AutoConfig
from model import Kronos, KronosTokenizer, KronosPredictor
import plotly.graph_objects as go
from datetime import timedelta

print("Loading Kronos-base...")
device = "cuda" if torch.cuda.is_available() else "cpu"

# 1. Load config dari HF dulu
config = AutoConfig.from_pretrained("THUDM/Kronos-base", trust_remote_code=True)
# 2. Init tokenizer pake config
tokenizer = KronosTokenizer.from_pretrained("THUDM/Kronos-base", config=config)
# 3. Init model pake config
model = Kronos.from_pretrained("THUDM/Kronos-base", config=config)
model.to(device)
model.eval()
# 4. Baru bikin predictor
predictor = KronosPredictor(model, tokenizer, device=device)
print("Model ready!")

def predict_xauusd(file, pred_len, lookback):
    if file is None:
        return None, "Upload CSV MT5 dulu"
    try:
        df = pd.read_csv(file.name)
        df.columns = df.columns.str.lower()
        if 'date' in df.columns: df.rename(columns={'date': 'time'}, inplace=True)
        if 'vol' in df.columns: df.rename(columns={'vol': 'volume'}, inplace=True)
        if 'tickvol' in df.columns: df.rename(columns={'tickvol': 'volume'}, inplace=True)
        
        required = ['time', 'open', 'high', 'low', 'close', 'volume']
        if not all(col in df.columns for col in required):
            return None, f"CSV harus ada: {required}. Punyamu: {list(df.columns)}"
        
        df = df.tail(int(lookback)).copy()
        df['time'] = pd.to_datetime(df['time'])
        df = df.sort_values('time').reset_index(drop=True)
        
        pred_df = predictor.predict(df=df, pred_len=int(pred_len))
        
        fig = go.Figure()
        hist_df = df.tail(100)
        fig.add_trace(go.Candlestick(x=hist_df['time'], open=hist_df['open'], high=hist_df['high'], low=hist_df['low'], close=hist_df['close'], name='Data Real'))
        
        last_time = df['time'].iloc[-1]
        tf_minutes = (df['time'].iloc[-1] - df['time'].iloc[-2]).total_seconds() / 60
        future_time = [last_time + timedelta(minutes=tf_minutes * (i+1)) for i in range(int(pred_len))]
        
        fig.add_trace(go.Candlestick(x=future_time, open=pred_df['open'], high=pred_df['high'], low=pred_df['low'], close=pred_df['close'], name='Prediksi'))
        fig.update_layout(template='plotly_dark', xaxis_rangeslider_visible=False, height=600, title=f'Prediksi XAUUSD {int(pred_len)} Candle')
        
        last_close = df['close'].iloc[-1]
        close_pred = pred_df['close'].iloc[-1]
        change = ((close_pred - last_close) / last_close) * 100
        summary = f"**Close Terakhir:** ${last_close:.2f}  \n**Prediksi Close:** ${close_pred:.2f} **({change:+.2f}%)**  \n**Range:** ${pred_df['low'].min():.2f} - ${pred_df['high'].max():.2f}"
        return fig, summary
    except Exception as e:
        return None, f"Error: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown("# 🪙 Kronos XAUUSD Predictor")
    file_input = gr.File(label="Upload CSV XAUUSD MT5")
    with gr.Row():
        lookback = gr.Slider(60, 500, value=120, step=10, label="Data Historis")
        pred_len = gr.Slider(6, 96, value=24, step=1, label="Candle Prediksi")
    btn = gr.Button("Prediksi", variant="primary")
    chart = gr.Plot()
    summary = gr.Markdown()
    btn.click(predict_xauusd, [file_input, pred_len, lookback], [chart, summary])

if __name__ == "__main__":
    demo.launch()