//***************************************************************************
//* Copyright (c) 2015 Saint Petersburg State University
//* Copyright (c) 2011-2014 Saint Petersburg Academic University
//* All Rights Reserved
//* See file LICENSE for details.
//***************************************************************************

#ifndef __HAMMER_HKMER_DISTANCE_HPP__
#define __HAMMER_HKMER_DISTANCE_HPP__

#include "hkmer.hpp"
#include "err_helper_table.hpp"

namespace hammer {

enum IonEventType {
  kIonEventMismatch,
  kIonEventBaseInsertion,
  kIonEventRunInsertion,
  kIonEventBaseDeletion,
  kIonEventRunDeletion
};

template <typename It1, typename It2>
struct IonPairAlignEvent {
  IonEventType type;
  It1 x_iter;
  It2 y_iter;
  unsigned length;
};

template<typename It1, typename It2>
class IonPairAligner {

  It1 x_it_;
  It1 x_end_;
  It2 y_it_;
  It2 y_end_;

  bool empty_;
  hammer::HomopolymerRun cx_, cy_;
  int end_diff_;
  bool at_the_start_; // turned off once we find a pair of runs with same nucleotide

  IonPairAlignEvent<It1, It2> front_;

  // true iff alignment process is not yet finished
  bool checkForZeroLengthRuns() {
    if (x_it_ == x_end_ || y_it_ == y_end_)
      return false;

    if (cx_.len > 0 && cy_.len > 0)
      return true;

    bool result = true;
    while (cx_.len == 0) {
      if (++x_it_ == x_end_) {
        result = false;
        break;
      }
      cx_ = *x_it_;
    }

    while (cy_.len == 0) {
      if (++y_it_ == y_end_) {
        result = false;
        break;
      }
      cy_ = *y_it_;
    }

    return result;
  }

  bool fetchNextX() {
    ++x_it_;
    if (x_it_ == x_end_)
      return false;
    cx_ = *x_it_;
    return true;
  }

  bool fetchNextY() {
    ++y_it_;
    if (y_it_ == y_end_)
      return false;
    cy_ = *y_it_;
    return true;
  }

  void updateEventOffset() {
    front_.x_iter = x_it_;
    front_.y_iter = y_it_;
  }

  void yieldBaseInsertion() {
    front_.type = kIonEventBaseInsertion;
    front_.length = cy_.len - cx_.len;
    updateEventOffset();
  }

  void yieldBaseDeletion() {
    front_.type = kIonEventBaseDeletion;
    front_.length = cx_.len - cy_.len;
    updateEventOffset();
  }

  void yieldRunEvent(IonEventType type) {
    front_.type = type;
    front_.length = 1;
    updateEventOffset();
  }

  void yieldMismatch() { yieldRunEvent(kIonEventMismatch); }
  void yieldRunInsertion() { yieldRunEvent(kIonEventRunInsertion); }
  void yieldRunDeletion() { yieldRunEvent(kIonEventRunDeletion); }

  void finishAlignmentProcess() {
    empty_ = true;
    if (x_it_ != x_end_) {
        end_diff_ += x_end_ - x_it_;
    }
    if (y_it_ != y_end_) {
        end_diff_ -= y_end_ - y_it_;
    }
  }

 public:
  IonPairAligner(const It1 &x_begin, const It1 &x_end,
                 const It2 &y_begin, const It2 &y_end)
    : x_it_(x_begin), x_end_(x_end), y_it_(y_begin), y_end_(y_end),
        empty_(false), cx_(*x_it_), cy_(*y_it_), end_diff_(0), at_the_start_(true)
  {
    popFront();
  }

  bool empty() const { return empty_; }

  IonPairAlignEvent<It1, It2> front() const { return front_; }

  // see comment to distanceHKMer function
  int endDiff() const { return end_diff_; }

  void popFront() {
    VERIFY(x_it_ <= x_end_);
    VERIFY(y_it_ <= y_end_);

    while (true) {
      if (x_it_ == x_end_ || y_it_ == y_end_ || !checkForZeroLengthRuns()) {
        if (x_it_ == x_end_ && y_it_ != y_end_ && cy_.len < y_it_->len) {
          cx_ = *(x_it_ - 1);
          cx_.len = 0;
          // end_diff_ -= 1;
          yieldBaseInsertion();
          fetchNextY();
          return;
        }

        if (y_it_ == y_end_ && x_it_ != x_end_ && cx_.len < x_it_->len) {
          cy_ = *(y_it_ - 1);
          cy_.len = 0;
          // end_diff_ += 1;
          yieldBaseDeletion();
          fetchNextX();
          return;
        }

        finishAlignmentProcess();
        return;
      }

      bool end = false;
      while (cx_.raw == cy_.raw) {
        // don't short-circuit for correct end_diff_ calculation!
        end = !fetchNextX();
        end |= !fetchNextY();
        if (end) break;
        at_the_start_ = false;
      }

      if (!end)
        break;
    }

    if (!checkForZeroLengthRuns()) {
      finishAlignmentProcess();
      return;
    }

    VERIFY(cx_.len > 0);
    VERIFY(cy_.len > 0);
    VERIFY(x_it_ < x_end_);
    VERIFY(y_it_ < y_end_);

    if (cx_.nucl == cy_.nucl) {

      if (cx_.len >= 4 && cy_.len >= 4) {
        if (cx_.len < cy_.len)
          yieldBaseInsertion();
        else
          yieldBaseDeletion();
        fetchNextX();
        fetchNextY();
        at_the_start_ = false;
        return;
      } else {
        --cx_.len;
        --cy_.len;
        at_the_start_ = false;
        popFront();
      }

    } else if (at_the_start_) {
      // alignment can't start with a deletion or insertion
      // unless it is in a homopolymer run
      --cx_.len, --cy_.len;
      yieldMismatch();
      return;
    } else {

      using namespace hammer::errHelper;
      auto hint = getHint(x_it_, x_end_, y_it_, y_end_, cx_.len, cy_.len);

      switch (hint) {
        case kMismatch:
          --cx_.len, --cy_.len;
          yieldMismatch();
          return;
        case kInsertion:
          --cy_.len;
          yieldRunInsertion();
          return;
        case kDeletion:
          --cx_.len;
          yieldRunDeletion();
          return;
      }
    }
  }
};


// returns distance between two homopolymer sequences;
// optionally, fills *end_diff:
//  [ --------- X ----------- ]
// [---------- Y -------]######
//                       \____/
//                       end_diff
template <int kMismatchCost=1,
          int kBaseInsertionCost=1, int kRunInsertionCost=1,
          int kBaseDeletionCost=1, int kRunDeletionCost=1,
          typename It1, typename It2>
inline unsigned distanceHKMer(const It1 &x_begin, const It1 &x_end,
                              const It2 &y_begin, const It2 &y_end,
                              unsigned tau = -1, int *end_diff=NULL) {
  unsigned dist = 0;

  IonPairAligner<It1, It2> aligner(x_begin, x_end, y_begin, y_end);

  while (!aligner.empty()) {
    auto event = aligner.front();
    switch (event.type) {
      case kIonEventMismatch:
        dist += kMismatchCost * event.length; break;
      case kIonEventBaseInsertion:
        dist += kBaseInsertionCost * event.length; break;
      case kIonEventBaseDeletion:
        dist += kBaseDeletionCost * event.length; break;
      case kIonEventRunInsertion:
        dist += kRunInsertionCost * event.length; break;
      case kIonEventRunDeletion:
        dist += kRunDeletionCost * event.length; break;
      default: break;
    }
    if (dist > tau && end_diff == NULL)
        break;
    aligner.popFront();
  }

  if (end_diff != NULL)
      *end_diff = aligner.endDiff();

  return dist;
}

#include <cassert>
#include <iostream>
namespace unittest {

  namespace detail {

    typedef hammer::HKMer::StorageType::const_iterator It;

    inline unsigned distanceHKMer(It beg1, It end1, It beg2, It end2) {
      unsigned dist = 0;

      IonPairAligner<It, It> aligner(beg1, end1, beg2, end2);

      while (!aligner.empty()) {
        auto event = aligner.front();
        switch (event.type) {
        case kIonEventMismatch:
          std::cerr << event.length << 'X';
          dist += event.length; break;
        case kIonEventBaseInsertion:
          std::cerr << event.length << 'I';
          dist += event.length; break;
        case kIonEventBaseDeletion:
          std::cerr << event.length << 'D';
          dist += event.length; break;
        case kIonEventRunInsertion:
          std::cerr << event.length << 'I';
          dist += event.length; break;
        case kIonEventRunDeletion:
          std::cerr << event.length << 'D';
          dist += event.length; break;
        default: break;
        }
        aligner.popFront();
      }

      std::cerr << " (end. diff. = " << aligner.endDiff() << ")" << std::endl;
      return dist;
    }

    inline unsigned distance(const std::string &s, const std::string &t) {
      using namespace hammer;
      HKMer k1, k2;
      for (size_t i = 0; i < s.size(); ++i)
        k1 <<= s[i];
      for (size_t i = 0; i < t.size(); ++i)
        k2 <<= t[i];

      return distanceHKMer(k1.begin(), k1.end(), k2.begin(), k2.end());
    }
  }

  inline void hkmer_distance() {
    using namespace detail;

    assert(distance("ACGTACGTACGTACGT",
                    "CGTACGTACGTACGTA") > 1);

    assert(distance("AACGTACGTACGTACGT",
                    "CGTACGTACGTACGTA") > 1);

    assert(distance("GATAGCGATTTGTTCGGTTTAGGGGGGG",
                    "GATAGCGATTTGTTCGTTTAG") >= 7);

    assert(distance("ATAGCGATTTGTTCGGTTTAGGGGGGGT",
                    "ATAGCGATTTGTTCGTTTAGA") >= 7);

    assert(distance("GATTTGTTCGGTTTAGGGGGGGTAGGGGGATTA",
                    "GATTTGTTCGTTTAGGGGGGGTAGGGGGATTA") == 1);

    assert(distance("TTAAGGCTTACAAAGACTGCGTTT",
                    "TTAAGGCTTACAAAGACTGCGTTTT") == 1);

    assert(distance("AAGGCTTACAAAGACTGCGTTTAA",
                    "AAGGCTTACAAAGACTGCGTA") >= 2);

    assert(distance("ACAAAGACTGCGTTTAAGAGC",
                    "ACAAAGACTGCGTTTTAAGAGC") == 1);

    assert(distance("CTAGGAATGAAAAAGAGAACAAGAA",
                    "CTAGGAATGAAAAAGAGAAAAAAGAATG") == 2);

    assert(distance("ACACACAGGGTTTTTGAACTGGATT",
                    "ACACACAGGGTTTTGAACTGGATT") == 1);

    assert(distance("ACATAAGCCTTTGTACTTAGC",
                    "ACATAAGCCTTTGACTTAGCA") == 1);
  }
}


};
#endif // __HAMMER_HKMER_DISTANCE_HPP__
