
/*******************************************************************************
 *
 * MIT License
 *
 * Copyright (c) 2025 Advanced Micro Devices, Inc.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 *
 *******************************************************************************/
#include <miopen/config.hpp>
#include <miopen/conv/heuristics/ai_heuristics.hpp>

#include <miopen/db_path.hpp>
#include <miopen/logger.hpp>

#if MIOPEN_ENABLE_AI_IMMED_MODE_FALLBACK

namespace miopen {
namespace ai {
namespace immed_mode {

// Helper function to safely load JSON with error handling
static std::optional<nlohmann::json> LoadJSONSafe(const std::string& arch)
{

    // const std::string archh = "gfx942";

    try
    {
        const auto file_path = GetSystemDbPath() / (arch + "_metadata.tn.model");
        return common::LoadJSON(file_path);
    }
    catch(const std::exception& e)
    {
        MIOPEN_LOG_I2("Failed to load JSON for " << arch << ": " << e.what());
        return std::nullopt;
    }
    catch(...)
    {
        MIOPEN_LOG_I2("Failed to load JSON for " << arch << ": unknown error");
        return std::nullopt;
    }
    MIOPEN_LOG_I2("Read metadata for: " << arch);
}

// Static helper functions for loading individual components
std::optional<int> MetadataND::LoadSpatialDim(const std::string& arch)
{
    auto json_opt = LoadJSONSafe(arch);
    if(!json_opt)
        return std::nullopt;

    try
    {
        return json_opt->at("redundant_columns").at("spatial_dim").get<int>();
    }
    catch(const std::exception& e)
    {
        MIOPEN_LOG_I2("Failed to load spatial_dim for " << arch << ": " << e.what());
        return std::nullopt;
    }
}

std::optional<std::vector<std::string>> MetadataND::LoadFeatures(const std::string& arch)
{
    auto json_opt = LoadJSONSafe(arch);
    if(!json_opt)
        return std::nullopt;

    try
    {
        return json_opt->at("conv_params_used_as_features").get<std::vector<std::string>>();
    }
    catch(const std::exception& e)
    {
        MIOPEN_LOG_I2("Failed to load features for " << arch << ": " << e.what());
        return std::nullopt;
    }
}

std::optional<size_t> MetadataND::LoadNumInputs(const std::string& arch)
{
    auto json_opt = LoadJSONSafe(arch);
    if(!json_opt)
        return std::nullopt;

    try
    {
        return json_opt->at("num_inputs").get<size_t>();
    }
    catch(const std::exception& e)
    {
        MIOPEN_LOG_I2("Failed to load num_inputs for " << arch << ": " << e.what());
        return std::nullopt;
    }
}

std::optional<size_t> MetadataND::LoadNumCu(const std::string& arch)
{
    auto json_opt = LoadJSONSafe(arch);
    if(!json_opt)
        return std::nullopt;

    try
    {
        return json_opt->at("gpu").at("num_cu").get<size_t>();
    }
    catch(const std::exception& e)
    {
        MIOPEN_LOG_I2("Failed to load gpu.num_cu for " << arch << ": " << e.what());
        return std::nullopt;
    }
}

std::optional<size_t> MetadataND::LoadNumOutputs(const std::string& arch)
{
    auto json_opt = LoadJSONSafe(arch);
    if(!json_opt)
        return std::nullopt;

    try
    {
        return json_opt->at("num_outputs").get<size_t>();
    }
    catch(const std::exception& e)
    {
        MIOPEN_LOG_I2("Failed to load num_outputs for " << arch << ": " << e.what());
        return std::nullopt;
    }
}

std::optional<size_t> MetadataND::LoadNumSolvers(const std::string& arch)
{
    auto json_opt = LoadJSONSafe(arch);
    if(!json_opt)
        return std::nullopt;

    try
    {
        return json_opt->at("num_solvers").get<size_t>();
    }
    catch(const std::exception& e)
    {
        MIOPEN_LOG_I2("Failed to load num_solvers for " << arch << ": " << e.what());
        return std::nullopt;
    }
}

std::optional<std::unordered_map<size_t, std::string>>
MetadataND::LoadSolverMap(const std::string& arch)
{
    auto json_opt = LoadJSONSafe(arch);
    if(!json_opt)
        return std::nullopt;

    try
    {
        return common::ReverseMap<std::string, size_t>(json_opt->at("encodings").at("solver"));
    }
    catch(const std::exception& e)
    {
        MIOPEN_LOG_I2("Failed to load solver_map for " << arch << ": " << e.what());
        return std::nullopt;
    }
}

std::optional<std::vector<float>>
MetadataND::LoadFeaturesMean([[maybe_unused]] const std::string& arch, size_t num_inputs)
{
    // For now, return default values (could be enhanced to load from JSON stats)
    // This is a simplified version that returns zeros for mean
    return std::vector<float>(num_inputs, 0.0f);
}

std::optional<std::vector<float>>
MetadataND::LoadFeaturesStd([[maybe_unused]] const std::string& arch, size_t num_inputs)
{
    // For now, return default values (could be enhanced to load from JSON stats)
    // This is a simplified version that returns ones for std
    return std::vector<float>(num_inputs, 1.0f);
}

std::optional<std::unordered_map<std::string, int>>
MetadataND::LoadDirectionEncodings(const std::string& arch)
{
    auto json_opt = LoadJSONSafe(arch);
    if(!json_opt)
        return std::nullopt;

    try
    {
        return json_opt->at("encodings")
            .at("direction")
            .get<std::unordered_map<std::string, int>>();
    }
    catch(const std::exception& e)
    {
        MIOPEN_LOG_I2("Failed to load direction_encodings for " << arch << ": " << e.what());
        return std::nullopt;
    }
}

std::optional<std::unordered_map<std::string, int>>
MetadataND::LoadPrecisionEncodings(const std::string& arch)
{
    auto json_opt = LoadJSONSafe(arch);
    if(!json_opt)
        return std::nullopt;

    try
    {
        return json_opt->at("encodings")
            .at("precision")
            .get<std::unordered_map<std::string, int>>();
    }
    catch(const std::exception& e)
    {
        MIOPEN_LOG_I2("Failed to load precision_encodings for " << arch << ": " << e.what());
        return std::nullopt;
    }
}

std::optional<std::unordered_map<std::string, int>>
MetadataND::LoadInLayoutEncodings(const std::string& arch)
{
    auto json_opt = LoadJSONSafe(arch);
    if(!json_opt)
        return std::nullopt;

    try
    {
        return json_opt->at("encodings")
            .at("in_layout")
            .get<std::unordered_map<std::string, int>>();
    }
    catch(const std::exception& e)
    {
        MIOPEN_LOG_I2("Failed to load in_layout_encodings for " << arch << ": " << e.what());
        return std::nullopt;
    }
}

std::optional<std::unordered_map<std::string, int>>
MetadataND::LoadFilLayoutEncodings(const std::string& arch)
{
    auto json_opt = LoadJSONSafe(arch);
    if(!json_opt)
        return std::nullopt;

    try
    {
        return json_opt->at("encodings")
            .at("fil_layout")
            .get<std::unordered_map<std::string, int>>();
    }
    catch(const std::exception& e)
    {
        MIOPEN_LOG_I2("Failed to load fil_layout_encodings for " << arch << ": " << e.what());
        return std::nullopt;
    }
}

std::optional<std::unordered_map<std::string, int>>
MetadataND::LoadOutLayoutEncodings(const std::string& arch)
{
    auto json_opt = LoadJSONSafe(arch);
    if(!json_opt)
        return std::nullopt;

    try
    {
        return json_opt->at("encodings")
            .at("out_layout")
            .get<std::unordered_map<std::string, int>>();
    }
    catch(const std::exception& e)
    {
        MIOPEN_LOG_I2("Failed to load out_layout_encodings for " << arch << ": " << e.what());
        return std::nullopt;
    }
}

// Constructor - loads all data immediately with error handling
MIOPEN_INTERNALS_EXPORT
MetadataND::MetadataND(const std::string& device, const int& dim)
    : is_valid(false),  // Initialize to false, will be set to true if all loads succeed
      spatial_dim(dim), // Store the dimension this model supports
      features(),
      num_inputs(0),
      num_outputs(0),
      num_solvers(0),
      solver_map(),
      features_mean(),
      features_std(),
      direction_encodings_3d(),
      precision_encodings_3d(),
      in_layout_encodings(),
      fil_layout_encodings(),
      out_layout_encodings()
{
    if(dim == 2)
    {
        model_prefix = device;
    }
    else if(dim == 3)
    {
        model_prefix = device + "_3d";
    }
    else
    {
        MIOPEN_LOG_I2("Unsupported dimension " << dim << " for MetadataND, expected 2 or 3");
        return;
    }

    // Optional: Load spatial_dim from JSON to validate it matches expected dimension
    auto spatial_dim_opt = LoadSpatialDim(model_prefix);
    if(spatial_dim_opt && *spatial_dim_opt != dim)
    {
        MIOPEN_LOG_W("MetadataND: Spatial dimension mismatch for "
                     << model_prefix << " - expected " << dim << ", found " << *spatial_dim_opt);
    }
    // Load all components using std::optional pattern (using full arch_name with "_3d")
    auto features_opt    = LoadFeatures(model_prefix);
    auto num_inputs_opt  = LoadNumInputs(model_prefix);
    auto num_outputs_opt = LoadNumOutputs(model_prefix);
    auto num_solvers_opt = LoadNumSolvers(model_prefix);
    auto solver_map_opt  = LoadSolverMap(model_prefix);

    // Check if basic components loaded successfully
    if(!features_opt || !num_inputs_opt || !num_outputs_opt || !num_solvers_opt || !solver_map_opt)
    {
        MIOPEN_LOG_I2("MetadataND: Failed to load basic components for " << model_prefix);
        return;
    }

    // Now load mean/std with the known num_inputs
    auto features_mean_opt = LoadFeaturesMean(model_prefix, *num_inputs_opt);
    auto features_std_opt  = LoadFeaturesStd(model_prefix, *num_inputs_opt);

    // Load encoding maps
    auto direction_encodings_opt  = LoadDirectionEncodings(model_prefix);
    auto precision_encodings_opt  = LoadPrecisionEncodings(model_prefix);
    auto in_layout_encodings_opt  = LoadInLayoutEncodings(model_prefix);
    auto fil_layout_encodings_opt = LoadFilLayoutEncodings(model_prefix);
    auto out_layout_encodings_opt = LoadOutLayoutEncodings(model_prefix);

    // Check if all components loaded successfully
    if(!features_mean_opt || !features_std_opt || !direction_encodings_opt ||
       !precision_encodings_opt || !in_layout_encodings_opt || !fil_layout_encodings_opt ||
       !out_layout_encodings_opt)
    {
        MIOPEN_LOG_I2("MetadataND: Failed to load encoding components for " << model_prefix);
        return;
    }

    // All components loaded successfully, now we can safely move the data
    // We need to const_cast because the members are const
    features               = std::move(*features_opt);
    num_inputs             = *num_inputs_opt;
    num_outputs            = *num_outputs_opt;
    num_solvers            = *num_solvers_opt;
    solver_map             = std::move(*solver_map_opt);
    features_mean          = std::move(*features_mean_opt);
    features_std           = std::move(*features_std_opt);
    direction_encodings_3d = std::move(*direction_encodings_opt);
    precision_encodings_3d = std::move(*precision_encodings_opt);
    in_layout_encodings    = std::move(*in_layout_encodings_opt);
    fil_layout_encodings   = std::move(*fil_layout_encodings_opt);
    out_layout_encodings   = std::move(*out_layout_encodings_opt);
    // num_cu is only consumed by the engineered 2D feature path. Absence leaves it 0, which zeroes
    // the hardware-aware derived features (a degraded 2D prediction), so warn rather than fail.
    if(auto num_cu_opt = LoadNumCu(model_prefix))
        num_cu_3d = *num_cu_opt;
    else
        MIOPEN_LOG_W("TunaNet ND metadata for "
                     << model_prefix
                     << " has no gpu.num_cu; engineered 2D features will use num_cu=0");
    // Mark as valid after successful loading
    is_valid = true;

    if(miopen::IsLogging(LoggingLevel::Info2))
    {
        MIOPEN_LOG_I2("MetadataND loaded successfully for arch: "
                      << model_prefix << ", num_inputs=" << num_inputs
                      << ", num_solvers=" << num_solvers);
    }
}

// Encoding methods with safe error handling
MIOPEN_INTERNALS_EXPORT
size_t MetadataND::EncodeDirection(miopen::conv::Direction dir) const
{
    if(!is_valid)
        return 0;

    try
    {
        if(dir == conv::Direction::BackwardWeights)
            return direction_encodings_3d.at("W");
        else if(dir == conv::Direction::BackwardData)
            return direction_encodings_3d.at("B");
        else
            return direction_encodings_3d.at("F");
    }
    catch(...)
    {
        MIOPEN_LOG_W("Direction encoding failed in ND metadata, returning 0");
        return 0;
    }
}

MIOPEN_INTERNALS_EXPORT
size_t MetadataND::EncodePrecision(miopenDataType_t data_type) const
{
    if(!is_valid)
        return 0;

    const char* key = common::DataTypeToEncodingKey(data_type);
    if(key != nullptr)
    {
        const auto it = precision_encodings_3d.find(key);
        if(it != precision_encodings_3d.end())
            return it->second;
    }
    // Unsupported datatype, or a precision this model wasn't trained on: throw so the caller falls
    // back to the non-AI heuristic. Returning an out-of-range index here would feed an all-zero
    // precision one-hot into the model -- a silently degraded prediction, which we want to avoid.
    MIOPEN_THROW("Precision not supported by this TunaNet ND model");
}

MIOPEN_INTERNALS_EXPORT
size_t MetadataND::EncodeLayout(const std::string& layout) const
{
    if(!is_valid)
        return 0;

    auto it = in_layout_encodings.find(layout);
    if(it != in_layout_encodings.end())
        return it->second;

    MIOPEN_LOG_W("Unsupported layout " << layout << " in ND metadata, returning 0");
    return 0;
}

MIOPEN_INTERNALS_EXPORT
size_t MetadataND::EncodeInLayout(const std::string& layout) const
{
    if(!is_valid)
        return 0;

    auto it = in_layout_encodings.find(layout);
    return (it != in_layout_encodings.end()) ? it->second : 0;
}

MIOPEN_INTERNALS_EXPORT
size_t MetadataND::EncodeFilLayout(const std::string& layout) const
{
    if(!is_valid)
        return 0;

    auto it = fil_layout_encodings.find(layout);
    return (it != fil_layout_encodings.end()) ? it->second : 0;
}

MIOPEN_INTERNALS_EXPORT
size_t MetadataND::EncodeOutLayout(const std::string& layout) const
{
    if(!is_valid)
        return 0;

    auto it = out_layout_encodings.find(layout);
    return (it != out_layout_encodings.end()) ? it->second : 0;
}

} // namespace immed_mode
} // namespace ai
} // namespace miopen

#endif // MIOPEN_ENABLE_AI_IMMED_MODE_FALLBACK
