#include <gecode/minimodel.hh>
#include "regular-extended.hh"
#include <algorithm>

namespace SigLex {
    using namespace Gecode;

    class CountingTransition {
        int n;
    public:
        CountingTransition(int n_) : n(n_) { }

        int operator() (int in_state, int symbol) {
            int d = in_state / 4 - n;
            int y = in_state % 4;
            int newd = d + (symbol==2) - (symbol==1);
            int newy = std::max(y, symbol);
            int ret;
            if (-n <= newd && newd <= n && symbol >= y)
                ret = (n+newd)*4+newy;
            else
                ret = -1;
            // std::cerr << "counting(" << n << ") " << in_state << " " << symbol << " -> " << ret
            //           << "\t (" << d << "," << y << ")->(" << newd << "," << newy << ")\n";
            return ret;
            //                        std::cerr << i << " " <<
            //                        t[i].i_state << " " <<
            //                        t[i].symbol << " " <<
            //                        t[i].o_state << "\n";
        }
    };

    class OrderingTransition {
    public:
        int operator() (int state, int symbol) {
            if (symbol <= 0 || state == 1)
                return std::max(state, static_cast<int>(symbol<0));
            else
                return -1;
        }
    };

    void
    siglex(Home home, int k, IntArgs p, const IntVarArgs& x, IntVarArray *dp, IntVarArray *yp, IntVarArray *finp) {
        // Order variables in each partition.
        int a = p.size();
        int n = x.size();
        for (int i = 0 ; i < a ; i++) {
            int stop = (i < a-1) ? (p[i+1]-1) : (n-1);
            for (int j = p[i] ; j < stop ; j++) {
                rel(home, x[j] <= x[j+1]);
                // std::cerr << "x[" << j << "] <= x[" << j+1 << "]\n";
            }
        }

        // Support::DynamicArray<DFA::Transition, Heap> t(heap);
        // int i = 0;
        // for (int d = -n ; d <= n ; d++)
        //     for (int y = 0 ; y <= 3 ; y++)
        //         for (int symbol = y ; symbol <= 3 ; symbol++) {
        //             t[i].i_state = (n+d)*4+y;
        //             t[i].symbol = symbol;
        //             int newd = d + (symbol==2) - (symbol==1);
        //             int newy = std::max(y, symbol);
        //             t[i].o_state = (n+newd)*4+newy;
        //             std::cerr << i << " " << t[i].i_state << " " << t[i].symbol << " " << t[i].o_state << "\n";
        //             if (-n <= newd && newd <= n)
        //                 i++;
        //         }
        // t[i].i_state = -1;
        int start_state = (n+0)*4+0;
        // Any state can be a final state.
        int n_states = (2*n+1)*4;
        int f[n_states+1];
        for (int fi = 0 ; fi < n_states ; fi++)
            f[fi] = fi;
        f[n_states] = -1;
        // for (int fi = 0 ; fi < n_states+1 ; fi++) std::cerr << f[fi] << " "; std::cerr << "\n";
        RegularExtended::ExtDFA<CountingTransition> dfa(start_state, n_states, 0, 3, CountingTransition(n));

        // Introduce Y_k^i variables.
        IntVarArray Y(home, n, 0, 3);
        IntVarArray D(home, a, -n, n);
        IntVarArray fin(home, a, 0, (n+n)*4+3);

        // For each partition, compute and order the signatures.
        for (int i = 0 ; i < a ; i++) {
            int stop = i < a-1 ? p[i+1]-1 : n-1;
            for (int j = p[i] ; j <= stop ; j++) {
                BoolVarArray B(home, 4, 0, 1);
                channel(home, B, Y[j], 0, ICL_DOM);
                rel(home, B[0] == (x[j] < k), ICL_DOM);
                rel(home, B[1] == (x[j] == k), ICL_DOM);
                rel(home, B[2] == (x[j] == k+1), ICL_DOM);
                rel(home, B[3] == (x[j] > k+1), ICL_DOM);
            }

            RegularExtended::regularExtended(home, Y.slice(p[i], 1, stop-p[i]+1), fin[i], dfa, ICL_DOM);
            // std::cerr << "Y: " << p[i] << " length " << stop-p[i]+1 << "\n";
            // Relate fin and D[i].
            rel(home, D[i] == (fin[i]/4)-n, ICL_DOM);
        }

        // Order the D variables.
        // Support::DynamicArray<DFA::Transition, Heap> t2(heap);
        // int i2 = 0;
        // for (int b = 0 ; b <= 1 ; b++)
        //     for (int d = -n ; d <= n ; d++) {
        //         t2[i2].i_state = b;
        //         t2[i2].symbol = d;
        //         t2[i2].o_state = std::max(b, static_cast<int>(d<0));
        //         if (d <= 0 || b == 1) {
        //             std::cerr << i2 << " " << t2[i2].i_state << " " << t2[i2].symbol << " " << t2[i2].o_state << "\n";
        //             i2++;
        //         }
        //     }
        // t2[i2].i_state = -1;
        int start_state2 = 0;
        // int final_states2[] = {0,1, -1};
        RegularExtended::ExtDFA<OrderingTransition> dfa2(start_state2, 2, -n, n, OrderingTransition());
        IntVar fin2(home, 0, 1);
        RegularExtended::regularExtended(home, D, fin2, dfa2, ICL_DOM);

        if (dp != NULL) *dp = D;
        if (yp != NULL) *yp = Y;
        if (finp != NULL) *finp = fin;
    }

    using namespace std;
    // http://stackoverflow.com/questions/279854/how-do-i-sort-a-vector-of-pairs-based-on-the-second-element-of-the-pair/280662#280662
    template<template <typename> class P = std::less >
    struct compare_pair_second {
        template<class T1, class T2> bool operator()(const std::pair<T1, T2>& left, const std::pair<T1, T2>& right) {
            return P<T2>()(left.second, right.second);
        }
    };

    template <class T, class U>
    void print_pair(pair<T,U> e) { cout << "(" << e.first << "," << e.second << ") "; }
    
    void
    siglexdec(Home home, int k, IntArgs p, const IntVarArgs& x, IntVarArray *d, IntVarArray *y, IntVarArray *finp) {
        // Compute partition sizes.
        vector<pair<int, int> > partitions;
        int a = p.size();
        for (int i = 0 ; i < a ; i++) {
            int sz = (i == a-1) ? (x.size() - p[i]) : (p[i+1] - p[i]);
            partitions.push_back(make_pair(i, sz));
        }

        // Sort partitions by size.
        // void (*pp)(pair<int,int>) = print_pair;
        // for_each(partitions.begin(), partitions.end(), pp);
        // cout << "\n";
        sort(partitions.begin(), partitions.end(), compare_pair_second<std::greater>());
        // for_each(partitions.begin(), partitions.end(), pp);
        // cout << "\n";

        // Reorder variables and partitions.
        IntArgs newp(a);
        IntVarArgs newx(x.size());
        int counter = 0;
        int pcounter = 0;
        for (vector<pair<int,int> >::iterator i = partitions.begin() ; i != partitions.end() ; ++i) {
            newp[pcounter++] = counter;
            int j = i->first;
            for (int k = 0 ; k < i->second ; k++) {
                newx[counter] = x[p[j]+k];
                // cout << p[j]+k << " -> " << counter << "\n";
                counter++;
            }
        }

        // cout << newp;

        // Post constraint on re-ordered partitions.
        siglex(home, k, newp, newx, d, y, finp);
    }
}
