import pandas as pd
import numpy as np


def compute_flow_features(flow_df: pd.DataFrame) -> dict:
    if flow_df is None or flow_df.empty:
        return {}

    features = {}
    cols = flow_df.columns.tolist()

    main_col = _find(cols, ["主力净流入-净额", "主力净流入净额"])
    big_col = _find(cols, ["大单净流入-净额", "大单净流入净额"])
    super_col = _find(cols, ["超大单净流入-净额", "超大单净流入净额"])
    main_pct_col = _find(cols, ["主力净流入-净占比(%)", "主力净流入净占比"])
    close_col = _find(cols, ["收盘价"])
    chg_col = _find(cols, ["涨跌幅(%)", "涨跌幅"])

    if main_col:
        vals = pd.to_numeric(flow_df[main_col], errors="coerce")
        features["main_flow_1d"] = vals.iloc[-1] if len(vals) > 0 else 0
        features["main_flow_3d"] = vals.tail(3).sum() if len(vals) >= 3 else 0
        features["main_flow_5d"] = vals.tail(5).sum() if len(vals) >= 5 else 0
        features["main_flow_trend"] = (
            vals.tail(3).mean() - vals.tail(10).mean() if len(vals) >= 10 else 0
        )

    if big_col:
        vals = pd.to_numeric(flow_df[big_col], errors="coerce")
        features["big_flow_1d"] = vals.iloc[-1] if len(vals) > 0 else 0
        features["big_flow_5d"] = vals.tail(5).sum() if len(vals) >= 5 else 0

    if super_col:
        vals = pd.to_numeric(flow_df[super_col], errors="coerce")
        features["super_flow_1d"] = vals.iloc[-1] if len(vals) > 0 else 0

    if main_pct_col:
        vals = pd.to_numeric(flow_df[main_pct_col], errors="coerce")
        features["main_flow_pct"] = vals.iloc[-1] if len(vals) > 0 else 0
        features["main_flow_pct_3d"] = vals.tail(3).mean() if len(vals) >= 3 else 0

    if main_col and close_col:
        flow_vals = pd.to_numeric(flow_df[main_col], errors="coerce")
        price_vals = pd.to_numeric(flow_df[close_col], errors="coerce")
        if len(flow_vals) >= 5 and price_vals.std() > 0:
            features["flow_price_corr_5d"] = flow_vals.tail(5).corr(price_vals.tail(5))

    return features


def compute_flow_rank_features(rank_df: pd.DataFrame, stock_code: str) -> dict:
    if rank_df is None or rank_df.empty:
        return {}

    features = {}
    code_col = _find(rank_df.columns.tolist(), ["代码", "股票代码"])
    main_col = _find(rank_df.columns.tolist(), ["主力净流入-净额", "主力净流入净额"])

    if not code_col or not main_col:
        return features

    stock_row = rank_df[rank_df[code_col] == stock_code]
    if stock_row.empty:
        return features

    all_main = pd.to_numeric(rank_df[main_col], errors="coerce").sort_values(ascending=False)
    stock_val = pd.to_numeric(stock_row[main_col].iloc[0], errors="coerce")

    features["main_flow_rank"] = (all_main > stock_val).sum() + 1
    features["main_flow_rank_pct"] = features["main_flow_rank"] / len(all_main)
    features["main_flow_is_top50"] = 1 if features["main_flow_rank_pct"] <= 0.5 else 0
    features["main_flow_is_top20"] = 1 if features["main_flow_rank_pct"] <= 0.2 else 0

    return features


def _find(cols: list, candidates: list) -> str:
    for c in candidates:
        if c in cols:
            return c
    return ""