def check_weights(data_ps):
    """Check if sum of propensity score weights match sample size
    Args:
        data_ps (pandas.DataFrame): dataframe with propensity score
    Return:
        tuple: sample size, treated size from weigths, untreated size froms weigths
    """
    weight_t = 1./data_ps.query("nudge==1")["pscore"]
    weight_nt = 1./(1.-data_ps.query("nudge==0")["pscore"])
    print(f"Original sample size {data_ps.shape[0]}")
    print("Original treatment sample size", data_ps.query("nudge==1").shape[0])
    print("Original control sample size", data_ps.query("nudge==0").shape[0])
    print(f"Weighted treatment sample size {round(sum(weight_t), 1)}")
    print(f"Weighted control sample size {round(sum(weight_nt), 1)}")

    return data_ps.shape[0], sum(weight_t), sum(weight_nt)