#pragma once
#include "MString.h"
#include "traits.h"
#include <algorithm>
#include <netcdf.h>
#include <utility>
#include <vector>

using michlib::MString;

class NCFileW
{
 enum Type
 {
  UNKNOWN,
  G1V1,
  G1V2,
  G2V2
 };
 struct Var
 {
  MString name;
  int     id;
 };

 static constexpr auto fill = std::numeric_limits<float>::max();

 int  ncid;
 Type type = UNKNOWN;
 union
 {
  struct
  {
   int xdimid, ydimid;
  };
  int dimid[2];
 };
 int                     xid, yid;
 int                     compress;
 std::vector<struct Var> vars;

 template<class D> static constexpr Type DetType()
 {
  if constexpr(ReadIs2DGeoRectArray<D>) return G1V2;
  if constexpr(ReadIs2DGeoArray<D>) return G2V2;
  if constexpr(ReadIs1DGeoArray<D>) return G1V1;
  return UNKNOWN;
 };

 template<class D> static constexpr bool Is1DType() { return Is1DType(DetType<D>()); }

 static constexpr bool Is1DType(Type t) { return t == G1V1; }

 MString CreateFile(Type stype, const MString& name, const MString& history, int compression, size_t nx, size_t ny);

 public:
 static constexpr auto Fill() { return fill; }

 ~NCFileW() { Close(); }

 template<class D> MString Create(const D& data, const MString& name, const MString& history, int compression)
 {
  if(type != UNKNOWN) return "File already created";
  if constexpr(Is1DType<D>())
   return CreateFile(DetType<D>(), name, history, compression, data.N(), 0);
  else
   return CreateFile(DetType<D>(), name, history, compression, data.Nx(), data.Ny());
 }

 size_t VarId(const MString& name) const
 {
  for(size_t i = 0; i < vars.size(); i++)
   if(vars[i].name == name) return i + 1;

  return 0;
 }

 bool HaveVar(const MString& name) const { return VarId(name) != 0; }

 void EndDef() const
 {
  if(type != UNKNOWN) nc_enddef(ncid);
 }
 void Close()
 {
  if(type != UNKNOWN) nc_close(ncid);
  type = UNKNOWN;
 }

 MString AddVariable(const MString& name, const MString& stname, const MString& lname, const MString& units, const MString& comment);

 template<class D, class Op> MString WriteVariable(const D& data, size_t varid, Op op) const
 {
  static constexpr auto dtype = DetType<D>();
  if(type == UNKNOWN) return "File not open";
  if(type != dtype) return "Incompatible data type";
  if(varid == 0) return "Incorrect variable";
  if constexpr(dtype == UNKNOWN) return "Unknown data type";

  size_t v = varid - 1;

  EndDef();

  int ret;
  if constexpr(Is1DType(dtype))
  {
   const size_t i = 0, c = data.N();
   float        buf[c];
   for(size_t ix = 0; ix < c; ix++) buf[ix] = data.IsFill(ix) ? fill : op(ix);

   ret = nc_put_vara_float(ncid, vars[v].id, &i, &c, buf);
   if(ret != NC_NOERR) return "Can't write " + vars[v].name + " variable in the netcdf file";
  }
  else
  {
   const size_t i[2] = {0, 0};
   const size_t c[2] = {data.Nx(), data.Ny()};
   float        buf[c[0] * c[1]];

   for(size_t ix = 0; ix < c[0]; ix++)
    for(size_t iy = 0; iy < c[1]; iy++) buf[ix * c[1] + iy] = data.IsFill(ix, iy) ? fill : op(ix, iy);

   ret = nc_put_vara_float(ncid, vars[v].id, i, c, buf);
   if(ret != NC_NOERR) return "Can't write " + vars[v].name + " variable in the netcdf file";
  }
  return "";
 }

 template<class D, class Op> MString WriteVariable(const D& data, const MString& name, Op op) const { return WriteVariable(data, VarId(name), op); }

 template<class D> MString WriteVariable(const D& data, size_t varid) const
 {
  if constexpr(Is1DType(DetType<D>()))
   return WriteVariable(data, varid, [&data = std::as_const(data)](size_t i) { return data(i); });
  else
   return WriteVariable(data, varid, [&data = std::as_const(data)](size_t i, size_t j) { return data(i, j); });
 }

 template<class D> MString WriteVariable(const D& data, const MString& name) const { return WriteVariable(data, VarId(name)); }

 template<class D> MString WriteGrid(const D& data) const
 {
  static constexpr auto dtype = DetType<D>();
  if(type == UNKNOWN) return "File not open";
  if(type != dtype) return "Incompatible data type";

  EndDef();

  int ret;
  if constexpr(dtype == G1V1)
  {
   const size_t i = 0, c = data.N();
   float        bufx[c];
   float        bufy[c];
   for(size_t ix = 0; ix < c; ix++)
   {
    bufx[ix] = data.Lon(ix);
    bufy[ix] = data.Lat(ix);
   }
   ret = nc_put_vara_float(ncid, xid, &i, &c, bufx);
   if(ret != NC_NOERR) return "Can't write longitude variable in the netcdf file";
   ret = nc_put_vara_float(ncid, yid, &i, &c, bufy);
   if(ret != NC_NOERR) return "Can't write latitude variable in the netcdf file";
  }
  else if constexpr(dtype == G1V2)
  {
   const size_t i = 0, cx = data.Nx(), cy = data.Ny();
   float        bufx[cx];
   float        bufy[cy];

   for(size_t ix = 0; ix < cx; ix++) bufx[ix] = data.Ix2Lon(ix);
   for(size_t iy = 0; iy < cy; iy++) bufy[iy] = data.Iy2Lat(iy);

   ret = nc_put_vara_float(ncid, xid, &i, &cx, bufx);
   if(ret != NC_NOERR) return "Can't write longitude variable in the netcdf file";
   ret = nc_put_vara_float(ncid, yid, &i, &cy, bufy);
   if(ret != NC_NOERR) return "Can't write latitude variable in the netcdf file";
  }
  else if constexpr(dtype == G2V2)
  {
   const size_t i[2] = {0, 0};
   const size_t c[2] = {data.Nx(), data.Ny()};
   float        bufx[c[0] * c[1]];
   float        bufy[c[0] * c[1]];
   for(size_t ix = 0; ix < c[0]; ix++)
    for(size_t iy = 0; iy < c[1]; iy++)
    {
     bufx[ix * c[1] + iy] = data.Lon(ix, iy);
     bufy[ix * c[1] + iy] = data.Lat(ix, iy);
    }
   ret = nc_put_vara_float(ncid, xid, i, c, bufx);
   if(ret != NC_NOERR) return "Can't write longitude variable in the netcdf file";
   ret = nc_put_vara_float(ncid, yid, i, c, bufy);
   if(ret != NC_NOERR) return "Can't write latitude variable in the netcdf file";
  }
  else
   return "Unknown data type";

  return "";
 }
};