/*
 * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
 *
 * 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 "contrib/nlohmann/json/json.hpp"

#include "../include/cudnn_frontend_Logging.h"

#include <cstdlib>
#include <fstream>
#pragma once

using json = nlohmann::json;

namespace cudnn_frontend {

// Loads the json handle from the json file 
// json file is defined by environment variable
// CUDNN_ERRATA_JSON_FILE. If the environment variable
// is not set the value set in the API is considered.
static bool
load_from_config(json &json_handle, const std::string & errata_json) {
    const char * err_json = get_environment("CUDNN_ERRATA_JSON_FILE");
    if (err_json == NULL && errata_json == "") {return false;}
    if (err_json == NULL) { err_json = errata_json.c_str();}
    std::ifstream ifs(err_json, std::ifstream::in);
    if (!ifs.is_open() || !ifs.good()) {return false;}
    ifs >> json_handle;
    return true;
}

/**
 * @brief Checks the shape of an operation to compare against errata filter height and width for kernel blocking
 * 
 * @param op The operation's tensors to check
 * @param shape_format The shape format of the tensor (NCHW vs NHWC)
 * @param tensor_attr The cudnnBackendAttributeName_t of the tensor's shape we want to check
 * @param blocked_height The height we want to filter out
 * @param blocked_width The width we want to filter out
 * @param blocked_channels The channels we want to filter out. Defaults to -1 (not filter out channels)
 * @return true The passed in operation shape matches the blocked shape
 * @return false The passed in operation shape does not match the blocked shape
 */
static bool
check_shape(cudnnBackendDescriptor_t &op,
            const std::string &shape_format,
            cudnnBackendAttributeName_t tensor_attr,
            const std::vector<int64_t> &blocked_shape) {

    // Get backend descriptor to individual tensor to be able to get shape
    ManagedOpaqueDescriptor tensor = make_shared_backend_pointer(CUDNN_BACKEND_TENSOR_DESCRIPTOR);
    cudnnBackendDescriptor_t tensor_ = tensor->get_backend_descriptor();
    int64_t count = 0;
    cudnnStatus_t status = cudnnBackendGetAttribute(op,
                                        tensor_attr,
                                        CUDNN_TYPE_BACKEND_DESCRIPTOR,
                                        1,
                                        &count,
                                        &tensor_);
    if (status != CUDNN_STATUS_SUCCESS) {
        #ifndef NV_CUDNN_DISABLE_EXCEPTION
        throw cudnnException(
            std::string("Error getting attribute. cudnn_status: " + to_string(status)).c_str(), status);
        #endif
    }

    // Get tensor dims
    std::array<int64_t, 5> tensor_dims;
    status = cudnnBackendGetAttribute(tensor_,
                                        CUDNN_ATTR_TENSOR_DIMENSIONS,
                                        CUDNN_TYPE_INT64,
                                        5,
                                        &count,
                                        tensor_dims.data());
    if (status != CUDNN_STATUS_SUCCESS) {
        #ifndef NV_CUDNN_DISABLE_EXCEPTION
        throw cudnnException(
            std::string("Error getting attribute. cudnn_status: " + to_string(status)).c_str(), status);
        #endif
    }
    // tensor_dims is 1 indexed
    int64_t first_dim = tensor_dims[1]; // batch size for input/output tensor, output channels for filter tensor
    int64_t blocked_first_dim = blocked_shape[0];

    // Defaults to true becuase -1 means we don't filter that out (Wildcard). If something later blocks, then the comparison will be correct
    bool blocked = (blocked_first_dim != -1) ? (first_dim == blocked_first_dim) : true;

    // Check for shape format to extract the right dimension. Filter shape will always be "NCHW" for convenience.
    int64_t channels = (shape_format == "NCHW") ? tensor_dims[2] : tensor_dims[4]; // channels
    int64_t blocked_channels = (shape_format == "NCHW") ? blocked_shape[1] : blocked_shape[3];
    blocked = (blocked_channels != -1) ? (blocked && channels == blocked_channels) : true;

    int64_t height = (shape_format == "NCHW") ? tensor_dims[3] : tensor_dims[2];
    int64_t blocked_height = (shape_format == "NCHW") ? blocked_shape[2] : blocked_shape[1];
    blocked = (blocked_height != -1) ? (blocked && height == blocked_height) : true;
    
    int64_t width = (shape_format == "NCHW") ? tensor_dims[4] : tensor_dims[3];
    int64_t blocked_width = (shape_format == "NCHW") ? blocked_shape[3] : blocked_shape[2];
    blocked = (blocked_width != -1) ? (blocked && width == blocked_width) : true;

    return blocked;
}

template <typename T>
static bool 
check_rule(const json &json_handle, const std::string &executionPlanTag,
    cudnnHandle_t handle, T fn) {
    std::string operation = json_handle["operation"];
    int64_t engine        =  json_handle["engine"];
    uint64_t cudnn_start     =  0;
    uint64_t cudnn_end       =  std::numeric_limits<uint64_t>::max();
    if (json_handle.contains("cudnn_version_start")) {
        cudnn_start   =  json_handle["cudnn_version_start"];
    }
    if (json_handle.contains("cudnn_version_end")) {
        cudnn_end     =  json_handle["cudnn_version_end"];
    }
    std::string tag_prefix = operation + "_eng" + std::to_string(engine) + "_"; 
    std::string mod_tag    = executionPlanTag + "_";
    bool blocked = 
        tag_prefix.size() <= mod_tag.size() &&
        std::equal(tag_prefix.begin(), tag_prefix.end(), mod_tag.begin()) &&
        CUDNN_VERSION >= cudnn_start &&
        CUDNN_VERSION < cudnn_end;

    if (blocked && json_handle.contains("knob")) { // Short circuit if operation and engine do not match
        for (auto& kv : json_handle["knob"]) {
            blocked = blocked &&
                (executionPlanTag.find(kv) != std::string::npos);
        }
    }
    blocked = blocked && fn(); 
    return blocked;

    CUDNN_FRONTEND_UNUSED(handle);
}

// Overload for check_rule to take in an operation graph for shape filtering
template <typename T>
static bool 
check_rule(const json &json_handle, const std::string &executionPlanTag,
    cudnnHandle_t handle, T fn, const OperationGraph& opGraph) {
    std::string operation = json_handle["operation"];
    int64_t engine        =  json_handle["engine"];
    uint64_t cudnn_start     =  0;
    uint64_t cudnn_end       =  std::numeric_limits<uint64_t>::max();
    if (json_handle.contains("cudnn_version_start")) {
        cudnn_start   =  json_handle["cudnn_version_start"];
    }
    if (json_handle.contains("cudnn_version_end")) {
        cudnn_end     =  json_handle["cudnn_version_end"];
    }
    std::string tag_prefix = operation + "_eng" + std::to_string(engine) + "_"; 
    std::string mod_tag    = executionPlanTag + "_";
    bool blocked = 
        tag_prefix.size() <= mod_tag.size() &&
        std::equal(tag_prefix.begin(), tag_prefix.end(), mod_tag.begin()) &&
        CUDNN_VERSION >= cudnn_start &&
        CUDNN_VERSION < cudnn_end;

    if (blocked && json_handle.contains("knob")) { // Short circuit if operation and engine do not match
        for (auto& kv : json_handle["knob"]) {
            blocked = blocked &&
                (executionPlanTag.find(kv) != std::string::npos);
        }
    }
    
    if (blocked && json_handle.contains("input_shape")) { // Check if user wants to block kernel for specific input shape
        if (!json_handle.contains("shape_format")) {
            std::string message = "ERROR: Please set a shape format (e.g. shape_format: \"NCWH\") for errata filters using input/kernel shape";
            #ifndef NV_CUDNN_DISABLE_EXCEPTION
            throw cudnnException(message.c_str(), CUDNN_STATUS_BAD_PARAM);
            #endif
            getLogger() << message << std::endl;
            return blocked;
        }

        std::array<ManagedOpaqueDescriptor, MAX_OPGRAPH_OPS> ops = opGraph.getOps();
        std::array<cudnnBackendDescriptor_t, MAX_OPGRAPH_OPS> ops_;
        for (unsigned int i = 0; i < opGraph.getOpCount(); i++) {
            ops_[i] = ops[i]->get_backend_descriptor();
        }

        std::string shape_format = json_handle["shape_format"];
        std::vector<int64_t> blocked_shape = json_handle["input_shape"];

        // Forward conv operation
        if (operation == "ConvFwd") {
            blocked = blocked && check_shape(ops_[0], shape_format, CUDNN_ATTR_OPERATION_CONVOLUTION_FORWARD_X, blocked_shape);

        // Operation is conv wgrad
        } else if (operation == "ConvBwdFilter") {
            blocked = blocked && check_shape(ops_[0], shape_format, CUDNN_ATTR_OPERATION_CONVOLUTION_BWD_FILTER_X, blocked_shape);

        // Operation is conv dgrad 
        } else if (operation == "ConvBwdData") {
            blocked = blocked && check_shape(ops_[0], shape_format, CUDNN_ATTR_OPERATION_CONVOLUTION_BWD_DATA_DX, blocked_shape);
        }
    }

    if (blocked && json_handle.contains("filter_shape")) { // Check if user wants to block kernel for specific filter shape
        std::array<ManagedOpaqueDescriptor, 50> ops = opGraph.getOps();
        std::array<cudnnBackendDescriptor_t, 50> ops_;
        for (unsigned int i = 0; i < opGraph.getOpCount(); i++) {
            ops_[i] = ops[i]->get_backend_descriptor();
        }

        std::vector<int64_t> blocked_shape = json_handle["filter_shape"];

        // Forward conv operation
        if (operation == "ConvFwd") {
            // Filter format is always [output channels, input channels, height, width] so we hardcode "NCHW" to match and not repeat code
            blocked = blocked && check_shape(ops_[0], "NCHW", CUDNN_ATTR_OPERATION_CONVOLUTION_FORWARD_W, blocked_shape);

        // Operation is conv wgrad
        } else if (operation == "ConvBwdFilter") {
            blocked = blocked && check_shape(ops_[0], "NCHW", CUDNN_ATTR_OPERATION_CONVOLUTION_BWD_FILTER_DW, blocked_shape);

        // Operation is conv dgrad 
        } else if (operation == "ConvBwdData") {
            blocked = blocked && check_shape(ops_[0], "NCHW", CUDNN_ATTR_OPERATION_CONVOLUTION_BWD_DATA_W, blocked_shape);
        }
    }
    blocked = blocked && fn(); 
    return blocked;

    CUDNN_FRONTEND_UNUSED(handle);
}

// Takes in an initialzed json handle and checks if it satisfies the 
// condition for running it. Returns true if the given executionPlanTag
// is faulty.
template <typename T>
static bool
check_errata(const json &json_handle, const std::string & executionPlanTag,
    cudnnHandle_t handle, T fn) {

    cudnn_frontend::getLogger() << "[cudnn_frontend] " << "Verifying " << executionPlanTag;
    for (auto const &rule : json_handle["rules"]) {
        if (check_rule<T>(rule, executionPlanTag, handle, fn)) {
            cudnn_frontend::getLogger() << ". Blocking." << std::endl;
            return true;
        }
    }

    cudnn_frontend::getLogger() << ". Passed." << std::endl;
    return false;
}

// Overload. Takes in an initialzed json handle, an execution plan tag, and a operation graph and checks if it satisfies the 
// condition for running it. Returns true if the given executionPlanTag + operation graph
// is faulty
template <typename T>
static bool
check_errata(const json &json_handle, const std::string & executionPlanTag,
    cudnnHandle_t handle, const OperationGraph &opGraph, T fn) {

    cudnn_frontend::getLogger() << "[cudnn_frontend] " << "Verifying " << executionPlanTag;
    for (auto const &rule : json_handle["rules"]) {
        if (check_rule<T>(rule, executionPlanTag, handle, fn, opGraph)) {
            cudnn_frontend::getLogger() << ". Blocking." << std::endl;
            return true;
        }
    }

    cudnn_frontend::getLogger() << ". Passed." << std::endl;
    return false;
}

}
