# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_image.ipynb.

# %% auto 0
__all__ = ['TAG_MAP', 'load_builtin_image', 'import_image', 'build_meta_canned_image', 'ehands_input_to_qcrank_input',
           'prep_image_qcrank_input', 'create_img']

# %% ../nbs/03_image.ipynb 1
import numpy as np
from PIL import Image
import os
import importlib.resources as pkg_resources
from pprint import pprint
from . import data  # qgear/data package
from toolbox.Util_H5io4 import write4_data_hdf5

# Mapping tags to filenames inside qgear/data
TAG_MAP = {
    'a1': 'alphabet_x32_y1.png',
    'b1': 'alphabet_x64_y16.png',
    'b2': 'high-heels_x32_y32.png',
    'b3': 'high-heels_x64_y64.png',
    'b4': 'high-heels_x128_y128.png',
    'c1': 'xray-finger-8y_x64_y80.png',
    'd1': 'igb_facade_x192_y128.png',
    'd2': 'zebra_x192_y128.png',
    'd3': 'xray-hand-4y_x128_y192.png',
    'd4': 'xray-foot-12y_x192_y128.png',
    'd5': 'e_coli_bacteria_x192_y128.png',
    'e1': 'zebra_x384_y256.png',
    'e2': 'cameraman_x256_y384.png',
    'e3': 'micro_germ_x384_y256.png'
}

def load_builtin_image(filename):
    """Load a PNG from qgear/data/ as numpy array."""
    img_path = pkg_resources.files(data) / filename
    img = Image.open(img_path)
    return np.array(img)

def import_image(tag, inp_path=None, verb=1):
    """Load image from external path or built-in package data."""
    if tag not in TAG_MAP:
        raise ValueError(f"Undefined tag: {tag}")
    filename = TAG_MAP[tag]

    # If no path provided → load from package data
    if inp_path is None:
        if verb:
            print(f"Loading built-in image '{filename}' from qgear/data/")
        return {"phys_image": load_builtin_image(filename), "image_name": filename.replace(".png", "")}

    # If path provided → try it, else fallback to built-in
    img_file = os.path.join(inp_path, filename)
    if not os.path.exists(img_file):
        if verb:
            print(f"Image not found at {img_file}, falling back to built-in")
        return {"phys_image": load_builtin_image(filename), "image_name": filename.replace(".png", "")}

    img = Image.open(img_file)
    return {"phys_image": np.array(img), "image_name": filename.replace(".png", "")}

# %% ../nbs/03_image.ipynb 2
# ---------------------------------
# Build metadata
# ---------------------------------
def build_meta_canned_image(tag: str, bigD: dict, nq_addr: int = 9, verb: int = 1):
    """Build metadata for canned image."""
    pixY, pixX = bigD['phys_image'].shape
    pd = {
        'num_sample': 1,
        'nq_addr': nq_addr,
        'seq_len': 1 << nq_addr
    }
    cad = {
        'image_name': bigD['image_name'],
        'image_shape_xy': [pixX, pixY],
        'image_pixels': pixX * pixY
    }
    assert cad['image_pixels'] % pd['seq_len'] == 0
    pd['nq_fdata'] = cad['image_pixels'] // pd['seq_len']
    pd['num_clbit'] = pd['nq_fdata'] + pd['nq_addr']
    assert pd['num_clbit'] <= 42
    pd['qcrank_max_fval'] = np.pi
    cad['canned_type'] = 'gray_image'
    md = {'payload': pd, 'canned': cad}
    md['short_name'] = f"canImg_{tag}_{pixX}_{pixY}"

    if verb:
        print("Metadata:", cad)
        if verb > 1:
            pprint(md)
    return md

# ---------------------------------
# Transform Escherhands input to QCrank format
# ---------------------------------
def ehands_input_to_qcrank_input(udata: np.ndarray):
    """Convert Escherhands-style data to QCrank format."""
    fdata = np.arccos(udata)
    fdata = np.transpose(fdata, (2, 1, 0))
    fdata = fdata[:, ::-1, :]
    return fdata

# ---------------------------------
# Prepare image for QCrank
# ---------------------------------
def prep_image_qcrank_input(md: dict, bigD: dict, verb: int = 1):
    """Prepare normalized image data for QCrank."""
    pmd = md['payload']
    n_addr = pmd['seq_len']
    nq_data = pmd['nq_fdata']
    n_img = pmd['num_sample']

    imgA = bigD['phys_image']
    min_val, max_val = imgA.min(), imgA.max()
    imgAN = 2 * ((imgA - min_val) / (max_val - min_val)) - 1
    if verb:
        print(f"Normalized image min/max: {imgAN.min()}, {imgAN.max()}")
    bigD['norm_image'] = imgAN.astype(np.float32)

    inpA = imgAN.flatten()
    if verb:
        print(f"Flattened image shape: {inpA.shape}, reshaped: {inpA.reshape(-1, n_addr).shape}")

    inp_udata = np.zeros((n_img, nq_data, n_addr), dtype=np.float32)
    inp_udata[0] = inpA.reshape(-1, n_addr)

    bigD['inp_udata'] = inp_udata
    bigD['inp_fdata'] = ehands_input_to_qcrank_input(inp_udata)


# %% ../nbs/03_image.ipynb 3
def create_img(tag="b2", nq_addr=9, inp_path=None, out_path="out", verb=1):
    """Create QCrank input from a canned image (built-in or external)."""
    os.makedirs(out_path, exist_ok=True)

    bigD = import_image(tag, inp_path=inp_path, verb=verb)
    md = build_meta_canned_image(tag, bigD, nq_addr=nq_addr, verb=verb)
    prep_image_qcrank_input(md, bigD, verb=verb)

    out_file = os.path.join(out_path, md['short_name'] + ".qcrank_inp.h5")
    write4_data_hdf5(bigD, out_file, metaD=md)

    if verb:
        print(f"Saved QCrank input: {out_file}")
    return md, bigD
