import re
from functions import Functions
import numpy as np


def eval_obj(x, samples):
    if isinstance(x, float):
        return x
    elif isinstance(x, str):
        return samples[x]
    else:
        return x.eval(samples)


def validate_obj(x, samples):
    if isinstance(x, float):
        return True
    elif isinstance(x, str):
        return x in samples
    else:
        return x.validate(samples)


class Expr:
    def __init__(self, func, args=[]):
        self.func = func
        self.args = args

    def __str__(self):
        args = [str(i) for i in self.args]
        return f'Expr({repr(self.func)} : {args})'

    def eval(self, samples):
        f = lambda i: eval_obj(i, samples)
        args = list(map(f, self.args))
        self.func(samples, "__aux__", args)
        return samples["__aux__"]

    def validate(self, samples):
        f = lambda i: validate_obj(i, samples)
        validation = list(map(f, self.args))
        return False not in validation

    def __parse(self, in_list):
        args = []
        expr = None
        tokens_consumed = 0
        i = 0
        while i < (len(in_list)):
            token = in_list[i]
            tokens_consumed += 1
            if token in Functions.operators:
                expr = Expr(Functions.get_function(token))
            elif token == "(":
                t, new_expr = expr.__parse(in_list[i + 1:])
                args.append(new_expr)
                i += t
                tokens_consumed += t
            elif token == ")":
                self.args = args
                return tokens_consumed, self
            elif token == ",":
                pass
            elif token[0] == "$":
                args.append(token[1:])
            else:
                args.append(float(token))
            i += 1
        return expr

    @staticmethod
    def _parse(in_list):
        aux = Expr(None)
        return aux.__parse(in_list)

    @staticmethod
    def parse_str(in_str):
        in_str = in_str.replace(' ', '')
        in_list = re.split(r'(,|\(|\))', in_str)
        in_list = list(filter(''.__ne__, in_list))
        try:
            return Expr._parse(in_list)
        except Exception:
            raise Exception("malformed expression")


if __name__ == "__main__":
    expr = Expr.parse_str("add(mul($foo, 0.01, 0.25), mul(div($bar, 1000), 0.5), mul(div($keke, $baba), 0.25))")
    # expr = Expr.parse_str("add($foo, $bar, mul($baba, $keke))")
    print(str(expr))
    samples = {
        "__aux__": None,
        "foo": np.asarray([1, 2, 3, 4, 5], dtype=np.float64),
        "bar": np.asarray([1, 2, 3, 4, 5], dtype=np.float64),
        "keke": np.asarray([1, 2, 3, 4, 5], dtype=np.float64),
        "baba": np.asarray([1, 2, 3, 4, 5], dtype=np.float64)
    }
    print(expr.eval(samples))
