module;

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

export module routemon:geo.utm.multizonal;

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

namespace routemon::geo::utm::multizonal {

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

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

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

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

class from_wgs84_transform : public from_wgs84_transform_base
{
  zone to_zone_;

public:
  explicit from_wgs84_transform(zone to_zone)
    : from_wgs84_transform_base{{}, bgeo::srs::epsg{to_zone.wgs84_proj_epsg()}},
      to_zone_{to_zone}
  {
  }

  from_wgs84_transform() : from_wgs84_transform{zone::min()} {}

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

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

public:
  static from_wgs84_transforms const& instance()
  {
    static from_wgs84_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);
}

// Distribute a WGS 84 linestring over the UTM zones that it finds
// itself in or near to.
//
// Since it is not at all unthinkable that some road works will end up
// crossing different UTM zones, we need some machinery to work with
// such situations. Since we want to report all situations within a
// specified distance of the planned route, we must also consider the
// situation where this is indeed the case, but the situation is
// situated at the opposite side of a delineating UTM zone meridian
// w.r.t. the planned route.
//
// To solve this issue, we create segments of the input line string
// for multiple UTM zones when crossing zone boundaries or when very
// close to zone boundaries.
auto distribute_linestring_over_zones(utm::zonable_wgs84_linestring const& ls)
    -> multi_zone<std::vector<utm::zone_local::prim::linestring>>
{
  auto result = multi_zone<std::vector<utm::zone_local::prim::linestring>>{};

  auto to_utm = from_wgs84_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)
  {
    result[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] = 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);
  }

  return result;
}

// Split a WGS 84 line string into a sequence of UTM-zone-local line
// strings, starting a new UTM-zone-local line string when the input
// line string crosses a UTM zone boundary. The line segment that
// crosses the zone boundary can be found in the UTM-zone-local line
// strings for both zones which its points are in (we assume that line
// segments will be short enough to not cause significant distortion
// here, so we do not put in the effort to clip at the
// zone-delineating meridian here).
auto split_linestring_across_zones(utm::zonable_wgs84_linestring ls)
    -> std::vector<utm::zone_local::linestring>
{
  auto splits = std::vector<utm::zone_local::linestring>{};

  auto to_utm = from_wgs84_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));
      splits.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())
  {
    splits.push_back(std::move(*mworking_seg));
  }

  return splits;
}

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 = distribute_linestring_over_zones(ls);
    for (auto z = utm::zone::min(); z != utm::zone::max(); z = z.next())
    {
      for (auto ls : zone_segments[z])
      {
        auto 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(std::vector<utm::zone_local::linestring> const& lss) const
      -> std::generator<T const&>
  {
    for (auto const& ls : lss)
    {
      for (auto it = local_rtrees_[ls.zone].qbegin(bgeo::index::intersects(ls));
           it != local_rtrees_[ls.zone].qend(); it++)
      {
        using multi_linestring =
            bgeo::model::multi_linestring<utm::zone_local::prim::linestring>;

        auto relevant_ls_parts = multi_linestring{};
        bgeo::intersection(
            std::get<0>(*it),
            static_cast<utm::zone_local::prim::linestring const&>(ls),
            relevant_ls_parts);
        // it is somewhat awkward and arbitrary that this measurement is
        // performed here
        if (bgeo::distance(relevant_ls_parts, std::get<1>(*it)) < 30.0 /* m */)
        {
          co_yield std::get<2>(*it);
        }
      }
    }
  }
};

template <class T>
class point_rtree
{
public:
  using index_value =
      std::tuple<utm::zone_local::prim::box, utm::zone_local::prim::point, 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_point const& p, U&& arg) -> void
  {
    auto to_utm = from_wgs84_transforms::instance();

    auto p_zone = utm::zone::for_wgs84_point(p);
    auto mp_alt_zone = std::optional<utm::zone>{};
    if (auto [neighbor_zone, neighbor_zone_dist] = neighbor_utm_zone(p);
        neighbor_zone_dist < 30.0 /* m */)
      mp_alt_zone = neighbor_zone;

    auto insert = [&](utm::zone z) -> void
    {
      auto p_utm = to_utm[z].apply(p);
      auto bpe = bgeo::return_buffer<utm::zone_local::prim::box>(
          bgeo::return_envelope<utm::zone_local::prim::box>(p_utm),
          30.0 /* m */);
      local_rtrees_[z].insert(index_value{bpe, p_utm, arg});
    };

    insert(p_zone);
    if (mp_alt_zone)
      insert(*mp_alt_zone);
  }

  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(std::vector<utm::zone_local::linestring> const& lss) const
      -> std::generator<T const&>
  {
    for (auto const& ls : lss)
    {
      for (auto it = local_rtrees_[ls.zone].qbegin(bgeo::index::intersects(ls));
           it != local_rtrees_[ls.zone].qend(); it++)
      {
        using multi_linestring =
            bgeo::model::multi_linestring<utm::zone_local::prim::linestring>;

        auto relevant_ls_parts = multi_linestring{};
        bgeo::intersection(
            std::get<0>(*it),
            static_cast<utm::zone_local::prim::linestring const&>(ls),
            relevant_ls_parts);
        // it is somewhat awkward and arbitrary that this measurement is
        // performed here
        if (bgeo::distance(relevant_ls_parts, std::get<1>(*it)) < 30.0 /* m */)
        {
          co_yield std::get<2>(*it);
        }
      }
    }
  }
};

} // namespace routemon::geo::utm::multizonal