import argparse
import glob
import json
import os
import pickle
from concurrent.futures import ProcessPoolExecutor, as_completed

import numpy as np
from sklearn.metrics import roc_auc_score
from tqdm import tqdm

# 定義設定檔與 Profile
PROFILE_NAMES = ("inst", "instr", "inst_and_instr")
SETUP_NAMES = ("all_roc", "user_roc", "all_prop", "user_prop")

_PROMPTS = None

def _init_worker(prompts):
    global _PROMPTS
    _PROMPTS = prompts

def iter_raw_pickles(input_path, input_dir):
    if input_path:
        yield input_path
        return
    if not input_dir:
        raise ValueError("Provide --input-path or --input-dir.")
    pattern = os.path.join(input_dir, "head_scoring_raw_split_*.pkl")
    paths = sorted(glob.glob(pattern))
    if not paths:
        raise ValueError(f"No split pickle files found in {input_dir}")
    for path in paths:
        yield path

def safe_roc_auc(y_true, y_score):
    # ROC AUC 需要至少兩個類別 (0 和 1)
    if len(set(y_true)) < 2:
        return None
    try:
        return float(roc_auc_score(y_true, y_score))
    except ValueError:
        return None

def calculate_metrics(attn_segment, target_segment):
    """
    計算單一片段的 ROC 和 Proportion
    """
    if not attn_segment:
        return None, None

    attn_sum = float(np.sum(attn_segment))
    if attn_sum <= 0:
        return None, None

    # 正規化 Attention (總和為 1)
    norm_attn = np.array([a / attn_sum for a in attn_segment], dtype=np.float64)
    target_segment = np.array(target_segment, dtype=np.float64)

    # 1. 計算 Proportion (比例)
    # Target (0/1) * Norm_Attn 的總和，即落在 Target 區域的 Attention 比例
    prop_score = float(np.sum(target_segment * norm_attn))

    # 2. 計算 ROC AUC
    roc_score = safe_roc_auc(target_segment, norm_attn)

    return roc_score, prop_score

def score_head(head_key, head_data, prompts):
    # 初始化結果容器
    # 結構: head_scores[setup][profile] = [list of scores]
    head_scores = {s: {p: [] for p in PROFILE_NAMES} for s in SETUP_NAMES}

    for idx, prompt in enumerate(prompts):
        # 1. 準備各種 Mask
        # 假設 attn_weight 的長度等於 mask 的長度
        attn_weight = head_data["attn_weight"][idx]
        
        user_mask = prompt["user_mask"]
        inst_mask = prompt["inst_mask"]
        instr_mask = prompt["instr_mask"]
        inst_and_instr_mask = [bool(a and b) for a, b in zip(inst_mask, instr_mask)]

        # Target Masks 字典
        targets = {
            "inst": inst_mask,
            "instr": instr_mask,
            "inst_and_instr": inst_and_instr_mask
        }

        # 2. 定義兩種 Scope 的 Indices (All vs User)
        # All: 全部 indices
        indices_all = range(len(attn_weight))
        # User: user_mask 為 True 的 indices
        indices_user = [i for i, m in enumerate(user_mask) if m]

        scopes = {
            "all": indices_all,
            "user": indices_user
        }

        # 3. 遍歷 Scope -> Profile -> 計算指標
        for scope_name, indices in scopes.items():
            if not indices:
                continue
            
            # 根據 indices 取出當前的 attn 和 target 片段
            # 這裡轉成 numpy array 可能比較快，但為了保持跟原 code 邏輯一致使用 list comprehension
            # 為了效能，先取出 attn slice
            cur_attn = [attn_weight[i] for i in indices]

            for profile_name, target_mask in targets.items():
                cur_target = [target_mask[i] for i in indices]

                # 計算該 Scope + Profile 下的 metrics
                roc, prop = calculate_metrics(cur_attn, cur_target)

                # 存入對應的 Setup
                # Setup 命名規則: {scope}_roc, {scope}_prop
                
                if roc is not None:
                    head_scores[f"{scope_name}_roc"][profile_name].append(roc)
                
                if prop is not None:
                    head_scores[f"{scope_name}_prop"][profile_name].append(prop)

    return head_key, head_scores

def score_heads_chunk(chunk_items):
    chunk_results = []
    for head_key, head_data in chunk_items:
        chunk_results.append(score_head(head_key, head_data, _PROMPTS))
    return chunk_results

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input-path", default=None)
    parser.add_argument("--input-dir", default=None)
    parser.add_argument("--output-path", default=None)
    parser.add_argument("--workers", type=int, default=os.cpu_count() or 4)
    args = parser.parse_args()

    # 初始化聚合容器
    # scores[setup][profile][head_key] = [score1, score2, ...]
    scores = {s: {p: {} for p in PROFILE_NAMES} for s in SETUP_NAMES}

    for path in iter_raw_pickles(args.input_path, args.input_dir):
        with open(path, "rb") as f:
            raw = pickle.load(f)
        
        prompts = raw["prompts"]
        heads = raw["heads"]
        head_items = list(heads.items())
        total_heads = len(head_items)
        
        # 簡單的分塊邏輯
        chunk_size = max(1, total_heads // (max(1, args.workers) * 4))
        chunks = [
            head_items[i : i + chunk_size]
            for i in range(0, total_heads, chunk_size)
        ]

        with ProcessPoolExecutor(
            max_workers=args.workers,
            initializer=_init_worker,
            initargs=(prompts,),
        ) as executor:
            futures = [executor.submit(score_heads_chunk, chunk) for chunk in chunks]
            
            progress = tqdm(
                total=total_heads,
                desc=f"Scoring {os.path.basename(path)}",
                unit="head",
            )
            
            for fut in as_completed(futures):
                chunk_results = fut.result()
                for head_key, head_res in chunk_results:
                    # head_res 結構: {setup: {profile: [scores]}}
                    for setup_name, profile_map in head_res.items():
                        for profile_name, val_list in profile_map.items():
                            if val_list:
                                scores[setup_name][profile_name].setdefault(head_key, []).extend(val_list)
                
                progress.update(len(chunk_results))
            progress.close()

    # 計算統計量 (Mean, Std, Count)
    # Output 結構: output[setup][profile][head_key] = {mean, std, count}
    final_output = {s: {p: {} for p in PROFILE_NAMES} for s in SETUP_NAMES}

    print("Aggregating statistics...")
    for setup_name in SETUP_NAMES:
        for profile_name in PROFILE_NAMES:
            head_map = scores[setup_name][profile_name]
            for head_key, values in head_map.items():
                if values:
                    arr = np.array(values, dtype=np.float64)
                    final_output[setup_name][profile_name][head_key] = {
                        "mean": float(arr.mean()),
                        "std": float(arr.std()),
                        "count": int(len(values)),
                    }
                else:
                    final_output[setup_name][profile_name][head_key] = {
                        "mean": 0.0,
                        "std": 0.0,
                        "count": 0,
                    }

    # 決定輸出路徑
    if args.output_path:
        output_path = args.output_path
    else:
        base_dir = args.input_dir or os.path.dirname(args.input_path) or "."
        output_path = os.path.join(base_dir, "head_scoring_combined.json")

    print(f"Saving results to {output_path}...")
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(final_output, f, indent=2)

if __name__ == "__main__":
    main()