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

parser = argparse.ArgumentParser(description="Read in the SASA data files and generate a histogram.")
parser.add_argument("--num-points", dest='numPoints', type=int, help="The number of data points in the sasa file.")
parser.add_argument("--output-name", dest='outputName', type=str, help="The basename of the generated figures. A pdf and png will be generated based on this name.")
parser.add_argument("-p", help="Don't plot the protein-all DNA sasa values.", dest='noBig', action='store_true')
parser.add_argument("-s", help="Don't plot the protein-binding-site sasa values.", dest='noSmall', action='store_true')
args = parser.parse_args()

numPoints = args.numPoints
outputName = args.outputName

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("proteinA.dat", "righthalf.dat", "righthalfprotein.dat", numPoints)

ymin = min([min(x) for x in [sasaP1, sasaP2, sasaPP]])
ymax = max([max(x) for x in [sasaP1, sasaP2, sasaPP]])

fig = plt.figure(figsize=(3,2), dpi=600)
axHist = fig.add_axes([0.3, 0.3, 0.6, 0.6])
axHist.set_xlim((ymin, ymax))
#axHist.yaxis.tick_right()
axHist.set_ylim((0, numPoints/50))
axHist.tick_params(axis='x', which='both', bottom=True, labelbottom=True)
def addPlot(dats, label, color, axHist):
    #axPlot.plot(times, dats, color+',', alpha=0.2)
    #axPlot.plot(times, scipy.ndimage.gaussian_filter(dats, 2000), color+'-', label=label)

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



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