"""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 argparse
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_graviton_executorch_optimized.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],
    output_path: Path | None,
    summary_output_path: Path | None,
) -> None:
    """Save explicitly requested inference outputs."""
    if output_path:
        Image.fromarray(sr_array).save(output_path)
        print(f"Super-resolved image saved to: {output_path}")

    if summary_output_path:
        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,
        }
        with summary_output_path.open("w") as f:
            json.dump(summary, f, indent=2)
        print(f"Summary saved to: {summary_output_path}")


# ── Main ───────────────────────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, help="Path for the super-resolved PNG")
    parser.add_argument("--summary-output", type=Path, help="Path for the JSON summary")
    args = parser.parse_args()

    script_dir = Path(__file__).resolve().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, args.output, args.summary_output)
    if not args.output and not args.summary_output:
        print("No output files written; use --output or --summary-output to save results.")


if __name__ == "__main__":
    main()