#!/usr/bin/python2

"""
Fetch and extract SSMI Version-04 QuikSCAT Ocean Surface Wind Vectors

Usage:
    python2 qscatv4.py start_year start_month end_year end_month

Where:
    start_year  is YYYY format.
    start_month is 1-12 for Jan-Dec.                         
    end_year    is YYYY format.              
    end_month   is 1-12 for Jan-Dec.
"""

__author__    = 'Chris Calloway'
__email__     = 'cbc@unc.edu'
__copyright__ = 'Copyright 2011 UNC-CH Dept. of Marine Sciences'
__license__   = 'GPL2'

import os
import sys
import traceback
import calendar
import ftplib
import gzip

ftp_server = 'ftp.ssmi.com'
ftp_user   = 'anonymous'
ftp_pass   = 'cbc@unc.edu'

work_dir           = os.path.join(os.path.sep,
                                  'afs',
                                  'isis.unc.edu',
                                  'depts',
                                  'marine',
                                  'workspace',
                                  'hseim',
                                  'volumes',
                                  'vol1',
                                  'nc_wind_2011',
                                  'satellite')
base_subdir        = os.path.join('qscat', 'bmaps_v04')
year_subdir_spec   = 'y%(year)4u'
month_subdir_spec  = 'm%(month)02u'

gz_spec    = 'qscat_%(year)4u%(month)02u%(day)02uv4'

def _fetch(f, gz_path):
    """
    Fetch a single gzipped file.
    """

    try:
        g = open(gz_path, 'wb')
    except KeyboardInterrupt:
        raise
    except:
        traceback.print_exc()
        return False
    try:
        gz = os.path.basename(gz_path)
        print 'Fetching:   %s.' % gz_path
        print f.retrbinary('RETR ' + gz, g.write)
    except KeyboardInterrupt:
        raise
    except:
        traceback.print_exc()
        g.close()
        return False
    else:
        g.close()
    return True

def _extract(gz_path):
    """
    Extract a single gzipped file.
    """

    gz_name = os.path.splitext(gz_path)[0]
    print 'Extracting: %s.' % gz_name
    try:
        dest = open(gz_name, 'wb')
    except KeyboardInterrupt:
        raise
    except:
        traceback.print_exc()
        return False
    try:
        src  = gzip.GzipFile(gz_path)
    except KeyboardInterrupt:
        raise
    except:
        dest.close()
        traceback.print_exc()
        return False
    try:
        content = src.read()
    except KeyboardInterrupt:
        raise
    except:
        dest.close()
        src.close()
        traceback.print_exc()
        return False
    else:
        src.close()
    try:
        dest.write(content)
    except KeyboardInterrupt:
        raise
    except:
        dest.close()
        traceback.print_exc()
        return False
    else:
        dest.close()
    try:
        os.remove(gz_path)
    except KeyboardInterrupt:
        raise
    except:
        traceback.print_exc()
        return False
    return True
                

def _process(f, year, month):
    """
    Fetch and extract all the files for a month.
    """

    format = {'year': year, 'month': month,}
    try:
        subdir = os.path.join(base_subdir,
                              year_subdir_spec % format,
                              month_subdir_spec % format)
        os.makedirs(subdir)
    except OSError, e:
        if not str(e).startswith("[Errno 17] File exists"):
            traceback.print_exc()
            return False
    try:
        print f.cwd('/' + subdir.replace(os.path.sep, '/'))
    except ftplib.error_perm, e:
        traceback.print_exc()
        return False
    day_range = range(1, calendar.monthrange(year, month)[1] + 1)
    for day in day_range:
        format.update({'day': day,})
        gz      = (gz_spec % format) + os.path.extsep + 'gz'
        gz_path = os.path.join(subdir, gz)
        if not _fetch(f, gz_path): continue
        if not _extract(gz_path): continue
    return True

def qscatv4(start_year, start_month, end_year, end_month):
    """
    Fetch and extract a range of files.

    qscatv4(start_year, start_month, end_year, end_month) -> None

    start_year  is inclusive.
    start_month is 1-12 inclusive.
    end_year    is exclusive.
    end_month   is 2-13 exclusive.
    """

    f = ftplib.FTP(ftp_server)
    print f.login(ftp_user, ftp_pass)
    year_range = range(start_year, end_year)
    for year in year_range:
        if year == year_range[0]:
            _start_month = start_month
        else:
            _start_month = 1
        if year == year_range[-1]:
            _end_month = end_month
        else:
            _end_month = 13
        month_range = range(_start_month, _end_month)
        for month in month_range:
            if not _process(f, year, month): continue
    print f.quit()

if __name__ == '__main__':
    import pdb; pdb.set_trace()
    try:
        if len(sys.argv) != 5: raise
        start_year  = int(sys.argv[1])
        start_month = int(sys.argv[2])
        end_year    = int(sys.argv[3]) + 1
        end_month   = int(sys.argv[4]) + 1 
    except KeyboardInterrupt:
        raise
    except:
        print __doc__
        sys.exit()
    os.chdir(work_dir)
    qscatv4(start_year, start_month, end_year, end_month)