#!/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 matplotlib import figure
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas

# 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,
  "figure.dpi": 300,
  "savefig.dpi": 300,
  "text.usetex": True
  })


# 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


# 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.Figure()
  FigureCanvas(fig)
  fig.set_size_inches(5, 5)
  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-')

  axes.fill(mass, minimum_mass(options, mass), facecolor='red', alpha=0.5)

  # plot equal mass 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
sys.exit(0)