"""B93Q1 self-contained, JSON-safe mixed-precision model storage.

93-level and 8649-level symmetric weight quantisation. Buffers and quantisation
scales are lossless IEEE float32 carried as five Base93 digits. The alphabet is
the user's exact Scratch alphabet. This is weight storage, not an integer-only
inference engine.
"""
from pathlib import Path
import json
import numpy as np

ALPHABET=''.join(chr(i) for i in range(32,127) if i not in (34,92))
LUT=np.frombuffer(ALPHABET.encode('ascii'),dtype=np.uint8)
INV=np.full(128,-1,dtype=np.int16)
INV[LUT]=np.arange(93,dtype=np.int16)
SEP='\\'
MAGIC='B93Q1'

def encode_uint(values,digits):
    values=np.asarray(values,dtype=np.uint64).reshape(-1).copy()
    if np.any(values>=93**digits):raise ValueError('Integer does not fit')
    out=np.empty((len(values),digits),dtype=np.uint8)
    for i in range(digits-1,-1,-1):
        out[:,i]=LUT[values%93];values//=93
    return out.tobytes().decode('ascii')

def decode_uint(text,digits):
    if digits<1 or len(text)%digits:raise ValueError('Invalid digit count')
    chars=np.frombuffer(text.encode('ascii'),dtype=np.uint8)
    if np.any(chars>=128):raise ValueError('Non-ASCII digit')
    codes=INV[chars]
    if np.any(codes<0):raise ValueError('Invalid Base93 digit')
    codes=codes.reshape(-1,digits).astype(np.uint64)
    values=np.zeros(len(codes),dtype=np.uint64)
    for i in range(digits):values=values*93+codes[:,i]
    return values

def pack_f32(values):
    a=np.asarray(values,dtype='<f4').reshape(-1)
    if not np.isfinite(a).all():raise ValueError('Non-finite float')
    return encode_uint(a.view('<u4'),5)

def unpack_f32(text):
    bits=decode_uint(text,5)
    if np.any(bits>0xffffffff):raise ValueError('Invalid float32 bits')
    a=bits.astype('<u4').view('<f4')
    if not np.isfinite(a).all():raise ValueError('Non-finite float')
    return a

def row_shape(shape):
    return (shape[0],int(np.prod(shape[1:]))) if len(shape)>=2 else (1,int(np.prod(shape)))

def quantize(values,digits,group_size):
    a=np.asarray(values,dtype=np.float32)
    if digits not in (1,2):raise ValueError('Weight precision must be 1 or 2')
    rows,width=row_shape(a.shape);groups=(width+group_size-1)//group_size
    padded=np.zeros((rows,groups*group_size),dtype=np.float32)
    padded[:,:width]=a.reshape(rows,width)
    blocks=padded.reshape(rows,groups,group_size)
    qmax=(93**digits-1)//2
    scales=np.max(np.abs(blocks),axis=-1)/qmax
    scales=np.where(scales>0,scales,1).astype(np.float32)
    q=np.rint(blocks/scales[...,None]).clip(-qmax,qmax).astype(np.int32)
    codes=(q.reshape(rows,-1)[:,:width]+qmax).astype(np.uint16).reshape(-1)
    restored=(q*scales[...,None]).reshape(rows,-1)[:,:width].reshape(a.shape).astype(np.float32)
    return restored,scales.reshape(-1),codes

def dequantize(codes,scales,shape,digits,group_size):
    rows,width=row_shape(shape);groups=(width+group_size-1)//group_size
    if len(codes)!=rows*width or len(scales)!=rows*groups:raise ValueError('Payload shape mismatch')
    scales=np.asarray(scales).reshape(rows,groups)
    per_value=scales[:,np.arange(width)//group_size]
    qmax=(93**digits-1)//2
    signed=np.asarray(codes,dtype=np.int32).reshape(rows,width)-qmax
    return (signed*per_value).astype(np.float32).reshape(shape)

def encode_state(state,parameter_names,config,provenance,plan):
    fields=[MAGIC,ALPHABET,json.dumps(config,separators=(',',':'),ensure_ascii=True),
            json.dumps(provenance,separators=(',',':'),ensure_ascii=True),str(len(state))]
    records=[];position=sum(len(f)+1 for f in fields)
    for name,t in state.items():
        a=t.detach().cpu().numpy() if hasattr(t,'detach') else np.asarray(t)
        shape=list(a.shape);start=position
        if name in parameter_names:
            digits,group_size=plan[name]
            restored,scales,codes=quantize(a,digits,group_size)
            kind='q';scale_text=pack_f32(scales);payload=encode_uint(codes,digits)
        elif a.dtype.kind=='f':
            kind='f';digits=5;group_size=0;scale_text='';payload=pack_f32(a)
        elif a.dtype.kind in 'iu':
            kind='i';digits=0;group_size=0;scale_text='';payload=','.join(map(str,a.reshape(-1)))
        else:raise ValueError('Unsupported dtype '+str(a.dtype))
        head=[name,','.join(map(str,shape)),kind,str(digits),str(group_size)]
        scale_start=position+sum(len(f)+1 for f in head)
        payload_start=scale_start+len(scale_text)+1
        record=head+[scale_text,payload]
        if any(SEP in x for x in record):raise ValueError('Separator inside field')
        fields+=record;position+=sum(len(f)+1 for f in record)
        records.append({'name':name,'shape':shape,'kind':kind,'chars_per_value':digits,
            'group_size':group_size,'num_values':int(a.size),'scale_offset':scale_start,
            'scale_count':len(scale_text)//5,'payload_offset':payload_start,
            'payload_length':len(payload),'record_offset':start})
    text=SEP.join(fields)+SEP
    manifest={'format':MAGIC,'alphabet':ALPHABET,'offset_base':0,'config':config,
        'provenance':provenance,'text_bytes':len(text.encode('ascii')),
        'json_string_bytes':len(json.dumps(text,ensure_ascii=True).encode('ascii')),
        'tensors':records}
    return text,manifest

def decode_state(text):
    fields=text.split(SEP)
    if len(fields)<6 or fields[0]!=MAGIC or fields[1]!=ALPHABET or fields[-1]!='':
        raise ValueError('Invalid B93Q1 header or terminator')
    config=json.loads(fields[2]);provenance=json.loads(fields[3]);n=int(fields[4])
    if len(fields)!=6+n*7:raise ValueError('Incorrect field count')
    result={}
    for i in range(n):
        name,shape,kind,digits,group_size,scale_text,payload=fields[5+i*7:12+i*7]
        if name in result:raise ValueError('Duplicate tensor')
        shape=tuple(int(x) for x in shape.split(',')) if shape else ()
        digits=int(digits);group_size=int(group_size)
        if kind=='q':
            if digits not in (1,2) or group_size<1:raise ValueError('Invalid quantisation')
            a=dequantize(decode_uint(payload,digits),unpack_f32(scale_text),shape,digits,group_size)
        elif kind=='f':a=unpack_f32(payload).reshape(shape)
        elif kind=='i':a=np.array([int(x) for x in payload.split(',')],dtype=np.int64).reshape(shape)
        else:raise ValueError('Unknown tensor encoding')
        result[name]=a
    return result,config,provenance

def load_model(path,device='cpu'):
    import torch
    from semantic_model import SemanticColorizer
    state,config,_=decode_state(Path(path).read_text(encoding='ascii'))
    model=SemanticColorizer(**{k:config[k] for k in ('head','width','queries')})
    model.load_state_dict({k:torch.from_numpy(v.copy()) for k,v in state.items()},strict=True)
    return model.to(device).eval()

if __name__=='__main__':
    import argparse
    p=argparse.ArgumentParser();p.add_argument('text');p.add_argument('--output',default='decoded_model')
    args=p.parse_args()
    from safetensors.torch import save_file
    import torch
    arrays,config,provenance=decode_state(Path(args.text).read_text(encoding='ascii'))
    dest=Path(args.output);dest.mkdir(exist_ok=True,parents=True)
    save_file({k:torch.from_numpy(v.copy()) for k,v in arrays.items()},str(dest/'model.safetensors'))
    (dest/'config.json').write_text(json.dumps(config,indent=2))
    print('Decoded',len(arrays),'tensors to',dest)