#!/usr/bin/env python
import argparse
import pysam
import sys

try:
	import matplotlib
	matplotlib.use('agg')
	import matplotlib.pyplot as plt
	import seaborn as sns
	sns.set_palette("husl", desat=.6)
except ImportError:
	print("Could not import either matplotlib or seaborn")
        sys.exit(1)
#import matplotlib.pyplot as plt

def is_proper_aligned_unique_innie(read):
    mapped_ok = not (read.is_unmapped or read.mate_is_unmapped)
    orientation_ok = (read.is_reverse and not read.mate_is_reverse and read.tlen < 0) or \
            (not read.is_reverse and read.mate_is_reverse and read.tlen > 0)
    quality_ok = read.mapq > 10 and not read.is_secondary
    same_ref = read.rname == read.mrnm

    return mapped_ok and orientation_ok and quality_ok and same_ref

def is_proper_aligned_unique_outie(read):
    mapped_ok = not (read.is_unmapped or read.mate_is_unmapped)
    orientation_ok = (read.is_reverse and not read.mate_is_reverse and read.tlen > 0) or \
            (not read.is_reverse and read.mate_is_reverse and read.tlen < 0)
    quality_ok = read.mapq > 10 and not read.is_secondary
    same_ref = read.rname == read.mrnm

    return mapped_ok and orientation_ok and quality_ok and same_ref

def ff_or_rr(read):
    mapped_ok = not (read.is_unmapped or read.mate_is_unmapped)
    orientation_ok = (read.is_reverse and read.mate_is_reverse) or \
            (read.is_reverse and read.mate_is_reverse)
    quality_ok = read.mapq > 10 and not read.is_secondary
    same_ref = read.rname == read.mrnm

    return mapped_ok and orientation_ok and quality_ok and same_ref


def plot_insert_sizes(bamfile,outfile_path, args):
    bamfile = pysam.Samfile(bamfile, 'rb')
    i_sizes_fr = []
    i_sizes_rf = []
    i_sizes_ff = []
    for read in bamfile:
        if abs(read.tlen) < 0 or abs(read.tlen) > args.x_upper:
            continue
	if is_proper_aligned_unique_outie(read):
            i_sizes_rf.append(abs(read.tlen))
            #print("cool", read.tlen, read.rlen, read.qlen, read.alen)
        if is_proper_aligned_unique_innie(read):
            i_sizes_fr.append(abs(read.tlen))
	if ff_or_rr(read):
            i_sizes_ff.append(abs(read.tlen))
	#if is_proper_aligned_unique_innie(read) and abs(read.tlen)> 1000:
        #    print("strange read", read.tlen, read.rlen, read.qlen, read.alen)
    print("Number of pairs ---> <----", len(i_sizes_fr))
    print("Number of pairs <--- ---->", len(i_sizes_rf))
    print("Number of pairs ----> ----> or <---- <-----", len(i_sizes_ff))
    print("contamine_rate:", len(i_sizes_fr)/ float(len(i_sizes_fr) + len(i_sizes_rf)) )
    n = float(len(i_sizes_rf))
    try:
        mu_rf = sum(i_sizes_rf)/n
        stddev_rf = (sum(list(map((lambda x: x ** 2 - 2 * x * mu_rf + mu_rf ** 2), i_sizes_rf))) / (n - 1)) ** 0.5
    except:
        mu_rf, stddev_rf = 0,0
    print("mean rf:", mu_rf)
    print("stddev rf:", stddev_rf)

    n = float(len(i_sizes_fr))
    try:
        mu_fr = sum(i_sizes_fr)/n
        stddev_fr = (sum(list(map((lambda x: x ** 2 - 2 * x * mu_fr + mu_fr ** 2), i_sizes_fr))) / (n - 1)) ** 0.5
    except:
        mu_fr, stddev_fr = 0,0
    print("mean fr:", mu_fr)
    print("stddev fr:", stddev_fr)

    n = float(len(i_sizes_ff))
    try:
        mu_ff = sum(i_sizes_ff)/n
        stddev_ff = (sum(list(map((lambda x: x ** 2 - 2 * x * mu_ff + mu_ff ** 2), i_sizes_ff))) / (n - 1)) ** 0.5
    except:
        mu_ff, stddev_ff = 0,0

    print("mean ff or rr:", mu_ff)
    print("stddev ff or rr", stddev_ff)

    labels = []
    labels.append("fr --> <--")
    labels.append("rf <-- -->")
    labels.append("ff/rr --> -->")
    if args.fr:
        plt.hist(i_sizes_rf,bins=100,color='b', alpha=0.4, label=labels[1])
        plt.hist(i_sizes_fr,bins=100,color='r', alpha=0.4,label=labels[0])
        plt.hist(i_sizes_ff,bins=100,color='g', alpha=0.4, label=labels[2])
    else: 
        plt.hist(i_sizes_fr,bins=100,color='r', alpha=0.4, label=labels[0])
        plt.hist(i_sizes_rf,bins=100,color='b', alpha=0.4, label=labels[1])
        plt.hist(i_sizes_ff,bins=100,color='g', alpha=0.4, label=labels[2])
    plt.ylabel('Frequency') 
    plt.xlabel('Fragment size')  
    plt.legend( )
    plt.savefig(outfile_path)

if __name__ == '__main__':

	parser = argparse.ArgumentParser()
	parser.add_argument('bam', type=str, help='bam file. ')
	parser.add_argument('outfile', type=str, help='outfolder path NEEDS TO BE FULL PATH with no slash at end!. e.g /Users/ksahlin/_tmp/MODIL/modil_out ')
        parser.add_argument('--fr', dest="fr", action="store_true", help='Plot order, default plots fr then rf. If --fr is set, it first plots rf, then fr as the upper layer.')
        parser.add_argument('--x_upper_limit', dest="x_upper", type=int, default = 10000, help='Upper limit on isize plot ')
	args = parser.parse_args()

	plot_insert_sizes( args.bam, args.outfile, args )

