#!/usr/bin/env python

"""Python script to plot tally data generated by OpenMC."""

import os
import sys

from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from matplotlib.backends.backend_tkagg import NavigationToolbar2TkAgg
from matplotlib.figure import Figure
import matplotlib.pyplot as plt
import numpy as np

from openmc.statepoint import StatePoint

if sys.version_info[0] < 3:
    import Tkinter as tk
    import tkFileDialog as filedialog
    import tkFont as font
    import tkMessageBox as messagebox
    import ttk as ttk
else:
    import tkinter as tk
    import tkinter.filedialog as filedialog
    import tkinter.font as font
    import tkinter.messagebox as messagebox
    import tkinter.ttk as ttk


class MeshPlotter(tk.Frame):
    def __init__(self, parent, filename):
        tk.Frame.__init__(self, parent)

        self.labels = {'cell': 'Cell:', 'cellborn': 'Cell born:',
                       'surface': 'Surface:', 'material': 'Material:',
                       'universe': 'Universe:', 'energyin': 'Energy in:',
                       'energyout': 'Energy out:'}

        self.filterBoxes = {}

        # Read data from source or leakage fraction file
        self.get_file_data(filename)

        # Set up top-level window
        top = self.winfo_toplevel()
        top.title('Mesh Tally Plotter: ' + filename)
        top.rowconfigure(0, weight=1)
        top.columnconfigure(0, weight=1)
        self.grid(sticky=tk.W+tk.N)

        # Create widgets and draw to screen
        self.create_widgets()
        self.update()

    def create_widgets(self):
        figureFrame = tk.Frame(self)
        figureFrame.grid(row=0, column=0)

        # Create the Figure and Canvas
        self.dpi = 100
        self.fig = Figure((5.0, 5.0), dpi=self.dpi)
        self.canvas = FigureCanvasTkAgg(self.fig, master=figureFrame)
        self.canvas.get_tk_widget().pack(side=tk.TOP, fill=tk.BOTH, expand=1)

        # Create the navigation toolbar, tied to the canvas
        self.mpl_toolbar = NavigationToolbar2TkAgg(self.canvas, figureFrame)
        self.mpl_toolbar.update()
        self.canvas._tkcanvas.pack(side=tk.TOP, fill=tk.BOTH, expand=1)

        # Create frame for comboboxes
        self.selectFrame = tk.Frame(self)
        self.selectFrame.grid(row=1, column=0, sticky=tk.W+tk.E)

        # Tally selection
        labelTally = tk.Label(self.selectFrame, text='Tally:')
        labelTally.grid(row=0, column=0, sticky=tk.W)
        self.tallyBox = ttk.Combobox(self.selectFrame, state='readonly')
        self.tallyBox['values'] = [self.datafile.tallies[i].id
                                   for i in self.meshTallies]
        self.tallyBox.current(0)
        self.tallyBox.grid(row=0, column=1, sticky=tk.W+tk.E)
        self.tallyBox.bind('<<ComboboxSelected>>', self.update)

        # Planar basis selection
        labelBasis = tk.Label(self.selectFrame, text='Basis:')
        labelBasis.grid(row=1, column=0, sticky=tk.W)
        self.basisBox = ttk.Combobox(self.selectFrame, state='readonly')
        self.basisBox['values'] = ('xy', 'yz', 'xz')
        self.basisBox.current(0)
        self.basisBox.grid(row=1, column=1, sticky=tk.W+tk.E)
        self.basisBox.bind('<<ComboboxSelected>>', self.update)

        # Axial level
        labelAxial = tk.Label(self.selectFrame, text='Axial level:')
        labelAxial.grid(row=2, column=0, sticky=tk.W)
        self.axialBox = ttk.Combobox(self.selectFrame, state='readonly')
        self.axialBox.grid(row=2, column=1, sticky=tk.W+tk.E)
        self.axialBox.bind('<<ComboboxSelected>>', self.redraw)

        # Option for mean/uncertainty
        labelMean = tk.Label(self.selectFrame, text='Mean/Uncertainty:')
        labelMean.grid(row=3, column=0, sticky=tk.W)
        self.meanBox = ttk.Combobox(self.selectFrame, state='readonly')
        self.meanBox['values'] = ('Mean', 'Absolute uncertainty',
                                  'Relative uncertainty')
        self.meanBox.current(0)
        self.meanBox.grid(row=3, column=1, sticky=tk.W+tk.E)
        self.meanBox.bind('<<ComboboxSelected>>', self.update)

        # Scores
        labelScore = tk.Label(self.selectFrame, text='Score:')
        labelScore.grid(row=4, column=0, sticky=tk.W)
        self.scoreBox = ttk.Combobox(self.selectFrame, state='readonly')
        self.scoreBox.grid(row=4, column=1, sticky=tk.W+tk.E)
        self.scoreBox.bind('<<ComboboxSelected>>', self.redraw)

        # Filter label
        boldfont = font.Font(weight='bold')
        labelFilters = tk.Label(self.selectFrame, text='Filters:',
                                font=boldfont)
        labelFilters.grid(row=5, column=0, sticky=tk.W)

    def update(self, event=None):
        if not event:
            widget = None
        else:
            widget = event.widget

        tally_id = self.meshTallies[self.tallyBox.current()]
        selectedTally = self.datafile.tallies[tally_id]

        # Get mesh for selected tally
        self.mesh = selectedTally.filters_by_name['mesh'].mesh

        # Get mesh dimensions
        self.nx, self.ny, self.nz = self.mesh.dimension

        # Repopulate comboboxes baesd on current basis selection
        text = self.basisBox.get()
        if text == 'xy':
            self.axialBox['values'] = [str(i+1) for i in range(self.nz)]
        elif text == 'yz':
            self.axialBox['values'] = [str(i+1) for i in range(self.nx)]
        else:
            self.axialBox['values'] = [str(i+1) for i in range(self.ny)]
        self.axialBox.current(0)

        # If update() was called by a change in the basis combobox, we don't
        # need to repopulate the filters
        if widget == self.basisBox:
            self.redraw()
            return

        # Update scores
        self.scoreBox['values'] = selectedTally.scores
        self.scoreBox.current(0)

        # Remove any filter labels/comboboxes that exist
        for row in range(6, self.selectFrame.grid_size()[1]):
            for w in self.selectFrame.grid_slaves(row=row):
                w.grid_forget()
                w.destroy()

        # create a label/combobox for each filter in selected tally
        count = 0
        for f in selectedTally.filters:
            filterType = f.type
            if filterType == 'mesh':
                continue
            count += 1

            # Create label and combobox for this filter
            label = tk.Label(self.selectFrame, text=self.labels[filterType])
            label.grid(row=count+6, column=0, sticky=tk.W)
            combobox = ttk.Combobox(self.selectFrame, state='readonly')
            self.filterBoxes[filterType] = combobox

            # Set combobox items
            if filterType in ['energyin', 'energyout']:
                combobox['values'] = ['{0} to {1}'.format(*f.bins[i:i+2])
                                      for i in range(f.length)]
            else:
                combobox['values'] = [str(i) for i in f.bins]

            combobox.current(0)
            combobox.grid(row=count+6, column=1, sticky=tk.W+tk.E)
            combobox.bind('<<ComboboxSelected>>', self.redraw)

        # If There are no filters, leave a 'None available' message
        if count == 0:
            count += 1
            label = tk.Label(self.selectFrame, text="None Available")
            label.grid(row=count+6, column=0, sticky=tk.W)

        self.redraw()

    def redraw(self, event=None):
        basis = self.basisBox.current() + 1
        axial_level = self.axialBox.current() + 1
        mbvalue = self.meanBox.get()

        # Get selected tally
        tally_id = self.meshTallies[self.tallyBox.current()]
        selectedTally = self.datafile.tallies[tally_id]

        # Create spec_list
        spec_list = []
        for f in selectedTally.filters:
            if f.type == 'mesh':
                mesh_filter = f
                continue
            index = self.filterBoxes[f.type].current()
            spec_list.append((f, index))

        text = self.basisBox.get()
        if text == 'xy':
            dims = (self.nx, self.ny)
        elif text == 'yz':
            dims = (self.ny, self.nz)
        else:
            dims = (self.nx, self.nz)
        matrix = np.zeros(dims)
        for i in range(dims[0]):
            for j in range(dims[1]):
                if text == 'xy':
                    meshtuple = (i + 1, j + 1, axial_level)
                elif text == 'yz':
                    meshtuple = (axial_level, i + 1, j + 1)
                else:
                    meshtuple = (i + 1, axial_level, j + 1)
                filters, filter_bins = zip(*spec_list + [
                    (mesh_filter, meshtuple)])
                mean = selectedTally.get_value(
                    self.scoreBox.get(), filters, filter_bins)
                stdev = selectedTally.get_value(
                    self.scoreBox.get(), filters, filter_bins,
                    value='std_dev')
                if mbvalue == 'Mean':
                    matrix[i, j] = mean
                elif mbvalue == 'Absolute uncertainty':
                    matrix[i, j] = stdev
                else:
                    if mean > 0.:
                        matrix[i, j] = stdev/mean
                    else:
                        matrix[i, j] = 0.

        # Clear the figure
        self.fig.clear()

        # Make figure, set up color bar
        self.axes = self.fig.add_subplot(111)
        cax = self.axes.imshow(matrix.transpose(), vmin=0.0, vmax=matrix.max(),
                               interpolation='none', origin='lower')
        self.fig.colorbar(cax)

        self.axes.set_xticks([])
        self.axes.set_yticks([])
        self.axes.set_aspect('equal')

        # Draw canvas
        self.canvas.draw()

    def get_file_data(self, filename):
        # Create StatePoint object and read in data
        self.datafile = StatePoint(filename)
        self.datafile.read_results()
        self.datafile.compute_stdev()

        # Find which tallies are mesh tallies
        self.meshTallies = []
        for itally, tally in self.datafile.tallies.iteritems():
            if any([f.type == 'mesh' for f in tally.filters]):
                self.meshTallies.append(itally)
                tally.filters_by_name = {f.type: f for f in tally.filters}

        if not self.meshTallies:
            messagebox.showerror("Invalid StatePoint File",
                                 "File does not contain mesh tallies!")
            sys.exit(1)


if __name__ == '__main__':
    # Hide root window
    root = tk.Tk()
    root.withdraw()

    # If no filename given as command-line argument, open file dialog
    if len(sys.argv) < 2:
        filename = filedialog.askopenfilename(title='Select statepoint file',
                                              initialdir='.')
    else:
        filename = sys.argv[1]

    if filename:
        # Check to make sure file exists
        if not os.path.isfile(filename):
            messagebox.showerror("File not found",
                                 "Could not find regular file: " + filename)
            sys.exit(1)

        app = MeshPlotter(root, filename)
        root.deiconify()
        root.mainloop()