# -*- coding: utf-8 -*-
"""
동일 득표쌍(collision) 분석 + 득표수 렌덤 노이즈(±1) 비교
- 관측치에 무작위로 +1 또는 -1의 노이즈를 주었을 때 동일쌍 수가 어떻게 변하는지 비교합니다.
"""
import json, sys, numpy as np
from collections import defaultdict, Counter
from itertools import combinations
sys.stdout.reconfigure(encoding='utf-8')

rng = np.random.default_rng(20260603)
recs = json.load(open('data/data_raw.json', encoding='utf-8'))
print(f'총 레코드: {len(recs)}')

# ----------------------------------------------------
# [추가] 득표수에 무작위로 +1 또는 -1을 더하는 함수
# ----------------------------------------------------
def apply_random_noise(records):
    perturbed_records = []
    for r in records:
        # 기존 record 복사
        new_r = r.copy()
        new_votes = []
        for v in r['votes']:
            if v is None:
                new_votes.append(None)
                continue

            # 무작위로 +1 또는 -1 선택
            noise = rng.choice([-1, 1])
            perturbed_v = v + noise

            # 득표수가 음수가 되지 않도록 방지
            if perturbed_v < 0:
                perturbed_v = 0
            new_votes.append(perturbed_v)

        new_r['votes'] = new_votes
        perturbed_records.append(new_r)
    return perturbed_records

# 1. 원본 데이터 그룹화
groups_original = defaultdict(list)
for r in recs:
    groups_original[(r['contest'], r['gubun'])].append(r)

# 2. 노이즈가 추가된 데이터 생성 및 그룹화
recs_perturbed = apply_random_noise(recs)
groups_perturbed = defaultdict(list)
for r in recs_perturbed:
    groups_perturbed[(r['contest'], r['gubun'])].append(r)

# ----------------------------------------------------
# 기존 함수들 (그대로 유지)
# ----------------------------------------------------
def pair_count(values):
    c = Counter(values)
    pairs = sum(n * (n - 1) // 2 for n in c.values() if n >= 2)
    groups = sum(1 for n in c.values() if n >= 2)
    return pairs, groups

def topk_key(vec, ranks):
    sv = sorted(vec, reverse=True)
    if len(sv) <= max(ranks):
        return None
    return tuple(sv[r] for r in ranks)

# target_groups를 인자로 받도록 수정된 분석 함수
def analyze_observed(gubun_filter, rankpairs, target_groups):
    out = {}
    for rp in rankpairs:
        pooled_vals = []
        per_contest_pairs = 0
        n_units = 0
        for (contest, gubun), rs in target_groups.items():
            if gubun != gubun_filter:
                continue
            vals = [topk_key(r['votes'], rp) for r in rs]
            vals = [v for v in vals if v is not None]
            n_units += len(vals)
            p, g = pair_count(vals)
            per_contest_pairs += p
            pooled_vals += vals
        pp, pg = pair_count(pooled_vals)
        out[rp] = dict(pooled_pairs=pp, pooled_groups=pg,
                        per_contest_pairs=per_contest_pairs, n_units=n_units)
    return out

def full_vector_collisions(gubun_filter, target_groups):
    per_contest = 0
    for (contest, gubun), rs in target_groups.items():
        if gubun != gubun_filter:
            continue
        p, g = pair_count([tuple(r['votes']) for r in rs])
        per_contest += p
    return per_contest

# ----------------------------------------------------
# 메인 실행부
# ----------------------------------------------------
if __name__ == '__main__':
    rankpairs = [(0, 1), (1, 2), (2, 3)]

    for gubun in ['관내사전투표', '선거일투표']:
        print('\n' + '=' * 70)
        print(f'[{gubun}] 결과 비교')
        print('=' * 70)

        # 원본 데이터 분석
        obs_orig = analyze_observed(gubun, rankpairs, groups_original)
        fv_orig = full_vector_collisions(gubun, groups_original)

        # 노이즈 데이터 분석
        obs_pert = analyze_observed(gubun, rankpairs, groups_perturbed)
        fv_pert = full_vector_collisions(gubun, groups_perturbed)

        for rp in rankpairs:
            o_orig = obs_orig[rp]
            o_pert = obs_pert[rp]

            print(f'📌 순위쌍 {tuple(r+1 for r in rp)} (단위수={o_orig["n_units"]}):')
            print(f'  - [원본] 지역구내 동일쌍: {o_orig["per_contest_pairs"]} | 전국풀링: {o_orig["pooled_pairs"]}')
            print(f'  - [노이즈] 지역구내 동일쌍: {o_pert["per_contest_pairs"]} | 전국풀링: {o_pert["pooled_pairs"]}')
            print(f'  - [변화량] 지역구내: {o_pert["per_contest_pairs"] - o_orig["per_contest_pairs"]} | 전국풀링: {o_pert["pooled_pairs"] - o_orig["pooled_pairs"]}')
            print('-' * 50)

        print(f'✨ 전체 득표벡터 완전일치 쌍 (지역구 내):')
        print(f'  - [원본]: {fv_orig} 쌍')
        print(f'  - [노이즈]: {fv_pert} 쌍 (변화량: {fv_pert - fv_orig})')