PPL TYPE 5

class Type(ABC): pass
class IntType(Type): pass
class FloatType(Type): pass
class BoolType(Type): pass
class NoType(Type): pass
class FunctionType(Type): pass
class Symbol:
    def __init__(self, name, typ):
        self.name = name
        if type(typ) is list:
            self.typ = typ
            return
        self.typ = [typ]
class Utils:
    def infer(env, name, typ):
        for symbol in env:
                if symbol.name == name:
                    symbol.typ = [typ]
                    return typ
class GetType:
    def checkType(s):
        if s == "IntType":
            return IntType()
        if s == "FloatType":
            return FloatType()
        if s == "BoolType":
            return BoolType()
        return NoType()

class UtilsInMoreEnvironments:
    def infer(env, name, typ):
        for symbol_list in env:
            for symbol in symbol_list:
                if symbol.name == name:
                    symbol.typ = [typ]
                    return typ
class GetEnv(Visitor):
    def visitProgram(self,ctx:Program,o):
        o = []
        for decl in ctx.decl:
            o = self.visit(decl,o)
        return o
    def visitVarDecl(self,ctx:VarDecl,o):
        for x in o:
            if ctx.name == x.name:
                raise Redeclared(ctx)
        o += [Symbol(ctx.name, NoType())]
        return o

    def visitFuncDecl(self,ctx:FuncDecl,o):
        for x in o:
            if ctx.name == x.name:
                raise Redeclared(ctx)
        env = []
        for decl in ctx.param:
            env = self.visit(decl,env)
        for decl in ctx.local:
            env = self.visit(decl,env)
        o += [Symbol(ctx.name,[FunctionType(),env])]
        return o
class StaticCheck(Visitor):
    def visitProgram(self,ctx:Program,o):
        env = GetEnv().visit(ctx,o)
        for decl in ctx.decl:
            if type(decl) is FuncDecl:
                self.visit(decl, env)
        for stmt in ctx.stmts:
            self.visit(stmt, env)
# class FuncDecl(Decl): #name:str,param:List[VarDecl],local:List[Decl],stmts:List[Stmt]
    def visitFuncDecl(self,ctx:FuncDecl,o):
        for x in o:
            if ctx.name == x.name:
                if type(x.typ[0]) is not FunctionType: return
                for stmt in ctx.stmts:
                    if type(stmt) is CallStmt:
                        self.visit(stmt,o)
                    if type(stmt) is Assign:
                        self.visit(stmt,x.typ[1])
                return
        return
    def visitAssign(self,ctx:Assign,o):
        rtype = self.visit(ctx.rhs, o)
        ltype = self.visit(ctx.lhs, o)

        if type(rtype) is NoType and type(ltype) is NoType:
            raise TypeCannotBeInferred(ctx)

        if type(ltype) is NoType:
            Utils.infer(o, ctx.lhs.name, rtype)
            return

        if type(rtype) is NoType:
            Utils.infer(o, ctx.rhs.name, ltype)
            return

        if type(ltype) is type(rtype): return
        raise TypeMismatchInStatement(ctx)
# class CallStmt(Stmt): #name:str,args:List[Exp]
    def visitCallStmt(self,ctx:CallStmt,o):
        #args la id: No o dau ? Global or Function Scope.
        #args ko co type. Chi co Literal -> Type hoac Id -> Type.
        args = ctx.args
        for x in o:
            if ctx.name == x.name:
                if type(x.typ[0]) is not FunctionType:
                    raise UndeclaredIdentifier (ctx.name)
                else:
                    if len(args) != len(x.typ[1]):
                        raise TypeMismatchInStatement (ctx)
                    for i in range(0,len(args)):
                    # x.typ[1] is Function Enviroment
                    # o is global environment, which contains Function
                    # up
                        env = o
                        type_of_args = type(self.visit(args[i],env))
                        type_of_param = type(x.typ[1][i].typ[0])
                        # raise UndeclaredIdentifier (str(type_of_args))
                        if type_of_args is NoType and type_of_param is NoType:
                            raise TypeCannotBeInferred (ctx)

                        if type_of_args is NoType and type_of_param is not NoType:

                            UtilsInMoreEnvironments.infer([x.typ[1],o],args[i].name,type_of_param)

                        if type_of_args is not NoType and type_of_param is NoType:

                            Utils.infer(x.typ[1],x.typ[1][i].name,type_of_args)

                        if type_of_args != type_of_param:

                            raise TypeMismatchInStatement (ctx)
                    return
        raise UndeclaredIdentifier(ctx.name)
# Case 1: FuncVar is NoType, while Call variable Type -> Suy dien for CallVariable

# Case 2: FuncVar have Type, call var have type -> Check if they are same type
# Case 3: Both is no type -> TypeCannotBeInferred

    def visitBinOp(self,ctx:BinOp,o):
        e1t = self.visit(ctx.e1, o)
        e2t = self.visit(ctx.e2, o)

        if ctx.op in ['+', '-', '*', '/']:
            if type(e1t) is NoType:
                e1t = Utils.infer(o, ctx.e1.name, IntType())
            if type(e2t) is NoType:
                e2t = Utils.infer(o, ctx.e2.name, IntType())
            if type(e1t) is IntType and type(e2t) is IntType:
                return IntType()
            raise TypeMismatchInExpression(ctx)

        if ctx.op in ['+.', '-.', '*.', '/.']:
            if type(e1t) is NoType:
                e1t = Utils.infer(o, ctx.e1.name, FloatType())
            if type(e2t) is NoType:
                e2t = Utils.infer(o, ctx.e2.name, FloatType())
            if type(e2t) is NoType:
                e1t = Utils.infer(o, ctx.e1.name, FloatType())
                e2t = Utils.infer(o, ctx.e2.name, FloatType())
            if type(e1t) is FloatType and type(e2t) is FloatType:
                return FloatType()
            raise TypeMismatchInExpression(ctx)
        if ctx.op in ['>', '=']:
            if type(e1t) is NoType:
                e1t = Utils.infer(o, ctx.e1.name, IntType())
            if type(e2t) is NoType and type(e1t) is IntType:
                e2t = Utils.infer(o, ctx.e2.name, IntType())
            if type(e1t) is IntType and type(e2t) is IntType:
                return BoolType()
        if ctx.op in ['>.', '=.']:
            if type(e1t) is NoType:
                e1t = Utils.infer(o, ctx.e1.name, FloatType())
            if type(e2t) is NoType:
                e2t = Utils.infer(o, ctx.e2.name, FloatType())
            if type(e1t) is FloatType and type(e2t) is FloatType:
                return BoolType()
            raise TypeMismatchInExpression(ctx)
        if ctx.op in ['&&', '||', '>b' , '=b']:
            if type(e1t) is NoType:
                e1t = Utils.infer(o, ctx.e1.name, BoolType())
            if type(e2t) is NoType:
                e2t = Utils.infer(o, ctx.e2.name, BoolType())
            if type(e1t) is BoolType and type(e2t) is BoolType:
                return BoolType()
            raise TypeMismatchInExpression(ctx)

    def visitUnOp(self,ctx:UnOp,o):
        e1t = self.visit(ctx.e, o)
        if ctx.op in ['-']:
            if type(e1t) is NoType:
                e1t = Utils.infer(o, ctx.e.name, IntType())
            if type(e1t) is IntType:
                return IntType()
            raise TypeMismatchInExpression(ctx)

        if ctx.op in ['-.']:
            if type(e1t) is NoType:
                e1t = Utils.infer(o, ctx.e.name, FloatType())
            if type(e1t) is FloatType:
                return FloatType()
            raise TypeMismatchInExpression(ctx)

        if ctx.op in ['i2f']:
            if type(e1t) is NoType:
                e1t = Utils.infer(o, ctx.e.name, IntType())
            if type(e1t) == IntType:
                return FloatType()
            raise TypeMismatchInExpression(ctx)

        if ctx.op in ['!']:
            if type(e1t) is NoType:
                e1t = Utils.infer(o, ctx.e.name, BoolType())
            if type(e1t) == BoolType:
                return BoolType()
            raise TypeMismatchInExpression(ctx)

        if ctx.op in ['floor']:
            if type(e1t) is NoType:
                e1t = Utils.infer(o, ctx.e.name, FloatType())
            if type(e1t) is FloatType:
                return IntType()
            raise TypeMismatchInExpression(ctx)

    def visitIntLit(self,ctx:IntLit,o): return IntType()

    def visitFloatLit(self,ctx,o): return FloatType()

    def visitBoolLit(self,ctx,o): return BoolType()

    def visitId(self,ctx,o):
        for symbol in o:
                if ctx.name == symbol.name:
                    return symbol.typ[0]
        raise UndeclaredIdentifier(ctx.name)