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


def kcd_counter(claims_df, kcd_patterns, udate, a, b,
                cancer_pattern=None, brain_pattern=None, heart_pattern=None,
                disease_pattern=None, injury_pattern=None):
    """
    claims_df에서 특정 kcd코드들을 udate 기준 과거 a년 내에 청구 기록(입원, 수술만)이 있었던 사람이,
    udate 기준 미래 b년에 cancer, brain, heart, disease_sur, injury_sur, disease_hos, injury_hos가
    각각 얼마나 나오는지를 kcd코드별로 정리하는 함수.
    
    성능 최적화:
    - 날짜 변환 최소화
    - 벡터화 연산 활용
    - n_calculator와 n_hos_calculator 재사용
    
    Parameters:
    -----------
    claims_df : pd.DataFrame
        청구 데이터프레임. 'ID', 'kcd', 'sdate', 'edate', 'sur', 'hosout', 'dis_inj' 컬럼이 있어야 함
    kcd_patterns : list of str
        조사하고 싶은 kcd 코드들의 정규표현식 리스트
        예: [r'^A00', r'^A01', r'^C[0-9]']
    udate : str or datetime
        기준일 (yyyy-mm-dd 형식의 문자열 또는 datetime 객체)
    a : int
        과거 a년 내의 데이터를 조회할 기간 (년)
    b : int
        미래 b년 내의 데이터를 카운트할 기간 (년)
    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
        kcd 코드별 결과 테이블
        컬럼: 'kcd', 'ID수', 'cancer', 'brain', 'heart', 'disease_sur', 'injury_sur', 'disease_hos', 'injury_hos'
    """
    if len(claims_df) == 0:
        # 빈 데이터프레임 반환
        return pd.DataFrame({
            'kcd': kcd_patterns if isinstance(kcd_patterns, list) else [],
            'ID수': [],
            'cancer': [],
            'brain': [],
            'heart': [],
            'disease_sur': [],
            'injury_sur': [],
            'disease_hos': [],
            'injury_hos': []
        })
    
    # 필수 컬럼 확인
    required_cols = ['ID', 'kcd', 'sdate', 'edate', 'sur', 'hosout']
    missing_cols = [col for col in required_cols if col not in claims_df.columns]
    if missing_cols:
        raise ValueError(f"필수 컬럼이 없습니다: {missing_cols}")
    
    # kcd_patterns가 리스트가 아닌 경우 리스트로 변환
    if not isinstance(kcd_patterns, list):
        kcd_patterns = [kcd_patterns]
    
    # 날짜 변환
    udate_dt = pd.to_datetime(udate)
    
    # 기간 설정
    past_start = udate_dt - relativedelta(years=a)
    future_end = udate_dt + relativedelta(years=b)
    
    # 전체 데이터의 날짜 변환 (한 번만)
    sdate_dt = pd.to_datetime(claims_df['sdate'])
    edate_dt = pd.to_datetime(claims_df['edate'])
    
    # kcd 코드를 문자열로 변환 (한 번만)
    kcd_str = claims_df['kcd'].astype(str)
    
    # 입원 또는 수술 조건 (한 번만 계산)
    hos_sur_mask = (claims_df['hosout'] == 1.0) | (claims_df['sur'] == 1.0)
    
    # 결과를 저장할 리스트
    results = []
    
    # 각 kcd 패턴에 대해 처리
    for kcd_pattern in kcd_patterns:
        # 1. 과거 a년 내에 해당 kcd 코드가 있고 입원/수술이 있었던 ID 목록 찾기
        # 기간 필터: past_start ~ udate (udate 미포함, udate 직전일까지만 포함)
        past_period_mask = (sdate_dt < udate_dt) & (edate_dt >= past_start)
        
        # kcd 코드 매칭
        kcd_re = re.compile(kcd_pattern)
        kcd_match = kcd_str.apply(lambda x: bool(kcd_re.search(str(x))))
        
        # 모든 조건 결합
        past_mask = past_period_mask & kcd_match & hos_sur_mask
        
        if not past_mask.any():
            # 조건을 만족하는 ID가 없으면 0으로 채운 행 추가
            results.append({
                'kcd': kcd_pattern,
                'ID수': 0,
                'cancer': 0,
                'brain': 0,
                'heart': 0,
                'disease_sur': 0,
                'injury_sur': 0,
                'disease_hos': 0,
                'injury_hos': 0
            })
            continue
        
        # 해당 ID 목록 추출 (고유 ID만)
        target_ids = claims_df.loc[past_mask, 'ID'].unique()
        target_id_set = set(target_ids)
        id_count = len(target_ids)
        
        # 2. 해당 ID들의 미래 b년 데이터만 필터링
        # 기간 필터: udate ~ future_end
        future_period_mask = (sdate_dt <= future_end) & (edate_dt >= udate_dt)
        
        # ID 필터
        id_mask = claims_df['ID'].isin(target_id_set)
        
        # 미래 기간 내의 해당 ID 데이터 추출
        future_df = claims_df[future_period_mask & id_mask].copy()
        
        if len(future_df) == 0:
            # 미래 데이터가 없으면 ID 수만 기록하고 나머지는 0
            results.append({
                'kcd': kcd_pattern,
                'ID수': id_count,
                'cancer': 0,
                'brain': 0,
                'heart': 0,
                'disease_sur': 0,
                'injury_sur': 0,
                'disease_hos': 0,
                'injury_hos': 0
            })
            continue
        
        # 3. n_calculator와 n_hos_calculator로 카운트
        n_result = n_calculator(
            future_df, udate, future_end,
            cancer_pattern, brain_pattern, heart_pattern,
            disease_pattern, injury_pattern
        )
        
        # dis_inj 컬럼이 있는지 확인 후 n_hos_calculator 실행
        if 'dis_inj' in future_df.columns:
            hos_result = n_hos_calculator(future_df, udate, future_end)
        else:
            # dis_inj 컬럼이 없으면 0으로 설정
            hos_result = {
                'dis_hos_days': 0,
                'inj_hos_days': 0
            }
        
        # 4. 결과 저장
        results.append({
            'kcd': kcd_pattern,
            'ID수': id_count,
            'cancer': n_result['cancer_id_count'],
            'brain': n_result['brain_id_count'],
            'heart': n_result['heart_id_count'],
            'disease_sur': n_result['disease_sur_count'],
            'injury_sur': n_result['injury_sur_count'],
            'disease_hos': hos_result['dis_hos_days'],
            'injury_hos': hos_result['inj_hos_days']
        })
    
    # 결과 데이터프레임 생성
    result_df = pd.DataFrame(results)
    
    return result_df

