/* Copyright (C) 2013 Ion Torrent Systems, Inc. All Rights Reserved */

#include <iostream>
#include <algorithm>
#include "DfcCompr.h"
//#include "Utils.h"
#include <malloc.h>
#include <cmath>
// DFC doesn't use anything from the Eigen library
//#define EIGEN_USE_MKL_ALL 1
//#include <Eigen/Dense>
//#include <Eigen/LU>

using namespace std;
//using namespace Eigen;

/***
 The DFT engine needs to be kept in sync with any changes to the number of frames.
*/
int DfcCompr::SetNumFrames(int frames)
{
    if ((frames > 0) && (frames != n_frames))
    {
        n_frames = frames;
        DFT.SetPts(frames);
    }

    return n_frames;
}

/***
 Range check the new window length value and generate the new window coefficients.
 Window lengths should be relatively short.
*/
int DfcCompr::SetWindowLength(int length)
{
    if ((length >= 0) && (length < (n_frames/4)) && (length != winParamLen))
    {
        winParamLen = length;
        GenerateWindow();
    }

    return winParamLen;
}

/***
 Range check the new window alpha value and generat the new window coefficients.
*/
float DfcCompr::SetWindowAlpha(float alpha)
{
    if ((alpha >= 0.0f) && (alpha <= 3.0f) && (alpha != winParamAlpha))
    {
        winParamAlpha = alpha;
        GenerateWindow();
    }

    return winParamAlpha;
}

/***
 Generate a new set of Gaussian window coefficients based on the current set of
 window length and alpha configuration parameters.
*/
void DfcCompr::GenerateWindow()
{
    // setup vectors to store compression and decompression windows
    windowCompress.resize(winParamLen*2);
    windowExpand.resize(winParamLen*2);

    float halfN = static_cast<float>(winParamLen);
    float n     = 0.5f-halfN;  // starting at -(N-1)/2
    int   idx, maxIdx;
    float an_sq, val, oneOnVal;

    maxIdx = (winParamLen * 2) - 1;
    for (idx = 0; idx < winParamLen; ++idx)
    {
        an_sq    = winParamAlpha * (n / halfN);
        val      = exp((-0.5f) * (an_sq * an_sq));
        oneOnVal = 1.0f / val;
        n       += 1.0f;  // setup n for next loop iteration
        windowCompress[idx]        = val;
        windowCompress[maxIdx-idx] = val;  // exploit window symmetry
        windowExpand[idx]          = oneOnVal;
        windowExpand[maxIdx-idx]   = oneOnVal;
    }

}

/***
  Decompose the image data into a set of basis vectors and each well's projection
  onto them.
  @param n_wells - number of wells in the image patch
  @param n_frame - number of frames in this image patch.
  @param image - Same order as RawImage structure. Individual wells of data in frame, row, col major order so
                 value(row_i,col_j,frame_k) = image[row_i * ncol + col_j + (nrow * ncol * frame_k)]
  @param n_sample_wells - number of wells in the sample of the patch
  @param image_sample - sample of image above for vectors and such
  @param compressed - output of a lossy compressed patch
*/
void DfcCompr::LossyCompress(float *image)
{
    //////////////////////////////////////////
    // Temporary scratch_pad memory allocation
    if (scratch_pad == NULL) scratch_pad = new float[n_wells*n_basis*2];
    //////////////////////////////////////////

    int wellIdx, frameIdx;
    float tmpWell[n_frames];
    register float tmpSum;
    register int offset, winIdx;
    float * __restrict dftPtrI = &scratch_pad[0];
    float * __restrict dftPtrQ = &scratch_pad[n_wells*n_basis];
    float deltaMagI[n_basis];
    float deltaMagQ[n_basis];

    // Initialize keyframe
    for (frameIdx = 0; frameIdx < n_basis; ++frameIdx)
    {
        keyFrameI[frameIdx] = 0.0f;
        keyFrameQ[frameIdx] = 0.0f;
    }

    for (wellIdx = 0, offset = 0; wellIdx < n_wells; ++wellIdx)
    {
        // Step 1a - extract the next well from the image cube
        tmpSum = 0.0f;
        winIdx = 0;
        for (frameIdx = 0; frameIdx < n_frames; ++frameIdx)
        {
            tmpWell[frameIdx] = image[(frameIdx*n_wells)+wellIdx];
            tmpSum           += image[(frameIdx*n_wells)+wellIdx];
        }
        // Step 1b - remove DC offset and apply windowing to each well
        tmpSum /= static_cast<float>(n_frames);  // DC offset
        // Apply front part of window
        for (frameIdx = 0; frameIdx < winParamLen; ++frameIdx)
        {
            tmpWell[frameIdx] -= tmpSum;
            tmpWell[frameIdx] *= windowCompress[winIdx++];
        }
        // Middle part (window is unity)
        for (; frameIdx < (n_frames-winParamLen); ++frameIdx)
            tmpWell[frameIdx] -= tmpSum;
        // Apply back part of window
        for (; frameIdx < n_frames; ++frameIdx)
        {
            tmpWell[frameIdx] -= tmpSum;
            tmpWell[frameIdx] *= windowCompress[winIdx++];
        }
        // Step 2 - Compute partial DFT
        // store for deltas and accumulate for keyframe
        winIdx = n_basis * wellIdx;
        DFT.PartialDFT(1, static_cast<unsigned int>(n_basis), &tmpWell[0], &dftPtrI[winIdx], &dftPtrQ[winIdx]);
        for (frameIdx = 0; frameIdx < n_basis; ++frameIdx, ++offset)
        {
            keyFrameI[frameIdx] += dftPtrI[offset];
            keyFrameQ[frameIdx] += dftPtrQ[offset];
        }
    }

    // Step 3 - Convert accumulated spectrum into mean keyframe and initialize
    // maximum delta magnitude vector
    for (frameIdx = 0; frameIdx < n_basis; ++frameIdx)
    {
        keyFrameI[frameIdx] /= static_cast<float>(n_wells);
        keyFrameQ[frameIdx] /= static_cast<float>(n_wells);
        deltaMagI[frameIdx]  = 0.0f;
        deltaMagQ[frameIdx]  = 0.0f;
    }

    // Step 4 - Calculate the correlation statistics and emphasis vector needed to
    // populate the bits per frequency element vector
    Emphasis();

    // Step 5 - Convert DFT data from absolute to raw delta, keeping maximum
    // component magnitude observed for each frequency component along the way
    for (wellIdx = 0, offset = 0; wellIdx < n_wells; ++wellIdx)
        // May be able to use SSE/AVX vector acceleration here to process multiple frequency elements in parallel
        for (frameIdx = 0; frameIdx < n_basis; ++frameIdx, ++offset)
        {
            dftPtrI[offset] -= keyFrameI[frameIdx];
            dftPtrQ[offset] -= keyFrameQ[frameIdx];
            if (deltaMagI[frameIdx] < abs(dftPtrI[offset])) deltaMagI[frameIdx] = abs(dftPtrI[offset]);  // update maximum observed real magnitude vector
            if (deltaMagQ[frameIdx] < abs(dftPtrQ[offset])) deltaMagQ[frameIdx] = abs(dftPtrQ[offset]);  // update maximum observed imaginary magnitude vector
        }

    // Populate the scale vectors
    for (frameIdx = 0; frameIdx < n_basis; ++frameIdx)
    {
        tmpSum = static_cast<float>((1<<static_cast<int>(bitsPerFreq[frameIdx]-1))-1);
        scaleVectorI[frameIdx] = deltaMagI[frameIdx] / tmpSum;
        scaleVectorQ[frameIdx] = deltaMagQ[frameIdx] / tmpSum;
    }

    // Step 6 - Quantize the spectrum delta values and pack into output vectors
    for (wellIdx = 0, offset = 0; wellIdx < n_wells; ++wellIdx)
        // May be able to use SSE/AVX vector acceleration here to process multiple frequency elements in parallel
        for (frameIdx = 0; frameIdx < n_basis; ++frameIdx, ++offset)
        {
            deltaI[offset] = static_cast<short>(round(dftPtrI[offset]/scaleVectorI[frameIdx]));
            deltaQ[offset] = static_cast<short>(round(dftPtrQ[offset]/scaleVectorQ[frameIdx]));
        }

    ////////////////////////////////////////////
    // Temporary scratch_pad memory deallocation
    delete[] scratch_pad; scratch_pad = NULL;
    ////////////////////////////////////////////
}

/***
*/
void DfcCompr::Emphasis()
{
    float kernelI[n_basis];
    float kernelQ[n_basis];
    float *dftPtrI = &scratch_pad[0];
    float *dftPtrQ = &scratch_pad[n_wells*n_basis];
    float *corrI = new float[n_wells*n_basis];  // probably a good idea to make this static or semi-static
    float *corrQ = new float[n_wells*n_basis];  // probably a good idea to make this static or semi-static
    register int offset, freqIdx;
    int   wellIdx;
    float meanMag[n_basis];
    float meanAng[n_basis];
    float varMag[n_basis];
    float varAng[n_basis];
    register float delta;
    float emphasisVector[n_basis];
    register int bitVal;

    // Generate correlation kernel from keyframe and initialize statistics vectors
    for (freqIdx = 0; freqIdx < n_basis; ++freqIdx)
    {
        kernelI[freqIdx] = keyFrameI[freqIdx];
        kernelQ[freqIdx] = -keyFrameQ[freqIdx];
        meanMag[freqIdx] = 0.0f;
        meanAng[freqIdx] = 0.0f;
        varMag[freqIdx]  = 0.0f;
        varAng[freqIdx]  = 0.0f;
    }

    // Correlate each well spectrum with the complex conjugate of the mean well spectrum
    for (wellIdx = 0, offset = 0; wellIdx < n_wells; ++wellIdx)
        // May be able to vectorize the correlation process
        for (freqIdx = 0; freqIdx < n_basis; ++freqIdx, ++offset)
        {
            // DFT data        ==> (a+ib)
            // kernel          ==> (x+ib)
            // correlated data ==> (a+ib)(x+iy) = ax+iay+ibx-by = (ax-by)+i(ay+bx)
            corrI[offset] = dftPtrI[offset]*kernelI[freqIdx]-dftPtrQ[offset]*kernelQ[freqIdx];  // (ax-by)
            corrQ[offset] = dftPtrI[offset]*kernelQ[freqIdx]+dftPtrQ[offset]*kernelI[freqIdx];  // (ay+bx)
        }

    // Find mean of magnitude and angle of each frequency element
    for (wellIdx = 0, offset = 0; wellIdx < n_wells; ++wellIdx)
        // May be able to vectorize the frequency elements
        for (freqIdx = 0; freqIdx < n_basis; ++freqIdx, ++offset)
        {
            meanMag[freqIdx] += sqrt((corrI[offset]*corrI[offset])+(corrQ[offset]*corrQ[offset]));
            meanAng[freqIdx] += atan2(corrQ[offset], corrI[offset]);
        }
    for (freqIdx = 0; freqIdx < n_basis; ++freqIdx)
    {
        meanMag[freqIdx] /= static_cast<float>(n_wells);
        meanAng[freqIdx] /= static_cast<float>(n_wells);
    }

    // Find standard deviation (now that we have the mean)
    for (wellIdx = 0, offset = 0; wellIdx < n_wells; ++wellIdx)
        // May be able to vectorize the frequency elements
        for (freqIdx = 0; freqIdx < n_basis; ++freqIdx, ++offset)
        {
            delta = sqrt((corrI[offset]*corrI[offset])+(corrQ[offset]*corrQ[offset])) - meanMag[freqIdx];
            varMag[freqIdx] += delta * delta;  // accumulate magnitude variance
            delta = atan2(corrQ[offset], corrI[offset]) - meanAng[freqIdx];
            varAng[freqIdx] += delta * delta;  // accumulate angle variance
        }

    // Finished with correlation data (don't need to release once we make this block of memory static/semi-static)
    delete[] corrI;
    delete[] corrQ;

    // Convert accumulated variance to raw emphasis (original Matlab calculation
    // uses standard deviation rather than variance, but that uses an additional
    // sqrt step that has been optimized out here)
    for (freqIdx = 0, delta = 0.0f; freqIdx < n_basis; ++freqIdx)
    {
        emphasisVector[freqIdx] = sqrt((varMag[freqIdx]*varAng[freqIdx])/static_cast<float>(n_wells));
        if (emphasisVector[freqIdx] > delta)
            delta = emphasisVector[freqIdx];  // keep the maximum value for normalization
    }

    // Normalize the emphasis vector and convert it to bits per frequency element
    for (freqIdx = 0; freqIdx < n_basis; ++freqIdx)
    {
        // calculate the number of magnitude bits to allocate (log will always return zero or a negative value)
        bitVal = static_cast<int>(static_cast<float>(n_maxBits)+ceil(log2(emphasisVector[freqIdx]/delta)));
        if (bitVal < n_minBits) bitVal = n_minBits;  // ensure at least the minimum number of magnitude bits are allocated
        ++bitVal;  // add the sign bit
        bitsPerFreq[freqIdx] = static_cast<unsigned char>(bitVal);
    }
}

/***
 Reconstruct the image data from the compressed frequency domain vectors.
 @param n_wells - number of wells in the image patch
 @param n_frame - number of frames in this image patch.
 @param image   - Same order as RawImage structure. Individual wells of data in frame, row, col major order so
                  value(row_i,col_j,frame_k) = image[row_i * ncol + col_j + (nrow * ncol * frame_k)]
*/
void DfcCompr::LossyUncompress(float *image)
{
    float wellBufI[n_basis];
    float wellBufQ[n_basis];
    float tmpBuf[n_frames];
    int   wellIdx, frameIdx;
    register int offset, winIdx;

    for (wellIdx = 0, offset = 0; wellIdx < n_wells; ++wellIdx)
    {
        // Reconstruct partial spectrum from keyframe, delta and scale information
        for (frameIdx = 0; frameIdx < n_basis; ++frameIdx, ++offset)
        {
            wellBufI[frameIdx] = keyFrameI[frameIdx] + scaleVectorI[frameIdx]*static_cast<float>(deltaI[offset]);
            wellBufQ[frameIdx] = keyFrameQ[frameIdx] + scaleVectorQ[frameIdx]*static_cast<float>(deltaQ[offset]);
        }
        // Reconstruct time domain signal from partial spectrum using IDFT
        DFT.PartialIDFT(1, n_basis, &tmpBuf[0], &wellBufI[0], &wellBufQ[0]);
        // Store reconstructed values into the image cube
        for (frameIdx = 0, winIdx = 0; frameIdx < winParamLen; ++frameIdx)
            image[(frameIdx*n_wells)+wellIdx] = tmpBuf[frameIdx] * windowExpand[winIdx++];
        for (; frameIdx < (n_frames-winParamLen); ++frameIdx)
            image[(frameIdx*n_wells)+wellIdx] = tmpBuf[frameIdx];
        for (; frameIdx < n_frames; ++frameIdx)
            image[(frameIdx*n_wells)+wellIdx] = tmpBuf[frameIdx] * windowExpand[winIdx++];
    }
}

/*****************************************************************************/

/***
 Map the memory blocks and call LossyCompress.
*/
void DfcComprWrapper::LossyCompress(float *image)
{
    SetupPointers();
    dfc.LossyCompress(image);
};

/***
 Map the memory blocks and call LossyUncompress.
*/
void DfcComprWrapper::LossyUncompress(float *image)
{
    SetupPointers();
    dfc.LossyUncompress(image);
};

/***
 Do some evil pointer arithmetic to map the DFC header and payload data vectors
 into the basis_vectors and coefficients memory blocks.
*/
void DfcComprWrapper::SetupPointers()
{
    void *hdrPtr = reinterpret_cast<void *>(basis_vectors);
    void *dtrPtr = reinterpret_cast<void *>(coefficients);
    int   offset = sizeof(float) * dfc.n_basis;

    // Map the header vectors into the basis_vectors memory block
    dfc.keyFrameI    = reinterpret_cast<float *>(hdrPtr);
    dfc.keyFrameQ    = reinterpret_cast<float *>(hdrPtr) + offset;
    dfc.scaleVectorI = reinterpret_cast<float *>(hdrPtr) + 2*offset;
    dfc.scaleVectorQ = reinterpret_cast<float *>(hdrPtr) + 3*offset;
    dfc.bitsPerFreq  = reinterpret_cast<unsigned char *>(hdrPtr) + 4*offset;

    // Map the delta data into the coefficients memory block
    offset     = sizeof(short) * dfc.n_basis * dfc.n_wells;
    dfc.deltaI = reinterpret_cast<short *>(dtrPtr);
    dfc.deltaQ = reinterpret_cast<short *>(dtrPtr) + offset;
}


/*******************************************************************************
*
* 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.
*
*******************************************************************************/

#include <miopen/solver.hpp>
#include <miopen/env.hpp>
#include <miopen/stringutils.hpp>
#include <miopen/sequences.hpp>
#include <miopen/kernel_build_params.hpp>
#include <miopen/generic_search.hpp>
#include <miopen/conv/data_invoke_params.hpp>
#include <miopen/conv/compiled_in_parameters.hpp>
#include <miopen/conv/wrw_invoke_params.hpp>

#include <boost/any.hpp>

#include <tuple>

MIOPEN_DECLARE_ENV_VAR(MIOPEN_DEBUG_AMD_WINOGRAD_RXS_F2X3)
MIOPEN_DECLARE_ENV_VAR(MIOPEN_DEBUG_AMD_WINOGRAD_RXS_F2X3_PERF_VALS)

#define WINODATA 2
#define WINOFILTER 3
#define MAX_CU_LIMIT 512

static inline size_t Ceil(const size_t v, const size_t m)
{
    assert(m > 0);
    return (v + m - 1) / m;
}

static inline size_t quantize_up(size_t val, size_t factor) { return Ceil(val, factor) * factor; }

static inline int GetBestNGroupParam(const int R,
                                     const int S,
                                     const int R_stride,
                                     const int S_stride,
                                     const int C,
                                     const int K,
                                     const int OH,
                                     const int OW,
                                     const int pad_H,
                                     const int pad_W,
                                     const int N,
                                     const int idilation_w,
                                     const int idilation_h,
                                     const int n_groups,
                                     const int G)
{
    int o_tile     = WINODATA;
    int f_tile     = WINOFILTER;
    int r_factor   = f_tile * 2;
    int s_factor   = r_factor;
    int c_factor   = 2;
    int k_factor   = 32;
    int nwh_factor = 32;
    int w_factor   = o_tile * idilation_w * S_stride;
    int h_factor   = o_tile * idilation_h * R_stride;

    if(S_stride == 1 && idilation_w == 1 && S <= f_tile)
        s_factor = f_tile;
    if((R_stride == 1 && idilation_h == 1) || (R % (f_tile * 2)) == 1)
        r_factor = f_tile;
    if(S_stride == 2 || R_stride == 2 || idilation_w == 2 || idilation_h == 2)
        c_factor = 1;

    size_t g_s = quantize_up(S, s_factor);
    size_t g_r = quantize_up(R, r_factor);
    size_t g_c = quantize_up(C, c_factor);
    size_t g_k = quantize_up(K, k_factor);
    size_t g_w = OW;
    size_t g_h = OH;

    if((pad_W % 2 == 0) && (idilation_w > 1 || S_stride > 1))
        g_w += 1;
    if((pad_H % 2 == 1) && (idilation_h > 1 || R_stride > 1))
        g_h += 1;

    g_w            = quantize_up(g_w, w_factor);
    g_h            = quantize_up(g_h, h_factor);
    size_t g_n_w_h = quantize_up(g_w * g_h * N, nwh_factor * w_factor * h_factor);

    int best_n_groups_cnt = 1;
    double min_param      = 0;
    for(auto i = 1; i < n_groups; ++i)
    {
        size_t g_n_w_h_k =
            quantize_up(g_n_w_h * g_k, nwh_factor * w_factor * h_factor * k_factor * i);
        size_t granulated_mac_count = g_n_w_h_k * g_c * g_s * g_r;
        size_t n_groups_per_cu      = Ceil(i * G, n_groups);
        double perf_metric = static_cast<double>(n_groups_per_cu) * granulated_mac_count / i;
        if(static_cast<double>(granulated_mac_count) / i > 1.0e+7)
            perf_metric *= (1 + i * 0.003);
        else
            perf_metric *= (1 + i * 0.04);
        if(i == 1)
            min_param = perf_metric;
        if(min_param > perf_metric)
        {
            best_n_groups_cnt = i;
            min_param         = perf_metric;
        }
    }
    return best_n_groups_cnt;
}

namespace miopen {
namespace solver {

namespace {
// clang-format off
    auto PerfFieldRules()
    {
        return seq::MakeRuleSet(
            std::make_tuple(seq::Span<int, 1, MAX_CU_LIMIT>{}, &PerformanceConfigConvBinWinogradRxSf2x3::n_groups)
        );
    }
// clang-format on

/// \todo Consider re-using code from RxS.
inline bool IsShaderContraintsMet(const int R,
                                  const int S,
                                  const int,
                                  const int,
                                  const int C,
                                  const int K,
                                  const int H,
                                  const int W,
                                  const int OH,
                                  const int OW,
                                  const int N,
                                  const ConvolutionContext& params)
{
    // Padding for bwd data shall not be negative.
    /// \todo Either remove WrW related code or re-use function from RxS
    if(params.direction.IsBackwardData() || params.direction.IsBackwardWrW())
    {
        if(!(0 <= params.GetBackwardPadW() && params.GetBackwardPadW() < std::pow(2, 16)))
            return false;
        if(!(0 <= params.GetBackwardPadH() && params.GetBackwardPadH() < std::pow(2, 16)))
            return false;
    }
    const auto grid_workgroup_count_x = params.GetStream().GetMaxComputeUnits();
    assert(params.weights_layout.length() == 0);
    // clang-format off
        // Check implementation limits.
        return N < std::pow(2, 16)
            && C < std::pow(2, 16)
            && K < std::pow(2, 16)
            && H < std::pow(2, 16)
            && W < std::pow(2, 16)
            && OH < std::pow(2, 16)
            && OW < std::pow(2, 16)
            && params.pad_w < std::pow(2, 16)
            && params.pad_h < std::pow(2, 16)
            && S < std::pow(2, 16)
            && R < std::pow(2, 16)
            && grid_workgroup_count_x < std::pow(2, 16)
            && (C * H * W) <= std::pow(2, 28)
            && (OH * OW) <= std::pow(2, 23)
            && (K * OH * OW) <= std::pow(2, 28)
            && (K * R * S) <= std::pow(2, 28)
            && (C * R * S) <= std::pow(2, 28);
    // clang-format on
}

} // namespace

PerformanceConfigConvBinWinogradRxSf2x3::PerformanceConfigConvBinWinogradRxSf2x3(int n_groups_)
    : n_groups(n_groups_)
{
}

void PerformanceConfigConvBinWinogradRxSf2x3::EuristicInit(const ConvolutionContext& config)
{
    const auto n_inputs_per_group  = config.n_inputs / config.group_counts,
               n_outputs_per_group = config.n_outputs / config.group_counts;
    if(config.group_counts == 1)
    {
        n_groups = config.GetStream().GetMaxComputeUnits();
        return;
    }

    if(config.direction.IsBackwardWrW())
    {
        n_groups = GetBestNGroupParam(config.in_height,
                                      config.in_width,
                                      config.kernel_dilation_h,
                                      config.kernel_dilation_w,
                                      config.batch_sz,    // N
                                      n_inputs_per_group, // K
                                      config.kernel_size_h,
                                      config.kernel_size_w,
                                      config.pad_w,
                                      config.pad_h,
                                      n_outputs_per_group, // C
                                      config.kernel_stride_h,
                                      config.kernel_stride_w,
                                      config.GetStream().GetMaxComputeUnits(),
                                      config.group_counts);
    }
    else
    {
        n_groups = GetBestNGroupParam(config.kernel_size_h, // RxS
                                      config.kernel_size_w,
                                      config.kernel_stride_h,
                                      config.kernel_stride_w,
                                      n_inputs_per_group,  // C
                                      n_outputs_per_group, // K
                                      config.out_height,   // OHxOW
                                      config.out_width,
                                      config.pad_w,
                                      config.pad_h,
                                      config.batch_sz, // N
                                      config.kernel_dilation_h,
                                      config.kernel_dilation_w,
                                      config.GetStream().GetMaxComputeUnits(),
                                      config.group_counts);
    }
}

bool PerformanceConfigConvBinWinogradRxSf2x3::SetNextValue()
{
    return !PerfFieldRules().Next(*this);
}

bool PerformanceConfigConvBinWinogradRxSf2x3::IsValidValue() const
{
    return PerfFieldRules().IsIn(*this);
}

bool PerformanceConfigConvBinWinogradRxSf2x3::IsValid(const ConvolutionContext& config) const
{
    if(config.GetStream().GetMaxComputeUnits() < n_groups)
        return false;

    if(!IsValidValue())
        return false;
    return true;
}

inline bool PerformanceConfigConvBinWinogradRxSf2x3::
operator==(const PerformanceConfigConvBinWinogradRxSf2x3& other) const
{
    return n_groups == other.n_groups;
}

std::string PerformanceConfigConvBinWinogradRxSf2x3::ToString() const
{
    std::ostringstream ss;
    Serialize(ss);
    return ss.str();
}

PerformanceConfigConvBinWinogradRxSf2x3
ConvBinWinogradRxSf2x3::GetPerformanceConfig(const ConvolutionContext& params) const
{
    PerformanceConfigConvBinWinogradRxSf2x3 pp;
    pp.EuristicInit(params);
    MIOPEN_LOG_I(pp.ToString());
    return pp;
}

bool ConvBinWinogradRxSf2x3::IsValidPerformanceConfig(
    const ConvolutionContext& problem, const PerformanceConfigConvBinWinogradRxSf2x3& c) const
{
    return c.IsValidValue() && c.IsValid(problem);
}

PerformanceConfigConvBinWinogradRxSf2x3
ConvBinWinogradRxSf2x3::Search(const ConvolutionContext& context) const
{
    if(context.direction.IsForward())
        return GenericSearchFwd(*this, context);
    else
        return GenericSearchBwd(*this, context);
}

inline void FillVarsFromConfig(int& H,
                               int& W,
                               int& R,
                               int& S,
                               int& R_stride,
                               int& S_stride,
                               int& C,
                               int& K,
                               int& out_H,
                               int& out_W,
                               int& pad_H,
                               int& pad_W,
                               int& N,
                               int& idilation_w,
                               int& idilation_h,
                               int& n_groups,
                               int& group_cnt,
                               const ConvolutionContext& config)
{
    group_cnt   = config.group_counts;
    n_groups    = config.GetStream().GetMaxComputeUnits();
    pad_H       = config.direction.IsForward() ? config.pad_h : config.GetBackwardPadH();
    pad_W       = config.direction.IsForward() ? config.pad_w : config.GetBackwardPadW();
    H           = config.in_height;
    W           = config.in_width;
    R           = config.kernel_size_h;
    S           = config.kernel_size_w;
    R_stride    = config.kernel_stride_h;
    S_stride    = config.kernel_stride_w;
    C           = config.n_inputs;
    K           = config.n_outputs;
    out_H       = config.out_height;
    out_W       = config.out_width;
    N           = config.batch_sz;
    idilation_w = config.kernel_dilation_h;
    idilation_h = config.kernel_dilation_w;
}

template <typename B, typename T, typename TW>
int ConvBinWinogradRxSf2x3::RunAndMeasureSolution(const miopen::Handle& profile_h,
                                                  B bot_ocl_buf,
                                                  T top_ocl_buf,
                                                  TW wei_ocl_buf,
                                                  ConstData_t bias_ocl_buf,
                                                  const ConvolutionContext& config,
                                                  const ConvSolution& solution,
                                                  float& elapsed_time) const
{
    assert(bias_ocl_buf == nullptr);
    (void)bias_ocl_buf;
    const KernelInfo k_info = solution.construction_params.back();
#ifdef NDEBUG
    try
#endif
    {
        elapsed_time = std::numeric_limits<float>::max();
        // ConvolutionContext::general_compile_options is for OpenCL kernels
        // and thus not applicable for assembly.
        auto kernel = profile_h.AddKernel("",
                                          "",
                                          k_info.kernel_file,
                                          k_info.kernel_name,
                                          k_info.l_wk,
                                          k_info.g_wk,
                                          k_info.comp_options);

        int reserved      = 0;
        int* reserved_ptr = nullptr;

        static const int F_REVERSE_R     = 1 << 0;
        static const int F_REVERSE_S     = 1 << 1;
        static const int F_FLIP_K_C      = 1 << 2;
        static const int F_NKC_STRIDES   = 1 << 9;
        static const int F_GROUP_STRIDES = 1 << 10;

        const auto is_forward = config.direction.IsForward();
        const auto is_wrw     = config.direction.IsBackwardWrW();

        int H, W, R, S, R_stride, S_stride, C, K, out_H, out_W, pad_H, pad_W, N, idilation_w,
            idilation_h, n_groups, group_cnt;
        // clang-format off
        FillVarsFromConfig(H, W, R, S, R_stride, S_stride, C, K, out_H, out_W, pad_H,
                           pad_W, N, idilation_w, idilation_h, n_groups, group_cnt, config);
        // clang-format on
        int flags = ((is_forward || is_wrw) ? 0 : F_REVERSE_R + F_REVERSE_S + F_FLIP_K_C) +
                    F_NKC_STRIDES + F_GROUP_STRIDES;

        C = C / group_cnt;
        K = K / group_cnt;
        // cppcheck-suppress unreadVariable

        const auto in_layout = (is_wrw) ? GetSwappedNCLayout(GetMemLayout_t(config.in_layout))
                                        : GetMemLayout_t(config.in_layout);
        const auto out_layout = (is_wrw) ? GetSwappedNCLayout(GetMemLayout_t(config.out_layout))
                                         : GetMemLayout_t(config.out_layout);
        const auto wei_layout =
            (is_wrw) ? GetSwappedNCLayout(MemLayout_t::NCHW)
                     : (is_forward ? (MemLayout_t::NCHW) : GetSwappedNCLayout(MemLayout_t::NCHW));

        // cppcheck-suppress unreadVariable
        BuffInfo d_buf(GetGroupConvLayout(in_layout, true),
                       N,
                       C,
                       H,
                       W,
                       group_cnt,
                       GetTypeSize(config.in_data_type)),
            // cppcheck-suppress unreadVariable
            o_buf(GetGroupConvLayout(out_layout, true),
                  N,
                  K,
                  out_H,
                  out_W,
                  group_cnt,
                  GetTypeSize(config.out_data_type)),
            // cppcheck-suppress unreadVariable
            f_buf(GetGroupConvLayout(wei_layout, false),
                  K,
                  C,
                  R,
                  S,
                  group_cnt,
                  GetTypeSize(config.weights_data_type));

        if(k_info.l_wk[0] != 0)
            n_groups = solver::ConvBinWinogradRxSf2x3::GetNGroups(config.group_counts,
                                                                  k_info.g_wk[0] / k_info.l_wk[0]);
        else
            n_groups = solver::ConvBinWinogradRxSf2x3::GetNGroups(
                config.group_counts,
                k_info.g_wk[0] / 512); // For OCL runtime. Issue #1724

        kernel(N,
               C,
               H,
               W,
               K,
               n_groups,
               flags,
               reserved,
               bot_ocl_buf,
               wei_ocl_buf,
               top_ocl_buf,
               reserved_ptr, // Unused return_addr.
               R,
               S,
               pad_H, // Like Fwd wino.
               pad_W,
               out_H,
               out_W,
               reserved_ptr, // Unused bias_addr.
               reserved,     // Unused relu_alpha.
               d_buf.byte_stride.nk,
               d_buf.byte_stride.c,
               d_buf.byte_stride.h,
               d_buf.byte_stride.w,
               f_buf.byte_stride.nk,
               f_buf.byte_stride.c,
               f_buf.byte_stride.h,
               f_buf.byte_stride.w,
               o_buf.byte_stride.nk,
               o_buf.byte_stride.c,
               o_buf.byte_stride.h,
               o_buf.byte_stride.w,
               group_cnt,
               d_buf.byte_stride.g,
               f_buf.byte_stride.g,
               o_buf.byte_stride.g);

        elapsed_time = profile_h.GetKernelTime();
    }
#ifdef NDEBUG
    catch(miopen::Exception& ex)
    {
        MIOPEN_LOG_WE(ex.what());
        return -1;
    }
#endif
    return 0;
}

bool ConvBinWinogradRxSf2x3::IsApplicable(const ConvolutionContext& params) const
{
    if(!params.Is2d())
        return false;
    if(!(params.IsFp32() || params.IsFp16()))
        return false;
    if(miopen::IsDisabled(MIOPEN_DEBUG_AMD_WINOGRAD_RXS_F2X3{}))
        return false;
    if(!params.use_asm_kernels)
        return false;
    if(!params.rmv.IsV3())
        return false;

    const auto name = params.GetStream().GetDeviceName();
    if(!(StartsWith(name, "gfx9")))
        return false;
    if(params.IsFp16() && !(StartsWith(name, "gfx906") || StartsWith(name, "gfx908")))
        return false;

    // clang-format off
    if (! ( (params.kernel_stride_w == 1 || params.kernel_stride_w == 2)
        && params.kernel_stride_w == params.kernel_stride_h
        && params.kernel_dilation_w == 1
        && params.kernel_dilation_h == 1
        && params.bias == 0
        && params.in_layout == "NCHW"))
        return false;
    // clang-format on

    const auto n_inputs_per_group  = params.n_inputs / params.group_counts,
               n_outputs_per_group = params.n_outputs / params.group_counts;

    if(params.direction.IsBackwardWrW())
    {
        if(params.kernel_stride_w == 2)
            return false;
        return IsShaderContraintsMet(params.in_height,
                                     params.in_width,
                                     params.kernel_dilation_h,
                                     params.kernel_dilation_w,
                                     params.batch_sz,    // N
                                     n_inputs_per_group, // K
                                     params.out_height,
                                     params.out_width,
                                     params.kernel_size_h,
                                     params.kernel_size_w,
                                     n_outputs_per_group, // C
                                     params);
    }
    else
    {
        return IsShaderContraintsMet(params.kernel_size_h, // RxS
                                     params.kernel_size_w,
                                     params.kernel_stride_h,
                                     params.kernel_stride_w,
                                     n_inputs_per_group,  // C
                                     n_outputs_per_group, // K
                                     params.in_height,    // HxW
                                     params.in_width,
                                     params.out_height, // OHxOW
                                     params.out_width,
                                     params.batch_sz, // N
                                     params);
    }
}

ConvSolution
ConvBinWinogradRxSf2x3::GetSolution(const ConvolutionContext& params,
                                    const PerformanceConfigConvBinWinogradRxSf2x3& config,
                                    const bool disableConfigOverrideFromEnv) const
{
    const auto n_groups = config.n_groups;
    static bool IsWarned;
    if(!IsWarned)
    {
        if(params.GetStream().GetMaxComputeUnits() > MAX_CU_LIMIT)
            MIOPEN_LOG_WE(SolverDbId(*this) << ": GPU has "
                                            << params.GetStream().GetMaxComputeUnits()
                                            << "CUs, but this solver supports max "
                                            << MAX_CU_LIMIT
                                            << "and thus may show sub-optimal performance.");
        IsWarned = true;
    }

    ConvSolution result;

    const PerformanceConfigConvBinWinogradRxSf2x3* pcfg = &config;
    PerformanceConfigConvBinWinogradRxSf2x3 fromEnv;

    if(!disableConfigOverrideFromEnv)
    {
        std::string s;
        const auto p_asciz = miopen::GetStringEnv(MIOPEN_DEBUG_AMD_WINOGRAD_RXS_F2X3_PERF_VALS{});
        if(p_asciz != nullptr)
        {
            s = std::string(p_asciz);
            if(!s.empty()) // else nothing to parse.
            {
                if(!fromEnv.Deserialize(s) || !fromEnv.IsValid(params))
                {
                    MIOPEN_LOG_E("MIOPEN_DEBUG_AMD_WINOGRAD_RXS_F2X3_PERF_VALS: "
                                 "Bad format or invalid for the problem config: "
                                 << s);
                }
                else
                {
                    MIOPEN_LOG_I("Overridden from env: " << fromEnv.ToString());
                    pcfg = &fromEnv;
                }
            }
        }
    }

    KernelInfo kernel;

    kernel.g_wk.push_back(512 * pcfg->GetNGroups() * params.group_counts);
    kernel.g_wk.push_back(1);
    kernel.g_wk.push_back(1);

    kernel.l_wk.push_back(512);
    kernel.l_wk.push_back(1);
    kernel.l_wk.push_back(1);

    KernelBuildParameters options{
        {"ROCM_METADATA_VERSION", 5},
    };
    kernel.comp_options = options.GenerateFor(kbp::GcnAsm{});

    std::string kernel_name    = "miopenSp3AsmConv";
    std::string kernel_file    = "Conv_Winograd";
    std::string kernel_postfix = "_v21_1_0_gfx9";

    if(params.IsFp32())
        kernel_postfix += "_fp32";
    else
    {
        kernel_postfix += "_fp16_dot2_edc";
    }
    if(params.kernel_stride_w == 1)
    {
        kernel_postfix += "_stride1";
    }
    else if(params.kernel_stride_w == 2 && !params.direction.IsBackwardData())
    {
        kernel_postfix += "_stride2";
    }
    else // if(params.kernel_dilation_h == 2)
    {
        kernel_postfix += "_dilation2";
    }

    if(params.group_counts != 1 || params.direction.IsBackwardWrW())
        kernel_postfix += "_group";

    kernel.kernel_name = kernel_name + kernel_postfix;
    kernel.kernel_file = kernel_file + kernel_postfix + ".s";

    result.construction_params.push_back(kernel);

    if(!params.direction.IsBackwardWrW())
    {
        const bool is_forward     = params.direction.IsForward();
        constexpr int F_REVERSE_R = 1 << 0;
        constexpr int F_REVERSE_S = 1 << 1;
        constexpr int F_FLIP_K_C  = 1 << 2;
        // These are not used yet. Nevertheless let's keep as a shader documentation.
        // constexpr int F_FLIP_DATA_N_C = 1 << 3; // Unsupported in f3x2.
        // constexpr int F_FLIP_OUT_N_K = 1 << 4; // Unsupported in f3x2.
        // constexpr int L_F_ADDR_INDIRECT  = 1 << 6;
        // constexpr int L_F_BIAS  = 1 << 7;
        // constexpr int L_F_LEAKY_RELU  = 1 << 8;
        constexpr int L_F_NKC_STRIDES   = 1 << 9;
        constexpr int L_F_GROUP_STRIDES = 1 << 10;
        int reserved                    = 0;
        int* reserved_ptr               = nullptr;
        int ignore;

        int N, C, H, W, K, out_H, out_W, R, S, pad_H, pad_W;
        GetCompiledInParameters(
            params, &N, &C, &H, &W, &K, &ignore, &out_H, &out_W, &R, &S, &pad_H, &pad_W);
        const auto group_cnt = params.group_counts;
        C                    = C / group_cnt;
        K                    = K / group_cnt;
        int flags            = is_forward ? 0 : F_REVERSE_R + F_REVERSE_S + F_FLIP_K_C;
        flags |= L_F_NKC_STRIDES + L_F_GROUP_STRIDES;

        // cppcheck-suppress unreadVariable
        BuffInfo d_buf(GetGroupConvLayout(GetMemLayout_t(params.in_layout), true),
                       N,
                       C,
                       H,
                       W,
                       group_cnt,
                       GetTypeSize(params.in_data_type)),
            // cppcheck-suppress unreadVariable
            o_buf(GetGroupConvLayout(GetMemLayout_t(params.out_layout), true),
                  N,
                  K,
                  out_H,
                  out_W,
                  group_cnt,
                  GetTypeSize(params.out_data_type)),
            // cppcheck-suppress unreadVariable
            f_buf(GetGroupConvLayout(is_forward ? (MemLayout_t::NCHW)
                                                : GetSwappedNCLayout(MemLayout_t::NCHW),
                                     false),
                  K,
                  C,
                  R,
                  S,
                  group_cnt,
                  GetTypeSize(params.weights_data_type));

        result.invoker_factory = [=](std::vector<Kernel> kernels) {
            return [=](const Handle& handle, const boost::any& primitive_params) {
                const auto k         = handle.Run(kernels[0]);
                const auto& data_ctx = boost::any_cast<conv::DataInvokeParams>(primitive_params);
                const auto& tensors  = data_ctx.tensors;

                // clang-format off
                MIOPEN_LOG_I2(" N=" << N << " G=" << group_cnt << " C=" << C << " H=" << H << " W=" << W << " K=" << K
                    << " n_groups=" << n_groups << " flags=" << flags << " R=" << R << " S=" << S
                    << " pad_H=" << pad_H << " pad_W=" << pad_W << " out_H=" << out_H << " out_W=" << out_W
                    << " d_buf.byte_stride.nk=" << d_buf.byte_stride.nk << " d_buf.byte_stride.c=" << d_buf.byte_stride.c
                    << " d_buf.byte_stride.h=" << d_buf.byte_stride.h << " d_buf.byte_stride.w=" << d_buf.byte_stride.w
                    << " f_buf.byte_stride.nk=" << f_buf.byte_stride.nk << " f_buf.byte_stride.c=" << f_buf.byte_stride.c
                    << " f_buf.byte_stride.h=" << f_buf.byte_stride.h << " f_buf.byte_stride.w=" << f_buf.byte_stride.w
                    << " o_buf.byte_stride.nk=" << o_buf.byte_stride.nk << " o_buf.byte_stride.c=" << o_buf.byte_stride.c
                    << " o_buf.byte_stride.h="  << o_buf.byte_stride.h <<  " o_buf.byte_stride.w=" << o_buf.byte_stride.w
                    << " d_buf.byte_stride.g=" << d_buf.byte_stride.g  << " o_buf.byte_stride.g="  << o_buf.byte_stride.g
                    << " f_buf.byte_stride.g=" << f_buf.byte_stride.g); // clang-format on

                k(N,
                  C,
                  H,
                  W,
                  K,
                  n_groups,
                  flags,
                  reserved,
                  tensors.in,
                  tensors.w,
                  tensors.out,
                  reserved_ptr, // Unused return_addr.
                  R,
                  S,
                  pad_H, // Like Fwd wino.
                  pad_W,
                  out_H,
                  out_W,
                  reserved_ptr, // Unused bias_addr.
                  reserved,     // Unused relu_alpha.
                  d_buf.byte_stride.nk,
                  d_buf.byte_stride.c,
                  d_buf.byte_stride.h,
                  d_buf.byte_stride.w,
                  f_buf.byte_stride.nk,
                  f_buf.byte_stride.c,
                  f_buf.byte_stride.h,
                  f_buf.byte_stride.w,
                  o_buf.byte_stride.nk,
                  o_buf.byte_stride.c,
                  o_buf.byte_stride.h,
                  o_buf.byte_stride.w,
                  group_cnt,
                  d_buf.byte_stride.g,
                  f_buf.byte_stride.g,
                  o_buf.byte_stride.g);
            };
        };
    }
    else
    {
        int unused = 0;
        int N, C, H, W, K, out_H, out_W, R, S;
        GetCompiledInParameters(
            params, &C, &K, &R, &S, &N, &unused, &H, &W, &out_H, &out_W, &unused, &unused);
        const auto group_cnt             = params.group_counts;
        static const int F_NKC_STRIDES   = 1 << 9;
        static const int F_GROUP_STRIDES = 1 << 10;
        int flags                        = F_NKC_STRIDES + F_GROUP_STRIDES;
        N                                = N / group_cnt;
        K                                = K / group_cnt;
        int pad_H                        = params.conv_problem.GetConv().GetConvPads()[0];
        int pad_W                        = params.conv_problem.GetConv().GetConvPads()[1];

        BuffInfo d_buf(
            GetGroupConvLayout(GetSwappedNCLayout(GetMemLayout_t(params.in_layout)), true),
            N,
            C,
            H,
            W,
            group_cnt,
            GetTypeSize(params.in_data_type)),
            o_buf(GetGroupConvLayout(GetSwappedNCLayout(GetMemLayout_t(params.out_layout)), false),
                  N,
                  K,
                  out_H,
                  out_W,
                  group_cnt,
                  GetTypeSize(params.out_data_type)),
            f_buf(GetGroupConvLayout(GetSwappedNCLayout(MemLayout_t::NCHW), true),
                  K,
                  C,
                  R,
                  S,
                  group_cnt,
                  GetTypeSize(params.weights_data_type));

        decltype(auto) batch_sz = params.batch_sz;
        decltype(auto) n_inputs = params.n_inputs;

        result.invoker_factory = [=](std::vector<Kernel> kernels) {
            return [=](const Handle& handle, const boost::any& primitive_params) {
                decltype(auto) invoke_params =
                    boost::any_cast<conv::WrWInvokeParams>(primitive_params);
                const auto& tensors = invoke_params.tensors;

                // clang-format off
                MIOPEN_LOG_I2(" N=" << N << " G=" << group_cnt << " C=" << C << " H=" << H << " W=" << W << " K=" << K
                    << " n_groups=" << n_groups << " flags=" << flags << " R=" << R << " S=" << S
                    << " pad_H=" << pad_H << " pad_W=" << pad_W << " out_H=" << out_H << " out_W=" << out_W
                    << " d_buf.byte_stride.nk=" << d_buf.byte_stride.nk << " d_buf.byte_stride.c=" << d_buf.byte_stride.c
                    << " d_buf.byte_stride.h=" << d_buf.byte_stride.h << " d_buf.byte_stride.w=" << d_buf.byte_stride.w
                    << " f_buf.byte_stride.nk=" << f_buf.byte_stride.nk << " f_buf.byte_stride.c=" << f_buf.byte_stride.c
                    << " f_buf.byte_stride.h=" << f_buf.byte_stride.h << " f_buf.byte_stride.w=" << f_buf.byte_stride.w
                    << " o_buf.byte_stride.nk=" << o_buf.byte_stride.nk << " o_buf.byte_stride.c=" << o_buf.byte_stride.c
                    << " o_buf.byte_stride.h="  << o_buf.byte_stride.h <<  " o_buf.byte_stride.w=" << o_buf.byte_stride.w
                    << " d_buf.byte_stride.g=" << d_buf.byte_stride.g  << " o_buf.byte_stride.g="  << o_buf.byte_stride.g
                    << " f_buf.byte_stride.g=" << f_buf.byte_stride.g); // clang-format on
                MIOPEN_LOG_I2(" ctx.batch_sz=" << batch_sz << "ctx.n_inputs=" << n_inputs);

                int reserved      = 0;
                int* reserved_ptr = nullptr;

                handle.Run(kernels[0])(N,
                                       C,
                                       H,
                                       W,
                                       K,
                                       n_groups,
                                       flags,
                                       reserved,
                                       tensors.x,
                                       tensors.dy,
                                       tensors.dw,
                                       reserved_ptr, // Unused return_addr.
                                       R,
                                       S,
                                       pad_H, // Like Fwd wino.
                                       pad_W,
                                       out_H,
                                       out_W,
                                       reserved_ptr, // Unused bias_addr.
                                       reserved,     // Unused relu_alpha.
                                       d_buf.byte_stride.nk,
                                       d_buf.byte_stride.c,
                                       d_buf.byte_stride.h,
                                       d_buf.byte_stride.w,
                                       f_buf.byte_stride.nk,
                                       f_buf.byte_stride.c,
                                       f_buf.byte_stride.h,
                                       f_buf.byte_stride.w,
                                       o_buf.byte_stride.nk,
                                       o_buf.byte_stride.c,
                                       o_buf.byte_stride.h,
                                       o_buf.byte_stride.w,
                                       group_cnt,
                                       d_buf.byte_stride.g,
                                       f_buf.byte_stride.g,
                                       o_buf.byte_stride.g);
            };
        };
    }

    return result;
}

} // namespace solver
} // namespace miopen


//
//=======================================================================
// Copyright 1997, 1998, 1999, 2000 University of Notre Dame.
// Authors: Andrew Lumsdaine, Lie-Quan Lee, Jeremy G. Siek
//
// Distributed under the Boost Software License, Version 1.0. (See
// accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
//=======================================================================
//
// $URL$
// $Id$
//
// NOTE: this file have been taken from boost 1.46.1 for using
//       with Modificable_priority_queue (to enhance the
//       non-documented mutable_queue).
//       original file is <boost/graph/detail/array_binary_tree.hpp>
//
#ifndef CGAL_INTERNAL_ARRAY_BINARY_TREE_HPP
#define CGAL_INTERNAL_ARRAY_BINARY_TREE_HPP

#include <iterator>
#include <functional>
#include <boost/config.hpp>

namespace CGAL { namespace internal {
namespace boost_ {

/*
 * Note: array_binary_tree is a completey balanced binary tree.
 */
#if !defined BOOST_NO_STD_ITERATOR_TRAITS
  template <class RandomAccessIterator, class ID>
#else
  template <class RandomAccessIterator, class ValueType, class ID>
#endif
class array_binary_tree_node {
public:
  typedef array_binary_tree_node ArrayBinaryTreeNode;
  typedef RandomAccessIterator rep_iterator;
#if !defined BOOST_NO_STD_ITERATOR_TRAITS
  typedef typename std::iterator_traits<RandomAccessIterator>::difference_type
          difference_type;
  typedef typename std::iterator_traits<RandomAccessIterator>::value_type
          value_type;
#else
  typedef int difference_type;
  typedef ValueType value_type;
#endif
  typedef difference_type size_type;

  struct children_type {
    struct iterator
        : ::boost::iterator<std::bidirectional_iterator_tag, ArrayBinaryTreeNode,
                       difference_type, array_binary_tree_node*, ArrayBinaryTreeNode&>
    { // replace with iterator_adaptor implementation -JGS

      inline iterator() : i(0), n(0) { }
      inline iterator(const iterator& x) : r(x.r), i(x.i), n(x.n), id(x.id) { }
      inline iterator& operator=(const iterator& x) {
        r = x.r; i = x.i; n = x.n;
        /*egcs generate a warning*/
        id = x.id;
        return *this;
      }
      inline iterator(rep_iterator rr,
                      size_type ii,
                      size_type nn,
                      const ID& _id) : r(rr), i(ii), n(nn), id(_id) { }
      inline array_binary_tree_node operator*() {
        return ArrayBinaryTreeNode(r, i, n, id); }
      inline iterator& operator++() { ++i; return *this; }
      inline iterator operator++(int)
        { iterator t = *this; ++(*this); return t; }
      inline bool operator==(const iterator& x) const { return i == x.i; }
      inline bool operator!=(const iterator& x) const
        { return !(*this == x); }
      rep_iterator r;
      size_type i;
      size_type n;
      ID id;
    };
    inline children_type() : i(0), n(0) { }
    inline children_type(const children_type& x)
      : r(x.r), i(x.i), n(x.n), id(x.id) { }
    inline children_type& operator=(const children_type& x) {
      r = x.r; i = x.i; n = x.n;
      /*egcs generate a warning*/
      id = x.id;
      return *this;
    }
    inline children_type(rep_iterator rr,
                         size_type ii,
                         size_type nn,
                         const ID& _id) : r(rr), i(ii), n(nn), id(_id) { }
    inline iterator begin() { return iterator(r, 2 * i + 1, n, id); }
    inline iterator end() { return iterator(r, 2 * i + 1 + size(), n, id); }
    inline size_type size() const {
      size_type c = 2 * i + 1;
      size_type s;
      if      (c + 1 < n) s = 2;
      else if (c < n)     s = 1;
      else                s = 0;
      return s;
    }
    rep_iterator r;
    size_type i;
    size_type n;
    ID id;
  };
  inline array_binary_tree_node() : i(0), n(0) { }
  inline array_binary_tree_node(const array_binary_tree_node& x)
    : r(x.r), i(x.i), n(x.n), id(x.id) { }
  inline ArrayBinaryTreeNode& operator=(const ArrayBinaryTreeNode& x) {
    r = x.r;
    i = x.i;
    n = x.n;
    /*egcs generate a warning*/
    id = x.id;
    return *this;
  }
  inline array_binary_tree_node(rep_iterator start,
                                rep_iterator end,
                                rep_iterator pos, const ID& _id)
    : r(start), i(pos - start), n(end - start), id(_id) { }
  inline array_binary_tree_node(rep_iterator rr,
                                size_type ii,
                                size_type nn, const ID& _id)
    : r(rr), i(ii), n(nn), id(_id) { }
  inline value_type& value() { return *(r + i); }
  inline const value_type& value() const { return *(r + i); }
  inline ArrayBinaryTreeNode parent() const {
    return ArrayBinaryTreeNode(r, (i - 1) / 2, n, id);
  }
  inline bool has_parent() const { return i != 0; }
  inline children_type children() { return children_type(r, i, n, id); }
  /*
  inline void swap(array_binary_tree_node x) {
    value_type tmp = x.value();
    x.value() = value();
    value() = tmp;
    i = x.i;
  }
  */
  template <class ExternalData>
  inline void swap(ArrayBinaryTreeNode x, ExternalData& edata ) {
    using ::boost::get;

    value_type tmp = x.value();

    /*swap external data*/
    edata[ get(id, tmp) ]     = i;
    edata[ get(id, value()) ] = x.i;

    x.value() = value();
    value() = tmp;
    i = x.i;
  }
   inline const children_type children() const {
    return children_type(r, i, n);
  }
  inline size_type index() const { return i; }
  rep_iterator r;
  size_type i;
  size_type n;
  ID id;
};

template <class RandomAccessContainer,
       class Compare = std::less<typename RandomAccessContainer::value_type> >
struct compare_array_node {
  typedef typename RandomAccessContainer::value_type value_type;
  compare_array_node(const Compare& x) : comp(x) {}
  compare_array_node(const compare_array_node& x) : comp(x.comp) {}

  template< class node_type >
  inline bool operator()(const node_type& x, const node_type& y) {
    return comp(x.value(), y.value());
  }

  template< class node_type >
  inline bool operator()(const node_type& x, const node_type& y) const {
    return comp(x.value(), y.value());
  }
  Compare comp;
};

} } } //namespace CGAL::internal::boost_

#endif /* CGAL_INTERNAL_ARRAY_BINARY_TREE_HPP */


#include <stan/math/rev.hpp>
#include <boost/mp11.hpp>
#include <gtest/gtest.h>
#include <test/unit/math/rev/functor/test_fixture_ode.hpp>
#include <test/unit/math/rev/functor/test_fixture_ode_fho.hpp>
#include <stan/math/torsten/test/unit/dsolve/ode_test_functors.hpp>

/**
 *
 * Use same solver functor type for both w & w/o tolerance control
 */
template <typename solve_type, typename... Ts>
using ode_test_tuple = std::tuple<solve_type, solve_type, Ts...>;

/**
 * Outer product of test types
 */
using forced_harm_osc_ts_test_types = boost::mp11::mp_product<
    ode_test_tuple,
    ::testing::Types<pmx_ode_adams_functor, pmx_ode_bdf_functor, pmx_ode_ckrk_functor,
                     pmx_ode_rk45_functor>,
    ::testing::Types<stan::math::var_value<double> >,          // time
    ::testing::Types<double, stan::math::var_value<double> >,  // y0
    ::testing::Types<double, stan::math::var_value<double> >   // theta
    >;

TYPED_TEST_SUITE_P(forced_harm_osc_ts_test);
TYPED_TEST_P(forced_harm_osc_ts_test, ts_ad) { this->test_ts_ad(); }
REGISTER_TYPED_TEST_SUITE_P(forced_harm_osc_ts_test, ts_ad);
INSTANTIATE_TYPED_TEST_SUITE_P(StanOde, forced_harm_osc_ts_test,
                               forced_harm_osc_ts_test_types);


//Copyright (c) 2015 Zachary Kann
//
//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.

// ---
// Author: Zachary Kann

// A set of functions for calculation of electric field / frequency maps
// as outlined in the papers of Skinner et al. See, for example, S.M. Gruenbaum,
// et al. J. Chem. Theory Comput. 2013 9 (7), 3109. These maps are necessary
// for calculation of vibrational spectra using the methodology  outline in the
// aforementioned papers.

// Abbreviations:
//   IR = infrared
//   SFG = sum-frequency generation

#include <armadillo>
#include "z_frequency_map.hpp"
#include "z_histogram.hpp"
#include "z_molecule_group.hpp"
#include "z_cx_tcf.hpp"

#ifndef _Z_SFG_MAP_HPP_
#define _Z_SFG_MAP_HPP_

class SfgMap:
  public FrequencyMap {
 public:
  inline SfgMap(const MoleculeGroup& group, const Chromophore chromophore,
                const double timestep, const int correlation_length = 200,
                const double slab_center_z = 0.0)
      : FrequencyMap(group, chromophore, timestep, correlation_length),
        slab_center_z_(slab_center_z) {
    mu_01_ = arma::zeros<arma::mat>(num_chromophores(), steps_guess_);
    alpha_01_ = arma::zeros<arma::mat>(num_chromophores(), steps_guess_);
  }

 private:
  arma::mat mu_01_;
  arma::mat alpha_01_;
  double slab_center_z_;
  static const double r_switch_ = 0.4;
  static const double r_switch_squared_ = 0.16;
  static const double r_switch_cubed_ = 0.064;

  inline void ResizeArrays() {
    omega_01_.resize(num_chromophores(), steps_guess_);
    mu_01_.resize(num_chromophores(), steps_guess_);
    alpha_01_.resize(num_chromophores(), steps_guess_);
  }

  inline double MuOrAlphaProduct(const int chromophore, const int corr_start,
                                 const int i_corr) {
    return mu_01_(chromophore,corr_start)*mu_01_(chromophore,corr_start+i_corr);
  }

  // Calculates switching function needed for SFG calculations in slab geometry.
  // Allows smooth switching between contribution to upper and lower surface as
  // one moves through the slab.
  double SwitchingFunction(const arma::rowvec& box);

  // Calculates omega and mu/alpha using the frequency maps.
  void UseMapping(const int chromophore,
                  const arma::rowvec& box);
};
#endif