import pandas as pd
import re
from datetime import datetime
from dateutil.relativedelta import relativedelta

# 미리 정의된 KCD 패턴 딕셔너리
_DEFAULT_KCD_PATTERNS = {
    'CANCER': r'^C[0-9]',
    'cancer': r'^C[0-9]',
    'BRAIN': r'^I6[0-9]',
    'brain': r'^I6[0-9]',
    'HEART': r'^I[0-5][0-9]|^I[7-9][0-9]',
    'heart': r'^I[0-5][0-9]|^I[7-9][0-9]',
}


def kcd_ith_mon(claims_df, udate, n_months, kcd_patterns, pattern_dict=None):
    """
    claims_df를 받아 최근 n개월 이내 해당 ID에 특정 kcd코드가 존재하는지 확인.
    ID별로 집계하여 반환.
    
    성능 최적화:
    - 날짜 변환 최소화 (한 번만 변환)
    - 벡터화 연산 활용
    - 정규표현식 매칭 최적화
    
    Parameters:
    -----------
    claims_df : pd.DataFrame
        청구 데이터프레임. 'ID', 'kcd', 'sdate', 'edate' 컬럼이 있어야 함
    udate : str or datetime
        기준 날짜 (yyyy-mm-dd 형식의 문자열 또는 datetime 객체)
    n_months : int or list of int
        확인할 개월 수. 예: 3 또는 [3, 6, 12]
    kcd_patterns : dict, str, or list
        kcd 코드 패턴. 다음 형태를 지원:
        - 딕셔너리: {'pattern_name': 'regex_pattern', ...}
        - 단일 문자열: 정규표현식 (컬럼명은 'kcd_pst_{n}m')
        - 리스트: ['CANCER', 'brain', 'heart'] (pattern_dict에서 찾아서 사용)
        예: {'brain': r'I6[0-9]', 'cancer': r'C[0-9]'} 또는 ['CANCER', 'brain']
    pattern_dict : dict, optional
        패턴 이름과 정규표현식을 매핑한 딕셔너리. None이면 기본 패턴 딕셔너리 사용.
        예: {'CANCER': r'^C[0-9]', 'brain': r'^I6[0-9]'}
    
    Returns:
    --------
    pd.DataFrame
        ID별 집계 결과. 컬럼: 'ID', '{kcd_name}_pst_{n}m' (kcd_name은 패턴 이름, n은 n_months 값)
        해당 kcd 코드가 있으면 1, 없으면 0
    """
    # n_months를 리스트로 변환
    if isinstance(n_months, int):
        n_months_list = [n_months]
    elif isinstance(n_months, (list, tuple)):
        n_months_list = list(n_months)
    else:
        raise ValueError("n_months는 int 또는 list of int여야 합니다.")
    
    # kcd_patterns 처리
    if isinstance(kcd_patterns, dict):
        # 딕셔너리 형태: {'name': 'pattern', ...}
        patterns_dict = kcd_patterns
    elif isinstance(kcd_patterns, str):
        # 단일 문자열: 기본 이름 'kcd' 사용
        patterns_dict = {'kcd': kcd_patterns}
    elif isinstance(kcd_patterns, (list, tuple)):
        # 리스트 형태: ['CANCER', 'brain', 'heart'] - pattern_dict에서 찾아서 사용
        pattern_lookup = pattern_dict if pattern_dict is not None else _DEFAULT_KCD_PATTERNS
        patterns_dict = {}
        for pattern_name in kcd_patterns:
            if pattern_name in pattern_lookup:
                patterns_dict[pattern_name] = pattern_lookup[pattern_name]
            else:
                raise ValueError(f"패턴 이름 '{pattern_name}'을 pattern_dict에서 찾을 수 없습니다. "
                               f"사용 가능한 패턴: {list(pattern_lookup.keys())}")
    else:
        raise ValueError("kcd_patterns는 dict, str, 또는 list여야 합니다.")
    
    # udate를 datetime으로 변환 (한 번만)
    if isinstance(udate, str):
        udate_dt = pd.to_datetime(udate)
    elif isinstance(udate, datetime):
        udate_dt = pd.to_datetime(udate)
    else:
        raise ValueError("udate는 문자열(yyyy-mm-dd) 또는 datetime 객체여야 합니다.")
    
    # 모든 고유 ID 가져오기 (claims_df의 모든 ID)
    all_ids = claims_df['ID'].unique()
    result_df = pd.DataFrame({'ID': all_ids})
    
    # 날짜 변환 (한 번만, 원본 DataFrame은 수정하지 않음)
    sdate_dt = pd.to_datetime(claims_df['sdate'])
    edate_dt = pd.to_datetime(claims_df['edate'])
    
    # kcd를 문자열로 변환 (한 번만, 재사용)
    kcd_str = claims_df['kcd'].astype(str)
    
    # 각 kcd 패턴별로 처리
    for kcd_name, kcd_pattern in patterns_dict.items():
        # kcd 코드 매칭 (벡터화 연산 - apply 대신 str.contains 사용)
        # str.contains는 내부적으로 최적화되어 apply보다 훨씬 빠름
        kcd_matched = kcd_str.str.contains(kcd_pattern, na=False, regex=True)
        
        # 각 n값에 대해 컬럼 생성
        for n in n_months_list:
            # n개월 이전 날짜 계산
            start_date = udate_dt - relativedelta(months=n)
            
            # 한 번에 필터링 (기존 코드 스타일처럼 boolean 마스크로 한 번에 처리)
            # sdate < udate AND edate >= start_date AND kcd 매칭
            # udate는 포함하지 않음 (udate 직전일까지만 포함)
            period_mask = (sdate_dt < udate_dt) & (edate_dt >= start_date) & kcd_matched
            
            # 해당 기간 내 해당 kcd 코드가 있는 ID들 (고유값) - 벡터화 연산
            ids_with_kcd = claims_df.loc[period_mask, 'ID'].unique()
            
            # 컬럼명 생성
            col_name = f'{kcd_name}_pst_{n}m'
            
            # ID별로 1 또는 0 할당 (벡터화 연산)
            result_df[col_name] = result_df['ID'].isin(ids_with_kcd).astype(int)
    
    return result_df

