"""Minimal inference example for ESRGAN 4x Super-Resolution using ExecuTorch.

Loads a quantized .pte model and runs super-resolution inference on a single image.
The model upscales the input image by 4x. For images larger than 128x128, the
input is automatically split into overlapping tiles, each tile is super-resolved,
and the results are blended and stitched into the final output.
"""

import json
from pathlib import Path

import numpy as np
import torch
from executorch.runtime import Runtime
from PIL import Image
from torchvision import transforms

# ── Configuration ──────────────────────────────────────────────────────────────
MODEL_PATH = "esrgan-x4-int8-executorch.pte"
IMAGE_PATH = "sample_input.jpg"
TILE_SIZE = 128       # model was exported with 128x128 input tiles
SCALE = 4             # 4x upscaling factor
TILE_OVERLAP = 8      # pixel overlap between tiles for seamless blending


# ── Model Loading ──────────────────────────────────────────────────────────────
def load_model(pte_path: str):
    """Load ExecuTorch .pte model and return the forward method."""
    runtime = Runtime.get()
    program = runtime.load_program(pte_path)
    return program.load_method("forward")


# ── Preprocessing ──────────────────────────────────────────────────────────────
def preprocess(image_path: str) -> tuple[torch.Tensor, tuple[int, int]]:
    """Load and preprocess image for model input.

    The model expects [0, 1] float32 RGB tensors — no normalization is applied.
    Returns the input tensor and the original image size for reference.
    """
    image = Image.open(image_path).convert("RGB")
    original_size = (image.width, image.height)
    to_tensor = transforms.ToTensor()  # converts [0-255] uint8 -> [0, 1] float32
    input_tensor = to_tensor(image).unsqueeze(0)  # [1, 3, H, W]
    return input_tensor, original_size


# ── Tiled Inference ────────────────────────────────────────────────────────────
def run_tiled_inference(method, input_tensor: torch.Tensor) -> torch.Tensor:
    """Run super-resolution inference with overlapping tile stitching.

    Splits the input into overlapping 128x128 tiles, runs each through the
    model, and blends overlapping regions using pixel-level averaging.
    """
    _, _, h, w = input_tensor.shape
    out_h, out_w = h * SCALE, w * SCALE

    output = torch.zeros(1, 3, out_h, out_w)
    weights = torch.zeros(1, 1, out_h, out_w)

    stride = TILE_SIZE - TILE_OVERLAP

    y_positions = list(range(0, max(1, h - TILE_SIZE + 1), stride))
    if not y_positions or y_positions[-1] + TILE_SIZE < h:
        y_positions.append(max(0, h - TILE_SIZE))

    x_positions = list(range(0, max(1, w - TILE_SIZE + 1), stride))
    if not x_positions or x_positions[-1] + TILE_SIZE < w:
        x_positions.append(max(0, w - TILE_SIZE))

    for y in y_positions:
        for x in x_positions:
            tile = input_tensor[:, :, y:y + TILE_SIZE, x:x + TILE_SIZE].contiguous()
            sr_tile = method.execute([tile])[0]

            oy, ox = y * SCALE, x * SCALE
            oh, ow = TILE_SIZE * SCALE, TILE_SIZE * SCALE
            output[:, :, oy:oy + oh, ox:ox + ow] += sr_tile
            weights[:, :, oy:oy + oh, ox:ox + ow] += 1.0

    return output / weights.clamp(min=1.0)


# ── Postprocessing ─────────────────────────────────────────────────────────────
def postprocess(raw_output: torch.Tensor) -> np.ndarray:
    """Clamp output to [0, 1] and convert to uint8 numpy array (H, W, 3)."""
    sr_tensor = raw_output.clamp(0.0, 1.0).squeeze(0)  # [3, H*4, W*4]
    sr_array = (sr_tensor.permute(1, 2, 0).numpy() * 255.0).round().astype(np.uint8)
    return sr_array


# ── Save Results ──────────────────────────────────────────────────────────────
def save_results(sr_array: np.ndarray, original_size: tuple[int, int]) -> None:
    """Save the super-resolved image and a JSON summary to the script directory."""
    script_dir = Path(__file__).parent

    # Save super-resolved image
    output_image = Image.fromarray(sr_array)
    output_path = script_dir / "sample_output.png"
    output_image.save(output_path)
    print(f"Super-resolved image saved to: {output_path}")

    # Save JSON summary
    summary = {
        "input_size": {"width": original_size[0], "height": original_size[1]},
        "output_size": {"width": sr_array.shape[1], "height": sr_array.shape[0]},
        "scale_factor": SCALE,
    }
    json_path = script_dir / "super_resolution.json"
    with open(json_path, "w") as f:
        json.dump(summary, f, indent=2)
    print(f"Summary saved to: {json_path}")


# ── Main ───────────────────────────────────────────────────────────────────────
def main():
    script_dir = Path(__file__).parent
    model_path = script_dir / MODEL_PATH
    image_path = script_dir / IMAGE_PATH

    print(f"Loading model from: {model_path}")
    method = load_model(str(model_path))

    print(f"Preprocessing image: {image_path}")
    input_tensor, original_size = preprocess(str(image_path))
    _, _, h, w = input_tensor.shape
    print(f"  Input size: {w}x{h} -> output will be {w * SCALE}x{h * SCALE}")

    print("Running tiled super-resolution inference...")
    raw_output = run_tiled_inference(method, input_tensor)

    sr_array = postprocess(raw_output)
    print(f"  Output shape: {sr_array.shape[1]}x{sr_array.shape[0]} (WxH)")

    save_results(sr_array, original_size)


if __name__ == "__main__":
    main()