import pandas as pd
import numpy as np
from numba import njit
from numba.typed import List as NumbaList
from joblib import Parallel, delayed
from .glue_code import glue_code_


# ============================================================
# 🔥 NUMBA ENGINE: 날짜는 "일(int)"로 처리 (timedelta 에러 방지)
# ============================================================
@njit
def merge_engine_numba_unique(sdays, edays, kcd_codes, add_values, interval):
    n = len(sdays)
    results = NumbaList()

    current_start = sdays[0]
    current_end = edays[0]

    current_kcds = NumbaList()
    current_kcds.append(kcd_codes[0])

    current_row_ids = NumbaList()
    current_row_ids.append(0)

    max_gap = interval + 1  # 허용 간격(일)

    for i in range(1, n):

        # additional_col 동일 여부
        same = True
        for j in range(add_values.shape[1]):
            if add_values[i, j] != add_values[i - 1, j]:
                same = False
                break

        # 병합 조건: 날짜 이어짐 + 추가 컬럼 동일
        if (sdays[i] <= current_end + max_gap) and same:
            if edays[i] > current_end:
                current_end = edays[i]
            current_kcds.append(kcd_codes[i])
            current_row_ids.append(i)
        else:
            # block 종료
            results.append((current_start, current_end, current_kcds, current_row_ids))

            # 새 block 시작
            current_start = sdays[i]
            current_end = edays[i]

            current_kcds = NumbaList()
            current_kcds.append(kcd_codes[i])

            current_row_ids = NumbaList()
            current_row_ids.append(i)

    # 마지막 block
    results.append((current_start, current_end, current_kcds, current_row_ids))

    return results


# ============================================================
# 🔥 merge_group: 문자열 → factorize → numba → 고유 날짜 stay + 복원
# ============================================================
def merge_group(group, additional_col, interval=0):

    group = group.sort_values(by=["sdate"] + additional_col)

    # ---- 1) kcd factorize (문자열 → int 코드) ----
    kcd_codes, kcd_unique = pd.factorize(group["kcd"])
    kcd_unique = np.array(kcd_unique, dtype=object)

    # ---- 2) additional_col factorize ----
    add_unique_map = {}
    add_values_list = []

    for col in additional_col:
        codes, uniques = pd.factorize(group[col])
        add_unique_map[col] = np.array(uniques, dtype=object)
        add_values_list.append(codes)

    # numba는 (n,0) shape 싫어해서, 추가 컬럼 없으면 더미 1컬럼 생성
    if len(add_values_list) == 0:
        add_values = np.zeros((len(group), 1), dtype=np.int64)
    else:
        add_values = np.vstack(add_values_list).T.astype(np.int64)  # (n, p)

    # ---- 3) 날짜를 "일 단위 int"로 변환 ----
    # datetime64[ns] → datetime64[D] → int64(일)
    s_arr = group["sdate"].values.astype("datetime64[D]")
    e_arr = group["edate"].values.astype("datetime64[D]")
    sdays = s_arr.astype("int64")
    edays = e_arr.astype("int64")

    # ---- 4) numba 엔진 실행 ----
    blocks = merge_engine_numba_unique(
        sdays,
        edays,
        kcd_codes.astype(np.int64),
        add_values,
        int(interval),
    )

    # ---- 5) 블록 후처리: 고유 날짜 stay + 문자열 복원 ----
    out = []

    for (start_day, end_day, kcd_code_list, row_ids) in blocks:

        # kcd 복원
        decoded_kcd_list = [kcd_unique[i] for i in kcd_code_list]
        merged_kcd = glue_code_(decoded_kcd_list)

        # additional_col 복원 (블록 내 모든 row 동일하므로 첫 row 기준)
        restored_add_cols = {}
        first_row_idx = row_ids[0]
        for idx, col in enumerate(additional_col):
            restored_add_cols[col] = add_unique_map[col][add_values[first_row_idx, idx]]

        # stay = 고유 날짜 개수
        unique_days = set()
        for rid in row_ids:
            s_day = sdays[rid]
            e_day = edays[rid]
            d = s_day
            while d <= e_day:
                unique_days.add(d)
                d += 1

        stay_unique = len(unique_days)

        # 대표 행 정보: 첫 row 기준
        first_row = group.iloc[first_row_idx]

        row = {
            "ID": first_row["ID"],
            "gender": first_row["gender"],
            "age": first_row["age"],
            "age_band": first_row["age_band"],
            "kcd": merged_kcd,
            "sdate": pd.Timestamp(start_day, unit="D"),
            "edate": pd.Timestamp(end_day, unit="D"),
            "stay": stay_unique,
        }
        row.update(restored_add_cols)
        out.append(row)

    return out


# ============================================================
# 🔥 merge_overlapping_date_range: ID 단위 병렬 처리
# ============================================================
def merge_overlapping_date_range(df, additional_col=[], interval=0, n_jobs=-1):

    tmp = df.copy()
    tmp["sdate"] = pd.to_datetime(tmp["sdate"])
    tmp["edate"] = pd.to_datetime(tmp["edate"])

    groups = [g for _, g in tmp.groupby("ID")]

    results = Parallel(n_jobs=n_jobs)(
        delayed(merge_group)(g, additional_col, interval)
        for g in groups
    )

    final = [row for grp in results for row in grp]
    return pd.DataFrame(final)

