import pandas as pd
import numpy as np
from typing import Optional, List
from .in_ith_mon import in_ith_mon
from .in_ith_year import in_ith_year
from .kcd_ith_year import kcd_ith_year
from .merge_overlapping_date_range import merge_overlapping_date_range
from .make_merge_ins import make_merge_ins


def prepare_segmentation_columns(claims_df, ins_df, udate, si_levels=None):
    """
    split_date_range 완료된 claims_df와 ins_df를 받아서 
    세그멘테이션에 필요한 모든 컬럼을 자동 생성
    
    Parameters:
    -----------
    claims_df : pd.DataFrame
        split_date_range까지 완료된 청구 데이터
    ins_df : pd.DataFrame
        가입자 데이터 (ID 컬럼 필수)
    udate : str or datetime
        기준 날짜 (yyyy-mm-dd)
    si_levels : list of int, optional
        SI 레벨 (예: [6,5,4])
        기본값: None
        
    Returns:
    --------
    pd.DataFrame
        세그멘테이션에 필요한 모든 컬럼이 추가된 ins_df
        
    Examples:
    ---------
    >>> # std만 사용하는 경우
    >>> ins_with_cols = prepare_segmentation_columns(claims_df, ins_df, '2018-01-01')
    >>> 
    >>> # std + si_666, si_555, si_444 사용하는 경우
    >>> ins_with_cols = prepare_segmentation_columns(
    ...     claims_df, ins_df, '2018-01-01', si_levels=[6,5,4]
    ... )
    """
    print(f"\n[prepare_segmentation_columns] 세그멘테이션 컬럼 생성 시작")
    print(f"  기준 날짜: {udate}")
    print(f"  SI 레벨: {si_levels}")
    
    result_df = ins_df.copy()
    
    # 1. std_non용: clm_pst_5y
    print("\n  [1/6] clm_pst_5y 생성 중...")
    clm_result = in_ith_year(claims_df, udate, [5], ['out', 'hos', 'sur'])
    clm_result['clm_pst_5y'] = (
        (clm_result.filter(regex='_pst_5y').sum(axis=1) > 0).astype(int)
    )
    clm_result = clm_result[['ID', 'clm_pst_5y']]
    result_df = make_merge_ins(result_df, clm_result)
    
    # 2. std_clm용: out_pst_3m
    print("  [2/6] out_pst_3m 생성 중...")
    out_3m_result = in_ith_mon(claims_df, udate, [3], ['out'])
    result_df = make_merge_ins(result_df, out_3m_result)
    
    # 3. std_clm용: hos_pst_5y, sur_pst_5y
    print("  [3/6] hos_pst_5y, sur_pst_5y 생성 중...")
    std_clm_result = in_ith_year(claims_df, udate, [5], ['hos', 'sur'])
    result_df = make_merge_ins(result_df, std_clm_result)
    
    # 4. std_clm용: trt7_pst_5y
    print("  [4/6] trt7_pst_5y 생성 중...")
    trt7_result = merge_overlapping_date_range(claims_df, interval=7, threshold=7, udate=udate, years=5)
    result_df = make_merge_ins(result_df, trt7_result)
    
    # 5. std_clm용: ci10_pst_5y, decl_pst_5y
    print("  [5/6] ci10_pst_5y, decl_pst_5y 생성 중...")
    ci10_pattern = r'^C[0-9]'  # 암
    decl_pattern = r'^I[25][0-9]'  # 심장질환
    kcd_std_result = kcd_ith_year(
        claims_df, udate, [5], 
        {'ci10': ci10_pattern, 'decl': decl_pattern}
    )
    result_df = make_merge_ins(result_df, kcd_std_result)
    
    # 6. si_nnn용: hos_pst_ny, sur_pst_ny, ci6_pst_ny
    if si_levels is not None and len(si_levels) > 0:
        print(f"  [6/6] SI 레벨 {si_levels}에 대한 컬럼 생성 중...")
        
        # hos, sur
        si_hos_sur_result = in_ith_year(claims_df, udate, si_levels, ['hos', 'sur'])
        result_df = make_merge_ins(result_df, si_hos_sur_result)
        
        # ci6 (뇌혈관질환)
        ci6_pattern = r'^I6[0-9]'
        si_ci6_result = kcd_ith_year(claims_df, udate, si_levels, {'ci6': ci6_pattern})
        result_df = make_merge_ins(result_df, si_ci6_result)
    else:
        print("  [6/6] SI 레벨 없음 - 건너뜀")
    
    print(f"\n  ✓ 세그멘테이션 컬럼 생성 완료: {len(result_df.columns)}개 컬럼")
    return result_df


def make_segmentation(ins_df, si_levels=None):
    """
    ins_df를 가지고 각 기준에 따라 가장 우량한 고객부터 불량한 고객까지 세그멘테이션을 수행하는 함수.
    
    우량한 고객부터 순서대로 라벨링을 하며, 한번 라벨링 되면 다시 라벨링 되지 않음.
    라벨링 되지 않은 나머지는 'resid'로 라벨링.
    세그멘테이션은 ins_df에 'cat'이라는 컬럼을 추가하는 방식.
    
    기준(우량 순):
    1. 'std_non' / clm_pst_5y == 0 (udate 이전 5년 이내 어떠한 청구도 없는 ID)
    2. 'std_clm' / out_pst_3m==0, hos_pst_5y==0, sur_pst_5y==0, trt7_pst_5y<7, ci10_pst_5y==0, decl_pst_5y==0
       - trt7_pst_5y < 7: 연속 치료 기간이 7일 미만이어야 함
    3. 'si_nnn' / out_pst_3m==0, hos_pst_ny==0, sur_pst_ny==0, ci6_pst_ny==0 (n은 si_levels에서 지정)
    
    Parameters:
    -----------
    ins_df : pd.DataFrame
        가입자 데이터프레임. 'ID' 컬럼과 세그멘테이션에 필요한 컬럼들이 있어야 함.
        필요한 컬럼은 si_levels에 따라 달라짐:
        - std_non: clm_pst_5y
        - std_clm: out_pst_3m, hos_pst_5y, sur_pst_5y, trt7_pst_5y, ci10_pst_5y, decl_pst_5y
        - si_nnn: out_pst_3m, hos_pst_ny, sur_pst_ny, ci6_pst_ny (n은 si_levels 값)
    si_levels : list of int, optional
        SI 세그멘테이션에 사용할 레벨 (예: [6,5,4] -> si_666, si_555, si_444)
        기본값: None (SI 세그멘테이션 없이 std_non, std_clm만 사용)
    
    Returns:
    --------
    pd.DataFrame
        'cat' 컬럼이 추가된 ins_df 복사본.
        각 ID는 다음 중 하나의 값을 가짐: 'std_non', 'std_clm', 'si_nnn', 'resid'
    
    Examples:
    ---------
    >>> # std만 사용
    >>> ins_df_seg = make_segmentation(ins_df)
    >>> print(ins_df_seg['cat'].value_counts())
    std_non    500
    std_clm    300
    resid      200
    
    >>> # std + si_666, si_555, si_444 사용
    >>> ins_df_seg = make_segmentation(ins_df, si_levels=[6, 5, 4])
    >>> print(ins_df_seg['cat'].value_counts())
    std_non    500
    std_clm    300
    si_666     100
    si_555      50
    si_444      30
    resid       20
    """
    # ins_df 복사 (원본 수정 방지)
    result_df = ins_df.copy()
    
    # 'cat' 컬럼 초기화
    result_df['cat'] = None
    
    # 세그멘테이션 규칙 정의 (우량 순서대로)
    segmentation_rules = []
    
    # 1. std_non (고정)
    segmentation_rules.append({
        'name': 'std_non',
        'description': 'clm_pst_5y == 0 (5년 이내 어떠한 청구도 없음)',
        'conditions': ['clm_pst_5y']
    })
    
    # 2. std_clm (고정)
    segmentation_rules.append({
        'name': 'std_clm',
        'description': 'out_pst_3m==0, hos_pst_5y==0, sur_pst_5y==0, trt7_pst_5y<7, ci10_pst_5y==0, decl_pst_5y==0',
        'conditions': ['out_pst_3m', 'hos_pst_5y', 'sur_pst_5y', 'trt7_pst_5y', 'ci10_pst_5y', 'decl_pst_5y']
    })
    
    # 3. si_nnn (동적 생성, 9부터 1까지 내림차순)
    if si_levels is not None and len(si_levels) > 0:
        # 내림차순 정렬 (9가 가장 우량)
        sorted_levels = sorted(si_levels, reverse=True)
        
        for n in sorted_levels:
            segmentation_rules.append({
                'name': f'si_{n}{n}{n}',
                'description': f'out_pst_3m==0, hos_pst_{n}y==0, sur_pst_{n}y==0, ci6_pst_{n}y==0',
                'conditions': [
                    'out_pst_3m',
                    f'hos_pst_{n}y',
                    f'sur_pst_{n}y',
                    f'ci6_pst_{n}y'
                ]
            })
    
    # 필요한 컬럼 확인
    all_required_cols = set()
    for rule in segmentation_rules:
        all_required_cols.update(rule['conditions'])
    
    missing_cols = all_required_cols - set(result_df.columns)
    if missing_cols:
        raise ValueError(f"필수 컬럼이 없습니다: {sorted(missing_cols)}")
    
    # 각 세그멘테이션 규칙을 우량 순서대로 적용
    print(f"\n[make_segmentation] 세그멘테이션 수행")
    print(f"  전체 ID 수: {len(result_df):,}")
    print(f"  규칙 수: {len(segmentation_rules)}")
    print()
    
    for rule in segmentation_rules:
        rule_name = rule['name']
        conditions = rule['conditions']
        
        # 조건 마스크 생성
        mask = pd.Series(True, index=result_df.index)
        
        for col in conditions:
            if col not in result_df.columns:
                raise ValueError(f"컬럼 '{col}'가 없습니다.")
            
            # trt7_pst_5y는 특별 처리: 연속 치료 기간이 7일 미만이어야 조건 만족
            if col == 'trt7_pst_5y':
                # 값이 7 미만이고 결측값(99)이 아니어야 함
                col_mask = (result_df[col] < 7) & (~result_df[col].isna()) & (result_df[col] != 99)
            else:
                # 다른 컬럼은 정확히 0이어야 조건 만족
                col_mask = (result_df[col] == 0) | (result_df[col] == 0.0)
            
            mask = mask & col_mask
        
        # 아직 라벨링되지 않은 ID만 해당
        unlabeled_mask = result_df['cat'].isna()
        
        # 최종 마스크: 조건을 만족하고 아직 라벨링되지 않은 ID
        final_mask = mask & unlabeled_mask
        
        # 라벨링 수행
        result_df.loc[final_mask, 'cat'] = rule_name
        
        # 세그멘테이션별 ID 개수 출력
        count = final_mask.sum()
        print(f"  [{rule_name}] {count:,}개 ID ({count/len(result_df)*100:.2f}%): {rule['description']}")
    
    # 라벨링되지 않은 나머지는 'resid'로 처리
    resid_mask = result_df['cat'].isna()
    resid_count = resid_mask.sum()
    result_df.loc[resid_mask, 'cat'] = 'resid'
    
    if resid_count > 0:
        print(f"  [resid] {resid_count:,}개 ID ({resid_count/len(result_df)*100:.2f}%): 나머지")
    
    # 최종 세그멘테이션 요약
    print(f"\n[세그멘테이션 완료]")
    print(f"세그멘테이션별 분포:")
    cat_counts = result_df['cat'].value_counts()
    for cat, count in cat_counts.items():
        print(f"  {cat}: {count:,}개 ({count/len(result_df)*100:.2f}%)")
    
    return result_df


def make_segmentation_from_scratch(claims_df, ins_df, udate, si_levels=None):
    """
    원시 데이터(claims_df, ins_df)부터 세그멘테이션까지 한 번에 수행
    
    이 함수는 split_date_range까지 완료된 claims_df와 ins_df를 받아서
    필요한 모든 컬럼을 자동으로 생성하고 세그멘테이션을 수행합니다.
    
    Parameters:
    -----------
    claims_df : pd.DataFrame
        청구 데이터 (split_date_range까지만 완료된 상태)
        필요한 컬럼: ID, sdate, edate, hosout, sur, kcd
    ins_df : pd.DataFrame
        가입자 데이터
        필요한 컬럼: ID
    udate : str or datetime
        기준 날짜 (yyyy-mm-dd)
    si_levels : list of int, optional
        SI 레벨 (예: [6,5,4])
        기본값: None (std_non, std_clm만 사용)
        
    Returns:
    --------
    pd.DataFrame
        'cat' 컬럼이 추가된 세그멘테이션 결과
        
    Examples:
    ---------
    >>> # std만 사용
    >>> result = make_segmentation_from_scratch(claims_df, ins_df, '2018-01-01')
    >>> print(result['cat'].value_counts())
    std_non    500
    std_clm    300
    resid      200
    
    >>> # std + si_666~si_444 사용
    >>> result = make_segmentation_from_scratch(
    ...     claims_df, ins_df, '2018-01-01', si_levels=[6,5,4]
    ... )
    >>> print(result['cat'].value_counts())
    std_non    500
    std_clm    300
    si_666     100
    si_555      50
    si_444      30
    resid       20
    
    >>> # std + si_999~si_111 사용 (모든 레벨)
    >>> result = make_segmentation_from_scratch(
    ...     claims_df, ins_df, '2018-01-01', si_levels=[9,8,7,6,5,4,3,2,1]
    ... )
    """
    print("=" * 80)
    print("make_segmentation_from_scratch: 전체 프로세스 시작")
    print("=" * 80)
    
    # 1. 컬럼 자동 생성
    ins_df_with_cols = prepare_segmentation_columns(claims_df, ins_df, udate, si_levels)
    
    # 2. 세그멘테이션 수행
    result = make_segmentation(ins_df_with_cols, si_levels)
    
    print("=" * 80)
    print("make_segmentation_from_scratch: 전체 프로세스 완료")
    print("=" * 80)
    
    return result
