/*******************************************************************************
 *
 * MIT License
 *
 * Copyright (c) 2020 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.
 *
 *******************************************************************************/

#pragma once

#include <miopen/conv_algo_name.hpp>
#include <miopen/names.hpp>
#include <miopen/scalar.hpp>

#include <miopen/problem_description_base.hpp>
#include <miopen/convolution.hpp>

#if MIOPEN_ENABLE_SQLITE
#include <miopen/sqlite_db.hpp>
#endif

namespace miopen {

struct ExecutionContext;

MIOPEN_INTERNALS_EXPORT std::string
EncodeDataTypesForKey(miopenDataType_t in, miopenDataType_t weights, miopenDataType_t out);

template <class TElement>
constexpr auto GetDHW(unsigned spatial_dims, const std::vector<TElement>& data)
{
    if(spatial_dims == 2)
        return std::make_tuple(0, data[0], data[1]);
    return std::make_tuple(data[0], data[1], data[2]);
}

template <class TElement>
constexpr TElement GetD3(unsigned spatial_dims, const std::vector<TElement>& data)
{
    return std::get<0>(GetDHW(spatial_dims, data));
}

template <class TElement>
constexpr TElement GetH3(unsigned spatial_dims, const std::vector<TElement>& data)
{
    return std::get<1>(GetDHW(spatial_dims, data));
}

template <class TElement>
constexpr TElement GetW3(unsigned spatial_dims, const std::vector<TElement>& data)
{
    return std::get<2>(GetDHW(spatial_dims, data));
}
template <class TElement>
constexpr auto GetCHWN(const std::vector<TElement>& data)
{
    return miopen::tien<4>(data, 1);
}

template <class TElement>
constexpr TElement GetNofCHWN(const std::vector<TElement>& data)
{
    return std::get<3>(GetCHWN(data));
}

template <class TElement>
constexpr TElement GetCofCHWN(const std::vector<TElement>& data)
{
    return std::get<0>(GetCHWN(data));
}

template <class TElement>
constexpr TElement GetHofCHWN(const std::vector<TElement>& data)
{
    return std::get<1>(GetCHWN(data));
}

template <class TElement>
constexpr TElement GetWofCHWN(const std::vector<TElement>& data)
{
    return std::get<2>(GetCHWN(data));
}

namespace conv {

MIOPEN_INTERNALS_EXPORT miopenAlphaBetaCase_t ClassifyAlphaBeta(const Scalar& alpha,
                                                                const Scalar& beta);

struct ProblemDescription : ProblemDescriptionBase
#if MIOPEN_ENABLE_SQLITE
    ,
                            SQLiteSerializable<ProblemDescription>
#endif
{
    ProblemDescription() = default;

    /// \todo Get rid of the swapping of x and y.
    ProblemDescription(const TensorDescriptor& in_, // x for Forward, y for Backward*
                       const TensorDescriptor& weights_,
                       const TensorDescriptor& out_, // y for Forward, x for Backward*
                       const ConvolutionDescriptor& conv_,
                       Direction direction_,
                       int bias_            = 0,
                       const Scalar& alpha_ = Scalar(1.0),
                       const Scalar& beta_  = Scalar(0.0))
        : in(in_),
          weights(weights_),
          out(out_),
          conv(conv_),
          in_layout(ComputeInLayout()),
          weights_layout(ComputeWeightsLayout()),
          out_layout(ComputeOutLayout()),
          direction(direction_),
          bias(bias_),
          alpha(alpha_),
          beta(beta_),
          alpha_beta_case(ClassifyAlphaBeta(alpha, beta))
    {
        HeuristicUpdateLayouts();
    }

    // Conv descriptor getters
    unsigned GetSpatialDims() const { return conv.GetSpatialDimension(); }
    int GetPadD() const { return GetD3(GetSpatialDims(), conv.GetConvPads()); }
    int GetPadH() const { return GetH3(GetSpatialDims(), conv.GetConvPads()); }
    int GetPadW() const { return GetW3(GetSpatialDims(), conv.GetConvPads()); }
    int GetKernelStrideD() const { return GetD3(GetSpatialDims(), conv.GetConvStrides()); }
    int GetKernelStrideH() const { return GetH3(GetSpatialDims(), conv.GetConvStrides()); }
    int GetKernelStrideW() const { return GetW3(GetSpatialDims(), conv.GetConvStrides()); }
    int GetDilationD() const { return GetD3(GetSpatialDims(), conv.GetConvDilations()); }
    int GetDilationH() const { return GetH3(GetSpatialDims(), conv.GetConvDilations()); }
    int GetDilationW() const { return GetW3(GetSpatialDims(), conv.GetConvDilations()); }
    int GetGroupCount() const { return conv.GetGroupCount(); }
    int GetVectorLength() const { return in.GetVectorLength(); }

    // In getters
    miopenDataType_t GetInDataType() const { return in.GetType(); }
    std::optional<miopenDataType_t> GetInCastType() const { return in.GetCastType(); }
    std::size_t GetInBatchSize() const { return GetN5(GetSpatialDims(), in.GetLengths()); }
    std::size_t GetBatchSize() const { return GetInBatchSize(); } // alias of GetInBatchSize()
    std::size_t GetInChannels() const { return GetC5(GetSpatialDims(), in.GetLengths()); }
    std::size_t GetInDepth() const { return GetD5(GetSpatialDims(), in.GetLengths()); }
    std::size_t GetInHeight() const { return GetH5(GetSpatialDims(), in.GetLengths()); }
    std::size_t GetInWidth() const { return GetW5(GetSpatialDims(), in.GetLengths()); }
    std::size_t GetInBatchStride() const { return GetN5(GetSpatialDims(), in.GetStrides()); }
    std::size_t GetInChannelStride() const { return GetC5(GetSpatialDims(), in.GetStrides()); }
    std::size_t GetInStrideD() const { return GetD5(GetSpatialDims(), in.GetStrides()); }
    std::size_t GetInStrideH() const { return GetH5(GetSpatialDims(), in.GetStrides()); }
    std::size_t GetInStrideW() const { return GetW5(GetSpatialDims(), in.GetStrides()); }
    std::string GetInLayout() const { return in_layout; }
    std::size_t GetInElementSize() const { return GetTypeSize(GetInDataType()); }
    std::size_t GetInSize() const { return in.GetNumBytes(); }

    // Out getters
    miopenDataType_t GetOutDataType() const { return out.GetType(); }
    std::optional<miopenDataType_t> GetOutCastType() const { return out.GetCastType(); }
    std::size_t GetOutBatchSize() const { return GetN5(GetSpatialDims(), out.GetLengths()); }
    std::size_t GetOutChannels() const { return GetC5(GetSpatialDims(), out.GetLengths()); }
    std::size_t GetOutDepth() const { return GetD5(GetSpatialDims(), out.GetLengths()); }
    std::size_t GetOutHeight() const { return GetH5(GetSpatialDims(), out.GetLengths()); }
    std::size_t GetOutWidth() const { return GetW5(GetSpatialDims(), out.GetLengths()); }
    std::size_t GetOutBatchStride() const { return GetN5(GetSpatialDims(), out.GetStrides()); }
    std::size_t GetOutChannelStride() const { return GetC5(GetSpatialDims(), out.GetStrides()); }
    std::size_t GetOutStrideD() const { return GetD5(GetSpatialDims(), out.GetStrides()); }
    std::size_t GetOutStrideH() const { return GetH5(GetSpatialDims(), out.GetStrides()); }
    std::size_t GetOutStrideW() const { return GetW5(GetSpatialDims(), out.GetStrides()); }
    std::string GetOutLayout() const { return out_layout; }
    std::size_t GetOutElementSize() const { return GetTypeSize(GetOutDataType()); }
    std::size_t GetOutSize() const { return out.GetNumBytes(); }

    // Weights getters
    miopenDataType_t GetWeightsDataType() const { return weights.GetType(); }
    std::optional<miopenDataType_t> GetWeightsCastType() const { return weights.GetCastType(); }
    std::size_t GetWeightsDepth() const { return GetD5(GetSpatialDims(), weights.GetLengths()); }
    std::size_t GetWeightsHeight() const
    {
        if(weights_layout == "CHWNc")
            return GetHofCHWN(weights.GetLengths());
        else
            return GetH5(GetSpatialDims(), weights.GetLengths());
    }
    std::size_t GetWeightsWidth() const
    {
        if(weights_layout == "CHWNc")
            return GetWofCHWN(weights.GetLengths());
        else
            return GetW5(GetSpatialDims(), weights.GetLengths());
    }
    std::size_t GetWeightsStrideK() const { return GetN5(GetSpatialDims(), weights.GetStrides()); }
    std::size_t GetWeightsStrideC() const { return GetC5(GetSpatialDims(), weights.GetStrides()); }
    std::size_t GetWeightsStrideD() const { return GetD5(GetSpatialDims(), weights.GetStrides()); }
    std::size_t GetWeightsStrideH() const { return GetH5(GetSpatialDims(), weights.GetStrides()); }
    std::size_t GetWeightsStrideW() const { return GetW5(GetSpatialDims(), weights.GetStrides()); }
    std::string GetWeightsLayout() const { return weights_layout; }
    std::size_t GetWeightsElementSize() const { return GetTypeSize(GetWeightsDataType()); }
    std::size_t GetWeightsSize() const { return weights.GetNumBytes(); }

    const TensorDescriptor& GetIn() const { return in; }
    const TensorDescriptor& GetWeights() const { return weights; }
    const TensorDescriptor& GetOut() const { return out; }
    const ConvolutionDescriptor& GetConv() const { return conv; }

    Direction GetDirection() const { return direction; }
    bool IsDirectionForward() const { return direction == conv::Direction::Forward; }
    bool IsDirectionBackwardData() const { return direction == conv::Direction::BackwardData; }
    bool IsDirectionBackwardWrW() const { return direction == conv::Direction::BackwardWeights; }
    MIOPEN_INTERNALS_EXPORT std::string GetDirectionStr() const;

    const Scalar& GetAlpha() const { return alpha; }
    const Scalar& GetBeta() const { return beta; }

    miopenAlphaBetaCase_t GetAlphaBetaCase() const { return alpha_beta_case; }

    std::string GetAlphaBetaCaseStr() const;

    int GetBias() const { return bias; }

    std::size_t GetBiasSize() const
    {
        return (GetBias() != 0) ? (GetOutChannels() * GetOutElementSize()) : 0;
    }

    int64_t GetBackwardPadW() const
    {
        return static_cast<int64_t>(GetWeightsWidth()) - GetPadW() - 1;
    }
    int64_t GetBackwardPadH() const
    {
        return static_cast<int64_t>(GetWeightsHeight()) - GetPadH() - 1;
    }

    bool IsAsymmetricPadH() const
    {
        return conv.paddingMode == miopenPaddingSame && (GetWeightsHeight() % 2) == 0;
    }
    bool IsAsymmetricPadW() const
    {
        return conv.paddingMode == miopenPaddingSame && (GetWeightsWidth() % 2) == 0;
    }

    bool Is2d() const { return GetSpatialDims() == 2; }
    bool Is3d() const { return GetSpatialDims() == 3; }

    bool IsFp32() const
    {
        return GetInDataType() == miopenFloat && GetWeightsDataType() == miopenFloat &&
               GetOutDataType() == miopenFloat;
    }
    bool IsFp16() const
    {
        return GetInDataType() == miopenHalf && GetWeightsDataType() == miopenHalf &&
               GetOutDataType() == miopenHalf;
    }
    bool IsBfp16() const
    {
        return GetInDataType() == miopenBFloat16 && GetWeightsDataType() == miopenBFloat16 &&
               GetOutDataType() == miopenBFloat16;
    }
    bool IsInt8() const
    {
        return GetInDataType() == miopenInt8 && GetWeightsDataType() == miopenInt8 &&
               (GetOutDataType() == miopenInt32 || GetOutDataType() == miopenInt8 ||
                GetOutDataType() == miopenFloat);
    }
    bool IsFp8() const
    {
        return GetInDataType() == miopenFloat8_fnuz || GetWeightsDataType() == miopenFloat8_fnuz ||
               GetOutDataType() == miopenFloat8_fnuz;
    }
    bool IsBfp8() const
    {
        return GetInDataType() == miopenBFloat8_fnuz ||
               GetWeightsDataType() == miopenBFloat8_fnuz || GetOutDataType() == miopenBFloat8_fnuz;
    }
    bool IsTensorsCasted() const
    {
        return GetInCastType() || GetWeightsCastType() || GetOutCastType();
    }

    bool UseTF32() const { return use_tf32; }

    // To be used in Solvers that do not implement ALT FP16 kernels.
    // Those Solvers must be non-applicable for gfx90a when this function returns true.
    bool IsGfx90aFp16altRequired() const
    {
        if(!IsFp16())
            return false;
        if(direction == conv::Direction::Forward)
            return conv.attribute.gfx90aFp16alt.GetFwd();
        if(direction == conv::Direction::BackwardData)
            return conv.attribute.gfx90aFp16alt.GetBwd();
        if(direction == conv::Direction::BackwardWeights)
            return conv.attribute.gfx90aFp16alt.GetWrW();
        MIOPEN_THROW("Direction must be known!");
    }

    MIOPEN_INTERNALS_EXPORT bool IsLayoutDefault() const;
    MIOPEN_INTERNALS_EXPORT bool IsLayoutNHWC() const;
    MIOPEN_INTERNALS_EXPORT bool IsLayoutNCHWc() const;
    bool IsNCHWc_NCHWc() const;
    bool IsNCHWc_CHWNc() const;

    bool HasNonPackedTensors() const
    {
        return !(in.IsPacked() && weights.IsPacked() && out.IsPacked());
    }

    bool HasMixedDataTypes() const
    {
        return !(GetInDataType() == GetWeightsDataType() &&
                 GetWeightsDataType() == GetOutDataType());
    }

    bool AllTensorsDimsFitIntoInt() const
    {
        return in.AllDimsFitIntoInt() && weights.AllDimsFitIntoInt() && out.AllDimsFitIntoInt();
    }

    bool AllTensorsLengthsFitIntoInt() const
    {
        return in.AllLengthsFitIntoInt() && weights.AllLengthsFitIntoInt() &&
               out.AllLengthsFitIntoInt();
    }

    MIOPEN_INTERNALS_EXPORT void HeuristicUpdateLayouts();

    MIOPEN_INTERNALS_EXPORT void MakeNetworkConfig(std::string& conf_key) const;

    NetworkConfig MakeNetworkConfig() const override
    {
        std::string ret;
        MakeNetworkConfig(ret);
        return NetworkConfig{ret};
    }

    // Todo: remove after fixing fin
    [[deprecated]] NetworkConfig BuildConfKey() const { return MakeNetworkConfig(); }

    MIOPEN_INTERNALS_EXPORT void Serialize(std::ostream& stream) const;

    friend std::ostream& operator<<(std::ostream& os, const ProblemDescription& obj)
    {
        obj.Serialize(os);
        return os;
    }

    template <class Self>
    static void Visit(Self&& self, std::function<void(int64_t, std::string)> f)
    {
        // The column names match the driver command line argument names
        f(self.GetSpatialDims(), "spatial_dim");
        f(self.GetInChannels(), "in_channels");
        f(self.GetInHeight(), "in_h");
        f(self.GetInWidth(), "in_w");
        f(self.GetInDepth(), "in_d");
        f(self.GetWeightsHeight(), "fil_h");
        f(self.GetWeightsWidth(), "fil_w");
        f(self.GetWeightsDepth(), "fil_d");
        f(self.GetOutChannels(), "out_channels");
        f(self.GetBatchSize(), "batchsize");
        f(self.GetPadH(), "pad_h");
        f(self.GetPadW(), "pad_w");
        f(self.GetPadD(), "pad_d");
        f(self.GetKernelStrideH(), "conv_stride_h");
        f(self.GetKernelStrideW(), "conv_stride_w");
        f(self.GetKernelStrideD(), "conv_stride_d");
        f(self.GetDilationH(), "dilation_h");
        f(self.GetDilationW(), "dilation_w");
        f(self.GetDilationD(), "dilation_d");
        f(self.GetBias(), "bias");
        f(self.GetGroupCount(), "group_count");
    }

    template <class Self>
    static void Visit(Self&& self, std::function<void(std::string, std::string)> f)
    {
        f(self.GetInLayout(), "layout");
        std::string data_type = EncodeDataTypesForKey(
            self.GetInDataType(), self.GetWeightsDataType(), self.GetOutDataType());
        f(data_type, "data_type");
        f(self.GetDirectionStr(), "direction");
        if(data_type == "FP32" && self.UseTF32())
            f("TF32", "compute_datatype");
    }

    template <class Self, class Visitor>
    static void VisitAll(Self&& self, const Visitor& f)
    {
        Visit(std::forward<Self>(self), [&](int64_t value, std::string name) { f(value, name); });
        Visit(std::forward<Self>(self),
              [&](std::string value, std::string name) { f(value, name); });
    }

    MIOPEN_INTERNALS_EXPORT void SetupFloats(ExecutionContext& ctx) const;
    MIOPEN_INTERNALS_EXPORT void SetupComputeType(const ExecutionContext& ctx) const;

private:
    std::string ComputeLayout(const TensorDescriptor& td) const;
    MIOPEN_INTERNALS_EXPORT std::string ComputeInLayout() const;
    MIOPEN_INTERNALS_EXPORT std::string ComputeOutLayout() const;
    MIOPEN_INTERNALS_EXPORT std::string ComputeWeightsLayout() const;

    TensorDescriptor in;
    TensorDescriptor weights;
    TensorDescriptor out;
    ConvolutionDescriptor conv;
    std::string in_layout;
    std::string weights_layout;
    std::string out_layout;
    Direction direction                   = Direction::Forward;
    int bias                              = 0;
    Scalar alpha                          = Scalar(1.0);
    Scalar beta                           = Scalar(0.0);
    miopenAlphaBetaCase_t alpha_beta_case = DEFAULT;
    mutable bool use_tf32                 = false;
};

inline bool IsPointOutput3dStrideEqFilter(const ProblemDescription& problem,
                                          Direction direction,
                                          bool require_input_spatial_eq_filter)
{
    if(problem.GetDirection() != direction)
        return false;

    const auto& conv = problem.GetConv();
    if(conv.GetSpatialDimension() != 3 || conv.group_count != 1)
        return false;

    const auto& in_desc    = problem.GetIn();
    const auto& out_desc   = problem.GetOut();
    const auto& point_desc = direction == Direction::Forward ? out_desc : in_desc;
    const auto& w_desc     = problem.GetWeights();

    const auto& point_lens = point_desc.GetLengths();
    if(point_lens.size() != 5 || point_lens[2] != 1 || point_lens[3] != 1 || point_lens[4] != 1)
        return false;

    const auto& in_lens = in_desc.GetLengths();

    const auto& pads      = conv.GetConvPads();
    const auto& strides   = conv.GetConvStrides();
    const auto& dilations = conv.GetConvDilations();
    const auto& w_lens    = w_desc.GetLengths();
    if(w_lens.size() != 5 || pads.size() != 3 || strides.size() != 3 || dilations.size() != 3)
        return false;

    for(int i = 0; i < 3; ++i)
    {
        if(pads[i] != 0 || dilations[i] != 1)
            return false;
        if(static_cast<int>(w_lens[2 + i]) != strides[i])
            return false;
        if(require_input_spatial_eq_filter &&
           static_cast<int>(in_lens[2 + i]) != static_cast<int>(w_lens[2 + i]))
            return false;
    }

    return true;
}

inline bool IsBwdDataPointOutput3dStrideEqFilter(const ProblemDescription& problem)
{
    return IsPointOutput3dStrideEqFilter(problem, Direction::BackwardData, false);
}

inline bool IsFwdDataPointOutput3dStrideEqFilter(const ProblemDescription& problem)
{
    // For direct GEMM without Im2Col, input spatial must equal filter spatial.
    return IsPointOutput3dStrideEqFilter(problem, Direction::Forward, true);
}

} // namespace conv
} // namespace miopen
