<commit_before># -*- coding: UTF-8 -*-
"""Definitions for `Fitter` class."""
import io
import json
import os
import re
import shutil
import sys
import time
import warnings
from collections import OrderedDict
from copy import deepcopy
from difflib import get_close_matches

import dropbox
import numpy as np
import scipy
from astrocats.catalog.entry import ENTRY, Entry
from astrocats.catalog.model import MODEL
from astrocats.catalog.photometry import PHOTOMETRY
from astrocats.catalog.quantity import QUANTITY
from astrocats.catalog.realization import REALIZATION
from emcee.autocorr import AutocorrError
from schwimmbad import MPIPool, SerialPool
from six import string_types

from mosfit.mossampler import MOSSampler
from mosfit.printer import Printer
from mosfit.utils import (calculate_WAIC, entabbed_json_dump,
                          entabbed_json_dumps, flux_density_unit,
                          frequency_unit, get_model_hash, get_url_file_handle,
                          is_number, listify, pretty_num)

from .model import Model

warnings.filterwarnings("ignore")


def draw_walker(test=True):
    """Draw a walker from the global model variable."""
    global model
    return model.draw_walker(test)  # noqa: F821


def likelihood(x):
    """Return a likelihood score using the global model variable."""
    global model
    return model.likelihood(x)  # noqa: F821


def prior(x):
    """Return a prior score using the global model variable."""
    global model
    return model.prior(x)  # noqa: F821


def frack(x):
    """Frack at the specified parameter combination."""
    global model
    return model.frack(x)  # noqa: F821


class Fitter(object):
    """Fit transient events with the provided model."""

    _MAX_ACORC = 5

    def __init__(self):
        """Initialize `Fitter`."""
        pass

    def fit_events(self,
                   events=[],
                   models=[],
                   plot_points='',
                   max_time='',
                   band_list=[],
                   band_systems=[],
                   band_instruments=[],
                   band_bandsets=[],
                   iterations=1000,
                   num_walkers=50,
                   num_temps=1,
                   parameter_paths=[''],
                   fracking=True,
                   frack_step=50,
                   wrap_length=100,
                   travis=False,
                   burn=None,
                   post_burn=None,
                   gibbs=False,
                   smooth_times=-1,
                   extrapolate_time=0.0,
                   limit_fitting_mjds=False,
                   exclude_bands=[],
                   exclude_instruments=[],
                   suffix='',
                   offline=False,
                   upload=False,
                   write=False,
                   quiet=False,
                   upload_token='',
                   check_upload_quality=False,
                   printer=None,
                   variance_for_each=[],
                   user_fixed_parameters=[],
                   run_until_converged=False,
                   save_full_chain=False,
                   draw_above_likelihood=False,
                   maximum_walltime=False,
                   start_time=False,
                   print_trees=False,
                   maximum_memory=np.inf,
                   **kwargs):
        """Fit a list of events with a list of models."""
        if start_time is False:
            start_time = time.time()
        self._start_time = start_time
        self._maximum_walltime = maximum_walltime
        self._maximum_memory = maximum_memory
        self._debug = False

        dir_path = os.path.dirname(os.path.realpath(__file__))
        self._travis = travis
        self._wrap_length = wrap_length
        self._draw_above_likelihood = draw_above_likelihood

        if not printer:
            self._printer = Printer(wrap_length=wrap_length, quiet=quiet)
        else:
            self._printer = printer

        prt = self._printer

        event_list = listify(events)
        model_list = listify(models)

        if len(model_list) and not len(event_list):
            event_list = ['']

        event_list = [x.replace('‑', '-') for x in event_list]

        self._catalogs = {
            'OSC': ('https://sne.space/astrocats/astrocats/'
                    'supernovae/output'),
            'OTC': ('https://tde.space/astrocats/astrocats/'
                    'tidaldisruptions/output')
        }

        if not len(event_list) and not len(model_list):
            self._printer.wrapped(
                'No events or models specified, initializing and then '
                'exiting.', warning=True)

        entries = [[] for x in range(len(event_list))]
        ps = [[] for x in range(len(event_list))]
        lnprobs = [[] for x in range(len(event_list))]

        data = {}

        self._event_name = 'Batch'
        self._event_catalog = ''
        for ei, event in enumerate(event_list):
            self._event_name = ''
            self._event_path = ''
            if event:
                try:
                    pool = MPIPool()
                except Exception:
                    pool = SerialPool()
                if pool.is_master():
                    path = ''
                    # If the event name ends in .json, assume a path
                    if event.endswith('.json'):
                        path = event
                        self._event_name = event.replace('.json',
                                                         '').split('/')[-1]

                    # If not (or the file doesn't exist), download from OSC
                    if not path or not os.path.exists(path):
                        names_paths = [
                            os.path.join(dir_path, 'cache', x +
                                         '.names.min.json') for x in
                            self._catalogs]
                        input_name = event.replace('.json', '')
                        prt.wrapped(
                            'Event `{}` interpreted as transient '
                            'name, downloading lists of transient '
                            'aliases...'.format(input_name))
                        if not offline:
                            for ci, catalog in enumerate(self._catalogs):
                                try:
                                    response = get_url_file_handle(
                                        self._catalogs[catalog] +
                                        '/names.min.json',
                                        timeout=10)
                                except Exception:
                                    prt.wrapped(
                                        'Warning: Could not download {} '
                                        'names (are you online?), using '
                                        'cached list.'.format(catalog))
                                    raise
                                else:
                                    with open(names_paths[ci], 'wb') as f:
                                        shutil.copyfileobj(response, f)
                        names = OrderedDict()
                        for ci, catalog in enumerate(self._catalogs):
                            if os.path.exists(names_paths[ci]):
                                with open(names_paths[ci], 'r') as f:
                                    names[catalog] = json.load(
                                        f, object_pairs_hook=OrderedDict)
                            else:
                                self._printer.wrapped(
                                    'Warning: Could not read list of {} '
                                    'names!'.format(catalog), warning=True)
                                if offline:
                                    self._printer.wrapped(
                                        'Try omitting the `--offline` flag.')
                                raise RuntimeError

                            if event in names[catalog]:
                                self._event_name = event
                                self._event_catalog = catalog
                            else:
                                for name in names[catalog]:
                                    if (event in names[catalog][name] or
                                            'SN' + event in
                                            names[catalog][name]):
                                        self._event_name = name
                                        self._event_catalog = catalog
                                        break

                        if not self._event_name:
                            for ci, catalog in enumerate(self._catalogs):
                                namekeys = []
                                for name in names[catalog]:
                                    namekeys.extend(names[catalog][name])
                                namekeys = list(sorted(set(namekeys)))
                                matches = get_close_matches(
                                    event, namekeys, n=5, cutoff=0.8)
                                # matches = []
                                if len(matches) < 5 and is_number(event[0]):
                                    prt.wrapped(
                                        'Could not find event, performing '
                                        'extended name search...')
                                    snprefixes = set(('SN19', 'SN20'))
                                    for name in names[catalog]:
                                        ind = re.search("\d", name)
                                        if ind and ind.start() > 0:
                                            snprefixes.add(name[:ind.start()])
                                    snprefixes = list(sorted(snprefixes))
                                    for prefix in snprefixes:
                                        testname = prefix + event
                                        new_matches = get_close_matches(
                                            testname, namekeys, cutoff=0.95,
                                            n=1)
                                        if (len(new_matches) and
                                                new_matches[0] not in matches):
                                            matches.append(new_matches[0])
                                        if len(matches) == 5:
                                            break
                                if len(matches):
                                    if travis:
                                        response = matches[0]
                                    else:
                                        response = prt.prompt(
                                            'No exact match to given '
                                            'transient '
                                            'found. Did you mean one of the '
                                            'following transients?',
                                            kind='select',
                                            options=matches,
                                            none_string=(
                                                'None of the above, ' +
                                                ('skip this event.' if
                                                 ci == len(self._catalogs) - 1
                                                 else
                                                 'try the next catalog.')))
                                    if response:
                                        for name in names[catalog]:
                                            if response in names[
                                                    catalog][name]:
                                                self._event_name = name
                                                self._event_catalog = catalog
                                                break

                        if not self._event_name:
                            prt.wrapped(
                                'Could not find event by that name, skipping!')
                            continue
                        urlname = self._event_name + '.json'
                        name_path = os.path.join(dir_path, 'cache', urlname)

                        if not offline:
                            prt.wrapped(
                                'Found event by primary name `{}` in the {}, '
                                'downloading data...'.format(
                                    self._event_name, self._event_catalog))
                            try:
                                response = get_url_file_handle(
                                    self._catalogs[self._event_catalog] +
                                    '/json/' + urlname,
                                    timeout=10)
                            except Exception:
                                prt.wrapped(
                                    'Warning: Could not download data for '
                                    ' `{}`, '
                                    'will attempt to use cached data.'.format(
                                        self._event_name))
                            else:
                                with open(name_path, 'wb') as f:
                                    shutil.copyfileobj(response, f)
                        path = name_path

                    if os.path.exists(path):
                        with open(path, 'r') as f:
                            data = json.load(f, object_pairs_hook=OrderedDict)
                        prt.wrapped('Event file:')
                        prt.wrapped('  ' + path)
                    else:
                        prt.wrapped(
                            'Error: Could not find data for `{}` locally or '
                            'in the {}.'.format(
                                self._event_name,
                                '/'.join(self._catalogs.keys())))
                        if offline:
                            prt.wrapped(
                                'Try omitting the `--offline` flag.')
                        raise RuntimeError

                    for rank in range(1, pool.size + 1):
                        pool.comm.send(self._event_name, dest=rank, tag=0)
                        pool.comm.send(path, dest=rank, tag=1)
                        pool.comm.send(data, dest=rank, tag=2)
                else:
                    self._event_name = pool.comm.recv(source=0, tag=0)
                    path = pool.comm.recv(source=0, tag=1)
                    data = pool.comm.recv(source=0, tag=2)
                    pool.wait()

                self._event_path = path

                if pool.is_master():
                    pool.close()

            if model_list:
                lmodel_list = model_list
            else:
                lmodel_list = ['']

            entries[ei] = [None for y in range(len(lmodel_list))]
            ps[ei] = [None for y in range(len(lmodel_list))]
            lnprobs[ei] = [None for y in range(len(lmodel_list))]

            for mi, mod_name in enumerate(lmodel_list):
                for parameter_path in parameter_paths:
                    try:
                        pool = MPIPool()
                    except Exception:
                        pool = SerialPool()
                    self._model = Model(
                        model=mod_name,
                        data=data,
                        parameter_path=parameter_path,
                        wrap_length=wrap_length,
                        fitter=self,
                        pool=pool,
                        print_trees=print_trees)

                    if not self._model._model_name:
                        self._printer.wrapped(
                            'Skipping `{}`, no models available to fit the '
                            'transient.'.format(self._event_name),
                            warning=True)
                        continue

                    if not event:
                        self._printer.wrapped(
                            'No event specified, generating dummy data.')
                        self._event_name = mod_name
                        gen_args = {
                            'name': mod_name,
                            'max_time': max_time,
                            'plot_points': plot_points,
                            'band_list': band_list,
                            'band_systems': band_systems,
                            'band_instruments': band_instruments,
                            'band_bandsets': band_bandsets
                        }
                        data = self.generate_dummy_data(**gen_args)

                    success = self.load_data(
                        data,
                        event_name=self._event_name,
                        iterations=iterations,
                        fracking=fracking,
                        burn=burn,
                        post_burn=post_burn,
                        smooth_times=smooth_times,
                        extrapolate_time=extrapolate_time,
                        limit_fitting_mjds=limit_fitting_mjds,
                        exclude_bands=exclude_bands,
                        exclude_instruments=exclude_instruments,
                        band_list=band_list,
                        band_systems=band_systems,
                        band_instruments=band_instruments,
                        band_bandsets=band_bandsets,
                        variance_for_each=variance_for_each,
                        user_fixed_parameters=user_fixed_parameters,
                        pool=pool)

                    if success:
                        entry, p, lnprob = self.fit_data(
                            event_name=self._event_name,
                            iterations=iterations,
                            num_walkers=num_walkers,
                            num_temps=num_temps,
                            fracking=fracking,
                            frack_step=frack_step,
                            gibbs=gibbs,
                            pool=pool,
                            suffix=suffix,
                            write=write,
                            upload=upload,
                            upload_token=upload_token,
                            check_upload_quality=check_upload_quality,
                            run_until_converged=run_until_converged,
                            save_full_chain=save_full_chain)
                        entries[ei][mi] = deepcopy(entry)
                        ps[ei][mi] = deepcopy(p)
                        lnprobs[ei][mi] = deepcopy(lnprob)

                    if pool.is_master():
                        pool.close()

        return (entries, ps, lnprobs)

    def load_data(self,
                  data,
                  event_name='',
                  iterations=2000,
                  fracking=True,
                  burn=None,
                  post_burn=None,
                  smooth_times=-1,
                  extrapolate_time=0.0,
                  limit_fitting_mjds=False,
                  exclude_bands=[],
                  exclude_instruments=[],
                  band_list=[],
                  band_systems=[],
                  band_instruments=[],
                  band_bandsets=[],
                  variance_for_each=[],
                  user_fixed_parameters=[],
                  pool=''):
        """Load the data for the specified event."""
        prt = self._printer
        fixed_parameters = []
        for task in self._model._call_stack:
            cur_task = self._model._call_stack[task]
            self._model._modules[task].set_event_name(event_name)
            if cur_task['kind'] == 'data':
                success = self._model._modules[task].set_data(
                    data,
                    req_key_values={'band': self._model._bands},
                    subtract_minimum_keys=['times'],
                    smooth_times=smooth_times,
                    extrapolate_time=extrapolate_time,
                    limit_fitting_mjds=limit_fitting_mjds,
                    exclude_bands=exclude_bands,
                    exclude_instruments=exclude_instruments,
                    band_list=band_list,
                    band_systems=band_systems,
                    band_instruments=band_instruments,
                    band_bandsets=band_bandsets)
                if not success:
                    return False
                fixed_parameters.extend(self._model._modules[task]
                                        .get_data_determined_parameters())

            # Fix user-specified parameters.
            for fi, param in enumerate(user_fixed_parameters):
                if (task == param or
                        self._model._call_stack[task].get(
                            'class', '') == param):
                    fixed_parameters.append(task)
                    if fi < len(user_fixed_parameters) - 1 and is_number(
                            user_fixed_parameters[fi + 1]):
                        value = float(user_fixed_parameters[fi + 1])
                        if value not in self._model._call_stack:
                            self._model._call_stack[task]['value'] = value
                    if 'min_value' in self._model._call_stack[task]:
                        del self._model._call_stack[task]['min_value']
                    if 'max_value' in self._model._call_stack[task]:
                        del self._model._call_stack[task]['max_value']
                    self._model._modules[task].fix_value(
                        self._model._call_stack[task]['value'])

        self._model.determine_free_parameters(fixed_parameters)

        self._model.exchange_requests()

        # Run through once to set all inits.
        for root in ['output', 'objective']:
            outputs = self._model.run_stack(
                [0.0 for x in range(self._model._num_free_parameters)],
                root=root)

        # Create any data-dependent free parameters.
        self._model.create_data_dependent_free_parameters(
            variance_for_each, outputs)

        # Determine free parameters again as above may have changed them.
        self._model.determine_free_parameters(fixed_parameters)

        self._model.determine_number_of_measurements()

        self._model.exchange_requests()

        # Reset modules
        for task in self._model._call_stack:
            self._model._modules[task].reset_preprocessed()

        # Run through inits once more.
        for root in ['output', 'objective']:
            outputs = self._model.run_stack(
                [0.0 for x in range(self._model._num_free_parameters)],
                root=root)

        # Collect observed band info
        if pool.is_master() and 'photometry' in self._model._modules:
            prt.wrapped('Bands being used for current transient:')
            bis = list(
                filter(lambda a: a != -1,
                       sorted(set(outputs['all_band_indices']))))
            ois = []
            for bi in bis:
                ois.append(
                    any([
                        y
                        for x, y in zip(outputs['all_band_indices'], outputs[
                            'observed']) if x == bi
                    ]))
            band_len = max([
                len(self._model._modules['photometry']._unique_bands[bi][
                    'SVO']) for bi in bis
            ])
            filts = self._model._modules['photometry']
            ubs = filts._unique_bands
            filterarr = [(ubs[bis[i]]['systems'], ubs[bis[i]]['bandsets'],
                          filts._average_wavelengths[bis[i]],
                          filts._band_offsets[bis[i]], ois[i], bis[i])
                         for i in range(len(bis))]
            filterrows = [(
                ' ' + (' ' if s[-2] else '*') + ubs[s[-1]]['SVO']
                .ljust(band_len) + ' [' + ', '.join(
                    list(
                        filter(None, (
                            'Bandset: ' + s[1] if s[1] else '',
                            'System: ' + s[0] if s[0] else '',
                            'AB offset: ' + pretty_num(s[3]))))) +
                ']').replace(' []', '') for s in list(sorted(filterarr))]
            if not all(ois):
                filterrows.append('  (* = Not observed in this band)')
            self._printer.prt('\n'.join(filterrows))

        self._event_name = event_name
        self._emcee_est_t = 0.0
        self._bh_est_t = 0.0
        self._fracking = fracking
        if burn is not None:
            self._burn_in = burn
            self._post_burn = max(iterations - burn, 0)
        elif post_burn is not None:
            self._post_burn = post_burn
            self._burn_in = max(iterations - post_burn, 0)
        else:
            self._burn_in = int(np.round(iterations / 2))
            self._post_burn = max(iterations - self._burn_in, 0)

        return True

    def fit_data(self,
                 event_name='',
                 iterations=2000,
                 frack_step=20,
                 num_walkers=50,
                 num_temps=1,
                 fracking=True,
                 gibbs=False,
                 pool='',
                 suffix='',
                 write=False,
                 upload=False,
                 upload_token='',
                 check_upload_quality=True,
                 run_until_converged=False,
                 save_full_chain=False):
        """Fit the data for a given event.

        Fitting performed using a combination of emcee and fracking.
        """
        from mosfit.__init__ import __version__
        global model
        model = self._model
        prt = self._printer

        upload_this = upload and iterations > 0

        if not pool.is_master():
            try:
                pool.wait()
            except (KeyboardInterrupt, SystemExit):
                pass
            return (None, None, None)

        ntemps, ndim, nwalkers = (num_temps, model._num_free_parameters,
                                  num_walkers)

        test_walker = iterations > 0
        lnprob = None
        lnlike = None
        pool_size = max(pool.size, 1)
        # Derived so only half a walker redrawn with Gaussian distribution.
        redraw_mult = 1.0 * np.sqrt(
            2) * scipy.special.erfinv(float(nwalkers - 1) / nwalkers)

        self._printer.prt(
            '{} measurements, {} free parameters.'.format(
                model._num_measurements, ndim))
        if model._num_measurements <= ndim:
            self._printer.wrapped(
                'Warning: Number of free parameters exceeds number of '
                'measurements. Please treat results with caution.',
                warning=True)
        if nwalkers < 10 * ndim:
            self._printer.wrapped(
                'Warning: While emcee accepts a number of walkers that is '
                'twice the number of dimensions or greater, simple tests '
                'with toy problems show poor convergence with this minimum. '
                'We suggest using at least 10x the number of dimensions '
                '(`-N {}` suggested for this fit), and preferably more, if '
                'feasible.'.format(10 * ndim),
                warning=True)
        self._printer.prt('\n\n')
        p0 = [[] for x in range(ntemps)]

        for i, pt in enumerate(p0):
            dwscores = []
            while len(p0[i]) <= nwalkers:
                prt.status(
                    self,
                    desc='Drawing initial walkers',
                    progress=[
                        i * nwalkers + len(p0[i]) + 1, nwalkers * ntemps])
                if len(p0[i]) == nwalkers:
                    break

                if pool.size == 0:
                    p, score = draw_walker(test_walker)
                    p0[i].append(p)
                    dwscores.append(score)
                else:
                    nmap = nwalkers - len(p0[i])
                    dws = pool.map(draw_walker, [test_walker] * nmap)
                    p0[i].extend([x[0] for x in dws])
                    dwscores.extend([x[1] for x in dws])

                if self._draw_above_likelihood is not False:
                    self._draw_above_likelihood = np.mean(dwscores)

        sampler = MOSSampler(
            ntemps, nwalkers, ndim, likelihood, prior, pool=pool)

        prt.inline('Initial draws completed!')
        self._printer.prt('\n\n')
        p = list(p0)

        sli = 1.0  # Keep track of how many times chain halved
        emi = 0
        tft = 0.0  # Total fracking time
        acor = None
        aacort = -1
        aa = 0
        psrf = np.inf
        s_exception = None
        all_chain = np.array([])

        try:
            st = time.time()

            max_chunk = 1000
            iter_chunks = int(np.ceil(float(iterations) / max_chunk))
            iter_arr = [max_chunk if xi < iter_chunks - 1 else
                        iterations - max_chunk * (iter_chunks - 1)
                        for xi, x in enumerate(range(iter_chunks))]
            # Make sure a chunk separation is located at self._burn_in
            chunk_is = sorted(set(
                np.concatenate(([0, self._burn_in], np.cumsum(iter_arr)))))
            iter_arr = np.diff(chunk_is)

            # The argument of the for loop runs emcee, after each iteration of
            # emcee the contents of the for loop are executed.
            converged = False
            exceeded_walltime = False
            ici = 0
            while run_until_converged or ici < len(iter_arr):
                slr = int(np.round(sli))
                ic = max_chunk if run_until_converged else iter_arr[ici]
                if exceeded_walltime:
                    break
                if run_until_converged and converged and emi > iterations:
                    break
                for li, (
                        p, lnprob, lnlike) in enumerate(
                            sampler.sample(
                                p, iterations=ic, gibbs=gibbs if
                                emi >= self._burn_in else True)):
                    if (self._maximum_walltime is not False and
                            time.time() - self._start_time >
                            self._maximum_walltime):
                        self._printer.wrapped(
                            'Walltime exceeded, exiting...', warning=True)
                        exceeded_walltime = True
                        break
                    emi = emi + 1
                    emim1 = emi - 1
                    messages = []

                    # Record then reset sampler proposal/acceptance counts.
                    accepts = list(
                        np.mean(sampler.nprop_accepted / sampler.nprop,
                                axis=1))
                    sampler.nprop = np.zeros(
                        (sampler.ntemps, sampler.nwalkers), dtype=np.float)
                    sampler.nprop_accepted = np.zeros(
                        (sampler.ntemps, sampler.nwalkers),
                        dtype=np.float)

                    # First, redraw any walkers with scores significantly
                    # worse than their peers (only during burn-in).
                    if emim1 <= self._burn_in:
                        pmedian = [np.median(x) for x in lnprob]
                        pmead = [np.mean([abs(y - pmedian) for y in x])
                                 for x in lnprob]
                        redraw_count = 0
                        bad_redraws = 0
                        for ti, tprob in enumerate(lnprob):
                            for wi, wprob in enumerate(tprob):
                                if (wprob <= pmedian[ti] -
                                    max(redraw_mult * pmead[ti],
                                        float(nwalkers)) or
                                        np.isnan(wprob)):
                                    redraw_count = redraw_count + 1
                                    dxx = np.random.normal(
                                        scale=0.01, size=ndim)
                                    tar_x = np.array(
                                        p[np.random.randint(ntemps)][
                                            np.random.randint(nwalkers)])
                                    new_x = np.clip(tar_x + dxx, 0.0, 1.0)
                                    new_like = likelihood(new_x)
                                    new_prob = new_like + prior(new_x)
                                    if new_prob > wprob or np.isnan(wprob):
                                        p[ti][wi] = new_x
                                        lnlike[ti][wi] = new_like
                                        lnprob[ti][wi] = new_prob
                                    else:
                                        bad_redraws = bad_redraws + 1
                        if redraw_count > 0:
                            messages.append(
                                '{:.0%} redraw, {}/{} success'.format(
                                    redraw_count / (nwalkers * ntemps),
                                    redraw_count - bad_redraws, redraw_count))

                    # Calculate the autocorrelation time.
                    low = 10
                    asize = 0.5 * (emim1 - self._burn_in) / low
                    if asize >= 0:
                        acorc = max(
                            1, min(self._MAX_ACORC,
                                   int(np.floor(0.5 * emi / low))))
                        aacort = -1.0
                        aa = 0
                        ams = self._burn_in
                        cur_chain = (np.concatenate(
                            (all_chain, sampler.chain[:, :, :li + 1:slr, :]),
                            axis=2) if len(all_chain) else
                            sampler.chain[:, :, :li + 1:slr, :])
                        for a in range(acorc, 1, -1):
                            ms = self._burn_in
                            if ms >= emi - low:
                                break
                            try:
                                acorts = sampler.get_autocorr_time(
                                    chain=cur_chain, low=low, c=a,
                                    min_step=int(np.round(float(ms) / sli)),
                                    max_walkers=5, fast=True)
                                acort = max([
                                    max(x)
                                    for x in acorts
                                ])
                            except AutocorrError:
                                continue
                            else:
                                aa = a
                                aacort = acort * slr
                                ams = ms
                                break
                        acor = [aacort, aa, ams]

                    # Calculate the PSRF (Gelman-Rubin statistic).
                    if li > 1 and emi > self._burn_in + 2:
                        cur_chain = (np.concatenate(
                            (all_chain, sampler.chain[:, :, :li + 1:slr, :]),
                            axis=2) if len(all_chain) else
                            sampler.chain[:, :, :li + 1:slr, :])
                        vws = np.zeros((ntemps, ndim))
                        for ti in range(ntemps):
                            for xi in range(ndim):
                                vchain = cur_chain[
                                    ti, :, int(np.round(
                                        self._burn_in / sli)):, xi]
                                m = len(vchain)
                                n = len(vchain[0])
                                mom = np.mean(np.mean(vchain, axis=1))
                                b = n / float(m - 1) * np.sum(<commit_msg>MAINT: slightly more accurate acort<commit_after># -*- coding: UTF-8 -*-
"""Definitions for `Fitter` class."""
import io
import json
import os
import re
import shutil
import sys
import time
import warnings
from collections import OrderedDict
from copy import deepcopy
from difflib import get_close_matches

import dropbox
import numpy as np
import scipy
from astrocats.catalog.entry import ENTRY, Entry
from astrocats.catalog.model import MODEL
from astrocats.catalog.photometry import PHOTOMETRY
from astrocats.catalog.quantity import QUANTITY
from astrocats.catalog.realization import REALIZATION
from emcee.autocorr import AutocorrError
from schwimmbad import MPIPool, SerialPool
from six import string_types

from mosfit.mossampler import MOSSampler
from mosfit.printer import Printer
from mosfit.utils import (calculate_WAIC, entabbed_json_dump,
                          entabbed_json_dumps, flux_density_unit,
                          frequency_unit, get_model_hash, get_url_file_handle,
                          is_number, listify, pretty_num)

from .model import Model

warnings.filterwarnings("ignore")


def draw_walker(test=True):
    """Draw a walker from the global model variable."""
    global model
    return model.draw_walker(test)  # noqa: F821


def likelihood(x):
    """Return a likelihood score using the global model variable."""
    global model
    return model.likelihood(x)  # noqa: F821


def prior(x):
    """Return a prior score using the global model variable."""
    global model
    return model.prior(x)  # noqa: F821


def frack(x):
    """Frack at the specified parameter combination."""
    global model
    return model.frack(x)  # noqa: F821


class Fitter(object):
    """Fit transient events with the provided model."""

    _MAX_ACORC = 5

    def __init__(self):
        """Initialize `Fitter`."""
        pass

    def fit_events(self,
                   events=[],
                   models=[],
                   plot_points='',
                   max_time='',
                   band_list=[],
                   band_systems=[],
                   band_instruments=[],
                   band_bandsets=[],
                   iterations=1000,
                   num_walkers=50,
                   num_temps=1,
                   parameter_paths=[''],
                   fracking=True,
                   frack_step=50,
                   wrap_length=100,
                   travis=False,
                   burn=None,
                   post_burn=None,
                   gibbs=False,
                   smooth_times=-1,
                   extrapolate_time=0.0,
                   limit_fitting_mjds=False,
                   exclude_bands=[],
                   exclude_instruments=[],
                   suffix='',
                   offline=False,
                   upload=False,
                   write=False,
                   quiet=False,
                   upload_token='',
                   check_upload_quality=False,
                   printer=None,
                   variance_for_each=[],
                   user_fixed_parameters=[],
                   run_until_converged=False,
                   save_full_chain=False,
                   draw_above_likelihood=False,
                   maximum_walltime=False,
                   start_time=False,
                   print_trees=False,
                   maximum_memory=np.inf,
                   **kwargs):
        """Fit a list of events with a list of models."""
        if start_time is False:
            start_time = time.time()
        self._start_time = start_time
        self._maximum_walltime = maximum_walltime
        self._maximum_memory = maximum_memory
        self._debug = False

        dir_path = os.path.dirname(os.path.realpath(__file__))
        self._travis = travis
        self._wrap_length = wrap_length
        self._draw_above_likelihood = draw_above_likelihood

        if not printer:
            self._printer = Printer(wrap_length=wrap_length, quiet=quiet)
        else:
            self._printer = printer

        prt = self._printer

        event_list = listify(events)
        model_list = listify(models)

        if len(model_list) and not len(event_list):
            event_list = ['']

        event_list = [x.replace('‑', '-') for x in event_list]

        self._catalogs = {
            'OSC': ('https://sne.space/astrocats/astrocats/'
                    'supernovae/output'),
            'OTC': ('https://tde.space/astrocats/astrocats/'
                    'tidaldisruptions/output')
        }

        if not len(event_list) and not len(model_list):
            self._printer.wrapped(
                'No events or models specified, initializing and then '
                'exiting.', warning=True)

        entries = [[] for x in range(len(event_list))]
        ps = [[] for x in range(len(event_list))]
        lnprobs = [[] for x in range(len(event_list))]

        data = {}

        self._event_name = 'Batch'
        self._event_catalog = ''
        for ei, event in enumerate(event_list):
            self._event_name = ''
            self._event_path = ''
            if event:
                try:
                    pool = MPIPool()
                except Exception:
                    pool = SerialPool()
                if pool.is_master():
                    path = ''
                    # If the event name ends in .json, assume a path
                    if event.endswith('.json'):
                        path = event
                        self._event_name = event.replace('.json',
                                                         '').split('/')[-1]

                    # If not (or the file doesn't exist), download from OSC
                    if not path or not os.path.exists(path):
                        names_paths = [
                            os.path.join(dir_path, 'cache', x +
                                         '.names.min.json') for x in
                            self._catalogs]
                        input_name = event.replace('.json', '')
                        prt.wrapped(
                            'Event `{}` interpreted as transient '
                            'name, downloading lists of transient '
                            'aliases...'.format(input_name))
                        if not offline:
                            for ci, catalog in enumerate(self._catalogs):
                                try:
                                    response = get_url_file_handle(
                                        self._catalogs[catalog] +
                                        '/names.min.json',
                                        timeout=10)
                                except Exception:
                                    prt.wrapped(
                                        'Warning: Could not download {} '
                                        'names (are you online?), using '
                                        'cached list.'.format(catalog))
                                    raise
                                else:
                                    with open(names_paths[ci], 'wb') as f:
                                        shutil.copyfileobj(response, f)
                        names = OrderedDict()
                        for ci, catalog in enumerate(self._catalogs):
                            if os.path.exists(names_paths[ci]):
                                with open(names_paths[ci], 'r') as f:
                                    names[catalog] = json.load(
                                        f, object_pairs_hook=OrderedDict)
                            else:
                                self._printer.wrapped(
                                    'Warning: Could not read list of {} '
                                    'names!'.format(catalog), warning=True)
                                if offline:
                                    self._printer.wrapped(
                                        'Try omitting the `--offline` flag.')
                                raise RuntimeError

                            if event in names[catalog]:
                                self._event_name = event
                                self._event_catalog = catalog
                            else:
                                for name in names[catalog]:
                                    if (event in names[catalog][name] or
                                            'SN' + event in
                                            names[catalog][name]):
                                        self._event_name = name
                                        self._event_catalog = catalog
                                        break

                        if not self._event_name:
                            for ci, catalog in enumerate(self._catalogs):
                                namekeys = []
                                for name in names[catalog]:
                                    namekeys.extend(names[catalog][name])
                                namekeys = list(sorted(set(namekeys)))
                                matches = get_close_matches(
                                    event, namekeys, n=5, cutoff=0.8)
                                # matches = []
                                if len(matches) < 5 and is_number(event[0]):
                                    prt.wrapped(
                                        'Could not find event, performing '
                                        'extended name search...')
                                    snprefixes = set(('SN19', 'SN20'))
                                    for name in names[catalog]:
                                        ind = re.search("\d", name)
                                        if ind and ind.start() > 0:
                                            snprefixes.add(name[:ind.start()])
                                    snprefixes = list(sorted(snprefixes))
                                    for prefix in snprefixes:
                                        testname = prefix + event
                                        new_matches = get_close_matches(
                                            testname, namekeys, cutoff=0.95,
                                            n=1)
                                        if (len(new_matches) and
                                                new_matches[0] not in matches):
                                            matches.append(new_matches[0])
                                        if len(matches) == 5:
                                            break
                                if len(matches):
                                    if travis:
                                        response = matches[0]
                                    else:
                                        response = prt.prompt(
                                            'No exact match to given '
                                            'transient '
                                            'found. Did you mean one of the '
                                            'following transients?',
                                            kind='select',
                                            options=matches,
                                            none_string=(
                                                'None of the above, ' +
                                                ('skip this event.' if
                                                 ci == len(self._catalogs) - 1
                                                 else
                                                 'try the next catalog.')))
                                    if response:
                                        for name in names[catalog]:
                                            if response in names[
                                                    catalog][name]:
                                                self._event_name = name
                                                self._event_catalog = catalog
                                                break

                        if not self._event_name:
                            prt.wrapped(
                                'Could not find event by that name, skipping!')
                            continue
                        urlname = self._event_name + '.json'
                        name_path = os.path.join(dir_path, 'cache', urlname)

                        if not offline:
                            prt.wrapped(
                                'Found event by primary name `{}` in the {}, '
                                'downloading data...'.format(
                                    self._event_name, self._event_catalog))
                            try:
                                response = get_url_file_handle(
                                    self._catalogs[self._event_catalog] +
                                    '/json/' + urlname,
                                    timeout=10)
                            except Exception:
                                prt.wrapped(
                                    'Warning: Could not download data for '
                                    ' `{}`, '
                                    'will attempt to use cached data.'.format(
                                        self._event_name))
                            else:
                                with open(name_path, 'wb') as f:
                                    shutil.copyfileobj(response, f)
                        path = name_path

                    if os.path.exists(path):
                        with open(path, 'r') as f:
                            data = json.load(f, object_pairs_hook=OrderedDict)
                        prt.wrapped('Event file:')
                        prt.wrapped('  ' + path)
                    else:
                        prt.wrapped(
                            'Error: Could not find data for `{}` locally or '
                            'in the {}.'.format(
                                self._event_name,
                                '/'.join(self._catalogs.keys())))
                        if offline:
                            prt.wrapped(
                                'Try omitting the `--offline` flag.')
                        raise RuntimeError

                    for rank in range(1, pool.size + 1):
                        pool.comm.send(self._event_name, dest=rank, tag=0)
                        pool.comm.send(path, dest=rank, tag=1)
                        pool.comm.send(data, dest=rank, tag=2)
                else:
                    self._event_name = pool.comm.recv(source=0, tag=0)
                    path = pool.comm.recv(source=0, tag=1)
                    data = pool.comm.recv(source=0, tag=2)
                    pool.wait()

                self._event_path = path

                if pool.is_master():
                    pool.close()

            if model_list:
                lmodel_list = model_list
            else:
                lmodel_list = ['']

            entries[ei] = [None for y in range(len(lmodel_list))]
            ps[ei] = [None for y in range(len(lmodel_list))]
            lnprobs[ei] = [None for y in range(len(lmodel_list))]

            for mi, mod_name in enumerate(lmodel_list):
                for parameter_path in parameter_paths:
                    try:
                        pool = MPIPool()
                    except Exception:
                        pool = SerialPool()
                    self._model = Model(
                        model=mod_name,
                        data=data,
                        parameter_path=parameter_path,
                        wrap_length=wrap_length,
                        fitter=self,
                        pool=pool,
                        print_trees=print_trees)

                    if not self._model._model_name:
                        self._printer.wrapped(
                            'Skipping `{}`, no models available to fit the '
                            'transient.'.format(self._event_name),
                            warning=True)
                        continue

                    if not event:
                        self._printer.wrapped(
                            'No event specified, generating dummy data.')
                        self._event_name = mod_name
                        gen_args = {
                            'name': mod_name,
                            'max_time': max_time,
                            'plot_points': plot_points,
                            'band_list': band_list,
                            'band_systems': band_systems,
                            'band_instruments': band_instruments,
                            'band_bandsets': band_bandsets
                        }
                        data = self.generate_dummy_data(**gen_args)

                    success = self.load_data(
                        data,
                        event_name=self._event_name,
                        iterations=iterations,
                        fracking=fracking,
                        burn=burn,
                        post_burn=post_burn,
                        smooth_times=smooth_times,
                        extrapolate_time=extrapolate_time,
                        limit_fitting_mjds=limit_fitting_mjds,
                        exclude_bands=exclude_bands,
                        exclude_instruments=exclude_instruments,
                        band_list=band_list,
                        band_systems=band_systems,
                        band_instruments=band_instruments,
                        band_bandsets=band_bandsets,
                        variance_for_each=variance_for_each,
                        user_fixed_parameters=user_fixed_parameters,
                        pool=pool)

                    if success:
                        entry, p, lnprob = self.fit_data(
                            event_name=self._event_name,
                            iterations=iterations,
                            num_walkers=num_walkers,
                            num_temps=num_temps,
                            fracking=fracking,
                            frack_step=frack_step,
                            gibbs=gibbs,
                            pool=pool,
                            suffix=suffix,
                            write=write,
                            upload=upload,
                            upload_token=upload_token,
                            check_upload_quality=check_upload_quality,
                            run_until_converged=run_until_converged,
                            save_full_chain=save_full_chain)
                        entries[ei][mi] = deepcopy(entry)
                        ps[ei][mi] = deepcopy(p)
                        lnprobs[ei][mi] = deepcopy(lnprob)

                    if pool.is_master():
                        pool.close()

        return (entries, ps, lnprobs)

    def load_data(self,
                  data,
                  event_name='',
                  iterations=2000,
                  fracking=True,
                  burn=None,
                  post_burn=None,
                  smooth_times=-1,
                  extrapolate_time=0.0,
                  limit_fitting_mjds=False,
                  exclude_bands=[],
                  exclude_instruments=[],
                  band_list=[],
                  band_systems=[],
                  band_instruments=[],
                  band_bandsets=[],
                  variance_for_each=[],
                  user_fixed_parameters=[],
                  pool=''):
        """Load the data for the specified event."""
        prt = self._printer
        fixed_parameters = []
        for task in self._model._call_stack:
            cur_task = self._model._call_stack[task]
            self._model._modules[task].set_event_name(event_name)
            if cur_task['kind'] == 'data':
                success = self._model._modules[task].set_data(
                    data,
                    req_key_values={'band': self._model._bands},
                    subtract_minimum_keys=['times'],
                    smooth_times=smooth_times,
                    extrapolate_time=extrapolate_time,
                    limit_fitting_mjds=limit_fitting_mjds,
                    exclude_bands=exclude_bands,
                    exclude_instruments=exclude_instruments,
                    band_list=band_list,
                    band_systems=band_systems,
                    band_instruments=band_instruments,
                    band_bandsets=band_bandsets)
                if not success:
                    return False
                fixed_parameters.extend(self._model._modules[task]
                                        .get_data_determined_parameters())

            # Fix user-specified parameters.
            for fi, param in enumerate(user_fixed_parameters):
                if (task == param or
                        self._model._call_stack[task].get(
                            'class', '') == param):
                    fixed_parameters.append(task)
                    if fi < len(user_fixed_parameters) - 1 and is_number(
                            user_fixed_parameters[fi + 1]):
                        value = float(user_fixed_parameters[fi + 1])
                        if value not in self._model._call_stack:
                            self._model._call_stack[task]['value'] = value
                    if 'min_value' in self._model._call_stack[task]:
                        del self._model._call_stack[task]['min_value']
                    if 'max_value' in self._model._call_stack[task]:
                        del self._model._call_stack[task]['max_value']
                    self._model._modules[task].fix_value(
                        self._model._call_stack[task]['value'])

        self._model.determine_free_parameters(fixed_parameters)

        self._model.exchange_requests()

        # Run through once to set all inits.
        for root in ['output', 'objective']:
            outputs = self._model.run_stack(
                [0.0 for x in range(self._model._num_free_parameters)],
                root=root)

        # Create any data-dependent free parameters.
        self._model.create_data_dependent_free_parameters(
            variance_for_each, outputs)

        # Determine free parameters again as above may have changed them.
        self._model.determine_free_parameters(fixed_parameters)

        self._model.determine_number_of_measurements()

        self._model.exchange_requests()

        # Reset modules
        for task in self._model._call_stack:
            self._model._modules[task].reset_preprocessed()

        # Run through inits once more.
        for root in ['output', 'objective']:
            outputs = self._model.run_stack(
                [0.0 for x in range(self._model._num_free_parameters)],
                root=root)

        # Collect observed band info
        if pool.is_master() and 'photometry' in self._model._modules:
            prt.wrapped('Bands being used for current transient:')
            bis = list(
                filter(lambda a: a != -1,
                       sorted(set(outputs['all_band_indices']))))
            ois = []
            for bi in bis:
                ois.append(
                    any([
                        y
                        for x, y in zip(outputs['all_band_indices'], outputs[
                            'observed']) if x == bi
                    ]))
            band_len = max([
                len(self._model._modules['photometry']._unique_bands[bi][
                    'SVO']) for bi in bis
            ])
            filts = self._model._modules['photometry']
            ubs = filts._unique_bands
            filterarr = [(ubs[bis[i]]['systems'], ubs[bis[i]]['bandsets'],
                          filts._average_wavelengths[bis[i]],
                          filts._band_offsets[bis[i]], ois[i], bis[i])
                         for i in range(len(bis))]
            filterrows = [(
                ' ' + (' ' if s[-2] else '*') + ubs[s[-1]]['SVO']
                .ljust(band_len) + ' [' + ', '.join(
                    list(
                        filter(None, (
                            'Bandset: ' + s[1] if s[1] else '',
                            'System: ' + s[0] if s[0] else '',
                            'AB offset: ' + pretty_num(s[3]))))) +
                ']').replace(' []', '') for s in list(sorted(filterarr))]
            if not all(ois):
                filterrows.append('  (* = Not observed in this band)')
            self._printer.prt('\n'.join(filterrows))

        self._event_name = event_name
        self._emcee_est_t = 0.0
        self._bh_est_t = 0.0
        self._fracking = fracking
        if burn is not None:
            self._burn_in = burn
            self._post_burn = max(iterations - burn, 0)
        elif post_burn is not None:
            self._post_burn = post_burn
            self._burn_in = max(iterations - post_burn, 0)
        else:
            self._burn_in = int(np.round(iterations / 2))
            self._post_burn = max(iterations - self._burn_in, 0)

        return True

    def fit_data(self,
                 event_name='',
                 iterations=2000,
                 frack_step=20,
                 num_walkers=50,
                 num_temps=1,
                 fracking=True,
                 gibbs=False,
                 pool='',
                 suffix='',
                 write=False,
                 upload=False,
                 upload_token='',
                 check_upload_quality=True,
                 run_until_converged=False,
                 save_full_chain=False):
        """Fit the data for a given event.

        Fitting performed using a combination of emcee and fracking.
        """
        from mosfit.__init__ import __version__
        global model
        model = self._model
        prt = self._printer

        upload_this = upload and iterations > 0

        if not pool.is_master():
            try:
                pool.wait()
            except (KeyboardInterrupt, SystemExit):
                pass
            return (None, None, None)

        ntemps, ndim, nwalkers = (num_temps, model._num_free_parameters,
                                  num_walkers)

        test_walker = iterations > 0
        lnprob = None
        lnlike = None
        pool_size = max(pool.size, 1)
        # Derived so only half a walker redrawn with Gaussian distribution.
        redraw_mult = 1.0 * np.sqrt(
            2) * scipy.special.erfinv(float(nwalkers - 1) / nwalkers)

        self._printer.prt(
            '{} measurements, {} free parameters.'.format(
                model._num_measurements, ndim))
        if model._num_measurements <= ndim:
            self._printer.wrapped(
                'Warning: Number of free parameters exceeds number of '
                'measurements. Please treat results with caution.',
                warning=True)
        if nwalkers < 10 * ndim:
            self._printer.wrapped(
                'Warning: While emcee accepts a number of walkers that is '
                'twice the number of dimensions or greater, simple tests '
                'with toy problems show poor convergence with this minimum. '
                'We suggest using at least 10x the number of dimensions '
                '(`-N {}` suggested for this fit), and preferably more, if '
                'feasible.'.format(10 * ndim),
                warning=True)
        self._printer.prt('\n\n')
        p0 = [[] for x in range(ntemps)]

        for i, pt in enumerate(p0):
            dwscores = []
            while len(p0[i]) <= nwalkers:
                prt.status(
                    self,
                    desc='Drawing initial walkers',
                    progress=[
                        i * nwalkers + len(p0[i]) + 1, nwalkers * ntemps])
                if len(p0[i]) == nwalkers:
                    break

                if pool.size == 0:
                    p, score = draw_walker(test_walker)
                    p0[i].append(p)
                    dwscores.append(score)
                else:
                    nmap = nwalkers - len(p0[i])
                    dws = pool.map(draw_walker, [test_walker] * nmap)
                    p0[i].extend([x[0] for x in dws])
                    dwscores.extend([x[1] for x in dws])

                if self._draw_above_likelihood is not False:
                    self._draw_above_likelihood = np.mean(dwscores)

        sampler = MOSSampler(
            ntemps, nwalkers, ndim, likelihood, prior, pool=pool)

        prt.inline('Initial draws completed!')
        self._printer.prt('\n\n')
        p = list(p0)

        sli = 1.0  # Keep track of how many times chain halved
        emi = 0
        tft = 0.0  # Total fracking time
        acor = None
        aacort = -1
        aa = 0
        psrf = np.inf
        s_exception = None
        all_chain = np.array([])

        try:
            st = time.time()

            max_chunk = 1000
            iter_chunks = int(np.ceil(float(iterations) / max_chunk))
            iter_arr = [max_chunk if xi < iter_chunks - 1 else
                        iterations - max_chunk * (iter_chunks - 1)
                        for xi, x in enumerate(range(iter_chunks))]
            # Make sure a chunk separation is located at self._burn_in
            chunk_is = sorted(set(
                np.concatenate(([0, self._burn_in], np.cumsum(iter_arr)))))
            iter_arr = np.diff(chunk_is)

            # The argument of the for loop runs emcee, after each iteration of
            # emcee the contents of the for loop are executed.
            converged = False
            exceeded_walltime = False
            ici = 0
            while run_until_converged or ici < len(iter_arr):
                slr = int(np.round(sli))
                ic = max_chunk if run_until_converged else iter_arr[ici]
                if exceeded_walltime:
                    break
                if run_until_converged and converged and emi > iterations:
                    break
                for li, (
                        p, lnprob, lnlike) in enumerate(
                            sampler.sample(
                                p, iterations=ic, gibbs=gibbs if
                                emi >= self._burn_in else True)):
                    if (self._maximum_walltime is not False and
                            time.time() - self._start_time >
                            self._maximum_walltime):
                        self._printer.wrapped(
                            'Walltime exceeded, exiting...', warning=True)
                        exceeded_walltime = True
                        break
                    emi = emi + 1
                    emim1 = emi - 1
                    messages = []

                    # Record then reset sampler proposal/acceptance counts.
                    accepts = list(
                        np.mean(sampler.nprop_accepted / sampler.nprop,
                                axis=1))
                    sampler.nprop = np.zeros(
                        (sampler.ntemps, sampler.nwalkers), dtype=np.float)
                    sampler.nprop_accepted = np.zeros(
                        (sampler.ntemps, sampler.nwalkers),
                        dtype=np.float)

                    # First, redraw any walkers with scores significantly
                    # worse than their peers (only during burn-in).
                    if emim1 <= self._burn_in:
                        pmedian = [np.median(x) for x in lnprob]
                        pmead = [np.mean([abs(y - pmedian) for y in x])
                                 for x in lnprob]
                        redraw_count = 0
                        bad_redraws = 0
                        for ti, tprob in enumerate(lnprob):
                            for wi, wprob in enumerate(tprob):
                                if (wprob <= pmedian[ti] -
                                    max(redraw_mult * pmead[ti],
                                        float(nwalkers)) or
                                        np.isnan(wprob)):
                                    redraw_count = redraw_count + 1
                                    dxx = np.random.normal(
                                        scale=0.01, size=ndim)
                                    tar_x = np.array(
                                        p[np.random.randint(ntemps)][
                                            np.random.randint(nwalkers)])
                                    new_x = np.clip(tar_x + dxx, 0.0, 1.0)
                                    new_like = likelihood(new_x)
                                    new_prob = new_like + prior(new_x)
                                    if new_prob > wprob or np.isnan(wprob):
                                        p[ti][wi] = new_x
                                        lnlike[ti][wi] = new_like
                                        lnprob[ti][wi] = new_prob
                                    else:
                                        bad_redraws = bad_redraws + 1
                        if redraw_count > 0:
                            messages.append(
                                '{:.0%} redraw, {}/{} success'.format(
                                    redraw_count / (nwalkers * ntemps),
                                    redraw_count - bad_redraws, redraw_count))

                    # Calculate the autocorrelation time.
                    low = 10
                    asize = 0.5 * (emim1 - self._burn_in) / low
                    if asize >= 0:
                        acorc = max(
                            1, min(self._MAX_ACORC,
                                   int(np.floor(0.5 * emi / low))))
                        aacort = -1.0
                        aa = 0
                        ams = self._burn_in
                        cur_chain = (np.concatenate(
                            (all_chain, sampler.chain[:, :, :li + 1:slr, :]),
                            axis=2) if len(all_chain) else
                            sampler.chain[:, :, :li + 1:slr, :])
                        for a in range(acorc, 1, -1):
                            ms = self._burn_in
                            if ms >= emi - low:
                                break
                            try:
                                acorts = sampler.get_autocorr_time(
                                    chain=cur_chain, low=low, c=a,
                                    min_step=int(np.round(float(ms) / sli)),
                                    max_walkers=5, fast=True)
                                acort = max([
                                    max(x)
                                    for x in acorts
                                ])
                            except AutocorrError:
                                continue
                            else:
                                aa = a
                                aacort = acort * sli
                                ams = ms
                                break
                        acor = [aacort, aa, ams]

                    # Calculate the PSRF (Gelman-Rubin statistic).
                    if li > 1 and emi > self._burn_in + 2:
                        cur_chain = (np.concatenate(
                            (all_chain, sampler.chain[:, :, :li + 1:slr, :]),
                            axis=2) if len(all_chain) else
                            sampler.chain[:, :, :li + 1:slr, :])
                        vws = np.zeros((ntemps, ndim))
                        for ti in range(ntemps):
                            for xi in range(ndim):
                                vchain = cur_chain[
                                    ti, :, int(np.round(
                                        self._burn_in / sli)):, xi]
                                m = len(vchain)
                                n = len(vchain[0])
                                mom = np.mean(np.mean(vchain, axis=1))
                                b = n / float(m - 1) * np.sum(
{'fraction_non_alphanumeric': 0.07909643357989386, 'fraction_numerical': 0.005169521819231721, 'mean_word_length': 0.9522505837829133, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 27, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 348, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2708009', 'n_tokens_mistral': 17910, 'n_tokens_neox': 16543, 'n_words': 3855}