"""Functions for the common machine learning utilities plus diffusion related methods."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/source/02_utils.ipynb.

# %% auto #0
__all__ = ['Loss', 'mix_gaussian_loss', 'kl_divergence', 'log_density_gaussian', 'matrix_log_density_gaussian',
           'log_importance_weight_matrix', 'get_log_pz_qz_prodzi_qzCx', 'BtcvaeLoss', 'GaussianMixtureMetric',
           'KLDMetric', 'ReconstructionMetric', 'DKLMetric', 'DWKLMetric', 'TCMetric', 'MIMetric', 'ShowLossCallback',
           'ShowKLDsCallback', 'ShowLatentsCallback', 'KLDsCallback', 'plot_klds', 'GMsCallback', 'save_model',
           'load_checkpoint', 'sig2D', 'D2sig']

# %% ../nbs/source/02_utils.ipynb #df667b48-20b3-4791-9313-1f7adf3bab4e
from fastai.vision.all import *
import andi_datasets
from .models import VAEWaveNet

# %% ../nbs/source/02_utils.ipynb #66cbf264-254f-4644-9533-81592221737b
from torch.distributions import Normal

# %% ../nbs/source/02_utils.ipynb #593db014-f461-4fbf-a791-a32de1429e80
# inspired from https://github.com/EtienneT/vae/blob/master/vae.ipynb
# little trick to input the beta parameter in the loss and make a schedule for fastai.
class Loss():
    """β-VAE loss function for a reconstruction term that is
    the negative log-likelihood of a mixture of Gaussians"""
    def __init__(self, RF:int, c:int, use_pad:bool=False, beta:float=0., reduction:str='sum', log=False): store_attr()
    def __call__(self, pred, target,):
        if   self.c > 0: y_hat, mu, logvar, c_, z = pred # VAE + WaveNet
        elif self.c== 0: y_hat                    = pred # WaveNet
        elif self.c==-1: y_hat, mu, logvar, z     = pred # VAE
        if self.use_pad:
            log_likelihood = mix_gaussian_loss(y_hat, target, reduction=self.reduction)
        else:
            log_likelihood = mix_gaussian_loss(y_hat, target[:,self.RF:], reduction=self.reduction)

        if self.reduction=='none':
            kl = kl_divergence(mu, logvar, reduction='none') if self.c else 0
            return log_likelihood, self.beta, kl
        else:
            kl = kl_divergence(mu, logvar, reduction='meansum') if self.c else 0
            return log_likelihood + self.beta*kl

        if self.log: self.last_components = (log_likelihood, kl)

    def set_hyper(self, k, v):
        self.k = v

# %% ../nbs/source/02_utils.ipynb #87336916-f3c1-489f-b6be-66fcadbdb67e
# we can easily define discretized version of the gaussian loss, however,
# use continuous version as same as the https://clarinet-demo.github.io/

def mix_gaussian_loss(y_hat:Tensor, # Predicted output (B x C x T)
                      y:Tensor,     # Target (B x T x 1)
                      log_scale_min:float=-12.0, # Log scale minimum value
                      reduction:str='sum'
                      )-> Tensor:   # Loss
    """Mixture of continuous Gaussian distributions loss.
    Note that it is assumed that input is scaled to [-1, 1]."""


    ##  # BxCxHxW -> BxCxT
    if y_hat.dim() == 4:
        y_hat = y_hat.reshape(*y_hat.shape[:2], -1)
    if y.dim() == 4:
        # BxCxHxW -> BxCxT
        y = y.reshape(*y.shape[:2], -1)
        # BxCxT -> BxTxC
        y = y.transpose(1, 2)


    assert y_hat.dim() == 3, f"y_hat {y_hat.dim()} " + f"{y_hat.shape}"
    C = y_hat.size(1)
    if C == 2:
        nr_mix = 1
    else:
        assert y_hat.size(1) % 3 == 0
        nr_mix = y_hat.size(1) // 3

    channel_dim = 1
    # BxCxT -> BxTxC
    y_hat = y_hat.transpose(1, 2)
    channel_dim = 2

    # unpack parameters.
    if C == 2:
        # special case for C == 2, just for compatibility
        logit_probs = None
        means = y_hat[:, :, 0:1]
        log_scales = torch.clamp(y_hat[:, :, 1:2], min=log_scale_min) # clamp removable?
    else:
        #  (B, T, num_mixtures) x 3
        logit_probs = y_hat[:, :, :nr_mix]
        means = y_hat[:, :, nr_mix:2 * nr_mix]
        log_scales = torch.clamp(y_hat[:, :, 2 * nr_mix:3 * nr_mix],
                                 min=log_scale_min)

    # repeat to match all the distributions in the mixture
    # B x T x 1 -> B x T x num_mixtures
    y = y.expand_as(means)

    centered_y = y - means
    dist = Normal(loc=0., scale=torch.exp(log_scales))

    log_probs = dist.log_prob(centered_y)

    if nr_mix > 1:
        log_probs = log_probs + F.log_softmax(logit_probs, -1) # mixture probabilities are added: log P(x_{t=1}^{(i=1)})+log P_{jt}^{(i)}

    if reduction == 'mean':
        if nr_mix == 1:
            # NLL
            return -torch.mean(log_probs) # sum in every dimension (B, T, num_mixtures) and divide by the number of elements, return an scalar
        else:
#             return -torch.sum(log_sum_exp(log_probs))
            return -torch.mean(torch.logsumexp(log_probs, channel_dim))
            # logsumexp in channel/mixtures dimension and then sum everything else
    elif reduction == 'sum':
        if nr_mix == 1:
            # NLL
            return -torch.sum(log_probs) # sum in every dimension (B, T, num_mixtures) and return an scalar
        else:
#             return -torch.sum(log_sum_exp(log_probs))
            return -torch.sum(torch.logsumexp(log_probs, channel_dim))
            # logsumexp in channel/mixtures dimension and then sum everything else
    elif reduction == 'sumtime':
        if nr_mix == 1:
            return -torch.sum(log_probs.mean(0)) # sum in time mean in batch (B, T, num_mixtures) and return an scalar
        else:
            return -torch.sum(torch.logsumexp(log_probs, channel_dim).mean(0))
    else:
        if nr_mix == 1:
            return -log_probs
        else:
#             return -log_sum_exp(log_probs).unsqueeze(-1)
            return -torch.logsumexp(log_probs, channel_dim).unsqueeze(-1)

# %% ../nbs/source/02_utils.ipynb #4c048512-6a8a-43c5-a3fc-59f152cee75e
def kl_divergence(mu:Tensor, logvar:Tensor, reduction:str='meansum'):
    """Divergence term of the VAE loss function."""
    klds = -0.5*(1. + logvar - mu.pow(2) - logvar.exp()) # batch_size, z_dim
    if reduction == 'meansum':
        # total_kld
        kld = klds.sum(1).mean(0)
        # sum of dimension_wise_kld # size 1
    elif reduction == 'mean':
        kld = klds.mean()
    elif reduction == 'sum':
        kld = klds.sum()
    elif reduction == 'meanb': # dimension_wise_kld
        kld = klds.mean(0)
    else:
        kld = klds

    # sanity check
    assert torch.all(torch.isfinite(kld)), f"Non finite value found on KLD loss computation: {kld.item()}."
    f'May {logvar=} be too big/small?'
    return kld

# %% ../nbs/source/02_utils.ipynb #1fc7b1cf-7e9d-4c4e-9220-564fbb048cd4
def log_density_gaussian(x, mu, logvar):
    """Calculates log density of a Gaussian."""
    normalization = -0.5 * (np.log(2 * np.pi) + logvar)
    inv_var = torch.exp(-logvar)
    log_density = normalization - 0.5 * ((x - mu) ** 2 * inv_var)
    return log_density

def matrix_log_density_gaussian(x, mu, logvar):
    """Calculates log density of a Gaussian for all batch pairs."""
    batch_size, dim = x.shape
    x = x[:, None, :]
    mu = mu[None, :, :]
    logvar = logvar[None, :, :]
    return log_density_gaussian(x, mu, logvar)

def log_importance_weight_matrix(batch_size: int, dataset_size: int):
    """Calculates a log importance weight matrix."""
    N = dataset_size
    M = batch_size - 1
    strat_weight = (N - M) / (N * M)
    W = torch.Tensor(batch_size, batch_size).fill_(1 / M)
    W.view(-1)[::M + 1] = 1 / N
    W.view(-1)[1::M + 1] = strat_weight
    W[M - 1, 0] = strat_weight
    return W.log()

def get_log_pz_qz_prodzi_qzCx(latent_sample, latent_dist, n_data, is_mss=True):
    batch_size, hidden_dim = latent_sample.shape

    # calculate log q(z|x)
    log_q_zCx = log_density_gaussian(latent_sample, *latent_dist).sum(dim=1)

    # calculate log p(z) (standard normal prior)
    zeros = torch.zeros_like(latent_sample)
    log_pz = log_density_gaussian(latent_sample, zeros, zeros).sum(1)

    mat_log_qz = matrix_log_density_gaussian(latent_sample, *latent_dist)

    if is_mss:
        # use stratification
        log_iw_mat = log_importance_weight_matrix(batch_size, n_data).to(latent_sample.device)
        mat_log_qz = mat_log_qz + log_iw_mat[:, :, None]

    log_qz = torch.logsumexp(mat_log_qz.sum(2), dim=1, keepdim=False)
    log_prod_qzi = torch.logsumexp(mat_log_qz, dim=1, keepdim=False).sum(1)

    return log_pz, log_qz, log_prod_qzi, log_q_zCx

# %% ../nbs/source/02_utils.ipynb #1c4a836f-a777-4a9e-864f-fc25bfaead85
class BtcvaeLoss():
    """Beta-TCVAE loss with modular gamma, KL, and MoG log-likelihood."""
    def __init__(self,
             RF, c, use_pad=False, alpha=1e-4, beta=3.0, gamma=0.1,
             reduction='sum', is_mss=True, log=True):
        store_attr()

    def __call__(self, pred, target):
        # Model output unpacking
        if   self.c > 0: y_hat, mu, logvar, c_, z = pred # VAE + WaveNet
        elif self.c== 0: y_hat                    = pred # WaveNet
        elif self.c==-1: y_hat, mu, logvar, z     = pred # VAE

        # Reconstruction loss
        if self.use_pad:
            recon_loss = mix_gaussian_loss(y_hat, target, reduction=self.reduction)
        else:
            recon_loss = mix_gaussian_loss(y_hat, target[:, self.RF:], reduction=self.reduction)

        # Total correlation loss (only for VAE)
        log_pz, log_qz, log_prod_qzi, log_q_zCx = get_log_pz_qz_prodzi_qzCx(latent_sample=z,
                                                                            latent_dist=(mu,logvar),
                                                                            n_data=z.shape[0],
                                                                            is_mss=self.is_mss,
                                                                           )
        # I[z;x] = KL[q(z,x)||q(x)q(z)] = E_x[KL[q(z|x)||q(z)]]
        mi_loss = (log_q_zCx - log_qz).mean()

        # TC[z] = KL[q(z)||\prod_i z_i]
        tc_loss = (log_qz - log_prod_qzi).mean()

        # dw_kl_loss is KL[q(z)||p(z)] instead of usual KL[q(z|x)||p(z))]
        dw_kl_loss = (log_prod_qzi - log_pz).mean()

        # total loss
        loss = recon_loss + (self.alpha * mi_loss +
                             self.gamma * tc_loss +
                             self.beta  * dw_kl_loss)

        # Attach components for Callback logging (optional)
        if self.log:  self.last_components = recon_loss.item(), dw_kl_loss.item(), tc_loss.item(), mi_loss.item()
        return loss

# %% ../nbs/source/02_utils.ipynb #1a9855c4-e59a-471c-a995-8a8663e3f15d
# inspired from https://github.com/EtienneT/vae/blob/master/vae.ipynb
# and https://github.com/fastai/fastai/blob/master/fastai/learner.py#L467
class GaussianMixtureMetric(AvgMetric):
    """Metric to log the Gaussian mixture loss"""
    def __init__(self, RF:int, c:int,
                 use_pad:bool=False, func=mix_gaussian_loss, reduction:str='mean'): store_attr()
    def accumulate(self, learn):
        bs = find_bs(learn.yb)
        if   self.c > 0: pred, mu, logvar, c_, z = to_detach(learn.pred, cpu=False) # VAE + WaveNet
        elif self.c== 0: pred                    = to_detach(learn.pred, cpu=False) # WaveNet
        elif self.c==-1: pred, mu, logvar, z     = to_detach(learn.pred, cpu=False) # VAE

        if self.use_pad:
            gaussian_mix_loss = to_detach(self.func(pred, learn.y, reduction=self.reduction).mean(), cpu=True)
        else:
            gaussian_mix_loss = to_detach(self.func(pred, learn.y[:,self.RF:], reduction=self.reduction).mean(), cpu=True)
        self.total += gaussian_mix_loss*bs
        self.count += bs

# %% ../nbs/source/02_utils.ipynb #259cec0e-7a14-445e-be9d-fd759617c6d9
class KLDMetric(AvgMetric):
    """Metric to log the Kullback-Leibler divergence term"""
    def __init__(self, c:int): store_attr()
    def accumulate(self, learn):
        bs = find_bs(learn.yb)
        if   self.c > 0: pred, mu, logvar, c_, z = to_detach(learn.pred, cpu=False) # VAE + WaveNet
        elif self.c==-1: pred, mu, logvar        = to_detach(learn.pred, cpu=False) # VAE
        KLD = -0.5*(1 + logvar - mu.pow(2) - logvar.exp()).sum(1).mean(0) if self.c and mu is not None else 0 # batch_size, z_dim
        KLD = to_detach(KLD, cpu=True)
        self.total += KLD*bs
        self.count += bs
    @property
    def name(self):  return 'kld'

# %% ../nbs/source/02_utils.ipynb #78c37d94-b965-44c7-8872-c0115498e142
class ReconstructionMetric(ValueMetric):
    def __init__(self, loss_func, metric_name='rec'): store_attr()
    @property
    def value(self):  return self.loss_func.last_components[0]
class DKLMetric(ValueMetric):
    def __init__(self, loss_func, metric_name='dkl'): store_attr()
    @property
    def value(self):  return self.loss_func.last_components[1]*self.loss_func.beta
class DWKLMetric(ValueMetric):
    def __init__(self, loss_func, metric_name='dw_kl'): store_attr()
    @property
    def value(self):  return self.loss_func.last_components[1]*self.loss_func.beta
class TCMetric(ValueMetric):
    def __init__(self, loss_func, metric_name='tc'): store_attr()
    @property
    def value(self):  return self.loss_func.last_components[2]*self.loss_func.gamma
class MIMetric(ValueMetric):
    def __init__(self, loss_func, metric_name='mi'): store_attr()
    @property
    def value(self):  return self.loss_func.last_components[3]*self.loss_func.alpha

# %% ../nbs/source/02_utils.ipynb #3a6132fc-983c-4071-8d26-4aeb75916a2f
class ShowLossCallback(Callback):
    "Update a graph of training and validation loss"
    order,run_valid=65,False
    def __init__(self, title:str=''): self.title=title
    def before_fit(self):
        self.run = not hasattr(self.learn, 'lr_finder') and not hasattr(self, "gather_preds")
        if not(self.run): return
        self.nb_batches = []
        assert hasattr(self.learn, 'progress')

    def after_train(self):  self.nb_batches.append(self.train_iter)

    def after_epoch(self):
        "Plot validation loss in the pbar graph"
        if not self.nb_batches: return
        rec = self.learn.recorder
        iters = range_of(rec.losses)
        val_losses = [v[1] for v in rec.values]
        # no bounds to see fluctuation in real time
        x_bounds = None # (0, (self.n_epoch - len(self.nb_batches)) * self.nb_batches[0] + len(rec.losses)) # None
        y_bounds = None # (min((min(Tensor(rec.losses)), min(Tensor(val_losses)))),
                    # max((max(Tensor(rec.losses)), max(Tensor(val_losses)))))
        self.progress.mbar.update_graph([(iters, rec.losses), (self.nb_batches, val_losses)], x_bounds, y_bounds)

        if IN_NOTEBOOK:
            self.progress.mbar.graph_ax.set_xlabel('Batches', x=0, horizontalalignment='right', verticalalignment='bottom')
            self.progress.mbar.graph_ax.set_ylabel('Loss')
            self.progress.mbar.graph_ax.set_title(self.title)
            # Epochs x axis
            # global batches_per_epoch
            if self.nb_batches:
                batches_per_epoch = self.nb_batches[0]# [-1]/self.n_epoch
                def epoch2batch(epoch):    return epoch * batches_per_epoch
                def batch2epoch(batch):    return batch / batches_per_epoch
                secax = self.progress.mbar.graph_ax.secondary_xaxis(-0.2, functions=(batch2epoch, epoch2batch))
                secax.set_xlabel('Epochs', x=0, horizontalalignment='right', verticalalignment='bottom');
                self.progress.mbar.graph_out.update(self.progress.mbar.graph_ax.figure)

# %% ../nbs/source/02_utils.ipynb #5f206a71-4b54-4638-9ba6-78d5c86b9025
class ShowKLDsCallback(Callback):
    "Update a graph of training and validation loss"
    order,run_valid=65,False
    names = ['']
    def __init__(self, title:str=''):    self.title=title

    def before_fit(self):
        self.run = not hasattr(self.learn, 'lr_finder') and not hasattr(self, "gather_preds")
        if not(self.run): return
        self.nb_batches = []
        assert hasattr(self.learn, 'progress')

    def after_validate(self):  self.on_iter_end()

    def after_fit(self):
        self.on_iter_end()
        if hasattr(self, 'graph_fig'): delattr(self, 'graph_fig')
        # needed to recreate the graph when fitting again
        # https://github.com/fastai/fastai/blob/master/fastai/callback/progress.py#L39

    def after_train(self):    self.on_iter_end();    self.nb_batches.append(self.train_iter)

    def on_iter_end(self):
        # from https://github.com/fastai/fastprogress/blob/master/fastprogress/fastprogress.py#L155
        if hasattr(self, 'imgs_fig'):
            plt.close()
            self.imgs_out.update(self.imgs_fig)
        if hasattr(self, 'graph_fig'):
            plt.close()
            self.graph_out.update(self.graph_fig)

    def after_epoch(self):
        "Plot validation loss in the pbar graph"
        if not self.nb_batches: return
        klds = np.stack(self.learn.kl_ds.preds)
        start_b = 0#17500
        # plt.yscale('symlog', linthresh=1e-8);
        # no bounds to see fluctuation in real time
        x_bounds = None # (0, (self.n_epoch - len(self.nb_batches)) * self.nb_batches[0] + len(rec.losses)) # None
        y_bounds = None # (min((min(Tensor(rec.losses)), min(Tensor(val_losses)))),
                    # max((max(Tensor(rec.losses)), max(Tensor(val_losses)))))
        graphs = [(range(start_b,len(klds)),klds[start_b:])]
        if IN_NOTEBOOK: self.update_graph(graphs, x_bounds, y_bounds)

    def update_graph(self, graphs, x_bounds=None, y_bounds=None, figsize=(6,4)):
        if not hasattr(self, 'graph_fig'):
            self.graph_fig, self.graph_ax = plt.subplots(1, figsize=figsize)
            self.graph_out = display(self.graph_ax.figure, display_id=True)
        self.graph_ax.clear()
        if len(self.names) < len(graphs): self.names += [''] * (len(graphs) - len(self.names))
        for g,n in zip(graphs,self.names): self.graph_ax.plot(*g, label=n)
        h,l = self.graph_ax.get_legend_handles_labels()
        if h: self.graph_ax.legend(h,l,loc='upper right')
        if x_bounds is not None: self.graph_ax.set_xlim(*x_bounds)
        if y_bounds is not None: self.graph_ax.set_ylim(*y_bounds)

        self.graph_ax.set_xlabel('Batches', x=0., horizontalalignment='right', verticalalignment='bottom')
        self.graph_ax.set_ylabel(r'DKL')
        self.graph_ax.grid()
        self.graph_ax.set_yscale('log')
        self.graph_ax.set_title(self.title)
        # Epochs x axis
        # global batches_per_epoch
        if self.nb_batches:
            batches_per_epoch = self.nb_batches[0]
            def epoch2batch(epoch):    return epoch * batches_per_epoch
            def batch2epoch(batch):    return batch / batches_per_epoch
            secax = self.graph_ax.secondary_xaxis(-0.2, functions=(batch2epoch, epoch2batch))
            secax.set_xlabel('Epochs', x=0., horizontalalignment='right', verticalalignment='bottom');
        self.graph_out.update(self.graph_ax.figure)

# %% ../nbs/source/02_utils.ipynb #d650c358-f166-4604-bf88-7c16d303de40
class ShowLatentsCallback(ShowKLDsCallback):
    "Update a graph of latent space"
    order,run_valid=65,False
    names = ['']
    def __init__(self, c, title:str=''): store_attr()

    def before_fit(self):
        self.run = not hasattr(self.learn, 'lr_finder') and not hasattr(self, "gather_preds")
        if not(self.run): return
        self.nb_batches = []
        assert hasattr(self.learn, 'progress')
        # get idx
        alphas_items = self.learn.dls.valid.items[:,0]
        Ds_items     = self.learn.dls.valid.items[:,1]
        u_a=np.unique(alphas_items, return_index=True, return_inverse=True)
        u_D=np.unique(Ds_items, return_index=True, return_inverse=True)
        alphas_idx = [np.flatnonzero(alphas_items==a) for a in u_a[0]]
        Ds_idx = [np.flatnonzero(Ds_items==D) for D in u_D[0]]

        intersect_idx = []
        for i,a in enumerate(u_a[0]):
            for j,D in enumerate(u_D[0]):
                intersect_idx.append(np.intersect1d(alphas_idx[i],Ds_idx[j]))

        self.intersect_idx = np.array(intersect_idx).reshape(len(u_a[0]),len(u_D[0]))
        self.alphas_idx_flat = [item for sublist in alphas_idx for item in sublist]
        self.Ds_idx_flat = [item for sublist in Ds_idx for item in sublist]

    def after_epoch(self):
        "Plot validation loss in the pbar graph"
        if not self.nb_batches: return

        cb = GatherPredsCallback(with_targs=False)
        ctx_mgrs = self.validation_context(cbs=L(cb), inner=True)
        dl = self.learn.dls[1].new(shuffle=False, drop_last=False) if self.learn.dls[1].drop_last else self.learn.dls[1]
        with ContextManagers(ctx_mgrs):
            self._do_epoch_validate(dl=dl)
            preds, _ = cb.all_tensors()
        if self.c >0: pred, mu, logvar, c, z = to_detach(preds)
        elif self.c==-1: pred, mu, logvar    = to_detach(preds)

        # plot
        # no bounds to see fluctuation in real time
        x_bounds = None; y_bounds = None
        graphs = [(range(len(self.alphas_idx_flat)),mu[self.alphas_idx_flat])]
        if IN_NOTEBOOK: self.update_graph(graphs, x_bounds, y_bounds)

    def update_graph(self, graphs, x_bounds=None, y_bounds=None, figsize=(6,4)):
        if not hasattr(self, 'graph_fig'):
            self.graph_fig, self.graph_ax = plt.subplots(1, figsize=figsize)
            self.graph_out = display(self.graph_ax.figure, display_id=True)
        self.graph_ax.clear()
        if len(self.names) < len(graphs): self.names += [''] * (len(graphs) - len(self.names))
        for g,n in zip(graphs,self.names): self.graph_ax.plot(*g, label=n)
        h,l = self.graph_ax.get_legend_handles_labels()
        if h: self.graph_ax.legend(h,l,loc='upper right')
        if x_bounds is not None: self.graph_ax.set_xlim(*x_bounds)
        if y_bounds is not None: self.graph_ax.set_ylim(*y_bounds)

        self.graph_ax.set_xlabel('ordered a index',)
        self.graph_ax.set_ylabel(r'\(\mu\) latent neuron activation')
        self.graph_ax.grid()
        # self.graph_ax.set_yscale('log')
        self.graph_ax.set_title(self.title)
        self.graph_out.update(self.graph_ax.figure)

# %% ../nbs/source/02_utils.ipynb #bd530d50-ef79-4a8e-b76b-c6ca6920a0ca
class KLDsCallback(Callback):
    "Record KLD per latent variable"
    order,run_valid=65,False
    def __init__(self,c:int): store_attr(); self.preds = []
    def after_batch(self):
        if   self.c > 0: pred, mu, logvar, c_, z = to_detach(self.learn.pred, cpu=False) # VAE + WaveNet
        elif self.c==-1: pred, mu, logvar        = to_detach(self.learn.pred, cpu=False) # VAE
        KLD = -0.5*(1 + logvar - mu.pow(2) - logvar.exp()).mean(0) if mu is not None else 0. # z_dim
        self.preds.append(KLD) # z_dim

# %% ../nbs/source/02_utils.ipynb #37587298-94ff-4345-b081-2d844230dcda
def plot_klds(learn, start_b:int=0, title:str=''):
    klds = learn.model.beta*np.stack(learn.kl_ds.preds) if learn.model.beta!=0. else np.stack(learn.kl_ds.preds)
    plt.semilogy(range(start_b,len(klds)),klds[start_b:]);
    plt.xlabel('Batches');plt.ylabel(r'\(\beta\) DKL');
    plt.title(title);
    plt.grid();
    if learn.model.beta>0. and hasattr(learn, 'loss_param_scheduler'): plt.plot(learn.loss_param_scheduler.hps['beta'], alpha=0.5);
    plt.yscale('symlog', linthresh=1e-8);
    return plt.gcf()

# %% ../nbs/source/02_utils.ipynb #b6fc733a-9bcc-4c4c-901a-40a34243cff8
class GMsCallback(Callback):
    "Record NLL gaussian mixture log-likelihood means per alpha and D during training"
    order,run_valid=65,False
    def __init__(self, RF:int, c:int, use_pad:bool=False, reduction:str='none'):
        store_attr(); self.preds = [];
    def after_pred(self):
        if   self.c > 0: pred, mu, logvar, c_, z = to_detach(self.learn.pred, cpu=False) # VAE + WaveNet
        elif self.c== 0: pred                    = to_detach(self.learn.pred, cpu=False) # WaveNet
        elif self.c==-1: pred, mu, logvar, z     = to_detach(self.learn.pred, cpu=False) # VAE
        if self.use_pad:
            gaussian_mix_loss = mix_gaussian_loss(pred, self.learn.y, reduction=self.reduction)
        else:
            gaussian_mix_loss = mix_gaussian_loss(pred, self.learn.y[:,self.RF:], reduction=self.reduction)
        gaussian_mix_loss = to_detach(gaussian_mix_loss, cpu=True)
        # get idx
        i = self.learn.iter # current iteration, batch number
        current_items = self.learn.dls.items[ self.learn.dls.train._DataLoader__idxs[i*self.learn.dls.bs:(i+1)*self.learn.dls.bs]]
        alphas_items, Ds_items = current_items[:,0], current_items[:,1]
        u_a=np.unique(alphas_items);  u_D=np.unique(Ds_items)
        alphas_idx = [np.flatnonzero(alphas_items==a) for a in u_a]
        Ds_idx = [np.flatnonzero(Ds_items==D) for D in u_D]

        intersect_idx = np.array([[np.intersect1d(alphas_idx[i],Ds_idx[j]) for j,D in enumerate(u_D)]
                                  for i,a in enumerate(u_a)]
                                 ,dtype=object)

        gaussian_mix_loss_mean_aD = np.array([
            [gaussian_mix_loss[intersect_idx[i,j]].mean()
             for j,D in enumerate(u_D)]
            for i,a in enumerate(u_a)] )

        self.preds.append(gaussian_mix_loss_mean_aD)

# %% ../nbs/source/02_utils.ipynb #5df5fc9a-bc45-406a-8eeb-44ed37f55bef
def save_model(fpath, model, model_args:dict, ds_args:dict):
    state = dict(model_args=model_args, ds_args=ds_args,
                 network_state_dict=model.state_dict(),)
    torch.save(state, f'{fpath}.tar')
    print('Saved at '+ f'{fpath}.tar')

# %% ../nbs/source/02_utils.ipynb #ceed4a9d-53b0-4f95-9948-a66424da4144
def load_checkpoint(fpath, model_class=VAEWaveNet, device:str='cuda'):
    try: # load into device
        checkpoint = torch.load(f'{fpath}.tar', weights_only=False, map_location=device)
    except: # use CPU
        device = 'cpu'
        checkpoint = torch.load(f'{fpath}.tar', weights_only=False, map_location=device)

    print("Loading checkpoint: "+f"{fpath}.tar" + f'\non device: {device}')

    network = model_class(**checkpoint['model_args']).to(device)
    try:
        network.load_state_dict(checkpoint['network_state_dict'])
    except:
        raise RuntimeError("Error loading the checkpoint")
    return checkpoint, network

# %% ../nbs/source/02_utils.ipynb #d8f5e2dd-6988-43eb-9c98-f74c19d6c79a
def sig2D(sigma):
    """Converts standard deviation into the associated diffusion coefficient
    $D=\sigma^2/2$"""
    return sigma**2*0.5
def D2sig(D):
    """Converts standard deviation into the associated diffusion coefficient
    $\sigma=\sqrt{2D}$"""
    return (2*D)**0.5
