import torch, struct, zlib, numpy as np, itertools, os, gc, sys, psutil
from diffusers import StableDiffusionXLPipeline

# CONFIGURATION
_ARTIFACT = "Marius_V18_SDXL.lzr2"
_BASE = "stabilityai/stable-diffusion-xl-base-1.0"

def _stat():
    process = psutil.Process(os.getpid())
    return process.memory_info().rss / (1024 ** 3)

def inject_solvay(path, pipe):
    if not os.path.exists(path): raise FileNotFoundError(f"Missing artifact: {path}")
    print(f"🌀 Initializing solvay V18 Protocol...")
    
    _opts, u = {}, pipe.unet
    _g_v = lambda d: np.array(list(itertools.product([-1, 0, 1], repeat=d)), dtype=np.float32)
    idx = 0
    
    with open(path, "rb") as f:
        if f.read(4) != b"LZR2": raise ValueError("Invalid Signature")
        while True:
            lkb = f.read(4)
            if not lkb: break
            key = f.read(struct.unpack('I', lkb)[0]).decode('utf-8')
            ls = struct.unpack('I', f.read(4))[0]
            sh = [struct.unpack('I', f.read(4))[0] for _ in range(ls)]
            tf = struct.unpack('B', f.read(1))[0]
            _w = None

            if tf == 1:
                dp, C = struct.unpack('I', f.read(4))[0], sh[0]
                _a = np.frombuffer(f.read(C*dp*4), dtype=np.float32).reshape(C, dp)
                _mn, _sc = np.frombuffer(f.read(C*4), dtype=np.float32), np.frombuffer(f.read(C*4), dtype=np.float32)
                lz = struct.unpack('I', f.read(4))[0]
                _ix_flat = np.frombuffer(zlib.decompress(f.read(lz)), dtype=np.uint16)
                n_blocks = _ix_flat.size // C
                _ix = _ix_flat.reshape(C, n_blocks)
                no = struct.unpack('I', f.read(4))[0]
                N_feat = int(np.prod(sh[1:])) if len(sh) > 1 else 1
                
                if dp not in _opts: _opts[dp] = _g_v(dp)
                rc = _opts[dp][_ix].reshape(C, -1) if n_blocks > 0 else np.zeros((C, 0), dtype=np.float32)
                
                fb = np.zeros((C, N_feat), dtype=np.float32)
                vw = min(rc.shape[1], N_feat)
                if vw > 0: fb[:, :vw] = rc[:, :vw]
                fb = (fb + _mn[:, None]) * _sc[:, None]

                if no > 0:
                    md = max(C, n_blocks) * dp
                    fmt, fsz = ('H', 8) if md < 65536 else ('I', 12)
                    dt = np.dtype([('r', np.uint16 if fmt=='H' else np.uint32), ('c', np.uint16 if fmt=='H' else np.uint32), ('v', np.float32)])
                    batch = np.frombuffer(f.read(no * fsz), dtype=dt)
                    m = (batch['r'] < C) & (batch['c'] < N_feat)
                    vb = batch[m]
                    fb[vb['r'], vb['c']] = vb['v']
                
                _w = torch.from_numpy(fb.reshape(sh).astype(np.float16))

            if _w is not None:
                try:
                    t = u
                    pts = key.split('.')
                    for p in pts[:-1]: t = getattr(t, p)
                    getattr(t, pts[-1]).data.copy_(_w.to(pipe.device, dtype=torch.float16))
                except: pass
                del _w
            
            idx += 1
            if idx % 10 == 0:
                sys.stdout.write(f"\r🌊 [STREAM] Synchronizing Module {idx:04d} | SysRAM: {_stat():.1f}GB")
                sys.stdout.flush()
            if idx % 200 == 0: gc.collect()

    print(f"\n✅ solvay Core Active ({idx} modules loaded).")

def get_pipe():
    print("⏳ Loading base architecture...")
    pipe = StableDiffusionXLPipeline.from_pretrained(_BASE, torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
    pipe.enable_model_cpu_offload()
    inject_solvay(_ARTIFACT, pipe)
    return pipe