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

using michlib::MString;

class NCFileWBase
{
 template<class T, class Dummy> struct NCTypeD
 {
  static constexpr nc_type nc = NC_NAT;
 };

 template<class Dummy> struct NCTypeD<michlib::int1, Dummy>
 {
  static constexpr nc_type nc = NC_BYTE;
 };
 template<class Dummy> struct NCTypeD<michlib::uint1, Dummy>
 {
  static constexpr nc_type nc = NC_UBYTE;
  static int               put_var(int nc, int vid, const michlib::uint1* data) { return nc_put_var_ubyte(nc, vid, data); }
  static int put_vara(int nc, int vid, const size_t* startp, const size_t* countp, const michlib::uint1* data) { return nc_put_vara_ubyte(nc, vid, startp, countp, data); }
 };
 template<class Dummy> struct NCTypeD<michlib::int2, Dummy>
 {
  static constexpr nc_type nc = NC_SHORT;
  static int               put_var(int nc, int vid, const michlib::int2* data) { return nc_put_var_short(nc, vid, data); }
  static int put_vara(int nc, int vid, const size_t* startp, const size_t* countp, const michlib::int2* data) { return nc_put_vara_short(nc, vid, startp, countp, data); }
 };
 template<class Dummy> struct NCTypeD<michlib::uint2, Dummy>
 {
  static constexpr nc_type nc = NC_USHORT;
  static int               put_var(int nc, int vid, const michlib::uint2* data) { return nc_put_var_ushort(nc, vid, data); }
  static int put_vara(int nc, int vid, const size_t* startp, const size_t* countp, const michlib::uint2* data) { return nc_put_vara_ushort(nc, vid, startp, countp, data); }
 };
 template<class Dummy> struct NCTypeD<michlib::int4, Dummy>
 {
  static constexpr nc_type nc = NC_INT;
  static int               put_var(int nc, int vid, const michlib::int4* data) { return nc_put_var_int(nc, vid, data); }
  static int put_vara(int nc, int vid, const size_t* startp, const size_t* countp, const michlib::int4* data) { return nc_put_vara_int(nc, vid, startp, countp, data); }
 };
 template<class Dummy> struct NCTypeD<michlib::uint4, Dummy>
 {
  static constexpr nc_type nc = NC_UINT;
  static int               put_var(int nc, int vid, const michlib::uint4* data) { return nc_put_var_uint(nc, vid, data); }
  static int put_vara(int nc, int vid, const size_t* startp, const size_t* countp, const michlib::uint4* data) { return nc_put_vara_uint(nc, vid, startp, countp, data); }
 };
 template<class Dummy> struct NCTypeD<michlib::int8, Dummy>
 {
  static constexpr nc_type nc = NC_INT64;
 };
 template<class Dummy> struct NCTypeD<michlib::uint8, Dummy>
 {
  static constexpr nc_type nc = NC_UINT64;
 };
 template<class Dummy> struct NCTypeD<float, Dummy>
 {
  static constexpr nc_type nc = NC_FLOAT;
  static int               put_var(int nc, int vid, const float* data) { return nc_put_var_float(nc, vid, data); }
  static int               put_vara(int nc, int vid, const size_t* startp, const size_t* countp, const float* data) { return nc_put_vara_float(nc, vid, startp, countp, data); }
 };
 template<class Dummy> struct NCTypeD<double, Dummy>
 {
  static constexpr nc_type nc = NC_DOUBLE;
  static int               put_var(int nc, int vid, const double* data) { return nc_put_var_double(nc, vid, data); }
  static int               put_vara(int nc, int vid, const size_t* startp, const size_t* countp, const double* data) { return nc_put_vara_double(nc, vid, startp, countp, data); }
 };
 template<class Dummy> struct NCTypeD<MString, Dummy>
 {
  static constexpr nc_type nc = NC_CHAR;
 };
 template<class Dummy> struct NCTypeD<char*, Dummy>
 {
  static constexpr nc_type nc = NC_CHAR;
 };

 template<class T> using NCType = NCTypeD<T, void>;

 public:
 template<class T> static constexpr nc_type Type2NCType = NCType<T>::nc;

 // Error class
 class Error
 {
  int err;

  public:
  constexpr Error(): err(NC_NOERR) {}
  constexpr Error(int n_err): err(n_err) {}
  explicit    operator bool() const { return err == NC_NOERR; }
  const char* ErrMessage() const { return nc_strerror(err); }
  int         Err() const { return err; }
  bool        IsErr() const { return err != NC_NOERR; }
  void        Reset() { err = NC_NOERR; }
  void        Reset(int n_err) { err = n_err; }
  void        Reset(const Error& e) { err = e.err; }
 };

 ~NCFileWBase() { Close(); }

 int         Id() const { return ncid; }
 int         Err() const { return err.Err(); }
 const char* ErrMessage() const { return err.ErrMessage(); }

 explicit operator bool() const { return Ok(); }
 bool     Ok() const { return opened && !err.IsErr(); }
 operator int() const { return ncid; }

 void Close()
 {
  if(opened) nc_close(ncid);
  err.Reset();
  opened = false;
 }

 void Open(const MString& path)
 {
  if(opened) return;
  err.Reset(nc_create(path.Buf(), NC_CLOBBER | NC_NETCDF4, &ncid));
  opened = !err.IsErr();
 }

 void AddDim(const MString& name, size_t len)
 {
  if(err.IsErr()) return;
  static int dimid;
  err.Reset(nc_def_dim(ncid, name.Buf(), len, &dimid));
 }

 MString AddAtts(const michlib_internal::ParameterListEx& pars)
 {
  for(const auto& p: pars.GetParameterList())
  {
   MString ret;
   if(p.first.Prefix() != "") continue; // Only parameters with empty prefix considered for now
   switch(p.second.valtype)
   {
   case michlib_internal::Parameter::UINT1:
   {
    ret = AddAtt(p.first.Name(), p.second.value.u1);
    break;
   }
   case michlib_internal::Parameter::UINT2:
   {
    ret = AddAtt(p.first.Name(), p.second.value.u2);
    break;
   }
   case michlib_internal::Parameter::UINT4:
   {
    ret = AddAtt(p.first.Name(), p.second.value.u4);
    break;
   }
   case michlib_internal::Parameter::UINT8:
   {
    ret = AddAtt(p.first.Name(), p.second.value.u8);
    break;
   }
   case michlib_internal::Parameter::INT1:
   {
    ret = AddAtt(p.first.Name(), p.second.value.i1);
    break;
   }
   case michlib_internal::Parameter::INT2:
   {
    ret = AddAtt(p.first.Name(), p.second.value.i2);
    break;
   }
   case michlib_internal::Parameter::INT4:
   {
    ret = AddAtt(p.first.Name(), p.second.value.i4);
    break;
   }
   case michlib_internal::Parameter::INT8:
   {
    ret = AddAtt(p.first.Name(), p.second.value.i8);
    break;
   }
   case michlib_internal::Parameter::FLOAT:
   {
    ret = AddAtt(p.first.Name(), p.second.value.r4);
    break;
   }
   case michlib_internal::Parameter::DOUBLE:
   {
    ret = AddAtt(p.first.Name(), p.second.value.r8);
    break;
   }
   case michlib_internal::Parameter::LDOUBLE:
   {
    ret = AddAtt(p.first.Name(), static_cast<double>(p.second.value.r10));
    break;
   }
   case michlib_internal::Parameter::BOOL:
   {
    ret = AddAtt(p.first.Name(), p.second.value.b ? 1 : 0);
    break;
   }
   case michlib_internal::Parameter::STRING:
   {
    ret = AddAtt(p.first.Name(), p.second.svalue);
    break;
   }
   case michlib_internal::Parameter::INVALID: break; // Silently skip
   }
   if(ret.Exist()) return ret;
  }
  return "";
 }

 template<class T>
 requires(Type2NCType<T> != NC_NAT)
 MString AddAtt(const MString& name, const T& val)
 {
  if(err.IsErr()) return "Can't write attribute " + name + " due to previous errors";
  static constexpr auto nct    = Type2NCType<T>;
  static constexpr bool ischar = std::is_same_v<T, char*>;

  if constexpr(nct != NC_CHAR)
   err.Reset(nc_put_att(ncid, NC_GLOBAL, name.Buf(), nct, 1, &val));
  else if constexpr(ischar)
   err.Reset(nc_put_att_text(ncid, NC_GLOBAL, name.Buf(), strlen(val) + 1, val));
  else
   err.Reset(nc_put_att_text(ncid, NC_GLOBAL, name.Buf(), val.Len() + 1, val.Buf()));

  return err.IsErr() ? "Can't write attribute " + name + ": " + ErrMessage() : "";
 }

 template<class T>
 requires(Type2NCType<T> != NC_NAT)
 MString AddAtt(const MString& vname, const MString& name, const T& val)
 {
  if(err.IsErr()) return "Can't write attribute " + name + " to the variable " + vname + " due to previous errors";
  static constexpr auto nct    = Type2NCType<T>;
  static constexpr bool ischar = std::is_same_v<T, char*>;
  int                   varid;

  err.Reset(nc_inq_varid(ncid, vname.Buf(), &varid));
  if(err.IsErr()) return "Can't find variable " + vname + ": " + ErrMessage();

  if constexpr(nct != NC_CHAR)
   err.Reset(nc_put_att(ncid, varid, name.Buf(), nct, 1, &val));
  else if constexpr(ischar)
   err.Reset(nc_put_att_text(ncid, varid, name.Buf(), strlen(val) + 1, val));
  else
   err.Reset(nc_put_att_text(ncid, varid, name.Buf(), val.Len() + 1, val.Buf()));

  return err.IsErr() ? "Can't write attribute " + name + " to the variable " + vname + ": " + ErrMessage() : "";
 }

 MString AddAtt(const MString& vname, const MString& name, const char* val)
 {
  if(err.IsErr()) return "Can't write attribute " + name + " to the variable " + vname + " due to previous errors";
  int varid;

  err.Reset(nc_inq_varid(ncid, vname.Buf(), &varid));
  if(err.IsErr()) return "Can't find variable " + vname + ": " + ErrMessage();
  err.Reset(nc_put_att_text(ncid, varid, name.Buf(), strlen(val) + 1, val));
  return err.IsErr() ? "Can't write attribute " + name + " to the variable " + vname + ": " + ErrMessage() : "";
 }

 template<class... T> void AddVar(const MString& vname, nc_type type, const T&... dnames) { AddVar(vname, type, std::vector<MString>({dnames...})); }

 void EndDef() const { nc_enddef(ncid); }

 void SetComp(const MString& vname, int comp)
 {
  if(err.IsErr()) return;
  int varid;

  err.Reset(nc_inq_varid(ncid, vname.Buf(), &varid));
  if(err.IsErr()) return;

  err.Reset(nc_def_var_deflate(ncid, varid, 1, 1, comp));
 }

 template<class T>
 requires requires(int nc, int vid, const T* d) {
  {
   NCTypeD<T, void>::put_var(nc, vid, d)
  } -> std::same_as<int>;
 }
 void WriteVar(const MString& vname, const T* data)
 {
  if(err.IsErr()) return;

  int varid;
  err.Reset(nc_inq_varid(ncid, vname.Buf(), &varid));
  if(err.IsErr()) return;

  err.Reset(NCTypeD<T, void>::put_var(ncid, varid, data));
 }

 template<class T>
 requires requires(int nc, int vid, const size_t* start, const size_t* count, const T* d) {
  {
   NCTypeD<T, void>::put_vara(nc, vid, start, count, d)
  } -> std::same_as<int>;
 }
 void WriteVar(const MString& vname, size_t ind, const T* data)
 {
  if(err.IsErr()) return;

  int varid;
  err.Reset(nc_inq_varid(ncid, vname.Buf(), &varid));
  if(err.IsErr()) return;

  int ndim;
  err.Reset(nc_inq_var(ncid, varid, nullptr, nullptr, &ndim, nullptr, nullptr));
  if(err.IsErr()) return;
  std::vector<int> dimids(ndim);
  err.Reset(nc_inq_var(ncid, varid, nullptr, nullptr, nullptr, dimids.data(), nullptr));
  if(err.IsErr()) return;

  std::vector<size_t> start(ndim), count(ndim);
  start[0] = ind;
  count[0] = 1;
  for(size_t i = 1; i < dimids.size(); i++)
  {
   start[i] = 0;
   err.Reset(nc_inq_dim(ncid, dimids[i], nullptr, &count[i]));
   if(err.IsErr()) return;
  }

  err.Reset(NCTypeD<T, void>::put_vara(ncid, varid, start.data(), count.data(), data));
 }

 private:
 // Members
 int   ncid;
 Error err    = NC_NOERR;
 bool  opened = false;

 void AddVar(const MString& vname, nc_type type, std::vector<MString>&& dnames)
 {
  if(err.IsErr()) return;
  std::vector<int> dimids(dnames.size());
  int              varid;

  for(size_t i = 0; i < dnames.size(); i++)
  {
   err.Reset(nc_inq_dimid(ncid, dnames[i].Buf(), &dimids[i]));
   if(err.IsErr()) return;
  }
  err.Reset(nc_def_var(ncid, vname.Buf(), type, dimids.size(), dimids.data(), &varid));
 }
};

class NCFileW: public NCFileWBase
{
 enum Type
 {
  UNKNOWN,
  PSET,
  GPSET,
  GRID,
  GGRID,
  RGRID,
  GRGRID
 };

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

 Type type = UNKNOWN;
 int  compress;

 template<class D> static constexpr Type DetType()
 {
  if constexpr(ReadIs2DGeoRectArray<D>) return GRGRID;
  if constexpr(ReadIs2DGeoArray<D>) return GGRID;
  if constexpr(ReadIs1DGeoArray<D>) return GPSET;
  if constexpr(ReadIs2DXYRectArray<D>) return RGRID;
  if constexpr(ReadIs2DXYArray<D>) return GRID;
  if constexpr(ReadIs1DArray<D>) return PSET;
  return UNKNOWN;
 };

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

 static constexpr bool IsGeoType(Type t) { return t == GPSET || t == GGRID || t == GRGRID; }
 static constexpr bool Is1DType(Type t) { return t == PSET || t == GPSET; }

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

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

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

 void Close()
 {
  NCFileWBase::Close();
  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, const MString& name, Op op)
 {
  static constexpr auto dtype = DetType<D>();
  if(type == UNKNOWN) return "File not open";
  if(type != dtype) return "Incompatible data type";
  if constexpr(dtype == UNKNOWN) return "Unknown data type";

  EndDef();

  if constexpr(Is1DType(dtype))
  {
   const size_t c = data.N();
   float        buf[c];
   for(size_t ix = 0; ix < c; ix++) buf[ix] = data.IsFill(ix) ? fill : op(ix);
   WriteVar(name, buf);
  }
  else
  {
   const size_t c[2] = {data.Ny(), data.Nx()};
   float        buf[c[0] * c[1]];
   for(size_t iy = 0; iy < c[0]; iy++)
    for(size_t ix = 0; ix < c[1]; ix++) buf[iy * c[1] + ix] = data.IsFill(ix, iy) ? fill : op(ix, iy);
   WriteVar(name, buf);
  }

  if(!*this) return "Can't write variable " + name + ": " + ErrMessage();

  return "";
 }

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

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

  EndDef();

  if constexpr(dtype == UNKNOWN)
   return "Unknown data type";
  else if constexpr(dtype == PSET)
  {
   const size_t c = data.N();
   float        bufx[c];
   float        bufy[c];
   for(size_t ix = 0; ix < c; ix++)
   {
    bufx[ix] = data.X(ix);
    bufy[ix] = data.Y(ix);
   }
   WriteVar("x", bufx);
   WriteVar("y", bufy);
  }
  else if constexpr(dtype == GPSET)
  {
   const size_t 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);
   }
   WriteVar("longitude", bufx);
   WriteVar("latitude", bufy);
  }
  else if constexpr(dtype == RGRID)
  {
   const size_t cx = data.Nx(), cy = data.Ny();
   float        bufx[cx];
   float        bufy[cy];
   for(size_t ix = 0; ix < cx; ix++) bufx[ix] = data.Ix2X(ix);
   for(size_t iy = 0; iy < cy; iy++) bufy[iy] = data.Iy2Y(iy);
   WriteVar("x", bufx);
   WriteVar("y", bufy);
  }
  else if constexpr(dtype == GRGRID)
  {
   const size_t 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);
   WriteVar("longitude", bufx);
   WriteVar("latitude", bufy);
  }
  else if constexpr(dtype == GRID)
  {
   const size_t c[2] = {data.Ny(), data.Nx()};
   float        bufx[c[0] * c[1]];
   float        bufy[c[0] * c[1]];
   for(size_t iy = 0; iy < c[0]; iy++)
    for(size_t ix = 0; ix < c[1]; ix++)
    {
     bufx[iy * c[1] + ix] = data.X(ix, iy);
     bufy[iy * c[1] + ix] = data.Y(ix, iy);
    }
   WriteVar("x", bufx);
   WriteVar("y", bufy);
  }
  else if constexpr(dtype == GGRID)
  {
   const size_t c[2] = {data.Ny(), data.Nx()};
   float        bufx[c[0] * c[1]];
   float        bufy[c[0] * c[1]];
   for(size_t iy = 0; iy < c[0]; iy++)
    for(size_t ix = 0; ix < c[1]; ix++)
    {
     bufx[iy * c[1] + ix] = data.Lon(ix, iy);
     bufy[iy * c[1] + ix] = data.Lat(ix, iy);
    }
   WriteVar("longitude", bufx);
   WriteVar("latitude", bufy);
  }
  else
   return "Unknown data type";

  if(!*this) return MString("Can't write grid: ") + ErrMessage();

  return "";
 }
};