#!/usr/bin/python
# -*- coding: utf-8 -*-

#
#% $Id$
#
#
# Copyright (C) 2002-2011
# The MeqTree Foundation &
# ASTRON (Netherlands Foundation for Research in Astronomy)
# P.O.Box 2, 7990 AA Dwingeloo, The Netherlands
#
# 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, see <http://www.gnu.org/licenses/>,
# or write to the Free Software Foundation, Inc.,
# 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#

import sys
import pyfits
import re
import os.path
import pyfits
import math
import numpy
import traceback
import fnmatch

DEG = math.pi/180;

NATIVE = "Tigger";

if __name__ == '__main__':
  import Kittens.utils
  from Kittens.utils import curry
  _verbosity = Kittens.utils.verbosity(name="convert-model");
  dprint = _verbosity.dprint;
  dprintf = _verbosity.dprintf;

  # find Tigger
  try:
    import Tigger
  except ImportError:
    dirname = os.path.dirname(os.path.realpath(__file__));
    # go up the directory tree looking for directory "Tigger"
    while len(dirname) > 1:
      if os.path.basename(dirname) == "Tigger":
        break;
      dirname = os.path.dirname(dirname);
    else:
      print "Unable to locate the Tigger directory, it is not a parent of %s. Please check your installation and/or PYTHONPATH."%os.path.realpath(__file__);
      sys.exit(1);
    sys.path.append(os.path.dirname(dirname));
    try:
      import Tigger
    except:
      print "Unable to import the Tigger package from %s. Please check your installation and PYTHONPATH."%dirname;
      sys.exit(1);

  Tigger.nuke_matplotlib();  # don't let the door hit you in the ass, sucka

  # setup some standard command-line option parsing
  #
  from optparse import OptionParser
  parser = OptionParser(usage="""%prog: sky_model [NAME or SELTAG<>SELVAL] [TAG=[TYPE:]VALUE or +TAG or !TAG or /TAG ...]""",
                        description=
"""Sets or changes tags of selected sources in the sky model.
Use NAME (with shell-style wildcards allowed) to select sources by name, or 
=SELTAG to select sources having the specified (non-zero) tag, or SELTAG<>SELVAL to
select sources by comparing a tag to a value, where '<>'  represents a comparison
operator, and can be one of == (or =),!=,<=,<,>,>= (or the FORTRAN-style
operators .eq.,.ne.,.le.,.lt.,.gt.,.ge.). SELVAL may also be followed by one of the characters 
'd', 'm' or 's', in which case it will be converted from degrees,
minutes or seconds into radians. This is useful for selections such as "r<5d".
Then, with a subset of sources selected, use TAG=TYPE:VALUE (where TYPE is one of: bool, int, float, str, complex)
to set a tag on the selected sources to a value of a specific type, or TAG=VALUE to determine type
automatically, or +TAG to set a bool True tag, !TAG to set a False tag, and /TAG to remove a tag."""
);

  parser.add_option("-l","--list",action="store_true",
                    help="Simply lists selected sources, does not apply any tags.");
  parser.add_option("-o","--output",metavar="FILENAME",type="string",
                    help="Saves changes to different output model. Default is to save in place.");
  parser.add_option("-f","--force",action="store_true",
                    help="Saves changes to model without prompting. Default is to prompt.");
  parser.add_option("-d", "--debug",dest="verbose",type="string",action="append",metavar="Context=Level",
                    help="(for debugging Python code) sets verbosity level of the named Python context. May be used multiple times.");

  parser.set_defaults();

  (options,rem_args) = parser.parse_args();

  # get filenames
  if len(rem_args) < 2:
    parser.error("Incorrect number of arguments. Use -h for help.");
  
  skymodel = rem_args[0];
  # load the model
  model = Tigger.load(skymodel);
  if not model.sources:
    print "Input model %s contains no sources"%skymodel;
    sys.exit(0);
  print "Input model contains %d sources"%len(model.sources);
  
  # comparison predicates for the SELTAG<>SELVAL option
  select_predicates = {
    '==':lambda x,y:x==y,
    '!=':lambda x,y:x!=y,
    '>=':lambda x,y:x>=y,
    '<=':lambda x,y:x<=y,
    '>' :lambda x,y:x>y,
    '<' :lambda x,y:x<y,
    '.eq.':lambda x,y:x==y,
    '.ne.':lambda x,y:x!=y,
    '.ge.':lambda x,y:x>=y,
    '.le.':lambda x,y:x<=y,
    '.gt.' :lambda x,y:x>y,
    '.lt.' :lambda x,y:x<y
  };
  # units for same
  select_units = dict(d=DEG,m=DEG/60,s=DEG/3600);

  # This is where we accumulate the result of selection arguments, until we hit the first tagging argument.
  # Initially None, meaning no explicit selection
  selected_ids = None;
  
  # This is where we put the selection when we hit the first tagging argument.
  selection = None;
  
  # this is set to true when the selection is listed
  listed = False;
  # set to true when the model is modified
  modified = False;
  
  def apply_selection (sel,selstr):
    global selection;
    global selected_ids;
    global listed;
    listed = False;
    """Helper function: applies selection argument""";
    # if selection is not None, then we've already selected and tagged something, so we need
    # to reset the selection to empty and start again. If selected_ids is None, this is the first selection
    if selection is not None or selected_ids is None:
      print "Selecting sources:";
      selected_ids = set();
      selection = None;
    # add to current selection
    selected_ids.update(map(id,sel));
    # print result
    if not len(sel):
      print '  %-16s: no sources selected'%selstr;
    elif len(sel) == 1:
      print '  %-16s: one source selected (%s)'%(selstr,sel[0].name);
    elif len(sel) <= 5:
      print '  %-16s: %d sources selected (%s)'%(selstr,len(sel)," ".join([src.name for src in sel]));
    else:
      print '  %-16s: %d sources selected'%(selstr,len(sel));
      
  def retrieve_selection ():
    global selection;
    global selected_ids;
    """Helper function: retrieves current selection in preparation for tagging""";
    # if selection is None, then we need to set it up based on selected_ids
    if selection is None:
      # no explicit selection: use entire model
      if selected_ids is None:
        selection = model.sources;
        print "No explicit selection, using all sources.";
      # else use selected set
      else:
        selection = [ src for src in model.sources if id(src) in selected_ids ];
        print "Using %d selected sources:"%len(selection);
    if options.list:
      print "Sources: %s"%(" ".join([x.name for x in selection]));
      global listed;
      listed = True;
    return selection;

  def getTagValue (src,tag):
    """Helper function: looks for the given tag in the source, or in its sub-objects""";
    for obj in src,src.pos,src.flux,getattr(src,'shape',None),getattr(src,'spectrum',None):
      if obj is not None and hasattr(obj,tag):
        return getattr(obj,tag);
    return None;

  def lookupObject (src,tagname):
    """helper function to look into sub-objects of a Source object.
    Given src and "a", returns src,"a"
    Given src and "a.b", returns src.a and "b"
    """;
    tags = tagname.split(".");
    for subobj in tags[:-1]:
      src = getattr(src,subobj,None);
      if src is None:
        print "Can't resolve attribute %s for source %s"%(tagname,src.name);
        sys.exit(1);
    return src,tags[-1];
    
  # loop over all arguments
  for arg in rem_args[1:]:
  # Match either the SELTAG<>SELVAL, or the TAG=[TYPE:]VALUE, or the [+!/]TAG forms
  # If none match, assume the NAME form
    mselcomp = re.match("^(?i)([^=<>!.]+)(%s)([^dms]+)([dms])?"%"|".join([ key.replace('.','\.') for key in select_predicates.keys()]),arg);
    mseltag = re.match("=(.+)$",arg); 
    mset = re.match("^(.+)=((bool|int|str|float|complex):)?(.+)$",arg);
    msetbool = re.match("^([+!/])(.+)$",arg);
    
    # SELTAG<>SELVAL selection
    if mselcomp:
      seltag,oper,selval,unit = mselcomp.groups();
      try:  
        selval = float(selval)*select_units.get(unit,1.);
      except:
        parser.error("Malformed selection string '%s': right-hand side is not a number."%arg);
      predicate = select_predicates[oper.lower()];
      # get tag value
      srctag = [ (src,getTagValue(src,seltag)) for src in model.sources ];
      apply_selection([ src for src,tag in srctag if tag is not None and predicate(tag,selval) ],arg);
    elif mseltag:
      seltag = mseltag.groups()[0];
      apply_selection([ src for src in model.sources if getTagValue(src,seltag) ],arg);
    elif not mseltag and not mselcomp and not mset and not msetbool:
      apply_selection([ src for src in model.sources if fnmatch.fnmatch(src.name,arg) ],arg);
    elif mset:
      sources = retrieve_selection();
      if options.list:
        print "--list in effect, ignoring tagging commands";
        continue;
      tagname,typespec,typename,value = mset.groups();
      # if type is specified, use it to explicitly convert the value
      # first bool: allow True/False/T/F
      if typename == "bool":
        val = value.lower();
        if val == "true" or val == "t":
          newval = True;
        elif val == "false" or val == "f":
          newval = False;
        else:
          try:
            newval = bool(int(value));
          except:
            print "Can't parse \"%s\" as a value of type bool"%value;
            sys.exit(2);
      # else some other type is specified -- use it to convert the value
      elif typename:
        try:
          newval = getattr(__builtin__,typename)(value);
        except:
          print "Can't parse \"%s\" as a value of type %s"%(value,typename);
          sys.exit(2);
      # else auto-convert
      else:
        newval = None;
        for tp in int,float,complex,str:
          try:
            newval = tp(value);
            break;
          except:
            pass;
      # ok, value determined
      if type(newval) is str:
        value = '"%s"'%value;
      if sources:
        print "  setting tag %s=%s (type '%s')"%(tagname,value,type(newval).__name__);
        for src in sources:
          obj,tag = lookupObject(src,tagname);
          obj.setAttribute(tag,newval);
        modified = True;
      else:
        print "No sources selected, ignoring tagging commands";
    elif msetbool:
      sources = retrieve_selection();
      if options.list:
        print "--list in effect, ignoring tagging commands";
        continue;
      if sources:
        op,tagname = msetbool.groups();
        if op == "+":
          print "  setting tag %s=True"%tagname;
          method = 'setAttribute';
          args = (tagname,True);
        elif op == "!":
          print "  setting tag %s=False"%tagname;
          method = 'setAttribute';
          args = (tagname,False);
        elif op == "/":
          print "  removing tag %s"%tagname;
          method = 'removeAttribute';
          args = (tagname,);
        for src in sources:
          obj,tag = lookupObject(src,tagname);
          getattr(obj,method)(*args);
        modified = True;
      else:
        print "No sources selected, ignoring tagging commands";

  if options.list:
    if not listed:
      retrieve_selection();
    
  if not modified:
    print "Model was not modified";
    sys.exit(0);
          
  # prompt
  if not options.force:
    try:
      raw_input("Press ENTER to save model or Ctrl+C to cancel: ");
    except:
      print "Cancelling";
      sys.exit(1);

  # save output
  if options.output:
    model.save(options.output);
    print "Saved updated model to %s"%options,output;
  else:
    model.save(skymodel);
    print "Saved updated model";
