#include "ofApp.h"

// some v4l2 global settings
int camWidth = 640;
int camHeight = 480;

void ofApp::setupKinect()
{
  ofLogNotice(__FUNCTION__) << "Found " << ofxAzureKinect::Device::getInstalledCount() << " installed devices.";

  auto kinectSettings = ofxAzureKinect::DeviceSettings();
  kinectSettings.updateIr = false;
  kinectSettings.updateColor = true;
  kinectSettings.colorResolution = K4A_COLOR_RESOLUTION_1080P;
  kinectSettings.updateVbo = false;

  auto deviceSettings = ofxAzureKinect::DeviceSettings();
  deviceSettings.syncImages = false;
  deviceSettings.depthMode = K4A_DEPTH_MODE_NFOV_UNBINNED;
  deviceSettings.updateIr = false;
  deviceSettings.updateColor = false;
  //deviceSettings.colorResolution = K4A_COLOR_RESOLUTION_1080P;
  deviceSettings.updateWorld = true;
  deviceSettings.updateVbo = false;
  auto bodyTrackingSettings = ofxAzureKinect::BodyTrackingSettings();
  //bodyTrackingSettings.processingMode = K4ABT_TRACKER_PROCESSING_MODE_CPU;
  bodyTrackingSettings.updateBodies = true;
  if (kinectDevice.open())
  {
    kinectDevice.startCameras(kinectSettings, bodyTrackingSettings);
  }
}

void ofApp::setupThermal()
{
  // this must be called before init (otherwise fprintf will tell you so)
  // note that high framerates will only function properly if the usb has enough bandwidth
  // for example, a ps3 eye cam at 60 fps will only function when it has full USB 2.0 bandwidth available
  v4l2Cam.setDesiredFramerate(60);

  // use this to set appropriate device and capture method
  std::string dev = captureDeviceName;
  v4l2Cam.initGrabber((std::string("/dev/") + dev).c_str(), IO_METHOD_MMAP, camWidth, camHeight);

  // some initial settings
  int set_gain = 2.0;
  bool set_autogain = true;

  // rudimentary settings implementation: each settings needs a seperate call to the settings method
  v4l2Cam.settings(ofxV4L2_AUTOGAIN, set_autogain);
  v4l2Cam.settings(ofxV4L2_GAIN, set_gain);

  // we use a texture because the ofxV4L2 class has no draw method (yet)
  // we use GL_LUMINANCE because the ofxV4L2 class supports only grayscale (for now)
  v4l2Tex.allocate(512, 512, GL_RGB);

  v4l2Pixels.allocate(512, 512, OF_PIXELS_RGB);
  for (int i = 0; i < 512; i++)
  {
    for (int j = 0; j < 512; j++)
    {
      v4l2Pixels.setColor(j, i, 0);
    }
  }
}

void ofApp::setupParticles()
{
  // 1,000,000 particles
  unsigned w = 512;
  unsigned h = 512;

  particles.init(w, h, OF_PRIMITIVE_POINTS, false, 4);

  particles.loadShaders("shaders/particles/update", "shaders/particles/draw");

  float *particlePosns = new float[w * h * 4];
  for (unsigned y = 0; y < h; ++y)
  {
    for (unsigned x = 0; x < w; ++x)
    {
      unsigned idx = y * w + x;
      particlePosns[idx * 4] = 400.f * x / (float)w - 200.f;     // particle x
      particlePosns[idx * 4 + 1] = 400.f * y / (float)h - 200.f; // particle y
      particlePosns[idx * 4 + 2] = 0.f;                          // particle z
      particlePosns[idx * 4 + 3] = 0.f;                          // dummy
    }
  }
  particles.loadDataTexture(ofxGpuParticles::POSITION, particlePosns);

  // initial velocities
  particles.zeroDataTexture(ofxGpuParticles::VELOCITY);
  for (unsigned y = 0; y < h; ++y)
  {
    for (unsigned x = 0; x < w; ++x)
    {
      unsigned idx = y * w + x;
      particlePosns[idx * 4] = 0;               //
      particlePosns[idx * 4 + 1] = ofRandomf(); // age
      particlePosns[idx * 4 + 2] = 0.f;         //
      particlePosns[idx * 4 + 3] = 0.f;         //
    }
  }
  particles.loadDataTexture(ofxGpuParticles::MISC, particlePosns);
  particles.zeroDataTexture(3);

  delete[] particlePosns;

  particles.whateverImages.insert({"u_depth", kinectDevice.getDepthTex()});
  particles.whateverImages.insert({"u_world", kinectDevice.getDepthToWorldTex()});
  particles.whateverImages.insert({"u_v4l2cam", v4l2Tex});

  for (auto i : gradients)
  {
    particles.whateverImages.insert({i.first, i.second.getTexture()});
  }

  // listen for update event to set additonal update uniforms
  ofAddListener(particles.updateEvent, this, &ofApp::onParticlesUpdate);
  ofAddListener(particles.drawEvent, this, &ofApp::onParticlesDraw);
}

void ofApp::setupGui()
{
  gui.setup("ZONE B", "settings.xml"); // most of the time you don't need a name
  gui.add(registrationXY.setup("XY", {-208, -257}, {-400, -400}, {400, 400}));
  gui.add(registrationScale.setup("scale", 2.0, 0.1, 3));
  gui.add(captureDeviceName.setup("device", "video3"));
  gui.loadFromFile("settings.xml");
  gui.setPosition(10, 10);

  debugGui.setup("DEBUG");
  debugGui.add(calibMode.setup("calib", false));
  debugGui.add(dummyMode.setup("dummy", true));                      // should be false
  debugGui.add(dummyXY.setup("XY", {256, 256}, {0, 0}, {511, 511})); // should be false
  debugGui.add(debugFps.setup("FPS", "0"));
  debugGui.setPosition(230, 10);
}

void ofApp::setup()
{
  // ofDisableArbTex();
  // ofSetVerticalSync(false);
  //ofSetLogLevel(OF_LOG_VERBOSE);

  setupGui();
  setupKinect();
  setupThermal();

  boundShader.allocate(ofGetWidth(), ofGetHeight());
  boundShader.load("shaders/bound.frag");

  fbos.insert({"ofcam", ofFbo()});
  fbos.at("ofcam").allocate(ofGetWidth(), ofGetHeight(), GL_RGBA32F_ARB);

  for (auto s : gradientNames)
  {
    gradients.insert({s, ofFloatImage()});
    gradients.at(s).load(s + ".png");
  }
  setupParticles();
}

void ofApp::exit()
{
  kinectDevice.close();
}

ofVec3f ofApp::getDepthAt(int x, int y)
{
  auto ray = kinectDevice.getDepthToWorldPix().getColor(x, y);
  auto depthShort = kinectDevice.getDepthPix().getColor(x, y).r;
  float depth = -depthShort;
  if (depthShort == 0)
    depth = -2000;
  ofVec3f pos(ray.r * depth, ray.g * depth, depth);
  return pos;
}

void ofApp::updateThermal()
{
  v4l2Cam.grabFrame();
  if (v4l2Cam.isNewFrame())
  {
    auto &body = kinectDevice.getBodyIndexPix();
    hotspots.clear();
    for (int i = 0; i < kinectDevice.getNumBodies(); i++)
    {
      hotspots.push_back(std::vector<ofVec3f>());
    }
    for (int i = 0; i < 512; i++)
    {
      for (int j = 0; j < 512; j++)
      {
        int x = (j - 256) * registrationScale + 256 + registrationXY->x;
        int y = (i - 256) * registrationScale + 256 + registrationXY->y;
        if (ofInRange(x, 0, camWidth - 1) == false || ofInRange(y, 0, camHeight - 1) == false)
        {
          continue;
        }
        int count = x + y * camWidth;
        int a = v4l2Cam.getPixels()[count];
        v4l2Pixels.setColor(j, i, a);
        if (i % 4 == 0 && j % 4 == 0)
        {
          auto c = body.getColor(j, i);
          if (c.r < hotspots.size())
          {
            hotspots.at(c.r).push_back(ofVec3f(j, i, a));
          }
        }
      }
    }
    v4l2Tex.allocate(v4l2Pixels);
    struct
    {
      bool operator()(ofVec3f a, ofVec3f b) const
      {
        return a.z > b.z;
      }
    } compareZThermal;
    for (int i = 0; i < hotspots.size(); i++)
    {
      std::sort(hotspots.at(i).begin(), hotspots.at(i).end(), compareZThermal);
    }

    if (hotspots.size() > 0 && hotspots.at(0).size() > 0)
    {
      hotspot0 = hotspots.at(0).at(0);
    }
    if (dummyMode == false)
    {
      if (hotspots.size() > 1 && hotspots.at(1).size() > 0)
      {
        hotspot1 = hotspots.at(1).at(0);
      }
      else if (hotspots.size() > 0 && hotspots.at(0).size() > 1)
      {
        hotspot1 = hotspots.at(0).at(1);
      }
    }
    else
    {
      hotspot1 = ofVec3f(dummyXY);
    }

    hotspot3d0 = getDepthAt(hotspot0.x, hotspot0.y);
    hotspot3d1 = getDepthAt(hotspot1.x, hotspot1.y);
  }
}

void ofApp::update()
{
  updateThermal();

  particles.update();

  debugFps = ofToString(ofGetFrameRate(), 2);
}

void ofApp::onParticlesUpdate(ofxShader &shader)
{
  ofVec3f mouse(ofGetMouseX() - .5f * ofGetWidth(), .5f * ofGetHeight() - ofGetMouseY(), 0.f);
  shader.setUniform3fv("mouse", mouse.getPtr());
  shader.setUniform1f("elapsed", ofGetLastFrameTime());
  shader.setUniform1f("radiusSquared", 200.f * 200.f);

  shader.setUniform2i("uFrameSize", kinectDevice.getDepthTex().getWidth(), kinectDevice.getDepthTex().getHeight());
  shader.setUniform2i("uDepthFrameSize", kinectDevice.getDepthTex().getWidth(), kinectDevice.getDepthTex().getHeight());

  shader.setUniform3f("uHottest0", hotspot0);
  shader.setUniform3f("uHottest1", hotspot1);
  shader.setUniform3f("uHottest3d0", hotspot3d0);
  shader.setUniform3f("uHottest3d1", hotspot3d1);
}

void ofApp::onParticlesDraw(ofxShader &shader)
{
  shader.setUniform3f("uHottest0", hotspot0);
  shader.setUniform3f("uHottest1", hotspot1);
  shader.setUniform3f("uHottest3d0", hotspot3d0);
  shader.setUniform3f("uHottest3d1", hotspot3d1);
}

void ofApp::drawMain()
{
  ofDisableDepthTest();
  auto tex = kinectDevice.getDepthTex();
  boundShader.setUniformTexture("u_depth", tex);
  boundShader.setUniformTexture("u_ofcam", fbos.at("ofcam"));
  boundShader.setUniformTexture("u_v4l2cam", v4l2Tex);
  boundShader.setUniform1i("u_calib", calibMode == true ? 1 : 0);
  boundShader.setUniform2f("u_calibXY", registrationXY);
  boundShader.setUniform1f("u_calibScale", registrationScale);
  boundShader.setUniform1i("u_init", 1);
  boundShader.render();
  boundShader.setUniform1i("u_init", 0);
  boundShader.render();
}

void ofApp::drawDebug()
{
  gui.draw();
  debugGui.draw();
}

void ofApp::draw()
{
  ofBackground(0);
  if (kinectDevice.isStreaming())
  {
    particles.whateverImages.at("u_depth") = kinectDevice.getDepthTex();
    particles.whateverImages.at("u_world") = kinectDevice.getDepthToWorldTex();
    particles.whateverImages.at("u_v4l2cam") = v4l2Tex;

    fbos.at("ofcam").begin();
    ofClear(0);
    cam.begin();
    particles.draw();
    cam.end();
    fbos.at("ofcam").end();

    drawMain();
    boundShader.draw(0, 0);
  }
  drawDebug();
}

void ofApp::keyPressed(int key)
{
}

void ofApp::keyReleased(int key)
{
}

void ofApp::mouseMoved(int x, int y)
{
}

void ofApp::mouseDragged(int x, int y, int button)
{
}

void ofApp::mousePressed(int x, int y, int button)
{
}

void ofApp::mouseReleased(int x, int y, int button)
{
}

void ofApp::mouseEntered(int x, int y)
{
}

void ofApp::mouseExited(int x, int y)
{
}

void ofApp::windowResized(int w, int h)
{
}

void ofApp::gotMessage(ofMessage msg)
{
}

void ofApp::dragEvent(ofDragInfo dragInfo)
{
}