#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
A kernel-density-based, embarrassingly parallel ensemble sampler.
"""

from __future__ import (division, print_function, absolute_import, unicode_literals)
from .pool import choose_pool
import numpy as np
import numpy.ma as ma
from scipy.stats import chisquare
from .clustered_kde import optimized_kde, TransdimensionalKDE

def dynamic_pbar_update(step_size, last_step_size, acc, pbar=None):
    if pbar is not None:
        pbar.set_postfix_str('| Single-Step Acceptence Rate: {0:.3f}'.format(acc))
        pbar.update(step_size-last_step_size)

def in_notebook():
    import __main__ as main
    return not hasattr(main, '__file__')

class _GetLnProbWrapper(object):
    """Convenience class for evaluating multiple probability densities at a single point."""
    def __init__(self, lnpost, kde, *args):
        self.lnpost = lnpost
        self.kde = kde
        self.args = args

    def update_kde(self, kde):
        self.kde = kde

    def lnprobs(self, p):
        """
        Evaluate the log probability density of the stored target distribution fuction
        and KDE at `p`.

        :param p: Location to evaluate probability densties at.

        :returns: ``lnpost(p)``, ``kde(p)``
        """
        result = self.lnpost(p, *self.args)
        kde = self.kde(p)

        # allow posterior function to optionally
        # return additional metadata
        try:
            lnpost = result[0]
            blob = result[1]
            return lnpost, kde, blob
        except (IndexError, TypeError):
            lnpost = result
            return lnpost, kde

    __call__ = lnprobs

_lnprob_wrapper = None

def _set_global_lnprob_wrapper(wrapper_instance):
    """Sets `_lnprob_wrapper` to a global variable equal to the provided
    instance of `_GetLnProbWrapper`. Making the wrapper a global variable
    allows chilren processes to access their own copy of the wrapper without
    needing to push out the same data on every call. For this to work it must
    be done *before* a multiprocessing pool is initialized. See
    http://stackoverflow.com/a/10118250 for details.
    """
    global _lnprob_wrapper
    _lnprob_wrapper = wrapper_instance

def _get_lnprob_from_wrapper(p):
    return _lnprob_wrapper.lnprobs(p)

def _update_wrapper_kde(kde):
    _lnprob_wrapper.update_kde(kde)

class Sampler(object):
    """
    An Ensemble sampler.

    The :attr:`chain` member of this object has the shape: `(nsteps, nwalkers, ndim)` where
    `nsteps` is the stored number of steps taken thus far.

    :param nwalkers:
        The number of individual MCMC chains to include in the ensemble.

    :param ndim:
        Number of dimensions in the parameter space.  If `transd` is ``True`` this is the number of
        unique dimensions across the parameter spaces.

    :param lnpostfn:
        A function that takes a vector in the parameter space as input and returns the natural
        logarithm of the posterior probability for that position.

    :param transd:
        If ``True``, the sampler will operate across parameter spaces using a
        :class:`.clustered_kde.TransdimensionalKDE` proposal distribution. In this mode a masked
        array with samples in each of the possible sets of dimensions must be given for the initial
        ensemble distribution.

    :param processes: (optional)
        The number of processes to use with :mod:`multiprocessing`.  If ``None``, all available
        cores are used.

    :param pool: (optional)
        A pre-constructed pool with a map method. If ``None`` a pool will be created using
        :mod:`multiprocessing`.

    """
    def __init__(self, nwalkers, ndim, lnpostfn, transd=False,
                 processes=None, pool=None, mpi=False, args=[]):
        self.nwalkers = nwalkers
        self.dim = ndim
        self._kde = None
        self._kde_size = self.nwalkers
        self._updates = []
        self._burnin_length = None

        self._get_lnpost = lnpostfn
        self._lnpost_args = args

        self.iterations = 0
        self.stored_iterations = 0
        self.mpi = mpi
        self.processes = processes
        if processes is None and pool is not None:
            raise ValueError("Please provide the number of processes"
                             " if also providing a pool instance.")
        self._managing_pool = False
        if pool is not None:
            self._pool = pool
        else:
            self._managing_pool = True
            self._pool = choose_pool(processes, mpi=mpi)

        if not hasattr(self._pool, 'map'):
            raise AttributeError("Pool object must have a map() method.")

        self._transd = transd
        if self._transd:
            self._chain = ma.masked_all((0, self.nwalkers, self.dim))
        else:
            self._chain = np.zeros((0, self.nwalkers, self.dim))

        self._lnpost = np.empty((0, self.nwalkers))
        self._lnprop = np.empty((0, self.nwalkers))
        self._acceptance = np.zeros((0, self.nwalkers))
        self._blobs = []
        self.tqdm = None

        self._last_run_mcmc_result = None
        self._burnin_spaces = None
        self._failed_p = None
        self._set_wrapper()

    def _set_wrapper(self):
        """ Determine the lnprob wrapper """
        # If the pool can guarantee a call to all processes we can update
        # the data explicitly, which is much faster
        if hasattr(self._pool, 'broadcast'):
            wrapper = _GetLnProbWrapper(self._get_lnpost, self._kde,

                                        *self._lnpost_args)
            self._pool.broadcast(_set_global_lnprob_wrapper, wrapper)

    def _get_wrapper(self):
        """ Retunr the lnprob wrapper function call """
        if hasattr(self._pool, 'broadcast'):
            return _get_lnprob_from_wrapper
        else:
            return _GetLnProbWrapper(self._get_lnpost, self._kde,

                                     *self._lnpost_args)

    def _get_dynamci_pbar(self, progress, t):
        pbar = None
        if progress:
            import tqdm
            self.tqdm = tqdm
            if in_notebook():
                pbar = self.tqdm.tqdm_notebook(desc="Burning in until Const Acc Rate over {} ACLs".format(t), position=0, leave=True,
                                               total=t)
            else:
                pbar = self.tqdm.tqdm(desc="Burning in until Const Acc Rate over {} ACLs".format(t), position=0, leave=True,
                                      total=t)
        return pbar, dynamic_pbar_update

    def _get_finite_pbar(self, progress, N):
        pbar = None
        if progress:
            if self.tqdm is None:
                import tqdm
                self.tqdm = tqdm
            if in_notebook():
                pbar = self.tqdm.tqdm_notebook(total=N, desc="Running MCMC")
            else:
                pbar = self.tqdm.tqdm(total=N, desc="Running MCMC")
        return pbar

    def burnin(self, p0=None, lnpost0=None, lnprop0=None, blob0=None,
               test_steps=16, critical_pval=0.05, max_steps=None,
               verbose=False, callback=None, progress=False, **kwargs):
        """
        Evolve an ensemble until the acceptance rate becomes roughly constant.  This is done by
        splitting acceptances in half and checking for statistical consistency.  This isn't
        guaranteed to return a fully burned-in ensemble, but usually does.

        :param p0: (optional)
            A list of the initial walker positions.  It should have the shape `(nwalkers, ndim)`.
            If ``None`` and the sampler has been run previously, it'll pick up where it left off.

        :param lnpost0: (optional)
            The list of log posterior probabilities for the walkers at positions `p0`. If ``lnpost0
            is None``, the initial values are calculated. It should have the shape `(nwalkers, )`.

        :param lnprop0: (optional)
            List of log proposal densities for walkers at positions `p0`. If ``lnprop0 is None``,
            the initial values are calculated. It should have the shape `(nwalkers, )`.

        :param blob0: (optional)
            The list of blob data for walkers at positions `p0`.

        :param test_steps: (optional)
            The (rough) number of accepted steps over which to check for acceptance rate
            consistency. If you find burnin repeatedly ending prematurely try increasing this.

        :param critical_pval: (optional)
            The critial p-value for considering the acceptance distribution
            consistent.  If the calculated p-value is over this, then the acceptance rate
            is considered to be stationary.  Lower this if you want burnin criteria to be
            less strict.

        :param max_steps: (optional)
            An absolute maximum number of steps to take, in case burnin is too painful.

        :param verbose: (optional)
            Print status messages each time a milestone is reached in the burnin.

        :param callback: (optional)
            Callback to use (default to None)

        :param progress: (optional)
            Boolean flag to report progress of burnin() / MCMC steps through a progress bar. Needs tqdm installed

        :param kwargs: (optional)
            The rest is passed to :meth:`run_mcmc`.

        After burnin...

        :returns:
            * ``p`` - A list of the current walker positions with shape `(nwalkers, ndim)`.

            * ``lnpost`` - Array of log posterior probabilities for walkers at positions `p`; has
              shape `(nwalkers, )`.

            * ``lnprop`` - Array of log proposal densities for walkers at positions `p`; has shape
              `(nwalkers, )`.

            * ``blob`` - (if `lnprobfn` returns blobs) The list of blob data for the walkers at
              positions `p`.
        """
        if p0 is not None:
            p0 = np.atleast_2d(p0)

        # Determine the maximum iteration to look for
        start = self.iterations
        # Confine walkers to their space during burnin
        freeze_transd = False
        if self._transd:
            freeze_transd = True
            self._burnin_spaces = ~p0.mask
        pbar, pbar_status = self._get_dynamci_pbar(progress, test_steps)
        max_iter = np.inf
        if max_steps is not None:
            max_iter = start + max_steps
        if progress:
            verbose = False

        sampling_ct = 0
        step_size = 2
        last_step_size = 0
        while step_size <= test_steps:
            # Update the proposal
            if p0 is not None:
                self.update_proposal(p0, max_samples=self.nwalkers)
                lnprop0 = self._kde(p0)
            if verbose:
                print('Updated proposal')

            # Take one step to estimate acceptance rate
            test_interval = 1
            results = self.run_mcmc(test_interval, p0, lnpost0, lnprop0, blob0,
                                    freeze_transd=freeze_transd, spaces=self._burnin_spaces,
                                    **kwargs)
            try:
                p, lnpost, lnprop, blob = results
            except ValueError:
                p, lnpost, lnprop = results
                blob = None

            # Use the fraction of acceptances in the last step to estimate acceptance rate
            #   Bottom out at 1% if acceptances are really bad
            last_acc_rate = max(np.mean(self.acceptance[-1]), 0.01)

            # Estimate ACT based on acceptance
            act = int(np.ceil(2.0/last_acc_rate - 1.0))

            if verbose:
                print('Single-step acceptance rate is ', last_acc_rate)
                print('Producing ACT of ', act)

            # Use the ACT to set the new test interval, but avoid
            # overstepping a specified max.  We throw away the first
            # 2*act worth of steps as an initial burnin when comparing
            # acceptance rates
            test_interval = min((step_size+2)*act, max_iter - self.iterations)

            # Make sure we're taking at least one step
            test_interval = max(test_interval, 1)
            sampling_ct += 1

            if verbose:
                if ~np.isinf(max_iter):
                    print("Time {} running mcmc during burnin with {}/{} Iterations and test_interval = {}:".format(
                        sampling_ct, self.iterations, max_iter, test_interval))
                else:
                    print("Time {} running mcmc during burnin with test_interval = {}:".format(
                        sampling_ct, test_interval))

            results = self.run_mcmc(test_interval, p, lnpost, lnprop, blob,
                                    freeze_transd=freeze_transd, spaces=self._burnin_spaces, progress=progress,
                                    **kwargs)
            try:
                p, lnpost, lnprop, blob = results
            except ValueError:
                p, lnpost, lnprop = results
                blob = None

            if callback is not None:
                callback(self)

            # Quit if we hit the max
            if self.iterations >= max_iter:
                print("Burnin hit {} iterations before completing.".format(max_iter))
                break

            # Only check for consistency past the burn-in stage of 2*act
            if self.consistent_acceptance_rate(window_size=step_size*act, critical_pval=critical_pval):
                if verbose:
                    print('Acceptance rate constant over ', step_size, ' ACTs')
                pbar_status(step_size, last_step_size, last_acc_rate, pbar)
                last_step_size = step_size
                step_size *= 2
            else:
                if verbose:
                    print('Acceptance rate varies, trying again')

            if verbose:
                print('') # Newline

            p0, lnpost0, lnprop0, blob0 = p, lnpost, lnprop, blob

        self._burnin_length = self.updates[-1]
        if pbar is not None:
            pbar.close()

        if blob is None:
            return p, lnpost, lnprop
        else:
            return p, lnpost, lnprop, blob

    def sample(self, p0=None, lnpost0=None, lnprop0=None, blob0=None,
               iterations=1, kde=None, update_interval=None, kde_size=None,
               freeze_transd=False, spaces=None, storechain=True, **kwargs):
        """
        Advance the ensemble `iterations` steps as a generator.

        :param p0: (optional)
            A list of the initial walker positions.  It should have the shape `(nwalkers, ndim)`.
            If ``None`` and a proposal distribution exists, walker positions will be drawn from the
            proposal.

        :param lnpost0: (optional)
            The list of log posterior probabilities for the walkers at positions `p0`. If ``lnpost0
            is None``, the initial values are calculated. It should have the shape `(nwalkers, )`.

        :param lnprop0: (optional)
            List of log proposal densities for walkers at positions `p0`. If ``lnprop0 is None``,
            the initial values are calculated. It should have the shape `(nwalkers, )`.

        :param blob0: (optional)
            The list of blob data for walkers at positions `p0`.

        :param iterations: (optional)
            The number of steps to run.

        :param kde: (optional)
            An already-constucted, evaluatable KDE with a ``draw`` method.

        :param update_interval: (optional)
            Number of steps between proposal updates.

        :param kde_size: (optional)
            Maximum sample size for KDE construction.  When the KDE is updated, existing samples are
            thinned by factors of two until there's enough room for `nwalkers` new samples.  The
            default is `nwalkers`, and must be greater than :math:`\geq``nwalkers` if specified.

        :param freeze_transd: (optional)
            If ``True`` when transdimensional sampling, walkers are confined to their parameter
            space.  This is helpful during burnin, and allows fox fixed-D burnin before
            transdimensional sampling.

        :param spaces: (optional)
            Confine walkers to the requested parameter spaces. Expects an inverted mask with shape
            `(nwalkers, ndim)`.

        :param storechain: (optional)
            Flag for disabling chain and probability density storage in :attr:`chain`,
            :attr:`lnpost`, and :attr:`lnprop`.

        :param kwargs: (optional)
            The rest is passed to :meth:`update_proposal`.

        After each iteration...

        :yields:
            * ``p`` - An array of current walker positions with shape `(nwalkers, ndim)`.

            * ``lnpost`` - The list of log posterior probabilities for the walkers at positions
              ``p``, with shape `(nwalkers, )`.

            * ``lnprop`` - The list of log proposal densities for the walkers at positions `p`, with
              shape `(nwalkers, )`.

            * ``blob`` - (if `lnprobfn` returns blobs) The list of blob data for the walkers at
              positions `p`.
        """
        if p0 is None:
            p = self.draw(self.nwalkers)
        else:
            try:
                # Try copying to preserve array type (i.e.masked or not)
                p = p0.copy()
            except AttributeError:
                # If not already an array, make it a non-masked array by default.
                #   Operations with masked arrays can be slow.
                p = np.array(p0, copy=True)

        if kde_size:
            self._kde_size = kde_size

        # Build a proposal if one doesn't already exist
        if kde is not None:
            self.set_kde(kde)
        elif self._kde is None:
            self.update_proposal(p, max_samples=self._kde_size, **kwargs)
            lnprop0 = self._kde(p)

        lnpost = lnpost0
        lnprop = lnprop0
        blob = blob0

        if lnpost is None or lnprop is None:
            results = list(self._pool.map(self._get_wrapper(), p))
            lnpost = np.array([r[0] for r in results]) if lnpost is None else lnpost
            lnprop = np.array([r[1] for r in results]) if lnprop is None else lnprop

            if blob is None:
                try:
                    blob = [r[2] for r in results]
                except IndexError:
                    blob = None

        # ensure blob has the right shape
        if blob is None:
            blob = [None]*self.nwalkers

        # Prepare arrays for storage ahead of time
        self._acceptance = np.concatenate((self._acceptance,
                                           np.zeros((iterations, self.nwalkers))))
        if storechain:
            # Make sure to mask things if the stored chain has a mask
            if hasattr(self._chain, "mask"):
                self._chain = ma.concatenate((self._chain,
                                              ma.masked_all((iterations, self.nwalkers, self.dim))))
            else:
                self._chain = np.concatenate((self._chain,
                                              np.zeros((iterations, self.nwalkers, self.dim))))

            self._lnpost = np.concatenate((self._lnpost, np.zeros((iterations, self.nwalkers))))
            self._lnprop = np.concatenate((self._lnprop, np.zeros((iterations, self.nwalkers))))

        for i in range(iterations):
            try:
                if freeze_transd and spaces is None:
                    spaces = ~p.mask
                # Draw new walker locations from the proposal
                p_p = self.draw(self.nwalkers, spaces=spaces)

                # Calculate the posterior probability and proposal density
                # at the proposed locations
                try:
                    results = list(self._pool.map(self._get_wrapper(), p_p))
                    lnpost_p = np.array([r[0] for r in results])
                    lnprop_p = np.array([r[1] for r in results])
                    try:
                        blob_p = [r[2] for r in results]
                    except IndexError:
                        blob_p = None

                # Catch any exceptions and exit gracefully
                except Exception as e:
                    self.rollback(self.stored_iterations)
                    self._failed_p = p_p

                    print("Offending samples stored in ``failed_p``.")
                    raise

                # Calculate the (ln) Metropolis-Hastings ratio
                ln_mh_ratio = lnpost_p - lnpost + lnprop - lnprop_p

                # Accept if ratio is greater than 1
                acc = ln_mh_ratio > 0

                # Decide which of the remainder will be accepted
                worse = ~acc
                nworse = np.count_nonzero(worse)
                uniform_draws = np.random.rand(nworse)
                acc[worse] = ln_mh_ratio[worse] > np.log(uniform_draws)

                # Update locations and probability densities
                if np.any(acc):
                    p[acc] = p_p[acc]
                    lnpost[acc] = lnpost_p[acc]
                    lnprop[acc] = lnprop_p[acc]

                    if blob_p is None:
                        blob = None
                    else:
                        blob = [blob_p[i] if a else blob[i] for i, a in enumerate(acc)]
                self._acceptance[self.iterations, :] = acc

                # Update the proposal at the requested interval
                if self.trigger_update(update_interval):
                    self.update_proposal(p, max_samples=self._kde_size, **kwargs)
                    lnprop = self._kde(p)

                if storechain:
                    # Store stuff
                    self._chain[self.stored_iterations, :, :] = p
                    self._lnpost[self.stored_iterations, :] = lnpost
                    self._lnprop[self.stored_iterations, :] = lnprop

                    if blob:
                        self._blobs.append(blob)

                    self.stored_iterations += 1
                self.iterations += 1

                # create generator for sampled points
                if blob:
                    yield p, lnpost, lnprop, blob
                else:
                    yield p, lnpost, lnprop

            except KeyboardInterrupt:
                # Resize arrays to remove allocated but unfilled elements
                if storechain:
                    self.rollback(self.stored_iterations)

                # If the sampler is handling the pool, reset it automatically
                if self._managing_pool:
                    self._pool.close()
                    self._pool = choose_pool(self.processes, mpi=self.mpi)

                raise

    def ln_ev(self, ndraws):
        """Produces a Monte-Carlo estimate of the evidence integral using the
        current propasal.

        :param ndraws: The number of draws to make from the proposal
          for the evidence estimate.

        :return: ``(lnZ, dlnZ)``.  Evidence estimate and associated
          uncertainty.
        """

        pts = self.draw(ndraws)
        results = list(self._pool.map(self._get_wrapper(), pts))
        lnpost = np.array([r[0] for r in results])
        lnprop = np.array([r[1] for r in results])

        lninteg = lnpost - lnprop
        lnZ = np.logaddexp.reduce(lninteg) - np.log(lninteg.shape[0])
        lnZ2 = np.logaddexp.reduce(2.0*lninteg) - np.log(lninteg.shape[0])

        # sigma^2 = <Z^2> - <Z>^2
        # log(sigma^2) = log(<Z^2>) + log(1 - <Z>^2/<Z^2>)
        # Standard error = sqrt(sigma^2/N)
        lnsZ = 0.5*(lnZ2 + np.log1p(-np.exp(2.0*lnZ - lnZ2)) - np.log(lninteg.shape[0]))

        # dlnZ = sigma / Z
        dlnZ = np.exp(lnsZ - lnZ)

        return lnZ, dlnZ

    def draw(self, size, spaces=None):
        """
        Draw `size` samples from the current proposal distribution.

        :param size:
            Number of samples to draw.

        :param spaces:
            If not ``None`` while transdimensional sampling, draws are confined to the requested
            spaces.  Such a thing is useful for burnin (e.g. ``spaces = ~p.mask``).

        :returns: `size` draws from the proposal distribution.
        """
        if self._transd:
            draws = self._kde.draw(size, spaces=spaces)
        else:
            draws = self._kde.draw(size)
        return draws

    def trigger_update(self, interval=None):
        """
        Decide whether to trigger a proposal update.

        :param interval:
            Interval between proposal updates.  If ``None``, no updates will be done.  If
            ``"auto"``, acceptances are split in half and checked for consistency (see
            :meth:`consistent_acceptance_rate`).  If an ``int``, the proposal will be updated
            every `interval` iterations.

        :returns: ``bool`` indicating whether a proposal update is due.
        """
        trigger = False
        if interval is None:
            trigger = False
        elif interval == 'auto':
            trigger = not self.consistent_acceptance_rate()
        elif isinstance(interval, int):
            if self.iterations % interval == 0:
                trigger = True
        else:
            raise RuntimeError("Unexpected `interval` in `trigger_update`.")

        return trigger

    def update_proposal(self, p, max_samples=None, **kwargs):
        """
        Update the proposal density with points `p`.

        :param p:
            Samples to update the proposal with.

        :param max_samples: (optional)
            The maximum number of samples to use for constructing or updating the kde.  If a KDE is
            supplied and adding the samples from it will go over this, old samples are thinned by
            factors of two until under the limit.

        :param kwargs: (optional)
            The rest is passed to the KDE constructor.
        """
        self._updates.append(self.iterations)

        if self._transd:
            self.set_kde(TransdimensionalKDE(p, pool=self._pool, kde=self._kde,
                                            max_samples=self._kde_size, **kwargs))
        else:
            self.set_kde(optimized_kde(p, pool=self._pool, kde=self._kde,
                                      max_samples=self._kde_size, **kwargs))

    def set_kde(self, kde):
        """Sets self's kde and re-creates the pool to use it."""
        self._kde = kde

        if hasattr(self._pool, 'broadcast'):
            self._pool.broadcast(_update_wrapper_kde, kde)

    @property
    def pool(self):
        """Returns the pool instance."""
        return self._pool

    @property
    def failed_p(self):
        """
        Sample that caused the last exception.
        """
        return self._failed_p

    @property
    def chain(self):
        """
        Ensemble's past samples, with shape `(iterations, nwalkers, ndim)`.
        """
        return self._chain

    @property
    def blobs(self):
        """
        Ensemble's past metadata.
        """
        return self._blobs

    @property
    def lnpost(self):
        """
        Ensemble's past posterior probabilities, with shape `(iterations, nwalkers)`.
        """
        return self._lnpost

    @property
    def lnprop(self):
        """
        Ensemble's past proposal probabilities, with shape `(iterations, nwalkers)`.
        """
        return self._lnprop

    @property
    def updates(self):
        """
        Step numbers where the proposal density was updated.
        """
        return self._updates

    @property
    def acceptance(self):
        """
        Boolean array of ensemble's past acceptances, with shape `(iterations, nwalkers)`.
        """
        return self._acceptance

    @property
    def acceptance_fraction(self):
        """
        A 1-D array of length :attr:`stored_iterations` of the fraction of walkers that accepted
        each step.
        """
        return np.mean(self.acceptance, axis=1)

    @property
    def acceptance_rate(self):
        """
        An `(nwalkers, )`-shaped array of the windowed acceptance rate for each walker.  The size of
        the window is chosen automatically based on the fraction of acceptances in the ensembles
        last step.  See :meth:`windowed_acceptance_rate` if you want more control.
        """
        return self.windowed_acceptance_rate()

    @property
    def autocorrelation_times(self, lookback=None):
        """
        An `nwalkers`-long vector of the estimated autocorrelation time of each walker, estimated
        using the number of step acceptances over the last `lookback` steps.

        This function leverages the convenience that the proposal density doesn't functionally depend
        on walkers' current locations, which means an accepted step *must* be independent.  This allows
        for an analytic estimate of the autocorrelation time :math:\tau that depends only on the
        acceptance rate

        .. math:: \tau = \frac{2}{r_\mathrm{acc}} - 1

        NOTE: This method can only be used if the chain is being stored.

        :param lookback:
            Number of previous steps to use for the autocorrelation time estimates.  If ``None``,
            all steps since last proposal update will be used.
        """
        if lookback is None:
            lookback = self.iterations - self.updates[-1]

        # Turn lookback into a negative index to count from the end of the array
        lookback *= -1

        acc_rates = np.mean(self.acceptance[lookback:], axis=0)

        return 2.0/acc_rates - 1.0

    def windowed_acceptance_rate(self, window=None):
        """
        An `(nwalkers, -1)`-shaped array of the windowed acceptance rate for each walker.

        :param window:
            Number of iterations to calculate acceptance rate over.  If ``None``, the fraction of
            accepances across the ensemble's last step are used for scale.
        """
        N = len(self.acceptance)

        # Use the mean acceptance rate of the last step to set the window
        if window is None:
            window = 20 * int(1/np.mean(self.acceptance[-1]))

        rates = np.empty((self.nwalkers, N - window + 1))
        weights = np.ones(window)/window

        for w in range(self.nwalkers):
            rates[w] = np.convolve(self.acceptance[:, w], weights, 'valid')

        return rates

    def consistent_acceptance_rate(self, window_size=None, critical_pval=0.05):
        """
        A convenience function for :meth:`burnin` and :meth:`trigger_update`.  Returns ``True``
        if the number of acceptances each step are consistent with the acceptance rate of the
        last step.  This is done using a chi-squared test.

        :param window_size:
            Number of iterations to look back for acceptances.  If ``None``, the iteration of the
            last proposal update (from :attr:`updates`) is used.

        :param critical_pval:
            The critial p-value for considering the distribution consistent.  If the calculated
            p-value is over this, then ``True`` is returned.
        """
        if window_size is None:
            if len(self.updates) == 0:
                return False
            else:
                window_start = self.updates[-1]
        else:
            window_start = self.iterations - window_size

        window_length = self.iterations - window_start

        # If window is really small, return `consistent` to avoid gratuitous updating
        consistent = True
        if window_length > 2:
            last_acc_rate = self.acceptance_fraction[-1]

            nacc = self.nwalkers * self.acceptance_fraction[window_start:self.iterations]
            expected_freqs = last_acc_rate * self.nwalkers * np.ones_like(nacc)

            _, p_val = chisquare(nacc, expected_freqs)

            if p_val < critical_pval:
                consistent = False

        return consistent

    def rollback(self, iteration):
        """
        Shrink internal arrays down to a length of `iteration` and reset the :attr:`pool` if there
        is one.  This is helpful for keeping things consistent after a :exc:`KeyboardInterrupt`.
        """
        self._chain = self._chain[:iteration]
        self._lnpost = self._lnpost[:iteration]
        self._lnprop = self._lnprop[:iteration]
        self._acceptance = self._acceptance[:iteration]
        self._blobs = self._blobs[:iteration]

    def run_mcmc(self, N, p0=None, lnpost0=None, lnprop0=None, blob0=None, progress=False, **kwargs):
        """
        Iterate :meth:`sample` for `N` iterations and return the result.

        :param N:
            The number of steps to take.

        :param p0: (optional)
            A list of the initial walker positions.  It should have the shape `(nwalkers, ndim)`.
            If ``None`` and the sampler has been run previously, it'll pick up where it left off.

        :param lnpost0: (optional)
            The list of log posterior probabilities for the walkers at positions `p0`. If ``lnpost0
            is None``, the initial values are calculated. It should have the shape `(nwalkers, )`.

        :param lnprop0: (optional)
            List of log proposal densities for walkers at positions `p0`. If ``lnprop0 is None``,
            the initial values are calculated. It should have the shape `(nwalkers, )`.

        :param blob0: (optional)
            The list of blob data for walkers at positions `p0`.

        :param progress: (optional)
            Boolean flag to turn on progress bar reporting through MCMC iterations

        :param kwargs: (optional)
            The rest is passed to :meth:`sample`.

        After `N` steps...

        :returns:
            * ``p`` - An array of current walker positions with shape `(nwalkers, ndim)`.

            * ``lnpost`` - The list of log posterior probabilities for the walkers at positions
              ``p``, with shape `(nwalkers, )`.

            * ``lnprop`` - The list of log proposal densities for the walkers at positions `p`, with
              shape `(nwalkers, )`.

            * ``blob`` - (if `lnprobfn` returns blobs) The list of blob data for the walkers at
              positions `p`.
        """

        if p0 is None:
            if self._last_run_mcmc_result is None:
                try:
                    p0 = self.chain[-1]
                    if lnpost0 is None:
                        lnpost0 = self.lnpost[-1]
                    if lnprop0 is None:
                        lnprop0 = self.lnprop[-1]
                except IndexError:
                    raise ValueError("Cannot have p0=None if the sampler hasn't been called.")
            else:
                p0 = self._last_run_mcmc_result[0]
                if lnpost0 is None:
                    lnpost0 = self._last_run_mcmc_result[1]
                if lnprop0 is None:
                    lnprop0 = self._last_run_mcmc_result[2]

        if self._kde is not None:
            if self._last_run_mcmc_result is None and (lnpost0 is None or lnprop0 is None):
                results = list(self._pool.map(self._get_wrapper(), p0))
                if lnpost0 is None:
                    lnpost0 = np.array([r[0] for r in results])
                if lnprop0 is None:
                    lnprop0 = np.array([r[1] for r in results])
                if blob0 is None:
                    try:
                        blob0 = [r[2] for r in results]
                    except IndexError:
                        blob0 = None

        pbar = self._get_finite_pbar(progress, N)
        it = self.iterations
        for results in self.sample(p0, lnpost0, lnprop0, blob0, N, **kwargs):
            if pbar is not None:
                pbar.update(1)
                pbar.set_postfix_str("| {0}/{1} Walkers Accepted | Last step Acc Rate: {2:.3f}".format(np.count_nonzero(self.acceptance[it]),
                                                                                self.nwalkers, self.acceptance_fraction[it]))
                it += 1

        # Store the results for later continuation and toss out the blob
        self._last_run_mcmc_result = results[:3]
        if pbar is not None:
            pbar.close()
        return results

    @property
    def burnin_length(self):
        """
        If :meth:`burnin` was not used, and ``burnin_length`` is
        ``None``, the iteration of the last proposal update will
        be used.
        """
        if self._burnin_length is None:
            return self.updates[-1]
        else:
            return self._burnin_length

    def get_samples(self, burnin_length=None):
        """
        Extract the independent samples collected after burnin.
        If :meth:`burnin` was not used, and ``burnin_length`` is
        ``None``, the iteration of the last proposal update will
        be used.

        :param burnin_length: (optional)
            The step after which to extract samples from.

        :returns: An `(nsamples, ndim)` array of independent samples.
        """
        if burnin_length is None:
            burnin_length = self.burnin_length
        ACTs = np.ceil(self.autocorrelation_times).astype(int)
        samples = [self.chain[burnin_length::ACT, walker]
                   for walker, ACT in zip(range(self.nwalkers), ACTs)]
        return np.vstack(samples)


# -*- coding: utf-8 -*-
import inspect
from datetime import timedelta
from typing import List, Union

import pandas as pd
from sqlalchemy import Column, Float, String, DateTime
from sqlalchemy.orm import Session

from zvt.api.data_type import Region, Provider
from zvt.contract import IntervalLevel
from zvt.utils.time_utils import date_and_time, is_same_time, now_pd_timestamp

class Mixin(object):
    id = Column(String, primary_key=True)
    # entity id for this model
    entity_id = Column(String)

    # the meaning could be different for different case,most of time it means 'happen time'
    timestamp = Column(DateTime)

    # __mapper_args__ = {
    #    "order_by": timestamp.desc()
    # }

    @classmethod
    def help(cls):
        print(inspect.getsource(cls))

    @classmethod
    def important_cols(cls):
        return []

    @classmethod
    def time_field(cls):
        return 'timestamp'

    @classmethod
    def register_recorder_cls(cls, region: Region, provider: Provider, recorder_cls):
        """
        register the recorder for the schema

        :param provider:
        :param recorder_cls:
        """
        # don't make provider_map_recorder as class field,it should be created for the sub class as need
        if not hasattr(cls, 'provider_map_recorder'):
            cls.provider_map_recorder = {}

        if region not in cls.provider_map_recorder:
            cls.provider_map_recorder[region] = {}
            cls.provider_map_recorder[region][provider] = recorder_cls

        elif provider not in cls.provider_map_recorder[region]:
            cls.provider_map_recorder[region][provider] = recorder_cls

    @classmethod
    def register_provider(cls, region: Region, provider: Provider):
        # dont't make providers as class field,it should be created for the sub class as need
        if not hasattr(cls, 'providers'):
            cls.providers = {}

        if region in cls.providers.keys():
            if provider not in cls.providers[region]:
                cls.providers[region].append(provider)
        else:
            cls.providers.update({region: [provider]})

    @classmethod
    def test_data_correctness(cls, region, provider, data_samples):
        for data in data_samples:
            item = cls.query_data(region=region, provider=provider, ids=[data['id']], return_type='dict')
            print(item)
            for k in data:
                if k == 'timestamp':
                    assert is_same_time(item[0][k], data[k])
                else:
                    assert item[0][k] == data[k]

    @classmethod
    def query_data(cls,
                   region: Region,
                   provider_index: int = 0,
                   ids: List[str] = None,
                   entity_ids: List[str] = None,
                   entity_id: str = None,
                   codes: List[str] = None,
                   code: str = None,
                   level: Union[IntervalLevel, str] = None,
                   provider: Provider = Provider.Default,
                   columns: List = None,
                   col_label: dict = None,
                   return_type: str = 'df',
                   start_timestamp: Union[pd.Timestamp, str] = None,
                   end_timestamp: Union[pd.Timestamp, str] = None,
                   filters: List = None,
                   session: Session = None,
                   order=None,
                   limit: int = None,
                   index: Union[str, list] = None,
                   time_field: str = 'timestamp',
                   func=None):
        from .api import get_data
        if provider == Provider.Default:
            provider = cls.providers[region][provider_index]
        return get_data(region=region, data_schema=cls, ids=ids, entity_ids=entity_ids, entity_id=entity_id, codes=codes,
                        code=code, level=level, provider=provider, columns=columns, col_label=col_label, return_type=return_type,
                        start_timestamp=start_timestamp, end_timestamp=end_timestamp, filters=filters, session=session,
                        order=order, limit=limit, index=index, time_field=time_field, fun=func)

    @classmethod
    def record_data(cls,
                    region: Region,
                    provider_index: int = 0,
                    provider: Provider = Provider.Default,
                    exchanges=None,
                    entity_ids=None,
                    codes=None,
                    batch_size=None,
                    force_update=None,
                    sleeping_time=None,
                    default_size=None,
                    real_time=None,
                    fix_duplicate_way=None,
                    start_timestamp=None,
                    end_timestamp=None,
                    close_hour=None,
                    close_minute=None,
                    one_day_trading_minutes=None,
                    **kwargs):
        if cls.provider_map_recorder:
            # print(f'{cls.__name__} registered recorders:{cls.provider_map_recorder}')
            if provider:
                recorder_class = cls.provider_map_recorder[region][provider]
            else:
                recorder_class = cls.provider_map_recorder[region][cls.providers[region][provider_index]]

            # get args for specific recorder class
            from zvt.contract.recorder import TimeSeriesDataRecorder
            if issubclass(recorder_class, TimeSeriesDataRecorder):
                args = [item for item in inspect.getfullargspec(cls.record_data).args if
                        item not in ('cls', 'region', 'provider_index', 'provider')]
            else:
                args = ['batch_size', 'force_update', 'sleeping_time']

            # just fill the None arg to kw,so we could use the recorder_class default args
            kw = {}
            for arg in args:
                tmp = eval(arg)
                if tmp is not None:
                    kw[arg] = tmp

            # FixedCycleDataRecorder
            from zvt.contract.recorder import FixedCycleDataRecorder
            if issubclass(recorder_class, FixedCycleDataRecorder):
                # contract:
                # 1)use FixedCycleDataRecorder to record the data with IntervalLevel
                # 2)the table of schema with IntervalLevel format is {entity}_{level}_[adjust_type]_{event}
                table: str = cls.__tablename__
                try:
                    items = table.split('_')
                    if len(items) == 4:
                        adjust_type = items[2]
                        kw['adjust_type'] = adjust_type
                    level = IntervalLevel(items[1])
                except:
                    # for other schema not with normal format,but need to calculate size for remaining days
                    level = IntervalLevel.LEVEL_1DAY

                kw['level'] = level

            # add other custom args
            for k in kwargs:
                kw[k] = kwargs[k]

            r = recorder_class(**kw)
            r.run()
        else:
            print(f'no recorders for {cls.__name__}')

class NormalMixin(Mixin):
    # the record created time in db
    created_timestamp = Column(DateTime, default=now_pd_timestamp(Region.CHN))
    # the record updated time in db, some recorder would check it for whether need to refresh
    updated_timestamp = Column(DateTime)

class Player(Mixin):
    # 参与者类型
    player_type = Column(String(length=64))
    # 所属国家
    country = Column(String(length=32))
    # 编码
    code = Column(String(length=64))
    # 名字
    name = Column(String(length=128))

class EntityMixin(Mixin):
    # 标的类型
    entity_type = Column(String(length=64))
    # 所属交易所
    exchange = Column(String(length=32))
    # 编码
    code = Column(String(length=64))
    # 名字
    name = Column(String(length=256))
    # 上市日
    list_date = Column(DateTime)
    # 退市日
    end_date = Column(DateTime)

    @classmethod
    def get_trading_dates(cls, start_date=None, end_date=None):
        """
        overwrite it to get the trading dates of the entity

        :param start_date:
        :param end_date:
        :return:
        """
        return pd.date_range(start_date, end_date, freq='B')

    @classmethod
    def get_trading_intervals(cls):
        """
        overwrite it to get the trading intervals of the entity

        :return:[(start,end)]
        """
        return [('09:30', '11:30'), ('13:00', '15:00')]

    @classmethod
    def get_interval_timestamps(cls, start_date, end_date, level: IntervalLevel):
        """
        generate the timestamps for the level

        :param start_date:
        :param end_date:
        :param level:
        """

        for current_date in cls.get_trading_dates(start_date=start_date, end_date=end_date):
            if level >= IntervalLevel.LEVEL_1DAY:
                yield current_date
            else:
                start_end_list = cls.get_trading_intervals()

                for start_end in start_end_list:
                    start = start_end[0]
                    end = start_end[1]

                    current_timestamp = date_and_time(the_date=current_date, the_time=start)
                    end_timestamp = date_and_time(the_date=current_date, the_time=end)

                    while current_timestamp <= end_timestamp:
                        yield current_timestamp
                        current_timestamp = current_timestamp + timedelta(minutes=level.to_minute())

    @classmethod
    def is_open_timestamp(cls, timestamp):
        timestamp = pd.Timestamp(timestamp)
        return is_same_time(timestamp, date_and_time(the_date=timestamp.date(),
                                                     the_time=cls.get_trading_intervals()[0][0]))

    @classmethod
    def is_close_timestamp(cls, timestamp):
        timestamp = pd.Timestamp(timestamp)
        return is_same_time(timestamp, date_and_time(the_date=timestamp.date(),
                                                     the_time=cls.get_trading_intervals()[-1][1]))

    @classmethod
    def is_finished_kdata_timestamp(cls, timestamp: pd.Timestamp, level: IntervalLevel):
        """
        :param timestamp: the timestamp could be recorded in kdata of the level
        :type timestamp: pd.Timestamp
        :param level:
        :type level: zvt.domain.common.IntervalLevel
        :return:
        :rtype: bool
        """
        timestamp = pd.Timestamp(timestamp)

        for t in cls.get_interval_timestamps(timestamp.date(), timestamp.date(), level=level):
            if is_same_time(t, timestamp):
                return True

        return False

    @classmethod
    def could_short(cls):
        """
        whether could be shorted

        :return:
        """
        return False

    @classmethod
    def get_trading_t(cls):
        """
        0 means t+0
        1 means t+1

        :return:
        """
        return 1

class NormalEntityMixin(EntityMixin):
    # the record created time in db
    created_timestamp = Column(DateTime, default=now_pd_timestamp(Region.CHN))
    # the record updated time in db, some recorder would check it for whether need to refresh
    updated_timestamp = Column(DateTime)

class Portfolio(EntityMixin):
    @classmethod
    def get_stocks(cls, region: Region, provider: Provider, timestamp, code=None, codes=None, ids=None):
        """
        the publishing policy of portfolio positions is different for different types,
        overwrite this function for get the holding stocks in specific date

        :param code: portfolio(etf/block/index...) code
        :param codes: portfolio(etf/block/index...) codes
        :param ids: portfolio(etf/block/index...) ids
        :param timestamp: the date of the holding stocks
        :param provider: the data provider
        :return:
        """
        from zvt.contract.api import get_schema_by_name
        schema_str = f'{cls.__name__}Stock'
        portfolio_stock = get_schema_by_name(schema_str)
        return portfolio_stock.query_data(region=region, provider=provider, code=code, codes=codes, timestamp=timestamp, ids=ids)

# 组合(Fund,Etf,Index,Block等)和个股(Stock)的关系 应该继承自该类
# 该基础类可以这样理解:
# entity为组合本身,其包含了stock这种entity,timestamp为持仓日期,从py的"你知道你在干啥"的哲学出发，不加任何约束
class PortfolioStock(Mixin):
    # portfolio标的类型
    entity_type = Column(String(length=64))
    # portfolio所属交易所
    exchange = Column(String(length=32))
    # portfolio编码
    code = Column(String(length=64))
    # portfolio名字
    name = Column(String(length=128))

    stock_id = Column(String)
    stock_code = Column(String(length=64))
    stock_name = Column(String(length=256))

# 支持时间变化,报告期标的调整
class PortfolioStockHistory(PortfolioStock):
    # 报告期,season1,half_year,season3,year
    report_period = Column(String(length=32))
    # 3-31,6-30,9-30,12-31
    report_date = Column(DateTime)

    # 占净值比例
    proportion = Column(Float)
    # 持有股票的数量
    shares = Column(Float)
    # 持有股票的市值
    market_cap = Column(Float)

__all__ = ['EntityMixin', 'Mixin', 'NormalMixin', 'NormalEntityMixin', 'Portfolio', 'PortfolioStock',
           'PortfolioStockHistory']


#!/usr/bin/env python
#Hennie Veldthuis

import sys
import xml.etree.ElementTree as ET

def main(argv):
    if len(argv) == 3:
        tree = ET.parse(argv[1])
        root = tree.getroot()

        for point in root:
            f0_start  = float(point[13].text)
            f0_end    = float(point[12].text)
            bottom_hz = float(point[5].text)
            top_hz    = float(point[7].text)
            if not((f0_start > bottom_hz) and (f0_start < top_hz) and (f0_end > bottom_hz) and (f0_end < top_hz)):
                # values are incorrect,
                # remove item
                root.remove(point)

        tree.write(argv[2])
    else:
        print("ERROR, not enough arguments, use: python3 spontal_test.py spontal.xml spontal_filtered.xml")

if __name__ == '__main__':
    main(sys.argv)


#----------------------------------------------------------------------------#
# Imports
#----------------------------------------------------------------------------#

import os
from sqlalchemy import Column, String, Integer, create_engine
from flask import Flask, render_template, request, Response, flash, redirect, url_for, jsonify, abort
from flask_sqlalchemy import SQLAlchemy
from flask_moment import Moment
from sqlalchemy.orm import backref
import psycopg2
import json

'''
setup_db(app)
    binds a flask application and a SQLAlchemy service
'''
#----------------------------------------------------------------------------#
# App Config.
#----------------------------------------------------------------------------#

database_name = "capstone"
database_path = os.environ['DATABASE_URL']
conn = psycopg2.connect(database_path, sslmode='require')

#database_name = "capstone"
# database_path = "postgres://{}:{}@{}/{}".format(
#    'postgres', '123', 'localhost:5432', database_name)

db = SQLAlchemy()

def setup_db(app, database_path=database_path):
    app.config["SQLALCHEMY_DATABASE_URI"] = database_path
    app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False
    db.app = app
    db.init_app(app)
    db.create_all()

'''
CREATE MODELS
MOVIES, ACTORS, POSTER
'''

#----------------------------------------------------------------------------#
# Models.
#----------------------------------------------------------------------------#

class Movie(db.Model):
    __tablename__ = 'movie'

    id = Column(Integer, primary_key=True)
    title = Column(String, nullable=False)
    release_data = Column(String, nullable=False)
    actors = db.relationship('Actor', secondary='poster',
                             backref='movie', lazy=True)

    def __init__(self, title, release_data):
        self.title = title
        self.release_data = release_data

    def insert(self):
        db.session.add(self)
        db.session.commit()

    def update(self):
        db.session.commit()

    def delete(self):
        db.session.delete(self)
        db.session.commit()

    def format(self):
        return {
            'id': self.id,
            'title': self.title,
            'release_data': self.release_data,
            'actors': [actor.name for actor in self.actors]
        }

class Actor(db.Model):
    __tablename__ = 'actor'

    id = Column(Integer, primary_key=True)
    name = Column(String, nullable=False)
    age = Column(String)
    gendar = Column(String)
    movies = db.relationship('Movie', secondary='poster',
                             backref='actor', lazy=True)

    def __init__(self, name, age, gendar):
        self.name = name
        self.age = age
        self.gendar = gendar

    def insert(self):
        db.session.add(self)
        db.session.commit()

    def update(self):
        db.session.commit()

    def delete(self):
        db.session.delete(self)
        db.session.commit()

    def format(self):
        return {
            'id': self.id,
            'name': self.name,
            'age': self.age,
            'gendar': self.gendar,
            'movies': [movie.title for movie in self.movies]
        }

class Poster(db.Model):
    __tablename__ = 'poster'

    poster_id = Column(Integer, primary_key=True)
    movie_id = Column(Integer, db.ForeignKey(
        'movie.id', ondelete='CASCADE'), primary_key=True)
    actor_id = Column(Integer, db.ForeignKey(
        'actor.id', ondelete='CASCADE'), primary_key=True)

    def insert(self):
        db.session.add(self)
        db.session.commit()

    def update(self):
        db.session.commit()

    def delete(self):
        db.session.delete(self)
        db.session.commit()


import time
from functools import partial
from rptest.services.cluster import cluster
from rptest.clients.rpk import RpkTool, RpkException
from rptest.services.kafka_cli_consumer import KafkaCliConsumer
from ducktape.mark import parametrize
from rptest.services.admin import Admin
from rptest.util import expect_exception
from ducktape.utils.util import wait_until
from rptest.clients.types import TopicSpec
from rptest.tests.redpanda_test import RedpandaTest
from rptest.services.redpanda import RESTART_LOG_ALLOW_LIST
from rptest.services.redpanda_installer import RedpandaInstaller, wait_for_num_versions, ver_string

class OffsetRetentionDisabledAfterUpgrade(RedpandaTest):
    """
    When upgrading to Redpanda v23 or later offset retention should be disabled
    by default. Offset retention did not exist pre-v23, so existing clusters
    should have to opt-in after upgrade in order to avoid suprises.

    When a cluster upgrades to v23 then only newly committed offsets are
    reclaimed automatically. Offsets written prior to the upgrade must be
    manually removed using the kafka delete offset api.

    Offsets committed prior to v23 that are normally non-reclaimable become
    reclaimable if they are updated following the upgrade.
    """
    topics = (TopicSpec(), TopicSpec(), TopicSpec(), TopicSpec(), TopicSpec())

    # pre-defined named topics used to simulate use of topics that are used only
    # pre-v23, both pre and post upgrade, and only after upgrade to v23. the
    # idle variants stop receiving updates prior to enabling legacy support.
    pre_v23_topic = topics[0].name
    prepost_v23_topic = topics[1].name
    prepost_v23_topic_idle = topics[2].name
    post_v23_topic = topics[3].name
    post_v23_topic_idle = topics[4].name

    feature_config_timing = {
        "group_offset_retention_sec",
        "group_offset_retention_check_ms",
    }
    feature_config_legacy = "legacy_group_offset_retention_enabled"
    feature_config_names = feature_config_timing | {feature_config_legacy}

    def __init__(self, test_context):
        super(OffsetRetentionDisabledAfterUpgrade,
              self).__init__(test_context=test_context, num_brokers=3)
        self.installer = self.redpanda._installer

    def setUp(self):
        # handled by test case to support parameterization
        pass

    def _validate_pre_upgrade(self, version):
        """
        1. verify starting version of cluster nodes
        2. verify expected configs of initial cluster
        """
        self.installer.install(self.redpanda.nodes, version)
        super(OffsetRetentionDisabledAfterUpgrade, self).setUp()

        # wait until all nodes are running the same version
        unique_versions = wait_for_num_versions(self.redpanda, 1)
        assert ver_string(version) in unique_versions

        # sanity check none of feature configs should exist
        admin = Admin(self.redpanda)
        config = admin.get_cluster_config()
        assert config.keys().isdisjoint(self.feature_config_names)

    def _perform_upgrade(self, initial_version, version):
        """
        1. verify upgrade of all cluster nodes
        2. verify expected configs after upgrade
        """
        # upgrade all nodes to target version
        self.installer.install(self.redpanda.nodes, version)
        self.redpanda.restart_nodes(self.redpanda.nodes)

        # wait until all nodes are running the same version
        unique_versions = wait_for_num_versions(self.redpanda, 1)
        assert ver_string(initial_version) not in unique_versions

        # sanity check that all new feature configs exist
        admin = Admin(self.redpanda)
        config = admin.get_cluster_config()
        assert config.keys() > self.feature_config_names

        # configs should all have positive values
        for name in self.feature_config_timing:
            assert config[name] > 0

        # after upgrade legacy support should be disabled
        assert config[self.feature_config_legacy] == False

    def _offset_removal_occurred(self, period, pre_upgrade,
                                 pre_legacy_support):
        rpk = RpkTool(self.redpanda)
        group = "hey_group"

        def offsets_exist(include_idle):
            """
            Note to future ci-failure investigators. If you were to ever observe
            a failure in which it appeared that prepost/post topic offsets were
            not removed it could be because the following occurred:

              1) the version fence was not written for some reason
              2) so some of the offsets were incorrectly marked as legacy
              3) something else happened like updates to commits also failing

            it's unlikely that this would happen, but if it does, the proper fix
            would be to update the test to check for this condition (may be
            non-trivial to automate this check) and then retry the test.

            in real-world scenarios if this were to happen then it is expected
            that offset delete api is used to clean up any offsets that lose
            this race condition that exists during an upgrade.
            """
            desc = rpk.group_describe(group)
            if len(desc.partitions) == 0:
                return False

            pre = any(p.topic == self.pre_v23_topic for p in desc.partitions)
            prepost = any(p.topic == self.prepost_v23_topic
                          for p in desc.partitions)
            prepost_idle = any(p.topic == self.prepost_v23_topic_idle
                               for p in desc.partitions)
            post = any(p.topic == self.post_v23_topic for p in desc.partitions)
            post_idle = any(p.topic == self.post_v23_topic_idle
                            for p in desc.partitions)

            assert pre, "pre-topic offsets should never be reclaimed"

            # prior to upgrade, the only offsets that should exist are for the
            # pre and prepost topics.
            if pre_upgrade:
                assert not post, "post-topic has not yet been written"
                return prepost and prepost_idle

            if include_idle:
                return prepost and post and prepost_idle and post_idle
            return prepost and post

        # consume from the group for twice as long as the retention period
        # setting and verify that group offsets exist for the entire time.
        start = time.time()
        while time.time() - start < (period * 2):
            if pre_upgrade:
                # pre-v23
                rpk.produce(self.pre_v23_topic, "k", "v")
                rpk.consume(self.pre_v23_topic, n=1, group=group)
            else:
                # post-v23
                rpk.produce(self.post_v23_topic, "k", "v")
                rpk.consume(self.post_v23_topic, n=1, group=group)
                if pre_legacy_support:
                    rpk.produce(self.post_v23_topic_idle, "k", "v")
                    rpk.consume(self.post_v23_topic_idle, n=1, group=group)

            # pre/post-v23
            rpk.produce(self.prepost_v23_topic, "k", "v")
            rpk.consume(self.prepost_v23_topic, n=1, group=group)
            if pre_legacy_support:
                rpk.produce(self.prepost_v23_topic_idle, "k", "v")
                rpk.consume(self.prepost_v23_topic_idle, n=1, group=group)

            wait_until(lambda: offsets_exist(pre_legacy_support),
                       timeout_sec=1,
                       backoff_sec=1)
            time.sleep(1)

        # after one half life the offset should still exist. prior to legacy
        # support being enabled include idle / be strict. after it is enabled
        # then it's possible the idle topics get reclaimed so relax this check.
        time.sleep(period / 2)
        assert offsets_exist(pre_legacy_support)

        # after waiting for twice the retention period, it should be gone.
        # always be strict on this check and include idle topics.
        time.sleep(period * 2 - period / 2)
        return not offsets_exist(True)

    @cluster(num_nodes=3, log_allow_list=RESTART_LOG_ALLOW_LIST)
    @parametrize(initial_version=(22, 2, 9))
    @parametrize(initial_version=(22, 3, 11))
    def test_upgrade_from_pre_v23(self, initial_version):
        """
        1. test that retention feature doesn't work on legacy version
        2. upgrade cluster and test that feature still doesn't work
        3. enable legacy support on the cluster
        4. test that retention works properly
        """
        period = 30

        # in old cluster offset retention should not be active
        self._validate_pre_upgrade(initial_version)
        assert not self._offset_removal_occurred(period, True, True)

        # in cluster upgraded from pre-v23 retention should not be active
        self._perform_upgrade(initial_version, (23, 1))
        rpk = RpkTool(self.redpanda)
        rpk.cluster_config_set("group_offset_retention_sec", str(period))
        rpk.cluster_config_set("group_offset_retention_check_ms", str(1000))
        assert not self._offset_removal_occurred(period, False, True)

        # enable legacy. enablng legacy support takes affect at the next
        # retention check. since that is configured above to happen every second
        # then the response time should be adequate.
        rpk.cluster_config_set("legacy_group_offset_retention_enabled",
                               str(True))
        assert self._offset_removal_occurred(period, False, False)

class OffsetRetentionTest(RedpandaTest):
    topics = (TopicSpec(), )
    period = 30

    def __init__(self, test_context):
        # retention time is set to 30 seconds and expired offset queries happen
        # every second. these are likely unrealistic in practice, but allow us
        # to build tests that are quick and responsive.
        super(OffsetRetentionTest,
              self).__init__(test_context=test_context,
                             extra_rp_conf=dict(
                                 group_offset_retention_sec=self.period,
                                 group_offset_retention_check_ms=1000,
                             ))

        self.rpk = RpkTool(self.redpanda)

    @cluster(num_nodes=3)
    def test_offset_expiration(self):
        group = "hey_group"

        def offsets_exist():
            desc = self.rpk.group_describe(group)
            return len(desc.partitions) > 0

        # consume from the group for twice as long as the retention period
        # setting and verify that group offsets exist for the entire time.
        start = time.time()
        while time.time() - start < (self.period * 2):
            self.rpk.produce(self.topic, "k", "v")
            self.rpk.consume(self.topic, n=1, group=group)
            wait_until(offsets_exist, timeout_sec=1, backoff_sec=1)
            time.sleep(1)

        # after one half life the offset should still exist
        time.sleep(self.period / 2)
        assert offsets_exist()

        # after waiting for twice the retention period, it should be gone
        time.sleep(self.period * 2 - self.period / 2)
        assert not offsets_exist()

class OffsetDeletionTest(RedpandaTest):
    topics = (TopicSpec(partition_count=3), )
    group = "hey_hey_group"

    def __init__(self, test_context):
        super(OffsetDeletionTest, self).__init__(test_context=test_context)

        self.rpk = RpkTool(self.redpanda)

    def make_consumer(self, topic, session_timeout_ms=60000):
        return KafkaCliConsumer(
            self.test_context,
            self.redpanda,
            topic=topic,
            group=self.group,
            from_beginning=True,
            consumer_properties={'session.timeout.ms': session_timeout_ms},
            instance_name=f'cli-consumer-offset-delete-test-{topic}')

    @cluster(num_nodes=5)
    def test_offset_deletion(self):
        def wait_for_partitions_in_group(n):
            desc = self.rpk.group_describe(self.group)
            return len(desc.partitions) == n

        def assert_status(output, expected_status, expected_topic):
            for response in output:
                assert response.topic == expected_topic, f"Expected: {expected_topic} Observed: {response.topic}"
                assert response.status == expected_status, response.status

        # Produce some data to a topic and consume assigning new consumer group
        for _ in range(0, 10):
            self.rpk.produce(self.topic, "k", "v", partition=0)
            self.rpk.produce(self.topic, "k", "v", partition=1)
            self.rpk.produce(self.topic, "k", "v", partition=2)
        self.rpk.consume(self.topic, n=3, group=self.group)
        wait_until(partial(wait_for_partitions_in_group, 3),
                   timeout_sec=30,
                   backoff_sec=1)

        # Assert offset-delete errors when request contains missing topic-partitions
        missing_topic = f"{self.topic}-foo"
        missing_topic_partitions = {missing_topic: [0, 1, 2]}

        def assert_unknown_topic_or_partition(e):
            output = e.parsed_output
            assert len(output) == 3
            assert_status(output, 'UNKNOWN_TOPIC_OR_PARTITION', missing_topic)
            return True

        with expect_exception(RpkException, assert_unknown_topic_or_partition):
            self.rpk.offset_delete(self.group, missing_topic_partitions)

        # Assert offset-delete errors when non-existent group is passed in
        topic_partitions = {self.topic: [0, 1, 2]}
        with expect_exception(RpkException,
                              lambda e: 'GROUP_ID_NOT_FOUND' in str(e)):
            self.rpk.offset_delete("missing", topic_partitions)

        # Assert offset-delete errors when attempting to delete offsets of topic
        # partitions still assigned to an active group
        consumer = self.make_consumer(self.topic)
        consumer.start()
        consumer.wait_for_messages(1)
        wait_until(partial(wait_for_partitions_in_group, 3),
                   timeout_sec=30,
                   backoff_sec=1)

        def assert_group_already_subscribed(e):
            output = e.parsed_output
            assert len(output) == 3
            assert_status(output, 'GROUP_SUBSCRIBED_TO_TOPIC', self.topic)
            return True

        with expect_exception(RpkException, assert_group_already_subscribed):
            self.rpk.offset_delete(self.group, topic_partitions)

        consumer.stop()
        consumer.wait()
        consumer.free()

        # Assert offset-delete removes offsets of dead groups
        output = self.rpk.offset_delete(self.group, topic_partitions)
        assert len(output) == 3
        assert_status(output, 'OK', self.topic)

        desc = self.rpk.group_describe(self.group)
        assert len(desc.partitions) == 0

        # Assert offset-delete removes offsets of unsubscribed topic/partitions
        # within an active group
        new_topic = "foo"
        self.rpk.create_topic(new_topic, partitions=3)
        for i in range(0, 3):
            self.rpk.produce(new_topic, "k", "v", partition=i)

        # Add two new consumers to the group
        consumer_a = self.make_consumer(self.topic)
        min_default_group_session_timeout_ms = 6000
        consumer_b = self.make_consumer(new_topic,
                                        min_default_group_session_timeout_ms)
        consumer_a.start()
        consumer_b.start()

        # Wait until both have joined the group
        wait_until(partial(wait_for_partitions_in_group, 6),
                   timeout_sec=30,
                   backoff_sec=1)

        # Have them both consume some data
        consumer_a.wait_for_messages(1)
        consumer_b.wait_for_messages(1)

        # Shutdown one consumer
        consumer_b.stop()
        consumer_b.wait()

        # After consumer_b shuts down wait for the group rebalance to occur
        # and unsubscription of topic/partitions it was assigned to
        wait_until(partial(wait_for_partitions_in_group, 3),
                   timeout_sec=30,
                   backoff_sec=1)

        # Remove the offsets that it had been keeping in redpanda
        def validate_output(e):
            output = e.parsed_output
            assert len(output) == 4
            good_responses = [x for x in output if x.partition != 3]
            bad_responses = [x for x in output if x.partition == 3]
            assert len(good_responses) == 3
            assert len(bad_responses) == 1
            assert_status(good_responses, "OK", new_topic)
            assert_status(bad_responses, "UNKNOWN_TOPIC_OR_PARTITION",
                          new_topic)
            return True

        with expect_exception(RpkException, validate_output):
            self.rpk.offset_delete(self.group, {new_topic: [0, 1, 2, 3]})

        # Shutdown other consumer
        consumer_a.stop()
        consumer_a.wait()


#!/usr/bin/env python3
# vim: set fileencoding=utf8 :

import types
from contextlib import contextmanager
from tornado import gen
from tornado.concurrent import Future, is_future
import inflection
import python_jsonschema_objects as pjs

from tornado_jsonapi.exceptions import MissingResourceSchemaError

class Resource:
    class ResourceObject:
        def id_(self):
            raise NotImplementedError

        def type_(self):
            raise NotImplementedError

        def attributes(self):
            raise NotImplementedError

        # TODO : relationships, links, meta

    def __init__(self, schema):
        self.schema = schema
        builder = pjs.ObjectBuilder(self.schema)
        classes = builder.build_classes()
        name = inflection.camelize(self.name())
        if name not in classes:
            raise MissingResourceSchemaError(name)
        self._schema = classes[name]

    def _on_request_end(self):
        pass

    def name(self):
        raise NotImplementedError

    def exists(self, id_):
        return self.read(id_) is not None

    def create(self, attributes):
        raise NotImplementedError

    def read(self, id_):
        raise NotImplementedError

    def update(self, id_, attributes):
        raise NotImplementedError

    def delete(self, id_):
        raise NotImplementedError

    def list_(self, limit=0, page=0):
        raise NotImplementedError

    def list_count(self):
        raise NotImplementedError

try:
    import sqlalchemy
    import alchemyjsonschema
    import alchemyjsonschema.dictify
except ImportError:
    pass

class SQLAlchemyResource(Resource):
    class ResourceObject:
        def __init__(self, resource, model, blacklist=None):
            self.resource = resource
            self.model = model
            self.blacklist = blacklist
            if self.blacklist is None:
                self.blacklist = []

        def id_(self):
            return str(
                self.model.__getattribute__(self.resource._primary_columns[0])
            )

        def type_(self):
            return self.resource.name()

        def attributes(self):
            attributes_ = alchemyjsonschema.dictify.jsonify(
                self.model, self.resource.schema
            )
            for key in self.blacklist:
                attributes_.pop(key, None)
            return attributes_

    def __init__(self, model_cls, sessionmaker):
        self._primary_columns = model_cls.__table__.primary_key.columns.keys()
        if len(self._primary_columns) > 1:
            raise NotImplementedError("Compound primary keys not supported")
        self.model_cls = model_cls
        self.model_primary_key = getattr(
            self.model_cls, self._primary_columns[0]
        )
        self.sessionmaker = sessionmaker
        self.session = sqlalchemy.orm.scoped_session(sessionmaker)
        self.blacklist = []
        factory = alchemyjsonschema.SchemaFactory(
            alchemyjsonschema.StructuralWalker
        )
        schema = factory(self.model_cls, excludes=self._primary_columns)
        super().__init__(schema)

    def _on_request_end(self):
        self.session.remove()

    def _id_filter(self, id_):
        return {self._primary_columns[0]: id_}

    def name(self):
        return inflection.camelize(
            self.model_cls.__name__, uppercase_first_letter=False
        )

    def exists(self, id_):
        return self.session.query(
            self.session.query(self.model_cls)
            .filter_by(**self._id_filter(id_))
            .exists()
        ).scalar()

    def create(self, attributes):
        model = self.model_cls(**attributes)
        self.session.add(model)
        self.session.commit()
        return SQLAlchemyResource.ResourceObject(
            self, model, blacklist=self.blacklist
        )

    def read(self, id_):
        model = (
            self.session.query(self.model_cls)
            .filter_by(**self._id_filter(id_))
            .one_or_none()
        )
        return (
            None
            if model is None
            else SQLAlchemyResource.ResourceObject(
                self, model, blacklist=self.blacklist
            )
        )

    def update(self, id_, attributes):
        model = (
            self.session.query(self.model_cls)
            .filter_by(**self._id_filter(id_))
            .one()
        )
        for k, v in attributes.items():
            setattr(model, k, v)
        self.session.merge(model)
        self.session.commit()
        return SQLAlchemyResource.ResourceObject(
            self, model, blacklist=self.blacklist
        )

    def delete(self, id_):
        # TODO shouldnt it be
        #  http://docs.sqlalchemy.org/en/latest/core/tutorial.html#deletes
        r = (
            self.session.query(self.model_cls)
            .filter_by(**self._id_filter(id_))
            .delete()
        )
        self.session.commit()
        return r

    def list_count(self):
        return self.session.query(
            sqlalchemy.func.count(self.model_primary_key)
        ).scalar()

    def list_(self, limit=0, page=0):
        if limit > 0:
            start = abs(page) * limit
            stop = start + limit
            models = self.session.query(self.model_cls).slice(start, stop)
        else:
            models = self.session.query(self.model_cls)
        res = []
        for model in models:
            res.append(
                SQLAlchemyResource.ResourceObject(
                    self, model, blacklist=self.blacklist
                )
            )
        return res

try:
    import dbapiext
except ImportError:
    pass

@gen.coroutine
def dbapi2Cursor(connection, transaction=False):
    @contextmanager
    def wrapper(connection, transaction):
        cursor = connection.cursor()
        try:
            yield cursor
            if transaction:
                connection.commit()
        except:
            if transaction:
                connection.rollback()
            raise
        finally:
            cursor.close()

    return wrapper(connection, transaction)

@gen.coroutine
def momokoCursor(pool, transaction=False):
    class wrapper:
        def __init__(self, pool, connection, transaction):
            self.pool = pool
            self.connection = connection
            self.transaction = transaction

        def __enter__(self):
            return self.connection

        @gen.coroutine
        def __exit__(self, type_, value, traceback):
            if self.transaction:
                if type_ is None:
                    yield self.connection.execute("COMMIT")
                else:
                    yield self.connection.execute("ROLLBACK")
            self.pool.putconn(self.connection)

    connection = yield pool.getconn(ping=False)
    if transaction:
        yield connection.execute("BEGIN")
    return wrapper(pool, connection, transaction)

class DBAPI2Resource(Resource):
    _types_mapping = {
        "boolean": "boolean",
        "integer": "integer",
        "number": "numeric",
        "string": "text",
    }

    class ResourceObject:
        def __init__(self, resource, row):
            self._resource = resource
            assert row is not None
            self.row = row

        def id_(self):
            return str(self.row[-1])

        def type_(self):
            return self._resource.name()

        def attributes(self):
            return {
                n: v for (n, v) in zip(self._resource.columns, self.row[:-1])
            }

    def __init__(self, schema, dbapi, connection):
        self.cursor = dbapi2Cursor
        self.connection = connection
        self.dbapi = dbapi
        if dbapi.__name__ == "momoko":
            self.cursor = momokoCursor
            self.dbapi = dbapi.psycopg2
        dbapiext.set_paramstyle(self.dbapi)
        self._tablename = inflection.pluralize(schema["title"])
        super().__init__(schema)
        self.columns = list(self._schema()._properties.keys())

    def _is_sqlite(self):
        return self.dbapi.__name__ == "sqlite3"

    def _is_postgresql(self):
        return self.dbapi.__name__ == "psycopg2"

    def _create_primary_key(self):
        id_type = "integer"
        if self._is_postgresql():
            id_type = "bigserial"
        key = "id {} not null primary key".format(id_type)
        if self._is_sqlite():
            key += " autoincrement"
        return key

    @gen.coroutine
    def _create_table(self):
        types = [
            self._types_mapping[self._schema.propinfo(c)["type"]] +
            (" not null" if c in self._schema.__required__ else "")
            for c in self.columns
        ]
        column_defs = [self._create_primary_key()] + [
            "{} {}".format(c, t) for c, t in zip(self.columns, types)
        ]
        with (yield self.cursor(self.connection, transaction=True)) as cursor:
            cur = dbapiext.execute_f(
                cursor,
                "create table if not exists %s (%s)",
                self._tablename,
                column_defs,
            )
            if is_future(cur):
                yield cur

    def name(self):
        return self.schema["title"]

    @gen.coroutine
    def exists(self, id_):
        with (yield self.cursor(self.connection)) as cursor:
            cur = dbapiext.execute_f(
                cursor, "select 1 from %s where id = %X", self._tablename, id_
            )
            if is_future(cur):
                cur = yield cur
            return cur.fetchone() is not None

    @gen.coroutine
    def create(self, attributes):
        with (yield self.cursor(self.connection, transaction=True)) as cursor:
            cur = dbapiext.execute_f(
                cursor,
                "insert into %s (%s) values (%X)",
                self._tablename,
                list(attributes.keys()),
                list(attributes.values()),
            )
            if is_future(cur):
                cur = yield cur
        with (yield self.cursor(self.connection)) as cursor:
            # XXX assuming id is monotonically increasing
            cur = dbapiext.execute_f(
                cursor,
                "select %s from %s order by id desc limit 1",
                self.columns + ["id"],
                self._tablename,
            )
            if is_future(cur):
                cur = yield cur
            return DBAPI2Resource.ResourceObject(self, cur.fetchone())

    @gen.coroutine
    def read(self, id_):
        with (yield self.cursor(self.connection)) as cursor:
            cur = dbapiext.execute_f(
                cursor,
                "select %s from %s where id = %X",
                self.columns + ["id"],
                self._tablename,
                id_,
            )
            if is_future(cur):
                cur = yield cur
            row = cur.fetchone()
            if not row:
                return None
            return DBAPI2Resource.ResourceObject(self, row)

    @gen.coroutine
    def update(self, id_, attributes):
        with (yield self.cursor(self.connection, transaction=True)) as cursor:
            cur = dbapiext.execute_f(
                cursor,
                "update %s set %X where id = %X",
                self._tablename,
                attributes,
                id_,
            )
            if is_future(cur):
                cur = yield cur
        with (yield self.cursor(self.connection)) as cursor:
            cur = dbapiext.execute_f(
                cursor,
                "select %s from %s where id = %X",
                self.columns + ["id"],
                self._tablename,
                id_,
            )
            if is_future(cur):
                cur = yield cur
            return DBAPI2Resource.ResourceObject(self, cur.fetchone())

    @gen.coroutine
    def delete(self, id_):
        with (yield self.cursor(self.connection, transaction=True)) as cursor:
            cur = dbapiext.execute_f(
                cursor, "delete from %s where id = %X", self._tablename, id_
            )
            if is_future(cur):
                cur = yield cur
            return cur.rowcount

    @gen.coroutine
    def list_(self, limit=0, page=0):
        with (yield self.cursor(self.connection)) as cursor:
            if limit > 0:
                start = abs(page) * limit
                cur = dbapiext.execute_f(
                    cursor,
                    "select %s from %s limit %d offset %d",
                    self.columns + ["id"],
                    self._tablename,
                    limit,
                    start,
                )
            else:
                cur = dbapiext.execute_f(
                    cursor,
                    "select %s from %s",
                    self.columns + ["id"],
                    self._tablename,
                )
            if is_future(cur):
                cur = yield cur
            rows = cur.fetchall()
            return [DBAPI2Resource.ResourceObject(self, row) for row in rows]

    @gen.coroutine
    def list_count(self):
        with (yield self.cursor(self.connection)) as cursor:
            cur = dbapiext.execute_f(
                cursor, "select count(1) from %s", self._tablename
            )
            if is_future(cur):
                cur = yield cur
            rows = cur.fetchone()
            return rows[0]


import json
import math
import os
import time
from concurrent.futures import ThreadPoolExecutor
from urllib import parse

import requests
from rich import print
from rich.console import Console
from rich.progress import Progress
from rich.progress import track
from selenium import webdriver

console = Console()

class qqzone_upload(object):
    def __init__(self, path):
        self.path = path

    def _login_and_wait(self):
        opt = webdriver.ChromeOptions()

        self.driver = webdriver.Chrome(chrome_options=opt)
        self.driver.get("https://i.qq.com")

        while True:
            url = self.driver.current_url
            if "user.qzone.qq.com" in url:
                time.sleep(0.2)
                self.my_uin = str(url[url.rfind("/") + 1 :])
                self.cookies = self.driver.get_cookies()
                self.g_tk = str(
                    self.driver.execute_script("return QZONE.FP.getACSRFToken()")
                )
                break
            else:
                time.sleep(1)

    def _init_session(self):
        self.session = requests.Session()
        for cookie in self.cookies:
            self.session.cookies.set(cookie["name"], cookie["value"])
        self.session.headers = {
            "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/66.0.3359.181 Safari/537.36"
        }

    def _add_ablum_and_get_id(self):
        ablum_name = self.path[self.path.rfind("\\") + 1 :]
        print(ablum_name)
        url = (
            "https://user.qzone.qq.com/proxy/domain/photo.qzone.qq.com/cgi-bin/common/cgi_add_album_v2?g_tk="
            + self.g_tk
        )
        data = (
            "inCharset=utf-8&outCharset=utf-8&hostUin="
            + self.my_uin
            + "&notice=0&callbackFun=_Callback&format=fs&plat=qzone&source=qzone&appid=4&uin="
            + self.my_uin
            + "&album_type=&birth_time=&degree_type=0&enroll_time=&albumname="
            + parse.quote(ablum_name.encode("utf-8"))
            + "&albumdesc=&albumclass=100&priv=3&queFstion=&answer=&whiteList=&bitmap=10000000&qzreferrer=https%3A%2F%2Fuser.qzone.qq.com%2F201001948%2F4"
        )
        r = self.session.post(url, data, verify=False)
        self.id = (r.text[r.text.find('"id" : ') + 8 :])[
            : (r.text[r.text.find('"id" : ') + 8 :]).find('",')
        ]
        print(self.id)

    def _add_the_photo(self, photo_name):
        photo_path = self.path + "\\" + photo_name
        self.driver.find_element_by_xpath('//*[@id="container"]/div/a/input').send_keys(
            photo_path
        )

    def _input_all_added(self):
        self.driver.find_element_by_xpath('//input[@_phototype="2"]').click()
        self.driver.find_element_by_xpath(
            '//a[@class="op-btn btn-upload j-btn-start j-uploading-hide"]'
        ).click()
        self.driver.switch_to.default_content()
        self.driver.switch_to.frame("tphoto")
        while True:
            try:
                self.driver.find_element_by_xpath(
                    '//span[contains(text(), "继续上传")]'
                ).click()
            except:
                try:
                    self.driver.find_element_by_xpath('//a[@id="back_upload"]').click()
                    time.sleep(3)
                    self.driver.switch_to.default_content()
                    self.driver.switch_to.frame("photoUploadDialog")
                    break
                except BaseException as e:
                    time.sleep(3)

    def _input_photos(self):
        self.driver.switch_to.frame("photoUploadDialog")
        # str = '//li[@_aid="' + self.id + '"]'
        # self.driver.find_element_by_xpath(str).click()

        """
        num=len(self.name_lists)
        num2=num%500
        num1=num//500

        for i in track(range(num1),"Uploading..."):
            for j in range(i*500,i*500+500,1):
                self._add_the_photo(self.name_lists[j])
            self._input_all_added()

        if num2!=0:
            for i in range(num2):
                self._add_the_photo(self.name_lists[i+num1])
            self._input_all_added()
        """

        for i in track(range(len(self.name_lists) - self.start_num), "Uploading..."):
            self._add_the_photo(self.name_lists[i + self.start_num])
            if (i + 1) % 500 == 0 or (i + 1 + self.start_num) >= len(self.name_lists):
                self._input_all_added()

    def _make_name_list(self):
        g = os.walk(self.path)
        for path, dir_list, file_list in g:
            self.name_lists = file_list
            # print(file_list)

    def start_at(self, num=0, make_ablum=False):
        self.start_num = num
        self._make_name_list()
        self._login_and_wait()
        if make_ablum or num == 0:
            self._init_session()
            self._add_ablum_and_get_id()
        wait = input("wait")
        self._input_photos()
        self.driver.quit()


from itertools import product
import numpy as np
import skimage.color

_FREQS = np.arange(0.05, 0.30, 0.1)
_SIGMAS = range(1, 21, 7)
_THETAS = np.arange(0, np.pi, 0.25*np.pi)

def add_channels(img, colorspace='rgb',
                 methods=['gabor', 'gaussian', 'sobel'],
                 methods_params=None):
    '''Add artificial channels to an image

    Parameters
    ----------
    img : np.ndarray
        NxMx3 array with colored image data
    colorspace : str, optional
        String indicating colorspace of *img* (rgb/hsv/etc.)
    methods : list, optional
        List of strings indicating channels to be added
    methods_params : dict, optional
        Dictionairy with named options for channel functions

    Notes
    -----
    Currently implemented channels are:
    * gabor, with options *frequencies* and *thetas*
    * gaussian, with option *sigmas*
    * sobel, without any options

    Returns
    -------
    np.ndarray
        NxMx(3+P) array with image data with extra channels where P
        is the number of channels added
    '''

    if colorspace.lower() == 'rgb':
        gray = img.mean(-1)
    elif colorspace.lower() == 'hsv':
        gray = img[..., -1]
    else:
        raise ValueError(
            'Unsupported colorspace [%s]' % colorspace)

    # default settings
    methods_params = _update_params(methods_params)

    # create extra channels
    channels = []

    # Create a set of gabor kernels and apply them
    if 'gabor' in methods:
        for i, frequency in enumerate(methods_params['frequencies']):
            response = []
            for theta in methods_params['thetas']:
                real, imag = skimage.filter.gabor_filter(
                    gray,
                    frequency=frequency,
                    #sigma_x=sigma,
                    #sigma_y=sigma,
                    theta=theta)
                response.append(np.sqrt(real ** 2 + imag ** 2))

            response = np.array(response).max(axis=0)

            channels.append(response)

    # Create relative greyishness
    if 'gaussian' in methods:
        for i, sigma1 in enumerate(methods_params['sigmas']):
            # blur
            response1 = skimage.filter.gaussian_filter(gray, sigma1)
            # blur on a higher level
            response2 = skimage.filter.gaussian_filter(gray, sigma1+1)
            # create relative darkness
            response = response2 - response1

            channels.append(response)

    # Create sobel filter channel
    if 'sobel' in methods:
        channels.append(skimage.filter.sobel(gray))

    # Combine all channels
    allchannels = np.concatenate([
        # greyscale
        np.uint8(gray)[..., np.newaxis],
        # color
        img
    ] + [
        # other channels
        channel[..., np.newaxis]
        for channel
        in channels
    ], -1)

    return allchannels

def get_channel_bounds(methods=['gabor', 'gaussian', 'sobel'],
                       methods_params=None):
    '''Get theoretical bounds of channel values

    Parameters
    ----------
    methods : list, optional
        List of strings indicating channels to be added
    methods_params : dict, optional
        Dictionairy with named options for channel functions

    Returns
    -------
    list
        List of dicts with keys *min* and *max* indicating the
        theoretical boundaries of the channel values
    '''

    # default settings
    methods_params = _update_params(methods_params)

    # initialize all bounds to 0 and 255
    stats = [{'max': 0., 'min': 255.}
             for i in range(get_number_channels(methods, methods_params))]

    # update gabor channel bounds, if requested
    if 'gabor' in methods:
        for i, frequency in enumerate(methods_params['frequencies']):
            k = skimage.filter.gabor_kernel(frequency=frequency, theta=0)
            resp = np.zeros(k.shape) + 255.
            resp[np.real(k) < 0] = 0.
            resp_real = np.sum(resp*np.real(k))
            resp_imag = np.sum(resp*np.imag(k))
            resp = np.sqrt(resp_real**2 + resp_imag**2)
            stats[i]['max'] = resp.max()

    return stats

def get_number_channels(methods=['gabor', 'gaussian', 'sobel'],
                        methods_params=None):
    '''Get number of artificial channels

    Parameters
    ----------
    methods : list, optional
        List of strings indicating channels to be added
    methods_params : dict, optional
        Dictionairy with named options for channel functions

    Returns
    -------
    int
        Number of channels added when using the specified settings
    '''

    methods_params = _update_params(methods_params)

    n = 0
    if 'gabor' in methods:
        n += len(methods_params['frequencies'])

    if 'gaussian' in methods:
        n += len(methods_params['sigmas'])

    if 'sobel' in methods:
        n += 1

    return n

def normalize_channel(channel, channelstats):
    '''Scale channel to uint8 based on maximum possible filter response

    Parameters
    ----------
    channel : np.ndarray
        Array with channel values
    channelstats : dict
        Dictionary with fields *min* and *max* indicating the channel
        value bounds of the dataset and used for normalization

    Returns
    -------
    np.ndarray
        Array with normalized channel values
    '''

    return np.uint8((channel-channelstats['min']) /
                    (channelstats['max']-channelstats['min'])*255)

def _update_params(params):
    '''Set channel options to module defaults'''

    methods_params = {'frequencies':_FREQS,
                      'sigmas':_SIGMAS,
                      'thetas':_THETAS}

    methods_params.update(params)

    return methods_params