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

#include <cmath>

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

#include "minus_renderer.h"

MinusRenderer::MinusRenderer(float near, float far, float fov,
                             float aspect_ratio)
    : near_(near), far_(far), fov_(fov), aspect_ratio_(aspect_ratio),
      rasterizer_(std::make_unique<Rasterizer>()) {
  spdlog::info("near: {}, far: {}, fov: {}, aspect ratio: {}", near_, far_,
               fov_, aspect_ratio_);
}

void MinusRenderer::set_meshes(const std::vector<Mesh> &meshes) {
  meshes_ = meshes;
}

void MinusRenderer::set_lamps(
    const std::vector<std::shared_ptr<PointLamp>> &lamps) {
  lamps_ = lamps;
}

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;
}

std::pair<TransformMatrix, TransformMatrix>
MinusRenderer::construct_transform(const int resolution_width,
                                   const int resolution_height) const {
  // Squish Transform
  const TransformMatrix &squish_transform{{near_, 0, 0, 0},
                                          {0, near_, 0, 0},
                                          {0, 0, near_ + far_, -near_ * far_},
                                          {0, 0, 1, 0}};
  const TransformMatrix &inv_squish_transform{
      {1, 0, 0, 0},
      {0, 1, 0, 0},
      {0, 0, 0, near_},
      {0, 0, -1. / far_, (near_ + far_) / far_}};

  // 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;

  const TransformMatrix &ortho_transform =
      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}} *
      TransformMatrix{{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}};
  const TransformMatrix &inv_ortho_transform{
      {(right - left) / 2., 0, 0, 0.5 * (right - left)},
      {0, (top - bottom) / 2., 0, 0.5 * (top - bottom)},
      {0, 0, std::fabs(far_ - near_) / 2., 0.5 * (far_ - near_)},
      {0, 0, 0, 1}};

  // View port transform
  const TransformMatrix &view_port_transform{
      {resolution_width / 2., 0, 0, resolution_width * 0.5},
      {0, resolution_height / 2., 0, resolution_height * 0.5},
      {0, 0, 1, 0},
      {0, 0, 0, 1}};
  const TransformMatrix &inv_view_port_transform =
      TransformMatrix{{2. / resolution_width, 0, 0, 0},
                      {0, 2. / resolution_height, 0, 0},
                      {0, 0, 1, 0},
                      {0, 0, 0, 1}} *
      TransformMatrix{{1, 0, 0, -resolution_width * 0.5},
                      {0, 1, 0, -resolution_height * 0.5},
                      {0, 0, 1, 0},
                      {0, 0, 0, 1}};

  const auto &ss_transform =
      view_port_transform * ortho_transform * squish_transform;
  const auto &inv_ss_transform =
      inv_squish_transform * inv_ortho_transform * inv_view_port_transform;
  return {ss_transform, inv_ss_transform};
}

std::tuple<double, double, double>
MinusRenderer::calculate_barycentric_coordinate(const Triangle &t,
                                                const Point3d &p) {
  const auto &points = t.get_vertex_position();
  const auto &pa = points[0];
  const auto &pb = points[1];
  const auto &pc = points[2];
  const auto &triangle_normal = t.normal_vector();

  const double squared_area = triangle_normal.norm() * triangle_normal.norm();
  const double alpha =
      (triangle_normal.dot((pc - pb).cross(p - pb))) / squared_area;
  const double beta =
      (triangle_normal.dot((pa - pc).cross(p - pc))) / squared_area;
  const double gamma =
      (triangle_normal.dot((pb - pa).cross(p - pa))) / squared_area;

  return {alpha, beta, gamma};
}

void MinusRenderer::model_transform(const TransformMatrix &mtx) {
  for (auto &m : meshes_) {
    for (auto &v : m.get_vertices()) {
      v.position = (mtx * v.position.homogeneous()).hnormalized();
      v.normal = (mtx * v.normal.homogeneous()).hnormalized();
    }
  }
}

void MinusRenderer::view_transform(const TransformMatrix &mtx) {
  for (auto &m : meshes_) {
    for (auto &v : m.get_vertices()) {
      v.position = (mtx * v.position.homogeneous()).hnormalized();
      v.normal = (mtx * v.normal.homogeneous()).hnormalized();
    }
  }
}

cv::Vec3b MinusRenderer::fragment_shader(const RasterizerResult &fragment,
                                         const Point3d &position,
                                         const Mesh &mesh) {
  const auto &triangle = mesh.get_primitive(fragment.triangle.triangle_index);
  const auto &[alpha, beta, gamma] =
      calculate_barycentric_coordinate(triangle, position);

  cv::Vec3b res{0, 0, 0};
  if (alpha >= 0. && alpha <= 1. && beta >= 0. && beta <= 1. && gamma >= 0. &&
      gamma <= 1.) {
    std::vector<Point2d> triangle_uv;
    std::vector<Point3d> triangle_normal;
    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));
    }

    Point2d uv{alpha * triangle_uv[0].x() + beta * triangle_uv[1].x() +
                   gamma * triangle_uv[2].x(),
               alpha * triangle_uv[0].y() + beta * triangle_uv[1].y() +
                   gamma * triangle_uv[2].y()};
    const auto &normal =
        -1. *
        Point3d(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())
            .normalized();

    const auto &material = mesh.get_material();
    res = material->shade(Vertex{position, normal, uv}, Point3d{}, lamps_);
  }
  return res;
}

cv::Mat MinusRenderer::render(const int resolution_width,
                              const int resolution_height) {
  spdlog::trace("Tick once.");

  // 构造 camera space -> screen space 的变换矩阵
  const auto [ss_transform, inv_ss_transform] =
      construct_transform(resolution_width, resolution_height);

  spdlog::trace("Rasterize.");
  rasterizer_->reset(resolution_width, resolution_height);
  for (decltype(meshes_)::size_type mesh_index = 0; mesh_index < meshes_.size();
       ++mesh_index) {
    auto primitives = (meshes_[mesh_index]).get_primitives();
    for (auto &t : primitives) {
      const auto [z_min, z_max] = t.get_depth_range();
      if (z_max >= far_ && z_min <= near_) {
        t.set_points(apply_transform(ss_transform, t.get_vertex_position()));
      }
    }
    rasterizer_->rasterize(mesh_index, primitives);
  }

  const auto &shading_points = rasterizer_->get_shading_points();
  assert(!shading_points.empty());
  cv::Mat color_image(shading_points.size(), (shading_points[0]).size(),
                      CV_8UC3);
  spdlog::trace("Shading.");
  for (size_t y = 0; y < shading_points.size(); ++y) {
    for (size_t 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 &fragment = shading_points[y][x];
      if (fragment.triangle.mesh_index >= 0 &&
          fragment.triangle.mesh_index < meshes_.size()) {
        const auto &mesh = meshes_[fragment.triangle.mesh_index];
        if (fragment.depth <= 1. && fragment.depth >= -1.) {
          const Point3d &position =
              (inv_ss_transform *
               (Point3d{x + 0.5, y + 0.5, fragment.depth}.homogeneous()))
                  .hnormalized();

          pixel_color = fragment_shader(fragment, position, mesh);
        }
      }
    }
  }
  return color_image;
}
