from typing import Iterable, Optional

import pandas as pd
import numpy as np
from scipy.special import expit


# def get_expanded_df(df, event_type_col='J', duration_col='X', pid_col='pid'):
#     """
#     This function gets a dataframe describing each sample the time of the observed events,
#     and returns an expanded dataframe as explained in [1]-[2].
#     Right censoring is allowed and must be marked as event type 0.
#
#     :param df: original dataframe (pd.DataFrame)
#     :param event_type_col: event type column name (str)
#     :param duration_col: time column name (str)
#     :param pid_col: patient id column name (str)
#
#     :return: result_df: expanded dataframe
#
#     References:
#         [1] Meir, Tomer and Gorfine, Malka, "Discrete-time Competing-Risks Regression with or without Penalization", https://arxiv.org/abs/2303.01186
#         [2] Meir, Tomer and Gutman, Rom and Gorfine, Malka "PyDTS: A Python Package for Discrete-Time Survival (Regularized) Regression with Competing Risks", https://arxiv.org/abs/2204.05731
#     """

def get_expanded_df(
        df: pd.DataFrame,
        event_type_col: str = 'J',
        duration_col: str = 'X',
        pid_col: str = 'pid') -> pd.DataFrame:
    """
    Expands a discrete-time survival dataset into a long-format dataframe suitable for modeling. This function receives a dataframe where each row corresponds to a subject with observed  event type and duration. It returns an expanded dataframe where each subject is represented  by multiple rows, one for each time point up to their observed time. Right censoring is allowed and should be indicated by event type 0.

    Args:
        df (pd.DataFrame): Original input dataframe containing one row per subject.
        event_type_col (str): Name of the column indicating event type. Censoring is marked by 0.
        duration_col (str): Name of the column indicating event or censoring time.
        pid_col (str): Name of the column indicating subject/patient ID.

    Returns:
        pd.DataFrame: Expanded dataframe in long format, with one row per subject-time pair.
    """
    unique_times = df[duration_col].sort_values().unique()
    result_df = df.reindex(df.index.repeat(df[duration_col]))
    result_df[duration_col] = result_df.groupby(pid_col).cumcount() + 1
    # drop times that didn't happen
    result_df.drop(index=result_df.loc[~result_df[duration_col].isin(unique_times)].index, inplace=True)
    result_df.reset_index(drop=True, inplace=True)
    last_idx = result_df.drop_duplicates(subset=[pid_col], keep='last').index
    events = sorted(df[event_type_col].unique())
    result_df.loc[last_idx, [f'j_{e}' for e in events]] = pd.get_dummies(
        result_df.loc[last_idx, event_type_col]).values
    result_df[[f'j_{e}' for e in events]] = result_df[[f'j_{e}' for e in events]].fillna(0)
    result_df[f'j_0'] = 1 - result_df[[f'j_{e}' for e in events if e > 0]].sum(axis=1)
    return result_df


def compare_models_coef_per_event(first_model: pd.Series,
                                  second_model: pd.Series,
                                  real_values: np.array,
                                  event: int,
                                  first_model_label:str = "first",
                                  second_model_label:str = "second"
                                  ) -> pd.DataFrame:
    event_suffix = f"_{event}"
    assert (first_model.index == second_model.index).all(), "All index should be the same"
    models = pd.concat([first_model.to_frame(first_model_label),
                        second_model.to_frame(second_model_label)], axis=1)
    models.index += event_suffix
    real_values_s = pd.Series(real_values, index=models.index)

    return pd.concat([models, real_values_s.to_frame("real")], axis=1)


def present_coefs(res_dict):
    from IPython.display import display
    for coef_type, events_dict in res_dict.items():
        print(f"for coef: {coef_type.capitalize()}")
        df = pd.concat([temp_df for temp_df in events_dict.values()])
        display(df)


def get_real_hazard(df, real_coef_dict, times, events):
    a_t = {event: {t: real_coef_dict['alpha'][event](t) for t in times} for event in events}
    b = pd.concat([df.dot(real_coef_dict['beta'][j]) for j in events], axis=1, keys=events)

    for j in events:
        df[[f'hazard_j{j}_t{t}' for t in times]] = pd.concat([expit(a_t[j][t] + b[j]) for t in times],
                                                             axis=1).values
    return df


def assert_fit(event_df, times, event_type_col='J', duration_col='X'):
    if not event_df['success'].all():
        problematic_times = event_df.loc[~event_df['success'], duration_col].tolist()
        event = event_df[event_type_col].max()  # all the events in the dataframe are the same
        raise RuntimeError(f"Number of observed events at some time points are too small. Consider collapsing neighbor time points."
                           f"\n See https://tomer1812.github.io/pydts/UsageExample-RegroupingData/ for more details.")
    if not np.all([t in event_df[duration_col].values for t in times]):
        event = event_df[event_type_col].max()  # all the events in the dataframe are the same
        problematic_times = pd.Index(event_df[duration_col]).symmetric_difference(times).tolist()
        raise RuntimeError(f"Number of observed events at some time points are too small. Consider collapsing neighbor time points."
                           f"\n See https://tomer1812.github.io/pydts/UsageExample-RegroupingData/ for more details.")


def create_df_for_cif_plots(df: pd.DataFrame, field: str,
                            covariates: Iterable,
                            vals: Optional[Iterable] = None,
                            quantiles: Optional[Iterable] = None,
                            zero_others: Optional[bool] = False
                            ) -> pd.DataFrame:
    """
    This method creates df for cif plot, where it zeros

    Args:
        df (pd.DataFrame): Dataframe which we yield the statiscal propetrics (means, quantiles, etc) and stacture
        field (str): The field which will represent the change
        covariates (Iterable): The covariates of the given model
        vals (Optional[Iterable]): The values to use for the field
        quantiles (Optional[Iterable]): The quantiles to use as values for the field
        zero_others (bool): Whether to zero the other covarites or to zero them

    Returns:
        df (pd.DataFrame): A dataframe that contains records per value for cif ploting
    """

    cov_not_fitted = [cov for cov in covariates if cov not in df.columns]
    assert len(cov_not_fitted) == 0, \
        f"Required covariates are missing from df: {cov_not_fitted}"

    df_for_ploting = df.copy()
    if vals is not None:
        pass
    elif quantiles is not None:
        vals = df_for_ploting[field].quantile(quantiles).values
    else:
        raise NotImplemented("Only Quantiles or specific values is supported")
    temp_series = []
    template_s = df_for_ploting.iloc[0][covariates].copy()
    if zero_others:
        impute_val = 0
    else:
        impute_val = df_for_ploting[covariates].mean().values
    for val in vals:
        temp_s = template_s.copy()
        temp_s[covariates] = impute_val
        temp_s[field] = val
        temp_series.append(temp_s)

    return pd.concat(temp_series, axis=1).T