#include <sli/fitscc.h>
#include <sli/mdarray_statistics.h>
#include <stdlib.h>
#include "mask.h"

using namespace sli;


static void repair_saturated(fits_image &image_hdu, const fits_image &mask_hdu) {
    const int extend = 1;

    mdarray_float &data = image_hdu.float_array();
    const mdarray_uchar &mask = mask_hdu.uchar_array();
    const int w = data.length(0),
              h = data.length(1);

    for (int y = 0;  y < h;  y++) for (int x = 0;  x < w;  x++) {
        if (mask(x, y) & SATURATED) {
            int start = x, end;
            for (end = start;  end < w && mask(end, y) & SATURATED;)
                end++;
            //       SATURATED
            // +-----xxxxxxxx---------+
            // |     |       |        |
            // 0   start    end     w - 1
            end--;
            start -= extend;
            end   += extend;
            double val_left  = data(start - 1, y),
                   val_right = data(end + 1, y);
            for (x = start; x <= end;  x++) {
                double t = (double)(x - start) / (end - start);
                data(x, y) = val_right * t + val_left * (1. - t);
            }
            x = end;
        }
    }
}


int main(int argc, char *argv[]) {
    if (argc != 3) {
        sli__eprintf("usage: %s INPUT OUTPUT\n", argv[0]);
        exit(1);
    }

    fitscc fits;
    fits.read_stream(argv[1]);
    repair_saturated(fits.image(0L), fits.image(1L));
    fits.write_stream(argv[2]);

    return 0;
}