#include "ofMain.h"
#include "ofxBox2d.h"
#include "ofxGui.h"
#include "ofxOscParameterSync.h"
#include "ofxAChaosLib.h"
#include "ofxMSATensorFlow.h"
#include "ofxOpenCv.h"
#include <map>
#include <set>
#include "PdBase.hpp"
#include "AChaosBase.h"

#define PLANETS 512

class AChaosChannel {
public:
    REAL min, max;
    vector<REAL> mystory;
    AChaosChannel(){
        min = 1e10;
        max = -1e10;
        mystory.clear();
    }
    void update(REAL val){
        if(val < min) min = val;
        if(val > max) max = val;
        mystory.push_back(val);
        if(mystory.size()>PLANETS){
            mystory.erase(mystory.begin());
        }
    }
};

//--------------------------------------------------------------
class AChaosViz {
public:

    AChaosBase *obj;
    vector<AChaosChannel> channels;
    int n;

    void setup(AChaosBase *ptr){
        obj = ptr;
        n=obj->ov.size();
        for(int i=0; i<n;i++){
            AChaosChannel ch;
            channels.push_back(ch);
        }
    }

    void update(float * data){
        for(int i=0; i<n;i++){
            channels[i].update(data[i]);
        }
    }

    void draw(){
        float w = ofGetWidth();
        float h = ofGetHeight();
        float chh = h / n;
        int s = channels[0].mystory.size();
        if(s<1){
            return;
        }
        float xg = w / s;

        for(int i=0; i<n;i++){

            float ya = i*chh;
            float yc = (i+1)*chh;
            float yb = yc - chh*0.5f;
            ofLine(0, yb, w, yb);
            ofDrawBitmapString(ofToString(channels[i].max), 5, ya + 0.1*chh);
            ofDrawBitmapString(ofToString(channels[i].min), 5, yc);
            ofDrawBitmapString(ofToString(channels[i].mystory[s-1]), 5, yb);

            for(int k=0; k<s-1; k++){
                float y0 = ofMap(channels[i].mystory[k], channels[i].min, channels[i].max, ya, yc);
                float y1 = ofMap(channels[i].mystory[k+1], channels[i].min, channels[i].max, ya, yc);

                ofLine(xg*k, y0, (xg+1)*k, y1);
            }

        }
    }
};

class ofApp : public ofBaseApp{

public:
    void setup();
    void update();
    void draw();

    void exit();

    void keyPressed(int key);
    void keyReleased(int key);
    void mouseMoved(int x, int y );
    void mouseDragged(int x, int y, int button);
    void mousePressed(int x, int y, int button);
    void mouseReleased(int x, int y, int button);
    void mouseEntered(int x, int y);
    void mouseExited(int x, int y);
    void windowResized(int w, int h);
    void dragEvent(ofDragInfo dragInfo);
    void gotMessage(ofMessage msg);

    void drawGravity(ofPoint p, ofPoint gravity);


    void circleResolutionChanged(int & circleResolution);
    void ringButtonPressed();

    void vSyncChanged(bool & vSync);

    bool bHide;

    ofxOscParameterSync sync;

    ofParameter<float> size;
    ofParameter<int> number;
    ofParameter<float> radius;
    ofParameter<bool> check;
    ofParameterGroup parameters;
    ofParameter<ofColor> color;
    ofParameter<ofVec2f> center;
    ofParameter<int> circleResolution;
    ofParameter<bool> filled;
    ofxButton twoCircles;
    ofxButton ringButton;
    ofParameter<string> screenSize;

    ofxPanel gui;
    ofXml settings;

    ofParameter<bool> vSync;
    ofTrueTypeFont font;
    ofSoundPlayer ring;

    AChaosNavierStokes  chaos;
    AChaosViz       viz;

    ofxBox2d                 box2dA;
    ofxBox2d                 box2dB;

    vector <shared_ptr<ofxBox2dCircle> >  circlesA;
    vector <shared_ptr<ofxBox2dCircle> >  circlesB;
    vector <shared_ptr<ofxBox2dRect> >    sharedRects;

    ofRectangle boundsA;
    ofRectangle boundsB;

    // shared pointer to tensorflow::Session
    msa::tf::Session_ptr session;

    // shared pointer to tensorflow::GraphDef
    msa::tf::GraphDef_ptr graph_def;

    // input tensors
    tensorflow::Tensor a, b;

    // output tensors
    vector<tensorflow::Tensor> outputs;

};

//--------------------------------------------------------------
void ofApp::setup(){
    parameters.setName("parameters");
    parameters.add(size.set("size",10,0,100));
    parameters.add(number.set("number",10,0,100));
    parameters.add(check.set("check",false));
    parameters.add(color.set("color",ofColor(127),ofColor(0,0),ofColor(255)));
    gui.setup(parameters);
    // by now needs to pass the gui parameter groups since the panel internally creates it's own group
    sync.setup((ofParameterGroup&)gui.getParameter(),6666,"localhost",6667);
    ofSetVerticalSync(true);
    chaos.setup(); viz.setup(&chaos);

    ofSetFrameRate(60);

    // the world bounds
    boundsA.set(0, 0, 500, 500);
    boundsB.set(500, 0, 500, 500);

    // setup world A
    box2dA.init();
    box2dA.setFPS(60);
    box2dA.setGravity(0, -10);
    box2dA.createBounds(boundsA);
    box2dA.registerGrabbing();

    // setup world B
    box2dB.init();
    box2dB.setFPS(60);
    box2dB.setGravity(0, 10);
    box2dB.createBounds(boundsB);
    box2dB.registerGrabbing();

    // add some cirlces to world A
    for(int i=0; i<10; i++) {
        shared_ptr<ofxBox2dCircle> c = shared_ptr<ofxBox2dCircle>(new ofxBox2dCircle);
        c.get()->setPhysics(1, 0.5, 1);
        c.get()->setup(box2dA.getWorld(), 250+ofRandom(-50, 50), 10, ofRandom(10,30));
        circlesA.push_back(c);
    }

    // add some cirlces to world B
    for(int i=0; i<10; i++) {
        shared_ptr<ofxBox2dCircle> c = shared_ptr<ofxBox2dCircle>(new ofxBox2dCircle);
        c.get()->setPhysics(1, 0.5, 1);
        c.get()->setup(box2dB.getWorld(), 750+ofRandom(-50, 50), 10, ofRandom(10,30));
        circlesB.push_back(c);
    }

    // we can also have a vector of any shape that is
    // shared between both worlds
    for(int i=0; i<20; i++) {

        shared_ptr<ofxBox2dRect> r = shared_ptr<ofxBox2dRect>(new ofxBox2dRect);
        r.get()->setPhysics(1, 0.7, 0.9);

        // add to world A
        if(i <= 9) {
            r.get()->setup(box2dA.getWorld(), 250+ofRandom(-50, 50), 10, ofRandom(10,30), ofRandom(10,30));
        }

        // add to world B
        else {
            r.get()->setup(box2dB.getWorld(), 750+ofRandom(-50, 50), 10, ofRandom(10,30), ofRandom(10,30));
        }

        // add to one vector
        sharedRects.push_back(r);
    }

    ofSetColor(255);
    ofBackground(0);
    ofSetVerticalSync(true);

    // Load graph (i.e. trained model) we exported from python, add to session, return if error
    graph_def = msa::tf::load_graph_def("models/model.pb");
    if(!graph_def) return;

    // initialize session with graph
    session = msa::tf::create_session_with_graph(graph_def);

    // initialize input tensor dimensions
    // (not sure what the best way to do this was as there isn't an 'init' method, just a constructor)
    a = tensorflow::Tensor(tensorflow::DT_FLOAT, tensorflow::TensorShape());
    b = tensorflow::Tensor(tensorflow::DT_FLOAT, tensorflow::TensorShape());
}

//--------------------------------------------------------------
void ofApp::update(){
    sync.update();
    box2dA.update();
    box2dB.update();
}

//--------------------------------------------------------------
void ofApp::draw(){
    gui.draw();
    ofSetColor(color);
    for(int i=0;i<number;i++){
        ofDrawCircle(ofGetWidth()*.5-size*((number-1)*0.5-i), ofGetHeight()*.5, size);
    }
    viz. draw();

    float angle = (atan2(gravity.y, gravity.x) * 180 / PI) - 90;
    float len   = MIN(200, gravity.length()*10); // scale it up a bit

    ofPushMatrix();
    ofTranslate(p.x, p.y);
    ofRotate(angle);
    ofDrawLine(0, 0, 0, len);
    ofDrawTriangle(0, len,
                   -5, len-10,
                   5, len-10);
    ofPopMatrix();

    // world A
    boundsA.inside(ofGetMouseX(), ofGetMouseY()) ? ofSetColor(80) : ofSetColor(100);
    ofFill();
    ofDrawRectangle(boundsA);

    // world A
    boundsB.inside(ofGetMouseX(), ofGetMouseY()) ? ofSetColor(180) : ofSetColor(200);
    ofFill();
    ofDrawRectangle(boundsB);

    ofFill();

    // A World Circles
    for (int i=0; i<circlesA.size(); i++) {
        ofSetHexColor(0xBFE364);
        circlesA[i].get()->draw();
    }

    // B World Circles
    for (int i=0; i<circlesB.size(); i++) {
        ofSetHexColor(0xE83AAB);
        circlesB[i].get()->draw();
    }

    // Shared Rects
    for (int i=0; i<sharedRects.size(); i++) {
        ofSetHexColor(0x2F9BA1);
        ofFill();
        sharedRects[i].get()->draw();
    }

    ofSetColor(255); ofDrawBitmapString("World A\nGravity "+ofToString(box2dA.getGravity().x, 1)+","+ofToString(box2dA.getGravity().y, 1), 250, ofGetHeight()/2);
    ofSetColor(90);  ofDrawBitmapString("World B\nGravity "+ofToString(box2dB.getGravity().x, 1)+","+ofToString(box2dB.getGravity().y, 1), 750, ofGetHeight()/2);

    ofPoint centerA(250, 250);
    ofPoint centerB(750, 250);
    ofSetColor(255);
    drawGravity(centerA, box2dA.getGravity());
    drawGravity(centerB, box2dB.getGravity());

    stringstream message;

    if(session) {
        // inputs are linked to mouse position, normalized to 0..10
        a.scalar<float>()() = round(ofMap(ofGetMouseX(), 0, ofGetWidth(), 0, 10));
        b.scalar<float>()() = round(ofMap(ofGetMouseY(), 0, ofGetHeight(), 0, 10));

        // Collect inputs into a vector
        // IMPORTANT: the string must match the name of the variable/node in the graph
        vector<pair<string, tensorflow::Tensor>> inputs = {
            { "a", a },
            { "b", b },
        };

        // desired outputs which we want processed and returned from the graph
        // IMPORTANT: the string must match the name of the variable/node in the graph
        vector<string> output_names = { "c" };

        // Run the graph, pass in our inputs and desired outputs, evaluate operation and return
        session->Run(inputs, output_names, {}, &outputs);

        // outputs is a vector of tensors, we're interested in only the first tensor
        auto &c = outputs[0];

        // get scalar values of each tensor (since they're 1D and single element it's easy)
        float val_a = a.scalar<float>()();
        float val_b = b.scalar<float>()();
        float val_c = c.scalar<float>()();

        // Print the results
        message << "MOVE MOUSE!" << endl << endl;
        message << val_a << " (" << a.DebugString() << ")" << endl;
        message << " * " << endl;
        message << val_b << " (" << b.DebugString() << ")" << endl;
        message << " = " << endl;
        message << val_c << " (" << c.DebugString() << ")" << endl;
        message << endl;
        message << "all this madness, just to calculate a * b" << endl;

    } else {
        message << "Error during initialization, check console for details.";

    }

    ofDrawBitmapString(message.str(), 30, 30);
}

}

//--------------------------------------------------------------
void ofApp::drawGravity(ofPoint p, ofPoint gravity) {

    float angle = (atan2(gravity.y, gravity.x) * 180 / PI) - 90;
    float len   = MIN(200, gravity.length()*10); // scale it up a bit

    ofPushMatrix();
    ofTranslate(p.x, p.y);
    ofRotate(angle);
    ofDrawLine(0, 0, 0, len);
    ofDrawTriangle(0, len,
                   -5, len-10,
                   5, len-10);
    ofPopMatrix();
}


//--------------------------------------------------------------
void ofApp::mouseDragged(int x, int y, int button){
    float maxGravity = 100;
    if(boundsA.inside(x, y)) {
        float gx = ofMap(x, 0, 500, -maxGravity, maxGravity);
        float gy = ofMap(y, 0, ofGetHeight(), -maxGravity, maxGravity);
        box2dA.setGravity(gx, gy);
    }
    else if(boundsB.inside(x, y)) {
        float gx = ofMap(x, 500, ofGetWidth(), -maxGravity, maxGravity);
        float gy = ofMap(y, 0, ofGetHeight(), -maxGravity, maxGravity);
        box2dB.setGravity(gx, gy);
    }
}

int main( ){
    ofSetupOpenGL(600,600,OF_WINDOW);           // <-------- setup the GL context

    // this kicks off the running of my app
    // can be OF_WINDOW or OF_FULLSCREEN
    // pass in width and height too:
    ofRunApp(new ofApp());

}