#!/usr/bin/python

__prog__ = "plot_inspiral_skymap"

import matplotlib
matplotlib.use('Agg')
import pylab
import math
from pylal import rate
import numpy
from optparse import *
from pylal import git_version
from pylal import webUtils
from pylal import InspiralUtils
import sys

def bin_and_smooth(X, Y, DATA, XYbin):
  # doesn't actually smooth yet
  out = rate.BinnedArray(XYbin)
  for i, d in enumerate(DATA):
    out[Y[i],X[i]] += d
  #out.to_pdf()
  wn = rate.gaussian_window(11, 11, 10)
  rate.filter_array(out.array, wn)
  return out.array
   
def get_2d_RA_DEC_bins(RAbins=600, DECbins=300, minRA=0, maxRA=2.0 * math.pi, minDEC = -math.pi/2.0, maxDEC = math.pi/2.0):
  """
  Given the component mass range low, high of the search it will
  return 2D bins with size bins in each direction
  """
  #print minRA, minDEC
  #Multiply the bin number by 1.5 because we'll smooth later
  RAbin = rate.LinearBins(minRA, maxRA, int(RAbins))
  DECbin = rate.LinearBins(minDEC, maxDEC, int(DECbins))
  twoDRADECBin = rate.NDBins( (DECbin, RAbin) )
  return RAbin, DECbin, twoDRADECBin

def injection(RAbin, DECbin, probs, injRA, injDEC):
  RAix = RAbin[injRA]
  DECix = DECbin[injDEC]
  prob = probs[DECix, RAix]
  ix_above_prob = (probs>=prob).nonzero()
  ix_below_prob = (probs<=prob).nonzero()
  
  # FIXME Check that this is correct, use cos when dec convention changes
  DEC = numpy.outer(DECbin.centres(), numpy.ones(len(RAbin.centres())))
  Area = numpy.ones(probs.shape) * RAbin.delta * DECbin.delta * numpy.cos(DEC)
  print Area[ix_above_prob].sum(), Area[ix_below_prob].sum()
  return Area[ix_above_prob].sum() / Area.sum() # should be 4 pi

def plotSkyMap(map_data_file,ra_res,dec_res,figname,injRA=None,injDEC=None, stat_value=None, mass_1=None, mass_2=None, spin_1_x=None, spin_1_y=None, spin_1_z=None, spin_2_x=None, spin_2_y=None, spin_2_z=None, distance=None ):

  M = numpy.loadtxt(map_data_file)
  RA = M[:,0]; DEC = M[:,1]; PROB = M[:,2]
  maxIX = PROB.argmax()

  # FIXME compute bin boundaries?
  RAbin, DECbin, twoDRADECBin = get_2d_RA_DEC_bins()

  Xi = RAbin.centres()
  Yi = DECbin.centres()
  #FIXME Z has the RA and DEC reversed from what you might expect
  #Z = pylab.griddata(RA, DEC, PROB, Xi, Yi)
  Z = bin_and_smooth(RA, DEC, PROB, twoDRADECBin)
  
  pylab.figure(1)

  pylab.subplot(2,1,1)
  pylab.contourf(RAbin.centres(), DECbin.centres(), Z, 250)
  pylab.xlabel('RA')
  pylab.ylabel('DEC')
  if injRA and injDEC: pylab.plot([injRA],[injDEC],'o',markersize=16,markeredgecolor='red',markerfacecolor='None', markeredgewidth=2,antialiased=True)
  pylab.xlim((min(RA),max(RA)))
  pylab.ylim((min(DEC),max(DEC)))
  pylab.gca().set_aspect(1)
  pylab.subplots_adjust(wspace=0.0,hspace=0.0)

  ofile = open(figname.replace(".png",".txt"),'w')

  if injRA and injDEC:
    injRank = injection(RAbin, DECbin, Z, injRA, injDEC)
    print >>ofile, '#Prob', 'RA', 'DEC, injRA, injDEC, areaRank, stat, mass_1, mass_2, spin_1_x, spin_1_y, spin_1_z, spin_2_x, spin_2_y', 'spin_2_z', 'distance'
    print >>ofile,  max(PROB), RA[maxIX], DEC[maxIX], injRA, injDEC, injRank, stat_value, mass_1, mass_2, spin_1_x, spin_1_y, spin_1_z, spin_2_x, spin_2_y, spin_2_z, distance
    ofile.close()
  else:
    injRank = None
    print >>ofile, '#Prob', 'RA', 'DEC, injRA, injDEC, areaRank, stat, mass_1, mass_2, spin_1_x, spin_1_y, spin_1_z, spin_2_x, spin_2_y', 'spin_2_z', 'distance'
    print >>ofile,  max(PROB), RA[maxIX], DEC[maxIX], injRA, injDEC, injRank, stat_value, mass_1, mass_2, spin_1_x, spin_1_y, spin_1_z, spin_2_x, spin_2_y, spin_2_z, distance
    ofile.close()


##### PRINTS PARAMETERS ONTO IMAGE #####

  if injRank: plttxt ="RA: %.3f DEC: %.3f \nInjRA %.3f InjDEC: %.3f injRank (sqrdeg): %.3f" % (RA[maxIX], DEC[maxIX], injRA, injDEC, injRank*41253)
  else: plttxt ="RA: %.3f DEC: %.3f" % (RA[maxIX], DEC[maxIX])
  pylab.title(plttxt)
 
########################################

  pylab.subplot(2,1,2)
  pylab.contourf(RAbin.centres(), DECbin.centres(), numpy.exp(Z), 250)
  pylab.xlabel('RA')
  pylab.ylabel('DEC')
  if injRA and injDEC: pylab.plot([injRA],[injDEC],'o',markersize=16,markeredgecolor='red',markerfacecolor='None', markeredgewidth=2,antialiased=True)
  pylab.xlim((min(RA),max(RA)))
  pylab.ylim((min(DEC),max(DEC)))
  pylab.gca().set_aspect(1)
  pylab.subplots_adjust(wspace=0.0,hspace=0.0)

  # Save and close we are done
  InspiralUtils.savefig_pylal(figname)
  pylab.close()

# MAIN PROGRAM


usage = """ %prog [options]
"""
parser = OptionParser(usage, version=git_version.verbose_msg)
parser.add_option("-f","--map-data-file",action="store",type="string",\
    metavar=" FILE",help="use map file FILE")
parser.add_option("-r","--ra-res",action="store",type="int",\
    metavar=" RA_RES",help="use gps time GPS")
parser.add_option("-d","--dec-res",action="store",type="int",\
    metavar=" DEC_RES",help="use DEC_RES pixels to display declination")
parser.add_option("-o","--output-path",action="store",type="string",\
    default="", metavar=" PATH",\
    help="path where the figures would be stored")
parser.add_option("-O","--enable-output",action="store_true",\
    default="false",  metavar="OUTPUT",\
    help="enable the generation of the html and cache documents")
parser.add_option("-T","--user-tag", action="store",type="string", \
    default=None, metavar=" USERTAG",help="user tag for the output file name")
parser.add_option("","--ifo-times",action="store",\
    type="string", default=None, metavar=" IFOTIMES",\
    help="provide ifo times for naming figure")
parser.add_option("","--ifo-tag",action="store",\
    type="string",  metavar=" IFOTAG",\
    help="ifo tag gives the information about ifo times and stage")
parser.add_option("","--gps-start-time", action="store",type="float", \
    metavar=" GPSSTARTTIME",help="gps start time (for naming figure and \
    output files)")
parser.add_option("","--gps-end-time", action="store",type="float", \
    metavar=" GPSENDTIME",help="gps end time (for naming figure and \
    output files)")
parser.add_option("-R","--injection-right-ascension",action="store",type="float",\
    metavar=" RA_INJ",help="injection has right asencion RA_INJ")
parser.add_option("-D","--injection-declination",action="store",type="float",\
    metavar=" DEC_INJ",help="injection has declination DEC_INJ")
parser.add_option("-S","--stat-value",action="store",type="float",default=0.0,\
    metavar=" STAT",help="injection has found stat value STAT")
parser.add_option("-z","--injection-mass1",action="store",type="float",\
    metavar=" MASS_1",help="mass of the first star")
parser.add_option("-y","--injection-mass2",action="store",type="float",\
    metavar=" MASS_2 ",help="mass of the second star")

parser.add_option("-x","--injection-spin1x",action="store",type="float",\
    metavar=" INJ_SPIN_1_X ",help="spin of the first star in x direction")
parser.add_option("-w","--injection-spin1y",action="store",type="float",\
    metavar=" INJ_SPIN_1_Y ",help="spin of first star in y direction")
parser.add_option("-e","--injection-spin1z",action="store",type="float",\
    metavar=" INJ_SPIN_1_Z ",help="spin of first star in Z direction")

parser.add_option("-a","--injection-spin2x",action="store",type="float",\
    metavar=" INJ_SPIN_1_X ",help="spin of the second star in x direction")
parser.add_option("-b","--injection-spin2y",action="store",type="float",\
    metavar=" INJ_SPIN_1_Y ",help="spin of second star in y direction")
parser.add_option("-c","--injection-spin2z",action="store",type="float",\
    metavar=" INJ_SPIN_1_Z ",help="spin of second star in Z direction")

parser.add_option("-g","--injection-distance",action="store",type="float",\
    metavar=" DISTANCE ",help="distance to system")


command_line = sys.argv[1:]
(opts,args) = parser.parse_args()

if not opts.output_path or not opts.dec_res or not opts.ra_res or not opts.map_data_file:
  print >> sys.stderr, "invalid command line try --help for usage"
  sys.exit(0)

opts = InspiralUtils.initialise(opts, __prog__, git_version.verbose_msg)
fnameList = []
tagList = []

figname = InspiralUtils.set_figure_name(opts,"skylocation")

plotSkyMap(opts.map_data_file,opts.ra_res,opts.dec_res,figname,opts.injection_right_ascension,opts.injection_declination, opts.stat_value, opts.injection_mass1, opts.injection_mass2, opts.injection_spin1x, opts.injection_spin1y, opts.injection_spin1z, opts.injection_spin2x, opts.injection_spin2y, opts.injection_spin2z, opts.injection_distance)

fnameList.append(figname)
tagList.append("Inspiral sky location probability")

if opts.enable_output:
  html_filename = InspiralUtils.write_html_output(opts, args, fnameList, \
    tagList)
  InspiralUtils.write_cache_output(opts, html_filename, fnameList)
