import pandas as pd
import re
from datetime import datetime


def n_calculator(claims_df, start, end, 
                 cancer_pattern=None, brain_pattern=None, heart_pattern=None,
                 disease_pattern=None, injury_pattern=None):
    """
    claims_df를 받아 시작일(start)~종료일(end) 기간 내의 데이터를 대상으로
    cancer, brain, heart의 발생 ID 개수와 disease_sur, injury_sur의 청구 횟수를 계산.
    
    성능 최적화:
    - 날짜 변환 최소화 (한 번만 변환)
    - 벡터화 연산 활용
    - 정규표현식 컴파일 최소화
    - 불필요한 데이터 복사 최소화
    
    Parameters:
    -----------
    claims_df : pd.DataFrame
        청구 데이터프레임. 'ID', 'kcd', 'sdate', 'edate', 'sur' 컬럼이 있어야 함
    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:
    --------
    dict
        {
            'cancer_id_count': int,      # cancer 발생한 고유 ID 개수
            'brain_id_count': int,       # brain 발생한 고유 ID 개수
            'heart_id_count': int,       # heart 발생한 고유 ID 개수
            'disease_sur_count': int,    # disease_sur 청구 횟수 (sur=1인 경우만)
            'injury_sur_count': int      # injury_sur 청구 횟수 (sur=1인 경우만)
        }
    """
    if len(claims_df) == 0:
        return {
            'cancer_id_count': 0,
            'brain_id_count': 0,
            'heart_id_count': 0,
            'disease_sur_count': 0,
            'injury_sur_count': 0
        }
    
    # 날짜 변환 (한 번만)
    start_dt = pd.to_datetime(start)
    end_dt = pd.to_datetime(end)
    
    if start_dt > end_dt:
        raise ValueError("start 날짜가 end 날짜보다 늦을 수 없습니다.")
    
    # 날짜 변환 (한 번만)
    sdate_dt = pd.to_datetime(claims_df['sdate'])
    edate_dt = pd.to_datetime(claims_df['edate'])
    
    # 기간 내 데이터 필터링 (sdate <= end AND edate >= start)
    period_mask = (sdate_dt <= end_dt) & (edate_dt >= start_dt)
    
    if not period_mask.any():
        return {
            'cancer_id_count': 0,
            'brain_id_count': 0,
            'heart_id_count': 0,
            'disease_sur_count': 0,
            'injury_sur_count': 0
        }
    
    # 기간 내 데이터만 추출
    filtered_df = claims_df[period_mask].copy()
    
    # kcd 코드를 문자열로 변환 (한 번만)
    kcd_str = filtered_df['kcd'].astype(str)
    
    # 기본 패턴 설정
    if cancer_pattern is None:
        cancer_pattern = r'^C[0-9]'  # C로 시작하는 코드 (암)
    if brain_pattern is None:
        brain_pattern = r'^I6[0-9]'  # I6으로 시작하는 코드 (뇌혈관 질환)
    if heart_pattern is None:
        heart_pattern = r'^I[25][0-9]'  # I2, I5로 시작하는 코드 (심장 질환)
    if disease_pattern is None:
        disease_pattern = r'^[CIFAEJMG][0-9]'  # 질병 관련 코드
    if injury_pattern is None:
        injury_pattern = r'^[ST][0-9]'  # S, T로 시작하는 코드 (외상)
    
    # 정규표현식 컴파일 (한 번만)
    cancer_re = re.compile(cancer_pattern)
    brain_re = re.compile(brain_pattern)
    heart_re = re.compile(heart_pattern)
    disease_re = re.compile(disease_pattern)
    injury_re = re.compile(injury_pattern)
    
    # 각 패턴별 매칭 (벡터화 연산)
    cancer_match = kcd_str.apply(lambda x: bool(cancer_re.search(x)))
    brain_match = kcd_str.apply(lambda x: bool(brain_re.search(x)))
    heart_match = kcd_str.apply(lambda x: bool(heart_re.search(x)))
    disease_match = kcd_str.apply(lambda x: bool(disease_re.search(x)))
    injury_match = kcd_str.apply(lambda x: bool(injury_re.search(x)))
    
    # cancer, brain, heart: 발생한 고유 ID 개수 계산
    cancer_ids = filtered_df.loc[cancer_match, 'ID'].unique()
    brain_ids = filtered_df.loc[brain_match, 'ID'].unique()
    heart_ids = filtered_df.loc[heart_match, 'ID'].unique()
    
    cancer_id_count = len(cancer_ids)
    brain_id_count = len(brain_ids)
    heart_id_count = len(heart_ids)
    
    # disease, injury: sur=1인 청구 횟수 계산
    sur_mask = filtered_df['sur'] == 1.0
    
    disease_sur_count = ((disease_match) & (sur_mask)).sum()
    injury_sur_count = ((injury_match) & (sur_mask)).sum()
    
    return {
        'cancer_id_count': cancer_id_count,
        'brain_id_count': brain_id_count,
        'heart_id_count': heart_id_count,
        'disease_sur_count': int(disease_sur_count),
        'injury_sur_count': int(injury_sur_count)
    }

