Print["This is Perturbation AIM Full Version 2.8"] ;
(*                                                                                                                               *)
(* Credits: original algorithm designed by Gary Anderson, 2002-3;                                                                *)
(*   corrected, optimized, and recoded by Eric Swanson, 2004.                                                                    *)
(*   Additional features and optimizations by Eric Swanson, 2008, 2009, 2010.                                                    *)
(* Code is currently maintained by Eric Swanson, see http://www.ericswanson.org or http://www.ericswanson.us for most recent     *)
(*   version and for background, instructions, capabilities, examples, and tips.                                                 *)
(* This version of the code has only been tested with Mathematica version 7.  Perturbation AIM version 2.4 will work in both     *)
(*   Mathematica 5 and Mathematica 6, but lacks some of the functionality and numerical robustness of the current perturbation   *)
(*   AIM version.  While it is not too hard to modify the code to work in Mathematica 4 and earlier, those earlier versions      *)
(*   are *much* slower at solving linear systems of equations, and a few of the function calls are also a bit more awkward.      *)
(* This is the "advanced" or "Full" version of the perturbationAIM code.  It is faster and more numerically robust than older    *)
(*   versions of the code, but the underlying algorithm is also a little bit more opaque as a result of these optimizations.     *)
(*   If you are trying to understand how the code works, there is a "simple" version of perturbation AIM 2.4 available at the    *)
(*   web site above, which does not feature many of these optimizations.  To understand the code, it is recommended that you     *)
(*   go through the "simple" version first, and then go through the code below.                                                  *)
(* No matter which version you are using, you should compute your final set of results using the arbitrary precision             *)
(*   capabilities of the code, as these are *much* more likely to be free of numerical problems.  As a general rule, I have      *)
(*   found that 50 digits of precision (AIMPrecision = 50) and a zero tolerance of 10^-10 (AIMZeroTol = 10^-10) provide          *)
(*   adequate robustness for most models.  However, very highly curved models, such as those that arise in finance, may require  *)
(*   even higher levels of precision to generate reliable answers.                                                               *)
(*                                                                                                                               *)

SetOptions["stdout",PageWidth->120] ; (* most people don't use 80 columns anymore; set the default to 120 *)
<<Combinatorica`                      (* we'll use the Compositions function *)

AIMZeroTol = 10^-10 ;
AIMPrecision = $MachinePrecision ;

(* Useful utilities that we will repeatedly call.                                                                                *)
(* The primary improvement here with respect to the "simple" version is the introduction of a new routine, AIMBackLookingVars,   *)
(*   that searches through the model for any variables that are clearly "backward-looking" (i.e., obviously a function of just   *)
(*   a few state variables, typically one or two, rather than the whole set of state variables of the model).  You could call    *)
(*   these clearly "predetermined" variables if you like, but I've avoided using that term.  This helps speed up the code by     *)
(*   reducing the amount of differentiation that must be done and imposing a number of zero restrictions on the solution.        *)

AIMGenericArgs[eqns_]:= AIMGenericArgs[eqns] = Table[Unique["arg"],{AIMNArgs[eqns]}] ;

AIMLagVars[eqns_]:= AIMLagVars[eqns] = Flatten[Map[Table[#[t+i],{i,Min[Cases[eqns,#[t+j_.]:>j,Infinity]],-1}]&, AIMVarNames[eqns]]] ;

AIMMaxLag[eqns_]:= AIMMaxLag[eqns] = -Min[Cases[eqns,x_Symbol[t+i_.]:>i,Infinity]] ;

AIMMaxLead[eqns_]:= AIMMaxLead[eqns] = Max[Cases[eqns,x_Symbol[t+i_.]:>i,Infinity]] ;

AIMNArgs[eqns_]:= AIMNArgs[eqns] = Length[AIMStateVars[eqns]] ;

AIMNEqns[eqns_]:= AIMNEqns[eqns] = Length[Flatten[{eqns}]] ;

AIMNVars[eqns_]:= AIMNVars[eqns] = Length[AIMVarNames[eqns]] ;

AIMShocks[eqns_]:= AIMShocks[eqns] = Union[Cases[eqns,eps[_][t],Infinity]] ;

AIMSSSubs={Sigma->0, eps[_][_]->0, x_[t+_.]:>Symbol[SymbolName[x]<>"AIMSS"]} ;

AIMSSVars[eqns_]:= AIMSSVars[eqns] = Map[Symbol[SymbolName[#]<>"AIMSS"]&, AIMVarNames[eqns]] ;

AIMStateVars[eqns_]:= AIMStateVars[eqns] = Join[AIMLagVars[eqns], AIMShocks[eqns], {Sigma}] ;

AIMVarNames[eqns_]:= AIMVarNames[eqns] = Union[Cases[eqns,x_Symbol[t+i_.]:>x,Infinity]] ;

(* Here is the AIMBackLookingVars routine.  It first finds all equations that have no leads.  For those equations that have      *)
(*   only one variable dated t, it is then obvious that the solution to this variable is entirely a function of the lagged       *)
(*   variables and shocks in that particular equation.  This particular variable can then be solved in terms of a reduced state  *)
(*   vector, so that its particular bFunc has fewer arguments, hence fewer derivatives and fewer undetermined coefficients that  *)
(*   must be solved.                                                                                                             *)
(* The routine then iterates, so that a variable is also counted as clearly backward-looking if the only other variables in      *)
(*   its equation are lagged variables, date t shocks, and other variables dated t that have already been shown to be clearly    *)
(*   backward-looking.                                                                                                           *)

AIMBackLookingVars[eqns_]:= AIMBackLookingVars[eqns] =
Module[{eqnsnoleads,symbols,nsv,nsvargs,possiblensv,nextpos},
eqnsnoleads = Select[Flatten[{eqns}],Cases[#,_[t+i_/;i>0],Infinity] =={}&] ;
symbols = Map[Join[AIMLagVars[#],Union[Cases[#,_[t],Infinity]]]&, eqnsnoleads] ;
{nsv,nsvargs} = FixedPoint[ Function[nsv,
			symbols = symbols /.Thread[nsv[[1]]->nsv[[2]]] ;
			possiblensv = Map[Cases[#,_Symbol[t]]&, symbols] ;
			nextpos = Flatten[Position[Map[Length,possiblensv], 1]] ;
			{Join[nsv[[1]], Flatten[possiblensv[[nextpos]]]], Join[nsv[[2]], DeleteCases[symbols[[nextpos]],_Symbol[t],Infinity]]}
			], {{},{}}] ;
{nsv /.x_[t]:>x, Map[Union[Flatten[#]]&,nsvargs]}
] ;

AIMBLVPos[eqns_]:= AIMBLVPos[eqns] = Flatten[Map[Position[AIMVarNames[eqns],#]&, AIMBackLookingVars[eqns][[1]]]] ;


(* Calculate the steady state *)

AIMSS[eqns_,parameters_List,opts___Rule]:= AIMSS[eqns,parameters,AIMZeroTol,AIMPrecision] =
Module[{ssguess,sseqns,symbols,ss,ss1,zeropos},
 ssguess = AIMSSGuess /.{opts} /.AIMSSGuess->AIMSSGuess[eqns] /.x_Real:>SetPrecision[x,AIMPrecision] ;
 sseqns = eqns /.x_Real:>SetPrecision[x,AIMPrecision] //.(parameters/.x_Real:>SetPrecision[x,AIMPrecision]) /.Equal->Subtract /.AIMSSSubs ;
 AIMModelDiagnostics[eqns] ;
 If[Length[ssguess]=!=AIMNVars[eqns], Print["Model has ",AIMNVars[eqns]," variables but AIMSSGuess has ",Length[ssguess]," elements"]; $Failed,
  Print["Finding steady state, time is ", Date[],", starting from initial guess: ", Thread[AIMVarNames[eqns]->N[ssguess]]] ;
  symbols = Complement[Union[Cases[sseqns,_Symbol,Infinity]], Join[AIMSSVars[eqns],{E}]] ;
  If[Length[symbols]>0, Print["Warning: found symbols ",symbols," in equations"]] ;
  ss = Check[ss1=FindRoot[sseqns,Transpose[{AIMSSVars[eqns],ssguess}],WorkingPrecision->Max[.96*AIMPrecision,$MachinePrecision],MaxIterations->5000],
	  Print["FindRoot has not converged, most recent candidate is: ", N[ss1]] ;
          Print[" yielding a sum of squared errors of ", N[sseqns.sseqns/.ss1]] ;
	  Print[" Try a different initial SS guess."] ; $Failed,
	{FindRoot::cvmit,FindRoot::lstol}] ;
  If[Head[ss]===List,
   Print["Found steady state, precision is estimated to be ",Precision[ss]," digits, sum of squared errors is ", N[sseqns.sseqns/.ss]] ;
   zeropos = Flatten[Position[Map[Last,ss=Chop[ss,AIMZeroTol]], 0]] ;
   Print["Found steady-state values < AIMZeroTol for ", Length[zeropos]," variables: ", AIMVarNames[eqns][[zeropos]]] ;
   Print[" Imposing hard zeros for these variables, sum of squared errors is ", N[sseqns.sseqns/.ss]] ;
   AIMSSGuess[eqns] = Map[Last,ss] ; (* update default SS guess *)
  ] ;
  ss
]] ;

AIMSS[eqns_,opts__Rule]:= AIMSS[eqns,{},opts] ; (* if AIMSS called without separate parameters, insert {} *)
AIMSS[eqns_]:= AIMSS[eqns,{}] ;
AIMSS[eqns_,parameters_List,zerotol_,precision_]:= AIMSS[eqns,parameters] ;

AIMSSGuess[eqns_]:= Table[0,{AIMNVars[eqns]}] ; (* sets the default steady-state initial guess *)


(* Front end that does error-checking and formats the output *)

AIMSeries[eqns_,parameters_List,deg_Integer]:=
Module[{soln,const,argsubs},
 If[AIMSS[eqns,parameters,AIMZeroTol,AIMPrecision]===$Failed || (soln=AIMSoln[eqns,parameters,deg,AIMZeroTol,AIMPrecision])===$Failed, $Failed,
  Print["Formatting output, time is ", Date[]] ;
  const = AIMSSVars[eqns] /.AIMSS[eqns,parameters,AIMZeroTol,AIMPrecision] ;
  argsubs = Thread[AIMGenericArgs[eqns]->AIMStateVars[eqns]] ;
  Thread[Through[AIMVarNames[eqns][t]] == const + (soln /.argsubs)]
]] ;

AIMSeries[eqns_,deg_Integer]:= AIMSeries[eqns,{},deg] ; (* if AIMSeries called without separate parameters, insert {} *)



(* Front end for Linear AIM *)

AIMSoln[eqns_,parameters_,1,zerotol_,precision_]:= AIMSoln[eqns,parameters,1,zerotol,precision] =
Module[{eqnseq0,allvars,hmat,epsmat,cofb,s0inv,alllagvars},
 eqnseq0 = Flatten[{eqns /.x_Real:>SetPrecision[x,precision]}] //.(parameters /.x_Real:>SetPrecision[x,precision]) /.Equal->Subtract ;
 allvars = Flatten[Map[Through[AIMVarNames[eqns] [t+#]]&, Range[-Max[AIMMaxLag[eqns],1],AIMMaxLead[eqns]]]] ;
 If[!ValueQ[AIMSS[eqns,parameters,zerotol,precision]], AIMSS[eqns,parameters]] ;
 Print["Linearizing model around steady state, time is ", Date[]] ;
 hmat = Chop[Outer[D,eqnseq0,allvars] /.AIMSSSubs /.AIMSS[eqns,parameters,zerotol,precision], zerotol] ;
 epsmat = Chop[Outer[D,eqnseq0,AIMShocks[eqns]] /.AIMSSSubs /.AIMSS[eqns,parameters,zerotol,precision], zerotol] ;
 If[({cofb,s0inv} = AIMLinearSoln[hmat,AIMMaxLead[eqns]]) === {{},{}}, $Failed,
 Print["Precision of linear solution is estimated to be ",Precision[{cofb,s0inv}]," digits, time is ", Date[]] ;
 alllagvars = allvars[[Range[Max[AIMMaxLag[eqns],1]*AIMNVars[eqns]]]] ;
 Chop[cofb .alllagvars - s0inv .epsmat .AIMShocks[eqns] /.Thread[AIMStateVars[eqns]->AIMGenericArgs[eqns]], zerotol]
]] ;


AIMSoln[eqns_,parameters_,0,zerotol_,precision_]:=
Table[0,{AIMNVars[eqns]}] ;  (* this is unnecessary but allows the user to call AIMSeries[eqns,0] or AIMSoln[eqns,0] *)


(* This is the heart of the program.  Derivatives are evaluated at steady state, the expectation is taken, and coefficients      *)
(*   are solved using the method of undetermined coefficients.                                                                   *)
(* As in the simple version, we solve a certainty-equivalent version of the problem first.  In addition, for the clearly         *)
(*   backward-looking variables, we delete terms and coefficients that are clearly zero, and thus solve only those coefficients  *)
(*   that are in general nonzero.                                                                                                *)

AIMSoln[eqns_,parameters_,deg_Integer,zerotol_,precision_]:= AIMSoln[eqns,parameters,deg,zerotol,precision] =
Module[{args,blvpos,bfuncargs,drvindxs,cedrvindxs,stdrvindxs,cecoeffs,stcoeffs,nextceTerms,nextstTerms,
	cesubs,cesystem,cesoln,cematrix,dum,stsubs,stsystem,stmatrix,stsoln1,stsoln2,stsoln},
 args = AIMGenericArgs[eqns] ;
 blvpos = AIMBLVPos[eqns] ;
 bfuncargs = Table[args,{AIMNVars[eqns]}] ;
 bfuncargs[[blvpos]] = AIMBackLookingVars[eqns][[2]] /.Thread[AIMStateVars[eqns]->AIMGenericArgs[eqns]] ;
 drvindxs = Compositions[deg,AIMNArgs[eqns]] ;
  cedrvindxs = Select[drvindxs, #[[-1]]==0 &] ;
  stdrvindxs = Select[drvindxs, #[[-1]] >1 &] ;
 cecoeffs = Table[Unique[coeff,{Temporary}],{AIMNVars[eqns]},{Length[cedrvindxs]}] ;
  stcoeffs = Table[Unique[coeff,{Temporary}],{AIMNVars[eqns]},{Length[stdrvindxs]}] ;
 nextceTerms = cecoeffs .Map[Apply[Times,Power[args,#]]&, cedrvindxs] ;
  nextstTerms = stcoeffs .Map[Apply[Times,Power[args,#]]&, stdrvindxs] ;
 nextceTerms[[blvpos]] = Map[ nextceTerms[[#]] /.Thread[Complement[args,bfuncargs[[#]]]->0]&, blvpos] ;
  nextstTerms[[blvpos]] = 0 ;
 cecoeffs[[blvpos]] = Map[Complement[Cases[nextceTerms[[#]],_Symbol,Infinity], bfuncargs[[#]]]&, blvpos] ;
  stcoeffs[[blvpos]] = {} ;
 cecoeffs = Flatten[cecoeffs] ;
  stcoeffs = Flatten[stcoeffs] ;
 cesubs = Thread[Map[bFunc,AIMVarNames[eqns]] -> MapThread[Apply[Function,{#1,#2}]&,
							{bfuncargs, AIMSoln[eqns,parameters,deg-1,zerotol,precision] + nextceTerms}]] ;
 Print["Calculating CE Derivatives, time is ", Date[]] ;
 cesoln = If[AIMNArgs[eqns]===1, PrependTo[args,dum]; nextceTerms,
   cesystem = Chop[Flatten[CoefficientArrays[AIMCEDerivatives[eqns,deg][0]
	//.(parameters /.x_Real:>SetPrecision[x,precision]) /.AIMSS[eqns,parameters,zerotol,precision] /.cesubs, Drop[args,-2]]], zerotol] ;
   Print["Undetermined CE coefficients to solve: ",Length[cecoeffs]] ;
   Print["Calculating CE solution, time is ", Date[]] ;
   cematrix = CoefficientArrays[DeleteCases[cesystem,0], cecoeffs] ;
   nextceTerms /.Thread[cecoeffs -> Chop[LinearSolve[cematrix[[2]],-cematrix[[1]]],zerotol]]] ;
 Print["Precision of solution is estimated to be ",Precision[cesoln]," digits, time is ", Date[]] ;
 If[stcoeffs==={}, stsoln=0,
  stsubs = Thread[Map[bFunc,AIMVarNames[eqns]] -> MapThread[Apply[Function,{#1,#2}]&,
							{bfuncargs, AIMSoln[eqns,parameters,deg-1,zerotol,precision] + cesoln + nextstTerms}]] ;
  Print["Calculating Stoch Derivatives, time is ", Date[]] ;
  stsystem = Chop[Expand[Flatten[CoefficientArrays[Chop[
	AIMStDerivatives[eqns,deg][0] //.(parameters/.x_Real:>SetPrecision[x,precision]) /.AIMSS[eqns,parameters,zerotol,precision]
		/.stsubs, zerotol],Drop[args,-1]]]], zerotol] /.eps[x_][_]^n_:>mom[x,n] /.eps[_][_]->0 ;
  Print["Undetermined Stoch coefficients to solve: ",Length[stcoeffs]] ;
  Print["Calculating Stoch solution, time is ", Date[]] ;
  stmatrix = CoefficientArrays[DeleteCases[stsystem,0], stcoeffs] ;
(* Note: normally, one would never use -Inverse[stmatrix[[2]]].stmatrix[[1]] to solve a linear system of equations, but
    there appears to be a subtle bug in Mathematica's LinearSolve (also inherited by NSolve and Solve) that causes
    occasional problems when there are symbols in the system (in this case, mom[_,n]).  The following is more robust: *)
  stsoln1 = nextstTerms /.(Flatten[NSolve[stsystem, stcoeffs]] /.x_/;Abs[N[x]]<zerotol ->0) ;
  stsoln2 = nextstTerms /.Thread[stcoeffs -> Chop[Expand[-Inverse[stmatrix[[2]]].stmatrix[[1]]], zerotol]] ;
  stsoln = If[Precision[stsoln1]>Precision[stsoln2],stsoln1,stsoln2] ;
 ] ;
 Apply[Remove,Join[cecoeffs,stcoeffs]] ;
 Print["Precision of solution is estimated to be ",Precision[stsoln]," digits, time is ", Date[]] ;
 AIMSoln[eqns,parameters,deg-1,zerotol,precision] + cesoln + stsoln
] ;

(* That's essentially it.  The following routine calculates derivatives of the equations composed with the (unknown) solution    *)
(*   functions b.	                                                                                                         *)
(* Like the simple version, we use univariate differentiation to compute the multivariate derivatives of the function Fob.       *)
(*   Unlike the simple version, we compute the certainty-equivalent derivatives (those not involving Sigma) first, and the       *)
(*   stochastic derivatives (those involving Sigma) second.  This allows us to avoid ever computing the first derivatives of     *)
(*   Fob with respect to Sigma, which we know from theory must be zero.                                                          *)

AIMCEDerivatives[eqns_,2]:= AIMCEDerivatives[eqns,2] =
Derivative[2][Function[tAIM, Evaluate[AIMSubBFuncsIntoEqns[eqns] /.
 Thread[AIMStateVars[eqns]->tAIM*Join[Drop[AIMGenericArgs[eqns],-2],{1,0}]]]]] ;

AIMCEDerivatives[eqns_,deg_Integer]:= AIMCEDerivatives[eqns,deg] = AIMCEDerivatives[eqns,deg-1]' ;


AIMStDerivatives[eqns_,2]:= AIMStDerivatives[eqns,2] =
Function[tAIM, Evaluate[Derivative[0,2] [Function[{tAIM1,tAIM2},
  Evaluate[AIMSubBFuncsIntoEqns[eqns] /.Thread[AIMStateVars[eqns]->Append[tAIM1*Drop[AIMGenericArgs[eqns],-1],tAIM2]]]]] [tAIM,tAIM]]] ;

AIMStDerivatives[eqns_,deg_Integer]:= AIMStDerivatives[eqns,deg] = AIMStDerivatives[eqns,deg-1]' ;


AIMSubBFuncsIntoEqns[eqns_]:= AIMSubBFuncsIntoEqns[eqns] =
With[{deveqns = eqns /.x_Symbol[t+i_.]:>Symbol[SymbolName[x]<>"AIMSS"] +x[t+i] /.Equal->Subtract},
 AIMSubOutLeadVars[eqns,deveqns,AIMMaxLead[eqns]]
] ;


AIMBFuncs[eqns_]:= AIMBFuncs[eqns] =
Module[{bfuncs},
 bfuncs = Map[Apply[bFunc[#],AIMStateVars[eqns]]&, AIMVarNames[eqns]] ;
 bfuncs[[AIMBLVPos[eqns]]] = MapThread[Apply[bFunc[#1],#2]&, AIMBackLookingVars[eqns]] ;
 bfuncs
] ;


AIMSubOutLeadVars[origeqns_,eqnssofar_,0]:= 
eqnssofar /.Thread[Through[AIMVarNames[origeqns][t]]->AIMBFuncs[origeqns]] /.eps[x_][t+i_]:>Sigma*eps[x][t+i] ;

AIMSubOutLeadVars[origeqns_,eqnssofar_,lead_Integer]:=
With[{reducedeqns=eqnssofar /.Thread[Through[AIMVarNames[origeqns][t+lead]] -> (AIMBFuncs[origeqns]/.t:>t+lead)]},
 AIMSubOutLeadVars[origeqns, reducedeqns, lead-1]
] ;


(* Print out model diagnostics (obviously) *)

AIMModelDiagnostics[eqns_]:= (
Print["Model Diagnostics:"] ;
Print["Number of equations: ", AIMNEqns[eqns]] ;
Print["Number of variables: ", AIMNVars[eqns]] ;
Print["Number of shocks:    ", Length[AIMShocks[eqns]]] ;
Print["Maximum lag:         ", AIMMaxLag[eqns]] ;
Print["Maximum lead:        ", AIMMaxLead[eqns]] ;
Print["Lagged variables: ", AIMLagVars[eqns]] ;
Print["Shocks: ", AIMShocks[eqns]] ;
Print[" together with Sigma, these yield ", AIMNArgs[eqns], " state variables"] ;
Print["Variables ", AIMBackLookingVars[eqns][[1]], " are clearly backward-looking"] ;
Print[" using the following reduced state variable sets for these variables:"] ;
Print[AIMBackLookingVars[eqns][[2]]] ;
Print["List of all variables: ", AIMVarNames[eqns]] ;
Print["Using internal working precison of ", AIMPrecision, " digits (AIMPrecision)"] ;
Print["Treating numbers < ", N[AIMZeroTol], " as zero (AIMZeroTol)"] ;
) ;


(* Everything that follows is Linear AIM.  See Anderson and Moore (1985) for a description of the algorithm.                      *)
(* In contrast to the simple version of the code, here we perform "exact" ShiftRights first (i.e., shift equations forward that   *)
(*   have a complete row of zeros in the lead matrix, which reduces the number of calls to the QR decomposition and thus reduces  *)
(*   potential numerical problems.)                                                                                               *)
(* Also, we compute the stability conditions of the model using the Schur decomposition if the eigenvectors corresponding to the  *)
(*   large eigenvalues of the model are linearly dependent.  The Schur decomposition is guaranteed to exist for all models, while *)
(*   linearly independent eigenvectors are not.  The disadvantage of using Schur vectors is that computing them is slow: first,   *)
(*   the computation is inherently slower, and second, the fact that the code must swap Schur blocks using a routine that I (ets) *)
(*   wrote, rather than internal optimized machine code, slows it down quite a bit.                                               *)

AIMLinearSoln[hmat_,nleads_Integer]:= AIMLinearSoln[hmat,nleads] =
Module[{hrows,hcols,shiftedh,qmat,stabconds,bmat,smat},
 {hrows,hcols} = Dimensions[hmat] ;
 {shiftedh,qmat} = FixedPoint[AIMExactShiftRight, {hmat,{}}, nleads] ;
 {shiftedh,qmat} = FixedPoint[AIMShiftRight, {shiftedh,qmat}, nleads*hrows] ;
 Print["Lead matrix is full rank; computing stability conditions, time is ", Date[]] ;
 stabconds = AIMLinearStabConds[Transpose[AIMAR1Form[shiftedh]]] ;
 Print["Model has ",Length[stabconds]," unstable roots, time is ", Date[]] ;
 qmat = Join[qmat,stabconds] ;
 If[Length[qmat]<hrows*nleads, Print["Multiple Linear Solutions"]; {{},{}},
 If[Length[qmat]>hrows*nleads, Print["No Linear Solutions"]; {{},{}},
 bmat = LinearSolve[AIMLastBlock[qmat],-AIMFirstBlocks[qmat]] ;
 smat = hmat .Join[IdentityMatrix[hcols-hrows*nleads], MapThread[Join,{Table[0,{hrows*nleads},{hrows}],bmat}]] ;
 Chop[{Take[bmat,hrows], Inverse[AIMLastBlock[smat]]}, AIMZeroTol]
]]] ;


AIMLinearSoln[hmat_,0]:= With[{cofb=LinearSolve[AIMLastBlock[hmat],-AIMFirstBlocks[hmat]]}, Chop[{cofb,Inverse[AIMLastBlock[hmat]]}, AIMZeroTol]] ;


AIMExactShiftRight[{hmatold_,qmatsofar_}]:=
Module[{hrows=Length[hmatold],zerorows,firstblocks},
If[ Length[zerorows = Position[Map[Max,Abs[AIMLastBlock[hmatold]]],0]] ==0, {hmatold,qmatsofar},
 firstblocks = AIMFirstBlocks[hmatold] ;
 Print["Shifting ",Length[zerorows]," equations ",zerorows," forward"] ;
 {ReplacePart[hmatold, ArrayFlatten[{{Table[0,{hrows},{hrows}],firstblocks}}], zerorows,zerorows], Join[qmatsofar,firstblocks[[Flatten[zerorows]]]]}
]] ;


AIMShiftRight[{hmatold_,qmatsofar_}]:=
Module[{hrows=Length[hmatold],q,r,zerorows,hmatnew,firstblocks},
 {q,r} = Take[QRDecomposition[AIMLastBlock[hmatold], Pivoting->True], 2] ;
If[ Length[zerorows = Position[Abs[Tr[r,List]], x_/;x<AIMZeroTol]] ==0, {hmatold,qmatsofar},
 Print["Shifting ",Length[zerorows]," linear combinations of equations forward"] ;
 firstblocks = AIMFirstBlocks[hmatnew = q .hmatold] ;
 Chop[{ReplacePart[hmatnew,ArrayFlatten[{{Table[0,{hrows},{hrows}],firstblocks}}], zerorows,zerorows],
										Join[qmatsofar,firstblocks[[Flatten[zerorows]]]]}, AIMZeroTol]
]] ;


AIMAR1Form[shiftedh_]:=
With[{hrows=Length[shiftedh],hcols=Length[First[shiftedh]]},
 Chop[Join[MapThread[Join,{Table[0,{hcols-2*hrows},{hrows}],IdentityMatrix[hcols-2*hrows]}],
	LinearSolve[AIMLastBlock[shiftedh],-AIMFirstBlocks[shiftedh]]], AIMZeroTol]
] ;


AIMLastBlock[matrix_]:= With[{dims=Dimensions[matrix]}, Take[matrix,{1,dims[[1]]},{dims[[2]]-dims[[1]]+1,dims[[2]]}]] ;

AIMFirstBlocks[matrix_]:= With[{dims=Dimensions[matrix]}, Take[matrix,{1,dims[[1]]},{1,dims[[2]]-dims[[1]]}]] ;


(* Mathematica's implementation of the Schur decomposition does not allow the user to specify any particular order for the       *)
(*   eigenvalues to appear along the diagonal.  To compute the stability conditions of the model, we require all eigenvalues     *)
(*   greater than one in absolute value to appear in the top part of the diagonal.  The following two routines swap adjacent     *)
(*   Schur blocks until this ordering is achieved.                                                                               *)
(* The routine is based on a 2002 paper by Jan Brandts entitled "Matlab Code for Sorted Real Schur Forms," available at          *)
(*   http://staff.science.uva.nl/~brandts.  One advantage of this particular algorithm is that it allows for swapping 1x1 and    *)
(*   2x2 Schur blocks, so that we can avoid complex numbers entirely.  This keeps annoying (albeit tiny) imaginary numbers       *)
(*   from cropping up in our stability conditions and hence our solution.                                                        *)
(* One can also use a Generalized Schur Decomposition instead of the plain vanilla Schur Decomposition, but in Linear AIM the    *)
(*   gains from this are small because of the efficient way that AIM handles leads (using ShiftRights to render the lead matrix  *)
(*   invertible, and avoiding stacking equations into the more numerically unstable AR1 form until the very last minute).        *)

AIMLinearStabConds[matrix_]:=
With[{n=Length[matrix]}, Module[{stabconds,q,t,subdiagelts,twoblockpos,blockpos,eigvals,bigrtpos,swaplist},
 stabconds = Chop[Eigenvectors[matrix, Length[Cases[Abs[Eigenvalues[matrix]],i_/;i>1+AIMZeroTol]]], AIMZeroTol] ;
 If[Min[Map[Max,Abs[stabconds]]]>0, stabconds,
  Print["Found degenerate Blanchard-Kahn eigenvector, using Schur decomposition instead, time is ", Date[]] ;
  {q,t} = Chop[SchurDecomposition[matrix], AIMZeroTol] ;
  subdiagelts = Map[t[[#+1,#]]&, Range[n-1]] ;
  twoblockpos = Flatten[Position[subdiagelts, i_/;i!=0]] ;
  blockpos = Complement[Range[n+1], twoblockpos+1] ;
  eigvals = Map[t[[#,#]]&, Drop[blockpos,-1]] ;
  eigvals[[Flatten[Map[Position[blockpos,#]&, twoblockpos]]]] += Map[Sqrt[t[[#,#+1]] *t[[#+1,#]]]&, twoblockpos] ;
  bigrtpos = Flatten[Position[Abs[eigvals], i_/;i>1+AIMZeroTol]] ;
  swaplist = Flatten[Map[Range[bigrtpos[[#]]-1,#,-1]&, Range[Length[bigrtpos]]]];
  {q,t,blockpos} = Fold[AIMSwapSchurBlocks, {Transpose[q],t,blockpos}, swaplist];
  stabconds = Take[q, blockpos[[Length[bigrtpos]+1]] -1]
]]] ;


AIMSwapSchurBlocks[{q_,t_,blockpos_},swappos_]:=
Module[{index1,index2,n1,n2,a11,a12,a22,xmat,ans,xq,index12,qswpd,tprime,tswpd,blockposswpd},
 index1 = Range[blockpos[[swappos]], blockpos[[swappos+1]]-1] ;
 index2 = Range[blockpos[[swappos+1]], blockpos[[swappos+2]]-1] ;
 {n1,n2} = {Length[index1],Length[index2]} ;
 {a11,a12,a22} = {Take[t,index1,index1], Take[t,index1,index2], Take[t,index2,index2]} ;
 xmat = Table[Unique[x,{Temporary}],{n1},{n2}] ;
 ans = Flatten[NSolve[Flatten[a11.xmat-xmat.a22]-Flatten[a12], Flatten[xmat]]];
 xq = First[QRDecomposition[ ArrayFlatten[{{-xmat /.ans,IdentityMatrix[n1]},{IdentityMatrix[n2],Table[0,{n2},{n1}]}}]]] ;
 index12 = Join[index1,index2] ;
 qswpd = q ; (* premultiply q by xq *)
  qswpd[[index12]] = xq .qswpd[[index12]] ;
 tprime = Transpose[t] ; (* transform t to xq .t .xq' *)
  tprime[[index12]] = xq .tprime[[index12]] ;
  tswpd = Transpose[tprime] ;
  tswpd[[index12]] = xq .tswpd[[index12]] ;
 blockposswpd = ReplacePart[blockpos, blockpos[[swappos+1]] +n2-n1, swappos+1] ;
 Apply[Remove,Flatten[xmat]] ;
 Chop[{qswpd,tswpd,blockposswpd}, AIMZeroTol]
] ;