import pandas as pd
import numpy as np

from config import STOCK_FILTER


def apply_risk_filters(df: pd.DataFrame) -> pd.DataFrame:
    if df.empty:
        return df
    df = df.copy()
    df = _remove_st(df)
    df = _remove_suspended(df)
    df = _remove_limit_stocks(df)
    df = _remove_low_liquidity(df)
    df = df.reset_index(drop=True)
    return df


def apply_model_risk_gates(scores_df: pd.DataFrame, config: dict = None) -> pd.DataFrame:
    cfg = config or {}
    min_score = cfg.get("min_score", 0.3)
    max_rank = cfg.get("max_rank", 100)

    scores_df = scores_df.copy()
    scores_df = scores_df[scores_df["ml_score"] >= min_score]
    scores_df = scores_df.head(max_rank)
    return scores_df


def rank_stocks(scored_df: pd.DataFrame, weights: dict = None) -> pd.DataFrame:
    w = weights or {"ml": 0.7, "flow": 0.2, "sector": 0.1}
    df = scored_df.copy()

    if "ml_score_norm" in df.columns:
        df["ml_weighted"] = df["ml_score_norm"] * w.get("ml", 0.7)
    else:
        df["ml_weighted"] = 0

    if "main_flow_pct" in df.columns:
        flow_vals = df["main_flow_pct"]
        flow_norm = (flow_vals - flow_vals.min()) / (flow_vals.max() - flow_vals.min() + 1e-10)
        df["flow_weighted"] = flow_norm * w.get("flow", 0.2)
    else:
        df["flow_weighted"] = 0

    if "stock_in_hot_sector" in df.columns:
        df["sector_weighted"] = df["stock_in_hot_sector"].astype(float) * w.get("sector", 0.1)
    else:
        df["sector_weighted"] = 0

    df["total_score"] = df["ml_weighted"] + df["flow_weighted"] + df["sector_weighted"]
    df = df.sort_values("total_score", ascending=False).reset_index(drop=True)
    df["rank"] = range(1, len(df) + 1)
    return df


def _remove_st(df: pd.DataFrame) -> pd.DataFrame:
    name_col = _find_col(df, ["名称", "股票名称", "name"])
    if name_col:
        return df[~df[name_col].str.contains(r"ST|退", na=False)].copy()
    return df


def _remove_suspended(df: pd.DataFrame) -> pd.DataFrame:
    price_col = _find_col(df, ["最新价", "收盘", "close", "price"])
    if price_col:
        return df[df[price_col] > 0].copy()
    return df


def _remove_limit_stocks(df: pd.DataFrame) -> pd.DataFrame:
    chg_col = _find_col(df, ["涨跌幅", "涨跌幅(%)", "change_pct"])
    if chg_col:
        return df[(df[chg_col] > -9.9) & (df[chg_col] < 9.9)].copy()
    return df


def _remove_low_liquidity(df: pd.DataFrame) -> pd.DataFrame:
    amt_col = _find_col(df, ["成交额", "成交额(元)", "amount"])
    if amt_col:
        return df[df[amt_col] >= STOCK_FILTER["min_turnover"]].copy()
    return df


def _find_col(df: pd.DataFrame, candidates: list) -> str:
    for c in candidates:
        if c in df.columns:
            return c
    return ""