// Copyright 2024 Bytedance Inc. All Rights Reserved.
// Author: tianlei.richard@qq.com (tianlei.richard)

#include <cmath>

#include "OBJ_Loader.h"
#include "spdlog/spdlog.h"

#include "minus_renderer.h"
#include "rasterizer.h"

MinusRenderer::MinusRenderer(float near, float far, float fov,
                             float aspect_ratio)
    : near_(near), far_(far), fov_(fov), aspect_ratio_(aspect_ratio),
      projection_matrix_(orthographic_tranform() * squish_tranform()) {
  spdlog::info("near: {}, far: {}, fov: {}, aspect ratio: {}", near_, far_,
               fov_, aspect_ratio_);
}

float MinusRenderer::calculate_height(const float fov, const float near) {
  return std::fabs(near) * std::tan(fov * 0.5) * 2;
}

float MinusRenderer::calculate_width(const float height, const float ratio) {
  return height * ratio;
}

TransformMatrix MinusRenderer::squish_tranform() {
  // Frustum to Cuboid
  return TransformMatrix{{near_, 0, 0, 0},
                         {0, near_, 0, 0},
                         {0, 0, near_ + far_, -near_ * far_},
                         {0, 0, 1, 0}};
}

TransformMatrix MinusRenderer::orthographic_tranform() {
  const float height = calculate_height(fov_, near_);
  const float width = calculate_width(height, aspect_ratio_);
  const float right = width * 0.5;
  const float left = -right;
  const float top = height * 0.5;
  const float bottom = -top;

  TransformMatrix m{{1, 0, 0, -0.5 * (right - left)},
                    {0, 1, 0, -0.5 * (top - bottom)},
                    {0, 0, 1, -0.5 * (far_ - near_)},
                    {0, 0, 0, 1}};

  m = TransformMatrix{{2 / (right - left), 0, 0, 0},
                      {0, 2 / (top - bottom), 0, 0},
                      {0, 0, 2 / std::fabs(far_ - near_), 0},
                      {0, 0, 0, 1}} *
      m;

  return m;
}

TransformMatrix MinusRenderer::view_port_transform(const float width,
                                                   const float height) {
  return TransformMatrix{{width * 0.5, 0, 0, width * 0.5},
                         {0, height * 0.5, 0, height * 0.5},
                         {0, 0, 1, 0},
                         {0, 0, 0, 1}};
}

void MinusRenderer::model_transform(const TransformMatrix &mtx) {
  for (auto &m : meshes_) {
    for (auto &p : m) {
      p.affine_transform(mtx);
    }
  }
}

void MinusRenderer::view_transform(const TransformMatrix &mtx) {
  for (auto &m : meshes_) {
    for (auto &p : m) {
      p.affine_transform(mtx);
    }
  }
}

cv::Mat MinusRenderer::render(const int resolution_width,
                              const int resolution_height) {
  std::vector<Mesh> meshes = meshes_;
  Rasterizer rasterizer{resolution_width, resolution_height};
  for (auto &primitives : meshes) {
    for (int i = 0; i < primitives.size(); ++i) {
      auto &t = primitives[i];
      t.projective_transform(projection_matrix_);
      t.affine_transform(
          view_port_transform(resolution_width, resolution_height));
    }
    rasterizer.rasterize(primitives);
  }
  const auto &pixels = rasterizer.get_picture();
  assert(!pixels.empty());
  cv::Mat color_image(pixels.size(), (pixels[0]).size(), CV_8UC3);
  cv::Mat depth_image(pixels.size(), (pixels[0]).size(), CV_8UC1);

  for (int i = 0; i < pixels.size(); ++i) {
    const auto &row = pixels[i];
    for (int j = 0; j < row.size(); ++j) {
      auto &pixel_color = color_image.at<cv::Vec3b>(pixels.size() - i - 1, j);
      pixel_color = cv::Vec3b{
          static_cast<unsigned char>(row[j].x() *
                                     std::numeric_limits<unsigned char>::max()),
          static_cast<unsigned char>(row[j].y() *
                                     std::numeric_limits<unsigned char>::max()),
          static_cast<unsigned char>(
              row[j].z() * std::numeric_limits<unsigned char>::max())};
      auto &pixel_depth =
          depth_image.at<unsigned char>(pixels.size() - i - 1, j);
      if (std::isinf(std::fabs(row[j].w()))) {
        pixel_depth = 0;
      } else {
        const float k =
            static_cast<float>(std::numeric_limits<unsigned char>::max()) * 0.5;
        const float b = k;
        pixel_depth = static_cast<unsigned char>(k * row[j].w() + b);
      }
    }
  }
  return depth_image;
}

bool MinusRenderer::load_mesh(const std::string &file_path) {
  objl::Loader loader{};
  bool load_status = loader.LoadFile(file_path.c_str());
  if (load_status) {
    for (const auto &mesh : loader.LoadedMeshes) {
      spdlog::info("Current mesh has {} vertices.", mesh.Vertices.size());
      Mesh m{};
      for (int i = 0; i < mesh.Vertices.size(); i += 3) {
        const auto objl_v0 = mesh.Vertices[mesh.Indices[i]].Position;
        Point3d v0(objl_v0.X, objl_v0.Y, objl_v0.Z);
        const auto objl_v1 = mesh.Vertices[mesh.Indices[i + 1]].Position;
        Point3d v1(objl_v1.X, objl_v1.Y, objl_v1.Z);
        const auto objl_v2 = mesh.Vertices[mesh.Indices[i + 2]].Position;
        Point3d v2(objl_v2.X, objl_v2.Y, objl_v2.Z);
        m.push_back(Triangle(v0, v1, v2));
      }
      meshes_.push_back(m);
    }
  }
  return load_status;
}