#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
天勤(TqSdk) 专业版 · 下载生猪 5 分钟全历史 -> 干净 CSV
=====================================================
- 用你自己的信易账号（专业版/试用）在本地运行；账号密码运行时当场输入，不写进文件、不外发。
- 下载生猪主连 KQ.m@DCE.lh 的 5 分钟线，整理成 lh_5min.csv，
  列名为 datetime,open,high,low,close,volume —— 可直接喂给 lh_expma_backtest.py。

准备：
  pip install -i https://pypi.tuna.tsinghua.edu.cn/simple tqsdk pandas
运行：
  python download_tq.py
"""
import getpass
from datetime import datetime
from contextlib import closing
import pandas as pd

try:
    from tqsdk import TqApi, TqAuth
    from tqsdk.tools import DataDownloader
except ImportError:
    raise SystemExit("未安装 tqsdk，请先运行：pip install tqsdk pandas")

# ---------- 可按需修改 ----------
SYMBOL   = "KQ.m@DCE.lh"          # 生猪主连（DCE 生猪 product=lh）。想下单一合约可改为 "DCE.lh2601" 等
DUR_SEC  = 300                    # 5 分钟 = 300 秒（1分钟=60，日线=86400，tick=0）
START    = datetime(2021, 1, 8)   # 生猪 2021-01-08 上市
END      = datetime.now()         # 到今天
RAW_CSV  = "lh_5min_raw.csv"      # DataDownloader 原始输出（列名带合约前缀）
OUT_CSV  = "lh_5min.csv"          # 整理后的干净文件（喂给回测系统）
# --------------------------------

def main():
    print("生猪 5 分钟历史下载（天勤专业版）")
    user = input("信易用户名: ").strip()
    pwd  = getpass.getpass("信易密码（输入时不显示）: ")

    api = TqApi(auth=TqAuth(user, pwd))
    dl = DataDownloader(api, symbol_list=SYMBOL, dur_sec=DUR_SEC,
                        start_dt=START, end_dt=END, csv_file_name=RAW_CSV)
    print(f"开始下载 {SYMBOL} {DUR_SEC//60} 分钟线 {START.date()} ~ {END.date()} ...")
    with closing(api):
        while not dl.is_finished():
            api.wait_update()
            print(f"\r进度 {dl.get_progress():.1f}%", end="", flush=True)
    print("\n下载完成，正在整理为干净 CSV ...")

    raw = pd.read_csv(RAW_CSV)
    cols = list(raw.columns)
    def pick(suffix):
        for c in cols:
            if c.endswith(suffix):
                return c
        return None
    o, h, l, c = pick(".open"), pick(".high"), pick(".low"), pick(".close")
    v = pick(".volume")
    dt_col = "datetime" if "datetime" in cols else cols[0]
    if not all([o, h, l, c]):
        raise SystemExit(f"原始文件列名不符合预期：{cols}")
    keep = [dt_col, o, h, l, c] + ([v] if v else [])
    out = raw[keep].copy()
    out.columns = ["datetime", "open", "high", "low", "close"] + (["volume"] if v else [])
    # datetime 若是纳秒时间戳则转为可读时间
    if pd.api.types.is_numeric_dtype(out["datetime"]):
        out["datetime"] = pd.to_datetime(out["datetime"], unit="ns")
    out.to_csv(OUT_CSV, index=False, encoding="utf-8-sig")
    print(f"完成！已生成 {OUT_CSV}，共 {len(out)} 根K线，"
          f"区间 {out['datetime'].iloc[0]} ~ {out['datetime'].iloc[-1]}")
    print(f"接下来即可运行：python lh_expma_backtest.py --data {OUT_CSV} --mode wf --outdir out")

if __name__ == "__main__":
    main()