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

#include <cmath>

#include "spdlog/spdlog.h"
#include <opencv2/highgui/highgui.hpp>

#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_transform() * squish_transform()) {
  spdlog::info("near: {}, far: {}, fov: {}, aspect ratio: {}", near_, far_,
               fov_, aspect_ratio_);
}

void MinusRenderer::set_meshes(const std::vector<Mesh> &meshes) {
  meshes_ = meshes;
  spdlog::info("Mesh size: {}, the primitives size of the first mesh: {}",
               meshes_.size(), meshes_.front().get_primitives().size());
}

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_transform() {
  // 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_transform() {
  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}};
}

std::tuple<double, double, double>
MinusRenderer::calculate_barycentric_coordinate(const Triangle &t,
                                                const Point2d &p) {
  const auto &points = t.get_vertex_position();
  const auto &A = (points[0]).head(2);
  const auto &B = (points[1]).head(2);
  const auto &C = (points[2]).head(2);
  double alpha =
      (-(p.x() - B.x()) * (C.y() - B.y()) + (p.y() - B.y()) * (C.x() - B.x())) /
      (-(A.x() - B.x()) * (C.y() - B.y()) + (A.y() - B.y()) * (C.x() - B.x()));
  double beta =
      (-(p.x() - C.x()) * (A.y() - C.y()) + (p.y() - C.y()) * (A.x() - C.x())) /
      (-(B.x() - C.x()) * (A.y() - C.y()) + (B.y() - C.y()) * (A.x() - C.x()));
  double gamma = 1. - alpha - beta;
  return {alpha, beta, gamma};
}

void MinusRenderer::model_transform(const TransformMatrix &mtx) {
  for (auto &m : meshes_) {
    for (auto &t : m.get_primitives()) {
      const auto &[res, _] = apply_transform(mtx, t.get_vertex_position());
      t.set_points(res);
    }
  }
}

void MinusRenderer::view_transform(const TransformMatrix &mtx) {
  for (auto &m : meshes_) {
    for (auto &t : m.get_primitives()) {
      const auto &[res, _] = apply_transform(mtx, t.get_vertex_position());
      t.set_points(res);
    }
  }
}

cv::Mat MinusRenderer::render(const int resolution_width,
                              const int resolution_height) {
  std::vector<Mesh> meshes = meshes_;
  std::vector<std::vector<double>> perspective_coeff(meshes.size());
  Rasterizer rasterizer{resolution_width, resolution_height};
  for (int mesh_index = 0; mesh_index < meshes.size(); ++mesh_index) {
    auto &mesh = meshes[mesh_index];
    auto &primitives = mesh.get_primitives();
    perspective_coeff[mesh_index] =
        std::vector<double>(mesh.get_vertex_size(), 1.);
    for (int i = 0; i < primitives.size(); ++i) {
      auto &t = primitives[i];
      const auto &triangle_vertices = t.get_vertex_position();
      const auto &[projective_res, w] =
          apply_transform(projection_matrix_, triangle_vertices);
      const auto &indices = t.get_vertex_index();
      for (int j = 0; j < indices.size(); ++j) {
        perspective_coeff[mesh_index][indices[j]] = w[j];
      }
      const auto &[res, _] = apply_transform(
          view_port_transform(resolution_width, resolution_height),
          projective_res);
      t.set_points(res);
    }
  }
  const auto &shading_points = rasterizer.rasterize(meshes);
  assert(!shading_points.empty());
  cv::Mat color_image(shading_points.size(), (shading_points[0]).size(),
                      CV_8UC3);
  for (int y = 0; y < shading_points.size(); ++y) {
    for (int x = 0; x < shading_points[y].size(); ++x) {
      auto &pixel_color =
          color_image.at<cv::Vec3b>(shading_points.size() - y - 1, x);
      pixel_color = cv::Vec3b(0, 0, 0);

      const auto &p = shading_points[y][x];
      if (p.triangle.mesh_index >= 0 && p.triangle.mesh_index < meshes.size()) {
        const auto &mesh = meshes[p.triangle.mesh_index];
        const auto &primitives = mesh.get_primitives();
        if (p.triangle.triangle_index >= 0 &&
            p.triangle.triangle_index < primitives.size()) {
          const auto &triangle = primitives[p.triangle.triangle_index];
          const auto &[alpha, beta, gamma] = calculate_barycentric_coordinate(
              triangle, Point2d{x + 0.5, y + 0.5});

          std::vector<Point2d> triangle_uv;
          std::vector<Point3d> triangle_normal;
          std::vector<double> w_coeff;
          for (const auto vertex_index : triangle.get_vertex_index()) {
            triangle_uv.push_back(mesh.get_texture_coordinate(vertex_index));
            triangle_normal.push_back(mesh.get_normal_vector(vertex_index));
            w_coeff.push_back(
                perspective_coeff[p.triangle.mesh_index][vertex_index]);
          }

          // TODO(tianlei): 推导一下为什么需要除以下面这个
          // 在 screen space 对 1/w 的插值就相当于在 world space 对 w 的插值
          auto w_reciprocal = alpha * (1. / w_coeff[0]) +
                              beta * (1. / w_coeff[1]) +
                              gamma * (1. / w_coeff[2]);
          Point2d uv{(alpha * (triangle_uv[0].x() / w_coeff[0]) +
                      beta * (triangle_uv[1].x() / w_coeff[1]) +
                      gamma * (triangle_uv[2].x() / w_coeff[2])) /
                         w_reciprocal,
                     (alpha * (triangle_uv[0].y() / w_coeff[0]) +
                      beta * (triangle_uv[1].y() / w_coeff[1]) +
                      gamma * (triangle_uv[2].y() / w_coeff[2])) /
                         w_reciprocal};

          Point3d normal{
              alpha * triangle_normal[0].x() + beta * triangle_normal[1].x() +
                  gamma * triangle_normal[2].x(),
              alpha * triangle_normal[0].y() + beta * triangle_normal[1].y() +
                  gamma * triangle_normal[2].y(),
              alpha * triangle_normal[0].z() + beta * triangle_normal[1].z() +
                  gamma * triangle_normal[2].z()};

          auto &material = mesh.get_material();
          pixel_color = material->shade(
              Vertex{Point3d{x, y, p.depth}, normal, uv}, Point3d{});
        }
      }
    }
  }
  return color_image;
}