"""Official PNG codec compatible weight reconstruction adapters."""

from __future__ import annotations

import hashlib
import json
from collections.abc import Mapping, Sequence
from pathlib import Path

import numpy as np
from PIL import Image


def decode_weight_image(image_path: str | Path, manifest_path: str | Path) -> tuple[np.ndarray, list[str]]:
    """公式manifest順でRGBのR/Gをfp16 bit patternとしてfloat32へ復元する。"""
    metadata = json.loads(Path(manifest_path).read_text())
    pixels = np.asarray(Image.open(image_path).convert("RGB"), dtype=np.uint8).reshape(-1, 3)
    total = int(metadata["total_parameters"])
    if len(pixels) < total:
        raise ValueError(f"weight image has {len(pixels)} pixels, needs {total}")
    bits = (pixels[:total, 0].astype(np.uint16) << 8) | pixels[:total, 1].astype(np.uint16)
    return bits.view(np.float16).astype(np.float32), [str(item["name"]) for item in metadata["params"]]


def flatten_manifest(arrays: Mapping[str, np.ndarray]) -> tuple[np.ndarray, list[tuple[str, int, int]]]:
    """mappingの挿入順を保持し、float16 quantization後のfloat32 flatを返す。"""
    parts: list[np.ndarray] = []
    manifest: list[tuple[str, int, int]] = []
    offset = 0
    for name, value in arrays.items():
        part = np.asarray(value).reshape(-1).astype(np.float16).astype(np.float32)
        parts.append(part)
        manifest.append((name, offset, offset + len(part)))
        offset += len(part)
    return np.concatenate(parts) if parts else np.empty(0, dtype=np.float32), manifest


def layer_offsets(names: Sequence[str], sizes: Sequence[int]) -> dict[str, tuple[int, int]]:
    """layer名とnumelからmetric用offsetを構築する。"""
    offsets: dict[str, tuple[int, int]] = {}
    offset = 0
    for name, size in zip(names, sizes):
        offsets[name] = (offset, offset + int(size))
        offset += int(size)
    return offsets


def sha256_file(path: str | Path) -> str:
    """ファイルのSHA-256を返す。"""
    digest = hashlib.sha256()
    with Path(path).open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def reconstruct_dit(image_path: str | Path, manifest_path: str | Path, device: str = "cpu", config_path: str | Path | None = None):
    """model_png.jsonのname/shape順でDiTを再構築し、copy後の全parameterをfinite gateする。"""
    import torch

    metadata = json.loads(Path(manifest_path).read_text(encoding="utf-8"))
    config = json.loads(Path(config_path).read_text(encoding="utf-8"))["dit"] if config_path is not None else metadata["cfg"]
    from dit import DiT

    model = DiT(dim=int(config["dim"]), depth=int(config["depth"]), heads=int(config["heads"]))
    named = dict(model.named_parameters())
    pixels = np.asarray(Image.open(image_path).convert("RGB"), dtype=np.uint8).reshape(-1, 3)
    total = int(metadata["total_parameters"])
    if pixels.shape[0] < total:
        raise ValueError("weight image is shorter than manifest")
    bits = (pixels[:total, 0].astype(np.uint16) << 8) | pixels[:total, 1].astype(np.uint16)
    flat = bits.view(np.float16)
    offset = 0
    with torch.no_grad():
        for item in metadata["params"]:
            name = str(item["name"])
            shape = tuple(int(value) for value in item["shape"])
            numel = int(item["numel"])
            if name not in named or tuple(named[name].shape) != shape or numel != int(np.prod(shape)):
                raise ValueError(f"manifest mismatch for {name}")
            named[name].copy_(torch.from_numpy(flat[offset:offset + numel].copy()).view(shape).float())
            offset += numel
    if offset != total or any(not bool(torch.isfinite(parameter).all()) for parameter in model.parameters()):
        raise ValueError("finite gate failed after model parameter copy")
    return model.to(device).eval()