import algorithms.heaps as heaps;
import algorithms.introsort as introsort;
import comparators.*;

alias log2 = introsort.log2;


//
// reduce
//

[F, A | Sequence?(A)]
reduce(f:F, a:A) {
    var iter = iterator(a);
    assert(hasNext?(iter));
    var result = next(iter);
    while (hasNext?(iter)) {
        forward v = next(iter);
        result = f(result, v);
    }
    return result;
}

[F, A, X | Sequence?(A) and CallDefined?(call, F, X, SequenceElementType(A))]
overload reduce(f:F, initial:X, a:A) {
    var result = initial;
    for (x in a)
        result = f(result, x);
    return result;
}



//
// reduceBy
//
[KeyF, F, A, X | Sequence?(A) and CallDefined?(call, F, X, SequenceElementType(A))]
reduceBy(keyF: KeyF, f: F, initial: X, a:A) {
    alias KeyType = Type(keyF(typeToLValue(SequenceElementType(A))));
    alias ValueType = X;
    var results = HashMap[KeyType, ValueType]();

    for (x in a) {
        var key = keyF(x);
        var resultPtr = lookup(results, key);
        if (null?(resultPtr))
            put(results, key, f(initial, x));
        else
            resultPtr^ = f(resultPtr^, x);
    }
    return results;
}



//
// sum, product, count
//

[A | Sequence?(A)]
sum(a:A) {
    alias T = SequenceElementType(A);
    return reduce(add, T(0), a);
}

[A | Sequence?(A)]
product(a:A) {
    alias T = SequenceElementType(A);
    return reduce(multiply, T(1), a);
}

[A | Sequence?(A)]
count(pred, a:A) {
    return reduce((n,elt) => if (pred(elt)) n+1 else n, SizeT(0), a);
}



//
// select
//

private define select;

[A | Sequence?(A)]
overload select(a:A, f) {
    var iter = iterator(a);
    assert(hasNext?(iter));
    var result = next(iter);
    while (hasNext?(iter)) {
        forward v = next(iter);
        if (f(v, result))
            result = v;
    }
    return result;
}

[A | LValueSequence?(A)]
overload select(a:A, f) {
    var iter = iterator(a);
    assert(hasNext?(iter));
    var result = &next(iter);
    while (hasNext?(iter)) {
        forward v = next(iter);
        if (f(v, result^))
            result = &v;
    }
    return ref result^;
}



//
// min, max for sequences
//

[A | Sequence?(A)]
overload min(a:A) = forward select(a, lesser?);

[A | Sequence?(A)]
overload max(a:A) = forward select(a, greaterEquals?);



//
// findIf, reverseFindIf
//

findIf(seq, pred) {
    return findIf(begin(seq), end(seq), pred);
}

overload findIf(_first, _last, pred) {
    var cur, last = _first, _last;
    while ((cur != last) and not pred(cur^))
        inc(cur);
    return cur;
}

reverseFindIf(seq, pred) {
    return reverseFindIf(begin(seq), end(seq), pred);
}

overload reverseFindIf(_first, _last, pred) {
    var first, cur = _first, _last;
    while (cur != first) {
        var prev = cur;
        dec(prev);
        if (pred(prev^))
            break;
        cur = prev;
    }
    return cur;
}

any?(seq, pred) = findIf(seq, pred) != end(seq);
all?(seq, pred) = findIf(seq, x -> not pred(x)) != end(seq);


//
// find
//

define find;

[A, T | Sequence?(A) and CallDefined?(equals?, SequenceElementType(A), T)]
overload find(a:A, x:T) {
    return find(begin(a), end(a), x);
}

[C, T | Coordinate?(C) and CallDefined?(equals?, CoordinateTargetType(C), T)]
overload find(begin:C, end:C, x:T) {
    var i = begin;
    while (i != end) {
        if (i^ == x)
            break;
        inc(i);
    }
    return i;
}


[A, B | SizedSequence?(A) and SizedSequence?(B)
        and (SequenceElementType(A) == SequenceElementType(B))]
overload find(a:A, b:B) {
    return find(begin(a), end(a), b);
}

[C, B | Coordinate?(C) and SizedSequence?(B)
        and (CoordinateTargetType(C) == SequenceElementType(B))]
overload find(begin_:C, end_:C, b:B) {
    // FIXME: use something like boyer-moore.
    var begin, end = begin_, end_;
    var bsize = size(b);
    for (c in range(begin, end)) {
        var n = end - c;
        if (n < bsize)
            break;
        var range = coordinateRange(c, c+bsize);
        if (range == b)
            return c;
    }
    return end;
}



//
// reverseFind
//

define reverseFind;

[A, T | Sequence?(A) and CallDefined?(equals?, SequenceElementType(A), T)]
overload reverseFind(a:A, x:T) {
    return reverseFind(begin(a), end(a), x);
}

[C, T | Coordinate?(C) and CallDefined?(equals?, CoordinateTargetType(C), T)]
overload reverseFind(begin:C, end:C, x:T) {
    var i = end;
    while (i != begin) {
        var j = i;
        dec(j);
        if (j^ == x)
            break;
        i = j;
    }
    return i;
}



//
// in?
//
in?(seq, elem) = find(seq, elem) != end(seq);



//
// beginsWith?, endsWith?
//

[A, B | Sequence?(A) and Sequence?(B) and (SequenceElementType(A) == SequenceElementType(B))]
beginsWith?(seq:A, prefix:B) =
    (size(seq) >= size(prefix)) and (slicedUpto(seq, size(prefix)) == prefix);

[A, T | Sequence?(A) and (T == SequenceElementType(A))]
overload beginsWith?(seq:A, x:T) =
    (size(seq) > 0) and (seq[0] == x);


[A, B | Sequence?(A) and Sequence?(B) and (SequenceElementType(A) == SequenceElementType(B))]
endsWith?(seq:A, suffix:B) =
    (size(seq) >= size(suffix))
    and (slicedFrom(seq, size(seq)-size(suffix)) == suffix);

[A, T | Sequence?(A) and (T == SequenceElementType(A))]
overload endsWith?(seq:A, x:T) =
    (size(seq) > 0) and (seq[size(seq)-1] == x);



//
// reverseInPlace
//

[A | RandomAccessSequence?(A) and LValueSequence?(A)]
reverseInPlace(a:A) {
    var n = size(a);
    for (i in range(n/2))
        swap(a[i], a[n-i-1]);
}



//
// split
//

[A | CoordinateSequence?(A)]
split(a:A, sep) {
    alias T = SequenceElementType(A);
    var result = Vector[Vector[T]]();
    var cur, last = begin(a), end(a);
    while (true) {
        var pos = find(cur, last, sep);
        push(result, Vector(coordinateRange(cur, pos)));
        if (pos == last)
            break;
        cur = advancePast(pos, sep);
    }
    return move(result);
}

private define advancePast;

[C, S | Coordinate?(C) and Sequence?(S)
        and (CoordinateTargetType(C) == SequenceElementType(S))]
overload advancePast(pos:C, x:S) {
    return pos + size(x);
}

[C, T | Coordinate?(C) and (CoordinateTargetType(C) == T)]
overload advancePast(pos:C, x:T) {
    return pos + 1;
}



//
// join
//

[S, A | Sequence?(A) and Sequence?(SequenceElementType(A))
        and ((S == SequenceElementType(SequenceElementType(A)))
             or (Sequence?(S)
                 and (SequenceElementType(S) == SequenceElementType(SequenceElementType(A)))))]
join(sep:S, a:A) {
    alias T = SequenceElementType(SequenceElementType(A));
    var result = Vector[T]();
    var first = true;
    for (x in a) {
        if (not first)
            push(result, sep);
        else
            first = false;
        push(result, x);
    }
    return move(result);
}



//
// binarySearch, binarySearchLowerBound, binarySearchUpperBound
//

[A | Sequence?(A)]
binarySearch(a:A, x) {
    return ..binarySearch(begin(a), end(a), x);
}

[A | Sequence?(A)]
binarySearchLowerBound(a:A, x) {
    return binarySearchLowerBound(begin(a), end(a), x);
}

[A | Sequence?(A)]
binarySearchUpperBound(a:A, x) {
    return binarySearchUpperBound(begin(a), end(a), x);
}

overload binarySearch(begin, end, x) {
    var i = binarySearchLowerBound(begin, end, x);
    return i, binarySearchUpperBound(i, end, x);
}

overload binarySearchLowerBound(begin_, end_, x) {
    var begin, end = begin_, end_;
    while (begin < end) {
        var mid = begin + (end - begin)/2;
        if (mid^ < x)
            begin = mid + 1;
        else
            end = mid;
    }
    return begin;
}

overload binarySearchUpperBound(begin_, end_, x) {
    var begin, end = begin_, end_;
    while (begin < end) {
        var mid = begin + (end - begin)/2;
        if (mid^ <= x)
            begin = mid + 1;
        else
            end = mid;
    }
    return end;
}



//
// sort - use quicksort with median-of-3 partitioning
//

[C | Comparator?(C)]
sortBy(a, comparator : C) {
    introsort.introSort(a, (x, y) -> lesser?(comparator, x, y));
}

sort(a){
  sortBy(a, natural());
}

overload sort(xs, f){
  sortBy(xs, comparing(f));
}
