"""
INFERENCE.py — run PixelModel from its model.safetensors weights.

This is the standalone entry point for anyone who just wants to load the
safetensors weights and generate an image, without needing model.png or
the rest of this repo's training code.

Usage:
  python INFERENCE.py "a red circle"
  python INFERENCE.py "a red circle" --model model.safetensors --out out.png --scale 8
"""

import argparse
import os
import sys

import numpy as np
import torch
from PIL import Image
from safetensors.torch import load_file

PROMPT_DIM = 32
OUT_SIZE = 32


def prompt_to_embedding(prompt: str) -> torch.Tensor:
    """Deterministic char-level embedding -> PROMPT_DIM vector."""
    vec = torch.zeros(PROMPT_DIM)
    for i, ch in enumerate(prompt.lower()):
        idx = i % PROMPT_DIM
        vec[idx] += ord(ch) / 127.0
    norm = vec.norm()
    if norm > 0:
        vec = vec / norm
    return vec


def forward(weights: dict, prompt: str) -> torch.Tensor:
    emb = prompt_to_embedding(prompt)
    x = torch.tanh(weights["W1"] @ emb)
    x = torch.tanh(weights["W2"] @ x)
    x = torch.sigmoid(weights["W3"] @ x)
    return x.reshape(OUT_SIZE, OUT_SIZE, 3)


def generate(prompt: str, model_path: str, out_path: str, scale: int = 8):
    if not os.path.exists(model_path):
        sys.exit(
            f"Model not found: {model_path}\n"
            f"Run: python convert_to_safetensors.py   to create one from model.png first."
        )

    weights = load_file(model_path)

    with torch.no_grad():
        result = forward(weights, prompt)

    arr = (result.numpy() * 255).clip(0, 255).astype(np.uint8)
    img = Image.fromarray(arr, mode="RGB")

    if scale > 1:
        img = img.resize((OUT_SIZE * scale, OUT_SIZE * scale), Image.NEAREST)

    img.save(out_path)
    print(f"prompt : '{prompt}'")
    print(f"model  : {model_path}  ({os.path.getsize(model_path)} bytes, safetensors)")
    print(f"output : {out_path}  ({OUT_SIZE * scale}x{OUT_SIZE * scale} px)")


if __name__ == "__main__":
    p = argparse.ArgumentParser(description="PixelModel inference (safetensors)")
    p.add_argument("prompt", help="Text prompt")
    p.add_argument("--model", default="model.safetensors", help="Path to safetensors weights")
    p.add_argument("--out", default="out.png", help="Output image path")
    p.add_argument("--scale", type=int, default=8, help="Upscale factor for output (default 8 -> 256x256)")
    args = p.parse_args()
    generate(args.prompt, args.model, args.out, args.scale)