#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np
import time

try:
	import cfunctions as c
except:
	c=None


def dd_func_lags_mult(panel,ll,g,AMAL,vname1,vname2,transpose=False, u_gradient=False):
	de2_zeta_xi_RE,de2_zeta_xi 	= dd_func_lags_mult_arima(panel,ll,g,AMAL,vname1,vname2,transpose, u_gradient)
	dd_re_variance 							= dd_func_re_variance(panel,ll,g,vname1,vname2,de2_zeta_xi_RE,u_gradient)
	d2LL_d2var_zeta_xi,d2LL_d2e_zeta_xi_RE 	= dd_func_garch(panel,ll,g,vname1,vname2,de2_zeta_xi_RE,de2_zeta_xi,dd_re_variance,u_gradient)
	
	return d2LL_d2var_zeta_xi,d2LL_d2e_zeta_xi_RE

def dd_func_lags_mult_arima(panel,ll,g,AMAL,vname1,vname2,transpose, u_gradient):
	#de_xi is "N x T x m", de_zeta is "N x T x k" and L is "T x T"
	de_xi=g.__dict__['de_'+vname1]
	de_zeta=g.__dict__['de_'+vname2]
	de2_zeta_xi_RE=None
	if de_xi is None or de_zeta is None:
		return None,None
	(N,T,m)=de_xi.shape
	(N,T,k)=de_zeta.shape	
	#ARIMA:
	if not AMAL is None:
		if u_gradient:
			de2_zeta_xi=-panel.arma_dot.dot(AMAL,g.X_RE,ll)#"T x N x s x m #for error beta-rho covariance, the u gradient must be used	
		else:
			de2_zeta_xi=panel.arma_dot.dot(AMAL,de_zeta,ll)#"T x N x s x m
		if transpose:#only happens if lags==k
			de2_zeta_xi=de2_zeta_xi+np.swapaxes(de2_zeta_xi,2,3)#adds the transpose
		de2_zeta_xi=de2_zeta_xi*panel.included[4]
	else:
		de2_zeta_xi=None
	de2_zeta_xi_RE=de2_zeta_xi
	
	return de2_zeta_xi_RE,de2_zeta_xi


def dd_func_re_variance(panel,ll,g,vname1,vname2,de2_zeta_xi_RE,u_gradient):
	if panel.N<=1 or panel.options.fixed_random_group_eff.value==0:
		return None
	#voaltility RE:

	de_xi_RE=g.__dict__['de_'+vname1+'_RE']
	de_zeta_RE=g.__dict__['de_'+vname2+'_RE']
	if de_xi_RE is None or de_zeta_RE is None:
		return None
	(N,T,m)=de_xi_RE.shape
	(N,T,k)=de_zeta_RE.shape		
	de_xi_RE_r=de_xi_RE.reshape((N,T,m,1))
	de_zeta_RE_r=de_zeta_RE.reshape((N,T,1,k))	
	incl=panel.included[4]
	
	dvRE_xi=g.__dict__['dvRE_'+vname1].reshape((N,T,m,1))
	dvRE_zeta=g.__dict__['dvRE_'+vname2].reshape((N,T,1,k))			
	
	if not de2_zeta_xi_RE is None:
		d_xi_input=g.__dict__['d_'+vname1+'_input']
		d_zeta_input=g.__dict__['d_'+vname2+'_input']		
		try:
			dd_e_RE_sq=(2*de_xi_RE_r*de_zeta_RE_r+2*ll.e_RE.reshape(N,T,1,1)*de2_zeta_xi_RE)*incl
		except:
			dd_e_RE_sq=(2*de_xi_RE_r*de_zeta_RE_r+2*ll.e_RE.reshape(N,T,1,1)*de2_zeta_xi_RE)*incl
		ddmeane2= panel.mean(dd_e_RE_sq,(0,1))*incl
		dd_input=(dd_e_RE_sq-ddmeane2)*incl

		ddvRE_d_xi_zeta=ddmeane2-add((ll.re_obj_i_v.ddRE(dd_input,d_xi_input,d_zeta_input,ll.varRE_input,vname1,vname2,panel),
									  ll.re_obj_t_v.ddRE(dd_input,d_xi_input,d_zeta_input,ll.varRE_input,vname1,vname2,panel)),True)
	else:
		ddvRE_d_xi_zeta=None
		
	dvarRE,		ddvarRE		=	ll.dvarRE.reshape((N,T,1,1)),		ll.ddvarRE.reshape((N,T,1,1))
	dd_re_variance	=	add((prod((dvarRE,ddvRE_d_xi_zeta))	,	ddvarRE*dvRE_xi*dvRE_zeta))
	return dd_re_variance


def dd_func_garch(panel,ll,g,vname1,vname2,de2_zeta_xi_RE,de2_zeta_xi,dd_re_variance,u_gradient):
	#GARCH: 
	incl=panel.included[4]
	if (g.__dict__['de_'+vname1] is None) or (g.__dict__['de_'+vname2] is None):
		return None, None	
	(N,T,m)=g.__dict__['de_'+vname1].shape
	(N,T,k)=g.__dict__['de_'+vname2].shape
	DLL_e=g.DLL_e.reshape(N,T,1,1)
	
	d2LL_d2e_zeta_xi_RE=None
	if not de2_zeta_xi_RE is None:	
		d2LL_d2e_zeta_xi_RE = de2_zeta_xi_RE * DLL_e	
		d2LL_d2e_zeta_xi_RE = np.sum(np.sum(d2LL_d2e_zeta_xi_RE*incl,0),0)
		
	RE_suffix=''

	de_xi_RE=g.__dict__['de_'+vname1+RE_suffix]
	de_zeta_RE=g.__dict__['de_'+vname2+RE_suffix]	

	
	dLL_var=g.dLL_var.reshape(N,T,1,1)
	if de2_zeta_xi_RE is None:
		de2_zeta_xi_RE=0
	d2var_zeta_xi=None
	d2LL_d2var_zeta_xi=None
	d_omega_e=None
	d2var_zeta_xi_h=None
	if panel.pqdkm[4]>0:
		if u_gradient:
			de_zeta_RE=g.__dict__['de_'+vname2+RE_suffix]
		h_e_de2_zeta_xi =  ll.h_e_val.reshape(N,T,1,1)  * de2_zeta_xi_RE
		h_2e_dezeta_dexi = ll.h_2e_val.reshape(N,T,1,1) * de_xi_RE.reshape((N,T,m,1)) * de_zeta_RE.reshape((N,T,1,k))

		d2var_zeta_xi_h = (h_e_de2_zeta_xi + h_2e_dezeta_dexi)
		
		d2var_zeta_xi_h = panel.arma_dot.dot(ll.GAR_1MA, d2var_zeta_xi_h,ll)
		
	d2var_zeta_xi = add((d2var_zeta_xi_h,  dd_re_variance,d_omega_e), True)
	if not d2var_zeta_xi is None:
		d2LL_d2var_zeta_xi = np.sum(prod((d2var_zeta_xi,dLL_var*incl),True),(0,1))
	
	return d2LL_d2var_zeta_xi,d2LL_d2e_zeta_xi_RE

def dd_func_lags(panel,ll,L,d,dLL,transpose=False):
	#d is "N x T x m" and L is "k x T x T"
	if panel.pqdkm[4]==0:
		return None
	if d is None:
		return None		
	(N,T,m)=d.shape
	if L is None:
		x=0
	elif len(L)==0:
		return None
	elif type(L)==tuple:#elif len(L.shape)==3:
		x=panel.arma_dot.dot(L,d,ll)#"T x N x k x m"
		if x is None:
			return None
	elif len(L.shape)==2:
		x=dot(L,d).reshape(N,T,1,m)
	dLL=dLL.reshape((N,T,1,1))
	return np.sum(np.sum(dLL*x,1),0)#and sum it	


def add(iterable,ignore=False):
	"""Sums iterable. If ignore=True all elements except those that are None are added. If ignore=False, None is returned if any element is None. """
	x=None
	for i in range(len(iterable)):
		if not iterable[i] is None:
			if x is None:
				x=iterable[i]
			else:
				x=x+iterable[i]
		else:
			if not ignore:
				return None
	return x

def prod(iterable,ignore=False):
	"""Takes the product sum of iterable. If ignore=True all elements except those that are None are multiplied. 
	If ignore=False, None is returned if any element is None. """
	x=None
	for i in iterable:
		if not i is None:
			if x is None:
				x=i
			else:
				x=x*i
		else:
			if not ignore:
				return None
	return x

def sumNT(nparray):
	if nparray is None:
		return None
	s=nparray.shape
	if len(s)<3:
		raise RuntimeError("Not enough dimensions")
	
	return np.sum(nparray.reshape(list(s)+[1]),(0,1))
	

def concat_matrix(block_matrix):
	m=[]
	for i in range(len(block_matrix)):
		r=block_matrix[i]
		C=[]
		for j in range(len(r)):
			if not r[j] is None:
				C.append(r[j])
		if len(C):
			m.append(np.concatenate(C,1))
	m=np.concatenate(m,0)
	return m

def concat_marray(matrix_array):
	arr=[]
	for i in matrix_array:
		if not i is None:
			arr.append(i)
	arr=np.concatenate(arr,2)
	return arr
		
def dd_func(d2LL_de2,d2LL_dln_de,d2LL_dln2,de_dh,de_dg,dln_dh,dln_dg,dLL_de2_dh_dg,dLL_dln2_dh_dg):
	a=[]
	a.append(dd_func_mult(de_dh,d2LL_de2,de_dg))

	a.append(dd_func_mult(de_dh,d2LL_dln_de,dln_dg))
	a.append(dd_func_mult(dln_dh,d2LL_dln_de,de_dg))

	a.append(dd_func_mult(dln_dh,d2LL_dln2,dln_dg))

	a.append(dLL_de2_dh_dg)
	a.append(dLL_dln2_dh_dg)
	return add(a,True)

def dd_func_mult(d0,mult,d1):
	#d0 is N x T x k and d1 is N x T x m
	if d0 is None or d1 is None or mult is None:
		return None
	(N,T,k)=d0.shape
	(N,T,m)=d1.shape
	if np.any(np.isnan(d0)) or np.any(np.isnan(d1)):
		x=np.empty((k,m))
		x[:]=np.nan
		return x
	d0=d0*mult
	d0=np.reshape(d0,(N,T,k,1))
	d1=np.reshape(d1,(N,T,1,m))
	try:
		x=np.sum(np.sum(d0*d1,0),0)#->k x m 
	except RuntimeWarning as e:
		if e.args[0]=='overflow encountered in multiply':
			d0=np.minimum(np.maximum(d0,-1e+100),1e+100)
			d1=np.minimum(np.maximum(d1,-1e+100),1e+100)
			x=np.sum(np.sum(d0*d1,0),0)#->k x m 
		else:
			raise RuntimeWarning(e)
	return x


def dot(a,b,reduce_dims=True):
	"""Matrix multiplication. Returns the dot product of a*b where either a or be or both to be
	arrays of matrices. Faster than mmult, less general and only used for special purpose.
	Todo: generalize and merge"""


	if len(a.shape)==3 and len(b.shape)==2:
		x = np.array([np.dot(a[i],b) for i in range(a.shape[0])])
	elif len(a.shape)==3 and len(b.shape)==3:
		x = np.sum([np.dot(a[i].T,b[i]) for i in range(a.shape[0])],0)
	elif len(a.shape)==2 and len(b.shape)==3:
		x = np.array([np.dot(a,b[i]) for i in range(b.shape[0])])
	return x


class arma_dot_obj:
	def __init__(self):
		pass
		
	def dotroll(self,aband,k,sign,b,ll):
		x = sign*self.fast_dot(aband, b)
		w=[]
		for i in range(k):
			w.append(np.roll(np.array(x),i+1,1))
			w[i][:,:i+1]=0
		x=np.array(w)
		x=np.moveaxis(x,0,2)
		return x
			
			
	def fast_dot(self, a, b):
		a_, name = a
		n = get_n(a_)
		if n is None:
			n = len(a_)
			
		r = a_[0]*b
		for i in range(1,n):
			r[:,i:] += a_[i]*b[:,:-i]
			
		return r

		
	def dot(self,a,b,ll):
		if len(a)>2:#then this is a proper matrix
			(aband,k,sgn)=a
			if k==0:
				return None
			return self.dotroll(aband, k, sgn, b, ll)
		x = self.fast_dot(a, b)
		return x


def get_n(a):
	minval = 0
	a_1 = np.abs(a[1:])
	max_a = np.max(a_1)
	if np.min(np.abs(a_1)) >= minval:
		return None
	if max_a == 0:
		return 1		
	else:
		nz = np.nonzero(a_1/max_a < minval)[0]
		if len(nz)>0:
			return nz[0]+1
		else:
			return None