#!/usr/bin/env python

import sys
import Image
from math import *
from scipy.special.orthogonal import genlaguerre

size = 320
I_0 = 1.
w = size/4.

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 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 = hypot(x,y)
    phi = atan2(x,y)
    rho = 2*r**2/w**2
    I = I_0 * rho**l * (genlaguerre(p,l)(rho))**2 \
      * (cos(l*phi))**2 * 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)
  print 'row ' + str((x+1)/2) + ' of ' + str(size/2) + ' complete'

# 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))
  print 'row ' + str(x+1) + ' of ' + str(size) + ' complete'