'''
This file contains classes that are use for visualizing the results generated by Digest module
Author: Yuancong Wang
Email: wangyuancong@163.com
version: 1.0
'''

# import pandas as pd
import sys
import matplotlib.pyplot as plt
import numpy as np
import os
import codecs
from PyPDF2 import PdfFileReader, PdfFileWriter, PdfFileMerger


class Drawer(object):
    '''
    This class is responsible for drawing bar plots according to the data generated by Digest module
    '''

    def __init__(self, dir_summary):
        # set the output dir for storing the plots.
        self.output_dir = os.path.join(dir_summary, 'output_plots')
        if not os.path.exists(self.output_dir):
            os.mkdir(self.output_dir)

        # get the directory that contains summary files
        self.dir_summary = dir_summary
        self.summary_files = list()
        self.summary_file_paths = list()
        # self.res_dir = os.path.join(dir_summary,'res')

    def get_the_summary_paths(self):
        '''
        Obtaining the paths of the summary files
        :return: None
        '''

        files = os.listdir(self.dir_summary)

        for file in files:
            if file.startswith('digestion_summary'):
                path = os.path.join(self.dir_summary, file)
                self.summary_file_paths.append(path)

    def iter_plot_loci(self):
        '''
        iter every plot of loci
        :return: None
        '''
        for file in self.summary_file_paths:
            file_size = os.path.getsize(file)
            if file_size == 0:
                print('The file is empty, quit!')
                continue
            file_name = os.path.basename(file)
            print('plotting ' + file_name + '...')

            with open(file, 'r') as f:
                contents = f.readlines()
                #     print(contents)

                int_num_loci = list()

                name_e = contents[0].split('\t')
                num_loci = contents[1].split('\t')

                for loc in num_loci:
                    int_num_loci.append(int(loc))

                self.plot_loci(file_name, name_e, int_num_loci)


    def plot_loci(self, file_name, x, y):
        '''
        make bar plot for each summary file
        :param file:
        :param x: The list of enzymes names
        :param y: The number of each loci recognized by enzyme.
        :return: None
        '''
        file = file_name
        x = x
        y = y
        fig = plt.figure()
        plt.bar(x, y, 0.4, color="green")
        plt.xlabel("")
        plt.ylabel("The number of loci")
        # plt.show()
        name_ends = '_'.join(file.split('_')[-3:])
        fig_path = os.path.join(self.output_dir, "num_loci_{}.pdf".format(name_ends))

        # Save the plot file
        fig.savefig(fig_path, dpi=600, format='pdf')

    def iter_frag_num(self):
        '''
        iterring every file and plotting them
        :return: None
        '''
        for file in self.summary_file_paths:
            file_size = os.path.getsize(file)
            if file_size == 0:
                print('The file is empty, quit!')
                continue

            file_name = os.path.basename(file)
            print('plotting ' + file_name + ' ......')

            with open(file, 'r') as f:
                contents = f.readlines()

                labels_all = contents[2].split('\t')
                values_all = contents[3].split('\t')

                value_all_math = [float(v) for v in values_all]

                labels_range = contents[4].split('\t')
                values_range = contents[5].split('\t')

                values_range_math = [float(v) for v in values_range]

                # labels_all_num = labels_all[0:3:2]
                values_all_num = value_all_math[0:3:2]

                # labels_range_num = labels_range[0:3:2]
                values_range_num = values_range_math[0:3:2]

                # labels_all_cover = labels_range[1:4:2]
                values_all_cover = value_all_math[1:4:2]

                # labels_range_cover = labels_all[1:4:2]
                values_range_cover = values_range_math[1:4:2]

                self.plot_frag_num('frag_num', file_name, values_all_num, values_range_num)


    def plot_frag_num(self, file_name_starts, file_name, values_all_num, values_range_num):

        file_name = file_name
        name_ends = '_'.join(file_name.split('_')[-3:])

        labels = ['All Three Types', 'FB_BF Only']
        num_all = [int(value) for value in values_all_num]
        num_range = [int(value) for value in values_range_num]

        x = np.arange(len(labels))  # the label locations
        width = 0.35  # the width of the bars

        fig, ax = plt.subplots()
        rects1 = ax.bar(x - width / 2, num_all, width, label='All Fragments')
        rects2 = ax.bar(x + width / 2, num_range, width, label='Fragments in range')

        # Add some text for labels, title and custom x-axis tick labels, etc.

        ax.set_ylabel('Numbers of Fragments')
        ax.set_title(name_ends)
        ax.set_xticks(x)
        ax.set_xticklabels(labels)
        ax.legend()

        def autolabel(rects):
            """Attach a text label above each bar in *rects*, displaying its height."""
            for rect in rects:
                height = rect.get_height()
                ax.annotate('{}'.format(height),
                            xy=(rect.get_x() + rect.get_width() / 2, height),
                            xytext=(0, 3),  # 3 points vertical offset
                            textcoords="offset points",
                            ha='center', va='bottom')

        autolabel(rects1)
        autolabel(rects2)

        fig.tight_layout()

        fig_path = os.path.join(self.output_dir, "{}_{}.pdf".format(file_name_starts, name_ends))

        # Save the plot file
        fig.savefig(fig_path, dpi=600, format='pdf')

        # plt.show()

###########################
    def iter_frag_cov(self):
        '''
        iterring every file and plotting them
        :return: None
        '''
        for file in self.summary_file_paths:
            file_size = os.path.getsize(file)
            if file_size == 0:
                print('The file is empty, quit!')
                continue

            file_name = os.path.basename(file)
            print('plotting ' + file_name + ' ......')

            with open(file, 'r') as f:
                contents = f.readlines()

                labels_all = contents[2].split('\t')
                values_all = contents[3].split('\t')

                value_all_math = [float(v) for v in values_all]

                labels_range = contents[4].split('\t')
                values_range = contents[5].split('\t')

                values_range_math = [float(v) for v in values_range]

                # labels_all_cover = labels_range[1:4:2]
                values_all_cover = value_all_math[1:4:2]

                # labels_range_cover = labels_all[1:4:2]
                values_range_cover = values_range_math[1:4:2]
                self.plot_frag_cov('frag_cov', file_name, values_all_cover, values_range_cover)

    def plot_frag_cov(self, file_name_starts, file_name, values_all_cover, values_range_cover):

        file_name = file_name
        name_ends = '_'.join(file_name.split('_')[-3:])

        labels = ['All Three Types', 'FB_BF Only']
        num_all = [round(value, 3) for value in values_all_cover]
        num_range = [round(value, 3) for value in values_range_cover]

        x = np.arange(len(labels))  # the label locations
        width = 0.35  # the width of the bars

        fig, ax = plt.subplots()
        rects1 = ax.bar(x - width / 2, num_all, width, label='All Fragments')
        rects2 = ax.bar(x + width / 2, num_range, width, label='Fragments in range')

        # Add some text for labels, title and custom x-axis tick labels, etc.

        ax.set_ylabel('Coverage')
        ax.set_title(name_ends)
        ax.set_xticks(x)
        ax.set_xticklabels(labels)
        ax.legend()

        def autolabel(rects):
            """Attach a text label above each bar in *rects*, displaying its height."""
            for rect in rects:
                height = rect.get_height()
                ax.annotate('{}'.format(height),
                            xy=(rect.get_x() + rect.get_width() / 2, height),
                            xytext=(0, 3),  # 3 points vertical offset
                            textcoords="offset points",
                            ha='center', va='bottom')

        autolabel(rects1)
        autolabel(rects2)

        fig.tight_layout()

        fig_path = os.path.join(self.output_dir, "{}_{}.pdf".format(file_name_starts, name_ends))

        # Save the plot file
        fig.savefig(fig_path, dpi=600, format='pdf')

        # plt.show()


    def merge_pdf(self, output_filename='merged_result_file.pdf'):
        # 遍历目录下的所有pdf将其合并输出到一个pdf文件中，输出的pdf文件默认带书签，书签名为之前的文件名
        # 默认情况下原始文件的书签不会导入，使用import_bookmarks=True可以将原文件所带的书签也导入到输出的pdf文件中
        merger = PdfFileMerger()
        dir_path = self.output_dir

        res_files = os.listdir(dir_path)

        res_files_pdf = list()

        # for file in res_files:
        #
        #     if '.pdf' in file:
        #         res_files_pdf.append(file)
        res_files_pdf = [file for file in res_files if '.pdf' in file]

        res_file_paths = list()
        for res_file in res_files_pdf:
            res_file_path = os.path.join(self.output_dir, res_file)
            res_file_paths.append(res_file_path)

        if not res_files_pdf:
            print("No pdf file in the current directory.")
            sys.exit()

        for filepath in res_file_paths:
            f = codecs.open(filepath, 'rb')
            file_rd = PdfFileReader(f)
            short_filename = os.path.basename(os.path.splitext(filepath)[0])
            merger.append(file_rd, bookmark=short_filename, import_bookmarks=False)
            print('merging files {}...'.format(filepath))
            f.close()
        out_filename = os.path.join(self.output_dir, output_filename)
        merger.write(out_filename)
        print('out put merged files...'.format(out_filename))
        merger.close()


