#!/usr/bin/env python3

import sys
import math
from PIL import Image
from scipy.special import genlaguerre

size = 320
I_0 = 1.0
w = size / 4.0

p = int(sys.argv[1])
l = int(sys.argv[2])

def linear_to_sRGB(l):
    # Formula from http://www.w3.org/Graphics/Color/sRGB
    if l <= 0.00304:
        l = 12.92 * l
    else:
        l = 1.055 * pow(l, 1.0 / 2.4) - 0.055
    return int(255.0 * l)

# First, let's make a floating-point image of the raw intensities.
raw = Image.new("F", (size, size))
high = 0
for x in range(1, size, 2):
    for y in range(1, size, 2):
        r = math.hypot(x, y)
        phi = math.atan2(x, y)
        rho = 2 * r**2 / w**2
        I = (
            I_0
            * rho**l
            * (genlaguerre(p, l)(rho)) ** 2
            * (math.cos(l * phi)) ** 2
            * math.exp(-rho)
        )
        if I > high:
            high = I
        raw.putpixel((size // 2 + (x - 1) // 2, size // 2 + (y - 1) // 2), I)
        raw.putpixel((size // 2 + (x - 1) // 2, size // 2 - (y + 1) // 2), I)
        raw.putpixel((size // 2 - (x + 1) // 2, size // 2 + (y - 1) // 2), I)
        raw.putpixel((size // 2 - (x + 1) // 2, size // 2 - (y + 1) // 2), I)

# Now, let's normalize them and export them as sRGB.
cooked = Image.new("L", (size, size))
for x in range(size):
    for y in range(size):
        I = raw.getpixel((x, y)) / high
        cooked.putpixel((x, y), linear_to_sRGB(I))

cooked.save(f"plot-{p}-{l}.png")