#!/usr/bin/env python
#
# parameter_space.py - black hole merger search
#                    - produce figure showing parameter space
#

# import required modules
import sys
import os
from optparse import OptionParser
import numpy

# import required modules from matplotlib
import matplotlib

from pylab import figure, show

# set matplotlib plot parameters
matplotlib.rcParams.update({
  "font.size": 8.0,
  "axes.titlesize": 10.0,
  "axes.labelsize": 10.0,
  "xtick.labelsize": 8.0,
  "ytick.labelsize": 8.0,
  "legend.fontsize": 8.0,
  })


# method to parse and check command line options
def parse_args():
  # usage information
  usage = "%prog [options]"

  # option parser
  parser = OptionParser(usage)
  parser.add_option("--min-mass", type="float", action="store",
      help="set minimum total black hole mass to MIN_MASS")
  parser.add_option("--max-mass", type="float", action="store",
      help="set maxiumm total black hole mass to MAX_MASS")
  parser.add_option("--output", type="string", action="store",
      help="save plot to file OUTPUT")

  # parse command line options
  options, args = parser.parse_args()

  # check for extra options
  if len(args) != 0:
    parser.error("extraneous command line options")

  # check for required options
  if options.min_mass == None:
    parser.error("--min-mass must be specified")
  if options.max_mass == None:
    parser.error("--max-mass must be specified")
  if options.output == None:
    parser.error("--output must be specified")

  # check for sensible options
  if options.min_mass < 0:
    parser.error("minimum mass must be positive")
  if options.max_mass < 0:
    parser.error("maximum mass must be positive")
  if options.min_mass > options.max_mass:
    parser.error("minimum mass is greater than maximum mass")
  if options.min_mass == options.max_mass:
    parser.error("minimum and maximum mass are equal")

  # return parsed options
  return options

def fill_between(x, y1, y2):
    # add x,y2 in reverse order for proper polygon filling
    return zip(x,y1) + [(thisx, thisy) for thisx, thisy in zip(x, y2)[::-1]]
    
# minimum mass line
def minimum_mass(options, mass_one):
  mass_two = options.min_mass - mass_one
  return mass_two


# maximum mass line
def maximum_mass(options, mass_one):
  mass_two = options.max_mass - mass_one
  return mass_two


# method to produce parameter space plot
def parameter_space_plot(options):
  # setup figure
  fig = figure()
  axes = fig.gca()
  axes.grid(True)
  axes.set_xlabel("Mass 1 / $M_\odot$")
  axes.set_ylabel("Mass 2 / $M_\odot$")

  # setup mass array
  mass = numpy.arange(0, options.max_mass + 1, 1)

  # plot min/max mass lines
  #axes.plot(mass, minimum_mass(options, mass), 'b-')
  #axes.plot(mass, maximum_mass(options, mass), 'b-')

  x, y = mass, minimum_mass(options, mass)
  xs, ys = zip(*fill_between(x, [0]*len(x), y))
  axes.fill(xs, ys, facecolor='red', alpha=0.5)
  print 
  # plot equal ass line
  #axes.plot(mass, mass, 'k--')

  # set axes limits
  axes.set_xlim([0, options.max_mass])
  axes.set_ylim([0, options.max_mass])

  # return plot
  return fig

#
# main program entry point
#

# parse and check command line options
options = parse_args()

# produce and save plot
parameter_space_plot(options).savefig(options.output)

# exit
show()