/*
 * Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <cstdint>
#include <optional>
#include <string>
#include <type_traits>

namespace tensorrt_llm {
namespace common {

enum class NVFP44Over6ErrMode : std::uint8_t {
  MAE = 0,
  MSE = 1,
};

template <int E4M3_MAX = 448, NVFP44Over6ErrMode ERR_MODE = NVFP44Over6ErrMode::MAE,
          bool ERR_USE_FAST_MATH = false>
struct NVFP44Over6Config {
  static_assert(E4M3_MAX == 448 || E4M3_MAX == 256, "Unsupported NVFP4 E4M3 max.");

  static constexpr int e4m3Max = E4M3_MAX;
  static constexpr NVFP44Over6ErrMode errMode = ERR_MODE;
  static constexpr bool errUseFastMath = ERR_USE_FAST_MATH;
};

template <typename T>
struct IsNVFP44Over6Config : std::false_type {};

template <int E4M3_MAX, NVFP44Over6ErrMode ERR_MODE, bool ERR_USE_FAST_MATH>
struct IsNVFP44Over6Config<NVFP44Over6Config<E4M3_MAX, ERR_MODE, ERR_USE_FAST_MATH>>
    : std::true_type {};

class QuantMode {
  // [WARNING] KEEP BELOW DEFINITION IN SYNC WITH tensorrt_llm/quantization/mode.py
 public:
  using BaseType = std::uint32_t;

  explicit constexpr QuantMode(BaseType value) noexcept : mValue{value} {}

  QuantMode() noexcept = default;

  constexpr QuantMode(QuantMode const&) noexcept = default;

  constexpr QuantMode& operator=(QuantMode const& other) noexcept = default;

  static constexpr QuantMode none() noexcept { return QuantMode(BaseType(0)); }

  static constexpr QuantMode int4Weights() noexcept { return QuantMode(BaseType(1u) << 0); }

  static constexpr QuantMode int8Weights() noexcept { return QuantMode(BaseType(1u) << 1); }

  static constexpr QuantMode activations() noexcept { return QuantMode(BaseType(1u) << 2); }

  static constexpr QuantMode perChannelScaling() noexcept { return QuantMode(BaseType(1u) << 3); }

  static constexpr QuantMode perTokenScaling() noexcept { return QuantMode(BaseType(1u) << 4); }

  static constexpr QuantMode perGroupScaling() noexcept { return QuantMode(BaseType(1u) << 5); }

  static constexpr QuantMode int8KvCache() noexcept { return QuantMode(BaseType(1u) << 6); }

  static constexpr QuantMode fp8KvCache() noexcept { return QuantMode(BaseType(1u) << 7); }

  static constexpr QuantMode fp8Qdq() noexcept { return QuantMode(BaseType(1u) << 8); }

  static constexpr QuantMode fp8RowWise() noexcept {
    return QuantMode(BaseType(1u) << 3 | BaseType(1u) << 4 | BaseType(1u) << 9);
  }

  static constexpr QuantMode fp8BlockScales() noexcept { return QuantMode(BaseType(1u) << 10); }

  static constexpr QuantMode w4a8QServe() noexcept { return QuantMode(BaseType(1u) << 11); }

  static constexpr QuantMode nvfp4() noexcept { return QuantMode(BaseType(1u) << 12); }

  static constexpr QuantMode fp4KvCache() noexcept { return QuantMode(BaseType(1u) << 13); }

  constexpr BaseType value() const noexcept { return mValue; }

  constexpr bool isSet(QuantMode const& mode) const noexcept {
    return (mValue & mode.value()) == mode.value();
  }

  constexpr bool hasInt4Weights() const noexcept { return isSet(int4Weights()); }

  constexpr bool hasInt8Weights() const noexcept { return isSet(int8Weights()); }

  constexpr bool hasActivations() const noexcept { return isSet(activations()); }

  constexpr bool hasPerChannelScaling() const noexcept { return isSet(perChannelScaling()); }

  constexpr bool hasPerTokenScaling() const noexcept { return isSet(perTokenScaling()); }

  constexpr bool hasPerGroupScaling() const noexcept { return isSet(perGroupScaling()); }

  constexpr bool hasStaticActivationScaling() const noexcept { return !hasPerTokenScaling(); }

  constexpr bool hasInt8KvCache() const noexcept { return isSet(int8KvCache()); }

  constexpr bool hasFp8KvCache() const noexcept { return isSet(fp8KvCache()); }

  constexpr bool hasFp4KvCache() const noexcept { return isSet(fp4KvCache()); }

  constexpr bool hasFp8Qdq() const noexcept { return isSet(fp8Qdq()); }

  constexpr bool hasFp8RowWise() const noexcept { return isSet(fp8RowWise()); }

  constexpr bool hasNvfp4() const noexcept { return isSet(nvfp4()); }

  constexpr bool hasKvCacheQuant() const noexcept {
    return hasInt8KvCache() || hasFp8KvCache() || hasFp4KvCache();
  }

  static constexpr QuantMode fromDescription(
      bool quantizeWeights = false, bool quantizeActivations = false, bool perToken = false,
      bool perChannel = false, bool perGroup = false, bool useInt4Weights = false,
      bool useInt8KvCache = false, bool useFp8KvCache = false, bool useFp8Qdq = false,
      bool useFp8RowWise = false, bool useW4a8QServe = false, bool useFp4Quant = false,
      bool useFp8BlockScales = false) {
    QuantMode quantMode{};
    if (quantizeWeights) {
      if (useInt4Weights)
        quantMode += int4Weights();
      else
        quantMode += int8Weights();
    }

    if (quantizeActivations) {
      quantMode += activations();
    }

    if (perChannel) {
      quantMode += QuantMode::perChannelScaling();
    }
    if (perToken) {
      quantMode += QuantMode::perTokenScaling();
    }
    if (perGroup) {
      quantMode += QuantMode::perGroupScaling();
    }

    if (useInt8KvCache) {
      quantMode += int8KvCache();
    }

    if (useFp8KvCache) {
      quantMode += fp8KvCache();
    }

    if (useFp8Qdq) {
      quantMode += fp8Qdq();
    }

    if (useFp8RowWise) {
      quantMode += fp8RowWise();
    }

    if (useFp8BlockScales) {
      quantMode += fp8BlockScales();
    }

    if (useW4a8QServe) {
      quantMode += w4a8QServe();
    }

    if (useFp4Quant) {
      quantMode += nvfp4();
    }

    return quantMode;
  }

  static constexpr QuantMode useSmoothQuant(bool perToken = false, bool perChannel = false) {
    return fromDescription(true, true, perToken, perChannel);
  }

  static constexpr QuantMode useQServe(bool perGroup) {
    return fromDescription(true, true, false, false, perGroup, true, false, false, false, false,
                           true);
  }

  static constexpr QuantMode useWeightOnly(bool useInt4Weights = false, bool perGroup = false) {
    return fromDescription(true, false, false, false, perGroup, useInt4Weights);
  }

  static QuantMode const fromQuantAlgo(std::optional<std::string> quantAlgo = std::nullopt,
                                       std::optional<std::string> kvCacheQuantAlgo = std::nullopt) {
    QuantMode quantMode{};
    if (quantAlgo == "W8A16") {
      quantMode = useWeightOnly(false, false);
    } else if (quantAlgo == "W4A16") {
      quantMode = useWeightOnly(true, false);
    } else if (quantAlgo == "W4A16_AWQ") {
      quantMode = useWeightOnly(true, true);
    } else if (quantAlgo == "W4A8_AWQ") {
      quantMode = useWeightOnly(true, true);
    } else if (quantAlgo == "W4A8_QSERVE_PER_GROUP") {
      quantMode = useQServe(false);
    } else if (quantAlgo == "W4A8_QSERVE_PER_CHANNEL") {
      quantMode = useQServe(true);
    } else if (quantAlgo == "W4A16_GPTQ") {
      quantMode = useWeightOnly(true, true);
    } else if (quantAlgo == "W8A8_SQ_PER_CHANNEL") {
      quantMode = useSmoothQuant(false, true);
    } else if (quantAlgo == "W8A8_SQ_PER_TENSOR_PLUGIN") {
      quantMode = useSmoothQuant(false, false);
    } else if (quantAlgo == "W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN") {
      quantMode = useSmoothQuant(true, true);
    } else if (quantAlgo == "W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN") {
      quantMode = useSmoothQuant(false, true);
    } else if (quantAlgo == "W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN") {
      quantMode = useSmoothQuant(true, false);
    } else if (quantAlgo == "FP8") {
      quantMode = fromDescription(false, false, false, false, false, false, false, false, true);
    } else if (quantAlgo == "FP8_ROWWISE") {
      quantMode =
          fromDescription(false, false, true, true, false, false, false, false, false, true);
    } else if (quantAlgo == "FP4") {
      quantMode = fromDescription(false, false, false, false, false, false, false, false, false,
                                  false, false, true);
    } else if (quantAlgo == "FP8_BLOCK_SCALES") {
      quantMode = fromDescription(false, false, false, false, false, false, false, false, false,
                                  false, false, false, true);
    }

    if (kvCacheQuantAlgo == "INT8") {
      quantMode += int8KvCache();
    } else if (kvCacheQuantAlgo == "FP8") {
      quantMode += fp8KvCache();
    } else if (kvCacheQuantAlgo == "NVFP4") {
      quantMode += fp4KvCache();
    }

    return quantMode;
  }

  constexpr QuantMode operator+(QuantMode const& other) const noexcept {
    return QuantMode(mValue | other.mValue);
  }

  constexpr QuantMode& operator+=(QuantMode const& other) noexcept { return *this = *this + other; }

  constexpr QuantMode operator-(QuantMode const& other) const noexcept {
    return QuantMode(mValue & ~other.mValue);
  }

  constexpr QuantMode& operator-=(QuantMode const& other) noexcept { return *this = *this - other; }

  constexpr bool operator==(QuantMode const& other) const noexcept {
    return mValue == other.mValue;
  }

  constexpr bool operator!=(QuantMode const& other) const noexcept { return !(*this == other); }

 private:
  BaseType mValue{0};
};

class GroupwiseQuantAlgo {
  // Keep align with tensorrt_llm/quantization/mode.py
 public:
  static constexpr int BIAS = int(1) << 0;
  static constexpr int ZERO = int(1) << 1;
  static constexpr int PRE_QUANT_SCALE = int(1) << 2;
  static constexpr int FP8_ALPHA = int(1) << 3;
  static constexpr int INT8_WEIGHT = int(1) << 4;
};

}  // namespace common
}  // namespace tensorrt_llm
