#pragma once
#include "geohelpers.h"
#include "traits.h"
#include <vector>

using michlib::M_PI;
using michlib::real;
using michlib::Sqrt;

enum class Metric
{
 EUCLID,
 SPHERIC
};

class UVDataStorage
{
 protected:
 static constexpr real fillval = 1.0e10;

 std::vector<real> u, v, u2;
 size_t            nx = 0, ny = 0;
 Metric            metric;
 MString           unit;

 UVDataStorage() = default;
 real D(real lon1, real lat1, real lon2, real lat2) const
 {
  switch(metric)
  {
  case(Metric::EUCLID): return michlib::Hypot(lon2 - lon1, lat2 - lat1);
  // Spheric distance in km
  case(Metric::SPHERIC): return michlib::GCD(lon1 * M_PI / 180.0, lat1 * M_PI / 180.0, lon2 * M_PI / 180.0, lat2 * M_PI / 180.0) * 6371.0;
  }
  return 0.0;
 }

 // For derivatives
 real Ud(size_t ix, size_t iy) const { return IsCoast(ix, iy) ? 0.0 : U(ix, iy); }
 real Vd(size_t ix, size_t iy) const { return IsCoast(ix, iy) ? 0.0 : V(ix, iy); }

 void SetUnit(const MString& str) { unit = str; }

 public:
 enum STPOINT
 {
  SADDLE = 0,
  SACICFOCUS,
  SKNOT,
  UACICFOCUS,
  UKNOT,
  SCICFOCUS,
  UCICFOCUS,
  NOPOINT
 };

 struct StPoint
 {
  real    x = 0.0, y = 0.0;
  STPOINT type = NOPOINT;
 };

 const real& U(size_t i) const { return u[i]; }
 const real& V(size_t i) const { return v[i]; }
 const real& U(size_t ix, size_t iy) const { return U(iy * nx + ix); }
 const real& V(size_t ix, size_t iy) const { return V(iy * nx + ix); }

 real& U(size_t i) { return u[i]; }
 real& V(size_t i) { return v[i]; }
 real& U(size_t ix, size_t iy) { return U(iy * nx + ix); }
 real& V(size_t ix, size_t iy) { return V(iy * nx + ix); }

 const real& U2(size_t i) const { return u2[i]; }
 const real& U2(size_t ix, size_t iy) const { return U2(iy * nx + ix); }

 real& U2(size_t i) { return u2[i]; }
 real& U2(size_t ix, size_t iy) { return U2(iy * nx + ix); }

 size_t N() const { return u.size(); }
 size_t Nx() const { return nx; }
 size_t Ny() const { return ny; }

 explicit operator bool() const { return N() != 0; }

 bool IsCoast(size_t i) const { return U(i) == fillval || V(i) == fillval; }
 bool IsCoast(size_t ix, size_t iy) const { return U(ix, iy) == fillval || V(ix, iy) == fillval; }

 bool IsFill(size_t i) const { return IsCoast(i); }
 bool IsFill(size_t ix, size_t iy) const { return IsCoast(ix, iy); }

 static real Fillval() { return fillval; }

 const MString& Unit() const { return unit; }

 const MString DUnit() const
 {
  switch(metric)
  {
  case(Metric::EUCLID): return "";
  case(Metric::SPHERIC): return "km";
  }
  return "";
 }
};

template<class OneVarData, bool isgrid> class UVDataDims;

template<class OneVarData> class UVDataDims<OneVarData, true>: public UVDataStorage
{
 real x0 = 0.0, y0 = 0.0;
 real xstep = 0.0, ystep = 0.0;

 void Resize(size_t sz)
 {
  UVDataStorage::u.resize(sz);
  UVDataStorage::v.resize(sz);
  UVDataStorage::u2.resize(sz);
 }

 protected:
 UVDataDims() = default;
 UVDataDims(const OneVarData& u, const OneVarData& v, Metric m = Metric::SPHERIC)
 {
  if(!(u && v))
  {
   *this = UVDataDims();
   return;
  }
  x0     = u.Lon(0);
  y0     = u.Lat(0);
  xstep  = u.XStep();
  ystep  = u.YStep();
  metric = m;
  nx     = u.Nx();
  ny     = u.Ny();

  Resize(u.N());
  for(size_t i = 0; i < u.N(); i++)
  {
   if(u(i) == u.Fillval() || v(i) == v.Fillval())
    U(i) = V(i) = U2(i) = Fillval();
   else
   {
    U(i)  = u(i);
    V(i)  = v(i);
    U2(i) = u(i) * u(i) + v(i) * v(i);
   }
  }
  SetUnit(u.Unit());
 }

 public:
 real Lon(size_t ix, [[maybe_unused]] size_t iy) const { return x0 + ix * xstep; }
 real Lat([[maybe_unused]] size_t ix, size_t iy) const { return y0 + iy * ystep; }
 real Lon(size_t i) const { return Lon(i % nx, i / nx); }
 real Lat(size_t i) const { return Lat(i % nx, i / nx); }

 real XStep() const { return xstep; }
 real YStep() const { return ystep; }

 real Ix2Lon(size_t ix) const { return x0 + ix * xstep; }
 real Iy2Lat(size_t iy) const { return y0 + iy * ystep; }
};

template<class OneVarData> class UVDataDims<OneVarData, false>: public UVDataStorage
{
 std::vector<real> lon, lat;

 void Resize(size_t sz)
 {
  UVDataStorage::u.resize(sz);
  UVDataStorage::v.resize(sz);
  UVDataStorage::u2.resize(sz);
  lon.resize(sz);
  lat.resize(sz);
 }

 protected:
 UVDataDims() = default;
 UVDataDims(const OneVarData& u, const OneVarData& v, Metric m = Metric::SPHERIC)
 {
  if(!(u && v))
  {
   *this = UVDataDims();
   return;
  }
  metric = m;
  nx     = u.Nx();
  ny     = u.Ny();

  Resize(u.N());
  for(size_t i = 0; i < u.N(); i++)
  {
   lon[i] = u.Lon(i);
   lat[i] = u.Lat(i);
   if(u(i) == u.Fillval() || v(i) == v.Fillval())
    U(i) = V(i) = U2(i) = Fillval();
   else
   {
    U(i)  = u(i);
    V(i)  = v(i);
    U2(i) = u(i) * u(i) + v(i) * v(i);
   }
  }
 }

 public:
 real Lon(size_t ix, [[maybe_unused]] size_t iy) const { return Lon(nx * iy + ix); }
 real Lat([[maybe_unused]] size_t ix, size_t iy) const { return Lat(nx * iy + ix); }
 real Lon(size_t i) const { return lon[i]; }
 real Lat(size_t i) const { return lat[i]; }
};

template<class OneVarData> class UVData: public UVDataDims<OneVarData, internal::HaveXYStep<OneVarData>>
{
 using B = UVDataDims<OneVarData, internal::HaveXYStep<OneVarData>>;
 using B::Ud, B::Vd, B::nx, B::ny, B::D;

 // For stationary points. Works only for meridian-parallel grids.
 real LonR(size_t i, size_t j, real x, [[maybe_unused]] real y) const { return B::Lon(i, j) + x * (B::Lon(i + 1, j) - B::Lon(i, j)); }
 real LatR(size_t i, size_t j, [[maybe_unused]] real x, real y) const { return B::Lat(i, j) + y * (B::Lat(i, j + 1) - B::Lat(i, j)); }

 real dUdX(size_t ix, size_t iy) const
 {
  if(IsCoast(ix, iy)) return Fillval();
  if(ix == 0) return (Ud(1, iy) - Ud(0, iy)) / D(Lon(0, iy), Lat(0, iy), Lon(1, iy), Lat(1, iy));
  if(ix == nx - 1) return (Ud(nx - 1, iy) - Ud(nx - 2, iy)) / D(Lon(nx - 2, iy), Lat(nx - 2, iy), Lon(nx - 1, iy), Lat(nx - 1, iy));
  return 0.5 * ((Ud(ix + 1, iy) - Ud(ix, iy)) / D(Lon(ix, iy), Lat(ix, iy), Lon(ix + 1, iy), Lat(ix + 1, iy)) +
                (Ud(ix, iy) - Ud(ix - 1, iy)) / D(Lon(ix - 1, iy), Lat(ix - 1, iy), Lon(ix, iy), Lat(ix, iy)));
 }
 real dVdX(size_t ix, size_t iy) const
 {
  if(IsCoast(ix, iy)) return Fillval();
  if(ix == 0) return (Vd(1, iy) - Vd(0, iy)) / D(Lon(0, iy), Lat(0, iy), Lon(1, iy), Lat(1, iy));
  if(ix == nx - 1) return (Vd(nx - 1, iy) - Vd(nx - 2, iy)) / D(Lon(nx - 2, iy), Lat(nx - 2, iy), Lon(nx - 1, iy), Lat(nx - 1, iy));
  return 0.5 * ((Vd(ix + 1, iy) - Vd(ix, iy)) / D(Lon(ix, iy), Lat(ix, iy), Lon(ix + 1, iy), Lat(ix + 1, iy)) +
                (Vd(ix, iy) - Vd(ix - 1, iy)) / D(Lon(ix - 1, iy), Lat(ix - 1, iy), Lon(ix, iy), Lat(ix, iy)));
 }
 real dUdY(size_t ix, size_t iy) const
 {
  if(IsCoast(ix, iy)) return Fillval();
  if(iy == 0) return (Ud(ix, 1) - Ud(ix, 0)) / D(Lon(ix, 0), Lat(ix, 0), Lon(ix, 1), Lat(ix, 1));
  if(iy == ny - 1) return (Ud(ix, ny - 1) - Ud(ix, ny - 2)) / D(Lon(ix, ny - 2), Lat(ix, ny - 2), Lon(ix, ny - 1), Lat(ix, ny - 1));
  return 0.5 * ((Ud(ix, iy + 1) - Ud(ix, iy)) / D(Lon(ix, iy), Lat(ix, iy), Lon(ix, iy + 1), Lat(ix, iy + 1)) +
                (Ud(ix, iy) - Ud(ix, iy - 1)) / D(Lon(ix, iy - 1), Lat(ix, iy - 1), Lon(ix, iy), Lat(ix, iy)));
 }
 real dVdY(size_t ix, size_t iy) const
 {
  if(IsCoast(ix, iy)) return Fillval();
  if(iy == 0) return (Vd(ix, 1) - Vd(ix, 0)) / D(Lon(ix, 0), Lat(ix, 0), Lon(ix, 1), Lat(ix, 1));
  if(iy == ny - 1) return (Vd(ix, ny - 1) - Vd(ix, ny - 2)) / D(Lon(ix, ny - 2), Lat(ix, ny - 2), Lon(ix, ny - 1), Lat(ix, ny - 1));
  return 0.5 * ((Vd(ix, iy + 1) - Vd(ix, iy)) / D(Lon(ix, iy), Lat(ix, iy), Lon(ix, iy + 1), Lat(ix, iy + 1)) +
                (Vd(ix, iy) - Vd(ix, iy - 1)) / D(Lon(ix, iy - 1), Lat(ix, iy - 1), Lon(ix, iy), Lat(ix, iy)));
 }

 public:
 using B::Fillval, B::IsCoast, B::Lon, B::Lat, B::U, B::V, B::Nx, B::Ny;

 UVData() = default;
 UVData(const OneVarData& u, const OneVarData& v, Metric m = Metric::SPHERIC): B(u, v, m) {}

 real Div(size_t i) const { return Div(i % nx, i / B::nx); }
 real Rot(size_t i) const { return Rot(i % nx, i / B::nx); }
 real Div(size_t ix, size_t iy) const
 {
  if(!*this) return 0.0;
  real ux = dUdX(ix, iy);
  real vy = dVdY(ix, iy);
  return (ux == Fillval() || vy == Fillval()) ? Fillval() : (ux + vy);
 }
 real Rot(size_t ix, size_t iy) const
 {
  if(!*this) return 0.0;
  real vx = dVdX(ix, iy);
  real uy = dUdY(ix, iy);
  return (vx == Fillval() || uy == Fillval()) ? Fillval() : (vx - uy);
 }
 // Okubo-Weiss parameter
 real OW(size_t i) const { return OW(i % nx, i / nx); }
 real OW(size_t ix, size_t iy) const
 {
  if(!*this) return 0.0;
  real ux = dUdX(ix, iy);
  real uy = dUdY(ix, iy);
  real vx = dVdX(ix, iy);
  real vy = dVdY(ix, iy);

  if(ux == Fillval() || uy == Fillval() || vx == Fillval() || vy == Fillval()) return Fillval();

  real sn = ux - vy;
  real ss = vx + uy;
  real w  = vx - uy;
  return sn * sn + ss * ss - w * w;
 }

 auto StablePoints(size_t ix, size_t iy) const
 {
  std::vector<struct B::StPoint> points;
  if(!*this) return points;
  if(ix >= Nx() - 1 || iy >= Ny() - 1) return points;
  if(IsCoast(ix, iy) || IsCoast(ix + 1, iy) || IsCoast(ix, iy + 1) || IsCoast(ix + 1, iy + 1)) return points;

  real au = U(ix, iy) - U(ix, iy + 1) - U(ix + 1, iy) + U(ix + 1, iy + 1);
  real bu = U(ix, iy + 1) - U(ix, iy);
  real cu = U(ix + 1, iy) - U(ix, iy);
  real du = U(ix, iy);

  real av = V(ix, iy) - V(ix, iy + 1) - V(ix + 1, iy) + V(ix + 1, iy + 1);
  real bv = V(ix, iy + 1) - V(ix, iy);
  real cv = V(ix + 1, iy) - V(ix, iy);
  real dv = V(ix, iy);

  real D   = au * bv - av * bu;
  real a   = (au * av * cu - au * au * cv) / D;
  real b   = (au * bv * cu - au * bu * cv + au * av * du - au * au * dv) / D;
  real c   = (au * bv * du - au * bu * dv) / D;
  real des = b * b - 4.0 * a * c;
  if(des < 0) return points;
  des = Sqrt(des);

  real p  = av * cu - au * cv;
  real q  = av * bu - au * bv;
  real r  = av * du - au * dv;
  real x1 = 0.5 * (-b + des) / a;
  real x2 = 0.5 * (-b - des) / a;
  real y1 = -(p * x1 + r) / q;
  real y2 = -(p * x2 + r) / q;

  auto PointType = [au = au, av = av, bu = bu, bv = bv, cu = cu, cv = cv](real x, real y) -> auto
  {
   real a   = 1.0;
   real b   = -y * au - x * av - bv - cu;
   real c   = y * (au * bv - av * bu) + x * (av * cu - au * cv) + bv * cu - bu * cv;
   real des = b * b - 4.0 * a * c;
   if(des > 0.0)
   {
    if(std::max(-b + Sqrt(des), -b - Sqrt(des)) > 0.0 && c > 0.0)
     return UVDataStorage::UKNOT;
    else if(std::max(-b + Sqrt(des), -b - Sqrt(des)) > 0.0)
     return UVDataStorage::SADDLE;
    else
     return UVDataStorage::SKNOT;
   }
   else
   {
    bool acyclcrit = (av * y + cv > 0.0);
    if(b < 0.0)
     return acyclcrit ? UVDataStorage::SACICFOCUS : UVDataStorage::SCICFOCUS;
    else
     return acyclcrit ? UVDataStorage::UACICFOCUS : UVDataStorage::UCICFOCUS;
   }
  };

  if(x1 >= 0.0 && x1 < 1.0 && y1 >= 0.0 && y1 < 1.0) points.emplace_back(LonR(ix, iy, x1, y1), LatR(ix, iy, x1, y1), PointType(x1, y1));
  if(x2 >= 0.0 && x2 < 1.0 && y2 >= 0.0 && y2 < 1.0) points.emplace_back(LonR(ix, iy, x2, y2), LatR(ix, iy, x2, y2), PointType(x2, y2));
  return points;
 }
};