"""
코호트 분석 모듈

연도별 또는 나이대별 코호트를 비교하여 
세대별/연령대별 위험도를 분석합니다.
"""

import pandas as pd
import numpy as np
from datetime import datetime
from ..calculator import n_calculator, n_hos_calculator


def analyze_cohort_risk(claims_df, ins_df, cohort_by='age_band',  # 'age_band' or 'year'
                       start_date=None, end_date=None,
                       cancer_pattern=None, brain_pattern=None, heart_pattern=None,
                       disease_pattern=None, injury_pattern=None):
    """
    코호트별 위험도를 분석.
    
    Parameters:
    -----------
    claims_df : pd.DataFrame
        청구 데이터프레임. 'ID', 'kcd', 'sdate', 'edate', 'sur', 'hosout' 컬럼 필요
    ins_df : pd.DataFrame
        가입자 데이터프레임. 'ID', cohort_by 컬럼 필요
    cohort_by : str
        코호트 구분 기준 ('age_band' 또는 'year')
        'year'의 경우 claims_df의 sdate에서 연도를 추출
    start_date : str or datetime, optional
        분석 시작일 (None이면 전체 기간)
    end_date : str or datetime, optional
        분석 종료일 (None이면 전체 기간)
    cancer_pattern : str, optional
        cancer 정규표현식 패턴
    brain_pattern : str, optional
        brain 정규표현식 패턴
    heart_pattern : str, optional
        heart 정규표현식 패턴
    disease_pattern : str, optional
        disease 정규표현식 패턴
    injury_pattern : str, optional
        injury 정규표현식 패턴
    
    Returns:
    --------
    pd.DataFrame
        코호트별 위험도 비교 결과
        컬럼: '코호트', 'ID수', 'cancer_risk', 'brain_risk', 'heart_risk', 
              'disease_sur_risk', 'injury_sur_risk', 
              'disease_hos_risk', 'injury_hos_risk'
    """
    # claims_df와 ins_df 병합
    if cohort_by == 'year':
        # sdate에서 연도 추출
        merged_df = claims_df.copy()
        merged_df['year'] = pd.to_datetime(merged_df['sdate']).dt.year
    else:
        # ins_df에서 코호트 정보 가져오기
        if cohort_by not in ins_df.columns:
            raise ValueError(f"ins_df에 '{cohort_by}' 컬럼이 없습니다.")
        # claims_df에 이미 해당 컬럼이 있으면 사용, 없으면 ins_df에서 가져오기
        if cohort_by in claims_df.columns:
            merged_df = claims_df.copy()
        else:
            merged_df = claims_df.merge(ins_df[['ID', cohort_by]], on='ID', how='left')
    
    # 기간 필터링
    if start_date is not None and end_date is not None:
        start_dt = pd.to_datetime(start_date)
        end_dt = pd.to_datetime(end_date)
        sdate_dt = pd.to_datetime(merged_df['sdate'])
        edate_dt = pd.to_datetime(merged_df['edate'])
        period_mask = (sdate_dt <= end_dt) & (edate_dt >= start_dt)
        merged_df = merged_df[period_mask].copy()
    
    # 코호트별 분석
    cohort_col = 'year' if cohort_by == 'year' else cohort_by
    cohorts = merged_df[cohort_col].dropna().unique()
    
    if len(cohorts) == 0:
        return pd.DataFrame()
    
    results = []
    
    for cohort in cohorts:
        cohort_df = merged_df[merged_df[cohort_col] == cohort].copy()
        
        if len(cohort_df) == 0:
            continue
        
        n_ids = cohort_df['ID'].nunique()
        
        if n_ids == 0:
            continue
        
        # 지표 계산
        if start_date is not None and end_date is not None:
            n_result = n_calculator(cohort_df, start_date, end_date,
                                   cancer_pattern, brain_pattern, heart_pattern,
                                   disease_pattern, injury_pattern)
            if 'dis_inj' in cohort_df.columns:
                hos_result = n_hos_calculator(cohort_df, start_date, end_date)
            else:
                hos_result = {'dis_hos_days': 0, 'inj_hos_days': 0}
        else:
            # 전체 기간 분석
            min_date = pd.to_datetime(cohort_df['sdate']).min()
            max_date = pd.to_datetime(cohort_df['edate']).max()
            n_result = n_calculator(cohort_df, min_date, max_date,
                                   cancer_pattern, brain_pattern, heart_pattern,
                                   disease_pattern, injury_pattern)
            if 'dis_inj' in cohort_df.columns:
                hos_result = n_hos_calculator(cohort_df, min_date, max_date)
            else:
                hos_result = {'dis_hos_days': 0, 'inj_hos_days': 0}
        
        # 위험도 계산
        results.append({
            '코호트': str(cohort),
            'ID수': n_ids,
            'cancer_risk': n_result['cancer_id_count'] / n_ids if n_ids > 0 else 0,
            'brain_risk': n_result['brain_id_count'] / n_ids if n_ids > 0 else 0,
            'heart_risk': n_result['heart_id_count'] / n_ids if n_ids > 0 else 0,
            'disease_sur_risk': n_result['disease_sur_count'] / n_ids if n_ids > 0 else 0,
            'injury_sur_risk': n_result['injury_sur_count'] / n_ids if n_ids > 0 else 0,
            'disease_hos_risk': hos_result['dis_hos_days'] / n_ids if n_ids > 0 else 0,
            'injury_hos_risk': hos_result['inj_hos_days'] / n_ids if n_ids > 0 else 0
        })
    
    if len(results) == 0:
        return pd.DataFrame()
    
    return pd.DataFrame(results)

