"""Pipeline Step 15 (BM25) — 정식 BM25 로 랭킹 계산해서 현재 가중치 방식과 비교.

BM25 공식 (표준):
  score(D, Q) = Σ_{t∈Q ∩ D}  IDF(t) × (tf(t,D) × (k1+1)) / (tf(t,D) + k1 × (1 - b + b × |D|/avg_dl))
  IDF(t) = log((N - n(t) + 0.5) / (n(t) + 0.5) + 1)

설정:
  D = 각 Figma 프레임 (키워드 멀티셋)
  Q = MDX 섹션 토큰 집합
  N = 32 프레임
  n(t) = t 를 포함한 프레임 수 (frame_df)
  tf(t, D) = 프레임 D 에서 t 의 local_count
  |D| = 프레임 D 의 총 토큰 수
  avg_dl = 모든 프레임 평균 길이
  k1 = 1.5, b = 0.75 (표준값)

비교: 현재 0.30/0.50/0.20 가중치 결과 vs BM25 결과.
"""
from collections import defaultdict
from math import log
from pathlib import Path
import yaml

HERE = Path(__file__).parent
N_FRAMES = 32
K1 = 1.5
B = 0.75


def main():
    auto = yaml.safe_load((HERE / 'auto_anchor_candidates.yaml').read_text(encoding='utf-8'))
    v1 = yaml.safe_load((HERE / 'mdx_matching_result.yaml').read_text(encoding='utf-8'))
    normalized = yaml.safe_load((HERE / 'normalized_text_tokens.yaml').read_text(encoding='utf-8'))

    # ─── 프레임별 토큰 multiset (tf = 등장 횟수) ───
    frame_tokens = defaultdict(lambda: defaultdict(int))
    for set_id, s in auto['source_text_sets'].items():
        fid = s['frame_id']
        for t in s.get('terms', []):
            tok = t['token']
            frame_tokens[fid][tok] += t.get('local_count_in_frame', 1)

    # 프레임 길이 |D| (= 모든 토큰 local_count 합)
    frame_len = {fid: sum(c.values()) for fid, c in frame_tokens.items()}
    avg_dl = sum(frame_len.values()) / max(len(frame_len), 1)

    # n(t): t 를 포함하는 프레임 수
    frame_df_map = defaultdict(int)
    for fid, counts in frame_tokens.items():
        for tok in counts:
            frame_df_map[tok] += 1

    def idf_bm25(t):
        n = frame_df_map.get(t, 0)
        return log((N_FRAMES - n + 0.5) / (n + 0.5) + 1)

    def bm25_score(frame_id, mdx_tokens):
        counts = frame_tokens[frame_id]
        dl = frame_len[frame_id]
        if dl == 0:
            return 0.0
        score = 0.0
        for t in mdx_tokens:
            if t in counts:
                tf = counts[t]
                idf_val = idf_bm25(t)
                denom = tf + K1 * (1 - B + B * dl / avg_dl)
                tf_norm = tf * (K1 + 1) / denom
                score += idf_val * tf_norm
        return score

    # frame_id → frame_number
    frame_num_map = {fid: info['frame_number'] for fid, info in auto['frame_stats'].items()}

    # ─── 각 MDX 섹션별 계산 ───
    print('=' * 90)
    print(f'{"섹션":<8} | {"순위":<3} | {"현재 (0.30/0.50/0.20)":<28} | {"BM25 (k1=1.5, b=0.75)":<28}')
    print('=' * 90)

    top1_agree = 0
    top3_agree = 0
    total = 0

    for sid, sec in v1['mdx_sections'].items():
        mdx_tokens = set(normalized['mdx'][sid].get('unique_tokens', []))
        bm25_scores = {
            fid: (bm25_score(fid, mdx_tokens), frame_num_map.get(fid))
            for fid in frame_tokens
        }
        bm25_ranking = sorted(bm25_scores.items(), key=lambda x: -x[1][0])

        current_ranking = sec['rank_by_matching_score']
        total += 1
        if current_ranking[0]['frame_id'] == bm25_ranking[0][0]:
            top1_agree += 1
        cur_top3 = {r['frame_id'] for r in current_ranking[:3]}
        bm25_top3 = {x[0] for x in bm25_ranking[:3]}
        if cur_top3 == bm25_top3:
            top3_agree += 1

        for i in range(5):
            cur = current_ranking[i]
            bm = bm25_ranking[i]
            cur_s = f'Frame {cur["frame_number"]:>2}  score {cur["matching_score"]:.3f}'
            bm_s = f'Frame {bm[1][1]:>2}  score {bm[1][0]:>6.2f}'
            prefix = sid if i == 0 else ''
            print(f'{prefix:<8} | {i+1:<3} | {cur_s:<28} | {bm_s:<28}')
        print('-' * 90)

    print()
    print('=' * 90)
    print(f'Top-1 일치: {top1_agree}/{total}')
    print(f'Top-3 집합 일치: {top3_agree}/{total}')
    print('=' * 90)
    print()
    print(f'N = {N_FRAMES} 프레임, avg_dl = {avg_dl:.1f} (평균 프레임 토큰 수)')
    print(f'k1 = {K1}, b = {B} (표준 BM25 파라미터)')


if __name__ == '__main__':
    main()