#!/usr/bin/env python3
import numpy as np
import matplotlib.pyplot as plt
import sys
import scipy.ndimage
import argparse

parser = argparse.ArgumentParser(description="Reads in the sasa files and generates a plot.")
parser.add_argument("-n", help="number of observations.", type=int, dest='numPoints')
parser.add_argument("-o", help="name of output file. A pdf and png will be generated based on this name.", dest='outFname')
parser.add_argument("-p", help="Don't plot the protein-nucleic tracks (the one with high values)", dest='noBig', action='store_true')
parser.add_argument("-s", help="Don't plot the left and right side binding site data.", dest='noSmall', action='store_true')
args = parser.parse_args()
numPoints = args.numPoints


def readSasaVals(file1, file2, fileComb, numPoints):
    fp1 = open(file1, "r")
    fp2 = open(file2, "r")
    fpC = open(fileComb, "r")
    ret = np.zeros((numPoints,))
    for i in range(numPoints):
        l1 = float(fp1.readline())
        l2 = float(fp2.readline())
        lC = float(fpC.readline())
        combSasa = l1 + l2
        Δsasa = (combSasa - lC) / 2
        ret[i] = Δsasa
    return ret

sasaP1 = readSasaVals("proteinA.dat", "nucleic.dat", "proteinAnucleic.dat", numPoints)
sasaP2 = readSasaVals("proteinB.dat", "nucleic.dat", "proteinBnucleic.dat", numPoints)
sasaPP = readSasaVals("proteinA.dat", "proteinB.dat", "proteinprotein.dat", numPoints)
sasaLP = readSasaVals("proteinA.dat", "lefthalf.dat", "lefthalfprotein.dat", numPoints)
sasaRP = readSasaVals("proteinB.dat", "righthalf.dat", "righthalfprotein.dat", numPoints)
times = np.arange(0, numPoints) / 100.0
if args.noBig:
    candidates = [sasaPP, sasaLP, sasaRP]
else:
    candidates = [sasaP1, sasaP2, sasaPP, sasaLP, sasaRP]
ymin = min([min(x) for x in candidates])
ymax = max([max(x) for x in candidates])

fig = plt.figure(figsize=(8,5), dpi=300)
axPlot = fig.add_axes([0.1, 0.1, 0.7, 0.8])
axHist = fig.add_axes([0.8, 0.1, 0.1, 0.8])
axPlot.set_ylim((ymin, ymax))
axHist.set_ylim((ymin, ymax))
axHist.yaxis.tick_right()
axPlot.set_xlim((0, max(times)))
axHist.set_xlim((0, numPoints/50))
axHist.tick_params(axis='x', which='both', bottom=False, labelbottom=False)

def addPlot(times, dats, label, color, axPlot, axHist):
    axPlot.plot(times, dats,marker=',', linestyle='None', color=color, alpha=0.2)
    axPlot.plot(times, scipy.ndimage.gaussian_filter(dats, 2000), color=color, linestyle='-', label=label)

    histOfVals = np.histogram(dats[dats>1], bins=200)#, range=(ymin, ymax))
    axHist.plot(histOfVals[0], histOfVals[1][:-1], color=color, linestyle='solid')
    medValue = np.median(dats)
    axHist.plot([0, np.max(histOfVals[0]) / 3], [medValue, medValue], color=color , linestyle='solid')



if not args.noBig:
    addPlot(times, sasaP1, "protein A-DNA", '#069AA8', axPlot, axHist)
    addPlot(times, sasaP2, "protein B-DNA", '#014D6D', axPlot, axHist)
addPlot(times, sasaPP, "protein-protein", '#CC6677', axPlot, axHist)
if not args.noSmall:
    addPlot(times, sasaLP, "protein-left site", '#46DAC8', axPlot, axHist)
    addPlot(times, sasaRP, "protein-right site", '#266DBD', axPlot, axHist)
axPlot.legend()
axPlot.set_ylabel("Buried surface (Ų)")
axPlot.set_xlabel("Time (ns)")
plt.savefig(args.outFname + ".png")
plt.savefig(args.outFname + ".pdf")
#plt.show()