import logging
import os
from multiprocessing import Pool
from tqdm import tqdm
from collections.abc import Callable
from cdwave.data import WaveformFull, Dataset
from cdwave.fnc import Waveform

logger = logging.getLogger(__name__)


def calc_parameter(waveform: WaveformFull) -> dict:
    """Calculate parameters of a waveform"""
    series = waveform.get_signal_series()
    wave = Waveform(series)
    if not wave.get_peaks():
        return None
    wave.analyse()
    try:
        r = wave.get_parameters()
    except Exception as e:
        logger.error('Cannot get parameters of %s', waveform)
        raise e
    return r


pbar_wrapper = {'x': None}


def default_process_fnc(status=-1, total=0):
    pbar = pbar_wrapper['x']
    if status == 0 and total != 0:
        pbar_wrapper['x'] = tqdm(total=total)
    elif status == 1:
        pbar.update()
    elif status == 2:
        pbar.close()


def calc_parameters_for_waveforms(dataset: Dataset,
                                  process_fnc: Callable = None,
                                  batch: int = 200,
                                  processes: int = None,
                                  custom_calculator: Callable = None):
    """Calculate parameter for waveforms

    Args:
        dataset: The waveform dataset
        process_fnc: A processing function used to send out the progress,
            see `default_process_fnc`, which uses tqdm
        batch: Batch size for multi-processing
        processes: Number of processors
        custom_calculator: A custom calculator which can setup the custom thresholds.
            If `None`, `calc_parameter` will be used by default.
    """
    if processes is None:
        try:
            processes = int(os.environ['NUMEXPR_MAX_THREADS'])
        except KeyError:
            processes = 10
        except ValueError:
            processes = 10
    if custom_calculator is None:
        custom_calculator = calc_parameter
    n = dataset.size
    if process_fnc:
        process_fnc(status=0, total=n)
    reses = []
    batch_waveforms = []
    for i, waveform in enumerate(dataset.waveforms):
        batch_waveforms.append(waveform)
        if i % batch == batch - 1 or i == n - 1:
            with Pool(processes=processes) as pool:
                for _waveform in batch_waveforms:
                    reses.append(pool.apply_async(
                        custom_calculator, (_waveform,)))
                for j, res in enumerate(reses):
                    r = res.get()
                    if r is not None:
                        batch_waveforms[j].parameters = r
            reses = []
            batch_waveforms = []
        if process_fnc:
            process_fnc(status=1)
    if process_fnc:
        process_fnc(status=2)
