// Experimental Paddle CPU operators for the fixed TrOCR encoder profile.
// Welford/reduction order follows pinned PyTorch; exp follows pinned SLEEF.
// Retain trocr_paddle_norm_probe.LICENSE and trocr_paddle_sleef_probe.LICENSE.
#include <array>
#include <cfenv>
#include <cmath>
#include <cstdint>
#include <vector>
#include "paddle/extension.h"

namespace {
constexpr int kWidth = 384;
constexpr int kAttention = 578;
struct Moments {
  int count = 0;
  std::array<float, 8> mean{};
  std::array<float, 8> m2{};
};
void Merge(Moments& dst, const Moments& src) {
  const int old = dst.count;
  const float factor = float(src.count) / float(old + src.count);
  for (int i = 0; i < 8; ++i) {
    const float delta = src.mean[i] - dst.mean[i];
    const float scaled = factor * delta;
    dst.mean[i] = dst.mean[i] + scaled;
    dst.m2[i] = std::fma(delta * float(old), scaled, dst.m2[i] + src.m2[i]);
  }
  dst.count += src.count;
}
void CheckCPU(const paddle::Tensor& value) {
  PD_CHECK(value.is_cpu(), "Source-order research operators require CPU inputs.");
  PD_CHECK(value.dtype() == paddle::DataType::FLOAT32, "Require FP32 inputs.");
  PD_CHECK(std::fegetround() == FE_TONEAREST, "Require round-to-nearest floating point.");
}
float NegativeExp(float d) {
  if (d < -104.f) return 0.f;
  const int q = int(std::nearbyint(d * 1.4426950408889634073599246810018921374f));
  float s = std::fma(float(q), -0.693145751953125f, d);
  s = std::fma(float(q), -1.428606765330187045e-6f, s);
  float u = 0.000198527617612853646278381f;
  u = std::fma(u, s, 0.00139304355252534151077271f);
  u = std::fma(u, s, 0.00833336077630519866943359f);
  u = std::fma(u, s, 0.0416664853692054748535156f);
  u = std::fma(u, s, 0.166666671633720397949219f);
  u = std::fma(u, s, 0.5f);
  u = 1.f + std::fma(s * s, u, s);
  const int half = int(std::floor(float(q) * .5f));
  return std::ldexp(std::ldexp(u, half), q - half);
}
}  // namespace

std::vector<paddle::Tensor> SourceNorm384(const paddle::Tensor& x,
                                         const paddle::Tensor& gamma,
                                         const paddle::Tensor& beta) {
  CheckCPU(x); CheckCPU(gamma); CheckCPU(beta);
  PD_CHECK(!x.shape().empty() && x.shape().back() == kWidth, "Require 384 channels.");
  PD_CHECK(gamma.numel() == kWidth && beta.numel() == kWidth, "Require 384 affine values.");
  auto out = paddle::empty(x.shape(), x.dtype(), x.place());
  const float* input = x.data<float>();
  const float* weight = gamma.data<float>();
  const float* bias = beta.data<float>();
  float* result = out.data<float>();
  for (int64_t row = 0; row < x.numel() / kWidth; ++row) {
    const float* values = input + row * kWidth;
    Moments first, second;
    for (int chunk = 0; chunk < 3; ++chunk) {
      Moments local; local.count = 16;
      for (int j = 0; j < 16; ++j) {
        const float factor = 1.f / float(j + 1);
        for (int lane = 0; lane < 8; ++lane) {
          const float value = values[chunk * 128 + j * 8 + lane];
          const float delta = value - local.mean[lane];
          local.mean[lane] = std::fma(factor, delta, local.mean[lane]);
          local.m2[lane] = std::fma(delta, value - local.mean[lane], local.m2[lane]);
        }
      }
      Merge(first, local);
      if (chunk == 1) { Merge(second, first); first = Moments{}; }
    }
    Merge(first, second);
    float mean = 0.f, m2 = 0.f;
    for (int lane = 0; lane < 8; ++lane) {
      const float factor = 48.f / float((lane + 1) * 48);
      const float delta = first.mean[lane] - mean;
      mean = std::fma(factor, delta, mean);
      const float correction = (delta * delta) * factor;
      m2 = m2 + std::fma(correction, float(lane * 48), first.m2[lane]);
    }
    const float variance = m2 / float(kWidth);
    const float scale = 1.f / std::sqrt(variance + 1e-12f);
    for (int i = 0; i < kWidth; ++i) {
      const float normalized = (values[i] - mean) * scale;
      result[row * kWidth + i] = std::fma(normalized, weight[i], bias[i]);
    }
  }
  return {out};
}

std::vector<paddle::Tensor> SourceSoftmax578(const paddle::Tensor& x) {
  CheckCPU(x);
  PD_CHECK(!x.shape().empty() && x.shape().back() == kAttention, "Require 578 attention scores.");
  auto out = paddle::empty(x.shape(), x.dtype(), x.place());
  const float* input = x.data<float>(); float* output = out.data<float>();
  for (int64_t row = 0; row < x.numel() / kAttention; ++row) {
    const float* values = input + row * kAttention;
    float* result = output + row * kAttention;
    float maximum = values[0];
    for (int i = 0; i < kAttention; ++i) {
      PD_CHECK(std::isfinite(values[i]), "Require finite scores.");
      if (values[i] > maximum) maximum = values[i];
    }
    for (int i = 0; i < kAttention; ++i) result[i] = NegativeExp(values[i] - maximum);
    std::array<float, 16> sum;
    for (int i = 0; i < 16; ++i) sum[i] = result[i];
    for (int block = 1; block < 36; ++block)
      for (int lane = 0; lane < 16; ++lane) sum[lane] = sum[lane] + result[block * 16 + lane];
    for (int lane = 0; lane < 2; ++lane) sum[lane] = sum[lane] + result[576 + lane];
    for (int width : {8, 4, 2, 1})
      for (int lane = 0; lane < width; ++lane) sum[lane] = sum[lane] + sum[lane + width];
    const float inverse = 1.f / sum[0];
    for (int i = 0; i < kAttention; ++i) result[i] = result[i] * inverse;
  }
  return {out};
}

std::vector<std::vector<int64_t>> NormShape(const std::vector<int64_t>& x,
                                           const std::vector<int64_t>& gamma,
                                           const std::vector<int64_t>& beta) { return {x}; }
std::vector<paddle::DataType> NormType(paddle::DataType x, paddle::DataType gamma,
                                     paddle::DataType beta) { return {x}; }
std::vector<std::vector<int64_t>> SoftmaxShape(const std::vector<int64_t>& x) { return {x}; }
std::vector<paddle::DataType> SoftmaxType(paddle::DataType x) { return {x}; }

PD_BUILD_OP(ocr_source_norm384_v1)
    .Inputs({"X", "Gamma", "Beta"}).Outputs({"Out"})
    .SetKernelFn(PD_KERNEL(SourceNorm384))
    .SetInferShapeFn(PD_INFER_SHAPE(NormShape)).SetInferDtypeFn(PD_INFER_DTYPE(NormType));
PD_BUILD_OP(ocr_source_softmax578_v1)
    .Inputs({"X"}).Outputs({"Out"})
    .SetKernelFn(PD_KERNEL(SourceSoftmax578))
    .SetInferShapeFn(PD_INFER_SHAPE(SoftmaxShape)).SetInferDtypeFn(PD_INFER_DTYPE(SoftmaxType));