"""Runs each model on the sample images on the CPU and compares its output with the first model's."""
import os
import sys
import time

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

samples_dir, out_dir, *models = sys.argv[1:]
os.makedirs(out_dir, exist_ok=True)
MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)


def pixels(path):
    image = Image.open(path).convert("RGB")
    resized = np.asarray(image.resize((1024, 1024), Image.BILINEAR), dtype=np.float32) / 255
    return image, ((resized - MEAN) / STD).transpose(2, 0, 1)[None]


def sigmoid(x):
    return 1 / (1 + np.exp(-x))


samples = [(name, *pixels(os.path.join(samples_dir, name))) for name in sorted(os.listdir(samples_dir))]
reference = {}
for index, path in enumerate(models):
    label = os.path.splitext(os.path.basename(path))[0]
    options = ort.SessionOptions()
    options.log_severity_level = 3
    started = time.time()
    session = ort.InferenceSession(path, options, providers=["CPUExecutionProvider"])
    print(f"{label}: loaded in {time.time() - started:.1f} s, {os.path.getsize(path) / 1e6:.0f} MB")
    input_name = session.get_inputs()[0].name
    for name, image, x in samples:
        started = time.time()
        logits = session.run(None, {input_name: x})[0].astype(np.float32)[0, 0]
        took = time.time() - started
        alpha = sigmoid(logits)
        Image.fromarray((alpha * 255).round().astype(np.uint8)).resize(image.size, Image.BILINEAR).save(os.path.join(out_dir, f"{os.path.splitext(name)[0]}-{label}.png"))
        if index == 0:
            reference[name] = (logits, alpha)
            print(f"  {name}: {took:.1f} s, logits {logits.min():.2f}..{logits.max():.2f}, {100 * (alpha > 0.5).mean():.1f}% foreground")
            continue
        ref_logits, ref_alpha = reference[name]
        diff = np.abs(alpha - ref_alpha)
        print(
            f"  {name}: {took:.1f} s, max logit diff {np.abs(logits - ref_logits).max():.2e}, "
            f"alpha diff max {diff.max():.2e} mean {diff.mean():.2e}, "
            f"pixels off by more than 1/255: {100 * (diff > 1 / 255).mean():.3f}%"
        )