# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import torch
from torch.nn import Module

from brevitas.nn import QuantLinear
from brevitas.quant import IntBias
from brevitas.quant_tensor import QuantTensor

OUTPUT_FEATURES = 10
INPUT_FEATURES = 5
BIT_WIDTH = 5


class TestQuantLinearInit:

    def test_module_init_defaults(self):
        mod = QuantLinear(out_features=OUTPUT_FEATURES, in_features=INPUT_FEATURES, bias=False)
        assert mod

    def test_module_init_bias_fp(self):
        mod = QuantLinear(out_features=OUTPUT_FEATURES, in_features=INPUT_FEATURES, bias=True)
        assert mod

    def test_module_init_bias_int(self):
        mod = QuantLinear(
            out_features=OUTPUT_FEATURES, in_features=INPUT_FEATURES, bias=True, bias_quant=IntBias)
        assert mod

    def test_module_init_scale_impl_type_override(self):
        mod = QuantLinear(
            out_features=OUTPUT_FEATURES,
            in_features=INPUT_FEATURES,
            bias=True,
            weight_scaling_impl_type='HE')
        assert mod.quant_weight_scale()


class TestQuantLinearFwd:

    def test_forward_defaults(self):
        mod = QuantLinear(out_features=OUTPUT_FEATURES, in_features=INPUT_FEATURES, bias=True)
        x = torch.rand(size=(3, INPUT_FEATURES))
        assert mod(x) is not None

    def test_forward_bias_fp(self):
        mod = QuantLinear(out_features=OUTPUT_FEATURES, in_features=INPUT_FEATURES, bias=True)
        x = torch.rand(size=(3, INPUT_FEATURES))
        assert mod(x) is not None

    def test_forward_bias_int(self):
        mod = QuantLinear(
            out_features=OUTPUT_FEATURES, in_features=INPUT_FEATURES, bias=True, bias_quant=IntBias)
        x = QuantTensor(
            torch.rand(size=(3, INPUT_FEATURES)),
            torch.tensor(1.0),
            torch.tensor(0.0),
            torch.tensor(3))
        assert mod(x) is not None
