"""Pipeline Step 15 (BM25 + 세트) — 개별 토큰 + 복합 토큰(세트) 모두 넣은 BM25.

이전 pipeline_15_bm25_comparison 은 개별 토큰만 썼음 → 세트 정보 누락.
이번엔 세트를 "복합 토큰(phrase token)" 으로 취급해 BM25 에 포함.

복합 토큰 정의:
  · 프레임의 각 source_text_line 에서 추출된 키워드 묶음 (예: [BIM, DX, 이해])
  · tf(compound, frame) = 1 (한 source line = 1 occurrence)
  · df(compound) = "그 모든 키워드를 다 가진 프레임 수"
  · IDF(compound) = log((N - df + 0.5) / (df + 0.5) + 1)

복합 토큰 매칭 (MDX 쪽):
  · MDX 에 compound 의 키워드가 얼마나 있나 → coverage (0~1)
  · 기여도 = IDF(compound) × tf_norm × coverage  (부분 매칭 허용)

최종 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'))

    # ─── 개별 토큰 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', []):
            frame_tokens[fid][t['token']] += t.get('local_count_in_frame', 1)

    # ─── 복합 토큰 (세트) 프레임별 ───
    frame_compounds = defaultdict(list)
    for set_id, s in auto['source_text_sets'].items():
        fid = s['frame_id']
        comp = tuple(sorted(set(s['term_values'])))
        if len(comp) >= 1:
            frame_compounds[fid].append(comp)

    # ─── 문서 길이 |D| ───
    # 개별 토큰 수 + 복합 토큰 수 (각 compound 는 1 unit)
    frame_len = {}
    for fid in frame_tokens:
        ind_len = sum(frame_tokens[fid].values())
        comp_len = len(frame_compounds.get(fid, []))
        frame_len[fid] = ind_len + comp_len
    avg_dl = sum(frame_len.values()) / max(len(frame_len), 1)

    # ─── 개별 토큰 df ───
    ind_df = defaultdict(int)
    for fid, counts in frame_tokens.items():
        for tok in counts:
            ind_df[tok] += 1

    # ─── 복합 토큰 df ───
    all_compounds = set()
    for comps in frame_compounds.values():
        all_compounds.update(comps)
    comp_df = {}
    for c in all_compounds:
        c_set = set(c)
        df = sum(1 for fid, tokens in frame_tokens.items() if c_set.issubset(tokens))
        comp_df[c] = max(df, 1)

    def idf(df):
        return log((N_FRAMES - df + 0.5) / (df + 0.5) + 1)

    # ─── BM25 점수 (개별 + 복합) ───
    def score_frame(fid, mdx_tokens):
        dl = frame_len[fid]
        if dl == 0:
            return 0.0
        score = 0.0
        # 개별 토큰
        for t, tf in frame_tokens[fid].items():
            if t in mdx_tokens:
                idf_val = idf(ind_df[t])
                tf_norm = tf * (K1 + 1) / (tf + K1 * (1 - B + B * dl / avg_dl))
                score += idf_val * tf_norm
        # 복합 토큰
        for c in frame_compounds.get(fid, []):
            c_set = set(c)
            if not c_set:
                continue
            hits = c_set & mdx_tokens
            coverage = len(hits) / len(c_set)
            if coverage > 0:
                idf_val = idf(comp_df[c])
                tf = 1
                tf_norm = tf * (K1 + 1) / (tf + K1 * (1 - B + B * dl / avg_dl))
                score += idf_val * tf_norm * coverage
        return score

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

    print('=' * 95)
    print(f'{"섹션":<8} | {"순위":<3} | {"현재 (0.30/0.50/0.20)":<28} | {"BM25 + 세트 (복합 토큰)":<30}')
    print('=' * 95)

    top1_agree = 0
    top3_agree = 0
    target_hits_bm25 = 0
    target_sids = ['01-2', '02-2.2', '03-1', '03-2']
    total = 0
    answer_map = v1['meta']['answer_map']

    for sid, sec in v1['mdx_sections'].items():
        mdx_tokens = set(normalized['mdx'][sid].get('unique_tokens', []))
        scores = {fid: (score_frame(fid, mdx_tokens), frame_num_map.get(fid)) for fid in frame_tokens}
        ranking = sorted(scores.items(), key=lambda x: -x[1][0])
        current = sec['rank_by_matching_score']
        total += 1
        if current[0]['frame_id'] == ranking[0][0]:
            top1_agree += 1
        if {r['frame_id'] for r in current[:3]} == {x[0] for x in ranking[:3]}:
            top3_agree += 1

        # TARGET 정답률
        if sid in target_sids:
            ans_num = answer_map[sid]
            if ranking[0][1][1] == ans_num:
                target_hits_bm25 += 1

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

    print()
    print('=' * 95)
    print(f'Top-1 일치: {top1_agree}/{total}')
    print(f'Top-3 집합 일치: {top3_agree}/{total}')
    print(f'TARGET 정답률: BM25+세트 = {target_hits_bm25}/4 (vs 현재 4/4)')
    print('=' * 95)


if __name__ == '__main__':
    main()