import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from pandas.api.types import is_datetime64_any_dtype


def kj_preprocessing(claims_df, ref_df, udate):
    """
    경증 코드 전처리 v2 - 4가지 조건을 모두 pass하는 청구건 식별 및 삭제
    
    4가지 조건 (모두 pass해야 삭제됨):
    1. 경증 코드 일치: ref_df에 있는 kcd코드이면 pass
    2. 경과일수 확인: udate 기준으로 경과월 이전에 발생한 입원/수술 청구건이면 pass (통원은 무조건 pass)
    3. 입원일수 확인: ID-kcd별 5년 이내 입원일수 합 <= ref_df 입원일수이면 pass (통원 및 5년 이상 지난 입원건은 무조건 pass)
    4. 수술횟수 확인: ID-kcd별 5년 이내 수술횟수 <= ref_df 수술횟수이면 pass (비수술 및 5년 이상 지난 수술건은 무조건 pass)
    
    Parameters:
    -----------
    claims_df : pd.DataFrame
        청구 데이터프레임. 'ID', 'kcd', 'sdate', 'edate', 'hosout', 'sur' 컬럼이 있어야 함
    ref_df : pd.DataFrame
        기준 데이터프레임. 'kcd', '경과월', '입원일수', '수술횟수' 컬럼이 있어야 함
    udate : str or datetime
        기준 날짜 (yyyy-mm-dd 형식의 문자열 또는 datetime 객체)
    
    Returns:
    --------
    pd.DataFrame
        원본 claims_df에 다음 컬럼들이 추가됨:
        - cond1_pass: 경증 코드 일치 여부
        - cond2_pass: 경과일수 조건 만족 여부
        - cond3_pass: 입원일수 조건 만족 여부
        - cond4_pass: 수술횟수 조건 만족 여부
        - to_delete: 4가지 조건을 모두 만족하여 삭제 대상인지 여부
    """
    if len(claims_df) == 0:
        result_df = claims_df.copy()
        result_df['cond1_pass'] = False
        result_df['cond2_pass'] = False
        result_df['cond3_pass'] = False
        result_df['cond4_pass'] = False
        result_df['to_delete'] = False
        return result_df
    
    # 필수 컬럼 확인
    required_claims_cols = ['ID', 'kcd', 'sdate', 'edate', 'hosout', 'sur']
    missing_cols = [col for col in required_claims_cols if col not in claims_df.columns]
    if missing_cols:
        raise ValueError(f"claims_df에 필수 컬럼이 없습니다: {missing_cols}")
    
    required_ref_cols = ['kcd', '경과월', '입원일수', '수술횟수']
    missing_cols = [col for col in required_ref_cols if col not in ref_df.columns]
    if missing_cols:
        raise ValueError(f"ref_df에 필수 컬럼이 없습니다: {missing_cols}")
    
    # 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 객체여야 합니다.")
    
    # 결과 데이터프레임 초기화 (원본 인덱스 보존)
    result_df = claims_df.copy()
    original_index = result_df.index
    
    # ref_df를 dict로 변환하여 빠른 조회
    ref_lookup = ref_df.set_index('kcd')[['경과월', '입원일수', '수술횟수']].to_dict('index')
    
    # 날짜 컬럼 변환
    sdate_dt = pd.to_datetime(result_df['sdate'])
    edate_dt = pd.to_datetime(result_df['edate'])
    
    # 5년 이전 날짜 계산 (조건3, 4에서 사용)
    five_years_ago = udate_dt - timedelta(days=365*5)
    
    # 기본 마스크 생성 (merge 전에 미리 생성)
    is_inpatient = (result_df['hosout'] == 1.0).values
    is_surgery = (result_df['sur'] == 1.0).values
    is_within_5years = (edate_dt >= five_years_ago).values
    
    # ==================== 조건1: 경증 코드 일치 ====================
    result_df['cond1_pass'] = result_df['kcd'].isin(ref_lookup.keys())
    
    # ==================== 조건2: 경과일수 확인 ====================
    # 통원은 무조건 pass, 입원/수술만 체크
    is_inpatient_or_surgery = is_inpatient | is_surgery
    
    # 각 kcd별 경과월 가져오기
    elapsed_months = result_df['kcd'].map(lambda x: ref_lookup.get(x, {}).get('경과월', np.inf))

    # udate에서 경과월만큼 빼기 (1개월 = 30일로 근사)
    threshold_dates = udate_dt - pd.to_timedelta(elapsed_months * 30, unit='D')
    
    # 청구건의 edate가 threshold_dates 이전이면 pass
    # 통원은 무조건 pass
    result_df['cond2_pass'] = (~is_inpatient_or_surgery) | (edate_dt < threshold_dates).values
    
    # ==================== 조건3: 입원일수 확인 ====================
    # 입원 청구건 필터링 (5년 이내)
    inpatient_within_5years = is_inpatient & is_within_5years
    
    if inpatient_within_5years.any():
        # ID-kcd별 입원일수 계산 (중복 제거 없이 단순 합산)
        inpatient_df = result_df[inpatient_within_5years].copy()
        inpatient_df['hos_days'] = (
            (pd.to_datetime(inpatient_df['edate']) - pd.to_datetime(inpatient_df['sdate'])).dt.days + 1
        )
        
        # ID-kcd별 입원일수 합산
        id_kcd_hos_days = inpatient_df.groupby(['ID', 'kcd'])['hos_days'].sum()
        
        # 벡터화 연산: merge를 사용하여 각 row에 입원일수 매핑
        temp_hos_days = pd.DataFrame({
            'ID': result_df['ID'].values,
            'kcd': result_df['kcd'].values,
            '_original_idx': original_index
        }).merge(
            id_kcd_hos_days.reset_index().rename(columns={'hos_days': '_id_kcd_hos_days'}),
            on=['ID', 'kcd'],
            how='left'
        ).set_index('_original_idx')['_id_kcd_hos_days'].fillna(0)
        
        # ref_df의 입원일수와 비교
        ref_hos_days = result_df['kcd'].map(
            lambda x: ref_lookup.get(x, {}).get('입원일수', np.inf)
        )
        
        # 입원 청구건: ID-kcd별 입원일수 <= ref 입원일수이면 pass
        # 통원 및 5년 이상 지난 입원건: 무조건 pass
        result_df['cond3_pass'] = (
            (~is_inpatient) |  # 통원은 무조건 pass
            (~is_within_5years) |  # 5년 이상 지난 입원건은 무조건 pass
            (temp_hos_days.values <= ref_hos_days.values)  # 입원일수 조건 만족
        )
    else:
        # 5년 이내 입원 청구건이 없으면 모두 pass
        result_df['cond3_pass'] = True
    
    # ==================== 조건4: 수술횟수 확인 ====================
    # 수술 청구건 필터링 (5년 이내)
    surgery_within_5years = is_surgery & is_within_5years
    
    if surgery_within_5years.any():
        # ID-kcd별 수술횟수 계산
        surgery_df = result_df[surgery_within_5years].copy()
        id_kcd_surgery_count = surgery_df.groupby(['ID', 'kcd']).size()
    
        # 벡터화 연산: merge를 사용하여 각 row에 수술횟수 매핑
        temp_surgery_count = pd.DataFrame({
            'ID': result_df['ID'].values,
            'kcd': result_df['kcd'].values,
            '_original_idx': original_index
        }).merge(
            id_kcd_surgery_count.reset_index().rename(columns={0: '_id_kcd_surgery_count'}),
            on=['ID', 'kcd'],
            how='left'
        ).set_index('_original_idx')['_id_kcd_surgery_count'].fillna(0)
        
        # ref_df의 수술횟수와 비교
        ref_surgery_count = result_df['kcd'].map(
            lambda x: ref_lookup.get(x, {}).get('수술횟수', np.inf)
        )
        
        # 수술 청구건: ID-kcd별 수술횟수 <= ref 수술횟수이면 pass
        # 비수술 및 5년 이상 지난 수술건: 무조건 pass
        result_df['cond4_pass'] = (
            (~is_surgery) |  # 비수술은 무조건 pass
            (~is_within_5years) |  # 5년 이상 지난 수술건은 무조건 pass
            (temp_surgery_count.values <= ref_surgery_count.values)  # 수술횟수 조건 만족
        )
    else:
        # 5년 이내 수술 청구건이 없으면 모두 pass
        result_df['cond4_pass'] = True
    
    # ==================== 최종 삭제 여부 판정 ====================
    result_df['to_delete'] = (
        result_df['cond1_pass'] &
        result_df['cond2_pass'] &
        result_df['cond3_pass'] &
        result_df['cond4_pass']
    )
    
    return result_df


def kj_preprocessing_and_remove(claims_df, ref_df, udate):
    """
    kj_preprocessing을 실행하고 to_delete=True인 row를 제거한 데이터프레임 반환
    
    Parameters:
    -----------
    claims_df : pd.DataFrame
        청구 데이터프레임
    ref_df : pd.DataFrame
        기준 데이터프레임
    udate : str or datetime
        기준 날짜
    
    Returns:
    --------
    pd.DataFrame
        to_delete=True인 row가 제거된 데이터프레임
        (cond1_pass, cond2_pass, cond3_pass, cond4_pass, to_delete 컬럼 포함)
    """
    result_df = kj_preprocessing(claims_df, ref_df, udate)
    return result_df[~result_df['to_delete']].reset_index(drop=True)
