#include <string.h>
#include <stdio.h>
#include <math.h>
#include "fitsio.h"

// Author: Kristopher Setnes
// Date: July 5, 2020

// This program computes the difference of the average pixel value between 
// auto focus pixels and non auto focus pixels.  A new file is written with
// the fudge amount subtracted from the auto focus pixels.

// This script was written by looking at example cfitsio programs.
// https://heasarc.gsfc.nasa.gov/docs/software/fitsio/cexamples.html

// A post on PixInsight forums was also of great help.
// https://pixinsight.com/forum/index.php?threads/strange-results-and-questions-related-to-imagecalibration-and-drizzleintegration.14612/


// This returns true if the pixel is a special auto focus pixel.
// I only have the Xtrans sensor used in a Fujifilm X-T20.
// Other cameras need to be added to this function.
int is_af_pixel(char *camera, int x, int y) { 
  // x=1, y=1 is the bottom left corner of the image.

  if ( strcmp(camera, "'Fujifilm X-T20'")==0 ) {
    // We know these pixels. 
    return ( x>=1511 && x<=4520 && y>=508 && y<=3524 ) &&
           ( (x-2)%3==0 && ( (y-4)%12==0 || (y-8)%12==0 ) );
  } else {
    // We don't know this camera. 
    printf("Unknown camera %s.\n",camera);
    exit(EXIT_FAILURE);
  }
}


int main(int argc, char *argv[])
{
    // Input and Output File Pointers.
    fitsfile *infptr, *outfptr;
 
    // CFITSIO status initialized to 0.
    int status = 0;

    // Unit Type, Number of Axis
    int hdutype, naxis;

    // axis[0] = count of X pixels.
    // axis[1] = count of Y pixels.
    // fpixel[0] = first X pixel (controls read and writes)
    // fpixel[1] = first Y pixel (controls read and writes)
    long naxes[2], fpixel[2];

    // This is for an entire row of pixels.
    double *pix;

    // non-focus pixels
    double nfsum = 0., nfmean = 0., nfmin = 1.E33, nfmax = -1.E33;
    long nfcount=0;

    // auto focus pixels
    double afsum = 0., afmean = 0., afmin = 1.E33, afmax = -1.E33;
    long afcount=0;

    // The fudge amount to apply to auto focus pixels. (computed)
    double fudge;

    // The target sample size (needs to be <= AF area)
    // We want to avoid any amp glow that might be visible in the AF area.
    // This value might need to be tweaked for different sensors.
    // If you want to sample the entire autofocus area on an X-T20, use {3020,3020}
    // There are 506016 AF pixels on a Fujifilm X-T20.
    // long sample_size[] = {3020,3020};
    long sample_size[] = {2048,2048};

    // First pixel (x,y) for sample box.
    long fsample[2];

    // Strings to store camera name and comment.
    char camera[FLEN_VALUE], camera_comment[FLEN_COMMENT];

    if (argc != 3) { 
      printf("Usage: fix_xtrans image_in image_out \n");
      printf("\n");
      printf("Fix auto focus pixels in callibration frames.\n");
      printf("\n");
      printf("Example: \n");
      printf("  fix_xtrans image.fits new_image.fits\n");
      return(0);
    }

    if ( !fits_open_image(&infptr, argv[1], READONLY, &status) )
    {
      if (fits_get_hdu_type(infptr, &hdutype, &status) || hdutype != IMAGE_HDU) { 
        printf("Error: this program only works on images, not tables\n");
        return(1);
      }

      fits_read_keyword(infptr, "INSTRUME", camera, camera_comment, &status);

      fits_get_img_dim(infptr, &naxis, &status);
      fits_get_img_size(infptr, 2, naxes, &status);

      if (status || naxis != 2) { 
        printf("Error: NAXIS = %d.  Only 2-D images are supported.\n", naxis);
        return(1);
      }

      // Memory for 1 row.
      pix = (double *) malloc(naxes[0] * sizeof(double));

      if (pix == NULL) {
        printf("Memory allocation error\n");
        return(1);
      }

      // Sample size must be <= image size.
      if (sample_size[0] > naxes[0] || sample_size[1] > naxes[1]) {
        printf( "Sample size (%ldx%ld) is larger than image size (%ldx%ld).\n"
              , sample_size[0], sample_size[1], naxes[0], naxes[1] );
        return(1);
      }

      // Define the sample box starting pixel.
      fsample[0] = floor((naxes[0]-sample_size[0])/2)+1;
      fsample[1] = floor((naxes[1]-sample_size[1])/2)+1;

      // Loop through sample rows.
      for (fpixel[0] = 1, fpixel[1] = fsample[1]; fpixel[1] < fsample[1]+sample_size[1]; fpixel[1]++)
      {  

        // x=1, y=1 is the bottom left corner of the image.
        int x=fpixel[0], y=fpixel[1];

        // We read all pixels for row Y
        if (fits_read_pix(infptr, TDOUBLE, fpixel, naxes[0],0, pix,0, &status))
           break;   /* jump out of loop on error */

        // Loop through sample pixels in row.
        for (int ii = fsample[0]-1; ii < fsample[0]+sample_size[0]-1; ii++) {

          // Set x. (ii starts at 0, x starts at 1)
          x = ii+1;     

          // We are in the sample box.
          if ( is_af_pixel( camera, x, y ) ) { 
            // We are an AF pixel.
            afcount++;
            afsum += pix[ii];
            if (pix[ii] < afmin) afmin = pix[ii];
            if (pix[ii] > afmax) afmax = pix[ii];
          } else {
            // We are not an AF pixel.
            nfcount++;
            nfsum += pix[ii];               
            if (pix[ii] < nfmin) nfmin = pix[ii];
            if (pix[ii] > nfmax) nfmax = pix[ii];
          }
        }
      }

      // Make sure we have a valid sample.
      if ( nfcount == 0 || afcount == 0 ) {
        printf("Failed to sample enough pixels.\n");
        return(1);
      }

      // Compute averages and fudge amount.
      nfmean = nfsum / nfcount;
      afmean = afsum / afcount;
      fudge = afmean - nfmean;


      // Now we loop again only this time we will adjust the values of the AF pixels.
      if ( !fits_create_file(&outfptr, argv[2], &status) ) {

        // copy all the header keywords from first image to new output file
        fits_copy_header(infptr, outfptr, &status); 

        // Loop through all rows.
        for (fpixel[0]=1, fpixel[1]=1; fpixel[1] <= naxes[1]; fpixel[1]++)
        {

           // We read all pixels for row Y
           if (fits_read_pix(infptr, TDOUBLE, fpixel, naxes[0],0, pix,0, &status))
              break;   /* jump out of loop on error */

           // Loop through all pixels in row.
           for (int ix = 0; ix < naxes[0]; ix++) {

             if ( is_af_pixel( camera, ix+1, fpixel[1] ) ) {
               // This is an auto focus pixel.  Subtract the fudge.
               pix[ix] = pix[ix] - fudge;
             }

           }

           // Write the line of pixels.
           fits_write_pix(outfptr, TDOUBLE, fpixel, naxes[0], pix, &status);

        } 

      }

      // Close files.  Free pix memory.
      fits_close_file(infptr, &status);
      fits_close_file(outfptr, &status);
      free(pix);

    }

    if (status)  {
        fits_report_error(stderr, status);
    } else {
        printf("\n");
        printf("          Camera Model  =  %s\n",        camera);
        printf("           Sensor Size  =  %ld x %ld\n", naxes[0], naxes[1]);
        printf("           Sample Size  =  %ld x %ld\n", sample_size[0], sample_size[1]);
        printf("      Sample Start X:Y  =  %ld:%ld\n",   fsample[0], fsample[1]);          
        printf("        Sampled Pixels  =  %ld\n\n",     nfcount+afcount); 

        printf(" Non Auto Focus Pixels  =  %ld\n",       nfcount);
        printf("               Average  = % .10f\n",     nfmean);
        printf("               Minimum  = % .10f\n",     nfmin);
        printf("               Maximum  = % .10f\n\n",   nfmax);

        printf("     Auto Focus Pixels  =  %ld\n",       afcount);
        printf("               Average  = % .10f\n",     afmean);
        printf("               Minimum  = % .10f\n",     afmin);
        printf("               Maximum  = % .10f\n\n",   afmax);

        printf("          Fudge Amount  = % .10f\n\n",   fudge);
    }

    return(status);
}