#!/usr/bin/env python

import numpy
import pylab

import urllib

#- Fetch the plain-text, fixed-column-width version of the Asiago Supernova
#- Catalog (ASC) from the Asiago Supernova Group website.  This group tries
#- to maintain a comprehensive list of all supernovae discovered and announced 
#- publically via IAU circulars.
#-
#- Literature reference for the ASC:
#-      [Barbon et al., 1999, A&AS, 139, 531]
#-
#- The Asiago Supernova Group website is at:
#-      [http://graspa.oapd.inaf.it/]
#-
#- Another resource is the Sternberg Supernova Catalog:
#-      [http://www.sai.msu.su/sn/sncat/]

stream = urllib.urlopen( "http://graspa.oapd.inaf.it/index.php?option=com_content&view=article&id=63&Itemid=82&dir=%2Fvar%2Fwww%2Fhtml%2Fjsmallfib_top%2Fcat&download_file=jsmallfib_top%2Fcat%2Fcat.txt" )
lines = [ line.rstrip( "\n\r" ) for line in stream ]
stream.close()

lines.pop( 0 )

#- We are interested in the supernova year, R.A. and Dec., and discoverer.  The 
#- ASC has remained consistent since 1999 in formatting its plain-text edition, 
#- so in the spirit of keeping it simple we adopt the fixed fields as we find 
#- them on this day.  Notice that we read both the coordinates of the host 
#- galaxy and the supernova in each case --- the host galaxy values are needed
#- in case the supernova coordinates are not listed.

bounds = { # The following column definitions are correct as of 2010-04-20.
        "year"      : (   2,    6 ),
        "gal_ra"    : (  32,   42 ),
        "gal_dec"   : (  42,   52 ),
        "sn_ra"     : (  52,   62 ),
        "sn_dec"    : (  62,   77 ),
        "disc"      : ( 197, None ) }

#- Process each row so that the supernova R.A. and Dec. fields are in radians.
#- This means making sure that if no supernova coordinates are recorded in the
#- ASC, we substitute in the host galaxy coordinates, which will be good enough
#- for the figure we are going to make --- the error is much smaller than the 
#- size of the symbols we will use for plotting.

years = []
ras   = []
decs  = []
discs = []

for line in lines :
    if not line.rstrip() :
        continue

    # Year and discoverer are easy.
    year = int( line[ bounds[ "year" ][ 0 ] : bounds[ "year" ][ 1 ] ] )
    disc = line[ bounds[ "disc" ][ 0 ] : ]

    # Try for supernova R.A. and Decl. and use galaxy values only if needed.
    ra  = line[ bounds[ "sn_ra"  ][ 0 ] : bounds[ "sn_ra"  ][ 1 ] ].strip( ": " )
    dec = line[ bounds[ "sn_dec" ][ 0 ] : bounds[ "sn_dec" ][ 1 ] ].strip( ": " )
    if not ( ra and dec ) :
        ra  = line[ bounds[ "gal_ra"  ][ 0 ] : bounds[ "gal_ra"  ][ 1 ] ].strip( ": " )
        dec = line[ bounds[ "gal_dec" ][ 0 ] : bounds[ "gal_dec" ][ 1 ] ].strip( ": " )

    # Convert R.A. to float degrees.
    h, m, s = ra[ 0 : 2 ], ra[ 2 : 4 ], ra[ 4 : ]
    try :
        test = float( s )
    except ValueError:
        s = "00"
    ra = 15.0 * ( float( h ) + ( float( m ) + float( s ) / 60.0 ) / 60.0 )

    # Convert Dec. to float degrees.
    p, d, m, s = dec[ 0 ], dec[ 1 : 3 ], dec[ 3 : 5 ], dec[ 5 : ]
    try :
        test = float( s )
    except ValueError:
        s = "00"
    dec = float( d ) + ( float( m ) + float( s ) / 60.0 ) / 60.0
    if p != "+": dec = -dec

    # Push onto each array.
    years.append( year )
    ras.append  ( ra   )
    decs.append ( dec  )
    discs.append( disc )

years = numpy.array( years, dtype = "int" )
discs = numpy.array( discs )
ras   = numpy.radians( numpy.array( ras   ) )
decs  = numpy.radians( numpy.array( decs  ) )
ras[ ras > numpy.pi ] = ras[ ras > numpy.pi ] - 2.0 * numpy.pi

#- Colors for some selected surveys --- not exhaustive by any stretch.

colors = []
for i, disc in enumerate( discs ) :
    if "sdss" in disc.lower() :
        discs[ i ] = "SDSS"
        colors.append( "#C28800" )
    elif "loss" in disc.lower() or "lotoss" in disc.lower() :
        discs[ i ] = "LOSS"
        colors.append( "#B20050" )
    elif "hzsst" in disc.lower() :
        discs[ i ] = "High-Z"
        colors.append( "#FF9C00" )
    elif "nsf" in disc.lower() :
        discs[ i ] = "SNfactory"
        colors.append( "#06D1AE" )
    elif "essence" in disc.lower() :
        discs[ i ] = "ESSENCE"
        colors.append( "#0F00B2" )
    elif "scp" in disc.lower() :
        discs[ i ] = "SCP"
        colors.append( "#8C8458" )
    elif "cfht" in disc.lower() :
        discs[ i ] = "SNLS"
        colors.append( "#3B6642" )
    else :
        discs[ i ] = "other" # could add others here
        colors.append( "#516587" )
colors = numpy.array( colors )

#- Generate a bunch of images to be run through ImageMagick convert to make an
#- animated GIF.

uniq_years = sorted( set( years ) )

discs_used = []
color_used = []

for i, year in enumerate( range( 1885, 2011 ) ) :
    fig = pylab.figure( figsize = ( 8, 4 ), dpi = 180 )
    axes = pylab.axes( projection = "hammer" )
    axes.grid( True )
    selected = years < year
    print year,
    print selected.sum(),
    if any( selected ) :
        axes.scatter( ras[ selected ], decs[ selected ], c = colors[ selected ], edgecolors = "none", alpha = 0.5 )
    selected = years == year
    print selected.sum()
    if any( selected ) :
        axes.scatter( ras[ selected ], decs[ selected ], c = colors[ selected ], edgecolors = "none", alpha = 1.0 )
        labels = discs[ selected ]
        for disc, color in zip( discs[ selected ], colors[ selected ] ) :
            if disc == "other" :
                continue
            if disc not in discs_used :
                discs_used.append( disc  )
                color_used.append( color )
    print discs_used
    for i, disc in enumerate( discs_used ) :
        pylab.text( 0.99, 0.32 - 0.06 * i, disc, color = color_used[ i ], transform = axes.transAxes )
    pylab.title( year )
    pylab.savefig( str( year ) + ".jpg" )
    pylab.clf()