#%%
import numpy as np
import librosa
import scipy as sp
import IPython.display as ipd
import matplotlib.pyplot as plt
import os

import json
import codecs

import argparse
import random
import time

from pythonosc import udp_client

################## insert here the name of the cover you want to segment ###############
song_name = "Florida"
######################################################################################

import DEF_features as DEF_features
#%%############## DEFINITION OF FEATURES FUNCTIONS ##########################
#ZCR
def compute_zcr(win, Fs):
    win_sign = np.sign(win)

    N = win.shape[0]

    sign_diff = np.abs(win_sign[:-1] - win_sign[1:])

    zcr = len(sign_diff[sign_diff != 0]) * Fs / N

    # equivalent to

    zcr2 = np.sum(sign_diff) * Fs / (2 * N)
    return zcr
#SPECTRAL CENTROID
def compute_speccentr(spec):
    k_axis = np.arange(1, spec.shape[0] + 1)
    centr = np.sum(np.transpose(k_axis) * np.abs(spec)) / (np.sum(np.abs(spec))+np.finfo(np.float32).eps)
    return centr

def compute_specdec(spec):
    mul_fact  = 1 / (np.sum(np.abs(spec[1:]))+np.finfo(np.float32).eps)
    num = np.abs(spec[1:]) - np.tile(np.abs(spec[0]), len(spec) - 1)
    den = np.arange(1, len(spec))
    spectral_decrease = mul_fact * np.sum(num / den)
    return spectral_decrease

#beat tracking
from librosa.beat import beat_track

def compute_beats(y, sr):
    return beat_track(y, sr=sr)

#not using this because imprecise
def compute_beats_samp(y, sr):
    return beat_track(y, sr=sr, units='samples')

#this function is not used
def first_beat(y, sample_beats):#   ASSUMING WE HAVE A 4/4
    size4= sample_beats.size-(sample_beats.size%4) # closest multiple of 4
    beats4 = np.reshape(sample_beats[:size4],(-1,4))
    energy4 = np.mean(np.power(y[beats4],2), axis=0)
    return np.argmax(energy4)

#I'm  using this because I can set beat_track to output samples directlyb but they are not precise
from librosa import frames_to_samples
def beats_to_sample(beats, y, sr):
    """
    Aligning supposed beats to the peak of energy in the y signal

    beats:  np.ndarray
        frames index where beats are supposed to be
    y: np.ndarray
        input signal
    sr: int
        samplerate

    Returns
    y_beat: np.ndarray
        array with 1 when there is a beat
    beats_indices: np.ndarray
        array with indices of the beats
    """

    y_beat = np.zeros(y.shape)
    margin=int(0.1 * sr)
    for beat in frames_to_samples(beats):
        bs_index = beat-margin +np.argmax(np.abs(y[beat-margin:beat+margin]))
        y_beat[bs_index]=1
    return y_beat, np.where(y_beat==1)[0]

def compute_rms(S):
  return np.sqrt(np.mean(np.abs(S)**2, axis=0, keepdims=True))

def compute_specroll(spec, perc, freq):
  spec= np.abs(spec)
  total_energy = np.cumsum(spec, 0)
  threshold = perc * total_energy[-1]
  ind = np.where(total_energy < threshold, np.nan, 1)
  return np.nanmin(ind*freq)

#%%

song_name_format=song_name + ".mp3"
DATA_DIR="data"
assert os.path.exists(DATA_DIR), "wrong data dir"

filename_in=os.path.join(DATA_DIR, song_name_format) # put whatever you like
audio, Fs = librosa.load(filename_in, sr=None)

ipd.Audio(audio, rate=Fs)

#%%
######################################SETTING PARAMETERS#
win_length = int(np.floor(0.01 * Fs))
hop_size = int(np.floor(0.075 * Fs))

window = sp.signal.get_window(window='hanning', Nx=win_length)

features_names = ['Zero Crossing Rate', 'Spectral Decrease', 'Spectral Centroid', 'Root mean square', 'Spectral Roll off']

#for beat tracking
t=np.arange(audio.size)/Fs

#%%
######################################COMPUTE FEATURES
win_number = int(np.floor((audio.shape[0] - win_length) / hop_size))

n_features = np.asarray(features_names).shape[0]

features = np.zeros((win_number, n_features))
for i in np.arange(win_number):
    frame = audio[i * hop_size : i * hop_size + win_length]
    frame_wind = frame * window

    spec = np.fft.fft(frame_wind)
    nyquist = int(np.floor(spec.shape[0] / 2))
    spec = spec[1:nyquist]
    freq = np.fft.fftfreq(n=win_length, d= 1/Fs)
    freq = freq[1:nyquist]

    features[i, 0] = compute_zcr(frame_wind, Fs)
    features[i, 1] = compute_specdec(spec)
    features[i, 2] = compute_speccentr(spec)
    features[i, 3] = compute_rms(spec)
    features[i, 4] = compute_specroll(spec, 0.85, freq)

#bpm, sample_beats = compute_beats_samp(audio, sr=Fs)

bpm, beats = compute_beats(audio, sr=Fs)
y_beats,sample_beats= beats_to_sample(beats, audio, sr=Fs)
print(bpm)

#%%
###################################### Plot the selected signal and features
plt.figure(1, figsize=(16, 8))

plt.subplot(n_features+1,1,1)
time_axis = np.arange(audio.shape[0]) / Fs
plt.plot(time_axis, audio)
plt.grid(True)
plt.title('audio')

plt.subplot(n_features+1,1,2)
feat_time_axis = np.arange(features.shape[0]) * hop_size / Fs
plt.title(features_names[0])
plt.plot(feat_time_axis, features[:, 0])
plt.grid(True)

plt.subplot(n_features+1,1,3)
feat_time_axis = np.arange(features.shape[0]) * hop_size / Fs
plt.title(features_names[1])
plt.plot(feat_time_axis, features[:, 1])
plt.grid(True)

plt.subplot(n_features+1,1,4)
feat_time_axis = np.arange(features.shape[0]) * hop_size / Fs
plt.title(features_names[2])
plt.plot(feat_time_axis, features[:, 2])
plt.grid(True);

plt.subplot(n_features+1,1,5)
feat_time_axis = np.arange(features.shape[0]) * hop_size / Fs
plt.title(features_names[3])
plt.plot(feat_time_axis, features[:, 3])
plt.grid(True);

plt.subplot(n_features+1,1,6)
feat_time_axis = np.arange(features.shape[0]) * hop_size / Fs
plt.title(features_names[4])
plt.plot(feat_time_axis, features[:, 4])
plt.grid(True);

#See if beats are correct
plt.figure(2, figsize=(16, 8))
plt.plot(t, audio, label="y")
plt.scatter(t[sample_beats], audio[sample_beats], label="beats", color="red")
plt.xlim([0, 5])
plt.xlabel("Time [s]")
plt.legend()

#%%
###################################### Smoothing process
#low pass
smooth_factor = 0.1
features_s=features
for i in range(win_number-1):
  for a in range(features.shape[1]):
    features_s[i+1][a]=smooth_factor*features_s[i+1][a]+(1-smooth_factor)*features_s[i][a]
features_s.shape

plt.figure(3, figsize=(16, 8))

plt.subplot(n_features+1,1,1)
time_axis = np.arange(audio.shape[0]) / Fs
plt.plot(time_axis, audio)
plt.grid(True)
plt.title('audio')

plt.subplot(n_features+1,1,2)
feat_time_axis = np.arange(features_s.shape[0]) * hop_size / Fs
plt.title(features_names[0])
plt.plot(feat_time_axis, features_s[:, 0])
plt.grid(True)

plt.subplot(n_features+1,1,3)
feat_time_axis = np.arange(features_s.shape[0]) * hop_size / Fs
plt.title(features_names[1])
plt.plot(feat_time_axis, features_s[:, 1])
plt.grid(True)

plt.subplot(n_features+1,1,4)
feat_time_axis = np.arange(features_s.shape[0]) * hop_size / Fs
plt.title(features_names[2])
plt.plot(feat_time_axis, features_s[:, 2])
plt.grid(True);

plt.subplot(n_features+1,1,5)
feat_time_axis = np.arange(features_s.shape[0]) * hop_size / Fs
plt.title(features_names[3])
plt.plot(feat_time_axis, features_s[:, 3])
plt.grid(True);

plt.subplot(n_features+1,1,6)
feat_time_axis = np.arange(features_s.shape[0]) * hop_size / Fs
plt.title(features_names[4])
plt.plot(feat_time_axis, features_s[:, 4])
plt.grid(True);

#%%
###################################### Mean, Max, min evaluation
features_mean = np.mean(features_s, 0)
features_Max= np.max(features_s, 0)
features_min = np.min(features_s, 0)
print(features_mean)
print(features_Max)
print(features_min)

print(features_s.shape)
# %%
for x in range(features_s.shape[1]):
    features_s[:, x] = (features_s[:, x]-features_min[x])/(features_Max[x]-features_min[x])

# %%
data0= features_s[:, 0].tolist()
data1= features_s[:, 1].tolist()
data2= features_s[:, 2].tolist()
data3= features_s[:, 3].tolist()
data4= features_s[:, 4].tolist()

sample_beats_s = sample_beats/Fs
data5 = sample_beats_s.tolist()

file_path = song_name+".json" ## your path variable
file_path=os.path.join(DATA_DIR, file_path)
json.dump({'zcr': data0, 'specdec': data1, 'speccentr': data2, 'rms': data3, 'specroll': data4, 'hopsize_s': hop_size/Fs, 'beats': data5}, codecs.open(file_path, 'w',
    encoding='utf-8'), separators=(',', ':'), sort_keys=False, indent=4)

# %%
#if __name__ == "__main__":
#
#    parser = argparse.ArgumentParser()
#    parser.add_argument("--ip", default="127.0.0.1",
#      help="The ip of the OSC server")
#    parser.add_argument("--port", type=int, default=5005,
#      help="The port the OSC server is listening on")
#    args = parser.parse_args()
#
#    client = udp_client.SimpleUDPClient(args.ip, args.port)
#    client.send_message("/Fs", Fs)
#    client.send_message("/t", t)
#
#    for x in range(features_s.shape[1]):
#        client.send_message("/features_s", features_s[:][x])
#    for x in range(sample_beats.shape):
#        client.send_message("/beats", sample_beats[x])

#plt.show()
# %%


from collections import defaultdict, Iterable
from copy import copy
import pickle
import logging
import warnings

import numpy as np
import theano.gradient as tg

from .backends.base import BaseTrace, MultiTrace
from .backends.ndarray import NDArray
from .distributions.distribution import draw_values
from .model import modelcontext, Point, all_continuous
from .step_methods import (NUTS, HamiltonianMC, Metropolis, BinaryMetropolis,
                           BinaryGibbsMetropolis, CategoricalGibbsMetropolis,
                           Slice, CompoundStep, arraystep, smc)
from .util import update_start_vals, get_untransformed_name, is_transformed_name, get_default_varnames
from .vartypes import discrete_types
from pymc3.step_methods.hmc import quadpotential
import pymc3 as pm
from tqdm import tqdm

import sys
sys.setrecursionlimit(10000)

__all__ = ['sample', 'iter_sample', 'sample_posterior_predictive',
           'sample_posterior_predictive_w', 'init_nuts',
           'sample_prior_predictive', 'sample_ppc', 'sample_ppc_w']

STEP_METHODS = (NUTS, HamiltonianMC, Metropolis, BinaryMetropolis,
                BinaryGibbsMetropolis, Slice, CategoricalGibbsMetropolis)

_log = logging.getLogger('pymc3')

def instantiate_steppers(model, steps, selected_steps, step_kwargs=None):
    """Instantiates steppers assigned to the model variables.

    This function is intended to be called automatically from `sample()`, but
    may be called manually.

    Parameters
    ----------
    model : Model object
        A fully-specified model object
    step : step function or vector of step functions
        One or more step functions that have been assigned to some subset of
        the model's parameters. Defaults to None (no assigned variables).
    selected_steps: dictionary of step methods and variables
        The step methods and the variables that have were assigned to them.
    step_kwargs : dict
        Parameters for the samplers. Keys are the lower case names of
        the step method, values a dict of arguments.

    Returns
    -------
    methods : list
        List of step methods associated with the model's variables.
    """
    if step_kwargs is None:
        step_kwargs = {}

    used_keys = set()
    for step_class, vars in selected_steps.items():
        if len(vars) == 0:
            continue
        args = step_kwargs.get(step_class.name, {})
        used_keys.add(step_class.name)
        step = step_class(vars=vars, **args)
        steps.append(step)

    unused_args = set(step_kwargs).difference(used_keys)
    if unused_args:
        raise ValueError('Unused step method arguments: %s' % unused_args)

    if len(steps) == 1:
        steps = steps[0]

    return steps

def assign_step_methods(model, step=None, methods=STEP_METHODS,
                        step_kwargs=None):
    """Assign model variables to appropriate step methods.

    Passing a specified model will auto-assign its constituent stochastic
    variables to step methods based on the characteristics of the variables.
    This function is intended to be called automatically from `sample()`, but
    may be called manually. Each step method passed should have a
    `competence()` method that returns an ordinal competence value
    corresponding to the variable passed to it. This value quantifies the
    appropriateness of the step method for sampling the variable.

    Parameters
    ----------
    model : Model object
        A fully-specified model object
    step : step function or vector of step functions
        One or more step functions that have been assigned to some subset of
        the model's parameters. Defaults to None (no assigned variables).
    methods : vector of step method classes
        The set of step methods from which the function may choose. Defaults
        to the main step methods provided by PyMC3.
    step_kwargs : dict
        Parameters for the samplers. Keys are the lower case names of
        the step method, values a dict of arguments.

    Returns
    -------
    methods : list
        List of step methods associated with the model's variables.
    """
    steps = []
    assigned_vars = set()

    if step is not None:
        try:
            steps += list(step)
        except TypeError:
            steps.append(step)
        for step in steps:
            try:
                assigned_vars = assigned_vars.union(set(step.vars))
            except AttributeError:
                for method in step.methods:
                    assigned_vars = assigned_vars.union(set(method.vars))

    # Use competence classmethods to select step methods for remaining
    # variables
    selected_steps = defaultdict(list)
    for var in model.free_RVs:
        if var not in assigned_vars:
            # determine if a gradient can be computed
            has_gradient = var.dtype not in discrete_types
            if has_gradient:
                try:
                    tg.grad(model.logpt, var)
                except (AttributeError,
                        NotImplementedError,
                        tg.NullTypeGradError):
                    has_gradient = False
            # select the best method
            selected = max(methods, key=lambda method,
                           var=var, has_gradient=has_gradient:
                           method._competence(var, has_gradient))
            selected_steps[selected].append(var)

    return instantiate_steppers(model, steps, selected_steps, step_kwargs)

def _print_step_hierarchy(s, level=0):
    if isinstance(s, (list, tuple)):
        _log.info('>' * level + 'list')
        for i in s:
            _print_step_hierarchy(i, level+1)
    elif isinstance(s, CompoundStep):
        _log.info('>' * level + 'CompoundStep')
        for i in s.methods:
            _print_step_hierarchy(i, level+1)
    else:
        varnames = ', '.join([get_untransformed_name(v.name) if is_transformed_name(v.name)
                              else v.name for v in s.vars])
        _log.info('>' * level + '{}: [{}]'.format(s.__class__.__name__, varnames))

def _cpu_count():
    """Try to guess the number of CPUs in the system.

    We use the number provided by psutil if that is installed.
    If not, we use the number provided by multiprocessing, but assume
    that half of the cpus are only hardware threads and ignore those.
    """
    try:
        import psutil
        cpus = psutil.cpu_count(False)
    except ImportError:
        import multiprocessing
        try:
            cpus = multiprocessing.cpu_count() // 2
        except NotImplementedError:
            cpus = 1
    if cpus is None:
        cpus = 1
    return cpus

def sample(draws=500, step=None, init='auto', n_init=200000, start=None, trace=None, chain_idx=0,
           chains=None, cores=None, tune=500, progressbar=True,
           model=None, random_seed=None, discard_tuned_samples=True,
           compute_convergence_checks=True, **kwargs):
    """Draw samples from the posterior using the given step methods.

    Multiple step methods are supported via compound step methods.

    Parameters
    ----------
    draws : int
        The number of samples to draw. Defaults to 500. The number of tuned samples are discarded
        by default. See discard_tuned_samples.
    step : function or iterable of functions
        A step function or collection of functions. If there are variables without a step methods,
        step methods for those variables will be assigned automatically.
    init : str
        Initialization method to use for auto-assigned NUTS samplers.

        * auto : Choose a default initialization method automatically.
          Currently, this is `'jitter+adapt_diag'`, but this can change in the future.
          If you depend on the exact behaviour, choose an initialization method explicitly.
        * adapt_diag : Start with a identity mass matrix and then adapt a diagonal based on the
          variance of the tuning samples. All chains use the test value (usually the prior mean)
          as starting point.
        * jitter+adapt_diag : Same as `adapt_diag`, but add uniform jitter in [-1, 1] to the
          starting point in each chain.
        * advi+adapt_diag : Run ADVI and then adapt the resulting diagonal mass matrix based on the
          sample variance of the tuning samples.
        * advi+adapt_diag_grad : Run ADVI and then adapt the resulting diagonal mass matrix based
          on the variance of the gradients during tuning. This is **experimental** and might be
          removed in a future release.
        * advi : Run ADVI to estimate posterior mean and diagonal mass matrix.
        * advi_map: Initialize ADVI with MAP and use MAP as starting point.
        * map : Use the MAP as starting point. This is discouraged.
        * nuts : Run NUTS and estimate posterior mean and mass matrix from the trace.
    n_init : int
        Number of iterations of initializer. Only works for 'nuts' and 'ADVI'.
        If 'ADVI', number of iterations, if 'nuts', number of draws.
    start : dict, or array of dict
        Starting point in parameter space (or partial point)
        Defaults to trace.point(-1)) if there is a trace provided and model.test_point if not
        (defaults to empty dict). Initialization methods for NUTS (see `init` keyword) can
        overwrite the default. For 'SMC' it should be a list of dict with length `chains`.
    trace : backend, list, or MultiTrace
        This should be a backend instance, a list of variables to track, or a MultiTrace object
        with past values. If a MultiTrace object is given, it must contain samples for the chain
        number `chain`. If None or a list of variables, the NDArray backend is used.
        Passing either "text" or "sqlite" is taken as a shortcut to set up the corresponding
        backend (with "mcmc" used as the base name). Ignored when using 'SMC'.
    chain_idx : int
        Chain number used to store sample in backend. If `chains` is greater than one, chain
        numbers will start here. Ignored when using 'SMC'.
    chains : int
        The number of chains to sample. Running independent chains is important for some
        convergence statistics and can also reveal multiple modes in the posterior. If `None`,
        then set to either `cores` or 2, whichever is larger. For SMC the number of chains is the
        number of draws.
    cores : int
        The number of chains to run in parallel. If `None`, set to the number of CPUs in the
        system, but at most 4 (for 'SMC' ignored if `pm.SMC(parallel=False)`. Keep in mind that
        some chains might themselves be multithreaded via openmp or BLAS. In those cases it might
        be faster to set this to 1.
    tune : int
        Number of iterations to tune, defaults to 500. Ignored when using 'SMC'. Samplers adjust
        the step sizes, scalings or similar during tuning. Tuning samples will be drawn in addition
        to the number specified in the `draws` argument, and will be discarded unless
        `discard_tuned_samples` is set to False.
    progressbar : bool
        Whether or not to display a progress bar in the command line. The bar shows the percentage
        of completion, the sampling speed in samples per second (SPS), and the estimated remaining
        time until completion ("expected time of arrival"; ETA).
    model : Model (optional if in `with` context)
    random_seed : int or list of ints
        A list is accepted if `cores` is greater than one.
    discard_tuned_samples : bool
        Whether to discard posterior samples of the tune interval. Ignored when using 'SMC'
    compute_convergence_checks : bool, default=True
        Whether to compute sampler statistics like gelman-rubin and effective_n.
        Ignored when using 'SMC'

    Returns
    -------
    trace : pymc3.backends.base.MultiTrace
        A `MultiTrace` object that contains the samples.

    Notes
    -----

    Optional keyword arguments can be passed to `sample` to be delivered to the
    `step_method`s used during sampling. In particular, the NUTS step method accepts
    a number of arguments. Common options are:

        * target_accept: float in [0, 1]. The step size is tuned such that we approximate this
          acceptance rate. Higher values like 0.9 or 0.95 often work better for problematic
          posteriors.
        * max_treedepth: The maximum depth of the trajectory tree.
        * step_scale: float, default 0.25
          The initial guess for the step size scaled down by `1/n**(1/4)`.

    You can find a full list of arguments in the docstring of the step methods.

    Examples
    --------
    .. code:: ipython

        >>> import pymc3 as pm
        ... n = 100
        ... h = 61
        ... alpha = 2
        ... beta = 2

    .. code:: ipython

        >>> with pm.Model() as model: # context management
        ...     p = pm.Beta('p', alpha=alpha, beta=beta)
        ...     y = pm.Binomial('y', n=n, p=p, observed=h)
        ...     trace = pm.sample(2000, tune=1000, cores=4)
        >>> pm.summary(trace)
               mean        sd  mc_error   hpd_2.5  hpd_97.5
        p  0.604625  0.047086   0.00078  0.510498  0.694774
    """
    model = modelcontext(model)

    nuts_kwargs = kwargs.pop('nuts_kwargs', None)
    if nuts_kwargs is not None:
        warnings.warn("The nuts_kwargs argument has been deprecated. Pass step "
                      "method arguments directly to sample instead",
                      DeprecationWarning)
        kwargs.update(nuts_kwargs)
    step_kwargs = kwargs.pop('step_kwargs', None)
    if step_kwargs is not None:
        warnings.warn("The step_kwargs argument has been deprecated. Pass step "
                      "method arguments directly to sample instead",
                      DeprecationWarning)
        kwargs.update(step_kwargs)

    if cores is None:
        cores = min(4, _cpu_count())

    if isinstance(step, pm.step_methods.smc.SMC):
        trace = smc.sample_smc(draws=draws,
                               step=step,
                               cores=cores,
                               progressbar=progressbar,
                               model=model,
                               random_seed=random_seed)
    else:
        if 'njobs' in kwargs:
            cores = kwargs['njobs']
            warnings.warn(
                "The njobs argument has been deprecated. Use cores instead.",
                DeprecationWarning)
        if 'nchains' in kwargs:
            chains = kwargs['nchains']
            warnings.warn(
                "The nchains argument has been deprecated. Use chains instead.",
                DeprecationWarning)
        if chains is None:
            chains = max(2, cores)
        if isinstance(start, dict):
            start = [start] * chains
        if random_seed == -1:
            random_seed = None
        if chains == 1 and isinstance(random_seed, int):
            random_seed = [random_seed]
        if random_seed is None or isinstance(random_seed, int):
            if random_seed is not None:
                np.random.seed(random_seed)
            random_seed = [np.random.randint(2 ** 30) for _ in range(chains)]
        if not isinstance(random_seed, Iterable):
            raise TypeError(
                'Invalid value for `random_seed`. Must be tuple, list or int')
        if 'chain' in kwargs:
            chain_idx = kwargs['chain']
            warnings.warn(
                "The chain argument has been deprecated. Use chain_idx instead.",
                DeprecationWarning)

        if start is not None:
            for start_vals in start:
                _check_start_shape(model, start_vals)

        # small trace warning
        if draws == 0:
            msg = "Tuning was enabled throughout the whole trace."
            _log.warning(msg)
        elif draws < 500:
            msg = "Only %s samples in chain." % draws
            _log.warning(msg)

        draws += tune

        if model.ndim == 0:
            raise ValueError('The model does not contain any free variables.')

        if step is None and init is not None and all_continuous(model.vars):
            try:
                # By default, try to use NUTS
                _log.info('Auto-assigning NUTS sampler...')
                start_, step = init_nuts(init=init, chains=chains, n_init=n_init,
                                         model=model, random_seed=random_seed,
                                         progressbar=progressbar, **kwargs)
                if start is None:
                    start = start_
            except (AttributeError, NotImplementedError, tg.NullTypeGradError):
                # gradient computation failed
                _log.info("Initializing NUTS failed. "
                          "Falling back to elementwise auto-assignment.")
                _log.debug('Exception in init nuts', exec_info=True)
                step = assign_step_methods(model, step, step_kwargs=kwargs)
        else:
            step = assign_step_methods(model, step, step_kwargs=kwargs)

        if isinstance(step, list):
            step = CompoundStep(step)
        if start is None:
            start = {}
        if isinstance(start, dict):
            start = [start] * chains

        sample_args = {'draws': draws,
                       'step': step,
                       'start': start,
                       'trace': trace,
                       'chain': chain_idx,
                       'chains': chains,
                       'tune': tune,
                       'progressbar': progressbar,
                       'model': model,
                       'random_seed': random_seed,
                       'cores': cores, }

        sample_args.update(kwargs)

        has_population_samplers = np.any([isinstance(m, arraystep.PopulationArrayStepShared)
                                          for m in (step.methods if isinstance(step, CompoundStep) else [step])])

        parallel = cores > 1 and chains > 1 and not has_population_samplers
        if parallel:
            _log.info('Multiprocess sampling ({} chains in {} jobs)'.format(chains, cores))
            _print_step_hierarchy(step)
            try:
                trace = _mp_sample(**sample_args)
            except pickle.PickleError:
                _log.warning("Could not pickle model, sampling singlethreaded.")
                _log.debug('Pickling error:', exec_info=True)
                parallel = False
            except AttributeError as e:
                if str(e).startswith("AttributeError: Can't pickle"):
                    _log.warning("Could not pickle model, sampling singlethreaded.")
                    _log.debug('Pickling error:', exec_info=True)
                    parallel = False
                else:
                    raise
        if not parallel:
            if has_population_samplers:
                _log.info('Population sampling ({} chains)'.format(chains))
                _print_step_hierarchy(step)
                trace = _sample_population(**sample_args)
            else:
                _log.info('Sequential sampling ({} chains in 1 job)'.format(chains))
                _print_step_hierarchy(step)
                trace = _sample_many(**sample_args)

        discard = tune if discard_tuned_samples else 0
        trace = trace[discard:]

        if compute_convergence_checks:
            if draws-tune < 100:
                warnings.warn("The number of samples is too small to check convergence reliably.")
            else:
                trace.report._run_convergence_checks(trace, model)

        trace.report._log_summary()

    return trace

def _check_start_shape(model, start):
    if not isinstance(start, dict):
        raise TypeError("start argument must be a dict or an array-like of dicts")
    e = ''
    for var in model.vars:
        if var.name in start.keys():
            var_shape = var.shape.tag.test_value
            start_var_shape = np.shape(start[var.name])
            if start_var_shape:
                if not np.array_equal(var_shape, start_var_shape):
                    e += "\nExpected shape {} for var '{}', got: {}".format(
                        tuple(var_shape), var.name, start_var_shape
                    )
            # if start var has no shape
            else:
                # if model var has a specified shape
                if var_shape.size > 0:
                    e += "\nExpected shape {} for var " \
                         "'{}', got scalar {}".format(
                             tuple(var_shape), var.name, start[var.name]
                         )

    if e != '':
        raise ValueError("Bad shape for start argument:{}".format(e))

def _sample_many(draws, chain, chains, start, random_seed, step, **kwargs):
    traces = []
    for i in range(chains):
        trace = _sample(draws=draws, chain=chain + i, start=start[i],
                        step=step, random_seed=random_seed[i], **kwargs)
        if trace is None:
            if len(traces) == 0:
                raise ValueError('Sampling stopped before a sample was created.')
            else:
                break
        elif len(trace) < draws:
            if len(traces) == 0:
                traces.append(trace)
            break
        else:
            traces.append(trace)
    return MultiTrace(traces)

def _sample_population(draws, chain, chains, start, random_seed, step, tune,
                       model, progressbar=None, parallelize=False, **kwargs):
    # create the generator that iterates all chains in parallel
    chains = [chain + c for c in range(chains)]
    sampling = _prepare_iter_population(draws, chains, step, start, parallelize,
                                        tune=tune, model=model, random_seed=random_seed)

    if progressbar:
        sampling = tqdm(sampling, total=draws)

    latest_traces = None
    for it, traces in enumerate(sampling):
        latest_traces = traces
    return MultiTrace(latest_traces)

def _sample(chain, progressbar, random_seed, start, draws=None, step=None,
            trace=None, tune=None, model=None, **kwargs):
    skip_first = kwargs.get('skip_first', 0)
    refresh_every = kwargs.get('refresh_every', 100)

    sampling = _iter_sample(draws, step, start, trace, chain,
                            tune, model, random_seed)
    if progressbar:
        sampling = tqdm(sampling, total=draws)
    try:
        strace = None
        for it, strace in enumerate(sampling):
            if it >= skip_first:
                trace = MultiTrace([strace])
    except KeyboardInterrupt:
        pass
    finally:
        if progressbar:
            sampling.close()
    return strace

def iter_sample(draws, step, start=None, trace=None, chain=0, tune=None,
                model=None, random_seed=None):
    """Generator that returns a trace on each iteration using the given
    step method.  Multiple step methods supported via compound step
    method returns the amount of time taken.

    Parameters
    ----------
    draws : int
        The number of samples to draw
    step : function
        Step function
    start : dict
        Starting point in parameter space (or partial point). Defaults to trace.point(-1)) if
        there is a trace provided and model.test_point if not (defaults to empty dict)
    trace : backend, list, or MultiTrace
        This should be a backend instance, a list of variables to track, or a MultiTrace object
        with past values. If a MultiTrace object is given, it must contain samples for the chain
        number `chain`. If None or a list of variables, the NDArray backend is used.
    chain : int
        Chain number used to store sample in backend. If `cores` is greater than one, chain numbers
        will start here.
    tune : int
        Number of iterations to tune, if applicable (defaults to None)
    model : Model (optional if in `with` context)
    random_seed : int or list of ints
        A list is accepted if more if `cores` is greater than one.

    Examples
    --------
    ::

        for trace in iter_sample(500, step):
            ...
    """
    sampling = _iter_sample(draws, step, start, trace, chain, tune,
                            model, random_seed)
    for i, strace in enumerate(sampling):
        yield MultiTrace([strace[:i + 1]])

def _iter_sample(draws, step, start=None, trace=None, chain=0, tune=None,
                 model=None, random_seed=None):
    model = modelcontext(model)
    draws = int(draws)
    if random_seed is not None:
        np.random.seed(random_seed)
    if draws < 1:
        raise ValueError('Argument `draws` must be greater than 0.')

    if start is None:
        start = {}

    strace = _choose_backend(trace, chain, model=model)

    if len(strace) > 0:
        update_start_vals(start, strace.point(-1), model)
    else:
        update_start_vals(start, model.test_point, model)

    try:
        step = CompoundStep(step)
    except TypeError:
        pass

    point = Point(start, model=model)

    if step.generates_stats and strace.supports_sampler_stats:
        strace.setup(draws, chain, step.stats_dtypes)
    else:
        strace.setup(draws, chain)

    try:
        step.tune = bool(tune)
        for i in range(draws):
            if i == tune:
                step = stop_tuning(step)
            if step.generates_stats:
                point, states = step.step(point)
                if strace.supports_sampler_stats:
                    strace.record(point, states)
                else:
                    strace.record(point)
            else:
                point = step.step(point)
                strace.record(point)
            yield strace
    except KeyboardInterrupt:
        strace.close()
        if hasattr(step, 'warnings'):
            warns = step.warnings()
            strace._add_warnings(warns)
        raise
    except BaseException:
        strace.close()
        raise
    else:
        strace.close()
        if hasattr(step, 'warnings'):
            warns = step.warnings()
            strace._add_warnings(warns)

class PopulationStepper:
    def __init__(self, steppers, parallelize):
        """Tries to use multiprocessing to parallelize chains.

        Falls back to sequential evaluation if multiprocessing fails.

        In the multiprocessing mode of operation, a new process is started for each
        chain/stepper and Pipes are used to communicate with the main process.

        Parameters
        ----------
        steppers : list
            A collection of independent step methods, one for each chain.
        parallelize : bool
            Indicates if chain parallelization is desired
        """
        self.nchains = len(steppers)
        self.is_parallelized = False
        self._master_ends = []
        self._processes = []
        self._steppers = steppers
        if parallelize:
            try:
                # configure a child process for each stepper
                _log.info('Attempting to parallelize chains.')
                import multiprocessing
                for c, stepper in enumerate(tqdm(steppers)):
                    slave_end, master_end = multiprocessing.Pipe()
                    stepper_dumps = pickle.dumps(stepper, protocol=4)
                    process = multiprocessing.Process(
                        target=self.__class__._run_slave,
                        args=(c, stepper_dumps, slave_end),
                        name='ChainWalker{}'.format(c)
                    )
                    # we want the child process to exit if the parent is terminated
                    process.daemon = True
                    # Starting the process might fail and takes time.
                    # By doing it in the constructor, the sampling progress bar
                    # will not be confused by the process start.
                    process.start()
                    self._master_ends.append(master_end)
                    self._processes.append(process)
                self.is_parallelized = True
            except Exception:
                _log.info('Population parallelization failed. '
                          'Falling back to sequential stepping of chains.')
                _log.debug('Error was: ', exec_info=True)
        else:
            _log.info('Chains are not parallelized. You can enable this by passing '
                      'pm.sample(parallelize=True).')
        return super().__init__()

    def __enter__(self):
        """Does nothing because processes are already started in __init__."""
        return

    def __exit__(self, exc_type, exc_val, exc_tb):
        if len(self._processes) > 0:
            try:
                for master_end in self._master_ends:
                    master_end.send(None)
                for process in self._processes:
                    process.join(timeout=3)
            except Exception:
                _log.warning('Termination failed.')
        return

    @staticmethod
    def _run_slave(c, stepper_dumps, slave_end):
        """Started on a separate process to perform stepping of a chain.

        Parameters
        ----------
        c : int
            number of this chain
        stepper : BlockedStep
            a step method such as CompoundStep
        slave_end : multiprocessing.connection.PipeConnection
            This is our connection to the main process
        """
        # re-seed each child process to make them unique
        np.random.seed(None)
        try:
            stepper = pickle.loads(stepper_dumps)
            # the stepper is not necessarily a PopulationArraySharedStep itself,
            # but rather a CompoundStep. PopulationArrayStepShared.population
            # has to be updated, therefore we identify the substeppers first.
            population_steppers = []
            for sm in (stepper.methods if isinstance(stepper, CompoundStep) else [stepper]):
                if isinstance(sm, arraystep.PopulationArrayStepShared):
                    population_steppers.append(sm)
            while True:
                incoming = slave_end.recv()
                # receiving a None is the signal to exit
                if incoming is None:
                    break
                tune_stop, population = incoming
                if tune_stop:
                    stop_tuning(stepper)
                # forward the population to the PopulationArrayStepShared objects
                # This is necessary because due to the process fork, the population
                # object is no longer shared between the steppers.
                for popstep in population_steppers:
                    popstep.population = population
                update = stepper.step(population[c])
                slave_end.send(update)
        except Exception:
            _log.exception('ChainWalker{}'.format(c))
        return

    def step(self, tune_stop, population):
        """Steps the entire population of chains.

        Parameters
        ----------
        tune_stop : bool
            Indicates if the condition (i == tune) is fulfilled
        population : list
            Current Points of all chains

        Returns
        -------
        update : Point
            The new positions of the chains
        """
        updates = [None] * self.nchains
        if self.is_parallelized:
            for c in range(self.nchains):
                self._master_ends[c].send((tune_stop, population))
            # Blockingly get the step outcomes
            for c in range(self.nchains):
                updates[c] = self._master_ends[c].recv()
        else:
            for c in range(self.nchains):
                if tune_stop:
                    self._steppers[c] = stop_tuning(self._steppers[c])
                updates[c] = self._steppers[c].step(population[c])
        return updates

def _prepare_iter_population(draws, chains, step, start, parallelize, tune=None,
                             model=None, random_seed=None):
    """Prepares a PopulationStepper and traces for population sampling.

    Returns
    -------
    _iter_population : generator
        The generator the yields traces of all chains at the same time
    """
    # chains contains the chain numbers, but for indexing we need indices...
    nchains = len(chains)
    model = modelcontext(model)
    draws = int(draws)
    if random_seed is not None:
        np.random.seed(random_seed)
    if draws < 1:
        raise ValueError('Argument `draws` should be above 0.')

    # The initialization of traces, samplers and points must happen in the right order:
    # 1. traces are initialized and update_start_vals configures variable transforms
    # 2. population of points is created
    # 3. steppers are initialized and linked to the points object
    # 4. traces are configured to track the sampler stats
    # 5. a PopulationStepper is configured for parallelized stepping

    # 1. prepare a BaseTrace for each chain
    traces = [_choose_backend(None, chain, model=model) for chain in chains]
    for c, strace in enumerate(traces):
        # initialize the trace size and variable transforms
        if len(strace) > 0:
            update_start_vals(start[c], strace.point(-1), model)
        else:
            update_start_vals(start[c], model.test_point, model)

    # 2. create a population (points) that tracks each chain
    # it is updated as the chains are advanced
    population = [Point(start[c], model=model) for c in range(nchains)]

    # 3. Set up the steppers
    steppers = [None] * nchains
    for c in range(nchains):
        # need indepenent samplers for each chain
        # it is important to copy the actual steppers (but not the delta_logp)
        if isinstance(step, CompoundStep):
            chainstep = CompoundStep([copy(m) for m in step.methods])
        else:
            chainstep = copy(step)
        # link population samplers to the shared population state
        for sm in (chainstep.methods if isinstance(step, CompoundStep) else [chainstep]):
            if isinstance(sm, arraystep.PopulationArrayStepShared):
                sm.link_population(population, c)
        steppers[c] = chainstep

    # 4. configure tracking of sampler stats
    for c in range(nchains):
        if steppers[c].generates_stats and traces[c].supports_sampler_stats:
            traces[c].setup(draws, c, steppers[c].stats_dtypes)
        else:
            traces[c].setup(draws, c)

    # 5. configure the PopulationStepper (expensive call)
    popstep = PopulationStepper(steppers, parallelize)

    # Because the preparations above are expensive, the actual iterator is
    # in another method. This way the progbar will not be disturbed.
    return _iter_population(draws, tune, popstep, steppers, traces, population)

def _iter_population(draws, tune, popstep, steppers, traces, points):
    """Generator that iterates a PopulationStepper.

    Parameters
    ----------
    draws : int
        number of draws per chain
    tune : int
        number of tuning steps
    popstep : PopulationStepper
        the helper object for (parallelized) stepping of chains
    steppers : list
        The step methods for each chain
    traces : list
        Traces for each chain
    points : list
        population of chain states
    """
    try:
        with popstep:
            # iterate draws of all chains
            for i in range(draws):
                updates = popstep.step(i == tune, points)

                # apply the update to the points and record to the traces
                for c, strace in enumerate(traces):
                    if steppers[c].generates_stats:
                        points[c], states = updates[c]
                        if strace.supports_sampler_stats:
                            strace.record(points[c], states)
                        else:
                            strace.record(points[c])
                    else:
                        points[c] = updates[c]
                        strace.record(points[c])
                # yield the state of all chains in parallel
                yield traces
    except KeyboardInterrupt:
        for c, strace in enumerate(traces):
            strace.close()
            if hasattr(steppers[c], 'report'):
                steppers[c].report._finalize(strace)
        raise
    except BaseException:
        for c, strace in enumerate(traces):
            strace.close()
        raise
    else:
        for c, strace in enumerate(traces):
            strace.close()
            if hasattr(steppers[c], 'report'):
                steppers[c].report._finalize(strace)

def _choose_backend(trace, chain, shortcuts=None, **kwds):
    if isinstance(trace, BaseTrace):
        return trace
    if isinstance(trace, MultiTrace):
        return trace._straces[chain]
    if trace is None:
        return NDArray(**kwds)

    if shortcuts is None:
        shortcuts = pm.backends._shortcuts

    try:
        backend = shortcuts[trace]['backend']
        name = shortcuts[trace]['name']
        return backend(name, **kwds)
    except TypeError:
        return NDArray(vars=trace, **kwds)
    except KeyError:
        raise ValueError('Argument `trace` is invalid.')

def _mp_sample(draws, tune, step, chains, cores, chain, random_seed,
               start, progressbar, trace=None, model=None, **kwargs):

    import pymc3.parallel_sampling as ps
    # We did draws += tune in pm.sample
    draws -= tune

    traces = []
    for idx in range(chain, chain + chains):
        if trace is not None:
            strace = _choose_backend(copy(trace), idx, model=model)
        else:
            strace = _choose_backend(None, idx, model=model)
        # for user supply start value, fill-in missing value if the supplied
        # dict does not contain all parameters
        update_start_vals(start[idx - chain], model.test_point, model)
        if step.generates_stats and strace.supports_sampler_stats:
            strace.setup(draws + tune, idx + chain, step.stats_dtypes)
        else:
            strace.setup(draws + tune, idx + chain)
        traces.append(strace)

    sampler = ps.ParallelSampler(
        draws, tune, chains, cores, random_seed, start, step,
        chain, progressbar)
    try:
        try:
            with sampler:
                for draw in sampler:
                    trace = traces[draw.chain - chain]
                    if (trace.supports_sampler_stats
                            and draw.stats is not None):
                        trace.record(draw.point, draw.stats)
                    else:
                        trace.record(draw.point)
                    if draw.is_last:
                        trace.close()
                        if draw.warnings is not None:
                            trace._add_warnings(draw.warnings)
        except ps.ParallelSamplingError as error:
            trace = traces[error._chain - chain]
            trace._add_warnings(error._warnings)
            for trace in traces:
                trace.close()

            multitrace = MultiTrace(traces)
            multitrace._report._log_summary()
            raise
        return MultiTrace(traces)
    except KeyboardInterrupt:
        traces, length = _choose_chains(traces, tune)
        return MultiTrace(traces)[:length]
    finally:
        for trace in traces:
            trace.close()

def _choose_chains(traces, tune):
    if tune is None:
        tune = 0

    if not traces:
        return []

    lengths = [max(0, len(trace) - tune) for trace in traces]
    if not sum(lengths):
        raise ValueError('Not enough samples to build a trace.')

    idxs = np.argsort(lengths)[::-1]
    l_sort = np.array(lengths)[idxs]

    final_length = l_sort[0]
    last_total = 0
    for i, length in enumerate(l_sort):
        total = (i + 1) * length
        if total < last_total:
            use_until = i
            break
        last_total = total
        final_length = length
    else:
        use_until = len(lengths)

    return [traces[idx] for idx in idxs[:use_until]], final_length + tune

def stop_tuning(step):
    """ stop tuning the current step method """

    step.stop_tuning()
    return step

def sample_posterior_predictive(trace, samples=None, model=None, vars=None, size=None,
                                random_seed=None, progressbar=True):
    """Generate posterior predictive samples from a model given a trace.

    Parameters
    ----------
    trace : backend, list, or MultiTrace
        Trace generated from MCMC sampling. Or a list containing dicts from
        find_MAP() or points
    samples : int
        Number of posterior predictive samples to generate. Defaults to the length of `trace`
    model : Model (optional if in `with` context)
        Model used to generate `trace`
    vars : iterable
        Variables for which to compute the posterior predictive samples.
        Defaults to `model.observed_RVs`.
    size : int
        The number of random draws from the distribution specified by the parameters in each
        sample of the trace.
    random_seed : int
        Seed for the random number generator.
    progressbar : bool
        Whether or not to display a progress bar in the command line. The bar shows the percentage
        of completion, the sampling speed in samples per second (SPS), and the estimated remaining
        time until completion ("expected time of arrival"; ETA).

    Returns
    -------
    samples : dict
        Dictionary with the variables as keys. The values corresponding to the
        posterior predictive samples.
    """
    len_trace = len(trace)
    try:
        nchain = trace.nchains
    except AttributeError:
        nchain = 1

    if samples is None:
        samples = sum(len(v) for v in trace._straces.values())

    model = modelcontext(model)

    if vars is None:
        vars = model.observed_RVs

    if random_seed is not None:
        np.random.seed(random_seed)

    indices = np.arange(samples)

    if progressbar:
        indices = tqdm(indices, total=samples)

    ppc_trace = defaultdict(list)
    try:
        for idx in indices:
            if nchain > 1:
                chain_idx, point_idx = np.divmod(idx, len_trace)
                param = trace._straces[chain_idx % nchain].point(point_idx)
            else:
                param = trace[idx % len_trace]

            values = draw_values(vars, point=param, size=size)
            for k, v in zip(vars, values):
                ppc_trace[k.name].append(v)

    except KeyboardInterrupt:
        pass

    finally:
        if progressbar:
            indices.close()

    return {k: np.asarray(v) for k, v in ppc_trace.items()}

def sample_ppc(*args, **kwargs):
    """This method is deprecated.  Please use :func:`~sampling.sample_posterior_predictive`"""
    message = 'sample_ppc() is deprecated.  Please use sample_posterior_predictive()'
    warnings.warn(message, DeprecationWarning, stacklevel=2)
    return sample_posterior_predictive(*args, **kwargs)

def sample_posterior_predictive_w(traces, samples=None, models=None, weights=None,
                                  random_seed=None, progressbar=True):
    """Generate weighted posterior predictive samples from a list of models and
    a list of traces according to a set of weights.

    Parameters
    ----------
    traces : list or list of lists
        List of traces generated from MCMC sampling, or a list of list
        containing dicts from find_MAP() or points. The number of traces should
        be equal to the number of weights.
    samples : int
        Number of posterior predictive samples to generate. Defaults to the
        length of the shorter trace in traces.
    models : list
        List of models used to generate the list of traces. The number of models should be equal to
        the number of weights and the number of observed RVs should be the same for all models.
        By default a single model will be inferred from `with` context, in this case results will
        only be meaningful if all models share the same distributions for the observed RVs.
    weights: array-like
        Individual weights for each trace. Default, same weight for each model.
    random_seed : int
        Seed for the random number generator.
    progressbar : bool
        Whether or not to display a progress bar in the command line. The bar shows the percentage
        of completion, the sampling speed in samples per second (SPS), and the estimated remaining
        time until completion ("expected time of arrival"; ETA).

    Returns
    -------
    samples : dict
        Dictionary with the variables as keys. The values corresponding to the
        posterior predictive samples from the weighted models.
    """
    np.random.seed(random_seed)

    if models is None:
        models = [modelcontext(models)] * len(traces)

    if weights is None:
        weights = [1] * len(traces)

    if len(traces) != len(weights):
        raise ValueError('The number of traces and weights should be the same')

    if len(models) != len(weights):
        raise ValueError('The number of models and weights should be the same')

    length_morv = len(models[0].observed_RVs)
    if not all(len(i.observed_RVs) == length_morv for i in models):
        raise ValueError(
            'The number of observed RVs should be the same for all models')

    weights = np.asarray(weights)
    p = weights / np.sum(weights)

    min_tr = min([len(i) * i.nchains for i in traces])

    n = (min_tr * p).astype('int')
    # ensure n sum up to min_tr
    idx = np.argmax(n)
    n[idx] = n[idx] + min_tr - np.sum(n)
    trace = []
    for i, j in enumerate(n):
        tr = traces[i]
        len_trace = len(tr)
        try:
            nchain = tr.nchains
        except AttributeError:
            nchain = 1

        indices = np.random.randint(0, nchain * len_trace, j)
        if nchain > 1:
            chain_idx, point_idx = np.divmod(indices, len_trace)
            for idx in zip(chain_idx, point_idx):
                trace.append(tr._straces[idx[0]].point(idx[1]))
        else:
            for idx in indices:
                trace.append(tr[idx])

    obs = [x for m in models for x in m.observed_RVs]
    variables = np.repeat(obs, n)

    lengths = list(set([np.atleast_1d(observed).shape for observed in obs]))

    if len(lengths) == 1:
        size = [None for i in variables]
    elif len(lengths) > 2:
        raise ValueError('Observed variables could not be broadcast together')
    else:
        size = []
        x = np.zeros(shape=lengths[0])
        y = np.zeros(shape=lengths[1])
        b = np.broadcast(x, y)
        for var in variables:
            shape = np.shape(np.atleast_1d(var.distribution.default()))
            if shape != b.shape:
                size.append(b.shape)
            else:
                size.append(None)
    len_trace = len(trace)

    if samples is None:
        samples = len_trace

    indices = np.random.randint(0, len_trace, samples)

    if progressbar:
        indices = tqdm(indices, total=samples)

    try:
        ppc = defaultdict(list)
        for idx in indices:
            param = trace[idx]
            var = variables[idx]
            # TODO sample_posterior_predictive_w is currently only work for model with
            # one observed.
            ppc[var.name].append(draw_values([var],
                                             point=param,
                                             size=size[idx]
                                             )[0])

    except KeyboardInterrupt:
        pass

    finally:
        if progressbar:
            indices.close()

    return {k: np.asarray(v) for k, v in ppc.items()}

def sample_ppc_w(*args, **kwargs):
    """This method is deprecated.  Please use :func:`~sampling.sample_posterior_predictive_w`"""
    message = 'sample_ppc() is deprecated.  Please use sample_posterior_predictive_w()'
    warnings.warn(message, DeprecationWarning, stacklevel=2)
    return sample_posterior_predictive_w(*args, **kwargs)

def sample_prior_predictive(samples=500, model=None, vars=None, random_seed=None):
    """Generate samples from the prior predictive distribution.

    Parameters
    ----------
    samples : int
        Number of samples from the prior predictive to generate. Defaults to 500.
    model : Model (optional if in `with` context)
    vars : iterable
        A list of names of variables for which to compute the posterior predictive
         samples.
        Defaults to `model.named_vars`.
    random_seed : int
        Seed for the random number generator.

    Returns
    -------
    dict
        Dictionary with variable names as keys. The values are numpy arrays of prior
         samples.
    """
    model = modelcontext(model)

    if vars is None:
        vars = set(model.named_vars.keys())

    if random_seed is not None:
        np.random.seed(random_seed)
    names = get_default_varnames(model.named_vars, include_transformed=False)
    # draw_values fails with auto-transformed variables. transform them later!
    values = draw_values([model[name] for name in names], size=samples)

    data = {k: v for k, v in zip(names, values)}

    prior = {}
    for var_name in vars:
        if var_name in data:
            prior[var_name] = data[var_name]
        elif is_transformed_name(var_name):
            untransformed = get_untransformed_name(var_name)
            if untransformed in data:
                prior[var_name] = model[untransformed].transformation.forward_val(
                    data[untransformed])
    return prior

def init_nuts(init='auto', chains=1, n_init=500000, model=None,
              random_seed=None, progressbar=True, **kwargs):
    """Set up the mass matrix initialization for NUTS.

    NUTS convergence and sampling speed is extremely dependent on the
    choice of mass/scaling matrix. This function implements different
    methods for choosing or adapting the mass matrix.

    Parameters
    ----------
    init : str
        Initialization method to use.

        * auto : Choose a default initialization method automatically.
          Currently, this is `'jitter+adapt_diag'`, but this can change in the future. If you
          depend on the exact behaviour, choose an initialization method explicitly.
        * adapt_diag : Start with a identity mass matrix and then adapt a diagonal based on the
          variance of the tuning samples. All chains use the test value (usually the prior mean)
          as starting point.
        * jitter+adapt_diag : Same as `adapt_diag`, but use uniform jitter in [-1, 1] as starting
          point in each chain.
        * advi+adapt_diag : Run ADVI and then adapt the resulting diagonal mass matrix based on the
          sample variance of the tuning samples.
        * advi+adapt_diag_grad : Run ADVI and then adapt the resulting diagonal mass matrix based
          on the variance of the gradients during tuning. This is **experimental** and might be
          removed in a future release.
        * advi : Run ADVI to estimate posterior mean and diagonal mass matrix.
        * advi_map: Initialize ADVI with MAP and use MAP as starting point.
        * map : Use the MAP as starting point. This is discouraged.
        * nuts : Run NUTS and estimate posterior mean and mass matrix from
          the trace.
    chains : int
        Number of jobs to start.
    n_init : int
        Number of iterations of initializer
        If 'ADVI', number of iterations, if 'nuts', number of draws.
    model : Model (optional if in `with` context)
    progressbar : bool
        Whether or not to display a progressbar for advi sampling.
    **kwargs : keyword arguments
        Extra keyword arguments are forwarded to pymc3.NUTS.

    Returns
    -------
    start : pymc3.model.Point
        Starting point for sampler
    nuts_sampler : pymc3.step_methods.NUTS
        Instantiated and initialized NUTS sampler object
    """
    model = modelcontext(model)

    vars = kwargs.get('vars', model.vars)
    if set(vars) != set(model.vars):
        raise ValueError('Must use init_nuts on all variables of a model.')
    if not all_continuous(vars):
        raise ValueError('init_nuts can only be used for models with only '
                         'continuous variables.')

    if not isinstance(init, str):
        raise TypeError('init must be a string.')

    if init is not None:
        init = init.lower()

    if init == 'auto':
        init = 'jitter+adapt_diag'

    _log.info('Initializing NUTS using {}...'.format(init))

    if random_seed is not None:
        random_seed = int(np.atleast_1d(random_seed)[0])
        np.random.seed(random_seed)

    cb = [
        pm.callbacks.CheckParametersConvergence(
            tolerance=1e-2, diff='absolute'),
        pm.callbacks.CheckParametersConvergence(
            tolerance=1e-2, diff='relative'),
    ]

    if init == 'adapt_diag':
        start = [model.test_point] * chains
        mean = np.mean([model.dict_to_array(vals) for vals in start], axis=0)
        var = np.ones_like(mean)
        potential = quadpotential.QuadPotentialDiagAdapt(
            model.ndim, mean, var, 10)
    elif init == 'jitter+adapt_diag':
        start = []
        for _ in range(chains):
            mean = {var: val.copy() for var, val in model.test_point.items()}
            for val in mean.values():
                val[...] += 2 * np.random.rand(*val.shape) - 1
            start.append(mean)
        mean = np.mean([model.dict_to_array(vals) for vals in start], axis=0)
        var = np.ones_like(mean)
        potential = quadpotential.QuadPotentialDiagAdapt(
            model.ndim, mean, var, 10)
    elif init == 'advi+adapt_diag_grad':
        approx = pm.fit(
            random_seed=random_seed,
            n=n_init, method='advi', model=model,
            callbacks=cb,
            progressbar=progressbar,
            obj_optimizer=pm.adagrad_window,
        )  # type: pm.MeanField
        start = approx.sample(draws=chains)
        start = list(start)
        stds = approx.bij.rmap(approx.std.eval())
        cov = model.dict_to_array(stds) ** 2
        mean = approx.bij.rmap(approx.mean.get_value())
        mean = model.dict_to_array(mean)
        weight = 50
        potential = quadpotential.QuadPotentialDiagAdaptGrad(
            model.ndim, mean, cov, weight)
    elif init == 'advi+adapt_diag':
        approx = pm.fit(
            random_seed=random_seed,
            n=n_init, method='advi', model=model,
            callbacks=cb,
            progressbar=progressbar,
            obj_optimizer=pm.adagrad_window,
        )  # type: pm.MeanField
        start = approx.sample(draws=chains)
        start = list(start)
        stds = approx.bij.rmap(approx.std.eval())
        cov = model.dict_to_array(stds) ** 2
        mean = approx.bij.rmap(approx.mean.get_value())
        mean = model.dict_to_array(mean)
        weight = 50
        potential = quadpotential.QuadPotentialDiagAdapt(
            model.ndim, mean, cov, weight)
    elif init == 'advi':
        approx = pm.fit(
            random_seed=random_seed,
            n=n_init, method='advi', model=model,
            callbacks=cb,
            progressbar=progressbar,
            obj_optimizer=pm.adagrad_window
        )  # type: pm.MeanField
        start = approx.sample(draws=chains)
        start = list(start)
        stds = approx.bij.rmap(approx.std.eval())
        cov = model.dict_to_array(stds) ** 2
        potential = quadpotential.QuadPotentialDiag(cov)
    elif init == 'advi_map':
        start = pm.find_MAP(include_transformed=True)
        approx = pm.MeanField(model=model, start=start)
        pm.fit(
            random_seed=random_seed,
            n=n_init, method=pm.KLqp(approx),
            callbacks=cb,
            progressbar=progressbar,
            obj_optimizer=pm.adagrad_window
        )
        start = approx.sample(draws=chains)
        start = list(start)
        stds = approx.bij.rmap(approx.std.eval())
        cov = model.dict_to_array(stds) ** 2
        potential = quadpotential.QuadPotentialDiag(cov)
    elif init == 'map':
        start = pm.find_MAP(include_transformed=True)
        cov = pm.find_hessian(point=start)
        start = [start] * chains
        potential = quadpotential.QuadPotentialFull(cov)
    elif init == 'nuts':
        init_trace = pm.sample(draws=n_init, step=pm.NUTS(),
                               tune=n_init // 2,
                               random_seed=random_seed)
        cov = np.atleast_1d(pm.trace_cov(init_trace))
        start = list(np.random.choice(init_trace, chains))
        potential = quadpotential.QuadPotentialFull(cov)
    else:
        raise ValueError(
            'Unknown initializer: {}.'.format(init))

    step = pm.NUTS(potential=potential, model=model, **kwargs)

    return start, step


import numpy as np
import math, sys, time
import numpy as np
import scipy.linalg, scipy.constants
import matplotlib.pyplot as plt
from scipy import signal
import pymatgen as pmg
import itertools
from pymatgen.core import Lattice
from pymatgen.core.tensors import Tensor
from pymatgen.analysis.elasticity.elastic import ElasticTensor
from pymatgen.analysis.elasticity.elastic import ElasticTensorExpansion
import pandas as pd

data = np.genfromtxt("./pcu_2.log", skip_footer=29, skip_header=17)

a_mean = []
b_mean = []
c_mean = []
alpha_mean = []
beta_mean = []
gamma_mean = []
vol_mean = []
mechanical_eig_max = []
mechanical_eig_min = []
temp = []

h = []
for frame in data:
    lattice = Lattice.from_lengths_and_angles([frame[14], frame[15], frame[16]], [frame[17], frame[18], frame[19]])
    hmatrix = lattice.matrix
    h.append(hmatrix)

a_mean.append(np.mean(data[:,14]))
b_mean.append(np.mean(data[:,15]))
c_mean.append(np.mean(data[:,16]))
alpha_mean.append(np.mean(data[:,17]))
beta_mean.append(np.mean(data[:,18]))
gamma_mean.append(np.mean(data[:,19]))
vol_mean.append(np.mean(data[:,13]))
temp.append(np.mean(data[:,11]))

h = np.array(h)
h0 = np.mean(h, axis=0)
h0_std = np.std(h, axis=0)
h0m1  = np.linalg.inv(h0)
h0m1t = np.linalg.inv(np.transpose(h0))
eps_mat = np.array(list(map(lambda hi : (np.dot(h0m1t, np.dot(hi.transpose(), np.dot(hi, h0m1))) - np.identity(3))/2.0, h)))
eps_voigt = np.empty([6,eps_mat.shape[0]], "d")
eps_voigt[0] = eps_mat[:,0,0]
eps_voigt[1] = eps_mat[:,1,1]
eps_voigt[2] = eps_mat[:,2,2]
eps_voigt[3] = eps_mat[:,2,1]
eps_voigt[4] = eps_mat[:,2,0]
eps_voigt[5] = eps_mat[:,1,0]
eps_voigt_mean = np.mean(eps_voigt, axis=1)
Smat = np.zeros([6,6], "d")
for i in range(6):
    for j in range(i,6):
        Smat[i,j] = np.mean(eps_voigt[i]*eps_voigt[j])-eps_voigt_mean[i]*eps_voigt_mean[j]
        if i != j: Smat[j,i] = Smat[i,j]
T=300
V = np.linalg.det(h0)
Cmat = np.linalg.inv(Smat)*((1.3806488E-23*T)/(V*1E-30))*1E-9
Cmat = np.around(Cmat, decimals=2)
print(Cmat)
Cmat_eig = np.linalg.eigvals(Cmat)
mechanical_eig_max.append(max(Cmat_eig))
mechanical_eig_min.append(min(Cmat_eig))
print(mechanical_eig_max)
print(mechanical_eig_min)

mat = Tensor.from_voigt(Cmat)
elastic = ElasticTensor(mat)
print(elastic.k_vrh)
#np.savetxt("pcu_1.cij", Cmat)


import os
import sys

# PostGreSQL
import psycopg2

instrument_path_base = '/var/r1/'
www_path_base = '/var/www/html'

ac_config_dict = {
	'aircraft.tailnumber' : os.getenv("AIRCRAFT"),
	'project.name' : os.getenv("PROJECT"),
	'project.path' : os.getenv("PROJ_DIR")
	}

def connect_db():
  return psycopg2.connect(database="real-time", user="ads", host="acserver")

# For ground testing.
def connect_db_ground():
  return psycopg2.connect(database="real-time-C130", user="ads",
                          host="eol-rt-data.ml-ext.ucar.edu")

def get_config(key):

  # Open real-time onboard database
  db = connect_db()
  cur = db.cursor()

  # Get flight number from database.
  cur.execute("SELECT value FROM global_attributes WHERE key='FlightNumber';")
  rows = cur.fetchall()
  ac_config_dict['project.flightnumber'] = rows[0][0]

  # Assemble camera path.
  cpath = instrument_path_base + ac_config_dict['project.name']
  cpath += '/camera_images/flight_number_'
  cpath += ac_config_dict['project.flightnumber']

  ac_config_dict['cameras.path'] = cpath

  ac_config_dict['toga.path'] = instrument_path_base + \
                                ac_config_dict['project.name'] + '/toga'
  ac_config_dict['mtp.path'] = instrument_path_base + \
                               ac_config_dict['project.name'] + '/mtp'
  ac_config_dict['dropsonde.raw_path'] = instrument_path_base + '/dropsondes'
  ac_config_dict['dropsonde.skewt_path'] = www_path_base + '/skewt'

  db.close()

  return ac_config_dict.get(key, None)


from sqlite3 import Cursor, Connection, Row
from typing import Type, TypeVar

from db.base.backends import Backend
from db.base.fields import Field
from db.base.model import Model
from db.base.query import GetQuery, CreateTableQuery, InsertQuery, UpdateQuery
from db.base.utils import camel_to_snake

QueryClass = TypeVar('QueryClass')

class SQLException(Exception):
    pass

class SQLFactory:
    def __init__(self, instance_class: Type[Model], db_backend: Backend, connection):
        self.instance_class = instance_class
        self.db_backend = db_backend
        self.connection: Connection = connection

        # Эта строчка означает, что методы fetch будут возвращать словари с данными, а не кортежи
        self.connection.row_factory = Row
        self.cursor: Cursor = connection.cursor()

        assert issubclass(self.instance_class, Model), (
            '{} должен быть типом {}'.format(self.instance_class, Model)
        )

        self.table_name: str = camel_to_snake(self.instance_class.__name__)
        self.table_fields: list[Field] = self.instance_class.get_fields()

        self.raw_sql: str = ''
        self.__instance = None

    @property
    def instance(self):
        return self.__instance

    @instance.setter
    def instance(self, value):
        assert isinstance(value, self.instance_class), '"instance" должен быть типом {}'.format(self.instance_class)
        self.__instance = value

    @property
    def raw_sql(self):
        return self.__raw_sql

    @raw_sql.setter
    def raw_sql(self, value):
        assert isinstance(value, str), '"raw_sql" атрибут должен быть строкой'
        self.__raw_sql = value

    def clear_query(self) -> None:
        """ Очищает строку запроса SQL """

        self.raw_sql = ''

    def _fetch_one(self, fetched: list[Row]):
        """ Срабатывает при использовании self.perform_fetch(), если не передан параметр many """

        assert len(fetched) < 2, 'Запрос возвратил больше одного объекта'
        if len(fetched) == 1:
            row = fetched[0]
            return self.instance_class(**row)

    def _fetch_many(self, fetched: list[Row]):
        """ Срабатывает при использовании self.perform_fetch(), если передан параметр many=True """

        result = []
        for elem in fetched:
            result.append(self.instance_class(**elem))
        return result

    def perform_fetch(self, many: bool = False):
        assert 'SELECT' in self.raw_sql, "Вызвать метод 'perform_fetch' можно только с SELECT запросом"
        self.cursor.execute(self.raw_sql)

        fetch: list[Row] = self.cursor.fetchall()
        self.clear_query()

        if not many:
            return self._fetch_one(fetched=fetch)
        else:
            return self._fetch_many(fetched=fetch)

    def perform_update(self, commit: bool = True):
        """ Выполняет SQL запрос, вносящий какие либо изменения в БД, составленный в self.raw_sql """

        assert 'CREATE' in self.raw_sql or \
               'UPDATE' in self.raw_sql or \
               'DELETE' in self.raw_sql or \
               'INSERT' in self.raw_sql, (
            'Вызвать метод "perform_update" можно только с запросом, вносящим изменения в БД'
        )

        self.cursor.execute(self.raw_sql)
        self.clear_query()

        if commit:
            self.connection.commit()

    def _init_query_class(self, query_class: Type[QueryClass]) -> QueryClass:
        return query_class(
            instance=self.instance_class,
            table_name=self.table_name,
            table_fields=self.table_fields
        )

    def create_table(self) -> None:
        """
        Создает SQL запрос для создания таблицы на основе текущей модели в БД
        и записывает его в self.raw_sql
        """

        query_ = self._init_query_class(CreateTableQuery)
        self.raw_sql = query_.get_sql()

    def create(self, **kwargs) -> None:
        """
        Создает запрос для создания объекта текущей модели в БД
        и записывает его в self.raw_sql
        """

        query_ = self._init_query_class(InsertQuery)
        query_.insert_values(**kwargs)
        self.raw_sql = query_.get_sql()

    def get(self, all_tables: bool = False, **kwargs) -> None:
        assert all_tables or kwargs, 'Не переданы параметры фильтрации'

        query_ = self._init_query_class(GetQuery)
        if not all_tables:
            query_.where(**kwargs)
        self.raw_sql = query_.get_sql()

    def update(self, **kwargs):
        """
        Создает запрос для обновления данных объекта текущей модели в БД
        и записывает его в self.raw_sql
        """

        assert self.instance is not None, 'Необходимо определить атрибут "instance" прежде чем вызывать метод update()'
        query_ = self._init_query_class(UpdateQuery)
        query_.where(**kwargs)
        query_.bind_obj(self.instance)
        self.raw_sql = query_.get_sql()

if __name__ == '__main__':
    pass


from core.renderers import BlogmeJSONRenderer

class UserJSONRenderer(BlogmeJSONRenderer):
    charset = 'utf-8'
    object_label = 'user'

    def render(self, data, accepted_media_type=None, renderer_context=None):
        token = data.get('token', None)

        if token is not None and isinstance(token, bytes):
            data['token'] = token.decode('utf-8')

        return super(UserJSONRenderer,self).render(data)


import sys
from PyQt5.QtCore import QUrl, QEventLoop
from PyQt5.QtWidgets import QApplication, QWidget, QPushButton, QVBoxLayout
from PyQt5.QtWebEngineWidgets import QWebEngineView

APP = QApplication(['', '--no-sandbox'])

class Widget(QWidget):
    def __init__(self, url):
        super().__init__()
        self.url = url
        self.initUI()

    def initUI(self):

        # set title
        self.setWindowTitle('Query')

        # create web view
        web_view = QWebEngineView()
        web_view.setUrl(QUrl(self.url))

        # create button
        button = QPushButton('Get Query', self)
        button.resize(button.sizeHint())
        button.clicked.connect(lambda: self.set_url(web_view))

        # create layout
        lay = QVBoxLayout(self)
        lay.addWidget(button)
        lay.addWidget(web_view)

    def set_url(self, web_view):
        self.url = web_view.url().toString()
        self.close()

    def get_url(self):
        return self.url

def webView(url):

    app = QApplication.instance()
    app.processEvents(QEventLoop.ExcludeUserInputEvents | QEventLoop.ExcludeSocketNotifiers | QEventLoop.WaitForMoreEvents)

    widget = Widget(url)
    widget.showMaximized()

    return_code = app.exec_()
    if return_code:
        print(f"Return code: {return_code}... exiting")
        exit(return_code)

    app.quit()

    return widget.get_url()


from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys
import csv
import re
import time
import os

from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from seleniumwire import webdriver as wiredriver
from selenium.webdriver.chrome.service import Service as ChromeService

TIMER_SLEEP = 2
# Configurações do Chrome
chrome_options = Options()
# Descomente a linha abaixo se você quiser executar o Chrome em modo headless (sem interface gráfica)
# chrome_options.add_argument("--headless")

# Configurações do Selenium Wire
# proxy_address = 'https://179.83.82.10'  # Substitua pelo endereço do seu proxy
# seleniumwire_options = {
#     'backend': 'seleniumwire',
#     'proxy': {
#         'http': proxy_address,
#         'https': proxy_address,
#         'no_proxy': 'localhost,127.0.0.1',  # Domínios que devem ser acessados diretamente sem usar o proxy
#     }
# }

# # # Configuração do serviço do ChromeDriver
# # service = webdriver.chrome.service.Service('chromedriver.exe')

# # # Inicialize o driver do Chrome com as opções do Chrome e do Selenium Wire
# driver = wiredriver.Chrome(service=service, chrome_options=chrome_options, seleniumwire_options=seleniumwire_options)

# # Inicialize o servidor proxy do Selenium Wire
# proxy = driver.request_interceptor

# # Abra a página de login
# url = 'https://sales.ticketsforfun.com.br/#/authentication/login'
# driver.get(url)
# # Defina a função JavaScript a ser executada no navegador
# js_script = """
# function interceptXHR() {
#     console.log("Executando script");
#     try {
#         // Interceptar as requisições XHR
#         var originalOpen = window.XMLHttpRequest.prototype.open;
#         var originalSend = window.XMLHttpRequest.prototype.send;

#         window.XMLHttpRequest.prototype.open = function(method, url) {
#             // Salvar a URL da requisição para uso posterior
#             this._url = url;
#             originalOpen.apply(this, arguments);
#         };

#         window.XMLHttpRequest.prototype.send = function(data) {
#             // Verificar se a URL da requisição corresponde ao padrão desejado
#             if (this._url && this._url.startsWith('https://services.tix.byinti.com/neofront-v3/authentication/send-pin-code?')) {

#                 // Modificar a URL da requisição
#                 var modifiedUrl = this._url.replace('email=jalmirrosa6%40gmail.com', 'email=igorbzreis%40gmail.com');

#                 // Chamar o método originalOpen para aplicar a modificação
#                 originalOpen.call(this, 'GET', modifiedUrl, true);

#                 // Agora, chamar o método originalSend com os dados modificados
#                 originalSend.call(this, data);

#                 console.log("Requisição modificada e enviada:", modifiedUrl);
#             } else {
#                 // Se não for a requisição desejada, chamar o método originalSend sem modificação
#                 originalSend.call(this, data);
#             }
#         };
#     } catch (error) {
#         console.error('Erro na interceptação XHR:', error.message);
#     }
# }

# function addClickListenersToButtons() {
#     try {
#         // Adicionar um ouvinte de evento de clique a todos os botões
#         var buttons = document.querySelectorAll('button');
#         buttons.forEach(function(button) {
#             button.addEventListener('click', function(event) {
#                 console.log("Botão clicado! Texto do botão: " + event.target.innerText);
#                 // Chamar a função de interceptação de requisições ao clicar em um botão
#                 interceptXHR();
#             });
#         });
#     } catch (error) {
#         console.error('Erro ao adicionar ouvinte de evento de clique:', error.message);
#     }
# }

# // Executar a função de interceptação de requisições ao carregar a página
# interceptXHR();

# // Executar a função de adicionar ouvinte de evento de clique ao carregar a página
# addClickListenersToButtons();

# """

# # Execute o script JavaScript no navegador
# driver.execute_script(js_script)

# # Caminho do seu driver do Selenium (exemplo usando o Chrome)
# # Certifique-se de baixar o driver correspondente à sua versão do navegador

# # Aguarde alguns segundos para a página carregar completamente
# time.sleep(5)

# # Permita que o usuário interaja com a página
# input("Interaja com a página. Pressione Enter quando terminar.")

# # Fechar o navegador
# driver.quit()

# # Configuração do serviço do ChromeDriver
TIMER_SLEEP = 2
chrome_options = Options()
chrome_service = ChromeService(executable_path='chromedriver.exe')
chrome_options = webdriver.ChromeOptions()
# chrome_options.add_argument("--headless")
driver = webdriver.Chrome(service=chrome_service, options=chrome_options)

# 1° Acesse o link -
url = "https://busca.inpi.gov.br/pePI/jsp/marcas/Pesquisa_num_processo.jsp"
driver.get(url)
wait = WebDriverWait(driver, 10)
continuar_link = wait.until(EC.visibility_of_element_located((By.XPATH, "//a[contains(text(), 'Continuar...')]")))
continuar_link.click()
url = "https://busca.inpi.gov.br/pePI/jsp/marcas/Pesquisa_num_processo.jsp"
driver.get(url)
time.sleep(5)

# CODIGO QUE VAI SE REPETIR
def buscar_marca(num_pedido):
    print(f"Esperando {TIMER_SLEEP} segundos...")
    time.sleep(TIMER_SLEEP)
    num_processo = ""
    nome_marca = ""
    situacao = ""
    apresentacao = ""
    natureza = ""
    classificacao_nice = ""
    data_de_deposito = ""
    classificacao_de_viena = ""
    texto_do_despacho = ""

    url = "https://busca.inpi.gov.br/pePI/jsp/marcas/Pesquisa_num_processo.jsp"
    driver.get(url)

    try:
        # 2° Busque o input com name="NumPedido" e insira o código
        # num_pedido = "927361949"  # Substitua pelo número de pedido desejado
        input_num_pedido = driver.find_element(By.NAME, "NumPedido")
        input_num_pedido.send_keys(num_pedido)
        input_num_pedido.send_keys(Keys.RETURN)
    except Exception as e:
        print("823 - Erro ao procurar por processo - ", e)

    try:
        # Encontre o elemento <a> com o atributo href contendo "/pePI/servlet/MarcasServletController?Action=detail&amp"
        elemento_a = driver.find_element(By.XPATH,f"//a[contains(text(), '{num_pedido}')]")
        href = elemento_a.get_attribute("href")
        driver.get(href)
    except Exception as e:
        print("753 - Não encontrado na busca - ", e)

    try:
    # 4° Na página aberta, faça um crawler das informações da penúltima tabela
    # Aguarde até que a penúltima tabela seja visível
        tabelas = wait.until(EC.presence_of_all_elements_located((By.TAG_NAME, "table")))

        informacoes = {}
        # Localize os elementos <td> que contêm os rótulos
        rotulos = driver.find_elements(By.XPATH,"//td[contains(font, 'Nº do Processo:') or contains(font, 'Marca:') or contains(font, 'Situação:') or contains(font, 'Apresentação:') or contains(font, 'Natureza:') or contains(font, 'Apostila :')]")
        # Itere sobre os rótulos e capture os valores correspondentes
        for rotulo in rotulos:
            chave = rotulo.find_element(By.XPATH, "font").text.strip().replace(":", "")
            valor = rotulo.find_element(By.XPATH, "following-sibling::td/font").text.strip()
            informacoes[chave] = valor

        # Exiba as informações coletadas
        for chave, valor in informacoes.items():
            num_processo = informacoes['Nº do Processo']
            nome_marca = informacoes['Marca']
            situacao = informacoes['Situação']
            apresentacao = informacoes['Apresentação']
            natureza = informacoes['Natureza']

            # print(f"{chave}: {valor}")
    except Exception as e:
        print(f"423 - Erro ao coletar informações da tabela: {num_pedido} - ", e)

    # CAPTURAR CLASSIFICAÇÃO NICE
    try:
        class_nice = driver.find_element(By.XPATH, '//*[@id="principal"]/div[3]/div/div/table/tbody/tr/td[1]/font/a')
        classificacao_nice = class_nice.text
        # print("Classificacao NICE: ", class_nice.text)
    except Exception as e:
        print("831 - Erro ao capturar classificação NICE - ", e)

    #CAPTURAR A DATA DE DEPOSITO
    try:
       # Encontre a tabela com o elemento font que contém 'Data de Depósito'
        table_element = driver.find_element(By.XPATH, "//table[contains(.//font, 'Data de Depósito')]")

        # Em seguida, encontre o primeiro elemento th dentro do tbody
        th_element = table_element.find_element(By.XPATH, ".//tbody/tr/th[1]")

        # Capture o texto do primeiro th
        data_de_deposito = th_element.text
    except Exception as e:
        print("712 - Erro ao capturar data de deposito - ", e)

    # CAPTURAR CLASSIFICAÇÃO DE VIENA
    try:
        # Encontre a tabela usando o XPath especificado
        tabela = driver.find_element(By.XPATH, "//*[@id='principal']/div[4]/div/div/table")
        # Encontre todas as linhas da tabela dentro do corpo (elementos <tr>)
        linhas = tabela.find_elements(By.XPATH, ".//tbody/tr")
        # Itere sobre as linhas para obter os valores
        valores_totais = []
        for linha in linhas:
            # Encontre as colunas em cada linha (elementos <td>)
            colunas = linha.find_elements(By.XPATH, ".//td")
            # Obtenha os valores de cada coluna
            valores = [coluna.text for coluna in colunas]

            valores_totais.append(valores)
            # Faça algo com os valores, por exemplo, imprima-os
            # print("Classificacao de VIENA: ", valores)

        classificacao_de_viena = valores_totais
    except Exception as e:
        print("483 - Erro ao capturar classificação de VIENA - ", e)

    try:
        # Definir o padrão de regex
        padrao = r"A marca reproduz ou imita os seguintes registros de terceiros[^;]*"
        correspondencias = re.search(padrao, driver.page_source)
        if correspondencias:
            texto_do_despacho = correspondencias.group(0)
            # print("Texto completo: ", texto_do_despacho)

        else:
            print("Padrão não encontrado na página.")

    except Exception as e:
        print("9123 - Erro ao capturar detalhe do despacho:", e)

    try:
        #SALVANDO IMAGEM
        time.sleep(1)
        url = driver.current_url
        match = re.search(r'CodPedido=(\d+)', url)
        cod_pedido = ""
        if match:
            cod_pedido = match.group(1)
            # print("CodPedido:", cod_pedido)

        linkimg = f"https://busca.inpi.gov.br/pePI/servlet/LogoMarcasServletController?Action=image&codProcesso={cod_pedido}".split("/")[-1]

        elemento_img = driver.find_element(By.CSS_SELECTOR, f"img[src*='{linkimg}']")
        # print(elemento_img, linkimg)

        with open(f'imgs/imgs_revista/{num_pedido}.png', 'wb') as file:
            file.write(elemento_img.screenshot_as_png)
    except Exception as e:
        print(f"923 - Processo {cod_pedido} sem imagem ")

    return num_processo, nome_marca, situacao, apresentacao, natureza, classificacao_nice, data_de_deposito, classificacao_de_viena, texto_do_despacho

cabecalho = [
                'N° Processo',
                'Marca Indeferida',
                'Situação',
                'Apresentação',
                'Natureza',
                'Classe de Nice',
                'Classificação de Viena',
                'Data de Deposito',

                'N° Processo',
                'Marca Registrada',
                'Situação',
                'Apresentação',
                'Natureza',
                'Classe de Nice',
                'Classificação de Viena',
                'Data de Deposito',

                'Despacho',
                'Revista',
                'Visual',
                'Escrita',
                'Fonética'
            ]

def processo_ja_existe(processo_indeferido, processo_deferido, nome_arquivo):
    with open(nome_arquivo, 'r', newline='', encoding='utf-8') as arquivo_csv:
        leitor = csv.reader(arquivo_csv, delimiter='\t')

        # Pula o cabeçalho
        next(leitor, None)

        # Verifica se o número do processo já existe nas colunas 1 e 9
        for linha in leitor:
            if linha[0] == processo_indeferido or linha[8] == processo_deferido:
                return True
    return False

for revista in os.listdir('revistas'):
    print(f"Analisando revista {revista}")

    total_analisado = 0
    dados = []

    try:
        with open(f'revistas/{revista}', newline='', encoding='utf-8') as arquivo_csv:
            leitor_csv = csv.reader(arquivo_csv, delimiter="\t")
            next(leitor_csv)
            for linha in leitor_csv:
                processo_indeferido = linha[1]

                texto_complementar = linha[3]
                numeros_de_processo = re.findall(r'Processo (\d+)', texto_complementar)
                print("Processo Indeferido: ", processo_indeferido, "Processos deferidos: ", numeros_de_processo)

                # busco pelo processo indeferido
                num_processo_in, nome_marca_in, situacao_in, apresentacao_in, natureza_in, classificacao_nice_in, data_de_deposito_in, classificacao_de_viena_in, texto_do_despacho_in = buscar_marca(processo_indeferido)
                print("Marca indeferida: ", num_processo_in, nome_marca_in, situacao_in, apresentacao_in, natureza_in, classificacao_nice_in, data_de_deposito_in, classificacao_de_viena_in )
                # agora busco pelos processo deferido (marcasregistradas)

                for processo_deferido in numeros_de_processo:
                #crawler para buscar processos indeferidos
                    if not processo_ja_existe(processo_indeferido, processo_deferido,   f'{revista}_analisada.csv'):
                        num_processo, nome_marca, situacao, apresentacao, natureza, classificacao_nice, data_de_deposito, classificacao_de_viena, texto_do_despacho = buscar_marca(processo_deferido)
                        print("Marca deferida: ", num_processo, nome_marca, situacao, apresentacao, natureza, classificacao_nice, data_de_deposito, classificacao_de_viena)

                        dados.append(
                            [num_processo_in, nome_marca_in, situacao_in, apresentacao_in, natureza_in, classificacao_nice_in, classificacao_de_viena_in, data_de_deposito_in,
                            num_processo, nome_marca, situacao, apresentacao, natureza, classificacao_nice, classificacao_de_viena, data_de_deposito,
                            texto_do_despacho_in, revista, 0, 0, 0]
                        )

                        novo_registro = [num_processo_in, nome_marca_in, situacao_in, apresentacao_in, natureza_in, classificacao_nice_in, classificacao_de_viena_in, data_de_deposito_in,
                            num_processo, nome_marca, situacao, apresentacao, natureza, classificacao_nice, classificacao_de_viena, data_de_deposito,
                            texto_do_despacho_in, revista, 0, 0, 0]

                        nome_arquivo = f'{revista}_analisada.csv'
                        arquivo_existe = os.path.exists(nome_arquivo)

                        with open(nome_arquivo, 'a', newline='', encoding='utf-8') as arquivo_csv:
                            escritor = csv.writer(arquivo_csv, delimiter='\t')

                            if not arquivo_existe or arquivo_csv.tell() == 0:
                                escritor.writerow(cabecalho)

                            escritor.writerow(novo_registro)

                    else:
                        print(f"Processo {processo_indeferido} e processo {processo_deferido} já analisado!")

                    total_analisado += 1

                print("Total Analisado: ", total_analisado)

    except Exception as e:
        print(f"Erro ao analisar revista {revista}. Total Analisado: {total_analisado} -- ", e)

    print("Salvando informações...")

    # cabecalho = [
    #             'N° Processo',
    #             'Marca Indeferida',
    #             'Situação',
    #             'Apresentação',
    #             'Natureza',
    #             'Classe de Nice',
    #             'Classificação de Viena',
    #             'Data de Deposito',

    #             'N° Processo',
    #             'Marca Registrada',
    #             'Situação',
    #             'Apresentação',
    #             'Natureza',
    #             'Classe de Nice',
    #             'Classificação de Viena',
    #             'Data de Deposito',

    #             'Despacho',
    #             'Revista',
    #             'Visual',
    #             'Escrita',
    #             'Fonética'
    #         ]

    # nome_arquivo = f'{revista}_analisada.csv'
    # # Crie e escreva o arquivo CSV
    # with open(nome_arquivo, 'w', newline='', encoding='utf-8') as arquivo_csv:
    #     escritor = csv.writer(arquivo_csv, delimiter='\t')
    #     escritor.writerow(cabecalho)
    #     for linha in dados:
    #         escritor.writerow(linha)


# -*- coding: utf-8 -*-
# __auther__ : "johnny leaf"
# __date__: "2018.8.9"
# __project name__ : "dianping_spider"
# __CODE__: Python3

import time
import datetime
import sys

from Edp_excel import report
from Edp_excel import DataAn
from Spider_edp import dpspider

LoginUrl = 'https://e.dianping.com'

AcounDic = {
    '测试': [],
    # '八大处亚运村': ['bjykyl', 'ykyl180225', '90030589'],
    # '八大处平安门诊部': ['13883278696', 'bjpa888', '96640607'],
    # '煤炭总医院': ['mtzyyzx', 'meitan000', '98380431'],
    # '北京炫美': ['BJxuanmei123456', 'xmzx123456777', '92742255'],
    # '成都健丽': ['cdjianli', 'cdjs12345', '73082729'],
    # '成都美瑞': ['MeiruiTF', '93963llq', '8352512'],
    # '成都僮颜': ['tongyan88888', 'tongyan61', '98951640'],
    # '德尔美客': ['deermeike', '30826rpo', '97312957'],
    # '瑷珊': ['ieshan23', '37878jmd', '69046169'],
    # '和谐同方': ['HX73357653', '65787488', '4255672'],
    # '时光': ['tjsgzx520', 'tjsgzx2015', '59241747'],
    # '凤凰怡美': ['fenghuangyimei', 'fhymfh1234', '90290461']
}

def GetStartEnddate(label):
    timelist = []
    today = datetime.date.today()
    if label == '半月报':
        day = 1
        timelist.append(str(
            datetime.date(
                today.year - (
                    today.month == 1
                    ), today.month - day or day+11, 1)))
        timelist.append(str(today))
    elif label == '月报':
        starttime = datetime.date(today.year - (today.month == 1), today.month - 2 or 13, 1)
        endtime = datetime.date(today.year - (today.month == 1), today.month - 1 or 12, 31 or 30)
        timelist.append(str(starttime))
        timelist.append(str(endtime))
    else:
        timelist = label.split(' ')
    return timelist

def main(TimeInfo):
    for acount in AcounDic:
        DataAn.main(TimeInfo, acount)
        # res = dpspider.Get_CookeandSession(LoginUrl, AcounDic[acount])
        # print('%s 已取得Cookies参数' % acount)
        # if res:
        #     flowresult = dpspider.Getflowdata(res, AcounDic[acount])
        #     print('%s 已爬取流量数据' % acount)
        #     chatreuslt = dpspider.Getchatdata(res, AcounDic[acount], TimeInfo)
        #     print('%s 已爬取数咨询据' % acount)
        #     appointmentresult = dpspider.GetAppointresult(res, AcounDic[acount])
        #     print('%s 已爬取预约数据' % acount)
        #     SaleOnlineresult = dpspider.GetSaleOnlineresult(res)
        #     print('%s 已爬取线上销售数据' % acount)
        #     CommentResult = dpspider.GetCommentResult(res, AcounDic[acount])
        #     print('%s 已爬取体验报告' % acount)
        #     report.Report_main(flowresult, chatreuslt, appointmentresult, SaleOnlineresult, CommentResult, acount)
        #     print('====%s报表制作完成====' % acount)
        #     print('==========================')
        #     DataAn.main(TimeInfo, acount)
        # else:
        #     time.sleep(1)

def GetUserinput(inputinfo):
    if inputinfo == '半月报':
        timeresult = GetStartEnddate('半月报')
        print('==========%s模式==========' % inputinfo)
    elif inputinfo == '月报':
        timeresult = GetStartEnddate('月报')
        print('==========%s模式==========' % inputinfo)
    else:
        timeresult = GetStartEnddate(inputinfo)
        print('======%s至%s=====' % (timeresult[0], timeresult[1]))
    return timeresult

if __name__ == '__main__':
    userInput = sys.argv[1]
    TimeInfo = GetUserinput(userInput)
    main(TimeInfo)


"""Functions to create loggers"""

import logging
from typing import Text
import sys

def get_console_handler() -> logging.StreamHandler:
    """Get console handler

    Returns:
        logging.StreamHandler which logs into stdout

    """
    console_handler = logging.StreamHandler(sys.stdout)
    formatter = logging.Formatter("%(asctime)s — %(name)s — %(levelname)s — %(message)s")
    console_handler.setFormatter(formatter)

    return console_handler

def get_logger(name: Text = __name__,
               log_level: Text or int = logging.DEBUG) -> logging.Logger:
    """Get logger

    Params:
        name {Text}: logger name
        log_level {Text or int}: logging level; can be string name or integer value

    Returns:
        logging.Logger instance

    """
    logger = logging.getLogger(name)
    logger.setLevel(log_level)

    # Prevent duplicate outputs in Jupyter Notebook
    if logger.hasHandlers():
        logger.handlers.clear()

    logger.addHandler(get_console_handler())
    logger.propagate = False

    return logger


"""
right_sett.py

Right side of the setting page of the app.
"""

#-------------------------------------------------------------------#

from src.utils.gui_utils import Frame
from src.interface.menus.settings.right.sett_header import SettHeader
from src.interface.menus.settings.right.sett_body import SettBody

#-------------------------------------------------------------------#

class RightSett(Frame):
    """
    Container of the header, body and footer of the shopping page.
    """
    def __init__(self, manager=None):
        super().__init__(manager)
        self.manager = manager
        self.grid_propagate(False)

        self.header = SettHeader(self)
        self.header.grid(row=0, column=0, sticky="nsew")
        self.body = SettBody(self)
        self.body.grid(row=1, column=0, sticky="nsew")

        self.columnconfigure(0, weight=1)
        self.rowconfigure(0, weight=1)
        self.rowconfigure(1, weight=4)