"""五大期货交易所前20会员成交持仓排名数据获取工具。

交易所对照:
  - CZCE  郑商所 → get_rank_table_czce   (玻璃FG、纯碱SA、甲醇MA等)
  - SHFE  上期所 → get_shfe_rank_table    (铜CU、铝AL、螺纹RB、原油SC等)
  - DCE   大商所 → get_dce_rank_table     (铁矿石I、豆粕M、棕榈P等)
  - CFFEX 中金所 → get_cffex_rank_table   (IF、IC、IH、T等)
  - GFEX  广期所 → futures_gfex_position_rank (工业硅SI、碳酸锂LC等)

交易所公布规则:
  - CZCE: 同时公布品种汇总 + 各合约明细
  - SHFE/CFFEX: 只公布合约明细，品种汇总由函数自行加总
  - DCE: 只公布品种排名
  - GFEX: 公布合约明细
"""

import akshare as ak
import pandas as pd
from datetime import datetime, timedelta

EXCHANGE_CODE = {
    "CZCE":  "CZCE",
    "SHFE":  "SHFE",
    "DCE":   "DCE",
    "CFFEX": "CFFEX",
    "GFEX":  "GFEX",
}

# 常见品种 -> 交易所映射
VARIETY_EXCHANGE = {
    # 郑商所 CZCE
    "FG": "CZCE", "SA": "CZCE", "MA": "CZCE", "TA": "CZCE",
    "RM": "CZCE", "OI": "CZCE", "ZC": "CZCE", "SF": "CZCE",
    "SM": "CZCE", "UR": "CZCE", "PF": "CZCE", "PX": "CZCE",
    "SH": "CZCE", "PR": "CZCE", "PL": "CZCE", "AP": "CZCE",
    "CJ": "CZCE", "CY": "CZCE", "PK": "CZCE", "CF": "CZCE",
    "SR": "CZCE",
    # 上期所 SHFE
    "CU": "SHFE", "AL": "SHFE", "ZN": "SHFE", "PB": "SHFE",
    "NI": "SHFE", "SN": "SHFE", "AU": "SHFE", "AG": "SHFE",
    "RB": "SHFE", "HC": "SHFE", "WR": "SHFE", "FU": "SHFE",
    "BU": "SHFE", "RU": "SHFE", "SC": "SHFE", "NR": "SHFE",
    "SP": "SHFE", "SS": "SHFE", "LU": "SHFE", "BC": "SHFE",
    "AO": "SHFE", "BR": "SHFE", "EC": "SHFE",
    # 大商所 DCE
    "C":  "DCE", "CS": "DCE", "A":  "DCE", "B":  "DCE",
    "M":  "DCE", "Y":  "DCE", "P":  "DCE", "FB": "DCE",
    "BB": "DCE", "JD": "DCE", "L":  "DCE", "V":  "DCE",
    "PP": "DCE", "J":  "DCE", "JM": "DCE", "I":  "DCE",
    "EG": "DCE", "RR": "DCE", "EB": "DCE", "PG": "DCE",
    "LH": "DCE", "LG": "DCE", "BZ": "DCE",
    # 中金所 CFFEX
    "IF": "CFFEX", "IC": "CFFEX", "IH": "CFFEX", "IM": "CFFEX",
    "T":  "CFFEX", "TF": "CFFEX", "TS": "CFFEX", "TL": "CFFEX",
    # 广期所 GFEX
    "SI": "GFEX", "LC": "GFEX", "PS": "GFEX",
}

# 各交易所的品种代码列表（用于 get_rank_sum）
ALL_VARS_CZCE  = ["WH","PM","CF","SR","TA","OI","RI","MA","FG","RS","RM","ZC",
                  "JR","LR","SF","SM","WT","TC","GN","RO","ER","SRX","SRY",
                  "WSX","WSY","CY","AP","UR","CJ","SA","PK","PF","PX","SH","PR","PL"]
ALL_VARS_DCE   = ["C","CS","A","B","M","Y","P","FB","BB","JD","L","V","PP",
                  "J","JM","I","EG","RR","EB","PG","LH","LG","BZ"]
ALL_VARS_SHFE  = ["CU","AL","ZN","PB","NI","SN","AU","AG","RB","WR","HC",
                  "FU","BU","RU","SC","NR","SP","SS","LU","BC","AO","BR","EC"]
ALL_VARS_CFFEX = ["IF","IC","IM","IH","T","TF","TS","TL"]
ALL_VARS_GFEX  = ["SI","LC","PS"]

# 没有统一字段名映射时用的 to_int 辅助
def _to_int(x):
    if isinstance(x, str):
        return int(x.replace(",", ""))
    return int(x)


# ---------------------------------------------------------------------------
# 底层接口，按交易所获取当日前20排名
# ---------------------------------------------------------------------------

def fetch_top20(date: str = None, exchange: str = None, variety: str = None) -> pd.DataFrame:
    """获取指定交易日该品种的前20会员成交持仓排名。

    主力合约检测调用 :func:`main_contract`，本函数自动定位。

    Parameters
    ----------
    date : str or None
        交易日，格式 YYYY-MM-DD 或 YYYYMMDD。None 表示最新。
    exchange : str or None
        交易所代码（CZCE/SHFE/DCE/CFFEX/GFEX）。None 时从 variety 推算。
    variety : str
        品种代码（如 "FG", "RB", "I", "IF"）。如果为 None 则返回全品种。

    Returns
    -------
    pd.DataFrame
        前20会员排名数据。
        列：rank, party_name, vol, vol_chg,
            long_party_name, long_open_interest, long_open_interest_chg,
            short_party_name, short_open_interest, short_open_interest_chg,
            symbol, exchange
    """
    if exchange is None and variety is not None:
        exchange = VARIETY_EXCHANGE.get(variety.upper())
        if exchange is None:
            raise ValueError(f"未知品种 {variety}，请指定 exchange 参数")

    if exchange == "CZCE":
        return _fetch_czce(date, variety)
    elif exchange == "SHFE":
        return _fetch_shfe(date, variety)
    elif exchange == "DCE":
        return _fetch_dce(date, variety)
    elif exchange == "CFFEX":
        return _fetch_cffex(date, variety)
    elif exchange == "GFEX":
        return _fetch_gfex(date, variety)
    else:
        raise ValueError(f"不支持的交易所: {exchange}")


def _fetch_czce(date=None, variety=None):
    rank = ak.get_rank_table_czce(date=date)
    if variety:
        variety = variety.upper()
        if variety in rank:
            # 品种级别汇总
            return _normalize_czce(rank[variety], variety, "CZCE")
        # 合约级别: 找到该品种的所有合约
        contract_keys = [k for k in rank if k.startswith(variety) and k != variety]
        if not contract_keys:
            raise ValueError(f"未找到品种 {variety} 的数据")
        # 返回汇总版
        return _normalize_czce(rank.get(variety, pd.DataFrame()), variety, "CZCE")
    else:
        # 全部品种 - 返回汇总
        return rank  # dict, 调用方自己处理


def _fetch_shfe(date=None, variety=None):
    vars_list = [variety.upper()] if variety else None
    rank = ak.get_shfe_rank_table(date=date, vars_list=vars_list)
    # SHFE 返回 dict: {合约代码: DataFrame}
    rows = []
    for symbol, df in rank.items():
        for _, r in df.iterrows():
            rows.append({
                "rank": int(r["名次"]),
                "party_name": r["会员简称"],
                "vol": r["成交量"],
                "vol_chg": r["成交量增减"],
                "long_party_name": r["持买单量-会员简称"],
                "long_open_interest": r["持买单量-合计"],
                "long_open_interest_chg": r["持买单量-增减"],
                "short_party_name": r["持卖单量-会员简称"],
                "short_open_interest": r["持卖单量-合计"],
                "short_open_interest_chg": r["持卖单量-增减"],
                "symbol": symbol,
            })
    df = pd.DataFrame(rows)
    if variety:
        df = df[df["symbol"].str.startswith(variety.upper())]
    return df


def _fetch_dce(date=None, variety=None):
    vars_list = [variety.upper()] if variety else None
    rank = ak.get_dce_rank_table(date=date, vars_list=vars_list)
    rows = []
    for symbol, df in rank.items():
        for _, r in df.iterrows():
            rows.append({
                "rank": int(r["名次"]),
                "party_name": r["会员简称"],
                "vol": r["成交量"],
                "vol_chg": r["成交量增减"],
                "long_party_name": r["持买单-会员简称"],
                "long_open_interest": r["持买单"],
                "long_open_interest_chg": r["持买单增减"],
                "short_party_name": r["持卖单-会员简称"],
                "short_open_interest": r["持卖单"],
                "short_open_interest_chg": r["持卖单增减"],
                "symbol": symbol,
            })
    return pd.DataFrame(rows)


def _fetch_cffex(date=None, variety=None):
    vars_list = [variety.upper()] if variety else None
    rank = ak.get_cffex_rank_table(date=date, vars_list=vars_list)
    rows = []
    for symbol, df in rank.items():
        for _, r in df.iterrows():
            rows.append({
                "rank": int(r["名次"]),
                "party_name": r["会员简称"],
                "vol": r["成交量"],
                "vol_chg": r["成交量增减"],
                "long_party_name": r["持买单-会员简称"],
                "long_open_interest": r["持买单-合计"],
                "long_open_interest_chg": r["持买单-增减"],
                "short_party_name": r["持卖单-会员简称"],
                "short_open_interest": r["持卖单-合计"],
                "short_open_interest_chg": r["持卖单-增减"],
                "symbol": symbol,
            })
    return pd.DataFrame(rows)


def _fetch_gfex(date=None, variety=None):
    vars_list = [variety.upper()] if variety else None
    rank = ak.futures_gfex_position_rank(date=date, vars_list=vars_list)
    rows = []
    if isinstance(rank, dict):
        for symbol, df in rank.items():
            for _, r in df.iterrows():
                rows.append({
                    "rank": int(r["名次"]),
                    "party_name": r["会员简称"],
                    "vol": r["成交量"],
                    "vol_chg": r["成交量增减"],
                    "long_party_name": r["持买单-会员简称"],
                    "long_open_interest": r["持买单-合计"],
                    "long_open_interest_chg": r["持买单-增减"],
                    "short_party_name": r["持卖单-会员简称"],
                    "short_open_interest": r["持卖单-合计"],
                    "short_open_interest_chg": r["持卖单-增减"],
                    "symbol": symbol,
                })
    return pd.DataFrame(rows)


def _normalize_czce(df, variety, exchange):
    """标准化 CZCE 品种汇总 DataFrame 字段名"""
    df = df.copy()
    df["symbol"] = df.get("symbol", variety)
    df["exchange"] = exchange
    return df


# ---------------------------------------------------------------------------
# 主力合约检测
# ---------------------------------------------------------------------------

def main_contract(variety: str, date: str = None) -> str:
    """检测该品种在指定交易日的主力合约代码。

    逻辑：成交量/持仓量最大的合约即为主力。
    """
    variety = variety.upper()
    exchange = VARIETY_EXCHANGE.get(variety)
    if exchange is None:
        raise ValueError(f"未知品种 {variety}")

    rank = ak.get_rank_table_czce(date=date) if exchange == "CZCE" else {}
    if exchange == "CZCE":
        # CZCE: 合约级别的 key 是 品种+年月 (FG609), 品种级是 FG
        contract_keys = [k for k in rank if k.startswith(variety) and k != variety]
        if not contract_keys:
            raise ValueError(f"未找到品种 {variety} 的合约数据")
        vols = {}
        for ck in contract_keys:
            df = rank[ck]
            vol_sum = df["vol"].apply(_to_int).sum()
            vols[ck] = vol_sum
        return max(vols, key=vols.get)
    else:
        # 其他交易所：通过成交额最大的标的判断
        raise NotImplementedError(f"{exchange} 主力合约检测待实现")


# ---------------------------------------------------------------------------
# 品种汇总对比（多交易所多品种）
# ---------------------------------------------------------------------------

def rank_summary(date: str = None, varieties: list = None) -> pd.DataFrame:
    """多品种前N会员成交量/持仓汇总对比（仅支持 CZCE 品种）。

    使用 get_rank_table_czce 的品种级别汇总数据，
    返回各品种前5/10/15/20成交量和多空持仓总和。

    Parameters
    ----------
    date : str
        交易日。
    varieties : list
        品种代码列表。默认 FG, SA, MA, CF, TA, UR, PF, SM。

    Returns
    -------
    pd.DataFrame
        每行一个品种，含品种级前N成交量和多空持仓总和。
    """
    if varieties is None:
        varieties = ["FG", "SA", "MA", "CF", "TA", "UR", "PF", "SM"]
    rank = ak.get_rank_table_czce(date=date)
    rows = []
    for var in varieties:
        var = var.upper()
        if var not in rank:
            continue
        df = rank[var]
        df["vol_int"] = df["vol"].apply(_to_int)
        df["long_int"] = df["long_open_interest"].apply(_to_int)
        df["short_int"] = df["short_open_interest"].apply(_to_int)
        total_vol = df["vol_int"].sum()
        vol_5 = df.head(5)["vol_int"].sum()
        vol_10 = df.head(10)["vol_int"].sum()
        long_sum = df["long_int"].sum()
        short_sum = df["short_int"].sum()
        rows.append({
            "var": var,
            "vol_total": total_vol,
            "vol_top5": vol_5,
            "vol_top10": vol_10,
            "vol_top20": df["vol_int"].sum(),
            "long_sum": long_sum,
            "short_sum": short_sum,
            "net": long_sum - short_sum,
        })
    return pd.DataFrame(rows).sort_values("vol_total", ascending=False)


# ---------------------------------------------------------------------------
# 主力合约前20成交量
# ---------------------------------------------------------------------------

def main_top20(variety: str, date: str = None) -> pd.DataFrame:
    """获取品种主力合约的前20会员成交量排名。

    Parameters
    ----------
    variety : str
        品种代码，如 "FG"。
    date : str
        交易日。None 表示最新。

    Returns
    -------
    pd.DataFrame
        主力合约前20会员成交量排名 + 多空持仓排名。
    """
    exchange = VARIETY_EXCHANGE.get(variety.upper())
    if exchange != "CZCE":
        raise NotImplementedError(f"目前仅支持 CZCE 品种的主力合约前20查询")

    rank = ak.get_rank_table_czce(date=date)
    contract_keys = [k for k in rank if k.startswith(variety.upper()) and k != variety.upper()]
    vols = {ck: rank[ck]["vol"].apply(_to_int).sum() for ck in contract_keys}
    main_code = max(vols, key=vols.get)

    df = rank[main_code].copy()
    df["symbol"] = main_code
    # 数值列处理（CZCE 返回字符串含逗号）
    for col in ["vol", "vol_chg", "long_open_interest", "long_open_interest_chg",
                "short_open_interest", "short_open_interest_chg"]:
        df[col] = df[col].apply(_to_int)
    return df.sort_values("rank")


# ---------------------------------------------------------------------------
# 输出格式化
# ---------------------------------------------------------------------------

def print_top20(df: pd.DataFrame, title: str = ""):
    """打印前20成交量排名 + 多空前10。"""
    if title:
        print(f"\n{'='*60}")
        print(f"  {title}")
        print(f"{'='*60}")

    print(f"\n【成交量前20】")
    print(f"{'排名':>4} {'会员':<20} {'成交量':>10} {'变化':>10}")
    print("-" * 50)
    for _, r in df.iterrows():
        print(f"{int(r['rank']):>4} {str(r['vol_party_name']):<20} "
              f"{int(r['vol']):>10,} {int(r['vol_chg']):>10,}")

    print(f"\n【持多单前10】")
    print(f"{'排名':>4} {'会员':<20} {'多单':>10} {'变化':>10}")
    print("-" * 50)
    for _, r in df.head(10).iterrows():
        print(f"{int(r['rank']):>4} {str(r['long_party_name']):<20} "
              f"{int(r['long_open_interest']):>10,} {int(r['long_open_interest_chg']):>10,}")

    print(f"\n【持空单前10】")
    print(f"{'排名':>4} {'会员':<20} {'空单':>10} {'变化':>10}")
    print("-" * 50)
    for _, r in df.head(10).iterrows():
        print(f"{int(r['rank']):>4} {str(r['short_party_name']):<20} "
              f"{int(r['short_open_interest']):>10,} {int(r['short_open_interest_chg']):>10,}")
