import Mrc
import datadoc
import numpy as np
import pylab
import scipy.stats.stats
import sys

points = []
imageShape = None
for filename in sys.argv[1:]:
    data = datadoc.DataDoc(Mrc.bindFile(filename)).imageArray[0,0,0]
    print filename, "%d" % data.mean()
    imageShape = data.shape
    points.append((data.mean(), data))
points.sort(lambda a, b: cmp(a[0], b[0]))

xVals = np.array([item[0] for item in points])
yVals = np.array([item[1] for item in points])
yVals.shape = len(yVals), np.product(imageShape)

result, residuals, rank, singular_values, rcond = np.polyfit(xVals, yVals, 1, full = True)
slopes = result[0]
intercepts = result[1]

slopes.shape = imageShape
intercepts.shape = imageShape
yVals.shape = (len(yVals),) + imageShape

result = np.zeros(imageShape, dtype = np.float32)
for i in xrange(slopes.shape[0]):
    print i
    for j in xrange(slopes.shape[1]):
        xen = xVals * slopes[i,j] + intercepts[i,j]
        yen = yVals[:,i,j]
        result[i,j] = scipy.stats.stats.pearsonr(xen, yen)[0]

result.shape = np.product(result.shape)
yVals.shape = len(xVals), result.shape

print result.shape, result.dtype
indices = np.argsort(result) # This fails

colors = ['r', 'g', 'b', 'c', 'm', 'y', 'k']
for i in indices[:len(colors)]:
    pylab.plot(xVals, yVals[:, i], colors[i])

pylab.grid()
pylab.xlabel('Mean image intensity')
pylab.ylabel('Specific pixel intensity')
pylab.title("Response curves for %d most nonlinear pixels" % len(colors))
pylab.show()