#pragma once
#include "BFileW.h"
#include "actiondep.h"
#include "ncfilew.h"
#include <utility>

using michlib::BFileW;

class UVMethods
{
 protected:
 template<class D> class SparserHolder
 {
  protected:
  const D& data;
  size_t   shiftx, shifty, skipx, skipy;
  SparserHolder(const D& d, size_t shx, size_t shy, size_t sx, size_t sy): data(d), shiftx(shx), shifty(shy), skipx(sx), skipy(sy) {}

  size_t I2Ix(size_t i) const { return i * skipx + shiftx; }
  size_t I2Iy(size_t i) const { return i * skipy + shifty; }

  public:
  size_t Nx() const
  {
   size_t n = data.Nx() - shiftx;
   return n / skipx + ((n % skipx > 0) ? 1 : 0);
  }
  size_t Ny() const
  {
   size_t n = data.Ny() - shifty;
   return n / skipy + ((n % skipy > 0) ? 1 : 0);
  }

  static constexpr auto Fillval() { return D::Fillval(); }

  real Lon(size_t ix, size_t iy) const { return data.Lon(I2Ix(ix), I2Iy(iy)); }
  real Lat(size_t ix, size_t iy) const { return data.Lat(I2Ix(ix), I2Iy(iy)); }
  real U(size_t ix, size_t iy) const { return data.U(I2Ix(ix), I2Iy(iy)); }
  real V(size_t ix, size_t iy) const { return data.V(I2Ix(ix), I2Iy(iy)); }
  bool IsFill(size_t ix, size_t iy) const { return data.IsFill(I2Ix(ix), I2Iy(iy)); }
 };

 template<class D, bool b> class Sparser2DGeoRectArray: public SparserHolder<D>
 {
  protected:
  Sparser2DGeoRectArray(const D& d, size_t shx, size_t shy, size_t sx, size_t sy): SparserHolder<D>(d, shx, shy, sx, sy) {}
 };
 template<class D> class Sparser2DGeoRectArray<D, true>: public SparserHolder<D>
 {
  protected:
  using SparserHolder<D>::data;
  using SparserHolder<D>::I2Ix;
  using SparserHolder<D>::I2Iy;

  Sparser2DGeoRectArray(const D& d, size_t shx, size_t shy, size_t sx, size_t sy): SparserHolder<D>(d, shx, shy, sx, sy) {}

  public:
  real Ix2Lon(size_t ix) const { return data.Ix2Lon(I2Ix(ix)); }
  real Iy2Lat(size_t iy) const { return data.Iy2Lat(I2Iy(iy)); }
 };

 template<class D> class Sparser: public Sparser2DGeoRectArray<D, ReadIs2DGeoRectArray<D>>
 {
  public:
  Sparser(const D& d, size_t shx, size_t shy, size_t sx, size_t sy): Sparser2DGeoRectArray<D, ReadIs2DGeoRectArray<D>>(d, shx, shy, sx, sy) {}
 };

 class StPoints
 {
  std::vector<real>           x, y;
  std::vector<michlib::uint1> t;
  bool                        lonlat;

  size_t N() const { return t.size(); }

  public:
  StPoints(bool ll): lonlat(ll) {}

  void Add(const std::vector<UVDataStorage::StPoint>& points)
  {
   for(const auto& p: points)
   {
    x.push_back(p.x);
    y.push_back(p.y);
    t.push_back(michlib::int_cast<michlib::uint1>(p.type));
   }
  }

  void WriteBinBile(const MString& name) const;

  MString WriteNcFile(const MString& name, const MString& history, int comp) const;
 };
};

ADD_ACTION(UV, uv, ReadPSupported<Source> || ReadSupported<Source>, UVMethods);

template<class D> MString ActionUV::DoAction(const CLArgs& args, D& ds)
{
 auto [reg, regerr] = GetRegion<D>(args);
 if(regerr.Exist()) return regerr;

 auto resop = ds.Open(args);
 if(resop.Exist()) return "Can't open source: " + resop;

 michlib_internal::ParameterListEx pars;
 pars.UsePrefix("");
 pars.SetParameter("source", args.at("source"));

 auto [tindexes, err] = GetTIndexes(ds, args, pars);
 if(err.Exist()) return err;

 std::unique_ptr<const BaseParameters> sourcepars;
 if constexpr(ParametersSupported<D>)
 {
  if constexpr(ParametersRequiredRegion<D>)
  {
   auto [p, err] = ds.Parameters(pars, args, reg);
   if(err.Exist()) return err;
   sourcepars.reset(p);
  }
  else
  {
   auto [p, err] = ds.Parameters(pars, args);
   if(err.Exist()) return err;
   sourcepars.reset(p);
  }
 }
 auto p = sourcepars.get();

 auto data = ReadUV(ds, p, tindexes);
 if(!data) return "Can't read data";

 // Main file
 MString name   = args.contains("out") ? args.at("out") : "";
 MString outfmt = args.contains("outformat") ? args.at("outformat") : (GetExt(name) == "nc" ? "nc" : "bin");

 MString u = data.Unit().Exist() ? data.Unit() : "unknown", d = data.DUnit().Exist() ? data.DUnit() : "unknown";

 if(name.Exist())
 {
  if(outfmt == "bin")
  {
   BFileW fw;
   fw.Create(name, 9);
   fw.SetColumnName(1, "Longitude");
   fw.SetColumnName(2, "Latitude");
   fw.SetColumnName(3, "u, " + u);
   fw.SetColumnName(4, "v, " + u);
   fw.SetColumnName(5, "Div, (" + u + ")/" + d);
   fw.SetColumnName(6, "Rot, (" + u + ")/" + d);
   fw.SetColumnName(7, "Okubo-Weiss parameter, (" + u + ")2/" + d + "2");
   fw.SetColumnName(8, "Kinetic energy, (" + u + ")2");
   fw.SetColumnName(9, "Eddy kinetic energy, (" + u + ")2");
   fw.SetParameters(pars);

   for(size_t i = 0; i < data.N(); i++)
   {
    fw.Write(data.Lon(i));
    fw.Write(data.Lat(i));
    fw.Write(data.U(i) == data.Fillval() ? NAN : data.U(i));
    fw.Write(data.V(i) == data.Fillval() ? NAN : data.V(i));
    fw.Write(data.Div(i) == data.Fillval() ? NAN : data.Div(i));
    fw.Write(data.Rot(i) == data.Fillval() ? NAN : data.Rot(i));
    fw.Write(data.OW(i) == data.Fillval() ? NAN : data.OW(i));
    fw.Write(data.U2(i) == data.Fillval() ? NAN : data.U2(i));
    fw.Write((data.U(i) == data.Fillval() || data.V(i) == data.Fillval()) ? NAN : (data.U2(i) - (data.U(i) * data.U(i) + data.V(i) * data.V(i))));
   }
   fw.Finalize();
   fw.Close();
  }
  else if(outfmt == "nc" || outfmt == "netcdf")
  {
   int     compress = 3;
   NCFileW fw;
   MString err;

   if(args.contains("compress")) compress = args.at("compress").ToInt();

   if(!err.Exist()) err = fw.Create(data, name, args.at("_cmdline"), compress);
   if(!err.Exist()) err = fw.AddVariable("u", "", "Eastward velocity", u, "");
   if(!err.Exist()) err = fw.AddVariable("v", "", "Northward velocity", u, "");
   if(!err.Exist()) err = fw.AddVariable("div", "", "Velocity divergence", "(" + u + ")/" + d, "");
   if(!err.Exist()) err = fw.AddVariable("rot", "", "Velocity rotor", "(" + u + ")/" + d, "");
   if(!err.Exist()) err = fw.AddVariable("ow", "", "Okubo-Weiss parameter", "(" + u + ")2/" + d + "2", "");
   if(!err.Exist()) err = fw.AddVariable("ke", "", "Squared velocity module, u^2+v^2", "(" + u + ")2", "");
   if(!err.Exist()) err = fw.AddVariable("eke", "", "Squared velocity dispersion aka eddy kinetic energy, <u^2+v^2>-<u>^2-<v>^2", "(" + u + ")2", "");

   if(!err.Exist()) err = fw.WriteGrid(data);

   if(!err.Exist()) err = fw.WriteVariable(data, "u", [&data = std::as_const(data)](size_t i, size_t j) { return data.U(i, j); });
   if(!err.Exist()) err = fw.WriteVariable(data, "v", [&data = std::as_const(data)](size_t i, size_t j) { return data.V(i, j); });
   if(!err.Exist()) err = fw.WriteVariable(data, "div", [&data = std::as_const(data)](size_t i, size_t j) { return data.Div(i, j); });
   if(!err.Exist()) err = fw.WriteVariable(data, "rot", [&data = std::as_const(data)](size_t i, size_t j) { return data.Rot(i, j); });
   if(!err.Exist()) err = fw.WriteVariable(data, "ow", [&data = std::as_const(data)](size_t i, size_t j) { return data.OW(i, j); });
   if(!err.Exist()) err = fw.WriteVariable(data, "ke", [&data = std::as_const(data)](size_t i, size_t j) { return data.U2(i, j); });
   if(!err.Exist())
    err = fw.WriteVariable(data, "eke", [&data = std::as_const(data)](size_t i, size_t j) { return data.U2(i, j) - (data.U(i, j) * data.U(i, j) + data.V(i, j) * data.V(i, j)); });
   if(err.Exist()) return err;

   fw.Close();
  }
  else
   return "Unknown format: " + outfmt;
 }

 // Filtered vectors file
 name   = args.contains("velout") ? args.at("velout") : "";
 outfmt = args.contains("veloutformat") ? args.at("veloutformat") : (GetExt(name) == "nc" ? "nc" : "bin");
 if(name.Exist())
 {
  size_t shiftx = args.contains("shiftx") ? args.at("shiftx").ToInteger<size_t>() : 0;
  size_t shifty = args.contains("shifty") ? args.at("shifty").ToInteger<size_t>() : 0;
  size_t skipx  = args.contains("skipx") ? args.at("skipx").ToInteger<size_t>() : 1;
  size_t skipy  = args.contains("skipy") ? args.at("skipy").ToInteger<size_t>() : 1;

  Sparser<decltype(data)> sdata(data, shiftx, shifty, skipx, skipy);
  if(outfmt == "bin")
  {
   BFileW vel;
   vel.Create(name, 4);
   vel.SetColumnName(1, "Longitude");
   vel.SetColumnName(2, "Latitude");
   vel.SetColumnName(3, "u, " + u);
   vel.SetColumnName(4, "v, " + u);

   for(size_t ix = 0; ix < sdata.Nx(); ix++)
    for(size_t iy = 0; iy < sdata.Ny(); iy++)
    {
     vel.Write(sdata.Lon(ix, iy));
     vel.Write(sdata.Lat(ix, iy));
     vel.Write(sdata.U(ix, iy) == sdata.Fillval() ? NAN : sdata.U(ix, iy));
     vel.Write(sdata.V(ix, iy) == sdata.Fillval() ? NAN : sdata.V(ix, iy));
    }
   vel.Finalize();
   vel.Close();
  }
  else if(outfmt == "nc" || outfmt == "netcdf")
  {
   int     compress = 3;
   NCFileW fw;
   MString err;

   if(args.contains("compress")) compress = args.at("compress").ToInt();

   if(!err.Exist()) err = fw.Create(sdata, name, args.at("_cmdline"), compress);
   if(!err.Exist()) err = fw.AddVariable("u", "", "Eastward velocity", u, "");
   if(!err.Exist()) err = fw.AddVariable("v", "", "Northward velocity", u, "");

   if(!err.Exist()) err = fw.WriteGrid(sdata);

   if(!err.Exist()) err = fw.WriteVariable(sdata, "u", [&data = std::as_const(sdata)](size_t i, size_t j) { return data.U(i, j); });
   if(!err.Exist()) err = fw.WriteVariable(sdata, "v", [&data = std::as_const(sdata)](size_t i, size_t j) { return data.V(i, j); });
   if(err.Exist()) return err;

   fw.Close();
  }
  else
   return "Unknown format: " + outfmt;
 }

 // Stationary points
 name   = args.contains("stpout") ? args.at("stpout") : "";
 outfmt = args.contains("stpoutformat") ? args.at("stpoutformat") : (GetExt(name) == "nc" ? "nc" : "bin");
 if(name.Exist())
 {
  StPoints p(true);
  for(size_t ix = 0; ix < data.Nx() - 1; ix++)
   for(size_t iy = 0; iy < data.Ny() - 1; iy++) p.Add(data.StablePoints(ix, iy));

  if(outfmt == "bin")
   p.WriteBinBile(name);
  else if(outfmt == "nc" || outfmt == "netcdf")
  {
   int     compress = args.contains("compress") ? args.at("compress").ToInt() : 3;
   MString err      = p.WriteNcFile(name, args.at("_cmdline"), compress);
   if(err.Exist()) return err;
  }
  else
   return "Unknown format: " + outfmt;
 }

 return "";
};