/***************************************************************************************************
 * Copyright (c) 2022 - 2026 Ali Hassani.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 *all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 *
 **************************************************************************************************/
/*! \file
    \brief Token Unpermute
*/

#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <torch/extension.h>

#include <natten/helpers.h>
#include <natten/natten.h>

#ifdef NATTEN_WITH_CUTLASS
#include <natten/cuda/tokperm/tokperm.hpp>

template <typename StdTuple>
auto std_tuple_to_cute_tuple(StdTuple a) {
  static_assert(
      std::tuple_size_v<StdTuple> > 0 && std::tuple_size_v<StdTuple> < 4);

  if constexpr (std::tuple_size_v<StdTuple> == 1) {
    return cute::make_tuple(static_cast<int>(std::get<0>(a)));
  } else if constexpr (std::tuple_size_v<StdTuple> == 2) {
    return cute::make_tuple(
        static_cast<int>(std::get<0>(a)), static_cast<int>(std::get<1>(a)));
  } else {
    return cute::make_tuple(
        static_cast<int>(std::get<0>(a)),
        static_cast<int>(std::get<1>(a)),
        static_cast<int>(std::get<2>(a)));
  }
}

#endif

namespace natten {

template <class StdNADim>
void token_unpermute_generic(
    at::Tensor& out,
    const at::Tensor& in,
    const StdNADim& token_layout,
    const StdNADim& tile_shape,
    const StdNADim& dilation,
    bool flip_tiled_dims) {
  static_assert(
      std::tuple_size_v<StdNADim> > 0 && std::tuple_size_v<StdNADim> < 4);
  static constexpr int kNADim = std::tuple_size_v<StdNADim>;

#ifdef NATTEN_WITH_CUTLASS
  CHECK_CONTIGUOUS(out);
  CHECK_CONTIGUOUS(in);
  CHECK_CUDA(out);
  CHECK_CUDA(in);

  CheckIfPropertiesMatch(in, out);

  at::cuda::OptionalCUDAGuard device_guard(in.device());

  TORCH_CHECK(
      out.dim() == 3 + kNADim,
      "Token unpermute ",
      kNADim,
      "-D expected rank-",
      3 + kNADim,
      " tensor output, got ",
      "output.dim()=",
      out.dim());

  TORCH_CHECK(
      in.size(0) == out.size(0),
      "Token unpermute input and output batches must match, got ",
      "input.shape[0]=",
      in.size(0),
      ", output.shape[0]=",
      out.size(0));

  TORCH_CHECK(
      out.size(1 + kNADim) == in.size(2),
      "Token unpermute input and output heads must match, got ",
      "output.shape[-2]=",
      out.size(1 + kNADim),
      ", input.shape[-2]=",
      in.size(2));

  TORCH_CHECK(
      out.size(2 + kNADim) == in.size(3),
      "Token unpermute input and output dims must match, got ",
      "output.shape[-1]=",
      out.size(2 + kNADim),
      ", input.shape[-2]=",
      in.size(3));

  int batch = in.size(0);
  int seqlen = in.size(1);
  int heads = in.size(2);
  int dim = in.size(3);

  auto token_layout_ = std_tuple_to_cute_tuple(token_layout);
  auto tile_shape_ = std_tuple_to_cute_tuple(tile_shape);
  auto dilation_ = std_tuple_to_cute_tuple(dilation);

  TORCH_CHECK(
      size(token_layout_) <= seqlen,
      "Token unpermute output's flattened token layout must be <= input sequence size.");

  int device_id = in.device().index();
  auto cuda_stream = at::cuda::getCurrentCUDAStream(device_id);

  cudaDeviceProp* device_props = at::cuda::getDeviceProperties(device_id);
  const int cc = device_props->major * 10 + device_props->minor;

  bool is_fp8_allowed = cc == 100 || cc == 103;

  TORCH_CHECK(
      in.scalar_type() == torch::kFloat16 ||
          in.scalar_type() == torch::kBFloat16 ||
          in.scalar_type() == torch::kFloat32 ||
          (is_fp8_allowed &&
           (in.scalar_type() == c10::ScalarType::Float8_e4m3fn ||
            in.scalar_type() == c10::ScalarType::Float8_e5m2)),
      "Only FP32, FP16, BF16, and FP8 operands are supported for now.");

  bool success = false;
  if (in.scalar_type() == torch::kFloat32) {
    success = natten::tokperm::token_unpermute_op(
        reinterpret_cast<float*>(in.data_ptr()),
        reinterpret_cast<float*>(out.data_ptr()),
        batch,
        heads,
        dim,
        token_layout_,
        tile_shape_,
        dilation_,
        flip_tiled_dims,
        cuda_stream);
  } else if (in.scalar_type() == torch::kFloat16) {
    success = natten::tokperm::token_unpermute_op(
        reinterpret_cast<cutlass::half_t*>(in.data_ptr()),
        reinterpret_cast<cutlass::half_t*>(out.data_ptr()),
        batch,
        heads,
        dim,
        token_layout_,
        tile_shape_,
        dilation_,
        flip_tiled_dims,
        cuda_stream);
  } else if (in.scalar_type() == torch::kBFloat16) {
    success = natten::tokperm::token_unpermute_op(
        reinterpret_cast<cutlass::bfloat16_t*>(in.data_ptr()),
        reinterpret_cast<cutlass::bfloat16_t*>(out.data_ptr()),
        batch,
        heads,
        dim,
        token_layout_,
        tile_shape_,
        dilation_,
        flip_tiled_dims,
        cuda_stream);
  }
#if defined(NATTEN_WITH_BLACKWELL_FNA)
  else if (
      is_fp8_allowed && in.scalar_type() == c10::ScalarType::Float8_e4m3fn) {
    success = natten::tokperm::token_unpermute_op(
        reinterpret_cast<cutlass::float_e4m3_t*>(in.data_ptr()),
        reinterpret_cast<cutlass::float_e4m3_t*>(out.data_ptr()),
        batch,
        heads,
        dim,
        token_layout_,
        tile_shape_,
        dilation_,
        flip_tiled_dims,
        cuda_stream);
  } else if (
      is_fp8_allowed && in.scalar_type() == c10::ScalarType::Float8_e5m2) {
    success = natten::tokperm::token_unpermute_op(
        reinterpret_cast<cutlass::float_e5m2_t*>(in.data_ptr()),
        reinterpret_cast<cutlass::float_e5m2_t*>(out.data_ptr()),
        batch,
        heads,
        dim,
        token_layout_,
        tile_shape_,
        dilation_,
        flip_tiled_dims,
        cuda_stream);
  }
#endif

  TORCH_CHECK(success, "Token Unpermute kernel launch failed.");

#else
  TORCH_CHECK(false, "libnatten not compiled with CUTLASS.");
#endif
}

void token_unpermute_1d(
    at::Tensor& out,
    const at::Tensor& in,
    const std::tuple<int32_t>& tile_shape,
    const std::tuple<int32_t>& dilation,
    bool flip_tiled_dims) {
  TORCH_CHECK(out.dim() == 4, "Output tensor must be 4-D.");

  token_unpermute_generic(
      out, in, {out.size(1)}, tile_shape, dilation, flip_tiled_dims);
}

void token_unpermute_2d(
    at::Tensor& out,
    const at::Tensor& in,
    const std::tuple<int32_t, int32_t>& tile_shape,
    const std::tuple<int32_t, int32_t>& dilation,
    bool flip_tiled_dims) {
  TORCH_CHECK(out.dim() == 5, "Output tensor must be 5-D.");

  token_unpermute_generic(
      out,
      in,
      {out.size(1), out.size(2)},
      tile_shape,
      dilation,
      flip_tiled_dims);
}

void token_unpermute_3d(
    at::Tensor& out,
    const at::Tensor& in,
    const std::tuple<int32_t, int32_t, int32_t>& tile_shape,
    const std::tuple<int32_t, int32_t, int32_t>& dilation,
    bool flip_tiled_dims) {
  TORCH_CHECK(out.dim() == 6, "Output tensor must be 6-D.");

  token_unpermute_generic(
      out,
      in,
      {out.size(1), out.size(2), out.size(3)},
      tile_shape,
      dilation,
      flip_tiled_dims);
}

} // namespace natten
