#!/usr/bin/env python3

# -*- coding: utf-8 -*-

"""

Created on Wed Nov 21 12:52:49 2018

@author: shenwanxiang

Multi process Run

"""

import time

import pandas as pd

from tqdm import tqdm

from concurrent.futures import ProcessPoolExecutor, wait, as_completed

from multiprocessing import Pool,cpu_count,current_process 

import subprocess

from aggmap.utils.logtools import print_info, print_error, pbar,print_warn

def RunCmd(cmd):

    '''

    input:

        cmd: str

    output:

        status: int, 0 for success

        stdout: str

        stderr: str

    '''

    print_info('run command : %s' % cmd)

    def swap_log(swap, error = True):

        sinfo = []

        for l in swap.split('\n'):

            if l == '':

                continue

            sinfo.append(l)

        for o in sinfo:

            if error:

                print_error(o) 

            else:

                print_info(o) 

        return            

    output = subprocess.run(cmd, 

                            shell=True, 

                            stdout=subprocess.PIPE, 

                            stderr=subprocess.PIPE, 

                            universal_newlines=True)    

    status = output.returncode

    stdout = output.stdout

    stderr = output.stderr

    if status != 0:

        if output.stdout:

             swap_log(output.stdout, error=True)

        if output.stderr:

             swap_log(output.stderr, error=True)

    else:

        if output.stdout:

            swap_log(output.stdout, error=False)

    #return status

    return status, stdout, stderr

def ImapUnorder(processor, iterator, max_workers=10, fail_in_file = './filed.lst'):

    '''

    processor: fuction

    iterator: list or iterator,each element should be a tuple or dict, so that data can be used as ordered 

    '''

    with ProcessPoolExecutor(max_workers=max_workers) as executor:

        with open(fail_in_file, 'w+') as f:

            futures = {executor.submit(processor, IdPlusSmile):IdPlusSmile for IdPlusSmile in iterator}

            success, _ = wait(futures)

            with pbar(total = len(futures)) as pb:

                for i in success:

                    IdPlusSmile = futures[i]

                    print_info('deal '+ str(IdPlusSmile))

                    try:

                        data_dict = i.result()

                        yield data_dict

                    except Exception as exc:

                        print_warn('because of the process is dead, input: %s is fialed when deal with %s: %s, so we will deal it automatically' % (IdPlusSmile, processor, exc))

                        try: 

                            yield processor(IdPlusSmile)

                        except:

                            f.write(str(IdPlusSmile)+'\n')

                            print_error(' input: %s is fialed when deal with %s: %s' % (IdPlusSmile, processor, exc))

                    pb.update(1)

def MultiProcessUnorderedBarRun(func, deal_list, n_cpus=None):

    if n_cpus ==None:

        N_CPUS = cpu_count()

    else:

        N_CPUS = int(n_cpus)

    print_info('the number of process is %s' % N_CPUS)

    p = Pool(N_CPUS)

    res_list = []

    with pbar(total = len(deal_list), ascii=True) as pb:

        for res in p.imap_unordered(func, deal_list):

            pb.update(1)

            res_list.append(res)

    p.close()

    p.join()

    return res_list

def MultiProcessRun(func, deal_list, n_cpus=None):

    '''

    input:

        func: function to do with each element in the deal_list

        deal_list: list to be done

        n_cpus: use the number of cpus

    output:

        list of the return result for each func

    '''

    #round_c = [deal_list[i:i+batch_size] for i  in range(0, len(deal_list), batch_size)]

    #mata thinking: https://my.oschina.net/leejun2005/blog/203148

    if n_cpus ==None:

        N_CPUS = cpu_count()

    else:

        N_CPUS = int(n_cpus)

    print_info('the number of process is %s' % N_CPUS)

    pool = Pool(N_CPUS)

    a = pool.map(func, deal_list)

    pool.close()

    pool.join()

    return a

########### ordered map reduce  ##############

def _decorate_func(func, i, j):

    return [i, func(j)]

def _executor(func, series, n_cpus = 4):

    with ProcessPoolExecutor(max_workers=n_cpus) as executor:

        futures = [executor.submit(_decorate_func, func, i, j) for i,j in series.iteritems()]

    return futures

def MultiExecutorRun(func, deal_list, n_cpus = 4, tqdm_args = {'unit':'one'}):

    '''

    input:

        func: function to do with each element in the deal_list

        deal_list: list to be done

        n_cpus: use the number of cpus

        tqdm_args: args for tqdm

    output:

        list of the return value for each func

    '''

    lst  =list(deal_list)

    series = pd.Series(lst)

    futures = _executor(func, series, n_cpus = n_cpus)

    args = {

        'total': len(deal_list),

        'unit': 'one',

        'ascii': True,

        'unit_scale': True,

        'leave': True

    }

    args.update(tqdm_args)

    print_info(args)

    results = []

    indexs = []

    for f in tqdm(as_completed(futures), **args):

        #print(f)

        idx, result = f.result()

        indexs.append(idx)

        results.append(result)

    res = pd.Series(results,index=indexs)

    #sort unordered result

    ordered_lst = res.sort_index().tolist()

    return ordered_lst