from .grammar import *
parse_string = parser.parse


import ast

def eval_scalar(tree):
    try:
        if isinstance(tree, ast.Num):
            return tree.n
        elif isinstance(tree, ast.UnaryOp):
            if isinstance(tree.op, ast.USub):
                return -tree.operand.n
            elif isinstance(tree.op, ast.UAdd):
                return tree.operand.n
        else:
            raise Exception("Don't know how to do that.")
    except:
        raise Exception("Don't know how to do that.")


# from dolang.language import functions as functions_dict
# from dolang.codegen import to_source
# from typing import Tuple, List, Dict, Set, TypeVar, Union
# from ast import NodeTransformer, Name, UnaryOp, UAdd, USub, Load, Call

# try:
#     from .dataclasses import dataclass
# except:
#     from dataclasses import dataclass

# import ast
# from ast import Expr
# from functools import wraps
# import re

# T = TypeVar('T')

# Expression = ast.Expr

# functions = list(functions_dict.keys())

# __regex_eq__ = re.compile("([^=]*)=([^=]*)")


# def parse_string(s: str) -> ast.Expr:
#     ss = s.replace("^", "**")
#     m = re.match(__regex_eq__, ss)
#     if m:
#         ss = str.join('==', m.groups())
#     return ast.parse(ss).body[0]


# def dedup(l: List[T]) -> List[T]:
#     return list(dict.fromkeys(l))


# def stringify_variable(arg: Tuple[str, int]) -> str:
#     s = arg[0]
#     date = arg[1]
#     if date == 0:
#         return '{}__0_'.format(s)
#     elif date <= 0:
#         return '{}_m{}_'.format(s, str(-date))
#     elif date >= 0:
#         return '{}__{}_'.format(s, str(date))


# def stringify_parameter(p: str) -> str:
#     return '{}_'.format(p)


# def stringify_symbol(arg) -> str:
#     if isinstance(arg, str):
#         return stringify_parameter(arg)
#     elif isinstance(arg, tuple):
#         if len(arg) == 2 and isinstance(arg[0], str) and isinstance(
#                 arg[1], int):
#             return stringify_variable(arg)
#     raise Exception("Unknown canonical form: {}".format(arg))


# def destringify_variable(s: str) -> Tuple[str, int]:
#     raise Exception("Not implemented.")


# def destringify_parameter(s: str) -> Tuple[str, int]:
#     raise Exception("Not implemented.")


# def destringify(s: str):
#     raise Exception("Not implemented.")


# ###
# # The following function can take a string or an expression as the first argument, and return an argument of the same type.
# ###


# def expression_or_string(f):
#     @wraps(f)
#     def wrapper(*args, **kwds):
#         if not isinstance(args[0], str):
#             return f(*args, **kwds)
#         else:
#             a = parse_string(args[0])
#             nargs = tuple([a]) + args[1:]
#             res = f(*nargs, **kwds)
#             return to_source(res)

#     return wrapper


# @expression_or_string
# def stringify(expr: Expression, variables: List[str] = []) -> Expression:
#     import copy
#     en = ExpressionStringifier(variables=variables)
#     return en.visit(copy.deepcopy(expr))


# # shift timing in equations
# # time_shift(:(a+b(1)+c),1,[:b,:c]) == :(a+b(2)+c(1))


# @expression_or_string
# def time_shift(expr: Expression, n) -> Expression:
#     import copy
#     eexpr = copy.deepcopy(expr)
#     variables = [e[0] for e in list_variables(eexpr)]
#     return TimeShiftTransformer(shift=n, variables=variables).visit(eexpr)


# @expression_or_string
# def steady_state(expr: Expression) -> Expression:
#     import copy
#     eexpr = copy.deepcopy(expr)
#     variables = [e[0] for e in list_variables(eexpr)]
#     return TimeShiftTransformer(shift='S', variables=variables).visit(eexpr)


# @expression_or_string
# def sanitize(expr: Expression, variables=None, functions=None):
#     # special functions ?
#     es = ExpressionSanitizer(variables=variables)
#     return es.visit(expr)


# def list_variables(expr: Expression,
#                    funs: List[str] = None) -> List[Tuple[str, int]]:

#     if funs is None:
#         funs = []
#     ll = ListSymbols(known_functions=functions + funs)
#     ll.visit(expr)
#     if ll.problems:
#         e = Exception('Symbolic error.')
#         e.problems = ll.problems
#         raise e
#     return dedup([v[0] for v in ll.variables])


# @dataclass
# class SymbolList(dict):
#     variables: Set[Tuple[str, int]]
#     parameters: Set[str]
#     functions: Set[str]


# def list_symbols(expr: Expression, funs: List[str] = None) -> SymbolList:
#     if funs is None:
#         funs = []
#     ll = ListSymbols(known_functions=functions + funs)
#     ll.visit(expr)
#     if ll.problems:
#         e = Exception('Symbolic error.')
#         e.problems = ll.problems
#         print(e.problems)
#         raise e

#     def head(v):
#         return [i[0] for i in v]

#     return SymbolList(
#         dedup(head(ll.variables)), dedup(head(ll.constants)),
#         dedup(head(ll.functions)))


# ### Tree manipulation functions


# def eval_scalar(tree: Expression) -> Union[int, float]:
#     try:
#         if isinstance(tree, ast.Num):
#             return tree.n
#         elif isinstance(tree, ast.UnaryOp):
#             if isinstance(tree.op, ast.USub):
#                 return -tree.operand.n
#             elif isinstance(tree.op, ast.UAdd):
#                 return tree.operand.n
#         else:
#             raise Exception("Don't know how to do that.")
#     except:
#         raise Exception("Don't know how to do that.")


# class ListSymbols(ast.NodeVisitor):
#     def __init__(self, known_functions=[], known_variables=[]):
#         self.known_functions = known_functions
#         self.known_variables = known_variables
#         self.functions = []
#         self.variables = []
#         self.constants = []
#         self.problems = []

#     def visit_Call(self, call):
#         name = call.func.id
#         colno = call.func.col_offset
#         if name in self.known_functions:
#             self.functions.append((name, colno))
#             [self.visit(e) for e in call.args]
#         else:
#             try:
#                 assert (len(call.args) == 1)
#                 n = int(eval_scalar(call.args[0]))
#                 self.variables.append(((name, n), colno))
#             except Exception as e:
#                 self.problems.append([name, 0, colno, 'incorrect subscript'])
#                 # [self.visit(e) for e in call.args]

#     def visit_Name(self, name):
#         # colno = name.colno
#         colno = name.col_offset
#         n = 0
#         name = name.id
#         if name in self.known_functions:
#             self.problems.append([name, colno, 'function_not_called'])
#         else:
#             self.constants.append((name, colno))


# class ExpressionStringifier(NodeTransformer):

#     # replaces calls to variables by time subscripts

#     def __init__(self, variables=None, functions=None, constants=None):

#         self.variables = variables if variables is not None else []
#         if functions is None:
#             self.functions = [e for e in functions_dict.keys()]
#         else:
#             self.functions = functions
#         if constants is None:
#             self.constants = ['pi', 'e', 'inf']
#         else:
#             self.constants = constants
#         # self.variables = tvariables # ???

#     def visit_Name(self, node):

#         name = node.id
#         # if name self.functions:
#         #     return node
#         if name in self.constants:
#             return node
#         elif name in self.variables:
#             return Name(id=stringify_variable((name, 0)), ctx=Load())
#         else:
#             return Name(id=stringify_parameter(name), ctx=Load())

#     def visit_Call(self, node):

#         name = node.func.id
#         args = node.args[0]

#         if name in self.variables or name not in self.functions:
#             try:
#                 date = eval_scalar(args)
#             except:
#                 raise Exception("Unrecognized subscript.")
#             newname = stringify_variable((name, date))
#             if newname is not None:
#                 return Name(newname, Load())
#         else:
#             return Call(
#                 func=node.func,
#                 args=[self.visit(e) for e in node.args],
#                 keywords=[])


# class TimeShiftTransformer(ast.NodeTransformer):
#     def __init__(self, shift=0, variables=None):

#         self.variables = variables if variables is not None else []
#         self.shift = shift

#     def visit_Call(self, node):

#         name = node.func.id
#         args = node.args[0]

#         if name in self.variables:
#             if isinstance(args, UnaryOp):
#                 # we have s(+1)
#                 if (isinstance(args.op, UAdd)):
#                     args = args.operand
#                     date = args.n
#                 elif (isinstance(args.op, USub)):
#                     args = args.operand
#                     date = -args.n
#                 else:
#                     raise Exception("Unrecognized subscript.")
#             else:
#                 date = args.n
#             if self.shift == 'S':
#                 new_date = 0
#             else:
#                 new_date = date + self.shift
#             return ast.parse('{}({})'.format(name, new_date)).body[0].value
#         else:

#             # , keywords=node.keywords,  kwargs=node.kwargs)
#             return Call(
#                 func=node.func,
#                 args=[self.visit(e) for e in node.args],
#                 keywords=[])


# class ExpressionSanitizer(NodeTransformer):

#     # replaces calls to variables by time subscripts
#     def __init__(self, variables=None):
#         self.variables = variables if variables is not None else []

#     def visit_Name(self, node):
#         name = node.id
#         if name in self.variables:
#             return ast.parse('{}(0)'.format(name)).body[0].value
#         else:
#             return node

#     def visit_Call(self, node):
#         name = node.func.id
#         if name in self.variables:
#             t = eval_scalar(node.args[0])
#             return ast.parse('{}({})'.format(name, t)).body[0].value
#         else:
#             return Call(
#                 func=node.func,
#                 args=[self.visit(e) for e in node.args],
#                 keywords=[])


# class NameSubstituter(ast.NodeTransformer):

#     # substitutes a name by an expression
#     def __init__(self, substitutions: Dict[str, Expr]):
#         self.substitutions = substitutions

#     def visit_Name(self, node):
#         name = node.id
#         if name in self.substitutions:
#             return self.substitutions[name]
#         else:
#             return node
