import pandas as pd
from datetime import datetime


def n_hos_calculator(claims_df, start, end):
    """
    claims_df를 받아 시작일(start)~종료일(end) 기간 내의 입원 데이터를 대상으로
    'dis_inj' 컬럼의 'dis'와 'inj'에 해당하는 입원일수를 각각 합산.
    
    성능 최적화:
    - 날짜 변환 최소화 (한 번만 변환)
    - 벡터화 연산 활용
    - 불필요한 데이터 복사 최소화
    - 중복 제거를 위한 효율적인 그룹화
    
    Parameters:
    -----------
    claims_df : pd.DataFrame
        청구 데이터프레임. 'ID', 'hosout', 'dis_inj', 'sdate', 'edate' 컬럼이 있어야 함
    start : str or datetime
        시작일 (yyyy-mm-dd 형식의 문자열 또는 datetime 객체)
    end : str or datetime
        종료일 (yyyy-mm-dd 형식의 문자열 또는 datetime 객체)
    
    Returns:
    --------
    dict
        {
            'dis_hos_days': int,  # dis에 해당하는 입원일수 총합
            'inj_hos_days': int   # inj에 해당하는 입원일수 총합
        }
    """
    if len(claims_df) == 0:
        return {
            'dis_hos_days': 0,
            'inj_hos_days': 0
        }
    
    # 필수 컬럼 확인
    required_cols = ['ID', 'hosout', 'dis_inj', 'sdate', 'edate']
    missing_cols = [col for col in required_cols if col not in claims_df.columns]
    if missing_cols:
        raise ValueError(f"필수 컬럼이 없습니다: {missing_cols}")
    
    # 날짜 변환 (한 번만)
    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'])
    
    # 1. 기간 내 데이터 필터링 (sdate <= end AND edate >= start)
    period_mask = (sdate_dt <= end_dt) & (edate_dt >= start_dt)
    
    if not period_mask.any():
        return {
            'dis_hos_days': 0,
            'inj_hos_days': 0
        }
    
    # 2. 입원 데이터만 필터링 (hosout == 1.0)
    hos_mask = claims_df['hosout'] == 1.0
    
    # 3. 'dis_inj'='none' 제거
    dis_inj_mask = claims_df['dis_inj'].isin(['dis', 'inj'])
    
    # 모든 조건 결합
    combined_mask = period_mask & hos_mask & dis_inj_mask
    
    if not combined_mask.any():
        return {
            'dis_hos_days': 0,
            'inj_hos_days': 0
        }
    
    # 필터링된 데이터 추출
    filtered_df = claims_df[combined_mask].copy()
    
    # 이미 변환된 날짜를 활용하여 입원일수 계산 (edate - sdate + 1)
    filtered_sdate_dt = sdate_dt[combined_mask]
    filtered_edate_dt = edate_dt[combined_mask]
    
    filtered_df['hos_days'] = (
        (filtered_edate_dt.values - filtered_sdate_dt.values) / pd.Timedelta(days=1) + 1
    ).astype(int)
    
    # 4. 중복 제거: 같은 ID 내에서 'dis_inj', 'sdate', 'edate'가 완전히 일치하는 경우 하나만 유지
    # 한 입원 건에 여러 KCD 코드가 있어서 청구건이 여러 개인 경우를 방지
    # drop_duplicates를 사용하여 효율적으로 처리
    deduplicated_df = filtered_df.drop_duplicates(
        subset=['ID', 'dis_inj', 'sdate', 'edate'],
        keep='first'
    )
    
    # 5. dis와 inj별로 입원일수 합산
    dis_mask = deduplicated_df['dis_inj'] == 'dis'
    inj_mask = deduplicated_df['dis_inj'] == 'inj'
    
    dis_hos_days = deduplicated_df.loc[dis_mask, 'hos_days'].sum()
    inj_hos_days = deduplicated_df.loc[inj_mask, 'hos_days'].sum()
    
    return {
        'dis_hos_days': int(dis_hos_days),
        'inj_hos_days': int(inj_hos_days)
    }

