# -------------------------------------------------------------------------------
# Create a 360 degree panoramic image by projection onto a cylinder.
#
# Make stereo left/right eye images where each vertical cylinder slice is
# rendered using cameras facing that slice (ie. line between cameras is
# perpendicular to slice direction).
#
# Make hemisphere dome stereo images by projecting onto 90 degree cone
# with axis along y using camera at center of cone base tilted 45
# degrees up with 90 degree vertical field of view.  Then project cone
# onto hemisphere, and hemisphere onto disk with radius proportional
# to angle from zenith.  This is for stereoscopic projection at the
# Imiloa planetarium in Hawaii.
#
def panorama(stereo = False, scale = 1):

    cp = circumference_pixels()
    rc = camera_center(stereo)

    import Matrix as M
    rtf = M.rotation_transform((0,1,0), 360.0 / cp, rc)
    rxf = M.chimera_xform(rtf)

    return cylinder_images(cp, rxf, stereo, scale)

# -------------------------------------------------------------------------------
# Compute radians per-pixels at center of field, then closest integer number of
# pixels in 360 degrees.
#
def circumference_pixels():

    from chimera import viewer as v
    w, h = v.windowSize                 # Window width and height in pixels
    c = v.camera
    from math import pi, tan, atan2
    fov = c.fieldOfView * pi / 180	# Horizontal field of view in radians
    d = 0.5*w/tan(0.5*fov)              # Eye to screen distance in pixels
    vfov = 2*atan2(d,0.5*h)             # Vertical field of view in radians
    rpp = 2*atan2(0.5, d)               # radians per pixel at equator.
    cp = int(round(2*pi/rpp))           # circumference pixels, nearest integer
    return cp

# -------------------------------------------------------------------------------
# Compute eye position or mid-point between stereo eye positions to use for
# rotating cameras or scene.
# 
def camera_center(stereo):

    from chimera import viewer as v
    c = v.camera
    if stereo:
        # Allow stereo left/right eye capture, but rotate about mid-point
        # between eyes.
        c.setMode('stereo left eye')
        le = c.eyePos(0)
        c.setMode('stereo right eye')
        re = c.eyePos(0)
        rc = tuple(0.5*(le[a]+re[a]) for a in (0,1,2))
    else:
        c.setMode('mono')
        rc = c.eyePos(0)
    return rc

# -------------------------------------------------------------------------------
#
def cylinder_images(width, rxf, stereo, scale = 1):

    from chimera import viewer as v

    # TODO: Move lights so they stay fixed with models.

    # Avoid depth cue lighting changes due to clip plane motion.
#    clip = v.clipping
#    v.clipping = True

    ilist = []
    modes = ('stereo left eye', 'stereo right eye') if stereo else ('mono',)
    for m in modes:
        v.camera.setMode(m)
        i = cylinder_image(width, rxf, scale)
        ilist.append(i)

#    v.clipping = clip

    return ilist

# -------------------------------------------------------------------------------
#
def cylinder_image(width, rxf, scale = 1):

    from chimera import openModels as om
    oslist = set(m.openState for m in om.list())

    from chimera import viewer as v
    w,h = v.windowSize
    w *= scale
    h *= scale
    from PIL import Image
    i = Image.new('RGB', (width,h))
    from time import time
    t0 = time()
    for k in range(width):
        ims = v.pilImages(w, h, supersample = 1)
        im = ims[0]
        vline = im.crop((w/2,0,w/2+1,h))
        col = (width/2 + k) % width
        i.paste(vline, (col,0))
        for os in oslist:
            os.globalXform(rxf)
    t1 = time()
    print 'Images per second', width / (t1 - t0)
    
    return i

# -------------------------------------------------------------------------------
# TODO: Add dome tilt parameter.
#
def dome_image(stereo, scale = 1):

    set_vertical_field_of_view(90)
    cp = circumference_pixels() * scale
    rc = camera_center(stereo)

    import Matrix as M
    tilt45tf = M.rotation_transform((1,0,0), 45, rc)
    txf = M.chimera_xform(tilt45tf)
    transform_models(txf)

    rsteptf = M.rotation_transform((0,1,-1), (360.0/cp), rc)
    rxf = M.chimera_xform(rsteptf)
    images = cylinder_images(cp, rxf, stereo, scale)

    transform_models(txf.inverse())

    dimages = [cone_to_dome(i) for i in images]
    return dimages

# -------------------------------------------------------------------------------
#
def set_vertical_field_of_view(degrees):

    from chimera import viewer as v
    w, h = v.windowSize                 # Window width and height in pixels
    from math import pi, tan, atan2
    d = (0.5*h)/tan(0.5 * degrees * pi / 180) # Eye to screen distance in pixels
    hfov = 2*atan2(0.5*w,d)             # Vertical field of view in radians
    v.camera.fieldOfView = hfov * 180 / pi

# -------------------------------------------------------------------------------
#
def transform_models(xf):

    from chimera import openModels as om
    for os in set(m.openState for m in om.list()):
        os.globalXform(xf)

# -------------------------------------------------------------------------------
# Convert rectangular image of unwrapped 90 degree cone to hemisphere image
# projected onto a circle with radius proportional to zenith angle.
#
def cone_to_dome(image):

    w, h = image.size
    from math import pi
    s = int(round(w / pi))
    from PIL import Image
    i = Image.new('RGB', (s,s))
    
    # TODO: Clamp and wrap pixels on cone.
    # TODO: Bilinear interpolate pixels.
    # TODO: Get half-pixel offsets right.
    # TODO: Vectorize calculation using numpy.
    from math import atan2, sqrt, pi, tan
    for di in range(s):
        for dj in range(s):
            x, y = di - 0.5*s, 0.5*s - dj
            r = sqrt(x*x + y*y)
            if r < 0.5*s:
                phi = pi * r / s
                a = (atan2(x,-y) + pi) / (2*pi)
                ci = int(round(w*a))
                cj = int(round(h * 0.5 * (1 + tan(phi - pi/4))))
                ci = (ci + w) % w
                cj = max(0, min(h-1, cj))
                i.putpixel((di,dj), image.getpixel((ci,cj)))

    return i
              
    

# -------------------------------------------------------------------------------
#
#images = panorama(stereo = True)
images = dome_image(stereo = True, scale = 4)
#images = dome_image(stereo = False)
w, h = images[0].size
#images[0].save('dome-%dx%d.jpg' % (w,h), "JPEG")
images[0].save('dome-%dx%d-left.jpg' % (w,h), "JPEG")
images[1].save('dome-%dx%d-right.jpg' % (w,h), "JPEG")
#images[0].save('cone-%dx%d-left.jpg' % (w,h), "JPEG")
#images[1].save('cone-%dx%d-right.jpg' % (w,h), "JPEG")
#images[0].save('cyl-%dx%d-left.jpg' % (w,h), "JPEG")
#images[1].save('cyl-%dx%d-right.jpg' % (w,h), "JPEG")