#pragma once

#include "../taskflow.hpp"

namespace tf {

// ----------------------------------------------------------------------------
// make_merge_task
// ----------------------------------------------------------------------------

/**
@brief creates a parallel merge task over two sorted ranges

@tparam B1 iterator type for the beginning of the first range  (satisfying tf::InputIteratorLike)  
@tparam E1 iterator type for the end of the first range        (satisfying tf::InputIteratorLike)  
@tparam B2 iterator type for the beginning of the second range (satisfying tf::InputIteratorLike) 
@tparam E2 iterator type for the end of the second range       (satisfying tf::InputIteratorLike) 
@tparam O  output iterator type
@tparam C  comparator type

@param first1  iterator to the beginning of the first sorted range
@param last1   iterator to the end of the first sorted range
@param first2  iterator to the beginning of the second sorted range
@param last2   iterator to the end of the second sorted range
@param cmp     comparator defining the sort order
@param d_first iterator to the beginning of the output range

Returns a callable suitable for @c tf::Executor::async,
@c tf::Executor::silent_async, or @c tf::Taskflow::emplace.

@par Algorithm

The algorithm parallelizes @c std::merge using the @em co-rank technique,
which avoids any inter-worker synchronization during the merge itself.

<b>Step 1 — Flat output partitioning.</b>
The merged output has @c N = n + m elements (where @c n and @c m are the
sizes of the two input ranges).  The output is divided into @c W equal
contiguous chunks, one per worker thread.  Worker @c w is responsible for
writing output positions <tt>[w*K, (w+1)*K)</tt> where @c K = N/W.

<b>Step 2 — Co-rank: finding the input sub-ranges.</b>
For each output chunk boundary @c rank, we need to know: "among the first
@c rank elements of the merged output, how many came from seq1 (call it
@c i) and how many from seq2 (call it @c j = rank - i)?"

Once @c (i, j) is known, the worker's input from seq1 is
<tt>seq1[i_start .. i_end)</tt> and from seq2 is
<tt>seq2[j_start .. j_end)</tt> — both independent slices that can be
merged in isolation.

Co-rank finds @c i by binary search over the range
<tt>[max(0, rank-m), min(n, rank)]</tt>.  For a candidate @c i with
@c j = rank - i, the partition is correct when:

@code
  seq1[i-1] <= seq2[j]    and    seq2[j-1] <= seq1[i]
@endcode

i.e. the last element of each slice does not exceed the first unused element
of the other sequence.  Because both sequences are sorted, only one of these
two conditions needs to be checked — as @c i increases, the correctness
condition transitions monotonically from "too small" to "correct", making
binary search applicable.

The special case @c j==0 (all @c rank elements come from seq1, none from
seq2) skips the @c seq2[j-1] dereference entirely, which would otherwise be
an out-of-bounds access.

<b>Step 3 — Independent parallel merge.</b>
Each worker calls @c std::merge on its identified input sub-slices and writes
directly to its non-overlapping portion of the output.  No mutex, no atomic
cursor, no synchronization.

<b>Why no partitioner?</b>
Unlike @c for_each or @c find, where dynamic/guided scheduling can adapt to
unequal per-element costs, @c std::merge on a chunk of size @c K always costs
exactly O(K) regardless of data values.  There is no load imbalance to
mitigate.  Static equal partitioning (W chunks, one per worker) is always
optimal.  Smaller chunks would only add unnecessary @c co_rank overhead
(O(log N) per boundary) without any benefit.

<b>Complexity.</b>
Each worker performs 2 co-rank calls at O(log N) each, plus one
@c std::merge at O(N/W).  Total: O(N + W log N), which is effectively O(N)
for typical W << N.

@note Both input ranges must be sorted with respect to @c cmp.
      The output range must not overlap either input range.
      Both input iterators must be random-access iterators.
*/
template <InputIteratorLike B1, InputIteratorLike E1, InputIteratorLike B2, InputIteratorLike E2, typename O, typename C>
auto make_merge_task(B1 first1, E1 last1, B2 first2, E2 last2, C cmp, O d_first) {

  using B1_t = std::decay_t<std::unwrap_ref_decay_t<B1>>;
  using E1_t = std::decay_t<std::unwrap_ref_decay_t<E1>>;
  using B2_t = std::decay_t<std::unwrap_ref_decay_t<B2>>;
  using E2_t = std::decay_t<std::unwrap_ref_decay_t<E2>>;
  using O_t  = std::decay_t<std::unwrap_ref_decay_t<O>>;

  return [=](Runtime& rt) mutable {

    B1_t beg1  = first1;
    E1_t end1  = last1;
    B2_t beg2  = first2;
    E2_t end2  = last2;
    O_t  d_beg = d_first;

    size_t n = static_cast<size_t>(std::distance(beg1, end1));
    size_t m = static_cast<size_t>(std::distance(beg2, end2));
    size_t N = n + m;
    size_t W = rt.executor().num_workers();

    if(N == 0) {
      return;
    }

    // sequential fallback for single worker or small N
    if(W <= 1 || N <= W) {
      std::merge(beg1, end1, beg2, end2, d_beg, cmp);
      return;
    }

    // co_rank: binary search for the unique split point i in [low, high)
    // such that taking i elements from seq1 and (rank-i) from seq2 gives
    // exactly the rank smallest merged elements.
    //
    // Invariant: seq2[j-1] <= seq1[i] means i is not too large (high = i),
    //            otherwise i is too small (low = i+1).
    //
    // Guard: when j==0, skip the seq2[j-1] dereference (would be seq2[-1])
    // and always increase low — if seq2 contributes nothing, i must grow.
    //
    // Captured by value so the lambda is safe to copy into async tasks.
    auto co_rank = [=](size_t rank) -> std::pair<size_t, size_t> {
      size_t low  = (rank > m) ? rank - m : 0;
      size_t high = (std::min)(n, rank);

      while(low < high) {
        size_t i = low + (high - low) / 2;
        size_t j = rank - i;

        if(j == 0 || !cmp(*std::next(beg2, j - 1), *std::next(beg1, i))) {
          low = i + 1;  // seq1[i] <= seq2[j-1]: i is too small, take more from seq1
        } else {
          high = i;     // seq1[i] >  seq2[j-1]: i is too large
        }
      }
      return {low, rank - low};
    };

    // merge_chunk: given an output range [part_b, part_e), use co_rank to
    // find the corresponding input sub-ranges and merge them independently.
    // Captured by value so it is safe to copy into rt.silent_async tasks —
    // all iterators and the co_rank callable are value-copied, no dangling
    // references to the outer stack frame.
    auto merge_chunk = [=](size_t part_b, size_t part_e) {
      auto [i_beg, j_beg] = co_rank(part_b);
      auto [i_end, j_end] = co_rank(part_e);
      std::merge(
        std::next(beg1, i_beg), std::next(beg1, i_end),
        std::next(beg2, j_beg), std::next(beg2, j_end),
        std::next(d_beg, part_b),
        cmp
      );
    };

    // static equal partition: W chunks, one per worker.
    // Each chunk boundary costs 2 co_rank calls O(log N) — smaller chunks
    // would add overhead with no load-balancing benefit since std::merge
    // always costs O(K) for a chunk of size K regardless of data.
    size_t chunk_size = (N + W - 1) / W;

    for(size_t w = 0; w < W; w++) {
      size_t curr_b = w * chunk_size;
      if(curr_b >= N) {
        break;
      }
      size_t curr_e = (std::min)(curr_b + chunk_size, N);
      if(w + 1 == W) {
        merge_chunk(curr_b, curr_e);  // last worker runs inline
      } else {
        rt.silent_async([=]() { merge_chunk(curr_b, curr_e); });
      }
    }
  };
}

// ----------------------------------------------------------------------------
// tf::Taskflow::merge
// ----------------------------------------------------------------------------

/**
@private
*/
template <InputIteratorLike B1, InputIteratorLike E1, InputIteratorLike B2, InputIteratorLike E2, typename O>
Task FlowBuilder::merge(B1 first1, E1 last1, B2 first2, E2 last2, O d_first) {
  return emplace(make_merge_task(
    first1, last1, first2, last2, std::less<>{}, d_first
  ));
}

/**
@private
*/
template <InputIteratorLike B1, InputIteratorLike E1, InputIteratorLike B2, InputIteratorLike E2,
          typename O, typename C>
Task FlowBuilder::merge(B1 first1, E1 last1, B2 first2, E2 last2, O d_first, C cmp) {
  return emplace(make_merge_task(
    first1, last1, first2, last2, cmp, d_first
  ));
}

}  // namespace tf
