import h5py
import pylab
import numpy
import matplotlib
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.figure import Figure

def get_min_max(filenames, fieldname, log=False, abs=False):
    def get_file_min_max(filename, fieldname, log=False, abs=False):
        Hfile = h5py.File(filename)
        field = numpy.array(Hfile['/'+fieldname])
        Hfile.close()
        if abs:
            field = numpy.abs(field)
        if log:
            field = numpy.log10(field)
        min = field.min()
        max = field.max()
        return min, max
    
    min, max = get_file_min_max(filenames[0], fieldname, log, abs)

    for filename in filenames[1:]:
        tmp_min, tmp_max = get_file_min_max(filename, fieldname, log, abs)
        min = numpy.min(tmp_min, min)
        max = numpy.max(tmp_max, max)

    return min, max

def imsave(filename, X, **kwargs):
    """ Homebrewed imsave to have nice colors... """
    figsize=(numpy.array(X.shape)/100.0)[::-1]
    pylab.rcParams.update({'figure.figsize':figsize})
    fig = Figure(figsize=figsize)
    canvas = FigureCanvas(fig)
    ax = fig.add_subplot(111)
    pylab.axes([0,0,1,1]) # Make the plot occupy the whole canvas
    pylab.axis('off')
    fig.set_size_inches(figsize)
    pylab.imshow(X,origin='lower', **kwargs)
    pylab.savefig(filename, facecolor='black', edgecolor='black', dpi=100)

def plot_field(filename,  fieldname, image_name, log=False, abs=False, min=None, max=None, cmap=pylab.cm.hot):
    Hfile = h5py.File(filename)
    field = numpy.array(Hfile['/'+fieldname])
    Hfile.close()
    if abs:
        field = numpy.abs(field)
    if log:
        field = numpy.log10(field)
    field_p = field.reshape((field.shape[1], field.shape[0]))
    imsave(image_name, field_p, vmin=min, vmax=max, cmap=cmap)

def plot_files(filenames, fieldname, options):
    if options.debug:
        print "Plotting %s field using color map %s" % (fieldname, options.cmap)
        if options.log:
            print "Using a log10 scale"
        if options.abs:
            print "Using the absolute values"

    min = None
    max = None

    if options.min == None or options.max == None:
        min, max = get_min_max(filenames, fieldname, log=options.log)
    if options.min != None:
        if options.log:
            min = numpy.log10(options.min)
        else:
            min = options.min
    if options.max != None:
        if options.log:
            max = numpy.log10(options.max)
        else:
            max = options.max
    if options.debug:
        print "Range to be plotted = (%8.6e, %8.6e)" % (min, max)

    if options.cmap == 'redshift':
        cdict = {'red': ((0.0, 1.0, 1.0),
                         (0.1, 0.0, 0.0),
                         (0.5, 0.0, 0.0),
                         (0.9, 1.0, 1.0),
                         (1.0, 1.0, 1.0)),
                 'green': ((0.0, 1.0, 1.0),
                           (0.1, 0.0, 0.0),
                           (0.5, 0.0, 0.0),
                           (0.9, 0.0, 0.0),
                           (1.0, 1.0, 1.0)),
                 'blue': ((0.0, 1.0, 1.0),
                         (0.1, 1.0, 1.0),
                          (0.5, 0.0, 0.0),
                          (0.9, 0.0, 0.0),
                          (1.0, 1.0, 1.0))}
        cmap = matplotlib.colors.LinearSegmentedColormap('my_colormap',cdict,256)
    else:
        cmap = eval("pylab.cm." + options.cmap)

    idx = 0
    for filename in filenames:
        if options.outputname:
            imagename = '%s.%d.png' % (options.outputname, idx)
        else:
            if filename[-2:] == 'h5':
                imagename = filename[:-3]+'.'+fieldname+'.png'
            elif filename[-3:] == 'hdf':
                imagename = filename[:-4]+'.'+fieldname+'.png'
            else:
                imagename = filename+'.'+fieldname+'.png'

        if options.debug:
                print "Plotting %s in %s to %s" % (fieldname, filename, imagename)
        plot_field(filename, fieldname, imagename,
                   log=options.log, abs=options.abs, min=min, max=max, cmap=cmap)
        idx += 1