from bio import *
import pickle

@tuple
class MyType:
    a: i32
    b: str
    c: float

class A:
    a: int
    v: list[str]

    def __eq__(self: A, other: A):
        return self.a == other.a and self.v == other.v

    def __ne__(self: A, other: A):
        return not (self == other)

    def __hash__(self: A):
        return self.a

    def __copy__(self: A):
        return A(self.a, copy(self.v))

@test
def test_pickle[T](x: T):
    import gzip
    path = 'build/testjar.bin'
    jar = gzip.open(path, 'wb')
    pickle.dump(x, jar)
    jar.close()

    jar = gzopen(path, 'rb')
    y = pickle.load(jar, T)
    jar.close()

    assert x == y

@test
def test_non_atomic_list_pickle[T](x: list[list[T]]):
    import gzip
    copy = [copy(a) for a in x]
    path = 'build/testjar.bin'
    jar = gzip.open(path, 'wb')
    pickle.dump(x, jar)
    jar.close()

    for v in x:
        v.clear()

    jar = gzopen(path, 'rb')
    y = pickle.load(jar, list[list[T]])
    jar.close()

    assert y == copy

@test
def test_non_atomic_dict_pickle[T](x: dict[str, list[T]]):
    import gzip
    copy = {k: copy(v) for k,v in x.items()}
    path = 'build/testjar.bin'
    jar = gzip.open(path, 'wb')
    pickle.dump(x, jar)
    jar.close()

    for v in x.values():
        v.clear()

    jar = gzopen(path, 'rb')
    y = pickle.load(jar, dict[str, list[T]])
    jar.close()

    assert y == copy

@test
def test_non_atomic_set_pickle(x: set[A]):
    import gzip
    copy = {copy(a) for a in x}
    path = 'build/testjar.bin'
    jar = gzip.open(path, 'wb')
    pickle.dump(x, jar)
    jar.close()

    for a in x:
        a.v.clear()

    jar = gzopen(path, 'rb')
    y = pickle.load(jar, set[A])
    jar.close()

    assert y == copy

K = Kmer[8]
test_pickle(42)
test_pickle(3.14)
test_pickle(True)
test_pickle(byte(90))
test_pickle(UInt[123](123123123))
test_pickle(Int[123](-123123123))
test_pickle([11, 22, 33, 44])
test_pickle({11, 22, 33, 44})
test_pickle({11: 1.1, 22: 2.2, 33: 3.3, 44: 4.4})
test_pickle('hello world')
test_pickle('')
test_pickle(s'ACGTAAGG')
test_pickle(~s'ACGTAAGG')
test_pickle(K(s'ACGTAAGG'))
test_pickle(MyType(i32(-1001), 'xyz', 5.55))
test_pickle((A(1, ['x', 'abc', '1.1.1.1']), 42, A(1000, ['foo'])))
test_pickle([K(s'ACGTAAGG'), K(s'TATCTGTT')])
test_pickle(list[K]())
test_pickle({K(s'ACGTAAGG'), K(s'CATTTTTA')})
test_pickle({~s'ACGTAAGG'})
test_pickle({K(s'ACGTAAGG'), K(s'TTTTGGTT')})
test_pickle(set[K]())
test_pickle({K(s'ACGTAAGG'): 99, K(s'TTATTCTT'): 42})
test_pickle(dict[K,K]())
test_pickle({~s'ACGTAAGG': ~s'ACGTAAGG'})
test_pickle((42, 3.14, True, byte(90), s'ACGTAAGG', K(s'ACGTAAGG')))
test_pickle({i32(42): [[{s'ACG', s'ACGTAGCG', ~s'ACGTAGCG'}, {s'ACG', s'ACGTAGCG', ~s'ACGTAGCG'}], list[set[seq]](), [set[seq]()], [{~s''}, {s'', s'GCGC'}]]})

test_non_atomic_list_pickle([[3,2,1], [-1,-2,-3], [111,999,888,777], list[int]()])
test_non_atomic_dict_pickle({'first': [3,2,1], 'second': [-1,-2,-3], 'third': [111,999,888,777], 'fourth:': list[int]()})
test_non_atomic_set_pickle({A(42, ['fourty', 'two']), A(0, list[str]()), A(-99, ['negative', 'ninety', 'nine'])})