import os
import glob
import pandas as pd
import numpy as np
from sklearn.preprocessing import MinMaxScaler, StandardScaler
import argparse

# === 参数解析 ===
parser = argparse.ArgumentParser(description="统一归一化评分所有 metrics_all.csv 文件并选择 top-N 样本")
parser.add_argument("--input_root", type=str, required=True, help="输入目录，例如 /data/tts/emogen/gens/hard/multi_times")
parser.add_argument("--normedcsv_output_dir", type=str, required=True, help="输出目录，例如 ./final_scored")
parser.add_argument("--finalcsv_output_dir", type=str, required=True, help="输出目录，例如 ./final_scored")
parser.add_argument("--top_k", type=int, default=1, help="选择第几高 final_score（默认 1 表示 top-1）")
parser.add_argument("--weights", type=str, default="1", help="选择第几高 final_score（默认 1 表示 top-1）")
args = parser.parse_args()

input_root = args.input_root
os.makedirs(args.normedcsv_output_dir, exist_ok=True)
os.makedirs(args.finalcsv_output_dir, exist_ok=True)

weights = args.weights.split("_")
# === 权重设置 ===
weights = {
    'cer': 0.0,
    'emo2vec_score': float(weights[0]),
    'audonnx_score': float(weights[1]),
    'resemb_score': 0.0,
    'wavlm_score': float(weights[2]),
    'dnsmos_var': float(weights[3]),
    'auto_pcp': float(weights[4]),
}
score_columns = list(weights.keys())
invert_columns = ['cer', 'dnsmos_var']  # 越小越好

# === 收集所有 metrics_all.csv ===
csv_paths = glob.glob(os.path.join(input_root, "*", "metrics_all.csv"))
grouped_data = {}

# === 合并数据，并记录来源文件 ===
for path in csv_paths:
    try:
        df = pd.read_csv(path)
        # if "zh" in path:
        #     continue
        df.columns = [c.strip().replace("\ufeff", "") for c in df.columns]
        if 'keys' in df.columns and 'line_i' in df.columns:
            df = df.sort_values(by=["keys", "line_i"]).reset_index(drop=True)
        df["__line_i"] = range(len(df))  # 保留原始顺序
        basename = os.path.basename(os.path.dirname(path))
        parts = basename.split("_")
        task_lang = "_".join(parts[:2])
        if task_lang not in grouped_data:
            grouped_data[task_lang] = []
        df["__source__"] = basename
        grouped_data[task_lang].append(df)
    except Exception as e:
        print(f"❌ 读取失败：{path}")

# === 按 group 统一归一化并筛选 top-k ===
for task_lang, df_list in grouped_data.items():
    merged = pd.concat(df_list, ignore_index=True)
    
    # 判断有效行
    def is_valid(row):
        return not any([
            row.get('cer') == 100.0,
            row.get('emo2vec_score') == -1,
            row.get('audonnx_score') == -1,
            row.get('resemb_score') == -1,
            row.get('wavlm_score') == -1,
            row.get('dnsmos_var') == -1,
            row.get('auto_pcp') == -1
        ])
    valid_mask = merged.apply(is_valid, axis=1)

    if valid_mask.sum() == 0:
        print(f"⚠️ 全为无效行：{task_lang}")
        continue

    scaler = StandardScaler()
    normed = scaler.fit_transform(merged.loc[valid_mask, score_columns])

    for i, col in enumerate(score_columns):
        if col in invert_columns:
            normed[:, i] = 1 - normed[:, i]
        merged[f"{col}_norm"] = np.nan
        merged.loc[valid_mask, f"{col}_norm"] = normed[:, i]

    weight_array = np.array([weights[c] for c in score_columns])
    final_score = np.dot(normed, weight_array)
    merged["final_score"] = -100.0
    merged.loc[valid_mask, "final_score"] = final_score

    # === 保存每个源文件（保留原顺序、不筛选）===
    new_df_list = []
    for basename in merged["__source__"].unique():
        sub_df = merged[merged["__source__"] == basename].copy()

        # 按 __line_i 排序并重设索引，确保 to_csv 保持正确顺序
        if "__line_i" in sub_df.columns:
            sub_df = sub_df.sort_values("__line_i").reset_index(drop=True)

        out_path = os.path.join(args.normedcsv_output_dir, f"{basename}.csv")
        sub_df.to_csv(out_path, index=False, encoding="utf-8")
        print(f"📄 已保存归一化文件：{out_path}")
        new_df_list.append(sub_df)

    df_list = new_df_list  # 这是归一化后、按原顺序保存的所有子文件数据列表
    result_rows = []  # 存储最终选中行的列表
    row_counts = [len(df) for df in df_list]  # 记录每个 DataFrame 的行数
    max_rows = max(row_counts)  # 获取所有文件中最大的行数（以此决定行索引遍历范围）

    for row_index in range(max_rows):  # 按照每一行的行号逐行处理
        row_variants = []  # 存储该行号在不同文件中的所有可选版本（来自不同模型）

        for df in df_list:  # 遍历所有模型/文件的 DataFrame
            if row_index < len(df):  # 如果该文件中当前行号存在
                row = df.iloc[row_index]  # 获取该行
                # 如果该行的 final_score 有效，则加入候选列表
                if row["final_score"] != -100.0 and pd.notna(row["final_score"]):
                    row_variants.append(row)

        if len(row_variants) >= args.top_k:
            # 正常选择第 top_k 高分
            topk_row = sorted(row_variants, key=lambda r: r["final_score"], reverse=True)[args.top_k - 1].copy()
        elif len(row_variants) > 0:
            # 候选不足，但至少有一个，就选最高分
            print(f"⚠️ {task_lang} row {row_index} 仅有 {len(row_variants)} 条，已选最高分候选")
            topk_row = sorted(row_variants, key=lambda r: r["final_score"], reverse=True)[0].copy()
        else:
            # 完全没有候选，填入空行（可定制）
            print(f"⚠️ {task_lang} row {row_index} 无有效候选，插入空行")
            topk_row = pd.Series(dtype=object)
            topk_row["final_score"] = -100.0
            result_rows.append(topk_row)
            continue  # 此处不再执行排名逻辑，直接加入空行

        # === 添加归一化字段的排名 ===
        for field in [f"{col}_norm" for col in score_columns]:
            scores = [r[field] for r in row_variants if pd.notna(r[field])]
            sorted_scores = sorted(scores, reverse=True)
            if pd.notna(topk_row.get(field)):
                try:
                    topk_row[field + "_rank"] = sorted_scores.index(topk_row[field]) + 1
                except ValueError:
                    topk_row[field + "_rank"] = -1
            else:
                topk_row[field + "_rank"] = -1
        result_rows.append(topk_row)

    # 将所有选中的行合并成一个最终 DataFrame（顺序即为行号顺序）
    final_df = pd.DataFrame(result_rows)

    # 构造输出路径，保存为 CSV 文件
    final_out_path = os.path.join(args.finalcsv_output_dir, f"top{args.top_k}_{task_lang}.csv")
    final_df.to_csv(final_out_path, index=False, encoding="utf-8")

    # 打印保存成功信息
    print(f"✅ 已保存：{task_lang}_top{args.top_k}.csv")