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

#This module contains the argument class for the panel object

import stat_functions as stat
import numpy as np
import loglikelihood as logl
import random_effects as re
import calculus_functions as cf




class arguments:
	"""Sets initial arguments and stores static properties of the arguments"""
	def __init__(self,panel):
		p, q, d, k, m=panel.pqdkm
		self.initial_user_defined = False
		self.categories=['beta','rho','lambda','gamma','psi','omega']
		if panel.z_active:
			self.categories+=['z']
		self.mu_removed=True
		if not self.mu_removed:
			self.categories+=['mu']
		self.make_namevector(panel,p, q, k, m)
		initargs=self.initargs(p, d, q, m, k, panel)
		self.position_defs(initargs)
		self.set_init_args(panel,initargs)
		self.get_user_constraints(panel)
		
		
	def get_user_constraints(self,panel):
		e="User contraints must be a dict of dicts or a string evaluating to that, on the form of ll.args.dict_string. User constraints not applied"
		if type(panel.options.user_constraints.value)==dict:
			self.user_constraints=panel.options.user_constraints.value
		else:
			if panel.options.user_constraints.value is None or panel.options.user_constraints.value=='':
				self.user_constraints={}
				return
			try:
				self.user_constraints=eval(panel.options.user_constraints.value)
			except SyntaxError:
				print(f"Syntax error: {e}")
				self.user_constraints={}
				return			
			except:
				print(e)
				self.user_constraints={}
				return
		if not panel.z_active and 'z' in self.user_constraints:
			self.user_constraints.pop('z')	
			
		for grp in self.user_constraints:
			try:
				self.user_constraints[grp] = np.array(self.user_constraints[grp]).flatten()
				test = 1.0 * self.user_constraints[grp]
			except:
				print(f'Failed to apply user constraints for {grp}')
				self.user_constraints.pop('grp')
				
		for grp in self.user_constraints:
			if not grp in self.caption_d:
				print(f"Constraint on {grp} not applied, {grp} not in arguments")
			if len(self.caption_d[grp]) < len(self.user_constraints[grp]):
				print(f"Not all constraints on {grp} applied, more constraints than arguments for {grp}")
				
		

	def initargs(self,p,d,q,m,k,panel):

		args=dict()
		args['beta']=np.zeros((panel.X.shape[2],1))
		args['omega']=np.zeros((panel.W.shape[2],1))
		args['rho']=np.zeros(p)
		args['lambda']=np.zeros(q)
		args['psi']=np.zeros(m)
		args['gamma']=np.zeros(k)
		args['omega'][0][0]=0
		args['mu']=np.array([])
		args['z']=np.array([])			

		
		if m>0 and panel.z_active:
			args['z']=np.array([1e-09])	

		if panel.N>1 and not self.mu_removed:
			args['mu']=np.array([0.0001])			
			

		return args

	def set_init_args(self,panel,initargs=None, default = True):
		p, q, d, k, m=panel.pqdkm
		
		if initargs is None:
			initargs = self.initargs(p, d, q, m, k, panel)
			
		#de2=np.roll(e**2,1)-e**2
		#c=stat.correl(np.concatenate((np.roll(de2,1),de2),2),panel)[0,1]
		beta,omega = self.set_init_regression(initargs,panel, default)
		self.args_start=self.create_args(initargs,panel)
		self.args_init=self.create_args(initargs,panel)
		self.set_restricted_args(p, d, q, m, k,panel,omega,beta)
		self.n_args=len(self.args_init.args_v)
		

	def set_restricted_args(self,p, d, q, m, k, panel,omega,beta):
		args_restricted=self.initargs(p, d, q, m, k, panel)
		args_restricted['beta'][0][0]=np.mean(panel.Y)
		args_restricted['omega'][0][0]=np.log(panel.var(panel.Y))
		self.args_restricted=self.create_args(args_restricted,panel)
		
		args_OLS=self.initargs(p, d, q, m, k, panel)	
		args_OLS['beta']=beta
		args_OLS['omega'][0][0]=omega
		self.args_OLS=self.create_args(args_OLS,panel)
		
	
	def create_null_ll(self,panel):
		if not hasattr(self,'LL_OLS'):
			self.LL_OLS=logl.LL(self.args_OLS,panel).LL
			self.LL_null=logl.LL(self.args_restricted,panel).LL	
		
	def position_defs(self,initargs):
		"""Defines positions in vector argument"""

		self.positions=dict()
		self.positions_map=dict()#a dictionary of indicies containing the string name and sub-position of index within the category
		k=0
		for i in self.categories:
			n=len(initargs[i])
			rng=range(k,k+n)
			self.positions[i]=rng
			for j in rng:
				self.positions_map[j]=[0,i,j-k]#equation,category,relative position
			k+=n

	
	def conv_to_dict(self,args):
		"""Converts a vector argument args to a dictionary argument. If args is a dict, it is returned unchanged"""
		if type(args)==dict:
			return args
		if type(args)==list:
			args=np.array(args)			
		d=dict()
		k=0
		for i in self.categories:
			n=len(self.positions[i])
			rng=range(k,k+n)
			d[i]=np.array(args[rng])
			if i=='beta' or i=='omega':
				d[i]=d[i].reshape((n,1))
			k+=n
		return d


	def conv_to_vector(self,args):
		"""Converts a dict argument args to vector argument. if args is a vector, it is returned unchanged.\n
		If args=None, the vector of self.args_init is returned"""
		if type(args)==list or type(args)==np.ndarray:
			return np.array(args)
		v=np.array([])
		for i in self.categories:
			s=np.array(args[i])
			if len(s.shape)==2:
				s=s.flatten()
			if len(s)>0:
				v=np.concatenate((v,s))
		return v


	def make_namevector(self,panel,p, q, k, m):
		"""Creates a vector of the names of all regression varaibles, 
		including variables, ARIMA and GARCH terms. This defines the positions
		of the variables througout the estimation."""
		d, names_d = {}, {}
		captions=list(panel.input.X.keys())#copy variable names
		d['beta']=list(captions)
		c=[list(captions)]
		names = [f'x{i}' for i in range(panel.n_beta)]
		names_d['beta'] = list(names)
		add_names(p,'rho%s    AR    p','rho',d,c,captions, names, names_d)
		add_names(q,'lambda%s MA    q','lambda',d,c,captions, names, names_d)
		add_names(k,'gamma%s  GARCH k','gamma',d,c,captions, names, names_d)
		add_names(m,'psi%s    ARCH  m','psi',d,c,captions, names, names_d)
		
		
		d['omega']=list(panel.input.W.keys())
		captions.extend(panel.input.W.keys())
		
		names_d['omega'] = [f'omega{i}' for i in range(panel.nW)]
		names.extend(names_d['omega'])
		
		c.append(d['omega'])
		
		if m>0:
			if panel.N>1 and not self.mu_removed:
				d['mu']=['mu (var.ID eff.)']
				captions.extend(d['mu'])
				names_d['mu']=['mu']
				names.extend(d['mu'])				
				c.append(d['mu'])
			if panel.z_active:
				d['z']=['z in h(e,z)']
				captions.extend(d['z'])
				names_d['z']=['z']
				names.extend(d['z'])					
				c.append(d['z'])
			
		self.caption_v=captions
		self.caption_d=d
		self.names_v = names
		self.names_d = names_d
		self.names_category_list=c

	def create_args(self,args,panel,constraints=None):
		if isinstance(args,arguments_set):
			self.test_consistency(args)
			return args
		args_v=self.conv_to_vector(args)
		if not constraints is None:
			constraints.within(args_v,True)	
			constraints.set_fixed(args_v)
		args_d=self.conv_to_dict(args_v)
		dict_string=[]
		for c in self.categories:
			s=[]
			captions=self.caption_d[c]
			a=args_d[c].flatten()
			for i in range(len(captions)):
				s.append(f"'{captions[i]}':{a[i]}")
			dict_string.append(f"'{c}':\n"+"{"+",\n".join(s)+"}")
		dict_string="{"+",\n".join(dict_string)+"}"
		return arguments_set(args_d, args_v, dict_string, self,panel)
	
	def test_consistency(self,args):
		#for debugging only
		m=self.positions_map
		for i in m:
			dict_arg=args.args_d[m[i][1]]
			if len(dict_arg.shape)==2:
				dict_arg=dict_arg[m[i][2]]
			if dict_arg[0]!=args.args_v[i]:
				raise RuntimeError("argument inconsistency")
			
	def get_name_ix(self,x,single_item=False):
		#returns name, list of indicies
		if x is None:
			return None, None
		if x in self.caption_v:
			if single_item:
				indicies=self.caption_v.index(x)
			else:
				indicies=[self.caption_v.index(x)]	
			return x,indicies
		elif x in self.positions and not single_item:
			indicies=list(self.positions[x])
			return x,indicies
		elif x in self.names_v:
			if single_item:
				indicies=self.names_v.index(x)
			else:
				indicies=[self.names_v.index(x)]	
			return x,indicies			
		try:
			name=self.caption_v[x]
		except Exception as e:
			raise RuntimeError(f"{e}. The identifier of an argument must be an integer or a string macthing a name in 'self.caption_v' or a category in 'self.positions'")
		if single_item:
			return name,x
		else:
			return name,[x]
		
	def set_init_regression(self, initargs,panel, default):
		usrargs =  panel.options.arguments.value
		if not usrargs is None:
			if type(usrargs)==str:
				try:
					usrargs = eval(usrargs.replace(" array"," np.array"))
				except NameError as e:
					if str(e)=="name 'array' is not defined":
						usrargs = eval(usrargs.replace("array"," np.array"))
			args = self.create_args(usrargs,panel)
			for c in args.args_d:
				initargs[c] = args.args_d[c]
			self.initial_user_defined = True
			return initargs['beta'], initargs['omega'][0,0]
		p, q, d, k, m=panel.pqdkm
		beta,rho,lmbda,u=ARMA_regression(panel)
		v=panel.var(u)
		if panel.options.EGARCH.value==0:
			omega0=panel.var(u)
			omega1 = 0
		else:
			omega0=np.log(panel.var(u))
			
		if panel.nW>1:
			omega1 = panel.var(panel.W[:,:,1:2])
			omega1 = omega0/omega1
			a=0.00
			initargs['omega'][1,0]=omega1*a
			initargs['omega'][0,0]=omega0*(1-a)
		else:
			initargs['omega'][0,0]=omega0
		initargs['beta']=beta
		
		if default: #turns out calculating ARMA parameters serves no purpose, default is therefore allways true. These coefficients are easily found by linesearch for given regression coefficients. 
			if q > 0:
				initargs['lambda'][0] = lmbda
				
			if p > 0:
				initargs['rho'][0] = rho	
	
		
		if panel.options.fixed_random_variance_eff.value==0:
			if v<1e-20:
				print('Warning, your model may be over determined. Check that you do not have the dependent among the independents')	
		return beta,omega0

def set_GARCH(panel,initargs,u,m):
	matrices=logl.set_garch_arch(panel,initargs)
	if matrices is None:
		e=u
	else:
		AMA_1,AMA_1AR,GAR_1,GAR_1MA=matrices
		e = cf.dot(AMA_1AR,u)*panel.included[3]		
	h=h_func(e, panel,initargs)
	if m>0:
		initargs['gamma'][0]=0
		initargs['psi'][0]=0


def h_func(e,panel,initargs):
	z=None
	if len(initargs['z'])>0:
		z=initargs['z'][0]	
	h_val,h_e_val,h_2e_val,h_z,h_2z,h_e_z=logl.h(e,z,panel)
	return h_val*panel.included[3]
	

	
def ARMA_regression(panel):
	gfre=panel.options.fixed_random_group_eff.value
	tfre=panel.options.fixed_random_time_eff.value
	re_obj_i=re.re_obj(panel,True,panel.T_i,panel.T_i,gfre)
	re_obj_t=re.re_obj(panel,False,panel.date_count_mtrx,panel.date_count,tfre)
	X=(panel.X+re_obj_i.RE(panel.X, panel)+re_obj_t.RE(panel.X, panel))*panel.included[3]
	Y=(panel.Y+re_obj_i.RE(panel.Y, panel)+re_obj_t.RE(panel.Y, panel))*panel.included[3]
	beta,u=stat.OLS(panel,X,Y,return_e=True)
	rho,lmbda=ARMA_process_calc(u,panel)
	return beta,rho,lmbda,u

def ARMA_process_calc(e,panel):
	c=stat.correlogram(panel,e,2,center=True)[1:]
	
	rho = max(min(c[1]/c[0],0.99), -0.99)

	
	lmbda = 0
	try:
		lmbda1 = (1 - 2*c[0]*rho + rho**2)/(2*(c[0]-rho))
		lmbda2 = (( (rho**2 - 1)*(rho**2 - 1 + 4*c[0]**2 - 4*c[0]*rho) )**0.5) / (2*(c[0]-rho))
		
		if abs(lmbda1+lmbda2)>abs(lmbda1-lmbda2):
			lmbda = max(min(lmbda1 - lmbda2,0.99), -0.99)
		else:
			lmbda = max(min(lmbda1 + lmbda2,0.99), -0.99)
	except:
		pass

	return rho,lmbda
		
			
def add_names(T,captionstr,category,d,c,captions, names, names_d):
	a=[]
	n=[]
	if ' ' in captionstr:
		namestr = captionstr.split(' ')[0]
	for i in range(T):
		a.append(captionstr %(i,))
		n.append(namestr %(i,))
	captions.extend(a)
	names.extend(n)
	d[category]=a
	names_d[category]=n
	c.append(a)


class arguments_set:
	"""A class that contains the numeric arguments used in the maximization
	in all shapes and forms needed."""
	def __init__(self,args_d,args_v,dict_string,arguments,panel):
		self.args_d=args_d#dictionary of arguments
		self.args_v=args_v#vector of arguments
		self.dict_string=dict_string#a string defining a dictionary of named arguments. For user input of initial arguments
		self.caption_v=arguments.caption_v#vector of captions
		self.caption_d=arguments.caption_d#dict of captions
		self.names_v=arguments.names_v#vector of names
		self.names_d=arguments.names_d#dict of names		
		self.n_args=len(self.args_v)
		self.pqdkm=panel.pqdkm
		self.positions=arguments.positions
		self.names_category_list=arguments.names_category_list
		