#!/usr/bin/env python

import numpy as np
import sys
import imageio.v3 as imageio  # aur: python-imageio
import scipy.optimize
import random


def srgb_to_linear(x):
    return np.fmax(x / 12.92, ((x + 0.055) / 1.055) ** 2.4)


def linear_to_srgb(x):
    return np.fmax(np.fmin(x * 12.92, 0.04045), 1.055 * x ** (1.0 / 2.4) - 0.055)


def solve_srgb(srgb_avg, light_avg):
    """
    Find Y in [max(0, 2*srgb_avg - 1), srgb_avg]

    srgb_to_lin(Y) + srgb_to_lin(2 * srgb_avg - Y) = 2 * light_avg
    """
    Y_min = max(0, 2 * srgb_avg - 1)
    Y_max = srgb_avg

    def fn(y):
        return srgb_to_linear(y) + srgb_to_linear(2 * srgb_avg - y) - 2 * light_avg

    x0 = scipy.optimize.brentq(fn, Y_min, Y_max)
    return x0


def max_range():
    scipy.optimize.maximize


if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("hide_image.py: src.png dst.png")
        print()
        print("Input should be an sRGB encoded PNG file with height divisible by two")

    src_path, dst_path = sys.argv[1:]
    print(src_path, ">", dst_path)

    im = imageio.imread(src_path)

    assert im.shape[0] % 2 == 0
    im = im / 255.0

    # First, shrink Y by a factor of 2
    im = srgb_to_linear(im)
    lin = 0.5 * (im[::2, :, :] + im[1::2, :, :])

    # the average srgb value we are targeting
    srgb_avg = 0.50

    # range of light values which we can encode in this way.
    min_light = srgb_to_linear(srgb_avg)
    srgb_min = max(0, 2 * srgb_avg - 1)
    max_light = 0.5 * (
        srgb_to_linear(srgb_min) + srgb_to_linear(2 * srgb_avg - srgb_min)
    )
    print(
        f"sRGB avg value: {srgb_avg:.6f} Light range: [{min_light:.6f}, {max_light:.6f}]"
    )

    # note: at pure gamma = 2.4, max light/min light ratio is 2**2.4/2 = 2.639
    # but due to srgb effects/sum clipping, srgb_avg = 0.5 maximizes ratio at 2.337
    # larger ratios are possible using more than 2 pixels -- up to ratio of 12.92

    db = np.zeros((lin.shape[0] * 2, lin.shape[1], lin.shape[2]))
    for i in range(lin.shape[0]):
        for j in range(lin.shape[1]):
            for k in range(lin.shape[2]):
                shifted_light = (max_light - min_light) * lin[i, j, k] + min_light
                srgb = solve_srgb(srgb_avg, shifted_light)
                
                # we can introduce arbitrary fine-grained patterns
                # by switching which of the two pixels gets the higher value
                # i.e, moving the color mass around
                if False:
                    # noisy split
                    if random.random() > 0.5:
                        srgb = 2 * srgb_avg - srgb
                else:
                    # pink/green split
                    if k == 1:
                        srgb = 2 * srgb_avg - srgb
                    
                db[2 * i, j, k] = srgb
                db[2 * i + 1, j, k] = 2 * srgb_avg - srgb

    # Save image
    out = np.array(np.round(255 * db), dtype=np.uint8)
    imageio.imwrite(dst_path, out)