import hashlib
import json
import random
import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
from .interaction_head import TriVectorInteractionHead
from .verifier_batching import prepare_batch_inputs

def identity(value):
    return hashlib.sha256(json.dumps(value,sort_keys=True,separators=(',',':'),allow_nan=False).encode()).hexdigest()

def seed_all(seed):
    random.seed(seed);np.random.seed(seed);torch.manual_seed(seed)
    if torch.cuda.is_available():torch.cuda.manual_seed_all(seed)
    torch.set_num_threads(8);torch.use_deterministic_algorithms(True)
    torch.backends.cuda.matmul.allow_tf32=False;torch.backends.cudnn.allow_tf32=False

class FullModel(nn.Module):
    def __init__(self,backbone):
        super().__init__();self.backbone=backbone
        self.head=TriVectorInteractionHead(backbone.config.hidden_size).to(next(backbone.parameters()).device)
    def forward(self,input_ids,attention_mask,position_ids,segments):
        hidden=self.backbone(input_ids=input_ids,attention_mask=attention_mask,
            position_ids=position_ids,use_cache=False,return_dict=True).last_hidden_state
        pools=self.head.pool_segments(hidden,segments,attention_mask)
        return self.head(*(F.layer_norm(p.float(),(p.shape[-1],)) for p in pools))

def collate(rows,pad,device):
    ids,mask=prepare_batch_inputs([torch.tensor(r['input_ids'],device=device) for r in rows],pad)
    seg=torch.zeros_like(ids)
    for i,row in enumerate(rows):seg[i,-len(row['segments']):]=torch.tensor(row['segments'],device=device)
    return dict(input_ids=ids,attention_mask=mask,position_ids=(mask.cumsum(-1)-1).clamp_min(0),segments=seg)

def infer(model,rows,pad,device='cuda',progress=None):
    model.eval();scores=[]
    with torch.inference_mode():
        for i,row in enumerate(rows):
            value=model(**collate([row],pad,device))[0].float().cpu()
            if not torch.isfinite(value).all():raise RuntimeError('Nonfinite inference')
            scores.append(value)
            if progress and (i%25==0 or i+1==len(rows)):progress(i+1,len(rows))
    return torch.stack(scores)

def tensor_hash(state):
    h=hashlib.sha256()
    for name,p in sorted(state.items()):
        t=p.detach().cpu().contiguous()
        h.update(name.encode());h.update(str(t.dtype).encode());h.update(str(tuple(t.shape)).encode())
        h.update(t.reshape(-1).view(torch.uint8).numpy().tobytes())
    return h.hexdigest()