from LanguagePrinter import LanguagePrinter

class PythonPrinter(LanguagePrinter):
    def __init__(self, outfile):
        super(PythonPrinter, self).__init__(outfile)
        self.expr = PythonExpressionPrinter(self)
        self.stmt = PythonStatementPrinter(self)
        self.comment('This Python file was autogenerated from Coq')
        self.writeln('from enum import Enum')
        self.writeln('import ZBitOps')
        self.writeln('from Utility import *')
        self.end_decl()

    def comment(self, s):
        self.writeln('# ' + s)

    def end_decl(self):
        self.writeln('')

    def type_alias(self, name, rhsName):
        pass # there are no types aliases in Python

    def enum(self, name, valueNames):
        self.writeln('class {}(Enum):'.format(name))
        for i, n in enumerate(valueNames, 1):
            self.writeln('    ' + n + ' = ' + str(i))
        self.end_decl()

    def variant(self, name, branches):
        '''
        name: str
        branches: list of (branchName, typesList) tuples
        '''
        self.writeln('class {}(object): pass'.format(name))
        self.end_decl()

        for branchName, argTypes in branches:
            self.writeln('class {}({}):'.format(branchName, name))
            self.increaseIndent()
            constructorArgs = ''.join([', f' + str(i) for i in range(len(argTypes))])
            self.writeln('def __init__(self{}):'.format(constructorArgs))
            self.increaseIndent()
            for i in range(len(argTypes)):
                self.writeln('self.f{} = f{}'.format(i, i))
            if len(argTypes) == 0:
                self.writeln('pass')
            self.decreaseIndent()
            self.decreaseIndent()
            self.end_decl()

    def constant_decl(self, name, typ, rhs):
        self.write(name + ' = ')
        self.writeln(rhs())
        self.end_decl()

    def fun_decl(self, name, argnamesWithTypes, returnType, body):
        self.writeln('def {}({}):'.format(name,
                ', '.join([argname for argname, tp in argnamesWithTypes])))
        self.increaseIndent()
        self.startln()
        self.write(body())
        self.decreaseIndent()
        self.end_decl()


class PythonExpressionPrinter:

    def __init__(self, context):
        self.context = context

    # private helper functions:

    def __binop(self, arg1, op, arg2):
        return "{} {} {}".format(arg1(), op, arg2())

    def __raw_function_call(self, func, args):
        return '{}({})'.format(func, ', '.join(args))


    # public functions:

    def alu_op_name(self, name):
        return 'Alu.' + name

    def function_call(self, func, args):
        return self.__raw_function_call(func(), [arg() for arg in args])

    def bit_literal(self, s):
        return '0b' + s

    def true_literal(self):
        return 'True'

    def false_literal(self):
        return 'False'

    def negate_bool(self, e):
        return 'not ' + e()

    def access_type_load(self):
        return 'AccessType.Load'

    def access_type_store(self):
        return 'AccessType.Store'

    def var(self, varName):
        return varName

    def if_expr(self, cond, ifyes, ifno):
        res = '(' + ifyes() + "\n"
        self.context.increaseIndent()
        res += self.context.indent + "if "
        res += cond() + "\n"
        res += self.context.indent + "else "
        res += ifno() + ')'
        self.context.decreaseIndent()
        return res

    def list(self, elems):
        return '[' + ', '.join([elem() for elem in elems]) + ']'

    def list_length(self, arg):
        return self.__raw_function_call('len', [arg()])

    def list_nth_default(self, index, l, default):
        return self.__raw_function_call("list_nth_default", [index(), l(), default()])

    def concat(self, first_arg, second_arg):
        return self.__binop(first_arg, '+', second_arg)

    def equality(self, first_arg, second_arg):
        return self.__binop(first_arg, '==', second_arg)

    def gt(self, first_arg, second_arg):
        return self.__binop(first_arg, '>', second_arg)

    def bigint_mul(self, first_arg, second_arg):
        return self.__binop(first_arg, '*', second_arg); # TODO

    def logical_or(self, first_arg, second_arg):
        return self.__binop(first_arg, '|', second_arg)

    def shift_left(self, first_arg, second_arg):
        return self.__binop(first_arg, '<<', second_arg)

    def boolean_and(self, first_arg, second_arg):
        return self.__binop(first_arg, 'and', second_arg)

    def boolean_or(self, first_arg, second_arg):
        return self.__binop(first_arg, 'or', second_arg)

    def silent_id(self, arg):
        return arg()


class PythonStatementPrinter:

    def __init__(self, context):
        self.context = context

    def if_stmt(self, cond, ifyes, ifno):
        res = 'if '
        res += cond()
        res += ':\n'
        self.context.increaseIndent()
        res += self.context.indent
        res += ifyes()
        res += '\n'
        self.context.decreaseIndent()
        res += self.context.indent
        res += "else:\n"
        self.context.increaseIndent()
        res += self.context.indent
        res += ifno()
        self.context.decreaseIndent()
        res += "\n" + self.context.indent
        return res

    def let_in(self, name, typ, rhs, body):
        res = name + ' = '
        res += rhs()
        res += '\n' + self.context.indent
        res += body()
        return res

    def nop(self):
        return 'pass'

    # match over Inductive (where constructors can take args)
    def match(self, discriminee, branches, default_branch):
        res = ''
        for constructorName, (argNames, branchBody) in branches.items():
            res += 'if isinstance({}, {}):\n'.format(discriminee, constructorName)
            self.context.increaseIndent()
            res += self.context.indent
            res += branchBody() # TODO make sure branch body can access fields
            self.context.decreaseIndent()
            res += '\n' + self.context.indent
        if default_branch:
            res += default_branch()
            res += '\n'
        return res

    # match over an enum (where constructors cannot take args)
    def switch(self, discriminee, enumName, branches, default_branch):
        res = ''
        for constructorName, branchBody in branches.items():
            res += 'if {} == {}.{}:\n'.format(discriminee, enumName, constructorName)
            self.context.increaseIndent()
            res += self.context.indent
            res += branchBody()
            self.context.decreaseIndent()
            res += '\n' + self.context.indent
        if default_branch:
            res += default_branch()
            res += '\n'
        return res

    def return_value(self, e):
        return "return " + e
