#pragma once
#include "timing/time_transform.h"
#include <algorithm>
#include <cassert>
#include <optional>
#include <vector>

namespace time_transform {

// Map for warping beats to seconds, domain is beats, codomain is seconds,
// parameterized on a list of at least two warp markers {beat_i, second_i} where
// all beat_i are distinct. Fulfills conditions:
// - Beats between beat_i and beat_{i+1} are warped linearly between
//   second_i and second_{i+1}
// - Beats outside the range [beat_0, beat_{n-1}) are warped with linear
//   extrapolation based on the slope of the first or last interval
// - With fewer than two markers, the map assumes bps = 1.0
// - With one marker, the map is a translation
// - With zero markers, the map is the identity map
template <TimeValue T = double> class WarpMap {
public:
  using DomainUnit = BeatsTag;
  using CodomainUnit = SecondsTag;
  using DomainTime = Beats<T>;
  using CodomainTime = Seconds<T>;
  using DomainDelta = BeatsDelta<T>;
  using CodomainDelta = SecondsDelta<T>;

  struct WarpMarker {
    DomainTime beat;
    CodomainTime second;

    bool operator<(const WarpMarker &other) const { return beat < other.beat; }
    // Used by std::unique to remove markers with identical domain times
    bool operator==(const WarpMarker &other) const {
      return beat == other.beat;
    }
  };

private:
  std::vector<WarpMarker> markers_;
  static constexpr T EPSILON = time_units::EPSILON;

  // Helper to calculate slope (seconds per beat)
  // Assumes m1.beat and m2.beat are distinct, which should be guaranteed
  // by the constructor logic for adjacent markers.
  static T calculate_slope(const WarpMarker &m1, const WarpMarker &m2) {
    DomainDelta domain_diff = m2.beat - m1.beat;
    CodomainDelta codomain_diff = m2.second - m1.second;

    // domain_diff.raw() should not be zero if markers are distinct (checked in
    // constructor for >=2 markers case)
    if (std::abs(domain_diff.raw()) < EPSILON) {
      if (std::abs(codomain_diff.raw()) < EPSILON) {
        return 0.0;
      } else {
        return codomain_diff.raw() > 0 ? std::numeric_limits<T>::infinity()
                                       : -std::numeric_limits<T>::infinity();
      }
    }
    return codomain_diff.raw() / domain_diff.raw();
  }

  int get_map_end_idx() const {
    size_t num_actual_markers = markers_.size();

    // 2 effective segments for identity or translation
    if (num_actual_markers < 2)
      return 2;

    // N+1 effective segments for N>=2 markers
    return static_cast<int>(num_actual_markers) + 1;
  }

public:
  explicit WarpMap(std::vector<WarpMarker> &&markers)
      : markers_(std::move(markers)) {
    // Sort and unique markers if there are any to process.
    std::sort(markers_.begin(), markers_.end());
    markers_.erase(std::unique(markers_.begin(), markers_.end()),
                   markers_.end());
  }

  WarpMap(std::initializer_list<WarpMarker> markers)
      : WarpMap(std::vector<WarpMarker>(markers)) {}

  class iterator {
  public:
    using iterator_category = std::bidirectional_iterator_tag;
    using value_type = MappedSegment<DomainUnit, CodomainUnit, T>;
    using difference_type = std::ptrdiff_t;
    using pointer = const value_type *;
    using reference = const value_type &;

  private:
    const WarpMap *map_ptr_;
    int effective_segment_idx_;
    mutable std::optional<value_type> current_mapped_segment_;

    int get_iterator_end_idx() const {
      if (!map_ptr_)
        return -1;
      return map_ptr_->get_map_end_idx();
    }

    void cache_current() const {
      if (current_mapped_segment_ || !map_ptr_)
        return;

      if (effective_segment_idx_ == get_iterator_end_idx()) {
        // This is the end iterator state, no segment to cache.
        return;
      }

      const auto &markers = map_ptr_->markers_;
      size_t num_actual_markers = markers.size();

      DomainTime source_start_time(0.0);
      DomainTime source_end_time(0.0);
      CodomainTime target_start_time(0.0);
      CodomainTime target_end_time(0.0);
      T slope_val = 1.0; // Default for 0/1 marker cases

      if (num_actual_markers == 0) {
        // Identity map
        assert(effective_segment_idx_ == 0 || effective_segment_idx_ == 1);
        slope_val = 1.0;
        if (effective_segment_idx_ == 0) {
          // Segment: (-inf, 0)
          source_start_time = DomainTime::neg_inf();
          source_end_time = DomainTime(0.0);
          target_start_time = CodomainTime::neg_inf();
          target_end_time = CodomainTime(0.0);
        } else {
          // Segment: [0, +inf)
          source_start_time = DomainTime(0.0);
          source_end_time = DomainTime::inf();
          target_start_time = CodomainTime(0.0);
          target_end_time = CodomainTime::inf();
        }
      } else if (num_actual_markers == 1) {
        // Translation map
        assert(effective_segment_idx_ == 0 || effective_segment_idx_ == 1);
        const auto &marker = markers[0];
        slope_val = 1.0;
        if (effective_segment_idx_ == 0) {
          // Segment: (-inf, marker.beat)
          source_start_time = DomainTime::neg_inf();
          source_end_time = marker.beat;
          target_start_time = CodomainTime::neg_inf();
          target_end_time = marker.second;
        } else {
          // Segment: [marker.beat, +inf)
          source_start_time = marker.beat;
          source_end_time = DomainTime::inf();
          target_start_time = marker.second;
          target_end_time = CodomainTime::inf();
        }
      } else {
        // Piecewise linear map with linear extrapolation
        assert(effective_segment_idx_ >= 0 &&
               effective_segment_idx_ <= static_cast<int>(num_actual_markers));

        if (effective_segment_idx_ == 0) {
          // Before first marker
          const auto &m0 = markers[0];
          const auto &m1 = markers[1];
          slope_val = WarpMap::calculate_slope(m0, m1);
          source_start_time = DomainTime::neg_inf();
          source_end_time = m0.beat;
          target_start_time = CodomainTime::neg_inf();
          target_end_time = m0.second;
        } else if (effective_segment_idx_ <
                   static_cast<int>(num_actual_markers)) {
          // Between markers
          const auto &prev_marker = markers[effective_segment_idx_ - 1];
          const auto &curr_marker = markers[effective_segment_idx_];
          slope_val = WarpMap::calculate_slope(prev_marker, curr_marker);
          source_start_time = prev_marker.beat;
          source_end_time = curr_marker.beat;
          target_start_time = prev_marker.second;
          target_end_time = curr_marker.second;
        } else {
          // After last marker (idx == num_actual_markers)
          const auto &m_last = markers.back();
          const auto &m_prev = markers[num_actual_markers - 2];
          slope_val = WarpMap::calculate_slope(m_prev, m_last);
          source_start_time = m_last.beat;
          source_end_time = DomainTime::inf();
          target_start_time = m_last.second;
          target_end_time = CodomainTime::inf();
        }
      }

      TimeRange<DomainUnit, T> src_r = {source_start_time, source_end_time};
      TimeRange<CodomainUnit, T> tgt_r = {target_start_time, target_end_time};
      current_mapped_segment_.emplace(src_r, tgt_r, slope_val,
                                      SegmentMarks::NONE);
    }

  public:
    iterator(const WarpMap *map, int idx)
        : map_ptr_(map), effective_segment_idx_(idx) {}

    iterator() : map_ptr_(nullptr), effective_segment_idx_(-1) {}

    reference operator*() const {
      cache_current();
      assert(current_mapped_segment_.has_value() &&
             "Dereferencing invalid or end WarpMap iterator");
      return *current_mapped_segment_;
    }
    pointer operator->() const {
      cache_current();
      assert(current_mapped_segment_.has_value() &&
             "Dereferencing invalid or end WarpMap iterator");
      return &(*current_mapped_segment_);
    }

    iterator &operator++() {
      if (!map_ptr_)
        return *this; // Default-constructed iterator
      current_mapped_segment_.reset();
      effective_segment_idx_++;
      int end_idx_for_this_map = get_iterator_end_idx();
      if (effective_segment_idx_ >= end_idx_for_this_map) {
        effective_segment_idx_ = end_idx_for_this_map; // Clamp to end state
      }
      return *this;
    }

    iterator operator++(int) {
      iterator tmp = *this;
      ++(*this);
      return tmp;
    }

    iterator &operator--() {
      if (!map_ptr_)
        return *this; // Default-constructed iterator
      current_mapped_segment_.reset();
      if (effective_segment_idx_ > 0) {
        effective_segment_idx_--;
      }
      return *this;
    }

    iterator operator--(int) {
      iterator tmp = *this;
      --(*this);
      return tmp;
    }

    bool operator==(const iterator &other) const {
      return map_ptr_ == other.map_ptr_ &&
             effective_segment_idx_ == other.effective_segment_idx_;
    }

    bool operator!=(const iterator &other) const { return !(*this == other); }
  };

  iterator begin() const { return iterator(this, 0); }
  iterator begin() { return iterator(this, 0); }

  iterator end() const { return iterator(this, get_map_end_idx()); }
  iterator end() { return iterator(this, get_map_end_idx()); }

  iterator getSegmentIteratorAt(const DomainTime &p) const {
    const size_t num_actual_markers = markers_.size();

    if (num_actual_markers == 0) {
      // Identity map, anchor at 0.0
      bool is_before_split =
          (p == DomainTime::neg_inf()) || (p.is_finite() && p.raw() < 0.0);
      return iterator(this, is_before_split ? 0 : 1);
    } else if (num_actual_markers == 1) {
      // Translation, anchor at marker
      const auto &marker = markers_[0];
      bool is_before_split = (p == DomainTime::neg_inf()) || (p < marker.beat);
      return iterator(this, is_before_split ? 0 : 1);
    } else {
      // Piecewise linear
      if (p == DomainTime::neg_inf() || p < markers_[0].beat) {
        return iterator(this, 0);
      }
      // markers_.back() is markers_[num_actual_markers - 1]
      if (p == DomainTime::inf() || p >= markers_.back().beat) {
        return iterator(this, static_cast<int>(num_actual_markers));
      }

      // p is between markers_[0].beat (inclusive) and markers_.back().beat
      // (exclusive)
      auto it_upper = std::upper_bound(
          markers_.begin(), markers_.end(), p,
          [](const DomainTime &val, const WarpMarker &marker_in_vec) {
            return val < marker_in_vec.beat;
          });

      assert(it_upper != markers_.begin() &&
             "Logic error: p < markers[0].beat should be handled earlier");
      assert(it_upper != markers_.end() &&
             "Logic error: p >= markers.back().beat should be handled earlier");

      int marker_idx_for_segment_start =
          static_cast<int>(std::distance(markers_.begin(), it_upper - 1));
      return iterator(this, marker_idx_for_segment_start + 1);
    }
  }
}; // End WarpMap class

static_assert(IsTimeTransformer<WarpMap<double>>,
              "WarpMap does not satisfy the IsTimeTransformer concept.");

} // namespace time_transform