/* See {dsm_segment2.h}  */
/* Last edited on 2016-04-09 20:09:38 by stolfilocal */

#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <assert.h>

#include <bool.h>

#include <dsm_image.h>
#include <dsm_segment2.h>

/* Background/vehicle segmentation. */
/* Version of Jan/2016 */

/* INTERNAL PROTOTYPES */
  
/* Computes the discrepancy between the frame pixel value {imp} and
  the current background image pixel value {bgp} (both high-pass
  filtered, with encoding {[-1 _ +1] --> {0 .. 255}}).  The background
  is adjusted by the current estimate of camera gain {mu} and the
  current shadow mask pixel {shp} (assumed to use the encoding {[0 _
  2] --> {0 .. 255}}).  The result is in the range {[-2 _ +2]}.  */  
double dsm_segment2_im_bg_discrepancy(uint8_t imp, double mu, uint8_t bgp, uint8_t shp);

/* Estimates the probability {cw} of a certain frame pixel being
  background, given the discrepancy {d} as computed by 
  {dsm_segment2_im_bg_discrepancy}. */  
double dsm_segment2_estimate_pixel_wt(double d, double lambda);

/* Adjusts the overall camera gain so that unshadowed parts of the backgrond 
  have shadow factor 1. */
void dsm_segment2_reestimate_mu
  ( dsm_image_t *wt, 
    dsm_image_t *sh, 
    double *muP, 
    int32_t ix[]      /* (WORK) Index table. */
  );

/* Computes the mask {up} where it seems safe to update the background: namely, where
  it was reasonably safe before, and the image looks like unshadowed background. */  
void dsm_segment2_where_to_update(dsm_image_t *wt, dsm_image_t *sh, dsm_image_t *up);

/* Updates the background {bg} from the image {img} where {up} is large, assuming 
  the probability of the pixel being background is {wt}. */
void dsm_update_bg
  ( dsm_image_t *img,  
    dsm_image_t *up,  
    dsm_image_t *wt,  
    dsm_image_t *sh,
    double mu, 
    dsm_image_t *bg );


/* Multiplies the reference background image {bg} by the estimated camera gain {mu}
  and the guessed shadow mask {sh} to obtain an adjusted background image {adj}. */
void dsm_segment2_adjust_background(dsm_image_t *bg, double mu, dsm_image_t *sh, dsm_image_t *adj);

/* Pixel comparison function for sorting shadow images. */
int dsm_segment2_cmp_sh_pixels(const void *pa, const void *pb, void *data);

/* Computes the bachground probability mask {wt} from the frame image {im_high} and the
  current reference background image {bg_high} (both high-pass
  filtered with encoding {[-1 _ +1] --> {0 .. 255}}.  The background 
  must have been adjusted by the current estimate of camera gain {mu} and the
  current shadow mask {sh} before the high-pass filtering.  The 
  weight {t} will use the encoding {[0 _ 1] --> {0..255}}.  

  The parameter {lambda} is the assumed deviation of the image-background
  discrepancy (see {dsm_segment2_im_bg_discrepancy}).

  Also stores the image-background discrepancy map in {disc}
  (with the encoding {[-3*lambda _ +3*lambda] --> {0 .. 255}}. */ 
void dsm_segment2_estimate_weight
  ( dsm_image_t *im_high,  /* High-pass-filtered frame image. */
    dsm_image_t *bg_high,  /* High-pass-filtered adjusted background image. */
    double lambda,
    dsm_image_t *disc,     /* (OUT) Discrepancy image. */
    dsm_image_t *wt        /* (OUT) Background probability mask. */
  );

/* Recomputes the shadow factor mask {sh} from the frame image {img},
   the current reference background image {bg}, and the current
   estimate of camera gain {mu}.  The input images are assumed to
   use the encoding {[0_1] --> {0..255}}.  The shadow masks will use
   the encoding {[0 _ 2] --> {0..255}}.  The parameter {rshblur} is the
   kernel window radius to be used when blurring the raw shadow mask
   {sh_sharp} to obtain {sh}. */
void dsm_segment2_estimate_shadow
  ( dsm_image_t *img,      /* Frame image. */
    dsm_image_t *bg,       /* Reference background image. */
    double mu,             /* Estimaged camera gain. */
    dsm_image_t *sh_sharp, /* (OUT) Shadow mask before blurring. */
    int rshblur,           /* Blur radius for shadow mask. */
    dsm_image_t *sh        /* (OUT) Shadow mask (blurred). */
  );

/* Compare procedure for sorting shadow mask pixels.  The {data} must be
  a pointer to the shadow image.  The arguments {pa,pb} must be pointers
  to {int32_t} vaiables that contain the linearized indices of the 
  shadow mask pixels to be compered. */
int cmp_sh_pixels(const void *pa, const void *pb, void *data);

/* IMPLEMENTATIONS */

void dsm_segment2_state_initialize(dsm_segment2_state_t *st, int nx, int ny, dsm_image_t *bg0)
  {
    st->nx = nx;
    st->ny = ny;

    st->bg = dsm_image_alloc(nx, ny, 0);
    dsm_image_copy(bg0, &(st->bg));

    st->mu = 1.0;                           /* Camera gain. */   
    st->sh = dsm_image_alloc(nx, ny, 127);  /* Last shadow factor times 127. */ 
   
    st->wt = dsm_image_alloc(nx, ny, 0);    /* Last background probability mask. */ 
     
    /* Classifier paramters: */
    st->lambda = 0.010;  /* Standard deviation of {im-bg} discrepancy. */
    st->rhigh = 9;       /* Blur radius for high-pass filtering. */
    st->rshblur = 3;     /* Blur radius for shadow mask. */
    st->rwterode1 = 25;  /* Structuring element radius to close holes in vehicles. */
    st->rwtdilate = 10;  /* Structuring element radius to clean small noise spots. */
    st->rwterode2 = 15;  /* Structuring element radius to expand vehicle spot. */

    /* Work images: */
    st->im_high            = dsm_image_alloc(nx, ny, 0);  /* High-pass filtered version of frame image. */
    st->bg_adj             = dsm_image_alloc(nx, ny, 0);  /* Bachground image adjusted for mu and shadow factor. */
    st->bg_high            = dsm_image_alloc(nx, ny, 0);  /* High-pass filtered version of background image. */
    st->up                 = dsm_image_alloc(nx, ny, 0);  /* Bachground update mask. */
    st->sh_sharp           = dsm_image_alloc(nx, ny, 0);  /* Shadow mask before blurring. */
    st->disc               = dsm_image_alloc(nx, ny, 0);  /* Image-background discrepancy image. */
    st->temp               = dsm_image_alloc(nx, ny, 0);  /* Generic work image. */

    /* Previous masks: */
    st->wt_prev = dsm_image_alloc(nx, ny,   0);      /* Background indicator image of previous frame. */
    st->up_prev = dsm_image_alloc(nx, ny,   0);      /* Background update mask of previous frame. */

    st->i32temp = (int32_t *)malloc(nx*ny*sizeof(int32_t));
    assert(st->i32temp != NULL);
  }

void dsm_segment2_identify_background
  ( dsm_image_t *img, 
    dsm_segment2_state_t *st,
    const char *out_prefix  /* Filename prefix for debugging images, or {NULL}. */
  )
  {
    bool_t debug = TRUE;

    if (debug) { fprintf(stderr, "----------------------------------------------------------------------\n"); }
    
    int nx = img->nx;
    int ny = img->ny;
    assert((nx == st->nx) && (ny == st->ny));
 
    /* Compute highpass version of {img}: */
    if (debug) { fprintf(stderr, "highpass-filtering image\n"); }
    dsm_image_log_high_pass (img, st->rhigh, st->rhigh, &(st->im_high), &(st->temp));

   if (out_prefix != NULL) {
      dsm_image_debug(&(st->im_high),  out_prefix, -1, -1,   "0_im_high");
    }

   if (out_prefix != NULL) {
     dsm_image_debug(&(st->bg),                out_prefix, -1, -1, "1_bg_in");
   }

   int maxiter = 2;
    for (int iter = 0; iter < maxiter; iter++) {
  
      if (debug) { fprintf(stderr, "start iteration %04d\n", iter); }

      if (out_prefix != NULL) {
        dsm_image_debug(&(st->sh),                out_prefix, -1, iter, "2_sh_ini");
      }
  
     /* Compute highpass version of current background image: */
      if (debug) { fprintf(stderr, "highpass-filtering reference bg (iter = %d)\n", iter); }
      dsm_segment2_adjust_background(&(st->bg), st->mu, &st->sh, &(st->bg_adj));
      dsm_image_log_high_pass (&(st->bg_adj), st->rhigh, st->rhigh, &(st->bg_high), &(st->temp));
      
      /* Estimate {wt} and {sh} from {im_high} and {bg_high}, {mu} and current {sh}: */
      if (debug) { fprintf(stderr, "estimating {wt} and {sh}\n"); }
      dsm_segment2_estimate_weight
        ( &(st->im_high), 
          &(st->bg_high),
          st->lambda,
          &(st->disc),  
          &(st->wt) );

      /* Estimate shadow mask and blur it a bit: */
      dsm_segment2_estimate_shadow(img, &(st->bg), st->mu, &(st->sh_sharp), st->rshblur, &(st->sh));

      if (out_prefix != NULL) {
        dsm_image_debug(&(st->wt),                out_prefix, -1, iter, "5_wt_a");
      }
  
      /* Close holes in vehicle spots: */
      dsm_image_erode(&(st->wt), st->rwterode1, &(st->temp));

      /* Remove small isolated noise spots: */
      dsm_image_dilate(&(st->wt), st->rwterode1+st->rwtdilate, &(st->temp));

      /* Expand vehicle spots for safety: */
      dsm_image_erode(&(st->wt), st->rwterode2, &(st->temp));

      /* Re-estimate {mu} and adjust {sh}: */
      if (debug) { fprintf(stderr, "reestimating {mu}\n"); }
      dsm_segment2_reestimate_mu(&(st->wt), &(st->sh), &(st->mu), st->i32temp);
      
      if (out_prefix != NULL) {
        dsm_image_debug(&(st->bg_adj),            out_prefix, -1, iter, "3_a_bg_adj");
        dsm_image_debug(&(st->bg_high),           out_prefix, -1, iter, "3_b_bg_high");
        dsm_image_debug(&(st->disc),              out_prefix, -1, iter, "4_disc");
        dsm_image_debug(&(st->wt),                out_prefix, -1, iter, "5_wt_b");
        dsm_image_debug(&(st->sh_sharp),          out_prefix, -1, iter, "6_sh_sharp");
        dsm_image_debug(&(st->sh),                out_prefix, -1, iter, "7_sh");
      }

      /* !!! Deveria testar convergencia aqui. !!! */
    }
  
    /* Updating the background where safe: */
    if (debug) { fprintf(stderr, "updating the background\n"); }
    dsm_segment2_where_to_update(&(st->wt), &(st->sh), &(st->up));
    dsm_update_bg(img, &(st->up), &(st->wt), &(st->sh), st->mu, &(st->bg));

    if (out_prefix != NULL) {
      dsm_image_debug(&(st->up),                out_prefix, -1, 9999, "8_up");
      dsm_image_debug(&(st->bg),                out_prefix, -1, 9999, "9_bg");
    }

    if (debug) { fprintf(stderr, "----------------------------------------------------------------------\n"); }
  }

void dsm_segment2_estimate_weight
  ( dsm_image_t *im_high,
    dsm_image_t *bg_high,
    double lambda,
    dsm_image_t *disc,     /* Discrepancy image. */
    dsm_image_t *wt        /* (OUT) Background probability mask. */
  )
  {
    int nx = im_high->nx;
    int ny = im_high->ny;
    assert((nx == bg_high->nx) && (ny == bg_high->ny));
    assert((nx == wt->nx) && (ny == wt->ny));

    /* Data for {lambda} checking: */
    double sumw = 0;
    double sumwd = 0.0;
    double sumwd2 = 0.0;

    for (int y = 0; y < ny; y++) {
      for (int x = 0; x < nx; x++) {
	/* Get relevant data for pixel {x,y}: */
	int id = y * nx + x;
	double imp = dsm_image_floatize(im_high->p[id], -1.0, +1.0);  /* Highpass image pixel in {[-1 _ +1]}. */
	double bgp = dsm_image_floatize(bg_high->p[id], -1.0, +1.0);  /* Highpass background pixel in {[-1 _ +1]}. */
	double d = imp - bgp; /* Discrepancy in the range {[-2 _ +2]}. */
	double cw = dsm_segment2_estimate_pixel_wt(d, lambda);
	wt->p[id] = dsm_image_quantize(cw, 0.0, 1.0); 
        
	/* Save discrepancy value: */
	disc->p[id] = dsm_image_quantize(d/(3*lambda), -1.0, +1.0); 
  
	/* Accumulate discrepancy statistics */
        sumw += cw;
        sumwd += cw*d;
        sumwd2 += cw*d*d;
      }
    }

    if (sumw > 0) {
      /* Compute the apparent {lambda} (deviation of discrepancy): */
      double davg = sumwd/sumw;
      double dvar = sumwd2/sumw; /* Assuming that average discrepancy is zero. */
      fprintf(stderr, "discrepancy average = %.3f deviation = %.3f\n", davg, sqrt(dvar));
    }
  }

  void dsm_segment2_estimate_shadow
  ( dsm_image_t *img,      /* Frame image. */
    dsm_image_t *bg,       /* Reference background image. */
    double mu,          /* Estimaged camera gain. */
    dsm_image_t *sh_sharp, /* (OUT) Shadow mask before blurring. */
    int rshblur,           /* Blur radius for shadow mask. */
    dsm_image_t *sh        /* (OUT) Shadow mask (blurred). */
  )
  {
    int nx = img->nx;
    int ny = img->ny;
    assert((nx == bg->nx) && (ny == bg->ny));
    assert((nx == sh_sharp->nx) && (ny == sh_sharp->ny));
    assert((nx == sh->nx) && (ny == sh->ny));

    for (int y = 0; y < ny; y++) {
      for (int x = 0; x < nx; x++) {
	/* Get relevant data for pixel {x,y}: */
	int id = y * nx + x;
	double imp =  img->p[id];       /* Image pixel in {[0 _ 255]}. */
	double bgp =  bg->p[id];        /* Background pixel in {[0 _ 255]}. */
        double shfac = imp/(mu*bgp); /* Shadow factor, typically in {[0_1]}, possibly {> 1}. */
	sh_sharp->p[id] = dsm_image_quantize(shfac, 0.0, 2.0);
      }
    }
    /* Blur the shadow mask: */ 
    /* !!! Should use a weighted blur !!! */
    dsm_image_blur (sh_sharp, rshblur, rshblur, sh);
  }

double dsm_segment2_im_bg_discrepancy(uint8_t imp, double mu, uint8_t bgp, uint8_t shp)
  {
    /* Pixel values in {[-1 _ +1]} scale: */
    double im_val = dsm_image_floatize(imp, 0.0, 1.0); /* Image pixel value. */
    double sh_val = dsm_image_floatize(shp, 0.0, 2.0); /* Shadow factor. */
    double bg_raw = dsm_image_floatize(bgp, 0.0, 1.0); /* Refrence background value. */
    double bg_adj = mu * bg_raw * sh_val; /* Corrected bg pixel value. */
    return im_val - bg_adj;
  }

double dsm_segment2_estimate_pixel_wt(double d, double lambda)
  {
    /* Pixel values in [0_1] scale: */
    double z = d/lambda;
    double Pbg = exp(-z*z/2)/lambda/sqrt(2*M_PI); /* Prob. dens. of {imp}, if background. */
    double Pve = 1/255.0;  /* Prob. dens. of {imp}, if vehicle. */
    /* Bayes's formula - a priori 50% prob of vehicle, 50% background: */
    double P = Pbg/(Pbg + Pve);
    return P;
  }

void dsm_segment2_reestimate_mu
  ( dsm_image_t *wt, 
    dsm_image_t *sh, 
    double *muP, 
    int32_t ix[]      /* (WORK) Index table. */
  )
  { 
    int nx = wt->nx;
    int ny = wt->ny;
    assert((nx == sh->nx) && (ny == sh->ny));

    int n = nx*ny; /* Total number of pixels. */

    /* Collect the indices of the pixels in {ix[0..n-1]: */
    for (int y = 0; y < ny; y++) {
      for (int x = 0; x < nx; x++) {
	/* Get relevant data for pixel {x,y}: */
	int id = y * nx + x;
        ix[id] = id;
      }
    }

    /* Sort {id[0..n-1]} by increasing shadow value: */
    qsort_r((void *)ix, n, sizeof(int32_t), &dsm_segment2_cmp_sh_pixels, (void*)sh);
    
    /* Compute total pixel weight {sumw}: */
    double sumw = 0;
    for (int ixw = 0; ixw < n; ixw++) {
      sumw += dsm_image_floatize(wt->p[ixw], 0.0, 1.0);
    }

    /* Compute 90% upper percentile of {sh} values weighted by {wt}: */
    double shperc = +HUGE;  /* The shadow factor greater than 90% of all factors. */
    double sumt = 0;
    for (int k = 0; (k < n) && (sumt < 0.10*sumw); k++)
      { int ixk = ix[n-1-k]; /* Index of pixel with rank {k} by shadow factor. */ 
        double shk = dsm_image_floatize(sh->p[ixk], 0.0, 2.0); /* Its shadow factor. */
        double wtk = dsm_image_floatize(wt->p[ixk], 0.0, 1.0); /* Its weight. */
        shperc = shk;
        sumt += wtk;
      }

    /* Adjust {mu} and shadow mask so that {shperc} becomes 1: */
    for (int y = 0; y < ny; y++) {
      for (int x = 0; x < nx; x++) {
	/* Get relevant data for pixel {x,y}: */
	int id = y * nx + x;
        double shfac_old = dsm_image_floatize(sh->p[id], 0.0, 2.0);
	double shfac_new = shfac_old / shperc;
        sh->p[id] = dsm_image_quantize(shfac_new, 0.0, 2.0);
      }
    }
    (*muP) *= shperc;
  }

void dsm_segment2_adjust_background(dsm_image_t *bg, double mu, dsm_image_t *sh, dsm_image_t *adj)
  { 
    int nx = bg->nx;
    int ny = bg->ny;
    assert((nx == sh->nx) && (ny == sh->ny));
    assert((nx == adj->nx) && (ny == adj->ny));

    for (int y = 0; y < ny; y++) {
      for (int x = 0; x < nx; x++) {
	int id = y * nx + x;
	double bgp = dsm_image_floatize(bg->p[id], 0.0, 1.0);  /* Reference pixel value in {[0 _ 1]}. */
	double shfac = dsm_image_floatize(sh->p[id], 0.0, 2.0); /* Shadow factor in {[0 _ 2]}. */
	double adjp = mu*bgp*shfac; /* Adjusted bg pixel in approx {[0 _ 1]}. */
        adj->p[id] = dsm_image_quantize(adjp, 0.0, 1.0);
      }
    }
  }

int dsm_segment2_cmp_sh_pixels(const void *pa, const void *pb, void *data)
  { 
    int32_t *ixa = (int32_t *)pa;
    int32_t *ixb = (int32_t *)pb;
    dsm_image_t * sh = (dsm_image_t *)data;
    
    if (sh->p[*ixa] < sh->p[*ixb]) 
      { return -1; }
    else if (sh->p[*ixa] > sh->p[*ixb]) 
      { return +1; }
    else 
      { return 0; }
  }
    
void dsm_segment2_where_to_update(dsm_image_t *wt, dsm_image_t *sh, dsm_image_t *up)
  { 
    /* !!! Should dilate stuff. !!! */
    int nx = wt->nx;
    int ny = wt->ny;
    assert((nx == sh->nx) && (ny == sh->ny));
    assert((nx == up->nx) && (ny == up->ny));

    for (int y = 0; y < ny; y++) {
      for (int x = 0; x < nx; x++) {
	int id = y * nx + x;
        double wtxy = dsm_image_floatize(wt->p[id], 0.0, 1.0); /* Bachground prob., [0_1]. */
	double shfac = dsm_image_floatize(sh->p[id], 0.0, 2.0); /* Shadowing factor. */
        double suxy = 1.0 - fabs(1.0 - shfac); /* Un-shadowed indicator [0_1]. */
        double upxy = suxy*suxy * wtxy*wtxy; /* Update factor {[0_1]} !!! Improve !!! */
        up->p[id] = dsm_image_quantize(upxy, 0.0, 1.0);
      }
    }
    
  }

void dsm_update_bg
  ( dsm_image_t *img,  
    dsm_image_t *up,  
    dsm_image_t *wt,  
    dsm_image_t *sh,
    double mu, 
    dsm_image_t *bg )
  { 
    int nx = img->nx;
    int ny = img->ny;
    assert((nx == sh->nx) && (ny == sh->ny));
    assert((nx == up->nx) && (ny == up->ny));
    assert((nx == wt->nx) && (ny == wt->ny));
    assert((nx == bg->nx) && (ny == bg->ny));

    for (int y = 0; y < ny; y++) {
      for (int x = 0; x < nx; x++) {
	int id = y * nx + x;
        double imxy = dsm_image_floatize(img->p[id], 0.0, 1.0);  /* Image pixel, [0 _ 1]. */
        double bgxy_old = dsm_image_floatize(bg->p[id], 0.0, 1.0);  /* Current background pixel, [0 _ 1]. */
        double upxy = dsm_image_floatize(up->p[id], 0.0, 1.0); /* Update mask, [0 _ 1]. */
        double wtxy = dsm_image_floatize(wt->p[id], 0.0, 1.0);  /* Bachground prob. [0 _ 1]. */
        double frac = 0.25*upxy*wtxy;  /* Background update fraction. */
        double shxy = dsm_image_floatize(sh->p[id], 0.0, 2.0); /* Shadow factor [0 _ 1]. */
        double bgxy_new = (1 - frac)*bgxy_old + frac*imxy/mu/shxy;
        bg->p[id] = dsm_image_quantize(bgxy_new, 0.0, 1.0);
      }
    }
  }