#   Copyright 2020 The PyMC Developers
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License 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.

# jms: This file is adapted from sample_nfmc.py / sample_smc.py

import logging
import multiprocessing as mp
import time
import warnings

from collections.abc import Iterable

import numpy as np
from scipy.special import logsumexp

from pymc3.backends.base import MultiTrace
from pymc3.model import modelcontext
from pymc3.parallel_sampling import _cpu_count
from pymc3.nfo.nfo import NFO
import torch
from scipy.stats import multivariate_normal



def opt_nfo(
    #Optimization parameters
    #initialization
    n0=10, #int, n0 the initial number of draws
    init_samples=None, #array, Whether to provide some pre-defined sequence or do pymc3 sampling
    #approximation
    k_trunc=np.inf, #IW clipping, not used by default
    eps_z=0.01, #float, tolerance on Z for q iter convergence (eps') #currently not used since not iterating SINF unless trainable
    nf_iter=1, #int, number of NF iters -should always be 1 in our implementation
    #annealing
    N=10, #int, N the TOTAL number of draws we want at each iteration - this is no longer used, is from when we used to run multiple fits
    t_ess=0.5, #float, ESS<t_ess*n0 t threshold on ESS for ESS3 (no longer temperature)
    g_AF=0, #float, size of gradient contribution to AF, not used now
    #exploration
    N_AF=1000, #int,number of points to use in q_w sampling for AF
    expl_top_AF=1, #int,cut for the top AF at a given temp level accepted at each beta
    expl_latent=0, #int,latent draw from around top IW1 or around random draw from q_w, accepted at each step
    expl_top_qw=0, #int,keep top q_w at this iteration
    beta_max=1, #float>0,highest exponent on tempered posterior, support >1 for exploitation
    rel_beta=1, #0<float<1, β2 = rel_beta*β, where β2 is the lower temp level used for sampling q_w, what we call 'X'
    frac_rel_beta_AF=1, #int, the modifier to the AF used to up/down-weight the w vs uw contribution, what we call "Y"
    latent_sigma=None, #float, the value of l
    use_latent_beta2=False, #whether to get the latent sample from q_w(β2) or from q_uw
    use_pq_beta_IW1=False, #whether to get the latent sample from near top IW1 or randomly from q_w
    bounds=None,#array, size 2xd, bounding box for samples FIXME make this more obvious, needed for prior
    N_temp=25, #int, cutoff on number of allowed temp iterations before giving up -> #FIXME eventually make this throw error
    #NF parameters
    model=None,
    frac_validate=0.0,
    iteration=None,
    alpha_w=(0,0),
    alpha_uw=(0,0),
    verbose=False,
    n_component=None,
    interp_nbin=None,
    KDE=True,
    bw_factor_min=1.0,
    bw_factor_max=1.0,
    bw_factor_num=1,
    rel_bw=1,
    edge_bins=None,
    ndata_wT=None,
    MSWD_max_iter=None,
    NBfirstlayer=True,
    logit=False,
    Whiten=False,
    trainable_qw=False, #whether to improve our q_w at each beta iteration with SGD
    sgd_steps=0, #number of steps used in Adam when training trainable q_w
    knots_trainable=5,
    batchsize=None,
    nocuda=False,
    patch=False,
    shape=[28,28,1],
    #Runtime
    random_seed=-1,
    parallel=False,
    cores=None
):
    r"""
    Normalizing flow-based Bayesian Optimization.

    Parameters
    ----------
    draws: int
        The number of samples to draw from the posterior (i.e. last stage). And also the number of
        independent chains. Defaults to 2000.
    norm_tol: float
        Fractional difference in the evidence estimate between two steps. If it falls below this we
        stop iterating over the NF fits.
    optim_iter: int
        Maximum number of optimization steps to run during the initialization.
    nf_iter: int
        Number of NF fit iterations to go through after the optimization step.
    model: Model (optional if in ``with`` context)).
    frac_validate: float
        Fraction of the live points at each NS iteration that we use for validation of the NF fit.
    alpha: tuple of floats
        Regularization parameters used for the NF fit.
    verbose: boolean
        Whether you want verbose output from the NF fit.
    random_seed: int
        random seed
    parallel: bool
        Distribute computations across cores if the number of cores is larger than 1.
        Defaults to False.
    cores : int
        Number of cores available for the optimization step. Defaults to None, in which case the CPU
        count is used.

    """

    _log = logging.getLogger("pymc3")
    _log.info("Initializing normalizing flow-based optimization...")

    model = modelcontext(model)
    if model.name:
        raise NotImplementedError(
            "The NS_NFO implementation currently does not support named models. "
            "See https://github.com/pymc-devs/pymc3/pull/4365."
        )
    if cores is None:
        cores = _cpu_count()
    chains=1

    _log.info(
        f"Sampling {chains} chain{'s' if chains > 1 else ''} "
        f"Cores available for optimization: {cores}"
    )
    if random_seed == -1:
        random_seed = None
    if chains == 1 and isinstance(random_seed, int):
        random_seed = [random_seed]
    if random_seed is None or isinstance(random_seed, int):
        if random_seed is not None:
            np.random.seed(random_seed)
        random_seed = [np.random.randint(2 ** 30) for _ in range(chains)]
    if not isinstance(random_seed, Iterable):
        raise TypeError("Invalid value for `random_seed`. Must be tuple, list or int")

    #we changed the name for end-user-facing readability, but internally more familiar with these names
    aN,bN,cN,dN = N_AF,expl_top_AF,expl_latent,expl_top_qw

    params = (
        n0,
        init_samples,
        k_trunc,
        eps_z,
        nf_iter,
        N,
        t_ess,
        g_AF,
        aN,
        bN,
        cN,
        dN,
        beta_max,
        rel_beta,
        frac_rel_beta_AF,
        latent_sigma,
        use_latent_beta2,
        use_pq_beta_IW1,
        bounds,
        N_temp,
        model,
        frac_validate,
        iteration,
        alpha_w,
        alpha_uw,
        cores,
        verbose,
        n_component,
        interp_nbin,
        KDE,
        bw_factor_min,
        bw_factor_max,
        bw_factor_num,
        rel_bw,
        edge_bins,
        ndata_wT,
        MSWD_max_iter,
        NBfirstlayer,
        logit,
        Whiten,
        trainable_qw,
        sgd_steps,
        knots_trainable,
        batchsize,
        nocuda,
        patch,
        shape,
        parallel,
    )

    t1 = time.time()

    results = []
    for i in range(chains):
        results.append(opt_nfo_int(*params, random_seed[i], i, _log))
    (
        traces,
        log_evidence,
        q_samples,
        importance_weights,
        logp,
        logq,
        train_logp,
        train_logq,
        logZ,
        q_models,
        q_ess,
        total_ess,
        min_var_bws,
        min_pq_bws,
        betas,
    ) = zip(*results)
    trace = MultiTrace(traces)
    trace.report.log_evidence = log_evidence
    trace.report.q_samples = q_samples
    trace.report.importance_weights = importance_weights
    trace.report.logp = logp
    trace.report.logq = logq
    trace.report.train_logp = train_logp
    trace.report.train_logq = train_logq
    trace.report.logZ = logZ
    trace.report.q_models = q_models
    trace.report.q_ess = q_ess
    trace.report.total_ess = total_ess
    trace.report.N = N
    trace.report.min_var_bws = min_var_bws
    trace.report.min_pq_bws = min_pq_bws
    trace.report._t_sampling = time.time() - t1
    trace.report.betas = betas
    return trace


def opt_nfo_int(
    n0,
    init_samples,
    k_trunc,
    eps_z,
    nf_iter,
    N,
    t_ess,
    g_AF,
    aN,
    bN,
    cN,
    dN,
    beta_max,
    rel_beta,
    frac_rel_beta_AF,
    latent_sigma,
    use_latent_beta2,
    use_pq_beta_IW1,
    bounds,
    N_temp,
    model,
    frac_validate,
    iteration,
    alpha_w,
    alpha_uw,
    cores,
    verbose,
    n_component,
    interp_nbin,
    KDE,
    bw_factor_min,
    bw_factor_max,
    bw_factor_num,
    rel_bw,
    edge_bins,
    ndata_wT,
    MSWD_max_iter,
    NBfirstlayer,
    logit,
    Whiten,
    trainable_qw,
    sgd_steps,
    knots_trainable,
    batchsize,
    nocuda,
    patch,
    shape,
    parallel,
    random_seed,
    chain,
    _log,
):
    """Run one NS_NFO instance."""
    nfo = NFO(
        n0=n0,
        init_samples=init_samples,
        k_trunc=k_trunc,
        eps_z=eps_z,
        nf_iter=nf_iter,
        N=N,
        t_ess=t_ess,
        beta_max=beta_max,
        model=model,
        frac_validate=frac_validate,
        iteration=iteration,
        alpha_w=alpha_w,
        alpha_uw=alpha_uw,
        verbose=verbose,
        n_component=n_component,
        interp_nbin=interp_nbin,
        KDE=KDE,
        bw_factor_min=bw_factor_min,
        bw_factor_max=bw_factor_max,
        bw_factor_num=bw_factor_num,
        rel_bw=rel_bw,
        edge_bins=edge_bins,
        ndata_wT=ndata_wT,
        MSWD_max_iter=MSWD_max_iter,
        NBfirstlayer=NBfirstlayer,
        logit=logit,
        Whiten=Whiten,
        trainable_qw=trainable_qw,
        sgd_steps=sgd_steps,
        knots_trainable=knots_trainable,
        batchsize=batchsize,
        nocuda=nocuda,
        bounds=bounds,
        patch=patch,
        shape=shape,
        chain=chain,
        random_seed=random_seed,
    )

    #dicts for storage #FIXME compactify
    iter_sample_dict = {}
    iter_weight_dict = {}
    iter_logp_dict = {}
    iter_logq_dict = {}
    iter_train_logp_dict = {}
    iter_train_logq_dict = {}
    iter_logZ_dict = {}
    iter_qmodel_dict = {}
    iter_q_ess_dict = {}
    iter_total_ess_dict = {}
    iter_min_var_bw_dict = {}
    iter_min_pq_bw_dict = {}

    '''///////////////////////////////'''
    '''/         INITIALIZE          /'''
    '''///////////////////////////////'''
    stage = 0
    print("stage: ", stage)
    betas = []
    nfo.initialize_var_info()
    nfo.setup_logp()
    nfo.initialize_population() #uses n0
    theta_0 = nfo.samples
    nfo.theta = theta_0
    nfo.latent_sigma=latent_sigma #for dict stuff
    betas.append(0)
    #need these in the first iteration to get weights for the fit in fit_nf
    #later replace prior with log_q in that function as we go
    nfo.sinf_logw = nfo.posterior_logp-nfo.prior_logp
    #this doesn't actually do anything except in prior
    nfo.importance_weights_pq = np.exp(nfo.posterior_logp-nfo.prior_logp-nfo.log_evidence_pq)

    save_dicts(nfo,stage,_log,iter_sample_dict,iter_weight_dict,iter_logp_dict,iter_logZ_dict,
    iter_qmodel_dict,iter_q_ess_dict,iter_total_ess_dict,iter_logq_dict,iter_train_logp_dict,iter_train_logq_dict,
    iter_min_var_bw_dict,iter_min_pq_bw_dict,betas=None)

    '''///////////////////////////////'''
    '''/           ANNEAL            /'''
    '''///////////////////////////////'''
    for n in range(N_temp):
        stage += 1
        print("stage: ", stage)
        #Fit unweighted approximator to samples
        nfo.fit_nf(weighted=False) #unweighted fit first to get IW
        quw = nfo.nf_model_uw
        #Select a temperature level beta based on ESS3
        nfo.logIW3 = nfo.posterior_logp - nfo.logq_uw
        nfo.update_weights_beta()

        #check if we want to sample from lower beta, gross
        if(stage>1):
            nearest_beta2_idx = (np.abs(np.array(betas)[:-1] - rel_beta*nfo.beta)).argmin()
        else:
            nearest_beta2_idx = 0

        '''///////////////////////////////'''
        '''/         APPROXIMATE         /'''
        '''///////////////////////////////'''
        nfo.fit_nf(weighted=True) #weighted fit, this knows about beta, also knows about trainable SINF
        qw = nfo.nf_model_w

        '''///////////////////////////////'''
        '''/           EXPLORE           /'''
        '''///////////////////////////////'''

        #sample from lower beta only if we are at high enough beta such that
        #rel_beta*beta >= betas[1] - usually this only excludes 1st step or so
        if(nearest_beta2_idx != 0 and stage>1 and rel_beta <1): #we can't use lower beta on first beta
            print("using beta2={0:.2f} at idx {1}, at stage-1 = {2}".format(betas[nearest_beta2_idx],nearest_beta2_idx,stage-1))
            qw_beta2 = iter_qmodel_dict['q{0}_w'.format(nearest_beta2_idx)]  #get the qw at the lower beta we need
            #FIXME lazy copy, make this a function, dropping gradient since sort of defunct now
            theta_w_beta2,_ = qw_beta2.sample(aN,device='cpu') #use aN for now
            nfo.beta_2 = betas[nearest_beta2_idx]
            nfo.tens_theta_w_beta2 = theta_w_beta2
            #evaluate fitted qs at this beta on samples from beta_2
            logq_uw_beta12 = t2a(quw.evaluate_density(theta_w_beta2))
            logq_w_beta12 = t2a(qw.evaluate_density(theta_w_beta2))

            theta_w_beta2 = t2a(theta_w_beta2) #has to happen after grad, grad wants tensor
            if(bounds is not None): #make sure we don't keep samples out of prior range, ow can get maxima out of bounds
                prior_mask_beta2 = np.all(theta_w_beta2>bounds[0],axis=-1) & np.all(theta_w_beta2<bounds[1],axis=-1)
                theta_w_beta2,logq_w_beta12,logq_uw_beta12 = theta_w_beta2[prior_mask_beta2],logq_w_beta12[prior_mask_beta2],logq_uw_beta12[prior_mask_beta2]
            logIW2_beta12 = np.squeeze(logq_w_beta12-logq_uw_beta12)
            logIW2 = logIW2_beta12
            theta_w = theta_w_beta2
            #define new acquisition function
            logIW2prime = logIW2 + (-1+rel_beta*frac_rel_beta_AF)*logq_w_beta12 #use downweighted version
            logIW2prime -= logsumexp(logIW2prime) #, normalize them
            nfo.AF = np.exp(logIW2prime)
            nfo.logIW2 = logIW2prime #just for output later
            #everything else outside this if statement is the same...

            #Draw a sample from q_uw via latent space near highest IW1
            if(latent_sigma is not None):
                max_latent=15#fixed for now, indicates a problem
                tries=0
                for i in range(max_latent):
                    tries+=1
                    if(use_pq_beta_IW1):
                        ranked_IW1 = np.argsort(nfo.importance_weights_pq_beta)[::-1]
                        top_IW1_sample = a2t(nfo.samples[ranked_IW1[:cN]]) #top pbeta/qbeta
                        # print("top IW sample ", top_IW1_sample)
                    else:
                        # ranked_IW1 = np.argsort(nfo.importance_weights_pq)[::-1]
                        #WARNING this is not actually IW1, just doing quick test
                        #replace IW1 with random draw from q_w at current beta level
                        qw_sample_rand,_ = qw.sample(cN,device='cpu')
                        top_IW1_sample = qw_sample_rand
                        # print("rand qw expl sample ", top_IW1_sample)

                    if(use_latent_beta2):
                        latent_z = np.squeeze(t2a(qw_beta2.inverse(top_IW1_sample)[0])) #go to latent space
                        latent_sigma = latent_sigma*np.ones_like(latent_z) #sphere width to sample from in latent space
                        z_prime = multivariate_normal.rvs(mean=latent_z,cov=latent_sigma,size=1)#,random_state=random_seed) #sample
                        theta_z = t2a(qw_beta2.forward(a2t(z_prime))[0]) #go back to data space
                    else:
                        latent_z = np.squeeze(t2a(quw.inverse(top_IW1_sample)[0])) #go to latent space
                        latent_sigma = latent_sigma*np.ones_like(latent_z) #sphere width to sample from in latent space
                        z_prime = multivariate_normal.rvs(mean=latent_z,cov=latent_sigma,size=1)#,random_state=random_seed) #sample
                        theta_z = t2a(quw.forward(a2t(z_prime))[0]) #go back to data space

                    # print("latent_z", latent_z)
                    # print("latent_sigma", latent_sigma)

                    z_prime = multivariate_normal.rvs(mean=latent_z,cov=latent_sigma,size=1)#,random_state=random_seed) #sample
                    theta_z = t2a(quw.forward(a2t(z_prime))[0]) #go back to data space
                    # print('z_prime: ', z_prime)

                    #take extra theta for iter from q_w(beta2) at highest q_w(beta)
                    if(dN>0):
                        idx_top_qw = np.argsort(np.squeeze(logq_w_beta12))[::-1]
                        top_qw = theta_w_beta2[idx_top_qw[:dN]]
                        theta_z = np.concatenate([theta_z,top_qw])
                    # print('theta_z: ', theta_z)

                    bound_check = np.all(theta_z>bounds[0],axis=-1) & np.all(theta_z<bounds[1],axis=-1)
                    if(np.all(bound_check)):
                        break #be done
                    #else we just keep going until max
                    if(tries==max_latent-1): raise ValueError("Too many OOB samples! Try decreasing latent_sigma")
                nfo.top_IW1_sample=t2a(top_IW1_sample)#for saving to dict later
                print("finished with latent draw, took {0} tries".format(i))

        else: #no sampling from lower beta levels
            theta_w,logq_w =qw.sample(aN,device='cpu')
            dlogq_w = qw.score(theta_w).numpy().astype(np.float64) #gradient dlogqw/dtheta

            nfo.tens_theta_w = theta_w
            logq_uw = t2a(quw.evaluate_density(theta_w))
            theta_w,logq_w = t2a(theta_w),t2a(logq_w) #has to happen after grad, grad wants tensor
            if(bounds is not None): #make sure we don't keep samples out of prior range, ow can get maxima out of bounds
                prior_mask = np.all(theta_w>bounds[0],axis=-1) & np.all(theta_w<bounds[1],axis=-1)
                theta_w,logq_w,logq_uw = theta_w[prior_mask],logq_w[prior_mask],logq_uw[prior_mask]
                dlogq_w = dlogq_w[prior_mask]
            logIW2 = np.squeeze(logq_w-logq_uw)#many (aN) values
            nfo.w_samples_beta = theta_w
            nfo.beta_2 = None
            #compute (scalar) acquisition function
            logIW2 -= logsumexp(logIW2) #, normalize them
            nfo.logIW2 = logIW2
            dim_sum_dlog_w = np.sum(dlogq_w,axis=-1) #trailing axis is size dims
            g_term = -g_AF*( (np.exp(np.squeeze(dim_sum_dlog_w) -logsumexp(np.squeeze(dim_sum_dlog_w))))**2 ) #normalize gradient term so g is not crazy
            nfo.AF = np.exp(logIW2) + g_term


        '''///////////////////////////////'''
        '''/           EXPLOIT           /'''
        '''///////////////////////////////'''

        nfo.w_samples = theta_w
        # ESS2 = len(logIW2)*nfo.calculate_ess(logIW2) not actually used?
        ranked_AF = np.argsort(nfo.AF)[::-1]
        nfo.ranked_topAF=ranked_AF[:bN]
        #rank order the AF results
        theta_new = theta_w[nfo.ranked_topAF] #check setting to see how many new points at each iteration

        #FIXME: this is very gross
        if(latent_sigma is not None and stage>1 and rel_beta <1 and nearest_beta2_idx != 0 ):
            #add new sample from latent space if want, this should get updated in the update call for p
            nfo.theta_no_z = theta_new
            theta_new = np.concatenate([theta_new,theta_z])
            #for saving to dicts later
            nfo.theta_z = theta_z

        new_sample_update(nfo,theta_new)
        #finish up
        save_dicts(nfo,stage,_log,iter_sample_dict,iter_weight_dict,iter_logp_dict,iter_logZ_dict,
        iter_qmodel_dict,iter_q_ess_dict,iter_total_ess_dict,iter_logq_dict,iter_train_logp_dict,iter_train_logq_dict,
        iter_min_var_bw_dict,iter_min_pq_bw_dict,betas=betas)
        #check if done
        if nfo.beta ==1: break
        if n==N_temp-1: print("WARNING: Failed to converge within N_temp temperature steps.")
    print("Finished")
    nfo.posterior = nfo.samples

    return (
        nfo.posterior_to_trace(),
        nfo.log_evidence,
        iter_sample_dict,
        iter_weight_dict,
        iter_logp_dict,
        iter_logq_dict,
        iter_train_logp_dict,
        iter_train_logq_dict,
        iter_logZ_dict,
        iter_qmodel_dict,
        iter_q_ess_dict,
        iter_total_ess_dict,
        iter_min_var_bw_dict,
        iter_min_pq_bw_dict,
        betas,
    )

#fix this, got grosser than I thought it would be
def save_dicts(nfo,stage,_log,
               iter_sample_dict,iter_weight_dict,
               iter_logp_dict,iter_logZ_dict,
               iter_qmodel_dict,iter_q_ess_dict,
               iter_total_ess_dict,iter_logq_dict,
               iter_train_logp_dict,iter_train_logq_dict,
               iter_min_var_bw_dict,iter_min_pq_bw_dict,betas=None):
    'Convenience function b/c dicts clog up the main loop'
    #iter
    nfo.nf_samples_to_trace()
    iter_sample_dict[f'q{int(stage)}'] = nfo.nf_trace
    iter_weight_dict[f'q{int(stage)}'] = nfo.weights
    iter_weight_dict[f'q{int(stage)}_pq'] = nfo.weights_pq
    iter_logp_dict[f'q{int(stage)}'] = nfo.posterior_logp
    iter_logZ_dict[f'q{int(stage)}'] = nfo.log_evidence
    iter_logZ_dict[f'q{int(stage)}_pq'] = nfo.log_evidence_pq
    iter_weight_dict[f'q{int(stage)}_pq'] = nfo.weights_pq

    # iter_qmodel_dict[f'q{int(stage)}'] = nfo.nf_model
    iter_q_ess_dict[f'q{int(stage)}'] = nfo.q_ess
    iter_total_ess_dict[f'q{int(stage)}'] = nfo.total_ess
    iter_total_ess_dict[f'q{int(stage)}_log_mean_loss'] = nfo.log_mean_loss
    print(f"logZ: {nfo.log_evidence:.3f}, ESS/N: {nfo.q_ess:.3f}, logZ_pq: {nfo.log_evidence_pq:.3f} log mean loss: {nfo.log_mean_loss:.3f}")
    if _log is not None:
        _log.info(f"Stage: {stage:3d}, logZ Estimate: {nfo.log_evidence:.3f},logZ_pq Estimate: {nfo.log_evidence_pq:.3f}")
        _log.info(f"Stage: {stage:3d}, q ESS/N: {nfo.q_ess:.3f}")

    if(stage>0): #these aren't defined in the initial
        if(betas is not None): betas.append(nfo.beta)
        iter_qmodel_dict[f'q{int(stage)}_w'] = nfo.nf_model_w
        iter_qmodel_dict[f'q{int(stage)}_uw'] = nfo.nf_model_uw
        iter_logq_dict[f'q{int(stage)}'] = nfo.logq
        iter_logq_dict[f'q{int(stage)}_w'] = nfo.logq_w
        iter_logq_dict[f'q{int(stage)}_uw'] = nfo.logq_uw
        iter_train_logp_dict[f'q{stage}'] = nfo.train_logp
        iter_train_logq_dict[f'q{stage}'] = nfo.train_logq
        iter_min_var_bw_dict[f'q{int(stage)}'] = nfo.min_var_bw
        iter_min_pq_bw_dict[f'q{int(stage)}'] = nfo.min_pq_bw
        iter_sample_dict[f'q{int(stage)}_w_samples'] = nfo.w_samples
        iter_weight_dict[f'q{int(stage)}_IW2'] = nfo.logIW2
        iter_weight_dict[f'q{int(stage)}_IW3'] = nfo.logIW3
        iter_weight_dict[f'q{int(stage)}_AF'] = nfo.AF
        iter_weight_dict[f'q{int(stage)}_IW2_p'] = nfo.logIW2_p
        iter_weight_dict[f'q{int(stage)}_AF_p'] = nfo.AF_p
        iter_logZ_dict[f'q{int(stage)}_pq_uw'] = nfo.log_evidence_pq_uw
        iter_logZ_dict[f'q{int(stage)}_pq_w'] = nfo.log_evidence_pq_w
        iter_weight_dict[f'q{int(stage)}_pq_uw'] = nfo.log_weight_pq_uw
        iter_weight_dict[f'q{int(stage)}_pq_w'] = nfo.log_weight_pq_w
        #gradient of q_w on samples where posterior is evaluated (regular samples)
        iter_sample_dict[f'grad_q{int(stage)}'] = nfo.grad_w_at_p
        if(nfo.latent_sigma is not None and stage>1 and nfo.beta_2 is not None):
            iter_sample_dict[f'topIW1_q{int(stage)}'] = nfo.top_IW1_sample
            iter_sample_dict[f'non_latent_q{int(stage)}'] = nfo.theta_no_z
            iter_sample_dict[f'latent_q{int(stage)}'] = nfo.theta_z

        if(nfo.beta_2 is not None):
            iter_sample_dict[f'beta2_{int(stage)}'] = nfo.beta_2

        if _log is not None:
            _log.info(f"Stage: {stage:3d}, Min variance BW factor: {nfo.min_var_bw}, Var(IW): {nfo.min_var_weights}, Min Zpq BW factor: {nfo.min_pq_bw}")
            if(betas is not None): _log.info(f"Stage: {stage:3d} Iter Beta: {nfo.beta:.3f}, True Beta: {nfo.beta*nfo.beta_max:.3f}")

def new_sample_update(nfo,theta_new):
    #update current samples and logp + other values
    tens_theta_new = torch.from_numpy(np.atleast_1d(theta_new).astype(np.float32))
    posterior_logp_new = np.array([nfo.posterior_logp_func(sample) for sample in theta_new]).squeeze()
    prior_logp_new = np.array([nfo.prior_logp_func(sample) for sample in theta_new]).squeeze()
    nfo.theta = np.concatenate([nfo.theta,theta_new])
    nfo.posterior_logp = np.concatenate([nfo.posterior_logp,np.atleast_1d(posterior_logp_new)])
    nfo.prior_logp = np.concatenate([nfo.prior_logp,np.atleast_1d(prior_logp_new)])
    #some of this is probably redundant, but need to update these for beta iteration
    #IW3 computed using tempered posterior - this is just for output, not actually used in fits.
    nfo.logIW3 = nfo.beta_max*nfo.beta*(nfo.posterior_logp-nfo.prior_logp) + nfo.prior_logp -np.concatenate([nfo.logq_uw,
                                            np.atleast_1d(nfo.nf_model_uw.evaluate_density(tens_theta_new).numpy().astype(np.float64))]) #this should be quw(theta)
    nfo.logIW2_p = nfo.logIW2[nfo.ranked_topAF] #keep the top one
    nfo.AF_p = nfo.AF[nfo.ranked_topAF]
    nfo.samples = nfo.theta
    nfo.nf_samples = np.copy(nfo.samples)
    #norm doesn't matter for IW1
    logqw_new = t2a(nfo.nf_model_w.evaluate_density(a2t(nfo.nf_samples)))
    nfo.importance_weights_pq = np.exp(nfo.posterior_logp- logqw_new)
    nfo.importance_weights_pq_beta = np.exp(nfo.posterior_logp*nfo.beta*nfo.beta_max- logqw_new)

    nfo.sinf_logw = nfo.posterior_logp-logqw_new #this is just used for ess (output and bw selection based on IW, which I don't use)
    #gradient of q_w on p points, need extra q because converting from log derivative, used for output
    nfo.grad_w_at_p = np.exp(t2a(nfo.nf_model_w.evaluate_density(a2t(nfo.nf_samples)))[:,np.newaxis]) * t2a(nfo.nf_model_w.score(a2t(nfo.nf_samples)))

    return 0

#Convenience functions, should replace repeated conversions with these as I do in notebook...
def t2a(tens): return tens.numpy().astype(np.float64)
def a2t(arr): return torch.from_numpy(arr.astype(np.float32))

#   Copyright 2020 The PyMC Developers
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License 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.

# jms: This file is adapted from sample_nfmc.py / sample_smc.py

import logging
import multiprocessing as mp
import time
import warnings

from collections.abc import Iterable

import numpy as np
from scipy.special import logsumexp

from pymc3.backends.base import MultiTrace
from pymc3.model import modelcontext
from pymc3.parallel_sampling import _cpu_count
from pymc3.nfo.nfo import NFO
import torch
from scipy.stats import multivariate_normal



def opt_nfo(
    #Optimization parameters
    #initialization
    n0=10, #int, n0 the initial number of draws
    init_samples=None, #array, Whether to provide some pre-defined sequence or do pymc3 sampling
    #approximation
    k_trunc=np.inf, #IW clipping, not used by default
    eps_z=0.01, #float, tolerance on Z for q iter convergence (eps') #currently not used since not iterating SINF unless trainable
    nf_iter=1, #int, number of NF iters -should always be 1 in our implementation
    #annealing
    N=10, #int, N the TOTAL number of draws we want at each iteration - this is no longer used, is from when we used to run multiple fits
    t_ess=0.5, #float, ESS<t_ess*n0 t threshold on ESS for ESS3 (no longer temperature)
    g_AF=0, #float, size of gradient contribution to AF, not used now
    #exploration
    N_AF=1000, #int,number of points to use in q_w sampling for AF
    expl_top_AF=1, #int,cut for the top AF at a given temp level accepted at each beta
    expl_latent=0, #int,latent draw from around top IW1 or around random draw from q_w, accepted at each step
    expl_top_qw=0, #int,keep top q_w at this iteration
    beta_max=1, #float>0,highest exponent on tempered posterior, support >1 for exploitation
    rel_beta=1, #0<float<1, β2 = rel_beta*β, where β2 is the lower temp level used for sampling q_w, what we call 'X'
    frac_rel_beta_AF=1, #int, the modifier to the AF used to up/down-weight the w vs uw contribution, what we call "Y"
    latent_sigma=None, #float, the value of l
    use_latent_beta2=False, #whether to get the latent sample from q_w(β2) or from q_uw
    use_pq_beta_IW1=False, #whether to get the latent sample from near top IW1 or randomly from q_w
    bounds=None,#array, size 2xd, bounding box for samples FIXME make this more obvious, needed for prior
    N_temp=25, #int, cutoff on number of allowed temp iterations before giving up -> #FIXME eventually make this throw error
    #NF parameters
    model=None,
    frac_validate=0.0,
    iteration=None,
    alpha_w=(0,0),
    alpha_uw=(0,0),
    verbose=False,
    n_component=None,
    interp_nbin=None,
    KDE=True,
    bw_factor_min=1.0,
    bw_factor_max=1.0,
    bw_factor_num=1,
    rel_bw=1,
    edge_bins=None,
    ndata_wT=None,
    MSWD_max_iter=None,
    NBfirstlayer=True,
    logit=False,
    Whiten=False,
    trainable_qw=False, #whether to improve our q_w at each beta iteration with SGD
    sgd_steps=0, #number of steps used in Adam when training trainable q_w
    knots_trainable=5,
    batchsize=None,
    nocuda=False,
    patch=False,
    shape=[28,28,1],
    #Runtime
    random_seed=-1,
    parallel=False,
    cores=None
):
    r"""
    Normalizing flow-based Bayesian Optimization.

    Parameters
    ----------
    draws: int
        The number of samples to draw from the posterior (i.e. last stage). And also the number of
        independent chains. Defaults to 2000.
    norm_tol: float
        Fractional difference in the evidence estimate between two steps. If it falls below this we
        stop iterating over the NF fits.
    optim_iter: int
        Maximum number of optimization steps to run during the initialization.
    nf_iter: int
        Number of NF fit iterations to go through after the optimization step.
    model: Model (optional if in ``with`` context)).
    frac_validate: float
        Fraction of the live points at each NS iteration that we use for validation of the NF fit.
    alpha: tuple of floats
        Regularization parameters used for the NF fit.
    verbose: boolean
        Whether you want verbose output from the NF fit.
    random_seed: int
        random seed
    parallel: bool
        Distribute computations across cores if the number of cores is larger than 1.
        Defaults to False.
    cores : int
        Number of cores available for the optimization step. Defaults to None, in which case the CPU
        count is used.

    """

    _log = logging.getLogger("pymc3")
    _log.info("Initializing normalizing flow-based optimization...")

    model = modelcontext(model)
    if model.name:
        raise NotImplementedError(
            "The NS_NFO implementation currently does not support named models. "
            "See https://github.com/pymc-devs/pymc3/pull/4365."
        )
    if cores is None:
        cores = _cpu_count()
    chains=1

    _log.info(
        f"Sampling {chains} chain{'s' if chains > 1 else ''} "
        f"Cores available for optimization: {cores}"
    )
    if random_seed == -1:
        random_seed = None
    if chains == 1 and isinstance(random_seed, int):
        random_seed = [random_seed]
    if random_seed is None or isinstance(random_seed, int):
        if random_seed is not None:
            np.random.seed(random_seed)
        random_seed = [np.random.randint(2 ** 30) for _ in range(chains)]
    if not isinstance(random_seed, Iterable):
        raise TypeError("Invalid value for `random_seed`. Must be tuple, list or int")

    #we changed the name for end-user-facing readability, but internally more familiar with these names
    aN,bN,cN,dN = N_AF,expl_top_AF,expl_latent,expl_top_qw

    params = (
        n0,
        init_samples,
        k_trunc,
        eps_z,
        nf_iter,
        N,
        t_ess,
        g_AF,
        aN,
        bN,
        cN,
        dN,
        beta_max,
        rel_beta,
        frac_rel_beta_AF,
        latent_sigma,
        use_latent_beta2,
        use_pq_beta_IW1,
        bounds,
        N_temp,
        model,
        frac_validate,
        iteration,
        alpha_w,
        alpha_uw,
        cores,
        verbose,
        n_component,
        interp_nbin,
        KDE,
        bw_factor_min,
        bw_factor_max,
        bw_factor_num,
        rel_bw,
        edge_bins,
        ndata_wT,
        MSWD_max_iter,
        NBfirstlayer,
        logit,
        Whiten,
        trainable_qw,
        sgd_steps,
        knots_trainable,
        batchsize,
        nocuda,
        patch,
        shape,
        parallel,
    )

    t1 = time.time()

    results = []
    for i in range(chains):
        results.append(opt_nfo_int(*params, random_seed[i], i, _log))
    (
        traces,
        log_evidence,
        q_samples,
        importance_weights,
        logp,
        logq,
        train_logp,
        train_logq,
        logZ,
        q_models,
        q_ess,
        total_ess,
        min_var_bws,
        min_pq_bws,
        betas,
    ) = zip(*results)
    trace = MultiTrace(traces)
    trace.report.log_evidence = log_evidence
    trace.report.q_samples = q_samples
    trace.report.importance_weights = importance_weights
    trace.report.logp = logp
    trace.report.logq = logq
    trace.report.train_logp = train_logp
    trace.report.train_logq = train_logq
    trace.report.logZ = logZ
    trace.report.q_models = q_models
    trace.report.q_ess = q_ess
    trace.report.total_ess = total_ess
    trace.report.N = N
    trace.report.min_var_bws = min_var_bws
    trace.report.min_pq_bws = min_pq_bws
    trace.report._t_sampling = time.time() - t1
    trace.report.betas = betas
    return trace


def opt_nfo_int(
    n0,
    init_samples,
    k_trunc,
    eps_z,
    nf_iter,
    N,
    t_ess,
    g_AF,
    aN,
    bN,
    cN,
    dN,
    beta_max,
    rel_beta,
    frac_rel_beta_AF,
    latent_sigma,
    use_latent_beta2,
    use_pq_beta_IW1,
    bounds,
    N_temp,
    model,
    frac_validate,
    iteration,
    alpha_w,
    alpha_uw,
    cores,
    verbose,
    n_component,
    interp_nbin,
    KDE,
    bw_factor_min,
    bw_factor_max,
    bw_factor_num,
    rel_bw,
    edge_bins,
    ndata_wT,
    MSWD_max_iter,
    NBfirstlayer,
    logit,
    Whiten,
    trainable_qw,
    sgd_steps,
    knots_trainable,
    batchsize,
    nocuda,
    patch,
    shape,
    parallel,
    random_seed,
    chain,
    _log,
):
    """Run one NS_NFO instance."""
    nfo = NFO(
        n0=n0,
        init_samples=init_samples,
        k_trunc=k_trunc,
        eps_z=eps_z,
        nf_iter=nf_iter,
        N=N,
        t_ess=t_ess,
        beta_max=beta_max,
        model=model,
        frac_validate=frac_validate,
        iteration=iteration,
        alpha_w=alpha_w,
        alpha_uw=alpha_uw,
        verbose=verbose,
        n_component=n_component,
        interp_nbin=interp_nbin,
        KDE=KDE,
        bw_factor_min=bw_factor_min,
        bw_factor_max=bw_factor_max,
        bw_factor_num=bw_factor_num,
        rel_bw=rel_bw,
        edge_bins=edge_bins,
        ndata_wT=ndata_wT,
        MSWD_max_iter=MSWD_max_iter,
        NBfirstlayer=NBfirstlayer,
        logit=logit,
        Whiten=Whiten,
        trainable_qw=trainable_qw,
        sgd_steps=sgd_steps,
        knots_trainable=knots_trainable,
        batchsize=batchsize,
        nocuda=nocuda,
        bounds=bounds,
        patch=patch,
        shape=shape,
        chain=chain,
        random_seed=random_seed,
    )

    #dicts for storage #FIXME compactify
    iter_sample_dict = {}
    iter_weight_dict = {}
    iter_logp_dict = {}
    iter_logq_dict = {}
    iter_train_logp_dict = {}
    iter_train_logq_dict = {}
    iter_logZ_dict = {}
    iter_qmodel_dict = {}
    iter_q_ess_dict = {}
    iter_total_ess_dict = {}
    iter_min_var_bw_dict = {}
    iter_min_pq_bw_dict = {}

    '''///////////////////////////////'''
    '''/         INITIALIZE          /'''
    '''///////////////////////////////'''
    stage = 0
    print("stage: ", stage)
    betas = []
    nfo.initialize_var_info()
    nfo.setup_logp()
    nfo.initialize_population() #uses n0
    theta_0 = nfo.samples
    nfo.theta = theta_0
    nfo.latent_sigma=latent_sigma #for dict stuff
    betas.append(0)
    #need these in the first iteration to get weights for the fit in fit_nf
    #later replace prior with log_q in that function as we go
    nfo.sinf_logw = nfo.posterior_logp-nfo.prior_logp
    #this doesn't actually do anything except in prior
    nfo.importance_weights_pq = np.exp(nfo.posterior_logp-nfo.prior_logp-nfo.log_evidence_pq)

    save_dicts(nfo,stage,_log,iter_sample_dict,iter_weight_dict,iter_logp_dict,iter_logZ_dict,
    iter_qmodel_dict,iter_q_ess_dict,iter_total_ess_dict,iter_logq_dict,iter_train_logp_dict,iter_train_logq_dict,
    iter_min_var_bw_dict,iter_min_pq_bw_dict,betas=None)

    '''///////////////////////////////'''
    '''/           ANNEAL            /'''
    '''///////////////////////////////'''
    for n in range(N_temp):
        stage += 1
        print("stage: ", stage)
        #Fit unweighted approximator to samples
        nfo.fit_nf(weighted=False) #unweighted fit first to get IW
        quw = nfo.nf_model_uw
        #Select a temperature level beta based on ESS3
        nfo.logIW3 = nfo.posterior_logp - nfo.logq_uw
        nfo.update_weights_beta()

        #check if we want to sample from lower beta, gross
        if(stage>1):
            nearest_beta2_idx = (np.abs(np.array(betas)[:-1] - rel_beta*nfo.beta)).argmin()
        else:
            nearest_beta2_idx = 0

        '''///////////////////////////////'''
        '''/         APPROXIMATE         /'''
        '''///////////////////////////////'''
        nfo.fit_nf(weighted=True) #weighted fit, this knows about beta, also knows about trainable SINF
        qw = nfo.nf_model_w

        '''///////////////////////////////'''
        '''/           EXPLORE           /'''
        '''///////////////////////////////'''

        #sample from lower beta only if we are at high enough beta such that
        #rel_beta*beta >= betas[1] - usually this only excludes 1st step or so
        if(nearest_beta2_idx != 0 and stage>1 and rel_beta <1): #we can't use lower beta on first beta
            print("using beta2={0:.2f} at idx {1}, at stage-1 = {2}".format(betas[nearest_beta2_idx],nearest_beta2_idx,stage-1))
            qw_beta2 = iter_qmodel_dict['q{0}_w'.format(nearest_beta2_idx)]  #get the qw at the lower beta we need
            #FIXME lazy copy, make this a function, dropping gradient since sort of defunct now
            theta_w_beta2,_ = qw_beta2.sample(aN,device='cpu') #use aN for now
            nfo.beta_2 = betas[nearest_beta2_idx]
            nfo.tens_theta_w_beta2 = theta_w_beta2
            #evaluate fitted qs at this beta on samples from beta_2
            logq_uw_beta12 = t2a(quw.evaluate_density(theta_w_beta2))
            logq_w_beta12 = t2a(qw.evaluate_density(theta_w_beta2))

            theta_w_beta2 = t2a(theta_w_beta2) #has to happen after grad, grad wants tensor
            if(bounds is not None): #make sure we don't keep samples out of prior range, ow can get maxima out of bounds
                prior_mask_beta2 = np.all(theta_w_beta2>bounds[0],axis=-1) & np.all(theta_w_beta2<bounds[1],axis=-1)
                theta_w_beta2,logq_w_beta12,logq_uw_beta12 = theta_w_beta2[prior_mask_beta2],logq_w_beta12[prior_mask_beta2],logq_uw_beta12[prior_mask_beta2]
            logIW2_beta12 = np.squeeze(logq_w_beta12-logq_uw_beta12)
            logIW2 = logIW2_beta12
            theta_w = theta_w_beta2
            #define new acquisition function
            logIW2prime = logIW2 + (-1+rel_beta*frac_rel_beta_AF)*logq_w_beta12 #use downweighted version
            logIW2prime -= logsumexp(logIW2prime) #, normalize them
            nfo.AF = np.exp(logIW2prime)
            nfo.logIW2 = logIW2prime #just for output later
            #everything else outside this if statement is the same...

            #Draw a sample from q_uw via latent space near highest IW1
            if(latent_sigma is not None):
                max_latent=15#fixed for now, indicates a problem
                tries=0
                for i in range(max_latent):
                    tries+=1
                    if(use_pq_beta_IW1):
                        ranked_IW1 = np.argsort(nfo.importance_weights_pq_beta)[::-1]
                        top_IW1_sample = a2t(nfo.samples[ranked_IW1[:cN]]) #top pbeta/qbeta
                        # print("top IW sample ", top_IW1_sample)
                    else:
                        # ranked_IW1 = np.argsort(nfo.importance_weights_pq)[::-1]
                        #WARNING this is not actually IW1, just doing quick test
                        #replace IW1 with random draw from q_w at current beta level
                        qw_sample_rand,_ = qw.sample(cN,device='cpu')
                        top_IW1_sample = qw_sample_rand
                        # print("rand qw expl sample ", top_IW1_sample)

                    if(use_latent_beta2):
                        latent_z = np.squeeze(t2a(qw_beta2.inverse(top_IW1_sample)[0])) #go to latent space
                        latent_sigma = latent_sigma*np.ones_like(latent_z) #sphere width to sample from in latent space
                        z_prime = multivariate_normal.rvs(mean=latent_z,cov=latent_sigma,size=1)#,random_state=random_seed) #sample
                        theta_z = t2a(qw_beta2.forward(a2t(z_prime))[0]) #go back to data space
                    else:
                        latent_z = np.squeeze(t2a(quw.inverse(top_IW1_sample)[0])) #go to latent space
                        latent_sigma = latent_sigma*np.ones_like(latent_z) #sphere width to sample from in latent space
                        z_prime = multivariate_normal.rvs(mean=latent_z,cov=latent_sigma,size=1)#,random_state=random_seed) #sample
                        theta_z = t2a(quw.forward(a2t(z_prime))[0]) #go back to data space

                    # print("latent_z", latent_z)
                    # print("latent_sigma", latent_sigma)

                    z_prime = multivariate_normal.rvs(mean=latent_z,cov=latent_sigma,size=1)#,random_state=random_seed) #sample
                    theta_z = t2a(quw.forward(a2t(z_prime))[0]) #go back to data space
                    # print('z_prime: ', z_prime)

                    #take extra theta for iter from q_w(beta2) at highest q_w(beta)
                    if(dN>0):
                        idx_top_qw = np.argsort(np.squeeze(logq_w_beta12))[::-1]
                        top_qw = theta_w_beta2[idx_top_qw[:dN]]
                        theta_z = np.concatenate([theta_z,top_qw])
                    # print('theta_z: ', theta_z)

                    bound_check = np.all(theta_z>bounds[0],axis=-1) & np.all(theta_z<bounds[1],axis=-1)
                    if(np.all(bound_check)):
                        break #be done
                    #else we just keep going until max
                    if(tries==max_latent-1): raise ValueError("Too many OOB samples! Try decreasing latent_sigma")
                nfo.top_IW1_sample=t2a(top_IW1_sample)#for saving to dict later
                print("finished with latent draw, took {0} tries".format(i))

        else: #no sampling from lower beta levels
            theta_w,logq_w =qw.sample(aN,device='cpu')
            dlogq_w = qw.score(theta_w).numpy().astype(np.float64) #gradient dlogqw/dtheta

            nfo.tens_theta_w = theta_w
            logq_uw = t2a(quw.evaluate_density(theta_w))
            theta_w,logq_w = t2a(theta_w),t2a(logq_w) #has to happen after grad, grad wants tensor
            if(bounds is not None): #make sure we don't keep samples out of prior range, ow can get maxima out of bounds
                prior_mask = np.all(theta_w>bounds[0],axis=-1) & np.all(theta_w<bounds[1],axis=-1)
                theta_w,logq_w,logq_uw = theta_w[prior_mask],logq_w[prior_mask],logq_uw[prior_mask]
                dlogq_w = dlogq_w[prior_mask]
            logIW2 = np.squeeze(logq_w-logq_uw)#many (aN) values
            nfo.w_samples_beta = theta_w
            nfo.beta_2 = None
            #compute (scalar) acquisition function
            logIW2 -= logsumexp(logIW2) #, normalize them
            nfo.logIW2 = logIW2
            dim_sum_dlog_w = np.sum(dlogq_w,axis=-1) #trailing axis is size dims
            g_term = -g_AF*( (np.exp(np.squeeze(dim_sum_dlog_w) -logsumexp(np.squeeze(dim_sum_dlog_w))))**2 ) #normalize gradient term so g is not crazy
            nfo.AF = np.exp(logIW2) + g_term


        '''///////////////////////////////'''
        '''/           EXPLOIT           /'''
        '''///////////////////////////////'''

        nfo.w_samples = theta_w
        # ESS2 = len(logIW2)*nfo.calculate_ess(logIW2) not actually used?
        ranked_AF = np.argsort(nfo.AF)[::-1]
        nfo.ranked_topAF=ranked_AF[:bN]
        #rank order the AF results
        theta_new = theta_w[nfo.ranked_topAF] #check setting to see how many new points at each iteration

        #FIXME: this is very gross
        if(latent_sigma is not None and stage>1 and rel_beta <1 and nearest_beta2_idx != 0 ):
            #add new sample from latent space if want, this should get updated in the update call for p
            nfo.theta_no_z = theta_new
            theta_new = np.concatenate([theta_new,theta_z])
            #for saving to dicts later
            nfo.theta_z = theta_z

        new_sample_update(nfo,theta_new)
        #finish up
        save_dicts(nfo,stage,_log,iter_sample_dict,iter_weight_dict,iter_logp_dict,iter_logZ_dict,
        iter_qmodel_dict,iter_q_ess_dict,iter_total_ess_dict,iter_logq_dict,iter_train_logp_dict,iter_train_logq_dict,
        iter_min_var_bw_dict,iter_min_pq_bw_dict,betas=betas)
        #check if done
        if nfo.beta ==1: break
        if n==N_temp-1: print("WARNING: Failed to converge within N_temp temperature steps.")
    print("Finished")
    nfo.posterior = nfo.samples

    return (
        nfo.posterior_to_trace(),
        nfo.log_evidence,
        iter_sample_dict,
        iter_weight_dict,
        iter_logp_dict,
        iter_logq_dict,
        iter_train_logp_dict,
        iter_train_logq_dict,
        iter_logZ_dict,
        iter_qmodel_dict,
        iter_q_ess_dict,
        iter_total_ess_dict,
        iter_min_var_bw_dict,
        iter_min_pq_bw_dict,
        betas,
    )

#fix this, got grosser than I thought it would be
def save_dicts(nfo,stage,_log,
               iter_sample_dict,iter_weight_dict,
               iter_logp_dict,iter_logZ_dict,
               iter_qmodel_dict,iter_q_ess_dict,
               iter_total_ess_dict,iter_logq_dict,
               iter_train_logp_dict,iter_train_logq_dict,
               iter_min_var_bw_dict,iter_min_pq_bw_dict,betas=None):
    'Convenience function b/c dicts clog up the main loop'
    #iter
    nfo.nf_samples_to_trace()
    iter_sample_dict[f'q{int(stage)}'] = nfo.nf_trace
    iter_weight_dict[f'q{int(stage)}'] = nfo.weights
    iter_weight_dict[f'q{int(stage)}_pq'] = nfo.weights_pq
    iter_logp_dict[f'q{int(stage)}'] = nfo.posterior_logp
    iter_logZ_dict[f'q{int(stage)}'] = nfo.log_evidence
    iter_logZ_dict[f'q{int(stage)}_pq'] = nfo.log_evidence_pq
    iter_weight_dict[f'q{int(stage)}_pq'] = nfo.weights_pq

    # iter_qmodel_dict[f'q{int(stage)}'] = nfo.nf_model
    iter_q_ess_dict[f'q{int(stage)}'] = nfo.q_ess
    iter_total_ess_dict[f'q{int(stage)}'] = nfo.total_ess
    iter_total_ess_dict[f'q{int(stage)}_log_mean_loss'] = nfo.log_mean_loss
    print(f"logZ: {nfo.log_evidence:.3f}, ESS/N: {nfo.q_ess:.3f}, logZ_pq: {nfo.log_evidence_pq:.3f} log mean loss: {nfo.log_mean_loss:.3f}")
    if _log is not None:
        _log.info(f"Stage: {stage:3d}, logZ Estimate: {nfo.log_evidence:.3f},logZ_pq Estimate: {nfo.log_evidence_pq:.3f}")
        _log.info(f"Stage: {stage:3d}, q ESS/N: {nfo.q_ess:.3f}")

    if(stage>0): #these aren't defined in the initial
        if(betas is not None): betas.append(nfo.beta)
        iter_qmodel_dict[f'q{int(stage)}_w'] = nfo.nf_model_w
        iter_qmodel_dict[f'q{int(stage)}_uw'] = nfo.nf_model_uw
        iter_logq_dict[f'q{int(stage)}'] = nfo.logq
        iter_logq_dict[f'q{int(stage)}_w'] = nfo.logq_w
        iter_logq_dict[f'q{int(stage)}_uw'] = nfo.logq_uw
        iter_train_logp_dict[f'q{stage}'] = nfo.train_logp
        iter_train_logq_dict[f'q{stage}'] = nfo.train_logq
        iter_min_var_bw_dict[f'q{int(stage)}'] = nfo.min_var_bw
        iter_min_pq_bw_dict[f'q{int(stage)}'] = nfo.min_pq_bw
        iter_sample_dict[f'q{int(stage)}_w_samples'] = nfo.w_samples
        iter_weight_dict[f'q{int(stage)}_IW2'] = nfo.logIW2
        iter_weight_dict[f'q{int(stage)}_IW3'] = nfo.logIW3
        iter_weight_dict[f'q{int(stage)}_AF'] = nfo.AF
        iter_weight_dict[f'q{int(stage)}_IW2_p'] = nfo.logIW2_p
        iter_weight_dict[f'q{int(stage)}_AF_p'] = nfo.AF_p
        iter_logZ_dict[f'q{int(stage)}_pq_uw'] = nfo.log_evidence_pq_uw
        iter_logZ_dict[f'q{int(stage)}_pq_w'] = nfo.log_evidence_pq_w
        iter_weight_dict[f'q{int(stage)}_pq_uw'] = nfo.log_weight_pq_uw
        iter_weight_dict[f'q{int(stage)}_pq_w'] = nfo.log_weight_pq_w
        #gradient of q_w on samples where posterior is evaluated (regular samples)
        iter_sample_dict[f'grad_q{int(stage)}'] = nfo.grad_w_at_p
        if(nfo.latent_sigma is not None and stage>1 and nfo.beta_2 is not None):
            iter_sample_dict[f'topIW1_q{int(stage)}'] = nfo.top_IW1_sample
            iter_sample_dict[f'non_latent_q{int(stage)}'] = nfo.theta_no_z
            iter_sample_dict[f'latent_q{int(stage)}'] = nfo.theta_z

        if(nfo.beta_2 is not None):
            iter_sample_dict[f'beta2_{int(stage)}'] = nfo.beta_2

        if _log is not None:
            _log.info(f"Stage: {stage:3d}, Min variance BW factor: {nfo.min_var_bw}, Var(IW): {nfo.min_var_weights}, Min Zpq BW factor: {nfo.min_pq_bw}")
            if(betas is not None): _log.info(f"Stage: {stage:3d} Iter Beta: {nfo.beta:.3f}, True Beta: {nfo.beta*nfo.beta_max:.3f}")

def new_sample_update(nfo,theta_new):
    #update current samples and logp + other values
    tens_theta_new = torch.from_numpy(np.atleast_1d(theta_new).astype(np.float32))
    posterior_logp_new = np.array([nfo.posterior_logp_func(sample) for sample in theta_new]).squeeze()
    prior_logp_new = np.array([nfo.prior_logp_func(sample) for sample in theta_new]).squeeze()
    nfo.theta = np.concatenate([nfo.theta,theta_new])
    nfo.posterior_logp = np.concatenate([nfo.posterior_logp,np.atleast_1d(posterior_logp_new)])
    nfo.prior_logp = np.concatenate([nfo.prior_logp,np.atleast_1d(prior_logp_new)])
    #some of this is probably redundant, but need to update these for beta iteration
    #IW3 computed using tempered posterior - this is just for output, not actually used in fits.
    nfo.logIW3 = nfo.beta_max*nfo.beta*(nfo.posterior_logp-nfo.prior_logp) + nfo.prior_logp -np.concatenate([nfo.logq_uw,
                                            np.atleast_1d(nfo.nf_model_uw.evaluate_density(tens_theta_new).numpy().astype(np.float64))]) #this should be quw(theta)
    nfo.logIW2_p = nfo.logIW2[nfo.ranked_topAF] #keep the top one
    nfo.AF_p = nfo.AF[nfo.ranked_topAF]
    nfo.samples = nfo.theta
    nfo.nf_samples = np.copy(nfo.samples)
    #norm doesn't matter for IW1
    logqw_new = t2a(nfo.nf_model_w.evaluate_density(a2t(nfo.nf_samples)))
    nfo.importance_weights_pq = np.exp(nfo.posterior_logp- logqw_new)
    nfo.importance_weights_pq_beta = np.exp(nfo.posterior_logp*nfo.beta*nfo.beta_max- logqw_new)

    nfo.sinf_logw = nfo.posterior_logp-logqw_new #this is just used for ess (output and bw selection based on IW, which I don't use)
    #gradient of q_w on p points, need extra q because converting from log derivative, used for output
    nfo.grad_w_at_p = np.exp(t2a(nfo.nf_model_w.evaluate_density(a2t(nfo.nf_samples)))[:,np.newaxis]) * t2a(nfo.nf_model_w.score(a2t(nfo.nf_samples)))

    return 0

#Convenience functions, should replace repeated conversions with these as I do in notebook...
def t2a(tens): return tens.numpy().astype(np.float64)
def a2t(arr): return torch.from_numpy(arr.astype(np.float32))

#   Copyright 2020 The PyMC Developers # #   Licensed under the Apache License, Version 2.0 (the "License"); #   you may not use this file except in compliance with the License. #   You may obtain a copy of the License at # #       http://www.apache.org/licenses/LICENSE-2.0 # #   Unless required by applicable law or agreed to in writing, software #   distributed under the License 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. # jms: This file is adapted from sample_nfmc.py / sample_smc.py #Optimization parameters #initialization #int, n0 the initial number of draws #array, Whether to provide some pre-defined sequence or do pymc3 sampling #approximation #IW clipping, not used by default #float, tolerance on Z for q iter convergence (eps') #currently not used since not iterating SINF unless trainable #int, number of NF iters -should always be 1 in our implementation #annealing #int, N the TOTAL number of draws we want at each iteration - this is no longer used, is from when we used to run multiple fits #float, ESS<t_ess*n0 t threshold on ESS for ESS3 (no longer temperature) #float, size of gradient contribution to AF, not used now #exploration #int,number of points to use in q_w sampling for AF #int,cut for the top AF at a given temp level accepted at each beta #int,latent draw from around top IW1 or around random draw from q_w, accepted at each step #int,keep top q_w at this iteration #float>0,highest exponent on tempered posterior, support >1 for exploitation #0<float<1, β2 = rel_beta*β, where β2 is the lower temp level used for sampling q_w, what we call 'X' #int, the modifier to the AF used to up/down-weight the w vs uw contribution, what we call "Y" #float, the value of l #whether to get the latent sample from q_w(β2) or from q_uw #whether to get the latent sample from near top IW1 or randomly from q_w #array, size 2xd, bounding box for samples FIXME make this more obvious, needed for prior #int, cutoff on number of allowed temp iterations before giving up -> #FIXME eventually make this throw error #NF parameters #whether to improve our q_w at each beta iteration with SGD #number of steps used in Adam when training trainable q_w #Runtime Normalizing flow-based Bayesian Optimization.      Parameters     ----------     draws: int         The number of samples to draw from the posterior (i.e. last stage). And also the number of         independent chains. Defaults to 2000.     norm_tol: float         Fractional difference in the evidence estimate between two steps. If it falls below this we         stop iterating over the NF fits.     optim_iter: int         Maximum number of optimization steps to run during the initialization.     nf_iter: int         Number of NF fit iterations to go through after the optimization step.     model: Model (optional if in ``with`` context)).     frac_validate: float         Fraction of the live points at each NS iteration that we use for validation of the NF fit.     alpha: tuple of floats         Regularization parameters used for the NF fit.     verbose: boolean         Whether you want verbose output from the NF fit.     random_seed: int         random seed     parallel: bool         Distribute computations across cores if the number of cores is larger than 1.         Defaults to False.     cores : int         Number of cores available for the optimization step. Defaults to None, in which case the CPU         count is used. #we changed the name for end-user-facing readability, but internally more familiar with these names Run one NS_NFO instance. #dicts for storage #FIXME compactify /////////////////////////////// /         INITIALIZE          / /////////////////////////////// #uses n0 #for dict stuff #need these in the first iteration to get weights for the fit in fit_nf #later replace prior with log_q in that function as we go #this doesn't actually do anything except in prior /////////////////////////////// /           ANNEAL            / /////////////////////////////// #Fit unweighted approximator to samples #unweighted fit first to get IW #Select a temperature level beta based on ESS3 #check if we want to sample from lower beta, gross /////////////////////////////// /         APPROXIMATE         / /////////////////////////////// #weighted fit, this knows about beta, also knows about trainable SINF /////////////////////////////// /           EXPLORE           / /////////////////////////////// #sample from lower beta only if we are at high enough beta such that #rel_beta*beta >= betas[1] - usually this only excludes 1st step or so #we can't use lower beta on first beta #get the qw at the lower beta we need #FIXME lazy copy, make this a function, dropping gradient since sort of defunct now #use aN for now #evaluate fitted qs at this beta on samples from beta_2 #has to happen after grad, grad wants tensor #make sure we don't keep samples out of prior range, ow can get maxima out of bounds #define new acquisition function #use downweighted version #, normalize them #just for output later #everything else outside this if statement is the same... #Draw a sample from q_uw via latent space near highest IW1 #fixed for now, indicates a problem #top pbeta/qbeta # print("top IW sample ", top_IW1_sample) # ranked_IW1 = np.argsort(nfo.importance_weights_pq)[::-1] #WARNING this is not actually IW1, just doing quick test #replace IW1 with random draw from q_w at current beta level # print("rand qw expl sample ", top_IW1_sample) #go to latent space #sphere width to sample from in latent space #,random_state=random_seed) #sample #go back to data space #go to latent space #sphere width to sample from in latent space #,random_state=random_seed) #sample #go back to data space # print("latent_z", latent_z) # print("latent_sigma", latent_sigma) #,random_state=random_seed) #sample #go back to data space # print('z_prime: ', z_prime) #take extra theta for iter from q_w(beta2) at highest q_w(beta) # print('theta_z: ', theta_z) #be done #else we just keep going until max #for saving to dict later #no sampling from lower beta levels #gradient dlogqw/dtheta #has to happen after grad, grad wants tensor #make sure we don't keep samples out of prior range, ow can get maxima out of bounds #many (aN) values #compute (scalar) acquisition function #, normalize them #trailing axis is size dims #normalize gradient term so g is not crazy /////////////////////////////// /           EXPLOIT           / /////////////////////////////// # ESS2 = len(logIW2)*nfo.calculate_ess(logIW2) not actually used? #rank order the AF results #check setting to see how many new points at each iteration #FIXME: this is very gross #add new sample from latent space if want, this should get updated in the update call for p #for saving to dicts later #finish up #check if done #fix this, got grosser than I thought it would be #iter # iter_qmodel_dict[f'q{int(stage)}'] = nfo.nf_model #these aren't defined in the initial #gradient of q_w on samples where posterior is evaluated (regular samples) #update current samples and logp + other values #some of this is probably redundant, but need to update these for beta iteration #IW3 computed using tempered posterior - this is just for output, not actually used in fits. #this should be quw(theta) #keep the top one #norm doesn't matter for IW1 #this is just used for ess (output and bw selection based on IW, which I don't use) #gradient of q_w on p points, need extra q because converting from log derivative, used for output #Convenience functions, should replace repeated conversions with these as I do in notebook...