// material colors
var
    //              50         100       200        300        400        500       
    pDeepPurple = ['#ede7f6', '#d1c4e9', '#b39ddb', '#9575cd', '#7e57c2', '#673ab7',
    //              600        700        800       900
                   '#5e35b1', '#512da8', '#4527a0', '#311b92'],
    pPink       = ['#fce4ec', '#f8bbd0', '#f48fb1', '#f06292', '#ec407a', '#e91e63',
                   '#d81b60', '#c2185b', '#ad1457', '#880e4f'],
    //               A100      A200        A400       A700
    pDeepPurpleA = ['#b388ff', '#7c4dff', '#651fff', '#6200ea'],
    pGreenA      = ['#b9f6ca', '#69f0ae', '#00e676', '#00c853'],
    pBlueA       = ['#82b1ff', '#448aff', '#2979ff', '#2962ff'],
    pPinkA       = ['#ff80ab', '#ff4081', '#f50057', '#c51162'],

    pDark  = ['#000000', '#212121', '#303030', '#424242'],
    pLight = ['#E0E0E0', '#F5F5F5', '#FAFAFA', '#FFFFFF'],

    oText     = [0.87, 0.54, 0.38, 0.12],
    textColor = ['#000000'],

    primaryColor   = [pDeepPurple[1], pDeepPurple[5], pDeepPurple[7]],
    secondaryColor = pBlueA[3],
    themeColor     = pLight,

    accentColor = secondaryColor,
    bgColor     = themeColor[3],

    linkColor             = pDeepPurple[2],
    linkColorAlt          = pPink[2],
    linkHighlightColor    = pDeepPurple[5],
    linkHighlightColorAlt = pPink[5],
    nodeColor             = pDeepPurple[4],
    nodeColorAlt          = pPink[4],
    nodeHighlightColor    = accentColor,
    nodeHighlightColorAlt = pPinkA[3]
;

function fLinkColor(d, i, highlight) {
    console.log(d.highlight, highlight);
    if (highlight === true || d.highlight) {
        if (d.altLayout) {
            return linkHighlightColorAlt;
        } else {
            return linkHighlightColor;
        }
    } else {
        if (d.altLayout) {
            return linkColorAlt;
        } else {
            return linkColor;
        }
    }
}
function fNodeColor(d, i, highlight) {
    if (highlight === true || d.highlight) {
        if (d.altLayout) {
            return nodeHighlightColorAlt;
        } else {
            return nodeHighlightColor;
        }
    } else {
        if (d.altLayout) {
            return nodeColorAlt;
        } else {
            return nodeColor;
        }
    }
}

/* network construction */

// TODO make multiple bias layers at positions > 0 possible?
var Network = (function() {
    function Network(layout) {
        this.layout = layout;
    }

    Network.prototype.init = function() {
        this._createNodes();
        this._createLinks();
    };

    Network.prototype.resize = function(env) {
        var largestLayerSize = 0;
        for (layer=0; layer<this.layout.net.layers; layer++) {
            if (this.layout.net[layer].size > largestLayerSize) largestLayerSize = this.layout.net[layer].size;
        }
        this.responsiveRadius = d3.min([env.width, env.height]) / d3.max([this.layout.net.layers, largestLayerSize]);
        this._updateNodes(env);
    };

    Network.prototype._createNodes = function() {
        // create layers
        this.nodes = [];
        for (layer=0; layer<this.layout.net.layers; layer++) {
            for (node=0; node<this.layout.net[layer].size; node++) {
                this.nodes.push({
                    id: layer.toString() + "_" + node.toString(),
                    layer: layer,
                    node: node,
                    altLayout: false,
                    highlight: false,
                    fixed: true
                });
            }
        }
    };

    Network.prototype._createLinks = function() {
        this.links = [];
        var nodesMap = d3.map();
        this.nodes.forEach(function (n) {
            return nodesMap.set(n.id, n);
        });
        // create links between adjacent layers
        for (layer=0; layer<this.layout.net.layers-1; layer++) { // for every layer except the last
            for (src=0; src<this.layout.net[layer].size; src++) { // get all nodes
                if (this.layout.net[layer].isBias) {
                    this.links.push({ // in case of a bias layer, connect 1:1
                        source: nodesMap.get(layer.toString() + "_" + src.toString()),
                        target: nodesMap.get((layer+1).toString() + "_" + src.toString()),
                        weight: 1,
                        altLayout: false,
                        highlight: false
                    });
                } else { // normal layer
                    for (tar=0; tar<this.layout.net[layer+1].size; tar++) { // get all nodes of the next layer
                        this.links.push({ // connect both nodes
                            source: nodesMap.get(layer.toString() + "_" + src.toString()),
                            target: nodesMap.get((layer+1).toString() + "_" + tar.toString()),
                            weight: 1,
                            altLayout: false,
                            highlight: false
                        });
                    }
                }
            }
        }
    };

    Network.prototype._updateNodes = function(env) {
        biasLayers = 0;
        for (layer=0; layer<this.layout.net.layers; layer++) {
            if (this.layout.net[layer].isBias) biasLayers++;
        }
        var cnt = 0;
        for (layer=0; layer<this.layout.net.layers; layer++) {
            for (node=0; node<this.layout.net[layer].size; node++) {
                // dynamic ("responsive") attributes for all window sizes
                if (this.layout.net[layer].isBias) {
                    this.nodes[cnt].x = (layer - biasLayers + 1.25) * (env.width / (this.layout.net.layers - biasLayers));
                    this.nodes[cnt].y = (node + 0.25) * (env.height / this.layout.net[layer].size);
                    this.nodes[cnt].r = this.responsiveRadius * 0.2;
                } else {
                    this.nodes[cnt].x = (layer - biasLayers + 0.5) * (env.width / (this.layout.net.layers - biasLayers));
                    this.nodes[cnt].y = (node + 0.5) * (env.height / this.layout.net[layer].size);
                    this.nodes[cnt].r = this.responsiveRadius * 0.3;
                }
                cnt++;
            }
        }
    };

    Network.prototype.updateWeights = function(weights, epochLayout) {
        var cnt = 0;
        for (layer=0; layer<weights.length; layer++) {
            for (layer2=0; layer2<weights[layer].length; layer2++) {
                for (node=0; node<weights[layer][layer2].length; node++) {
                    this.links[cnt].altLayout = (weights[layer][layer2][node] < 0);
                    // scaled = (x - min) / (max - min)
                    this.links[cnt++].weight = (Math.abs(weights[layer][layer2][node]) - epochLayout.min) / (epochLayout.max - epochLayout.min);
                }
            }
        }
    };

    return Network;
})();


/* data handling */

var Loader = (function() {
    function Loader() {
    }

    Loader.prototype.load = function(url, callback) {
        var self = this;
        d3.json(url)
            .get(function(e, json) {
                if (e) console.log(e);
                self._parse(json);
        });
        this._finished = callback;
    };

    Loader.prototype._parse = function(json) {
        // json format:
        // json.
        //   [batch][epoch] [[sourceIndex x targetIndex]] (int)
        // careful - root and first level nodes are objects and have no '.length'

        var layout = {};
        // layout format:
        // layout.
        //   net.layers (int)
        //   net.[net.layers].
        //      isBias (bool)
        //      size (int)
        //   layers.[batch][epoch].
        //      min (int)
        //      max (int)
        //   data.
        //      batches (int)
        //      epochs  (int)
        layout.data = {};
        layout.data.batches = Object.keys(json).length;
        layout.data.epochs  = Object.keys(json[0]).length;
        layout.net = {};
        // network structure
        layout.net.layers = json[0][0].length;
        for (layer=0; layer<layout.net.layers; layer++) {
            layerObj = {};
            if (json[0][0][layer].length == 1) { // bias
                layerObj.isBias = true;
                layerObj.size = json[0][0][0][layer].length; // workaround - Matrix should be rotatet serverside
            } else {
                layerObj.isBias = false;
                layerObj.size = json[0][0][layer].length;
            }
            layout.net[layer] = layerObj;
        }
        layout.net[layout.net.layers] = {'isBias': false, 'size': json[0][0][layout.net.layers-1][0].length};
        layout.net.layers += 1; // output layer
        // layer properties
        layout.layers = {};
        for (batch=0; batch<layout.data.batches; batch++) {
            layout.layers[batch] = {};
            for (epoch=0; epoch<layout.data.epochs; epoch++) {
                layout.layers[batch][epoch] = {};
                layout.layers[batch][epoch].max = 0;
                layout.layers[batch][epoch].min = Infinity;
                for (layer=0; layer<json[batch][epoch].length; layer++) {
                    var layerMin = d3.min(json[batch][epoch][layer], Math.abs);
                    var layerMax = d3.max(json[batch][epoch][layer], Math.abs);
                    if (layerMin < layout.layers[batch][epoch].min) layout.layers[batch][epoch].min = layerMin;
                    if (layerMax > layout.layers[batch][epoch].max) layout.layers[batch][epoch].max = layerMax;
                }
            }
        }
        if (this._finished) this._finished(layout, json);
    };

    return Loader;
})();

var Mocker = (function() {
    function Mocker() {
    }

    Mocker.prototype.load = function(shape, callback) {
        // shape is an array [layer1size, layer2size, …]
        callback(this._createLayout(shape), this._createData(shape));
    };

    Mocker.prototype._createLayout = function(shape) {
        // json format:
        // see Loader class
        var layout = {};
        layout.data = {};
        layout.data.batches = 1;
        layout.data.epochs  = 1;
        layout.net = {};
        // network structure
        layout.net.layers = shape.length;
        for (layer=0; layer<layout.net.layers; layer++) {
            layerObj = {};
            layerObj.isBias = false;
            layerObj.size = shape[layer];
            layout.net[layer] = layerObj;
        }
        // layer properties
        layout.layers = {};
        layout.layers[0] = {};
        layout.layers[0][0] = {};
        layout.layers[0][0].max = 0;
        layout.layers[0][0].min = 0;
        return layout;
    };

    Mocker.prototype._createData = function(shape) {
        var data = {"0": {"0": {}}};
        for (layer=0; layer<shape.length; layer++) {
            data[0][0][layer] = [];
            for (weight=0; weight<shape[layer]; weight++) {
                data[0][0][layer].push(0);
            }
        }
        return data;
    };

    return Mocker;
})();


/* animation stepping */
// TODO move out of this file
var Animator = (function() {
    function Animator(layout, time, step) {
        this.layout = layout;
        this.step = step;
        this.frameDuration = time / (this.layout.data.batches * this.layout.data.epochs / this.step);
    }

    Animator.prototype.start = function(bat, ep, callback, finishedCb) {
        var self = this;
        this.epoch = ep;
        this.batch = bat;
        this.finishedCb = finishedCb;
        wrapper = function() {
            callback(self.batch, self.epoch, self.frameDuration);
            self.animate();
        };
        this.pid = setInterval(wrapper, this.frameDuration);
    };

    Animator.prototype.stop = function() {
        if (this.finishedCb) this.finishedCb();
        clearInterval(this.pid);
    };

    Animator.prototype.animate = function() {
        this.epoch += this.step;
        if (this.epoch >= this.layout.data.epochs) {
            this.epoch = 0;
            if (++this.batch>= this.layout.data.batches) {
                this.batch = 0;
                this.stop();
            }
        }
    };

    return Animator;
})();


/* main */

var NetVisualizer = (function() {
    var _weightScaler = 10; // size of thickest line

    function NetVisualizer(svg) {
        this.loader = new Loader();
        this.mocker = new Mocker();
        this.svg = svg;
        this.stepHooks = [];
        this.stopHooks = [];
    }

    NetVisualizer.prototype.getSteps = function() {
        return this.layout.data.batches * this.layout.data.epochs;
    };
    NetVisualizer.prototype.getDatashape = function() {
        return this.layout.data;
    };

    NetVisualizer.prototype.load = function(url, callback) {
        var self = this;
        function wrapper(layout, data) {
            self.layout = layout;
            self.data = data;
            self._setup();
            callback();
        }
        this.loader.load(url, wrapper);
    };

    NetVisualizer.prototype.mock = function(shape, callback) {
        var self = this;
        function wrapper(layout, data) {
            self.layout = layout;
            self.data = data;
            self._setup();
            callback();
        }
        this.mocker.load(shape, wrapper);
    };

    NetVisualizer.prototype.stopAnimation = function() {
        if (this.animator) this.animator.stop();
    };

    NetVisualizer.prototype.animate = function(animationTime, resolution, step) {
        // animationTime in seconds per frame
        // resolution in epochs per frame
        if (!step) step = 0;
        var self = this;
        var batch = Math.floor(step / this.layout.data.epochs),
            epoch = step % this.layout.data.epochs;
        this.animator = new Animator(this.layout, animationTime * 1000, resolution);
        this.animator.start(batch, epoch, function(cbatch, cepoch) {
            self.show(cbatch, cepoch);
        }, function() {
            self.stopHooks.forEach(function(hook){ hook(); });
        });
    };

    NetVisualizer.prototype.hookStepper = function(hook) {
        this.stepHooks.push(hook);
    };

    NetVisualizer.prototype.hookStop = function(hook) {
        this.stopHooks.push(hook);
    };

    NetVisualizer.prototype.showStep = function(step) {
        var epoch = step % this.layout.data.epochs,
            batch = Math.floor(step / this.layout.data.epochs);
        this.show(batch, epoch);
    };

    NetVisualizer.prototype.show = function(batch, epoch, transitionDuration) {
        var self = this;
        this.network.updateWeights(this.data[batch][epoch], this.layout.layers[batch][epoch]);
        if (transitionDuration) {
            // 10% extra time per frame for the code execution
            this.domLink
                .transition()
                .duration(transitionDuration / 1.1);
        }
        this.domLink.style("stroke-width", function(d) { return d.weight * _weightScaler; })
            .style("stroke", fLinkColor);
        this.stepHooks.forEach(function(hook){ hook(batch * self.layout.data.epochs + epoch); });
    };

    NetVisualizer.prototype.reposition = function() {
        var width = parseInt(this.svg.style("width"))   || 800,
            height = parseInt(this.svg.style("height")) || 600;
        var node = this.svg.selectAll(".node");

        this.network.resize({"width": width, "height": height});

        node
            .attr("cx", function(d) { return d.x; })
            .attr("cy", function(d) { return d.y; })
            .attr("r",  function(d) { return d.r; });
        this.domLink
            .attr("x1", function(d) { return d.source.x; })
            .attr("y1", function(d) { return d.source.y; })
            .attr("x2", function(d) { return d.target.x; })
            .attr("y2", function(d) { return d.target.y; });
    };

    NetVisualizer.prototype._setup = function() {
        var self = this;
        // create empty elements and determine the visualisation's size
        var width = parseInt(this.svg.style("width"))   || 800,
            height = parseInt(this.svg.style("height")) || 600;

        this.network = new Network(this.layout);
        this.network.init();
        this.network.resize({"width": width, "height": height});
        this.network.updateWeights(this.data[0][0], this.layout.layers[0][0]);

        this.domLink = this.svg.selectAll(".link")
            .data(this.network.links)
            .enter().append("line")
            .attr("class", function(d) { return "link source-" + d.source.id + " target-" + d.target.id; })
            .style("stroke", fLinkColor)
            .style("stroke-width", function(d) { return d.weight * _weightScaler; })
            .on("mouseover", function(d) {
                d3.select("#node-" + d.source.id).style("fill", function(d, i) { return fNodeColor(d, i, true); });
                d3.select("#node-" + d.target.id).style("fill", function(d, i) { return fNodeColor(d, i, true); });
                d3.select(this).style("stroke", function(d, i) { return fLinkColor(d, i, true); });
            })
            .on("mouseout", function(d) {
                d3.select("#node-" + d.source.id).style("fill", fNodeColor);
                d3.select("#node-" + d.target.id).style("fill", fNodeColor);
                d3.select(this).style("stroke", fLinkColor);
            })
            .on("click", function(d) {
                d.highlight = !d.highlight;
            });
        var node = this.svg.selectAll(".node")
            .data(this.network.nodes)
            .enter().append("circle")
            .attr("class", function(d) { return "node layer-" + d.layer; })
            .attr("id", function(d) { return "node-" + d.id; })
            .style("fill", fNodeColor)
            .on("mouseover", function(n) {
                d3.selectAll(".source-" + n.id).style("stroke", function(d, i) { return fLinkColor(d, i, true); });
                d3.selectAll(".target-" + n.id).style("stroke", function(d, i) { return fLinkColor(d, i, true); });
                d3.select(this).style("fill", function(d, i) { return fNodeColor(d, i, true); });
            })
            .on("mouseout", function(n) {
                d3.selectAll(".source-" + n.id).style("stroke", fLinkColor);
                d3.selectAll(".target-" + n.id).style("stroke", fLinkColor);
                d3.select(this).style("fill", fNodeColor);
            })
            .on("click", function(d) {
                var me = d3.select(this);
                if (d.highlight) {
                    d3.selectAll(".target-" + d.id).each(function(d) { d.highlight=false; });
                    //d3.selectAll(".target-" + d.id).attr("data-highlight", "false");
                } else {
                    d3.selectAll(".target-" + d.id).each(function(d) { d.highlight=true; });
                    //d3.selectAll(".target-" + d.id).attr("data-highlight", "true");
                }
                d.highlight = !d.highlight;
            });
        this.reposition();
    };

    NetVisualizer.prototype.visualActivate = function(inputs) {
        var activations = this.activate(inputs);
        var flatActivations = [].concat.apply([], activations);
        var max = d3.max(flatActivations, Math.abs),
            min = d3.min(flatActivations, Math.abs);
        var node = this.svg.selectAll(".node");
        node.each(function(d) { d.altLayout = activations[d.layer][d.node]>0; })
            .attr("r",  function(d) { return 2 * d.r * (Math.abs(activations[d.layer][d.node]) - min) / (max-min); })
            .style("fill", fNodeColor);
    };

    NetVisualizer.prototype.activate = function(inputs) {
        // assumes the sigmoid function is used as activation
        function activation(x) { return 1 / (1 + Math.exp(-x)); }
        var weights = this.data[this.layout.data.batches-1][this.layout.data.epochs-1];
        var lastLayerOut = inputs;
        var activations = [inputs];
        var hasBias = false;
        if (this.layout.net[0].isBias) {
            for (bias=0; bias<inputs.length; bias++) {
                lastLayerOut[bias] += weights[0][0][bias];
            }
            hasBias = true;
            activations.push(lastLayerOut);
        }
        for (layer=hasBias?1:0; layer<this.layout.net.layers-1; layer++) {
            // manual vector matrix multiplication
            var thisLayerOut = [];
            for (tar=0; tar<this.layout.net[layer+1].size; tar++) {
                thisLayerOut[tar] = 0;
                for (src=0; src<this.layout.net[layer].size; src++) {
                    thisLayerOut[tar] += weights[layer][src][tar] * lastLayerOut[src];
                }
                thisLayerOut[tar] = activation(thisLayerOut[tar]);
            }
            lastLayerOut = thisLayerOut;
            activations.push(lastLayerOut);
        }
        return activations;
    };

    return NetVisualizer;
})();