#!/usr/bin/env python3
# @Author   : Yang Liu
# @FileName : plot_figure.py
# @Software : NANOME project
# @Organization : JAX Li Lab
# @Website  : https://github.com/TheJacksonLaboratory/nanome

"""
Plots and export data for nanome paper
"""
import argparse
import glob
import gzip
import os.path
import pickle
import re
from collections import defaultdict

import matplotlib.pyplot as plt
#### Additional libraries:
import numpy as np
import pandas as pd
import pybedtools
import seaborn as sns
from scipy import stats
from scipy.stats import pearsonr
from sklearn import metrics
from sklearn.metrics import precision_recall_curve

from nanocompare.eval_common import logger, set_log_debug_level, pic_base_dir
from nanocompare.global_settings import *


def set_style(font_scale=1.2):
    # sns.set()
    # sns.set(font_scale=1.15)
    # sns.set_style("white")

    # This sets reasonable defaults for font size for
    # a figure that will go in a paper
    sns.set_context("paper")
    sns.set(font_scale=font_scale)

    # Set the font to be serif, rather than sans
    # sns.set(font='serif')

    # Make the background white, and specify the
    # specific font family
    # sns.set_style("white", {
    #         "font.family": "serif",
    #         "font.serif" : ["Times", "Palatino", "serif"]
    #         })
    sns.set_style("white")


def grid_plot_correlation_matrix_for_fig5a(infn, removeDeepMod=False):
    """
    Plot the grid of corr COE, distribution and scatter plot based on input files
    :return:
    """
    ## Load data into df
    logger.debug(infn)
    df = pd.read_csv(infn, sep=',')

    sel_col = []
    rename_dict = {}
    for col in df.columns:
        if str(col).endswith('_freq'):
            sel_col.append(col)

            new_col_name = get_tool_name(str(col).replace('_freq', ''))
            rename_dict.update({str(col): new_col_name})
    df = df[sel_col]
    df = df.rename(columns=rename_dict)

    if removeDeepMod:
        # rename DeepMod column here
        df = df.drop("DeepMod", axis=1)

    num_col = len(df.columns)

    df = df.iloc[:, list(range(1, num_col)) + [0]]
    logger.debug(f'Load data = {len(df)}')

    # Rearrange columns by correlation from high to low
    corrdf = df.corr()
    # logger.info(corrdf)

    coe_series = corrdf.iloc[:, -1]
    # logger.info(coe_series[:-1].sort_values(ascending=False))

    orderedColumns = coe_series[:-1].sort_values(ascending=False).index.tolist()

    df = df.loc[:, ['BGTruth'] + orderedColumns]

    # Rename to BS-seq name based on WR's suggestions
    df = df.rename(columns={'BGTruth': 'BS-seq'})

    ## Plot correlation grid figure
    basefn = os.path.basename(infn)
    outfileName = "{}.jpg".format("fig.5a.s5abc." + basefn.replace(".csv", ""))

    if args.o:
        outdir = args.o
    else:
        outdir = pic_base_dir

    outfn = os.path.join(outdir, outfileName)

    plt.clf()

    fig, ax = plt.subplots()

    fig.set_size_inches(12, 12)

    left, width = .25, .5
    bottom, height = .25, .5
    right = left + width
    top = bottom + height

    position = 1
    for yrow in range(1, num_col + 1):  # row
        for xcol in range(1, num_col + 1):  # column, each tool
            # adding this so that one would be able to estimate how much time one has for a coffe
            logger.debug(f"Please wait, processing column xcol {xcol}, yrow {yrow}.")
            if xcol == yrow:
                # Diagonal, distribution lines:
                plt.subplot(num_col, num_col, position)
                params = {'legend.fontsize': 50,
                          'legend.handlelength': 0,
                          'legend.handletextpad': 0,
                          'legend.fancybox': False,
                          'legend.loc': 'upper right',
                          'legend.framealpha': 0,
                          'legend.borderaxespad': 0}
                plt.rcParams.update(params)
                ax = sns.kdeplot(df.iloc[:, xcol - 1], shade=False, color="black", legend=True)

                tool_fontsize = 10
                tool_fontsize = 14
                leg = ax.legend(labels=[str(df.columns[xcol - 1])], fontsize=tool_fontsize)
                for item in leg.legendHandles:
                    item.set_visible(False)

                # ax.annotate("abcd", xy=(0.2, 0.8), xycoords='axes fraction',
                #             fontsize=14)
                ax.set_yticklabels([])
                ax.set_xticklabels([])
                ax.set_ylabel('')
                ax.set_xlabel('')
                ax.tick_params(axis=u'both', which=u'both', length=0)
            elif xcol > yrow:
                # upper triangle, COE number:
                ax2 = plt.subplot(num_col, num_col, position)

                corrValue = pearsonr(df.iloc[:, xcol - 1], df.iloc[:, yrow - 1])
                corrValueStr = "{0:.3f}".format(corrValue[0])

                coe_font_size = 19
                coe_font_size = 22

                if yrow == 1:
                    ax2.text(0.5 * (left + right), 0.5 * (bottom + top), corrValueStr,
                             horizontalalignment='center',
                             verticalalignment='center',
                             fontsize=coe_font_size, color='#7B2323', weight='bold',
                             transform=ax2.transAxes)
                else:
                    ax2.text(0.5 * (left + right), 0.5 * (bottom + top), corrValueStr,
                             horizontalalignment='center',
                             verticalalignment='center',
                             fontsize=coe_font_size, color='black',
                             transform=ax2.transAxes)

                ax2.set_yticklabels([])
                ax2.set_xticklabels([])
                ax2.tick_params(axis=u'both', which=u'both', length=0)

            elif xcol < yrow:
                # lower triangle, scatter plot, using hexbin, x-column label, y-row label:
                ax3 = plt.subplot(num_col, num_col, position)

                ### Kernel density plot:
                BGres = 300j  # 150j # 75j # 300j ## for debugging change this parameter to "50j" or less

                m1_sign = np.array(list(df.iloc[:, xcol - 1]))
                m2_sign = np.array(list(df.iloc[:, yrow - 1]))
                xmin = ymin = -0.05
                xmax = ymax = 1.05

                X, Y = np.mgrid[xmin:xmax:BGres, ymin:ymax:BGres]
                positions = np.vstack([X.ravel(), Y.ravel()])
                values = np.vstack([m1_sign, m2_sign])
                kernel = stats.gaussian_kde(values)
                Z = np.reshape(kernel(positions).T, X.shape)

                ax3.imshow(np.rot90(Z), cmap=plt.cm.Spectral_r, extent=[xmin, xmax, ymin, ymax], aspect="auto")
                # ax3.imshow(np.rot90(Z), cmap=plt.cm.gist_stern, extent=[xmin, xmax, ymin, ymax], aspect="auto") # An example how to switch from Spectral_r to gist_stern colormap
                # ax3.plot(m1_sign, m2_sign, marker=',', color="white", linestyle='None', markersize=0.25, alpha=0.25) # uncomment this line if you want to generate version with points marker very lightly on the figure

                ### Scatter plot (previous implementation for the plot - i am keeping this for now but should be cleaned if we decide to go with KDE version)
                # if len(df) < 100:
                #                     gridRes = 15  # for HL60
                #                 else:
                #                     gridRes = 50
                #
                #                 mincnt = 1
                #
                #                 if len(df) < 100:  # for HL60 plot
                #                     plt.hexbin(df.iloc[:, xcol - 1], df.iloc[:, yrow - 1], gridsize=(gridRes, gridRes), cmap='Blues', bins='log', mincnt=None)
                #                     pass
                #                 else:
                #                     plt.hexbin(df.iloc[:, xcol - 1], df.iloc[:, yrow - 1], gridsize=(gridRes, gridRes), cmap='Blues', bins='log', mincnt=mincnt)
                tick_label_fontsize = 10
                tick_label_fontsize = 12
                if yrow == num_col:  # last row scatter plot shows x ticks
                    plt.xticks([0, 1], fontsize=tick_label_fontsize)
                    ax3.set_xticklabels(["0%", "100%"], fontsize=tick_label_fontsize)
                    for label, alnType in zip(ax3.get_xticklabels(), ['left', 'center']):
                        label.set_horizontalalignment(alnType)
                else:
                    plt.xticks([], fontsize=10)

                if xcol == 1:  # first column scatter plot shows y ticks
                    plt.yticks([0, 1], fontsize=tick_label_fontsize)
                    ax3.set_yticklabels(["0%", "100%"], fontsize=tick_label_fontsize)
                    for label, alnType in zip(ax3.get_yticklabels(), ['bottom', 'top']):
                        label.set_verticalalignment(alnType)
                else:
                    plt.yticks([], fontsize=10)

            position += 1

    fig.savefig(outfn, dpi=300, bbox_inches='tight')
    fig.savefig(outfn.replace(".jpg", ".pdf"), dpi=300, bbox_inches='tight')  # Generate also PDF version
    # plt.show()
    plt.close()
    logger.info(f"save to {outfn}")


def gen_figure_5a(infn):
    """
    Generate figure 5-a in paper
    Correlation plots
    :return:
    """
    grid_plot_correlation_matrix_for_fig5a(infn)
    pass


def plot_ROC_PR_curves(ret, outdir, tagname="tagname", removeDeepMod=False):
    """
    Plot ROC AUC and PR curves
    :param ret:
    :param outdir:
    :param tagname:
    :return:
    """
    figure_size = (4, 4)

    # title_font_size = 18
    label_font_size = 14

    ## Plot ROC curve
    plt.clf()
    plt.figure(figsize=figure_size)

    csfont = {'fontsize': label_font_size}

    conf_matrix = {}
    for toolname, toolcolor in zip(ToolNameList, ToolsColorList):
        if toolname == 'DeepMod' and removeDeepMod:
            continue
        ytrue = ret[f'{toolname}_true']
        ypred = ret[f'{toolname}_pred']
        yscore = ret[f'{toolname}_score']

        ## Construct comfusion matrix
        from sklearn.metrics import confusion_matrix
        conf_matrix[toolname] = confusion_matrix(ytrue, ypred, labels=[0, 1])

        # logger.debug(f'ytrue={len(ytrue)}, ypred={len(ypred)}, yscore={len(yscore)}')
        # TODO: plot points in curves
        fpr, tpr, threshold = metrics.roc_curve(ytrue, yscore, drop_intermediate=False)
        roc_auc = metrics.auc(fpr, tpr)
        plt.plot(fpr, tpr, toolcolor, label=f'{toolname}={roc_auc:.2f}')
    leg = plt.legend(loc='lower right', title="AUC")
    leg._legend_box.align = "left"
    # plt.plot([0, 1], [0, 1], 'r--')
    plt.xlim([0, 1])
    plt.ylim([0, 1])

    plt.ylabel('True Positive Rate', **csfont)
    plt.xlabel('False Positive Rate', **csfont)

    # plt.title('ROC Curves', fontsize=title_font_size)
    outfn = os.path.join(outdir, f'fig.3b.{tagname}.roc.curves.jpg')
    plt.savefig(outfn, format='jpg', bbox_inches='tight', dpi=300)
    plt.savefig(outfn.replace(".jpg", ".pdf"), dpi=300, bbox_inches='tight')  # Generate also PDF version

    # print(plt.rcParams["font.family"])
    # print(plt.rcParams['font.sans-serif'])

    # plt.show()
    plt.close()

    outfn = os.path.join(outdir, f'confusion.matrix.{tagname}.each.tool.csv.gz')
    outf = gzip.open(outfn, 'wt')
    for toolname in ToolNameList:
        if toolname not in conf_matrix:
            continue
        # logger.info(f"{toolname} = {conf_matrix[toolname]}")
        outf.write(f"{toolname}\n")
        mat = conf_matrix[toolname]
        outf.write(f"{mat[0, 0]},{mat[0, 1]}\n")
        outf.write(f"{mat[1, 0]},{mat[1, 1]}\n")
    outf.close()
    return

    ## PR curves
    plt.clf()
    plt.figure(figsize=figure_size)
    for toolname, toolcolor in zip(ToolNameList, ToolsColorList):
        if toolname == 'DeepMod' and removeDeepMod:
            continue
        ytrue = ret[f'{toolname}_true']
        yscore = ret[f'{toolname}_score']

        precision, recall, _ = precision_recall_curve(ytrue, yscore)
        average_precision = metrics.average_precision_score(ytrue, yscore)
        plt.plot(recall, precision, toolcolor, label=f'{toolname}={average_precision:.2f}')
    plt.legend(loc='lower left')
    plt.xlim([0, 1])
    plt.xlabel('Recall', fontsize=label_font_size)
    plt.ylabel('Precision', fontsize=label_font_size)
    plt.ylim([0.0, 1.05])
    # plt.title('Precision-Recall Curves', fontsize=title_font_size)

    outfn = os.path.join(outdir, f'fig.3b.{tagname}.pr.curves.jpg')
    plt.savefig(outfn, format='png', bbox_inches='tight', dpi=600)

    plt.show()
    plt.close()


def collect_singleton_vs_nonsingleton_df(runPrefix, pattern="*.summary.bsseq.singleton.nonsingleton.cov1.csv"):
    dflist = []
    # logger.debug(runPrefix)
    for run1 in runPrefix:
        filepat = os.path.join(runPrefix[run1], pattern)
        # logger.debug(run1)
        flist = glob.glob(filepat)
        if len(flist) != 1:
            logger.error(
                f"Too much/No summary of singleton vs non-singleton for {run1} in folder {runPrefix[run1]} with pattern={pattern}, len={len(flist)}")
        # logger.debug(f'Get file:{flist[0]}')
        if len(flist) >= 1:
            logger.info(f"read file: {flist[0]}")
            df = pd.read_csv(flist[0], index_col=0)
            dflist.append(df)
    retdf = pd.concat(dflist)
    retdf.index.name = 'Dataset'
    return retdf


def collect_distribution_genomic_df(runPrefix, pattern=None):
    dflist = []
    logger.debug(runPrefix)
    for run1 in runPrefix:
        filepat = os.path.join(runPrefix[run1], pattern)
        # logger.debug(run1)
        flist = glob.glob(filepat)
        if len(flist) != 1:
            logger.error(
                f"Too much/No summary of singleton vs non-singleton for {run1} in folder {runPrefix[run1]} with pattern={pattern}, len={len(flist)}")

        # logger.debug(f'Get file:{flist[0]}')
        if len(flist) >= 1:
            logger.info(f"read file: {flist[0]}")
            df = pd.read_excel(flist[0], index_col=0, engine='openpyxl')  #
            dflist.append(df)
    retdf = pd.concat(dflist)

    retdf['Dataset'] = pd.Categorical(retdf['Dataset'], datasets_order)
    retdf['Coord'] = pd.Categorical(retdf['Coord'], locations_order)
    retdf = retdf.sort_values(by=['Dataset', 'Coord'], ascending=[True, True])
    retdf = retdf.dropna()

    return retdf


def collect_performance_report_as_df(runPrefix):
    """
    create report from list of runPrefix, return specified columns
    :return:
    """
    dflist = []
    for runKey in runPrefix:
        logger.debug(f'runPrefix={runKey}')
        pattern = os.path.join(runPrefix[runKey], "performance?results", "*.performance.report.csv")
        files = glob.glob(pattern)
        for infn in files:
            df = pd.read_csv(infn, index_col=0, sep=",")
            dflist.append(df)
            logger.debug(f'Collect data from {runKey}:{os.path.basename(infn)} = {len(df)}')
    if len(dflist) == 0:
        raise Exception(f"Can not find report at {runKey} in folder {runPrefix[runKey]} with pattern={pattern}")
    combdf = pd.concat(dflist, ignore_index=True)
    logger.info(f'We collected total {len(dflist)} files with {len(combdf)} records')
    return combdf


def load_wide_format_performance_results(runPrefix,
                                         sel_locations=locations_category + locations_singleton + locations_new):
    """
    Collect the currently new performance of exp results for paper
    :return:
    """
    df = collect_performance_report_as_df(runPrefix)
    retdf = df[df['Location'].isin(sel_locations)]

    logger.debug(
        f"load_wide_format_performance_results, wide-format retdf={len(retdf)} using locations={sel_locations}")

    return retdf


def save_wide_format_performance_results(runPrefix, outdir, tagname):
    """
    Save all performance report results into a csv
    :return:
    """
    df = load_wide_format_performance_results(runPrefix)

    ## Rename Tool name to starndard name
    df["Tool"].replace({"Megalodon_ZW": "Megalodon", "DeepMod_C": "DeepMod"}, inplace=True)

    ## arrange df
    df['Dataset'] = pd.Categorical(df['Dataset'], datasets_order)
    df['Tool'] = pd.Categorical(df['Tool'], ToolNameList)
    df['Location'] = pd.Categorical(df['Location'], locations_order)
    df = df.sort_values(by=['Dataset', 'Location', 'Accuracy'], ascending=[True, True, False])

    outfn = os.path.join(outdir, f'performance-results{f"-{tagname}" if tagname else ""}.csv')
    df.to_csv(outfn)
    logger.info(f'save to {outfn}')


def parse_arguments():
    parser = argparse.ArgumentParser(prog='plot_figure (NANOME)', description='Plot and export data for Nanocompare paper.')
    parser.add_argument('-v', '--version', action='version', version=f'%(prog)s v{nanome_version}')
    parser.add_argument("cmd", help="name of command, fig5a, export-data, etc.")
    parser.add_argument('-i', nargs='+', help='list of input files', default=[])
    parser.add_argument('-o', type=str, help="output dir",
                        default=None)  # TODO: check all correct when change this to None
    parser.add_argument('--tagname', type=str, help="tagname of files", default=None)
    parser.add_argument('--beddir', type=str, help="bed file dir used for finding Concordant and Discordant",
                        default=None)
    parser.add_argument('--cutoff', type=int, help="the cutoff used on bed file", default=1)
    parser.add_argument('--signal-col', type=int, help="the column number of signals extract from bed to bedGraph",
                        default=6)
    parser.add_argument('--cov-col', type=int, help="the column number of signals extract from bed to bedGraph",
                        default=7)  # -1 means no cov col

    args = parser.parse_args()
    return args


if __name__ == '__main__':
    set_log_debug_level()

    ## Set tmp dir for bedtools
    bedtool_tmp_dir = "/fastscratch/liuya/nanocompare/bedtools_tmp"
    os.makedirs(bedtool_tmp_dir, exist_ok=True)
    pybedtools.helpers.set_tempdir(bedtool_tmp_dir)

    ## Use same font Arial for all plotting
    import matplotlib.font_manager as font_manager

    # font_manager._rebuild() # some times failed, must be sure change font
    plt.rcParams['font.sans-serif'] = ['Arial']

    args = parse_arguments()
    logger.debug(args)

    if args.cmd == 'fig5a':
        ## find ${resultDir} -name 'Meth_corr_plot_data*.csv' \
        ##      -type f -exec python plot_figure.py fig5a -i {} \;
        ## python plot_figure.py fig5a -i /projects/li-lab/Nanopore_compare/result/MethCorr-NA19240_RRBS/Meth_corr_plot_data_joined-NA19240_RRBS-bsCov5-minToolCov3-baseFormat1.csv

        if not args.o:
            outdir = pic_base_dir
        for fn in args.i:
            gen_figure_5a(fn)
    elif args.cmd == 'figs5ab-data':
        # Fig. S5A-B data: PCC in each regions
        if not args.o:
            args.o = pic_base_dir
        dflist = []
        for indir in args.i:
            # file like: NA12878.corrdata.coe.pvalue.each.regions.xlsx
            fnlist = glob.glob(os.path.join(indir, '*.corrdata.coe.pvalue.each.regions.xlsx'))
            if len(fnlist) != 1:
                raise Exception(f'Found more or none fnlist={fnlist}, for indir={indir}')
            logger.info(f'Find file: {fnlist[0]}')

            df = pd.read_excel(fnlist[0], index_col=0, engine='openpyxl')
            dflist.append(df)
        outdf = pd.concat(dflist)
        outfn = os.path.join(args.o,
                             f'All.corrdata.coe.pvalue.each.regions{f".{args.tagname}" if args.tagname else ""}.xlsx')
        outdf.to_excel(outfn)
        logger.info(f'save to {outfn}')
    elif args.cmd == 'export-data':
        # Table S3 data
        # using global setting variable runPrefixDict to locate all data
        ## python plot_figure.py export-data -i /projects/li-lab/Nanopore_compare/result/MethPerf-APL_RRBS_CPG /projects/li-lab/Nanopore_compare/result/MethPerf-HL60_RRBS_CPG /projects/li-lab/Nanopore_compare/result/MethPerf-K562_WGBS_CPG --tagname CPG

        ## python plot_figure.py export-data -i /projects/li-lab/yang/results/2021-02-22/Meth
        if len(args.i) == 0:
            raise Exception(f'No run dir specified')
        else:
            run_prefix = defaultdict()
            for dirname in args.i:
                run_prefix[os.path.basename(dirname)] = dirname
        if not args.o:
            args.o = pic_base_dir

        logger.info(run_prefix)
        save_wide_format_performance_results(run_prefix, args.o, args.tagname)

        # pattern1 = "*.summary.bsseq.singleton.nonsingleton.cov1.csv"
        # df = collect_singleton_vs_nonsingleton_df(run_prefix, pattern=pattern1)
        # ## arrange df
        # df = df.reindex(datasets_order)
        #
        # outfn = os.path.join(args.o,
        #                      f'dataset.singleton.vs.non-singleton{f".{args.tagname}" if args.tagname else ""}.cov1.csv')
        # df.to_csv(outfn)

        pattern2 = "*.summary.bsseq.singleton.nonsingleton.cov5.table.s2.csv"
        df = collect_singleton_vs_nonsingleton_df(run_prefix, pattern=pattern2)
        ## arrange df
        df = df.reindex(datasets_order)

        outfn = os.path.join(args.o,
                             f'dataset.singleton.vs.non-singleton{f".{args.tagname}" if args.tagname else ""}.cov5.table.s2.xlsx')
        df.to_excel(outfn)
        logger.info(f'save stats of singleton and non-singleton to {outfn}')

        ## concat all singleton/nonsingleton for each datasets
        ## sample: HL60.bgtruth.certain.sites.distribution.sing.nonsing.each.genomic.cov5.xlsx
        pattern3 = "*.bgtruth.certain.sites.distribution.sing.nonsing.each.genomic.cov5.table.s6.xlsx"
        df = collect_distribution_genomic_df(run_prefix, pattern=pattern3)
        ## arrange df
        # df['Dataset'] = pd.Categorical(df['Dataset'], datasets_order)
        # df['Coord'] = pd.Categorical(df['Coord'], locations_order)
        # df = df.sort_values(by='Dataset')

        outfn = os.path.join(args.o,
                             f'all.certain.sites.distribution.each.genomic.region{f".{args.tagname}" if args.tagname else ""}.cov5.table.s6.csv')
        df.to_csv(outfn)

        logger.info(f'save distribution of each genomic region to {outfn}')
    elif args.cmd == 'export-curve-data':
        # Figure 3B data
        ## python plot_figure.py export-curve-data -i /projects/li-lab/Nanopore_compare/result/MethPerf-HL60_RRBS /projects/li-lab/Nanopore_compare/result/MethPerf-K562_WGBS

        ## python plot_figure.py export-curve-data -i /projects/li-lab/Nanopore_compare/result/MethPerf-APL_RRBS_CPG /projects/li-lab/Nanopore_compare/result/MethPerf-HL60_RRBS_CPG /projects/li-lab/Nanopore_compare/result/MethPerf-K562_WGBS_CPG --tagname CPG
        if not args.o:
            args.o = pic_base_dir
        outdir = os.path.join(args.o, f'plot-curve-data{f"-{args.tagname}" if args.tagname else ""}')
        os.makedirs(outdir, exist_ok=True)

        if len(args.i) == 0:
            raise Exception("No run dir specified")

        for bdir in args.i:
            runPrefix = os.path.basename(bdir)
            cnt = 0
            for coordinate_bed_name in location_filename_to_abbvname.keys():
                curve_data = defaultdict(list)
                success = True
                for toolname in ToolNameList:
                    pattern_str = os.path.join(bdir, 'performance?results', 'curve_data',
                                               f'*.{toolname}*.*{coordinate_bed_name}.curve_data.pkl')
                    fnlist = glob.glob(pattern_str)
                    if len(fnlist) != 1:
                        logger.error(
                            f'Can not locate curve_data for Tool={toolname}, at Coord={coordinate_bed_name}, with pattern={pattern_str}, find results={fnlist}. Please check if MethPerf results folder={bdir} specified is correct.')
                        success = False
                        break
                    cnt += 1
                    with open(fnlist[0], 'rb') as f:
                        ret = pickle.load(f)
                        # logger.debug(ret)
                        curve_data[f'{toolname}_true'] = ret['yTrue']  # Ground truth label
                        curve_data[f'{toolname}_pred'] = ret['yPred']  # Prediction label
                        curve_data[f'{toolname}_score'] = ret['yScore']  # Prediction score, used for roc curve plotting
                if success:
                    outfn = os.path.join(outdir,
                                         f'{runPrefix}.plot.curve.data.ytrue.ypred.yscore.{location_filename_to_abbvname[coordinate_bed_name]}.pkl')
                    with open(outfn, 'wb') as f:
                        pickle.dump(curve_data, f)

                    outfn = os.path.join(outdir,
                                         f'{runPrefix}.plot.curve.data.ytrue.ypred.yscore.{location_filename_to_abbvname[coordinate_bed_name]}.dat')
                    with open(outfn, "w") as f:
                        for toolname in ToolNameList:
                            f.write(f"{toolname}_true:")
                            outstr = ','.join([str(value) for value in curve_data[f'{toolname}_true']])
                            f.write(outstr)
                            f.write("\n")

                            f.write(f"{toolname}_pred:")
                            outstr = ','.join([str(value) for value in curve_data[f'{toolname}_pred']])
                            f.write(outstr)
                            f.write("\n")

                            f.write(f"{toolname}_score:")
                            outstr = ','.join([f'{value:.4f}' for value in curve_data[f'{toolname}_score']])
                            f.write(outstr)
                            f.write("\n")
            logger.info(f'For runPrefix={runPrefix} at dir={bdir}, total files={cnt}')
    elif args.cmd == 'plot-curve-data':
        # Fig. 3B
        ## find /projects/li-lab/Nanopore_compare/result/plot-curve-data -name '*.pkl' -exec python plot_figure.py plot-curve-data -i {} \;
        if not args.o:
            args.o = pic_base_dir

        outdir = os.path.join(args.o, f'curves-figures{f"-{args.tagname}" if args.tagname else ""}')
        os.makedirs(outdir, exist_ok=True)
        for fn in args.i:
            logger.debug(f'Plot data from fn={fn}')
            with open(fn, 'rb') as infn:
                basename = os.path.basename(fn)
                bn = os.path.splitext(basename)[0]
                ret = pickle.load(infn)
                # logger.debug(ret.keys())
                plot_ROC_PR_curves(ret, outdir, tagname=bn)
    elif args.cmd == 'guppy-qos':
        ## collect basecall output for summary results, used for qos
        ## python plot_figure.py guppy-qos -i /fastscratch/liuya/nanocompare/HL60-Runs/HL60-N50-basecall /fastscratch/liuya/nanocompare/K562-Runs/K562-N50-basecall /projects/li-lab/Nanopore_compare/nanopore_fast5/NA19240-N300-basecall /fastscratch/liuya/nanocompare/APL-Runs/APL-N50-basecall
        for basedir in args.i:
            fpatstr = os.path.join(basedir, '*', 'sequencing_summary.txt')
            flist = glob.glob(fpatstr)
            # logger.debug(flist)
            dflist = []
            for fn in flist:
                df = pd.read_csv(fn, sep='\t')
                dflist.append(df)
            outdf = pd.concat(dflist)
            logger.info(outdf)

            bfn = os.path.basename(basedir)
            outfn = os.path.join(pic_base_dir, f'{bfn}.sequencing_summary.txt')
            outdf.to_csv(outfn, index=False, sep='\t')
    elif args.cmd == 'bed-to-bedGraph':
        # Convert bed into bedGraph, extract specific column signals (5mC, or 5hmC)
        # Input is 0-based, output is same 0-based for TSS analysis
        # python plot_figure.py bed-to-bedGraph -i $fn --cutoff $CUTOFF -o $outdir --cov_col 7  --signal-col 6
        # python plot_figure.py bed-to-bedGraph -i $fn --cutoff 5 --cov_col -1  --signal-col 3
        ## Default: cov_col=7, signal_col=6
        infn = args.i[0]
        cutoff = args.cutoff
        df = pd.read_csv(infn, sep='\t', header=None)
        logger.info(len(df))

        if args.cov_col != -1:  # filter with coverage
            df = df[df.iloc[:, args.cov_col] >= cutoff]
            logger.info(f'After cutoff={cutoff}, len(df)={len(df)}')

        if not args.o:  # output dir params
            outdir = pic_base_dir
        else:
            outdir = args.o

        baseInfn = os.path.basename(infn)
        baseOutfn = re.sub('cov..bed.gz', '', baseInfn) + f"cov{cutoff}.bedGraph"

        outfn = os.path.join(outdir, baseOutfn)
        df = df.iloc[:, [0, 1, 2, args.signal_col]]
        df.to_csv(outfn, sep='\t', header=False, index=None)
        logger.info(f'save to {outfn}')

        from pybedtools import BedTool

        bed1 = BedTool(outfn)
        bed1 = bed1.sort()
        outfn2 = outfn.replace('.bedGraph', '.sorted.bedGraph')
        bed1.saveas(outfn2)
        logger.info(f'after sort bed, save to {outfn2}')
    elif args.cmd == 'view-outmatrix':  # View the out marix for TSS, such as /projects/li-lab/yang/results/2021-04-25/tss-plots/
        ### python plot_figure.py view-outmatrix -i /projects/li-lab/yang/results/2021-04-25/tss-plots/NA19240.bin50.outMatrix.tsv
        infn = args.i[0]
        df = pd.read_csv(infn, sep='\t', compression='gzip', header=None, skiprows=[0], na_values='nan')
        logger.info(df)

        if args.o:
            outdir = args.o
        else:
            outdir = pic_base_dir
        outfn = os.path.basename(infn).replace('outMatrix.tsv', 'outMatrix.for.R.csv.gz')
        outfn = os.path.join(outdir, outfn)

        df.to_csv(outfn, index=False, header=False, sep=',', compression='gzip')
        logger.info(f'save to {outfn}')

        pass
    else:
        raise Exception(f'Command={args.cmd} is not support')

    logger.info('Plot data script DONE')
