import functools

import numpy as np
import pytest

import megengine
from megengine.autodiff.grad_manager import GradManager
from megengine.core.ops.builtin import GetVarShape, Reduce, TypeCvt
from megengine.core.tensor.utils import subgraph_fn
from megengine.device import CompNode, get_default_device
from megengine.jit import trace

_assert_allclose = functools.partial(np.testing.assert_allclose, atol=5e-6, rtol=5e-6)


@functools.lru_cache(maxsize=None)
def _get_batch_norm_fn(dtype, device, channels, ndim, interpret, gopt_level):
    @subgraph_fn(
        "BatchNormNd",
        dtype=dtype,
        device=device,
        nr_inputs=4,
        interpret=interpret,
        gopt_level=gopt_level,
    )
    def batch_norm_nd(inputs, f, c):
        input, eps, weight, bias = inputs[0:4]
        reduce_shape = c(
            (1, channels) + (1,) * (ndim - 2), dtype="int32", device=device
        )
        input_shape = f(GetVarShape(), input)
        input_elems = f(Reduce(mode="product", axis=0), input_shape)
        reduce_elems = f(Reduce(mode="product", axis=0), reduce_shape)
        reduce_size = f("//", input_elems, reduce_elems)
        reduce_size = f(TypeCvt(dtype=dtype), reduce_size)
        channel_x1s = f(Reduce(mode="sum"), input, reduce_shape)
        channel_x2s = f(Reduce(mode="sum_sqr"), input, reduce_shape)
        channel_mean = f("/", channel_x1s, reduce_size)
        channel_var = f(
            "-", f("/", channel_x2s, reduce_size), f("*", channel_mean, channel_mean),
        )
        invsqrt_channel_var = f("**", f("+", channel_var, eps), c(-0.5))
        inv_var_wt = f("*", invsqrt_channel_var, weight)
        neg_channel_mean = f("-", channel_mean)
        outvar = f(
            "fma3", input, inv_var_wt, f("fma3", neg_channel_mean, inv_var_wt, bias),
        )
        return (outvar,), (True,)

    return batch_norm_nd


@pytest.mark.parametrize("device", [
get_default_device()