#!/usr/bin/env python

import argparse
import csv
import os
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import scipy.stats


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("corr_coeff_file")
    parser.add_argument("translation_table")
    parser.add_argument("operon_list_folder")
    parser.add_argument("--output_pdf", default="regulons.pdf")
    args = parser.parse_args()
    corr_value_visualizer = CorrValueVisualizer(
        args.corr_coeff_file, args.operon_list_folder, args.translation_table,
        args.output_pdf)
    corr_value_visualizer._read_corr_coeff_file()
    corr_value_visualizer._read_translation_table()
    corr_value_visualizer._collect_values_of_regulons()
    corr_value_visualizer._plot_box_plots()
    corr_value_visualizer._run_stat_test()


class CorrValueVisualizer(object):

    def __init__(self, corr_coeff_file, operon_list_folder, translation_table,
                 output_pdf):
        self._corr_coeff_file = corr_coeff_file
        self._operon_list_folder = operon_list_folder
        self._translation_table = translation_table
        self._output_pdf = output_pdf

    def _read_corr_coeff_file(self):
        self._genes_and_coefficients = {}
        for row in csv.reader(open(self._corr_coeff_file), delimiter="\t"):
            if row[0].startswith("Gene"):
                continue
            gene, coefficient = row[0], row[1]
            self._genes_and_coefficients[gene] = float(coefficient)

    def _read_translation_table(self):
        self._gene_name_translation = {}
        for row in csv.reader(open(self._translation_table), delimiter="\t"):
            self._gene_name_translation[row[0]] = row[1]

    def _collect_values_of_regulons(self):
        self._regulons_and_values = {}
        for operon_list_file in os.listdir(self._operon_list_folder):
            operon_regulator = operon_list_file.split("_")[0]
            self._regulons_and_values[operon_regulator] = []
            for row in csv.reader(
                    open(self._operon_list_folder + "/" + operon_list_file),
                    delimiter="\t"):
                if (row[0].startswith("gene") or len(row[0]) == 0 or
                   len(row[0]) > 10):
                    continue
                else:
                    gene = row[0]
                    corr_value = self._get_corr_value(gene)
                    if corr_value is not None:
                        self._regulons_and_values[operon_regulator].append(
                            corr_value)

    def _plot_box_plots(self):
        box_color = "#FFB3B3"
        median_color = "#CC0000"
        fig = plt.figure()
        ax = fig.add_subplot(111)
        # Sort by median
        medians_and_names = sorted([
            (np.median(np.array(self._regulons_and_values[name])), name)
            for name in self._regulons_and_values.keys()])
        regulon_names = [name for median, name in medians_and_names]
        regulon_values = [
            self._regulons_and_values[regulon] for regulon in regulon_names]
        reguoln_sizes = [len(values) for values in regulon_values]
        plot = plt.boxplot(regulon_values, patch_artist=True,
                           widths=[0.8]*len(regulon_values))
        names_with_sizes = ["%s (%s)" % (name, size) for name, size in zip(
            regulon_names, reguoln_sizes)]
        plt.xticks(np.array(range(1, len(regulon_names)+1)),
                   names_with_sizes, rotation=45, ha="right")
        plt.title("Distributions of Pearson correlation "
                  "coefficients (compared to pinT)", fontsize=10)
        # Clean up
        ax.spines["top"].set_visible(True)
        ax.spines["bottom"].set_visible(True)
        ax.spines["left"].set_visible(True)
        ax.spines["right"].set_visible(True)
        ax.xaxis.set_ticks_position("none")
        ax.yaxis.set_ticks_position("left")
        plt.setp(plot["medians"], color=median_color, linewidth=2.0)
        plt.setp(plot["fliers"], color=box_color, marker=".",
                 alpha=0.8, markersize=5)
        plt.setp(plot["whiskers"], color=box_color, linestyle="solid",
                 linewidth=2.0)
        plt.setp(plot["caps"], color=box_color, linestyle="solid",
                 linewidth=2.0)
        plt.setp(plot["boxes"], color=box_color)
        for patch in plot["boxes"]:
            patch.set_facecolor(box_color)

        ax.grid(False)
        ax.patch.set_facecolor("white")

        plt.tight_layout()
        rect = plt.Rectangle(
            (0, -0.5), len(regulon_names)+1, 1,
            linewidth=1, color="#EEEEEE")
        plt.gca().add_patch(rect)
        fig.savefig(self._output_pdf)
        
    def _run_stat_test(self):
        phop_values = self._regulons_and_values["PhoP"]
        print("Regulon name\t"
              "Kolmogorov-Smirnov test comparison with PhoP p-value"
              "\tMedian correlation score")
        for regulon_name, regulon_values in self._regulons_and_values.items():
            print("%s\t%s\t%s" % (
                regulon_name,
                scipy.stats.ks_2samp(
                    np.array(phop_values), np.array(regulon_values))[1],
                np.median(np.array(regulon_values))))

    def _get_corr_value(self, gene):
        if gene in self._genes_and_coefficients:
            return self._genes_and_coefficients[gene]
        elif gene in self._gene_name_translation:
            return self._genes_and_coefficients[
                self._gene_name_translation[gene]]
    
if __name__ == "__main__":
    main()

