#!/usr/bin/env python3
"""Experimento: ¿aprende un modelo a predecir la salida de SHA-256 según el número de rondas?
Usa la cabecera real del bloque 974.323 y solo cambia el nonce (el caso más favorable al modelo).

Uso: sha_exp.py a 4,5,6,7,8,9,10,11,12,14,16,20,32,64 | sha_exp.py b | sha_exp.py c

Límites conocidos, comprobados en sha_exp_verif.py:
- 'a' recorta solo la compresión del bloque que contiene el nonce, sin la segunda pasada de SHA-256.
- 'a' mide a esta red, no a SHA-256: unos árboles por bit aciertan h6 a 9 rondas (50,645%, 18 desviaciones)
  y un test clásico de diferencias sin IA aún ve estructura a 10 rondas.
- Los nonces se sortean con repetición: el 0,016% de la prueba de 'a' y el 0,043% de la de 'b' están
  también en entrenamiento. No se ha quitado para conservar las cifras; favorecería al modelo y es
  demasiado poco para mover los resultados."""
import sys, json, hashlib, time
import numpy as np
from sklearn.neural_network import MLPClassifier
from sklearn.ensemble import HistGradientBoostingClassifier

K = np.array([
0x428a2f98,0x71374491,0xb5c0fbcf,0xe9b5dba5,0x3956c25b,0x59f111f1,0x923f82a4,0xab1c5ed5,0xd807aa98,0x12835b01,0x243185be,0x550c7dc3,0x72be5d74,0x80deb1fe,0x9bdc06a7,0xc19bf174,
0xe49b69c1,0xefbe4786,0x0fc19dc6,0x240ca1cc,0x2de92c6f,0x4a7484aa,0x5cb0a9dc,0x76f988da,0x983e5152,0xa831c66d,0xb00327c8,0xbf597fc7,0xc6e00bf3,0xd5a79147,0x06ca6351,0x14292967,
0x27b70a85,0x2e1b2138,0x4d2c6dfc,0x53380d13,0x650a7354,0x766a0abb,0x81c2c92e,0x92722c85,0xa2bfe8a1,0xa81a664b,0xc24b8b70,0xc76c51a3,0xd192e819,0xd6990624,0xf40e3585,0x106aa070,
0x19a4c116,0x1e376c08,0x2748774c,0x34b0bcb5,0x391c0cb3,0x4ed8aa4a,0x5b9cca4f,0x682e6ff3,0x748f82ee,0x78a5636f,0x84c87814,0x8cc70208,0x90befffa,0xa4506ceb,0xbef9a3f7,0xc67178f2],dtype=np.uint32)
IV = np.array([0x6a09e667,0xbb67ae85,0x3c6ef372,0xa54ff53a,0x510e527f,0x9b05688c,0x1f83d9ab,0x5be0cd19],dtype=np.uint32)
def rr(x,n): return (x>>np.uint32(n))|(x<<np.uint32(32-n))
def compress(state, W16, rounds=64, feedforward=True):
    """state: 8 arrays uint32 o escalares; W16: lista de 16 arrays uint32."""
    n = max(np.size(w) for w in W16)
    W = [np.broadcast_to(np.asarray(w,dtype=np.uint32),(n,)).copy() for w in W16]
    for t in range(16,64):
        s0 = rr(W[t-15],7)^rr(W[t-15],18)^(W[t-15]>>np.uint32(3))
        s1 = rr(W[t-2],17)^rr(W[t-2],19)^(W[t-2]>>np.uint32(10))
        W.append(W[t-16]+s0+W[t-7]+s1)
    a,b,c,d,e,f,g,h = [np.broadcast_to(np.asarray(s,dtype=np.uint32),(n,)).copy() for s in state]
    for t in range(rounds):
        S1 = rr(e,6)^rr(e,11)^rr(e,25); ch=(e&f)^(~e&g)
        t1 = h+S1+ch+K[t]+W[t]
        S0 = rr(a,2)^rr(a,13)^rr(a,22); mj=(a&b)^(a&c)^(b&c)
        t2 = S0+mj
        h,g,f,e,d,c,b,a = g,f,e,d+t1,c,b,a,t1+t2
    out=[a,b,c,d,e,f,g,h]
    if feedforward: out=[np.asarray(s,dtype=np.uint32)+o for s,o in zip(state,out)]
    return out
HDR = bytes.fromhex("000001208fb06bd1c96469b9dde6d35cb5a1184c78207f66ce2c01000000000000000000d250e41c5debe5619ced21d356eb155ee2cc3c0d989bdda7a0a77e78a4cff6bbe02de96aace50117806bc8f6")
def be(b): return int.from_bytes(b,'big')
W0 = [np.uint32(be(HDR[i:i+4])) for i in range(0,64,4)]
MID = compress(IV, W0)
MID = [np.uint32(m[0]) for m in MID]
TAIL = [np.uint32(be(HDR[64+i:64+i+4])) for i in range(0,12,4)]
def bswap(x): return ((x&0xff)<<np.uint32(24))|((x&0xff00)<<np.uint32(8))|((x>>np.uint32(8))&0xff00)|(x>>np.uint32(24))
def block2(nonce_le, rounds=64):
    W=TAIL+[bswap(nonce_le),np.uint32(0x80000000)]+[np.uint32(0)]*10+[np.uint32(640)]
    return compress(MID,W,rounds)
def dsha(nonce_le):
    o=block2(nonce_le)
    W=o+[np.uint32(0x80000000)]+[np.uint32(0)]*6+[np.uint32(256)]
    return compress(IV,W)
# Comprobación contra hashlib y contra el hash real del bloque
n=np.array([int.from_bytes(HDR[76:80],'little')],dtype=np.uint32)
h=dsha(n); dig=b''.join(int(x[0]).to_bytes(4,'big') for x in h)
assert dig==hashlib.sha256(hashlib.sha256(HDR).digest()).digest(); print('hash real del bloque:',dig[::-1].hex())
def bits(x, nb=32): return ((x[:,None]>>np.arange(nb,dtype=np.uint32))&1).astype(np.float32)
rng=np.random.default_rng(974323)
class Red:
    """Perceptrón de dos capas ocultas (512 y 256) con ReLU, salida logística por bit y Adam. Solo numpy."""
    def __init__(s,ni,no,sem=1):
        r=np.random.default_rng(sem);s.W=[];t=[ni,512,256,no]
        for a,b in zip(t,t[1:]):
            s.W+=[(r.normal(0,np.sqrt(2/a),(a,b))).astype(np.float32),np.zeros(b,dtype=np.float32)]
        s.m=[np.zeros_like(w) for w in s.W];s.v=[np.zeros_like(w) for w in s.W];s.t=0
    def adelante(s,X):
        h1=np.maximum(X@s.W[0]+s.W[1],0);h2=np.maximum(h1@s.W[2]+s.W[3],0);return h1,h2,h2@s.W[4]+s.W[5]
    def paso(s,X,Y,lr=1e-3):
        h1,h2,o=s.adelante(X);p=1/(1+np.exp(-np.clip(o,-30,30)));e=(p-Y)/Y.shape[1]/X.shape[0]
        g4=h2.T@e;g5=e.sum(0);d2=e@s.W[4].T;m2=h2>0;d2=d2*m2
        g2=h1.T@d2;g3=d2.sum(0);d1=(d2@s.W[2].T)*(h1>0);g0=X.T@d1;g1=d1.sum(0)
        s.t+=1
        for i,g in enumerate((g0,g1,g2,g3,g4,g5)):
            s.m[i]=.9*s.m[i]+.1*g;s.v[i]=.999*s.v[i]+.001*g*g
            s.W[i]-=lr*(s.m[i]/(1-.9**s.t))/(np.sqrt(s.v[i]/(1-.999**s.t))+1e-8)
    def fit(s,X,Y,ep=25,lote=512):
        r=np.random.default_rng(7)
        for _ in range(ep):
            o=r.permutation(len(X))
            for k in range(0,len(X)-lote+1,lote):
                b=o[k:k+lote];s.paso(X[b],Y[b])
        return s
def expA(rounds_list,ntr=300000,nte=50000,out='expA.json'):
    res=[]
    for r in rounds_list:
        t0=time.time()
        nonce=rng.integers(0,2**32,ntr+nte,dtype=np.uint32)
        o=block2(nonce,rounds=r)
        Y=np.concatenate([bits(x) for x in o],axis=1)
        X=bits(nonce)*2-1
        frac=Y[:ntr].mean(0);var=(frac>0.02)&(frac<0.98)
        Yv=Y[:,var]
        m=Red(32,int(var.sum())).fit(X[:ntr],Yv[:ntr])
        p=m.adelante(X[ntr:])[2]>0
        acc=(p==(Yv[ntr:]>.5)).mean(0)
        tr=((m.adelante(X[:20000])[2]>0)==(Yv[:20000]>.5)).mean(0)
        umbral=3.9*0.5/np.sqrt(nte)
        rec=dict(rounds=r,bits_variables=int(var.sum()),acc_media=float(acc.mean()),acc_max=float(acc.max()),acc_min=float(acc.min()),acc_train=float(tr.mean()),
                 bits_con_senal=int((acc>.5+umbral).sum()),bits_por_debajo=int((acc<.5-umbral).sum()),umbral=float(umbral),seg=time.time()-t0)
        print(rec,flush=True); res.append(rec); json.dump(res,open(out,'w'),indent=1)
def expB(zero_bits=6,ntr=2000000,nte=1000000):
    nonce=rng.integers(0,2**32,ntr+nte,dtype=np.uint32)
    h=dsha(nonce)
    # El hash se lee al revés: los ceros de cabeza son los últimos bytes del digest
    top=bswap(h[7])
    y=(top>>np.uint32(32-zero_bits))==0
    X=bits(nonce)
    print('positivos',y.mean(),1/2**zero_bits)
    out={}
    g=HistGradientBoostingClassifier(max_iter=300,learning_rate=0.1,max_leaf_nodes=63,random_state=1)
    g.fit(X[:ntr],y[:ntr]); s=g.predict_proba(X[ntr:])[:,1]
    m=MLPClassifier((512,256),batch_size=1024,max_iter=15,random_state=1,tol=1e-6,n_iter_no_change=5)
    m.fit(X[:ntr]*2-1,y[:ntr]); s2=m.predict_proba(X[ntr:]*2-1)[:,1]
    for name,sc in (('gbdt',s),('mlp',s2)):
        idx=np.argsort(-sc)
        yt=y[ntr:]
        for frac in (0.01,0.1,0.3,0.5):
            k=int(frac*nte); cap=yt[idx[:k]].sum()/yt.sum()
            se=np.sqrt(frac*(1-frac)/yt.sum())
            out[f'{name}_{frac}']=dict(captura=float(cap),esperado=frac,error_tipico=float(se))
        print(name,{k:v for k,v in out.items() if k.startswith(name)},flush=True)
    json.dump(out,open('expB.json','w'),indent=1)
def avalanche(n=200000):
    nonce=rng.integers(0,2**32,n,dtype=np.uint32)
    b=rng.integers(0,32,n,dtype=np.uint32)
    h1=dsha(nonce); h2=dsha(nonce^(np.uint32(1)<<b))
    d=np.zeros(n,dtype=np.int32)
    for x,y in zip(h1,h2):
        d+=np.unpackbits(np.frombuffer((x^y).tobytes(),dtype=np.uint8)).reshape(n,32).sum(1).astype(np.int32)
    print('avalancha media',d.mean(),'desv',d.std(),'min',d.min(),'max',d.max())
    json.dump(dict(media=float(d.mean()),desv=float(d.std()),min=int(d.min()),max=int(d.max())),open('avalancha.json','w'))
if __name__=='__main__':
    c=sys.argv[1]
    if c=='a': expA([int(x) for x in sys.argv[2].split(',')])
    if c=='b': expB()
    if c=='c': avalanche()