#include "toonzqt/lutcalibrator.h"

// Tnzlib includes

#include "toonz/preferences.h"

// TnzCore includes

#include "tmsgcore.h"

#include <qopenglshader></qopenglshader>

#include <qopenglshaderprogram></qopenglshaderprogram>

#include <qopenglframebufferobject></qopenglframebufferobject>

#include <qopengltexture></qopengltexture>

#include <qopenglcontext></qopenglcontext>

#include <qoffscreensurface></qoffscreensurface>

#include <qfile></qfile>

#include <qcolor></qcolor>

namespace {

inline bool execWarning(const QString& s) {

  DVGui::MsgBox(DVGui::WARNING, s);

  return false;

}

};  // namespace

#ifdef WIN32

#include <qstringlist></qstringlist>

#include <qsettings></qsettings>

#include <qbytearray></qbytearray>

namespace {

// obtain monitor information from registry

QStringList getMonitorNames() {

  QStringList subPathSet;

  // QSettings regSys("SYSTEM\\CurrentControlSet\\Enum\\DISPLAY",

  // QSettings::NativeFormat);

  QSettings regSys(

      "HKEY_LOCAL_MACHINE\\SYSTEM\\CurrentControlSet\\Enum\\DISPLAY",

      QSettings::NativeFormat);

  QStringList children = regSys.childGroups();

  // return value

  QStringList nameList;

  if (children.isEmpty()) {

    std::cout << "getMonitorNames : Failed to Open Registry" << std::endl;

    return nameList;

  }

  for (int c = 0; c < children.size(); c++) {

    // Parent Key : DISPLAY

    // Child Keys : AVO0000, ENC2174, etc.

    // Grandchild Key : 5&388..

    // Find grandchild key which contains a great-grandchild key named "Control"

    regSys.beginGroup(children.at(c));  // Child keys : AVO0000, ENC2174, etc.

    QStringList grandChildren = regSys.childGroups();

    for (int gc = 0; gc < grandChildren.size(); gc++) {

      regSys.beginGroup(grandChildren.at(gc));  // Grandchild key : 5&388..

      QStringList greatGrandChildren = regSys.childGroups();

      if (greatGrandChildren.contains(

              "Control"))  // If the key "Control" is found

      {

        // Obtain variable "EDID" from the key "Device Parameters"

        regSys.beginGroup("Device Parameters");

        // the key may be not "EDID", but "BAD_EDID"

        if (regSys.contains("EDID")) {

          QString subPath = regSys.group().replace("/", "\\").prepend(

              "SYSTEM\\CurrentControlSet\\Enum\\DISPLAY\\");

          subPathSet.push_back(subPath);

        }

        regSys.endGroup();

      }

      regSys.endGroup();

    }

    // subPath may not be one...?

    // if(!subPath.isEmpty())

    //	break;

    regSys.endGroup();

  }

  if (subPathSet.isEmpty()) {

    std::cout << "getMonitorNames : Failed to Find Current EDID" << std::endl;

    return nameList;

  }

  // for each subPath ( it may become more than one when using submonitor )

  for (int sp = 0; sp < subPathSet.size(); sp++) {

    QString subPath = subPathSet.at(sp);

    HKEY handle = 0;

    LONG res = RegOpenKeyExW(HKEY_LOCAL_MACHINE,

                             reinterpret_cast<const wchar_t*="">(subPath.utf16()),</const>

                             0, KEY_READ, &handle);

    if (res == ERROR_SUCCESS && handle) {

      QString keyStr("EDID");

      // get the size and type of the value

      DWORD dataType;

      DWORD dataSize;

      LONG res = RegQueryValueExW(

          handle, reinterpret_cast<const wchar_t*="">(keyStr.utf16()), 0,</const>

          &dataType, 0, &dataSize);

      if (res != ERROR_SUCCESS) {

        RegCloseKey(handle);

        continue;

      }

      // get the value

      QByteArray ba(dataSize, 0);

      res = RegQueryValueExW(

          handle, reinterpret_cast<const wchar_t*="">(keyStr.utf16()), 0, 0,</const>

          reinterpret_cast<unsigned char*="">(ba.data()), &dataSize);</unsigned>

      if (res != ERROR_SUCCESS) {

        RegCloseKey(handle);

        continue;

      }

      QString s;

      if (dataSize) {

        s = QString::fromUtf16((const ushort*)ba.constData(), ba.size() / 2);

      }

      QString valStr;

      QList<int> valArray;</int>

      for (int b = 0; b < s.length(); b++) {

        QChar c1((int)s[b].unicode() % 256);

        QChar c2((int)s[b].unicode() / 256);

        valStr.append(c1.toLatin1());

        valStr.append(c2.toLatin1());

        valArray.append((int)s[b].unicode() % 256);

        valArray.append((int)s[b].unicode() / 256);

      }

      // machine name starts from "FC 00", end with "0A"

      int index1 = valArray.indexOf(252);         // FC

      int index2 = valArray.indexOf(10, index1);  // 0A

      if (index1 > 0 && index2 > 0) {

        QString machineName = valStr.mid(index1 + 2, index2 - index1 - 2);

        nameList.push_back(machineName);

        // std::wcout << "machine name = " << machineName.toStdWString() <<

        // std::endl;

      }

      RegCloseKey(handle);

    } else {

      std::cout << "getMonitorNames : failed to get handle" << std::endl;

      continue;

    }

  }

  if (nameList.isEmpty())

    std::cout << "getMonitorNames : No Monitor Name Found" << std::endl;

  return nameList;

}

};  // namespace

#endif

//-----------------------------------------------------------------------------

LutCalibrator::LutCalibrator() {

  LutManager::instance()->registerCalibrator(this);

}

//-----------------------------------------------------------------------------

LutCalibrator::~LutCalibrator() {

  LutManager::instance()->removeCalibrator(this);

}

//-----------------------------------------------------------------------------

void LutCalibrator::initialize() {

  initializeOpenGLFunctions();

  m_isInitialized = true;

  if (!LutManager::instance()->isValid()) return;

  // create shader

  if (!initializeLutTextureShader()) {

    if (m_shader.program) delete m_shader.program;

    if (m_shader.vert) delete m_shader.vert;

    if (m_shader.frag) delete m_shader.frag;

    return;

  }

  createViewerVBO();

  // input 3dlut data to the shader

  assignLutTexture();

  m_isValid = true;

}

//-----------------------------------------------------------------------------

void LutCalibrator::cleanup() {

  m_isInitialized = false;

  if (!isValid()) return;

  // release shader

  if (m_shader.program) {

    delete m_shader.program;

    m_shader.program = NULL;

  }

  if (m_shader.vert) {

    delete m_shader.vert;

    m_shader.vert = NULL;

  }

  if (m_shader.frag) {

    delete m_shader.frag;

    m_shader.frag = NULL;

  }

  // release VBO

  if (m_viewerVBO.isCreated()) m_viewerVBO.destroy();

  // release LUT texture

  if (m_lutTex && m_lutTex->isCreated()) {

    m_lutTex->destroy();

    delete m_lutTex;

    m_lutTex = NULL;

  }

  m_isValid = false;

}

//-----------------------------------------------------------------------------

bool LutCalibrator::initializeLutTextureShader() {

  m_shader.vert = new QOpenGLShader(QOpenGLShader::Vertex);

  const char* simple_vsrc =

      "#version 330 core\n"

      "// Input vertex data, different for all executions of this shader.\n"

      "layout(location = 0) in vec3 vertexPosition;\n"

      "layout(location = 1) in vec2 texCoord;\n"

      "// Output data ; will be interpolated for each fragment.\n"

      "out vec2 UV;\n"

      "// Values that stay constant for the whole mesh.\n"

      "void main() {\n"

      "  // Output position of the vertex, in clip space : MVP * position\n"

      "  gl_Position = vec4(vertexPosition, 1);\n"

      "  // UV of the vertex. No special space for this one.\n"

      "  UV = texCoord;\n"

      "}\n";

  bool ret = m_shader.vert->compileSourceCode(simple_vsrc);

  if (!ret)

    return execWarning(

        QObject::tr("Failed to compile m_textureShader.vert.", "gl"));

  m_shader.frag = new QOpenGLShader(QOpenGLShader::Fragment);

  const char* simple_fsrc =

      "#version 330 core \n"

      "// Interpolated values from the vertex shaders \n"

      "in vec2 UV; \n"

      "// Output data \n"

      "out vec4 color; \n"

      "// Values that stay constant for the whole mesh. \n"

      "uniform sampler2D tex; \n"

      "uniform sampler3D lut; \n"

      "uniform vec3 lutSize; \n"

      "void main() { \n"

      "  vec3 rawColor = texture(tex,UV).rgb; \n"

      "  vec3 scale = (lutSize - 1.0) / lutSize; \n"

      "  vec3 offset = 1.0 / (2.0 * lutSize); \n"

      "  color = vec4(texture(lut, scale * rawColor + offset).rgb, 1.0); \n"

      "} \n";

  ret = m_shader.frag->compileSourceCode(simple_fsrc);

  if (!ret)

    return execWarning(QObject::tr("Failed to compile m_shader.frag.", "gl"));

  m_shader.program = new QOpenGLShaderProgram();

  // add shaders

  ret = m_shader.program->addShader(m_shader.vert);

  if (!ret)

    return execWarning(QObject::tr("Failed to add m_shader.vert.", "gl"));

  ret = m_shader.program->addShader(m_shader.frag);

  if (!ret)

    return execWarning(QObject::tr("Failed to add m_shader.frag.", "gl"));

  // link shaders

  ret = m_shader.program->link();

  if (!ret)

    return execWarning(QObject::tr("Failed to link simple shader: %1", "gl")

                           .arg(m_shader.program->log()));

  // obtain parameter locations

  m_shader.vertexAttrib = m_shader.program->attributeLocation("vertexPosition");

  if (m_shader.vertexAttrib == -1)

    return execWarning(

        QObject::tr("Failed to get attribute location of %1", "gl")

            .arg("vertexPosition"));

  m_shader.texCoordAttrib = m_shader.program->attributeLocation("texCoord");

  if (m_shader.texCoordAttrib == -1)

    return execWarning(

        QObject::tr("Failed to get attribute location of %1", "gl")

            .arg("texCoord"));

  m_shader.texUniform = m_shader.program->uniformLocation("tex");

  if (m_shader.texUniform == -1)

    return execWarning(

        QObject::tr("Failed to get uniform location of %1", "gl").arg("tex"));

  m_shader.lutUniform = m_shader.program->uniformLocation("lut");

  if (m_shader.lutUniform == -1)

    return execWarning(

        QObject::tr("Failed to get uniform location of %1", "gl").arg("lut"));

  m_shader.lutSizeUniform = m_shader.program->uniformLocation("lutSize");

  if (m_shader.lutSizeUniform == -1)

    return execWarning(QObject::tr("Failed to get uniform location of %1", "gl")

                           .arg("lutSize"));

  return true;

}

//-----------------------------------------------------------------------------

void LutCalibrator::createViewerVBO() {

  GLfloat vertex[]   = {-1.0f, -1.0f, 1.0f, -1.0f, 1.0f, 1.0f, -1.0f, 1.0f};

  GLfloat texCoord[] = {0.0f, 0.0f, 1.0f, 0.0f, 1.0f, 1.0f, 0.0f, 1.0f};

  m_viewerVBO.create();

  m_viewerVBO.bind();

  m_viewerVBO.allocate(4 * 4 * sizeof(GLfloat));

  m_viewerVBO.write(0, vertex, sizeof(vertex));

  m_viewerVBO.write(sizeof(vertex), texCoord, sizeof(texCoord));

  m_viewerVBO.release();

}

//-----------------------------------------------------------------------------

void LutCalibrator::onEndDraw(QOpenGLFramebufferObject* fbo) {

  assert((glGetError()) == GL_NO_ERROR);

  fbo->release();

  GLuint textureId = fbo->texture();

  glEnable(GL_TEXTURE_2D);

  glActiveTexture(GL_TEXTURE1);

  glBindTexture(GL_TEXTURE_2D, textureId);

  glActiveTexture(GL_TEXTURE2);

  m_lutTex->bind();

  glPushMatrix();

  glLoadIdentity();

  m_shader.program->bind();

  m_shader.program->setUniformValue(m_shader.texUniform,

                                    1);  // use texture unit 1

  m_shader.program->setUniformValue(m_shader.lutUniform,

                                    2);  // use texture unit 2

  GLfloat size = (GLfloat)LutManager::instance()->meshSize();

  m_shader.program->setUniformValue(m_shader.lutSizeUniform, size, size, size);

  m_shader.program->enableAttributeArray(m_shader.vertexAttrib);

  m_shader.program->enableAttributeArray(m_shader.texCoordAttrib);

  m_viewerVBO.bind();

  m_shader.program->setAttributeBuffer(m_shader.vertexAttrib, GL_FLOAT, 0, 2);

  m_shader.program->setAttributeBuffer(m_shader.texCoordAttrib, GL_FLOAT,

                                       4 * 2 * sizeof(GLfloat), 2);

  m_viewerVBO.release();

  glDrawArrays(GL_TRIANGLE_FAN, 0, 4);

  m_shader.program->disableAttributeArray(m_shader.vertexAttrib);

  m_shader.program->disableAttributeArray(m_shader.texCoordAttrib);

  m_shader.program->release();

  glPopMatrix();

  glActiveTexture(GL_TEXTURE0);  // reset the active texture unit to 0

  glDisable(GL_TEXTURE_2D);

  assert((glGetError()) == GL_NO_ERROR);

}

//-----------------------------------------------------------------------------

void LutCalibrator::assignLutTexture() {

  assert(glGetError() == GL_NO_ERROR);

  if (m_lutTex) delete m_lutTex;

  int meshSize = LutManager::instance()->meshSize();

  m_lutTex     = new QOpenGLTexture(QOpenGLTexture::Target3D);

  m_lutTex->setSize(meshSize, meshSize, meshSize);

  m_lutTex->setFormat(QOpenGLTexture::RGB32F);

  // m_lutTex->setLayers(1);

  m_lutTex->setMipLevels(1);

  m_lutTex->allocateStorage();

  m_lutTex->setMinMagFilters(QOpenGLTexture::Linear, QOpenGLTexture::Linear);

  m_lutTex->setWrapMode(QOpenGLTexture::ClampToEdge);

  m_lutTex->setData(QOpenGLTexture::RGB, QOpenGLTexture::Float32,

                    LutManager::instance()->data());

  assert(glGetError() == GL_NO_ERROR);

}

//-----------------------------------------------------------------------------

void LutCalibrator::update(bool textureChanged) {

  m_isValid = LutManager::instance()->isValid();

  if (textureChanged) assignLutTexture();

}

//=============================================================================

LutManager* LutManager::instance() {

  static LutManager _instance;

  return &_instance;

}

//-----------------------------------------------------------------------------

LutManager::LutManager() : m_isValid(false), m_currentLutPath() {

  // check whether preference enables color calibration

  if (!Preferences::instance()->isColorCalibrationEnabled()) return;

  // obtain current monitor name

  QString monitorName = getMonitorName();

  // obtain 3dlut path associated to the monitor name

  QString lutPath =

      Preferences::instance()->getColorCalibrationLutPath(monitorName);

  if (lutPath.isEmpty()) return;

  // check existence of the 3dlut file

  // load 3dlut data

  if (!loadLutFile(lutPath)) return;

  m_currentLutPath = lutPath;

  m_isValid        = true;

}

//-----------------------------------------------------------------------------

LutManager::~LutManager() {

  if (m_lut.data) delete[] m_lut.data;

}

//-----------------------------------------------------------------------------

QString& LutManager::getMonitorName() const {

  static QString monitorName;

  if (!monitorName.isEmpty()) return monitorName;

#ifdef WIN32

  QStringList list = getMonitorNames();

  if (list.isEmpty())

    monitorName = "Any Monitor";  // this should not be translated

  else

    monitorName = list.at(0);  // for now only the first monitor is handled

#else

  monitorName = "Any Monitor";  // this should not be translated

#endif

  return monitorName;

}

//-----------------------------------------------------------------------------

bool LutManager::loadLutFile(const QString& fp) {

  struct locals {

    // skip empty or comment lines

    static inline QString readDataLine(QTextStream& stream) {

      while (1) {

        if (stream.atEnd()) return QString();

        QString ret = stream.readLine();

        if (!ret.isEmpty() && ret[0] != QChar('#')) return ret;

      }

    }

    static inline int lutCoords(int r, int g, int b, int meshSize) {

      return b * meshSize * meshSize * 3 + g * meshSize * 3 + r * 3;

    }

  };

  QFile file(fp);

  if (!file.open(QIODevice::ReadOnly | QIODevice::Text))

    return execWarning(QObject::tr("Failed to Open 3DLUT File."));

  QTextStream stream(&file);

  QString line;

  //---- read the 3DLUT files

  // The first line should start from "3DMESH" keyword (case sensitive)

  line = locals::readDataLine(stream);

  if (line != "3DMESH") {

    file.close();

    return execWarning(

        QObject::tr("Failed to Load 3DLUT File.\nIt should start with "

                    "\"3DMESH\" keyword."));

  }

  // The second line is "Mesh [Input bit depth] [Output bit depth]"

  // "Mesh" is a keyword (case sensitive) 

  line             = locals::readDataLine(stream);

  QStringList list = line.split(" ");

  if (list.size() != 3 || list.at(0) != "Mesh") {

    file.close();

    return execWarning(

        QObject::tr("Failed to Load 3DLUT File.\nThe second line should be "

                    "\"Mesh [Input bit depth] [Output bit depth]\""));

  }

  int inputBitDepth  = list.at(1).toInt();

  int outputBitDepth = list.at(2).toInt();

  m_lut.meshSize = (int)pow(2.0, inputBitDepth) + 1;

  float maxValue = pow(2.0, outputBitDepth) - 1.0f;

  // The third line is corrections of values at each LUT grid

  line = locals::readDataLine(stream);

  list = line.split(" ", QString::SkipEmptyParts);

  if (list.size() != m_lut.meshSize) {

    file.close();

    return execWarning(QObject::tr("Failed to Load 3DLUT File."));

  }

  if (m_lut.data) delete[] m_lut.data;

  m_lut.data = new float[m_lut.meshSize * m_lut.meshSize * m_lut.meshSize * 3];

  for (int k = 0; k < m_lut.meshSize; ++k)  // r

  {

    for (int j = 0; j < m_lut.meshSize; ++j)  // g

    {

      for (int i = 0; i < m_lut.meshSize; ++i)  // b

      {

        line = locals::readDataLine(stream);

        list = line.split(" ", QString::SkipEmptyParts);

        if (list.size() != 3) {

          file.close();

          delete[] m_lut.data;

          return execWarning(QObject::tr("Failed to Load 3DLUT File."));

        }

        float* lut = m_lut.data + locals::lutCoords(k, j, i, m_lut.meshSize);

        *lut       = (float)(list.at(0).toInt()) / maxValue;

        lut++;

        *lut = (float)(list.at(1).toInt()) / maxValue;

        lut++;

        *lut = (float)(list.at(2).toInt()) / maxValue;

      }

    }

  }

  file.close();

  return true;

}

//-----------------------------------------------------------------------------

// input : 0-1

void LutManager::convert(float& r, float& g, float& b) {

  struct locals {

    static inline float lerp(float val1, float val2, float ratio) {

      return val1 * (1.0f - ratio) + val2 * ratio;

    }

    static inline int getCoord(int r, int g, int b, int meshSize) {

      return b * meshSize * meshSize * 3 + g * meshSize * 3 + r * 3;

    }

  };

  if (!m_isValid) return;

  float ratio[3];   // RGB軸

  int index[3][2];  // rgb インデックス

  float rawVal[3] = {r, g, b};

  // clamp values (for HDR image)

  for (int c = 0; c < 3; c++)

    rawVal[c] = (rawVal[c] < 0.f) ? 0.f : (rawVal[c] > 1.f) ? 1.f : rawVal[c];

  float vertex_color[2][2][2][3];  // 補間用の１ボクセルの頂点色

  for (int c = 0; c < 3; c++) {

    float val   = rawVal[c] * (float)(m_lut.meshSize - 1);

    index[c][0] = (int)val;

    // boundary condition: if rawVal == 1 the value will not be interpolated

    index[c][1] = (rawVal[c] >= 1.0f) ? index[c][0] : index[c][0] + 1;

    ratio[c]    = val - (float)index[c][0];

  }

  for (int rr = 0; rr < 2; rr++)

    for (int gg = 0; gg < 2; gg++)

      for (int bb = 0; bb < 2; bb++) {

        float* val = &m_lut.data[locals::getCoord(

            index[0][rr], index[1][gg], index[2][bb], m_lut.meshSize)];

        for (int chan = 0; chan < 3; chan++, val++)

          vertex_color[rr][gg][bb][chan] = *val;

      }

  float result[3];

  for (int chan = 0; chan < 3; chan++) {

    result[chan] = locals::lerp(

        locals::lerp(locals::lerp(vertex_color[0][0][0][chan],

                                  vertex_color[0][0][1][chan], ratio[2]),

                     locals::lerp(vertex_color[0][1][0][chan],

                                  vertex_color[0][1][1][chan], ratio[2]),

                     ratio[1]),

        locals::lerp(locals::lerp(vertex_color[1][0][0][chan],

                                  vertex_color[1][0][1][chan], ratio[2]),

                     locals::lerp(vertex_color[1][1][0][chan],

                                  vertex_color[1][1][1][chan], ratio[2]),

                     ratio[1]),

        ratio[0]);

  }

  r = result[0];

  g = result[1];

  b = result[2];

}

//-----------------------------------------------------------------------------

void LutManager::convert(QColor& col) {

  if (!m_isValid) return;

  float r = col.redF();

  float g = col.greenF();

  float b = col.blueF();

  convert(r, g, b);

  col = QColor::fromRgbF(r, g, b, col.alphaF());

}

//-----------------------------------------------------------------------------

void LutManager::convert(TPixel32& col) {

  if (!m_isValid) return;

  float r = (float)col.r / 255.0;

  float g = (float)col.g / 255.0;

  float b = (float)col.b / 255.0;

  convert(r, g, b);

  col = TPixel32((int)(r * 255.0 + 0.5), (int)(g * 255.0 + 0.5),

                 (int)(b * 255.0 + 0.5), col.m);

}

//-----------------------------------------------------------------------------

void LutManager::registerCalibrator(LutCalibrator* calibrator) {

  assert(!m_calibrators.contains(calibrator));

  m_calibrators.insert(calibrator);

}

//-----------------------------------------------------------------------------

void LutManager::removeCalibrator(LutCalibrator* calibrator) {

  assert(m_calibrators.contains(calibrator));

  m_calibrators.remove(calibrator);

}

//-----------------------------------------------------------------------------

void LutManager::update() {

  m_isValid           = false;

  bool textureChanged = false;

  if (Preferences::instance()->isColorCalibrationEnabled()) {

    // obtain current monitor name

    QString monitorName = getMonitorName();

    // obtain 3dlut path associated to the monitor name

    QString lutPath =

        Preferences::instance()->getColorCalibrationLutPath(monitorName);

    if (m_currentLutPath == lutPath)

      m_isValid = true;

    else if (loadLutFile(lutPath)) {

      m_isValid        = true;

      m_currentLutPath = lutPath;

      textureChanged   = true;

    }

  }

  // update textures for all calibrators

  for (auto calibrator : m_calibrators) calibrator->update(textureChanged);

}