import argparse
import os

import numpy as np
import onnxruntime as ort
from PIL import Image


HEIGHT = 512
WIDTH = 512
IMG_EXTENSIONS = ('.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP')


def parse_args():
    parser = argparse.ArgumentParser(description='Run GCANet ONNX inference with fixed 512x512 input.')
    parser.add_argument('--task', default='dehaze', choices=['dehaze', 'derain'])
    parser.add_argument('--onnx', default=None, help='Path to ONNX model. Default: onnx/gcanet_{task}_512x512_sim.onnx')
    parser.add_argument('--indir', default='examples')
    parser.add_argument('--outdir', default='onnx_output')
    return parser.parse_args()


def make_dataset(image_dir):
    images = []
    assert os.path.isdir(image_dir), '%s is not a valid directory' % image_dir
    for root, _, fnames in sorted(os.walk(image_dir)):
        for fname in fnames:
            if fname.endswith(IMG_EXTENSIONS):
                images.append(os.path.join(root, fname))
    return images


def edge_compute_np(img_chw):
    x_diffx = np.abs(img_chw[:, :, 1:] - img_chw[:, :, :-1])
    x_diffy = np.abs(img_chw[:, 1:, :] - img_chw[:, :-1, :])

    edge = np.zeros_like(img_chw, dtype=np.float32)
    edge[:, :, 1:] += x_diffx
    edge[:, :, :-1] += x_diffx
    edge[:, 1:, :] += x_diffy
    edge[:, :-1, :] += x_diffy
    edge = np.sum(edge, axis=0, keepdims=True) / 3.0
    edge = edge / 4.0
    return edge.astype(np.float32)


def preprocess(img_path):
    img = Image.open(img_path).convert('RGB')
    img = img.resize((WIDTH, HEIGHT), Image.BICUBIC)
    img_np = np.array(img).astype(np.float32)
    img_chw = np.transpose(img_np, (2, 0, 1))
    edge = edge_compute_np(img_chw)
    model_input = np.concatenate((img_chw, edge), axis=0)[None, :, :, :] - 128.0
    return img_chw, model_input.astype(np.float32)


def postprocess(pred, img_chw, only_residual):
    out = pred[0]
    if only_residual:
        out = out + img_chw
    out = np.round(out).clip(0, 255).astype(np.uint8)
    out = np.transpose(out, (1, 2, 0))
    return out


def main():
    args = parse_args()
    onnx_path = args.onnx or os.path.join('onnx', 'gcanet_%s_512x512_sim.onnx' % args.task)
    only_residual = args.task == 'dehaze'
    os.makedirs(args.outdir, exist_ok=True)

    session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
    input_name = session.get_inputs()[0].name

    for img_path in make_dataset(args.indir):
        img_chw, model_input = preprocess(img_path)
        pred = session.run(None, {input_name: model_input})[0]
        out_img = postprocess(pred, img_chw, only_residual)
        save_name = os.path.splitext(os.path.basename(img_path))[0] + '_%s_onnx.png' % args.task
        Image.fromarray(out_img).save(os.path.join(args.outdir, save_name))
        print('Saved:', os.path.join(args.outdir, save_name))


if __name__ == '__main__':
    main()