/**
 * Created by David DeTomaso on 6/24/2015.
 */

function HeatMap(parent, width, height) {
    var self = this;
    this.h = 1; //height of row

    if (width === undefined) {
        this.width = $(parent).width();
    } else {
        this.width = width;
    }

    if (height === undefined) {
        this.height = $(parent).height();
    } else {
        this.height = height;
    }

    this.svg = d3
        .select(parent)
        .append("svg")
        .attr("width", self.width)
        .attr("height", self.height);

    var offset = 0;

    this.labels = this.svg.append("g");

    this.labels
        .append("text")
        .classed("row_label", true)
        .attr("x", 0)
        .attr("y", offset + 28)
        .attr("font-size", "20px");

    this.labels
        .append("rect")
        .attr("x", this.width - 40)
        .attr("y", offset)
        .attr("width", 40)
        .attr("height", 30)
        .style("fill", "white");

    this.labels
        .append("text")
        .classed("rect_label", true)
        .attr("text-anchor", "middle")
        .attr("x", this.width - 20)
        .attr("y", offset + 17)
        .attr("font-size", "10px");

    offset += 30;
    offset += 10; //Some margin

    this.xlabel_margin = 30;
    this.heat_height = this.height - offset - this.xlabel_margin;
    this.grid_gap = 10; //# of pixels between the plus and minus heat maps
    this.grid_start = offset;
    this.grid_xoffset = 50; //# of pixels to the left of the grids for the + and - labels
    this.grid_xoffsetr = 50; //# of pixels to the right of the grids for the ylabel

    this.grid_plus = this.svg.append("g");
    this.grid_minus = this.svg.append("g");

    this.grid_plus_label = this.svg
        .append("text")
        .attr("text-anchor", "middle")
        .attr("x", this.grid_xoffset / 2)
        .attr("y", 0)
        .attr("font-size", "25px")
        .text("");

    this.grid_minus_label = this.svg
        .append("text")
        .attr("text-anchor", "middle")
        .attr("x", this.grid_xoffset / 2)
        .attr("y", 0)
        .attr("font-size", "25px")
        .text("");

    this.grid_x_label = this.svg
        .append("text")
        .attr("text-anchor", "middle")
        .attr("x", this.width / 2)
        .attr(
            "y",
            this.grid_start + this.heat_height + this.xlabel_margin * 0.8,
        )
        .attr("font-size", "25px")
        .text("Sample Clusters");

    this.grid_y_label = this.svg
        .append("text")
        .attr("text-anchor", "middle")
        .attr("x", this.width - this.grid_xoffsetr / 2)
        .attr("y", this.height / 2)
        .attr("font-size", "25px")
        .style("writing-mode", "tb")
        .text("Genes");

    //define a color scale using the min and max expression values
    this.colorScale = d3.scale
        .linear()
        .domain([0, 4])
        .range(["white", "lightcoral"]);

    this.data_plus = [];
    this.data_minus = [];
    this.selected = -1;
    this.selected_links = [];

    this.hovered_col = -1;

    this.hovered_row = "";

    this.last_event = 0;

    this.cell_labels = [];

    this.cluster_assignments = []; //Used in hoverCol.  Denotes which cluster each sample is assigned to

    this.click = undefined; // To be overwritten and called when clicking
}

HeatMap.prototype.setLabels = function (xlabel, ylabel) {
    this.grid_x_label.text(xlabel);
    this.grid_y_label.text(ylabel);
};

function dist_mat(data) {
    //creates a matrix of size data.length x data.length
    //distance between entries(rows) of data
    //each item in data (which is an array) is an array of numbers

    //Create the array
    var output = [];
    for (var i = 0; i < data.length; i++) {
        output.push([]);
    }

    //Calculate distances
    for (var i = 0; i < data.length; i++) {
        for (var j = i; j < data.length; j++) {
            var dd = dist(data[i], data[j]);
            output[i][j] = dd;
            output[j][i] = dd;
        }
    }

    return output;
}

function dist(x, y) {
    //x and y are arrays of numbers of the same length

    //Euclidean
    var total = 0;
    for (var i = 0; i < x.length; i++) {
        var temp = x[i] - y[i];
        total = total + temp * temp;
    }
    return Math.sqrt(total);

    //1-Pearsons
    var stdx = d3.deviation(x);
    var stdy = d3.deviation(y);
    var meanx = d3.mean(x);
    var meany = d3.mean(y);
    var total = 0;
    for (var i = 0; i < x.length; i++) {
        total = total + (x[i] - meanx) * (y[i] - meany);
    }
    total = total / (stdx * stdy);
    return 1 - total;
}

function cluster_order(data) {
    var distances = dist_mat(data);
    var clusters = [];

    //Convert array of rows into array of objects
    //row_data = row vector
    //index = index in the original matrix

    for (var i = 0; i < data.length; i++) {
        clusters.push([{ index: i, row_data: data[i] }]);
    }

    //clusters is a list of clusters
    //each cluster is a list of row vectors

    while (clusters.length > 1) {
        clusters = merge_closest_cluster_pair(clusters, distances);
    }

    var leaves_indices = [];
    for (var i = 0; i < clusters[0].length; i++) {
        leaves_indices.push(clusters[0][i].index);
    }

    return leaves_indices;
}

function merge_closest_cluster_pair(clusters, distances) {
    var best_pair = [0, 1];
    var closest_distance = 1e99;

    for (var i = 0; i < clusters.length; i++) {
        for (var j = i + 1; j < clusters.length; j++) {
            var distance = distance_between_clusters(
                clusters[i],
                clusters[j],
                distances,
            );
            if (distance < closest_distance) {
                closest_distance = distance;
                best_pair = [i, j];
            }
        }
    }

    var clusterA = clusters[best_pair[0]];
    var clusterB = clusters[best_pair[1]];
    var merged_cluster = merge_clusters(clusterA, clusterB, distances);

    var new_clusters = [];
    for (var i = 0; i < clusters.length; i++) {
        if (i != best_pair[0] && i != best_pair[1]) {
            new_clusters.push(clusters[i]);
        }
    }

    new_clusters.push(merged_cluster);
    return new_clusters;
}

function merge_clusters(clusterA, clusterB, distances) {
    //Merge can happen in two orders
    //Compare first/last clusters in either order to determine the best order

    //AB
    var a_last = clusterA[clusterA.length - 1];
    var b_first = clusterB[0];
    var distAB = distances[a_last.index][b_first.index];

    //BA
    var a_first = clusterA[0];
    var b_last = clusterB[clusterB.length - 1];
    var distBA = distances[a_first.index][b_last.index];

    var new_cluster;
    if (distBA < distAB) {
        new_cluster = clusterB.concat(clusterA);
    } else {
        new_cluster = clusterA.concat(clusterB);
    }

    return new_cluster;
}

function distance_between_clusters(clusterA, clusterB, distances) {
    //Compute the average distance
    var total_dist = 0;
    for (var i = 0; i < clusterA.length; i++) {
        for (var j = 0; j < clusterB.length; j++) {
            var rowA = clusterA[i];
            var rowB = clusterB[j];
            total_dist = total_dist + distances[rowA.index][rowB.index];
        }
    }
    total_dist = total_dist / clusterA.length / clusterB.length;
    return total_dist;

    //    var min_dist = 1e99;
    //    for(var i = 0; i < clusterA.length; i++)
    //    {
    //        for(var j = 0; j < clusterB.length; j++)
    //        {
    //            var rowA = clusterA[i];
    //            var rowB = clusterB[j];
    //            var curr_dist = distances[rowA.index][rowB.index];
    //            if(curr_dist < min_dist)
    //            {
    //                min_dist = curr_dist;
    //            }
    //        }
    //    }
    //
    //    return min_dist;
}

function order_cluster_group_rows(cluster_group) {
    if (cluster_group[0].data.length === 0) {
        return;
    }
    //create data matrix from cluster_group
    var data = [];
    for (var i = 0; i < cluster_group.length; i++) {
        var col_obj = cluster_group[i];
        var row = col_obj.data.map(function (x) {
            return x.value;
        });
        data.push(row);
    }

    data = d3.transpose(data);

    // Get optimal leaf orderings from heirarchical clustering
    var leaf_order = cluster_order(data);

    // Reorder the cluster_group in place
    for (var i = 0; i < cluster_group.length; i++) {
        var col_obj = cluster_group[i];
        var old_data = col_obj.data;
        var new_data = [];
        for (var j = 0; j < leaf_order.length; j++) {
            new_data.push(old_data[leaf_order[j]]);
        }
        col_obj.data = new_data;
    }
}

function order_cluster_group_cols(
    self,
    cluster_group_plus,
    cluster_group_minus,
    TOTAL_SAMPLES,
) {
    // Cluster the columns in the heatmap
    // First, merge the + and - sets into a single column
    // Then, use that column to calculate a new column order
    // Lastly, use the new column order to modify the cluster_group_plus(minus) objects

    if (
        cluster_group_plus[0].data.length === 0 &&
        cluster_group_minus[0].data.length === 0
    ) {
        return [cluster_group_plus, cluster_group_minus];
    }
    //create data matrix from cluster_group
    var data = [];
    for (var i = 0; i < cluster_group_plus.length; i++) {
        var col_obj_p = cluster_group_plus[i];
        var col_obj_m = cluster_group_minus[i];
        var row = col_obj_p.data.map(function (x) {
            return x.value;
        });
        row = row.concat(
            col_obj_m.data.map(function (x) {
                return x.value;
            }),
        );
        data.push(row);
    }

    // Get optimal leaf orderings from heirarchical clustering
    var leaf_order = cluster_order(data);

    cluster_group_plus = leaf_order.map(function (e) {
        return cluster_group_plus[e];
    });
    cluster_group_minus = leaf_order.map(function (e) {
        return cluster_group_minus[e];
    });

    // Change the x_offset for each cluster
    var x_offset = self.grid_xoffset;
    for (var j = 0; j < cluster_group_plus.length; j++) {
        var clust_p = cluster_group_plus[j];
        var width =
            (clust_p.weight / TOTAL_SAMPLES) *
            (self.width - self.grid_xoffset - self.grid_xoffsetr);

        for (var k = 0; k < clust_p.data.length; k++) {
            clust_p.data[k].x = x_offset;
        }

        x_offset = x_offset + width;
    }

    // Same thing for the minus clusters
    var x_offset = self.grid_xoffset;
    for (var j = 0; j < cluster_group_minus.length; j++) {
        var clust_m = cluster_group_minus[j];
        var width =
            (clust_m.weight / TOTAL_SAMPLES) *
            (self.width - self.grid_xoffset - self.grid_xoffsetr);

        for (var k = 0; k < clust_m.data.length; k++) {
            clust_m.data[k].x = x_offset;
        }

        x_offset = x_offset + width;
    }

    return [cluster_group_plus, cluster_group_minus];
}

HeatMap.prototype.setData = function (
    data,
    cluster_assignments,
    gene_labels,
    gene_signs,
    sample_labels,
) {
    //Data is an array of rows, each containing an array of values for each col
    //cluster_assignments is an array of numbers indicating assignment

    this.cell_labels = sample_labels;

    this.cluster_assignments = cluster_assignments;
    var dataT = d3.transpose(data);
    var TOTAL_SAMPLES = cluster_assignments.length;
    var clusters = {};
    for (var j = 0; j < cluster_assignments.length; j++) {
        if (clusters[cluster_assignments[j]] === undefined) {
            clusters[cluster_assignments[j]] = { weight: 0, data: [] };
        }

        var clust = clusters[cluster_assignments[j]];
        clust.weight = clust.weight + 1;
        clust.data.push(dataT[j]);
    }

    //Now clusters is a dict of "cluster index" -> ('weight' -> # of cols in cluster, 'data' -> N_cols x N_row for cluster)
    for (var clust_i in clusters) {
        var clust = clusters[clust_i];
        var cdT = d3.transpose(clust.data);
        clust.data = cdT.map(function (e) {
            return d3.mean(e);
        });
    }
    //Now, each clusters data is just a length N_genes array of values for that cluster.

    //Convert to a list for easier d3 binding
    var cluster_list = [];
    for (var clust_i in clusters) {
        var clust = clusters[clust_i];
        clust.index = clust_i;
        cluster_list.push(clust);
    }

    //Add in information on the width of each cluster
    var x_offset = this.grid_xoffset;
    for (var j = 0; j < cluster_list.length; j++) {
        var clust = cluster_list[j];
        var width =
            (clust.weight / TOTAL_SAMPLES) *
            (this.width - this.grid_xoffset - this.grid_xoffsetr);

        clust.data = clust.data.map(function (e, i) {
            return {
                value: e,
                x: x_offset,
                width: width,
                index: clust.index,
                gene: gene_labels[i],
            };
        });

        x_offset = x_offset + width;
    }

    var N_ROWS = cluster_list[0].data.length;

    //this.h = Math.floor((this.heat_height - this.grid_gap)/N_ROWS);
    this.h = (this.heat_height - this.grid_gap) / N_ROWS;
    if (this.h === 0) {
        this.h = (this.heat_height - this.grid_gap) / N_ROWS;
    }

    //Split data into data_plus and data_minus
    //Unsigned sigs (sign = 0) go in data_plus

    var cluster_list_plus = [];
    var cluster_list_minus = [];
    for (var i = 0; i < cluster_list.length; i++) {
        var clust = cluster_list[i];

        //get positive rows
        var clust_plus = {};
        clust_plus.data = clust.data.filter(function (e, i) {
            return gene_signs[i] === 0 || gene_signs[i] === 1;
        });

        //get negative rows
        var clust_minus = {};
        clust_minus.data = clust.data.filter(function (e, i) {
            return gene_signs[i] === -1;
        });

        //copy over all other properties
        for (var key in clust) {
            if (key !== "data") {
                clust_plus[key] = clust[key];
                clust_minus[key] = clust[key];
            }
        }

        cluster_list_plus.push(clust_plus);
        cluster_list_minus.push(clust_minus);
    }

    // Cluster within data_minus and data_plus

    order_cluster_group_rows(cluster_list_plus);
    order_cluster_group_rows(cluster_list_minus);

    var out = order_cluster_group_cols(
        this,
        cluster_list_plus,
        cluster_list_minus,
        TOTAL_SAMPLES,
    );
    cluster_list_plus = out[0];
    cluster_list_minus = out[1];

    this.data_plus = cluster_list_plus;
    this.data_minus = cluster_list_minus;

    this.setColormap();

    this.redraw()();
};

HeatMap.prototype.setColormap = function () {
    var self = this;
    var vals_p = self.data_plus.map(function (x) {
        return x.data.map(function (y) {
            return y.value;
        });
    });
    var vals_m = self.data_minus.map(function (x) {
        return x.data.map(function (y) {
            return y.value;
        });
    });
    var vals = vals_p.concat(vals_m);

    vals = d3.merge(vals);
    // Rows were z-normalized before getting here
    // So values shoudl be symmetric and zero-centered
    vals = vals.sort(d3.ascending);
    var upper = d3.quantile(vals, 0.99);
    var low = d3.quantile(vals, 0.01);
    var scale = Math.max(Math.abs(low), upper);

    self.colorScale = d3.scale
        .linear()
        .domain([scale * -1, 0, scale])
        .range(["steelblue", "white", "lightcoral"]);
};

HeatMap.prototype.setSelected = function (selected_index, event_id) {
    if (event_id === undefined) {
        event_id = Math.random();
    }

    //Needed to prevent infinite loops with linked hover and select events
    if (this.last_event !== event_id) {
        this.last_event = event_id;
        this.selected = selected_index;
        this.redraw()();
        this.selected_links.forEach(function (e) {
            e.setSelected(selected_index, event_id);
        });
    }
};

HeatMap.prototype.setHovered = function (hovered_indices, event_id) {};

HeatMap.prototype.setHoveredCol = function (hovered_col_index) {
    if (this.hovered_col === hovered_col_index) {
        return;
    }

    this.hovered_col = hovered_col_index;

    var hovered_cells = _.zip(this.cluster_assignments, this.cell_labels)
        .filter((x) => x[0] == hovered_col_index)
        .map((x) => x[1]);

    var event = new CustomEvent("hover-cells", {
        detail: hovered_cells,
        bubbles: true,
    });

    this.svg[0][0].dispatchEvent(event);
};

HeatMap.prototype.setHoveredRowLabel = function (row_label) {
    this.hovered_row = row_label;
    this.labels.select(".row_label").text(this.hovered_row);
};

//Sets the text and color for the square upper-right indicator
HeatMap.prototype.setHoveredIndicator = function (data_val) {
    if (data_val !== undefined) {
        this.labels.select("rect").style("fill", this.colorScale(data_val));
        this.labels.select(".rect_label").text(data_val.toFixed(3));
    } else {
        this.labels.select("rect").style("fill", "white");
        this.labels.select(".rect_label").text("");
    }
};

HeatMap.prototype.redraw = function () {
    var self = this;
    return function () {
        //generate heatmap columns for plus grid

        var pos_grid_start = self.grid_start;
        self.grid_plus.attr("transform", "translate(0," + pos_grid_start + ")");

        self.grid_plus.on("mouseleave", function () {
            self.setHoveredCol(-1);
        });

        var heatmapCols_plus = self.grid_plus
            .selectAll("g")
            .data(self.data_plus);

        heatmapCols_plus.enter().append("g");

        heatmapCols_plus.exit().remove();

        //generate heatmap rows
        var heatmaRects_plus = heatmapCols_plus
            .selectAll("rect")
            .data(function (d) {
                return d.data;
            });

        heatmaRects_plus
            .enter()
            .append("rect")
            .on("mouseover", function (d, i) {
                self.setHoveredRowLabel(d.gene);
                self.setHoveredCol(d.index);
                self.setHoveredIndicator(d.value);
            })
            .on("click", function (d, i) {
                if (self.click !== undefined) {
                    self.click(d.gene, d.index, d.value);
                }
            });

        heatmaRects_plus
            .style("fill", function (d) {
                return self.colorScale(d.value);
            })
            .attr("width", function (d) {
                return d.width;
            })
            .attr("height", self.h)
            .attr("y", function (d, i) {
                return i * self.h;
            })
            .attr("x", function (d) {
                return d.x;
            });

        heatmaRects_plus.exit().remove();

        var num_pos_rects;
        if (self.data_plus.length === 0) {
            num_pos_rects = 0;
        } else {
            num_pos_rects = self.data_plus[0].data.length;
        }

        var neg_grid_start = 0;
        neg_grid_start += pos_grid_start; // Start location for all grids
        neg_grid_start += self.h * num_pos_rects; // Offset for space taken up by positive grid
        if (num_pos_rects > 0) {
            neg_grid_start += self.grid_gap; // If there's a positive grid, introduce a gap between them
        }

        self.grid_minus.attr(
            "transform",
            "translate(0," + neg_grid_start + ")",
        );

        self.grid_minus.on("mouseleave", function () {
            self.setHoveredCol(-1);
        });

        //generate heatmap columns for minus grid
        var heatmapCols_minus = self.grid_minus
            .selectAll("g")
            .data(self.data_minus);

        heatmapCols_minus.enter().append("g");

        heatmapCols_minus.exit().remove();

        //generate heatmap rows
        var heatmapRects_minus = heatmapCols_minus
            .selectAll("rect")
            .data(function (d) {
                return d.data;
            });

        heatmapRects_minus
            .enter()
            .append("rect")
            .on("mouseover", function (d, i) {
                self.setHoveredRowLabel(d.gene);
                self.setHoveredCol(d.index);
                self.setHoveredIndicator(d.value);
            })
            .on("click", function (d, i) {
                if (self.click !== undefined) {
                    self.click(d.gene, d.index, d.value);
                }
            });

        self.svg.on("mouseleave", function (d) {
            self.setHoveredRowLabel("");
            self.setHoveredIndicator();
        });

        heatmapRects_minus
            .style("fill", function (d) {
                return self.colorScale(d.value);
            })
            .attr("width", function (d) {
                return d.width;
            })
            .attr("height", self.h)
            .attr("y", function (d, i) {
                return i * self.h;
            })
            .attr("x", function (d) {
                return d.x;
            });

        heatmapRects_minus.exit().remove();

        var num_neg_rects;
        if (self.data_minus.length === 0) {
            num_neg_rects = 0;
        } else {
            num_neg_rects = self.data_minus[0].data.length;
        }

        //Plus and minus labels for the grids
        if (num_pos_rects > 0) {
            var pos_grid_center = Math.floor(
                pos_grid_start + (self.h * num_pos_rects) / 2,
            );
            self.grid_plus_label
                .attr("y", pos_grid_center + 9) // Y lines up with baseline, need to offset to vertically center
                .text("+"); // Some browsers have an attribute to do this, doesn't work on FireFox
        } else {
            self.grid_plus_label.text("");
        }

        if (num_neg_rects > 0) {
            var neg_grid_center = Math.floor(
                neg_grid_start + (self.h * num_neg_rects) / 2,
            );
            self.grid_minus_label
                .attr("y", neg_grid_center + 9) // See note for plus label above
                .text("-");
        } else {
            self.grid_minus_label.text("");
        }
    };
};
