#!/usr/bin/env python
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter
from createPlotUtils import *

initOptions("Script for plotting MAP-Elites heatmaps.", "[input_directory] [OPTIONS]", "1.0")

#Derived defaults
def def_output_dir(): return base(getStr("config_file")) + ".pdf"
def def_title_font_size(): return getInt("font_size")+4
def def_tick_font_size(): return getInt("font_size")-6

#Directory settings
addPositionalOption("file", nargs="?")
addOption("templates", ".*")
addOption("output_file", def_output_dir, nargs=1)

#General plot settings
addOption("x_label", nargs=1)
addOption("y_label", nargs=1)
addOption("title", nargs=1)
addOption("input_directory", nargs=1)
addOption("map_size", 100, nargs=1)
addOption("fig_size", [8, 6], nargs=2)
addOption("x_scale", 1, nargs=1)
addOption("y_scale", 1, nargs=1)

#Font settings
addOption("has_title", True, nargs=1)
addOption("font_size", 18, nargs=1)
addOption("title_size", def_title_font_size, nargs=1)
addOption("tick_font_size", def_tick_font_size, nargs=1)

#Color map
cdict = {'red': [(0.0,  0.0, 0.0),
                 (0.33, 0.0, 0.0),
                 (0.66,  1.0, 1.0),
                 (1.0,  1.0, 1.0)],
         'blue': [(0.0,  0.0, 0.0),
                  (0.33, 1.0, 1.0),
                  (0.66,  0.0, 0.0),
                  (1.0,  0.0, 0.0)],
         'green': [(0.0,  0.0, 0.0),
                   (0.33, 0.0, 0.0),
                   (0.66,  0.0, 0.0),
                   (1.0,  1.0, 1.0)]}
my_cmap = matplotlib.colors.LinearSegmentedColormap('my_colormap', cdict, 256)


###################
##### CLASSES #####
###################
class DataOfInterest:
    def __init__(self, files):
        self.files = files
        self.data = None

    def get_files(self):
        return self.files

    def get_data(self):
        if not self.data:
            self.init_data()
        return self.data

    def init_data(self):
        size = getInt("map_size")
        self.data = np.zeros((size, size))
        for file_name in self.files:
            print "Reading:", file_name
            try:
                x, y, z = np.loadtxt(file_name).T
            except ValueError:
                continue
            for i in range(0, len(z)):
                x_index = max(min(round(x[i] * size), size-1), 0)
                y_index = max(min(round(y[i] * size), size-1), 0)
                current_value = self.data[x_index, y_index]
                if(z[i] > current_value):
                    self.data[x_index, y_index] = z[i]


######################
# PLOTTING FUNCTIONS #
######################
def create_plot(data_of_interest):
    data = data_of_interest.get_data()
    data = np.ma.masked_where(data == 0, data)
    im = plt.imshow(data.T, origin='lower', cmap=my_cmap)
    im.set_interpolation('nearest')


######################
### CONFIGURE PLOTS ##
######################
def scale(x, pos):
    'The two args are the value and tick position'
    y_scale = getFloat("y_scale")
    return '%1.1f' % ((x*y_scale) / 100.0)

def scale2(x, pos):
    'The two args are the value and tick position'
    x_scale = getFloat("x_scale")
    return '%1.1f' % ((x*x_scale) / 100.0)

def setup_plot():
    """A setup for the different plots (both the main plot and the small bar at the bottom)."""

    #Setup the matplotlib params
    preamble=[r'\usepackage[T1]{fontenc}',
              r'\usepackage{amsmath}',
              r'\usepackage{txfonts}',
              r'\usepackage{textcomp}']
    matplotlib.rc('font', **{'family':'sans-serif', 'sans-serif':['Helvetica']})
    matplotlib.rc('text.latex', preamble=preamble)
    params = {'backend': 'pdf',
              'axes.labelsize': getInt("font_size"),
              'text.fontsize': getInt("font_size"),
              'xtick.labelsize': getInt("tick_font_size"),
              'ytick.labelsize': getInt("tick_font_size"),
              'text.usetex': latex_available()}
    matplotlib.rcParams.update(params)

    #Set all labels and limits for the main window
    fig = plt.figure(figsize=getFloatPair("fig_size"))
    
    ax = plt.gca()
    ax.yaxis.set_major_formatter(FuncFormatter(scale))
    ax.xaxis.set_major_formatter(FuncFormatter(scale2))
    ax.set_ylabel(getStr("y_label"))
    ax.set_xlabel(getStr("x_label"))
    if getBool("has_title"):
        plt.title(getStr("title"), fontsize=getInt("title_size"))


def write_plot():
    print("Writing plot..")
    output_file = getStr("output_file")
    plt.colorbar()
    plt.tight_layout()
    plt.savefig(output_file)


######################
#### PARSE OPTIONS ###
######################
def parse_options():
    parse_global_options()
    if getExists("file"):
        files = getList("file")
    else:
        templates = getList("templates")
        input_directory = getStr("input_directory")
        files = get_files(templates, input_directory)
    data_of_interest = DataOfInterest(files)
    return data_of_interest



######################
######## MAIN ########
######################
def main():
    #Parse options
    data_of_interest = parse_options()

    #Setup plot
    setup_plot()

    #Plot all treatments
    create_plot(data_of_interest)

    #Writes the plots to disk
    write_plot()


if __name__ == '__main__':
    main()
