//
// SPDX-FileCopyrightText: Copyright 2023-2025 Arm Limited and/or its affiliates <open-source-office@arm.com>
//
// SPDX-License-Identifier: Apache-2.0
//
// 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
//
// 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.
//
// Partially generated by tosaValidationGenerator for TOSA Specification 1.0.0draft
// TODO: Implement the constraints.

#include "tosa/tosa_argument_checks.hpp"

#include "common/data_type.hpp"
#include "compiler/attributes.hpp"
#include "compiler/operation.hpp"
#include "compiler/tensor_properties.hpp"

#include <optional>
#include <unordered_map>
#include <unordered_set>

using namespace tosa::validator::checks;
using regor::DataType;
using regor::Operation;
using regor::OpType;
using regor::TensorConnection;
using regor::TensorUsage;
namespace
{

const std::unordered_set<std::string_view> tosaSupportedTypes = {
    "bool_t",
    "i4_t",
    "int4_t",
    "uint6_t",
    "i8_t",
    "int8_t",
    "uint8_t",
    "i16_t",
    "int16_t",
    "uint16_t",
    "i32_t",
    "int32_t",
    "i48_t",
    "int48_t",
    "fp16_t",
    "bf16_t",
    "fp32_t",
    "fp64_t",
    "index_t",
    "mode_t",
    "mul_t",
    "tensor_list_t",
    "tosa_graph_t",
    "shape_t",
    "acc_size_t",
    "acc_type_t",
    "resize_mode_t",
    "nan_propagation_mode_t",
    "rounding_mode_t",
};

std::optional<regor::DataType> MapType(const std::string_view &type)
{

    static const std::unordered_map<std::string_view, DataType> typeMap{
        {"bool_t", DataType::Bool8},
        {"i4_t", DataType::Int4Packed8},
        {"int4_t", DataType::Int4Packed8},
        {"i8_t", DataType::Int8},
        {"int8_t", DataType::Int8},
        {"uint8_t", DataType::UInt8},
        {"i16_t", DataType::Int16},
        {"int16_t", DataType::Int16},
        {"uint16_t", DataType::UInt16},
        {"i32_t", DataType::Int32},
        {"int32_t", DataType::Int32},
        {"i48_t", DataType::Int48},
        {"int48_t", DataType::Int48},
        {"fp16_t", DataType::Float16},
        {"bf16_t", DataType::BFloat16},
        {"fp32_t", DataType::Float32},
        {"index_t", DataType::Int32},
        {"shape_t", DataType::Int64},
        {"acc_size_t", DataType::Int32},
        {"resize_mode_t", DataType::Int32},
        // The actual mul_t type is determined by an attribute
        {"mul_t", DataType::Int16},
    };

    if ( auto p = typeMap.find(type); p != typeMap.end() ) return p->second;
    return std::nullopt;
}

const TensorConnection *GetTensorForArgument(const Operation *op, const Argument &argument)
{

    static const std::unordered_map<std::string_view, TensorUsage> inputUsageMap{
        {"input", TensorUsage::IFM},
        {"input_real", TensorUsage::IFM},
        {"input_imag", TensorUsage::IFM1},
        {"input1", TensorUsage::IFM},
        {"input2", TensorUsage::IFM1},
        {"input3", TensorUsage::IFM2},
        {"A", TensorUsage::IFM},
        {"B", TensorUsage::IFM1},
        {"weight", TensorUsage::Weights},
        {"bias", TensorUsage::Scales},
        {"values", TensorUsage::IFM},
        {"indices", TensorUsage::IFM1},
        {"values_in", TensorUsage::IFM},
        {"multiplier", TensorUsage::Params},
        {"multiples", TensorUsage::Params},
        {"shift", TensorUsage::Params1},
        {"scale", TensorUsage::Params},
        {"offset", TensorUsage::Params1},
        {"border", TensorUsage::Params2},
        {"padding", TensorUsage::Params},
        {"pad_const", TensorUsage::Params1},
        {"start", TensorUsage::Params},
        {"size", TensorUsage::Params1},
        {"input_zp", TensorUsage::Params},
        {"input1_zp", TensorUsage::Params},
        {"A_zp", TensorUsage::Params},
        {"B_zp", TensorUsage::Params1},
        {"weight_zp", TensorUsage::Params1},
        {"output_zp", TensorUsage::Params1},
        {"shape", TensorUsage::Params},
        {"table", TensorUsage::Params},
    };

    switch ( argument.category )
    {
        case Category::Input:  // input, weight, bias, input_real, input_imag, A, B, input1, input2, input3
            if ( op->Type() == OpType::Rescale )  // Rescale does not follow the pattern the other ops use
            {
                if ( argument.name == "input_zp" ) return op->Input(TensorUsage::Params2);
                if ( argument.name == "output_zp" ) return op->Input(TensorUsage::Params3);
            }
            // Mul shift parameter does not have the same pattern as Rescale
            if ( op->Type() == OpType::Mul && argument.name == "shift" ) return op->Input(TensorUsage::Params);
            // Scatter input parameter does not have the same pattern as other operations
            if ( op->Type() == OpType::Scatter && argument.name == "input" ) return op->Input(TensorUsage::IFM2);
            if ( auto p = inputUsageMap.find(argument.name); p != inputUsageMap.end() )
            {
                return op->Input(p->second);
            }
            break;
        case Category::Output:
        {
            auto usage = regor::TensorUsage::OFM;
            return op->Output(usage);
        }
        break;
        case Category::Attribute:  // kernel stride pad dilation inverse weight bias out_pad out_shape padding new_shape
                                   // start size multiples perms scale offset border multiplier shift values then_graph
                                   // else_graph cond_graph body_graph
            if ( auto p = inputUsageMap.find(argument.name); p != inputUsageMap.end() )
            {
                return op->Input(p->second);
            }
            break;
        case Category::ScalarAttribute:  // axis acc_size input_zp output_zp weight_zp A_zp B_zp min_val max_val round
                                         // pad_const mode scale32 double_round per_channel
            break;
        default:
            break;
    }
    return nullptr;
}

std::optional<int> GetRankForArgument(const Operation *op, const Argument &argument)
{
    auto *tensorConnection = GetTensorForArgument(op, argument);
    if ( tensorConnection ) return tensorConnection->shape.Size();
    return std::nullopt;
}

std::optional<regor::DataType> GetTypeForArgument(const Operation *op, const Argument &argument)
{
    static const std::unordered_map<std::string_view, DataType> attributeTypeMap{
        {"kernel", DataType::Int32},
        {"stride", DataType::Int32},
        {"pad", DataType::Int32},
        {"dilation", DataType::Int32},
        {"shape", DataType::Int32},
    };

    auto *tensorConnection = GetTensorForArgument(op, argument);
    if ( tensorConnection ) return tensorConnection->tensor->Type();
    switch ( argument.category )
    {
        case Category::Input:  // input, weight, bias, input_real, input_imag, A, B, input1, input2, input3
            break;
        case Category::Output:
            break;
        case Category::Attribute:  // kernel stride pad dilation inverse weight bias out_pad out_shape padding new_shape
                                   // start size multiples perms scale offset border multiplier shift values then_graph
                                   // else_graph cond_graph body_graph
            if ( attributeTypeMap.count(argument.name) > 0 ) return attributeTypeMap.at(argument.name);
            break;
        case Category::ScalarAttribute:  // axis acc_size input_zp output_zp weight_zp A_zp B_zp min_val max_val round
                                         // pad_const mode scale32 double_round per_channel
            // TODO
            break;
        default:
            break;
    }
    return std::nullopt;
}

bool IsReduceOp(OpType op)
{
    static const std::unordered_set<OpType> reduceOps{
        OpType::ReduceAll,
        OpType::ReduceAny,
        OpType::ReduceMax,
        OpType::ReduceMin,
        OpType::ReduceProduct,
        OpType::ReduceSum,
    };
    return reduceOps.count(op) > 0;
}

void ValidateArgumentShape(const regor::Operation *op, const Argument &argument, const tosa::validator::Context &)
{
    if ( argument.rank_bounds.first >= 0 || argument.rank_bounds.second >= 0 )
    {
        auto rank = GetRankForArgument(op, argument);
        if ( rank && (rank < argument.rank_bounds.first || rank > argument.rank_bounds.second) )
            throw std::invalid_argument("Tensor rank out of bounds");
    }
}

void ValidateArgumentShapes(const regor::Operation *op, const std::vector<const Argument *> &arguments, const tosa::validator::Context &context)
{
    if ( IsReduceOp(op->Type()) )
    {
        auto inRank = op->Input(regor::TensorUsage::IFM)->shape.Size();
        auto outRank = op->Output(regor::TensorUsage::OFM)->shape.Size();
        if ( inRank != outRank ) throw std::invalid_argument("Tensor ranks different");
        if ( inRank < 1 ) throw std::invalid_argument("Tensor rank < 1");
    }
    for ( const auto &argument : arguments )
        ValidateArgumentShape(op, *argument, context);
}

bool ValidateArgument(const regor::Operation *op, const Argument *argument, DataType expectedType)
{
    auto actualType = GetTypeForArgument(op, *argument);
    if ( !actualType ) return false;
    if ( *actualType == expectedType ) return true;
    return false;
}

bool ResolveAndValidateArgument(const regor::Operation *op, const Argument *argument, const Typesupport *typesupport)
{
    if ( argument == nullptr ) return false;
    if ( argument->category == Category::Attribute ) return true;        // For now
    if ( argument->category == Category::ScalarAttribute ) return true;  // For now
    std::string_view typeName = argument->element_type;
    if ( typesupport )
    {
        if ( auto p = typesupport->find(typeName); p != typesupport->end() )
        {
            typeName = p->second;
        }
    }
    auto expectedType = MapType(typeName);
    if ( !expectedType ) return false;
    if ( op->Type() == OpType::Rescale )
    {
        auto *attr = op->Attribute<regor::sign_attr_t>();
        if ( ((argument->category == Category::Input && argument->name == "input" && attr->input_unsigned) ||
                 (argument->category == Category::Output && argument->name == "output" && attr->output_unsigned)) )
        {

            if ( attr->output_unsigned && DataTypeSizeBits(*expectedType) > 16 )
            {
                // Error to make unsigned with anything but i8_t and i16_t
                return false;
            }
            // Signedness of IFM/OFM for Rescale depends on the input_unsigned/output_unsigned attributes
            // Actual input/output type can be either signed or unsigned
            auto unsignedType = *expectedType & ~unsigned(DataType::Signed);

            // Check unsigned type
            if ( ValidateArgument(op, argument, unsignedType) ) return true;

            // Unsgined check failed, signed check will be done below
        }
        if ( argument->category == Category::Input && argument->name == "multiplier" )
        {
            auto *r_attr = op->Attribute<regor::rescale_attr_t>();
            *expectedType = r_attr->scale32 ? DataType::Int32 : DataType::Int16;
            // Validate argument below
        }
    }
    return ValidateArgument(op, argument, *expectedType);
}

bool ResolveAndValidateArguments(
    const regor::Operation *op, const std::vector<const Argument *> arguments, const Typesupport *typesupport = nullptr)
{
    for ( const auto argument : arguments )
    {
        if ( !ResolveAndValidateArgument(op, argument, typesupport) ) return false;
    }
    return true;
}

bool ArgumentsCanBeResolvedAndValidated(const regor::Operation *op, const std::vector<const Argument *> &arguments,
    const std::vector<Typesupport> &typesupports)
{
    if ( op->Type() == regor::OpType::If )
    {
        auto cond = op->Input(TensorUsage::IFM);
        if ( cond == nullptr ) return false;
        auto type = cond->tensor->Type();
        if ( type != DataType::Bool8 && type != DataType::Bool ) return false;
        auto *attr = op->Attribute<regor::cond_attr_t>();
        if ( attr->then_branch.empty() ) return false;
        if ( attr->else_branch.empty() ) return false;
        return true;
    }
    if ( op->Type() == regor::OpType::While )
    {
        auto *attr = op->Attribute<regor::while_attr_t>();
        if ( attr->body_branch.empty() ) return false;
        if ( attr->cond_branch.empty() ) return false;
        return true;
    }
    if ( ResolveAndValidateArguments(op, arguments) ) return true;
    for ( const auto &typesupport : typesupports )
    {
        if ( ResolveAndValidateArguments(op, arguments, &typesupport) ) return true;
    }
    return false;
}

bool CanResolveArgument(const Argument *argument, const Typesupport *typesupport)
{
    if ( argument == nullptr ) return false;
    std::string_view typeName = argument->element_type;
    if ( typesupport )
    {
        if ( auto p = typesupport->find(typeName); p != typesupport->end() )
        {
            typeName = p->second;
        }
    }
    return tosaSupportedTypes.count(typeName) > 0;
}

bool CanResolveArguments(const std::vector<const Argument *> &arguments, const Typesupport *typesupport = nullptr)
{
    for ( const auto argument : arguments )
    {
        if ( !CanResolveArgument(argument, typesupport) ) return false;
    }
    return true;
}
bool ArgumentsCanBeResolved(const std::vector<const Argument *> &arguments, const std::vector<Typesupport> &typesupports)
{
    if ( CanResolveArguments(arguments) ) return true;
    for ( const auto &typesupport : typesupports )
    {
        if ( CanResolveArguments(arguments, &typesupport) ) return true;
    }
    return false;
}
}  // namespace
namespace tosa
{
namespace validator
{
namespace checks
{

void ValidateArguments(const regor::Operation *op, const std::vector<const Argument *> &arguments,
    const std::vector<Typesupport> &typesupports, const Context &context)
{
    if ( !ArgumentsCanBeResolved(arguments, typesupports) ) throw std::invalid_argument("Unsupported operation");
    if ( !ArgumentsCanBeResolvedAndValidated(op, arguments, typesupports) )
        throw std::invalid_argument("Invalid arguments");
    ValidateArgumentShapes(op, arguments, context);
}

}  // namespace checks
}  // namespace validator
}  // namespace tosa
