
// =================================================================================================
// This file is part of the CLBlast project. Author(s):
//   Cedric Nugteren <www.cedricnugteren.nl>
//
// This file implements data-prepration routines for proper input for the TRSM routine. Note: The
// data-preparation routines are taken from clBLAS
//
// =================================================================================================

#ifndef CLBLAST_TEST_ROUTINES_XTRSM_DATA_H_
#define CLBLAST_TEST_ROUTINES_XTRSM_DATA_H_

#include <random>
#include <string>
#include <vector>

#include "utilities/utilities.hpp"

namespace clblast {
// =================================================================================================

// Limits to prepare proper input data
template <typename T>
double TrsmLimitMatA();
template <>
double TrsmLimitMatA<float>() {
  return pow(2.0, 7);
}
template <>
double TrsmLimitMatA<double>() {
  return pow(2.0, 5);
}
template <>
double TrsmLimitMatA<float2>() {
  return TrsmLimitMatA<float>();
}
template <>
double TrsmLimitMatA<double2>() {
  return TrsmLimitMatA<double>();
}
template <typename T>
double TrsmLimitMatB();
template <>
double TrsmLimitMatB<float>() {
  return pow(2.0, 16);
}
template <>
double TrsmLimitMatB<double>() {
  return pow(2.0, 47);
}
template <>
double TrsmLimitMatB<float2>() {
  return TrsmLimitMatB<float>();
}
template <>
double TrsmLimitMatB<double2>() {
  return TrsmLimitMatB<double>();
}

// Matrix element setter
template <typename T>
void SetElement(const clblast::Layout layout, const size_t row, const size_t column, T* mat, const size_t ld,
                const T value) {
  if (layout == clblast::Layout::kRowMajor) {
    mat[column + ld * row] = value;
  } else {
    mat[row + ld * column] = value;
  }
}

// Matrix element getter
template <typename T>
T GetElement(const clblast::Layout layout, const size_t row, const size_t column, const T* mat, const size_t ld) {
  if (layout == clblast::Layout::kRowMajor) {
    return mat[column + ld * row];
  } else {
    return mat[row + ld * column];
  }
}

// Bounds a value between 'left' and 'right'. The random value is assumed to be between -1 and +1.
template <typename T>
T BoundRandom(const double rand_val, const double left, const double right) {
  const auto value = Constant<T>(rand_val * (right - left));
  if (AbsoluteValue<T>(value) < 0.0) {
    return value - Constant<T>(left);
  } else {
    return value + Constant<T>(left);
  }
}

// The clBLAS function to generate proper input matrices for matrices A & B. Note that this routine
// should remain deterministic. Random values are therefore taken from the existing input, which
// is scaled between -1 and +1.
template <typename T>
void GenerateProperTrsmMatrices(const Arguments<T>& args, const int seed, T* mat_a, T* mat_b) {
  // Random number generator
  std::mt19937 mt(seed);
  std::uniform_real_distribution<double> dist(-1.0, 1.0);

  const auto k = (args.side == Side::kLeft) ? args.m : args.n;

  // Determines: max(|a_{ii}|) and  min(|a_{ii}|)
  // Generates: a_{ii} which are constrainted by min/max
  auto min = ConstantZero<T>();
  if (args.diagonal == clblast::Diagonal::kUnit) {
    for (auto i = size_t{0}; i < k; ++i) {
      SetElement<T>(args.layout, i, i, mat_a, args.a_ld, ConstantOne<T>());  // must not be accessed
    }
  } else {
    auto max = Constant<T>(dist(mt) * TrsmLimitMatA<T>());
    if (AbsoluteValue(max) < 1.0) {
      max += Constant<T>(3.0);
    }  // no zero's on the diagonal
    min = max / Constant<T>(100.0);
    SetElement<T>(args.layout, 0, 0, mat_a, args.a_ld, max);
    for (auto i = size_t{1}; i < k; ++i) {
      auto value = BoundRandom<T>(dist(mt), AbsoluteValue(min), AbsoluteValue(max));
      if (AbsoluteValue(value) == 0) {
        value = max;
      }
      SetElement<T>(args.layout, i, i, mat_a, args.a_ld, value);
    }
  }

  // Generates a_{ij} for all j <> i.
  for (auto i = size_t{0}; i < k; ++i) {
    auto sum = (args.diagonal == clblast::Diagonal::kUnit)
                   ? AbsoluteValue(ConstantOne<T>())
                   : AbsoluteValue(GetElement<T>(args.layout, i, i, mat_a, args.a_ld));
    for (auto j = size_t{0}; j < k; ++j) {
      if (j == i) {
        continue;
      }
      auto value = ConstantZero<T>();
      if (((args.triangle == clblast::Triangle::kUpper) && (j > i)) ||
          ((args.triangle == clblast::Triangle::kLower) && (j < i))) {
        if (sum >= 1.0) {
          const auto limit = sum / std::sqrt(static_cast<double>(k) - static_cast<double>(j));
          value = Constant<T>(dist(mt) * limit);
          sum -= AbsoluteValue(value);
        }
      }
      SetElement<T>(args.layout, i, j, mat_a, args.a_ld, value);
    }
  }

  // Generate matrix B
  if (args.side == clblast::Side::kLeft) {
    for (auto j = size_t{0}; j < args.n; ++j) {
      auto sum = TrsmLimitMatB<T>();
      for (auto i = size_t{0}; i < args.m; ++i) {
        const auto a_value = GetElement<T>(args.layout, i, i, mat_a, args.a_ld);
        auto value = ConstantZero<T>();
        if (sum >= 0.0) {
          const auto limit =
              sum * AbsoluteValue(a_value) / std::sqrt(static_cast<double>(args.m) - static_cast<double>(i));
          value = Constant<T>(dist(mt) * limit);
          sum -= AbsoluteValue(value) / AbsoluteValue(a_value);
        }
        SetElement<T>(args.layout, i, j, mat_b, args.b_ld, value);
        if ((i == 0 && j == 0) || (AbsoluteValue(value) < AbsoluteValue(min))) {
          min = value;
        }
      }
    }
  } else {
    for (auto i = size_t{0}; i < args.m; ++i) {
      auto sum = TrsmLimitMatB<T>();
      for (auto j = size_t{0}; j < args.n; ++j) {
        const auto a_value = GetElement<T>(args.layout, j, j, mat_a, args.a_ld);
        auto value = ConstantZero<T>();
        if (sum >= 0.0) {
          const auto limit =
              sum * AbsoluteValue(a_value) / std::sqrt(static_cast<double>(args.n) - static_cast<double>(j));
          value = Constant<T>(dist(mt) * limit);
          sum -= AbsoluteValue(value) / AbsoluteValue(a_value);
        }
        SetElement<T>(args.layout, i, j, mat_b, args.b_ld, value);
        if ((i == 0 && j == 0) || (AbsoluteValue(value) < AbsoluteValue(min))) {
          min = value;
        }
      }
    }
  }
  if (args.diagonal == clblast::Diagonal::kUnit) {
    for (auto i = size_t{0}; i < k; ++i) {
      SetElement<T>(args.layout, i, i, mat_a, args.a_ld, ConstantOne<T>());  // must not be accessed
    }
  }

  // Calculate a proper alpha
  if (AbsoluteValue(min) > AbsoluteValue(args.alpha)) {
    // Not implemented
  }

  // Adjust matrix B according to the value of alpha
  if (AbsoluteValue(args.alpha) != 1.0 && AbsoluteValue(args.alpha) != 0.0) {
    for (auto i = size_t{0}; i < args.m; ++i) {
      for (auto j = size_t{0}; j < args.n; ++j) {
        auto value = GetElement<T>(args.layout, i, j, mat_b, args.b_ld);
        value /= args.alpha;
        SetElement<T>(args.layout, i, j, mat_b, args.b_ld, value);
      }
    }
  }
}

// =================================================================================================
}  // namespace clblast

// CLBLAST_TEST_ROUTINES_XTRSM_DATA_H_
#endif
