#!/bin/python

from util import log_error
import numpy as np
from enum import Enum


def noop(samples, target, args):
    pass


def delta(samples, target, args):
    source = args[0]
    for i, _ in enumerate(samples[target]):
        if i == 0 or samples[source][i] == np.nan:
            samples[target][i] = np.nan
        else:
            samples[target][i] = samples[source][i] - samples[source][i - 1]


def average(samples, target, args):
    source = args[0]
    samples[target] = np.nanmean(samples[source])


def max(samples, target, args):
    source = args[0]
    samples[target] = np.nanmax(samples[source])


def min(samples, target, args):
    source = args[0]
    samples[target] = np.nanmin(samples[source])


def variance(samples, target, args):
    source = args[0]
    samples[target] = np.nanvar(samples[source])


def stdev(samples, target, args):
    source = args[0]
    samples[target] = np.nanstd(samples[source])


def percentil(samples, target, args):
    source = args[0]
    perc = args[1]
    samples[target] = np.nanpercentile(samples[source], perc)


def average2sigmas(samples, target, args):
    av = args[0]
    std = args[1]
    samples[target] = samples[av] + (2 * samples[std])


def average_2sigmas(samples, target, args):
    av = args[0]
    std = args[1]
    samples[target] = samples[av] - (2 * samples[std])


def _binary_operation(samples, target, args, op):
    operand_1 = args[0]
    operand_2 = args[1]
    for i, _ in enumerate(samples[target]):
        samples[target][i] = op(samples[operand_1][i], samples[operand_2][i])


def get_arg(samples, arg):
    if isinstance(arg, str):
        return samples[arg]
    else:
        return arg


def mul(samples, target, args):
    samples[target] = get_arg(samples, args[0])
    for arg in args[1:]:
        samples[target] *= get_arg(samples, arg)


def add(samples, target, args):
    samples[target] = get_arg(samples, args[0])
    for arg in args[1:]:
        samples[target] += get_arg(samples, arg)


def sub(samples, target, args):
    samples[target] = get_arg(samples, args[0])
    for arg in args[1:]:
        samples[target] -= get_arg(samples, arg)


def div(samples, target, args):
    samples[target] = get_arg(samples, args[0])
    for arg in args[1:]:
        samples[target] /= get_arg(samples, arg)


class ArityType(Enum):
    AT_LEAST = 1
    AT_MOST = 2
    EXACTLY = 3


class Operator:
    def __init__(self, f, arity, arity_type=ArityType.EXACTLY):
        self.op = f
        self.arity = arity
        self.arity_type = arity_type

    def __repr__(self):
        return f'{self.op.__name__}'


class Functions:
    operators = {
        "delta": Operator(delta, 1),
        "average": Operator(average, 1),
        "mul": Operator(mul, 2, ArityType.AT_LEAST),
        "add": Operator(add, 2, ArityType.AT_LEAST),
        "sub": Operator(sub, 2, ArityType.AT_LEAST),
        "div": Operator(div, 2, ArityType.AT_LEAST),
        "noop": Operator(noop, 0),
        "max": Operator(max, 1),
        "min": Operator(min, 1),
        "variance": Operator(variance, 1),
        "standard_deviation": Operator(stdev, 1),
        "percentil": Operator(percentil, 2),
        "average2sigmas": Operator(average2sigmas, 2),
        "average_2sigmas": Operator(average_2sigmas, 2),
    }

    @staticmethod
    def get_function(name):
        if name not in Functions.operators:
            log_error('when getting operator, no function with name "{}" found.'.format(name))
            return Functions.operators["noop"].op
        return Functions.operators[name].op

    @staticmethod
    def get_function_arity(name):
        if name not in Functions.operators:
            log_error('when getting arity of operator, no function with name "{}" found.'.format(name))
            operator = Functions.operators["noop"]
        else:
            operator = Functions.operators[name]
        return (operator.arity, operator.arity_type)

    @staticmethod
    def get_function_arity_type(name):
        if name not in Functions.operators:
            log_error('when getting arity of operator, no function with name "{}" found.'.format(name))
            return Functions.operators["noop"].arity_type
        return Functions.operators[name].arity_type
