#include "lib/Utils/Polynomial/RNSPolynomial.h"

#include <cassert>
#include <utility>
#include <vector>

#include "lib/Dialect/ModArith/IR/ModArithAttributes.h"
#include "lib/Dialect/RNS/IR/RNSAttributes.h"
#include "lib/Utils/MathUtils.h"
#include "lib/Utils/Polynomial/NTT.h"
#include "llvm/include/llvm/ADT/APInt.h"        // from @llvm-project
#include "llvm/include/llvm/ADT/SmallVector.h"  // from @llvm-project

namespace mlir {
namespace heir {
namespace polynomial {

RNSPolynomial::RNSPolynomial(llvm::SmallVector<uint64_t> data,
                             llvm::SmallVector<uint64_t> moduli,
                             Form representation)
    : data(std::move(data)),
      moduli(std::move(moduli)),
      representation(representation) {
  assert(!this->moduli.empty() && "moduli cannot be empty");
  assert(this->data.size() % this->moduli.size() == 0 &&
         "numLimbs must divide data size");
  numCoeffs = this->data.size() / this->moduli.size();
}

RNSPolynomial RNSPolynomial::add(const RNSPolynomial& other) const {
  assert(representation == other.representation &&
         "Representations must match for arithmetic");
  assert(moduli == other.moduli && "Moduli must match for addition");
  assert(numCoeffs == other.numCoeffs && "Number of coefficients must match");

  llvm::SmallVector<uint64_t> resultData;
  resultData.reserve(data.size());

  for (size_t limbIdx = 0; limbIdx < getNumLimbs(); ++limbIdx) {
    uint64_t modulus = moduli[limbIdx];
    for (size_t coeffIdx = 0; coeffIdx < numCoeffs; ++coeffIdx) {
      uint64_t a = getElement(limbIdx, coeffIdx);
      uint64_t b = other.getElement(limbIdx, coeffIdx);
      uint64_t sum = a + b;
      if (sum >= modulus) {
        sum -= modulus;
      }
      resultData.push_back(sum);
    }
  }

  return RNSPolynomial(std::move(resultData), moduli, representation);
}

RNSPolynomial RNSPolynomial::sub(const RNSPolynomial& other) const {
  assert(representation == other.representation &&
         "Representations must match for arithmetic");
  assert(moduli == other.moduli && "Moduli must match for subtraction");
  assert(numCoeffs == other.numCoeffs && "Number of coefficients must match");

  llvm::SmallVector<uint64_t> resultData;
  resultData.reserve(data.size());

  for (size_t limbIdx = 0; limbIdx < getNumLimbs(); ++limbIdx) {
    uint64_t modulus = moduli[limbIdx];
    for (size_t coeffIdx = 0; coeffIdx < numCoeffs; ++coeffIdx) {
      uint64_t a = getElement(limbIdx, coeffIdx);
      uint64_t b = other.getElement(limbIdx, coeffIdx);
      uint64_t diff = a;
      if (diff < b) {
        diff += modulus;
      }
      diff -= b;
      resultData.push_back(diff);
    }
  }

  return RNSPolynomial(std::move(resultData), moduli, representation);
}

std::optional<RNSPolynomial> RNSPolynomial::scalarMul(
    llvm::ArrayRef<uint64_t> scalarValues) const {
  if (scalarValues.size() != getNumLimbs()) return std::nullopt;

  llvm::SmallVector<uint64_t> resultData;
  resultData.reserve(data.size());
  for (size_t limbIndex = 0; limbIndex < getNumLimbs(); ++limbIndex) {
    uint64_t modulus = moduli[limbIndex];
    uint64_t scalar = scalarValues[limbIndex] % modulus;
    for (size_t coeffIndex = 0; coeffIndex < numCoeffs; ++coeffIndex) {
      unsigned __int128 product =
          (unsigned __int128)getElement(limbIndex, coeffIndex) * scalar;
      resultData.push_back(product % modulus);
    }
  }

  return RNSPolynomial(std::move(resultData), moduli, representation);
}

std::optional<RNSPolynomial> RNSPolynomial::mul(
    const RNSPolynomial& other) const {
  assert(moduli == other.moduli && "Moduli must match for multiplication");
  assert(numCoeffs == other.numCoeffs && "Number of coefficients must match");
  assert(representation == other.representation &&
         "Mismatched representations or unsupported conversion");

  if (representation == Form::EVAL && other.representation == Form::EVAL) {
    llvm::SmallVector<uint64_t> resultData;
    resultData.reserve(data.size());

    for (size_t limbIdx = 0; limbIdx < getNumLimbs(); ++limbIdx) {
      uint64_t modulus = moduli[limbIdx];
      for (size_t coeffIdx = 0; coeffIdx < numCoeffs; ++coeffIdx) {
        uint64_t a = getElement(limbIdx, coeffIdx);
        uint64_t b = other.getElement(limbIdx, coeffIdx);
        unsigned __int128 prod = (unsigned __int128)a * b;
        uint64_t res = prod % modulus;
        resultData.push_back(res);
      }
    }
    return RNSPolynomial(std::move(resultData), moduli, representation);
  }

  if (representation == Form::COEFF && other.representation == Form::COEFF) {
    auto lhsNtt = toNtt();
    if (!lhsNtt) return std::nullopt;
    auto rhsNtt = other.toNtt();
    if (!rhsNtt) return std::nullopt;
    auto prodNtt = lhsNtt->mul(*rhsNtt);
    if (!prodNtt) return std::nullopt;
    return prodNtt->toCoefficient();
  }

  return std::nullopt;
}

std::optional<RNSPolynomial> RNSPolynomial::toNtt(
    llvm::ArrayRef<uint64_t> rootsOfUnity) const {
  if (numCoeffs > 0 && (numCoeffs & (numCoeffs - 1)) != 0) {
    return std::nullopt;
  }
  assert(representation == Form::COEFF && "Already in NTT representation");

  llvm::SmallVector<uint64_t> resultData;
  resultData.reserve(data.size());

  if (!rootsOfUnity.empty()) {
    assert(rootsOfUnity.size() == getNumLimbs() &&
           "mismatch in number of limbs for roots");
  }

  for (size_t limbIdx = 0; limbIdx < getNumLimbs(); ++limbIdx) {
    uint64_t modulus = moduli[limbIdx];
    uint64_t rootOfUnity;

    if (!rootsOfUnity.empty()) {
      rootOfUnity = rootsOfUnity[limbIdx];
    } else {
      llvm::APInt qAp(64, modulus);
      std::optional<llvm::APInt> rootOpt =
          findPrimitive2nthRoot(qAp, numCoeffs);
      if (!rootOpt.has_value()) return std::nullopt;
      rootOfUnity = rootOpt->getZExtValue();
    }

    std::vector<uint64_t> limbCoeffs;
    limbCoeffs.reserve(numCoeffs);
    for (size_t coeffIdx = 0; coeffIdx < numCoeffs; ++coeffIdx) {
      limbCoeffs.push_back(getElement(limbIdx, coeffIdx));
    }

    nttInPlace(limbCoeffs, modulus, rootOfUnity);
    resultData.insert(resultData.end(), limbCoeffs.begin(), limbCoeffs.end());
  }

  return RNSPolynomial(std::move(resultData), moduli, Form::EVAL);
}

std::optional<RNSPolynomial> RNSPolynomial::toNtt(rns::RNSAttr rootAttr) const {
  if (!rootAttr) return toNtt(llvm::ArrayRef<uint64_t>());

  llvm::SmallVector<uint64_t> roots;
  roots.reserve(rootAttr.getValues().size());
  for (auto root : rootAttr.getValues()) {
    auto rootMA = llvm::cast<mod_arith::ModArithAttr>(root);
    roots.push_back(rootMA.getValue().getValue().getZExtValue());
  }
  return toNtt(roots);
}

std::optional<RNSPolynomial> RNSPolynomial::toCoefficient(
    llvm::ArrayRef<uint64_t> rootsOfUnity) const {
  if (numCoeffs > 0 && (numCoeffs & (numCoeffs - 1)) != 0) {
    return std::nullopt;
  }
  assert(representation == Form::EVAL &&
         "Already in Coefficient representation");

  llvm::SmallVector<uint64_t> resultData;
  resultData.reserve(data.size());

  if (!rootsOfUnity.empty()) {
    assert(rootsOfUnity.size() == getNumLimbs() &&
           "mismatch in number of limbs for roots");
  }

  for (size_t limbIdx = 0; limbIdx < getNumLimbs(); ++limbIdx) {
    uint64_t modulus = moduli[limbIdx];
    uint64_t rootOfUnity;

    if (!rootsOfUnity.empty()) {
      rootOfUnity = rootsOfUnity[limbIdx];
    } else {
      llvm::APInt qAp(64, modulus);
      std::optional<llvm::APInt> rootOpt =
          findPrimitive2nthRoot(qAp, numCoeffs);
      if (!rootOpt.has_value()) return std::nullopt;
      rootOfUnity = rootOpt->getZExtValue();
    }

    std::vector<uint64_t> limbCoeffs;
    limbCoeffs.reserve(numCoeffs);
    for (size_t coeffIdx = 0; coeffIdx < numCoeffs; ++coeffIdx) {
      limbCoeffs.push_back(getElement(limbIdx, coeffIdx));
    }

    inttInPlace(limbCoeffs, modulus, rootOfUnity);
    resultData.insert(resultData.end(), limbCoeffs.begin(), limbCoeffs.end());
  }

  return RNSPolynomial(std::move(resultData), moduli, Form::COEFF);
}

std::optional<RNSPolynomial> RNSPolynomial::toCoefficient(
    rns::RNSAttr rootAttr) const {
  if (!rootAttr) return toCoefficient(llvm::ArrayRef<uint64_t>());

  llvm::SmallVector<uint64_t> roots;
  roots.reserve(rootAttr.getValues().size());
  for (auto root : rootAttr.getValues()) {
    auto rootMA = llvm::cast<mod_arith::ModArithAttr>(root);
    roots.push_back(rootMA.getValue().getValue().getZExtValue());
  }
  return toCoefficient(roots);
}

RNSPolynomial RNSPolynomial::slice(size_t start, size_t size) const {
  assert(start + size <= getNumLimbs() && "slice out of bounds");

  llvm::SmallVector<uint64_t> slicedData;
  slicedData.reserve(size * numCoeffs);

  auto dataRef = getData();
  auto slicedDataRef = dataRef.slice(start * numCoeffs, size * numCoeffs);
  slicedData.append(slicedDataRef.begin(), slicedDataRef.end());

  llvm::SmallVector<uint64_t> slicedModuli;
  slicedModuli.reserve(size);
  auto moduliRef = getModuli();
  auto slicedModuliRef = moduliRef.slice(start, size);
  slicedModuli.append(slicedModuliRef.begin(), slicedModuliRef.end());

  return RNSPolynomial(std::move(slicedData), std::move(slicedModuli),
                       representation);
}

}  // namespace polynomial
}  // namespace heir
}  // namespace mlir
