#include <stp/c_interface.h>
#include <vector>
#include <iostream>
#include <sstream>
#include <string>
#include <stdlib.h>

using namespace std;

int main(int argc, char** argv) {
    int colors = 3;
    int n = 0;

    if(argc <= 1) {
        cout << "Usage: " << argv[0] << " n" << endl;
        exit(1);
    }
    else {
        n = atoi(argv[1]);
    }

    VC handle = vc_createValidityChecker();

    //create binary variable of length "colors" for each vertex
    cout << "Stage 1: Generating variables for each vertex" << endl;
    std::vector<Expr> color_map;
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j <= i; ++j) {
            stringstream name;
            name << "col_" << i+1 << "_" << j+1;
            string name_st = name.str();

            Expr x = vc_varExpr(handle, name_st.c_str(), vc_bvType(handle, colors));
            color_map.push_back(x);
        }
    }

    //now we generate inequalities
    cout << "Stage 2: Generating basic inequalities" << endl;
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j <= i; ++j) {
            int index = i*(i + 1)/2+j;
            Expr cond = NULL;
            std::vector<Expr> feasibility;

            //every vertex gets exactly one
            for(int k = 0; k < colors; ++k) {
                Expr result = NULL;
                for(int l = 0; l < colors; ++l) {
                    Expr is_color_l = vc_notExpr(handle,vc_bvBoolExtract_One(handle,color_map[index],l));
                    if(l == k) {
                        is_color_l = vc_notExpr(handle,is_color_l);
                    }
                    if(result == NULL) {
                        result = is_color_l;
                    }
                    else {
                        result = vc_andExpr(handle,result,is_color_l);
                    }
                }

                feasibility.push_back(result);
            }

            vc_assertFormula(handle,vc_orExprN(handle,&feasibility[0],feasibility.size()));
        }
    }

    //no triangle has uniform colors
    cout << "Stage 3: Generating coloring inequalities" << endl;
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j <= i; ++j) {
            //(i,j) is the tip
            for( int i2 = i+1; i2 < n; ++i2) {
                int j2 = j;
                int j3 = j2+(i2-i);
                if(j3 > i2) continue;

                //we now consider the triangle (i,j) (i2,j2) (i2,j3)
                int v = i *(i  + 1)/2+j ;
                int w = i2*(i2 + 1)/2+j2;
                int q = i2*(i2 + 1)/2+j3;

                //v w and q do not all have the same color
                for(int k = 0; k < colors; ++k) {
                    Expr has_v_col_k = vc_bvBoolExtract_One(handle,color_map[v],k); //extract bit k
                    Expr has_w_col_k = vc_bvBoolExtract_One(handle,color_map[w],k);
                    Expr has_q_col_k = vc_bvBoolExtract_One(handle,color_map[q],k);

                    //if v and w have the same color, q has a different color, etc
                    Expr cond_vw = vc_impliesExpr(handle,vc_andExpr(handle,has_v_col_k,has_w_col_k),vc_notExpr(handle,vc_andExpr(handle,has_v_col_k,has_q_col_k)));

                    vc_assertFormula(handle,cond_vw);
                }
            }
        }
    }

    cout << "Stage 4: Solving" << endl;
    int result = vc_query(handle,vc_falseExpr(handle)); //find counterexample
    switch(result)  {
    case 0:
        cout << "Found a valid coloring with " << colors << " colors for n=" << n << endl;
        vc_printCounterExample(handle);
        break;
    case 1:
        cout << "Could NOT find a coloring with " << colors << " colors for n=" << n << endl;
        break;
    default:
        cout << "Unexpected result from vc_query (" << result << ")" << endl;
    }

    //clean up
    vc_Destroy(handle);
    color_map.clear();

    return 0;
}