# Copyright (c) OpenMMLab. All rights reserved.
# Modified from:
# https://github.com/pytorch/pytorch/blob/9ade03959392e5a90b74261012de1d806cab2253/torch/onnx/symbolic_opset9.py

from torch.onnx.symbolic_helper import parse_args

from juxtapose.mmdeploy.core import SYMBOLIC_REWRITER
from juxtapose.mmdeploy.utils import Backend


@SYMBOLIC_REWRITER.register_symbolic(
    "layer_norm", is_pytorch=True, arg_descriptors=["v", "is", "v", "v", "f", "i"]
)
def layer_norm__default(g, input, normalized_shape, weight, bias, eps, cudnn_enable):
    """Symbolic function for `layer_norm`

    Layer norm with torch<=1.12 might lead to wrong output shapes. Add
    keepdims=1 to each ReduceMean node to correct the shape.
    """
    import torch.onnx.symbolic_helper as sym_help
    from torch.onnx.symbolic_opset9 import add, mul, pow, sqrt, sub

    axes = [-i for i in range(len(normalized_shape), 0, -1)]

    two_cst = sym_help._generate_wrapped_number(g, 2.0)
    eps_cst = sym_help._generate_wrapped_number(g, eps)

    mean = g.op("ReduceMean", input, axes_i=axes, keepdims_i=1)
    numerator = sub(g, input, mean)
    # variance = e((x - e(x))^2), and (x - e(x)) is the numerator in the
    # layer_norm formula
    variance = g.op("ReduceMean", pow(g, numerator, two_cst), axes_i=axes, keepdims_i=1)
    denominator = sqrt(g, add(g, variance, eps_cst))

    layer_norm = g.op("Div", numerator, denominator)

    if not (weight is None or sym_help._is_none(weight)):
        layer_norm = mul(g, layer_norm, weight)
    if not (bias is None or sym_help._is_none(bias)):
        layer_norm = add(g, layer_norm, bias)

    return layer_norm


@parse_args("v", "is", "v", "v", "f", "i")
def _layer_norm_ncnn(g, input, normalized_shape, weight, bias, eps, cudnn_enable):
    """Symbolic function for `layer_norm`.

    PyTorch does not support export layer_norm to ONNX by default. We add the
    support here. `layer_norm` will be exported as ONNX node
    'mmdeploy::layer_norm'
    """
    weight.setDebugName("layernorm_weight")
    bias.setDebugName("layernorm_bias")
    return g.op("mmdeploy::LayerNorm", input, weight, bias, affine_i=1, epsilon_f=eps)


@SYMBOLIC_REWRITER.register_symbolic(
    "layer_norm", is_pytorch=True, backend=Backend.NCNN.value
)
def layer_norm__ncnn(*args):
    """Register default symbolic function for `layer_norm`.

    Add support to layer_norm to ONNX.
    """
    return _layer_norm_ncnn(*args)
