from typing import Dict, List, Tuple
import numpy as np
import codon

@codon.convert
class Foo:
    __slots__ = 'a', 'b', 'c'

    def __init__(self, n):
        self.a = n
        self.b = n**2
        self.c = n**3

    def __eq__(self, other):
        return (self.a == other.a and
                self.b == other.b and
                self.c == other.c)

    def __hash__(self):
        return hash((self.a, self.b, self.c))

    @codon.jit
    def total(self):
        return self.a + self.b + self.c

def test_convertible():
    assert Foo(10).total() == 1110

def test_many():
    @codon.jit
    def is_prime(n):
        if n <= 1:
            return False
        for i in range(2, n):
            if n % i == 0:
                return False
        return True

    assert sum(1 for i in range(100000, 200000) if is_prime(i)) == 8392

def test_roundtrip():
    @codon.jit
    def roundtrip(x):
        return x

    for _ in range(5):
        assert roundtrip(None) == None
        assert roundtrip(42) == 42
        assert roundtrip(3.14) == 3.14
        assert roundtrip(False) == False
        assert roundtrip(True) == True
        assert roundtrip('hello') == 'hello'
        assert roundtrip('') == ''
        assert roundtrip(2+3j) == 2+3j
        assert roundtrip(slice(1,2,3)) == slice(1,2,3)
        assert roundtrip([11,22,33]) == [11,22,33]
        assert roundtrip([[[42]]]) == [[[42]]]
        assert roundtrip({11,22,33}) == {11,22,33}
        assert roundtrip({11: 'one', 22: 'two', 33: 'three'}) == {11: 'one', 22: 'two', 33: 'three'}
        assert roundtrip((11,22,33)) == (11,22,33)
        assert Foo(roundtrip(Foo(123))[0]) == Foo(123)
        assert roundtrip(roundtrip) is roundtrip

def test_return_type():
    @codon.jit
    def run() -> Tuple[int, str, float, List[int], Dict[str, int]]:
        return (1, "str", 2.45, [1, 2, 3], {"a": 1, "b": 2})

    r = run()
    assert type(r) == tuple
    assert type(r[0]) == int
    assert type(r[1]) == str
    assert type(r[2]) == float
    assert type(r[3]) == list
    assert len(r[3]) == 3
    assert type(r[3][0]) == int
    assert type(r[4]) == dict
    assert len(r[4].items()) == 2
    assert type(next(iter(r[4].keys()))) == str
    assert type(next(iter(r[4].values()))) == int

def test_param_types():
    @codon.jit
    def run(a: int, b: Tuple[int, int], c: List[int], d: Dict[str, int]) -> int:
        s = 0
        for v in [a, *b, *c, *d.values()]:
            s += v
        return s

    r = run(1, (2, 3), [4, 5, 6], dict(a=7, b=8, c=9))
    assert type(r) == int
    assert r == 45

def test_error_handling():
    @codon.jit
    def type_error():
        return 1 + '1'

    try:
        type_error()
    except codon.JITError:
        pass
    except:
        assert False
    else:
        assert False

test_convertible()
test_many()
test_roundtrip()
test_return_type()
test_param_types()
test_error_handling()


@codon.jit
def foo(y):
    return f"{y.__class__.__name__}; {y}"


@codon.jit(debug=True)
def foo2(y):
    return f"{y.__class__.__name__}; {y}"

class Foo:
    def __init__(self):
        self.x = 1

@codon.jit
def a(x):
    return x+1

def b(x, z):
    y = a(x)
    return y * z

@codon.jit(pyvars=['b'])
def c(x, y):
    n = b(x,y) ** a(1)
    return n

def test_cross_calls():
    assert foo([None, 1]) == "List[Optional[int]]; [None, 1]"
    assert foo([1, None, 1]) == "List[Optional[int]]; [1, None, 1]"
    assert foo([1, None, 1.2]) == "List[pyobj]; [1, None, 1.2]"
    assert foo({None: 1}) == "Dict[pyobj,int]; {None: 1}"
    assert foo2([None, Foo()]).startswith("List[pyobj]; [None, <__main__.Foo object at")

    assert a(3) == 4
    assert b(3, 4) == 16
    assert round(c(5, 6.1), 2) == 1339.56

test_cross_calls()

@codon.jit
def d(x, y):
    z = x**2 + y
    return z

def test_ndarray():
    x = np.array([1, 2, 3, 4])
    y = np.array([10., 100., 1000., 10000.])
    z = d(x, y)
    assert np.array_equal(z, [11., 104., 1009., 10016.])
    assert z.dtype == np.float64

    x = np.array([1j])
    y = np.array([2j], np.complex64)
    z = d(x, y)
    assert np.array_equal(z, [-1+2j])
    assert z.dtype == np.complex128

    x = np.array([3, 4, 5], 'timedelta64[2s]')
    y = a(x)
    assert np.array_equal(y, np.array([4, 5, 6], 'timedelta64[2s]'))
    assert np.datetime_data(y.dtype) == ('s', 2)

    x = np.array([3, 4, 5], 'datetime64[2s]')
    y = a(x)
    assert np.array_equal(y, np.array([4, 5, 6], 'datetime64[2s]'))
    assert np.datetime_data(y.dtype) == ('s', 2)

test_ndarray()

@codon.jit
def e(x=2, y=99):
    return 2*x + y

def test_arg_order():
    assert e(1, 2) == 4
    assert e(1) == 101
    assert e(y=10, x=1) == 12
    assert e(x=1) == 101
    assert e() == 103

test_arg_order()