#!/usr/bin/env python3
# Copyright (c) the JPEG XL Project Authors. All rights reserved.
#
# Use of this source code is governed by a BSD-style
# license that can be found in the LICENSE file.
"""JPEG XL conformance test runner.

Tool to perform a conformance test for a decoder.
"""

import argparse
import json
import numpy
import os
import shutil
import subprocess
import sys
import tempfile

import lcms2

BRIGHT_RED = "91"
BRIGHT_BLUE = "94"

def ColoredMessage(color, message):
    print(f"\033[{color}m{message}\033[0m", flush=True)

def Failure(message):
    ColoredMessage(BRIGHT_RED, message)
    return False

def CompareNPY(ref, ref_icc, dec, dec_icc, frame_idx, rmse_limit, peak_error):
    """Compare a decoded numpy against the reference one."""
    if ref.shape != dec.shape:
        return Failure(f'Expected shape {ref.shape} but found {dec.shape}')
    ref_frame = ref[frame_idx]
    dec_frame = dec[frame_idx]
    num_channels = ref_frame.shape[2]

    if ref_icc != dec_icc and peak_error > 0:
        # Transform colors before comparison.
        # Skip this if we expect fully lossless, since it introduces tiny errors
        # (even in the case where ref_icc and dec_icc are equivalent so it's a no-op)
        if num_channels < 3:
            return Failure(f"Only RGB images are supported")
        dec_clr = dec_frame[:, :, 0:3]
        dec_frame[:, :, 0:3] = lcms2.convert_pixels(dec_icc, ref_icc, dec_clr)

    error = numpy.abs(ref_frame - dec_frame)
    actual_peak_error = error.max()
    error_by_channel = [error[:, :, ch] for ch in range(num_channels)]
    actual_rmses = [numpy.sqrt(numpy.mean(error_ch * error_ch)) for error_ch in error_by_channel]
    actual_rmse = max(actual_rmses)

    print(f"RMSE: {actual_rmses}, actual peak: {actual_peak_error}", flush=True)

    if actual_rmse > rmse_limit:
        return Failure(f"RMSE too large: {actual_rmse} > {rmse_limit}")

    if actual_peak_error > peak_error:
        return Failure(
            f"Peak error too large: {actual_peak_error} > {peak_error}")
    return True


def CompareBinaries(ref_bin, dec_bin):
    """Compare a decoded binary file against the reference for exact contents."""
    with open(ref_bin, 'rb') as reff:
        ref_data = reff.read()

    with open(dec_bin, 'rb') as decf:
        dec_data = decf.read()

    if ref_data != dec_data:
        return Failure(
            f'Binary files mismatch: {ref_bin} {dec_bin}')
    return True


TEST_KEYS = set(
    ['reconstructed_jpeg', 'original_icc', 'rms_error', 'peak_error'])


def CheckMeta(dec, ref):
    if isinstance(ref, dict):
        if not isinstance(dec, dict):
            return Failure("Malformed metadata file")
        for k, v in ref.items():
            if k in TEST_KEYS:
                continue
            if k not in dec:
                return Failure(
                    f"Malformed metadata file: key {k} not found")
            vv = dec[k]
            return CheckMeta(vv, v)
    elif isinstance(ref, list):
        if not isinstance(dec, list) or len(dec) != len(ref):
            return Failure("Malformed metadata file")
        for vv, v in zip(dec, ref):
            return CheckMeta(vv, v)
    elif isinstance(ref, float):
        if not isinstance(dec, float):
            return Failure("Malformed metadata file")
        if abs(dec - ref) > 0.0001:
            return Failure(
                f"Metadata: Expected {ref}, found {dec}")
    elif dec != ref:
        return Failure(f"Metadata: Expected {ref}, found {dec}")
    return True


def RunTest(context, test_id, work_dir):
    ok = True
    test_dir = os.path.join(context.corpus_dir, test_id)

    with open(os.path.join(test_dir, 'test.json'), 'r') as f:
        descriptor = json.load(f)
        if 'sha256sums' in descriptor:
            del descriptor['sha256sums']

    exact_tests = []

    input_filename = os.path.join(test_dir, 'input.jxl')
    pixel_prefix = os.path.join(work_dir, 'decoded')
    output_filename = pixel_prefix + '_image.npy'
    cmd = context.decoder_cmd + [input_filename, output_filename]
    cmd_jpeg = []
    if 'preview' in descriptor:
        preview_filename = os.path.join(work_dir,
                                        'decoded_preview.npy')
        cmd.extend(['--preview_out', preview_filename])
    if 'reconstructed_jpeg' in descriptor:
        jpeg_filename = os.path.join(work_dir, 'reconstructed.jpg')
        cmd_jpeg = context.decoder_cmd + [input_filename, jpeg_filename]
        exact_tests.append(('reconstructed.jpg', jpeg_filename))
    if 'original_icc' in descriptor:
        decoded_original_icc = os.path.join(
            work_dir, 'decoded_org.icc')
        cmd.extend(['--orig_icc_out', decoded_original_icc])
        exact_tests.append(('original.icc', decoded_original_icc))
    meta_filename = os.path.join(work_dir, 'meta.json')
    cmd.extend(['--metadata_out', meta_filename])
    cmd.extend(['--icc_out', pixel_prefix + '.icc'])
    cmd.extend(['--norender_spotcolors'])

    print(f"Running: {cmd}", flush=True)
    if subprocess.call(cmd) != 0:
        return Failure('Running the decoder (%s) returned error' %
                        ' '.join(cmd))
    if cmd_jpeg:
        print(f"Running: {cmd_jpeg}", flush=True)
        if subprocess.call(cmd_jpeg) != 0:
            return Failure(
                'Running the decoder (%s) returned error' %
                ' '.join(cmd_jpeg))

    # Run validation of exact files.
    for reference_basename, decoded_filename in exact_tests:
        reference_filename = os.path.join(test_dir,
                                            reference_basename)
        binary_ok = CompareBinaries(reference_filename,
                                    decoded_filename)
        if not binary_ok and context.update_on_failure:
            os.unlink(reference_filename)
            shutil.copy2(decoded_filename, reference_filename)
            binary_ok = True
        ok = ok & binary_ok

    # Validate metadata.
    with open(meta_filename, 'r') as f:
        meta = json.load(f)

    ok = ok & CheckMeta(meta, descriptor)

    # Pixel data.
    decoded_icc = pixel_prefix + '.icc'
    with open(decoded_icc, 'rb') as f:
        decoded_icc = f.read()
    reference_icc = os.path.join(test_dir, "reference.icc")
    with open(reference_icc, 'rb') as f:
        reference_icc = f.read()

    reference_npy_fn = os.path.join(test_dir, 'reference_image.npy')
    decoded_npy_fn = os.path.join(work_dir, 'decoded_image.npy')

    if not os.path.exists(decoded_npy_fn):
        return Failure('File not decoded: decoded_image.npy')

    reference_npy = numpy.load(reference_npy_fn)
    decoded_npy = numpy.load(decoded_npy_fn)

    frames_ok = True
    for i, fd in enumerate(descriptor['frames']):
        frames_ok = frames_ok & CompareNPY(
            reference_npy, reference_icc, decoded_npy,
            decoded_icc, i, fd['rms_error'],
            fd['peak_error'])

    if not frames_ok and context.update_on_failure:
        os.unlink(reference_npy_fn)
        shutil.copy2(decoded_npy_fn, reference_npy_fn)
        frames_ok = True
    ok = ok & frames_ok

    if 'preview' in descriptor:
        reference_npy_fn = os.path.join(test_dir,
                                        'reference_preview.npy')
        decoded_npy_fn = os.path.join(work_dir,
                                        'decoded_preview.npy')

        if not os.path.exists(decoded_npy_fn):
            ok = Failure(
                'File not decoded: decoded_preview.npy')

        reference_npy = numpy.load(reference_npy_fn)
        decoded_npy = numpy.load(decoded_npy_fn)
        preview_ok = CompareNPY(reference_npy, reference_icc,
                                decoded_npy, decoded_icc, 0,
                                descriptor['preview']['rms_error'],
                                descriptor['preview']['peak_error'])
        if not preview_ok & args.update_on_failure:
            os.unlink(reference_npy_fn)
            shutil.copy2(decoded_npy_fn, reference_npy_fn)
            preview_ok = True
        ok = ok & preview_ok

    return ok

def ConformanceTestRunner(args):
    context = type('Expando', (object,), {})()
    context.update_on_failure = args.update_on_failure
    # We can pass either the .txt file or the directory which defaults to the
    # full corpus. This is useful to run a subset of the corpus in other .txt
    # files.
    if os.path.isdir(args.corpus):
        context.corpus_dir = args.corpus
        corpus_txt = os.path.join(args.corpus, 'corpus.txt')
    else:
        context.corpus_dir = os.path.dirname(args.corpus)
        corpus_txt = args.corpus
    context.decoder_cmd = args.decoder.strip().split(' ')

    all_tests = []
    failed_tests = []

    with open(corpus_txt, 'r') as f:
        for test_id in f:
            all_tests.append(test_id)
            test_id = test_id.rstrip('\n')
            ColoredMessage(BRIGHT_BLUE, f"Testing {test_id}")
            with tempfile.TemporaryDirectory(prefix=test_id) as work_dir:
                test_ok = RunTest(context, test_id, work_dir)
                if not test_ok:
                    failed_tests.append(test_id)

    all = len(all_tests)
    failed = len(failed_tests)
    passed = all - failed
    ColoredMessage(BRIGHT_BLUE, f"{passed} of {all} tests passed")
    if failed_tests:
        ColoredMessage(BRIGHT_RED, "Failed tests:")
        for test in failed_tests:
            ColoredMessage(BRIGHT_RED, f"  {test}")
        return False
    return True

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--decoder',
                        metavar='DECODER',
                        required=True,
                        help='path to the decoder binary under test.')
    parser.add_argument(
        '--corpus',
        metavar='CORPUS',
        required=True,
        help=('path to the corpus directory or corpus descriptor'
              ' text file.'))
    parser.add_argument(
        '--update_on_failure', action='store_true',
        help='If set, updates reference files on failing checks.')
    args = parser.parse_args()
    if not ConformanceTestRunner(args):
        sys.exit(1)


if __name__ == '__main__':
    main()
