#!/usr/bin/env python
#
#  Copyright (C) 2007  Smithsonian Astrophysical Observatory
#
#
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 2 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License along
#  with this program; if not, write to the Free Software Foundation, Inc.,
#  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#
#
# Usage:
#    ./filename tgr0 tgd0 offset outfile
#
# Aim:
#   Creates a new copy of the CALDB HRC region file, named outfile,
#   where:
#      The inflection point in-dispersion position, TG_R[1], background         #      values are set to +/- tgr0
#      The inflection point cross-dispersion position, TG_D[1], background 
#      values are set to +/- tgd0
#      The other TG_D background elements are modified by +/- offset
#   (the sign depends on a combination of BACKGROUND_UP/DOWN and tg_m values)
#

import os
import sys

import pycrates as pyc
import numpy as np

def edit_tg_region_from_caldb(outfile, tgr0, tgd0, offset, verbose=False):
    """Edit the CALDB grating region file to adjust the background regions.

    outfile is the name of the new file
    tgr0 is the new value of TG_R[1] (should be positive)
    tgd0 is the new value of TG_D[1] (should be positive)
    offset is the value added to the other TG_D values.

    verbose=True causes some screen messages to be displayed
    """

    if os.getenv("CALDB") == None:
        raise RuntimeError, "Unable to find value of CALDB env. variable"
    cfile = os.getenv("CALDB") + "/data/chandra/hrc/tgmask2/letgD1999-07-22regN0002.fits"
    edit_tg_region_file(cfile, outfile, tgr0, tgd0, offset, verbose)

def edit_tg_region_file(infile, outfile, tgr0, tgd0, offset, verbose=False):
    """Edit the grating region file to adjust the background regions.

    infile should be the name of the region file to edit.
    outfile is the name of the new file
    tgr0 is the new value of TG_R[1] (should be positive)
    tgd0 is the new value of TG_D[1] (should be positive)
    offset is the value added to the other TG_D values.

    verbose=True causes some screen messages to be displayed
    """

    # Read in file
    #
    cr = pyc.read_file(infile)
    if cr == None:
        raise RuntimeError, "Unable to open '%s'" % infile
    if verbose:
        print "Read in data from %s" % infile

    # Check columns exist and minor validation of their size/shapes
    #
    for cname in ["rowid", "tg_r", "tg_d", "tg_m"]:
        if not pyc.col_exists(cr, cname):
            raise RuntimeError, "No column called %s in %s" % (cname,infile)

    rid = pyc.copy_colvals(cr, "rowid")
    tgm = pyc.copy_colvals(cr, "tg_m")

    tgr = pyc.get_colvals(cr, "tg_r")
    tgd = pyc.get_colvals(cr, "tg_d")

    if len(tgr.shape) != 2:
        raise RuntimeError, "Expected TG_R to be 2d, but found %d" % tgr.shape
    if len(tgd.shape) != 2:
        raise RuntimeError, "Expected TG_D to be 2d, but found %d" % tgd.shape

    # Add the offset to the TG_D values
    #
    if verbose:
        print "Adding %g to BACKGROUND_UP/DOWN columns" % offset
    idx = rid == "BACKGROUND_UP"
    tgd[idx,:] += offset

    idx = rid == "BACKGROUND_DOWN"
    tgd[idx,:] -= offset

    # Adjust the TG_R/D[1] values
    #
    if verbose:
        print "Adding %g to TG_R and %g to TG_D of BACKGROUND_UP/DOWN columns" % (tgr0,tgd0)
    idx = (rid == "BACKGROUND_UP") & (tgm == 1)
    tgr[idx,1] = tgr0
    tgd[idx,1] = tgd0

    idx = (rid == "BACKGROUND_UP") & (tgm == -1)
    tgr[idx,1] = -tgr0
    tgd[idx,1] = tgd0

    idx = (rid == "BACKGROUND_DOWN") & (tgm == -1)
    tgr[idx,1] = -tgr0
    tgd[idx,1] = -tgd0

    idx = (rid == "BACKGROUND_DOWN") & (tgm == 1)
    tgr[idx,1] = tgr0
    tgd[idx,1] = -tgd0

    # Write out the file
    #
    if verbose:
        print "Writing out file to '%s'" % outfile
    pyc.write_file(cr, outfile)


def usage(progname):
    print "Usage: %s tgr0 tgd0 offset outfile" % progname
    
def main(argv):                         

    # We assume argv can not be empty here
    #
    if len(argv) != 5:
        usage(argv[0])
        sys.exit(1)

    tgr0 = float(argv[1])
    tgd0 = float(argv[2])
    offset = float(argv[3])
    ofile = argv[4]

    edit_tg_region_from_caldb(ofile, tgr0, tgd0, offset)
    print "Created: %s" % ofile

if __name__ == '__main__':
    main(sys.argv)

#