import argparse
import json
import os
import sys

# 預設的 Lambda 值列表
DEFAULT_LAMBDAS = [0, 0.1, 0.5, 1, 1.5, 2]

def _load_json(path: str):
    if not os.path.isfile(path):
        raise FileNotFoundError(f"JSON not found: {path}")
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)

def _save_json(data, folder, filename):
    os.makedirs(folder, exist_ok=True)
    path = os.path.join(folder, filename)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=None) # 不縮排以節省空間
    return path

def process_head_map(head_map, lcb_lambda):
    """
    輸入: {head_key: {mean, std, count}}
    輸出: [[head_key1, score1], [head_key2, score2], ...] (按分數由高到低排序)
    """
    scored = []
    for head_name, stats in head_map.items():
        mean = float(stats.get("mean", 0.0))
        std = float(stats.get("std", 0.0))
        
        # LCB 分數計算: Mean - Lambda * Std
        score = mean - (lcb_lambda * std)
        
        scored.append((head_name, score))

    # 排序: 分數高的在前面
    scored.sort(key=lambda x: x[1], reverse=True)
    
    return [[name, float(score)] for name, score in scored]

def main():
    parser = argparse.ArgumentParser(description="Batch sort heads based on combined scoring results.")
    parser.add_argument("--input-json", required=True, help="Path to head_scoring_combined.json")
    parser.add_argument("--output-dir", default=None, help="Directory to save sorted JSON files. Defaults to {input_dir}/heads_sorted")
    args = parser.parse_args()

    # -----------------------------------------------------------
    # 自動決定 Output Directory
    # -----------------------------------------------------------
    if args.output_dir:
        out_dir = args.output_dir
    else:
        # 取得 input json 所在的資料夾路徑
        base_dir = os.path.dirname(os.path.abspath(args.input_json))
        out_dir = os.path.join(base_dir, "heads_sorted")
    
    print(f"📂 Loading data from {args.input_json}...")
    try:
        payload = _load_json(args.input_json)
    except Exception as e:
        print(f"❌ Error loading JSON: {e}")
        return

    count = 0
    
    # 1. 遍歷 Setup (Range + Metric)
    for setup_key, profiles_map in payload.items():
        # 解析 Range 和 Setup (Metric)
        # 假設 key 格式固定為 "{range}_{metric}"
        try:
            parts = setup_key.split('_', 1)
            if len(parts) != 2:
                print(f"⚠️ Skipping key with unexpected format: {setup_key}")
                continue
            range_val, metric_val = parts
        except ValueError:
            continue

        # 2. 遍歷 Profile
        for profile_name, head_map in profiles_map.items():
            if not head_map:
                continue

            # 3. 遍歷 Lambda
            for lcb_lambda in DEFAULT_LAMBDAS:
                # 計算並排序
                sorted_heads = process_head_map(head_map, lcb_lambda)
                
                # 4. 生成檔案名稱
                # 格式: {range}_{setup}_{profile}_{lambda}.json
                # lambda 格式化去掉多餘的 .0 (例如 0.0 -> 0)
                lambda_str = f"{lcb_lambda:g}" 
                filename = f"{range_val}_{metric_val}_{profile_name}_{lambda_str}.json"
                
                # 儲存
                _save_json(sorted_heads, out_dir, filename)
                count += 1

    print(f"✅ Successfully generated {count} sorted files in '{out_dir}/'")

if __name__ == "__main__":
    main()
