import pandas as pd
import numpy as np
from datetime import timedelta
from joblib import Parallel, delayed
from numba import njit

from .glue_code import glue_code_


# =========================================================
# Numba JIT 엔진 : 날짜 병합 로직만 NumPy 기반으로 최적화
# =========================================================
@njit
def merge_engine_numba(sdates, edates, kcds, add_values, interval):
    n = len(sdates)

    results = []

    # 첫 row 초기화
    current_start = sdates[0]
    current_end = edates[0]
    kcd_buffer = [kcds[0]]
    day_count = (current_end - current_start).days + 1

    for i in range(1, n):
        # additional_col 값 비교
        same_value = True
        for j in range(add_values.shape[1]):
            if add_values[i, j] != add_values[i-1, j]:
                same_value = False
                break

        # merge 조건
        if sdates[i] <= current_end + np.timedelta64(interval + 1, 'D') and same_value:
            if edates[i] > current_end:
                current_end = edates[i]
            kcd_buffer.append(kcds[i])
            day_count += (edates[i] - sdates[i]).days + 1
        else:
            # merge 된 구간 저장 (Numba에서는 object append 불가 → tuple로 저장)
            results.append((current_start, current_end, list(kcd_buffer), day_count))

            # restart
            current_start = sdates[i]
            current_end = edates[i]
            kcd_buffer = [kcds[i]]
            day_count = (edates[i] - sdates[i]).days + 1

    # 마지막 구간 저장
    results.append((current_start, current_end, list(kcd_buffer), day_count))

    return results


def merge_group(group, additional_col, interval=0):

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

    sdates = group['sdate'].to_numpy(dtype='datetime64[D]')
    edates = group['edate'].to_numpy(dtype='datetime64[D]')
    kcds = group['kcd'].to_numpy(dtype=np.int64) if group['kcd'].dtype == 'int' else group['kcd'].to_numpy()

    add_values = group[additional_col].to_numpy()

    # Numba JIT 호출
    merged_blocks = merge_engine_numba(sdates, edates, kcds, add_values, interval)

    output = []
    for idx, (s, e, merged_kcd_list, stay) in enumerate(merged_blocks):

        last_row_index = max(0, min(idx, len(group)-1))

        row = {
            'ID': group['ID'].iloc[last_row_index],
            'gender': group['gender'].iloc[last_row_index],
            'age': group['age'].iloc[last_row_index],
            'age_band': group['age_band'].iloc[last_row_index],
            'kcd': glue_code_(merged_kcd_list),
            'sdate': pd.Timestamp(s),
            'edate': pd.Timestamp(e),
            'stay': stay
        }

        for col in additional_col:
            row[col] = group[col].iloc[last_row_index]

        output.append(row)

    return output


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)(group, additional_col, interval)
        for group in groups
    )

    merged_rows = [row for group_result in results for row in group_result]

    return pd.DataFrame(merged_rows)
