import argparse
import os

import numpy as np
import axengine as axe
from PIL import Image, ImageDraw, ImageFont


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 axmodel inference with fixed 512x512 input.')
    parser.add_argument('--task', default='dehaze', choices=['dehaze', 'derain'])
    parser.add_argument('--axmodel', default='GCANet_u16.axmodel', help='Path to axmodel model.')
    parser.add_argument('--indir', default='examples')
    parser.add_argument('--outdir', default='axmodel_output')
    parser.add_argument(
        '--input-mode',
        default='raw_u8',
        choices=['compat_centered_u8', 'raw_u8'],
        help='compat_centered_u8 keeps the historical centered input encoding for the current axmodel; raw_u8 follows the declared U8 input processor literally.',
    )
    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, input_mode):
    img = Image.open(img_path).convert('RGB')
    img = img.resize((WIDTH, HEIGHT), Image.BICUBIC)
    img_rgb = np.array(img).astype(np.float32)
    img_rgb_chw = np.transpose(img_rgb, (2, 0, 1))
    edge = edge_compute_np(img_rgb_chw)

    # This 4-channel model is calibrated from RGB+edge NCHW tensors.
    # The build config marks src_format=BGR/tensor_format=RGB, but with
    # csc_mode=NoCSC there is no runtime color conversion, so feeding BGR here
    # would silently swap channels and introduce artifacts.
    model_input = np.concatenate((img_rgb_chw, edge), axis=0)[None, :, :, :]
    if input_mode == 'compat_centered_u8':
        # The original float model consumes (rgb+edge-128). The current axmodel
        # was compiled as U8 input without explicit dequant params, so we keep
        # the historical wraparound encoding here for compatibility.
        model_input = np.mod(np.round(model_input - 128.0), 256.0).astype(np.uint8)
    else:
        model_input = np.clip(np.round(model_input), 0, 255).astype(np.uint8)
    return img_rgb_chw, model_input


def postprocess(pred, img_chw, only_residual):
    out = pred[0].astype(np.float32)
    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 make_comparison_image(input_chw, output_hwc, task):
    input_hwc = np.transpose(input_chw, (1, 2, 0)).clip(0, 255).astype(np.uint8)
    input_img = Image.fromarray(input_hwc, mode='RGB')
    output_img = Image.fromarray(output_hwc, mode='RGB')

    title_height = 32
    comparison = Image.new('RGB', (WIDTH * 2, HEIGHT + title_height), color='white')
    comparison.paste(input_img, (0, title_height))
    comparison.paste(output_img, (WIDTH, title_height))

    draw = ImageDraw.Draw(comparison)
    font = ImageFont.load_default()
    right_label = 'dehaze' if task == 'dehaze' else task
    draw.text((8, 8), 'hazy', fill='black', font=font)
    draw.text((WIDTH + 8, 8), right_label, fill='black', font=font)
    draw.line((WIDTH, 0, WIDTH, HEIGHT + title_height), fill='black', width=1)
    return comparison


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

    session = axe.InferenceSession(axmodel_path, providers=['AxEngineExecutionProvider'])
    input_name = session.get_inputs()[0].name

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


if __name__ == '__main__':
    main()