/*
 * expand-map.c
 *
 * PIK/S. Petri/April 2013
 *
 * For GRASS r.watershed, the world is a flat disc (or rectangle).
 * For calculation of gradients and flow directions, it does not look
 * across the map boundaries, although in a world map the boundaries are periodical.
 * Thus, here we expand a given input map.
 * Append 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.
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <netcdf.h>

static char *myname = NULL;

static void
usage(char *myname, char *msg) {
  if (msg) fprintf(stderr, "%s\n", msg);
  fprintf(stderr, "usage: %s input-file input-var 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, maskvarid, londimid, latdimid, lonvarid, latvarid, ndims, *dimids = NULL;
  size_t nlons, nlats;
  char buf[NC_MAX_NAME];

  float *topo = NULL, topomissval = -999.0; /* NC_FILL_FLOAT; */
  float *mask = NULL, maskmissval = -999.0; /* NC_FILL_FLOAT; */

  double *lons = NULL, *lats = NULL;

  float *outtopo = NULL, *outmask = NULL;
  int outnlons, outnlats;

  int outncid, outvarid, outmaskvarid, outdimids[2], outlonvarid, outlatvarid;
  /* sigh. we have to make up a fake coordinate system for GRASS */
  double *outlons = NULL, *outlats = NULL;

  myname = *argv;

  if (argc != 5) usage(myname, "need 4 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_create(argv[4], NC_CLASSIC_MODEL, &outncid); handle(err);

  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");

  outnlats = 2*nlats;
  err = nc_def_dim(outncid, buf, outnlats, &(outdimids[0])); handle(err);
  err = nc_def_var(outncid, buf, NC_DOUBLE, 1, outdimids+0, &outlatvarid); handle(err);
  err = nc_put_att_text(outncid, outlatvarid, "units", strlen("degrees_north"), "degrees_north"); handle(err);
  err = nc_put_att_text(outncid, outlatvarid, "long_name", strlen("latitude"), "latitude"); handle(err);

  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");

  outnlons = 2*nlons;
  err = nc_def_dim(outncid, buf, outnlons, &(outdimids[1])); handle(err);
  err = nc_def_var(outncid, buf, NC_DOUBLE, 1, outdimids+1, &outlonvarid); handle(err);
  err = nc_put_att_text(outncid, outlonvarid, "units", strlen("degrees_east"), "degrees_east"); handle(err);
  err = nc_put_att_text(outncid, outlonvarid, "long_name", strlen("longitude"), "longitude"); handle(err);

  topo = (float *)malloc(nlats*nlons*sizeof(*topo));
  err = nc_get_var_float(ncid, varid, topo); handle(err);

  err = nc_inq_varid (ncid, argv[3], &maskvarid); handle(err);
  mask = (float *)malloc(nlats*nlons*sizeof(*mask));
  err = nc_get_var_float(ncid, maskvarid, mask); handle(err);

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

  /* everything looks OK so far, now lets start working */

  outtopo = (float *)malloc(outnlons*outnlats*sizeof(outtopo));
  memset(outtopo, '\0', outnlons*outnlats*sizeof(outtopo));
  outmask = (float *)malloc(outnlons*outnlats*sizeof(outmask));
  memset(outmask, '\0', outnlons*outnlats*sizeof(outmask));

  /* 1. put original map into center of output */
  for (j = 0; j < nlats; j++) {
    for (i = 0; i < nlons; i++) {
      outtopo[(j+nlats/2)*outnlons+i+nlons/2] = topo[j*nlons+i];
      outmask[(j+nlats/2)*outnlons+i+nlons/2] = mask[j*nlons+i];
    }
  }

  /* 2. copy west half of original map onto east eighth of output map */
  for (j = 0; j < nlats; j++) {
    for (i = 0; i < nlons/2; i++) {
      outtopo[(j+nlats/2)*outnlons+i+nlons+nlons/2] = topo[j*nlons+i];
      outmask[(j+nlats/2)*outnlons+i+nlons+nlons/2] = mask[j*nlons+i];
    }
  }

  /* 3. copy east half of original map onto west eighth of output map */
  for (j = 0; j < nlats; j++) {
    for (i = nlons/2; i < nlons; i++) {
      outtopo[(j+nlats/2)*outnlons+i-nlons/2] = topo[j*nlons+i];
      outmask[(j+nlats/2)*outnlons+i-nlons/2] = mask[j*nlons+i];
    }
  }

  /* 4. rotate and mirror south half of original map onto south eighth of output map */
  for (j = 0; j < nlats/2; j++) {
    for (i = 0; i < nlons; i++) {
      int outi = nlons/2 + (i+nlons/2)%nlons;
      int outj = nlats/2 - j - 1;
      outtopo[outj*outnlons+outi] = topo[j*nlons+i];
      outmask[outj*outnlons+outi] = mask[j*nlons+i];
    }
  }

  /* 5. rotate and mirror north half of original map onto north eighth of output map */
  for (j = nlats/2; j < nlats; j++) {
    for (i = 0; i < nlons; i++) {
      int outi = nlons/2 + (i+nlons/2)%nlons;
      int outj = outnlats - j + nlats/2 - 1;
      outtopo[outj*outnlons+outi] = topo[j*nlons+i];
      outmask[outj*outnlons+outi] = mask[j*nlons+i];
    }
  }
  /* 6. fill sw corner 16th with a copy of rotated and mirrored sw original map */
  for (j = 0; j < nlats/2; j++) {
    for (i = nlons; i < nlons + nlons/2; i++) {
      int outi = i - nlons;
      int outj = j;
      //printf("%d,%d to %d,%d\n", i,j,outi, outj);
      outtopo[outj*outnlons+outi] = outtopo[j*outnlons+i];
      outmask[outj*outnlons+outi] = outmask[j*outnlons+i];
    }
  }
  /* 7. fill se corner 16th with a copy of rotated and mirrored se original map */
  for (j = 0; j < nlats/2; j++) {
    for (i = nlons/2; i < nlons; i++) {
      int outi = i + nlons;
      int outj = j;
      //printf("%d,%d to %d,%d\n", i,j,outi, outj);
      outtopo[outj*outnlons+outi] = outtopo[j*outnlons+i];
      outmask[outj*outnlons+outi] = outmask[j*outnlons+i];
    }
  }
  /* 8. fill nw corner 16th with a copy of rotated and mirrored nw original map */
  for (j = nlats+nlats/2; j < outnlats; j++) {
    for (i = nlons; i < nlons + nlons/2; i++) {
      int outi = i - nlons;
      int outj = j;
      //printf("%d,%d to %d,%d\n", i,j,outi, outj);
      outtopo[outj*outnlons+outi] = outtopo[j*outnlons+i];
      outmask[outj*outnlons+outi] = outmask[j*outnlons+i];
    }
  }
  /* 9. fill ne corner 16th with a copy of rotated and mirrored ne original map */
  for (j = nlats+nlats/2; j < outnlats; j++) {
    for (i = nlons/2; i < nlons; i++) {
      int outi = i + nlons;
      int outj = j;
      //printf("%d,%d to %d,%d\n", i,j,outi, outj);
      outtopo[outj*outnlons+outi] = outtopo[j*outnlons+i];
      outmask[outj*outnlons+outi] = outmask[j*outnlons+i];
    }
  }
#if 0
#endif

  /* write to file */
  err = nc_def_var(outncid, argv[2], NC_FLOAT, 2, outdimids, &outvarid); handle(err);
  err = nc_put_att_text(outncid, outvarid, "units", strlen("m"), "m"); handle(err);
  err = nc_put_att_text(outncid, outvarid, "long_name", strlen("topography"), "topography"); handle(err);
  err = nc_put_att_float(outncid, outvarid, "_FillValue", NC_FLOAT, 1, &topomissval); handle(err);
  err = nc_put_att_float(outncid, outvarid, "missing_value", NC_FLOAT, 1, &topomissval); handle(err);

  err = nc_def_var(outncid, argv[3], NC_FLOAT, 2, outdimids, &outmaskvarid); handle(err);
  err = nc_put_att_text(outncid, outvarid, "units", strlen("1 = land, 0 = sea"), "1 = land, 0 = sea"); handle(err);
  err = nc_put_att_text(outncid, outvarid, "long_name", strlen("land-sea mask"), "land-sea mask"); handle(err);
  err = nc_put_att_float(outncid, outvarid, "_FillValue", NC_FLOAT, 1, &maskmissval); handle(err);

  err = nc_enddef(outncid); handle(err);

  err = nc_put_var_float(outncid, outvarid, outtopo); handle(err);
  err = nc_put_var_float(outncid, outmaskvarid, outmask); handle(err);

  outlons = (double *)malloc(outnlons*sizeof(double));
  outlons[0] = 0;
  for (i = 1; i < outnlons; i++) outlons[i] = outlons[0] + i*(360.0/outnlons);
  err = nc_put_var_double(outncid, outlonvarid, outlons); handle(err);

  outlats = (double *)malloc(outnlats*sizeof(double));
  outlats[0] = -90.0 + (180.0/outnlats/2.0);
  for (j = 1; j < outnlats; j++) outlats[j] = outlats[0] + j*(180.0/outnlats);
  err = nc_put_var_double(outncid, outlatvarid, outlats); handle(err);

  err = nc_close(outncid); handle(err);

  printf("done\n");

  return 0;
}