#!/usr/bin/python2
################################################################################
#                                                                              #
# spice_plot.py Simple python script to read in data from ngspice simulations  #
# and plot them using matplotlib, a very nice plotting library.                #
#                                                                              #
# @author       Brenton Gillis <bgillisb@gmail.com>                            #
# @date         2014-09-07                                                     #
# @version      1.1.0                                                          #
#                                                                              #
################################################################################

# The MIT License (MIT)

# Copyright (c) 2014 Brenton Gillis

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

# ------------------------------------------------------------------------------
# Imports
# ------------------------------------------------------------------------------

import pandas as pd
import numpy as np
import getopt
import sys
import matplotlib.pyplot as plt
from matplotlib.widgets import CheckButtons
from matplotlib.widgets import SpanSelector
from matplotlib.widgets import Button
from matplotlib.ticker import AutoMinorLocator

class Spice_Plot(object):

  # ----------------------------------------------------------------------------
  # Initialize Instance Variables
  # ----------------------------------------------------------------------------
  
  def __init__(self):
    super(Spice_Plot, self).__init__()

    # initialize some variables
    self.YPADDING = 0.15
    self.traces = {}
    self.states = []
    self.state_names = []


  # ----------------------------------------------------------------------------
  # read_data - Read in simulation data
  # ----------------------------------------------------------------------------

  def read_data(self, _filename):

    # read in pandas object and delimit on whitespace
    self.sim_data = pd.read_csv(_filename, delim_whitespace=True)

    # the pandas column names, in this case our point names
    self.point_names = list(self.sim_data.columns.values)

    # how many point names
    self.num_point_names = len(self.point_names)

    # plot the data
    self.plot_data(self.sim_data, self.point_names, self.num_point_names)


  # ----------------------------------------------------------------------------
  # plot_data - plots simulation data
  # ----------------------------------------------------------------------------

  def plot_data(self, sim_data, point_names, num_point_names):

    # initialize min/max variables
    self.min_y = self.max_y = self.min_x = self.max_x = 0
    time = sim_data[point_names[0]]

    # set min/max x axis
    self.current_min_x = self.min_x = min(sim_data[point_names[0]])
    self.current_max_x = self.max_x = max(sim_data[point_names[0]])

    # Create subplots object
    fig, ax = plt.subplots()

    # Create curve for plot based on number of point names
    previous_max = tmp_max = 0
    for i in range(num_point_names):

      # ignore first one because its usually just the time column
      if i != 0:

        # for the rest, create the plot
        self.traces[point_names[i]] = ax.plot(time, sim_data[point_names[i]], 
          label=point_names[i])

        # and determine the max and min for the y axis
        if self.max_y < max(sim_data[point_names[i]]):
          self.max_y = max(sim_data[point_names[i]])
          y = sim_data[point_names[i]]
        if self.min_y > min(sim_data[point_names[i]]):
          self.min_y = min(sim_data[point_names[i]])

        # update the self.states for use in the checkboxes
        self.states.append(True)
        self.state_names.append(point_names[i])

    self.current_max = y.max()
    self.current_min = y.min()

    # set plot axis based off of min and max values. Use a small buffer of 15%
    # to pad with whitespace
    plt.axis([self.min_x, self.max_x, self.min_y + (self.min_y*self.YPADDING), 
      self.max_y+(self.max_y* self.YPADDING)])

    # Add a legend, title, and labels
    plt.legend(bbox_to_anchor=(1, 1), loc=1, borderaxespad=0.1)
    plt.title("NGSpice Plot")
    plt.ylabel("Amplitude (V)")
    plt.xlabel("Time (s)")

    # make the checkboxes
    plt.subplots_adjust(left=0.2)
    checkax = plt.axes([0.03, 0.75, 0.1, 0.15])
    check = CheckButtons(checkax, self.state_names, self.states)

    # --------------------------------------------------------------------------
    # span selector function
    # --------------------------------------------------------------------------

    def onselect(xmin, xmax):
      indmin, indmax = np.searchsorted(time.values, (xmin, xmax))
      indmax = min(len(time)-1, indmax)
      thisx = time.values[indmin:indmax]
      thisy = y[indmin:indmax]
      ax.set_xlim(thisx[0], thisx[-1])
      self.current_min_x = thisx[0]
      self.current_max_x = thisx[-1]
      fig.canvas.draw()

    # set useblit True on gtkagg for enhanced performance
    span = SpanSelector(ax, onselect, 'horizontal', useblit=True,
                        rectprops=dict(alpha=0.5, facecolor='red') )

    # show or hide a trace
    def check_function(label):

      # Using the label, hide or show trace
      self.traces[label][0].set_visible(not self.traces[label][0].get_visible())

      # force redraw
      plt.draw()

    # --------------------------------------------------------------------------
    # scale up button
    # --------------------------------------------------------------------------

    def scale_up(event):

      difference = self.current_max - self.current_min

      # Scale up by YPADDING percentage of the current size
      self.current_min += difference * self.YPADDING
      self.current_max -= difference * self.YPADDING

      # Set the new y-axis and force redraw
      ax.set_ylim(self.current_min, self.current_max)
      plt.draw()

    # --------------------------------------------------------------------------
    # scale down button
    # --------------------------------------------------------------------------

    def scale_down(event):

      difference = self.current_max - self.current_min

      # Scale down by YPADDING percentage of the current size
      self.current_min -= difference * self.YPADDING
      self.current_max += difference * self.YPADDING

      # Set the new y-axis and force redraw
      ax.set_ylim(self.current_min, self.current_max)
      plt.draw()

    # --------------------------------------------------------------------------
    # pan up button
    # --------------------------------------------------------------------------

    def pan_up(event):

      difference = self.current_max - self.current_min

      # pan up by YPADDING percentage of the current size
      self.current_min += difference * self.YPADDING
      self.current_max += difference * self.YPADDING

      # Set the new y-axis and force redraw
      ax.set_ylim(self.current_min, self.current_max)
      plt.draw()

    # --------------------------------------------------------------------------
    # pan down button
    # --------------------------------------------------------------------------

    def pan_down(event):

      difference = self.current_max - self.current_min

      # pan down by YPADDING percentage of the current size
      self.current_min -= difference * self.YPADDING
      self.current_max -= difference * self.YPADDING

      # Set the new y-axis and force redraw
      ax.set_ylim(self.current_min, self.current_max)
      plt.draw()

    # --------------------------------------------------------------------------
    # pan right button
    # --------------------------------------------------------------------------

    def pan_right(event):

      difference = self.current_max_x - self.current_min_x

      # pan right by YPADDING percentage of the current size
      self.current_min_x += difference * self.YPADDING
      self.current_max_x += difference * self.YPADDING

      # Set the new x-axis and force redraw
      ax.set_xlim(self.current_min_x, self.current_max_x)
      plt.draw()

    # --------------------------------------------------------------------------
    # pan left button
    # --------------------------------------------------------------------------

    def pan_left(event):

      difference = self.current_max_x - self.current_min_x

      # pan right by YPADDING percentage of the current size
      self.current_min_x -= difference * self.YPADDING
      self.current_max_x -= difference * self.YPADDING

      # Set the new x-axis and force redraw
      ax.set_xlim(self.current_min_x, self.current_max_x)
      plt.draw()

    # --------------------------------------------------------------------------
    # reset to original state
    # --------------------------------------------------------------------------

    def reset(event):

      # If reset is called, reset the current min/max to also reset the scale
      # functions values
      self.current_min = y.min()
      self.current_max = y.max()

      # reset the x and y axes to the original states
      ax.set_xlim(time.min(), time.max())
      ax.set_ylim(self.current_min + self.current_min * self.YPADDING, 
        self.current_max + self.current_max * self.YPADDING)

      # force redraw
      plt.draw()

    # --------------------------------------------------------------------------
    # grid toggle
    # --------------------------------------------------------------------------

    def grid_toggle(event):

      if self.grid_state:
        ax.grid(b=True, which='major', color='0.5', linestyle='--')
        ax.grid(b=True, which='minor', color='0.75', linestyle='-')
        self.grid_state = False
      else:
        ax.grid(False, which='minor')
        ax.grid(False, which='major')
        self.grid_state = True

    # scan for a check box click
    check.on_clicked(check_function)

    # Add reset button
    ax_reset = plt.axes([0.045, 0.65, 0.07, 0.05])
    b_reset = Button(ax_reset, "Reset")
    b_reset.on_clicked(reset)

    # Add scale_up button
    ax_y_scale_up = plt.axes([0.045, 0.5, 0.04, 0.05])
    b_y_scale_up = Button(ax_y_scale_up, "+Y")
    b_y_scale_up.on_clicked(scale_up)

    # Add scale_down button
    ax_y_scale_down = plt.axes([0.045, 0.43, 0.04, 0.05])
    b_y_scale_down = Button(ax_y_scale_down, "-Y")
    b_y_scale_down.on_clicked(scale_down)

    # Add pan up button
    ax_y_pan_up = plt.axes([0.06, 0.3, 0.04, 0.05])
    b_y_pan_up = Button(ax_y_pan_up, "Up")
    b_y_pan_up.on_clicked(pan_up)

    # Add pan down button
    ax_y_pan_down = plt.axes([0.06, 0.16, 0.04, 0.05])
    b_y_pan_down = Button(ax_y_pan_down, "Dn")
    b_y_pan_down.on_clicked(pan_down)

    # Add pan right button
    ax_x_pan_right = plt.axes([0.09, 0.23, 0.04, 0.05])
    b_x_pan_right = Button(ax_x_pan_right, ">>")
    b_x_pan_right.on_clicked(pan_right)

    # Add pan left button
    ax_x_pan_left = plt.axes([0.03, 0.23, 0.04, 0.05])
    b_x_pan_left = Button(ax_x_pan_left, "<<")
    b_x_pan_left.on_clicked(pan_left)

    # Add grid toggle button
    ax_x_grid_toggle = plt.axes([0.045, 0.58, 0.05, 0.05])
    b_x_grid_toggle = Button(ax_x_grid_toggle, "Grid")
    b_x_grid_toggle.on_clicked(grid_toggle)

    # Finally, show the plots
    ax.yaxis.set_minor_locator(AutoMinorLocator(8))
    ax.xaxis.set_minor_locator(AutoMinorLocator(8))
    # ax.grid(b=True, which='major', color='0.5', linestyle='--')
    # ax.grid(b=True, which='minor', color='0.75', linestyle='-')
    self.grid_state = True
    grid_toggle(None)
    plt.show()

# ------------------------------------------------------------------------------
# getopts usage message
# ------------------------------------------------------------------------------

def get_opts_usage():

  usage = """
  This script will read in ngspice (or gnucap) data and plot using matplotlib.

  USAGE spice_plot.py -f [--file] <filename>

  The plots can be zoomed to observe areas of interest. Use the mouse to click
  and highlight a horizontal region to zoom in on the x-axis. Use the "+Y" and
  "-Y" buttons to scale the y-axis.

  There are four pan buttons: up(Up), right(>>), down(Dn), and left(<<).

  There is also a reset button to return to the default state when the plot
  was first generated.
  """
  print usage

# ------------------------------------------------------------------------------
# Main Function
# ------------------------------------------------------------------------------

def main():

  # attempt parsing and exit if error
  try:
    opts, args = getopt.getopt(sys.argv[1:], "hf:", ["help", "filename="])
    if len(opts) == 0:
      get_opts_usage()
      sys.exit()

  except getopt.GetoptError, err:
    print str(err)
    sys.exit(2)

  # Loop through the arguments
  for opt, arg in opts:

    # display help if need be and exit
    if opt in ("-h", "--help"):
      get_opts_usage()
      sys.exit()

    # read in filename and plot
    elif opt in ("-f", "--filename"):
      filename = arg
      SP = Spice_Plot()
      SP.read_data(filename)
    
    else:
      assert False, "unhandled option"

# Call Directly
if __name__ == '__main__':
  main()

# EOF

