//===-- lib/runtime/dot-product.h -------------------------------*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#ifndef FLANG_RT_RUNTIME_DOT_PRODUCT_H_
#define FLANG_RT_RUNTIME_DOT_PRODUCT_H_

#include "float.h"
#include "flang-rt/runtime/descriptor.h"
#include "flang-rt/runtime/terminator.h"
#include "flang-rt/runtime/tools.h"
#include "flang/Common/float128.h"
#include "flang/Runtime/cpp-type.h"
#include "flang/Runtime/reduction.h"
#include <cfloat>
#include <cinttypes>

namespace Fortran::runtime {

// Beware: DOT_PRODUCT of COMPLEX data uses the complex conjugate of the first
// argument; MATMUL does not.

// General accumulator for any type and stride; this is not used for
// contiguous numeric vectors.
template <TypeCategory RCAT, int RKIND, typename XT, typename YT>
class Accumulator {
public:
  using Result = AccumulationType<RCAT, RKIND>;
  RT_API_ATTRS Accumulator(const Descriptor &x, const Descriptor &y)
      : x_{x}, y_{y} {}
  RT_API_ATTRS void AccumulateIndexed(SubscriptValue xAt, SubscriptValue yAt) {
    if constexpr (RCAT == TypeCategory::Logical) {
      sum_ = sum_ ||
          (IsLogicalElementTrue(x_, &xAt) && IsLogicalElementTrue(y_, &yAt));
    } else {
      const XT &xElement{*x_.Element<XT>(&xAt)};
      const YT &yElement{*y_.Element<YT>(&yAt)};
      if constexpr (RCAT == TypeCategory::Complex) {
        sum_ += rtcmplx::conj(static_cast<Result>(xElement)) *
            static_cast<Result>(yElement);
      } else {
        sum_ += static_cast<Result>(xElement) * static_cast<Result>(yElement);
      }
    }
  }
  RT_API_ATTRS Result GetResult() const { return sum_; }

private:
  const Descriptor &x_, &y_;
  Result sum_{};
};

template <TypeCategory RCAT, int RKIND, typename XT, typename YT>
static inline RT_API_ATTRS CppTypeFor<RCAT, RKIND> DoDotProduct(
    const Descriptor &x, const Descriptor &y, Terminator &terminator) {
  using Result = CppTypeFor<RCAT, RKIND>;
  RUNTIME_CHECK(terminator, x.rank() == 1 && y.rank() == 1);
  SubscriptValue n{x.GetDimension(0).Extent()};
  if (SubscriptValue yN{y.GetDimension(0).Extent()}; yN != n) {
    terminator.Crash(
        "DOT_PRODUCT: SIZE(VECTOR_A) is %jd but SIZE(VECTOR_B) is %jd",
        static_cast<std::intmax_t>(n), static_cast<std::intmax_t>(yN));
  }
  if constexpr (RCAT != TypeCategory::Logical) {
    if (x.GetDimension(0).ByteStride() == sizeof(XT) &&
        y.GetDimension(0).ByteStride() == sizeof(YT)) {
      // Contiguous numeric vectors
      if constexpr (std::is_same_v<XT, YT>) {
        // Contiguous homogeneous numeric vectors
        if constexpr (std::is_same_v<XT, float>) {
          // TODO: call BLAS-1 SDOT or SDSDOT
        } else if constexpr (std::is_same_v<XT, double>) {
          // TODO: call BLAS-1 DDOT
        } else if constexpr (std::is_same_v<XT, rtcmplx::complex<float>>) {
          // TODO: call BLAS-1 CDOTC
        } else if constexpr (std::is_same_v<XT, rtcmplx::complex<double>>) {
          // TODO: call BLAS-1 ZDOTC
        }
      }
      XT *xp{x.OffsetElement<XT>(0)};
      YT *yp{y.OffsetElement<YT>(0)};
      using AccumType = AccumulationType<RCAT, RKIND>;
      AccumType accum{};
      if constexpr (RCAT == TypeCategory::Complex) {
        for (SubscriptValue j{0}; j < n; ++j) {
          // conj() may instantiate its argument twice,
          // so xp has to be incremented separately.
          // This is a workaround for an alleged bug in clang,
          // that shows up as:
          //   warning: multiple unsequenced modifications to 'xp'
          accum += rtcmplx::conj(static_cast<AccumType>(*xp)) *
              static_cast<AccumType>(*yp++);
          xp++;
        }
      } else {
        for (SubscriptValue j{0}; j < n; ++j) {
          accum +=
              static_cast<AccumType>(*xp++) * static_cast<AccumType>(*yp++);
        }
      }
      return static_cast<Result>(accum);
    }
  }
  // Non-contiguous, heterogeneous, & LOGICAL cases
  SubscriptValue xAt{x.GetDimension(0).LowerBound()};
  SubscriptValue yAt{y.GetDimension(0).LowerBound()};
  Accumulator<RCAT, RKIND, XT, YT> accumulator{x, y};
  for (SubscriptValue j{0}; j < n; ++j) {
    accumulator.AccumulateIndexed(xAt++, yAt++);
  }
  return static_cast<Result>(accumulator.GetResult());
}

template <TypeCategory RCAT, int RKIND> struct DotProduct {
  using Result = CppTypeFor<RCAT, RKIND>;
  template <TypeCategory XCAT, int XKIND> struct DP1 {
    template <TypeCategory YCAT, int YKIND> struct DP2 {
      RT_API_ATTRS Result operator()(const Descriptor &x, const Descriptor &y,
          Terminator &terminator) const {
        if constexpr (constexpr auto resultType{
                          GetResultType(XCAT, XKIND, YCAT, YKIND)}) {
          if constexpr (resultType->first == RCAT &&
              (resultType->second <= RKIND || RCAT == TypeCategory::Logical)) {
            return DoDotProduct<RCAT, RKIND, CppTypeFor<XCAT, XKIND>,
                CppTypeFor<YCAT, YKIND>>(x, y, terminator);
          }
        }
        terminator.Crash(
            "DOT_PRODUCT(%d(%d)): bad operand types (%d(%d), %d(%d))",
            static_cast<int>(RCAT), RKIND, static_cast<int>(XCAT), XKIND,
            static_cast<int>(YCAT), YKIND);
      }
    };
    RT_API_ATTRS Result operator()(const Descriptor &x, const Descriptor &y,
        Terminator &terminator, TypeCategory yCat, int yKind) const {
      return ApplyType<DP2, Result>(yCat, yKind, terminator, x, y, terminator);
    }
  };
  RT_API_ATTRS Result operator()(const Descriptor &x, const Descriptor &y,
      const char *source, int line) const {
    Terminator terminator{source, line};
    if (RCAT != TypeCategory::Logical && x.type() == y.type()) {
      // No conversions needed, operands and result have same known type
      return typename DP1<RCAT, RKIND>::template DP2<RCAT, RKIND>{}(
          x, y, terminator);
    } else {
      auto xCatKind{x.type().GetCategoryAndKind()};
      auto yCatKind{y.type().GetCategoryAndKind()};
      RUNTIME_CHECK(terminator, xCatKind.has_value() && yCatKind.has_value());
      return ApplyType<DP1, Result>(xCatKind->first, xCatKind->second,
          terminator, x, y, terminator, yCatKind->first, yCatKind->second);
    }
  }
};

} // namespace Fortran::runtime

#endif // FLANG_RT_RUNTIME_DOT_PRODUCT_H_
