#! /usr/bin/env python

import argparse, sys, os, errno

import logging

logging.basicConfig(level=logging.INFO, format='[%(asctime)s] [%(levelname)s] %(name)s: %(message)s')

import numpy as np

'''

import matplotlib

matplotlib.use('Agg')

import matplotlib.pyplot as plt

from matplotlib.backends.backend_pdf import PdfPages

import seaborn as sns

sns.set()

'''

import pandas as pd

from pandas import DataFrame, Series

from scipy.fftpack import fft

from scipy.signal import convolve

import numba

command_handlers = {}

def command_handler(f):

    command_handlers[f.__name__] = f

    return f

def read_coverage(filename):

    coverage = []

    gene_ids = []

    with open(filename, 'r') as f:

        for line in f:

            c = line.strip().split('\t')

            gene_id = c[0]

            values = np.array(c[1:]).astype(np.float64)

            gene_ids.append(gene_id)

            coverage.append(values)

    return gene_ids, coverage

@numba.jit('int64(int32[:], int32[:], float64, float64, float64)')

def icm_update(x, y, h=0.0, beta=1.0, eta=2.1):

    n_changes = 0

    N = x.shape[0]

    for i in range(N):

        dx = -2*x[i]

        dE = 0

        if i > 0:

            dE += h*dx - beta*dx*x[i - 1] - eta*dx*y[i]

        if i < (N - 1):

            dE += h*dx - beta*dx*x[i + 1] - eta*dx*y[i]

        if dE < 0:

            x[i] = -x[i]

            n_changes += 1

    return n_changes

def icm_smooth(x, h=0.0, beta=1.0, eta=2.1):

    '''Smooth signals using iterated conditional modes

    Args:

        x: 1D signal

    Returns:

        Smoothed signal of the same length of x

    '''

    x = x*2 - 1

    y = x.copy()

    #E = h*np.sum(x) - beta*x[:-1]*x[1:] - eta*x*y

    n_updates = icm_update(x, y, h=h, beta=beta, eta=eta)

    while n_updates > 0:

        n_updates = icm_update(x, y, h=h, beta=beta, eta=eta)

    x = (x > 0).astype(np.int32)

    return x

def call_peak_gene(sig, local_bg_width=3, local_bg_weight=0.5, bg_global=None, smooth=False):

    if bg_global is None:

        bg_global = np.mean(sig)

    filter = np.full(local_bg_width, 1.0/local_bg_width)

    bg_local = convolve(sig, filter, mode='same')

    bg = local_bg_weight*bg_local + (1.0 - local_bg_weight)*bg_global

    bg[np.isclose(bg, 0)] = 1

    snr = sig/bg

    peaks = (snr > 1.0).astype(np.int32)

    if smooth:

        peaks = icm_smooth(peaks, h=-2.0, beta=4.0, eta=2.0)

    x = np.zeros(len(peaks) + 2, dtype=np.int32)

    x[1:-1] = peaks

    starts = np.nonzero(x[1:] > x[:-1])[0]

    ends = np.nonzero(x[:-1] > x[1:])[0]

    peaks = np.column_stack([starts, ends])

    return peaks

def estimate_bg_global(signals):

    '''signals

    '''

    signals_mean = np.asarray([np.mean(s) for s in signals])

    bg = np.median(signals_mean)

    return bg

def call_peaks(signals, min_length=2):

    #bg_global = estimate_bg_global(signals)

    bg_global = None

    peaks = []

    for i, signal in enumerate(signals):

        peak_locations = call_peak_gene(signal, bg_global=bg_global, smooth=True)

        #print(signal)

        for start, end in peak_locations:

            #print(peak_locations)

            if (min_length is None) or ((end - start) >= min_length):

                peaks.append((i, start, end, signal[start:end].mean()))

    return peaks

@command_handler

def call_peak(args):

    from tqdm import trange

    logger.info('read input file: ' + args.input_file)

    gene_ids, signals = read_coverage(args.input_file)

    if args.use_log:

        signals = [np.log10(np.maximum(1e-3, a)) + 3 for a in signals]

    signals_mean = np.asarray([np.mean(a) for a in signals])

    bg_global = np.median(signals_mean)

    logger.info('create output plot file: ' + args.output_file)

    with open(args.output_file, 'w') as fout:

        for i in trange(len(signals), unit='gene'):

            peaks_locations = call_peak_gene(signals[i], bg_global=bg_global, local_bg_weight=args.local_bg_weight,

                local_bg_width=args.local_bg_width, smooth=args.smooth)

            for j in range(peaks_locations.shape[0]):

                fout.write('{}\t{}\t{}\n'.format(gene_ids[i], peaks_locations[j, 0], peaks_locations[j, 1]))

@command_handler

def refine_peaks(args):

    import pandas as pd

    import numpy as np

    from tqdm import tqdm

    import re

    import h5py

    #from bx.bbi.bigwig_file import BigWigFile

    import pyBigWig

    from ioutils import open_file_or_stdout

    logger.info('read input matrix file: ' + args.peaks)

    #matrix = pd.read_table(args.matrix, sep='\t', index_col=0)

    #feature_info = matrix.index.to_series().str.split('|', expand=True)

    #feature_info.columns = ['gene_id', 'gene_type', 'gene_name', 'domain_id', 'transcript_id', 'start', 'end']

    #feature_info['start'] = feature_info['start'].astype('int')

    #feature_info['end'] = feature_info['end'].astype('int')

    peaks = pd.read_table(args.peaks, sep='\t', header=None, dtype=str)

    if peaks.shape[1] < 6:

        raise ValueError('less than 6 columns in peak file')

    peaks.columns = ['chrom', 'start', 'end', 'name', 'score', 'strand'] + ['c%d'%i for i in range(6, peaks.shape[1])]

    peaks['start'] = peaks['start'].astype('int')

    peaks['end'] = peaks['end'].astype('int')

    logger.info('read chrom sizes: ' + args.chrom_sizes)

    chrom_sizes = pd.read_table(args.chrom_sizes, sep='\t', header=None, names=['chrom', 'size'])

    chrom_sizes = chrom_sizes.drop_duplicates('chrom')

    chrom_sizes = chrom_sizes.set_index('chrom').iloc[:, 0]

    logger.info('read input genomic bigwig file: ' + args.tbigwig)

    tbigwig = pyBigWig.open(args.tbigwig)

    #chrom_sizes.update(dict(tbigwig.get_chrom_sizes()))

    gbigwig = {}

    logger.info('read input genomic bigwig (+) file: ' + args.gbigwig_plus)

    gbigwig['+'] = pyBigWig.open(args.gbigwig_plus)

    logger.info('read input genomic bigwig (-) file: ' + args.gbigwig_minus)

    gbigwig['-'] = pyBigWig.open(args.gbigwig_minus)

    #chrom_sizes.update(dict(gbigwig['+'].get_chrom_sizes()))

    flanking = args.flanking

    signals = []

    signals_mean = []

    windows = []

    #pat_gene_id = re.compile('^(.*)_([0-9]+)_([0-9]+)_([+-])$')

    for _, peak in peaks.iterrows():

        if peak['chrom'].startswith('chr'):

        #if feature['gene_type'] == 'genomic':

            #chrom, start, end, strand = pat_gene_id.match(feature['gene_id']).groups()

            #start = int(start)

            #end = int(end)

            window_start = max(0, peak['start'] - flanking)

            window_end = min(peak['end'] + flanking, chrom_sizes[peak['chrom']])

            data = np.nan_to_num(gbigwig[peak['strand']].values(peak['chrom'], window_start, window_end))

        else:

            strand = '+'

            window_start = max(0, peak['start'] - flanking)

            window_end = min(peak['end'] + flanking, chrom_sizes[peak['chrom']])

            data = np.nan_to_num(tbigwig.values(peak['chrom'], window_start, window_end))

        if data is None:

            data = np.zeros((window_end - window_start))

            #logger.info('no coverage data found for peak: {}'.format(feature['domain_id']))

        if args.use_log:

            data = np.log2(np.maximum(data, 0.25)) + 2

        signals.append(data)

        signals_mean.append(np.mean(data))

        windows.append((peak['chrom'], window_start, window_end, peak['start'], peak['end'], peak['strand']))

    tbigwig.close()

    gbigwig['+'].close()

    gbigwig['-'].close()

    windows = pd.DataFrame.from_records(windows)

    windows.columns = ['chrom', 'window_start', 'window_end', 'start', 'end', 'strand']

    logger.info('call peaks')

    refined_peaks = call_peaks(signals, min_length=args.min_length)

    with open_file_or_stdout(args.output_file) as fout:

        for i, start, end, mean_signal in refined_peaks:

            # map peak coordinates from window to original

            peak = [windows['chrom'][i], 

                start + windows['start'][i],

                end + windows['start'][i],

                'peak_%d'%(i + 1),

                '%.4f'%mean_signal,

                strand

            ]

            # remove peaks not overlapping with the window

            if (peak[1] > windows['end'][i]) or (peak[2] < windows['start'][i]):

                continue

            fout.write('\t'.join(map(str, peak)) + '\n')

        #print('%s\t%d\t%d => %s\t%d\t%d'%(

        #    windows['chrom'][i], windows['start'][i], windows['end'][i],

        #    windows['chrom'][i], start, end))

def _call_peaks_localmax(x, min_peak_length=10, bin_width=10, min_cov=5, decay=0.5):

    '''Call peaks by extending from local maxima

    Parameters:

    ----------

    x: array-like, (length,)

        Input signal values

    min_peak_length: int

        Minimum length required for each peak

    bin_width: int

        Bin width for searching bins with mean coverage higher than min_cov

    min_cov: float

        Minimum coverage to define a peak

    decay: float

        Stops extending a peak after signal values fall below decay*peak_summit

    Returns:

    -----------

    peaks: list of list

        Peaks found

        Each element of the list is a list: [start, end, local_max]

    '''

    # average signal over bins with 50% overlap

    half_bin_width = bin_width//2

    length = x.shape[0]

    n_bins = max(1, length//half_bin_width)

    bin_cov = np.zeros(n_bins)

    for i in range(n_bins):

        bin_cov[i] = np.mean(x[(i*half_bin_width):min((i + 2)*half_bin_width, length)])

    cand_bins = np.nonzero(bin_cov > min_cov)[0]

    n_cand_bins = cand_bins.shape[0]

    cand_bin_index = 0

    left_bound = 0

    peaks = []

    while cand_bin_index < n_cand_bins:

        i = cand_bins[cand_bin_index]*half_bin_width

        start = i

        end = i

        # find local max

        while (start > left_bound) and (x[start - 1] >= x[start]) and (x[start - 1] >= min_cov):

            start -= 1

        while (end < (length - 1)) and (x[end + 1] >= x[end]) and (x[end + 1] >= min_cov):

            end += 1

        max_index = 0

        if x[start] >= x[end]:

            local_max = x[start]

            max_index = start

        else:

            local_max = x[end]

            max_index = end

        if local_max > min_cov:

            # find bounds when input signal drops below 0.5*local_max

            start = max_index

            while (start > left_bound) and (x[start - 1] >= decay*local_max):

                start -= 1

                local_max = max(local_max, x[start])

            end = max_index

            while (end < (len(x) - 1)) and (x[end + 1] >= decay*local_max):

                end += 1

                local_max = max(local_max, x[end])

            # add current peak to results

            if (end - start) >= min_peak_length:

                #print((start, end, local_max))

                peaks.append([start, end, local_max])

            # find next candidate bin

            left_bound = end

            next_cand_bin = end//half_bin_width

            while (cand_bin_index < n_cand_bins) and (cand_bins[cand_bin_index] < next_cand_bin):

                cand_bin_index += 1

        cand_bin_index += 1

    return peaks

@command_handler

def call_peaks_localmax(args):

    import pyBigWig

    import numpy as np

    logger.info('read input file: ' + args.input_file)

    bigwig = pyBigWig.open(args.input_file)

    logger.info('write output file: ' + args.output_file)

    bed = open(args.output_file, 'w')

    chroms = bigwig.chroms()

    n_peaks = 0

    for chrom, size in chroms.items():

        if chrom.startswith('chr'):

            continue

        x = np.nan_to_num(bigwig.values(chrom, 0, size))

        peaks_chrom = _call_peaks_localmax(x, 

            min_peak_length=args.min_peak_length, bin_width=args.bin_width,

            min_cov=args.min_cov, decay=args.decay)

        for peak in peaks_chrom:

            n_peaks += 1

            #peaks.append([chrom, peak[0], peak[1], 'peak_%d'%n_peaks, peak[2], '+'])

            bed.write('%s\t%d\t%d\tpeak_%d\t%d\t+\n'%(chrom, peak[0], peak[1], n_peaks, peak[2]))

    bigwig.close()

    bed.close()

if __name__ == '__main__':

    main_parser = argparse.ArgumentParser(description='Call peaks from exRNA signals')

    subparsers = main_parser.add_subparsers(dest='command')

    parser = subparsers.add_parser('call_peak')

    parser.add_argument('--input-file', '-i', type=str, required=True,

        help='input file of exRNA signals for each transcript')

    parser.add_argument('--use-log', action='store_true', 

        help='use log10 instead raw signals')

    parser.add_argument('--smooth', action='store_true',

        help='merge adjacent peaks')

    parser.add_argument('--local-bg-width', type=int, default=3,

        help='number of nearby bins for estimation of local background')

    parser.add_argument('--local-bg-weight', type=float, default=0.5, 

        help='weight for local background (0.0-1.0)')

    parser.add_argument('--output-file', '-o', type=str, required=True,

        help='output plot file BED format')

    parser = subparsers.add_parser('call_peaks_localmax')

    parser.add_argument('--input-file', '-i', type=str, required=True,

        help='input BigWig file of raw reads coverage')

    parser.add_argument('--min-peak-length', type=int, default=10,

        help='minimum length required for a peak')

    parser.add_argument('--decay', type=float, default=0.5,

        help='decay factor of peak summit to define peak boundary')

    parser.add_argument('--min-cov', type=float, default=5,

        help='minimum coverage required to define a peak')

    parser.add_argument('--bin-width', type=int, default=10,

        help='bin width to search enriched bins')

    parser.add_argument('--output-file', '-o', type=str, required=True,

        help='output peaks in BED format')

    parser = subparsers.add_parser('refine_peaks')

    parser.add_argument('--peaks', type=str, required=True,

        help='input count matrix with feature names as the first column')

    parser.add_argument('--tbigwig', type=str, required=True,

        help='transcript BigWig file')

    parser.add_argument('--gbigwig-plus', type=str, required=True,

        help='genomic BigWig (+) file')

    parser.add_argument('--gbigwig-minus', type=str, required=True,

        help='genomic BigWig (-) file')

    parser.add_argument('--chrom-sizes', type=str, required=True,

        help='chrom sizes')

    parser.add_argument('--output-file', '-o', type=str, default='-',

        help='output refined peaks')

    parser.add_argument('--use-log', action='store_true', 

        help='use log10 instead raw signals')

    parser.add_argument('--smooth', action='store_true',

        help='merge adjacent peaks')

    parser.add_argument('--local-bg-width', type=int, default=3,

        help='number of nearby bins for estimation of local background')

    parser.add_argument('--local-bg-weight', type=float, default=0.5, 

        help='weight for local background (0.0-1.0)')

    parser.add_argument('--flanking', type=int, default=20)

    parser.add_argument('--min-length', type=int, default=10)

    args = main_parser.parse_args()

    if args.command is None:

        raise ValueError('empty command')

    logger = logging.getLogger('call_peak.' + args.command)

    command_handlers.get(args.command)(args)