# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/methods/06_cchvae.ipynb.

# %% ../../nbs/methods/06_cchvae.ipynb 3
from __future__ import annotations
from ..import_essentials import *
from .base import CFModule, ParametricCFModule
from ..base import BaseConfig
from ..utils import auto_reshaping, grad_update, validate_configs, get_config
from ..data_utils import Feature, FeaturesList
from ..ml_model import MLP, MLPBlock
from ..data_module import DataModule
from keras.random import SeedGenerator

# %% auto 0
__all__ = ['CHVAE', 'CCHVAEConfig', 'CCHVAE']

# %% ../../nbs/methods/06_cchvae.ipynb 5
@keras.saving.register_keras_serializable()
class CHVAE(keras.Model):
    def __init__(
        self, 
        layers: list, 
        dropout_rate: float = 0., 
        **kwargs
    ):
        super().__init__()
        self.n_layers = layers
        self.dropout_rate = dropout_rate
        self.seed_generator = SeedGenerator(get_config().global_seed)

    def set_apply_constraints_fn(self, fn=None):
        if fn is None:
            fn = lambda x: x
        self.apply_constraints = fn

    def build(self, input_shape):
        encoder = keras.Sequential([
            MLPBlock(size, dropout_rate=self.dropout_rate) for size in self.n_layers[:-1]
        ])
        decoder = keras.Sequential([
            MLPBlock(size, dropout_rate=self.dropout_rate) for size in self.n_layers[::-1][1:]
        ])
        # Encoder
        self.mu_enc = keras.Sequential([encoder, keras.layers.Dense(self.n_layers[-1])])
        self.log_var_enc = keras.Sequential([encoder, keras.layers.Dense(self.n_layers[-1])])
        # Decoder
        self.mu_dec = keras.Sequential([
            decoder, keras.layers.Dense(input_shape[-1]), 
            keras.layers.BatchNormalization(),
        ])
        self.log_var_dec = keras.Sequential([
            decoder, keras.layers.Dense(input_shape[-1]), 
            keras.layers.BatchNormalization(),
        ])

    def encode(self, x, training=None):
        return self.mu_enc(x, training=training), self.log_var_enc(x, training=training)
    
    def decode(self, z, training=None):
        return self.mu_dec(z, training=training), self.log_var_dec(z, training=training)
    
    def reparameterize(self, mu, log_var):
        std = keras.ops.exp(0.5 * log_var)
        eps = keras.random.normal(std.shape, seed=self.seed_generator)
        return mu + eps * std
    
    def forward(self, x, training=None):
        mu_z, log_var_z = self.encode(x, training=training)
        z = self.reparameterize(mu_z, log_var_z)
        mu_x, log_var_x = self.decode(z, training=training)
        return mu_z, log_var_z, z, mu_x, log_var_x

    def call(self, inputs, training=None):
        x = inputs
        mu_z, log_var_z, z, mu_x, log_var_x = self.forward(x, training=training)
        # compute loss
        loss = self.compute_kl_loss(mu_z, log_var_z)
        self.add_loss(loss)
        reconstructed_x = mu_x
        # reconstructed_x = self.apply_constraints(inputs, mu_x, hard=not training)
        return reconstructed_x
    
    def regenerate(self, z):
        mu, log_var = self.decode(z)
        return mu
    
    def compute_kl_loss(self, mu, logvar):
        kl_loss = -0.5 * keras.ops.sum(1 + logvar - mu**2 - keras.ops.exp(logvar))
        return kl_loss

# %% ../../nbs/methods/06_cchvae.ipynb 6
def _hyper_sphere_coordindates(
    rng_key: jrand.PRNGKey, # Random number generator key
    x: Array, # Input instance with only continuous features. Shape: (1, n_features)
    n_samples: int,
    high: float, # Upper bound
    low: float, # Lower bound
    p_norm: int = 2 # Norm
):
    key_1, key_2 = jrand.split(rng_key)
    delta = jrand.normal(key_1, shape=(n_samples, x.shape[-1]))
    dist = jrand.normal(key_2, shape=(n_samples,)) * (high - low) + low
    norm_p = jnp.linalg.norm(delta, ord=p_norm, axis=1)
    d_norm = jnp.divide(dist, norm_p).reshape(-1, 1)  # rescale/normalize factor
    delta = jnp.multiply(delta, d_norm)
    candidates = x + delta
    return candidates
     

@ft.partial(jit, static_argnums=(3, 4, 5, 7))
def _cchvae(
    x: Array,
    rng_key: jrand.PRNGKey,
    y_target: Array,
    pred_fn: Callable[[Array], Array],
    max_steps: int,
    n_search_samples: int,
    step_size: float,
    chvae: CHVAE,
):
    """Counterfactual generation using CCHVAE."""
        
    @loop_tqdm(max_steps)
    def body_fn(i, state):
        count, candidate_cf, rng = state
        rng_key, subkey_1, subkey_2 = jrand.split(rng, num=3)
        low, high = step_size * count, step_size * (count + 1)
        # STEP 1 -- SAMPLE POINTS on hyper sphere around instance
        latent_neighbors = _hyper_sphere_coordindates(
            subkey_1, z_rep, n_search_samples, high=high, low=low, p_norm=1
        )
        x_ce = chvae.regenerate(latent_neighbors)
        x_ce = chvae.apply_constraints(x, x_ce.reshape(1, -1), hard=True)
        
        # STEP 2 -- COMPUTE l1 norms
        distances = jnp.abs(x_ce - x).sum(axis=1)

        # STEP 3 -- SELECT POINT with MINIMUM l1 norm
        y_candidates = pred_fn(x_ce).argmax(axis=1)
        indices = jnp.where(y_candidates == y_target, 1, 0).astype(bool)
        distances = jnp.where(indices, distances, jnp.inf)

        best_candidate_cf = x_ce[jnp.argmin(distances)].reshape(1, -1)
        
        candidate_cf = lax.cond(
            distances.min() < jnp.abs(x - candidate_cf).sum(axis=1).min(),
            lambda _: best_candidate_cf,
            lambda _: candidate_cf,
            None
        )

        count += 1
        return count, candidate_cf, rng_key
    
    y_target = y_target.reshape(1, -1).argmax(axis=1)
    z, _ = chvae.encode(x)
    # z_rep = jnp.repeat(z.reshape(1, -1), n_search_samples, axis=0)
    z_rep = z.reshape(1, -1)
    rng_key, _ = jrand.split(rng_key)
    # candidate_cf = jnp.array(x, copy=True)
    candidate_cf = jnp.ones_like(x) * jnp.inf
    state = (0, candidate_cf, rng_key) # (count, candidate_cf, rng_key)
    # count, candidate_cf, rng_key = jax.lax.while_loop(cond_fn, body_fn, state)
    count, candidate_cf, rng_key = lax.fori_loop(0, max_steps, body_fn, state)
    # if `inf` is found, return the original input
    candidate_cf = jnp.where(jnp.isinf(candidate_cf), x, candidate_cf)
    return candidate_cf


# %% ../../nbs/methods/06_cchvae.ipynb 7
class CCHVAEConfig(BaseConfig):
    vae_layers: List[int] = Field(
        [20, 16, 14, 12], description="List of hidden layer sizes for VAE."
    )
    opt_name: str = Field("adam", description="Optimizer name of VAE.")
    vae_lr: float = Field(0.001, description="Learning rate of VAE.")
    max_steps: int = Field(100, description="Max steps")
    n_search_samples: int = Field(100, description="Number of generated candidate counterfactuals.")
    step_size: float = Field(0.1, description="Step size")    

# %% ../../nbs/methods/06_cchvae.ipynb 8
class CCHVAE(ParametricCFModule):
    
    def __init__(self, config: Dict | CCHVAEConfig = None, chvae=None, name: str = 'cchvae'):
        if config is None:
            config = CCHVAEConfig()
        config = validate_configs(config, CCHVAEConfig)
        name = "CCHVAE" if name is None else name
        self.vae = chvae
        super().__init__(config, name=name)

    def _init_model(self, config: CCHVAEConfig, model: keras.Model):
        if model is None:
            model = CHVAE(self.config.vae_layers)
            model.compile(
                optimizer=keras.optimizers.get({
                    'class_name': config.opt_name, 
                    'config': {'learning_rate': config.vae_lr}
                }),
                loss=keras.losses.MeanSquaredError()
            )
        return model

    def train(
        self, 
        data: DataModule, 
        pred_fn: Callable = None,
        batch_size: int = 128,
        epochs: int = 10,
        **fit_kwargs
    ):
        if not isinstance(data, DataModule):
            raise ValueError(f"Expected `data` to be `DataModule`, "
                             f"got type=`{type(data).__name__}` instead.")
        X_train, y_train = data['train'] 
        self.vae = self._init_model(self.config, self.vae)
        self.vae.set_apply_constraints_fn(data.apply_constraints)
        self.vae.fit(
            X_train, X_train, 
            batch_size=batch_size, 
            epochs=epochs,
            **fit_kwargs
        )
        self._is_trained = True
        return self
    
    @auto_reshaping('x')
    def generate_cf(
        self,
        x: Array,
        pred_fn: Callable = None,
        y_target: Array = None,
        rng_key: jrand.PRNGKey = None,
        **kwargs
    ) -> Array:
        # TODO: Currently assumes binary classification.
        if y_target is None:
            y_target = 1 - pred_fn(x)
        else:
            y_target = jnp.array(y_target, copy=True)
        if rng_key is None:
            raise ValueError("`rng_key` must be provided, but got `None`.")
        
        return _cchvae(
            x,
            rng_key=rng_key,
            y_target=y_target,
            pred_fn=pred_fn,
            max_steps=self.config.max_steps,
            n_search_samples=self.config.n_search_samples,
            step_size=self.config.step_size,
            chvae=self.vae,
        )
