# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/methods/08_clue.ipynb.

# %% ../../nbs/methods/08_clue.ipynb 3
from __future__ import annotations
from ..import_essentials import *
from .base import CFModule, ParametricCFModule
from ..base import BaseConfig
from ..utils import auto_reshaping, validate_configs, get_config, grad_update
from ..ml_model import MLP, MLPBlock
from ..data_module import DataModule
from jax.scipy.stats.norm import logpdf as gaussian_logpdf
from keras.random import SeedGenerator


# %% auto 0
__all__ = ['Encoder', 'Decoder', 'kl_divergence', 'VAEGaussCat', 'CLUEConfig', 'get_reconstruction_loss_fn', 'CLUE']

# %% ../../nbs/methods/08_clue.ipynb 5
class Encoder(keras.layers.Layer):
    def __init__(self, sizes: List[int], dropout: float = 0.1):
        super().__init__()
        assert sizes[-1] % 2 == 0, f"sizes[-1] must be even, but got {sizes[-1]}"
        self.encoder = keras.Sequential([
            MLPBlock(size, dropout_rate=dropout) for size in sizes
        ])
    
    def call(self, x: Array, training: bool):
        params = self.encoder(x, training=training)
        d = params.shape[-1] // 2
        mu, sigma = params[:, :d], params[:, d:]
        sigma = jax.nn.softplus(sigma)
        sigma = jnp.clip(sigma, 1e-3)
        return mu, sigma

class Decoder(keras.layers.Layer):
    def __init__(self, sizes: List[int], output_size: int, dropout: float = 0.1):
        super().__init__()
        self.decoder = MLP(
            sizes, output_size=output_size, 
            dropout_rate=dropout, last_activation='sigmoid'
        )
    
    def __call__(self, z: Array, training: bool):
        mu_dec = self.decoder(z, training=training)
        return mu_dec

# %% ../../nbs/methods/08_clue.ipynb 7
@jit
def kl_divergence(p: Array, q: Array, eps: float = 2 ** -17) -> Array:
    loss_pointwise = p * (jnp.log(p + eps) - jnp.log(q + eps))
    return loss_pointwise

# %% ../../nbs/methods/08_clue.ipynb 8
class VAEGaussCat(keras.Model):
    def __init__(
        self, 
        enc_sizes: List[int] = [20, 16, 14, 12],
        dec_sizes: List[int] = [12, 14, 16, 20],
        dropout_rate: float = 0.1,    
    ):
        super().__init__()
        self.enc_sizes = enc_sizes
        self.dec_sizes = dec_sizes
        self.dropout_rate = dropout_rate
        self.seed_generator = SeedGenerator(get_config().global_seed)
        # default reconstruction loss to l2 loss
        self.reconstruction_loss = lambda x, y: optax.l2_loss(x, y).mean(-1)

    def set_reconstruction_loss(self, fn):
        self.reconstruction_loss = fn

    def build(self, input_shape):
        self.encoder = Encoder(self.enc_sizes, self.dropout_rate)
        self.decoder = Decoder(self.dec_sizes, input_shape[-1], self.dropout_rate)
    
    def encode(self, x, training=True):
        mu_z, var_z = self.encoder(x, training=training)
        return mu_z, var_z
    
    def sample_latent(self, rng_key, mean, var):
        key, _ = jax.random.split(rng_key)
        std = jnp.exp(0.5 * var)
        eps = jrand.normal(key, var.shape)
        return mean + eps * std
    
    def decode(self, z, training=True):
        reconstruct_x = self.decoder(z, training=training)
        return reconstruct_x        
    
    def sample_step(self, rng_key, mean, var, training=True):
        z = self.sample_latent(rng_key, mean, var)
        mu_x = self.decode(z, training=training)
        return mu_x
    
    def sample(self, x, mc_samples, training=True): # Shape: (mc_samples, batch_size, input_size)
        mean, var = self.encode(x, training=training)
        rng_key = self.seed_generator.next()
        keys = jax.random.split(rng_key, mc_samples)
        
        partial_sample_step = partial(
            self.sample_step, mean=mean, var=var, training=training
        )
        reconstruct_x = jax.vmap(partial_sample_step)(keys)
        return (mean, var, reconstruct_x)
        
    def compute_vae_loss(self, inputs, mu_z, logvar_z, reconstruct_x):
        kl_loss = -0.5 * (1 + logvar_z - jnp.power(mu_z, 2) - jnp.exp(logvar_z)).sum(-1)
        
        rec = self.reconstruction_loss(inputs, reconstruct_x.reshape(inputs.shape)).sum(-1)
        batchwise_loss = (rec + kl_loss) / inputs.shape[0]
        return batchwise_loss.mean()

    def call(self, inputs, training=True):
        mu_z, logvar_z, reconstruct_x = self.sample(inputs, mc_samples=1, training=training)
        loss = self.compute_vae_loss(inputs, mu_z, logvar_z, reconstruct_x)
        self.add_loss(loss)
        return reconstruct_x

# %% ../../nbs/methods/08_clue.ipynb 10
@ft.partial(jit, static_argnums=(3, 4, 6, 9, 12, 13))
def _clue_generate(
    x: Array,
    rng_key: jrand.PRNGKey,
    y_target: Array,
    pred_fn: Callable,
    max_steps: int,
    step_size: float,
    vae_module: VAEGaussCat,
    uncertainty_weight: float,
    aleatoric_weight: float,
    prior_weight: float,
    distance_weight: float,
    validity_weight: float,
    validity_fn: Callable,
    apply_fn: Callable
) -> Array:
    
    @jit
    def sample_latent_from_x(x: Array, rng_key: jrand.PRNGKey):
        key_1, key_2 = jrand.split(rng_key)
        mean, var = vae_module.encode(x, training=False)
        z = vae_module.sample_latent(key_2, mean, var)
        return z
    
    @ft.partial(jit, static_argnums=(1))
    def generate_from_z(z: Array, hard: bool):
        cf = vae_module.decode(z, training=False)
        cf = apply_fn(x, cf, hard=hard)
        return cf

    @jit
    def uncertainty_from_z(z: Array):
        cfs = generate_from_z(z, hard=False)
        pred_cfs = pred_fn(cfs)
        prob = pred_cfs[:, 1]
        total_uncertainty = -(prob * jnp.log(prob + 1e-10)).sum(-1)
        return total_uncertainty, cfs, pred_cfs
    
    @jit
    def compute_loss(z: Array):
        uncertainty, cfs, pred_cfs = uncertainty_from_z(z)
        loglik = gaussian_logpdf(z).sum(-1)
        dist = jnp.abs(cfs - x).mean()
        validity = validity_fn(y_target, pred_cfs).mean()
        loss = (
            (uncertainty_weight + aleatoric_weight) * uncertainty 
            + prior_weight * loglik
            + distance_weight * dist
            + validity_weight * validity
        )
        return loss.mean()
    
    @loop_tqdm(max_steps)
    def step(i, z_opt_state):
        z, opt_state = z_opt_state
        z_grad = jax.grad(compute_loss)(z)
        z, opt_state = grad_update(z_grad, z, opt_state, opt)
        return z, opt_state
    
    key_1, _ = jax.random.split(rng_key)
    z = sample_latent_from_x(x, key_1)
    opt = optax.adam(step_size)
    opt_state = opt.init(z)

    # Write a loop to optimize z using lax.fori_loop
    z, opt_state = lax.fori_loop(0, max_steps, step, (z, opt_state))
    cf = generate_from_z(z, hard=True)
    return cf

# %% ../../nbs/methods/08_clue.ipynb 12
class CLUEConfig(BaseConfig):
    enc_sizes: List[int] = Field(
        [20, 16, 14, 12], description="Sequence of Encoder layer sizes."
    )
    dec_sizes: List[int] = Field(
        [12, 14, 16, 20], description="Sequence of Decoder layer sizes."
    )
    dropout_rate: float = Field(0.1, description="Dropout rate")
    encoded_size: int = Field(5, description="Encoded size")
    lr: float = Field(0.001, description="Learning rate")
    max_steps: int = Field(500, description="Max steps")
    step_size: float = Field(0.01, description="Step size")
    vae_n_epochs: int = Field(10, description="Number of epochs for VAE")
    vae_batch_size: int = Field(128, description="Batch size for VAE")
    seed: int = Field(0, description="Seed for random number generator")

# %% ../../nbs/methods/08_clue.ipynb 13
def get_reconstruction_loss_fn(dm: DataModule):
    def reconstruction_loss(xs, cfs):
        losses = []
        for feat, (start, end) in dm.features.features_and_indices:
            if feat.is_categorical:
                losses.append(
                    optax.softmax_cross_entropy(cfs[:, start:end], xs[:, start:end]).reshape(-1, 1)
                )
            else:
                losses.append(optax.l2_loss(cfs[:, start:end], xs[:, start:end]))
        return jnp.concatenate(losses, axis=-1)
    
    return reconstruction_loss

# %% ../../nbs/methods/08_clue.ipynb 15
class CLUE(ParametricCFModule):

    def __init__(self, config: Dict | CLUEConfig = None, vae=None, name: str = 'CLUE'):
        if config is None:
            config = CLUEConfig()
        config = validate_configs(config, CLUEConfig)
        self.vae = vae
        super().__init__(config, name=name)

    def _init_model(self, config: CLUEConfig, model: VAEGaussCat):
        if model is None:
            model = VAEGaussCat(
                enc_sizes=config.enc_sizes, dec_sizes=config.dec_sizes, 
                dropout_rate=config.dropout_rate
            )
            model.compile(optimizer=keras.optimizers.Adam(config.lr), loss=None)
        return model
    
    def train(
        self, 
        data: DataModule, # data module
        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.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,
        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 _clue_generate(
            x, 
            rng_key=rng_key, 
            y_target=y_target,
            pred_fn=pred_fn,
            max_steps=self.config.max_steps,
            step_size=self.config.step_size,
            vae_module=self.vae,
            uncertainty_weight=.0,
            aleatoric_weight=0.0,
            prior_weight=0.0,
            distance_weight=.1,
            validity_weight=1.0,
            validity_fn=keras.losses.get({'class_name': 'KLDivergence', 'config': {'reduction': None}}),
            apply_fn=self.apply_constraints,
        )
