<filename>estimators/wald_estimator.py
from sklearn.linear_model import LinearRegression, LogisticRegression
import numpy as np
from sklearn.metrics import r2_score


def WaldEstimator(X, Y, Z, mode="bxby",return_predictive_score=False):

    # # adjust the probability range of Z to real
    # # if Z is a probability not binary label
    # # this will avoid adjusting extreme UAS
    # if len(set(Z))>2 and max(Z)<1 and min(Z)>0:
    #     Z = np.log(Z)-np.log(1-Z)

    if "bx" in mode.lower():
        clf1 = LogisticRegression(
            random_state=0, max_iter=10000, C=1e16, solver="lbfgs"
        )
    else:
        clf1 = LinearRegression()

    Z = Z.reshape(-1, 1)
    clf1.fit(Z, X)

    if "by" in mode.lower():
        clf2 = LogisticRegression(
            random_state=0, solver="lbfgs", tol=1e-10, max_iter=10000, C=1e16
        )
    else:
        clf2 = LinearRegression()

    clf2.fit(Z, Y)

    coef1 = clf1.coef_[0][0] if "bx" in mode.lower() else clf1.coef_[0]
    coef2 = clf2.coef_[0][0] if "by" in mode.lower() else clf2.coef_[0]
    
    if len(set(X))>2:
        if return_predictive_score is True:
            # conitnuous
            predictive_score = r2_score(X,clf1.predict(Z))
            return [coef2/coef1, predictive_score]
        else:
            return coef2/coef1
    elif np.all(X==Z.reshape(-1)):
        # observational association
        return coef2
    else:
        return coef2/coef1

import statsmodels.api as sm

def WaldEstimatorStats(X, Y, Z, mode="bxby", detail=False):

    # add intercept to the first column of Z
    Z = sm.add_constant(Z)

    # change the encoding of Y
    Y = (Y+1)/2

    # change the encoding of X if it is binary:
    if(len(set(X))==2):
        X = (X+1)/2

    if "bx" in mode.lower():
        clf1 = sm.Logit(X,Z)
    else:
        clf1 = sm.OLS(X,Z)

    clf1_results = clf1.fit(disp=0)

    if "by" in mode.lower():
        clf2 = sm.Logit(Y,Z)
    else:
        clf2 = sm.OLS(Y,Z)

    clf2_results = clf2.fit(disp=0)

    # coeffiicents
    coef1 = clf1_results.params[1]
    coef2 = clf2_results.params[1]

    # standard error
    se1 = clf1_results.bse[1]
    se2 = clf2_results.bse[1]

    if detail:
        return {
        "beta_x_z": coef1, 
        "beta_y_z": coef2,
        "se_x_z": se1,
        "se_y_z": se2
        }
    
    if np.all(X==Z.reshape(-1)):
        # observational association
        return coef2
    else:
        # wald ratio
        return coef2/coef1
    