#include <opencv2/opencv.hpp>

#include <iostream>
#include <vector>
#include <cmath>
#include <assert.h>

using namespace std;
using namespace cv;

#define MASK_BACKGROUND (1<<7)
#define MAX_ITERATION 300

inline uchar clampValue (float val) {
  if (val > 255) return 255;
  else if (val < 0) return 0;
  else return (uchar)val;
}

void montage (IplImage* srcImage, IplImage* destImage, IplImage* maskImage) {
  if (!srcImage || !destImage || !maskImage) {
    cerr << "cannot find image file" << endl;
    exit(-1);
  }

  int w = srcImage->width;
  int h = srcImage->height;

  if (w != maskImage->width || h != maskImage->height) {
    cerr << "mask size doesn't match src size" << endl;
    exit(-1);
  }

  vector<int> destPoints;
  vector<int> constants;

  int size=0;
  for (int y=0; y<h; y++) {
    for (int x=0; x<w; x++) {
      int maskIndex = maskImage->widthStep * y + x;
      uchar c = maskImage->imageData[maskIndex];
      //printf("%d\n", (int)c);
      //if (!(c & MASK_BACKGROUND)) {
      if (c & MASK_BACKGROUND) {
        //printf("input (%d,%d)\n", x, y);
        int destIndex = destImage->widthStep * y + (x * 3);
        destPoints.push_back(destIndex);

        int srcIndex = srcImage->widthStep * y + (x * 3);

        int constant[3] = {0};
        int sum1[3]={0};
        int sum2[3]={0};
        // right
        if (x < srcImage->width-1) {
          for (int i=0; i<3; i++) {
            int val1 = (uchar)(destImage->imageData[destIndex+3+i]) - (uchar)(destImage->imageData[destIndex+i]);
            int val2 = (uchar)(srcImage->imageData[srcIndex+3+i]) - (uchar)(srcImage->imageData[srcIndex+i]);
            //printf("%d,%d\n", val1, val2);
            //constant[i] += min(val1, val2);
            sum1[i]+=val1;
            sum2[i]+=val2;
            //constant[i] += val2;
          }
        }
        // left
        if (x > 0) {
          for (int i=0; i<3; i++) {
            int val1 = (uchar)(destImage->imageData[destIndex-3+i]) - (uchar)(destImage->imageData[destIndex+i]);
            int val2 = (uchar)(srcImage->imageData[srcIndex-3+i]) - (uchar)(srcImage->imageData[srcIndex+i]);
            //printf("%d,%d\n", val1, val2);
            //constant[i] += min(val1, val2);
            sum1[i]+=val1;
            sum2[i]+=val2;
            //constant[i] += val2;
          }
        }
        // top
        if (y > 0) {
          for (int i=0; i<3; i++) {
            int val1 = (uchar)(destImage->imageData[destImage->widthStep*(y-1)+(x*3)+i]) - (uchar)(destImage->imageData[destIndex+i]);
            int val2 = (uchar)(srcImage->imageData[srcImage->widthStep*(y-1)+(x*3)+i]) - (uchar)(srcImage->imageData[srcIndex+i]);
            //printf("%d,%d\n", val1, val2);
            //constant[i] += min(val1, val2);
            sum1[i]+=val1;
            sum2[i]+=val2;
            //constant[i] += val2;
          }
        }
        // bottom
        if (y < srcImage->height-1) {
          for (int i=0; i<3; i++) {
            int val1 = (uchar)(destImage->imageData[destImage->widthStep*(y+1)+(x*3)+i]) - (uchar)(destImage->imageData[destIndex+i]);
            int val2 = (uchar)(srcImage->imageData[srcImage->widthStep*(y+1)+(x*3)+i]) - (uchar)(srcImage->imageData[srcIndex+i]);
            //printf("%d,%d\n", val1, val2);
            //constant[i] += min(val1, val2);
            sum1[i]+=val1;
            sum2[i]+=val2;
            //constant[i] += val2;
          }
        }
        for (int i=0; i<3; i++) {
          //constants.push_back(constant[i]);
          //if (abs(sum1[i]) > abs(sum2[i])){
          //  constants.push_back(sum1[i]);
          //} else {
          constants.push_back(sum2[i]);
          //}
        }
      }
      // int offset = srcImage->widthStep * y + (x * 3);
      size++;
    }
  }

  printf("destPoints size=%d\n", (int)destPoints.size());
  printf("constants size=%d\n", (int)constants.size());
  uchar* destImageData = (uchar*)destImage->imageData;
  cv::Mat final(3*destImage->width, destImage->height, DataType<float>::type);
  for (int x=0; x<destImage->width; x++) for (int y=0; y<destImage->height; y++)
    for (int i=0; i<3; i++)
      final.at<float>(x*3+i, y) = ((uchar)destImage->imageData[destImage->widthStep*y+3*x+i]);

  //for (int x=0; x<destImage->width; x++) for (int y=0; y<destImage->height; y++)
  //  printf("(%f,%f,%f)\n", final.at<float>(x*3+0, y), final.at<float>(x*3+1, y), final.at<float>(x*3+2, y));

  // ヤコビ法で連立一次方程式を解く
  for (int loop=0; loop<MAX_ITERATION; loop++) {
    int n = destPoints.size();
    for (int p=0; p<n; p++) {
      int destIndex = destPoints[p];
      int y = destIndex / (destImage->widthStep);
      int x = (destIndex % (destImage->widthStep)) / 3;
      //printf("check (%d,%d)\n", x, y);
      float values[3] = {0};
      // right
      int count = 0;
      if (x < destImage->width-1) {
        count++;
        for (int i=0; i<3; i++) {
          //values[i] += (uchar)(destImageData[destIndex+3+i]);
          values[i] += final.at<float>((3*(x+1))+i, y);
        }
      }
      // left
      if (x > 0) {
        count++;
        for (int i=0; i<3; i++) {
          //values[i] += (uchar)(destImageData[destIndex-3+i]);
          values[i] += final.at<float>((3*(x-1))+i, y);
        }
      }
      // top
      if (y > 0) {
        count++;
        for (int i=0; i<3; i++) {
          //values[i] += (uchar)(destImageData[destImage->widthStep*(y-1)+(x*3)+i]);
          values[i] += final.at<float>((3*x)+i, y-1);
        }
      }
      // bottom
      if (y < destImage->height-1) {
        count++;
        for (int i=0; i<3; i++) {
          //values[i] += (uchar)(destImageData[destImage->widthStep*(y+1)+(x*3)+i]);
          values[i] += final.at<float>((3*x)+i, y+1);
        }
      }
      //for (int i=0; i<3; i++) {
      //  values[i] -= count*((uchar)(destImageData[destIndex+i]));
      //}
      // 更新
      for (int j=0; j<3; j++) {
        float newval = (values[j] - constants[p*3+j]) / count;
        float oldval = final.at<float>((3*x)+j, y);
        //if (newval >= 256) newval = 255;
        //if (newval < 0) newval = 0;
        //destImageData[destIndex+j] = (uchar)newval;
        //printf("%f->%f\n", oldval, newval);
        final.at<float>((3*x)+j, y) = newval;
      }
    }
  }

  {
    int n = destPoints.size();
    for (int p=0; p<n; p++) {
      int destIndex = destPoints[p];
      int y = destIndex / (destImage->widthStep);
      int x = (destIndex % (destImage->widthStep)) / 3;
      //printf("set (%d,%d)\n", x, y);
      for (int i=0; i<3; i++) {
        //printf(" %d->%d\n", x, y, (uchar)destImage->imageData[destIndex + i], clampValue(final.at<float>(x*3+i, y)));
        destImage->imageData[destIndex + i] = clampValue(final.at<float>(x*3+i, y));
      }
    }
  }
  return;
}

int main(int argc, char** argv) {
  const char* srcFile = argc == 4 ? argv[1] : "black1.jpg";
  const char* destFile = argc == 4 ? argv[2] : "takahiro.jpg";
  //const char* filename2 = argc == 3 ? argv[2] : "photo2.jpg";
  const char* maskFile = argc == 4 ? argv[3] : "mask_black.jpg";

  IplImage* srcImage = cvLoadImage(srcFile, CV_LOAD_IMAGE_ANYCOLOR|CV_LOAD_IMAGE_ANYDEPTH);
  IplImage* destImage = cvLoadImage(destFile, CV_LOAD_IMAGE_ANYCOLOR|CV_LOAD_IMAGE_ANYDEPTH);
  IplImage* maskImage = cvLoadImage(maskFile, CV_LOAD_IMAGE_GRAYSCALE);

  if (!srcImage || !destImage) {
    cerr << "cannot find image file" << endl;
    exit(-1);
  }

  int w = srcImage->width;
  int h = srcImage->height;

  if (w != maskImage->width || h != maskImage->height) {
    cerr << "mask size doesn't match src size" << endl;
    exit(-1);
  }

  montage(srcImage, destImage, maskImage);

  cvShowImage("Poisson Image Editing", destImage);
  cvSaveImage("poisson.jpeg", destImage);

  cvWaitKey(0);

  // 後始末
  cvReleaseImage(&maskImage);
  cvReleaseImage(&srcImage);
  cvReleaseImage(&destImage);

  return 0;
}