/*********************************************************

 Shapes the vector containing all initial values into 
 the needed vectors and matrices for the log-likelihood

*********************************************************/

function scalar HIP_params_shape(matrix *theta, matrix *beta, 
 				 matrix *alpha,  matrix *Pi,
				 matrix *psi, matrix *C, 
				 scalar k1, scalar k2, 
				 scalar p, scalar q) 

    scalar h = k1 + p + q
    matrix beta = theta[1: k1+p]

    if q > 0
	matrix alpha = theta[k1+p+1:h]
    endif
	
    if (p!=0)
	scalar k = k1+k2
	Pi = mshape(theta[h+1 : h+p*k],k,p)
	matrix psi = theta[h+p*k+1:h + p*(k+1),]
	matrix C = lower(unvech(theta[h+p*(k+1)+1:,]))
    else
	Pi = {}
	psi = {}
	C = {}
    endif

    return 0
end function

function matrix EndoVarNormalize(matrix E)
    matrix s = diag(mcov(E))
    scalar h = 0.5
    s = s .^ h
    return s
end function

function matrix ExoVarNormalize(matrix X)

    # we exclude from the treatment the constant
    # and the dummy variables

    k = cols(X)
    s = ones(k,1)
    matrix chk = minc((X .= 0) + (X .= 1))
    matrix   s = chk + (1-chk) .* sdc(X)
    return s'
end function

function scalar CraggDonald(matrix Sigma, matrix Y, matrix X, 
			    matrix Z)

    p = cols(Y)
    matrix U = {}
    mols(Z, X, &U)
    B = mols(Y, U)
    matrix l = eigsolve((Y'U) * B, Sigma)
    return minc(l)/cols(Z)
end function

function matrix InitProbit(series y, list W, scalar *err)

    scalar k = nelem(W)
    catch probit y W --quiet	# here we can run into the perfect-prediction problem
    err = $error
    if err == 0
	list DROPPED = W - $xlist
	# handle some failures
        if nelem(DROPPED) >0
            printf "The variable(s) %s were dropped\n", varname(DROPPED)
            printf "Initializing to zeros and hoping for the best.\n"
            matrix beta = zeros(k,1)
        else
            matrix beta = $coeff
        endif
	ret = $lnl | beta
    else
	printf "Warning: probit returned error %d (%s)\n", err, errmsg(err)
	printf "This is weird and disturbing.\n"
	ret = {NA} | zeros(k,1)
    endif

    return ret

end function

/*********************************************************

 Two-step estimation to compute initial
 values. Parameters estimated covariance matrix is not
 computed.

*********************************************************/

function matrix InitParm(bundle *b)

    matrix HETVAR = b["mHETVAR"]
    matrix ENDOG = b["mENDOG"]
    matrix Z = b["mZ"] # total regressors
    matrix X = b["mX"] # total instruments
    series y = b["depvar"]
    scalar k1 = b["mk1"]
    scalar k = b["mk"]
    scalar p = b["mp"]
    scalar q = b["mq"]
    scalar h = b["mh"]
    
    scalar het = b["het"]
    scalar iv = b["iv"]
    list W
    loop i=1..h -q
        series z$i = Z[,$i]
        W += z$i
    endloop

    RESCALEX = 0 # Experimental; not for now

    if RESCALEX
	matrix rescale = ExoVarNormalize(X[,1:k1])	
	b["rescaleX"] = rescale
	X[,1:k1] = X[,1:k1] ./ rescale'
	b["mX"] = X
	b["mEXOG"] = b["mEXOG"] ./ rescale'
    else
	matrix rescale = ones(k1,1)
	b["rescaleX"] = rescale
    endif

    matrix E = { }
    matrix V_Pi = { }
    if iv
	matrix Pi = mols(ENDOG, X, &E, &V_Pi)
	matrix Sigma = mcov(E)

	lambda = CraggDonald(Sigma, ENDOG, b["mEXOG"], b["mADDIN"])
	b["CraggDonald"] = lambda
	if lambda < 1.0e-2
	    printf "Weak instruments (lambda = %g); expect trouble.\n", lambda
	endif

	matrix rescale = EndoVarNormalize(E)
	b["rescaleY"] = rescale

	if rows(rescale) > 0
	    ENDOG = ENDOG ./ rescale'
	    b["mENDOG"] = ENDOG
	    Pi = Pi ./ rescale'
	    E  = E ./ rescale'
	    Sigma = Sigma ./ (rescale .* rescale')
	endif

	b["uhat"] = E
	loop i=1..p -q
	    series e$i = E[,$i]
	    W += e$i
	endloop

	matrix C = cholesky(invpd(Sigma))
    else
	matrix Pi = {}
	matrix psi = {}
    	matrix C = {}
    endif	    

    scalar err = 0
    ##BUG: cambiato il nome da beta a inipar
    # (conflitto con il vecchio script)
    matrix inipar = InitProbit(y, W, &err)

    if err > 0
        print "Initial probit failed (%s). Aborting.", errmsg(error)
	beta = {}
	scalar b["lnl0"] = NA
	return beta
    else
	scalar b["lnl0"] = inipar[1]
	matrix beta = inipar[2:h+1]
    endif

    if het
	series ndx = Z * beta
	series mills = y ? invmills(-ndx) : -invmills(ndx)
	series s2 = -ndx*mills
	alpha = mols(s2 , 1 ~ HETVAR)
	alpha = alpha[2:]
    else
	alpha = {}
    endif

    if iv
        ##BUG
	#gamma = $coeff[h+1: ]
	matrix gamma = inipar[h+2:]
 	scalar scale = sqrt(1 + qform(gamma', Sigma))
	matrix beta = beta ./ scale
	matrix psi = inv(C')*gamma ./ scale
    endif

    theta = beta | alpha | vec(Pi) | psi | vech(C)
    return theta

end function

function scalar rescale_results(bundle *mod, matrix *theta, matrix *vcv,
				matrix *G)

    # used to undo the scaling of the endogenous variables once estimation
    # is done; while we're at it, we also compute the Jacobian term to correct
    # the loglikelihood by

    matrix sY = mod["rescaleY"]
    matrix sX = mod["rescaleX"]
    matrix isY = 1 ./ sY
    matrix isX = 1 ./ sX
    scalar p = rows(sY)
    scalar k1 = mod["mk1"]
    scalar k2 = mod["mk2"]

    if mod["het"]
	s_a = ones(mod["mq"],1)
    else
	s_a = {}
    endif

    if mod["iv"]
	a_Pi = vec( (isX .* sY') | mshape(sY, p, k2)' )
	a_end = ones(p,1) | vech(mshape(isY,p,p)') 
    else
	a_Pi = {}
	a_end = {}
    endif    

    matrix a = isX | isY | s_a | a_Pi | a_end

    theta = theta .* a
    vcv   = vcv .* (a .* a')
    G     = G ./ a'
#   mod["mENDOG"] = mod["mENDOG"] .* sY'

    return sumc(ln(sY))

end function

  /******************************
       Log-likelihood

   ****************************/

function series loglik_m(matrix *Y, matrix *X, matrix *Pi, matrix *C,
			 matrix *ScaledRes)
    # computes the marginal loglikelihood: the argument ScaledRes will 
    # hold in output the re-scaled residuals (used in the conditional 
    # loglik)

    matrix ScaledRes = (Y - X*Pi) * C
    scalar J = sumc(ln(diag(C))) - cols(Y)*.91893853320467274178
    return J - 0.5*sumr(ScaledRes .^2)
end function

function scalar HIP_loglik(series y, matrix *EXOG, matrix *ENDOG[null],
			   matrix *ADDIN[null], matrix *HETVAR[null],
			   matrix *beta, matrix *alpha[null], 
			   matrix *Pi[null], matrix *psi[null], 
			   matrix *C[null], series *ll, matrix *omega, 
			   series *sigma)

    # computes the total loglikelihood: the arguments omega and sigma will 
    # hold in output the re-scaled residual and the conditional variance 
    # to avoid re-computing them later (eg for the score matrix); note that
    # sigma holds the conditional standard error (NOT the variance)

    scalar het = (rows(alpha) > 0)
    scalar iv = (rows(Pi) > 0)
    series ll = NA
    scalar err = 0
 
    # do checks first
    if iv
	scalar cVar = 1-psi'psi
	scalar pdCheck = (minc(diag(C)) > 1.0e-8) && \
	    (cVar > 1.0e-20)
	err = pdCheck ? 0 : 1
    endif

    if het && !err
	series ndxv = HETVAR * alpha
	scalar cvCheck = (max(ndxv) < 100)
	err = cvCheck ? 0 : 2
    endif

    if err
	return err
    endif

    # ok, we should be within 'nice' bounds by now

    if iv # compute marginal loglikelihood
	matrix X = EXOG ~ ADDIN
	matrix omega = {}
	series llm = loglik_m(&ENDOG, &X, &Pi, &C, &omega)
	if sum(missing(llm)) > 0
	    err = 3
	    printf "Problems with marginal loglik!\n"
	    return err
	endif
    endif

    # now compute conditional loglikelihood    
    # compute 'plain' nu first
    series nu = (EXOG ~ ENDOG) * beta

    # next, adjust as needed 
    if het
	series sigma = exp(ndxv)
	series nu = nu / sigma
    else
	series sigma = 1
    endif

    if iv
	series nu = (nu + (omega * psi)) / sqrt(cVar)
    endif    

    # form loglikelihood and add marginal if needed
    series ndx = y ? nu : -nu 

    check = min(ndx)
    if check < -35.0
	err = 4
#	printf "Problems with conditional loglik! min(ndx = %g)\n", check
	series ll = NA
    else
	series ll = ln(cnorm(ndx))
    endif

    if !err && iv
	ll += llm
    endif

    #done
    return err
     
end function

function matrix wkron(matrix *X, matrix *Z)

    k = cols(X)
    ret = {}
  
    if ($version < 10908)
	loop i=1..k --quiet	
            matrix ret ~= X[,i] .* Z
	endloop
    else
	ret = hdprod(X,Z)
    endif
	    	 
    return ret
end function


  /******************************
    Score matrix by observation
   ****************************/

function matrix HIP_Score(series y, matrix *EXOG, matrix *ENDOG[null],
			  matrix *ADDIN[null], matrix *HETVAR[null],
			  matrix *beta, matrix *alpha[null], 
			  matrix *Pi[null], matrix *psi[null], 
			  matrix *C[null], matrix *omega, series *sigma)

    scalar het = (rows(alpha) > 0)
    scalar p = cols(ENDOG)
    scalar iv = (p > 0)
    matrix Z = EXOG ~ ENDOG

    # build nu first 
    series ndxm = Z * beta

    if het
	series ndxm = ndxm/sigma
    endif

    if iv
	scalar sigmacond = sqrt(1 - psi'psi)
	series nu = (ndxm + (omega * psi)) / sigmacond 
    else
	scalar sigmacond = 1
	series nu = ndxm
    endif

    # the mills ratio
    series mills = y ? invmills(-nu) : -invmills(nu)

    # now the derivatives wrt nu
    series s_beta = mills/(sigmacond * sigma)
    matrix S_beta = s_beta .* Z

    if het
	series s_alpha = -(ndxm*mills)/sigmacond
	matrix S_alpha =  {s_alpha} .* HETVAR
    else
	matrix S_alpha = { }
    endif

    if iv
	matrix scaledpsi = psi ./ sigmacond
	matrix s_Pi = omega - (mills .* scaledpsi')
	matrix tmp = s_Pi * C'
	matrix X = EXOG ~ ADDIN
	matrix S_Pi = wkron(&tmp,&X)

	matrix S_psi = omega + nu .* scaledpsi'
	matrix S_psi = (mills/sigmacond) .* S_psi

	matrix tmp = zeros(p,p)
	tmp[diag] = 1 ./ diag(C)
	tmp = vech(tmp)'
	sel = vech(mshape(seq(1,p*p),p,p)')
	matrix omiC = omega*inv(C)
	matrix S_dc = wkron(&s_Pi,&omiC)
	matrix S_dc = tmp .- S_dc[,sel]
    else
	matrix S_psi = { }
	matrix S_dc = { }
	matrix S_Pi = { }
    endif

    ret = (S_beta ~ S_alpha ~ S_Pi ~ S_psi ~ S_dc)
    return ret
end function



function matrix Wald_test(bundle *b, scalar ini, scalar df)
    # it is assumed that the model has already been estimated
    # and that the coefficients and the vcv matrix exist

    scalar fin = ini + df - 1

    theta = b["theta"]
    vtheta = b["VCVtheta"]

    vtheta = vtheta[ini:fin, ini:fin]

    matrix WT = qform(theta[ini:fin]', invpd(vtheta))
    return (WT  ~ df ~ pvalue(X, df, WT[1]))
end function

function matrix LM_test(bundle *b)

    theta = b["theta"]
    k1 = b["mk1"]
    k2 = b["mk2"]
    k = b["mk"]
    p = b["mp"]
    q = b["mq"]
    scalar het = (q>0)

    scalar g = k1 + p + q
    
    alpha = theta[1:k1]
    delta = theta[k1+1:k1+p]
    if het
	gamma = theta[k1+p+1:g]
    endif

    Pi = mshape(theta[g+1 : g+p*k],k,p)
    P1 = Pi[1:k1,]
    P2 = Pi[k1+1:k,]
    psi = (alpha + P1*delta) | (P2*delta) | delta

    X = b["mX"]
    W = b["mHETVAR"]
    V = b["uhat"]    
    y = b["depvar"]
    
    if het
	series ex  = exp(-W*gamma)
	series ndx = {ex} .* ((X ~ V) * psi) 
    else
	series ndx = (X ~ V) * psi
    endif

    series mills = y ? invmills(-ndx) : -invmills(ndx)

    if het		  
	matrix S_psi = (mills * ex) .* (X ~ V)		  
	matrix S_gamma = - mills*ndx .* W
    else 
	matrix S_psi = mills .* (X ~ V)		  
	matrix S_gamma = { }
    endif

    matrix G = S_psi ~ S_gamma
		  
    scalar LM = qform(sumc(G),invpd(G'G))
    df = k2 - p
    return (LM ~ df ~ pvalue(X,df,LM))
		  
end function

function matrix LR_test(bundle *b, scalar df)
    
    LR = 2*(b["lnl1"] - b["lnl0"])
    return (LR ~ df ~ pvalue(X,df,LR))

end function

function matrix CI_test(bundle *b)

    scalar het = b["het"]
    theta = b["theta"]
    Z = b["mZ"]
    h = b["mh"]
    q = b["mq"]
    W = b["mHETVAR"]

    beta = theta[1:h]
    series ndxm = Z * beta
    if het 
	alpha = theta[h+1:h+q]
	series ndxv = W * alpha
    else
	series ndxv = 0
    endif

    y = b["depvar"]
    series mills = y ? invmills(-ndxm) : -invmills(ndxm)
    
    matrix e3 = exp(3*ndxv)*mills*(2 + ndxm*ndxm)
    matrix e4 = -exp(4*ndxv)*mills*ndxm*(3 + ndxm*ndxm)

    G = b["SCORE"]
    G = G[,1:h+q] ~ e3 ~ e4
    scalar T = qform(sumc(G),invpd(G'G))
    return (T ~ 2 ~ pvalue(X,2,T))
end function


function void HIP_diagnostics(bundle *b)

    k1 = b["mk1"]
    p = b["mp"]
    h = b["mh"]
    q = b["mq"]
    k = b["mk"]
    het = b["het"]
    iv = b["iv"]

    # Overall test (Wald)
    scalar df = k1 + p - 1
    scalar ini = 2
    WT = Wald_test(&b, ini, df)
    b["WaldAll"] = WT

    if iv
	# Endogeneity test (Wald)
	scalar df  = p
	scalar ini = h + q + k*p + 1
	WT = Wald_test(&b, ini, df)
	b["WaldEnd"] = WT

	if (k>h) # Test of overidentifying restrictions (LM)
	    LM = LM_test(&b)
	    b["LMOverid"] = LM
	endif
    else
	# Chesher&Irish CM test for normality (unsuitable
	# for IV but ok if HET)
	CI = CI_test(&b)
	b["normtest"] = CI
    endif

    # Heteroskedasticity test: Wald if iv, else LR 
    if het
	if iv
	     scalar ini = h + 1
	     T = Wald_test(&b, ini, q)
	 else
	     T = LR_test(&b, q)	     
	 endif
	 b["HETtest"] = T
     endif

end function