/*
 * watershed-to-DDM30.c
 *
 * PIK/S. Petri/July 2020
 *
 * Convert output from GRASS r.watershed into a format that 
 * is compatible with DDM30.asc and can be used by
 * LPJml/bin/drainage
 */

/* drainage
 * Output map: drainage direction. Provides the "aspect" for each cell
 * measured CCW from East. Multiplying positive values by 45 will give
 * the direction in degrees that the surface runoff will travel from that
 * cell. The value 0 (zero) indicates that the cell is a depression area
 * (defined by the depression input map). Negative values indicate that
 * surface runoff is leaving the boundaries of the current geographic
 * region. The absolute value of these negative cells indicates the
 * direction of flow.
 *
 *  That gives:
 *  0  1  2  3  4  5  6  7  8
 *  -  NE N  NW W  SW S  SE E 
 *
 * This now needs to be converted into something that looks like DDM30.asc
 *
 * The map is generated with a cellsize of 30 arc
 * minutes and is distributed as a zipped ASCII-grid that can be easily
 * imported in most GIS-software that support rasters or grids. The
 * direction of outflow from cell is according to the ArcInfo/ArcView
 * convention as follows: ``0'' is assigned to internal sinks and to
 * cells draining into the ocean while ``1'', ``2'', ``4'', ``8'', ``16'', ``32'',
 * ``64'' and ``128'' standing for the outflow directions E, SE, S, SW, W,
 * NW, N or NE.
 * -9 is missing_value, for ocean cells.
 * http://www.geo.uni-frankfurt.de/ipg/ag/dl/datensaetze/3_drainage_direction_map/index.html
 * http://www.uni-frankfurt.de/45217896/3_drainage_direction_map?
 * http://www.uni-frankfurt.de/45218101/DDM30
 *
 *
 * Read the GRASS-generated drainage map and the original land/sea mask.
 *
 *        -  NE  N  NW  W  SW  S  SE  E miss
 * GRASS  0   1  2   3  4   5  6   7  8
 * DDM30    128 64  32 16   8  4   2  1 -9
 *
 * Possible problems: along the map edges, r.watershed places negative
 * drainage direction values to indicate that the drainage leaves the
 * region. Apparently, that means that r.watershed does not ``wrap'' the
 * world around the longitude?
 * Yes. And, more important, not across the poles.
 * That hits us with a bone crushing sound in Antarctica.
 *
 * Above suspicion turns out to be true. Moreover, tracing drainage across
 * the map edges results in circular flow in several places, especially in antarctica,
 * but also across the east/west boundary in some places.
 *
 * Approach:
 * Extend the original map, by appending the western half to the eastern edge,
 * the eastern half to the western edge,
 * south half rotated and mirrored to southern edge,
 * northern half rotated and mirrored to northern edge.
 * Ignore lat / lon axis labels.
 * Then import it into GRASS, run r.watershed on it, export, and run
 * the inverse procedure, which just cuts out the original area.
 *
 * In that cut-out area, the edges should now contain only positive values,
 * because r.watershed did not perceive them as edges in its work.
 * Thus we must check for over-the-edge wrap-around here.
 * 
 */

/* If projection and region were set correctly, the generated NetCDF file has
 * 1440 columns longitudes 0 < .. < 360, and 720 rows.
 * If there are 721, then most probably the projection and/or region was not set
 * correctly, and the last column is either all-zeros or a duplicate of the first one.
 * In that case, it is just ignored here.
 * If there occurs any other number of rows, an error is raised.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <netcdf.h>

/*                         -  NE   N  NW   W  SW   S  SE   E */
//static char  *dname[] =  { "-","NE", "N","NW", "W","SW", "S","SE", "E"};
//static int grassval[] =  {  0,   1,    2,   3,   4,   5,   6,   7,  8 };
static int     ddmval[] =  {  0 ,128,   64,  32,  16,   8,   4,   2,  1 };

static void
usage(char *myname, char *msg) {
  if (msg) fprintf(stderr, "%s\n", msg);
  fprintf(stderr, "usage: %s drainage-file drainage-var mask-file mask-var output-file\n",
          myname);
  exit(1);
}

static void
handle(int err) {
  if (err != NC_NOERR) {
    fprintf(stderr, "%s\n", nc_strerror(err));
    exit(1);
  }
}


int main(int argc, char **argv) {

  int i, j, err;
  int ncid, varid, londimid, latdimid, lonvarid, latvarid, ndims, *dimids = NULL;
  size_t nlons, nlats;
  char buf[NC_MAX_NAME];

  int *drainage_direction_raw = NULL; /* expanded map as read from file */
  int *drainage_direction = NULL; /* cut-out center, the original map region, on same coordinates as mask */
  double *lons = NULL, *lats = NULL;

  int maskncid, maskvarid;
  size_t masknlons, masknlats;
  int *landmask = NULL;

  int samesize = 0;
  int is, ie, js, je;

  FILE *outfile;

  char *myname = *argv;

  if (argc != 6) usage(myname, "need 5 parameters");

  err = nc_open(argv[1], NC_NOWRITE, &ncid); handle(err);
  err = nc_inq_varid (ncid, argv[2], &varid); handle(err);
  err = nc_inq_varndims(ncid, varid, &ndims); handle(err);
  printf("%s: variable %s has %d dimensions\n", argv[1], argv[2], ndims);
  if (2 != ndims) {
    fprintf(stderr, "Error: %s: variable %s: expected exact 2 dimensions, but got %d\n",
            argv[1], argv[2], ndims);
    exit(1);
  }
  dimids = (int *)malloc(ndims*sizeof(*dimids));
  err = nc_inq_vardimid(ncid, varid, dimids); handle(err);
  latdimid = dimids[0];
  londimid = dimids[1];
  free(dimids);
  dimids = NULL;

  err = nc_inq_dim(ncid, latdimid, buf, &nlats); handle(err);
  printf("dimension 0 name %s len %lu\n", buf, nlats);
  err = nc_inq_varid(ncid, buf, &latvarid); handle(err);
  lats = (double *)malloc(nlats*sizeof(*lats));
  err = nc_get_var_double(ncid, latvarid, lats); handle(err);
  for (i = 0; i < nlats; i++) printf("%g ", lats[i]); printf("\n");

  err = nc_inq_dim(ncid, londimid, buf, &nlons); handle(err);
  printf("dimension 1 name %s len %lu\n", buf, nlons);
  err = nc_inq_varid(ncid, buf, &lonvarid); handle(err);
  lons = (double *)malloc(nlons*sizeof(*lons));
  err = nc_get_var_double(ncid, lonvarid, lons); handle(err);
  for (i = 0; i < nlons; i++) printf("%g ", lons[i]); printf("\n");

  drainage_direction_raw = (int *)malloc(nlats*nlons*sizeof(*drainage_direction_raw));
  err = nc_get_var_int(ncid, varid, drainage_direction_raw); handle(err);

  for (j = 0; j < nlats; j++) {
    for (i = 0; i < nlons; i++) {
      int d = drainage_direction_raw[j*nlons+i];
      if (d != (int)d || d < -8 || d > 8) {
        fprintf(stderr, "Error: drainage_direction must consist of integers -8 to 8 only, but has %d at position %d, %d\n",
                d, i, j);
        exit(1);
      }
    }
  }

  
  err = nc_open(argv[3], NC_NOWRITE, &maskncid); handle(err);
  err = nc_inq_varid (maskncid, argv[4], &maskvarid); handle(err);
  err = nc_inq_varndims(maskncid, maskvarid, &ndims); handle(err);
  printf("%s: variable %s has %d dimensions\n", argv[3], argv[4], ndims);
  if (2 != ndims) {
    fprintf(stderr, "Error: %s: variable %s: expected exact 2 dimensions, but got %d\n",
            argv[3], argv[4], ndims);
    exit(1);
  }
  dimids = (int *)malloc(ndims*sizeof(*dimids));
  err = nc_inq_vardimid(maskncid, maskvarid, dimids); handle(err);

  err = nc_inq_dim(maskncid, dimids[0], buf, &masknlats); handle(err);
  printf("dimension 0 name %s len %lu\n", buf, masknlats);
  err = nc_inq_dim(maskncid, dimids[1], buf, &masknlons); handle(err);
  printf("dimension 1 name %s len %lu\n", buf, masknlons);

  if (nlats == masknlats && nlons == masknlons) {
    /* both are expanded maps */
    samesize = 1;
  } else {
    if (nlats != 2*masknlats) {
      fprintf(stderr, "number of latitudes of drainage_direction and land mask do not match: %lu != 2*%lu\n",
              nlats, masknlats);
      exit(1);
    }

    if (nlons == 2*masknlons+1) {
      /* examine last column. should be all zeros */
      i = nlons-1;
      for (j = 0; j < nlats; j++) {
        if (drainage_direction_raw[j*nlons+i] != 0) {
          fprintf(stderr, "drainage_direction[%d, %d] = %d != 0\n", i, j, drainage_direction_raw[i*nlons+j]);
          exit(1);
        }
      }
    } else if (nlons != 2*masknlons) {
      fprintf(stderr, "number of longitudes of drainage_direction and land mask do not match: %lu != 2*%lu\n",
              nlons, masknlons);
      exit(1);
    }
  }

  landmask = (int *)malloc(masknlats*masknlons*sizeof(*landmask));
  err = nc_get_var_int(maskncid, maskvarid, landmask); handle(err);

  for (j = 0; j < masknlats; j++) {
    for (i = 0; i < masknlons; i++) {
      if (landmask[j*masknlons+i] != 0 && landmask[j*masknlons+i] != 1) {
        fprintf(stderr, "Error: land mask must consist of 0 and 1 values only, but has %d at position %d, %d\n",
                landmask[j*masknlons+i], i, j);
        exit(1);
      }
    }
  }

  /* so far, everything looks OK, now lets start the real work */

  if (samesize) {
    printf("drainage map and mask have same size\n");
    drainage_direction = drainage_direction_raw;
  } else {
    /* cut out original region from drainage_direction_raw */
    printf("cutting out original region from drainage map\n");
    drainage_direction = (int *)malloc(masknlats*masknlons*sizeof(*drainage_direction));
    for (j = 0; j < masknlats; j++) {
      for (i = 0; i < masknlons; i++) {
        drainage_direction[j*masknlons+i] = drainage_direction_raw[(j+masknlons/2)*nlons+i+masknlons/2];
      }
    }
    free(drainage_direction_raw);
  }
  drainage_direction_raw = NULL;
  nlons = -424242;
  nlats = -434343;

  /* scan edges of cut-out drainage-direction. off-map directions are most probably positive now */
  for (i = 0; i < masknlons; i++) {
    int d;
    d = drainage_direction[0*masknlons+i]; /* south edge */
    if (5 == d || 6 == d || 7 == d) { /* SW, S, SE */ drainage_direction[0*masknlons+i] = -d; }
    d = drainage_direction[(masknlats-1)*masknlons+i]; /* north edge */
    if (1 == d || 2 == d || 3 == d) { /* NE, N, NW */ drainage_direction[(masknlats-1)*masknlons+i] = -d; }
  }
  
  for (j = 0; j < masknlats; j++) {
    int d;
    d = drainage_direction[j*masknlons+0]; /* west edge */
    if (3 == d || 4 == d || 5 == d) { /* NW, W, SW */ drainage_direction[j*masknlons+0] = -d; }
    d = drainage_direction[j*masknlons+(masknlons-1)]; /* east edge */
    if (1 == d || 7 == d || 8 == d) { /* NE, SE, E */ drainage_direction[j*masknlons+(masknlons-1)] = -d; }
  }

  /* now lets report the problematic cells */
  for (j = 0; j < masknlats; j++) {
    for (i = 0; i < masknlons; i++) {
      int d = drainage_direction[j*masknlons+i];
      if (d < 0 && landmask[j*masknlons+i] > 0) {
        printf("drainage_direction off-map %d at position %d, %d\n", d, i, j);
      }
    }
  }

  /* produce the output map */
  outfile = fopen(argv[5], "w");
  if (NULL == outfile) perror(argv[5]);

  if (samesize) {
    /* now we need to cut out the map center to get back the original map */
    /* n_lon, n_lat, western boundary in degrees */
    printf("cut out at samesize\n");
    is = masknlons/4; ie = masknlons*3/4;
    js = masknlats/4; je = masknlats*3/4;
  } else {
    is = 0; ie = masknlons;
    js = 0; je = masknlats;
  }
  printf("is %d ie %d js %d je %d\n", is, ie, js, je);
  //ncols         720
  //nrows         360
  //xllcorner     -180
  //yllcorner     -90
  //cellsize      0.5
  //NODATA_value  -9
  fprintf(outfile, "ncols      %6d\n", is-ie);
  fprintf(outfile, "nrows      %6d\n", js-je);
  fprintf(outfile, "xllcorner  %6d\n", -180);
  fprintf(outfile, "yllcorner  %6d\n", -90);
  fprintf(outfile, "cellsize   0.5\n");
  fprintf(outfile, "NODATA_value  -9\n");
  for (j = je-1; j >= js; j--) { // reverse N-S axis
    for (i = is; i < ie; i++) {
      if (0 == landmask[j*masknlons+i]) {
        fprintf(outfile, " -9");
      } else {
        int d = drainage_direction[j*masknlons+i];
        fprintf(outfile, " %d", ddmval[abs(d)]);
      }
    }
    fprintf(outfile, "\n");
  }

  printf("done\n");
  return 0;
}