"""Matlab(tm) compatibility functions.

This will hopefully become a complete set of the basic functions available in
matlab.  The syntax is kept as close to the matlab syntax as possible.  One 
fundamental change is that the first index in matlab varies the fastest (as in 
FORTRAN).  That means that it will usually perform reductions over columns, 
whereas with this object the most natural reductions are over rows.  It's perfectly
possible to make this work the way it does in matlab if that's desired.
"""
from Numeric import *

# Elementary Matrices

# zeros is from matrixmodule in C
# ones is from Numeric.py

import RandomArray
def rand(*args):
	"""rand(d1,...,dn) returns a matrix of the given dimensions
	which is initialized to random numbers from a uniform distribution
	in the range [0,1).
	"""
	return RandomArray.random(args)

def eye(N, M=None, k=0, typecode=None):
	"""eye(N, M=N, k=0, typecode=None) returns a N-by-M matrix where the 
	k-th diagonal is all ones, and everything else is zeros.
	"""
	if M == None: M = N
	if type(M) == type('d'): 
		typecode = M
		M = N
	m = equal(subtract.outer(arange(N), arange(M)),-k)
	return asarray(m,typecode=typecode)

def tri(N, M=None, k=0, typecode=None):
	"""tri(N, M=N, k=0, typecode=None) returns a N-by-M matrix where all
	the diagonals starting from lower left corner up to the k-th are all ones.
	"""
	if M == None: M = N
	if type(M) == type('d'): 
		typecode = M
		M = N
	m = greater_equal(subtract.outer(arange(N), arange(M)),-k)
	return m.astype(typecode)
	
# Matrix manipulation

def diag(v, k=0):
	"""diag(v,k=0) returns the k-th diagonal if v is a matrix or
	returns a matrix with v as the k-th diagonal if v is a vector.
	"""
	v = asarray(v)
	s = v.shape
	if len(s)==1:
		n = s[0]+abs(k)
		if k > 0:
			v = concatenate((zeros(k, v.typecode()),v))
		elif k < 0:
			v = concatenate((v,zeros(-k, v.typecode())))
		return eye(n, k=k)*v
	elif len(s)==2:
		v = add.reduce(eye(s[0], s[1], k=k)*v)
		if k > 0: return v[k:]
		elif k < 0: return v[:k]
		else: return v
	else:
	        raise ValueError, "Input must be 1- or 2-D."
	

def fliplr(m):
	"""fliplr(m) returns a 2-D matrix m with the rows preserved and
	columns flipped in the left/right direction.  Only works with 2-D
	arrays.
	"""
	m = asarray(m)
	if len(m.shape) != 2:
		raise ValueError, "Input must be 2-D."
	return m[:, ::-1]

def flipud(m):
	"""flipud(m) returns a 2-D matrix with the columns preserved and
	rows flipped in the up/down direction.  Only works with 2-D arrays.
	"""
	m = asarray(m)
	if len(m.shape) != 2:
		raise ValueError, "Input must be 2-D."
	return m[::-1]
	
# reshape(x, m, n) is not used, instead use reshape(x, (m, n))

def rot90(m, k=1):
	"""rot90(m,k=1) returns the matrix found by rotating m by k*90 degrees
	in the counterclockwise direction.
	"""
	m = asarray(m)
	if len(m.shape) != 2:
		raise ValueError, "Input must be 2-D."
	k = k % 4
	if k == 0: return m
	elif k == 1: return transpose(fliplr(m))
	elif k == 2: return fliplr(flipud(m))
	elif k == 3: return fliplr(transpose(m))

def tril(m, k=0):
	"""tril(m,k=0) returns the elements on and below the k-th diagonal of
	m.  k=0 is the main diagonal, k > 0 is above and k < 0 is below the main
	diagonal.
	"""
	return tri(m.shape[0], m.shape[1], k=k, typecode=m.typecode())*m

def triu(m, k=0):
	"""triu(m,k=0) returns the elements on and above the k-th diagonal of
	m.  k=0 is the main diagonal, k > 0 is above and k < 0 is below the main
	diagonal.
	"""	
	return (1-tri(m.shape[0], m.shape[1], k-1, m.typecode()))*m 

# Data analysis

# Basic operations
def max(m):
	"""max(m) returns the maximum along the first dimension of m.
	"""
	return maximum.reduce(m)

def min(m):
	"""min(m) returns the minimum along the first dimension of m.
	"""
	return minimum.reduce(m)

# Actually from BASIS, but it fits in so naturally here...

def ptp(m):
	"""ptp(m) returns the maximum - minimum along the first dimension of m.
	"""
	return max(m)-min(m)

def mean(m):
	"""mean(m) returns the mean along the first dimension of m.  Note:  if m is
	an integer array, integer division will occur.
	"""
	return add.reduce(m)/len(m)

# sort is done in C but is done row-wise rather than column-wise
def msort(m):
	"""msort(m) returns a sort along the first dimension of m as in MATLAB.
	"""
	return transpose(sort(transpose(m)))

def median(m):
	"""median(m) returns a mean of m along the first dimension of m.
	"""
	return msort(m)[m.shape[0]/2]

def std(m):
	"""std(m) returns the standard deviation along the first
	dimension of m.  The result is unbiased meaning division by len(m)-1.
	"""
	mu = mean(m)
	return sqrt(add.reduce(pow(m-mu,2)))/sqrt(len(m)-1)

def sum(m):
	"""sum(m) returns the sum of the elements along the first
	dimension of m.
	"""
	return add.reduce(m)

def cumsum(m):
	"""cumsum(m) returns the cumulative sum of the elements along the
	first dimension of m.
	"""
	return add.accumulate(m)

def prod(m):
	"""prod(m) returns the product of the elements along the first
	dimension of m.
	"""
	return multiply.reduce(m)

def cumprod(m):
	"""cumprod(m) returns the cumulative product of the elments along the
	first dimension of m.
	"""
	return multiply.accumulate(m)

def trapz(y, x=None):
	"""trapz(y,x=None) integrates y = f(x) using the trapezoidal rule.
	"""
	if x == None: d = 1
	else: d = diff(x)
	return sum(d * (y[1:]+y[0:-1])/2)

def diff(x, n=1):
	"""diff(x,n=1) calculates the first-order, discrete difference
	approximation to the derivative.
	"""
	if n > 1:
	    return diff(x[1:]-x[:-1], n-1)
	else:
	    return x[1:]-x[:-1]
	
def corrcoef(x, y=None):
	"""The correlation coefficients
	"""
	c = cov(x, y)
	d = diag(c)
	return c/sqrt(multiply.outer(d,d))

def cov(m,y=None):
	m = asarray(m)
	mu = mean(m)
	if y != None: m = concatenate((m,y))
	sum_cov = 0.0
	for v in m:
		sum_cov = sum_cov+multiply.outer(v,v)
	return (sum_cov-len(m)*multiply.outer(mu,mu))/(len(m)-1.0)

# Added functions supplied by Travis Oliphant
import LinearAlgebra
def squeeze(a):
    "squeeze(a) removes any ones from the shape of a"
    b = asarray(a.shape)
    reshape (a, tuple (compress (not_equal (b, 1), b)))
    return

def kaiser(M,beta):
    """kaiser(M, beta) returns a Kaiser window of length M with shape parameter
    beta. It depends on the cephes module for the modified bessel function i0.
    """
    import cephes
    n = arange(0,M)
    alpha = (M-1)/2.0
    return cephes.i0(beta * sqrt(1-((n-alpha)/alpha)**2))/cephes.i0(beta)

def blackman(M):
    """blackman(M) returns the M-point Blackman window.
    """
    n = arange(0,M)
    return 0.42-0.5*cos(2*pi*n/M) + 0.08*cos(4*pi*n/M)


def bartlett(M):
    """bartlett(M) returns the M-point Bartlett window.
    """
    n = arange(0,M)
    return where(less_equal(n,M/2.0),2.0*n/M,2-2.0*n/M)

def hanning(M):
    """hanning(M) returns the M-point Hanning window.
    """
    n = arange(0,M)
    return 0.5-0.5*cos(2*pi*n/M)

def hamming(M):
    """hamming(M) returns the M-point Hamming window.
    """
    n = arange(0,M)
    return 0.54-0.46*cos(2*pi*n/M)

def sinc(x):
    """sinc(x) returns sin(pi*x)/(pi*x) at all points of array x.
    """
    return where(equal(x,0.0),1,sin(pi*x)/(pi*x))


def eig(v):
    """[x,v] = eig(m) returns the the eigenvalues of m in x and the corresponding
    eigenvectors in the rows of v.
    """
    return LinearAlgebra.eigenvectors(v)

def svd(v):
    """[u,x,v] = svd(m) return the singular value decomposition of m.
    """
    return LinearAlgebra.singular_value_decomposition(v)