import numpy as np
from pycuda import gpuarray  # type:ignore
from pycuda.compiler import SourceModule  # type:ignore
from pycuda.elementwise import ElementwiseKernel  # type:ignore

from pydtnn.backends.pycuda.optimizers.optimizer import OptimizerPycuda
from pydtnn.optimizers.adam import Adam

from pydtnn.backends.pycuda.layers.layer import LayerPycuda
from pydtnn.backends.pycuda.utils.tensor_array import TensorArray
from pydtnn.utils.constants import DTYPE2CTYPE


class AdamPycuda(Adam[TensorArray], OptimizerPycuda):
    """
    AdamPycuda optimizer
    """

    def __init__(self, learning_rate=1e-2, beta1=0.99, beta2=0.999, epsilon=1e-7, decay=0.0):
        super().__init__(learning_rate, beta1, beta2, epsilon, decay)

    def _kernel_init(self) -> None:
        func_pow = {np.dtype(np.float32): "powf", np.dtype(np.float64): "pow"}

        # --- GPU ---
        parameters_gpu = "{T} *w, {T} *dw, {T} *m, {T} *v, float it, " \
                         "float lr, float decay, float beta1, float beta2, float epsilon".format(T=DTYPE2CTYPE[self.model.dtype])
        operations_gpu = """
            m[i] = beta1 * m[i] + (1 - beta1) * dw[i];
            v[i] = beta2 * v[i] + (1 - beta2) * {func}(dw[i], 2);
            w[i] -= lr * (decay * w[i] + ((m[i] / (1 - {func}(beta1, it))) / sqrt(v[i] / (1 - {func}(beta2, it)) + epsilon)));
        """.format(func=func_pow[self.model.dtype])

        self.update_kernel = ElementwiseKernel(parameters_gpu, operations_gpu, "Adam_kernel")
        # -----------

        # GPU DIRECT-
        _name = "Adam_kernel_gpudirect"
        code = """
            __global__ void {name}({T} *w, {T} *dw, {T} *m, {T} *v,
                                   float it, float lr, float decay,
                                   float beta1, float beta2, float epsilon, int N)
            {{
                int i = blockIdx.x * blockDim.x + threadIdx.x;
                if (i < N)
                {{
                    m[i] = beta1 * m[i] + (1 - beta1) * dw[i];
                    v[i] = beta2 * v[i] + (1 - beta2) * {func}(dw[i], 2);
                    w[i] -= lr * (decay * w[i] + ((m[i] / (1 - {func}(beta1, it))) /
                                              sqrt(v[i] / (1 - {func}(beta2, it)) + epsilon)));
                }}
            }}"""
        code = code.format(T=DTYPE2CTYPE[self.model.dtype],
                           func=func_pow[self.model.dtype],
                           name=_name)

        self.update_gpudirect = SourceModule(code).get_function(_name)
        # -----------

    def _model_init(self, list_layers: list[LayerPycuda]) -> None:
        super()._model_init(list_layers)  # type: ignore (The type is correct: LayerPycuda extends LayerBase)

        for layer in list_layers:
            self.context[layer.id] = dict[str, int | gpuarray.GPUArray]()
            self.context[layer.id]["it"] = 0

            for w_ in layer.grad_vars.keys():
                w = getattr(layer, w_)
                self.context[layer.id]["m_%s" % w_] = gpuarray.zeros_like(w.ary, dtype=layer.model.dtype)
                self.context[layer.id]["v_%s" % w_] = gpuarray.zeros_like(w.ary, dtype=layer.model.dtype)

                self.memory_used += self.context[layer.id]["m_%s" % w_].nbytes + self.context[layer.id]["v_%s" % w_].nbytes  # type: ignore (They are both "gpuarray" and not "int")

    def update(self, layer: LayerPycuda):
        self.context[layer.id]["it"] += 1  # type: ignore ("it" is an "int".)
        it: int = self.context[layer.id]["it"]  # type: ignore ("it" is an "int".)

        for w_, dw_ in layer.grad_vars.items():
            w, dw = getattr(layer, w_), getattr(layer, dw_)
            m = self.context[layer.id]["m_%s" % w_]
            v = self.context[layer.id]["v_%s" % w_]
            w: TensorArray
            dw: TensorArray
            m: gpuarray.GPUArray
            v: gpuarray.GPUArray

            if self.gpudirect:
                n = self.get_batch_size(w)
                threads, blocks = self.get_threads_and_blocks()

                self.update_gpudirect(w.ary.gpudata, dw.ptr_intp, m.gpudata, v.gpudata,
                                      np.float32(it), np.float32(self.learning_rate),
                                      np.float32(self.decay), np.float32(self.beta1),
                                      np.float32(self.beta2), np.float32(self.epsilon),
                                      np.int32(n),
                                      grid=(int(blocks), 1, 1), block=(int(threads), 1, 1),
                                      stream=layer.stream_2)
            else:
                self.update_kernel(w.ary, dw.ary, m, v,
                                   np.float32(it), np.float32(self.learning_rate),
                                   np.float32(self.decay), np.float32(self.beta1),
                                   np.float32(self.beta2), np.float32(self.epsilon),
                                   stream=layer.stream_2)
