#include "calib/hmd_config.h"

#include <fstream>
#include <sstream>

#include "nlohmann/json.hpp"

namespace sauna {
namespace {

using nlohmann::json;

void readVec3(const json& j, double out[3]) {
  for (int i = 0; i < 3; i++) out[i] = j.at(i).get<double>();
}

void readMat3(const json& j, double out[3][3]) {
  for (int r = 0; r < 3; r++)
    for (int c = 0; c < 3; c++) out[r][c] = j.at(r).at(c).get<double>();
}

void readFrame(const json& j, FrameCalib* out) {
  readVec3(j.at("plus_x"), out->plus_x);
  readVec3(j.at("plus_z"), out->plus_z);
  readVec3(j.at("position"), out->position);
}

void readChannel(const json& j, DistortionChannel* out) {
  const std::string type = j.at("type").get<std::string>();
  if (type != "DISTORT_POLY3")
    throw std::runtime_error("unsupported distortion type " + type +
                             " (only DISTORT_POLY3 is S3-verified)");
  out->center_x = j.at("center_x").get<double>();
  out->center_y = j.at("center_y").get<double>();
  const json& k = j.at("coeffs");
  for (int i = 0; i < 3; i++) out->k[i] = k.at(i).get<double>();
}

}  // namespace

bool LoadHmdConfigFromString(const std::string& text, HmdConfig* out,
                             std::string* err) {
  try {
    json cfg = json::parse(text);

    out->serial = cfg.at("device_serial_number").get<std::string>();
    out->model = cfg.value("model_number", std::string());

    const json& dev = cfg.at("device");
    out->eye_width_px = dev.at("eye_target_width_in_pixels").get<int>();
    out->eye_height_px = dev.at("eye_target_height_in_pixels").get<int>();

    out->parallel_render_cameras = cfg.value("parallel_render_cameras", false);
    out->ipd_default_mm = cfg.at("ipd").at("default_mm").get<double>();
    out->seconds_from_vsync_to_photons =
        cfg.value("seconds_from_vsync_to_photons", 0.0);
    out->seconds_from_photons_to_vblank =
        cfg.value("seconds_from_photons_to_vblank", 0.0);

    readFrame(cfg.at("head"), &out->head);

    const json& imu = cfg.at("imu");
    readFrame(imu, &out->imu.frame);
    readVec3(imu.at("acc_bias"), out->imu.acc_bias);
    readVec3(imu.at("acc_scale"), out->imu.acc_scale);
    readVec3(imu.at("gyro_bias"), out->imu.gyro_bias);
    out->imu.has_gyro_scale = imu.contains("gyro_scale");
    if (out->imu.has_gyro_scale)
      readVec3(imu.at("gyro_scale"), out->imu.gyro_scale);

    const json& tte = cfg.at("tracking_to_eye_transform");
    if (tte.size() != 2) throw std::runtime_error("expected 2 eyes");
    const json& colorMult = cfg.at("display_color_mult");
    static const char* kChannelKeys[3] = {"distortion_red", "distortion",
                                          "distortion_blue"};
    for (int e = 0; e < 2; e++) {
      EyeCalib* eye = &out->eye[e];
      const json& je = tte.at(e);
      for (int ch = 0; ch < 3; ch++)
        readChannel(je.at(kChannelKeys[ch]), &eye->rgb[ch]);
      eye->grow_for_undistort = je.at("grow_for_undistort").get<double>();
      eye->undistort_r2_cutoff = je.at("undistort_r2_cutoff").get<double>();
      readMat3(je.at("intrinsics"), eye->intrinsics);
      readMat3(je.at("eye_to_head"), eye->eye_to_head);
      readVec3(colorMult.at(e), eye->color_mult);
    }
  } catch (const std::exception& ex) {
    if (err) *err = std::string("config parse error: ") + ex.what();
    return false;
  }
  return true;
}

bool LoadHmdConfig(const std::string& path, HmdConfig* out, std::string* err) {
  std::ifstream f(path);
  if (!f) {
    if (err) *err = "cannot open " + path;
    return false;
  }
  std::stringstream ss;
  ss << f.rdbuf();
  if (!LoadHmdConfigFromString(ss.str(), out, err)) {
    if (err) *err += " (" + path + ")";
    return false;
  }
  return true;
}

void ProjectionRawTangents(const EyeCalib& eye, double* L, double* R,
                           double* T, double* B) {
  const double G = 1.0 + eye.grow_for_undistort;
  const double cx = -eye.intrinsics[0][2];  // S3 sign convention
  const double cy = eye.intrinsics[1][2];
  const double fx = eye.intrinsics[0][0];
  const double fy = eye.intrinsics[1][1];
  const double tphiX = eye.eye_to_head[0][2] / eye.eye_to_head[2][2];
  const double tphiY = eye.eye_to_head[1][2] / eye.eye_to_head[2][2];
  auto fold = [](double t, double tphi) {
    return (t - tphi) / (1.0 + t * tphi);
  };
  *L = fold((-G - cx) / fx, tphiX);
  *R = fold((G - cx) / fx, tphiX);
  *T = -fold((G - cy) / fy, -tphiY);
  *B = -fold((-G - cy) / fy, -tphiY);
}

void EvalDistortion(const EyeCalib& eye, double u, double v,
                    double out_uv[3][2]) {
  const double scale = 0.5 / (1.0 + eye.grow_for_undistort);
  for (int ch = 0; ch < 3; ch++) {
    const DistortionChannel& c = eye.rgb[ch];
    const double tx = 2.0 * u - 1.0 - c.center_x;
    const double ty = 2.0 * v - 1.0 - c.center_y;
    // No r2 clamp: ComputeDistortion ignores undistort_r2_cutoff (S3).
    const double r2 = tx * tx + ty * ty;
    const double d = 1.0 + r2 * (c.k[0] + r2 * (c.k[1] + r2 * c.k[2]));
    out_uv[ch][0] = 0.5 + (tx * d + c.center_x) * scale;
    out_uv[ch][1] = 0.5 + (ty * d + c.center_y) * scale;
  }
}

}  // namespace sauna