import sys
from pathlib import Path

import numba as nb
import numpy as np
from numba.core.dispatcher import Dispatcher
from numba.core.registry import CPUDispatcher
from numba.pycc import CC


spec = nb.from_dtype(np.dtype([("y", np.int32), ("x", np.int32), ("phase", np.int32)]))


@nb.njit(nb.float64[:, :](spec[:], nb.float64))
def aot_func(srcs, t):
    m = np.zeros((200, 200), nb.float64)
    H, W = m.shape

    for y in range(H):
        for x in range(W):
            for src in srcs:
                dist = np.sqrt((y - src.y) ** 2 + (x - src.x) ** 2)
                m[y, x] += np.cos(2.0 * np.pi * (dist + t) + t * src.phase)
    return m


@nb.njit(nb.float64[:, :](spec[:], nb.float64), parallel=True)
def aot_func_parallel(srcs, t):
    m = np.zeros((200, 200), nb.float64)
    H, W = m.shape

    for y in nb.prange(H):
        for x in nb.prange(W):
            for src in srcs:
                dist = np.sqrt((y - src.y) ** 2 + (x - src.x) ** 2)
                m[y, x] += np.cos(2.0 * np.pi * (dist + t) + t * src.phase)
    return m


if __name__ == "__main__":

    for p in Path().glob("**/*.pyd"):
        p.unlink()

    mod_name = Path(__file__).stem
    cc = CC(mod_name)

    # Uncomment the following line to print out the compilation steps
    # cc.verbose = True

    mod = __import__(mod_name)
    mod = sys.modules[mod_name]
    for name, func in mod.__dict__.items():
        if isinstance(func, CPUDispatcher):
            print(f"compile: {name}")
            # signature = func.signatures[0]
            # signature = func.get_function_type().signature
            signature = func.nopython_signatures[0]
            deco = cc.export(func.__name__, signature)
            deco(func.py_func)

    cc.compile()