from __future__ import annotations

import argparse
from pathlib import Path

import caffe
import cv2


BASE_DIR = Path(__file__).resolve().parents[1]
SUPPORTED_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"}


def edit_fcn_proto(template_file: Path, output_file: Path, height: int, width: int) -> None:
    template = template_file.read_text()
    output_file.write_text(template.format(height=height, width=width))


def iter_images(input_dir: Path) -> list[Path]:
    return [
        path
        for path in sorted(input_dir.iterdir(), key=lambda p: p.name.lower())
        if path.is_file() and path.suffix.lower() in SUPPORTED_EXTENSIONS
    ]


def run(input_dir: Path, output_dir: Path, max_dim: int) -> None:
    caffe.set_mode_cpu()
    output_dir.mkdir(parents=True, exist_ok=True)

    template_file = BASE_DIR / "test" / "test_template.prototxt"
    deploy_file = BASE_DIR / "test" / "DeployT.prototxt"
    model_file = BASE_DIR / "AOD_Net.caffemodel"

    images = iter_images(input_dir)
    print(f"image numbers: {len(images)}")

    for index, img_path in enumerate(images, start=1):
        npstore = caffe.io.load_image(str(img_path))
        orig_h, orig_w = npstore.shape[0], npstore.shape[1]

        if orig_h > max_dim or orig_w > max_dim:
            scale = max_dim / float(max(orig_h, orig_w))
            new_h = int(orig_h * scale)
            new_w = int(orig_w * scale)
            npstore = cv2.resize(npstore, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
            print(f"Resized {img_path.name} from {orig_w}x{orig_h} to {new_w}x{new_h}")

        height, width = npstore.shape[0], npstore.shape[1]
        edit_fcn_proto(template_file, deploy_file, height, width)

        net = caffe.Net(str(deploy_file), str(model_file), caffe.TEST)
        data = npstore.transpose((2, 0, 1))
        net.blobs["data"].data[...] = [data]
        net.forward()

        result = net.blobs["sum"].data[0].transpose((1, 2, 0))
        result = result[:, :, ::-1]

        if height != orig_h or width != orig_w:
            result = cv2.resize(result, (orig_w, orig_h), interpolation=cv2.INTER_CUBIC)

        save_path = output_dir / f"{img_path.stem}_AOD-Net.png"
        cv2.imwrite(str(save_path), result * 255.0, [cv2.IMWRITE_JPEG_QUALITY, 100])
        print(f"[{index}/{len(images)}] saved: {save_path}")


def main() -> None:
    parser = argparse.ArgumentParser(description="Run AOD-Net on a folder of images.")
    parser.add_argument("input_dir", nargs="?", default=str(BASE_DIR / "data" / "img"))
    parser.add_argument("output_dir", nargs="?", default=str(BASE_DIR / "data" / "result"))
    parser.add_argument("--max-dim", type=int, default=1920)
    args = parser.parse_args()

    run(Path(args.input_dir), Path(args.output_dir), args.max_dim)


if __name__ == "__main__":
    main()