#include <cstdio>

#include "autostitch.h"

#include <opencv/cv.h>
#include <opencv/highgui.h>

#include "meowpp/Usage.h"

#include "meowpp/colors/RGB_Space.h"

#include "meowpp/dsa/DisjointSet.h"

#include "meowpp/geo/Vectors.h"

#include "meowpp/gra/Bitmap.h"
#include "meowpp/gra/Photo.h"
#include "meowpp/gra/Camera.h"

#include "meowpp/math/utility.h"
#include "meowpp/math/methods.h"


extern "C"{
#include <sys/types.h>
#include <dirent.h>
}

#include <vector>
#include <algorithm>
#include <string>
#include <cstdlib>


using namespace meow;

//////////////////////////////////////////////////////////////////////

Usage usg("autostitch");

double p0 = 0.07, P = 0.99;
double q = 0.7, r = 0.01, Q = 0.97;
double stop = 1;
double o_radius = 500;
double angle_t = PI / 4.0;
double aspect_t = 2.0;

std::vector<std::string> input_name;

MyK_Match match;
std::vector<Bitmap<RGBf_Space> >  input_bitmap;
std::vector<Bitmap<RGBf_Space> > output_bitmap;

std::vector<std::vector<FeaturePoint<double, double> > > fps;
std::vector<std::vector<Vector<double>               > > fpsv;
std::vector<std::vector<FeaturePointIndexPairs       > > pairs;

struct OutputSet {
  struct Edge {
    std::vector<Vector<double> > v1;
    std::vector<Vector<double> > v2;
    size_t i1, i2;
    bool done;
    Edge(size_t ii1, size_t ii2): i1(ii1), i2(ii2) {
      done = false;
    }
    bool operator<(Edge const& e) const {
      return (v1.size() < e.v1.size());
    }
  };
  std::vector<Camera<RGBf_Space> > cameras;
  std::vector<Edge>                edges;
};

std::vector<OutputSet> outputs;

//////////////////////////// **# setup #** ///////////////////////////
bool setup(int argc, char** argv) {
  usg.optionAdd("h", "Display this help document.");
  usg.optionAdd("i",
                "Specify the input images are in <type> "
                "instead of specifying from arguments",
                "<dirname>", "",
                false);
  usg.optionAdd("o",
                "Output file name, (not include '.jpg' suffix)",
                "<filename>",
                "output",
                false);
  usg.optionAdd("f",
                "File name for output the text data",
                "<filename>",
                "output.txt",
                false);
  usg.optionAdd("d",
                "Specify which Feature-Point-Detect algorithm to use",
                "<algorithm>",
                "",
                true);
  usg.optionAdd("ransac-p0",
                "Pribabilicity for RANSAC to choose a right feature point",
                "<floating point>", stringPrintf("%.10f", p0),
                false);
  usg.optionAdd("ransac-ok",
                "Pribabilicity for RANSAC access",
                "<floating point>", stringPrintf("%.10f", P),
                false);
  usg.optionAdd("prob-p1",
                "p1 for Prob. Model",
                "<floationg Point>", stringPrintf("%.10f", q),
                false);
  usg.optionAdd("prob-p0",
                "p0 for Prob. Model",
                "<floationg Point>", stringPrintf("%.10f", r),
                false);
  usg.optionAdd("prob-min",
                "p_min for Prob. Model",
                "<floationg Point>", stringPrintf("%.10f", Q),
                false);
  usg.optionAdd("s",
                "stop threshold for boundle adjustment",
                "<floationg Point>", stringPrintf("%.10f", stop),
                false);
  usg.optionAdd("output-radius",
                "output ball radius",
                "<floationg Point>", stringPrintf("%.10f", o_radius),
                false);
  usg.optionAdd("match-angle",
                "angle threshold for matching",
                "<floating point>", stringPrintf("%.10f", angle_t / PI * 180),
                false);
  usg.optionAdd("match-aspect",
                "aspect threshold for matching",
                "<floating point>", stringPrintf("%.10f", aspect_t),
                false);
  std::vector<std::string> fpsd_algorithm_list = ObjSelector<FPSD_ID>::names();
  for (size_t i = 0, I = fpsd_algorithm_list.size(); i < I; i++) {
    const ObjBase* tmp = ObjSelector<FPSD_ID>::get(fpsd_algorithm_list[i]);
    usg.optionValueAcceptAdd("d",
                             fpsd_algorithm_list[i],
                             tmp->type());
    usg.import(((MyFeaturePointsDetector*)tmp)->usage());
  }
  usg.import(match.usage());
  usg.import(MyRansacCheck::usage());
  // set arg
  std::string err_msg;
  bool ok = usg.arguments(argc, argv, &err_msg);
  if (usg.hasOptionSetup("h")) {
    printf("%s\n", usg.usage().c_str());
    exit(0);
  }
  if (!ok) {
    fprintf(stderr, "%s\n", err_msg.c_str());
    exit(-1);
  }
  return true;
}


//////////////// **# Input images and convert it #** /////////////////
bool input() {
  if (!usg.hasOptionSetup("i")) {
    input_name = usg.procArgs();
  }
  else {
    std::string base = usg.optionValue("i", 0);
    if (base.length() == 0 || base[base.length() - 1] != '/') {
      base += "/";
    }
    DIR* dir = opendir(base.c_str());
    if (!dir) {
      fprintf(stderr, "can't open dir '%s'\n", base.c_str());
      return -1;
    }
    for (dirent* ent; (ent = readdir(dir)) != NULL; ) {
      if (!cstringEndWith(ent->d_name, 4, ".jpeg", ".jpg", ".JPG", ".JPEG")) {
        continue;
      }
      input_name.push_back(base + std::string(ent->d_name));
    }
  }
  messagePrintf(1, "Loading images");
  for (size_t i = 0; i < input_name.size(); i++) {
    messagePrintf(1, "%s", input_name[i].c_str());
    cv::Mat img = cv::imread(input_name[i], CV_LOAD_IMAGE_COLOR);
    if (!img.data) {
      messagePrintf(-1, "opencv read error!, ignore");
      continue;
    }
    size_t width  = img.size().width ;
    size_t height = img.size().height;
    size_t index = input_bitmap.size();
    input_bitmap.resize(index + 1);
    input_bitmap[index].size(height, width, RGBf_Space(0));
    for (size_t x = 0; x < width; x++) {
      for (size_t y = 0; y < height; y++) {
        RGBi_Space tmp(Vector3D<int>(
            img.at<cv::Vec3b>(y, x)[2],
            img.at<cv::Vec3b>(y, x)[1],
            img.at<cv::Vec3b>(y, x)[0]));
        RGBf_Space p;
        colorTransformate(tmp, &p);
        input_bitmap[index].pixel(y, x, p);
      }
    }
    messagePrintf(-1, "%lux%lu, ok", width, height);
  }
  messagePrintf(-1, "ok");
  return true;
}

//////////////////////// **# FeaturePoint #** ////////////////////////
bool detect() {
  std::string fpsd_algo_name = usg.optionValue("d", 0);
  MyFeaturePointsDetector* detector(
    (MyFeaturePointsDetector*)ObjSelector<FPSD_ID>::create(fpsd_algo_name));
  detector->usage(usg);
  fps .resize(input_bitmap.size());
  fpsv.resize(input_bitmap.size());
  for (size_t i = 0, I = input_bitmap.size(); i < I; i++) {
    messagePrintf(1, "Detect the feature points for %lu-th pic", i);
    fps[i] = detector->detect(input_bitmap[i]);
    messagePrintf(-1, "ok, %lu", fps[i].size());
    for (size_t j = 0, J = fps[i].size(); j < J; j++) {
      fpsv[i].push_back(fps[i][j].position());
    }
  }
  delete detector;
  return true;
}


//////////////////////////// **# k-match #** /////////////////////////
bool kmatch() {
  match.usage(usg);
  messagePrintf( 1, "run k-match");
  FeaturePointIndexPairs mat(match.match(fps));
  pairs.resize(input_bitmap.size());
  for (size_t i = 0, I = input_bitmap.size(); i < I; i++) {
    pairs[i].resize(I);
  }
  for (size_t i = 0, I = mat.size(); i < I; ++i) {
    pairs[mat[i].from.first][mat[i].to.first].push_back(mat[i]);
  }
  messagePrintf(-1, "ok");
  return true;
}

//////////////////////////// **# RANSAC #** //////////////////////////
bool ransac() {
  messagePrintf( 1, "RANSAC");
  MyRansacCheck::usage(usg);
  // tmp output
  p0 = inRange(0.00001, 0.9999, atof(usg.optionValue("ransac-p0", 0).c_str()));
  P  = inRange(0.00001, 0.9999, atof(usg.optionValue("ransac-ok", 0).c_str()));
  for (size_t i = 0, I = input_bitmap.size(); i < I; i++) {
    for (size_t j = 0, J = input_bitmap.size(); j < J; j++) {
      size_t num = 4u; // !!!!!!!!!!!!!!!!!!!
      messagePrintf( 1, "ransac %lu --- %lu", i, j);
      MyRansacCheck chk(&(fpsv[i]), &(fpsv[j]));
      FeaturePointIndexPairs ret = ransac(pairs[i][j], chk, num, p0, P);
      if (!ret.empty()) {
        chk.rememberVCalc(ret);
        FeaturePointIndexPairs ok(ret);
        for (size_t k = 0, K = pairs[i][j].size(); k < K; k++) {
          bool chk_again = true;
          for (size_t l = 0, L = ret.size(); chk_again && l < L; l++) {
            if (ret[l] == pairs[i][j][k]) {
              chk_again = false;
            }
          }
          if (chk_again && chk.ok(pairs[i][j][k])) {
            ok.push_back(pairs[i][j][k]);
          }
        }
        if (ok.size() >= num) pairs[i][j] = ok;
        else                  pairs[i][j].clear();
        messagePrintf(-1, "ok(%lu)", pairs[i][j].size());
      }
      else {
        pairs[i][j].clear();
        messagePrintf(-1, "empty");
      }
    }
  }
  messagePrintf(-1, "ok");
  return true;
}

//////////////////// **# checking match again #** ////////////////////
bool match_check() {
  q = inRange(0.00001, 0.99999, atof(usg.optionValue("prob-p1", 0).c_str()));
  r = inRange(0.00001, 0.99999, atof(usg.optionValue("prob-p0", 0).c_str()));
  Q = inRange(0.00001, 0.99999, atof(usg.optionValue("prob-min", 0).c_str()));
  aspect_t = inRange(0.000001, 9999.0, atof(usg.optionValue("match-aspect", 0).c_str()));
  angle_t  = inRange(0.1,999.0, atof(usg.optionValue("match-angle", 0).c_str()))/180 * PI;
  double m_ni = log(q * (1 - r)) - log(r * (1 - q));
  double c    = log(Q) - log(1 - Q);
  double m_nf = log(1 - r) - log(1 - q);
  messagePrintf(1, "run prob_mod, ni * %.7f > %.7f + nf * %.7f ???",
                m_ni, c, m_nf);
  for (size_t i = 0, I = input_bitmap.size(); i < I; i++) {
    for (size_t j = 0, J = input_bitmap.size(); j < J; j++) {
      if (pairs[i][j].empty()) {
        continue;
      }
      double ni = pairs[i][j].size(), nf = 0;
      MyRansacCheck chk(&(fpsv[i]), &(fpsv[j]));
      chk.rememberVCalc(pairs[i][j]);
      if (chk.check(aspect_t, angle_t)) {
        for (size_t k = 0, K = fpsv[i].size(); k < K; k++) {
          Vector2D<double> to(chk.to(Vector2D<double>(fpsv[i][k](0),
                                                      fpsv[i][k](1))));
          if (0 <= to.x() && to.x() <= (double)input_bitmap[j].width() &&
              0 <= to.y() && to.y() <= (double)input_bitmap[j].height()) {
            nf++;
          }
        }
        if (ni * m_ni > c + m_nf * nf) {
          messagePrintf(0, "accept %lu --- %lu", i, j);
          messagePrintf(0,
                        "%.0f * %.3f = %.3f ?? %.3f = %.3f + %.3f * %.0f",
                        ni, m_ni, ni * m_ni, c + m_nf * nf, c, m_nf, nf);
          continue;
        }
      }
      pairs[i][j].clear();
    }
  }
  messagePrintf(-1, "ok");
  return true;
}

////////////////////// **# Write to output file #** //////////////////

bool output() {
  messagePrintf(1, "Write images");
  for (size_t i = 0; i < output_bitmap.size(); i++) {
    size_t width  = output_bitmap[i].width ();
    size_t height = output_bitmap[i].height();
    cv::Mat img(height, width, CV_8UC3);
    for (size_t x = 0; x < width; x++) {
      for (size_t y = 0; y < height; y++) {
        RGBi_Space tmp;
        colorTransformate(output_bitmap[i].pixel(y, x), &tmp);
        img.at<cv::Vec3b>(y, x)[0] = tmp.b();
        img.at<cv::Vec3b>(y, x)[1] = tmp.g();
        img.at<cv::Vec3b>(y, x)[2] = tmp.r();
      }
    }
    std::string output_name(usg.optionValue("o", 0)
                            + (output_bitmap.size() > 1
                               ? stringPrintf("%lu", i)
                               : "")
                            + ".jpg");
    messagePrintf(1, "Write to file '%s'", output_name.c_str());
    if (imwrite(output_name, img) == false) {
      messagePrintf(-1, "opencv fail, ignore");
    }
    else {
      messagePrintf(-1, "%lux%lu, ok", width, height);
    }
  }
  messagePrintf(-1, "ok");
  return true;
}

bool pair_output(){
  for(size_t i = 0, I = input_bitmap.size(); i < I; i++){
    for(size_t j = 0, J = input_bitmap.size(); j < J; j++){
      if(pairs[i][j].empty()) continue;
      MyRansacCheck chk(&(fpsv[i]), &(fpsv[j]));
      chk.rememberVCalc(pairs[i][j]);
      size_t index = output_bitmap.size();
      output_bitmap.push_back(input_bitmap[i]);
      for(ssize_t x = 0, X = input_bitmap[i].width(); x < X; x++) {
        for(ssize_t y = 0, Y = input_bitmap[i].height(); y < Y; y++) {
          Vector2D<double> to(chk.to(Vector2D<double>(x, y)));
          ssize_t x2 = to.x(), y2 = to.y();
          if (0 <= x2 && x2 <= (ssize_t)input_bitmap[j].width() &&
              0 <= y2 && y2 <= (ssize_t)input_bitmap[j].height()) {
            output_bitmap[index].pixel(y, x, (input_bitmap[i].pixel(y, x) +
                                              input_bitmap[j].pixel(y2,x2)) / 2
                                       );
          }
        }
      }
      for (size_t k = 0, K = pairs[i][j].size(); k < K; ++k) {
        ssize_t x0 = fpsv[i][pairs[i][j][k].from.second](0);
        ssize_t y0 = fpsv[i][pairs[i][j][k].from.second](1);
        for (ssize_t d = -10; d <= 10; ++d) {
          if (0 <= x0 + d && x0 + d < (ssize_t)input_bitmap[i].width() - 1) {
            output_bitmap[index].pixel(y0, x0 + d,
                                       RGBf_Space(Vector3D<double>(
                                           1.0, 1.0, 0.0
                                       )));
          }
          if (0 <= y0 + d && y0 + d < (ssize_t)input_bitmap[i].height() - 1) {
            output_bitmap[index].pixel(y0 + d, x0,
                                       RGBf_Space(Vector3D<double>(
                                           1.0, 1.0, 0.0
                                       )));
          }
        }
      }
    }
  }
  return output();
}

bool text_output() {
  std::string s = usg.optionValue("f", 0);
  FILE* f = fopen(s.c_str(), "w");
  fprintf(f, "%lu\n", input_bitmap.size());
  for (size_t i = 0, I = input_bitmap.size(); i < I; ++i) {
    fprintf(f, "%s\n", input_name[i].c_str());
    fprintf(f, "%lu %lu %lu   ",
            input_bitmap[i].height(), input_bitmap[i].width(), fpsv[i].size());
    for (size_t j = 0, J = fpsv[i].size(); j < J; ++j) {
      fprintf(f, "%.10f %.10f  ", fpsv[i][j](0), fpsv[i][j](1));
    }
    fprintf(f, "\n");
    for (size_t j = 0; j < I; ++j) {
      fprintf(f, "%lu   ", pairs[i][j].size());
      for (size_t k = 0, K = pairs[i][j].size(); k < K; ++k) {
        fprintf(f, "%lu %lu %lu %lu   ",
                pairs[i][j][k].from.first,
                pairs[i][j][k].from.second,
                pairs[i][j][k].to.first,
                pairs[i][j][k].to.second);
      }
      fprintf(f, "\n");
    }
  }
  fclose(f);
  return true;
}

int main(int argc, char** argv){
  setup(argc, argv);
  input();
  detect();
  kmatch();
  ransac();
  match_check();
  pair_output();
  text_output();
  return 0;
}