import pandas as pd
import numpy as np
import statsmodels.api as sm
from statsmodels.stats.mediation import Mediation

"""
Explanation: Mediation analysis with duration data
This notebook demonstrates mediation analysis when the
mediator and outcome are duration variables, modeled
using proportional hazards regression.  These examples
are based on simulated data.
End of explanation
"""


np.random.seed(3424)

"""
Explanation: Make the notebook reproducible.
End of explanation
"""


n = 1000

"""
Explanation: Specify a sample size.
End of explanation
"""


exp = np.random.normal(size=n)

"""
Explanation: Generate an exposure variable.
End of explanation
"""


def gen_mediator():
    mn = np.exp(exp)
    mtime0 = -mn * np.log(np.random.uniform(size=n))
    ctime = -2 * mn * np.log(np.random.uniform(size=n))
    mstatus = (ctime >= mtime0).astype(np.int)
    mtime = np.where(mtime0 <= ctime, mtime0, ctime)
    return mtime0, mtime, mstatus

"""
Explanation: Generate a mediator variable.
End of explanation
"""


def gen_outcome(otype, mtime0):
    if otype == "full":
        lp = 0.5*mtime0
    elif otype == "no":
        lp = exp
    else:
        lp = exp + mtime0
    mn = np.exp(-lp)
    ytime0 = -mn * np.log(np.random.uniform(size=n))
    ctime = -2 * mn * np.log(np.random.uniform(size=n))
    ystatus = (ctime >= ytime0).astype(np.int)
    ytime = np.where(ytime0 <= ctime, ytime0, ctime)
    return ytime, ystatus

"""
Explanation: Generate an outcome variable.
End of explanation
"""


def build_df(ytime, ystatus, mtime0, mtime, mstatus):
    df = pd.DataFrame({"ytime": ytime, "ystatus": ystatus,
                       "mtime": mtime, "mstatus": mstatus,
                       "exp": exp})
    return df

"""
Explanation: Build a dataframe containing all the relevant variables.
End of explanation
"""


def run(otype):

    mtime0, mtime, mstatus = gen_mediator()
    ytime, ystatus = gen_outcome(otype, mtime0)
    df = build_df(ytime, ystatus, mtime0, mtime, mstatus)

    outcome_model = sm.PHReg.from_formula("ytime ~ exp + mtime", status="ystatus", data=df)
    mediator_model = sm.PHReg.from_formula("mtime ~ exp", status="mstatus", data=df)

    med = Mediation(outcome_model, mediator_model, "exp", "mtime",
                    outcome_predict_kwargs={"pred_only": True})
    med_result = med.fit(n_rep=20)
    print(med_result.summary())

"""
Explanation: Run the full simulation and analysis, under a particular
population structure of mediation.
End of explanation
"""


run("full")

"""
Explanation: Run the example with full mediation
End of explanation
"""


run("partial")

"""
Explanation: Run the example with partial mediation
End of explanation
"""


run("no")

"""
Explanation: Run the example with no mediation
End of explanation
"""


