#!/usr/bin/python
#
# flam3cuda, one of a surprisingly large number of ports of the fractal flame
# algorithm to NVIDIA GPUs.
#
# This one is copyright 2010 Steven Robertson <steven@strobe.cc>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 or later
# as published by the Free Software Foundation.

import os
import sys
import argparse
import multiprocessing
from subprocess import Popen
from ctypes import *

import numpy as np
import Image
import scipy
import pycuda.autoinit

import cuburn._pyflam3_hacks
from fr0stlib import pyflam3
from cuburn import render
from cuburn.code.mwc import MWCTest

np.set_printoptions(precision=5, edgeitems=20)

real_stdout = sys.stdout

def fmt_time(time):
    # Format time in a lexically-ordered way that doesn't interfere with the
    # typical case of ascending natural numbers
    atime = abs(time)
    dcml = ('-' if time < 0 else '') + ('%05d' % np.floor(atime))
    frac = np.round((atime - np.floor(atime)) * 1000)
    if frac:
        return '%s_%03d' % (dcml, frac)
    return dcml

def save(args, time, raw):
    noalpha = raw[:,:,:3]
    if args.raw:
        real_stdout.write(buffer(np.uint8(noalpha * 255.0)))
        return

    name = os.path.join(args.dir, '%s_%s' % (args.name, fmt_time(time)))
    img = scipy.misc.toimage(noalpha, cmin=0, cmax=1)
    img.save(name+'.png')

    if args.jpg is not None:
        img.save(name+'.jpg', quality=args.jpg)
    print 'saved', name

def error(msg):
    print "Error:", msg
    sys.exit(1)

def main(args):
    if args.test:
        MWCTest.test_mwc()
        return

    if args.raw:
        sys.stdout = sys.stderr

    genome_ptr, ngenomes = pyflam3.Genome.from_string(args.flame.read())
    genomes = cast(genome_ptr, POINTER(pyflam3.Genome*ngenomes)).contents

    if not np.all([g.width == genomes[0].width for g in genomes]):
        error("Inconsistent width. Force with --width.")
    if not np.all([g.height == genomes[0].height for g in genomes]):
        error("Inconsistent height. Force with --height.")

    if args.qs:
        for g in genomes:
            g.sample_density *= args.qs
            g.ntemporal_samples = max(1, int(g.ntemporal_samples * args.qs))

    if args.width:
        if not args.scale:
            args.scale = float(args.width) / genomes[0].width
        for g in genomes:
            g.width = args.width
    if args.height:
        for g in genomes:
            g.height = args.height
    if args.scale:
        for g in genomes:
            g.pixels_per_unit *= args.scale

    if args.skip and not args.tempscale:
        args.tempscale = float(args.skip)
    if args.tempscale:
        for g in genomes:
            g.temporal_filter_width *= args.tempscale

    if not args.name:
        if args.flame.name == '<stdin>':
            args.name = genomes[0].flame_name
        else:
            args.name = os.path.splitext(os.path.basename(args.flame.name))[0]

    cp_times = [cp.time for cp in genomes]
    if args.renumber:
        for t, g in enumerate(genomes):
            g.time = t - args.renumber
        cp_times = [cp.time for cp in genomes]
    elif np.any(np.diff(cp_times) <= 0):
        error("Genome times are non-monotonic. Try using --renumber.")

    if args.skip:
        if args.start is None:
            args.start = cp_times[0]
        if args.end is None:
            args.end = cp_times[-1]
        times = np.arange(args.start, args.end, args.skip)
    else:
        times = [t for t in cp_times
                 if (args.start is None or t >= args.start)
                 and (args.end is None or t < args.end)]

    render._AnimRenderer.sync = args.sync or args.sleep
    if args.sleep:
        render._AnimRenderer.sleep = args.sleep / 1000.
    anim = render.Animation(genomes)
    if args.debug:
        anim.cmp_options.append('-G')
    anim.keep = args.keep or args.debug
    anim.compile()
    anim.load()

    if args.gfx:
        import pyglet
        window = pyglet.window.Window(anim.features.width, anim.features.height)
        image = pyglet.image.CheckerImagePattern().create_image(
                anim.features.width, anim.features.height)
        label = pyglet.text.Label('Rendering first frame', x=5, y=5,
                font_size=24, bold=True)

        @window.event
        def on_draw():
            window.clear()
            image.texture.blit(0, 0)
            label.draw()

        @window.event
        def on_key_press(sym, mod):
            if sym == pyglet.window.key.Q:
                pyglet.app.exit()

        @window.event
        def on_mouse_motion(x, y, dx, dy):
            pass

        frames = anim.render_frames(times, block=False)
        def poll(dt):
            out = next(frames, False)
            if out is False:
                label.text = "Done. ('q' to quit)"
                pyglet.clock.unschedule(poll)
            elif out is not None:
                time, buf = out
                save(args, time, buf)
                imgbuf = np.uint8(buf.flatten() * 255)
                image.set_data('RGBA', -anim.features.width*4, imgbuf.tostring())
                label.text = '%s %4g' % (args.name, time)

        pyglet.clock.set_fps_limit(30)
        pyglet.clock.schedule_interval(poll, 1/30.)
        pyglet.app.run()

    else:
        for time, out in anim.render_frames(times):
            save(args, time, out)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Render fractal flames.')

    parser.add_argument('flame', metavar='FILE', type=file,
        help="Path to genome file ('-' for stdin)")
    parser.add_argument('-g', action='store_true', dest='gfx',
        help="Show output in OpenGL window")
    parser.add_argument('-j', metavar='QUALITY', nargs='?',
        action='store', type=int, dest='jpg', const=90,
        help="Write .jpg in addition to .png (default quality 90)")
    parser.add_argument('-n', metavar='NAME', type=str, dest='name',
        help="Prefix to use when saving files (default is basename of input)")
    parser.add_argument('-o', metavar='DIR', type=str, dest='dir',
        help="Output directory", default='.')
    parser.add_argument('--raw', action='store_true', dest='raw',
        help="Do not write files; instead, send raw RGBA data to stdout.")

    time = parser.add_argument_group('Sequence options', description="""
        Control which frames are rendered from a genome sequence. If '-k' is
        not given, '-s' and '-e' act as limits, and any control point with a
        time in bounds is rendered at its central time. If '-k' is given,
        a list of times to render is given according to the semantics of
        Python's range operator, as in range(start, end, skip).

        If no options are given, all control points are rendered.""")
    time.add_argument('-s', dest='start', metavar='TIME', type=float,
        help="Start time of image sequence (inclusive)")
    time.add_argument('-e', dest='end', metavar='TIME', type=float,
        help="End time of image sequence (exclusive)")
    time.add_argument('-k', dest='skip', metavar='TIME', type=float,
        help="Skip time between frames in image sequence. Auto-sets "
             "--tempscale, use '--tempscale 1' to override.")
    time.add_argument('--renumber', metavar="TIME", type=float,
        dest='renumber', nargs='?', const=0,
        help="Renumber frame times, counting up from the supplied start time "
             "(default is 0).")

    genome = parser.add_argument_group('Genome options')
    genome.add_argument('--qs', type=float, metavar='SCALE',
        help="Scale quality and number of temporal samples")
    genome.add_argument('--scale', type=float, metavar='SCALE',
        help="Scale pixels per unit (camera zoom)")
    genome.add_argument('--tempscale', type=float, metavar='SCALE',
        help="Scale temporal filter width")
    genome.add_argument('--width', type=int, metavar='PIXELS',
        help="Use this width. Auto-sets scale, use '--scale 1' to override.")
    genome.add_argument('--height', type=int, metavar='PIXELS',
        help="Use this height (does *not* auto-set scale)")


    debug = parser.add_argument_group('Debug options')
    debug.add_argument('--test', action='store_true', dest='test',
        help='Run some internal tests')
    debug.add_argument('--keep', action='store_true', dest='keep',
        help='Keep compilation directory (disables kernel caching)')
    debug.add_argument('--debug', action='store_true', dest='debug',
        help='Compile kernel with debugging enabled (implies --keep)')
    debug.add_argument('--sync', action='store_true', dest='sync',
        help='Use synchronous launches whenever possible')
    debug.add_argument('--sleep', metavar='MSEC', type=int, dest='sleep',
            nargs='?', const='5',
            help='Sleep between invocations. Keeps a single-card system '
                 'usable. Implies --sync.')
    args = parser.parse_args()

    main(args)