# Copyright (C) 2022-2023 Exaloop Inc. <https://exaloop.io>
# dict implementation based on klib's khash

import internal.khash as khash
import internal.gc as gc

def _dict_hash(key) -> int:
    k = key.__hash__()
    return (k >> 33) ^ k ^ (k << 11)

class Dict:
    _n_buckets: int
    _size: int
    _n_occupied: int
    _upper_bound: int

    _flags: Ptr[u32]
    _keys: Ptr[K]
    _vals: Ptr[V]

    K: type
    V: type

    # Magic methods

    def _init(self):
        self._n_buckets = 0
        self._size = 0
        self._n_occupied = 0
        self._upper_bound = 0
        self._flags = Ptr[u32]()
        self._keys = Ptr[K]()
        self._vals = Ptr[V]()

    def _init_from(self, other):
        n = other._n_buckets

        if n == 0:
            self._init()
            return

        f = khash.__ac_fsize(n)
        self._n_buckets = n
        self._size = other._size
        self._n_occupied = other._n_occupied
        self._upper_bound = other._upper_bound

        flags_copy = Ptr[u32](f)
        keys_copy = Ptr[K](n)
        vals_copy = Ptr[V](n)
        str.memcpy(flags_copy.as_byte(), other._flags.as_byte(), f * gc.sizeof(u32))
        str.memcpy(keys_copy.as_byte(), other._keys.as_byte(), n * gc.sizeof(K))
        str.memcpy(vals_copy.as_byte(), other._vals.as_byte(), n * gc.sizeof(V))

        self._flags = flags_copy
        self._keys = keys_copy
        self._vals = vals_copy

    def __init__(self):
        self._init()

    def __init__(self, g: Generator[Tuple[K, V]]):
        self._init()
        for k, v in g:
            self[k] = v

    def __init__(self, other: Dict[K, V]):
        self._init_from(other)

    def __getitem__(self, key: K) -> V:
        x = self._kh_get(key)
        if x != self._kh_end():
            return self._vals[x]
        raise KeyError(str(key))

    def __setitem__(self, key: K, val: V):
        ret, x = self._kh_put(key)
        self._vals[x] = val

    def __delitem__(self, key: K):
        x = self._kh_get(key)
        if x != self._kh_end():
            self._kh_del(x)
        else:
            raise KeyError(str(key))

    def __contains__(self, key: K) -> bool:
        return self._kh_get(key) != self._kh_end()

    def __eq__(self, other: Dict[K, V]) -> bool:
        if self.__len__() != other.__len__():
            return False
        for k, v in self.items():
            if k not in other or other[k] != v:
                return False
        return True

    def __ne__(self, other: Dict[K, V]) -> bool:
        return not (self == other)

    def __iter__(self) -> Generator[K]:
        return self.keys()

    def __len__(self) -> int:
        return self._size

    def __or__(self, other):
        new = self.__copy__()
        new.update(other)
        return new

    def __ior__(self, other):
        self.update(other)
        return self

    def __copy__(self):
        if self.__len__() == 0:
            return Dict[K, V]()
        n = self._n_buckets
        f = khash.__ac_fsize(n)
        flags_copy = Ptr[u32](f)
        keys_copy = Ptr[K](n)
        vals_copy = Ptr[V](n)
        str.memcpy(flags_copy.as_byte(), self._flags.as_byte(), f * gc.sizeof(u32))
        str.memcpy(keys_copy.as_byte(), self._keys.as_byte(), n * gc.sizeof(K))
        str.memcpy(vals_copy.as_byte(), self._vals.as_byte(), n * gc.sizeof(V))
        return Dict[K, V](
            n,
            self._size,
            self._n_occupied,
            self._upper_bound,
            flags_copy,
            keys_copy,
            vals_copy,
        )

    def __deepcopy__(self) -> Dict[K, V]:
        return {k.__deepcopy__(): v.__deepcopy__() for k, v in self.items()}

    def __repr__(self) -> str:
        n = self.__len__()
        if n == 0:
            return "{}"
        else:
            buf = _strbuf()
            buf.append("{")
            first = True
            for k, v in self.items():
                if not first:
                    buf.append(", ")
                else:
                    first = False
                buf.append(k.__repr__())
                buf.append(": ")
                buf.append(v.__repr__())
            buf.append("}")
            return buf.__str__()

    # Helper methods

    def resize(self, new_n_buckets: int):
        self._kh_resize(new_n_buckets)

    def get(self, key: K, s: V) -> V:
        x = self._kh_get(key)
        return self._vals[x] if x != self._kh_end() else s

    def setdefault(self, key: K, val: V) -> V:
        ret, x = self._kh_put(key)
        if ret != 0:  # i.e. key not present
            self._vals[x] = val
            return val
        return self._vals[x]

    def increment(self, key: K, by: T = 1, T: type):
        ret, x = self._kh_put(key)
        if ret != 0:  # i.e. key not present
            self._vals[x] = by
        else:
            self._vals[x] += by

    def __dict_do_op_throws__(self, key: K, other: Z, op: F, F: type, Z: type):
        x = self._kh_get(key)
        if x == self._kh_end():
            raise KeyError(str(key))
        else:
            self._vals[x] = op(self._vals[x], other)

    def __dict_do_op__(self, key: K, other: Z, dflt: V, op: F, F: type, Z: type):
        ret, x = self._kh_put(key)
        self._vals[x] = op(dflt if ret != 0 else self._vals[x], other)

    def update(self, other):
        if isinstance(other, Dict[K, V]):
            for k, v in other.items():
                self[k] = v
        else:
            for k, v in other:
                self[k] = v

    def pop(self, key: K) -> V:
        x = self._kh_get(key)
        if x != self._kh_end():
            v = self._vals[x]
            self._kh_del(x)
            return v
        raise KeyError(str(key))

    def popitem(self) -> Tuple[K, V]:
        for k in self:
            return (k, self.pop(k))
        raise KeyError("dictionary is empty")

    def clear(self):
        self._kh_clear()

    def items(self) -> Generator[Tuple[K, V]]:
        i = self._kh_begin()
        while i < self._kh_end():
            if self._kh_exist(i):
                yield self._keys[i], self._vals[i]
            i += 1

    def keys(self) -> Generator[K]:
        for k, v in self.items():
            yield k

    def values(self) -> Generator[V]:
        for k, v in self.items():
            yield v

    def copy(self):
        return self.__copy__()

    def fromkeys(ks: Generator[K], v: V, K: type, V: type) -> Dict[K, V]:
        return {k: v for k in ks}

    # Internal helpers

    def _kh_clear(self):
        if self._flags:
            i = 0
            n = khash.__ac_fsize(self._n_buckets)
            while i < n:
                self._flags[i] = u32(0xAAAAAAAA)
                i += 1
            self._size = 0
            self._n_occupied = 0

    def _kh_get(self, key: K) -> int:
        if self._n_buckets:
            step = 0
            mask = self._n_buckets - 1
            k = _dict_hash(key)
            i = k & mask
            last = i
            while not khash.__ac_isempty(self._flags, i) and (
                khash.__ac_isdel(self._flags, i) or self._keys[i] != key
            ):
                step += 1
                i = (i + step) & mask
                if i == last:
                    return self._n_buckets
            return self._n_buckets if khash.__ac_iseither(self._flags, i) else i
        else:
            return 0

    def _kh_resize(self, new_n_buckets: int):
        HASH_UPPER = 0.77
        new_flags = Ptr[u32]()
        j = 1

        # round up to next power of 2
        new_n_buckets -= 1
        new_n_buckets |= new_n_buckets >> 1
        new_n_buckets |= new_n_buckets >> 2
        new_n_buckets |= new_n_buckets >> 4
        new_n_buckets |= new_n_buckets >> 8
        new_n_buckets |= new_n_buckets >> 16
        new_n_buckets |= new_n_buckets >> 32
        new_n_buckets += 1

        if new_n_buckets < 4:
            new_n_buckets = 4

        if self._size >= int(new_n_buckets * HASH_UPPER + 0.5):
            j = 0
        else:
            fsize = khash.__ac_fsize(new_n_buckets)
            new_flags = Ptr[u32](fsize)
            i = 0
            while i < fsize:
                new_flags[i] = u32(0xAAAAAAAA)
                i += 1

            if self._n_buckets < new_n_buckets:
                self._keys = Ptr[K](
                    gc.realloc(self._keys.as_byte(),
                               new_n_buckets * gc.sizeof(K),
                               self._n_buckets * gc.sizeof(K))
                )
                self._vals = Ptr[V](
                    gc.realloc(self._vals.as_byte(),
                               new_n_buckets * gc.sizeof(V),
                               self._n_buckets * gc.sizeof(V))
                )

        if j:
            j = 0
            while j != self._n_buckets:
                if khash.__ac_iseither(self._flags, j) == 0:
                    key = self._keys[j]
                    val = self._vals[j]
                    new_mask = new_n_buckets - 1
                    khash.__ac_set_isdel_true(self._flags, j)

                    while True:
                        step = 0
                        k = _dict_hash(key)
                        i = k & new_mask

                        while not khash.__ac_isempty(new_flags, i):
                            step += 1
                            i = (i + step) & new_mask

                        khash.__ac_set_isempty_false(new_flags, i)
                        if (
                            i < self._n_buckets
                            and khash.__ac_iseither(self._flags, i) == 0
                        ):
                            self._keys[i], key = key, self._keys[i]
                            self._vals[i], val = val, self._vals[i]
                            khash.__ac_set_isdel_true(self._flags, i)
                        else:
                            self._keys[i] = key
                            self._vals[i] = val
                            break
                j += 1

            if self._n_buckets > new_n_buckets:
                self._keys = Ptr[K](
                    gc.realloc(self._keys.as_byte(),
                               new_n_buckets * gc.sizeof(K),
                               self._n_buckets * gc.sizeof(K))
                )
                self._vals = Ptr[V](
                    gc.realloc(self._vals.as_byte(),
                               new_n_buckets * gc.sizeof(V),
                               self._n_buckets * gc.sizeof(V))
                )

            self._flags = new_flags
            self._n_buckets = new_n_buckets
            self._n_occupied = self._size
            self._upper_bound = int(self._n_buckets * HASH_UPPER + 0.5)

    def _kh_put(self, key: K) -> Tuple[int, int]:
        if self._n_occupied >= self._upper_bound:
            if self._n_buckets > (self._size << 1):
                self._kh_resize(self._n_buckets - 1)
            else:
                self._kh_resize(self._n_buckets + 1)

        mask = self._n_buckets - 1
        step = 0
        site = self._n_buckets
        x = site
        k = _dict_hash(key)
        i = k & mask
        if khash.__ac_isempty(self._flags, i):
            x = i
        else:
            last = i
            while not khash.__ac_isempty(self._flags, i) and (
                khash.__ac_isdel(self._flags, i) or self._keys[i] != key
            ):
                if khash.__ac_isdel(self._flags, i):
                    site = i
                step += 1
                i = (i + step) & mask
                if i == last:
                    x = site
                    break

            if x == self._n_buckets:
                if khash.__ac_isempty(self._flags, i) and site != self._n_buckets:
                    x = site
                else:
                    x = i

        ret = 0
        if khash.__ac_isempty(self._flags, x):
            self._keys[x] = key
            khash.__ac_set_isboth_false(self._flags, x)
            self._size += 1
            self._n_occupied += 1
            ret = 1
        elif khash.__ac_isdel(self._flags, x):
            self._keys[x] = key
            khash.__ac_set_isboth_false(self._flags, x)
            self._size += 1
            ret = 2

        return (ret, x)

    def _kh_del(self, x: int):
        if x != self._n_buckets and not khash.__ac_iseither(self._flags, x):
            khash.__ac_set_isdel_true(self._flags, x)
            self._size -= 1

    def _kh_begin(self) -> int:
        return 0

    def _kh_end(self) -> int:
        return self._n_buckets

    def _kh_exist(self, x: int) -> bool:
        return not khash.__ac_iseither(self._flags, x)

dict = Dict
