(* ::Package:: *)

(* ::Section:: *)
(*Choice of symbols*)


(* ::Text:: *)
(*Sample a 0=BAD with probability p*)
(*Sample a 1=GOOD with probability 1-p = q*)
(*Pick a random 0 and resample.*)
(*All-1 state is the final good state.*)


(* ::Text:: *)
(*Goal: compute P(Z^n | start in 011..11) on the finite chain {1,...,n}*)
(*Method: sum over (k>=1) of P(vertex n gets zero for the first time at step k)*)
(*Let M be the Markov Chain on n-1 vertices such that if site n becomes zero, then it goes to a special state with no outgoing arrows*)
(*Furthermore the all-good state has no self-loop (but has incoming paths). This last thing seems like it is not required but removes the 1 eigenvalue so I guess it is good for inverting. Let si be the initial state, let sf be the final special state.*)
(*Goal is to compute sum over (k>=0) of  si . M^k . sf*)
(*This is equal to si . (1 - M)^(-1) . sf*)


(* ::Chapter:: *)
(*Markov Chain Matrix*)


(* Note that the matrix function requires at least n\[GreaterEqual]3 *)
generateInitialState[n_]:=SparseArray[{(1+FromDigits[Join[{0},ConstantArray[1,n-2]],2])->1},{2^(n-1)+1}]
generateFinalState[n_]:=SparseArray[{(2^(n-1)+1)->1},{2^(n-1)+1}]
(* Matrix as described above! *)
generateLeakingMatrix[nReal_,p_]:=Module[{n,mcM,bits,numBads,j1,j2,j3,newbits,prob},
n = nReal-1; (* See description above. When vertex n is resampled it is a special state. *)
mcM=SparseArray[{},{2^n+1,2^n+1},0];

Do[
    (* Compute neighbours of i *)
    bits=IntegerDigits[i,2,n];
    numBads=Count[bits,0];
    (* Special cases: j==1 and j==n *)
    If[bits[[1]]==0, (* j==1 *)
        newbits=bits;
        Do[
            newbits[[{1,2}]]={b1,b2};
            prob=((1-p)^(b1+b2)*(p)^(2-(b1+b2)))/numBads;
            (* Set the matrix entry *)
            mcM[[i+1,FromDigits[newbits,2]+1]]+=prob;
        ,{b1,0,1},{b2,0,1}];
    ];
    If[bits[[n]]==0, (* j==n *)
        (* This 'n' is actually the n-1 vertex. So we check if it sets the real n vertex to zero. *)
        (* Potentially go into special state *)
        newbits=bits;
        Do[
            prob=((1-p)^(b1+b2+b3)*(p)^(3-(b1+b2+b3)))/numBads;
            If[b3==0,
                (* Special state *)
                mcM[[i+1,-1]]+=prob;
            ,
                (* Normal state *)
                newbits[[{n-1,n}]]={b1,b2};
                (* Set the matrix entry *)
                mcM[[i+1,FromDigits[newbits,2]+1]]+=prob;
            ];
        ,{b1,0,1},{b2,0,1},{b3,0,1}];
    ];
    If[n>=3,
    Do[ (* 2 <= j <= n-1 *)
        (* For all sites, check if it is BAD *)
        If[bits[[j]]==0,
            newbits=bits;
            Do[
                newbits[[{j-1,j,j+1}]]={b1,b2,b3};
                prob=((1-p)^(b1+b2+b3)*(p)^(3-(b1+b2+b3)))/numBads;
                (* Set the matrix entry *)
                mcM[[i+1,FromDigits[newbits,2]+1]]+=prob;
            ,{b1,0,1},{b2,0,1},{b3,0,1}];
        ]; (* endif bits[[j]] *)
    ,{j,2,n-1}];];
,{i,0,2^n-1}];

mcM
]/;(nReal>=3)


(* ::Section:: *)
(*Symbolic evaluation*)


(* ::Subsection:: *)
(*Expected number of resamples using inverse method*)


(* This function definition causes automatic memoization of results *)
getIdMinusM[n_]:=getIdMinusM[n]=Module[{M,IdMinusM},
PrintTemporary["Generating matrix (Id-M) for n = ",n," and general p."];
M=generateLeakingMatrix[n,p];
IdMinusM=IdentityMatrix[Length[M],SparseArray]-M;
IdMinusM
]

(* Replacement rules do not work on sparsearrays. Instead we use this, taken shamelessly from stackoverflow *)
ReplaceElement[s_SparseArray, rule_]:=SparseArray[ArrayRules[s]/.rule,Dimensions[s]]

computeFullExpression[n_,pValue_]:=Module[{nSites,si,sf,IdMinusM,inv1},
(*PrintTemporary["Generating initial state for n = ",n,", p = ",pValue];*)
(*PrintTemporary["Getting (Id-M) matrix."];*)
IdMinusM=getIdMinusM[n];
(*PrintTemporary["Entering p into matrix"];*)
IdMinusM=Simplify[ReplaceElement[IdMinusM,p->pValue]];

si=generateInitialState[n];
sf=generateFinalState[n];
PrintTemporary["Solving (Id-M).x = b for n = ",n," , p = ",pValue];
inv1=LinearSolve[IdMinusM,sf];
(*PrintTemporary["Computing final inner product"];*)
Simplify[si.inv1]
]


generateLeakingMatrix[2,p]


(* ::Subsection:: *)
(*Compute and save coefficients*)


(* For n up to 10 it can be done symbolically in p *)
n2function[p_]=computeFullExpression[2,p];
(*
n3function[p_]=computeFullExpression[3,p];
n4function[p_]=computeFullExpression[4,p];
n5function[p_]=computeFullExpression[5,p];
n6function[p_]=computeFullExpression[6,p];
n7function[p_]=computeFullExpression[7,p];
n8function[p_]=computeFullExpression[8,p];
n9function[p_]=computeFullExpression[9,p];
n10function[p_]=computeFullExpression[10,p];
*)





For[i=1,i<500,i++,
pValue=0/1000+i/500;
Rvalue=computeFullExpression[12,pValue];
PutAppend[{12,pValue,Rvalue},"functionValuesN12.m"]
PutAppend[",","functionValuesN12.m"]
]