# app.py
import os
import gradio as gr
import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import einops

from huggingface_hub import snapshot_download
from visionts import VisionTSpp, freq_to_seasonality_list

# ========================
# 配置
# ========================
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
REPO_ID = "Lefei/VisionTSpp"
LOCAL_DIR = "./hf_models/VisionTSpp"
CKPT_PATH = os.path.join(LOCAL_DIR, "visiontspp_model.ckpt")

ARCH = 'mae_base'  # 可选: 'mae_base', 'mae_large', 'mae_huge'

# 下载模型（Space 构建时执行一次）
if not os.path.exists(CKPT_PATH):
    os.makedirs(LOCAL_DIR, exist_ok=True)
    print("Downloading model from Hugging Face Hub...")
    snapshot_download(repo_id=REPO_ID, local_dir=LOCAL_DIR, local_dir_use_symlinks=False)

# 加载模型（全局加载一次）
model = VisionTSpp(
    ARCH,
    ckpt_path=CKPT_PATH,
    quantile=True,
    clip_input=True,
    complete_no_clip=False,
    color=True
).to(DEVICE)
print(f"Model loaded on {DEVICE}")

# ========================
# 核心预测与可视化函数
# ========================

def visual_ts(true, preds=None, lookback_len_visual=300, pred_len=96):
    """
    可视化真实值 vs 预测值
    true: [T, nvars]
    preds: [T, nvars]，与 true 对齐
    """
    if isinstance(true, torch.Tensor):
        true = true.cpu().numpy()
    if isinstance(preds, torch.Tensor):
        preds = preds.cpu().numpy()

    nvars = true.shape[1]

    FIG_WIDTH = 12
    FIG_HEIGHT_PER_VAR = 1.8
    FONT_S = 10

    fig, axes = plt.subplots(
        nrows=nvars, ncols=1, figsize=(FIG_WIDTH, nvars * FIG_HEIGHT_PER_VAR), sharex=True,
        gridspec_kw={'height_ratios': [1] * nvars}
    )
    if nvars == 1:
        axes = [axes]

    lookback_len = true.shape[0] - pred_len

    for i, ax in enumerate(axes):
        ax.plot(true[:, i], label='Ground Truth', color='gray', linewidth=1.8)
        if preds is not None:
            ax.plot(np.arange(lookback_len, len(true)), preds[lookback_len:, i],
                    label='Prediction (Median)', color='blue', linewidth=1.8)

        # 分隔线
        y_min, y_max = ax.get_ylim()
        ax.vlines(x=lookback_len, ymin=y_min, ymax=y_max,
                  colors='gray', linestyles='--', alpha=0.7, linewidth=1)

        ax.set_yticks([])
        ax.set_xticks([])
        ax.text(0.005, 0.8, f'Var {i+1}', transform=ax.transAxes, fontsize=FONT_S, weight='bold')

    # 图例
    if preds is not None:
        handles, labels = axes[0].get_legend_handles_labels()
        fig.legend(handles, labels, loc='upper right', bbox_to_anchor=(0.9, 0.9), prop={'size': FONT_S})

    # 计算 MSE/MAE
    if preds is not None:
        true_eval = true[-pred_len:]
        pred_eval = preds[-pred_len:]
        mse = np.mean((true_eval - pred_eval) ** 2)
        mae = np.mean(np.abs(true_eval - pred_eval))
        fig.suptitle(f'MSE: {mse:.4f}, MAE: {mae:.4f}', fontsize=12, y=0.95)

    plt.subplots_adjust(hspace=0)
    return fig  # 返回 matplotlib figure


def predict_and_visualize(df, context_len=960, pred_len=394, freq="15Min"):
    """
    输入: df (pandas.DataFrame)，必须包含 'date' 列和其他数值列
    输出: matplotlib 图像
    """
    if 'date' in df.columns:
        df['date'] = pd.to_datetime(df['date'])
        df = df.set_index('date')
    else:
        # 如果没有 date 列，假设是纯数值序列
        df = df.copy()

    data = df.values  # [T, nvars]
    nvars = data.shape[1]

    if data.shape[0] < context_len + pred_len:
        raise ValueError(f"数据太短，至少需要 {context_len + pred_len} 行，当前只有 {data.shape[0]} 行。")

    # 归一化（使用训练集前 70% 的统计量）
    train_len = int(len(data) * 0.7)
    x_mean = data[:train_len].mean(axis=0, keepdims=True)
    x_std = data[:train_len].std(axis=0, keepdims=True) + 1e-8
    data_norm = (data - x_mean) / x_std

    # 取最后一段作为测试窗口
    end_idx = len(data_norm)
    start_idx = end_idx - (context_len + pred_len)
    x = data_norm[start_idx:start_idx + context_len]  # [context_len, nvars]
    y_true = data_norm[start_idx + context_len:end_idx]  # [pred_len, nvars]

    # 设置周期性
    periodicity_list = freq_to_seasonality_list(freq)
    periodicity = periodicity_list[0] if periodicity_list else 1
    color_list = [i % 3 for i in range(nvars)]  # RGB 循环着色

    # 更新模型配置
    model.update_config(
        context_len=context_len,
        pred_len=pred_len,
        periodicity=periodicity,
        num_patch_input=7,
        padding_mode='constant'
    )

    # 转为 tensor
    x_tensor = torch.FloatTensor(x).unsqueeze(0).to(DEVICE)  # [1, T, N]
    y_true_tensor = torch.FloatTensor(y_true).unsqueeze(0).to(DEVICE)

    # 预测
    with torch.no_grad():
        y_pred, _, _, _, _ = model.forward(x_tensor, export_image=True, color_list=color_list)
        y_pred_median = y_pred[0]  # median prediction

    # 反归一化
    y_true_original = y_true * x_std + x_mean
    y_pred_original = y_pred_median[0].cpu().numpy() * x_std + x_mean

    # 构造完整序列用于可视化
    full_true = np.concatenate([x * x_std + x_mean, y_true_original], axis=0)
    full_pred = np.concatenate([x * x_std + x_mean, y_pred_original], axis=0)

    # 可视化
    fig = visual_ts(true=full_true, preds=full_pred, lookback_len_visual=context_len, pred_len=pred_len)
    return fig


# ========================
# 默认数据加载
# ========================
def load_default_data():
    data_path = "./datasets/ETTm1.csv"
    if not os.path.exists(data_path):
        os.makedirs("./datasets", exist_ok=True)
        url = "https://raw.githubusercontent.com/zhouhaoyi/ETDataset/main/ETT-small/ETTm1.csv"
        df = pd.read_csv(url)
        df.to_csv(data_path, index=False)
    else:
        df = pd.read_csv(data_path)
    return df


# ========================
# Gradio 界面
# ========================
def run_forecast(file_input, context_len, pred_len, freq):
    if file_input is not None:
        df = pd.read_csv(file_input.name)
        title = "Uploaded Data Prediction"
    else:
        df = load_default_data()
        title = "Default ETTm1 Dataset Prediction"

    try:
        fig = predict_and_visualize(df, context_len=int(context_len), pred_len=int(pred_len), freq=freq)
        fig.suptitle(title, fontsize=14, y=0.98)
        plt.close(fig)  # 防止重复显示
        return fig
    except Exception as e:
        # 返回错误信息图像
        fig, ax = plt.subplots()
        ax.text(0.5, 0.5, f"Error: {str(e)}", ha='center', va='center', wrap=True)
        ax.axis('off')
        plt.close(fig)
        return fig


# Gradio UI
with gr.Blocks(title="VisionTS++ 时间序列预测") as demo:
    gr.Markdown("# 🕰️ VisionTS++ 时间序列预测平台")
    gr.Markdown("上传你的多变量时间序列 CSV 文件，或使用默认 ETTm1 数据进行预测。")

    with gr.Row():
        file_input = gr.File(label="上传 CSV 文件（含 date 列或纯数值）", file_types=['.csv'])
        with gr.Column():
            context_len = gr.Number(label="历史长度 (context_len)", value=960)
            pred_len = gr.Number(label="预测长度 (pred_len)", value=394)
            freq = gr.Textbox(label="时间频率 (如 15Min, H)", value="15Min")

    btn = gr.Button("🚀 开始预测")

    output_plot = gr.Plot(label="预测结果")

    btn.click(
        fn=run_forecast,
        inputs=[file_input, context_len, pred_len, freq],
        outputs=output_plot
    )

    # 示例：使用默认数据
    gr.Examples(
        examples=[
            [None, 960, 394, "15Min"]
        ],
        inputs=[file_input, context_len, pred_len, freq],
        outputs=output_plot,
        fn=run_forecast,
        label="点击运行默认示例"
    )

# 启动
if __name__ == "__main__":
    demo.launch()