#!/usr/local/CDAT-4.1.2/bin/python

#REQUIRES INSTALLATION OF CDAT http://www2-pcmdi.llnl.gov/cdat
#FIRST LINE MUST REFERENCE PYTHON INSTALLED AS PART OF CDAT

#extraction script originally written by T. Pruitt 7/2007
#Modified by E. Maurer 1/2009

#hardcoded for the global 0.5 degree resolution data
#this is NOT meant to be anything more than a start for people.
#it is NOT a general use utility.
#no one involved in creating or distributing this script has any
#responsibility for its function or output.

import sys
import getopt

import os

import cdms
import cdutil
import cdtime


class NC2TXT:

    def __init__(self):
        """ """

    def setSpatialSelector(self, lat1, lat2, lon1, lon2):
        """ spatial domain """
        res = 0.5
    
        if lat1 <= lat2:
            la1 = lat1-res/2.
            la2 = lat2+res/2.
        else:
            la1 = lat2-res/2.
            la2 = lat1+res/2.
    
        if lon1 <= lon2:
            lo1 = lon1-res/2.
            lo2 = lon2+res/2.
        else:
            lo1 = lon2-res/2.
            lo2 = lon1+res/2.

        # valid only for pts
        self.lat = lat1
        self.lon = lon1

        self.space_selector = cdutil.region.domain(latitude=(la1, la2), longitude=(lo1, lo2))


    def setNetCDFData(self, prcpfn, tavgfn):
        """ """
        self.prcp = self.getNetCDFData(prcpfn, 'Prcp') 
        self.tavg = self.getNetCDFData(tavgfn, 'Tavg')


    def getNetCDFData(self, fn, varname):
        """ """
        fp = cdms.open(fn)
        data = fp(varname, self.space_selector)
        fp.close()
        return data


    def write_ASCII(self, dir):
        """ """
        fn = '%s/data_%s_%s' % (dir, self.lat, self.lon)
        fp = open(fn, 'w')
        #print fn

        # times
        time_axis = self.prcp.getTime()
        times = time_axis.getData()

        Prcp = self.prcp.getValue()
        Tavg = self.tavg.getValue()
        flag=0
        for i in range(len(times)):
            time = times[i]
            t = cdtime.reltime(time, time_axis.units)
            tc = t.tocomponent()
            mo = tc.month
            yr = tc.year

            prcp = Prcp[i, 0, 0]
            tavg = Tavg[i, 0, 0]

            if prcp > 10000. or tavg > 10000.:
                s = '%s %2s %7s %7s\n' % (yr, mo, 'NA', 'NA')
                flag=1
            else:
                s = '%s %2s %7.3f %7.3f\n' % (yr, mo, tavg, prcp)
                fp.write(s)

        fp.close()
        if (flag == 1):
            os.unlink(fn)

    def getNetCDFDimensions(self, fn=''):
        """ """
        if not fn:
            fn = self.fn
        fp = cdms.open(fn)
        lat = fp.axes['latitude'].getData()
        lon = fp.axes['longitude'].getData()
        fp.close()
        return (lat,lon)


def twoFiles2Ascii(latx, latn, lonx, lonn, prcp_fn='', tavg_fn=''):

    dir = './out_ascii'
    if not os.path.isdir(dir):
        os.mkdir(dir)

    nc = NC2TXT()
    lats, lons = nc.getNetCDFDimensions(prcp_fn)

    print len(lats), len(lons)
    #lats = lats[100:110]
    #lons = lons[100:110]

    for lon in lons:
        for lat in lats:
            if lat < latx and lat > latn and lon < lonx and lon > lonn:
                print lat, lon

                # single point extraction
                nc.setSpatialSelector(lat,lat,lon,lon)
                nc.setNetCDFData(prcp_fn, tavg_fn)
                nc.write_ASCII(dir)


############################################
#    #            END OF CLASS                #
############################################
use = '''
Usage: %s -[h] PrcpFile TavgFile <max_lat> <min_lat> <max_lon> <min_lon>

        -h  help

        Default:  Writes to stdout
'''
if __name__ == '__main__':

    def usage():
        sys.stderr.write(use % sys.argv[0])
        sys.exit(1)

    try:
        (opts, args) = getopt.getopt(sys.argv[1:], 'h')
    except getopt.error:
        usage()

    verbose = False
    grid_info = False
    proj_info = True
    for (opt,val) in opts:
        if opt == '-h':
            usage()
        elif opt == '-v':
            verbose = True
        else:
            raise OptionError, opt
            usage()

    if len(args) == 6:
        prcp_fn = args[0]
        tavg_fn = args[1]
        latx = float(args[2])
        latn = float(args[3])
        lonx = float(args[4])
        lonn = float(args[5])
        twoFiles2Ascii(latx, latn, lonx, lonn, prcp_fn, tavg_fn)
    else:
        usage()

    #s = NC2TXT()