# Exact zero-state in-control ARL of Shewhart X chart with Western Electric run rules
# (distribution propagation over a finite Markov chain; pure Python).
import math, sys
Phi = lambda x: 0.5*(1+math.erf(x/math.sqrt(2)))
# zones: +3,+2,+1,-1,-2,-3 meaning (2,3],(1,2],(0,1],(-1,0],(-2,-1],(-3,-2]; beyond 3 = signal (rule 1)
def probs(shift=0.0):
    c=[-3,-2,-1,0,1,2,3]
    p=[Phi(c[i+1]-shift)-Phi(c[i]-shift) for i in range(6)]  # -3,-2,-1,+1,+2,+3
    z=[-3,-2,-1,1,2,3]
    out=1-sum(p)
    return list(zip(z,p)), out
def arl(rules, shift=0.0, run4=8, tol=1e-8):
    zp,out = probs(shift)
    if 1 not in rules: raise SystemExit("rule1 always on here")
    dist={((),0):1.0}; total=0.0; n=0
    while True:
        total+=sum(dist.values()); n+=1
        nd={}
        for (hist,run),m in dist.items():
            for z,p in zp:
                h=(hist+(z,))[-5:]
                side=1 if z>0 else -1
                r = side if (run==0 or (run>0)!=(side>0)) else run+side
                sig=False
                if 2 in rules:
                    l3=h[-3:]
                    if sum(1 for x in l3 if x>=3)>=2 or sum(1 for x in l3 if x<=-3)>=2: sig=True
                if 3 in rules:
                    l5=h[-5:]
                    if sum(1 for x in l5 if x>=2)>=4 or sum(1 for x in l5 if x<=-2)>=4: sig=True
                if 4 in rules and abs(r)>=run4: sig=True
                if sig: continue
                key=(h[-4:], max(-(run4-1),min(run4-1,r)))
                nd[key]=nd.get(key,0)+m*p
        dist=nd
        if sum(dist.values())<tol: break
    return total
if __name__=="__main__":
    for rs in [(1,),(1,2),(1,3),(1,4),(1,2,3),(1,2,3,4)]:
        print(rs, round(arl(set(rs)),2))
    print("rule1+2+3+4, 9-run variant:", round(arl({1,2,3,4},run4=9),2))
    for d in (0.5,1.0,2.0):
        print("shift",d,"rule1",round(arl({1},d),2),"WE all",round(arl({1,2,3,4},d),2))