# @GPL_HEADER_START
# Use of the RAVE and related source code is subject to the terms
# of the following licenses:
# 
# GNU Public License (GPL) Rights pursuant to Version 2, June 1991
# Government Purpose License Rights (GPLR) pursuant to DFARS 252.225-7013
# 
# NO WARRANTY
# 
# ANY INFORMATION, MATERIALS, SERVICES, INTELLECTUAL PROPERTY OR OTHER
# PROPERTY OR RIGHTS GRANTED OR PROVIDED BY CARNEGIE MELLON UNIVERSITY
# PURSUANT TO THIS LICENSE (HEREINAFTER THE "DELIVERABLES") ARE ON AN
# "AS-IS" BASIS. CARNEGIE MELLON UNIVERSITY MAKES NO WARRANTIES OF ANY
# KIND, EITHER EXPRESS OR IMPLIED AS TO ANY MATTER INCLUDING, BUT NOT
# LIMITED TO, WARRANTY OF FITNESS FOR A PARTICULAR PURPOSE,
# MERCHANTABILITY, INFORMATIONAL CONTENT, NONINFRINGEMENT, OR ERROR-FREE
# OPERATION. CARNEGIE MELLON UNIVERSITY SHALL NOT BE LIABLE FOR INDIRECT,
# SPECIAL OR CONSEQUENTIAL DAMAGES, SUCH AS LOSS OF PROFITS OR INABILITY
# TO USE SAID INTELLECTUAL PROPERTY, UNDER THIS LICENSE, REGARDLESS OF
# WHETHER SUCH PARTY WAS AWARE OF THE POSSIBILITY OF SUCH DAMAGES.
# LICENSEE AGREES THAT IT WILL NOT MAKE ANY WARRANTY ON BEHALF OF
# CARNEGIE MELLON UNIVERSITY, EXPRESS OR IMPLIED, TO ANY PERSON
# CONCERNING THE APPLICATION OF OR THE RESULTS TO BE OBTAINED WITH THE
# DELIVERABLES UNDER THIS LICENSE.
# 
# Licensee hereby agrees to defend, indemnify, and hold harmless Carnegie
# Mellon University, its trustees, officers, employees, and agents from
# all claims or demands made against them (and any related losses,
# expenses, or attorney's fees) arising out of, or relating to Licensee's
# and/or its sub licensees' negligent use or willful misuse of or
# negligent conduct or willful misconduct regarding the Software,
# facilities, or other rights or assistance granted by Carnegie Mellon
# University under this License, including, but not limited to, any
# claims of product liability, personal injury, death, damage to
# property, or violation of any laws or regulations.
# 
# Carnegie Mellon University Software Engineering Institute authored
# documents are sponsored by the U.S. Department of Defense under
# Contract F19628-00-C-0003. Carnegie Mellon University retains
# copyrights in all material produced under this contract. The U.S.
# Government retains a non-exclusive, royalty-free license to publish or
# reproduce these documents, or allow others to do so, for U.S.
# Government purposes only pursuant to the copyright license under the
# contract clause at 252.227.7013.
# 
# @GPL_HEADER_END
from rave.plugins.dataset import Dataset
from PIL import Image, ImageDraw, ImageFont, ImageColor
from subprocess import Popen, PIPE
from tempfile import TemporaryFile
import sys, math, os
from rave.plugins.render_mpl import render_blank

default_font = None

def set_font_filename(filename):
    default_font = filename

# Squish 32 bits down to however many by merging entries
def merge_ranges(d, bits, addr_col, value_col):
    non_bits = 32 - bits
    d = d.sort_col(addr_col)
    addr_data = d[addr_col]
    value_data = d[value_col]
    result = Dataset([addr_col, value_col])
    result_addrs = []
    result_values = []
    current_addr = None
    current_value = 0
    for i in xrange(len(addr_data)):
        addr = addr_data[i]
        if isinstance(addr, (str, unicode)):
            addr = str(addr)
            octets = map((lambda x: int(x.strip())), addr.split('.'))
            v = 0
            for o in octets:
                v = v * 256
                v = v + o
            for i in xrange(4 - len(octets)):
                v = v * 256
            addr = v
        addr = (addr >> non_bits) << non_bits
        if addr == current_addr:
            current_value = current_value + value_data[i]
        else:
            if current_addr <> None:
                result.add_row([current_addr, current_value])
            current_addr = addr
            current_value = value_data[i]
    if current_addr <> None:
        result.add_row([current_addr, current_value])
    return result

def render_hilbert_im(addr_data, value_data, overlay,
                      hilbertize_program, bgcolor='#ffffff',
                      label_color='#ff0000', label_font=default_font,
                      min_color='#808080', max_color='#000000',
                      min_value=None, max_value=None,
                      address_bits=24, log_scale=True, no_labels=False,
                      width=None, height=None):
    if width or height:
        if width is None:
            width = height
        if height is None:
            height = width
        address_bits = int(math.ceil(math.log(max(width, height),2)))*2
    d = Dataset(addr=addr_data, val=value_data)
    d = merge_ranges(d, address_bits, 'addr', 'val')
    addr_data = d['addr']
    value_data = d['val']
    height_multiplier = 1
    address_nonbits = 32 - address_bits
    flip_xy = True
    if address_bits % 2 == 1:
        address_bits = address_bits + 1
        height_multiplier = 2
        flip_xy = False
    data_width = 2**(address_bits/2)
    data_height = 2**(address_bits/2)/height_multiplier
    background_color = ImageColor.getrgb(bgcolor)
    label_color = ImageColor.getrgb(label_color)
    min_color = ImageColor.getrgb(min_color)
    max_color = ImageColor.getrgb(max_color)
    # First, send input file to hilbertize and find min/max values
    temp_file = TemporaryFile()
    hilbertize = Popen([hilbertize_program, `address_bits/2`],
                       stdin=PIPE, stdout=temp_file)
    # find the min and max values if not given, and write to temporary file
    computed_min_value = 2**31
    computed_max_value = 0
    for i in xrange(len(addr_data)):
        a = addr_data[i]
        v = value_data[i]
        if v < computed_min_value:
            computed_min_value = v
        if v > computed_max_value:
            computed_max_value = v
        if address_nonbits > 0:
            a = a >> address_nonbits
        hilbertize.stdin.write("%d %d\n" % (a, v))
    hilbertize.stdin.close()
    hilbertize.wait()
    temp_file.seek(0)
    if min_value == None:
        min_value = computed_min_value
    if max_value == None:
        max_value = computed_max_value
    im = Image.new('RGB', (data_width, data_height), background_color)
    im.load()

    # Write overlay over the top
    if overlay:
        im_over = Image.open(overlay).resize(im.size).convert(im.mode)
        im = Image.blend(im_over, im, 0.0)
    
    putpixel = im.im.putpixel
    if max_value == min_value:
        min_value = max_value - 1
    if log_scale:
        log_adjust = 1 - min_value
        min_value = min_value + log_adjust
        max_value = max_value + log_adjust
        min_value = math.log(min_value)
        max_value = math.log(max_value)
    value_scale = max_value - min_value
    r_scale = float(max_color[0] - min_color[0])
    g_scale = float(max_color[1] - min_color[1])
    b_scale = float(max_color[2] - min_color[2])
    for line in temp_file:
        x, y, v = (int(i) for i in line.split())
        if log_scale:
            v = v + log_adjust
            v = math.log(v)
        r = int(min_color[0] + (r_scale * (v - min_value)) / value_scale)
        g = int(min_color[1] + (g_scale * (v - min_value)) / value_scale)
        b = int(min_color[2] + (b_scale * (v - min_value)) / value_scale)
        if flip_xy:
            putpixel((y, x), (r, g, b))
        else:
            putpixel((x, y), (r, g, b))
    temp_file.close()
    if (not no_labels) and (label_font is not None):
        base_font_size = 60
        if address_bits < 20:
            base_font_size = 100
        temp_file = TemporaryFile()
        hilbertize = Popen([hilbertize_program, `address_bits/2`],
                           stdin=PIPE, stdout=temp_file)
        for i in xrange(256):
            i_pos = i << (address_bits - 8)
            i_pos = i_pos + (1 << (address_bits - 9))
            print >>hilbertize.stdin, "%d %d" % (i_pos, i)
        hilbertize.stdin.close()
        hilbertize.wait()
        temp_file.seek(0)
        if flip_xy:
            font_size = base_font_size * data_width / 4096
        else:
            font_size = base_font_size * data_height / 4096
        x_off = 110 * data_width / 4096
        y_off = 128 * data_height * height_multiplier / 4096
        font = ImageFont.truetype(label_font, font_size)
        draw = ImageDraw.Draw(im)
        for line in temp_file:
            (x, y, l) = line.split()
            x = int(x)
            y = int(y)
            if flip_xy:
                draw.text((y-x_off,x-y_off), l, fill=label_color, font=font)
            else:
                draw.text((x-x_off,y-y_off), l, fill=label_color, font=font)
        temp_file.close()
    if width:
        im = im.resize((width, height), Image.BICUBIC)

    return im

def render_hilbert(out_file_name, addr_data, value_data, overlay,
                   hilbertize_program, bgcolor='#ffffff',
                   label_color='#ff0000', label_font=default_font,
                   min_color='#808080', max_color='#000000',
                   min_value=None, max_value=None,
                   address_bits=24, log_scale=True, no_labels=False,
                   width=None, height=None):
    if len(addr_data) == 0:
        render_blank(out_file_name, width or 512, height or 512, bgcolor)
        return
    im = render_hilbert_im(addr_data, value_data, overlay,
                           hilbertize_program, bgcolor, label_color,
                           label_font, min_color, max_color,
                           min_value, max_value, address_bits, log_scale,
                           no_labels, width, height)
    im.save(out_file_name, 'PNG')

__all__ = [ 'render_hilbert', 'render_hilbert_im' ]
