"""

Some utils for experiments

"""

# Author: Alicia Curth

from typing import Callable, Dict, Optional, Union

import jax.numpy as jnp

from catenets.models.jax import (

    DRNET_NAME,

    PSEUDOOUT_NAME,

    RANET_NAME,

    RNET_NAME,

    SNET1_NAME,

    SNET2_NAME,

    SNET3_NAME,

    SNET_NAME,

    T_NAME,

    XNET_NAME,

    PseudoOutcomeNet,

    get_catenet,

)

from catenets.models.jax.base import check_shape_1d_data

from catenets.models.jax.transformation_utils import (

    DR_TRANSFORMATION,

    PW_TRANSFORMATION,

    RA_TRANSFORMATION,

)

SEP = "_"

def eval_mse_model(

    inputs: jnp.ndarray,

    targets: jnp.ndarray,

    predict_fun: Callable,

    params: jnp.ndarray,

) -> jnp.ndarray:

    # evaluate the mse of a model given its function and params

    preds = predict_fun(params, inputs)

    return jnp.mean((preds - targets) ** 2)

def eval_mse(preds: jnp.ndarray, targets: jnp.ndarray) -> jnp.ndarray:

    preds = check_shape_1d_data(preds)

    targets = check_shape_1d_data(targets)

    return jnp.mean((preds - targets) ** 2)

def eval_root_mse(cate_pred: jnp.ndarray, cate_true: jnp.ndarray) -> jnp.ndarray:

    cate_true = check_shape_1d_data(cate_true)

    cate_pred = check_shape_1d_data(cate_pred)

    return jnp.sqrt(eval_mse(cate_pred, cate_true))

def eval_abs_error_ate(cate_pred: jnp.ndarray, cate_true: jnp.ndarray) -> jnp.ndarray:

    cate_true = check_shape_1d_data(cate_true)

    cate_pred = check_shape_1d_data(cate_pred)

    return jnp.abs(jnp.mean(cate_pred) - jnp.mean(cate_true))

def get_model_set(

    model_selection: Union[str, list] = "all", model_params: Optional[dict] = None

) -> Dict:

    """Helper function to retrieve a set of models"""

    # get model selection

    if type(model_selection) is str:

        if model_selection == "snet":

            models = get_all_snets()

        elif model_selection == "pseudo":

            models = get_all_pseudoout_models()

        elif model_selection == "twostep":

            models = get_all_twostep_models()

        elif model_selection == "all":

            models = dict(**get_all_snets(), **get_all_pseudoout_models())

        else:

            models = {model_selection: get_catenet(model_selection)()}  # type: ignore

    elif type(model_selection) is list:

        models = {}

        for model in model_selection:

            models.update({model: get_catenet(model)()})

    else:

        raise ValueError("model_selection should be string or list.")

    # set hyperparameters

    if model_params is not None:

        for model in models.values():

            existing_params = model.get_params()

            new_params = {

                key: val

                for key, val in model_params.items()

                if key in existing_params.keys()

            }

            model.set_params(**new_params)

    return models

ALL_SNETS = [T_NAME, SNET1_NAME, SNET2_NAME, SNET3_NAME, SNET_NAME]

ALL_PSEUDOOUT_MODELS = [DR_TRANSFORMATION, PW_TRANSFORMATION, RA_TRANSFORMATION]

ALL_TWOSTEP_MODELS = [DRNET_NAME, RANET_NAME, XNET_NAME, RNET_NAME]

def get_all_snets() -> Dict:

    model_dict = {}

    for name in ALL_SNETS:

        model_dict.update({name: get_catenet(name)()})

    return model_dict

def get_all_pseudoout_models() -> Dict:  # DR, RA, PW learner

    model_dict = {}

    for trans in ALL_PSEUDOOUT_MODELS:

        model_dict.update(

            {PSEUDOOUT_NAME + SEP + trans: PseudoOutcomeNet(transformation=trans)}

        )

    return model_dict

def get_all_twostep_models() -> Dict:  # DR, RA, R, X learner

    model_dict = {}

    for name in ALL_TWOSTEP_MODELS:

        model_dict.update({name: get_catenet(name)()})

    return model_dict