#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/gra/WatchBall.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;

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('d',
                "Specify which Feature-Point-Detect algorithm to use",
                "<algorithm>",
                "",
                true);
  usg.optionAdd('p',
                "Pribabilicity for RANSAC to choose a right feature point",
                "<floating point>", stringPrintf("%.10f", p0),
                false);
  usg.optionAdd('P',
                "Pribabilicity for RANSAC access",
                "<floating point>", stringPrintf("%.10f", P),
                false);
  usg.optionAdd('q',
                "p1 for Prob. Model",
                "<floationg Point>", stringPrintf("%.10f", q),
                false);
  usg.optionAdd('r',
                "p0 for Prob. Model",
                "<floationg Point>", stringPrintf("%.10f", r),
                false);
  usg.optionAdd('Q',
                "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('O',
                "output ball radius",
                "<floationg Point>", stringPrintf("%.10f", o_radius),
                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() {
  std::vector<std::string> input_name;
  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('p', 0).c_str()));
  P  = inRange(0.00001, 0.9999, atof(usg.optionValue('P', 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 = 3u;
      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;
}

///////////////////////// **# prob module #** ////////////////////////
bool prob_mod() {
  q = inRange(0.00001, 0.99999, atof(usg.optionValue('q', 0).c_str()));
  r = inRange(0.00001, 0.99999, atof(usg.optionValue('r', 0).c_str()));
  Q = inRange(0.00001, 0.99999, atof(usg.optionValue('Q', 0).c_str()));
  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]);
      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;
      }
      else {
        pairs[i][j].clear();
      }
    }
  }
  messagePrintf(-1, "ok");
  return true;
}

///////////////////// **# group them together #** ////////////////////
bool group() {
  messagePrintf(1, "group");
  // union
  DisjointSet dsj(input_bitmap.size());
  for (size_t i = 0, I = input_bitmap.size(); i < I; i++) {
    for (size_t j = 0; j < I; j++) {
      if(pairs[i][j].empty()) continue;
      dsj.merge(i, j);
    }
  }
  std::vector<size_t> root;
  for (size_t i = 0, I = input_bitmap.size(); i < I; i++) {
    if (dsj.root(i) == i) {
      root.push_back(i);
    }
  }
  // split into groups
  outputs.resize(root.size());
  for (size_t i = 0, I = root.size(); i < I; i++) {
    messagePrintf(1, "Group %d", i);
    std::vector<size_t> ids;
    for (size_t j = 0, J = input_bitmap.size(); j < J; j++) {
      if (dsj.root(j) != root[i]) continue;
      outputs[i].cameras.push_back(Camera<RGBf_Space>());
      outputs[i].cameras[outputs[i].cameras.size() - 1].photo(
        Photo<RGBf_Space>(input_bitmap[j])
      );
      ids.push_back(j);
      messagePrintf(0, "camera %lu from bitmap %lu",
                    outputs[i].cameras.size() - 1, j);
    }
    for (size_t j = 0, J = ids.size(); j < J; ++j) {
      for (size_t k = 0; k < J; ++k) {
        if (j == k) continue;
        size_t i1 = ids[j], i2 = ids[k];
        if (pairs[i1][i2].empty()) continue;
        outputs[i].edges.push_back(OutputSet::Edge(j, k));
        size_t index = outputs[i].edges.size() - 1;
        for (size_t n = 0, N = pairs[i1][i2].size(); n < N; ++n) {
          outputs[i].edges[index].v1.push_back(
            fpsv[i1][pairs[i1][i2][n].from.second]
          );
          outputs[i].edges[index].v2.push_back(
            fpsv[i1][pairs[i1][i2][n].to.second]
          );
        }
        messagePrintf(0, "Edge %lu---%lu, size = %lu",
                      i1, i2, outputs[i].edges[index].v1.size());
      }
    }
    std::sort(outputs[i].edges.begin(), outputs[i].edges.end());
    messagePrintf(-1, "");
  }
  messagePrintf(-1, "ok");
  return true;
}

////////////////////// **# boundle adjustment #** ////////////////////
bool boundle() {
  stop = inRange(0.01, 100000.0, atof(usg.optionValue('s', 0).c_str()));
  messagePrintf(1, "boundle adjustment");
  for (size_t i = 0, I = outputs.size(); i < I; i++) {
    int id = 0;
    std::set<size_t> in;
    size_t i1 = outputs[i].edges[0].i1;
    size_t i2 = outputs[i].edges[0].i2;
    for (size_t j = 0, J = outputs[i].edges[0].v1.size(); j < J; ++j) {
      outputs[i].cameras[i1].fixedPoints2DGet().identityPointAdd(
        id, outputs[i].edges[0].v1[j]);
      outputs[i].cameras[i2].fixedPoints2DGet().identityPointAdd(
        id, outputs[i].edges[0].v2[j]);
      id++;
    }
    in.insert(i1);
    in.insert(i2);
    double r_lst = Camera<RGBf_Space>::boundleAdjustment2D(
      &(outputs[i].cameras),
      stop
    );
    for (size_t j = 1, J = outputs[i].edges.size(); j < J; ++j) {
      size_t best;
      for (best = 0; best < J; ++best) {
        if (in.find(outputs[i].edges[best].i1) == in.end() &&
            in.find(outputs[i].edges[best].i2) == in.end()) continue;
        break;
      }
      i1 = outputs[i].edges[best].i1;
      i2 = outputs[i].edges[best].i2;
      for (size_t j = 0, J = outputs[i].edges[best].v1.size(); j < J; ++j) {
        outputs[i].cameras[i1].fixedPoints2DGet().identityPointAdd(
          id, outputs[i].edges[best].v1[j]);
        outputs[i].cameras[i2].fixedPoints2DGet().identityPointAdd(
          id, outputs[i].edges[best].v2[j]);
        id++;
      }
      in.insert(i1);
      in.insert(i2);
      std::vector<Camera<RGBf_Space> > tmp(outputs[i].cameras);
      double r = Camera<RGBf_Space>::boundleAdjustment2D(&tmp, stop);
      if (r > r_lst * 1.5) continue;
      outputs[i].cameras = tmp;
    }
  }
  messagePrintf(-1, "ok");
  return true;
}

bool expand() {
  o_radius = inRange(100.0, 1000000.0, atof(usg.optionValue('O', 0).c_str()));
  output_bitmap.resize(outputs.size());
  for (size_t i = 0, I = outputs.size(); i < I; ++i) {
    WatchBall<RGBf_Space> wb;
    wb.cameras(outputs[i].cameras);
    output_bitmap[i] = wb.expand(o_radius);
  }
  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 tmp_output() {
  output_bitmap = input_bitmap;
  for (size_t i = 0, I = input_bitmap.size(); i < I; i++) {
    for (size_t j = 0, J = fpsv[i].size(); j < J; j++) {
      ssize_t x = fpsv[i][j](0);
      ssize_t y = fpsv[i][j](1);
      ssize_t dx[2] = {0, 1}, x0[2] = {0, -10};
      ssize_t dy[2] = {1, 0}, y0[2] = {-10, 0};
      for(size_t k = 0; k < 2; k++){
        for(size_t count = 0; count < 20; count++){
          ssize_t xx = x + dx[k] * count + x0[k];
          ssize_t yy = y + dy[k] * count + y0[k];
          if(0 <= xx && xx < (ssize_t)input_bitmap[i].width() &&
             0 <= yy && yy < (ssize_t)input_bitmap[i].height()){
            output_bitmap[i].pixel(yy, xx, Vector3D<double>(1.0, 1.0, 0.0));
          }
        }
      }
    }
  }
  return output();
}
// */

/*
bool g_output(){
  output_bitmap.resize(input_bitmap.size() * 2);
  for(size_t i = 0, I = input_bitmap.size(); i < I; i++){
    output_bitmap[i * 2    ] = input_bitmap[i];
    output_bitmap[i * 2 + 1] = input_bitmap[i];
    output_bitmap[i * 2    ].gradiancedX(3, 3);
    output_bitmap[i * 2 + 1].gradiancedY(3, 3);
    for(size_t x = 0, X = output_bitmap[i * 2].width(); x < X; x++){
      for(size_t y = 0, Y = output_bitmap[i * 2].height(); y < Y; y++){
        Vector3D<double> v;
        v = output_bitmap[i * 2](y, x);
        output_bitmap[i * 2](y, x) = Vector3D<double>(v.length() / sqrt(2.0));
        v = output_bitmap[i * 2 + 1](y, x);
        output_bitmap[i * 2 + 1](y, x) = Vector3D<double>(v.length() / sqrt(2.0));
      }
    }
  }
  return output();
}

// */

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
                                       );
          }
        }
      }
    }
  }
  return output();
}
/*
bool pair_output2(){
  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((i + 1) % I != j && (j + 1) % J != i) continue;
      messagePrintf(0, "%3lu--%3lu: %lu", i, j, pairs[i][j].size());
      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 = fpsv[i].size(); k < K; k++){
        ssize_t dy[2] = {0, 1}, dx[2] = {1, 0};
        ssize_t y0[2] = {0, -10}, x0[2] = {-10, 0};
        ssize_t x = fpsv[i][k](0), y = fpsv[i][k](1);
        for(ssize_t m = 0; m < 2; m++){
          for(ssize_t n = 0; n < 20; n++){
            ssize_t xx = x + x0[m] + dx[m] * n;
            ssize_t yy = y + y0[m] + dy[m] * n;
            if(0 <= xx && xx <= (ssize_t)input_bitmap[j].width() &&
               0 <= yy && yy <= (ssize_t)input_bitmap[j].height()){
              output_bitmap[index].pixel(yy, xx,
                                         Vector3D<double>(1.0, 1.0, 0));
            }
          }
        }
      }
      for(size_t k = 0, K = pairs[i][j].size(); k < K; k++){
        ssize_t dy[2] = {0, 1}, dx[2] = {1, 0};
        ssize_t y0[2] = {0, -10}, x0[2] = {-10, 0};
        ssize_t x = fpsv[i][pairs[i][j][k].from.second](0);
        ssize_t y = fpsv[i][pairs[i][j][k].from.second](1);
        for(ssize_t m = 0; m < 2; m++){
          for(ssize_t n = 0; n < 20; n++){
            ssize_t xx = x + x0[m] + dx[m] * n;
            ssize_t yy = y + y0[m] + dy[m] * n;
            if(0 <= xx && xx <= (ssize_t)input_bitmap[j].width() &&
               0 <= yy && yy <= (ssize_t)input_bitmap[j].height()){
              output_bitmap[index].pixel(yy, xx,
                                         Vector3D<double>(1.0, 0.0, 0));
            }
          }
        }
      }
    }
  }
  return output();
}
// */


int main(int argc, char** argv){
  setup(argc, argv);
  input();
  detect();
  kmatch();
  ransac();
  prob_mod();
  group();
  pair_output(); return 0;
  boundle();
  expand();
  output();
  return 0;
}