"""Version 2 editing protocol: independent noise and inspectable RGBA diagnostics.

Diagnostics are descriptive, not quality scores or ground-truth matting metrics.
No filtering, sharpening, alpha thresholding or cleanup is applied to outputs.
"""
import math
from pathlib import Path

import numpy as np
from PIL import Image, ImageOps

PROTOCOL = 'editing-v2'
DEFAULT_SEEDS = (1000042, 1000123)


def validate_seeds(cases, seeds):
    if (not seeds or len(set(seeds)) != len(seeds)
            or any(type(s) is not int or not 0 <= s < 2**63 for s in seeds)):
        raise ValueError('Use distinct integer seeds in [0, 2**63)')
    source_seeds = {c.get('source_seed') for c in cases} - {None}
    if source_seeds.intersection(seeds):
        raise ValueError('Editing seeds must differ from every known source seed (noise replay)')
    return list(seeds)


def load_input(path):
    with Image.open(path) as image:
        image = ImageOps.exif_transpose(image)
        has_alpha = 'A' in image.getbands() or 'transparency' in image.info
        return image.convert('RGBA' if has_alpha else 'RGB')


def edit_dimensions(size, resolution):
    if resolution < 32 or resolution % 32 or min(size) <= 0:
        raise ValueError('Positive image size and resolution divisible by 32 required')
    ratio = size[0] / size[1]
    width = math.sqrt(resolution**2 * ratio)
    return max(32, round(width / 32)*32), max(32, round(width / ratio / 32)*32)


def image_diagnostics(image):
    rgba = np.asarray(image.convert('RGBA'))
    alpha = rgba[:, :, 3]
    # The thresholds are explicitly recorded; a white RGB background is not transparent.
    rgb = rgba[:, :, :3].astype(np.float32) / 255
    return dict(image_mode=image.mode, actual_size=list(image.size),
                alpha_min=int(alpha.min()), alpha_max=int(alpha.max()),
                alpha_transparent_fraction=float((alpha <= 5).mean()),
                alpha_opaque_fraction=float((alpha >= 250).mean()),
                alpha_soft_fraction=float(((alpha > 5) & (alpha < 250)).mean()),
                alpha_thresholds=[5, 250],
                rgb_dx_mean=float(np.abs(np.diff(rgb, axis=1)).mean()) if image.width > 1 else 0.,
                rgb_dy_mean=float(np.abs(np.diff(rgb, axis=0)).mean()) if image.height > 1 else 0.)


def save_diagnostics(image, output, stem):
    output = Path(output)
    output.mkdir(parents=True, exist_ok=True)
    rgba = image.convert('RGBA')
    rgba.getchannel('A').save(output/f'{stem}-alpha.png')
    for name, color in [('white', 'white'), ('black', 'black')]:
        background = Image.new('RGBA', image.size, color)
        Image.alpha_composite(background, rgba).convert('RGB').save(output/f'{stem}-{name}.png')
    yy, xx = np.indices((image.height, image.width))
    grid = np.where((xx//16 + yy//16) % 2, 192, 240).astype(np.uint8)
    background = Image.fromarray(np.repeat(grid[:, :, None], 3, axis=2)).convert('RGBA')
    Image.alpha_composite(background, rgba).convert('RGB').save(output/f'{stem}-checker.png')