def getMatrix(AllData):
  # create matrices for all the data
  numQs = len(AllData.keys())
  subjects = 12    #len(AllData[1]['RT'])
  correct = np.array(py.zeros([numQs, subjects]))
  confA = np.array(py.zeros([numQs, subjects]))
  confB = np.array(py.zeros([numQs, subjects]))
  RTs = np.array(py.zeros([numQs, subjects]))
  #print(AllData)
  for i in xrange(subjects):
    # rows
    for j in xrange(1,17):
      # columns
      correct[j-1,i] = AllData[j]['correct'][i]
      
  for i in xrange(subjects):
    for j in xrange(1,17):
      confA[j-1,i] = AllData[j]['confA'][i]
  for i in xrange(subjects):
    for j in xrange(1,17):
      confB[j-1,i] = AllData[j]['confB'][i]
  for i in xrange(subjects):
    for j in xrange(1,17):
      RTs[j-1,i] = AllData[j]['RT'][i]
  
  print(py.shape(correct), py.shape(confA), py.shape(confB), py.shape(RTs))
  return correct, confA, confB, RTs

def filter2d(x, y, axes=['y'], algos=['2sigma']):
    """
    Perform 2D data filtration by selected exes.
    In:
        x : ndarray, X vector
        y : ndarray, Y vector
        axes : list, axes names which are used to choose filtered values. x, y or any combination
    Out:
        xnew : ndarray, filtered X
        ynew : ndarray, filtered Y
    """
    xnew = pl.array(x, dtype='float')
    ynew = pl.array(y, dtype='float')
    mask_x = pl.ones(len(x), dtype='bool')
    mask_y = pl.ones(len(y), dtype='bool')
    if 'y' in axes:
        mask_y = filter1d(y,algos=algos)        
    if 'x' in axes:
        mask_x = filter1d(x,algos=algos)
    mask = mask_x * mask_y
    xnew *= mask
    ynew *= mask
    
    xnew = pl.ma.masked_equal(xnew,0)
    xnew = pl.ma.compressed(xnew)
    ynew = pl.ma.masked_equal(ynew,0)
    ynew = pl.ma.compressed(ynew)

    assert pl.shape(xnew) == pl.shape(ynew)
    return xnew, ynew

def loadMNISTImages(filename):
  f = open(filename, 'rb')

  # Verify Magic Number
  s = f.read(4)
  magic = int(s.encode('hex'),16)
  assert(magic == 2051)

  # Get Number of Images
  s = f.read(4)
  numImages = int(s.encode('hex'),16)
  s = f.read(4)
  numRows = int(s.encode('hex'),16)
  s = f.read(4)
  numCols = int(s.encode('hex'),16)

  # Get Data
  s = f.read()
  a = frombuffer(s, uint8)

  # Use 'F' to ensure that we read by column
  a = reshape(a, (numCols , numRows, numImages), order='F');
  images = transpose(a, (1, 0, 2))
  f.close()

  # Reshape to #pixels * #examples
  images  = reshape(a, (shape(images)[0] * shape(images)[1], numImages),
          order='F');
  images = double(images)/255
  return images

        def residuals(params,x,y,z):
            xo = params[0]
            xs = params[1]
            yo = params[2]
            ys = params[3]
            zo = params[4]
            zs = params[5]
            xys = params[6]
            xzs = params[7]
            yzs = params[8]

            xc = empty(shape(x))
            yc = empty(shape(y))
            zc = empty(shape(z))
            for i in range(len(x)):
                _x = x[i] - xo
                _y = y[i] - yo
                _z = z[i] - zo

                xc[i] = _x * (xs + _y * xys + _z * xzs)
                yc[i] = _y * (ys + _z * yzs)
                zc[i] = _z * (zs)

            res = []
            for i in range(len(xc)):
                norm = l2norm(array([xc[i],yc[i],zc[i]])) - 1.0
                res.append(norm)

            return array(res)

def g_l2_wd(x0, S, I, gamma):
    M = shape(S)[0]
    L, batch_size = shape(I)
    A = matrix(reshape(x0,[L,M]))
    E = I - A*S
    g = -E*S.T/batch_size + gamma*A
    return g.A1

def f_l2_wd(x0, S, I, gamma):
    M = shape(S)[0]
    L, batch_size = shape(I)
    A = matrix(reshape(x0,[L,M]))
    E = I - A*S
    f = 0.5*(E.A**2).sum()/batch_size + 0.5*gamma*(A.A**2).sum()
    return f

        def residuals(params,x,y,z):
            xo = params[0]
            xs = params[1]
            yo = params[2]
            ys = params[3]
            zo = params[4]
            zs = params[5]

            xc = empty(shape(x))
            for i in range(len(x)):
                xc[i] = (x[i] - xo) * xs

            yc = empty(shape(y))
            for i in range(len(y)):
                yc[i] = (y[i] - yo) * ys

            zc = empty(shape(z))
            for i in range(len(z)):
                zc[i] = (z[i] - zo) * zs

            res = []
            for i in range(len(xc)):
                norm = l2norm(array([xc[i],yc[i],zc[i]])) - 1.0
                res.append(norm)

            return array(res)

def plotEnsemble2D(ens,v1,v2,colordata=None,hess=None,\
		   size=50,labelBest=True,ensembleAlpha=0.75,contourAlpha=1.0):
	"""
	Plots a 2-dimensional projection of a given parameter
	ensemble, along given directions:
	     -- If v1 and v2 are scalars, project onto plane given by
		those two bare parameter directions.
	     -- If v1 and v2 are vectors, project onto those two vectors.
	
	When given colordata (either a single color, or an array
	of different colors the length of ensemble size), each point
	will be assigned a color based on the colordata.
	
	With labelBest set, the first point in the ensemble is
	plotted larger (to show the 'best fit' point for a usual 
	parameter ensemble).
	
	If a Hessian is given, cost contours will be plotted
	using plotContours2D.
	"""
	if pylab.shape(v1) is ():
		xdata = pylab.transpose(ens)[v1]
		ydata = pylab.transpose(ens)[v2]
		
		# label axes
		param1name, param2name = '',''
		try:
		    paramLabels = ens[0].keys()
		except:
		    paramLabels = None
		if paramLabels is not None:
		    param1name = ' ('+paramLabels[param1]+')'
		    param2name = ' ('+paramLabels[param2]+')'
		pylab.xlabel('Parameter '+str(v1)+param1name)
		pylab.ylabel('Parameter '+str(v2)+param2name)
	else:
		xdata = pylab.dot(ens,v1)
		ydata = pylab.dot(ens,v2)

	if colordata==None:
		colordata = pylab.ones(len(xdata))
		
	if labelBest: # plot first as larger circle
		if pylab.shape(colordata) is (): # single color
		    colordata0 = colordata
		    colordataRest = colordata
		else: # specified colors
		    colordata0 = [colordata[0]]
		    colordataRest = colordata[1:]
		scatterColors(xdata[1:],ydata[1:],colordataRest,		\
				size,alpha=ensembleAlpha)
		scatterColors([xdata[0]],[ydata[0]],colordata0,			\
				size*4,alpha=ensembleAlpha)
	else:
		scatterColors(xdata,ydata,colordata,size,alpha=ensembleAlpha)
		
	if hess is not None:
		plotApproxContours2D(hess,param1,param2,pylab.array(ens[0]),	\
			alpha=contourAlpha)

 def __mul__(self, V):
     if not self.TR:
         batchsize, v = shape(V)
         return self[0]*V + self[1] 
     else:
         batchsize, h = shape(V)
         assert(h==self.h)            
         return self[0].transpose()*V

def objfun_l2_wd(x0, S, I, gamma):
    M = shape(S)[0]
    L, batch_size = shape(I)
    A = matrix(reshape(x0,[L,M]))
    E = I - A*S 
    f = 0.5*(E.A**2).sum()/batch_size + 0.5*gamma*(A.A**2).sum()
    g = -E*S.T/batch_size + gamma*A   
    return (f,g.A1)

def correctBias(AllData):
  # correct for difficulty and plot each subject %correct vs confidence
  corrmatrix, confmatrix = returnConfMatrix(AllData)
  Qs, subjects = py.shape(corrmatrix)
  copts = [1,2,3,4,5]
  datamat = np.array(py.zeros([len(copts), subjects]))
  print(datamat)
  fig = py.figure()
  ax15 = fig.add_subplot(111) 
  i = 0
 
  while i < subjects:
    c1, c2, c3, c4, c5 = [],[],[],[],[]
    # get confidences for each subject
    j = 0
    while j < Qs:
      # get confidences and correct for each question
      if confmatrix[j][i] == 1:
        c1.append(corrmatrix[j][i])
      elif confmatrix[j][i] == 2:
        c2.append(corrmatrix[j][i])
      elif confmatrix[j][i] == 3:
        c3.append(corrmatrix[j][i])
      elif confmatrix[j][i] == 4:
        c4.append(corrmatrix[j][i])
      elif confmatrix[j][i] == 5:
        c5.append(corrmatrix[j][i])
      else:
        print('bad num encountered')
        
      j += 1
    print('i is %d' %i)
    minconf = ([py.mean(c1), py.mean(c2), py.mean(c3), 
                   py.mean(c4), py.mean(c5)])
    pmin = 10
    for p in minconf:
      if p < pmin and p != 0 and math.isnan(p) is not True:
        pmin = p
    
    print(pmin)
    datamat[0][i] = py.mean(c1)/pmin
    datamat[1][i] = py.mean(c2)/pmin
    datamat[2][i] = py.mean(c3)/pmin
    datamat[3][i] = py.mean(c4)/pmin
    datamat[4][i] = py.mean(c5)/pmin
    # print(datamat)
    print( py.shape(datamat))
    print(len(datamat[:,i]))
    ax15.plot(range(1,6), datamat[:,i], alpha=0.4, linewidth=4)
    i += 1
  
  ax15.set_ylabel('Modified Correct')
  ax15.set_xlabel('Confidence')
  ax15.set_title('All responses')
  ax15.set_xticks(np.arange(1,6))
  ax15.set_xticklabels( [1, 2, 3, 4, 5] )
  ax15.set_xlim(0,6)

    def _generate_feature_development_plots(self, important_features):
        """ This function generates the actual histogram plot"""
        # Everything is done class-wise
        for label in important_features.keys():
            # Axis limits are determined by the global maxima
            (minVal, maxVal) = (important_features[label].min(0).min(0),
                                important_features[label].max(0).max(0))                                
            nr_chans = pylab.shape(important_features[label])[0]
                                
            myFig = pylab.figure()
            myFig.set_size_inches((40,nr_chans))
            
            for i_chan in range(nr_chans):
                ax = myFig.add_subplot(nr_chans, 1, i_chan+1)
                
                # cycle line colors
                if (pylab.mod(i_chan,2) == 0): myCol = '#000080'
                else: myCol = '#003EFF'
                # plot features and black zero-line
                pylab.plot(important_features[label][i_chan,:],color=myCol)
                pylab.plot(range(len(important_features[label][i_chan,:])),
                    pylab.zeros(len(important_features[label][i_chan,:])),
                    'k--')
                pylab.ylim((minVal,maxVal))
                xmax = pylab.shape(important_features[label])[1]
                pylab.xlim((0,xmax))
                
                # draw vertical dashed line every 20 epochs
                for vertical_line_position in range(0,xmax+1,20):
                    pylab.axvline(x=vertical_line_position,
                                  ymin=0, ymax=1, color='k', linestyle='--')
                    
                # write title above uppermost subplot
                if i_chan+1 == 1:
                    pylab.title('Feature development: Amplitudes of %s Epochs'
                                % label, fontsize=40)
                # adjust the axes, i.e. remove upper and right,
                # shift the left to avoid overlaps,
                # and lower axis only @ bottom subplot
                if i_chan+1 < nr_chans:
                    self._adjust_spines(ax,['left'],i_chan)
                else:
                    self._adjust_spines(ax,['left', 'bottom'],i_chan)
                    pylab.xlabel('Number of Epoch', fontsize=36)
                # Write feature name next to the axis
                pylab.ylabel(self.corr_important_feat_names[i_chan],
                             fontsize=20, rotation='horizontal')
            # remove whitespace between subplots etc.
            myFig.subplots_adjust(bottom=0.03,left=0.08,right=0.97,
                                  top=0.94,wspace=0,hspace=0)

            self.feature_development_plot[label] = myFig

def spatio_temporal(ancl):

    os.mkdir('SpatioTemporalVels')

    print('RIGHT NOW THIS IS ONLY FOR VX!!!!!!!')

    p_arr = pl.arange(0,ancl.N)
    

    # How many cycles do we want to look at?
    how_many = 10

    var_arr = pl.array([])
    for i,j in enumerate(os.listdir('.')):
        if 'poindat.txt' not in j:
            continue
        print('working on file ' + j)
        poin_num = int(j[:j.find('p')])
        cur_file = open(j,'r')
        cur_sweep_var = float(cur_file.readline().split()[-1])
        cur_data=pl.genfromtxt(cur_file)
        cur_file.close()

        var_arr = pl.append(var_arr,cur_sweep_var)
        
        count = 0
        grid = cur_data[-int(how_many*2.0*pl.pi/ancl.dt):,:ancl.N]

        # in 1D because particles never cross eachother we can order them in the images to mathch
        # their physical order.
        grid_ordered = pl.zeros(pl.shape(grid))
        # can just use the initial conditions to figure out where each is
        init_x = cur_data[0,ancl.N:2*ancl.N]
        sorted_x = sorted(init_x)
        for a,alpha in enumerate(sorted_x):
            for b,beta in enumerate(init_x):
                if alpha == beta:
                    grid_ordered[:,a]=grid[:,b]
        
    
        print('shape of grid_ordered: ' + str(pl.shape(grid_ordered)))
        
        fig = pl.figure()
        ax = fig.add_subplot(111)
        # form of errorbar(x,y,xerr=xerr_arr,yerr=yerr_arr)
        ax.imshow(grid_ordered,interpolation="nearest", aspect='auto')
        ax.set_xlabel('Particle',fontsize=30)
        #ax.set_aspect('equal')
        ax.set_ylabel(r'$ t $',fontsize=30)
        fig.tight_layout()
        fig.savefig('SpatioTemporalVels/%(number)04d.png'%{'number':poin_num})
        pl.close(fig)

def approximate(x,y):
    """
    Linear approximation of y=f(x) using least square estimator.
    In:
        x : ndarray
        y : ndarray
    Out:
        a, b : float, as in a*x+b=y
    """
    assert pl.shape(x) == pl.shape(y)
    A = pl.vstack([x, pl.ones(len(x))]).T
    a, b = pl.lstsq(A, y)[0]
    return a, b

def choose_patches(IMAGES, L, batch_size=1000):
    sz = int(sqrt(L))
    imsz = shape(IMAGES)[0]
    num_images = shape(IMAGES)[2]
    BUFF = 4

    X = matrix(zeros([L,batch_size],'d'))
    for i in range(batch_size):
        j = int(floor(num_images * rand()))
        r = sz/2+BUFF+int(floor((imsz-sz-2*BUFF)*rand()))
        c = sz/2+BUFF+int(floor((imsz-sz-2*BUFF)*rand()))
        X[:,i] = reshape(IMAGES[r-sz/2:r+sz/2,c-sz/2:c+sz/2,j],[L,1])
    return X

def decoder(data,code,iflip):
    """convolves each data profile with the code sequence"""
    times=py.shape(data)[0]
    hts=py.shape(data)[1]
    codelength=py.shape(code)[0]
    code_rev=code[::-1] #decoding requires using the inverse of the code
    deflip=1
    #pdb.set_trace()
    for i in range (times):
        temp=py.convolve(data[i,:],code_rev)
        data[i,:]=deflip*temp[codelength-1:codelength+hts]
        deflip=deflip*iflip #call with iflip=-1 if tx has flip
        #pdb.set_trace()
    return data

def calcaV(W,method = "ratio"):
    """Calculate aV"""
    if method == "ratio":
        return pl.log(pl.absolute(W/pl.roll(W,-1,axis=1)))
    else:
        aVs = pl.zeros(pl.shape(W))
        n = pl.arange(1,pl.size(W,axis=1)+1)
        f = lambda b,t,W: W - b[0] * pl.exp(-b[1] * t)
        
        for i in xrange(pl.size(W,axis=0)):
            params,result = optimize.leastsq(f,[1.,1.],args=(n,W[i]))
            aVs[i] = params[1] * pl.ones(pl.shape(W[i]))
            
        return aVs

def pnccd_to_image(infile, outfile):
    try:
        f = h5py.File(infile)
    except:
        raise IOError("Can't read %s. It may not be a pnCCD file." % filename)

    i1 = f.keys().index("data")
    i2 = f.values()[i1].keys().index("data1")

    data = f.values()[i1].values()[i2].value

    img = spimage.sp_image_alloc(pylab.shape(data)[0], pylab.shape(data)[1], 1)
    img.image[:, :] = data[:, :]
    spimage.sp_image_write(img, outfile, 0)
    spimage.sp_image_free(img)

    def __init__(self, r_value, filename_coupling, amplitude, omega, cycles):
	"""
	Ode_function_call()

	Constructor.

	Parameters
	----------
	filename_coupling : string, path to the HDF5 file that contains the 
	    dipole couplings of the electronic H2+ problem.
	"""
	#Laser info.
	self.amplitude = amplitude
	self.omega = omega
	self.cycles = cycles
	self.pulse_duration = 2* pi /(self.omega) * self.cycles

	#Open files
	f = tables.openFile(filename_coupling)

	#Retrieve r value.
	r_grid = f.root.R_grid[:]
	self.r_index = argmin(abs(r_grid - r_value))	
	self.index_array = f.root.index_array[:]
	self.r = r_grid[self.r_index]

	#Retrieve Hamiltonian.
	self.H_0 = diag(f.root.E[:,self.r_index])
	self.H_1 = f.root.couplings[:,:,self.r_index]	
	
	#Close file.
	f.close()

	#Basis sizes.
	self.basis_size = shape(self.H_0)[0]

def gadget_merge_ics( outfile, filename1, filename2, offset1, offset2, voffset1=[0.,0.,0.], voffset2=[0.,0.,0.] ):
	snap1 = gadget_readsnapname( filename1 )
	snap2 = gadget_readsnapname( filename2 )

	for i in range(3):
		snap1.pos[:,i] += offset1[i]
		snap2.pos[:,i] += offset2[i]
		
	for i in range(3):
		snap1.vel[:,i] += voffset1[i]
		snap2.vel[:,i] += voffset2[i]

	npart = snap1.npart + snap2.npart
	data = {}
	data[ 'count' ] = npart
	data[ 'pos' ] = pylab.zeros( [npart, 3] )
	data[ 'pos' ][ 0:snap1.npart, : ] = snap1.pos
	data[ 'pos' ][ snap1.npart:npart, : ] = snap2.pos
	data[ 'vel' ] = pylab.zeros( [npart, 3] )
	data[ 'vel' ][ 0:snap1.npart, : ] = snap1.vel
	data[ 'vel' ][ snap1.npart:npart, : ] = snap2.vel
	data[ 'mass' ] = pylab.zeros( npart )
	data[ 'mass' ][ 0:snap1.npart ] = snap1.data["mass"]
	data[ 'mass' ][ snap1.npart:npart ] = snap2.data["mass"]
	data[ 'u' ] = pylab.zeros( npart )
	data[ 'u' ][ 0:snap1.npart ] = snap1.data["u"]
	data[ 'u' ][ snap1.npart:npart ] = snap2.data["u"]
	nxnuc = pylab.shape( snap1.data["xnuc"] )[1]
	data[ 'xnuc' ] = pylab.zeros( [npart, nxnuc] )
	data[ 'xnuc' ][ 0:snap1.npart, : ] = snap1.data["xnuc"]
	data[ 'xnuc' ][ snap1.npart:npart, : ] = snap2.data["xnuc"]

	gadget_write_ics( outfile, data, transpose=False )
	return