import pandas as pd
from .n_calculator import n_calculator
from .n_hos_calculator import n_hos_calculator


def RR_calculator(claims_df_A, claims_df_B, start, end,
                  cancer_pattern=None, brain_pattern=None, heart_pattern=None,
                  disease_pattern=None, injury_pattern=None):
    """
    claims_df_A와 claims_df_B를 받아 각 그룹의 위험도와 상대위험도를 계산.
    
    성능 최적화:
    - n_calculator와 n_hos_calculator 재사용
    - 벡터화 연산 활용
    
    Parameters:
    -----------
    claims_df_A : pd.DataFrame
        그룹 A의 청구 데이터프레임. 'ID', 'kcd', 'sdate', 'edate', 'sur', 'hosout', 'dis_inj' 컬럼이 있어야 함
    claims_df_B : pd.DataFrame
        그룹 B의 청구 데이터프레임. 'ID', 'kcd', 'sdate', 'edate', 'sur', 'hosout', 'dis_inj' 컬럼이 있어야 함
    start : str or datetime
        시작일 (yyyy-mm-dd 형식의 문자열 또는 datetime 객체)
    end : str or datetime
        종료일 (yyyy-mm-dd 형식의 문자열 또는 datetime 객체)
    cancer_pattern : str, optional
        cancer를 판별하는 정규표현식 패턴. 기본값: r'^C[0-9]'
    brain_pattern : str, optional
        brain을 판별하는 정규표현식 패턴. 기본값: r'^I6[0-9]'
    heart_pattern : str, optional
        heart를 판별하는 정규표현식 패턴. 기본값: r'^I[25][0-9]'
    disease_pattern : str, optional
        disease를 판별하는 정규표현식 패턴. 기본값: r'^[CIFAEJMG][0-9]'
    injury_pattern : str, optional
        injury를 판별하는 정규표현식 패턴. 기본값: r'^[ST][0-9]'
    
    Returns:
    --------
    pd.DataFrame
        위험도와 상대위험도를 정리한 테이블
        컬럼: '지표', 'A발생수', 'A모수', 'A위험도', 'B발생수', 'B모수', 'B위험도', 
              '차이', '상대위험도', '변화율'
        
        - A발생수/B발생수: 각 지표의 분자값 (ID수, 청구건수, 또는 입원일수)
        - A모수/B모수: 각 그룹의 고유 ID 수 (분모값)
        - A위험도/B위험도: 발생수 / 모수
        - 차이: B위험도 - A위험도
        - 상대위험도: B위험도 / A위험도
        - 변화율: (상대위험도 - 1) * 100
    """
    # 빈 데이터프레임 처리
    if len(claims_df_A) == 0 and len(claims_df_B) == 0:
        return pd.DataFrame({
            '지표': ['cancer', 'brain', 'heart', 'disease_sur', 'injury_sur', 'disease_hos', 'injury_hos'],
            'A발생수': [0] * 7,
            'A모수': [0] * 7,
            'A위험도': [0.0] * 7,
            'B발생수': [0] * 7,
            'B모수': [0] * 7,
            'B위험도': [0.0] * 7,
            '차이': [0.0] * 7,
            '상대위험도': [0.0] * 7,
            '변화율': [0.0] * 7
        })
    
    # 1. 각 그룹의 ID 개수 계산 (고유 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
    
    # 2. 각 그룹에 대해 n_calculator와 n_hos_calculator 실행
    # 그룹 A
    if len(claims_df_A) > 0:
        n_result_A = n_calculator(
            claims_df_A, start, end,
            cancer_pattern, brain_pattern, heart_pattern,
            disease_pattern, injury_pattern
        )
        hos_result_A = n_hos_calculator(claims_df_A, start, end)
    else:
        n_result_A = {
            'cancer_id_count': 0,
            'brain_id_count': 0,
            'heart_id_count': 0,
            'disease_sur_count': 0,
            'injury_sur_count': 0
        }
        hos_result_A = {
            'dis_hos_days': 0,
            'inj_hos_days': 0
        }
    
    # 그룹 B
    if len(claims_df_B) > 0:
        n_result_B = n_calculator(
            claims_df_B, start, end,
            cancer_pattern, brain_pattern, heart_pattern,
            disease_pattern, injury_pattern
        )
        hos_result_B = n_hos_calculator(claims_df_B, start, end)
    else:
        n_result_B = {
            'cancer_id_count': 0,
            'brain_id_count': 0,
            'heart_id_count': 0,
            'disease_sur_count': 0,
            'injury_sur_count': 0
        }
        hos_result_B = {
            'dis_hos_days': 0,
            'inj_hos_days': 0
        }
    
    # 3. 위험도 계산 (각 지표 / 해당 그룹의 ID 개수)
    # 그룹 A 위험도
    risk_A = {
        'cancer': n_result_A['cancer_id_count'] / n_A if n_A > 0 else 0.0,
        'brain': n_result_A['brain_id_count'] / n_A if n_A > 0 else 0.0,
        'heart': n_result_A['heart_id_count'] / n_A if n_A > 0 else 0.0,
        'disease_sur': n_result_A['disease_sur_count'] / n_A if n_A > 0 else 0.0,
        'injury_sur': n_result_A['injury_sur_count'] / n_A if n_A > 0 else 0.0,
        'disease_hos': hos_result_A['dis_hos_days'] / n_A if n_A > 0 else 0.0,
        'injury_hos': hos_result_A['inj_hos_days'] / n_A if n_A > 0 else 0.0
    }
    
    # 그룹 B 위험도
    risk_B = {
        'cancer': n_result_B['cancer_id_count'] / n_B if n_B > 0 else 0.0,
        'brain': n_result_B['brain_id_count'] / n_B if n_B > 0 else 0.0,
        'heart': n_result_B['heart_id_count'] / n_B if n_B > 0 else 0.0,
        'disease_sur': n_result_B['disease_sur_count'] / n_B if n_B > 0 else 0.0,
        'injury_sur': n_result_B['injury_sur_count'] / n_B if n_B > 0 else 0.0,
        'disease_hos': hos_result_B['dis_hos_days'] / n_B if n_B > 0 else 0.0,
        'injury_hos': hos_result_B['inj_hos_days'] / n_B if n_B > 0 else 0.0
    }
    
    # 4. 차이, 상대위험도, 변화율 계산
    difference = {}  # B위험도 - A위험도
    relative_risk = {}  # B위험도 / A위험도
    change_rate = {}  # (상대위험도 - 1) * 100
    
    for key in risk_A.keys():
        # 차이 계산
        difference[key] = risk_B[key] - risk_A[key]
        
        # 상대위험도 계산
        if risk_A[key] > 0:
            relative_risk[key] = risk_B[key] / risk_A[key]
        else:
            # A의 위험도가 0인 경우 처리
            if risk_B[key] > 0:
                relative_risk[key] = float('inf')  # 무한대
            else:
                relative_risk[key] = 0.0  # 둘 다 0
        
        # 변화율 계산 (상대위험도 - 1) * 100
        if relative_risk[key] == float('inf'):
            change_rate[key] = float('inf')
        else:
            change_rate[key] = (relative_risk[key] - 1) * 100
    
    # 5. 발생수 정리 (RR 계산에 사용된 분자값)
    count_A = {
        'cancer': n_result_A['cancer_id_count'],
        'brain': n_result_A['brain_id_count'],
        'heart': n_result_A['heart_id_count'],
        'disease_sur': n_result_A['disease_sur_count'],
        'injury_sur': n_result_A['injury_sur_count'],
        'disease_hos': hos_result_A['dis_hos_days'],
        'injury_hos': hos_result_A['inj_hos_days']
    }
    
    count_B = {
        'cancer': n_result_B['cancer_id_count'],
        'brain': n_result_B['brain_id_count'],
        'heart': n_result_B['heart_id_count'],
        'disease_sur': n_result_B['disease_sur_count'],
        'injury_sur': n_result_B['injury_sur_count'],
        'disease_hos': hos_result_B['dis_hos_days'],
        'injury_hos': hos_result_B['inj_hos_days']
    }
    
    # 6. 테이블 형식으로 정리
    indicators = ['cancer', 'brain', 'heart', 'disease_sur', 'injury_sur', 'disease_hos', 'injury_hos']
    
    result_df = pd.DataFrame({
        '지표': indicators,
        'A발생수': [count_A[ind] for ind in indicators],
        'A모수': [n_A] * 7,
        'A위험도': [risk_A[ind] for ind in indicators],
        'B발생수': [count_B[ind] for ind in indicators],
        'B모수': [n_B] * 7,
        'B위험도': [risk_B[ind] for ind in indicators],
        '차이': [difference[ind] for ind in indicators],
        '상대위험도': [relative_risk[ind] for ind in indicators],
        '변화율': [change_rate[ind] for ind in indicators]
    })
    
    # 무한대 값을 문자열로 표시 (보기 좋게)
    result_df['상대위험도'] = result_df['상대위험도'].apply(
        lambda x: 'Inf' if x == float('inf') else x
    )
    result_df['변화율'] = result_df['변화율'].apply(
        lambda x: 'Inf' if x == float('inf') else x
    )
    
    return result_df

