# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from akg.utils import kernel_exec as utils
import numpy as np
from tests.common.tensorio import compare_tensor
from tests.common.gen_random import random_gaussian
from tests.common.base import get_rtol_atol
from akg.ms.cce import prod_force_se_a

import random


def gen_data(input_shapes):
    input1 = random_gaussian(input_shapes[0], miu=1, sigma=0.1)
    input2 = random_gaussian(input_shapes[1], miu=1, sigma=0.1)
    input3 = random_gaussian(input_shapes[2], miu=1, sigma=0.1) * 96

    input1 = input1.astype(np.float32)
    input2 = input2.astype(np.float32)
    input3 = input3.astype(np.int32)

    expect = prod_force_cpu(input1, input2, input3)

    out_shape = expect.shape
    output = np.full(out_shape, np.nan, "float32")
    args = [input1, input2, input3, output]
    return args, expect, input1, input2, input3


def prod_force_cpu(net_deriv_tensor, in_deriv_tensor, nlist_tensor, natoms=192):
    net_deriv_tensor_shape = net_deriv_tensor.shape
    nlist_tensor_shape = nlist_tensor.shape
    natoms = net_deriv_tensor_shape[1]
    nframes = net_deriv_tensor_shape[0]
    ndescript = net_deriv_tensor_shape[2]
    nnei = nlist_tensor_shape[2]

    output_shape = [nframes, natoms, 3]

    force = np.full(output_shape, 0, "float32")
    for kk in range(nframes):
        for ii in range(natoms):
            for cc in range(3):
                force[kk, ii, cc] = 0.0
        for ii in range(natoms):
            for aa in range(ndescript):
                for cc in range(3):
                    force[kk, ii, cc] -= net_deriv_tensor[kk, ii, aa] * in_deriv_tensor[kk, ii, aa, cc]
            
            for jj in range(nnei):
                j_idx = nlist_tensor[kk, ii, jj]
                if j_idx > -1:
                    for aa in range(jj*4, jj*4 + 4):
                        for cc in range(3):
                            force[kk, j_idx, cc] += net_deriv_tensor[kk, ii, aa] * in_deriv_tensor[kk, ii, aa, cc]
    return force

def prod_force_se_a_run(input_shapes, input_dtype, attrs = {}):
    mod = utils.op_build_test(prod_force_se_a.ProdForceSeA, input_shapes, input_dtype, kernel_name = "force", attrs = attrs)
    args, expect, input1, input2, input3 = gen_data(input_shapes)
    output = utils.mod_launch(mod, args, expect=expect)
    rtol, atol = get_rtol_atol("prod_force_se_a", input_dtype)
    compare_res = compare_tensor(output, expect, rtol=rtol, atol=atol, equal_nan=True)
    return (input1, input2, input3), output, expect, compare_res
