# -*- coding: utf-8 -*-
"""Base pyro model class"""

import os
from collections import defaultdict, Iterable

import numpy as np
import matplotlib.pyplot as plt
from pyro import poutine
from tqdm.auto import tqdm

import pyro
import torch
from pyro.infer import SVI, JitTrace_ELBO
from pyro.infer import Predictive
from pyro.infer.autoguide import AutoDelta, AutoGuideList
from cell2location.distributions.AutoNormal import AutoNormal
from pyro.infer.autoguide import init_to_mean
import pyro.optim as optim
from sklearn.model_selection import train_test_split

from torch.utils.data import Dataset, DataLoader

from cell2location.models.base.base_model import BaseModel

def flatten_iterable(iterable):
    flattened_list = []
    for i in iterable:
        if isinstance(i, Iterable):
            flattened_list += flatten_iterable(i)
        else:
            flattened_list.append(i)
    return flattened_list

class MiniBatchDataset(Dataset):

    def __init__(self, x_data, extra_data, return_idx=False):
        self.x_data = x_data
        self.extra_data = extra_data
        self.return_idx = return_idx

    def __len__(self):
        return len(self.x_data)

    def __getitem__(self, idx):
        if self.return_idx:
            return self.x_data[idx], {**{'idx': idx}, **{k: v[idx] for k, v in self.extra_data.items()}}
        return self.x_data[idx], {k: v[idx] for k, v in self.extra_data.items()}

# base model class - defining shared methods but not the model itself
class PyroModel(BaseModel):
    r"""Base class for pyro models.
    :param n_fact: number of factors
    :param X_data: Numpy array of gene expression (cols) in spatial locations (rows)
    :param learning_rate: ADAM learning rate for optimising Variational inference objective
    :param n_iter: number of training iterations
    :param total_grad_norm_constraint: gradient constraints in optimisation
    """

    def __init__(
            self,
            X_data: np.ndarray,
            n_fact: int = 10,
            data_type: str = 'float32',
            n_iter: int = 200000,
            learning_rate=0.001,
            total_grad_norm_constraint=200,
            use_cuda=True,
            verbose=True,
            var_names=None, var_names_read=None,
            obs_names=None, fact_names=None, sample_id=None,
            minibatch_size=None,
            minibatch_seed=42,
            point_estim=[],
            custom_guides={}
    ):

        ############# Initialise parameters ################
        super().__init__(X_data, n_fact,
                         data_type, n_iter,
                         learning_rate, total_grad_norm_constraint,
                         verbose, var_names, var_names_read,
                         obs_names, fact_names, sample_id)

        self.extra_data = {}
        self.init_vals = {}
        self.minibatch_size = minibatch_size
        self.minibatch_seed = minibatch_seed
        self.MiniBatchDataset = MiniBatchDataset
        self.use_cuda = use_cuda
        self.device = 'cuda' if self.use_cuda else 'cpu'
        self.point_estim = point_estim
        self.custom_guides = custom_guides
        self.guide_type = 'AutoNormal'

        if self.use_cuda:
            if data_type == 'float32':
                torch.set_default_tensor_type(torch.cuda.FloatTensor)
            elif data_type == 'float16':
                torch.set_default_tensor_type(torch.cuda.HalfTensor)
            elif data_type == 'float64':
                torch.set_default_tensor_type(torch.cuda.DoubleTensor)
            else:
                raise ValueError('Only 32, 16 and 64-bit tensors can be used (data_type)')
        else:
            if data_type == 'float32':
                torch.set_default_tensor_type(torch.FloatTensor)
            elif data_type == 'float16':
                torch.set_default_tensor_type(torch.HalfTensor)
            elif data_type == 'float64':
                torch.set_default_tensor_type(torch.DoubleTensor)
            else:
                raise ValueError('Only 32, 16 and 64-bit tensors can be used (data_type)')

    def sample_prior(self):
        r""" Take one sample from the prior
        :return: self.prior_trace dictionary with an element for each parameter of the model.
        """
        print(".sample_prior() not implemented yet")

    def fit_advi_iterative_simple(self, n: int = 3, method='advi', n_type='restart',
                                  n_iter=None, learning_rate=None, progressbar=True, ):
        r""" Find posterior using ADVI (deprecated)
        (maximising likehood of the data and minimising KL-divergence of posterior to prior)
        :param n: number of independent initialisations
        :param method: which approximation of the posterior (guide) to use?.
            * ``'advi'`` - Univariate normal approximation (pyro.infer.autoguide.AutoDiagonalNormal)
            * ``'custom'`` - Custom guide using conjugate posteriors
        :return: self.svi dictionary with svi pyro objects for each n, and sefl.elbo dictionary storing training history.
        """

        # Pass data to pyro / pytorch
        self.x_data = torch.tensor(self.X_data.astype(self.data_type))  # .double()

        # initialise parameter store
        self.svi = {}
        self.hist = {}
        self.guide_i = {}
        self.samples = {}
        self.node_samples = {}

        self.n_type = n_type

        if n_iter is None:
            n_iter = self.n_iter

        if learning_rate is None:
            learning_rate = self.learning_rate

        if np.isin(n_type, ['bootstrap']):
            if self.X_data_sample is None:
                self.bootstrap_data(n=n)
        elif np.isin(n_type, ['cv']):
            self.generate_cv_data()  # cv data added to self.X_data_sample

        init_names = ['init_' + str(i + 1) for i in np.arange(n)]

        for i, name in enumerate(init_names):

            # initialise Variational distributiion = guide
            try:
                create_plates = getattr(self, 'create_plates')
            except Exception as e:
                create_plates = None

            if method is 'advi':
                self.guide_i[name] = AutoGuideList(self.model)
                self.guide_i[name].append(
                    AutoNormal(poutine.block(self.model, expose_all=True, hide_all=False, hide=self.point_estim),
                               init_loc_fn=init_to_mean, create_plates=create_plates))
                self.guide_i[name].append(
                    AutoDelta(poutine.block(self.model, hide_all=True, expose=self.point_estim,
                                            create_plates=create_plates)))
            elif method is 'custom':
                self.guide_i[name] = self.guide

            # pick dataset depending on the training mode and move to GPU
            if np.isin(n_type, ['cv', 'bootstrap']):
                self.x_data = torch.tensor(self.X_data_sample[i].astype(self.data_type))
            else:
                self.x_data = torch.tensor(self.X_data.astype(self.data_type))

            if self.use_cuda:
                # move tensors and modules to CUDA
                self.x_data = self.x_data.cuda()

            pyro.clear_param_store()

            self.guide_i[name](self.x_data)

            # initialise SVI inference method
            self.svi[name] = SVI(self.model, self.guide_i[name],
                                 optim.ClippedAdam({'lr': learning_rate,
                                                    # limit the gradient step from becoming too large
                                                    'clip_norm': self.total_grad_norm_constraint}),
                                 loss=JitTrace_ELBO())

            # record ELBO Loss history here
            self.hist[name] = []

            # train for n_iter
            it_iterator = tqdm(range(n_iter))
            for it in it_iterator:

                hist = self.svi[name].step(self.x_data)
                it_iterator.set_description('ELBO Loss: ' + str(np.round(hist, 3)))
                self.hist[name].append(hist)

                # if it % 50 == 0 & self.verbose:
                # logging.info("Elbo loss: {}".format(hist))
                if it % 500 == 0:
                    torch.cuda.empty_cache()

    def set_initial_values(self):
        r"""Method for setting initial values on covariate effect (gene_factors parameter)
        :return: nothing
        """
        if self.guide_type == 'AutoGuideList':
            def prefix(i):
                return f'AutoGuideList.{i}.'
        elif self.guide_type == 'AutoNormal':
            def prefix(i):
                return f'AutoNormal.'

        for k in list(self.init_vals.keys()):

            if k in self.point_estim:
                pyro.param(f'{prefix(1)}{k}',
                           torch.Tensor(self.init_vals[k][1](self.init_vals[k][0])))
            else:
                pyro.param(f'{prefix(0)}locs.{k}',
                           torch.Tensor(self.init_vals[k][1](self.init_vals[k][0])))

    def fit_advi_iterative(self, n=3, method='advi', n_type='restart',
                           n_iter=None, learning_rate=None,
                           progressbar=True, num_workers=2,
                           train_proportion=None, stratify_cv=None,
                           l2_weight=False, sample_scaling_weight=0.5,
                           checkpoints=None,
                           checkpoint_dir='./checkpoints',
                           tracking=False):

        r""" Train posterior using ADVI method.
        (maximising likehood of the data and minimising KL-divergence of posterior to prior)
        :param n: number of independent initialisations
        :param method: to allow for potential use of SVGD or MCMC (currently only ADVI implemented).
        :param n_type: type of repeated initialisation:
                                  'restart' to pick different initial value,
                                  'cv' for molecular cross-validation - splits counts into n datasets,
                                         for now, only n=2 is implemented
                                  'bootstrap' for fitting the model to multiple downsampled datasets.
                                         Run `mod.bootstrap_data()` to generate variants of data
        :param n_iter: number of iterations, supersedes self.n_iter
        :param train_proportion: if not None, which proportion of cells to use for training and which for validation.
        :param checkpoints: int, list of int's or None, number of checkpoints to save while model training or list of
            iterations to save checkpoints on
        :param checkpoint_dir: str, directory to save checkpoints in
        :param tracking: bool, track all latent variables during training - if True makes training 2 times slower
        :return: None
        """

        # initialise parameter store
        self.svi = {}
        self.hist = {}
        self.guide_i = {}
        self.trace_elbo_i = {}
        self.samples = {}
        self.node_samples = {}

        if tracking:
            self.logp_hist = {}

        if n_iter is None:
            n_iter = self.n_iter

        if type(checkpoints) is int:
            if n_iter < checkpoints:
                checkpoints = n_iter
            checkpoints = np.linspace(0, n_iter, checkpoints + 1, dtype=int)[1:]
            self.checkpoints = list(checkpoints)
        else:
            self.checkpoints = checkpoints

        self.checkpoint_dir = checkpoint_dir

        self.n_type = n_type
        self.l2_weight = l2_weight
        self.sample_scaling_weight = sample_scaling_weight
        self.train_proportion = train_proportion

        if stratify_cv is not None:
            self.stratify_cv = stratify_cv

        if train_proportion is not None:
            self.validation_hist = {}
            self.training_hist = {}
            if tracking:
                self.logp_hist_val = {}
                self.logp_hist_train = {}

        if learning_rate is None:
            learning_rate = self.learning_rate

        if np.isin(n_type, ['bootstrap']):
            if self.X_data_sample is None:
                self.bootstrap_data(n=n)
        elif np.isin(n_type, ['cv']):
            self.generate_cv_data()  # cv data added to self.X_data_sample

        init_names = ['init_' + str(i + 1) for i in np.arange(n)]

        for i, name in enumerate(init_names):
            ################### Initialise parameters & optimiser ###################
            # initialise Variational distribution = guide
            try:
                create_plates = getattr(self, 'create_plates')
            except Exception as e:
                create_plates = None

            if method is 'advi':
                if len(self.point_estim + flatten_iterable(self.custom_guides.keys())) > 0:
                    self.guide_i[name] = AutoGuideList(self.model)
                    normal_guide_block = poutine.block(self.model, expose_all=True, hide_all=False,
                                                       hide=self.point_estim + flatten_iterable(
                                                           self.custom_guides.keys()))
                    self.guide_i[name].append(AutoNormal(normal_guide_block, init_loc_fn=init_to_mean,
                                                         create_plates=create_plates))
                    self.guide_i[name].append(
                        AutoDelta(poutine.block(self.model, hide_all=True, expose=self.point_estim,
                                                create_plates=create_plates)))
                    for k, v in self.custom_guides.items():
                        self.guide_i[name].append(v)
                else:
                    self.guide_i[name] = AutoNormal(self.model, init_loc_fn=init_to_mean,
                                                    create_plates=create_plates)

                self.guide_type = type(self.guide_i[name]).__name__

            elif method is 'custom':
                self.guide_i[name] = self.guide

            def initialise_svi(x_data, extra_data):

                pyro.clear_param_store()

                self.set_initial_values()

                self.init_guide(name, x_data, extra_data)

                self.trace_elbo_i[name] = JitTrace_ELBO()  # JitTrace_ELBO()

                # initialise SVI inference method
                self.svi[name] = SVI(self.model, self.guide_i[name],
                                     optim.ClippedAdam({'lr': learning_rate,
                                                        # limit the gradient step from becoming too large
                                                        'clip_norm': self.total_grad_norm_constraint}),
                                     loss=self.trace_elbo_i[name])

            # record ELBO Loss history here
            self.hist[name] = []
            if tracking:
                self.logp_hist[name] = defaultdict(list)

            if train_proportion is not None:
                self.validation_hist[name] = []
                if tracking:
                    self.logp_hist_val[name] = defaultdict(list)

            ################### Select data for this iteration ###################
            if np.isin(n_type, ['cv', 'bootstrap']):
                X_data = self.X_data_sample[i].astype(self.data_type)
            else:
                X_data = self.X_data.astype(self.data_type)

            ################### Training / validation split ###################
            # split into training and validation
            if train_proportion is not None:
                idx = np.arange(len(X_data))
                train_idx, val_idx = train_test_split(idx, train_size=train_proportion,
                                                      shuffle=True, stratify=self.stratify_cv)

                extra_data_val = {k: torch.FloatTensor(v[val_idx]).to(self.device) for k, v in self.extra_data.items()}
                extra_data_train = {k: torch.FloatTensor(v[train_idx]) for k, v in self.extra_data.items()}

                x_data_val = torch.FloatTensor(X_data[val_idx]).to(self.device)
                x_data = torch.FloatTensor(X_data[train_idx])
            else:
                # just convert data to CPU tensors
                x_data = torch.FloatTensor(X_data)
                extra_data_train = {k: torch.FloatTensor(v) for k, v in self.extra_data.items()}

            ################### Move data to cuda - FULL data ###################
            # if not minibatch do this:
            if self.minibatch_size is None:
                # move tensors to CUDA
                x_data = x_data.to(self.device)
                for k in extra_data_train.keys():
                    extra_data_train[k] = extra_data_train[k].to(self.device)
                # extra_data_train = {k: v.to(self.device) for k, v in extra_data_train.items()}

            ################### MINIBATCH data ###################
            else:
                # create minibatches
                dataset = MiniBatchDataset(x_data, extra_data_train, return_idx=True)
                loader = DataLoader(dataset, batch_size=self.minibatch_size,
                                    num_workers=0, shuffle=True, drop_last=True)

            ################### Training the model ###################
            if self.minibatch_size is None:
                initialise_svi(x_data, extra_data_train)
            else:
                i = 0
                for batch in loader:
                    i = i + 1
                    if i == 1:
                        x_data_batch, extra_data_batch = batch
                        x_data_batch = x_data_batch.to(self.device)
                        extra_data_batch = {k: v.to(self.device) for k, v in extra_data_batch.items()}

                initialise_svi(x_data, extra_data_batch)

            # start training in epochs
            epochs_iterator = tqdm(range(n_iter))
            for epoch in epochs_iterator:

                if self.minibatch_size is None:
                    ################### Training FULL data ###################
                    iter_loss = self.step_train(name, x_data, extra_data_train)

                    self.hist[name].append(iter_loss)
                    # save data for posterior sampling
                    self.x_data = x_data
                    self.extra_data_train = extra_data_train

                    if tracking:
                        guide_tr, model_tr = self.step_trace(name, x_data, extra_data_train)
                        self.logp_hist[name]['guide'].append(guide_tr.log_prob_sum().item())
                        self.logp_hist[name]['model'].append(model_tr.log_prob_sum().item())

                        for k, v in model_tr.nodes.items():
                            if "log_prob_sum" in v:
                                self.logp_hist[name][k].append(v["log_prob_sum"].item())

                else:
                    ################### Training MINIBATCH data ###################
                    aver_loss = []
                    if tracking:
                        aver_logp_guide = []
                        aver_logp_model = []
                        aver_logp = defaultdict(list)

                    for batch in loader:

                        x_data_batch, extra_data_batch = batch
                        x_data_batch = x_data_batch.to(self.device)
                        extra_data_batch = {k: v.to(self.device) for k, v in extra_data_batch.items()}

                        loss = self.step_train(name, x_data_batch, extra_data_batch)

                        if tracking:
                            guide_tr, model_tr = self.step_trace(name, x_data_batch, extra_data_batch)
                            aver_logp_guide.append(guide_tr.log_prob_sum().item())
                            aver_logp_model.append(model_tr.log_prob_sum().item())

                            for k, v in model_tr.nodes.items():
                                if "log_prob_sum" in v:
                                    aver_logp[k].append(v["log_prob_sum"].item())

                        aver_loss.append(loss)

                    iter_loss = np.sum(aver_loss)

                    # save data for posterior sampling
                    self.x_data = x_data_batch
                    self.extra_data_train = extra_data_batch

                    self.hist[name].append(iter_loss)

                    if tracking:
                        iter_logp_guide = np.sum(aver_logp_guide)
                        iter_logp_model = np.sum(aver_logp_model)
                        self.logp_hist[name]['guide'].append(iter_logp_guide)
                        self.logp_hist[name]['model'].append(iter_logp_model)

                        for k, v in aver_logp.items():
                            self.logp_hist[name][k].append(np.sum(v))

                if self.checkpoints is not None:
                    if (epoch + 1) in self.checkpoints:
                        self.save_checkpoint(epoch + 1, prefix=name)

                ################### Evaluating cross-validation loss ###################
                if train_proportion is not None:

                    iter_loss_val = self.step_eval_loss(name, x_data_val, extra_data_val)

                    if tracking:
                        guide_tr, model_tr = self.step_trace(name, x_data_val, extra_data_val)
                        self.logp_hist_val[name]['guide'].append(guide_tr.log_prob_sum().item())
                        self.logp_hist_val[name]['model'].append(model_tr.log_prob_sum().item())

                        for k, v in model_tr.nodes.items():
                            if "log_prob_sum" in v:
                                self.logp_hist_val[name][k].append(v["log_prob_sum"].item())

                    self.validation_hist[name].append(iter_loss_val)
                    epochs_iterator.set_description(f'ELBO Loss: ' + '{:.4e}'.format(iter_loss) \
                                                    + ': Val loss: ' + '{:.4e}'.format(iter_loss_val))
                else:
                    epochs_iterator.set_description('ELBO Loss: ' + '{:.4e}'.format(iter_loss))

                if epoch % 20 == 0:
                    torch.cuda.empty_cache()

            if train_proportion is not None:
                # rescale loss
                self.validation_hist[name] = [i / (1 - train_proportion)
                                              for i in self.validation_hist[name]]
                self.hist[name] = [i / train_proportion for i in self.hist[name]]

                # reassing the main loss to be displayed
                self.training_hist[name] = self.hist[name]
                self.hist[name] = self.validation_hist[name]

                if tracking:
                    for k, v in self.logp_hist[name].items():
                        self.logp_hist[name][k] = [i / train_proportion for i in self.logp_hist[name][k]]
                        self.logp_hist_val[name][k] = [i / (1 - train_proportion) for i in self.logp_hist_val[name][k]]

                    self.logp_hist_train[name] = self.logp_hist[name]
                    self.logp_hist[name] = self.logp_hist_val[name]

            if self.verbose:
                print(plt.plot(np.log10(self.hist[name][0:])));

    def init_guide(self, name, x_data, extra_data):

        self.guide_i[name](x_data)

    def step_train(self, name, x_data, extra_data):

        return self.svi[name].step(x_data)

    def step_eval_loss(self, name, x_data, extra_data):

        return self.svi[name].evaluate_loss(x_data)

    def step_predictive(self, predictive, x_data, extra_data):

        return predictive(x_data)

    def step_trace(self, name, x_data, extra_data):

        guide_tr = poutine.trace(self.guide_i[name]).get_trace(x_data)
        model_tr = poutine.trace(poutine.replay(self.model,
                                                trace=guide_tr)).get_trace(x_data)
        return guide_tr, model_tr

    def fit_nuts(self, n_samples: int = 1000, warmup_steps: int = 1000, save_samples=False):

        self.samples = {}
        self.n_samples = n_samples

        # create sampler and run MCMC
        self.nuts_kernel = NUTS(self.model, jit_compile=True)
        self.mcmc = MCMC(self.nuts_kernel, num_samples=n_samples,
                         warmup_steps=warmup_steps)
        self.mcmc.run(self.x_data)

        post_samples = {k: v.detach().cpu().numpy() for k, v in self.mcmc.get_samples().items()}

        # summarise samples
        self.samples['post_sample_means'] = {v: post_samples[v].mean(axis=0) for v in post_samples.varnames}
        self.samples['post_sample_q05'] = {v: np.quantile(post_samples[v], 0.05, axis=0) for v in post_samples.varnames}
        self.samples['post_sample_q95'] = {v: np.quantile(post_samples[v], 0.95, axis=0) for v in post_samples.varnames}
        self.samples['post_sample_sds'] = {v: post_samples[v].std(axis=0) for v in post_samples.varnames}

        if (save_samples):
            self.samples['post_samples'] = post_samples

    def plot_history_old(self, iter_start: int = 15000, iter_end=-1):
        r""" Plot training history
        :param iter_start: omit initial iterations from the plot
        :param iter_end: omit last iterations from the plot
        """
        for i in self.hist.keys():
            print(plt.plot(np.log10(np.array(self.hist[i])[iter_start:iter_end])));

    def plot_history_1(self, iter_start=0, iter_end=-1,
                       mean_field_slot=None, log_y=True, ax=None):
        r""" Plot training history

        :param iter_start: omit initial iterations from the plot
        :param iter_end: omit last iterations from the plot
        """

        if ax is None:
            ax = plt
            ax.set_xlabel = plt.xlabel
            ax.set_ylabel = plt.ylabel

        if mean_field_slot is None:
            mean_field_slot = self.hist.keys()

        for i in mean_field_slot:

            if iter_end == -1:
                iter_end = np.array(self.hist[i]).flatten().shape[0]

            y = np.array(self.hist[i]).flatten()[iter_start:iter_end]
            if log_y:
                y = np.log10(y)
            ax.plot(np.arange(iter_start, iter_end), y, label='train')
            ax.set_xlabel('Training epochs')
            ax.set_ylabel('Reconstruction accuracy (ELBO loss)')
            ax.legend()
            plt.tight_layout()

    def sample_node1(self, node, init, batch_size: int = 10):

        predictive = Predictive(self.model, guide=self.guide_i[init],
                                num_samples=batch_size)

        post_samples = {k: v.detach().cpu().numpy()
                        for k, v in self.step_predictive(predictive, self.x_data, self.extra_data_train).items()
                        if k == node}

        return (post_samples[node])

    def sample_node(self, node, init, n_sampl_iter,
                    batch_size: int = 10, suff=''):

        # sample first batch
        self.samples[node + suff][init] = self.sample_node1(node, init, batch_size=batch_size)

        for it in tqdm(range(n_sampl_iter - 1)):
            # sample remaining batches
            post_node = self.sample_node1(node, init, batch_size=batch_size)

            # concatenate batches
            self.samples[node + suff][init] = np.concatenate((self.samples[node + suff][init], post_node), axis=0)

        # compute mean across samples
        self.samples[node + suff][init] = self.samples[node + suff][init].mean(0)

    def sample_all1(self, init='init_1', batch_size: int = 10):

        predictive = Predictive(self.model, guide=self.guide_i[init],
                                num_samples=batch_size)

        post_samples = {k: v.detach().cpu().numpy()
                        for k, v in self.step_predictive(predictive, self.x_data, self.extra_data_train).items()
                        if k != "data_target"}

        return (post_samples)

    def sample_all(self, n_sampl_iter, init='init_1', batch_size: int = 10):

        # sample first batch
        self.samples['post_samples'] = self.sample_all1(init, batch_size=batch_size)

        for it in tqdm(range(n_sampl_iter - 1)):
            # sample remaining batches
            post_samples = self.sample_all1(init, batch_size=batch_size)

            # concatenate batches
            self.samples['post_samples'] = {k: np.concatenate((self.samples['post_samples'][k],
                                                               post_samples[k]), axis=0)
                                            for k in post_samples.keys()}

    def b_evaluate_stability(self, node, n_samples: int = 1000, batch_size: int = 10,
                             align=True, transpose=True):
        r""" Evaluate stability of posterior samples between training initialisations
        (takes samples and correlates the values of factors between training initialisations)
        :param node: which pymc3 node to sample? Factors should be in columns.
        :param n_samples: the number of samples.
        :param batch_size: generate samples in batches of size `batch_size`. Necessary for the computation to fit in the GPU memory
        :return: self.samples[node_name+_stab] dictionary with an element for each training initialisation.
        """

        self.n_samples = n_samples
        self.n_sampl_iter = int(np.ceil(n_samples / batch_size))
        self.n_sampl_batch = batch_size

        self.samples[node + '_stab'] = {}

        for i in self.guide_i.keys():
            self.sample_node(node, i, self.n_sampl_iter,
                             batch_size=self.n_sampl_batch, suff='_stab')

        # plot correlations of posterior mean between training initialisations
        for i in range(len(self.samples[node + '_stab'].keys()) - 1):
            x = self.samples[node + '_stab']['init_' + str(1)]
            y = self.samples[node + '_stab']['init_' + str(i + 2)]
            if transpose:
                x = x.T
                y = y.T
            print(self.align_plot_stability(x, y,
                                            str(1), str(i + 2), align=align))

    def sample_posterior(self, node='all',
                         n_samples: int = 1000, batch_size: int = 10,
                         save_samples=False,
                         mean_field_slot='init_1'):
        r""" Sample posterior distribution of parameters - either all or single parameter
        :param node: pyro parameter to sample (e.g. default "all", self.spot_factors)
        :param n_samples: number of posterior samples to generate (1000 is recommended, reduce if you get GPU memory error)
        :param save_samples: save samples in addition to sample mean, 5% quantile, SD.
        :param return_samples: return summarised samples in addition to saving them in `self.samples`
        :param mean_field_slot: string, which mean_field slot to sample? 'init_1' by default
        :return: dictionary of dictionaries (mean, 5% quantile, SD, optionally all samples) with numpy arrays for each parameter.
        Optional dictionary of all samples contains parameters as numpy arrays of shape ``(n_samples, ...)``
        """

        self.n_samples = n_samples
        self.n_sampl_iter = int(np.ceil(n_samples / batch_size))
        self.n_sampl_batch = batch_size

        if (node == 'all'):
            # Sample all parameters - might use a lot of GPU memory

            self.sample_all(self.n_sampl_iter, init=mean_field_slot, batch_size=self.n_sampl_batch)

            self.param_names = list(self.samples['post_samples'].keys())

            self.samples['post_sample_means'] = {v: self.samples['post_samples'][v].mean(axis=0)
                                                 for v in self.param_names}
            self.samples['post_sample_q05'] = {v: np.quantile(self.samples['post_samples'][v], 0.05, axis=0)
                                               for v in self.param_names}
            self.samples['post_sample_q95'] = {v: np.quantile(self.samples['post_samples'][v], 0.95, axis=0)
                                               for v in self.param_names}
            self.samples['post_sample_sds'] = {v: self.samples['post_samples'][v].std(axis=0)
                                               for v in self.param_names}

            if not save_samples:
                self.samples['post_samples'] = None

        else:
            self.sample_node(node, mean_field_slot, self.n_sampl_iter,
                             batch_size=self.n_sampl_batch, suff='')

        return (self.samples)

    def save_checkpoint(self, n, prefix=''):
        r""" Save pyro parameter store (current status of Variational parameters) to disk
        :param n: epoch number
        :param prefix: filename prefix (e.g. init number)
        """

        if not os.path.exists(self.checkpoint_dir):
            os.makedirs(self.checkpoint_dir)

        filename = f'{self.checkpoint_dir}/{prefix}_{n}.ckp'
        pyro.get_param_store().save(filename)

    def load_checkpoint(self, filename):
        r""" Load pyro parameter store (current status of Variational parameters) from disk
        :param filename: checkpoint filename
        """

        if filename in os.listdir(self.checkpoint_dir):
            pyro.get_param_store().load(filename)
        else:
            checkpoints = os.listdir(self.checkpoint_dir)
            checkpoints = '\n'.join(checkpoints)
            checkpoint_dir_abspath = os.path.abspath(self.checkpoint_dir)
            raise ValueError(f'No such filename in {checkpoint_dir_abspath}, available filenames : \n'
                             f'{checkpoints}')


# -*- coding: utf-8 -*-
# =============================================================================
# Copyright (c) 2012, Lawrence Livermore National Security, LLC.
# Produced at the Lawrence Livermore National Laboratory.
# Written by Joel Bernier <bernier2@llnl.gov> and others.
# LLNL-CODE-529294.
# All rights reserved.
#
# This file is part of HEXRD. For details on dowloading the source,
# see the file COPYING.
#
# Please also see the file LICENSE.
#
# This program is free software; you can redistribute it and/or modify it under
# the terms of the GNU Lesser General Public License (as published by the Free
# Software Foundation) version 2.1 dated February 1999.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the IMPLIED WARRANTY OF MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE. See the terms and conditions of the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this program (see file LICENSE); if not, write to
# the Free Software Foundation, Inc., 59 Temple Place, Suite 330,
# Boston, MA 02111-1307 USA or visit <http://www.gnu.org/licenses/>.
# =============================================================================
"""
Module for XRD material class

Use the Material class directly for new materials.  Known
materials are defined by name in materialDict.
"""
from configparser import SafeConfigParser as Parser
import numpy

from hexrd.crystallography import PlaneData as PData
from hexrd.valunits import valWUnit
from hexrd import unitcell
from hexrd.constants import ptable
import copy

from os import path
from pathlib import Path
from CifFile import ReadCif
import  h5py
from warnings import warn
from hexrd.mksupport import Write2H5File
from hexrd.symbols import xtal_sys_dict
from hexrd.symbols import Hall_to_sgnum, HM_to_sgnum

__all__ = ['Material', 'loadMaterialList']

#
# ================================================== Module Data
#

def _angstroms(x):
    return valWUnit('lp', 'length',  x, 'angstrom')

def _degrees(x):
    return valWUnit('lp', 'angle',   x, 'degrees')

def _key(x):
    return x.name

#
# ---------------------------------------------------CLASS:  Material
#

class Material(object):
    """Simple class for holding lattice parameters, accessible by name.

    The class references materials by name and contains lattice and
    space group data.
    default data is for nickel, but name is material
    """
    DFLT_NAME = 'material.xtal'
    DFLT_XTAL = 'Ni'
    DFLT_SGNUM = 225

    DFLT_LPARMS = [_angstroms(3.61), _angstroms(3.61), _angstroms(3.61),
                   _degrees(90.0), _degrees(90.0), _degrees(90.0)]
    DFLT_SSMAX = 100

    DFLT_KEV = valWUnit('wavelength', 'energy', 80.725e0, 'keV')
    DFLT_STR = 0.0025
    DFLT_TTH = numpy.radians(0.25)
    DFLT_TTHMAX = numpy.radians(160.0)
    """
    ATOMINFO    Fractional Atom Position of an atom in the unit cell followed by the
    site occupany and debye waller (U) factor in A^(-2)
    B is related to U by B = 8 pi^2 U

    ATOMTYPE    atomic number of all the different species in the unitcell
    """
    DFLT_ATOMINFO = numpy.array([[0., 0., 0., 1.]])
    DFLT_U = numpy.array([4.18e-7])
    DFLT_ATOMTYPE = numpy.array([28])

    '''
    the dmin parameter is used to figure out the maximum sampling for g-vectors
    this parameter is in angstroms
    '''
    DFLT_DMIN = _angstroms(1.0)

    '''
    some materials have more than one space group setting. for ex
    the diamond cubic system has two settings with the origin either
    at (0,0,0) or at (1/4,1/4,1/4) etc. this handle takes care of these
    cases. but the defaiult is always 0

    default space group setting
    '''
    DFLT_SGSETTING = 0

    def __init__(self, name=None, material_file=None, dmin=DFLT_DMIN, kev=DFLT_KEV, sgsetting=DFLT_SGSETTING):
        """Constructor for Material

        name -- (str) name of crystal
        material_file -- (str) name of the material file
        which contains the crystal. this could be either cif
        or hdf5
        """
        self.name = name
        self.description = ''

        self._dmin = dmin

        self._beamEnergy = kev

        self.sgsetting = sgsetting

        if material_file:
            # Get values from configuration
            # self._readCfg(material_file)
            # >> @ date 08/20/2020 SS removing dependence on hklmax
            #self._hklMax = Material.DFLT_SSMAX
            # self._beamEnergy = Material.DFLT_KEV
            form = Path(material_file).suffix[1:]

            if(form == 'cif'):
                self._readCif(material_file)
            elif(form in ['h5', 'hdf5', 'xtal']):
                self._readHDFxtal(fhdf=material_file, xtal=name)
        else:
            # Use default values
            self._lparms = Material.DFLT_LPARMS
            # self._hklMax = Material.DFLT_SSMAX
            #
            self.description = ''
            #
            self.sgnum = Material.DFLT_SGNUM
            self._sgsetting = Material.DFLT_SGSETTING
            #
            self._atominfo = Material.DFLT_ATOMINFO
            #
            self._U = Material.DFLT_U
            #
            self._atomtype = Material.DFLT_ATOMTYPE
            #

        self.unitcell = unitcell.unitcell(
            self._lparms, self.sgnum, self._atomtype, self._atominfo, self._U,
            self._dmin.getVal('nm'), self._beamEnergy.value,
            self._sgsetting)

        self._newPdata()
        self.update_structure_factor()

    def __str__(self):
        """String representation"""
        s = 'Material:  %s\n' % self.name
        if self.description:
            s += '   description:  %s\n' % self.description
            pass
        s += '   plane Data:  %s' % str(self.planeData)
        return s

    def _newPdata(self):
        """Create a new plane data instance"""
        # spaceGroup module calulates forbidden reflections
        '''
        >> @date 08/20/2020 SS removing dependence of planeData
        initialization on the spaceGroup module. everything is
        initialized using the unitcell module now
        '''
        hkls = self.unitcell.getHKLs(self._dmin.getVal('nm')).T
        lprm = [self._lparms[i] for i in unitcell._rqpDict[self.unitcell.latticeType][0]]
        laue = self.unitcell._laueGroup
        self._pData = PData(hkls, lprm, laue,
                            self._beamEnergy, Material.DFLT_STR,
                            tThWidth=Material.DFLT_TTH,
                            tThMax=Material.DFLT_TTHMAX)
        '''
          Set default exclusions
          all reflections with two-theta smaller than 90 degrees
        '''
        tth = numpy.array([hkldata['tTheta'] for hkldata in self._pData.hklDataList])

        dflt_excl = numpy.ones(tth.shape,dtype=numpy.bool)
        dflt_excl[~numpy.isnan(tth)] = ~( (tth[~numpy.isnan(tth)] >= 0.0) & \
                                     (tth[~numpy.isnan(tth)] <= numpy.pi/2.0) )
        dflt_excl[0] = False
        self._pData.exclusions = dflt_excl

        return

    def update_structure_factor(self):
        hkls = self.planeData.getHKLs(allHKLs=True)
        sf = numpy.zeros([hkls.shape[0],])
        for i,g in enumerate(hkls):
            sf[i] = self.unitcell.CalcXRSF(g)

        self.planeData.set_structFact(sf[~self.planeData.exclusions])

    def _readCif(self, fcif=DFLT_NAME+'.cif'):
        """
        >> @AUTHOR:     Saransh Singh, Lawrence Livermore National Lab, saransh1@llnl.gov
        >> @DATE:       10/16/2019 SS 1.0 original
        >> @DETAILS:    hexrd3 will have real structure factors and will require the overhaul
                        of the crystallography. In this effort, we will have a cif reader and
                        also the HDF5 format reader in the material class. We will be using
                        pycifrw for i/o
        """

        # make sure file exists etc.
        if(fcif == Material.DFLT_NAME+'.cif'):
            try:
                cif = ReadCif(fcif)
            except(OSError):
                raise RuntimeError('OS Error: No file name supplied and default file name not found.')
        else:
            try:
                cif = ReadCif(fcif)
            except(OSError):
                raise RuntimeError('OS Error: File not found')

        # read the file
        for k in cif.keys():
            if('_cell_length_a' in cif[k]):
                m = k
                break
        cifdata = cif[m]
        # cifdata = cif[cif.keys()[0]]

        # make sure the space group is present in the cif file, either as
        # international table number, hermann-maguain or hall symbol
        sgkey = ['_space_group_IT_number',
                 '_symmetry_space_group_name_h-m',
                 '_symmetry_space_group_name_hall',
                 '_symmetry_Int_Tables_number']

        sgdata = False
        for key in sgkey:
            sgdata = sgdata or (key in cifdata)
            if(sgdata):
                skey = key
                break

        if(not(sgdata)):
            raise RuntimeError(' No space group information in CIF file! ')

        sgnum = 0
        if skey is sgkey[0]:
            sgnum = int(cifdata[sgkey[0]])
        elif (skey is sgkey[1]):
            HM = cifdata[sgkey[1]]
            HM = HM.replace(" ", "")
            sgnum = HM_to_sgnum[HM]
        elif (skey is sgkey[2]):
            hall = cifdata[sgkey[2]]
            hall = hall.replace(" ", "")
            sgnum = Hall_to_sgnum[HM]
        elif(skey is sgkey[3]):
            sgnum = int(cifdata[sgkey[3]])

        # lattice parameters
        lparms = []
        lpkey = ['_cell_length_a', '_cell_length_b', \
                 '_cell_length_c', '_cell_angle_alpha', \
                 '_cell_angle_beta', '_cell_angle_gamma']

        for key in lpkey:
            n = cifdata[key].find('(')
            if(n != -1):
                lparms.append(float(cifdata[key][:n]))
            else:
                lparms.append(float(cifdata[key]))

        for i in range(6):
                if(i < 3):
                    lparms[i] = _angstroms(lparms[i])
                else:
                    lparms[i] = _degrees(lparms[i])

        self._lparms = lparms
        self.sgnum   = sgnum

        # fractional atomic site, occ and vibration amplitude
        fracsitekey = ['_atom_site_fract_x', '_atom_site_fract_y',\
                        '_atom_site_fract_z',]

        occ_U       = ['_atom_site_occupancy',\
                        '_atom_site_u_iso_or_equiv','_atom_site_U_iso_or_equiv']

        sitedata = True
        for key in fracsitekey:
            sitedata = sitedata and (key in cifdata)

        if(not(sitedata)):
            raise RuntimeError(' fractional site position is not present or incomplete in the CIF file! ')

        atompos = []
        for key in fracsitekey:
            slist = cifdata[key]
            pos = []

            for p in slist:
                n = p.find('(')

                if(n != -1):
                    pos.append(p[:n])
                else:
                    pos.append(p)

            '''
            sometimes cif files have negative values so need to
            bring them back to fractional coordinates between 0-1
            '''
            pos = numpy.asarray(pos).astype(numpy.float64)
            pos,_ = numpy.modf(pos+100.0)
            atompos.append(pos)

        """note that the vibration amplitude, U is just the amplitude (in A)
            to convert to the typical B which occurs in the debye-waller factor,
            we will use the following formula
            B = 8 * pi ^2 * < U_av^2 >
            this will be done here so we dont have to worry about it later
        """

        pocc = (occ_U[0] in cifdata.keys())
        pU   = (occ_U[1] in cifdata.keys()) or (occ_U[2] in cifdata.keys())

        if(not pocc):
            warn('occupation fraction not present. setting it to 1')
            occ = numpy.ones(atompos[0].shape)
            atompos.append(occ)
        else:
            slist = cifdata[occ_U[0]]
            occ = []
            for p in slist:
                n = p.find('(')

                if(n != -1):
                    occ.append(p[:n])
                else:
                    occ.append(p)

            atompos.append(numpy.asarray(occ).astype(numpy.float64))

        if(not pU):
            warn('Debye-Waller factors not present. setting to same values for all atoms.')
            U = 1.0/numpy.pi/2./numpy.sqrt(2.) * numpy.ones(atompos[0].shape)
            self._U = U
        else:
            if(occ_U[1] in cifdata.keys()):
                k = occ_U[1]
            else:
                k = occ_U[2]

            slist = cifdata[k]
            U = []
            for p in slist:
                n = p.find('(')

                if(n != -1):
                    U.append(p[:n])
                else:
                    U.append(p)

            self._U = numpy.asarray(U).astype(numpy.float64)
        '''
        format everything in the right shape etc.
        '''
        self._atominfo = numpy.asarray(atompos).T

        '''
        get atome types here i.e. the atomic number of atoms at each site
        '''
        atype = '_atom_site_type_symbol'
        patype = (atype in cifdata)
        if(not patype):
            raise RuntimeError('atom types not defined in cif file.')

        satype = cifdata[atype]
        atomtype = []

        for s in satype:
            atomtype.append(ptable[s])

        self._atomtype  = numpy.asarray(atomtype).astype(numpy.int32)
        self._sgsetting = 0

    def _readHDFxtal(self, fhdf=DFLT_NAME, xtal=DFLT_NAME):
        """
        >> @AUTHOR:     Saransh Singh, Lawrence Livermore National Lab, saransh1@llnl.gov
        >> @DATE:       10/17/2019 SS 1.0 original
        >> @DETAILS:    hexrd3 will have real structure factors and will require the overhaul
                        of the crystallography. In this effort, we will have a HDF file reader.
                        the file will be the same as the EMsoft xtal file. h5py will be used for
                        i/o
        """

        fexist = path.exists(fhdf)
        if(fexist):
            fid = h5py.File(fhdf, 'r')
            xtal = "/"+xtal
            if xtal not in fid:
                raise IOError('crystal doesn''t exist in material file.')
        else:
            raise IOError('material file does not exist.')

        gid         = fid.get(xtal)

        sgnum       = numpy.asscalar(numpy.array(gid.get('SpaceGroupNumber'), \
                                    dtype = numpy.int32))
        """
            IMPORTANT NOTE:
            note that the latice parameters is nm by default
            hexrd on the other hand uses A as the default units, so we
            need to be careful and convert it right here, so there is no
            confusion later on
        """
        lparms      = list(gid.get('LatticeParameters'))

        for i in range(6):
            if(i < 3):
                lparms[i] = _angstroms(lparms[i]*10.0)
            else:
                lparms[i] = _degrees(lparms[i])

        self._lparms    = lparms
        #self._lparms    = self._toSixLP(sgnum, lparms)
        # fill space group and lattice parameters
        self.sgnum      = sgnum

        # the U factors are related to B by the relation B = 8pi^2 U
        self._atominfo  = numpy.transpose(numpy.array(gid.get('AtomData'), dtype = numpy.float64))
        self._U         = numpy.transpose(numpy.array(gid.get('U'), dtype = numpy.float64))

        # read atom types (by atomic number, Z)
        self._atomtype = numpy.array(gid.get('Atomtypes'), dtype = numpy.int32)
        self._atom_ntype = self._atomtype.shape[0]

        self._sgsetting = numpy.asscalar(numpy.array(gid.get('SpaceGroupSetting'), \
                                        dtype = numpy.int32))

        fid.close()

    def dump_material(self, filename):
        '''
        get the atominfo dictionaary aand the lattice parameters
        '''
        AtomInfo = {}

        AtomInfo['file'] = filename
        AtomInfo['xtalname'] = self.name
        AtomInfo['xtal_sys'] = xtal_sys_dict[self.unitcell.latticeType.lower()]
        AtomInfo['Z'] = self.unitcell.atom_type
        AtomInfo['SG'] = self.unitcell.sgnum
        AtomInfo['SGsetting'] = self.unitcell.sgsetting
        AtomInfo['APOS'] = self.unitcell.atom_pos
        AtomInfo['U'] = self.unitcell.U

        '''
        lattice parameters
        '''
        lat_param = {'a': self.unitcell.a,
                     'b': self.unitcell.b,
                     'c': self.unitcell.c,
                     'alpha': self.unitcell.alpha,
                     'beta': self.unitcell.beta,
                     'gamma': self.unitcell.gamma}

        Write2H5File(AtomInfo, lat_param)

    # ============================== API
    #
    #  ========== Properties
    #

    # property:  spaceGroup

    @property
    def spaceGroup(self):
        """(read only) Space group"""
        return self._spaceGroup

    @property
    def vol(self):
        return self.unitcell.vol

    # property:  sgnum

    def _get_sgnum(self):
        """Get method for sgnum"""
        return self._sgnum

    def _set_sgnum(self, v):
        """Set method for sgnum
        >> @date 08/20/2020 SS removed planedata initialization
            everytime sgnum is updated singe everything is initialized
            using unitcell now
        """
        self._sgnum = v

        # Update the unit cell if there is one
        if hasattr(self, 'unitcell'):
            self.unitcell.sgnum = v

    sgnum = property(_get_sgnum, _set_sgnum, None,
                     "Space group number")
    # property:  beamEnergy

    def _get_beamEnergy(self):
        """Get method for beamEnergy"""
        return self._beamEnergy

    def _set_beamEnergy(self, keV):
        """
        Set method for beamEnergy

        * note that units are assumed to be keV for
          float arguments.  Also can take a valWUnit
          instance
        """
        self._beamEnergy = keV
        self.planeData.wavelength = keV

        return

    beamEnergy = property(_get_beamEnergy, _set_beamEnergy, None,
                          "Beam energy in keV")

    #>> @date 08/20/2020 removing dependence on hklmax
    # property:  hklMax

    # def _get_hklMax(self):
    #     """Get method for hklMax"""
    #     return self._hklMax

    # def _set_hklMax(self, v):
    #     """Set method for hklMax"""
    #     self._hklMax = v
    #     self._newPdata()  # update planeData
    #     return

    # hklMax = property(_get_hklMax, _set_hklMax, None,
    #                   "Max sum of squares for HKLs")
    # property:  planeData

    @property
    def planeData(self):
        """(read only) Return the planeData attribute (lattice parameters)"""
        return self._pData

    # property:  latticeParameters

    def _get_latticeParameters(self):
        """Get method for latticeParameters"""
        return self._lparms

    def _set_latticeParameters(self, v):
        """Set method for latticeParameters"""
        if(len(v) != 6):
            v = unitcell._rqpDict[self.unitcell.latticeType][1](v)
        lp = [_angstroms(v[i]) for i in range(3)]
        for i in range(3,6):
            lp.append(_degrees(v[i]))
        self._lparms = lp

        rq_lp = unitcell._rqpDict[self.unitcell.latticeType][0]
        for i,vv in enumerate(lp):
            if(vv.isLength()):
                val = vv.value / 10.0
            else:
                val = vv.value
            setattr(self.unitcell, unitcell._lpname[i], val)
        v2 = [lp[x].value for x in rq_lp]
        self.planeData.lparms = v2

        return

    lpdoc = r"""Lattice parameters

On output, all six paramters are returned.

On input, either all six or a minimal set is accepted.

The values have units attached, i.e. they are valWunit instances.
"""
    latticeParameters = property(
            _get_latticeParameters, _set_latticeParameters,
            None, lpdoc)

    # property:  "name"

    def _get_name(self):
        """Set method for name"""
        return self._name

    def _set_name(self, v):
        """Set method for name"""
        self._name = v

        return

    name = property(_get_name, _set_name, None,
                    "Name of material")

    @property
    def dmin(self):
        return self._dmin

    @dmin.setter
    def dmin(self, v):
        if self._dmin == v:
            return

        self._dmin = v

        # Update the unit cell
        self.unitcell.dmin = v.getVal('nm')

        self._newPdata()
        self.update_structure_factor()

    # property: "atominfo"
    def _get_atominfo(self):
        """Set method for name"""
        return self._atominfo

    def _set_atominfo(self, v):
        """Set method for name"""
        if v.shape[1] == 4:
            self._atominfo = v
        else:
            print("Improper syntax, array must be n x 4")

        return

    atominfo = property(
        _get_atominfo, _set_atominfo, None,
        "Information about atomic positions and electron number")

    #
    #  ========== Methods
    #
    #
    pass  # end class

#
#  -----------------------------------------------END CLASS:  Material
#
#  Utility Functions
#

def loadMaterialList(cfgFile):
    """Load a list of materials from a file

    The file uses the config file format.  See ConfigParser module.
"""
    p = Parser()
    p.read(cfgFile)
    #
    #  Each section defines a material
    #
    names = p.sections()
    matList = [Material(n, p) for n in names]
    # Sort the list
    matList = sorted(matList, key=_key)

    return matList

def load_materials_hdf5(f, dmin=Material.DFLT_DMIN, kev=Material.DFLT_KEV,
                        sgsetting=Material.DFLT_SGSETTING):
    """Load materials from an HDF5 file

    The file uses the HDF5 file format.
    """
    with h5py.File(f, 'r') as rf:
        names = list(rf)

    return {
        name: Material(name, f, dmin=dmin, kev=kev, sgsetting=sgsetting)
        for name in names
    }

def save_materials_hdf5(f, materials):
    """Save a dict of materials into an HDF5 file"""
    for material in materials.values():
        material.dump_material(f)

#
#  ============================== Executable section for testing
#

if __name__ == '__main__':
    #
    #  For testing
    #
    import sys

    if len(sys.argv) == 1:
        print("need argument:  materials.cfg")
        sys.exit()
        pass

    ml = loadMaterialList(sys.argv[1])

    print('MATERIAL LIST\n')
    print(('   from file:  ', sys.argv[1]))
    for m in ml:
        print(m)
        pass
    pass


# Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# A copy of the License is located at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# or in the "license" file accompanying this file. This file is distributed
# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
# express or implied. See the License for the specific language governing
# permissions and limitations under the License.
import numpy as np
import autograd.numpy as anp
from autograd import grad
from typing import Tuple, Dict, List, Optional
from numpy.random import RandomState
import logging

from syne_tune.optimizer.schedulers.searchers.bayesopt.gpautograd.kernel \
    import KernelFunction
from syne_tune.optimizer.schedulers.searchers.bayesopt.gpautograd.mean \
    import MeanFunction
from syne_tune.optimizer.schedulers.searchers.bayesopt.gpautograd.learncurve.issm \
    import issm_likelihood_slow_computations, posterior_computations, \
    sample_posterior_marginals, _inner_product, issm_likelihood_computations, \
    issm_likelihood_precomputations, decode_features, _rowvec, \
    update_posterior_state, update_posterior_pvec, _flatvec, _colvec
from syne_tune.optimizer.schedulers.searchers.bayesopt.gpautograd.learncurve.issm \
    import sample_posterior_joint as sample_posterior_joint_issm, \
    predict_posterior_marginals_extended as predict_posterior_marginals_issm
from syne_tune.optimizer.schedulers.searchers.bayesopt.gpautograd.learncurve.freeze_thaw \
    import sample_posterior_joint as sample_posterior_joint_expdecay, \
    predict_posterior_marginals_extended as predict_posterior_marginals_expdecay
from syne_tune.optimizer.schedulers.searchers.bayesopt.gpautograd.learncurve.model_params \
    import ISSModelParameters
from syne_tune.optimizer.schedulers.searchers.bayesopt.gpautograd.learncurve.freeze_thaw \
    import resource_kernel_likelihood_slow_computations, \
    ExponentialDecayBaseKernelFunction, logdet_cholfact_cov_resource, \
    resource_kernel_likelihood_computations, \
    resource_kernel_likelihood_precomputations
from syne_tune.optimizer.schedulers.searchers.bayesopt.datatypes.common \
    import Configuration

logger = logging.getLogger(__name__)

__all__ = ['GaussProcAdditivePosteriorState',
           'IncrementalUpdateGPAdditivePosteriorState',
           'GaussProcISSMPosteriorState',
           'GaussProcExpDecayPosteriorState']

class GaussProcAdditivePosteriorState(object):
    """
    Represent posterior state for joint Gaussian model of learning curves over
    a number of configurations. The (additive) model is the sum of a Gaussian
    process model for function values at r_max and independent Gaussian models
    over r only.

    Importantly, inference scales cubically only in the number of
    configurations, not in the number of observations.

    """

    def __init__(
            self, data: Optional[Dict], mean: MeanFunction,
            kernel: KernelFunction, noise_variance, **kwargs):
        """
        `data` contains input points and targets, as obtained from
        `issm.prepare_data`. `iss_model` maintains the ISSM parameters.

        :param data: Input points and targets
        :param mean: Mean function m(X)
        :param kernel: Kernel function k(X, X')
        :param noise_variance: Noise variance
        """
        self.mean = mean
        self.kernel = kernel
        self.noise_variance = noise_variance
        self.poster_state = None
        if data is not None:
            self.r_min = data['r_min']
            self.r_max = data['r_max']
            # Compute posterior state
            self._compute_posterior_state(data, noise_variance, **kwargs)
        else:
            # Copy constructor, used by `IncrementalUpdateGPISSMPosteriorState`
            # subclass
            self.poster_state = kwargs['poster_state']
            self.r_min = kwargs['r_min']
            self.r_max = kwargs['r_max']

    def _compute_posterior_state(
            self, data: Dict, noise_variance, **kwargs):
        raise NotImplementedError()

    def neg_log_likelihood(self):
        assert 'criterion' in self.poster_state, \
            "neg_log_likelihood not defined for fantasizing posterior state"
        return self.poster_state['criterion']

    def predict(
            self, test_features: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        We compute marginals over f(x, r), where `test_features` are extended
        features.
        Note: The test configs must not overlap with any in the training set.
        Otherwise, at least if `r != r_max`, the predictive distributions
        computed here may be wrong.

        :param test_features: Extended features for test configs
        :return: posterior_means, posterior_variances
        """
        raise NotImplementedError()

    def sample_marginals(
            self, test_features: np.ndarray, num_samples: int = 1,
            random_state: Optional[RandomState] = None) -> np.ndarray:
        """
        See comments of `predict`.

        :param test_features: Input points for test configs
        :param num_samples: Number of samples
        :param random_state: PRNG
        :return: Marginal samples, (num_test, num_samples)
        """
        if random_state is None:
            random_state = np.random
        return sample_posterior_marginals(
            self.poster_state, self.mean, self.kernel, test_features,
            random_state=random_state, num_samples=num_samples)

    def backward_gradient(
            self, input: np.ndarray,
            head_gradients: Dict[str, np.ndarray],
            mean_data: float, std_data: float) -> np.ndarray:
        """
        Implements SurrogateModel.backward_gradient, see comments there.
        This is for a single posterior state. If the SurrogateModel uses
        MCMC, have to call this for every sample.

        :param input: Single input point x, shape (d,)
        :param head_gradients: See SurrogateModel.backward_gradient
        :param mean_data: Mean used to normalize targets
        :param std_data: Stddev used to normalize targets
        :return:
        """
        test_feature = np.reshape(input, (1, -1))

        def diff_test_feature(test_feature_array):
            norm_mean, norm_variance = self.predict(test_feature_array)
            # De-normalize, and variance -> stddev
            pred_mean = norm_mean * std_data + mean_data
            pred_std = anp.sqrt(norm_variance) * std_data
            head_gradients_mean = anp.reshape(head_gradients['mean'], pred_mean.shape)
            head_gradients_std = anp.reshape(head_gradients['std'], pred_std.shape)
            # Added to mimic mxnet.autograd.backward
            pred_mean_sum = _inner_product(pred_mean, head_gradients_mean)
            pred_std_sum = _inner_product(pred_std, head_gradients_std)
            return pred_mean_sum + pred_std_sum

        test_feature_gradient = grad(diff_test_feature)
        return np.reshape(test_feature_gradient(test_feature), input.shape)

    def _sample_curves_internal(
            self, data: Dict, poster_state: Dict, num_samples: int = 1,
            random_state: Optional[RandomState] = None) -> List[Dict]:
        raise NotImplementedError()

    def sample_curves(
            self, data: Dict, num_samples: int = 1,
            random_state: Optional[RandomState] = None) -> List[Dict]:
        """
        Given data from one or more configurations (as returned by
        `issm.prepare_data`), for each config, sample a curve from the
        joint posterior (predictive) distribution over latent targets.
        The curve for each config in `data` may be partly observed, but
        must not be fully observed. Samples for the different configs are
        independent. None of the configs in `data` must appear in the dataset
        used to compute the posterior state.

        The result is a list of Dict, one for each config. If for a config,
        targets in `data` are given for resource values range(r_min, r_obs),
        the dict entry `y` is a joint sample [y_r], r in range(r_obs, r_max+1).
        For some subclasses (e.g., ISSM), there is also an entry `f` with a
        joint sample [f_r], r in range(r_obs-1, r_max+1), the latent function
        values before noise. These entries are matrices with `num_samples`
        columns, which are independent (the joint dependence is along the rows).

        :param data: Data for configs to predict at
        :param num_samples: Number of samples to draw from each curve
        :param random_state: PRNG state to be used for sampling
        :return: See above
        """
        return self._sample_curves_internal(
            data=data, poster_state=self.poster_state, num_samples=num_samples,
            random_state=random_state)

    def has_precomputations(self, data: Dict) -> bool:
        raise NotImplementedError()

class IncrementalUpdateGPAdditivePosteriorState(GaussProcAdditivePosteriorState):
    """
    Extension of :class:`GaussProcAdditivePosteriorState` which allows for
    incremental updating (single config added to the dataset).
    This is required for simulation-based scoring, and for support of
    fantasizing.

    """
    def __init__(
            self, data: Optional[Dict], mean: MeanFunction,
            kernel: KernelFunction, noise_variance, **kwargs):
        super().__init__(
            data, mean, kernel, noise_variance, **kwargs)

    def _prepare_update(
            self, feature: np.ndarray, targets: np.ndarray) \
            -> Tuple[float, float, float]:
        """
        Helper method for `update`. Returns new entries for d, s, r2 vectors.

        :param feature: See `update`
        :param targets: See `update`
        :return: (d_new, s_new, r2_new)
        """
        raise NotImplementedError()

    def _update_internal(
            self, feature: np.ndarray, targets: np.ndarray) -> Dict:
        """
        Update posterior state, given a single new datapoint. `feature`,
        `targets` are like one entry of `data`. The method returns a new
        object with the updated state.

        :param feature: See above
        :param targets: See above
        :return: Arguments to create new posterior state
        """
        # Update posterior state
        feature = _rowvec(feature, _np=np)
        d_new, s_new, r2_new = self._prepare_update(feature, targets)
        new_poster_state = update_posterior_state(
            self.poster_state, self.kernel, feature, d_new, s_new, r2_new)
        # Return args to create new object by way of "copy constructor"
        return dict(
            data=None,
            mean=self.mean,
            kernel=self.kernel,
            noise_variance=self.noise_variance,
            poster_state=new_poster_state,
            r_min=self.r_min, r_max=self.r_max)

    def update(
            self, feature: np.ndarray, targets: np.ndarray) \
            -> 'IncrementalUpdateGPAdditivePosteriorState':
        raise NotImplementedError()

    def update_pvec(
            self, feature: np.ndarray, targets: np.ndarray) -> np.ndarray:
        """
        Part of `update`: Only update prediction vector p. This cannot be used
        to update p for several new datapoints.

        :param feature:
        :param targets:
        :return: New p vector
        """
        feature = _rowvec(feature, _np=np)
        d_new, s_new, r2_new = self._prepare_update(feature, targets)
        return update_posterior_pvec(
            self.poster_state, self.kernel, feature, d_new, s_new, r2_new)

    def _sample_posterior_joint_for_config(
            self, poster_state: Dict, config: Configuration, feature: np.ndarray,
            targets: np.ndarray, random_state: RandomState) -> np.ndarray:
        data = {'features': [feature],
                'targets': [targets],
                'configs': [config]}
        results = self._sample_curves_internal(
            data, poster_state, num_samples=1, random_state=random_state)
        return results[0]['y']

    def sample_and_update_for_pending(
            self, data_pending: Dict, sample_all_nonobserved: bool = False,
            random_state: Optional[RandomState] = None) \
            -> (List[np.ndarray], 'IncrementalUpdateGPAdditivePosteriorState'):
        """
        This function is needed for sampling fantasy targets, and also to
        support simulation-based scoring.

        `issm.prepare_data_with_pending` creates two data dicts `data_nopending`,
        `data_pending`, the first for configs with observed data, but no
        pending evals, the second for configs with pending evals.
        You create the state with `data_nopending`, then call this method with
        `data_pending`.

        This method is iterating over configs (or trials) in `data_pending`.
        For each config, it draws a joint sample from some non-observed
        targets, then updates the state conditioned on observed and sampled
        targets (by calling `update`). If `sample_all_nonobserved` is False,
        the number of targets sampled is the entry in
        `data_pending['num_pending']`. Otherwise, targets are sampled for all
        non-observed positions.

        The method returns the list of sampled target vectors, and the state
        at the end (like `update` does as well).

        :param data_pending: See above
        :param sample_all_nonobserved: See above
        :param random_state: PRNG
        :return: pending_targets, final_state
        """
        if random_state is None:
            random_state = np.random
        curr_poster_state = self.poster_state
        targets_lst = []
        final_state = self
        for config, feature, targets, num_pending in zip(
                data_pending['configs'], data_pending['features'],
                data_pending['targets'], data_pending['num_pending']):
            # Draw joint sample
            fantasies = _flatvec(self._sample_posterior_joint_for_config(
                curr_poster_state, config, feature, targets, random_state),
                _np=np)
            if not sample_all_nonobserved:
                fantasies = fantasies[:num_pending]
            fantasies = _colvec(fantasies, _np=np)
            targets_lst.append(fantasies)
            # Update state
            full_targets = np.vstack((targets, fantasies))
            final_state = self.update(feature, full_targets)
            curr_poster_state = final_state.poster_state
        return targets_lst, final_state

class GaussProcISSMPosteriorState(IncrementalUpdateGPAdditivePosteriorState):
    """
    Represent posterior state for joint Gaussian model of learning curves over
    a number of configurations. The model is the sum of a Gaussian process
    model for function values at r_max and independent Gaussian linear
    innovation state space models (ISSMs) of a particular power law decay
    form.

    """
    def __init__(
            self, data: Optional[Dict], mean: MeanFunction,
            kernel: KernelFunction, iss_model: ISSModelParameters,
            noise_variance, **kwargs):
        """
        `data` contains input points and targets, as obtained from
        `issm.prepare_data`. `iss_model` maintains the ISSM parameters.

        :param data: Input points and targets
        :param mean: Mean function m(X)
        :param kernel: Kernel function k(X, X')
        :param iss_model: ISS model
        :param noise_variance: Innovation and noise variance
        """
        self.iss_model = iss_model
        super().__init__(
            data, mean, kernel, noise_variance=noise_variance, **kwargs)

    def has_precomputations(self, data: Dict) -> bool:
        return all(k in data for k in ('ydims', 'num_configs', 'deltay', 'logr'))

    def _compute_posterior_state(
            self, data: Dict, noise_variance, **kwargs):
        profiler = kwargs.get('profiler')
        # Compute posterior state
        issm_params = self.iss_model.get_issm_params(data['features'])
        if self.has_precomputations(data):
            issm_likelihood = issm_likelihood_computations(
                precomputed=data,
                issm_params=issm_params,
                r_min=self.r_min, r_max=self.r_max,
                profiler=profiler)
        else:
            issm_likelihood = issm_likelihood_slow_computations(
                targets=data['targets'],
                issm_params=issm_params,
                r_min=self.r_min, r_max=self.r_max,
                profiler=profiler)
        if profiler is not None:
            profiler.start('poster_comp')
        self.poster_state = posterior_computations(
            features=data['features'],
            mean=self.mean, kernel=self.kernel,
            issm_likelihood=issm_likelihood,
            noise_variance=noise_variance)
        if profiler is not None:
            profiler.stop('poster_comp')

    def predict(
            self, test_features: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        resource_attr_range = (self.r_min, self.r_max)
        features, resources = decode_features(
            test_features, resource_attr_range=resource_attr_range)
        issm_params = self.iss_model.get_issm_params(features)
        return predict_posterior_marginals_issm(
            poster_state=self.poster_state,
            mean=self.mean,
            kernel=self.kernel,
            test_features=features,
            resources=resources,
            issm_params=issm_params,
            r_min=self.r_min, r_max=self.r_max)

    @staticmethod
    def data_precomputations(data: Dict):
        logger.info("Enhancing data dictionary by precomputed variables")
        data.update(issm_likelihood_precomputations(
            targets=data['targets'], r_min=data['r_min']))

    def _sample_curves_internal(
            self, data: Dict, poster_state: Dict, num_samples: int = 1,
            random_state: Optional[RandomState] = None) -> List[Dict]:
        if random_state is None:
            random_state = np.random
        results = []
        for feature, targets, config in zip(
                data['features'], data['targets'], data['configs']):
            issm_params = self.iss_model.get_issm_params(
                feature.reshape((1, -1)))
            results.append(sample_posterior_joint_issm(
                poster_state=poster_state,
                mean=self.mean,
                kernel=self.kernel,
                feature=feature,
                targets=targets,
                issm_params=issm_params,
                r_min=self.r_min, r_max=self.r_max,
                random_state=random_state,
                num_samples=num_samples))
        return results

    def _prepare_update(
            self, feature: np.ndarray, targets: np.ndarray) \
            -> Tuple[float, float, float]:
        issm_params = self.iss_model.get_issm_params(feature.reshape((1, -1)))
        issm_likelihood = issm_likelihood_slow_computations(
            targets=[_colvec(targets, _np=np)],
            issm_params=issm_params,
            r_min=self.r_min, r_max=self.r_max)
        d_new = issm_likelihood['d'].item()
        vtv = issm_likelihood['vtv'].item()
        wtv = issm_likelihood['wtv'].item()
        s_sq = vtv / self.noise_variance
        s_new = np.sqrt(s_sq)
        muhat = _flatvec(self.mean(feature)).item() - issm_likelihood['c'].item()
        r2_new = wtv / self.noise_variance - s_sq * muhat
        return d_new, s_new, r2_new

    def update(
            self, feature: np.ndarray, targets: np.ndarray) \
            -> 'IncrementalUpdateGPAdditivePosteriorState':
        create_kwargs = self._update_internal(feature, targets)
        return GaussProcISSMPosteriorState(
            **create_kwargs, iss_model=self.iss_model)

class GaussProcExpDecayPosteriorState(IncrementalUpdateGPAdditivePosteriorState):
    """
    Represent posterior state for joint Gaussian model of learning curves over
    a number of configurations. The model is the sum of a Gaussian process
    model for function values at r_max and independent Gaussian processes over
    r, using an exponential decay covariance function. The latter is shared
    between all configs.

    This is essentially the model from the Freeze Thaw paper (see also
    :class:`ExponentialDecayResourcesKernelFunction`).

    """
    def __init__(
            self, data: Optional[Dict],
            mean: MeanFunction,
            kernel: KernelFunction,
            res_kernel: ExponentialDecayBaseKernelFunction,
            noise_variance, **kwargs):
        """
        `data` contains input points and targets, as obtained from
        `issm.prepare_data`.

        :param data: Input points and targets
        :param mean: Mean function m(X)
        :param kernel: Kernel function k(X, X')
        :param res_kernel: Kernel function k_r(r, r'), of exponential decay
            type
        :param noise_variance: Innovation and noise variance
        """
        self.res_kernel = res_kernel
        super().__init__(
            data, mean, kernel, noise_variance=noise_variance, **kwargs)
        assert self.r_min == res_kernel.r_min and self.r_max == res_kernel.r_max, \
            ((self.r_min, self.r_max), (res_kernel.r_min, res_kernel.r_max))

    def has_precomputations(self, data: Dict) -> bool:
        return all(k in data for k in ('ydims', 'num_configs', 'yflat'))

    def _compute_posterior_state(
            self, data: Dict, noise_variance, **kwargs):
        profiler = kwargs.get('profiler')
        # Compute posterior state
        if profiler is not None:
            profiler.start('likelihood')
        if self.has_precomputations(data):
            issm_likelihood = resource_kernel_likelihood_computations(
                precomputed=data,
                res_kernel=self.res_kernel,
                noise_variance=noise_variance)
        else:
            issm_likelihood = resource_kernel_likelihood_slow_computations(
                targets=data['targets'],
                res_kernel=self.res_kernel,
                noise_variance=noise_variance)
        if profiler is not None:
            profiler.stop('likelihood')
            profiler.start('poster_comp')
        self.poster_state = posterior_computations(
            features=data['features'],
            mean=self.mean, kernel=self.kernel,
            issm_likelihood=issm_likelihood,
            noise_variance=noise_variance)
        if profiler is not None:
            profiler.stop('poster_comp')
        # Add missing term to criterion value
        if 'criterion' in self.poster_state:
            part3 = logdet_cholfact_cov_resource(issm_likelihood)
            self.poster_state['criterion'] += part3
        # Extra terms required in `sample_curves`
        self.poster_state['lfact_all'] = issm_likelihood['lfact_all']
        self.poster_state['means_all'] = issm_likelihood['means_all']
        self.poster_state['noise_variance'] = noise_variance

    def predict(
            self, test_features: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        resource_attr_range = (self.r_min, self.r_max)
        features, resources = decode_features(
            test_features, resource_attr_range=resource_attr_range)
        return predict_posterior_marginals_expdecay(
            poster_state=self.poster_state,
            mean=self.mean,
            kernel=self.kernel,
            test_features=features,
            resources=resources,
            res_kernel=self.res_kernel)

    @staticmethod
    def data_precomputations(data: Dict):
        data.update(resource_kernel_likelihood_precomputations(
            targets=data['targets']))

    def _sample_curves_internal(
            self, data: Dict, poster_state: Dict, num_samples: int = 1,
            random_state: Optional[RandomState] = None) -> List[Dict]:
        assert 'lfact_all' in poster_state
        if random_state is None:
            random_state = np.random
        lfact_all = poster_state['lfact_all']
        means_all = poster_state['means_all']
        noise_variance = poster_state['noise_variance']
        results = []
        for feature, targets in zip(data['features'], data['targets']):
            results.append(sample_posterior_joint_expdecay(
                poster_state=poster_state,
                mean=self.mean,
                kernel=self.kernel,
                feature=feature,
                targets=targets,
                res_kernel=self.res_kernel,
                noise_variance=noise_variance,
                lfact_all=lfact_all,
                means_all=means_all,
                random_state=random_state,
                num_samples=num_samples))
        return results

    def _prepare_update(
            self, feature: np.ndarray, targets: np.ndarray) \
            -> Tuple[float, float, float]:
        issm_likelihood = resource_kernel_likelihood_slow_computations(
            targets=[_colvec(targets, _np=np)],
            res_kernel=self.res_kernel,
            noise_variance=self.noise_variance)
        vtv = issm_likelihood['vtv'].item()
        wtv = issm_likelihood['wtv'].item()
        s_sq = vtv / self.noise_variance
        s_new = np.sqrt(s_sq)
        muhat = _flatvec(self.mean(feature)).item()
        r2_new = wtv / self.noise_variance - s_sq * muhat
        return 0.0, s_new, r2_new

    def update(
            self, feature: np.ndarray, targets: np.ndarray) \
            -> 'IncrementalUpdateGPAdditivePosteriorState':
        create_kwargs = self._update_internal(feature, targets)
        # Extra terms required in `sample_curves`
        new_poster_state = create_kwargs['poster_state']
        for k in ('lfact_all', 'means_all', 'noise_variance'):
            new_poster_state[k] = self.poster_state[k]
        return GaussProcExpDecayPosteriorState(
            **create_kwargs, res_kernel=self.res_kernel)