kMin = 1;   // minimum value of capital allowed
kMax = 40; // maximum value of capital allowed
kStates = kMax-kMin+1;  // number of possible values of state variable (capital)
beta = 0.9;
kVec = seqa(kMin,(kMax-kMin+1)/kStates,kStates);  // a kStates vector with the value of capital at each statements

investmentCost = 3;

// ********* Transition Matrices **********************

// create transition matrix for capital if the firm does not invest, which depends on depreciation.
transMatNoInvest = zeros(kStates,kStates);   // transition matrix from one capital level to another.  It is mostly zeros.
depreciationRate = 0.15|0.35|0.5;  // capital stays the same with probability 3, decreases by 1 with prob 2, and by 2 states with prob 1
for i(3,kStates,1);				// Fill in band of depreciation -- capital can go down by 1 or 2 states.
	transMatNoInvest[i,i-2:i] = depreciationRate';
endfor;

transMatNoInvest[1,1]=1;  			// In state 1, capital cannot go down.
transmatNoInvest[2,1:2] = sumc(depreciationRate[1:2])~depreciationRate[3]; // In state 2, capital cannot go down by 2 states.

// create transition matrix for capital if firm does invest
transMatInvest = zeros(kStates,kStates);
for i(3,kStates-5,1);
	transMatInvest[i,i+3:i+5]=depreciationRate';
endfor;

transMatInvest[1,6]=1;  			// In state 1, capital cannot go down, so you will definitely go to state 6
transmatInvest[2,6:7] = sumc(depreciationRate[1:2])~depreciationRate[3]; // In state 2, capital cannot go down by 2 states, you will definitely be in 6 or 7.
transMatInvest[kStates-4,kStates-1:kStates]=depreciationRate[1]~sumc(depreciationRate[2:3]);  // In state 36, capital can go to 39 or 40
transMatInvest[kStates-3:kStates,kStates]=ones(4,1); // Above 36, and investment automatically takes you to 40.



// ******** Now solve the Bellman equation via a fixed point algorithm *************** 

profitVec = profit(kVec);

// starting values
vNew = profitVec/(1-beta);  
v = vNew - 1;

// Fixed point algorithm
for i(1,1000,1);
	if maxc(abs(vNew-v))<1e-7; 
		Print "convergence at " i " iterations";
		break; 
	endif;
	v = vNew;
	vNew = getVnew(v);
endfor;

// Now graph the probability of investment at each level of capital.
probInvestment = exp(profitVec -investmentCost + beta*transMatInvest*v)./exp(v);

plotxy(kVec,probInvestment);
end;

/*********** Compute new V from old V and transition matrices **************/
proc getVnew(v);
local vNew;
vNew = profitVec+ ln(exp(-investmentCost + beta*transMatInvest*v) + exp(beta*transMatNoInvest*v));
retp(vNew);
endp;


/***************** Compute profit from a vector of capital levels ****************/
proc profit(k);
	retp(ln(k));
endp;