# Path: sampyl/state.py
# class State(collections.OrderedDict):
#     """ State object for storing parameter values.
#         
#         Inherits from OrderedDict.
# 
#     """
# 
#     def tovector(self):
#         """ Return the parameter values as a flat vector. """
#         return np.hstack(self.values())
# 
#     def fromvector(self, vec):
#         """ Update the state using a numpy array. 
# 
#             :param vec: np.array for updating the state.
#         """
#         var_sizes = self.size()
#         i = 0
#         for var in self:
#             self[var] = np.squeeze(vec[i:(i+var_sizes[var])])
#             i += var_sizes[var]
#         return self
# 
#     def freeze(self):
#         """ Return a immutable tuple of the state values."""
#         return tuple(self.tovector())
# 
#     @staticmethod
#     def init_fromvector(vec, state):
#         ""
#         vals = []
#         var_sizes = state.size()
#         i = 0
#         for var in state:
#             vals.append(np.squeeze(vec[i:(i+var_sizes[var])]))
#             i += var_sizes[var]
#         return State(zip(state.keys(), vals))
# 
#     @staticmethod
#     def fromfunc(func):
#         """ Initialize a State from the arguments of a function """
#         var_names = func_var_names(func)
#         return State.fromkeys(var_names)
# 
#     def size(self):
#         return State([(var, np.size(self[var])) for var in self])
# 
#     def __add__(self, other):
#         if isinstance(other, int) or isinstance(other, float):
#             return handle_number(self, other, '__add__')
#         elif isinstance(other, collections.Iterable):
#             return handle_iterable(self, other, '__add__')
#         else:
#             raise TypeError("Addition not supported for State and {}".format(other))
# 
#     def __sub__(self, other):
#         if isinstance(other, int) or isinstance(other, float):
#             return handle_number(self, other, '__sub__')
#         elif isinstance(other, collections.Iterable):
#             return handle_iterable(self, other, '__sub__')
#         else:
#             raise TypeError("Subtraction not supported for State and {}".format(other))
# 
#     def __mul__(self, other):
#         if isinstance(other, int) or isinstance(other, float):
#             return handle_number(self, other, '__mul__')
#         else:
#             raise TypeError("Multiplication not supported for State and {}".format(other))
# 
#     def __truediv__(self, other):
#         if isinstance(other, int) or isinstance(other, float):
#             return handle_number(self, other, '__truediv__')
#         else:
#             raise TypeError("Division not supported for State and {}".format(other))
# 
#     def __radd__(self, other):
#         if isinstance(other, int) or isinstance(other, float):
#             # Commutative, so nothing changes
#             return self + other
#         else:
#             raise TypeError("Can only broadcast from the left.")
# 
#     def __rmul__(self, other):
#         if isinstance(other, int) or isinstance(other, float):
#             # Commutative, so nothing changes
#             return self * other
#         else:
#             raise TypeError("Can only broadcast from the left.")
# 
#     def __rsub__(self, other):
#         if isinstance(other, int) or isinstance(other, float):
#             return handle_number(self, other, '__rsub__')
#         elif isinstance(other, collections.Iterable):
#             return handle_iterable(self, other, '__rsub__')
#         else:
#             raise TypeError("Subtraction not supported for State and {}".format(other))
# 
#     def __rtruediv__(self, other):
#         if isinstance(other, int) or isinstance(other, float):
#             return handle_number(self, other, '__truediv__')
#         else:
#             raise TypeError("Division not supported for State and {}".format(other))
# 
# Path: sampyl/progressbar.py
# def update_progress(current, total, width=30, end=False):
#     bar_width = width
#     block = int(round(bar_width * current/total))
#     text = "\rProgress: [{0}] {1} of {2} samples".\
#              format("#"*block + "-"*(bar_width-block), current, total)
#     if end:
#         text = text +'\n'
#     sys.stdout.write(text)
#     sys.stdout.flush()


import copy
from functools import partial
from multiprocessing import Pool
from .core import np, AUTOGRAD
from .samplers import *
from .distributions import *
from .state import State
from .progressbar import update_progress


def parallel(sampler, n_chains, samples, progress_bar=True, **kwargs):

    samplers = init_samplers(sampler, n_chains)
    chains = [copy.copy(samples) for _ in range(n_chains)]

    N_batches = 10
    n_samples = len(samples)
    batches = [n_samples//N_batches]*N_batches
    batches.append(n_samples % N_batches)

    pool = Pool(processes=n_chains)
    for i, N in enumerate(batches):
        func = partial(f, N)

        if i != 0:
            # Reinitialize samplers where the previous batch left off so that
            # the new batch starts at the correct state
            samplers = init_samplers(sampler, n_chains, chains=new_chains)
        new_chains = pool.map(func, samplers)

        for new, chain in zip(new_chains, chains):
            for j in range(N):
                chain[i*N + j] = new[j]

        if progress_bar:
                update_progress(N*(i+1), n_samples)

    if progress_bar:

# Path: sampyl/state.py
# class State(collections.OrderedDict):
#     """ State object for storing parameter values.
#         
#         Inherits from OrderedDict.
# 
#     """
# 
#     def tovector(self):
#         """ Return the parameter values as a flat vector. """
#         return np.hstack(self.values())
# 
#     def fromvector(self, vec):
#         """ Update the state using a numpy array. 
# 
#             :param vec: np.array for updating the state.
#         """
#         var_sizes = self.size()
#         i = 0
#         for var in self:
#             self[var] = np.squeeze(vec[i:(i+var_sizes[var])])
#             i += var_sizes[var]
#         return self
# 
#     def freeze(self):
#         """ Return a immutable tuple of the state values."""
#         return tuple(self.tovector())
# 
#     @staticmethod
#     def init_fromvector(vec, state):
#         ""
#         vals = []
#         var_sizes = state.size()
#         i = 0
#         for var in state:
#             vals.append(np.squeeze(vec[i:(i+var_sizes[var])]))
#             i += var_sizes[var]
#         return State(zip(state.keys(), vals))
# 
#     @staticmethod
#     def fromfunc(func):
#         """ Initialize a State from the arguments of a function """
#         var_names = func_var_names(func)
#         return State.fromkeys(var_names)
# 
#     def size(self):
#         return State([(var, np.size(self[var])) for var in self])
# 
#     def __add__(self, other):
#         if isinstance(other, int) or isinstance(other, float):
#             return handle_number(self, other, '__add__')
#         elif isinstance(other, collections.Iterable):
#             return handle_iterable(self, other, '__add__')
#         else:
#             raise TypeError("Addition not supported for State and {}".format(other))
# 
#     def __sub__(self, other):
#         if isinstance(other, int) or isinstance(other, float):
#             return handle_number(self, other, '__sub__')
#         elif isinstance(other, collections.Iterable):
#             return handle_iterable(self, other, '__sub__')
#         else:
#             raise TypeError("Subtraction not supported for State and {}".format(other))
# 
#     def __mul__(self, other):
#         if isinstance(other, int) or isinstance(other, float):
#             return handle_number(self, other, '__mul__')
#         else:
#             raise TypeError("Multiplication not supported for State and {}".format(other))
# 
#     def __truediv__(self, other):
#         if isinstance(other, int) or isinstance(other, float):
#             return handle_number(self, other, '__truediv__')
#         else:
#             raise TypeError("Division not supported for State and {}".format(other))
# 
#     def __radd__(self, other):
#         if isinstance(other, int) or isinstance(other, float):
#             # Commutative, so nothing changes
#             return self + other
#         else:
#             raise TypeError("Can only broadcast from the left.")
# 
#     def __rmul__(self, other):
#         if isinstance(other, int) or isinstance(other, float):
#             # Commutative, so nothing changes
#             return self * other
#         else:
#             raise TypeError("Can only broadcast from the left.")
# 
#     def __rsub__(self, other):
#         if isinstance(other, int) or isinstance(other, float):
#             return handle_number(self, other, '__rsub__')
#         elif isinstance(other, collections.Iterable):
#             return handle_iterable(self, other, '__rsub__')
#         else:
#             raise TypeError("Subtraction not supported for State and {}".format(other))
# 
#     def __rtruediv__(self, other):
#         if isinstance(other, int) or isinstance(other, float):
#             return handle_number(self, other, '__truediv__')
#         else:
#             raise TypeError("Division not supported for State and {}".format(other))
# 
# Path: sampyl/progressbar.py
# def update_progress(current, total, width=30, end=False):
#     bar_width = width
#     block = int(round(bar_width * current/total))
#     text = "\rProgress: [{0}] {1} of {2} samples".\
#              format("#"*block + "-"*(bar_width-block), current, total)
#     if end:
#         text = text +'\n'
#     sys.stdout.write(text)
#     sys.stdout.flush()


# Path: sampyl/parallel.py
import copy
from functools import partial
from multiprocessing import Pool
from .core import np, AUTOGRAD
from .samplers import *
from .distributions import *
from .state import State
from .progressbar import update_progress



def parallel(sampler, n_chains, samples, progress_bar=True, **kwargs):

    samplers = init_samplers(sampler, n_chains)
    chains = [copy.copy(samples) for _ in range(n_chains)]

    N_batches = 10
    n_samples = len(samples)
    batches = [n_samples//N_batches]*N_batches
    batches.append(n_samples % N_batches)

    pool = Pool(processes=n_chains)
    for i, N in enumerate(batches):
        func = partial(f, N)

        if i != 0:
            # Reinitialize samplers where the previous batch left off so that
            # the new batch starts at the correct state
            samplers = init_samplers(sampler, n_chains, chains=new_chains)
        new_chains = pool.map(func, samplers)

        for new, chain in zip(new_chains, chains):
            for j in range(N):
                chain[i*N + j] = new[j]

        if progress_bar:
                update_progress(N*(i+1), n_samples)

    if progress_bar:

        update_progress(n_samples, n_samples, end=True)