"""
통계적 유의성 검정 모듈

RR_calculator 결과에 대한 통계적 검정을 수행하여 
그룹 간 차이의 유의성을 검증합니다.
"""

import pandas as pd
import numpy as np
from scipy import stats
from ..calculator import n_calculator, n_hos_calculator


def test_rr_significance(claims_df_A, claims_df_B, start, end,
                         cancer_pattern=None, brain_pattern=None, heart_pattern=None,
                         disease_pattern=None, injury_pattern=None,
                         alpha=0.05, method='ztest'):
    """
    RR_calculator 결과에 대한 통계적 유의성 검정을 수행.
    
    Parameters:
    -----------
    claims_df_A : pd.DataFrame
        그룹 A의 claims 데이터. 'ID', 'kcd', 'sdate', 'edate', 'sur', 'hosout', 'dis_inj' 컬럼 필요
    claims_df_B : pd.DataFrame
        그룹 B의 claims 데이터. 'ID', 'kcd', 'sdate', 'edate', 'sur', 'hosout', 'dis_inj' 컬럼 필요
    start : str or datetime
        분석 시작일
    end : str or datetime
        분석 종료일
    cancer_pattern : str, optional
        cancer 정규표현식 패턴
    brain_pattern : str, optional
        brain 정규표현식 패턴
    heart_pattern : str, optional
        heart 정규표현식 패턴
    disease_pattern : str, optional
        disease 정규표현식 패턴
    injury_pattern : str, optional
        injury 정규표현식 패턴
    alpha : float
        유의수준 (기본값: 0.05)
    method : str
        검정 방법 ('ztest' 또는 'chi2')
    
    Returns:
    --------
    pd.DataFrame
        각 지표별 통계적 검정 결과
        컬럼: '지표', '그룹A_발생수', '그룹A_전체수', '그룹B_발생수', '그룹B_전체수',
              '그룹A_위험도', '그룹B_위험도', 'p_value', '신뢰구간_하한', '신뢰구간_상한', '유의함'
    """
    # 각 그룹의 ID 개수
    n_A = claims_df_A['ID'].nunique() if len(claims_df_A) > 0 else 0
    n_B = claims_df_B['ID'].nunique() if len(claims_df_B) > 0 else 0
    
    if n_A == 0 or n_B == 0:
        raise ValueError("그룹 A 또는 B의 데이터가 없습니다.")
    
    # 각 그룹의 지표 계산
    n_result_A = n_calculator(claims_df_A, start, end,
                              cancer_pattern, brain_pattern, heart_pattern,
                              disease_pattern, injury_pattern)
    n_result_B = n_calculator(claims_df_B, start, end,
                              cancer_pattern, brain_pattern, heart_pattern,
                              disease_pattern, injury_pattern)
    
    # 입원일수 계산 (dis_inj 컬럼이 있는 경우만)
    if 'dis_inj' in claims_df_A.columns and 'dis_inj' in claims_df_B.columns:
        hos_result_A = n_hos_calculator(claims_df_A, start, end)
        hos_result_B = n_hos_calculator(claims_df_B, start, end)
    else:
        hos_result_A = {'dis_hos_days': 0, 'inj_hos_days': 0}
        hos_result_B = {'dis_hos_days': 0, 'inj_hos_days': 0}
    
    # 지표별 검정 결과 저장
    results = []
    indicators = [
        ('cancer', n_result_A['cancer_id_count'], n_result_B['cancer_id_count']),
        ('brain', n_result_A['brain_id_count'], n_result_B['brain_id_count']),
        ('heart', n_result_A['heart_id_count'], n_result_B['heart_id_count']),
        ('disease_sur', n_result_A['disease_sur_count'], n_result_B['disease_sur_count']),
        ('injury_sur', n_result_A['injury_sur_count'], n_result_B['injury_sur_count']),
        ('disease_hos', hos_result_A['dis_hos_days'], hos_result_B['dis_hos_days']),
        ('injury_hos', hos_result_A['inj_hos_days'], hos_result_B['inj_hos_days'])
    ]
    
    for indicator, count_A, count_B in indicators:
        if indicator in ['disease_hos', 'injury_hos']:
            # 입원일수의 경우 평균 차이 검정은 복잡하므로 
            # 여기서는 비율 검정 대신 단순 비교만 제공
            p_value = np.nan
            ci_lower = np.nan
            ci_upper = np.nan
            risk_A = count_A / n_A if n_A > 0 else 0
            risk_B = count_B / n_B if n_B > 0 else 0
        else:
            # 비율 검정 (Z-test 또는 Chi-square test)
            risk_A = count_A / n_A if n_A > 0 else 0
            risk_B = count_B / n_B if n_B > 0 else 0
            
            if method == 'ztest':
                # 두 비율의 차이에 대한 Z-test
                # 수동으로 Z-test 계산
                if count_A == 0 and count_B == 0:
                    zstat = 0
                    p_value = 1.0
                elif n_A == 0 or n_B == 0:
                    zstat = 0
                    p_value = 1.0
                else:
                    # 합동 비율 (pooled proportion)
                    p_pool = (count_A + count_B) / (n_A + n_B)
                    # 표준 오차
                    se_pool = np.sqrt(p_pool * (1 - p_pool) * (1/n_A + 1/n_B))
                    if se_pool > 0:
                        # Z-통계량
                        zstat = (risk_B - risk_A) / se_pool
                        # p-value (양측 검정)
                        p_value = 2 * (1 - stats.norm.cdf(abs(zstat)))
                    else:
                        zstat = 0
                        p_value = 1.0
                
                # 신뢰구간 계산
                if n_A > 0 and n_B > 0:
                    se = np.sqrt(risk_A * (1 - risk_A) / n_A + risk_B * (1 - risk_B) / n_B)
                    if se > 0:
                        z_critical = stats.norm.ppf(1 - alpha / 2)
                        diff = risk_B - risk_A
                        ci_lower = diff - z_critical * se
                        ci_upper = diff + z_critical * se
                    else:
                        ci_lower = np.nan
                        ci_upper = np.nan
                else:
                    ci_lower = np.nan
                    ci_upper = np.nan
            else:
                # Chi-square test
                if count_A == 0 and count_B == 0:
                    chi2 = 0
                    p_value = 1.0
                else:
                    contingency_table = np.array([[count_A, n_A - count_A],
                                                  [count_B, n_B - count_B]])
                    chi2, p_value, dof, expected = stats.chi2_contingency(contingency_table)
                ci_lower = np.nan
                ci_upper = np.nan
        
        results.append({
            '지표': indicator,
            '그룹A_발생수': count_A,
            '그룹A_전체수': n_A,
            '그룹B_발생수': count_B,
            '그룹B_전체수': n_B,
            '그룹A_위험도': risk_A,
            '그룹B_위험도': risk_B,
            'p_value': p_value,
            '신뢰구간_하한': ci_lower,
            '신뢰구간_상한': ci_upper,
            '유의함': p_value < alpha if not np.isnan(p_value) else False
        })
    
    return pd.DataFrame(results)

