#include <iostream>
#include <array>
#include <pybind11/numpy.h>
#include <pybind11/stl.h>
#include <tbb/tbb.h>

#include <ceres/ceres.h>
#include <glog/logging.h>

#include "cie1931.h"

namespace py = pybind11;

using Array = py::array_t<double>;

#define NDIM        3
#define BILERP_RES  6

/// Discretization of quadrature scheme
#define CIE_FINE_SAMPLES ((CIE_SAMPLES - 1) * 3 + 1)

/// Precomputed tables for fast spectral -> RGB conversion
double lambda_tbl[CIE_FINE_SAMPLES],
       rgb_tbl[3][CIE_FINE_SAMPLES],
       rgb_to_xyz[3][3],
       xyz_to_rgb[3][3],
       xyz_whitepoint[3];

/// Currently supported gamuts
enum Gamut {
    SRGB,
    ProPhotoRGB,
    ACES2065_1,
    REC2020,
    ERGB,
    XYZ,
};

template <typename T> T sigmoid(T x) {
#if SIGMOID_EXP
    return 1.0 / (1.0 + exp(-x));
#else
    return 0.5 * x / sqrt(1.0 + x*x) + 0.5;
#endif
}

template <typename T> void cie_lab(T *p) {
    T X = T(0.0), Y = T(0.0), Z = T(0.0),
      Xw = T(xyz_whitepoint[0]),
      Yw = T(xyz_whitepoint[1]),
      Zw = T(xyz_whitepoint[2]);

    for (int j = 0; j < 3; ++j) {
        X += p[j] * T(rgb_to_xyz[0][j]);
        Y += p[j] * T(rgb_to_xyz[1][j]);
        Z += p[j] * T(rgb_to_xyz[2][j]);
    }

    auto f = [](T t) -> T {
        T delta = T(6.0 / 29.0);
        if (t > delta*delta*delta)
            return pow(t, 1.0/3.0);//cbrt(t);
        else
            return t / (delta*delta * T(3.0)) + T(4.0 / 29.0);
    };

    p[0] = T(116.0) * f(Y / Yw) - T(16.0);
    p[1] = T(500.0) * (f(X / Xw) - f(Y / Yw));
    p[2] = T(200.0) * (f(Y / Yw) - f(Z / Zw));
}

/**
 * This function precomputes tables used to convert arbitrary spectra
 * to RGB (either sRGB or ProPhoto RGB)
 *
 * A composite quadrature rule integrates the CIE curves, reflectance, and
 * illuminant spectrum over each 5nm segment in the 360..830nm range using
 * Simpson's 3/8 rule (4th-order accurate), which evaluates the integrand at
 * four positions per segment. While the CIE curves and illuminant spectrum are
 * linear over the segment, the reflectance could have arbitrary behavior,
 * hence the extra precations.
 */
void init_tables(Gamut gamut) {
    memset(rgb_tbl, 0, sizeof(rgb_tbl));
    memset(xyz_whitepoint, 0, sizeof(xyz_whitepoint));

    double h = (CIE_LAMBDA_MAX - CIE_LAMBDA_MIN) / (CIE_FINE_SAMPLES - 1);

    const double *illuminant = nullptr;

    switch (gamut) {
        case SRGB:
            illuminant = cie_d65;
            memcpy(xyz_to_rgb, xyz_to_srgb, sizeof(double) * 9);
            memcpy(rgb_to_xyz, srgb_to_xyz, sizeof(double) * 9);
            break;

        case ERGB:
            illuminant = cie_e;
            memcpy(xyz_to_rgb, xyz_to_ergb, sizeof(double) * 9);
            memcpy(rgb_to_xyz, ergb_to_xyz, sizeof(double) * 9);
            break;

        case XYZ:
            illuminant = cie_e;
            memcpy(xyz_to_rgb, xyz_to_xyz, sizeof(double) * 9);
            memcpy(rgb_to_xyz, xyz_to_xyz, sizeof(double) * 9);
            break;

        case ProPhotoRGB:
            illuminant = cie_d50;
            memcpy(xyz_to_rgb, xyz_to_prophoto_rgb, sizeof(double) * 9);
            memcpy(rgb_to_xyz, prophoto_rgb_to_xyz, sizeof(double) * 9);
            break;

        case ACES2065_1:
            illuminant = cie_d60;
            memcpy(xyz_to_rgb, xyz_to_aces2065_1, sizeof(double) * 9);
            memcpy(rgb_to_xyz, aces2065_1_to_xyz, sizeof(double) * 9);
            break;

        case REC2020:
            illuminant = cie_d65;
            memcpy(xyz_to_rgb, xyz_to_rec2020, sizeof(double) * 9);
            memcpy(rgb_to_xyz, rec2020_to_xyz, sizeof(double) * 9);
            break;
    }

    for (int i = 0; i < CIE_FINE_SAMPLES; ++i) {
        double lambda = CIE_LAMBDA_MIN + i * h;

        double xyz[3] = { cie_interp(cie_x, lambda),
                          cie_interp(cie_y, lambda),
                          cie_interp(cie_z, lambda) },
               I = cie_interp(illuminant, lambda);

        double weight = 3.0 / 8.0 * h;
        if (i == 0 || i == CIE_FINE_SAMPLES - 1)
            ;
        else if ((i - 1) % 3 == 2)
            weight *= 2.f;
        else
            weight *= 3.f;

        lambda_tbl[i] = lambda;
        for (int k = 0; k < 3; ++k)
            for (int j = 0; j < 3; ++j)
                rgb_tbl[k][i] += xyz_to_rgb[k][j] * xyz[j] * I * weight;

        for (int i = 0; i < 3; ++i)
            xyz_whitepoint[i] += xyz[i] * I * weight;
    }
}

class PointObjective {
public:
    PointObjective(const double *rgb, bool lab) : m_cie_lab(lab) {
        memcpy(m_rgb, rgb, sizeof(double) * 3);
        if (m_cie_lab)
            cie_lab(m_rgb);
    }

    template <typename T> bool operator()(const T *const coeffs, T *residuals) const {
        for (int j = 0; j < 3; ++j)
            residuals[j] = T(0.0);

        for (int i = 0; i < CIE_FINE_SAMPLES; ++i) {
            /* Scale lambda to 0..1 range */
            T lambda = T((lambda_tbl[i] - CIE_LAMBDA_MIN) /
                         (CIE_LAMBDA_MAX - CIE_LAMBDA_MIN));

            /* Polynomial */
            T x = T(0.f);
            for (int i = 0; i < NDIM; ++i)
                x = x * lambda + coeffs[i];

            /* Sigmoid */
            T s = sigmoid(x);

            /* Integrate against precomputed curves */
            for (int j = 0; j < 3; ++j)
                residuals[j] += rgb_tbl[j][i] * s;
        }

        if (m_cie_lab)
            cie_lab(residuals);

        for (int j = 0; j < 3; ++j)
            residuals[j] -= T(m_rgb[j]);

        return true;
    }
private:
    double m_rgb[3];
    bool m_cie_lab;
};

class SimilarityObjective {
public:
    SimilarityObjective(double weight, const double *values) : m_weight(weight) {
        memcpy(m_values, values, sizeof(double) * 3);
    }

    template <typename T> bool operator()(const T *const coeffs, T *residuals) const {
        for (int i = 0; i < 3; ++i)
            residuals[i] = (T(m_values[i]) - coeffs[i]) * T(m_weight);
        return true;
    }
private:
    double m_values[3];
    double m_weight;
};

/**
 * Optimizes wrt. the above objectives
 *
 * 'input' is a I1 x I2 x 3 tensor containing target RGB values
 *
 * 'output' is a I1 x I2 x NDIM tensor containing starting guesses initially,
 * and which will be overwritten with the resulting fits
 *
 * If 'lerp' is true, the implementation will optimize values in a local region,
 * taking into account that they will be bilinearly interpolated.
 *
 * If 'bound' is nonzero, the optimization is constrained to find
 * polynomial coefficients int the range [-bound..bound]
 */
Array optimize(const Array &input_, Array output_, double bound = 0,
               bool cie_lab = false, double similarity_constraint = 0.0) {
    using namespace ceres;

    auto input  = input_.unchecked<3>();
    auto output = output_.mutable_unchecked<3>();

    if (input.shape(0) != output.shape(0) ||
        input.shape(1) != output.shape(1) ||
        input.shape(2) != 3 ||
        output.shape(2) != NDIM)
        throw std::runtime_error("Input/output array shape mismatch!");

    Array error_({input.shape(0), input.shape(1)});
    auto error = error_.mutable_unchecked<2>();

    tbb::parallel_for(
        tbb::blocked_range<ssize_t>(0, input.shape(0), 1),
        [&](const tbb::blocked_range<ssize_t> &range) {
            for (ssize_t i = range.begin(); i < range.end(); ++i) {
                for (ssize_t j = 0; j < input.shape(1); ++j) {
                    double *x = output_.mutable_data(i, j, 0);
                    const double *y = input.data(i, j, 0);

                    Problem problem;

                    PointObjective *fct = new PointObjective(y, cie_lab);
                    {
                        CostFunction *cost = new AutoDiffCostFunction<PointObjective, 3, NDIM>(fct);
                        problem.AddResidualBlock(cost, nullptr, x);
                    }
                    
                    if (similarity_constraint != 0.0) {
                        SimilarityObjective *fct2 = new SimilarityObjective(similarity_constraint, x);
                        CostFunction *cost = new AutoDiffCostFunction<SimilarityObjective, 3, NDIM>(fct2);
                        problem.AddResidualBlock(cost, nullptr, x);
                    }

                    if (bound != 0) {
                        for (size_t k = 0; k < NDIM; ++k) {
                            problem.SetParameterLowerBound(x, k, -bound);
                            problem.SetParameterUpperBound(x, k,  bound);
                        }
                    }

                    Solver::Options options;
                    Solver::Summary summary;
                    Solve(options, &problem, &summary);
                    if (!summary.IsSolutionUsable())
                        throw std::runtime_error("Optimization failed: " + summary.BriefReport());

                    double r[3];
                    (*fct)(x, r);
                    error(i, j) = std::sqrt(r[0]*r[0] + r[1]*r[1] + r[2]*r[2]);
                }
            }
        }
    );
    return error_;
}

/// Return the RGB value associated with a given set of coefficients
Array eval_rgb(const Array &coeffs) {
    if (coeffs.ndim() != 1)
        throw std::runtime_error("Expected a vector as input!");

    Array result(3);
    memset(result.mutable_data(), 0, sizeof(double) * 3);

    for (int i = 0; i < CIE_FINE_SAMPLES; ++i) {
        /* Scale lambda to 0..1 range */
        double lambda = (lambda_tbl[i] - CIE_LAMBDA_MIN) /
                        (CIE_LAMBDA_MAX - CIE_LAMBDA_MIN);

        /* Polynomial */
        double x = 0.f;
        for (int i = 0; i < NDIM; ++i)
            x = x * lambda + coeffs.at(i);

        /* Sigmoid */
        double s = sigmoid(x);

        /* Integrate against precomputed curves */
        for (int j = 0; j < 3; ++j)
            result.mutable_at(j) += rgb_tbl[j][i] * s;
    }

    return result;
}

/// Return the RGB value associated with a given spectral distribution
Array spec_to_rgb(const Array &coeffs) {
    if (coeffs.ndim() != 1)
        throw std::runtime_error("Expected a vector as input!");

    if (coeffs.shape(0) != CIE_FINE_SAMPLES)
        throw std::runtime_error("Unsupported input size!");

    auto data = coeffs.unchecked<1>();

    Array result(3);
    memset(result.mutable_data(), 0, sizeof(double) * 3);

    for (int i = 0; i < CIE_FINE_SAMPLES; ++i) {
        double value = data[i];
        /* Integrate against precomputed curves */
        for (int j = 0; j < 3; ++j)
            result.mutable_at(j) += rgb_tbl[j][i] * value;
    }

    return result;
}

/// Evaluate the spectrum associated with a given set of coefficients
Array eval_spec(const Array &coeffs, const Array &wavelengths) {
    if (wavelengths.ndim() != 1 || coeffs.ndim() != 1)
        throw std::runtime_error("Expected vectors as input!");

    Array result(wavelengths.shape(0));

    for (int i = 0; i < wavelengths.shape(0); ++i) {
        /* Scale lambda to 0..1 range */
        double lambda = (wavelengths.at(i) - CIE_LAMBDA_MIN) /
                        (CIE_LAMBDA_MAX - CIE_LAMBDA_MIN);

        /* Polynomial */
        double x = 0.f;
        for (int i = 0; i < NDIM; ++i)
            x = x * lambda + coeffs.at(i);

        /* Sigmoid */
        double s = sigmoid(x);

        result.mutable_at(i) = s;
    }

    return result;
}

Array spec_wavelengths() {
    Array result(CIE_FINE_SAMPLES);

    for (int i = 0; i < CIE_FINE_SAMPLES; ++i) {
        double lambda =
            CIE_LAMBDA_MIN +
            i * ((CIE_LAMBDA_MAX - CIE_LAMBDA_MIN) / (CIE_FINE_SAMPLES - 1));
        result.mutable_at(i) = lambda;
    }

    return result;
}

/* Python bindings for the above functions */
PYBIND11_MODULE(spectra, m) {
    google::InitGoogleLogging("python");

    py::enum_<Gamut>(m, "Gamut")
        .value("SRGB", Gamut::SRGB)
        .value("ProPhotoRGB", Gamut::ProPhotoRGB)
        .value("ACES2065_1", Gamut::ACES2065_1)
        .value("REC2020", Gamut::REC2020)
        .value("ERGB", Gamut::ERGB)
        .value("XYZ", Gamut::XYZ);

    m.def("init_tables", &init_tables);

    m.def("optimize", &optimize, py::arg("input"), py::arg("output"),
          py::arg("bound") = 500, py::arg("cie_lab") = false,
          py::arg("similarity_constraint") = 0);

    m.def("eval_rgb", &eval_rgb, py::arg("coeffs"));
    m.def("spec_to_rgb", &spec_to_rgb, py::arg("coeffs"));
    m.def("wavelengths", &spec_wavelengths);

    m.def("eval_spec", &eval_spec, py::arg("coeffs"), py::arg("wavelengths"));

    // Return the currently selected RGB<->XYZ conversion matrices
    m.def("xyz_to_rgb", [] {
        Array result({ 3, 3 });
        memcpy(result.mutable_data(), xyz_to_rgb, sizeof(double) * 9);
        return result;
    });
    m.def("rgb_to_xyz", [] {
        Array result({ 3, 3 });
        memcpy(result.mutable_data(), rgb_to_xyz, sizeof(double) * 9);
        return result;
    });
    m.def("xyz_whitepoint", [] {
        Array result(3);
        memcpy(result.mutable_data(), xyz_whitepoint, sizeof(double) * 3);
        return result;
    });
}
