from functools import wraps
import time
import concurrent.futures as futures
from tqdm import tqdm
import multiprocessing
from multiprocessing import Pipe, Process, Manager

def time_count(function):
    @wraps(function)
    def decorated(*args, **kwargs):
        T1 = time.perf_counter()
        results = function(*args, **kwargs)
        T2 =time.perf_counter()
        print(f'{function} Time cost:{(T2 - T1)*1000} ms')
        return results
    return decorated


def inputlines(prompt=None):
    if prompt is not None: print(prompt)
    lines=[]
    while True:
        try:
            line=input()
            if line=='': break
            lines.append(line)
        except:
            break
    return lines

def threadMapTqdm(function, my_iter,description=None,workers=10,single_arg=False,closebar=False):
    def updatethread(function,arg,pbar,single_arg):
        if type(arg) not in [list,tuple] or single_arg:
            result=function(arg)
        else:
            result=function(*arg)
        pbar.update(1)
        return result
    l = len(my_iter)
    with tqdm(total=l) as pbar:
        if description is not None: pbar.set_description(description)
        with futures.ThreadPoolExecutor(max_workers=workers) as executor:
            jobs = {executor.submit(lambda p: updatethread(*p),[function,arg,pbar,single_arg]): idx for idx,arg in enumerate(my_iter)}
            results = {}
            for job in futures.as_completed(jobs):
                idx = jobs[job]
                results[idx] = job.result()
    results_sort=[results[idx] for idx in sorted(results)]
    if closebar:
        pbar.close()
    return results_sort


def updateprocess(function,arg,single_arg):
    if type(arg) not in [list,tuple] or single_arg:
        result=function(arg)
    else:
        result=function(*arg)
    return result

def processMapTqdm(function, my_iter,description=None,workers=4,single_arg=False,closebar=False):
    pool = multiprocessing.Pool(processes = workers)
    job_num=len(my_iter)
    results={}
    pbar=tqdm(total=job_num,leave=not closebar, desc=description)
    def updatebar(msg):
        pbar.update(1)
    for i, item in enumerate(my_iter):
        results[i]=pool.apply_async(updateprocess,[function, item, single_arg],callback=updatebar)
    pool.close()
    results_sort=[results[idx].get() for idx in sorted(results)]
    if closebar:
        pbar.close()
    return results_sort


def block_loop(function,arg,single_arg, process_index, pipe, return_dict):
    _out_pipe, _in_pipe = pipe
    _out_pipe.close()
    results=[]
    for item in arg:
        if type(item) not in [list,tuple] or single_arg:
            result=function(item)
        else:
            result=function(*item)
        results.append(result)
        _in_pipe.send(1)
    return_dict[process_index]=results
    return results

def processBlockTqdm(function, my_iter,description=None,workers=4,single_arg=False,closebar=False):
    manager=Manager()
    results=manager.dict()

    job_num=len(my_iter)
    my_iter_list=list(my_iter)
    block_size=job_num // workers + int(job_num % workers > 0)
    iter_blocks=[my_iter_list[block_size*i:block_size*i+block_size] for i in range(workers)]
    
    out_pipe, in_pipe = Pipe(True)
    process=[]
    pbar=tqdm(total=job_num,leave=not closebar, desc=description)
    for i in range(workers):
        new_process=Process(target=block_loop, args=[function, iter_blocks[i], single_arg, i, (out_pipe, in_pipe), results])
        process.append(new_process)
        new_process.start()
    in_pipe.close()
    while True:
        try:
            msg = out_pipe.recv()
            if msg == 1: pbar.update(1)
        except EOFError as e:
            print(e)
            break
    results_sort=[]

    for idx in sorted(results):
        results_sort+=results[idx]
    if closebar:
        pbar.close()
    return results_sort