"""Mild synthetic controls; no model benchmark and no autoplay assets."""
import argparse
import hashlib
import json
import subprocess
from pathlib import Path

import imageio_ffmpeg
import numpy as np
from PIL import Image, ImageDraw
from flicker_probe import measure

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--output', type=Path, default=Path('flicker-controls'))
parser.add_argument('--implicit-range', action='store_true', help='Reproduce the earlier failed color-conversion trial')
args = parser.parse_args()
ROOT = args.output.resolve()
if ROOT.exists():
    parser.error('Choose a new output directory; existing data is never overwritten')
ROOT.mkdir(parents=True)
FFMPEG = imageio_ffmpeg.get_ffmpeg_exe()
WIDTH, HEIGHT, FPS, COUNT = 480, 270, 24, 96
ROI = (180, 100, 121, 81)
BACKGROUND = (20, 20, 80, 80)
image = Image.new('RGB', (WIDTH, HEIGHT), (110, 110, 110))
draw = ImageDraw.Draw(image)
draw.rectangle((160, 45, 320, 235), fill=(190, 190, 190))
draw.rectangle((180, 100, 300, 180), fill=(55, 55, 55))
draw.rectangle((195, 120, 215, 160), fill=(220, 220, 220))
original = np.array(image)
shifted = original.copy()
shifted[120:161, 195:216] = 55
shifted[120:161, 265:286] = 220


def run(arguments, **kwargs):
    return subprocess.run([FFMPEG, '-hide_banner', '-loglevel', 'error', *arguments], check=True, **kwargs)


results = []
for name in ['clean', 'exposure', 'geometry', 'mixed']:
    frames = []
    for index in range(COUNT):
        frame = shifted if name in ['geometry', 'mixed'] and (index // 6) % 2 else original
        gain = 1.04 if (index // 4) % 2 else 0.96
        if name not in ['exposure', 'mixed']:
            gain = 1.0
        frames.append((frame.astype(float) * gain).round().clip(0, 255).astype('uint8'))
    source = ROOT / f'{name}.mkv'
    output = ROOT / f'{name}-deflicker.mkv'
    run(['-y', '-f', 'rawvideo', '-pix_fmt', 'rgb24', '-s', f'{WIDTH}x{HEIGHT}',
         '-r', str(FPS), '-i', 'pipe:0', '-c:v', 'ffv1', str(source)], input=np.stack(frames).tobytes())
    filter_chain = 'deflicker=size=5:mode=am'
    color_options = []
    if not args.implicit_range:
        filter_chain = 'scale=in_range=full:out_range=full,format=yuv444p,' + filter_chain
        color_options = ['-color_range', 'pc']
    run(['-y', '-i', str(source), '-vf', filter_chain, *color_options, '-c:v', 'ffv1', str(output)])
    pair = {'control': name}
    for key, filename in [('before', source), ('after', output)]:
        rows, summary = measure(filename, ROI, BACKGROUND)
        pair[key] = {'file': filename.name, 'sha256': hashlib.sha256(filename.read_bytes()).hexdigest(), **summary}
        (ROOT / f'{name}-{key}-pairs.json').write_text(json.dumps(rows, indent=2))
    results.append(pair)
    strip = Image.new('RGB', (WIDTH * 3, HEIGHT))
    for position, index in enumerate([0, 4, 6]):
        strip.paste(Image.fromarray(frames[index]), (position * WIDTH, 0))
    strip.save(ROOT / f'{name}-frames.png')

by_name = {row['control']: row for row in results}
exposure, geometry, mixed = [by_name[name] for name in ['exposure', 'geometry', 'mixed']]
checks = {
    'clean_static': by_name['clean']['after']['mean_residual_rgb'] == 0,
    'exposure_ratio_reduced': exposure['after']['mean_abs_ratio_minus_one'] < exposure['before']['mean_abs_ratio_minus_one'] * 0.65,
    'exposure_residual_small': exposure['after']['mean_residual_rgb'] < 0.5,
    'geometry_residual_retained': geometry['after']['mean_residual_rgb'] > geometry['before']['mean_residual_rgb'] * 0.9,
    'mixed_ratio_reduced': mixed['after']['mean_abs_ratio_minus_one'] < mixed['before']['mean_abs_ratio_minus_one'] * 0.65,
    'mixed_residual_retained': mixed['after']['mean_residual_rgb'] > mixed['before']['mean_residual_rgb'] * 0.85,
    'mixed_separated_from_exposure': mixed['after']['mean_residual_rgb'] > exposure['after']['mean_residual_rgb'] * 5,
}
record = {'method': ('Implicit FFmpeg color-range conversion (failed control). ' if args.implicit_range else 'Explicit full-range conversion. ') + 'Fixed regions in synthetic stationary 480x270 RGB clips, 96 frames at 24fps. Mild +/-4% gain every four frames; mark shifts every six frames. FFmpeg deflicker size5 arithmetic mean; decoded RGB measurements. See implicit_range_trial for conversion mode.',
          'implicit_range_trial': args.implicit_range,
          'ffmpeg': subprocess.check_output([FFMPEG, '-version'], text=True).splitlines()[0],
          'roi': ROI, 'background': BACKGROUND, 'checks': checks, 'results': results}
(ROOT / 'measurements.json').write_text(json.dumps(record, indent=2))
print(json.dumps(record, indent=2))
if not all(checks.values()):
    raise SystemExit('Experiment failed: do not draft a page claiming separation.')