from anyBSM import anyBSM as anyBSM_single
from numpy import array
from collections import defaultdict
from typing import get_type_hints
from tqdm import tqdm
from multiprocessing import Process, Queue
import logging

__doc__ = """
# Basic Parallelism
Example usage:
```python
from anyBSM.parallel import anyBSM
import numpy as np

# initialize 8 parallel worker models and one master-model
SM=anyBSM('SM',num_workers=8);

# set variable for all workers
SM.to_workers('warnSSSS', False)

# read variable from all workers
SM.from_workers('warnSSSS')

# prepare (dict-typed) input
inputs = [{'parameters': {'Mh': mh}} for mh in np.linspace(20,1000,50)]

# compute lambdahhh for all inputs, distributed across all workers
results = SM.lambdahhh(worker_kwargs=inputs,transpose=True)

# compute max SS->SS eigenvalue for all inputs, distributed across all workers
results = SM.eigSSSS(worker_kwargs=inputs)

# prepare another (list-typed) input and compute differential HH production xs
inputs = [mh**2 for mh in np.linspace(250,1001,20)]
SM.sigmahat_ggHH(worker_args=inputs,pdf_scale=250,transpose=True)

```
"""

def list2dict(ll, fast = False):
    """ convert list-of-dicts to dict-of-lists """
    if fast:
        result = defaultdict(list)
        [ result[k].append(v) for k,v in ll.items() for entry in ll ]
        return result

    result = defaultdict(list)
    keys = set()
    for entry in ll:
        if type(entry) is dict:
            keys = keys.union(set(entry.keys()))
    for k in keys:
        result[k] = []
        for entry in ll:
            if type(entry) is dict and k in entry:
                result[k].append(entry[k])
            else:
                result[k].append(entry)
    return result

class anyBSMWorker(Process):
    """ anyBSM worker process
    Spawns an anyBSM instance with given setup_kwargs.
    Listens on task_queue and writes results into result_queue
    """
    def __init__(self,
                 task_queue: Queue,
                 result_queue: Queue,
                 model : str,
                 setup_kwargs: dict = {}):
        super().__init__()
        self.setup_kwargs = {'quiet' : True, 'progress': False, 'ask': False}
        self.setup_kwargs.update(setup_kwargs)

        self.task_queue = task_queue
        self.result_queue = result_queue

        self.model = model

    def run(self):
        model = anyBSM_single(self.model, **self.setup_kwargs)

        while True:
            job = self.task_queue.get()
            if job == "STOP":
                break
            if len(job) == 2:
                if job[0] == "SET":
                    try:
                        setattr(model, job[1][0], job[1][1])
                        self.result_queue.put((job[1][0], job[1][1]))
                    except Exception as e:
                        self.result_queue.put((job, str(e)))
                    continue
                if job[0] == "GET":
                    try:
                        self.result_queue.put((f'{self.name}: {job[1]} =', getattr(model,job[1])))
                    except Exception as e:
                        self.result_queue.put((job, str(e)))
                    continue

            if len(job) != 3:
                self.result_queue.put((job, 'ERROR: job requires exactly three arguments (id, func_name, (*args, **kwargs))'))
                continue

            if job[1] == "EVAL":
                try:
                    result = model._eval(job[1][0],**job[1][1])
                    self.result_queue.put((job[0], job[1][1], result))
                except Exception as e:
                    self.result_queue.put((job, str(e)))
                continue

            eval_func = job[1]
            if isinstance(job[2], dict):
                kwargs = job[2]
                args = []
            else:
                args, kwargs = job[2]
            try:
                method = getattr(model, eval_func)
                result = method(*args, **kwargs)
                # print(self.name, eval_func, args,kwargs)
                self.result_queue.put((job[0], result))
            except Exception as e:
                self.result_queue.put((job[0], str(e)))


class anyBSM(anyBSM_single):
    def __init__(self, model : str, num_workers : int = 4, **kwargs):
        super().__init__(model, **kwargs)
        self._master_manager  = True
        self._model = model
        self._num_workers = num_workers
        self._setup_kwargs = kwargs
        self._task_queues = []
        self._result_queue = Queue()
        self._workers = []

        self.respawn_workers()

    def respawn_workers(self):
        self._shutdown()
        for i in range(self._num_workers):
            tq = Queue()
            worker = anyBSMWorker(tq, self._result_queue, self._model, self._setup_kwargs)
            worker.start()
            self._task_queues.append(tq)
            self._workers.append(worker)

    # def __setattr__(self, name, value):
    #     super().__setattr__(name, value)

    #     # Skip local-only attributes or private ones
    #     if name.startswith('_') or not hasattr(self, '_num_workers'):
    #         return

    #     # Broadcast to all workers
    #     for tq in self._task_queues:
    #         tq.put(("SET",[name, value]))
    #     for _ in range(self._num_workers):
    #         self._result_queue.get()

    def __getattribute__(self, name):
        if name.startswith("_") or not hasattr(self, '_master_manager'):
            return object.__getattribute__(self, name)

        attr = super().__getattribute__(name)

        # Wrap callable methods for optional parallelism
        if callable(attr) and not name.startswith("_"):
            def wrapper(*args,worker_args=[], worker_kwargs=[], transpose=False, **kwargs):
                inputs = []
                if worker_args and worker_kwargs:
                    if len(worker_args) != len(worker_kwargs):
                        logging.error(f'worker_args and worker_kwargs should be of same length: {len(worker_args)}!={len(worker_kwargs)}')
                    inputs = [[a,{**k, **kwargs}] for a,k in zip(worker_args,worker_kwargs)]
                elif worker_args:
                    shape = array(worker_args).shape
                    if len(shape) == 1:
                        inputs = [[[a],kwargs] for a in worker_args]
                    else:
                        inputs = [[a,kwargs] for a in worker_args]
                elif worker_kwargs:
                    inputs = [[args,{**k, **kwargs}] for k in worker_kwargs]
                if inputs:
                    if args and not worker_kwargs:
                        logging.warning('additional arguments ignored')
                    return self.parallel_map(name, inputs, transpose)
                return attr(*args, **kwargs)
            signature = f'type hints of non-paralel function: {get_type_hints(attr)}'
            if attr.__doc__:
                wrapper.__doc__ = attr.__doc__ + signature
            return wrapper

        return attr

    def parallel_map(self, func_name, inputs,transpose = False):
        li = len(inputs)
        for i, item in enumerate(inputs):
            self._task_queues[i % self._num_workers].put((i, func_name, item))

        results = []
        for _ in tqdm(range(li), desc=f"Parallel {func_name}"):
            results.append(self._result_queue.get())

        results = dict(results)
        results = [results[i] for i in range(li)]
        if transpose:
            hints = get_type_hints(getattr(anyBSM_single,func_name))
            htype = hints.get('return',None)
            if htype is dict:
                results = list2dict(results)
            elif htype is tuple or htype is list:
                results = array(results).transpose()
        return results

    def parallel_eval(self, string, inputs):
        """ runs anyBSM._eval(string, **inputs) with inputs in parallel.
        Args:
           * string: a str to evaluate; any string compatible with the current model
           * inputs: a list of dictionaries. Entries are computed in parallel
        """
        li = len(inputs)
        for i, item in enumerate(inputs):
            self._task_queues[i % self._num_workers].put(('EVAL', (i, string, item)))
        results = []
        for _ in tqdm(range(li), desc="Evaluating"):
            results.append(self._result_queue.get())
        results = dict(results)
        return [results[i] for i in range(li)]

    def from_workers(self, attr):
        for tq in self._task_queues:
            tq.put(("GET", attr))

        return [self._result_queue.get() for _ in self._task_queues]

    def to_workers(self, name, value):
        for tq in self._task_queues:
            tq.put(("SET",[name, value]))
        for _ in range(self._num_workers):
            self._result_queue.get()

    def _shutdown(self):
        for tq in self._task_queues:
            tq.put("STOP")
        for w in self._workers:
            w.join()

    def __exit__(self):
        self._shutdown()

    def __del__(self):
        self._shutdown()

    def __enter__(self):
        return self
