import sys
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy import stats
import numpy as np


def main():
    mapping_comparer = MappingComparer(
        sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4], sys.argv[5],
        sys.argv[6])
    mapping_comparer.read_and_filter_data()
    mapping_comparer.generate_scatter_plot()
    mapping_comparer.write_rrna_table()
    mapping_comparer.write_decreased_non_rrna()
    mapping_comparer.plot_histogram()


class MappingComparer(object):

    def __init__(self, input_file, organism, scatter_plot_file, histogram_file,
                 rrna_file, top_changed_non_rrna_file):
        self._multiplier = 1000000
        self._input_file = input_file
        assert organism in ["Salmonella", "Human"]
        self._organism = organism
        self._scatter_plot_file = scatter_plot_file
        self._histogram_file = histogram_file
        self._rrna_file = rrna_file
        self._top_changed_non_rrna_file = top_changed_non_rrna_file

    def read_and_filter_data(self):
        self._normalized_gene_quanti = pd.read_table(self._input_file)
        # Calculate the ratio
        self._normalized_gene_quanti["ratio"] = self._normalized_gene_quanti[
            "HeLa-S3:WT:16_h:total_RNA"] / self._normalized_gene_quanti[
            "HeLa-S3:WT:16_h:rRNA_free"]
        if self._organism == "Salmonella":
            self._normalized_gene_quanti_without_rrna = self._normalized_gene_quanti[
                self._normalized_gene_quanti["Feature"] != "rRNA"]
            self._normalized_gene_quanti_rrna = self._normalized_gene_quanti[
                self._normalized_gene_quanti["Feature"] == "rRNA"]
        else:
            self._normalized_gene_quanti_without_rrna = self._normalized_gene_quanti[
                ~self._normalized_gene_quanti["Attributes"].str.contains("=rRNA|=Mt_rRNA")]
            self._normalized_gene_quanti_rrna = self._normalized_gene_quanti[
                self._normalized_gene_quanti["Attributes"].str.contains("=rRNA|=Mt_rRNA")]

    def write_rrna_table(self):
        self._normalized_gene_quanti_rrna.to_csv(
            self._rrna_file, sep="\t", index=False)

    def plot_histogram(self):
        count_inf = len(self._normalized_gene_quanti_rrna["ratio"][
            np.isinf(self._normalized_gene_quanti_rrna["ratio"])])
        count_nan = len(self._normalized_gene_quanti_rrna["ratio"][
            self._normalized_gene_quanti_rrna["ratio"].isnull()])
        non_inf = self._normalized_gene_quanti_rrna["ratio"][
            ~np.isinf(self._normalized_gene_quanti_rrna["ratio"])]
        non_inf_non_nan = non_inf[~non_inf.isnull()]
        count_below_1 = len(non_inf_non_nan[non_inf_non_nan < 1.0])
        count_1_to_10 = len(non_inf_non_nan[
            (non_inf_non_nan >= 1.0) & (non_inf_non_nan < 10.0)])
        count_10_and_more = len(non_inf_non_nan[non_inf_non_nan >= 10.0])
        with open(self._histogram_file, "w") as output_file:
            output_file.write(
                "Counts inf: {}\n"
                "Counts NaN (0.0 vs. 0.0): {}\n"
                "Counts < 1.0: {}\n"
                "Counts >= 1.0 and < 10: {}\n"
                "Counts <= 10: {}\n".format(
                    count_inf, count_nan, count_below_1, count_1_to_10,
                    count_10_and_more))

    def write_decreased_non_rrna(self):
        non_nan = self._normalized_gene_quanti_without_rrna[
            ~self._normalized_gene_quanti_without_rrna["ratio"].isnull()]
        non_nan = non_nan[non_nan["ratio"] > 1.0]
        non_nan_non_inf = non_nan[~np.isinf(non_nan["ratio"])]
        top_changed_non_rrna = non_nan_non_inf.sort(["ratio"], ascending=False)
        inf = non_nan[np.isinf(non_nan["ratio"])]
        inf.to_csv(self._top_changed_non_rrna_file, sep="\t")
        top_changed_non_rrna.to_csv(
            self._top_changed_non_rrna_file, mode="a", sep="\t", index=False)

    def generate_scatter_plot(self):
        x_without_rrna = self._normalized_gene_quanti_without_rrna[
            "HeLa-S3:WT:16_h:total_RNA"] * self._multiplier
        y_without_rrna = self._normalized_gene_quanti_without_rrna[
            "HeLa-S3:WT:16_h:rRNA_free"] * self._multiplier
        x_rrna = self._normalized_gene_quanti_rrna[
            "HeLa-S3:WT:16_h:total_RNA"] * self._multiplier
        y_rrna = self._normalized_gene_quanti_rrna[
            "HeLa-S3:WT:16_h:rRNA_free"] * self._multiplier

        x = self._normalized_gene_quanti[
            "HeLa-S3:WT:16_h:total_RNA"] * self._multiplier
        y = self._normalized_gene_quanti[
            "HeLa-S3:WT:16_h:rRNA_free"] * self._multiplier
        min_value = min([min(x), min(y)])
        max_value = max([max(x), max(y)]) * 1.1

        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.set_aspect(1)
        font = {'family': 'sans-serif', 'weight': 'normal', 'size': 7}
        matplotlib.rc('font', **font)
        plt.plot(x_without_rrna, y_without_rrna, "k.", alpha=0.4)
        plt.plot(x_rrna, y_rrna, "r.", alpha=0.5)

        ax.set_yscale('log')
        ax.set_xscale('log')
        spearman = stats.spearmanr(
            self._normalized_gene_quanti_without_rrna[
                "HeLa-S3:WT:16_h:total_RNA"],
            self._normalized_gene_quanti_without_rrna[
                "HeLa-S3:WT:16_h:rRNA_free"])

        plt.xlabel("16 h WT total")
        plt.ylabel("16 h WT RiboZero treated")
        plt.title("{} (r = {})".format(self._organism, round(spearman[0], 4)))

        # The ones that have infinitive enrichtment in the non-treated
        # sample won't be plotted. For those we generate representative random
        # dot.
        np.random.seed = 1
        no_of_inf = len(self._normalized_gene_quanti_rrna[
            np.isinf(self._normalized_gene_quanti_rrna["ratio"])])
        if no_of_inf != 0:
            x_inf = np.random.uniform(max(x)*0.001, max(x), no_of_inf)
            y_inf = np.random.uniform(0.001, 0.0001, no_of_inf)
            x_inf = np.random.uniform(max(x)*0.9, max(x), no_of_inf)
            y_inf = np.random.uniform(1000000, 100000, no_of_inf)
            ax2 = plt.axes([0.65, 0.15, 0.15, 0.15])
            plt.plot(x_inf, y_inf, "r.", alpha=0.5)
            ax2.set_xticklabels([])
            ax2.set_yticklabels([])
            ax2.grid(False)
            plt.setp(ax2, xticks=[], yticks=[])
        fig.savefig(self._scatter_plot_file)
        
main()
