from type_def import *

types = {
    'int32': PrimitiveType('int32'),
    'int64': PrimitiveType('int64'),
    'float': PrimitiveType('float'),
    'bool': PrimitiveType('bool')
}

i32 = types['int32']
i64 = types['int64']
f32 = types['float']
b = types['bool']
I = TypeVariable('I', [i32, i64, f32])
A = TypeVariable('A', [])

def impl_math(ty):
    ty.methods['__add__'] = ([SelfType(), ty], ty, set())
    ty.methods['__sub__'] = ([SelfType(), ty], ty, set())
    ty.methods['__mul__'] = ([SelfType(), ty], ty, set())
    ty.methods['__neg__'] = ([SelfType()], ty, set())
    ty.methods['__truediv__'] = ([SelfType(), ty], f32, set())
    ty.methods['__floordiv__'] = ([SelfType(), ty], ty, set())
    ty.methods['__mod__'] = ([SelfType(), ty], ty, set())
    ty.methods['__pow__'] = ([SelfType(), ty], ty, set())

def impl_bits(ty):
    ty.methods['__lshift__'] = ([SelfType(), i32], ty, set())
    ty.methods['__rshift__'] = ([SelfType(), i32], ty, set())
    ty.methods['__xor__'] = ([SelfType(), ty], ty, set())

def impl_logic(ty):
    ty.methods['__and__'] = ([SelfType(), ty], ty, set())
    ty.methods['__or__'] = ([SelfType(), ty], ty, set())

def impl_eq(ty):
    ty.methods['__eq__'] = ([SelfType(), ty], b, set())
    ty.methods['__ne__'] = ([SelfType(), ty], b, set())

def impl_order(ty):
    ty.methods['__lt__'] = ([SelfType(), ty], b, set())
    ty.methods['__gt__'] = ([SelfType(), ty], b, set())
    ty.methods['__le__'] = ([SelfType(), ty], b, set())
    ty.methods['__ge__'] = ([SelfType(), ty], b, set())

impl_logic(b)
impl_eq(b)

impl_math(i32)
impl_bits(i32)
impl_logic(i32)
impl_eq(i32)
impl_order(i32)

impl_math(i64)
impl_bits(i64)
impl_logic(i64)
impl_eq(i64)
impl_order(i64)

impl_math(f32)
impl_eq(f32)
impl_order(f32)

i32.methods['__init__'] = ([SelfType(), I], None, {'I'})
i64.methods['__init__'] = ([SelfType(), I], None, {'I'})
f32.methods['__init__'] = ([SelfType(), I], None, {'I'})

ParametricType.method_table['list'] = {
    'get': ([SelfType(), i32], TypeVariable('T1', []), set()),
    'head': ([SelfType()], TypeVariable('T1', []), set())
}

simplest_ctx = Context({}, types)
simplest_ctx.functions['len'] = ([ListType(A)], i32, {'A'})
simplest_ctx.functions['range'] = ([i32], ListType(i32), set())