module;

#include <boost/geometry.hpp>
#include <boost/geometry/srs/epsg.hpp>
#include <boost/geometry/srs/transformation.hpp>

export module routemon:geo.multizonal;

import std;
import :geo;
import :geo.utm;
import :geo.utm.zone_local;
import :geo.wgs84;

namespace routemon::geo::multizonal {

template<std::default_initializable T>
class multi_zone
{
  std::array<T, utm::zone::max().as_index() + 1> zones_;

public:
  auto operator[](utm::zone z) -> T&
  {
    return zones_[z.as_index()];
  }

  auto operator[](utm::zone z) const -> T const&
  {
    return zones_[z.as_index()];
  }
};

namespace wgs84
{

// Transformations from WGS 84 (EPSG:4326)
using transform_from_t = bgeo::srs::transformation<bgeo::srs::static_epsg<4326>>;

class utm_transform : public transform_from_t
{
  utm::zone to_zone_;

public:
  explicit utm_transform(utm::zone to_zone)
    : transform_from_t{{}, bgeo::srs::epsg{to_zone.wgs84_proj_epsg()}},
      to_zone_{to_zone}
  {
  }

  utm_transform()
    : utm_transform{utm::zone::min()}
  {}

  auto apply(utm::zonable_wgs84_point p) -> utm::zone_local::point
  {
    auto local_p = utm::zone_local::point{to_zone_, 0.0, 0.0};
    forward(p, local_p);
    return local_p;
  }
};

class utm_transforms : public multi_zone<utm_transform>
{
  utm_transforms()
  {
    for (auto z = utm::zone::min(); z != utm::zone::max(); z = z.next())
    {
      (*this)[z] = utm_transform{z};
    }
  }

public:
  static utm_transforms const& instance()
  {
    static utm_transforms inst;
    return inst;
  }
};

auto neighbor_utm_zone(utm::zonable_wgs84_point p) -> std::pair<utm::zone, double>
{
  auto separating_meridian_lon = std::round(p.lon() / 6.0) * 6.0;
  auto closest_zone_middle = separating_meridian_lon < p.lon()
    ? separating_meridian_lon - 3.0
    : separating_meridian_lon + 3.0;
  auto closest_zone = utm::zone::for_wgs84_point(utm::zonable_wgs84_point::from(geo::wgs84::normalized_point{geo::wgs84::from_lat_lon(p.lat(), closest_zone_middle)}).value());

  auto separating_meridian_ls = geo::wgs84::linestring{
    geo::wgs84::from_lat_lon(90.0, 0.0),
    geo::wgs84::from_lat_lon(0.0, separating_meridian_lon),
    geo::wgs84::from_lat_lon(-90.0, 0.0),
  };
  auto closest_zone_dist = bgeo::distance(p, separating_meridian_ls, geo::wgs84::vincenty_strategy{});

  return std::make_pair(closest_zone, closest_zone_dist);
}

} // namespace wgs84

struct multizone_linestring
{
  multi_zone<std::vector<utm::zone_local::prim::linestring>> segments;

  // TODO: consider taking an input range instead
  explicit multizone_linestring(utm::zonable_wgs84_linestring const& ls)
  {
    auto to_utm = wgs84::utm_transforms::instance();
    auto working = multi_zone<utm::zone_local::prim::linestring>{};
    auto mprev_p_zone = std::optional<utm::zone>{};
    auto mprev_p_alt_zone = std::optional<utm::zone>{};

    auto push = [&](utm::zonable_wgs84_point p, utm::zone z)
    { working[z].push_back(to_utm[z].apply(p)); };
    auto flush = [&](utm::zone z)
    {
      segments[z].push_back(std::move(working[z]));
      working[z] = {};
    };

    for (auto const& p : ls)
    {
      auto p_zone = utm::zone::for_wgs84_point(p);
      auto mp_alt_zone = std::optional<utm::zone>{};
      if (auto [neighbor_zone, neighbor_zone_dist] = wgs84::neighbor_utm_zone(p);
          neighbor_zone_dist < 30.0 /* m */)
        mp_alt_zone = neighbor_zone;
      assert(!mp_alt_zone || *mp_alt_zone != p_zone);

      push(p, p_zone);
      if (mp_alt_zone)
        push(p, *mp_alt_zone);
      if (mprev_p_zone && *mprev_p_zone != p_zone && mprev_p_zone != mp_alt_zone)
      {
        push(p, *mprev_p_zone);
        flush(*mprev_p_zone);
      }
      if (mprev_p_alt_zone && *mprev_p_alt_zone != p_zone && mprev_p_alt_zone != mp_alt_zone)
      {
        push(p, *mprev_p_alt_zone);
        flush(*mprev_p_alt_zone);
      }

      mprev_p_zone = p_zone;
      mprev_p_alt_zone = mp_alt_zone;
    }

    if (mprev_p_zone && !working[*mprev_p_zone].empty())
    {
      flush(*mprev_p_zone);
    }
    if (mprev_p_alt_zone && !working[*mprev_p_alt_zone].empty())
    {
      flush(*mprev_p_alt_zone);
    }
  }
};

struct zoned_linestring_seg_seq
{
  std::vector<utm::zone_local::linestring> segments;

  explicit zoned_linestring_seg_seq(utm::zonable_wgs84_linestring ls)
  {
    auto to_utm = wgs84::utm_transforms::instance();
    auto mworking_seg = std::optional<utm::zone_local::linestring>{};

    for (auto const& p : ls)
    {
      auto p_zone = utm::zone::for_wgs84_point(p);
      if (mworking_seg && mworking_seg->zone != p_zone)
      {
        mworking_seg->push_back(to_utm[mworking_seg->zone].apply(p));
        segments.push_back(std::move(*mworking_seg));
        mworking_seg = std::nullopt;
      }
      if (!mworking_seg)
        mworking_seg = utm::zone_local::linestring{p_zone};
      mworking_seg->push_back(to_utm[p_zone].apply(p));
    }

    if (mworking_seg && !mworking_seg->empty())
    {
      segments.push_back(std::move(*mworking_seg));
    }
  }
};

template<class T>
class linestring_rtree
{
public:
  using index_value = std::tuple<utm::zone_local::prim::box, utm::zone_local::prim::linestring, T>;

private:
  multi_zone<bgeo::index::rtree<index_value, bgeo::index::quadratic<16>>> local_rtrees_;

public:
  template<std::convertible_to<T> U>
  auto insert(utm::zonable_wgs84_linestring const& ls, U&& arg) -> void
  {
    auto zone_segments = multizone_linestring{ls}.segments;
    for (auto z = utm::zone::min(); z != utm::zone::max(); z = z.next())
    {
      for (auto ls : zone_segments[z])
      {
        utm::zone_local::prim::box blse =
          bgeo::return_buffer<utm::zone_local::prim::box>(bgeo::return_envelope<utm::zone_local::prim::box>(ls), 30.0 /* m */);
        local_rtrees_[z].insert(index_value{blse, std::move(ls), std::forward<U>(arg)});
      }
    }
  }

  auto size() const -> std::size_t
  {
    auto total_size = 0uz;
    for (auto z = utm::zone::min(); z != utm::zone::max(); z = z.next())
      total_size += local_rtrees_[z].size();
    return total_size;
  }

  // Note: the same T may be generated more than once!
  auto intersection(zoned_linestring_seg_seq const& lss) const -> std::generator<T const&>
  {
    for (auto const& ls : lss.segments)
    {
      // auto ls_box = bgeo::return_envelope<utm::zone_local::prim::box>(ls);
      for (auto it = local_rtrees_[ls.zone].qbegin(bgeo::index::intersects(ls));
           it != local_rtrees_[ls.zone].qend(); it++)
      {
        if (bgeo::intersects(ls, std::get<1>(*it)))
        {
          co_yield std::get<2>(*it);
        }
      }
    }
  }
};

} // namespace routemon::geo::multizonal