#!/usr/bin/python
#  Copyright (C) 2012 by Carnegie Mellon University.
#  
#  @OPENSOURCE_HEADER_START@
#  Use of Rayon 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.227.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 FA8721-05-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.
#  @OPENSOURCE_HEADER_END@

from __future__ import division

import unittest
import sys
import math

from rayon.rytools import *
from rayon import plots, toolbox
from rayon import RayonException

# Commentary:
#
#    *sigh*


import datetime
def debug(msg):
    if False:
        tstamp = datetime.datetime.now()
        print "%s: %s" % (tstamp.strftime("%H:%M:%S"), msg)



def mk_proportional_scale(in_col):
    total = sum(in_col)
    if total == 0:
        return lambda in_val: 0.0
    else:
        return lambda in_val: in_val / total

# A set of ten palettes to use as the default. (Yeah, I could use an
# array, but conceptually this is a mapping of a number of colors to a
# palette. Don't judge me.)
palettes = {
    0 : ((0, 0, 0, 0),),
    1 : ((0, 149, 201, 255),),
    2 : ((0, 149, 201, 255), (255, 84, 31, 255)),
    3 : ((0, 149, 201, 255), (255, 184, 61, 255), (174, 33, 73, 255)),
    4 : ((0, 149, 201, 255), (255, 84, 31, 255),
         (134, 42, 168, 255), (255, 248, 96, 255)),
    5 : ((0, 149, 201, 255), (255, 84, 31, 255),
         (48, 38, 157, 255), (255, 184, 61, 255),
         (174, 33, 73, 255)),
    6 : ((0, 149, 201, 255), (255, 84, 31, 255), (48, 38, 157, 255),
         (255, 184, 61, 255), (174, 33, 73, 255), (212, 229, 86, 255)),
    7 : ((0, 149, 201, 255), (255, 84, 31, 255), (48, 38, 157, 255),
         (255, 184, 61, 255), (174, 33, 73, 255), (212, 229, 86, 255),
         (134, 42, 168, 255)),
    8 : ((0, 149, 201, 255), (255, 84, 31, 255),
         (48, 38, 157, 255), (255, 184, 61, 255),
         (174, 33, 73, 255), (212, 229, 86, 255),
         (134, 42, 168, 255), (255, 248, 96, 255)),
    9 : ((0, 149, 201, 255), (255, 84, 31, 255), (48, 38, 157, 255),
         (255, 184, 61, 255), (174, 33, 73, 255), (212, 229, 86, 255),
         (134, 42, 168, 255), (255, 248, 96, 255), (0, 93, 244, 255)),
    10: ((0, 149, 201, 255), (255, 84, 31, 255), (48, 38, 157, 255),
         (255, 184, 61, 255), (174, 33, 73, 255), (212, 229, 86, 255),
         (134, 42, 168, 255), (255, 248, 96, 255), (0, 93, 244, 255),
         (255, 150, 50, 255))
}
         
    

def mk_color_scale(vals, palette=None):
    def recycle(num, p):
        ct = 0
        while ct < num:
            for color in p:
                yield color
                ct += 1
                
    if palette is None:
        if len(vals) == 0:
            # Huh? Um, okay...(TODO: I may want to raise an error
            # at this point, not sure.)
            return lambda in_val: (0, 0, 0, 0)
        try:
            palette = palettes[len(vals)]
        except KeyError:
            # More than 10 values. Recycle as necessary
            palette = list(recycle(len(vals), palettes[10]))
    else:
        # Ensure palette has enough colors
        palette = list(recycle(len(vals), palette))

    rlookup = dict((k, v) for v, k in enumerate(vals))
    
    def s(in_val):
        return palette[rlookup[in_val]]
    return s
        

class PiePlot(plots.Plot):
    axes = ("color", "value", "stroke_width", "stroke_color")
    default_scales = {}
    def draw_(self, ctx, width, height):
        center_x = width / 2
        center_y = height / 2
        radius = min(width, height) / 2

        def perimeter_point_at(degrees):
            # Given an angle, return a point on the perimeter of the
            # pie chart corresponding to that angle, in device space
            # coordinates.
            theta = math.radians(degrees)
            return ((radius * math.cos(theta)) + center_x,
                    height - ((radius * math.sin(theta)) + center_y))
        

        def pie_piece(a_start, a_end, fill, s_width, s_color):
            # Draw a pie piece for the angle starting at a_start and
            # ending at a_end.
            pt1 = perimeter_point_at(a_start)
            pt2 = perimeter_point_at(a_end)
            debug("a_start: %s degrees (%s*pi radians)" %
                  (a_start, math.radians(a_start)/math.pi))
            debug("a_end: %s degrees (%s*pi radians)" %
                  (a_end, math.radians(a_end)/math.pi))
            debug("pt1: %s" % str(pt1))
            debug("pt2: %s" % str(pt2))
            debug("ctr: %s" % str((center_x, center_y)))
            segments = [
                ("line", (center_x, center_y)),
                ("line", pt1),
                #("line", pt2)
                ("arc", (center_x, center_y, radius, a_start, a_end, True))
                ]
            ctx.draw_shape(segments,
                           pt2,
                           stroke_width,
                           stroke_color,
                           fill)

        acc = 0 # degrees
        for (color,
             value,
             stroke_width,
             stroke_color) in self.get_scaled_points():
            angle = value * 360
            angle_end = angle + acc
            pie_piece(acc, angle_end, color, stroke_width, stroke_color)
            acc = angle_end

class ColorLegend(plots.Plot):
    axes = ("names","text_attrs", "stroke_color")
    default_scales = {}
    def draw_(self, ctx, width, height):
        # This works a little differently. We don't want to occupy all
        # available space; we want to pack ourselves efficiently
        # within this space.
        #
        # (No clipping. If there is to be clipping, the Chart and/or
        # Plot should do it. There is currently no clipping.)
        y = 0
        canvas_w, canvas_h = ctx.canvas_dimensions()
        def get_lbl_height(lbl, name):
            (_, tl_y), (_, bl_y), _, _, _ = lbl.get_bounding_box(ctx, name)
            return bl_y - tl_y
        
        lbl_height = max(get_lbl_height(lbl, name)
                         for ((_, lbl, _), (name, _, _))
                              in self.get_all_points())
        whitespace = lbl_height * .25
        rect_height = lbl_height + whitespace
        rect_width = lbl_height * 1.62
        line_height = rect_height + (rect_height * .25)
        start_of_text = (rect_height - lbl_height) * .5
        
        for ((color, lbl, stroke), (name, _, _)) in self.get_all_points():
            ctx.draw_rectangle(0, y + rect_height,
                               rect_width, rect_height, 1, stroke, color)
            lbl.draw(ctx, rect_width + 5, y + (rect_height * .5),
                     canvas_w, canvas_h, name)
            y += line_height
    

def extract_important_data(in_data, opts):
    if in_data.get_num_columns() < 2:
        raise RuntimeError("Too few columns in data to generate chart")
    
    cat_input, val_input, color_input = tuple(
        opts[x] for x in ('cat_input', 'val_input', 'color_input'))

    if cat_input is None and val_input is None:
        cat_input = 0
        val_input = 1
    elif cat_input is None or val_input is None:
        # ...meaning one of them *isn't* None
        raise RuntimeError("Must specify both cat-input and val-input")

    # Coerce categories to string, every time
    cat_col = tuple(str(s) for s in in_data.get_column(cat_input))
    val_col = in_data.get_column(val_input)

    if cat_col == val_col:
        raise RuntimeError("Duplicate inputs: cat-input and val-input")

    if color_input is None:
        color_col = None
    else:
        color_col = in_data.get_column(color_input)
        if color_col == cat_col:
            raise RuntimeError(
                "Duplicate inputs: cat-input and color-input")
        elif color_col == val_col:
            raise RuntimeError(
                "Duplicate inputs: color-input and val-input")

    return cat_col, val_col, color_col
        
            


    



def runtests():
    suite = unittest.TestLoader().loadTestsFromModule(sys.modules[__name__])
    unittest.TextTestRunner(verbosity=2).run(suite)

def main(opts):
    if opts['run_tests']:
        runtests()
        return 0

    tools = toolbox.Toolbox.for_file()

    input_data = tools.new_dataset_from_stream(
        istream_from_str(opts['input_path']),
        first_line_is_colnames=opts['first_line_colnames'])

    cat_col, val_col, color_col = extract_important_data(input_data, opts)

    color_scale = mk_color_scale(cat_col, color_col)
    prop_scale = mk_proportional_scale(val_col)
    
    

    chart = tools.new_chart('tiled_adv', col_weights=[3,1])
    
    pieplot = PiePlot(
        color_data=cat_col,
        color_scale=color_scale,
        value_data=val_col,
        value_scale=prop_scale,
        stroke_width_scale=1,
        stroke_color_scale=(0, 0, 0, 255))

    chart.add_plot(pieplot)
    
    legend = ColorLegend(
        names_data=cat_col,
        names_scale=color_scale,
        text_attrs_scale=tools.new_labeler(font_size="large", valign="center"),
        stroke_color_scale=(0, 0, 0, 255))

    chart.add_plot(legend, tpad=.1)

    if opts['title'] is not None:
        border = tools.new_border('hlabel',
                                  label=opts['title'],
                                  font_size="x-large")
        chart.add_top_border(border, .15)

    chart.set_padding(allpad="15px")
    chart.set_chart_background(opts['background_color'])
    page = tools.new_page_from_filename(
        opts['output_path'], opts['width'], opts['height'])
    page.write(chart)

    
    return 0

# Tests

class t_mk_color_scale(unittest.TestCase):
    red = (255, 51, 102, 255)
    orange = (205, 102, 51, 255)
    yellow = (205, 204, 51, 255)
    green = (102, 255, 51, 255)
    blue = (0, 61, 245, 255)
    indigo = (51, 102, 255, 255)
    violet = (138, 0, 184, 255)
    rainbow = (red, orange, yellow, green, blue, indigo, violet)
    
    def test_custom_palette_1(self):
        vals = [1,2,3,4,5]
        s = mk_color_scale(vals, self.rainbow[:5])
        for v, c in zip(vals, self.rainbow[:5]):
            self.assertEquals(s(v), c)

    def test_custom_palette_2(self):
        # colors recycle if len(palette) < len(values)
        vals = xrange(20)
        s = mk_color_scale(vals, self.rainbow)
        for v in vals:
            self.assertEquals(s(v), self.rainbow[v % len(self.rainbow)])

    def test_default_palette_1(self):
        vals = [1,2,3,4,5]
        s = mk_color_scale(vals)
        for v, c in zip(vals, palettes[5]):
            self.assertEquals(s(v), c)

    def test_empty_vals_1(self):
        s = mk_color_scale([])
        self.assertEquals(s(123), (0, 0, 0, 0))

    def test_10_vals_1(self):
        vals = list(xrange(10))
        s = mk_color_scale(vals)
        for v in vals:
            self.assertEquals(s(v), palettes[10][v])

    def test_11plus_vals_1(self):
        # values recycle after 
        vals = list(xrange(20))
        s = mk_color_scale(vals)
        for v in vals:
            self.assertEquals(s(v), palettes[10][v % 10])
        

class t_mk_proportional_scale(unittest.TestCase):
    def test_1(self):
        s = mk_proportional_scale([10, 20, 30, 40, 50])
        self.assertAlmostEquals(s(10), 10 / 150)

    def test_exceeds_proportion_1(self):
        # This scale happily handles values larger than its
        # proportion. (Doesn't need to, but it does.)
        s = mk_proportional_scale([10, 20, 30, 40, 50])
        self.assertAlmostEquals(s(300), 2.0)

    def test_negative_proportion_1(self):
        s = mk_proportional_scale([10, 20, 30, 40, 50])
        self.assertAlmostEquals(s(-10), -10 / 150)

    def test_non_elements_1(self):
        # This scale also processes elements that weren't in the input
        # set.
        s = mk_proportional_scale([10, 20, 30, 40, 50])
        self.assertAlmostEquals(s(12), 12 / 150)

    def test_zero_elements_1(self):
        s = mk_proportional_scale([])
        self.assertEquals(s(123), 0.0)

import rayon.data

def eid_succeeds(cols_in, opts_in, cols_out):
    opts = dict(zip(('cat_input', 'val_input', 'color_input'), opts_in))
    in_data = rayon.data.Dataset.from_columns(cols_in)
    def test(self):
        rc = extract_important_data(in_data, opts)
        for expected, actual in zip(cols_out, rc):
            if expected is None:
                self.assertEqual(expected, actual)
            else:
                self.assertEqual(rayon.data.Column(expected), actual)
    return test

def eid_fails(cols_in, opts_in, err=RuntimeError):
    opts = dict(zip(('cat_input', 'val_input', 'color_input'), opts_in))
    in_data = rayon.data.Dataset.from_columns(cols_in)
    def test(self):
        self.assertRaises(
            err,
            lambda: extract_important_data(in_data, opts))
    return test

                    
class t_extract_important_data(unittest.TestCase):
    cats = ['a', 'b', 'c']
    vals = [10, 20, 30]
    colors = [(255, 0, 0, 255),
              (0, 255, 0, 255),
              (0, 0, 255, 255)]

    # The rules:
    # 1 column: error
    # 2 column: must be cat and val
    #   * default: cat/val
    #   * If you specify one of cat/value, you must specify the other
    #   * Color is default
    # 3 column:
    #   * default: cat/val/color
    #   * If you specify one of cat/value, you must specify the other
    #   * If you specify cat/value, color is the remaining column
    # 4+ columns:
    #   * no default, specify everything
    
    def test_1column_1(self):
        opts = {'cat_input': None, 'val_input': None}
        in_data = rayon.data.Dataset.from_columns([['a', 'b', 'c']])
        def bad():
            extract_important_data(in_data, opts)
        self.assertRaises(RuntimeError, bad)

    # 2 columns:
        
    #  -- no columns specified
    test_2col_1 = eid_succeeds(
        # Input columns
        [cats, vals],
        # cat, val, color options (from command-line)
        [None, None, None],
        # expected output columns (as lists, or None if not specified)
        [cats, vals, None])

    # -- one column specified
    #                       Input cols  | options
    test_2col_2 = eid_fails([cats, vals], [0, None, None])
    test_2col_3 = eid_fails([cats, vals], [None, 1, None])

    # -- both columns specified
    test_2col_4 = eid_succeeds([cats, vals],
                               [0, 1, None],
                               [cats, vals, None])
    test_2col_5 = eid_succeeds([vals, cats],
                               [1, 0, None],
                               [cats, vals, None])

    # -- both columns specified, same column (ERROR)
    test_2col_6 = eid_fails([cats, vals], [0, 0, None])
    test_2col_7 = eid_fails([cats, vals], [1, 1, None])

    # -- invalid column index (ERROR)
    test_2col_7 = eid_fails([cats, vals], [0, 2, None], err=RayonException)
    test_2col_8 = eid_fails([cats, vals], [2, 0, None], err=RayonException)

    # # 3 columns:

    # -- no columns specified (default; notice, color is not
    #    automatically picked up)
    test_3col_1 = eid_succeeds([cats, vals, colors],
                               [None, None, None],
                               [cats, vals, None])

    # -- all columns specified
    test_3col_2 = eid_succeeds([cats, colors, vals],
                               [0, 2, 1],
                               [cats, vals, colors])

    # -- cat and value specified
    test_3col_3 = eid_succeeds([cats, colors, vals],
                               [0, 2, None],
                               [cats, vals, None])

    # -- only one of cat/val specified (ERROR)
    test_3col_4 = eid_fails([cats, vals, colors], [0, None, None])
    test_3col_5 = eid_fails([cats, vals, colors], [None, 1, None])

    # -- only color specified
    test_3col_6 = eid_succeeds([cats, vals, colors],
                               [None, None, 2],
                               [cats, vals, colors])

    # -- same column specified twice (or three times) (ERROR)
    test_3col_7 = eid_fails([cats, vals, colors], [0, 0, 1])
    test_3col_8 = eid_fails([cats, vals, colors], [0, 1, 0])
    test_3col_9 = eid_fails([cats, vals, colors], [1, 0, 0])

    def test_3col_10(self):
        # If we use names for cat and val, the color column inference
        # thing still works
        in_data = rayon.data.Dataset.from_columns(
            [self.colors, self.vals, self.cats],
            colnames=('a', 'b', 'c'))
        opts = {'cat_input': 'c', 'val_input': 'b', 'color_input' : None}
        cat_col, val_col, color_col = extract_important_data(in_data, opts)
        self.assertEqual(cat_col, rayon.data.Column(self.cats))
        self.assertEqual(val_col, rayon.data.Column(self.vals))
        self.assertEqual(color_col, None)

        
# Option parsing and execution

desc = "Visualize proportional data using a pie chart"

options = [
    Flag('run-tests'),
    ColnameOption(
        'cat-input', default_val=None,
        hlp="Name or index (from left, starting at 0) of "
        "column containing category information (Default: 0)"),
    ColnameOption(
        'val-input', default_val=None,
        hlp="Name or index (from left, starting at 0) of "
        "column containing value information (Default: 1)"),
    ColnameOption(
        'color-input', default_val=None,
        hlp="Name or index (from left, starting at 0) of "
        "column containing value information (Default: auto-select)"),
    StringOption(
        'title', default_val=None,
        hlp="Title of chart"),
    ColorOption(
        'background-color', default_val=(255, 255, 255, 255),
        hlp="Color of chart background"),
]

if __name__ == '__main__':
    execute(cmd_func=main, cmd_options=options, cmd_desc=desc)

