import numpy as np
import scipy.stats as sps

import kalman
from nipy.neurospin.utils import mahalanobis


DEF_TINY = 1e-50
DEF_DOFMAX = 1e10

models = {'spherical':['ols', 'kalman'], 
	  'ar1':['kalman']}

class glm:
	def __init__(self, Y=None, X=None, formula=None, axis=0, 
		     model='spherical', method=None, niter=2):

		# Check dimensions
		if Y == None:
			return
		else:
			self.fit(Y, X, formula, axis, model, method, niter)

	def fit(self, Y, X, formula=None, axis=0, model='spherical', method=None, niter=2):
		
		if Y.shape[axis] != X.shape[0]:
			raise ValueError, 'Response and predictors are inconsistent'
		
		# Find model type
		self._axis = axis
		if isinstance(formula, str):
			model = 'mfx'
		if models.has_key(model):
			self.model = model
			if method == None: 
				self.method = models[model][0]
			elif models[model].count(method):
				self.method = method
			else: 
				raise ValueError, 'Unknown method'
		else:
			raise ValueError, 'Unknown model'
			
		# Initialize fields
		constants = []
		a = 0
		
		# Switch on models / methods 
		if self.model == 'spherical': 
			constants = ['nvbeta', 'a']
			if self.method == 'ols':
				out = ols(Y, X, axis=axis)
			elif self.method == 'kalman':
				out = kalman.ols(Y, X, axis=axis)
                elif self.model == 'ar1':
			constants = ['a']
			out = kalman.ar1(Y, X, axis=axis, niter=niter)
			a = out[4]
			out = out[0:4]
			

		# Finalize
		self.beta, self.nvbeta, self.s2, self.dof = out 
		self.s2 = self.s2.squeeze()
		self.a = a
		self._constants = constants

	"""
	Save fit into a .npz file 
	"""
	def save(self, file): 
		np.savez(file, 
			 beta=self.beta, 
			 nvbeta=self.nvbeta, 
			 s2=self.s2, 
			 dof=self.dof, 
			 a=self.a, 
			 model=self.model,
			 method=self.method,
			 axis=self._axis, 
			 constants=self._constants)


        """
        c must be a numpy.ndarray (or anything that numpy.asarray
        can cast to a ndarray).
        For a F contrast, c must be q x p where q is the number of contrast vectors and
        p is the total number of regressors. 
        """
	def contrast(self, c, type='t', tiny=DEF_TINY, dofmax=DEF_DOFMAX):
		c = np.asarray(c)
		#dim = len(c.shape)
		if c.ndim == 1:
			dim=1
		else:
			dim = c.shape[0]
		axis = self._axis
		ndims = len(self.beta.shape)
                
		# Compute the contrast estimate: c*B
		B = np.rollaxis(self.beta, axis, ndims)
		con = np.inner(c, B) # q, X
                
		# Compute the variance of the contrast estimate: s2 * (c' * nvbeta * c)
		# Two cases are considered: either the input effect variance
		# is position-dependent (output by RKF_fit), or it is a global
		# one (output by KF_fit)
		s2 = self.s2.squeeze()
		nvbeta = self.nvbeta
		if not 'nvbeta' in self._constants: 
			nvbeta = np.rollaxis(nvbeta, axis, ndims+1)
			nvbeta = np.rollaxis(nvbeta, axis, ndims+1) # X, p, p
		if dim == 1:
			vcon = np.inner(c, np.inner(c, nvbeta))
			vcon = vcon.squeeze()*s2
		else: 
			vcon = np.dot(c, np.inner(nvbeta, c)) # q, X, q or q, q
			if not 'nvbeta' in self._constants: 
				vcon = np.rollaxis(vcon, ndims, 1)*s2 # q, q, X
			else:
				aux = vcon.shape # q, q
				vcon = np.resize(vcon, s2.shape+aux) # X, q, q
				vcon = vcon.transpose().reshape(aux+(s2.size,))*s2.reshape((s2.size,)) # q, q, Xflat
				vcon = vcon.reshape(aux+s2.shape) # q, q, X
		# Create contrast instance
		c = contrast(dim, type, tiny, dofmax)
		c.effect = con
		c.variance = vcon
		c.dof = self.dof
		return c
                


class contrast:

	def __init__(self, dim, type='t', tiny=DEF_TINY, dofmax=DEF_DOFMAX):
		"""
		tiny is a numerical constant for computations. 
		"""
		self.dim = dim
		self.effect = None
		self.variance = None
		self.dof = None
		if dim > 1: 
			if type is 't':
				type = 'F'
		self.type = type
		self._stat = None
		self._pvalue = None
		self._baseline = 0 
		self._tiny = tiny
		self._dofmax = dofmax

	def summary(self):
		"""
		Return a dictionary containing the estimated contrast effect,
		the associated ReML-based estimation variance, and the estimated
		degrees	of freedom (variance of the variance).
		"""
		return {'effect':self.effect, 'variance':self.variance, 'dof':self.dof}
	     
	def stat(self, baseline=0.0):
		"""
		Return the decision statistic associated with the test of the
		null hypothesis: (H0) 'contrast equals baseline'
	        """
		self._baseline = baseline
		
		# Case: one-dimensional contrast ==> t or t**2
		if self.dim == 1: 
			# avoids division by zero
			t = (self.effect-baseline) / np.sqrt(np.maximum(self.variance, self._tiny)) 
			if self.type == 'F':
				t = t**2
		# Case: F contrast 
		elif self.type == 'F':
			# F = |t|^2/q ,  |t|^2 = e^t v-1 e  
			t = mahalanobis(self.effect-baseline, np.maximum(self.variance, self._tiny))/self.dim
		# Case: tmin (conjunctions)
		elif self.type == 'tmin':
			vdiag = self.variance.reshape([self.dim**2]+list(self.variance.shape[2:]))[::self.dim+1]
			t = (self.effect-baseline) / np.sqrt(np.maximum(vdiag, self._tiny))
			t = t.min(0)
			
		# Unknwon stat
		else:
			raise ValueError, 'Unknown statistic type'

		self._stat = t
		return t


	def pvalue(self, baseline=0.0):
		"""
		Return a parametric approximation of the p-value associated
		with the null hypothesis: (H0) 'contrast equals baseline'
	        """
		if self._stat == None or not self._baseline == baseline: 
			self._stat = self.stat(baseline)
		# Valid conjunction as in Nichols et al, Neuroimage 25, 2005. 
		if self.type in ['t', 'tmin']:
			p = sps.t.sf(self._stat, np.minimum(self.dof, self._dofmax))
		elif self.type == 'F':
			p = sps.f.sf(self._stat, self.dim, np.minimum(self.dof, self._dofmax))
		else:
			raise ValueError, 'Unknown statistic type'
		self._pvalue = p 
		return p 

	def zscore(self, baseline=0.0):
		"""
		Return a parametric approximation of the z-score associated
		with the null hypothesis: (H0) 'contrast equals baseline'
	        """  
		if self._pvalue == None or not self._baseline == baseline: 
			self._pvalue = self.pvalue(baseline)

		# Avoid inf values kindly supplied by scipy. 
		"""
		z = sps.norm.isf(self._pvalue)
		th = z[np.where(z<np.inf)].max()
		z = np.minimum(z, int(th+1))
		"""
		from nipy.neurospin.utils.zscore import zscore
		z = zscore(self._pvalue)		
		return z


		
	def __add__(self, other):
		if self.dim != other.dim:
			return None
		con = contrast(self.dim)
		con.type = self.type
		con.effect = self.effect + other.effect
		con.variance = self.variance + other.variance
		con.dof = self.dof + other.dof 
		return con

	def __rmul__(self, other):
		k = float(other)
		con = contrast(self.dim)
		con.type = self.type
		con.effect = k * self.effect
		con.variance = k**2 * self.variance
		con.dof = self.dof
		return con

	__mul__ = __rmul__

	def __div__(self, other):
		return self.__rmul__(1/float(other))


def ols(Y, X, axis=0):
	"""
	beta, nvbeta, s2, dof = ols(Y, X, axis=0)

        Essentially, compute pinv(X)*Y
	"""
	ndims = len(Y.shape)
	pX = np.linalg.pinv(X)
	beta = np.rollaxis(np.inner(pX, np.rollaxis(Y, axis, ndims)), 0, axis+1)
	nvbeta = np.inner(pX, pX)
	res = Y - np.rollaxis(np.inner(X, np.rollaxis(beta, axis, ndims)), 0, axis+1)
	n = res.shape[axis]
	s2 = (res**2).sum(axis) / float(n-X.shape[1])
	dof = float(X.shape[0] - X.shape[1])
	return beta, nvbeta, s2, dof



def load(file): 
	"""
	Load a fitted glm
	"""
	from os.path import splitext
	if splitext(file)[1]=='':
		file = file + '.npz'
	fmod = np.load(file)
	mod = glm()
	mod.beta = fmod['beta']
	mod.nvbeta = fmod['nvbeta']
	mod.s2 = fmod['s2']
	mod.dof = fmod['dof']
	mod.a = fmod['a']
	mod.model = str(fmod['model'])
	mod.method = str(fmod['method'])
	mod._axis = int(fmod['axis'])
	mod._constants = list(fmod['constants'])
	return mod


