#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# ------------------------------------------------------------------------------
# 
# MIT License
# 
# Copyright (c) 2017 Gabriele Girelli
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
# 
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
# 
# Author: Gabriele Girelli
# Email: gigi.ga90@gmail.com
# Project: GPSeq
# Description: estimate region centrality from GPSeq sequencing data.
# 
# ------------------------------------------------------------------------------



# DEPENDENCIES =================================================================

import matplotlib
matplotlib.use('ps')

import argparse
import datetime
from matplotlib import pyplot as pp
from matplotlib.backends.backend_pdf import PdfPages
import numpy as np
import os
import pickle
import sys
from tqdm import tqdm

from ggc.args import check_threads, export_settings
from ggc.prompt import ask

import gpseqc.compare as compare
from gpseqc.compare import RankTable, MetricTable
from gpseqc.compare import dKT_iter, dKTw_iter
from gpseqc.compare import plot_comparison, plot_heatmap

# PARAMETERS ===================================================================

parser = argparse.ArgumentParser(description = """
Description:
 Centrality ranking comparison: calculate distance between centrality
 rankings. Please, note that the comparison is performed only on the
 intersection of the two rankings. If they do not intersect, an error message is
 displayed.

Notes:

 # Distance --------------------------------------------------------------------

  The Kendall tau (kt) distance counts the number of swaps required to convert
  one ranking into the other. The Kendall tau weighted (ktw) distance does the
  same, while taking into consideration also the weights used to rank the items.
  The Earth Mover's Distance (EMD) is an alternative approach that takes into
  account both the order and the weights. For more details, refer to the online
  documentation: https://github.com/ggirelli/gpseqc

 # -----------------------------------------------------------------------------
""", formatter_class = argparse.RawDescriptionHelpFormatter)

# Add mandatory arguments
parser.add_argument('rank1', type = str,
    help = '''Path to first ranking table ("estimated" table generated with
    gpseqc_estimate).''')
parser.add_argument('rank2', type = str,
    help = '''Path to second ranking table ("estimated" table generated with
    gpseqc_estimate).''')

# Add arguments with default value
parser.add_argument('-o', '--outdir', type = str, metavar = 'outdir',
    default = '.', help = """Path to output directory, created if missing.
    Defaults to current directory.""")
parser.add_argument('-d', '--dist', type = str, metavar = 'dist',
    help = """Whether 'kt' (Kendall tau), 'ktw' (weighted Kendall tau), or 'emd'
    (Earth Mover's Distance). Default: 'ktw'""",
    choices = list(compare.DISTANCE_FUNS.keys()), default = 'ktw')
parser.add_argument('-n', '--niter', type = int, metavar = 'niter',
    default = 5000, help = """Number of iterations to build the random
    distribution. Default: 5000""")
parser.add_argument('-t', '--threads', type = int, metavar = "threads",
    default = 1,
    help = """Number of threads for parallelization. Default: 1""")
parser.add_argument('-D', '--delim', type = str, metavar = 'delim',
    help = """Ranking file field delimiter. Default: TAB""", default = '\t')
parser.add_argument('-p', '--prefix', type = str, metavar = 'text',
    default = '', help = """Text for output file name prefix.""")
parser.add_argument('-s', '--suffix', type = str, metavar = 'text',
    default = '', help = """Text for output file name suffix.""")
parser.add_argument('-S', '--seed', type = int, metavar = 'seed',
    default = None, help = """Seed for random number generator.""")

# Flags
parser.add_argument('-y', '--do-all', action = 'store_const',
    help = """Do not ask for settings confirmation and proceed.""",
    const = True, default = False)
parser.add_argument('-C', '--custom-estimates', action = 'store_const',
    help = """Allow for custom estimate metrics.""",
    const = True, default = False)
parser.add_argument('--no-test', action = 'store_const',
    help = """Do not calculate p-values, only distances.""",
    const = True, default = False)
parser.add_argument('--no-plot', action = 'store_const',
    help = """Do not generate any plots.""",
    const = True, default = False)

# Version flag
version = "2.0.4"
parser.add_argument('--version', action = 'version',
    version = '%s v%s' % (sys.argv[0], version,))

# Parse arguments
args = parser.parse_args()

# Check input ------------------------------------------------------------------

args.threads = check_threads(args.threads)

assert os.path.isfile(args.rank1), "file not found: %s" % args.rank1
assert os.path.isfile(args.rank2), "file not found: %s" % args.rank2

assert_msg = "file expected, folder found: %s" % args.outdir
assert not os.path.isfile(args.outdir), assert_msg

assert args.niter >= 1, "at least 1 iteration required, got %d" % args.niter

if not os.path.isdir(args.outdir): os.mkdir(args.outdir)

# Adjust prefix/suffix if needed
if 0 != len(args.prefix):
    if '.' != args.prefix[-1]: args.prefix += '.'
if 0 != len(args.suffix):
    if '.' != args.suffix[0]: args.suffix = '.' + args.suffix

if not type(None) == type(args.seed):
    np.random.seed(args.seed)

# FUNCTION =====================================================================

def print_settings(args, clear = True):
    '''Show input settings, for confirmation.

    Args:
        args (Namespace): arguments parsed by argparse.
        clear (bool): clear screen before printing.
    '''
    s = " # GPSeqC - Ranking comparison v%s\n\n" % version

    s += " 1st rank : %s\n" % args.rank1
    s += " 2nd rank : %s\n" % args.rank2

    s += "\n Distance : %s\n" % args.dist
    s += "    nIter : %d\n" % args.niter
    s += "  Threads : %d\n" % args.threads

    s += "\n    Delim : '%s'\n" % args.delim
    if 0 != len(args.prefix): s += "   Prefix : '%s'\n" % args.prefix
    if 0 != len(args.suffix): s += "   Suffix : '%s'\n" % args.suffix

    s += "\n  Custom. : %r\n" % args.custom_estimates

    if type(None) != type(args.seed):
        s += "\n  Seed : %d\n" % args.seed

    if clear: print("\033[H\033[J%s" % s)
    else: print(s)
    return(s)

# RUN ==========================================================================

ssettings = print_settings(args)
if not args.do_all: ask("Confirm settings and proceed?")

args.suffix = ".rank_cmp.%s.%dniter%s" % (args.dist, args.niter, args.suffix)

# Save confirmed settings
with open(os.path.join(args.outdir, "%ssettings%s.txt" % (
    args.prefix, args.suffix)), "w+") as OH:
    export_settings(OH, ssettings)

# Read ranking tables
print("Reading and parsing rank tables...")
r1 = RankTable(args.rank1, allow_custom = args.custom_estimates)
r2 = RankTable(args.rank2, allow_custom = args.custom_estimates)

# Make intersection
print("Subsetting rank tables...")
r1 = r1.intersection(r2)
assert 0 != len(r1), "empty table left after subsetting."
r2 = r2.intersection(r1)
assert 0 != len(r2), "empty table left after subsetting."

# Distance label and function
dlabel = compare.DISTANCE_LABELS[args.dist]

if args.no_test:
    print("Calculating distance...")
    dtab = r1.compare(r2, args.dist, args.niter, skipSubset = True,
        progress = True, threads = args.threads)

    # Export DataFrames
    print("Exporting tsv tables...")
    path_tuple = (args.outdir, args.prefix, args.suffix)
    dtab.to_csv("%s/%sdist_table%s.tsv" % path_tuple, sep = args.delim)

    if not args.no_plot:
        # Plot
        plot_heatmap(dtab, [0, .5, 1], dlabel,
            outpath = "%s/%sdist%s.pdf" % path_tuple)

else: # Run comparison test
    print("Running comparison test...")
    dout = r1.test_comparison(r2, args.dist, args.niter, skipSubset = True,
        progress = True, threads = args.threads)

    # Export DataFrames --------------------------------------------------------
    print("Exporting tsv tables...")

    path_tuple = (args.outdir, args.prefix, args.suffix)
    dout["dist"].to_csv(
        "%s/%sdist_table%s.tsv" % path_tuple, sep = args.delim)
    dout["pval"].to_csv(
        "%s/%spval_table%s.tsv" % path_tuple, sep = args.delim)
    with open("%s/%soutput%s.pickle" % path_tuple, "wb") as OH:
        pickle.dump(dout, OH)

    if not args.no_plot:
        # Plot ---------------------------------------------------------------------
        print("Plotting...")

        # Recap heatmaps
        plot_heatmap(dout["dist"], [0, .5, 1], dlabel,
            outpath = "%s/%sdist%s.pdf" % path_tuple)
        plot_heatmap(dout["pval"], [0, 0.01, 0.05], "Distance p-value",
            outpath = "%s/%spval%s.pdf" % path_tuple)

        # Single comparison plots with Gaussian fit
        outpdf = PdfPages("%s/%snorm%s.pdf" % path_tuple)
        pgen = ((i, j) for i in dout["dist"].index
            for j in dout["dist"].columns)
        pgen = tqdm(pgen, total = dout["dist"].shape[0] * dout["dist"].shape[1])
        xlim = (0, 1) if args.dist != "emd" else None
        for (i, j) in pgen:
            title  = "Dist [%s]: %.4f\n" % (args.dist, dout["dist"].loc[i, j])
            title += "R1: %s; R2: %s" % (i, j)
            title += "\nn.perm=%d;" % args.niter
            title += " p.val: %.2E;" % dout["pval"].loc[i, j]
            plot_comparison(dout["dist"].loc[i, j],
                [d.loc[i, j] for d in dout["random_distribution"]],
                title, dlabel, xlim).savefig(outpdf, format = 'pdf')
            pp.close('all')
        outpdf.close()

# End --------------------------------------------------------------------------

################################################################################
