/*

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!                                                                   !!
!!                   GNU General Public License                      !!
!!                                                                   !!
!! This file is part of the Flexible Modeling System (FMS).          !!
!!                                                                   !!
!! FMS is free software; you can redistribute it and/or modify       !!
!! it and are expected to follow the terms of the GNU General Public !!
!! License as published by the Free Software Foundation.             !!
!!                                                                   !!
!! FMS is distributed in the hope that it will be useful,            !!
!! but WITHOUT ANY WARRANTY; without even the implied warranty of    !!
!! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the     !!
!! GNU General Public License for more details.                      !!
!!                                                                   !!
!! You should have received a copy of the GNU General Public License !!
!! along with FMS; if not, write to:                                 !!
!!          Free Software Foundation, Inc.                           !!
!!          59 Temple Place, Suite 330                               !!
!!          Boston, MA  02111-1307  USA                              !!
!! or see:                                                           !!
!!          http://www.gnu.org/licenses/gpl.txt                      !!
!!                                                                   !!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

*/

/*******************************************************************************
  contact: Zhi Liang (Zhi.Liang@noaa.gov)
*******************************************************************************/

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <getopt.h>
#include "constant.h"
#include "mpp.h"
#include "mpp_io.h"
#include "read_mosaic.h"
#include "tool_util.h"

char *usage[] = {
  "",
  "  edit_landXocean --mosaic mosaic_grid                                                ",
  "                                                                                      ",
  "edit_landXocean will adjust the landXocean exchange grid cell area to put the flux    ",
  "onto ocean points that are adjacent to land points.                                   ",
  "                                                                                      ",
  "edit_landXocean takes the following flags:                                            ",
  "                                                                                      ",
  "REQUIRED:                                                                             ",
  "                                                                                      ",
  "--mosaic_grid mosaic_grid   specify the mosaic grid file to be edited. This mosaic    ",
  "                            file should be a coupler mosaic file, which contains link ",
  "                            to ocean topog file and the exchange grid file.           ",
  "                                                                                      ",
  NULL
};  

int main(int argc, char* argv[])
{
  unsigned int opcode = 0;
  int          option_index, c, n;
  char         *mosaic_file     = NULL;
  char         **lXo_file = NULL;
  char         **lXo_base_file = NULL;
  char         history[1024];
  char         mosaic_dir[STRING], ocn_mosaic[STRING], lnd_mosaic[STRING];
  int          **ocn_mask;
  int          *nxgrid;
  double       **area_lXo;
  int          **il_lXo, **jl_lXo, **io_lXo, **jo_lXo;
    
  
  int          *nx_ocn, *ny_ocn, *nx_lnd, *ny_lnd;
  int          nlXo, ntiles_ocn, ntiles_lnd;
  int          x_cyclic, y_cyclic, folded_north;
  
  int errflg = (argc == 1);
  
  static struct option long_options[] = {
    {"mosaic_grid",       required_argument, NULL, 'a'},
    {0, 0, 0, 0},
  };

  mpp_init(&argc, &argv);

  /* This tool will be run on one processor */
  if(mpp_npes() > 1) mpp_error("edit_landXocean: parallel is not supported yet, try running one single processor");
  while ((c = getopt_long(argc, argv, "", long_options, &option_index)) != -1) {
    switch (c) {
    case 'a':
      mosaic_file  = optarg;
      break;
    case '?':
      errflg++;
      break;
    }
  }
  
  if (errflg) {
    char **u = usage;
    while (*u) { fprintf(stderr, "%s\n", *u); u++; }
    exit(2);
  }
  /* check the arguments */
  if( !mosaic_file    ) mpp_error("fregrid: mosaic_grid is not specified");

  strcpy(history,argv[0]);
  for(n=1;n<argc;n++)  {
    strcat(history, " ");
    strcat(history, argv[n]); 
  }

  /* read the mosaic file to get the landxocean exchange grid file name */
  {
    char tmp_file[STRING];
    char topog_file[STRING];
    size_t start[4], nread[4];
    int fid, vid;
    
    get_file_path(mosaic_file, mosaic_dir);
    fid = mpp_open(mosaic_file, MPP_READ);
    nlXo = mpp_get_dimlen( fid, "nfile_lXo");
    vid = mpp_get_varid(fid, "lXo_file");
    lXo_file = (char **)malloc(nlXo*sizeof(char *));
    lXo_base_file = (char **)malloc(nlXo*sizeof(char *));
    for(n=0; n<4; n++) {start[n]=0; nread[n]=1;}
    nread[1] = STRING;
    for(n=0; n<nlXo; n++) {
      start[0] = n;
      lXo_file[n] = (char *)malloc(STRING*sizeof(char));
      lXo_base_file[n] = (char *)malloc(STRING*sizeof(char));
      mpp_get_var_value_block(fid, vid, start, nread, lXo_base_file[n]);
      sprintf(lXo_file[n],  "%s/%s", mosaic_dir, lXo_base_file[n]);
    }
    
    /* read the ocean mosaic and land mosaic */
    vid = mpp_get_varid(fid, "ocn_mosaic_file");
    mpp_get_var_value(fid, vid, tmp_file);
    sprintf(ocn_mosaic, "%s/%s", mosaic_dir, tmp_file);
    vid = mpp_get_varid(fid, "lnd_mosaic_file");
    mpp_get_var_value(fid, vid, tmp_file);
    sprintf(lnd_mosaic, "%s/%s", mosaic_dir, tmp_file);    
    /* read the ocean topog file name */
    vid = mpp_get_varid(fid, "ocn_topog_file");
    mpp_get_var_value(fid, vid, tmp_file);
    sprintf(topog_file,  "%s/%s", mosaic_dir, tmp_file);
    mpp_close(fid);

  }
	 
  /*  get ntiles for ocean mosaic and land mosaic. currently we are assuming
      ntiles_ocean = 1 and ntiles_land = nlXo
      Also read the grid size for ocn_mosaic and land_mosaic.
  */
  ntiles_ocn = read_mosaic_ntiles(ocn_mosaic);
  ntiles_lnd = read_mosaic_ntiles(lnd_mosaic); 
  nx_ocn = (int *)malloc(ntiles_ocn*sizeof(int));
  ny_ocn = (int *)malloc(ntiles_ocn*sizeof(int));
  read_mosaic_grid_sizes(ocn_mosaic, nx_ocn, ny_ocn);
  nx_lnd = (int *)malloc(ntiles_lnd*sizeof(int));
  ny_lnd = (int *)malloc(ntiles_lnd*sizeof(int));
  read_mosaic_grid_sizes(lnd_mosaic, nx_lnd, ny_lnd);

  {
    char   **ocn_grid, **lnd_grid;
    char   tmp_file[STRING];
    size_t start[4], nread[4];
    int    fid, vid;
    int    ncontact;
    int   *tile1, *istart1, *iend1, *jstart1, *jend1;
    int   *tile2, *istart2, *iend2, *jstart2, *jend2;

    /* get the boundary condition for ocean mosaic. Since we are assuming there is only one
     tile in ocean mosaic, we could assume the boundary will be cyclic in x-direction,
     cyclic or folded-north for y-direction
    */
    x_cyclic     = 0;
    y_cyclic     = 0;
    folded_north = 0;    
    ncontact = read_mosaic_ncontacts(ocn_mosaic);
    tile1   = (int *)malloc(ncontact*sizeof(int));
    tile2   = (int *)malloc(ncontact*sizeof(int));
    istart1 = (int *)malloc(ncontact*sizeof(int));
    iend1   = (int *)malloc(ncontact*sizeof(int));
    jstart1 = (int *)malloc(ncontact*sizeof(int));
    jend1   = (int *)malloc(ncontact*sizeof(int));
    istart2 = (int *)malloc(ncontact*sizeof(int));
    iend2   = (int *)malloc(ncontact*sizeof(int));
    jstart2 = (int *)malloc(ncontact*sizeof(int));
    jend2   = (int *)malloc(ncontact*sizeof(int));
    if(ncontact >0) {
      read_mosaic_contact(ocn_mosaic, tile1, tile2, istart1, iend1, jstart1, jend1, istart2, iend2, jstart2, jend2);
      if(ncontact <= 2) { /* x-cyclic of y-cyclic */
	for(n=0; n<ncontact; n++) {
	  if(istart1[n] == iend1[n] && istart2[n] == iend2[n] ) /* x_cyclic */
	    x_cyclic = 1;
	  else if(jstart1[n] == jend1[n] && jstart2[n] == jend2[n] ) {
	    if( jstart1[n] == jend1[n] )
	      folded_north = 1;
	    else
	      y_cyclic = 1;
	  }
	  else
	    mpp_error("edit_landXocean: for one-tile mosaic, the boundary condition should be either x-cyclic, y-cyclic or folded-north");
	}
      }
      else
	mpp_error("edit_landXocean: the number of contact should be either 0, 1 or 2");    
    }
      
    free(tile1);
    free(tile2);
    free(istart1);
    free(iend1);
    free(jstart1);
    free(jend1);
    free(istart2);
    free(iend2);
    free(jstart2);
    free(jend2);
  }
  
  /* read the landXocean exchange grid information */
  {
    area_lXo = (double **)malloc(nlXo*sizeof(double *));
    il_lXo   = (int    **)malloc(nlXo*sizeof(int    *));
    jl_lXo   = (int    **)malloc(nlXo*sizeof(int    *));
    io_lXo   = (int    **)malloc(nlXo*sizeof(int    *));
    jo_lXo   = (int    **)malloc(nlXo*sizeof(int    *));
    nxgrid   = (int     *)malloc(nlXo*sizeof(int     ));
    for(n=0; n<nlXo; n++) {
      nxgrid[n]   = read_mosaic_xgrid_size(lXo_file[n]);
      area_lXo[n] = (double *)malloc(nxgrid[n]*sizeof(double));
      il_lXo[n]   = (int    *)malloc(nxgrid[n]*sizeof(int   ));  
      jl_lXo[n]   = (int    *)malloc(nxgrid[n]*sizeof(int   ));
      io_lXo[n]   = (int    *)malloc(nxgrid[n]*sizeof(int   ));
      jo_lXo[n]   = (int    *)malloc(nxgrid[n]*sizeof(int   ));
      read_mosaic_xgrid_order1(lXo_file[n], il_lXo[n], jl_lXo[n], io_lXo[n], jo_lXo[n], area_lXo[n]);    
    }
  }

  /*calculate ocean grid mask based on landXocean exchange grid information
    We also assume ocean cell is always full-cell. so the mask value will be
    either 0 (land) or 1 (ocean) */
  {
    int nx, ny, nxp1, nyp1, nxp2, nyp2;
    int n, m, i, j;
    
    ocn_mask = (int **)malloc(ntiles_ocn*sizeof(int *));
    nx   = nx_ocn[0];
    ny   = ny_ocn[0];
    nxp1 = nx + 1;
    nyp1 = ny + 1;
    nxp2 = nx + 2;
    nyp2 = ny + 2;
    for(n=0; n<ntiles_ocn; n++) {
      ocn_mask[n] = (int *)malloc(nxp2*nyp2*sizeof(int) );
      for(m=0; m<nxp2*nyp2; m++) ocn_mask[n][m] = 0;
    }

    for(n=0; n<nlXo; n++) {
      for(m=0; m<nxgrid[n]; m++) {
	i = io_lXo[n][m];   
	j = jo_lXo[n][m];
	ocn_mask[n][j*nxp2+i] = 1;
      }
    }
    /* fill the halo */
    if(x_cyclic) {
      for(j=1; j<nyp1; j++) {
	ocn_mask[0][j*nxp2     ] = ocn_mask[0][j*nxp2+nx];      /* West */
	ocn_mask[0][j*nxp2+nxp1] = ocn_mask[0][j*nxp2+ 1];      /* East */
      }
    }
    if(y_cyclic) {
      for(i=1; i<nxp1; i++) {
	ocn_mask[0][          i] = ocn_mask[0][ny*nxp2+i];      /* south */
	ocn_mask[0][nyp1*nxp2+i] = ocn_mask[0][   nxp2+i];      /* north */
      }
      if(x_cyclic) {
	ocn_mask[0][0]              = ocn_mask[0][ny*nxp2+nx];   /* southwest */
	ocn_mask[0][nxp1]           = ocn_mask[0][ny*nxp2+1];    /* southeast */    
	ocn_mask[0][nyp1*nxp2+nxp1] = ocn_mask[0][nxp2+1];       /* northeast */
	ocn_mask[0][nyp1*nxp2]      = ocn_mask[0][nxp2+nx];      /* northwest */
      }
    }
    else if(folded_north) {
      if( !x_cyclic) mpp_error("Error from program edit_landXocean: when the y-boundary condition "
			       "is folded-north, the x-boundary condition must be cyclic");
      for(i=1; i<nxp1; i++) {
	ocn_mask[0][nyp1*nxp2+i] = ocn_mask[0][ny*nxp2+nxp1-i];
      }
      ocn_mask[0][nyp1*nxp2]   = ocn_mask[0][ny*nxp2+1];
      ocn_mask[0][nyp1*nxp2+nxp1] = ocn_mask[0][ny*nxp2+nx];
    } 
  }
  /* figure out the exchange grid cells that the runoff will put to */
  /* set such point with value 1 and other points with value 0 */
  {
    double *scale, *tot_area, *set_area;
    int nxl, nyl, n, m, io, jo, il, jl, ioff, joff, ii, jj;
    int is_selected, nxp2;
    char orig_file[STRING], cmd[512];
    int fid, id_scale, id_ncell;
    
    nxp2 = nx_ocn[0] + 2;
    for(n=0; n<nlXo; n++) {
      nxl = nx_lnd[n];
      nyl = ny_lnd[n];
      tot_area = (double *)malloc(nxl*nyl*sizeof(double));
      set_area = (double *)malloc(nxl*nyl*sizeof(double));
      scale    = (double *)malloc(nxgrid[n]*sizeof(double));
      for(m=0; m<nxl*nyl; m++) {
	tot_area[m] = 0;
	set_area[m] = 0;
      }
      for(m=0; m<nxgrid[n]; m++) {
	il = il_lXo[n][m] - 1;
	jl = jl_lXo[n][m] - 1;  
	io = io_lXo[n][m];
	jo = jo_lXo[n][m];
	is_selected = 0;
	for(joff=-1; joff<=1; joff++) for(ioff=-1; ioff<=1; ioff++) {
	  if(ioff == 0 && joff == 0) continue;
	  ii = io+ioff;
	  jj = jo+joff;
	  if( ocn_mask[0][jj*nxp2+ii] == 0) {
	    is_selected = 1;
	  }
	}
	tot_area[jl*nxl+il] += area_lXo[n][m];
	if(is_selected) {
	  scale[m] = 1;
	  set_area[jl*nxl+il] += area_lXo[n][m];
	}
	else {
	  scale[m] = 0;
	}
      }

      for(m=0; m<nxgrid[n]; m++) {
	il = il_lXo[n][m] - 1;
	jl = jl_lXo[n][m] - 1;
	if(set_area[jl*nxl+il] == 0)
	  scale[m] = 1;
	else if(set_area[jl*nxl+il] > 0) 
	  scale[m] *= (tot_area[jl*nxl+il]/set_area[jl*nxl+il]);
	else {
	  printf("at point (%d, %d), set_area = %15.12f is negative\n", il, jl, set_area[jl*nxl+il]);
	  mpp_error("program edit_landXocean: set_area is negative for some point");
	}
      }
      free(tot_area);
      free(set_area);
      /* write the scale information into the xgrid file */
      if(!strcmp(lXo_file[n], lXo_base_file[n])) {
	sprintf(orig_file, "%s.orig", lXo_file[n]);
	sprintf(cmd, "cp %s %s", lXo_file[n], orig_file);
	system(cmd);
      }
      else {
	sprintf(cmd, "cp %s %s", lXo_file[n], lXo_base_file[n]);
        system(cmd);
      }
 
      fid = mpp_open(lXo_base_file[n], MPP_APPEND);
      mpp_redef(fid);
      id_ncell = mpp_get_dimid(fid, "ncells");
      id_scale = mpp_def_var(fid, "scale", MPP_DOUBLE, 1, &id_ncell, 0);
      mpp_end_def(fid);
      mpp_put_var_value(fid, id_scale, scale);
      mpp_close(fid);
    }
  }
      
  /* free the memory */
  for(n=0; n<nlXo; n++) {
    free(area_lXo[n]);
    free(il_lXo[n]);
    free(jl_lXo[n]);
    free(io_lXo[n]);
    free(jo_lXo[n]);
  }
  free(area_lXo);
  free(il_lXo);
  free(jl_lXo);
  free(io_lXo);
  free(jo_lXo);

  printf("Successfully running river_regrid and the following output file are generated.\n");

  mpp_end();

  return 0;  
} /* end of main */