/*
  This file is part of CDO. CDO is a collection of Operators to manipulate and analyse Climate model Data.

  Author: Uwe Schulzweida

*/

/*
   This module contains the following operators:

      Maskbox    masklonlatbox   Mask lon/lat box
      Maskbox    maskindexbox    Mask index box
      Maskbox    maskregion      Mask regions
*/

#include <cdi.h>

#include <utility>

#include "cdo_options.h"
#include "process_int.h"
#include "readline.h"
#include <mpim_grid.h>
#include "selboxinfo.h"
#include "util_string.h"
#include "dcw_reader.h"

struct Polygons
{
  std::vector<double> x, y;
  std::vector<size_t> segmentSize;
  std::vector<size_t> segmentOffset;
  size_t numSegments = 0;
};

static int
read_coords(Varray<double> &xvals, Varray<double> &yvals, const char *polyfile, FILE *fp)
{
  const auto maxVals = xvals.size();
  constexpr size_t MAX_LINE = 256;
  char line[MAX_LINE];

  size_t number = 0, jumpedlines = 0;
  while (cdo::readline(fp, line, MAX_LINE))
    {
      if (line[0] == '#' || line[0] == '\0')
        {
          jumpedlines++;
          continue;
        }

      if (number == 0 && line[0] == '>') continue;  // Dump of DCW-GMT
      if (line[0] == '&' || line[0] == '>') break;

      const auto lineNo = number + jumpedlines + 1;

      double xcoord, ycoord;
      const auto nread = sscanf(line, "%lf %lf", &xcoord, &ycoord);
      if (nread != 2) cdo_abort("Wrong value format in file %s at line %d", polyfile, lineNo);

      if (number >= maxVals) cdo_abort("Too many polygon points (max=%zu)!", maxVals);
      xvals[number] = xcoord;
      yvals[number] = ycoord;
      number++;
    }

  if ((number != 0) && (!(IS_EQUAL(xvals[0], xvals[number - 1]) && IS_EQUAL(yvals[0], yvals[number - 1]))))
    {
      xvals[number] = xvals[0];
      yvals[number] = yvals[0];
      number++;
    }

  if (Options::cdoVerbose)
    for (size_t i = 0; i < number; i++) fprintf(stderr, "%zu %g %g\n", i + 1, xvals[i], yvals[i]);

  return number;
}

static void
read_polygons_from_file(const char *filename, Polygons &polygons)
{
  auto fp = fopen(filename, "r");
  if (fp == nullptr) cdo_abort("Open failed on %s", filename);

  constexpr size_t maxVals = 1048576;
  Varray<double> xcoords(maxVals), ycoords(maxVals);

  size_t n = 0;
  while (true)
    {
      const auto segmentSize = read_coords(xcoords, ycoords, filename, fp);
      if (segmentSize == 0) break;
      if (segmentSize < 3) cdo_abort("Too few point for polygon in file %s (Min=3)!", filename);

      auto &x = polygons.x;
      auto &y = polygons.y;
      const auto offset = x.size();
      polygons.segmentSize.push_back(segmentSize);
      polygons.segmentOffset.push_back(offset);
      polygons.numSegments++;
      n += segmentSize;
      x.resize(n);
      y.resize(n);
      for (int i = 0; i < segmentSize; ++i) x[offset + i] = xcoords[i];
      for (int i = 0; i < segmentSize; ++i) y[offset + i] = ycoords[i];
    }

  fclose(fp);
}

static void
read_polygons_from_dcw(const char *codeNames, Polygons &polygons)
{
  if (*codeNames == 0) cdo_abort("DCW parameter missing!");

  DCW_Lists dcw_lists;
  if (dcw_load_lists(dcw_lists)) cdo_abort("dcw_load_lists failed!");

  auto codeList = split_string(codeNames, "\\+");
  dcw_sort_countries(dcw_lists);

  codeList = dcw_expand_code_list(dcw_lists, codeList);

  if (codeList.size() == 0) cdo_abort("Empty country code list!");

  auto &lon = polygons.x;
  auto &lat = polygons.y;
  if (dcw_get_lonlat(dcw_lists, codeList, lon, lat)) cdo_abort("Reading DCW data failed!");

  const auto n = lon.size();
  for (size_t i = 0; i < n; ++i)
    {
      if (is_equal(lon[i], 0.0) && is_equal(lat[i], 0.0))
        {
          polygons.segmentOffset.push_back(i + 1);
          polygons.numSegments++;
        }
    }
  auto numSegments = polygons.numSegments;
  for (size_t i = 0; i < numSegments - 1; ++i)
    {
      auto segmentSize = polygons.segmentOffset[i + 1] - polygons.segmentOffset[i] - 1;
      polygons.segmentSize.push_back(segmentSize);
    }
  polygons.segmentSize.push_back(polygons.x.size() - polygons.segmentOffset[numSegments - 1]);
}

static void
maskbox(std::vector<char> &mask, const int gridID, const SelboxInfo &sbox)
{
  const auto &lat1 = sbox.lat1;
  const auto &lat2 = sbox.lat2;
  const auto &lon11 = sbox.lon11;
  const auto &lon12 = sbox.lon12;
  const auto &lon21 = sbox.lon21;
  const auto &lon22 = sbox.lon22;
  const long nlon = gridInqXsize(gridID);
  const long nlat = gridInqYsize(gridID);

  for (long ilat = 0; ilat < nlat; ilat++)
    for (long ilon = 0; ilon < nlon; ilon++)
      if ((lat1 <= ilat && ilat <= lat2 && ((lon11 <= ilon && ilon <= lon12) || (lon21 <= ilon && ilon <= lon22))))
        mask[nlon * ilat + ilon] = false;
}

void getlonlatparams(int argc_offset, double &xlon1, double &xlon2, double &xlat1, double &xlat2);

static void
maskbox_cell(std::vector<char> &mask, const int gridID)
{
  double xlon1 = 0, xlon2 = 0, xlat1 = 0, xlat2 = 0;
  getlonlatparams(0, xlon1, xlon2, xlat1, xlat2);

  const auto gridsize = gridInqSize(gridID);

  Varray<double> xvals(gridsize), yvals(gridsize);
  gridInqXvals(gridID, xvals.data());
  gridInqYvals(gridID, yvals.data());

  // Convert lat/lon units if required
  cdo_grid_to_degree(gridID, CDI_XAXIS, gridsize, xvals.data(), "grid center lon");
  cdo_grid_to_degree(gridID, CDI_YAXIS, gridsize, yvals.data(), "grid center lat");

  if (xlon1 > xlon2) cdo_abort("The second longitude have to be greater than the first one!");

  if (xlat1 > xlat2) std::swap(xlat1, xlat2);

  for (size_t i = 0; i < gridsize; ++i)
    {
      mask[i] = true;

      const auto xval = xvals[i];
      const auto yval = yvals[i];
      if (yval >= xlat1 && yval <= xlat2)
        {
          if (((xval >= xlon1 && xval <= xlon2) || (xval - 360 >= xlon1 && xval - 360 <= xlon2)
               || (xval + 360 >= xlon1 && xval + 360 <= xlon2)))
            {
              mask[i] = false;
            }
        }
    }
}

static inline bool
is_point_inside(double xval, double yval, double xi, double xj, double yi, double yj)
{
  return (((yval > yi && yval < yj) || (yval > yj && yval < yi)) && (xval < ((xj - xi) * (yval - yi) / (yj - yi) + xi)));
}

static bool
point_is_inside(double xval, double yval, size_t n, const double *xcoords, const double *ycoords)
{
  bool c = false;

  for (size_t i = 0, j = n - 1; i < n; j = i++)
    {
      if (is_point_inside(xval, yval, xcoords[i], xcoords[j], ycoords[i], ycoords[j])) c = !c;
    }

  return c;
}

static bool
point_is_inside(double xval, double yval, double xmin, double xmax, const double *xcoords, const double *ycoords, size_t nofcoords)
{
  bool c = false;

  // clang-format off
  if (xval >= xmin && xval <= xmax)
    c = point_is_inside(xval,         yval, nofcoords, xcoords, ycoords);
  else if (xval > 180.0 && xval - 360.0 >= xmin && xval - 360.0 <= xmax)
    c = point_is_inside(xval - 360.0, yval, nofcoords, xcoords, ycoords);
  else if (xval <   0.0 && xval + 360.0 >= xmin && xval + 360.0 <= xmax)
    c = point_is_inside(xval + 360.0, yval, nofcoords, xcoords, ycoords);
  // clang-format on

  return c;
}

static void
mask_region_regular(std::vector<char> &mask, size_t nlon, size_t nlat, const Varray<double> &xvals, const Varray<double> &yvals,
                    const double *xcoords, const double *ycoords, size_t segmentSize)
{
  auto xmm = varray_min_max(segmentSize, xcoords);
  auto ymm = varray_min_max(segmentSize, ycoords);

  auto gridsize = nlon * nlat;
#ifdef _OPENMP
#pragma omp parallel for schedule(static) default(shared)
#endif
  for (size_t i = 0; i < gridsize; i++)
    {
      const auto ilat = i / nlon;
      const auto yval = yvals[ilat];
      if (yval >= ymm.min && yval <= ymm.max)
        {
          if (point_is_inside(xvals[i - ilat * nlon], yval, xmm.min, xmm.max, xcoords, ycoords, segmentSize)) mask[i] = false;
        }
    }
}

static void
mask_region_cell(std::vector<char> &mask, size_t gridsize, const Varray<double> &xvals, const Varray<double> &yvals,
                 const double *xcoords, const double *ycoords, size_t segmentSize)
{
  auto xmm = varray_min_max(segmentSize, xcoords);
  auto ymm = varray_min_max(segmentSize, ycoords);

#ifdef _OPENMP
#pragma omp parallel for schedule(static) default(shared)
#endif
  for (size_t i = 0; i < gridsize; i++)
    {
      const auto yval = yvals[i];
      if (yval >= ymm.min && yval <= ymm.max)
        {
          if (point_is_inside(xvals[i], yval, xmm.min, xmm.max, xcoords, ycoords, segmentSize)) mask[i] = false;
        }
    }
}

void *
Maskbox(void *process)
{
  int index, gridtype = -1;
  SelboxInfo sbox;

  cdo_initialize(process);

  // clang-format off
  const auto MASKLONLATBOX = cdo_operator_add("masklonlatbox", 0, 0, "western and eastern longitude and southern and northern latitude");
  const auto MASKINDEXBOX  = cdo_operator_add("maskindexbox",  0, 0, "index of first and last longitude and index of first and last latitude");
  const auto MASKREGION    = cdo_operator_add("maskregion",    0, 0, "limiting coordinates of the region");
  // clang-format on

  const auto operatorID = cdo_operator_id();

  operator_input_arg(cdo_operator_enter(operatorID));

  const auto streamID1 = cdo_open_read(0);

  const auto vlistID1 = cdo_stream_inq_vlist(streamID1);
  const auto vlistID2 = vlistDuplicate(vlistID1);

  const auto taxisID1 = vlistInqTaxis(vlistID1);
  const auto taxisID2 = taxisDuplicate(taxisID1);
  vlistDefTaxis(vlistID2, taxisID2);

  const auto nvars = vlistNvars(vlistID1);
  std::vector<bool> vars(nvars, false);

  const auto ngrids = vlistNgrids(vlistID1);
  int ndiffgrids = 0;
  for (index = 1; index < ngrids; index++)
    if (vlistGrid(vlistID1, 0) != vlistGrid(vlistID1, index)) ndiffgrids++;

  int gridID = -1;
  for (index = 0; index < ngrids; index++)
    {
      gridID = vlistGrid(vlistID1, index);
      gridtype = gridInqType(gridID);

      if (gridtype == GRID_LONLAT || gridtype == GRID_GAUSSIAN || gridtype == GRID_CURVILINEAR || gridtype == GRID_UNSTRUCTURED)
        break;
      if (operatorID == MASKINDEXBOX && gridtype == GRID_GENERIC && gridInqXsize(gridID) > 0 && gridInqYsize(gridID) > 0) break;
    }

  if (gridtype == GRID_GAUSSIAN_REDUCED) cdo_abort("Gaussian reduced grid found. Use option -R to convert it to a regular grid!");

  if (index == ngrids) cdo_abort("No regular lon/lat grid found!");
  if (ndiffgrids > 0) cdo_abort("Too many different grids!");

  operator_input_arg(cdo_operator_enter(operatorID));

  for (int varID = 0; varID < nvars; varID++) vars[varID] = (gridID == vlistInqVarGrid(vlistID1, varID));

  const auto streamID2 = cdo_open_write(1);
  cdo_def_vlist(streamID2, vlistID2);

  const auto gridsize = gridInqSize(gridID);
  Varray<double> array(gridsize);
  std::vector<char> mask(gridsize, true);

  if (operatorID == MASKLONLATBOX)
    {
      if (gridtype == GRID_CURVILINEAR || gridtype == GRID_UNSTRUCTURED)
        {
          maskbox_cell(mask, gridID);
        }
      else
        {
          genlonlatbox(0, gridID, sbox);
          maskbox(mask, gridID, sbox);
        }
    }
  else if (operatorID == MASKINDEXBOX)
    {
      genindexbox(0, gridID, sbox);
      maskbox(mask, gridID, sbox);
    }
  else if (operatorID == MASKREGION)
    {
      const auto nlon = gridInqXsize(gridID);
      const auto nlat = gridInqYsize(gridID);
      const auto fullGrid = (gridtype == GRID_CURVILINEAR || gridtype == GRID_UNSTRUCTURED);
      Varray<double> xvals(fullGrid ? gridsize : nlon), yvals(fullGrid ? gridsize : nlat);

      gridInqXvals(gridID, xvals.data());
      gridInqYvals(gridID, yvals.data());

      // Convert lat/lon units if required
      cdo_grid_to_degree(gridID, CDI_XAXIS, xvals.size(), xvals.data(), "grid center lon");
      cdo_grid_to_degree(gridID, CDI_YAXIS, yvals.size(), yvals.data(), "grid center lat");

      const auto nfiles = cdo_operator_argc();

      for (int i = 0; i < nfiles; i++)
        {
          Polygons polygons;
          const auto param = cdo_operator_argv(i).c_str();
          if (strncmp(param, "dcw:", 4) == 0)
            read_polygons_from_dcw(param + 4, polygons);
          else
            read_polygons_from_file(param, polygons);

          for (size_t k = 0; k < polygons.numSegments; ++k)
            {
              const auto segmentSize = polygons.segmentSize[k];
              if (segmentSize < 3) continue;
              const auto offset = polygons.segmentOffset[k];
              const auto xcoords = &polygons.x[offset];
              const auto ycoords = &polygons.y[offset];
              if (fullGrid)
                mask_region_cell(mask, gridsize, xvals, yvals, xcoords, ycoords, segmentSize);
              else
                mask_region_regular(mask, nlon, nlat, xvals, yvals, xcoords, ycoords, segmentSize);
            }
        }
    }

  int tsID = 0;
  while (true)
    {
      const auto nrecs = cdo_stream_inq_timestep(streamID1, tsID);
      if (nrecs == 0) break;

      cdo_taxis_copy_timestep(taxisID2, taxisID1);
      cdo_def_timestep(streamID2, tsID);

      for (int recID = 0; recID < nrecs; recID++)
        {
          int varID, levelID;
          cdo_inq_record(streamID1, &varID, &levelID);

          if (vars[varID])
            {
              size_t nmiss;
              cdo_read_record(streamID1, array.data(), &nmiss);

              const auto missval = vlistInqVarMissval(vlistID1, varID);
              for (size_t i = 0; i < gridsize; i++)
                if (mask[i]) array[i] = missval;

              nmiss = varray_num_mv(gridsize, array, missval);
              cdo_def_record(streamID2, varID, levelID);
              cdo_write_record(streamID2, array.data(), nmiss);
            }
        }

      tsID++;
    }

  cdo_stream_close(streamID2);
  cdo_stream_close(streamID1);

  cdo_finish();

  return nullptr;
}
