
import warnings
import numpy as np
import designs,models
from .. import _datachecks, _reml, _spm


def aov(model, contrasts, f_terms):
	'''
	This code is borrowed and modified from statsmodels.stats.anova_lm
	'''
	effects = np.asarray( np.dot(model.QT, model.Y) )
	if model.dim==0:
		effects = effects.flatten()
	SS      = np.dot(contrasts.C, effects**2)
	DF      = np.asarray(contrasts.C.sum(axis=1), dtype=int)
	F       = []
	for term0,term1 in f_terms:
		i       = contrasts.term_labels.index(term0)
		ss0,df0 = SS[i], DF[i]
		ms0     = ss0 / df0
		if term1 == 'Error':
			ss1,df1,ms1  = model._SSE, model._dfE, model._MSE
		else:
			i       = contrasts.term_labels.index(term1)
			ss1,df1 = SS[i], DF[i]
			ms1     = ss1 / df1
		f           = ms0 / ms1
		if model.dim == 0:
			F.append( _spm.SPM0D_F(f, (df0,df1), (ss0,ss1), (ms0,ms1), model.eij, model.QT) )
		else:
			F.append( _spm.SPM_F(f, (df0,df1), model.fwhm, model.resels, model.X, model._beta, model.eij, model.QT) )
	return F


### ONE-WAY DESIGNS ##############

def anova1(Y, A=None, equal_var=False):
	'''
	One-way ANOVA.
	
	:Parameters (Option 1):
		- *Y* --- A list or tuple of (J x Q) numpy arrays
		- *equal_var* --- If *True*, equal group variance will be assumed


	:Parameters (Option 2):
		- *Y* --- (J x Q) numpy array
		- *A* --- (J x 1) vector of integer group labels
		- *equal_var* --- If *True*, equal group variance will be assumed

	
	:Returns:
		- F : An **spm1d._spm.SPM_F** instance
	
	:Example:
	
	>>> F = spm1d.stats.anova1((Y0,Y1,Y2))
	>>> Fi = F.inference(alpha=0.05)
	>>> Fi.plot()
	'''
	if isinstance(Y, (list,tuple)):
		_datachecks.check('anova1list', Y)
		A   = np.hstack([[i]*y.shape[0] for i,y in enumerate(Y)])
		Y   = np.hstack(Y) if Y[0].ndim==1 else np.vstack(Y)
	else:
		_datachecks.check('anova1', Y, A)
	design  = designs.ANOVA1(A)
	model   = models.LinearModel(Y, design.X)
	model.fit()
	F       = aov(model, design.contrasts, design.f_terms)[0]
	if not equal_var:
		warnings.warn('\nWARNING:  Non-sphericity corrections for one-way ANOVA are currently approximate and have not been verified.\n', UserWarning, stacklevel=2)
		Y,X,r = model.Y, model.X, model.eij
		Q,C   = design.A.get_Q(), design.contrasts.C.T
		F.df  = _reml.estimate_df_anova1(Y, X, r, Q, C)
	return F



def anova1rm(Y, A, SUBJ, equal_var=True):
	'''
	One-way repeated-measures ANOVA.
	
	:Parameters:
		- *Y* --- (J x Q) numpy array
		- *A* --- (J x 1) vector of integer group labels
		- *SUBJ* --- (J x 1) vector of subject labels
		- *equal_var* --- If *True*, equal group variance will be assumed

	
	:Returns:
		- F : An **spm1d._spm.SPM_F** instance
	
	:Example:
	
	>>> Y = np.random.randn(9, 101)
	>>> A = np.array([1,1,1, 2,2,2, 3,3,3])
	>>> SUBJ = np.array([1,2,3, 1,2,3, 1,2,3])
	>>> F = spm1d.stats.anova1(Y, A, SUBJ)
	>>> Fi = F.inference(alpha=0.05)
	>>> Fi.plot()
	'''
	if not equal_var:
		raise( NotImplementedError( 'Non-sphericity corrections are not yet implemented. Set "equal_var" to "True" to force an assumption of equal variance.' ) )
	design  = designs.ANOVA1rm(A, SUBJ)
	model   = models.LinearModel(Y, design.X)
	model.fit()
	F       = aov(model, design.contrasts, design.f_terms)[0]
	return F






### TWO-WAY DESIGNS ##############


def anova2(Y, A, B, equal_var=True):
	'''
	Two-way ANOVA.
	
	:Parameters:
		- *Y* --- (J x Q) numpy array
		- *A* --- (J x 1) vector of integer labels for Factor A
		- *B* --- (J x 1) vector of integer labels for Factor B
		- *equal_var* --- If *True*, equal group variance will be assumed

	:Returns:
		- List of three **spm1d._spm.SPM_F** instances in the following order:
			1. Main effect A
			2. Main effect B
			3. Interaction AB
	'''
	if not equal_var:
		raise( NotImplementedError( 'Non-sphericity corrections are not yet implemented. Set "equal_var" to "True" to force an assumption of equal variance.' ) )
	design  = designs.ANOVA2(A, B)
	model   = models.LinearModel(Y, design.X)
	model.fit()
	F       = aov(model, design.contrasts, design.f_terms)
	# if not equal_var:
		# Y,X,r   = model.Y, model.X, model.eij
		# QA,QB,C = design.A.get_Q(), design.B.get_Q(), design.contrasts.C.T
		# Q       = QA + QB
		# u1,u2   = _reml.estimate_df_anova2(Y, X, r, Q, C)
		# for ff,u in zip(f,u1):
		# 	ff.df = u,u2
	return F


def anova2nested(Y, A, B, equal_var=True):
	'''
	Two-way nested ANOVA.
	
	:Parameters:
		- *Y* --- (J x Q) numpy array
		- *A* --- (J x 1) vector of integer labels for Factor A
		- *B* --- (J x 1) vector of integer labels for Factor B (nested in A)
		- *equal_var* --- If *True*, equal group variance will be assumed

	:Returns:
		- List of two **spm1d._spm.SPM_F** instances in the following order:
			1. Main effect A
			2. Main effect B

	:Note:
		- there is no interaction term in nested designs.
	'''
	if equal_var is not True:
		raise( NotImplementedError('Non-sphericity correction not implemented. To continue you must assume equal variance and set "equal_var=True".') )
	design  = designs.ANOVA2nested(A, B)
	model   = models.LinearModel(Y, design.X)
	model.fit()
	F       = aov(model, design.contrasts, design.f_terms)
	return F



def anova2rm(Y, A, B, SUBJ, equal_var=True):
	'''
	Two-way repeated-measures ANOVA.
	
	:Parameters:
		- *Y* --- (J x Q) numpy array
		- *A* --- (J x 1) vector of integer labels for Factor A
		- *B* --- (J x 1) vector of integer labels for Factor B
		- *SUBJ* --- (J x 1) vector of integer subject labels
		- *equal_var* --- If *True*, equal group variance will be assumed

	:Returns:
		- List of three **spm1d._spm.SPM_F** instances in the following order:
			1. Main effect A
			2. Main effect B
			3. Interaction AB
			
	:Note:
		- Non-sphericity correction not implemented. Equal variance must be assumed by setting "equal_var=True".
	'''
	if equal_var is not True:
		raise( NotImplementedError('Non-sphericity correction not implemented. To continue you must assume equal variance and set "equal_var=True".') )
	design  = designs.ANOVA2rm(A, B, SUBJ)
	model   = models.LinearModel(Y, design.X)
	if (model.dim == 1) and ( design.check_for_single_responses() ):
		model.fit( approx_residuals=design.contrasts.C[:5] )
	else:
		model.fit( )
	F       = aov(model, design.contrasts, design.f_terms)
	# if not equal_var:
	# 	Y,X,r   = solver.Y, solver.X, solver.eij
	# 	QA,QB,C = design.A.get_Q(), design.B.get_Q(), [c.C.T for c in design.contrasts]
	# 	Q       = QA + QB
	# 	u1,u2   = _reml.estimate_df_anova2(Y, X, r, Q, C)
	# 	for ff,u in zip(f,u1):
	# 		ff.df = u,u2
	return F


def anova2onerm(Y, A, B, SUBJ, equal_var=True):
	'''
	Two-way ANOVA with repeated-measures on one factor.
	
	:Parameters:
		- *Y* --- (J x Q) numpy array
		- *A* --- (J x 1) vector of integer labels for Factor A
		- *B* --- (J x 1) vector of integer labels for Factor B (the repeated-measures factor)
		- *SUBJ* --- (J x 1) vector of integer subject labels
		- *equal_var* --- If *True*, equal group variance will be assumed

	:Returns:
		- List of three **spm1d._spm.SPM_F** instances in the following order:
			1. Main effect A
			2. Main effect B
			3. Interaction AB
			
	:Note:
		- Non-sphericity correction not implemented. Equal variance must be assumed by setting "equal_var=True".
	'''
	if equal_var is not True:
		raise( NotImplementedError('Non-sphericity correction not implemented. To continue you must assume equal variance and set "equal_var=True".') )
	design  = designs.ANOVA2onerm(A, B, SUBJ)
	model   = models.LinearModel(Y, design.X)
	if (model.dim == 1) and ( design.check_for_single_responses() ):
		model.fit( approx_residuals=design.contrasts.C[:5] )
	else:
		model.fit( )
	F       = aov(model, design.contrasts, design.f_terms)
	return F




### THREE-WAY DESIGNS ##############


def anova3(Y, A, B, C, equal_var=True):
	'''
	Three-way ANOVA.
	
	:Parameters:
		- *Y* --- (J x Q) numpy array
		- *A* --- (J x 1) vector of integer labels for Factor A
		- *B* --- (J x 1) vector of integer labels for Factor B
		- *C* --- (J x 1) vector of integer labels for Factor C
		- *equal_var* --- If *True*, equal group variance will be assumed

	:Returns:
		- List of seven **spm1d._spm.SPM_F** instances in the following order:
			1. Main effect A
			2. Main effect B
			3. Main effect C
			4. Interaction AB
			5. Interaction AC
			6. Interaction BC
			7. Interaction ABC
	'''
	if equal_var is not True:
		raise( NotImplementedError('Non-sphericity correction not implemented. To continue you must assume equal variance and set "equal_var=True".') )
	design  = designs.ANOVA3(A, B, C)
	model   = models.LinearModel(Y, design.X)
	model.fit()
	F       = aov(model, design.contrasts, design.f_terms)
	return F
	# if not equal_var:
	# 	Y,X,r   = solver.Y, solver.X, solver.eij
	# 	QA,QB,QC,C = design.A.get_Q(), design.B.get_Q(), design.C.get_Q(), [c.C.T for c in design.contrasts]
	# 	Q       = QA + QB + QC
	# 	u1,u2   = _reml.estimate_df_anova2(Y, X, r, Q, C)
	# 	for ff,u in zip(f,u1):
	# 		ff.df = u,u2
	# return f

def anova3nested(Y, A, B, C, equal_var=True):
	'''
	Three-way fully nested ANOVA.
	
	:Parameters:
		- *Y* --- (J x Q) numpy array
		- *A* --- (J x 1) vector of integer labels for Factor A
		- *B* --- (J x 1) vector of integer labels for Factor B (nested in A)
		- *C* --- (J x 1) vector of integer labels for Factor C (nested in B)
		- *equal_var* --- If *True*, equal group variance will be assumed

	:Returns:
		- List of three **spm1d._spm.SPM_F** instances in the following order:
			1. Main effect A
			2. Main effect B
			3. Main effect C

	:Note:
		- there are no interaction terms in fully-nested designs.
	'''
	if equal_var is not True:
		raise( NotImplementedError('Non-sphericity correction not implemented. To continue you must assume equal variance and set "equal_var=True".') )
	design  = designs.ANOVA3nested(A, B, C)
	model   = models.LinearModel(Y, design.X)
	model.fit()
	F       = aov(model, design.contrasts, design.f_terms)
	return F


def anova3rm(Y, A, B, C, SUBJ, equal_var=True):
	'''
	Three-way ANOVA (repeated measures on all factors).
	
	:Parameters:
		- *Y* --- (J x Q) numpy array
		- *A* --- (J x 1) vector of integer labels for Factor A
		- *B* --- (J x 1) vector of integer labels for Factor B
		- *C* --- (J x 1) vector of integer labels for Factor C
		- *equal_var* --- If *True*, equal group variance will be assumed

	:Returns:
		- List of seven **spm1d._spm.SPM_F** instances in the following order:
			1. Main effect A
			2. Main effect B
			3. Main effect C
			4. Interaction AB
			5. Interaction AC
			6. Interaction BC
			7. Interaction ABC
	'''
	if equal_var is not True:
		raise( NotImplementedError('Non-sphericity correction not implemented. To continue you must assume equal variance and set "equal_var=True".') )
	design  = designs.ANOVA3rm(A, B, C, SUBJ)
	model   = models.LinearModel(Y, design.X)
	if (model.dim == 1) and ( design.check_for_single_responses() ):
		model.fit( approx_residuals=design.contrasts.C[:8] )
	else:
		model.fit( )
	F       = aov(model, design.contrasts, design.f_terms)
	return F
	


def anova3onerm(Y, A, B, C, SUBJ, equal_var=True):
	'''
	Three-way ANOVA with repeated-measures on one factor.
	
	:Parameters:
		- *Y* --- (J x Q) numpy array
		- *A* --- (J x 1) vector of integer labels for Factor A
		- *B* --- (J x 1) vector of integer labels for Factor B
		- *C* --- (J x 1) vector of integer labels for Factor C (the repeated-measures factor)
		- *SUBJ* --- (J x 1) vector of integer subject labels
		- *equal_var* --- If *True*, equal group variance will be assumed

	:Returns:
		- List of seven **spm1d._spm.SPM_F** instances in the following order:
			1. Main effect A
			2. Main effect B
			3. Main effect C
			4. Interaction AB
			5. Interaction AC
			6. Interaction BC
			7. Interaction ABC
			
	:Note:
		- Non-sphericity correction not implemented. Equal variance must be assumed by setting "equal_var=True".
	'''
	if equal_var is not True:
		raise( NotImplementedError('Non-sphericity correction not implemented. To continue you must assume equal variance and set "equal_var=True".') )
	design  = designs.ANOVA3onerm(A, B, C, SUBJ)
	model   = models.LinearModel(Y, design.X)
	if (model.dim == 1) and ( design.check_for_single_responses() ):
		model.fit( approx_residuals=design.contrasts.C[:8] )
	else:
		model.fit( )
	model.fit()
	F       = aov(model, design.contrasts, design.f_terms)
	return F



def anova3tworm(Y, A, B, C, SUBJ, equal_var=True):
	'''
	Three-way ANOVA with repeated-measures on two factors.
	
	:Parameters:
		- *Y* --- (J x Q) numpy array
		- *A* --- (J x 1) vector of integer labels for Factor A
		- *B* --- (J x 1) vector of integer labels for Factor B (a repeated-measures factor)
		- *C* --- (J x 1) vector of integer labels for Factor C (a repeated-measures factor)
		- *SUBJ* --- (J x 1) vector of integer subject labels
		- *equal_var* --- If *True*, equal group variance will be assumed

	:Returns:
		- List of seven **spm1d._spm.SPM_F** instances in the following order:
			1. Main effect A
			2. Main effect B
			3. Main effect C
			4. Interaction AB
			5. Interaction AC
			6. Interaction BC
			7. Interaction ABC
			
	:Note:
		- Non-sphericity correction not implemented. Equal variance must be assumed by setting "equal_var=True".
	'''
	if equal_var is not True:
		raise( NotImplementedError('Non-sphericity correction not implemented. To continue you must assume equal variance and set "equal_var=True".') )
	design  = designs.ANOVA3tworm(A, B, C, SUBJ)
	model   = models.LinearModel(Y, design.X)
	if (model.dim == 1) and ( design.check_for_single_responses() ):
		model.fit( approx_residuals=design.contrasts.C[:8] )
	else:
		model.fit( )
	F       = aov(model, design.contrasts, design.f_terms)
	return F


