"""
고위험 고객 식별 모듈

다중 청구, 고위험 KCD 조합 등을 가진 
고위험 고객을 식별합니다.
"""

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


def identify_high_risk_customers(claims_df, ins_df, 
                                 udate, lookback_years=5,
                                 min_claims_threshold=5,
                                 high_risk_kcd_patterns=None):
    """
    고위험 고객을 식별.
    
    Parameters:
    -----------
    claims_df : pd.DataFrame
        청구 데이터프레임. 'ID', 'kcd', 'sdate', 'edate', 'sur', 'hosout', 'death' 컬럼 필요
    ins_df : pd.DataFrame
        가입자 데이터프레임. 'ID', 'gender', 'age', 'age_band' 컬럼 필요
    udate : str or datetime
        기준일
    lookback_years : int
        과거 몇 년간의 데이터를 볼지
    min_claims_threshold : int
        고위험으로 판단할 최소 청구 횟수
    high_risk_kcd_patterns : list of str, optional
        고위험 KCD 패턴 리스트 (정규표현식)
        예: [r'^C[0-9]', r'^I6[0-9]'] (암, 뇌질환)
    
    Returns:
    --------
    pd.DataFrame
        고위험 고객 목록 및 특성
        컬럼: 'ID', 'claim_count', 'surgery_count', 'hospitalization_count', 
              'has_death', 'high_risk_kcd_count', 'risk_score',
              'gender', 'age', 'age_band' (ins_df에서 가져온 정보)
        risk_score로 내림차순 정렬
    """
    udate_dt = pd.to_datetime(udate)
    lookback_start = udate_dt - relativedelta(years=lookback_years)
    
    # 기준일 이전 데이터만 필터링 (udate 미포함, udate 직전일까지만 포함)
    sdate_dt = pd.to_datetime(claims_df['sdate'])
    edate_dt = pd.to_datetime(claims_df['edate'])
    period_mask = (sdate_dt < udate_dt) & (edate_dt >= lookback_start)
    filtered_df = claims_df[period_mask].copy()
    
    if len(filtered_df) == 0:
        return pd.DataFrame()
    
    # ID별 집계
    id_stats = filtered_df.groupby('ID').agg({
        'kcd': 'count',  # 청구 횟수
        'sur': lambda x: (x == 1.0).sum(),  # 수술 횟수
        'hosout': lambda x: (x == 1.0).sum(),  # 입원 횟수
        'death': lambda x: (x == 1.0).any()  # 사망 여부
    }).rename(columns={
        'kcd': 'claim_count',
        'sur': 'surgery_count',
        'hosout': 'hospitalization_count',
        'death': 'has_death'
    })
    
    # 고위험 KCD 패턴 매칭
    if high_risk_kcd_patterns is not None:
        high_risk_count = pd.Series(0, index=id_stats.index, dtype=int)
        kcd_str = filtered_df['kcd'].astype(str)
        
        for pattern in high_risk_kcd_patterns:
            pattern_re = re.compile(pattern)
            high_risk_mask = kcd_str.apply(lambda x: bool(pattern_re.search(x)))
            high_risk_ids = filtered_df.loc[high_risk_mask, 'ID'].unique()
            # 고유 ID에 대해 카운트 증가
            for id_val in high_risk_ids:
                if id_val in high_risk_count.index:
                    high_risk_count[id_val] += 1
        
        id_stats['high_risk_kcd_count'] = high_risk_count
    else:
        id_stats['high_risk_kcd_count'] = 0
    
    # 고위험 고객 필터링
    high_risk_mask = id_stats['claim_count'] >= min_claims_threshold
    high_risk_customers = id_stats[high_risk_mask].copy()
    
    if len(high_risk_customers) == 0:
        return pd.DataFrame()
    
    # ins_df와 병합하여 고객 정보 추가
    customer_cols = ['ID']
    if 'gender' in ins_df.columns:
        customer_cols.append('gender')
    if 'age' in ins_df.columns:
        customer_cols.append('age')
    if 'age_band' in ins_df.columns:
        customer_cols.append('age_band')
    
    high_risk_customers = high_risk_customers.reset_index().merge(
        ins_df[customer_cols], 
        on='ID', 
        how='left'
    )
    
    # 위험도 점수 계산 (정규화)
    # 각 지표를 0-1로 정규화 후 가중 평균
    max_claim = high_risk_customers['claim_count'].max()
    max_surgery = high_risk_customers['surgery_count'].max()
    max_hosp = high_risk_customers['hospitalization_count'].max()
    max_kcd = high_risk_customers['high_risk_kcd_count'].max() if high_risk_kcd_patterns else 1
    
    if max_claim > 0:
        norm_claim = high_risk_customers['claim_count'] / max_claim
    else:
        norm_claim = 0
    
    if max_surgery > 0:
        norm_surgery = high_risk_customers['surgery_count'] / max_surgery
    else:
        norm_surgery = 0
    
    if max_hosp > 0:
        norm_hosp = high_risk_customers['hospitalization_count'] / max_hosp
    else:
        norm_hosp = 0
    
    if max_kcd > 0:
        norm_kcd = high_risk_customers['high_risk_kcd_count'] / max_kcd
    else:
        norm_kcd = 0
    
    high_risk_customers['risk_score'] = (
        norm_claim * 0.4 +
        norm_surgery * 0.3 +
        norm_hosp * 0.2 +
        norm_kcd * 0.05 +
        high_risk_customers['has_death'].astype(int) * 0.05
    )
    
    # 위험도 점수로 정렬
    high_risk_customers = high_risk_customers.sort_values('risk_score', ascending=False)
    
    return high_risk_customers

