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] = 0.5~0.5; // 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] = 0.5~0.5; // In state 2, capital cannot go down by 2 states, you will definitely be in 6 or 7.
transMatInvest[kStates-4,kStates-1:kStates]=0.15~0.85;  // In state 96, capital can go to 99 or 100
transMatInvest[kStates-3:kStates,kStates]=ones(4,1); // Above 96, and investment automatically takes you to 100.





v = zeros(kStates,1);

profitVec = profit(kVec);
vNew = profitVec/(1-beta);  // a starting value
v = vNew - 1;
for i(1,1000,1);
	if maxc(abs(vNew-v))<1e-6; 
		Print "convergence at " i " iterations";
		break; 
	endif;
	v = vNew;
	
	// Now compute elements of value function in parallel.  Break up value function vector into four parts, and solve for each one simultaneously.
	ThreadStat	vNew[1:10] = getVnewi(1,10,v);
	ThreadStat  vNew[11:20] = getVnewi(11,20,v);
	ThreadStat	vNew[21:30] = getVnewi(21,30,v);
	ThreadStat  vNew[31:40] = getVnewi(31,40,v);
	ThreadJoin;
	
endfor;


probInvestment = exp(profitVec -investmentCost + beta*transMatInvest*v)./exp(v);

probInvestment;


proc getvNewi(kStart,kStop, v);
	local vnewi;
	vnewi = profitVec[kStart:kStop] + ln(exp(-investmentCost + beta*transMatInvest[kStart:kStop,.]*v) + exp(beta*transMatNoInvest[kStart:kStop,.]*v));
retp(vnewi);
endp;

proc profit(k);
	retp(ln(k));
endp;