#include "color.h"
#include <algorithm>
#include <cmath>

constexpr double cie_epsilon = 216. / 24389.;
constexpr double cie_kappa = 24389. / 27.;

constexpr double D65_X = 0.95047;
constexpr double D65_Y = 1.;
constexpr double D65_Z = 1.08883;

static std::array<double, 3> linrgb_to_srgb(std::array<double, 3> linrgb);
static std::array<double, 3> srgb_to_linrgb(std::array<double, 3> srgb);

static std::array<double, 3> linrgb_to_xyz(std::array<double, 3> linrgb);
static std::array<double, 3> xyz_to_linrgb(std::array<double, 3> xyz);

static std::array<double, 3> cielab_to_xyz(std::array<double, 3> cielab);
static std::array<double, 3> xyz_to_cielab(std::array<double, 3> xyz);

static std::array<double, 3> cieluv_to_xyz(std::array<double, 3> cieluv);
static std::array<double, 3> xyz_to_cieluv(std::array<double, 3> xyz);

static std::array<double, 3> linrgb_to_oklab(std::array<double, 3> linrgb);
static std::array<double, 3> oklab_to_linrgb(std::array<double, 3> oklab);

static std::array<double, 3> xyzmatrix_mult(const std::array<std::array<double, 3>, 3> M, const std::array<double, 3> v);

ColorSpace Color::default_colorspace = ColorSpace::CIELAB;

std::array<double, 3> Color::values() const {
  return {v1, v2, v3};
}

std::array<uint8_t, 3> Color::to_rgb() const {
  std::array<double, 3> normalised = this->to_normalized_rgb();
  return {
      static_cast<uint8_t>(std::round(normalised[0]*255.)),
      static_cast<uint8_t>(std::round(normalised[1]*255.)),
      static_cast<uint8_t>(std::round(normalised[2]*255.)),
  };
}

std::array<double, 3> Color::to_normalized_rgb() const {
  std::array<double, 3> linrgb{};
  switch (colorspace) {
  case ColorSpace::CIELAB:
    linrgb = xyz_to_linrgb(cielab_to_xyz({v1, v2, v3}));
    break;
  case ColorSpace::CIELUV:
    linrgb = xyz_to_linrgb(cieluv_to_xyz({v1, v2, v3}));
    break;
  case ColorSpace::OKLAB:
    linrgb = oklab_to_linrgb({v1, v2, v3});
    break;
  }
  return linrgb_to_srgb(linrgb);
}

Color Color::from_rgb(const std::array<uint8_t, 3> rgb, ColorSpace colorspace) {
  return from_rgb(rgb[0], rgb[1], rgb[2], colorspace);
}

Color Color::from_rgb(double red, double green, double blue, ColorSpace colorspace) {
  return from_normalized_rgb(red / 255., green / 255., blue / 255., colorspace);
}

Color Color::from_rgb(const std::array<double, 3> rgb, ColorSpace colorspace) {
  return from_rgb(rgb[0], rgb[1], rgb[2], colorspace);
}

Color Color::from_normalized_rgb(double red, double green, double blue, ColorSpace colorspace) {
  return from_normalized_rgb({red, green, blue}, colorspace);
}

Color Color::from_normalized_rgb(const std::array<double, 3> rgb, ColorSpace colorspace) {
  auto linrgb = srgb_to_linrgb(rgb);
  std::array<double, 3> values{};
  switch (colorspace) {
  case ColorSpace::CIELAB:
    values = xyz_to_cielab(linrgb_to_xyz(linrgb));
    break;
  case ColorSpace::CIELUV:
    values = xyz_to_cieluv(linrgb_to_xyz(linrgb));
    break;
  case ColorSpace::OKLAB:
    values = linrgb_to_oklab(linrgb);
    break;
  }

  return {values, colorspace};
}

Color Color::from_rgb(uint8_t red, uint8_t green, uint8_t blue) {
  return from_rgb(red, green, blue, default_colorspace);
}

Color Color::from_rgb(const std::array<uint8_t, 3> rgb) {
  return from_rgb(rgb, default_colorspace);
}

Color Color::from_rgb(double red, double green, double blue) {
  return from_rgb(red, green, blue, default_colorspace);
}

Color Color::from_rgb(const std::array<double, 3> rgb) {
  return from_rgb(rgb, default_colorspace);
}

Color Color::from_normalized_rgb(double red, double green, double blue) {
  return from_normalized_rgb(red, green, blue, default_colorspace);
}

Color Color::from_normalized_rgb(const std::array<double, 3> rgb) {
  return from_normalized_rgb(rgb, default_colorspace);
}

void Color::set_default_colorspace(ColorSpace colorspace) {
  default_colorspace = colorspace;
}

std::array<double, 3> linrgb_to_srgb(const std::array<double, 3> linrgb) {
  std::array<double, 3> srgb{};
  std::transform(linrgb.begin(), linrgb.end(), srgb.begin(), [](double v) {
    if (v <= 0.0031308) {
      return v * 12.92;
    } else {
      return 1.055 * pow(v, 1. / 2.4) - 0.055;
    }
  });
  return srgb;
}

std::array<double, 3> srgb_to_linrgb(const std::array<double, 3> srgb) {
  std::array<double, 3> linrgb{};
  std::transform(srgb.begin(), srgb.end(), linrgb.begin(), [](double v) {
    if (v <= 0.04045) {
      return v / 12.92;
    } else {
      return pow(((v + 0.055) / 1.055), 2.4);
    }
  });
  return linrgb;
}

static std::array<double, 3> xyzmatrix_mult(const std::array<std::array<double, 3>, 3> M, const std::array<double, 3> v) {
  return {
      M[0][0] * v[0] + M[0][1] * v[1] + M[0][2] * v[2],
      M[1][0] * v[0] + M[1][1] * v[1] + M[1][2] * v[2],
      M[2][0] * v[0] + M[2][1] * v[1] + M[2][2] * v[2],
  };
}

constexpr std::array<std::array<double, 3>, 3> RGB_TO_XYZ_M = {
    std::array<double, 3>{0.4124564, 0.3575761, 0.1804375},
    std::array<double, 3>{0.2126729, 0.7151522, 0.0721750},
    std::array<double, 3>{0.0193339, 0.1191920, 0.9503041},
};

constexpr std::array<std::array<double, 3>, 3> XYZ_TO_RGB_M = {
    std::array<double, 3>{3.2404542, -1.5371385, -0.4985314},
    std::array<double, 3>{-0.9692660, 1.8760108, 0.0415560},
    std::array<double, 3>{0.0556434, -0.2040259, 1.0572252},
};

std::array<double, 3> linrgb_to_xyz(const std::array<double, 3> linrgb) {
  return xyzmatrix_mult(RGB_TO_XYZ_M, linrgb);
}
std::array<double, 3> xyz_to_linrgb(const std::array<double, 3> xyz) {
  return xyzmatrix_mult(XYZ_TO_RGB_M, xyz);
}
std::array<double, 3> cielab_to_xyz(const std::array<double, 3> cielab) {
  auto L = cielab[0];
  auto a = cielab[1];
  auto b = cielab[2];

  auto fy = (L + 16.) / 116.;
  auto fz = fy - b / 200.;
  auto fx = a / 500. + fy;

  double xr, yr, zr;
  if (pow(fx, 3.) > cie_epsilon) {
    xr = pow(fx, 3.);
  } else {
    xr = (116. * fx - 16.) / cie_kappa;
  }

  if (L > cie_kappa * cie_epsilon) {
    yr = pow((L + 16.) / 116., 3.);
  } else {
    yr = L / cie_kappa;
  }

  if (pow(fz, 3.) > cie_epsilon) {
    zr = pow(fz, 3.);
  } else {
    zr = (116. * fz - 16.) / cie_kappa;
  }

  auto X = xr * D65_X;
  auto Y = yr * D65_Y;
  auto Z = zr * D65_Z;

  return {X, Y, Z};
}
std::array<double, 3> xyz_to_cielab(const std::array<double, 3> xyz) {

  double xr = xyz[0] / D65_X;
  double yr = xyz[1] / D65_Y;
  double zr = xyz[2] / D65_Z;

  double fx, fy, fz;
  if (xr > cie_epsilon) {
    fx = cbrt(xr);
  } else {
    fx = (cie_kappa * xr + 16.) / 116.;
  }
  if (yr > cie_epsilon) {
    fy = cbrt(yr);
  } else {
    fy = (cie_kappa * yr + 16.) / 116.;
  }
  if (zr > cie_epsilon) {
    fz = cbrt(zr);
  } else {
    fz = (cie_kappa * zr + 16.) / 116.;
  }

  double L = 116. * fy - 16.;
  double a = 500. * (fx - fy);
  double b = 200. * (fy - fz);

  return {L, a, b};
}
std::array<double, 3> cieluv_to_xyz(const std::array<double, 3> cieluv) {
  constexpr double ur_ = (4. * D65_X) / (D65_X + 15. * D65_Y + 3. * D65_Z);
  constexpr double vr_ = (9. * D65_Y) / (D65_X + 15. * D65_Y + 3. * D65_Z);

  double L = cieluv[0];
  double u = cieluv[1];
  double v = cieluv[2];

  double Y;
  if (L > cie_kappa * cie_epsilon) {
    Y = pow((L + 16.) / 116., 3.);
  } else {
    Y = L / cie_kappa;
  }

  double a_ = u + 13. * L * ur_;
  double a = a_ != 0. ? 1. / 3. * (52. * L / a_ - 1.) : 0.;
  double b = -5. * Y;
  double c = -1. / 3.;
  double d_ = v + 13. * L * vr_;
  double d = d_ != 0. ? Y * (39. * L / d_ - 5.) : 0.;

  double X = (d - b) / (a - c);
  double Z = X * a + b;
  return {X, Y, Z};
}
std::array<double, 3> xyz_to_cieluv(const std::array<double, 3> xyz) {
  constexpr double ur_ = (4. * D65_X) / (D65_X + 15. * D65_Y + 3. * D65_Z);
  constexpr double vr_ = (9. * D65_Y) / (D65_X + 15. * D65_Y + 3. * D65_Z);

  double X = xyz[0];
  double Y = xyz[1];
  double Z = xyz[2];

  double yr = Y / D65_Y;
  double D = (X + 15. * Y + 3. * Z);
  double u_, v_;
  if (D != 0) {
    u_ = (4. * X) / D;
    v_ = (9. * Y) / D;
  } else {
    u_ = 0;
    v_ = 0;
  }

  double L;
  if (yr > cie_epsilon) {
    L = 116. * cbrt(yr) - 16.;
  } else {
    L = cie_kappa * yr;
  }
  double u = 13. * L * (u_ - ur_);
  double v = 13. * L * (v_ - vr_);

  return {L, u, v};
}
std::array<double, 3> linrgb_to_oklab(const std::array<double, 3> linrgb) {
  // https://bottosson.github.io/posts/oklab/#converting-from-linear-srgb-to-oklab
  double l = 0.4122214708 * linrgb[0] + 0.5363325363 * linrgb[1] + 0.0514459929 * linrgb[2];
  double m = 0.2119034982 * linrgb[0] + 0.6806995451 * linrgb[1] + 0.1073969566 * linrgb[2];
  double s = 0.0883024619 * linrgb[0] + 0.2817188376 * linrgb[1] + 0.6299787005 * linrgb[2];

  double l_ = cbrt(l);
  double m_ = cbrt(m);
  double s_ = cbrt(s);

  return {
      0.2104542553 * l_ + 0.7936177850 * m_ - 0.0040720468 * s_,
      1.9779984951 * l_ - 2.4285922050 * m_ + 0.4505937099 * s_,
      0.0259040371 * l_ + 0.7827717662 * m_ - 0.8086757660 * s_,
  };
}
std::array<double, 3> oklab_to_linrgb(const std::array<double, 3> oklab) {
  double l_ = oklab[0] + 0.3963377774 * oklab[1] + 0.2158037573 * oklab[2];
  double m_ = oklab[0] - 0.1055613458 * oklab[1] - 0.0638541728 * oklab[2];
  double s_ = oklab[0] - 0.0894841775 * oklab[1] - 1.2914855480 * oklab[2];

  double l = l_ * l_ * l_;
  double m = m_ * m_ * m_;
  double s = s_ * s_ * s_;

  return {
      +4.0767416621 * l - 3.3077115913 * m + 0.2309699292 * s,
      -1.2684380046 * l + 2.6097574011 * m - 0.3413193965 * s,
      -0.0041960863 * l - 0.7034186147 * m + 1.7076147010 * s,
  };
}