from __future__ import division, absolute_import, print_function
from .. import affinitymat
from .. import nearest_neighbors
from .. import cluster
from .. import aggregator
from .. import core
from .. import util
from .. import seqlet_embedding
from .. import pattern_filterer as pattern_filterer_module
from joblib import Parallel, delayed
from collections import defaultdict, OrderedDict, Counter
import numpy as np
import time
import sys
import gc
import json
from ..util import print_memory_use

def get_seqlet_neighbors_with_initcluster(
        nearest_neighbors_to_compute,
        coarse_affmat, initclusters):
    if (initclusters is not None):
        assert len(initclusters)==len(coarse_affmat)
    #get the argsort for coarse_affmat
    coarse_affmat_argsort = np.argsort(-coarse_affmat, axis=-1)
    nearest_neighbors = []
    for row_idx,argsort_row in enumerate(coarse_affmat_argsort):
        combined_neighbor_row = []
        neighbor_row_topnn = argsort_row[:nearest_neighbors_to_compute+1]
        neighbor_set_topnn = set(neighbor_row_topnn)
        #combined_neighbor_row ends up being the union of the standard nearest
        # neighbors plus the nearest neighbors if focusing on the initclusters
        combined_neighbor_row.extend(neighbor_row_topnn)
        if (initclusters is not None):
            combined_neighbor_row.extend([
                y for y in ([x for x in argsort_row
                             if initclusters[x]==initclusters[row_idx]][
                             :nearest_neighbors_to_compute+1])
                if y not in neighbor_set_topnn])
        nearest_neighbors.append(combined_neighbor_row)
    return nearest_neighbors

def fish_out_kwargs(orig_kwargs, to_fish_out):
    fished_out = {}
    for kwarg_name in to_fish_out:
        if kwarg_name in orig_kwargs:
            fished_out[kwarg_name] = orig_kwargs[kwarg_name]
            del orig_kwargs[kwarg_name]
    return fished_out

#adds backwargs compatibility re gapped kmer arguments
def legacy_tfmodiscoseqletstopatternsfactory(current_constructor):

    def new_constructor(*args, **kwargs):
        gapped_kmer_kwargs = fish_out_kwargs(
            orig_kwargs=kwargs,
            to_fish_out=['kmer_len', 'num_gaps',
                         'num_mismatches', 'gpu_batch_size'])
        if (len(gapped_kmer_kwargs) > 0):
            assert 'embedder_factory' not in kwargs,\
                ("Cannot both specify embedder_factory and "
                 +str(gapped_kmer_kwargs))
            from modisco.seqlet_embedding import gapped_kmer
            kwargs['embedder_factory'] = (
                  seqlet_embedding.gapped_kmer
                  .GappedKmerEmbedderFactory(**gapped_kmer_kwargs))
        return current_constructor(*args, **kwargs)
    return new_constructor

##legacy
#alphabet_size=None,
#kmer_len=None, num_gaps=3, num_mismatches=2,
#gpu_batch_size=20,
class TfModiscoSeqletsToPatternsFactory(object):

    @legacy_tfmodiscoseqletstopatternsfactory
    def __init__(self, n_cores=4,
                       min_overlap_while_sliding=0.7,

                       #init clusterer factory
                       initclusterer_factory=None,

                       embedder_factory=(
                        seqlet_embedding.advanced_gapped_kmer
                                        .AdvancedGappedKmerEmbedderFactory()),

                       nearest_neighbors_to_compute=500,
                       use_pynnd=False,

                       affmat_correlation_threshold=0.15,
                       filter_beyond_first_round=False,
                       skip_fine_grained=False,

                       tsne_perplexity=10,
                       use_louvain=False,
                       louvain_initclusters_weight=1.0,
                       n_leiden_iterations_r1=-1,
                       n_leiden_iterations_r2=-1,
                       louvain_num_runs_and_levels_r1=[(200,-1)],
                       louvain_num_runs_and_levels_r2=[(200,-1)],
                       contin_runs_r1=50,
                       contin_runs_r2=50,
                       final_louvain_level_to_return=1,

                       frac_support_to_trim_to=0.2,
                       min_num_to_trim_to=30,
                       trim_to_window_size=30,
                       initial_flank_to_add=10,

                       prob_and_pertrack_sim_merge_thresholds=[
                       (0.8,0.8), (0.5, 0.85), (0.2, 0.9)],

                       prob_and_pertrack_sim_dealbreaker_thresholds=[
                        (0.4, 0.75), (0.2,0.8), (0.1, 0.85), (0.0,0.9)],

                       subcluster_perplexity=50,
                       merging_max_seqlets_subsample=300,
                       #threshold_for_spurious_merge_detection=0.8,

                       #min_similarity_for_seqlet_assignment=0.2,
                       final_min_cluster_size=30,
                       min_ic_in_window=0.6,#total IC in some windowsize window
                       min_ic_windowsize=6,
                       ppm_pseudocount=0.001,

                       final_flank_to_add=0,

                       verbose=True, seed=1234):

        self.initclusterer_factory = initclusterer_factory
        if (use_louvain==True):
            assert self.initclusterer_factory==None,\
                    ("Louvain doesn't support cluster initialization;"
                     +" set use_louvain to False")

        #affinity_mat calculation
        self.n_cores = n_cores
        self.min_overlap_while_sliding = min_overlap_while_sliding

        self.embedder_factory = embedder_factory

        self.nearest_neighbors_to_compute = nearest_neighbors_to_compute
        self.use_pynnd = use_pynnd

        self.affmat_correlation_threshold = affmat_correlation_threshold
        self.filter_beyond_first_round = filter_beyond_first_round
        self.skip_fine_grained = skip_fine_grained

        #affinity mat to tsne dist mat setting
        self.tsne_perplexity = tsne_perplexity

        #clustering settings
        self.use_louvain = use_louvain
        self.louvain_initclusters_weight = louvain_initclusters_weight
        self.n_leiden_iterations_r1 = n_leiden_iterations_r1
        self.n_leiden_iterations_r2 = n_leiden_iterations_r2
        self.contin_runs_r1 = contin_runs_r1
        self.contin_runs_r2 = contin_runs_r2
        self.louvain_num_runs_and_levels_r1 = louvain_num_runs_and_levels_r1
        self.louvain_num_runs_and_levels_r2 = louvain_num_runs_and_levels_r2
        self.final_louvain_level_to_return = final_louvain_level_to_return

        #postprocessor1 settings
        self.frac_support_to_trim_to = frac_support_to_trim_to
        self.min_num_to_trim_to = min_num_to_trim_to
        self.trim_to_window_size = trim_to_window_size
        self.initial_flank_to_add = initial_flank_to_add

        #subclustering
        self.subcluster_perplexity=subcluster_perplexity

        #merging similar patterns
        self.prob_and_pertrack_sim_merge_thresholds =\
            prob_and_pertrack_sim_merge_thresholds
        self.prob_and_pertrack_sim_dealbreaker_thresholds =\
            prob_and_pertrack_sim_dealbreaker_thresholds

        self.merging_max_seqlets_subsample = merging_max_seqlets_subsample
        #self.threshold_for_spurious_merge_detection =\
        #    threshold_for_spurious_merge_detection

        #reassignment settings
        #self.min_similarity_for_seqlet_assignment =\
        #    min_similarity_for_seqlet_assignment
        self.final_min_cluster_size = final_min_cluster_size
        self.min_ic_in_window = min_ic_in_window
        self.min_ic_windowsize = min_ic_windowsize
        self.ppm_pseudocount = ppm_pseudocount

        #final postprocessor settings
        self.final_flank_to_add=final_flank_to_add

        #other settings
        self.verbose = verbose
        self.seed = seed

    def get_jsonable_config(self):
        to_return =  OrderedDict([
                ('class_name', type(self).__name__),
                ('n_cores', self.n_cores),
                ('initclusterer_factory',
                 self.initclusterer_factory.get_jsonable_config()),
                ('min_overlap_while_sliding', self.min_overlap_while_sliding),
                ('embedder_factory',
                 self.embedder_factory.get_jsonable_config()),
                ('num_mismatches', self.num_mismatches),
                ('nearest_neighbors_to_compute',
                 self.nearest_neighbors_to_compute),
                ('affmat_correlation_threshold',
                 self.affmat_correlation_threshold),
                ('filter_beyond_first_round', filter_beyond_first_round),
                ('tsne_perplexity', self.tsne_perplexity),
                ('use_louvain', self.use_louvain),
                ('louvain_num_runs_and_levels_r1',
                 self.louvain_num_runs_and_levels_r1),
                ('louvain_num_runs_and_levels_r2',
                 self.louvain_num_runs_and_levels_r2),
                ('final_louvain_level_to_return',
                 self.final_louvain_level_to_return),
                ('contin_runs_r1',
                 self.contin_runs_r1),
                ('contin_runs_r2',
                 self.contin_runs_r2),
                ('frac_support_to_trim_to', self.frac_support_to_trim_to),
                ('min_num_to_trim_to', self.min_num_to_trim_to),
                ('trim_to_window_size', self.trim_to_window_size),
                ('initial_flank_to_add', self.initial_flank_to_add),
                ('subcluster_perplexity', self.subcluster_perplexity),
                ('prob_and_pertrack_sim_merge_thresholds',
                 self.prob_and_pertrack_sim_merge_thresholds),
                ('prob_and_pertrack_sim_dealbreaker_thresholds',
                 self.prob_and_pertrack_sim_dealbreaker_thresholds),
                ('merging_max_seqlets_subsample',
                 self.merging_max_seqlets_subsample),
                #('threshold_for_spurious_merge_detection',
                # self.threshold_for_spurious_merge_detection),
                ('min_similarity_for_seqlet_assignment',
                 self.min_similarity_for_seqlet_assignment),
                ('final_min_cluster_size', self.final_min_cluster_size),
                ('min_ic_in_window', self.min_ic_in_window),
                ('min_ic_windowsize', self.min_ic_windowsize),
                ('ppm_pseudocount', self.ppm_pseudocount),
                ('final_flank_to_add', self.final_flank_to_add),
                ])
        return to_return

    def __call__(self, track_set, onehot_track_name,
                       contrib_scores_track_names,
                       hypothetical_contribs_track_names,
                       track_signs,
                       other_comparison_track_names=[]):

        bg_freq = np.mean(
            track_set.track_name_to_data_track[onehot_track_name].fwd_tracks,
            axis=(0,1))
        assert len(bg_freq.shape)==1

        assert len(track_signs)==len(hypothetical_contribs_track_names)
        assert len(track_signs)==len(contrib_scores_track_names)

        seqlets_sorter = (lambda arr:
                          sorted(arr,
                                 key=lambda x:
                                  -np.sum([np.sum(np.abs(x[track_name].fwd))
                                     for track_name
                                     in contrib_scores_track_names])))

        if (self.initclusterer_factory is not None):
            self.initclusterer_factory.set_onehot_track_name(onehot_track_name)
        initclusterer_factory = self.initclusterer_factory

        pattern_comparison_settings =\
            affinitymat.core.PatternComparisonSettings(
                track_names=hypothetical_contribs_track_names
                            +contrib_scores_track_names
                            +other_comparison_track_names,
                track_transformer=affinitymat.L1Normalizer(),
                min_overlap=self.min_overlap_while_sliding)

        #coarse_grained 1d embedder
        seqlets_to_1d_embedder = self.embedder_factory(
                onehot_track_name=onehot_track_name,
                toscore_track_names_and_signs=list(
                zip(hypothetical_contribs_track_names,
                    [np.sign(x) for x in track_signs])),
                n_jobs=self.n_cores)

        #affinity matrix from embeddings
        if (self.use_pynnd):
            sparse_affmat_from_fwdnrev1dvecs = (
             affinitymat.core.PynndSparseNumpyCosineSimFromFwdAndRevOneDVecs(
                n_neighbors=self.nearest_neighbors_to_compute,
                verbose=self.verbose,
                n_jobs=self.n_cores))
        else:
            sparse_affmat_from_fwdnrev1dvecs =\
                affinitymat.core.SparseNumpyCosineSimFromFwdAndRevOneDVecs(
                        n_neighbors=self.nearest_neighbors_to_compute,
                        verbose=self.verbose)
        coarse_affmat_computer =\
          affinitymat.core.SparseAffmatFromFwdAndRevSeqletEmbeddings(
            seqlets_to_1d_embedder=seqlets_to_1d_embedder,
            sparse_affmat_from_fwdnrev1dvecs=sparse_affmat_from_fwdnrev1dvecs,
            verbose=self.verbose)

        affmat_from_seqlets_with_nn_pairs =\
            affinitymat.core.AffmatFromSeqletsWithNNpairs(
                pattern_comparison_settings=pattern_comparison_settings,
                sim_metric_on_nn_pairs=\
                    affinitymat.core.ParallelCpuCrossMetricOnNNpairs(
                        n_cores=self.n_cores,
                        cross_metric_single_region=
                            affinitymat.core.CrossContinJaccardSingleRegion()))

        filter_mask_from_correlation =\
            affinitymat.core.FilterMaskFromCorrelation(
                correlation_threshold=self.affmat_correlation_threshold,
                verbose=self.verbose)

        aff_to_dist_mat = affinitymat.transformers.AffToDistViaInvLogistic()
        #density_adapted_affmat_transformer =\
        #    affinitymat.transformers.NNTsneConditionalProbs(
        #        perplexity=self.tsne_perplexity,
        #        aff_to_dist_mat=aff_to_dist_mat)

        #prepare the clusterers for the different rounds
        # No longer a need for symmetrization because am symmetrizing by
        # taking the geometric mean elsewhere
        affmat_transformer_r1 =\
            affinitymat.transformers.AdhocAffMatTransformer(lambda x: x)
        #affinitymat.transformers.SymmetrizeByAddition(
        #                            probability_normalize=True)
        print("TfModiscoSeqletsToPatternsFactory: seed=%d" % self.seed)
        if (self.use_louvain):
            for n_runs, level_to_return in self.louvain_num_runs_and_levels_r1:
                affmat_transformer_r1 = affmat_transformer_r1.chain(
                    affinitymat.transformers.LouvainMembershipAverage(
                        n_runs=n_runs,
                        level_to_return=level_to_return,
                        parallel_threads=self.n_cores, seed=self.seed))
            clusterer_r1 = cluster.core.LouvainCluster(
                level_to_return=self.final_louvain_level_to_return,
                affmat_transformer=affmat_transformer_r1,
                contin_runs=self.contin_runs_r1,
                verbose=self.verbose, seed=self.seed)
        else:
            clusterer_r1 = cluster.core.LeidenClusterParallel(
                n_jobs=self.n_cores,
                affmat_transformer=affmat_transformer_r1,
                numseedstotry=self.contin_runs_r1,
                n_leiden_iterations=self.n_leiden_iterations_r1,
                verbose=self.verbose)

        #No longer a need for symmetrization because am symmetrizing by
        # taking the geometric mean elsewhere
        affmat_transformer_r2 =\
            affinitymat.transformers.AdhocAffMatTransformer(lambda x: x)
        #affmat_transformer_r2 = affinitymat.transformers.SymmetrizeByAddition(
        #                        probability_normalize=True)
        if (self.use_louvain):
            for n_runs, level_to_return in self.louvain_num_runs_and_levels_r2:
                affmat_transformer_r2 = affmat_transformer_r2.chain(
                    affinitymat.transformers.LouvainMembershipAverage(
                        n_runs=n_runs,
                        level_to_return=level_to_return,
                        parallel_threads=self.n_cores, seed=self.seed))
            clusterer_r2 = cluster.core.LouvainCluster(
                level_to_return=self.final_louvain_level_to_return,
                affmat_transformer=affmat_transformer_r2,
                contin_runs=self.contin_runs_r2,
                verbose=self.verbose, seed=self.seed,
                initclusters_weight=self.louvain_initclusters_weight)
        else:
            clusterer_r2 = cluster.core.LeidenClusterParallel(
                n_jobs=self.n_cores,
                affmat_transformer=affmat_transformer_r2,
                numseedstotry=self.contin_runs_r2,
                n_leiden_iterations=self.n_leiden_iterations_r2,
                verbose=self.verbose)

        clusterer_per_round = [clusterer_r1, clusterer_r2]

        #prepare the seqlet aggregator
        expand_trim_expand1 =\
            aggregator.ExpandSeqletsToFillPattern(
                track_set=track_set,
                flank_to_add=self.initial_flank_to_add).chain(
            aggregator.TrimToBestWindowByIC(
                window_size=self.trim_to_window_size,
                onehot_track_name=onehot_track_name,
                bg_freq=bg_freq)).chain(
            aggregator.ExpandSeqletsToFillPattern(
                track_set=track_set,
                flank_to_add=self.initial_flank_to_add))
        postprocessor1 =\
            aggregator.TrimToFracSupport(
                        min_frac=self.frac_support_to_trim_to,
                        min_num=self.min_num_to_trim_to,
                        verbose=self.verbose)\
                      .chain(expand_trim_expand1)
        seqlet_aggregator = aggregator.GreedySeqletAggregator(
            pattern_aligner=core.CrossContinJaccardPatternAligner(
                pattern_comparison_settings=pattern_comparison_settings),
                seqlet_sort_metric=
                    lambda x: -sum([np.sum(np.abs(x[track_name].fwd)) for
                               track_name in contrib_scores_track_names]),
            track_set=track_set, #needed for seqlet expansion
            postprocessor=postprocessor1)

        def sign_consistency_func(motif):
            motif_track_signs = [
                np.sign(np.sum(motif[contrib_scores_track_name].fwd)) for
                contrib_scores_track_name in contrib_scores_track_names]
            return all([(x==y) for x,y in zip(motif_track_signs, track_signs)])

        #prepare the similar patterns collapser
        pattern_to_seqlet_sim_computer =\
            affinitymat.core.AffmatFromSeqletsWithNNpairs(
                pattern_comparison_settings=pattern_comparison_settings,
                sim_metric_on_nn_pairs=\
                    affinitymat.core.ParallelCpuCrossMetricOnNNpairs(
                        n_cores=self.n_cores,
                        cross_metric_single_region=\
                            affinitymat.core.CrossContinJaccardSingleRegion(),
                        verbose=False))

        #similarity settings for merging
        prob_and_sim_merge_thresholds =\
            self.prob_and_pertrack_sim_merge_thresholds
        prob_and_sim_dealbreaker_thresholds =\
            self.prob_and_pertrack_sim_dealbreaker_thresholds

        similar_patterns_collapser =\
            aggregator.DynamicDistanceSimilarPatternsCollapser2(
                pattern_comparison_settings=pattern_comparison_settings,
                track_set=track_set,
                pattern_aligner=core.CrossCorrelationPatternAligner(
                    pattern_comparison_settings=
                        affinitymat.core.PatternComparisonSettings(
                            track_names=(
                                contrib_scores_track_names+
                                other_comparison_track_names),
                            track_transformer=
                                affinitymat.MeanNormalizer().chain(
                                affinitymat.MagnitudeNormalizer()),
                            min_overlap=self.min_overlap_while_sliding)),
                collapse_condition=(lambda prob, aligner_sim:
                    any([(prob >= x[0] and aligner_sim >= x[1])
                         for x in prob_and_sim_merge_thresholds])),
                dealbreaker_condition=(lambda prob, aligner_sim:
                    any([(prob <= x[0] and aligner_sim <= x[1])
                         for x in prob_and_sim_dealbreaker_thresholds])),
                postprocessor=postprocessor1,
                verbose=self.verbose,
                max_seqlets_subsample=self.merging_max_seqlets_subsample,
                n_cores=self.n_cores)

        subcluster_settings = {
            "pattern_comparison_settings": pattern_comparison_settings,
            "perplexity": self.subcluster_perplexity,
            "n_jobs": self.n_cores,
        }
        spurious_merge_detector = aggregator.DetectSpuriousMerging2(
            subcluster_settings=subcluster_settings,
            verbose=self.verbose,
            min_in_subcluster=max(self.final_min_cluster_size,
                                  self.subcluster_perplexity),
            similar_patterns_collapser=similar_patterns_collapser)

        #spurious_merge_detector = aggregator.DetectSpuriousMerging(
        #    track_names=contrib_scores_track_names,
        #    track_transformer=affinitymat.core.L1Normalizer(),
        #    affmat_from_1d=affinitymat.core.ContinJaccardSimilarity(
        #                    make_positive=True, verbose=False),
        #    diclusterer=cluster.core.LouvainCluster(
        #                    level_to_return=1,
        #                    max_clusters=2, contin_runs=20,
        #                    verbose=False, seed=self.seed),
        #    is_dissimilar_func=aggregator.PearsonCorrIsDissimilarFunc(
        #                threshold=self.threshold_for_spurious_merge_detection,
        #                verbose=self.verbose),
        #    min_in_subcluster=self.final_min_cluster_size)

        #similar_patterns_collapser =\
        #    aggregator.DynamicDistanceSimilarPatternsCollapser(
        #        pattern_to_pattern_sim_computer=
        #            pattern_to_seqlet_sim_computer,
        #        aff_to_dist_mat=aff_to_dist_mat,
        #        pattern_aligner=core.CrossCorrelationPatternAligner(
        #            pattern_comparison_settings=
        #                affinitymat.core.PatternComparisonSettings(
        #                    track_names=(
        #                        contrib_scores_track_names+
        #                        other_comparison_track_names),
        #                    track_transformer=
        #                        affinitymat.MeanNormalizer().chain(
        #                        affinitymat.MagnitudeNormalizer()),
        #                    min_overlap=self.min_overlap_while_sliding)),
        #        collapse_condition=(lambda prob, aligner_sim:
        #            any([(prob > x[0] and aligner_sim > x[1])
        #                 for x in prob_and_sim_merge_thresholds])),
        #        dealbreaker_condition=(lambda prob, aligner_sim:
        #            any([(prob < x[0] and aligner_sim < x[1])
        #                 for x in prob_and_sim_dealbreaker_thresholds])),
        #        postprocessor=postprocessor1,
        #        verbose=self.verbose)

        pattern_filterer = pattern_filterer_module.MinSeqletSupportFilterer(
            min_seqlet_support=self.final_min_cluster_size).chain(
                pattern_filterer_module.MinICinWindow(
                    window_size=self.min_ic_windowsize,
                    min_ic_in_window=self.min_ic_in_window,
                    background=bg_freq,
                    sequence_track_name=onehot_track_name,
                    ppm_pseudocount=self.ppm_pseudocount
                )
            )

        #seqlet_reassigner =\
        #   aggregator.ReassignSeqletsFromSmallClusters(
        #    seqlet_assigner=aggregator.AssignSeqletsByBestMetric(
        #        pattern_comparison_settings=pattern_comparison_settings,
        #        individual_aligner_metric=
        #            core.get_best_alignment_crosscontinjaccard,
        #        matrix_affinity_metric=
        #            affinitymat.core.CrossContinJaccardMultiCoreCPU(
        #                verbose=self.verbose, n_cores=self.n_cores),
        #        min_similarity=self.min_similarity_for_seqlet_assignment,
        #        track_set=track_set),
        #    min_cluster_size=self.final_min_cluster_size,
        #    postprocessor=expand_trim_expand1,
        #    verbose=self.verbose)

        final_postprocessor = aggregator.ExpandSeqletsToFillPattern(
                                        track_set=track_set,
                                        flank_to_add=self.final_flank_to_add)

        return TfModiscoSeqletsToPatterns(
                seqlets_sorter=seqlets_sorter,
                initclusterer_factory=initclusterer_factory,
                coarse_affmat_computer=coarse_affmat_computer,
                nearest_neighbors_to_compute=self.nearest_neighbors_to_compute,
                affmat_from_seqlets_with_nn_pairs=
                    affmat_from_seqlets_with_nn_pairs,
                filter_mask_from_correlation=filter_mask_from_correlation,
                filter_beyond_first_round=self.filter_beyond_first_round,
                skip_fine_grained=self.skip_fine_grained,
                aff_to_dist_mat=aff_to_dist_mat,
                tsne_perplexity=self.tsne_perplexity,
                #density_adapted_affmat_transformer=
                #    density_adapted_affmat_transformer,
                clusterer_per_round=clusterer_per_round,
                seqlet_aggregator=seqlet_aggregator,
                sign_consistency_func=sign_consistency_func,
                subcluster_settings=subcluster_settings,
                spurious_merge_detector=spurious_merge_detector,
                similar_patterns_collapser=similar_patterns_collapser,
                #seqlet_reassigner=seqlet_reassigner,
                pattern_filterer=pattern_filterer,
                final_postprocessor=final_postprocessor,
                verbose=self.verbose,
                n_cores=self.n_cores,
                other_config={
                 'onehot_track_name': onehot_track_name,
                 'contrib_scores_track_names': contrib_scores_track_names,
                 'hypothetical_contribs_track_names':
                    hypothetical_contribs_track_names,
                 'track_signs': track_signs,
                 'other_comparison_track_names': other_comparison_track_names},
                )

    def save_hdf5(self, grp):
        grp.attrs['jsonable_config'] =\
            json.dumps(self.jsonable_config, indent=4, separators=(',', ': '))

class SeqletsToPatternsResults(object):

    def __init__(self,
                 each_round_initcluster_motifs,
                 patterns,
                 remaining_patterns,
                 pattern_merge_hierarchy,
                 cluster_results,
                 total_time_taken,
                 other_config={},
                 success=True,
                 **kwargs):
        self.each_round_initcluster_motifs = each_round_initcluster_motifs
        self.other_config = other_config
        self.success = success
        self.patterns = patterns
        self.remaining_patterns = remaining_patterns
        self.pattern_merge_hierarchy = pattern_merge_hierarchy
        self.cluster_results = cluster_results
        self.total_time_taken = total_time_taken
        self.__dict__.update(**kwargs)

    def save_each_round_initcluster_motifs(self, grp):
        all_round_names = []
        for (round_idx,initcluster_motifs)\
            in enumerate(self.each_round_initcluster_motifs):
            round_name = "round_"+str(round_idx)
            util.save_patterns(patterns=initcluster_motifs,
                               grp=grp.create_group(round_name))
        util.save_string_list(
            string_list=all_round_names,
            dset_name="all_round_names",
            grp=grp)

    @classmethod
    def load_each_round_initcluster_motifs(cls, grp, track_set):
        all_round_names = util.load_string_list(dset_name="all_round_names",
                                                grp=grp)
        each_round_initcluster_motifs = []
        for round_name in all_round_names:
            round_grp = grp[round_name]
            initcluster_motifs = load_patterns(grp=round_grp,
                                               track_set=track_set)
            each_round_initcluster_motifs.append(initcluster_motifs)
        return each_round_initcluster_motifs

    @classmethod
    def from_hdf5(cls, grp, track_set):
        success = grp.attrs.get("success", False)
        if (success):
            if ("each_round_initcluster_motifs" not in grp):
                each_round_initcluster_motifs = None
            else:
                each_round_initcluster_motifs =\
                    cls.load_each_round_initcluster_motifs(
                        grp=grp["each_round_initcluster_motifs"],
                        track_set=track_set)
            patterns = util.load_patterns(grp=grp["patterns"],
                                          track_set=track_set)
            if "remaining_patterns" in grp:
                remaining_patterns = util.load_patterns(
                    grp=grp["remaining_patterns"],
                    track_set=track_set)
            else: #backwards compatibility
                remaining_patterns = []
            cluster_results = None
            total_time_taken = grp.attrs["total_time_taken"]
            if ("pattern_merge_hierarchy" in grp):
                pattern_merge_hierarchy =\
                    aggregator.PatternMergeHierarchy.from_hdf5(
                        grp=grp["pattern_merge_hierarchy"],
                        track_set=track_set)
            else:
                pattern_merge_hierarchy = None
            return cls(
                each_round_initcluster_motifs=each_round_initcluster_motifs,
                patterns=patterns,
                remaining_patterns=remaining_patterns,
                pattern_merge_hierarchy=pattern_merge_hierarchy,
                cluster_results=cluster_results,
                total_time_taken=total_time_taken)
        else:
            return cls(success=False, patterns=None, cluster_results=None,
                       total_time_taken=None,
                       each_round_initcluster_motifs=None,
                       remaining_patterns=None,
                       pattern_merge_hierarchy=None)

    def save_hdf5(self, grp):
        grp.attrs["success"] = self.success
        grp.attrs["other_config"] =\
            json.dumps(self.other_config, indent=4, separators=(',', ': '))
        if (self.success):
            grp.attrs["total_time_taken"] = self.total_time_taken
            if (self.each_round_initcluster_motifs is not None):
                self.save_each_round_initcluster_motifs(
                    grp=grp.create_group("each_round_initcluster_motifs"))
            util.save_patterns(self.patterns,
                               grp.create_group("patterns"))
            util.save_patterns(
                self.remaining_patterns,
                grp.create_group("remaining_patterns"))
            self.cluster_results.save_hdf5(grp.create_group("cluster_results"))
            grp.attrs['total_time_taken'] = self.total_time_taken
            self.pattern_merge_hierarchy.save_hdf5(
                    grp=grp.create_group("pattern_merge_hierarchy"))

class AbstractSeqletsToPatterns(object):

    def __call__(self, seqlets):
        raise NotImplementedError()

class TfModiscoSeqletsToPatterns(AbstractSeqletsToPatterns):

    def __init__(self, seqlets_sorter,
                       initclusterer_factory,
                       coarse_affmat_computer,
                       nearest_neighbors_to_compute,
                       affmat_from_seqlets_with_nn_pairs,
                       filter_mask_from_correlation,
                       filter_beyond_first_round,
                       skip_fine_grained,
                       aff_to_dist_mat,
                       tsne_perplexity,
                       #density_adapted_affmat_transformer,
                       clusterer_per_round,
                       seqlet_aggregator,
                       sign_consistency_func,
                       spurious_merge_detector,
                       similar_patterns_collapser,
                       pattern_filterer,
                       #seqlet_reassigner,
                       final_postprocessor,
                       subcluster_settings,
                       n_cores,
                       other_config={},
                       verbose=True):

        self.seqlets_sorter = seqlets_sorter
        self.initclusterer_factory = initclusterer_factory
        self.coarse_affmat_computer = coarse_affmat_computer
        self.nearest_neighbors_to_compute = nearest_neighbors_to_compute
        self.affmat_from_seqlets_with_nn_pairs =\
            affmat_from_seqlets_with_nn_pairs
        self.filter_mask_from_correlation = filter_mask_from_correlation
        self.filter_beyond_first_round = filter_beyond_first_round
        self.skip_fine_grained = skip_fine_grained
        self.aff_to_dist_mat = aff_to_dist_mat
        self.tsne_perplexity = tsne_perplexity
        #self.density_adapted_affmat_transformer =\
        #    density_adapted_affmat_transformer
        self.clusterer_per_round = clusterer_per_round
        self.seqlet_aggregator = seqlet_aggregator
        self.sign_consistency_func = sign_consistency_func

        self.spurious_merge_detector = spurious_merge_detector
        self.similar_patterns_collapser = similar_patterns_collapser
        #self.seqlet_reassigner = seqlet_reassigner
        self.pattern_filterer = pattern_filterer
        self.final_postprocessor = final_postprocessor

        self.verbose = verbose
        self.subcluster_settings = subcluster_settings
        self.n_cores = n_cores
        self.other_config = other_config

    def get_cluster_to_aggregate_motif(self, seqlets, cluster_indices,
                                       sign_consistency_check,
                                       min_seqlets_in_motif):
        num_clusters = max(cluster_indices+1)
        cluster_to_seqlets = defaultdict(list)
        assert len(seqlets)==len(cluster_indices)
        for seqlet,idx in zip(seqlets, cluster_indices):
            cluster_to_seqlets[idx].append(seqlet)

        cluster_to_motif = OrderedDict()
        cluster_to_eliminated_motif = OrderedDict()
        for i in range(num_clusters):
            if (len(cluster_to_seqlets[i]) >= min_seqlets_in_motif):
                if (self.verbose):
                    print("Aggregating for cluster "+str(i)+" with "
                          +str(len(cluster_to_seqlets[i]))+" seqlets")
                    print_memory_use()
                    sys.stdout.flush()
                motifs = self.seqlet_aggregator(cluster_to_seqlets[i])
                assert len(motifs)<=1
                if (len(motifs) > 0):
                    motif = motifs[0]
                    if (sign_consistency_check==False or
                        self.sign_consistency_func(motif)):
                        cluster_to_motif[i] = motif
                    else:
                        if (self.verbose):
                            print("Dropping cluster "+str(i)+
                                  " with "+str(motif.num_seqlets)
                                  +" seqlets due to sign disagreement")
                        cluster_to_eliminated_motif[i] = motif
        return cluster_to_motif, cluster_to_eliminated_motif

    def do_density_adaptation(self, new_rows_distmat_nn, new_rows_nn,
                                    new_rows_betas, new_rows_normfactors):
        new_rows_densadapted_affmat_nn = []
        for i in range(len(new_rows_distmat_nn)):
            densadapted_row = []
            for j,distance in zip(new_rows_nn[i], new_rows_distmat_nn[i]):
                densadapted_row.append(np.sqrt(
                  (np.exp(-distance/new_rows_betas[i])/new_rows_normfactors[i])
                 *(np.exp(-distance/new_rows_betas[j])/
                   new_rows_normfactors[j])))
            new_rows_densadapted_affmat_nn.append(densadapted_row)
        return new_rows_densadapted_affmat_nn

    def __call__(self, seqlets):

        seqlets = self.seqlets_sorter(seqlets)
        if (self.initclusterer_factory is not None):
            initclusterer = self.initclusterer_factory(seqlets=seqlets)
        else:
            initclusterer = None

        start = time.time()

        #seqlets_sets = []
        #coarse_affmats = []
        #nn_affmats = []
        #filtered_seqlets_sets = []
        #filtered_affmats = []
        #density_adapted_affmats = []
        #cluster_results_sets = []
        #cluster_to_motif_sets = []
        #cluster_to_eliminated_motif_sets = []

        if (initclusterer is not None):
            each_round_initcluster_motifs = []
        else:
            each_round_initcluster_motifs = None

        for round_idx, clusterer in enumerate(self.clusterer_per_round):
            import gc
            gc.collect()

            round_num = round_idx+1

            if (initclusterer is not None):
                initclusters = initclusterer(seqlets=seqlets)
                initcluster_motifs =\
                    list(self.get_cluster_to_aggregate_motif(
                            seqlets=seqlets,
                            cluster_indices=initclusters,
                            sign_consistency_check=False,
                            min_seqlets_in_motif=2)[0].values())
                each_round_initcluster_motifs.append(initcluster_motifs)
            else:
                initclusters = None
                initcluster_motifs = None

            if (len(seqlets)==0):
                if (self.verbose):
                    print("len(seqlets) is 0 - bailing!")
                return SeqletsToPatternsResults(
                        each_round_initcluster_motifs=None,
                        patterns=None,
                        remaining_patterns=None,
                        pattern_merge_hierarchy=None,
                        cluster_results=None,
                        total_time_taken=None,
                        success=False,
                        seqlets=None,
                        affmat=None)

            if (self.verbose):
                print("(Round "+str(round_num)+
                      ") num seqlets: "+str(len(seqlets)))
                print("(Round "+str(round_num)+") Computing coarse affmat")
                print_memory_use()
                sys.stdout.flush()
            #coarse_affmat = self.coarse_affmat_computer(seqlets)
            #coarse_affmats.append(coarse_affmat)
            coarse_affmat_nn, seqlet_neighbors =\
                self.coarse_affmat_computer(seqlets, initclusters=initclusters)
            gc.collect()

            if (self.verbose):
                print("(Round "+str(round_num)+") Computed coarse affmat")
                print_memory_use()
                sys.stdout.flush()

            if (self.skip_fine_grained==False):
                #nn_start = time.time()
                #if (self.verbose):
                #    print("(Round "+str(round_num)+") Compute nearest neighbors"
                #          +" from coarse affmat")
                #    print_memory_use()
                #    sys.stdout.flush()

                #seqlet_neighbors = get_seqlet_neighbors_with_initcluster(
                #    nearest_neighbors_to_compute=
                #     self.nearest_neighbors_to_compute,
                #    coarse_affmat=coarse_affmat,
                #    initclusters=initclusters)

                #if (self.verbose):
                #    print("Computed nearest neighbors in",
                #          round(time.time()-nn_start,2),"s")
                #    print_memory_use()
                #    sys.stdout.flush()

                nn_affmat_start = time.time()
                if (self.verbose):
                    print("(Round "+str(round_num)+") Computing affinity matrix"
                          +" on nearest neighbors")
                    print_memory_use()
                    sys.stdout.flush()
                #nn_affmat = self.affmat_from_seqlets_with_nn_pairs(
                #                    seqlet_neighbors=seqlet_neighbors,
                #                    seqlets=seqlets)
                #nn_affmats.append(nn_affmat)

                fine_affmat_nn = self.affmat_from_seqlets_with_nn_pairs(
                                    seqlet_neighbors=seqlet_neighbors,
                                    seqlets=seqlets,
                                    return_sparse=True)
                #get the fine_affmat_nn reorderings
                reorderings = np.array([np.argsort(-finesimsinrow)
                                        for finesimsinrow in fine_affmat_nn])

                #reorder fine_affmat_nn, coarse_affmat_nn and seqlet_neighbors
                # according to reorderings
                fine_affmat_nn = [finesimsinrow[rowreordering]
                                      for (finesimsinrow, rowreordering)
                                      in zip(fine_affmat_nn, reorderings)]
                coarse_affmat_nn = [coarsesimsinrow[rowreordering]
                                      for (coarsesimsinrow, rowreordering)
                                      in zip(coarse_affmat_nn, reorderings)]
                seqlet_neighbors = [nnrow[rowreordering]
                                      for (nnrow, rowreordering)
                                      in zip(seqlet_neighbors, reorderings)]

                del reorderings
                gc.collect()

                if (self.verbose):
                    print("(Round "+str(round_num)+") Computed affinity matrix"
                          +" on nearest neighbors in",
                          round(time.time()-nn_affmat_start,2),"s")
                    print_memory_use()
                    sys.stdout.flush()

                #filter by correlation
                if (round_idx == 0 or self.filter_beyond_first_round==True):
                    #the filter_mask_from_correlation function only operates
                    # on columns in which np.abs(main_affmat) > 0
                    #filtered_rows_mask = self.filter_mask_from_correlation(
                    #                        main_affmat=nn_affmat,
                    #                        other_affmat=coarse_affmat)
                    filtered_rows_mask = self.filter_mask_from_correlation(
                                            main_affmat=fine_affmat_nn,
                                            other_affmat=coarse_affmat_nn)
                    if (self.verbose):
                        print("(Round "+str(round_num)+") Retained "
                              +str(np.sum(filtered_rows_mask))
                              +" rows out of "+str(len(filtered_rows_mask))
                              +" after filtering")
                        print_memory_use()
                        sys.stdout.flush()
                else:
                    filtered_rows_mask = np.array([True for x in seqlets])
                    if (self.verbose):
                        print("Not applying filtering for "
                              +"rounds above first round")
                        print_memory_use()
                        sys.stdout.flush()

                del coarse_affmat_nn
                gc.collect()

                filtered_seqlets = [x[0] for x in
                    zip(seqlets, filtered_rows_mask) if (x[1])]
                if (initclusters is not None):
                    filtered_initclusters = initclusters[filtered_rows_mask]
                else:
                    filtered_initclusters = None
                #filtered_seqlets_sets.append(filtered_seqlets)

                #filtered_affmat =\
                #    nn_affmat[filtered_rows_mask][:,filtered_rows_mask]
                #del coarse_affmat
                #del nn_affmat

                #figure out a mapping from pre-filtering to the
                # post-filtering indices
                new_idx_mapping = (
                    np.cumsum(1.0*(filtered_rows_mask)).astype("int")-1)
                retained_indices = set(np.arange(len(filtered_rows_mask))[
                                                  filtered_rows_mask])
                del filtered_rows_mask
                filtered_neighbors = []
                filtered_affmat_nn = []
                for old_row_idx, (old_neighbors,affmat_row) in enumerate(
                                    zip(seqlet_neighbors, fine_affmat_nn)):
                    if old_row_idx in retained_indices:
                        filtered_old_neighbors = [
                            neighbor for neighbor in old_neighbors if neighbor
                            in retained_indices]
                        filtered_affmat_row = [
                            affmatval for affmatval,neighbor
                            in zip(affmat_row,old_neighbors)
                            if neighbor in retained_indices]
                        filtered_neighbors_row = [
                            new_idx_mapping[neighbor] for neighbor
                            in filtered_old_neighbors]
                        filtered_neighbors.append(filtered_neighbors_row)
                        filtered_affmat_nn.append(filtered_affmat_row)

                #overwrite seqlet_neighbors...should be ok if the rows are
                # not all the same length
                seqlet_neighbors = filtered_neighbors
                del (filtered_neighbors, retained_indices, new_idx_mapping)
            else:
                filtered_affmat = coarse_affmat_nn
                filtered_seqlets = seqlets
                if (initclusters is not None):
                    filtered_initclusters = initclusters
                else:
                    filtered_initclusters = None

            if (self.verbose):
                print("(Round "+str(round_num)+") Computing density "
                      +"adapted affmat")
                print_memory_use()
                sys.stdout.flush()

            #density_adapted_affmat =\
            #    self.density_adapted_affmat_transformer(filtered_affmat)
            #del filtered_affmat
            #density_adapted_affmats.append(density_adapted_affmat)

            #apply aff_to_dist_mat one row at a time
            distmat_nn = [self.aff_to_dist_mat(affinity_mat=x)
                          for x in filtered_affmat_nn]
            del filtered_affmat_nn

            if (self.verbose):
                print("Symmetrizing nearest neighbors")
            #Note: the fine-grained similarity metric isn't actually symmetric
            # because a different input will get padded with zeros depending
            # on which seqlets are specified as the filters and which seqlets
            # are specified as the 'thing to scan'. So explicit symmetrization
            # is worthwhile
            sym_seqlet_neighbors, sym_distmat_nn = util.symmetrize_nn_distmat(
                distmat_nn=distmat_nn, nn=seqlet_neighbors,
                average_with_transpose=True)
            del distmat_nn
            del seqlet_neighbors

            if (self.verbose):
                print("Computing betas for density adaptation")

            #Compute beta values for the density adaptation. *store it*
            betas_and_ps = Parallel(n_jobs=self.n_cores)(
                     delayed(util.binary_search_perplexity)(
                          self.tsne_perplexity, distances)
                     for distances in sym_distmat_nn)
            betas = np.array([x[0] for x in betas_and_ps])

            if (self.verbose):
                print("Computing normalizing denominators")

            #also compute the normalization factor needed to get probs to sum to 1
            #note: sticking to lists here because different rows of
            # sym_distmat_nn may have different lengths after adding in
            # the symmetric pairs
            densadapted_affmat_nn_unnorm = [np.exp(-np.array(distmat_row)/beta)
                for distmat_row, beta in zip(sym_distmat_nn, betas)]
            normfactors = np.array([max(np.sum(x),1e-8) for x in
                                    densadapted_affmat_nn_unnorm])

            sym_densadapted_affmat_nn = self.do_density_adaptation(
                new_rows_distmat_nn=sym_distmat_nn,
                new_rows_nn=sym_seqlet_neighbors,
                new_rows_betas=betas,
                new_rows_normfactors=normfactors)

            util.verify_symmetric_nn_affmat(
                affmat_nn=sym_densadapted_affmat_nn,
                nn=sym_seqlet_neighbors)

            #Make csr matrix
            csr_density_adapted_affmat =\
                util.coo_matrix_from_neighborsformat(
                    entries=sym_densadapted_affmat_nn,
                    neighbors=sym_seqlet_neighbors,
                    ncols=len(sym_densadapted_affmat_nn)).tocsr()

            if (self.verbose):
                print("(Round "+str(round_num)+") Computing clustering")
                print_memory_use()
                sys.stdout.flush()

            #cluster_results = clusterer(density_adapted_affmat,
            #                            initclusters=filtered_initclusters)
            #del density_adapted_affmat
            #cluster_results_sets.append(cluster_results)
            cluster_results = clusterer(csr_density_adapted_affmat,
                                        initclusters=filtered_initclusters)
            del csr_density_adapted_affmat

            num_clusters = max(cluster_results.cluster_indices+1)
            cluster_idx_counts = Counter(cluster_results.cluster_indices)
            if (self.verbose):
                print("Got "+str(num_clusters)
                      +" clusters after round "+str(round_num))
                print("Counts:")
                print(dict([x for x in cluster_idx_counts.items()]))
                print_memory_use()
                sys.stdout.flush()

            if (self.verbose):
                print("(Round "+str(round_num)+") Aggregating seqlets"
                      +" in each cluster")
                print_memory_use()
                sys.stdout.flush()

            cluster_to_motif, cluster_to_eliminated_motif =\
                self.get_cluster_to_aggregate_motif(
                    seqlets=filtered_seqlets,
                    cluster_indices=cluster_results.cluster_indices,
                    sign_consistency_check=True,
                    min_seqlets_in_motif=0)

            #obtain unique seqlets from adjusted motifs
            seqlets = list(dict([(y.exidx_start_end_string, y)
                             for x in cluster_to_motif.values()
                             for y in x.seqlets]).values())

        if (self.verbose):
            print("Got "+str(len(cluster_to_motif.values()))+" clusters")
            print("Splitting into subclusters...")
            print_memory_use()
            sys.stdout.flush()

        split_patterns = self.spurious_merge_detector(
                            cluster_to_motif.values())

        if (len(split_patterns)==0):
            if (self.verbose):
                print("No more surviving patterns - bailing!")
            return SeqletsToPatternsResults(
                    each_round_initcluster_motifs=None,
                    patterns=None,
                    remaining_patterns=None,
                    pattern_merge_hierarchy=None,
                    cluster_results=None,
                    total_time_taken=None,
                    success=False,
                    seqlets=None,
                    affmat=None)

        #Now start merging patterns
        if (self.verbose):
            print("Merging on "+str(len(split_patterns))+" clusters")
            print_memory_use()
            sys.stdout.flush()
        merged_patterns, pattern_merge_hierarchy =\
            self.similar_patterns_collapser(
                patterns=split_patterns)
        merged_patterns = sorted(merged_patterns, key=lambda x: -x.num_seqlets)
        if (self.verbose):
            print("Got "+str(len(merged_patterns))+" patterns after merging")
            print_memory_use()
            sys.stdout.flush()

        if (self.verbose):
            print("Performing filtering")
            print_memory_use()
            sys.stdout.flush()

        final_patterns, remaining_patterns = self.pattern_filterer(
                                                    merged_patterns)
        #reassigned_patterns = self.seqlet_reassigner(merged_patterns)
        final_patterns = self.final_postprocessor(final_patterns)
        remaining_patterns =\
            self.final_postprocessor(remaining_patterns)

        if (self.verbose):
            print("Got "+str(len(final_patterns))
                  +" patterns after filtering")
            print_memory_use()
            sys.stdout.flush()

        total_time_taken = round(time.time()-start,2)
        if (self.verbose):
            print("Total time taken is "
                  +str(total_time_taken)+"s")
            print_memory_use()
            sys.stdout.flush()

        #apply subclustering procedure on the final patterns
        print("Applying subclustering to the final motifs")
        for patternidx, pattern in enumerate(final_patterns):
            print("On pattern",patternidx)
            pattern.compute_subclusters_and_embedding(
                verbose=self.verbose,
                **self.subcluster_settings)

        results = SeqletsToPatternsResults(
            each_round_initcluster_motifs=each_round_initcluster_motifs,
            patterns=final_patterns,
            remaining_patterns=remaining_patterns,
            seqlets=filtered_seqlets, #last stage of filtered seqlets
            #affmat=filtered_affmat,
            cluster_results=cluster_results,
            total_time_taken=total_time_taken,

            #seqlets_sets=seqlets_sets,
            #coarse_affmats=coarse_affmats,
            #nn_affmats=nn_affmats,
            #filtered_seqlets_sets=filtered_seqlets_sets,
            #filtered_affmats=filtered_affmats,
            #density_adapted_affmats=density_adapted_affmats,
            #cluster_results_sets=cluster_results_sets,
            #cluster_to_motif_sets=cluster_to_motif_sets,
            #cluster_to_eliminated_motif_sets=cluster_to_eliminated_motif_sets,

            merged_patterns=merged_patterns,
            pattern_merge_hierarchy=pattern_merge_hierarchy,
            #reassigned_patterns=reassigned_patterns
        )

        return results


# Copyright Contributors to the Pyro project.
# SPDX-License-Identifier: Apache-2.0

from collections import namedtuple
from contextlib import contextmanager
from functools import partial
from typing import Callable, Dict, List, Optional
import warnings

import numpy as np

import jax
from jax import device_get, jacfwd, lax, random, tree_flatten, value_and_grad
from jax.flatten_util import ravel_pytree
from jax.lax import broadcast_shapes
import jax.numpy as jnp
from jax.tree_util import tree_map

import numpyro
from numpyro.distributions import constraints
from numpyro.distributions.transforms import biject_to
from numpyro.distributions.util import is_identically_one, sum_rightmost
from numpyro.handlers import condition, replay, seed, substitute, trace
from numpyro.infer.initialization import init_to_uniform, init_to_value
from numpyro.util import (
    _validate_model,
    find_stack_level,
    not_jax_tracer,
    soft_vmap,
    while_loop,
)

__all__ = [
    "find_valid_initial_params",
    "get_potential_fn",
    "log_density",
    "log_likelihood",
    "potential_energy",
    "initialize_model",
    "Predictive",
]

ModelInfo = namedtuple(
    "ModelInfo", ["param_info", "potential_fn", "postprocess_fn", "model_trace"]
)
ParamInfo = namedtuple("ParamInfo", ["z", "potential_energy", "z_grad"])

def log_density(model, model_args, model_kwargs, params):
    """
    (EXPERIMENTAL INTERFACE) Computes log of joint density for the model given
    latent values ``params``.

    :param model: Python callable containing NumPyro primitives.
    :param tuple model_args: args provided to the model.
    :param dict model_kwargs: kwargs provided to the model.
    :param dict params: dictionary of current parameter values keyed by site
        name.
    :return: log of joint density and a corresponding model trace
    """
    model = substitute(model, data=params)
    model_trace = trace(model).get_trace(*model_args, **model_kwargs)
    log_joint = jnp.zeros(())
    for site in model_trace.values():
        if site["type"] == "sample":
            value = site["value"]
            intermediates = site["intermediates"]
            scale = site["scale"]
            if intermediates:
                log_prob = site["fn"].log_prob(value, intermediates)
            else:
                guide_shape = jnp.shape(value)
                model_shape = tuple(
                    site["fn"].shape()
                )  # TensorShape from tfp needs casting to tuple
                try:
                    broadcast_shapes(guide_shape, model_shape)
                except ValueError:
                    raise ValueError(
                        "Model and guide shapes disagree at site: '{}': {} vs {}".format(
                            site["name"], model_shape, guide_shape
                        )
                    )
                log_prob = site["fn"].log_prob(value)

            if (scale is not None) and (not is_identically_one(scale)):
                log_prob = scale * log_prob

            log_prob = jnp.sum(log_prob)
            log_joint = log_joint + log_prob
    return log_joint, model_trace

class _without_rsample_stop_gradient(numpyro.primitives.Messenger):
    """
    Stop gradient for samples at latent sample sites for which has_rsample=False.
    """

    def postprocess_message(self, msg):
        if (
            msg["type"] == "sample"
            and (not msg["is_observed"])
            and (not msg["fn"].has_rsample)
        ):
            msg["value"] = lax.stop_gradient(msg["value"])
            # TODO: reconsider this logic
            # here we clear all the cached value so that gradients of log_prob(value) w.r.t.
            # all parameters of the transformed distributions match the behavior of
            # TransformedDistribution(d, transform) in Pyro with transform.cache_size == 0
            msg["intermediates"] = None

def get_importance_trace(model, guide, args, kwargs, params):
    """
    (EXPERIMENTAL) Returns traces from the guide and the model that is run against it.
    The returned traces also store the log probability at each site.

    .. note:: Gradients are blocked at latent sites which do not have reparametrized samplers.
    """
    guide = substitute(guide, data=params)
    with _without_rsample_stop_gradient():
        guide_trace = trace(guide).get_trace(*args, **kwargs)
    model = substitute(replay(model, guide_trace), data=params)
    model_trace = trace(model).get_trace(*args, **kwargs)
    for tr in (guide_trace, model_trace):
        for site in tr.values():
            if site["type"] == "sample":
                if "log_prob" not in site:
                    value = site["value"]
                    intermediates = site["intermediates"]
                    scale = site["scale"]
                    if intermediates:
                        log_prob = site["fn"].log_prob(value, intermediates)
                    else:
                        log_prob = site["fn"].log_prob(value)

                    if (scale is not None) and (not is_identically_one(scale)):
                        log_prob = scale * log_prob
                    site["log_prob"] = log_prob
    return model_trace, guide_trace

def transform_fn(transforms, params, invert=False):
    """
    (EXPERIMENTAL INTERFACE) Callable that applies a transformation from the `transforms`
    dict to values in the `params` dict and returns the transformed values keyed on
    the same names.

    :param transforms: Dictionary of transforms keyed by names. Names in
        `transforms` and `params` should align.
    :param params: Dictionary of arrays keyed by names.
    :param invert: Whether to apply the inverse of the transforms.
    :return: `dict` of transformed params.
    """
    if invert:
        transforms = {k: v.inv for k, v in transforms.items()}
    return {k: transforms[k](v) if k in transforms else v for k, v in params.items()}

def constrain_fn(model, model_args, model_kwargs, params, return_deterministic=False):
    """
    (EXPERIMENTAL INTERFACE) Gets value at each latent site in `model` given
    unconstrained parameters `params`. The `transforms` is used to transform these
    unconstrained parameters to base values of the corresponding priors in `model`.
    If a prior is a transformed distribution, the corresponding base value lies in
    the support of base distribution. Otherwise, the base value lies in the support
    of the distribution.

    :param model: a callable containing NumPyro primitives.
    :param tuple model_args: args provided to the model.
    :param dict model_kwargs: kwargs provided to the model.
    :param dict params: dictionary of unconstrained values keyed by site
        names.
    :param bool return_deterministic: whether to return the value of `deterministic`
        sites from the model. Defaults to `False`.
    :return: `dict` of transformed params.
    """

    def substitute_fn(site):
        if site["name"] in params:
            if site["type"] == "sample":
                with helpful_support_errors(site):
                    return biject_to(site["fn"].support)(params[site["name"]])
            else:
                return params[site["name"]]

    substituted_model = substitute(model, substitute_fn=substitute_fn)
    model_trace = trace(substituted_model).get_trace(*model_args, **model_kwargs)
    return {
        k: v["value"]
        for k, v in model_trace.items()
        if (k in params) or (return_deterministic and (v["type"] == "deterministic"))
    }

def _unconstrain_reparam(params, site):
    name = site["name"]
    if name in params:
        p = params[name]
        support = site["fn"].support
        with helpful_support_errors(site):
            t = biject_to(support)
        # in scan, we might only want to substitute an item at index i, rather than the whole sequence
        i = site["infer"].get("_scan_current_index", None)
        if i is not None:
            event_dim_shift = t.codomain.event_dim - t.domain.event_dim
            expected_unconstrained_dim = len(site["fn"].shape()) - event_dim_shift
            # check if p has additional time dimension
            if jnp.ndim(p) > expected_unconstrained_dim:
                p = p[i]

        if support in [constraints.real, constraints.real_vector]:
            return p
        value = t(p)

        log_det = t.log_abs_det_jacobian(p, value)
        log_det = sum_rightmost(
            log_det, jnp.ndim(log_det) - jnp.ndim(value) + len(site["fn"].event_shape)
        )
        if site["scale"] is not None:
            log_det = site["scale"] * log_det
        numpyro.factor("_{}_log_det".format(name), log_det)
        return value

def potential_energy(model, model_args, model_kwargs, params, enum=False):
    """
    (EXPERIMENTAL INTERFACE) Computes potential energy of a model given unconstrained params.
    Under the hood, we will transform these unconstrained parameters to the values
    belong to the supports of the corresponding priors in `model`.

    :param model: a callable containing NumPyro primitives.
    :param tuple model_args: args provided to the model.
    :param dict model_kwargs: kwargs provided to the model.
    :param dict params: unconstrained parameters of `model`.
    :param bool enum: whether to enumerate over discrete latent sites.
    :return: potential energy given unconstrained parameters.
    """
    if enum:
        from numpyro.contrib.funsor import log_density as log_density_
    else:
        log_density_ = log_density

    substituted_model = substitute(
        model, substitute_fn=partial(_unconstrain_reparam, params)
    )
    # no param is needed for log_density computation because we already substitute
    log_joint, model_trace = log_density_(
        substituted_model, model_args, model_kwargs, {}
    )
    return -log_joint

def _init_to_unconstrained_value(site=None, values={}):
    if site is None:
        return partial(_init_to_unconstrained_value, values=values)

def find_valid_initial_params(
    rng_key,
    model,
    *,
    init_strategy=init_to_uniform,
    enum=False,
    model_args=(),
    model_kwargs=None,
    prototype_params=None,
    forward_mode_differentiation=False,
    validate_grad=True,
):
    """
    (EXPERIMENTAL INTERFACE) Given a model with Pyro primitives, returns an initial
    valid unconstrained value for all the parameters. This function also returns
    the corresponding potential energy, the gradients, and an
    `is_valid` flag to say whether the initial parameters are valid. Parameter values
    are considered valid if the values and the gradients for the log density have
    finite values.

    :param jax.random.PRNGKey rng_key: random number generator seed to
        sample from the prior. The returned `init_params` will have the
        batch shape ``rng_key.shape[:-1]``.
    :param model: Python callable containing Pyro primitives.
    :param callable init_strategy: a per-site initialization function.
    :param bool enum: whether to enumerate over discrete latent sites.
    :param tuple model_args: args provided to the model.
    :param dict model_kwargs: kwargs provided to the model.
    :param dict prototype_params: an optional prototype parameters, which is used
        to define the shape for initial parameters.
    :param bool forward_mode_differentiation: whether to use forward-mode differentiation
        or reverse-mode differentiation. Defaults to False.
    :param bool validate_grad: whether to validate gradient of the initial params.
        Defaults to True.
    :return: tuple of `init_params_info` and `is_valid`, where `init_params_info` is the tuple
        containing the initial params, their potential energy, and their gradients.
    """
    model_kwargs = {} if model_kwargs is None else model_kwargs
    init_strategy = (
        init_strategy if isinstance(init_strategy, partial) else init_strategy()
    )
    # handle those init strategies differently to save computation
    if init_strategy.func is init_to_uniform:
        radius = init_strategy.keywords.get("radius")
        init_values = {}
    elif init_strategy.func is _init_to_unconstrained_value:
        radius = 2
        init_values = init_strategy.keywords.get("values")
    else:
        radius = None

    def cond_fn(state):
        i, _, _, is_valid = state
        return (i < 100) & (~is_valid)

    def body_fn(state):
        i, key, _, _ = state
        key, subkey = random.split(key)

        if radius is None or prototype_params is None:
            # XXX: we don't want to apply enum to draw latent samples
            model_ = model
            if enum:
                from numpyro.contrib.funsor import enum as enum_handler

                if isinstance(model, substitute) and isinstance(model.fn, enum_handler):
                    model_ = substitute(model.fn.fn, data=model.data)
                elif isinstance(model, enum_handler):
                    model_ = model.fn

            # Wrap model in a `substitute` handler to initialize from `init_loc_fn`.
            seeded_model = substitute(seed(model_, subkey), substitute_fn=init_strategy)
            model_trace = trace(seeded_model).get_trace(*model_args, **model_kwargs)
            constrained_values, inv_transforms = {}, {}
            for k, v in model_trace.items():
                if (
                    v["type"] == "sample"
                    and not v["is_observed"]
                    and not v["fn"].support.is_discrete
                ):
                    constrained_values[k] = v["value"]
                    with helpful_support_errors(v):
                        inv_transforms[k] = biject_to(v["fn"].support)
            params = transform_fn(
                inv_transforms,
                {k: v for k, v in constrained_values.items()},
                invert=True,
            )
        else:  # this branch doesn't require tracing the model
            params = {}
            for k, v in prototype_params.items():
                if k in init_values:
                    params[k] = init_values[k]
                else:
                    params[k] = random.uniform(
                        subkey, jnp.shape(v), minval=-radius, maxval=radius
                    )
                    key, subkey = random.split(key)

        potential_fn = partial(
            potential_energy, model, model_args, model_kwargs, enum=enum
        )
        if validate_grad:
            if forward_mode_differentiation:
                pe = potential_fn(params)
                z_grad = jacfwd(potential_fn)(params)
            else:
                pe, z_grad = value_and_grad(potential_fn)(params)
            z_grad_flat = ravel_pytree(z_grad)[0]
            is_valid = jnp.isfinite(pe) & jnp.all(jnp.isfinite(z_grad_flat))
        else:
            pe = potential_fn(params)
            is_valid = jnp.isfinite(pe)
            z_grad = None

        return i + 1, key, (params, pe, z_grad), is_valid

    def _find_valid_params(rng_key, exit_early=False):
        init_state = (0, rng_key, (prototype_params, 0.0, prototype_params), False)
        if exit_early and not_jax_tracer(rng_key):
            # Early return if valid params found. This is only helpful for single chain,
            # where we can avoid compiling body_fn in while_loop.
            _, _, (init_params, pe, z_grad), is_valid = init_state = body_fn(init_state)
            if not_jax_tracer(is_valid):
                if device_get(is_valid):
                    return (init_params, pe, z_grad), is_valid

        # XXX: this requires compiling the model, so for multi-chain, we trace the model 2-times
        # even if the init_state is a valid result
        _, _, (init_params, pe, z_grad), is_valid = while_loop(
            cond_fn, body_fn, init_state
        )
        return (init_params, pe, z_grad), is_valid

    # Handle possible vectorization
    if rng_key.ndim == 1:
        (init_params, pe, z_grad), is_valid = _find_valid_params(
            rng_key, exit_early=True
        )
    else:
        (init_params, pe, z_grad), is_valid = lax.map(_find_valid_params, rng_key)
    return (init_params, pe, z_grad), is_valid

def _get_model_transforms(model, model_args=(), model_kwargs=None):
    model_kwargs = {} if model_kwargs is None else model_kwargs
    model_trace = trace(model).get_trace(*model_args, **model_kwargs)
    inv_transforms = {}
    # model code may need to be replayed in the presence of deterministic sites
    replay_model = False
    has_enumerate_support = False
    for k, v in model_trace.items():
        if v["type"] == "sample" and not v["is_observed"]:
            if v["fn"].support.is_discrete:
                enum_type = v["infer"].get("enumerate")
                if enum_type is not None and (enum_type != "parallel"):
                    raise RuntimeError(
                        "This algorithm might only work for discrete sites with"
                        f" enumerate marked 'parallel'. But the site {k} is marked"
                        f" as '{enum_type}'."
                    )
                has_enumerate_support = True
                if not v["fn"].has_enumerate_support:
                    dist_name = type(v["fn"]).__name__
                    raise RuntimeError(
                        "This algorithm might only work for discrete sites with"
                        f" enumerate support. But the {dist_name} distribution at"
                        f" site {k} does not have enumerate support."
                    )
                if enum_type is None:
                    warnings.warn(
                        "Some algorithms will automatically enumerate the discrete"
                        f" latent site {k} of your model. In the future,"
                        " enumerated sites need to be marked with"
                        " `infer={'enumerate': 'parallel'}`.",
                        FutureWarning,
                        stacklevel=find_stack_level(),
                    )
            else:
                support = v["fn"].support
                with helpful_support_errors(v, raise_warnings=True):
                    inv_transforms[k] = biject_to(support)
                # XXX: the following code filters out most situations with dynamic supports
                args = ()
                if isinstance(support, constraints._GreaterThan):
                    args = ("lower_bound",)
                elif isinstance(support, constraints._Interval):
                    args = ("lower_bound", "upper_bound")
                for arg in args:
                    if not isinstance(getattr(support, arg), (int, float)):
                        replay_model = True
        elif v["type"] == "deterministic":
            replay_model = True
    return inv_transforms, replay_model, has_enumerate_support, model_trace

def _partial_args_kwargs(fn, *args, **kwargs):
    """Returns a partial function of `fn` and args, kwargs."""
    return partial(fn, args, kwargs)

def _drop_args_kwargs(fn, *args, **kwargs):
    """Returns the input function `fn`, ignoring args and kwargs."""
    return fn

def get_potential_fn(
    model,
    inv_transforms,
    *,
    enum=False,
    replay_model=False,
    dynamic_args=False,
    model_args=(),
    model_kwargs=None,
):
    """
    (EXPERIMENTAL INTERFACE) Given a model with Pyro primitives, returns a
    function which, given unconstrained parameters, evaluates the potential
    energy (negative log joint density). In addition, this returns a
    function to transform unconstrained values at sample sites to constrained
    values within their respective support.

    :param model: Python callable containing Pyro primitives.
    :param dict inv_transforms: dictionary of transforms keyed by names.
    :param bool enum: whether to enumerate over discrete latent sites.
    :param bool replay_model: whether we need to replay model in
        `postprocess_fn` to obtain `deterministic` sites.
    :param bool dynamic_args: if `True`, the `potential_fn` and
        `constraints_fn` are themselves dependent on model arguments.
        When provided a `*model_args, **model_kwargs`, they return
        `potential_fn` and `constraints_fn` callables, respectively.
    :param tuple model_args: args provided to the model.
    :param dict model_kwargs: kwargs provided to the model.
    :return: tuple of (`potential_fn`, `postprocess_fn`). The latter is used
        to constrain unconstrained samples (e.g. those returned by HMC)
        to values that lie within the site's support, and return values at
        `deterministic` sites in the model.
    """
    if dynamic_args:
        potential_fn = partial(
            _partial_args_kwargs, partial(potential_energy, model, enum=enum)
        )
        if replay_model:
            # XXX: we seed to sample discrete sites (but not collect them)
            model_ = seed(model.fn, 0) if enum else model
            postprocess_fn = partial(
                _partial_args_kwargs,
                partial(constrain_fn, model, return_deterministic=True),
            )
        else:
            postprocess_fn = partial(
                _drop_args_kwargs, partial(transform_fn, inv_transforms)
            )
    else:
        model_kwargs = {} if model_kwargs is None else model_kwargs
        potential_fn = partial(
            potential_energy, model, model_args, model_kwargs, enum=enum
        )
        if replay_model:
            model_ = seed(model.fn, 0) if enum else model
            postprocess_fn = partial(
                constrain_fn,
                model_,
                model_args,
                model_kwargs,
                return_deterministic=True,
            )
        else:
            postprocess_fn = partial(transform_fn, inv_transforms)

    return potential_fn, postprocess_fn

def _guess_max_plate_nesting(model_trace):
    """
    Guesses max_plate_nesting by using model trace.
    This optimistically assumes static model
    structure.
    """
    sites = [site for site in model_trace.values() if site["type"] == "sample"]

    dims = [
        frame.dim
        for site in sites
        for frame in site["cond_indep_stack"]
        if frame.dim is not None
    ]
    max_plate_nesting = -min(dims) if dims else 0
    return max_plate_nesting

def initialize_model(
    rng_key,
    model,
    *,
    init_strategy=init_to_uniform,
    dynamic_args=False,
    model_args=(),
    model_kwargs=None,
    forward_mode_differentiation=False,
    validate_grad=True,
):
    """
    (EXPERIMENTAL INTERFACE) Helper function that calls :func:`~numpyro.infer.util.get_potential_fn`
    and :func:`~numpyro.infer.util.find_valid_initial_params` under the hood
    to return a tuple of (`init_params_info`, `potential_fn`, `postprocess_fn`, `model_trace`).

    :param jax.random.PRNGKey rng_key: random number generator seed to
        sample from the prior. The returned `init_params` will have the
        batch shape ``rng_key.shape[:-1]``.
    :param model: Python callable containing Pyro primitives.
    :param callable init_strategy: a per-site initialization function.
        See :ref:`init_strategy` section for available functions.
    :param bool dynamic_args: if `True`, the `potential_fn` and
        `constraints_fn` are themselves dependent on model arguments.
        When provided a `*model_args, **model_kwargs`, they return
        `potential_fn` and `constraints_fn` callables, respectively.
    :param tuple model_args: args provided to the model.
    :param dict model_kwargs: kwargs provided to the model.
    :param bool forward_mode_differentiation: whether to use forward-mode differentiation
        or reverse-mode differentiation. By default, we use reverse mode but the forward
        mode can be useful in some cases to improve the performance. In addition, some
        control flow utility on JAX such as `jax.lax.while_loop` or `jax.lax.fori_loop`
        only supports forward-mode differentiation. See
        `JAX's The Autodiff Cookbook <https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html>`_
        for more information.
    :param bool validate_grad: whether to validate gradient of the initial params.
        Defaults to True.
    :return: a namedtupe `ModelInfo` which contains the fields
        (`param_info`, `potential_fn`, `postprocess_fn`, `model_trace`), where
        `param_info` is a namedtuple `ParamInfo` containing values from the prior
        used to initiate MCMC, their corresponding potential energy, and their gradients;
        `postprocess_fn` is a callable that uses inverse transforms
        to convert unconstrained HMC samples to constrained values that
        lie within the site's support, in addition to returning values
        at `deterministic` sites in the model.
    """
    model_kwargs = {} if model_kwargs is None else model_kwargs
    substituted_model = substitute(
        seed(model, rng_key if jnp.ndim(rng_key) == 1 else rng_key[0]),
        substitute_fn=init_strategy,
    )
    (
        inv_transforms,
        replay_model,
        has_enumerate_support,
        model_trace,
    ) = _get_model_transforms(substituted_model, model_args, model_kwargs)
    # substitute param sites from model_trace to model so
    # we don't need to generate again parameters of `numpyro.module`
    model = substitute(
        model,
        data={
            k: site["value"]
            for k, site in model_trace.items()
            if site["type"] in ["param"]
        },
    )
    constrained_values = {
        k: v["value"]
        for k, v in model_trace.items()
        if v["type"] == "sample"
        and not v["is_observed"]
        and not v["fn"].support.is_discrete
    }

    if has_enumerate_support:
        from numpyro.contrib.funsor import config_enumerate, enum

        if not isinstance(model, enum):
            max_plate_nesting = _guess_max_plate_nesting(model_trace)
            _validate_model(model_trace, plate_warning="error")
            model = enum(config_enumerate(model), -max_plate_nesting - 1)
    else:
        _validate_model(model_trace, plate_warning="loose")

    potential_fn, postprocess_fn = get_potential_fn(
        model,
        inv_transforms,
        replay_model=replay_model,
        enum=has_enumerate_support,
        dynamic_args=dynamic_args,
        model_args=model_args,
        model_kwargs=model_kwargs,
    )

    init_strategy = (
        init_strategy if isinstance(init_strategy, partial) else init_strategy()
    )
    if (init_strategy.func is init_to_value) and not replay_model:
        init_values = init_strategy.keywords.get("values")
        unconstrained_values = transform_fn(inv_transforms, init_values, invert=True)
        init_strategy = _init_to_unconstrained_value(values=unconstrained_values)
    prototype_params = transform_fn(inv_transforms, constrained_values, invert=True)
    (init_params, pe, grad), is_valid = find_valid_initial_params(
        rng_key,
        substitute(
            model,
            data={
                k: site["value"]
                for k, site in model_trace.items()
                if site["type"] in ["plate"]
            },
        ),
        init_strategy=init_strategy,
        enum=has_enumerate_support,
        model_args=model_args,
        model_kwargs=model_kwargs,
        prototype_params=prototype_params,
        forward_mode_differentiation=forward_mode_differentiation,
        validate_grad=validate_grad,
    )

    if not_jax_tracer(is_valid):
        if device_get(~jnp.all(is_valid)):
            with numpyro.validation_enabled(), trace() as tr:
                # validate parameters
                substituted_model(*model_args, **model_kwargs)
                # validate values
                for site in tr.values():
                    if site["type"] == "sample":
                        with warnings.catch_warnings(record=True) as ws:
                            site["fn"]._validate_sample(site["value"])
                        if len(ws) > 0:
                            for w in ws:
                                # at site information to the warning message
                                w.message.args = (
                                    "Site {}: {}".format(
                                        site["name"], w.message.args[0]
                                    ),
                                ) + w.message.args[1:]
                                warnings.showwarning(
                                    w.message,
                                    w.category,
                                    w.filename,
                                    w.lineno,
                                    file=w.file,
                                    line=w.line,
                                )
            raise RuntimeError(
                "Cannot find valid initial parameters. Please check your model again."
            )
    return ModelInfo(
        ParamInfo(init_params, pe, grad), potential_fn, postprocess_fn, model_trace
    )

def _predictive(
    rng_key,
    model,
    posterior_samples,
    batch_shape,
    return_sites=None,
    infer_discrete=False,
    parallel=True,
    model_args=(),
    model_kwargs={},
):
    masked_model = numpyro.handlers.mask(model, mask=False)
    if infer_discrete:
        # inspect the model to get some structure
        rng_key, subkey = random.split(rng_key)
        batch_ndim = len(batch_shape)
        prototype_sample = tree_map(
            lambda x: jnp.reshape(x, (-1,) + jnp.shape(x)[batch_ndim:])[0],
            posterior_samples,
        )
        prototype_trace = trace(
            seed(substitute(masked_model, prototype_sample), subkey)
        ).get_trace(*model_args, **model_kwargs)
        first_available_dim = -_guess_max_plate_nesting(prototype_trace) - 1

    def single_prediction(val):
        rng_key, samples = val
        if infer_discrete:
            from numpyro.contrib.funsor import config_enumerate
            from numpyro.contrib.funsor.discrete import _sample_posterior

            model_trace = prototype_trace
            temperature = 1
            pred_samples = _sample_posterior(
                config_enumerate(condition(model, samples)),
                first_available_dim,
                temperature,
                rng_key,
                *model_args,
                **model_kwargs,
            )
        else:
            model_trace = trace(
                seed(substitute(masked_model, samples), rng_key)
            ).get_trace(*model_args, **model_kwargs)
            pred_samples = {name: site["value"] for name, site in model_trace.items()}

        if return_sites is not None:
            if return_sites == "":
                sites = {
                    k for k, site in model_trace.items() if site["type"] != "plate"
                }
            else:
                sites = return_sites
        else:
            sites = {
                k
                for k, site in model_trace.items()
                if (site["type"] == "sample" and k not in samples)
                or (site["type"] == "deterministic")
            }
        return {name: value for name, value in pred_samples.items() if name in sites}

    num_samples = int(np.prod(batch_shape))
    if num_samples > 1:
        rng_key = random.split(rng_key, num_samples)
    rng_key = rng_key.reshape((*batch_shape, 2))
    chunk_size = num_samples if parallel else 1
    return soft_vmap(
        single_prediction, (rng_key, posterior_samples), len(batch_shape), chunk_size
    )

class Predictive(object):
    """
    This class is used to construct predictive distribution. The predictive distribution is obtained
    by running model conditioned on latent samples from `posterior_samples`.

    .. warning::
        The interface for the `Predictive` class is experimental, and
        might change in the future.

    :param model: Python callable containing Pyro primitives.
    :param dict posterior_samples: dictionary of samples from the posterior.
    :param callable guide: optional guide to get posterior samples of sites not present
        in `posterior_samples`.
    :param dict params: dictionary of values for param sites of model/guide.
    :param int num_samples: number of samples
    :param list return_sites: sites to return; by default only sample sites not present
        in `posterior_samples` are returned.
    :param bool infer_discrete: whether or not to sample discrete sites from the
        posterior, conditioned on observations and other latent values in
        ``posterior_samples``. Under the hood, those sites will be marked with
        ``site["infer"]["enumerate"] = "parallel"``. See how `infer_discrete` works at
        the `Pyro enumeration tutorial <https://pyro.ai/examples/enumeration.html>`_.
        Note that this requires ``funsor`` installation.
    :param bool parallel: whether to predict in parallel using JAX vectorized map :func:`jax.vmap`.
        Defaults to False.
    :param batch_ndims: the number of batch dimensions in posterior samples or parameters. If `None` defaults
        to 0 if guide is set (i.e. not `None`) and 1 otherwise. Usages for batched posterior samples:

        + set `batch_ndims=0` to get prediction for 1 single sample

        + set `batch_ndims=1` to get prediction for `posterior_samples`
          with shapes `(num_samples x ...)` (same as`batch_ndims=None` with `guide=None`)

        + set `batch_ndims=2` to get prediction for `posterior_samples`
          with shapes `(num_chains x N x ...)`. Note that if `num_samples`
          argument is not None, its value should be equal to `num_chains x N`.

        Usages for batched parameters:

        + set `batch_ndims=0` to get 1 sample from the guide and parameters (same as `batch_ndims=None` with guide)

        + set `batch_ndims=1` to get predictions from a one dimensional batch of the guide and parameters
          with shapes `(num_samples x batch_size x ...)`

    :return: dict of samples from the predictive distribution.

    **Example:**

    Given a model::

        def model(X, y=None):
            ...
            return numpyro.sample("obs", likelihood, obs=y)

    you can sample from the prior predictive::

        predictive = Predictive(model, num_samples=1000)
        y_pred = predictive(rng_key, X)["obs"]

    If you also have posterior samples, you can sample from the posterior predictive::

        predictive = Predictive(model, posterior_samples=posterior_samples)
        y_pred = predictive(rng_key, X)["obs"]

    See docstrings for :class:`~numpyro.infer.svi.SVI` and :class:`~numpyro.infer.mcmc.MCMCKernel`
    to see example code of this in context.
    """

    def __init__(
        self,
        model: Callable,
        posterior_samples: Optional[Dict] = None,
        *,
        guide: Optional[Callable] = None,
        params: Optional[Dict] = None,
        num_samples: Optional[int] = None,
        return_sites: Optional[List[str]] = None,
        infer_discrete: bool = False,
        parallel: bool = False,
        batch_ndims: Optional[int] = None,
    ):
        if posterior_samples is None and num_samples is None:
            raise ValueError(
                "Either posterior_samples or num_samples must be specified."
            )
        if posterior_samples is not None and guide is not None:
            raise ValueError(
                "Only one of guide or posterior_samples can be provided, not both."
            )

        batch_ndims = (
            batch_ndims if batch_ndims is not None else 1 if guide is None else 0
        )

        posterior_samples = {} if posterior_samples is None else posterior_samples

        prototype_site = batch_shape = batch_size = None
        for name, sample in posterior_samples.items():
            if batch_shape is not None and sample.shape[:batch_ndims] != batch_shape:
                raise ValueError(
                    f"Batch shapes at site {name} and {prototype_site} "
                    f"should be the same, but got "
                    f"{sample.shape[:batch_ndims]} and {batch_shape}"
                )
            else:
                prototype_site = name
                batch_shape = sample.shape[:batch_ndims]
                batch_size = int(np.prod(batch_shape))
                if (num_samples is not None) and (num_samples != batch_size):
                    warnings.warn(
                        "Sample's batch dimension size {} is different from the "
                        "provided {} num_samples argument. Defaulting to {}.".format(
                            batch_size, num_samples, batch_size
                        ),
                        UserWarning,
                        stacklevel=find_stack_level(),
                    )
                num_samples = batch_size

        if num_samples is None:
            raise ValueError(
                "No sample sites in posterior samples to infer `num_samples`."
            )

        if batch_shape is None:
            batch_shape = (1,) * (batch_ndims - 1) + (num_samples,)

        if return_sites is not None:
            assert isinstance(return_sites, (list, tuple, set))

        self.model = model
        self.posterior_samples = {} if posterior_samples is None else posterior_samples
        self.num_samples = num_samples
        self.guide = guide
        self.params = {} if params is None else params
        self.infer_discrete = infer_discrete
        self.return_sites = return_sites
        self.parallel = parallel
        self.batch_ndims = batch_ndims
        self._batch_shape = batch_shape

    def _call_with_params(self, rng_key, params, args, kwargs):
        posterior_samples = self.posterior_samples
        if self.guide is not None:
            rng_key, guide_rng_key = random.split(rng_key)
            # use return_sites='' as a special signal to return all sites
            guide = substitute(self.guide, params)
            posterior_samples = _predictive(
                guide_rng_key,
                guide,
                posterior_samples,
                self._batch_shape,
                return_sites="",
                parallel=self.parallel,
                model_args=args,
                model_kwargs=kwargs,
            )
        model = substitute(self.model, self.params)
        return _predictive(
            rng_key,
            model,
            posterior_samples,
            self._batch_shape,
            return_sites=self.return_sites,
            infer_discrete=self.infer_discrete,
            parallel=self.parallel,
            model_args=args,
            model_kwargs=kwargs,
        )

    def __call__(self, rng_key, *args, **kwargs):
        """
        Returns dict of samples from the predictive distribution. By default, only sample sites not
        contained in `posterior_samples` are returned. This can be modified by changing the
        `return_sites` keyword argument of this :class:`Predictive` instance.

        :param jax.random.PRNGKey rng_key: random key to draw samples.
        :param args: model arguments.
        :param kwargs: model kwargs.
        """
        if self.batch_ndims == 0 or self.params == {} or self.guide is None:
            return self._call_with_params(rng_key, self.params, args, kwargs)
        elif self.batch_ndims == 1:  # batch over parameters
            batch_size = jnp.shape(tree_flatten(self.params)[0][0])[0]
            rng_keys = random.split(rng_key, batch_size)
            return jax.vmap(
                partial(self._call_with_params, args=args, kwargs=kwargs),
                in_axes=0,
                out_axes=1,
            )(rng_keys, self.params)
        else:
            raise NotImplementedError

def log_likelihood(
    model, posterior_samples, *args, parallel=False, batch_ndims=1, **kwargs
):
    """
    (EXPERIMENTAL INTERFACE) Returns log likelihood at observation nodes of model,
    given samples of all latent variables.

    :param model: Python callable containing Pyro primitives.
    :param dict posterior_samples: dictionary of samples from the posterior.
    :param args: model arguments.
    :param batch_ndims: the number of batch dimensions in posterior samples. Some usages:

        + set `batch_ndims=0` to get log likelihoods for 1 single sample

        + set `batch_ndims=1` to get log likelihoods for `posterior_samples`
          with shapes `(num_samples x ...)`

        + set `batch_ndims=2` to get log likelihoods for `posterior_samples`
          with shapes `(num_chains x num_samples x ...)`

    :param kwargs: model kwargs.
    :return: dict of log likelihoods at observation sites.
    """

    def single_loglik(samples):
        substituted_model = (
            substitute(model, samples) if isinstance(samples, dict) else model
        )
        model_trace = trace(substituted_model).get_trace(*args, **kwargs)
        return {
            name: site["fn"].log_prob(site["value"])
            for name, site in model_trace.items()
            if site["type"] == "sample" and site["is_observed"]
        }

    prototype_site = batch_shape = None
    for name, sample in posterior_samples.items():
        if batch_shape is not None and jnp.shape(sample)[:batch_ndims] != batch_shape:
            raise ValueError(
                f"Batch shapes at site {name} and {prototype_site} "
                f"should be the same, but got "
                f"{sample.shape[:batch_ndims]} and {batch_shape}"
            )
        else:
            prototype_site = name
            batch_shape = jnp.shape(sample)[:batch_ndims]

    if batch_shape is None:  # posterior_samples is an empty dict
        batch_shape = (1,) * batch_ndims
        posterior_samples = np.zeros(batch_shape)

    batch_size = int(np.prod(batch_shape))
    chunk_size = batch_size if parallel else 1
    return soft_vmap(single_loglik, posterior_samples, len(batch_shape), chunk_size)

@contextmanager
def helpful_support_errors(site, raise_warnings=False):
    name = site["name"]
    support = getattr(site["fn"], "support", None)
    if isinstance(support, constraints.independent):
        support = support.base_constraint

    # Warnings
    if raise_warnings:
        if support is constraints.circular:
            msg = (
                f"Continuous inference poorly handles circular sample site '{name}'. "
                + "Consider using VonMises distribution together with "
                + "a reparameterizer, e.g. "
                + f"numpyro.handlers.reparam(config={{'{name}': CircularReparam()}})."
            )
            warnings.warn(msg, UserWarning, stacklevel=find_stack_level())

    # Exceptions
    try:
        yield
    except NotImplementedError as e:
        support_name = repr(support).lower()
        if "integer" in support_name or "boolean" in support_name:
            # TODO: mention enumeration when it is supported in SVI
            raise ValueError(
                f"Continuous inference cannot handle discrete sample site '{name}'."
            )
        if "sphere" in support_name:
            raise ValueError(
                f"Continuous inference cannot handle spherical sample site '{name}'. "
                "Consider using ProjectedNormal distribution together with "
                "a reparameterizer, e.g. "
                f"numpyro.handlers.reparam(config={{'{name}': ProjectedNormalReparam()}})."
            )
        raise e from None

import numpy,from numpy,import jax,from jax