#!/usr/bin/python

r"""
Fetch and extract Ifremer ASCAT Daily NetCDF files.

Usage:
    python ascat.py start_year start_month start_day \
                    end_year   end_month   end_day

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

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

import os
import sys
import traceback
import calendar
import ftplib
import bz2

ftp_server = 'ftp.ifremer.fr'
ftp_user   = 'anonymous'
ftp_pass   = 'cbc@unc.edu'
ftp_dir    = 'ifremer/cersat/products/gridded/MWF/L3/ASCAT/Daily/Netcdf/'

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('ascat', 'daily')
year_spec   = '%(year)4u'
month_spec  = '%(month)02u'
day_spec    = '%(day)02u'

def _fetch(f):
    """
    Fetch a single bzip2 file.
    """

    try:
        bz = f.nlst()[0]
        b = open(bz, 'wb')
    except KeyboardInterrupt:
        raise
    except:
        traceback.print_exc()
        return None
    try:
        print 'Fetching:   %s.' % bz
        print f.retrbinary('RETR ' + bz, b.write)
    except KeyboardInterrupt:
        raise
    except:
        traceback.print_exc()
        b.close()
        return None
    else:
        b.close()
    return bz

def _extract(bz):
    """
    Extract a single bzip2 file.
    """

    b = os.path.split(bz)
    b = os.path.join(b[0],b[1][:8] + os.extsep + "nc")
    print 'Extracting: %s.' % b
    try:
        dest = open(b, 'wb')
    except KeyboardInterrupt:
        raise
    except:
        traceback.print_exc()
        return False
    try:
        src  = bz2.BZ2File(bz)
    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(bz)
    except KeyboardInterrupt:
        raise
    except:
        traceback.print_exc()
        return False
    return True
                

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

    day_range = date['range']
    for day in day_range:
        date.update({'day': day,})
        try:
            print f.cwd(os.path.join(os.path.sep,
                                     ftp_dir,
                                     year_spec % date,
                                     month_spec % date,
                                     day_spec % date,
                                     "").replace(os.path.sep, '/'))
        except KeyboardInterrupt:
            raise
        except ftplib.error_perm, e:
            traceback.print_exc()
            continue
        bz = _fetch(f)
        if not bz: continue
        if not _extract(bz): continue
    return True

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

    ascat(start_year, start_month, start_day,
          end_year,   end_month,   end_day   ) -> None

    start_year  is inclusive,
    start_month is 1-12 inclusive,
    start_day   is 1-31 inclusive,
    end_year    is exclusive,
    end_month   is 2-13 exclusive,
    end_day     is 2-32 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 year == year_range[0] and month == month_range[0]:
                _start_day = start_day
            else:
                _start_day = 1
            if year == year_range[-1] and month == month_range[-1]:
                _end_day = end_day
            else:
                _end_day = calendar.monthrange(year, month)[1] + 1
            day_range = range(_start_day, _end_day)
            date = {'year': year,
                    'month': month,
                    'range': day_range,
                   }
            try:
                os.chdir(work_dir)
                sub_dir = os.path.join(base_subdir,
                                       year_spec % date,
                                       month_spec % date)
                os.makedirs(sub_dir)
                os.chdir(sub_dir)
            except OSError, e:
                if not str(e).startswith("[Errno 17] File exists"):
                    traceback.print_exc()
                    continue
            if not _process(f, date): continue
    print f.quit()

if __name__ == '__main__':
    try:
        if len(sys.argv) != 7: raise
        start_year  = int(sys.argv[1])
        start_month = int(sys.argv[2])
        start_day   = int(sys.argv[3])
        end_year    = int(sys.argv[4]) + 1
        end_month   = int(sys.argv[5]) + 1
        end_day     = int(sys.argv[6]) + 1 
    except KeyboardInterrupt:
        raise
    except:
        print __doc__
        sys.exit()
    ascat(start_year, start_month, start_day,
          end_year,   end_month,   end_day)