// Copyright (c) 2008 Jakub Travnik
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.


// Note that this code produces by no means accurate
// images of what the camera sensor and its software does.

PImage img;

int imagesize=400; // expected size of source image side
int windowsizew=3*imagesize, windowsizeh=2*imagesize;
int fullcolorpixelsize=16; // side of square area in source image that makes full color pixel
float fullcolorfillfactor=0.8; // square root of fill factor of full color pixel, 1.0 means full fill, 0.707 means half 50% area fill
int cfapixelsize=9; // side of square area in source image that makes cfa subpixel
float cfaaafilter=2.0; // 1.0 means no aafilter (1x1 pixel area), 2.0 means moderate filter (2x2 pixel area is used)

boolean strongcfaaafilter=true;
// if true, aa filter area will be used for subpixel color
// if false, aa filter area color blended with subpixel area color are used for subpixel color (center is prefered)

void setup() {
//fullcolorfillfactor=1.0;
//cfaaafilter=1.0;
  size(windowsizew,windowsizeh);
  colorMode(RGB, 1.0);
  frameRate(10);
  img = loadImage("input.jpg");
  img.loadPixels();
}

void drawGridCFA(int ofsx, int ofsy, int gridsize) {
  pushMatrix();
  image(img, 0, 0);
  rectMode(CORNER);
  int xlen=imagesize/gridsize-2;
  int ylen=imagesize/gridsize-2;
  PImage fimg = createImage(xlen, ylen, RGB);
  PImage fimgr = createImage(xlen, ylen, RGB);
  PImage fimgg = createImage(xlen, ylen, RGB);
  PImage fimgb = createImage(xlen, ylen, RGB);
  int samplesizeaa = int(cfaaafilter*gridsize);
  int sampleofsaa = (samplesizeaa-gridsize)/2;
  int samplesizecore = gridsize;
  int sampleofscore = 0;
  for (int x=0; x<xlen; x++) {
    for (int y=0; y<ylen; y++) {
      int gx=x*gridsize+ofsx;
      int gy=y*gridsize+ofsy;
      noFill();
      boolean bx=x%2==1;
      boolean by=y%2==1;
      color caafilter = sample(img, gx-sampleofsaa, gy-sampleofsaa, samplesizeaa, samplesizeaa);
      color c;
      if (strongcfaaafilter) {
        c=caafilter;
      } else {
        color ccore = sample(img, gx-sampleofscore, gy-sampleofscore, samplesizecore, samplesizecore);
        c = blendColors2(ccore, caafilter);
      }
      color c2;
      noStroke();
      if (bx==by) {
        //green
        c2=color(0.0, green(c), 0.0);
        fimgg.pixels[y*xlen+x]=c2;
      } else if (bx) {
        //red
        c2=color(red(c), 0.0, 0.0);
        fimgr.pixels[y*xlen+x]=c2;
      } else {
        // blue
        c2=color(0.0, 0.0, blue(c));
        fimgb.pixels[y*xlen+x]=c2;
    }
      fimg.pixels[y*xlen+x]=c2;
      fill(c2);
      rect(gx, gy, gridsize-1, gridsize-1);
    }
  }
  fimg.updatePixels();
  fimgr.updatePixels();
  fimgg.updatePixels();
  fimgb.updatePixels();
  noSmooth();
  translate(imagesize, 0);
  scale(4);
  image(fimgr,0,0);
  image(fimgg,fimg.width,0);
  image(fimgb,2*fimg.width,0);
  image(fimg,3*fimg.width,0);
  translate(0, fimg.height);
  // interpolate colors in each color plane
  fimg.loadPixels();
  fimgr.loadPixels();
  fimgg.loadPixels();
  fimgb.loadPixels();
  for(int y=0; y<ylen; y++) {
    for(int x=0; x<xlen; x++) {
      boolean bx=x%2==1;
      boolean by=y%2==1;
      // red
      if (bx==by) {
        if (bx) {
          fimgr.pixels[y*xlen+x]=blendColors2(getcolor(fimgr, x, y-1), getcolor(fimgr, x, y+1));
        } else {
          fimgr.pixels[y*xlen+x]=blendColors2(getcolor(fimgr, x-1, y), getcolor(fimgr, x+1, y));
        }
      }
      if (by && !bx) {
        fimgr.pixels[y*xlen+x]=blendColors4(getcolor(fimgr, x-1, y-1), getcolor(fimgr, x-1, y+1),
                                            getcolor(fimgr, x+1, y-1), getcolor(fimgr, x+1, y+1));
      }
      // blue
      if (bx==by) {
        if (bx) {
          fimgb.pixels[y*xlen+x]=blendColors2(getcolor(fimgb, x-1, y), getcolor(fimgb, x+1, y));
        } else {
          fimgb.pixels[y*xlen+x]=blendColors2(getcolor(fimgb, x, y-1), getcolor(fimgb, x, y+1));
        }
      }
      if (!by && bx) {
        fimgb.pixels[y*xlen+x]=blendColors4(getcolor(fimgb, x-1, y-1), getcolor(fimgb, x-1, y+1),
                                            getcolor(fimgb, x+1, y-1), getcolor(fimgb, x+1, y+1));
      }
      //green
      if (bx!=by) {
        fimgg.pixels[y*xlen+x]=blendColors4(getcolor(fimgg, x, y-1), getcolor(fimgg, x, y+1),
                                            getcolor(fimgg, x-1, y), getcolor(fimgg, x+1, y));
      }
    }
  }
  for(int y=0; y<ylen; y++) {
    for(int x=0; x<xlen; x++) {
      fimg.pixels[y*xlen+x]=cfaMergeColors3(fimgr.pixels[y*xlen+x], fimgg.pixels[y*xlen+x], fimgb.pixels[y*xlen+x]);
    }
  }
  fimg.updatePixels();
  fimgr.updatePixels();
  fimgg.updatePixels();
  fimgb.updatePixels();
  image(fimgr,0,0);
  image(fimgg,fimg.width,0);
  image(fimgb,2*fimg.width,0);
  image(fimg,3*fimg.width,0);
  scale(1);
  popMatrix();
}

color getcolor(PImage im, int x, int y) {
  if (x<0 || y<0 || x>=im.width || y>=im.height) {
    return color(0.0, 0.0, 0.0, 0.0);
  }
  return im.pixels[y*im.width+x];
}

color blendColors4(color c1, color c2, color c3, color c4) {
  float r=0.0,g=0.0,b=0.0,a=0.0;
  float a1=alpha(c1);
  r+=red(c1)*a1;
  g+=green(c1)*a1;
  b+=blue(c1)*a1;
  a+=a1;
  float a2=alpha(c2);
  r+=red(c2)*a2;
  g+=green(c2)*a2;
  b+=blue(c2)*a2;
  a+=a2;
  float a3=alpha(c3);
  r+=red(c3)*a3;
  g+=green(c3)*a3;
  b+=blue(c3)*a3;
  a+=a3;
  float a4=alpha(c4);
  r+=red(c4)*a4;
  g+=green(c4)*a4;
  b+=blue(c4)*a4;
  a+=a4;
  return color(r/a,g/a,b/a);
}

color blendColors2(color c1, color c2) {
  float r=0.0,g=0.0,b=0.0,a=0.0;
  float a1=alpha(c1);
  r+=red(c1)*a1;
  g+=green(c1)*a1;
  b+=blue(c1)*a1;
  a+=a1;
  float a2=alpha(c2);
  r+=red(c2)*a2;
  g+=green(c2)*a2;
  b+=blue(c2)*a2;
  a+=a2;
  return color(r/a,g/a,b/a);
}

color cfaMergeColors3(color c1, color c2, color c3) {
  float r=0.0,g=0.0,b=0.0;
  r+=red(c1);
  g+=green(c1);
  b+=blue(c1);
  r+=red(c2);
  g+=green(c2);
  b+=blue(c2);
  r+=red(c3);
  g+=green(c3);
  b+=blue(c3);
  return color(r,g,b);
}



void drawGridFullcolor(int ofsx, int ofsy, int gridsize, int fillsize) {
  pushMatrix();
  image(img, 0, 0);
  rectMode(CORNER);
  int xlen=imagesize/gridsize-1;
  int ylen=imagesize/gridsize-1;
  PImage fimg = createImage(xlen, ylen, RGB);
  for (int x=0; x<xlen; x++) {
    for (int y=0; y<ylen; y++) {
      int gx=x*gridsize+ofsx;
      int gy=y*gridsize+ofsy;
      noFill();
      stroke(#cccccc);
      rect(gx, gy, gridsize, gridsize);
      int fillshift = (gridsize-fillsize)/2;
      rect(gx+fillshift, gy+fillshift, fillsize, fillsize);
      color c=sample(img, gx+fillshift, gy+fillshift, fillsize, fillsize);
      fimg.pixels[y*xlen+x]=c;
      fill(c);
      noStroke();
      rect(gx+fillshift, gy+fillshift, fillsize, fillsize);
    }
  }
  noSmooth();
  translate(imagesize, 0);
  image(fimg,0,0);
  image(fimg,fimg.width,0,fimg.width*2, fimg.height*2);
  image(fimg,fimg.width*3,0,fimg.width*4, fimg.height*4);
  image(fimg,fimg.width*7,0,fimg.width*8, fimg.height*8);
  popMatrix();
}

color sample(PImage im, int x, int y, int w, int h) {
  int count = 0;
  float r=0.0,g=0.0,b=0.0;
  if (x<0) x=0;
  if (y<0) y=0;
  if (x+w>im.width) w=im.width-x;
  if (y+h>im.height) h=im.height-y;
  for(int yy=y; yy<y+h; yy++) {
    for(int xx=x; xx<x+w; xx++) {
      color c=im.pixels[yy*im.width+xx];
      r+=red(c);
      g+=green(c);
      b+=blue(c);
      count++;
    }
  }
  return color(r/count,g/count,b/count);
}

void draw() {
  textFont(createFont("SanSerif", 12));
  fill(#cccccc);
  noStroke();
  rect(0,0, windowsizew, windowsizeh);
  int fullcolorx=mouseX % fullcolorpixelsize;
  int fullcolory=mouseY % fullcolorpixelsize;
  drawGridFullcolor(fullcolorx, fullcolory, fullcolorpixelsize, int(fullcolorfillfactor*fullcolorpixelsize));
  fill(#000000);
  text("Full color sensor settings", imagesize*2, 20);
  text("Grid origin: x="+fullcolorx+" y="+fullcolory, imagesize*2, 40);
  text("Fill factor (area): "+round(fullcolorfillfactor*fullcolorfillfactor*100)+"%", imagesize*2, 60);
  text("Source pixels per pixel (area): "+round(fullcolorpixelsize*fullcolorpixelsize), imagesize*2, 80);
  translate(0,imagesize);
  int cfax=mouseX % (2*cfapixelsize);
  int cfay=mouseY % (2*cfapixelsize);
  drawGridCFA(cfax, cfay, cfapixelsize);
  fill(#000000);
  text("Bayer CFA sensor settings", imagesize*2, -90);
  text("Grid origin: x="+cfax+" y="+cfay, imagesize*2, -70);
  text("AA Filter blur (area): "+round(cfaaafilter*cfaaafilter), imagesize*2, -50);
  text("AA Filter mode (true=flat, false=center weigthed): "+strongcfaaafilter, imagesize*2, -30);
  text("Source pixels per subpixel (area): "+round(cfapixelsize*cfapixelsize), imagesize*2, -10);
}