# nwisextractor.py
#
# David J. Lampert (djlampert@gmail.com)
#
# Last updated: 07/26/2014
#
# contains a class to download and import data on daily streamflow gage 
# from the National Water Information Service (NWIS). the example below
# will download the source metadata file for the US, then download all
# the daily flow and water quality data for the Patuxent watershed.

import shutil, os, pickle, zipfile, datetime

from shapefile import Reader, Writer
from urllib    import request

from .gagestation import GageStation

class NWISExtractor:
    """A class to download the NWIS station source metadata for the USGS 
    daily discharge stations, and then retrieve information for a particular
    watershed (8-digit HUC).
    """

    def __init__(self, 
                 destination, 
                 url = 'http://water.usgs.gov/GIS/dsdl',
                 ):

        self.url         = url
        self.destination = destination

        if not os.path.isdir(destination):

            print('destination directory does not exist\n')
            
            try: os.mkdir(destination)
            except:
                print('warning, unable to create directory ' +
                      '{}'.format(destination))
                raise

    def report(self, 
               n, 
               block, 
               size,
               ):

        if n % 100 == 0:
            it = block * n / 10**6, size / 10**6
            print('{:.1f} MB of {:.1f} MB transferred'.format(*it))

    def download_metadata(self,
                          webfile = 'USGS_Streamgages-NHD_Locations_Shape.zip', 
                          sfile   = 'USGS_Streamgages-NHD_Locations',
                          verbose = True
                          ):
        """Downloads the source data files."""
        
        # directory for output

        if not os.path.isdir(self.destination):

            if verbose:
 
                print(self.destination)
                print('The NWIS metadata are not present in the destination\n')
                os.mkdir(self.destination)

        else: print('NWIS directory {} exists\n'.format(self.destination))

        # source zip file download

        zfile = '{}/{}'.format(self.destination, webfile)
        if not os.path.isfile(zfile):

            url = '{}/{}'.format(self.url, webfile)
            print('the source zip file for the NWIS metadata is not present\n')
            request.urlretrieve(url, zfile, self.report)

        elif verbose: 

            print('NWIS source metadata file {} is present\n'.format(zfile))

        # unzip 

        self.NWIS = '{}/{}'.format(self.destination, sfile)
        if not os.path.isfile(self.NWIS + '.shp'):
            
            print('extracting the gage shapefile from the archive')

            zf = zipfile.ZipFile(zfile)
            zf.extractall(self.destination)
            print('')

        elif verbose: print('gage metadata {} is present\n'.format(self.NWIS))

    def extract_HUC8(self, HUC8, output, gagefile = 'gagestations', 
                     verbose = True):
        """Extracts the USGS gage stations for a watershed from the gage 
        station shapefile into a shapefile for the 8-digit hydrologic unit 
        code of interest. 
        """

        # make sure the metadata exist locally

        self.download_metadata()

        # make sure the output destination exists

        if not os.path.isdir(output): os.mkdir(output)

        sfile = '{}/{}'.format(output, gagefile)
        if not os.path.isfile(sfile + '.shp'):

            # copy the projection

            shutil.copy(self.NWIS + '.prj', sfile + '.prj')

            # read the file

            gagereader  = Reader(self.NWIS, shapeType = 1)
            gagerecords = gagereader.records()

            # pull out the HUC8 record to parse the dataset

            HUC8_index  = gagereader.fields.index(['HUC',  'C', 8, 0]) - 1

            # iterate through the field and find gages in the watershed

            its = HUC8, sfile
            print('extracting gage stations in {} to {}\n'.format(*its))

            gage_indices = []

            i = 0
            for record in gagerecords:
                if record[HUC8_index] == HUC8: gage_indices.append(i)
                i+=1

            # write the data from the HUC8 to a new shapefile

            w = Writer(shapeType = 1)

            for field in gagereader.fields:  w.field(*field)

            for i in gage_indices:
                point = gagereader.shape(i).points[0]
                w.point(*point)
                w.record(*gagerecords[i])

            w.save(sfile)

            if verbose: 
                print('successfully extracted NWIS gage stations\n')

        elif verbose: 

            print('gage station file {} exists\n'.format(sfile))

        self.set_metadata(sfile)

    def set_metadata(self, gagefile):
        """Opens the gage file for the station metadata and saves it."""

        # metadata for stations

        self.gages  = []
        self.day1s  = []
        self.dayns  = []
        self.drains = []
        self.states = []
        self.sites  = []
        self.nwiss  = []
        self.aves   = []
        self.names  = []

        gagereader = Reader(gagefile, shapeType = 1)

        # get the fields with pertinent info

        day1_index  = gagereader.fields.index(['DAY1',       'N', 19, 0]) - 1
        dayn_index  = gagereader.fields.index(['DAYN',       'N', 19, 0]) - 1
        drain_index = gagereader.fields.index(['DA_SQ_MILE', 'N', 19, 2]) - 1
        HUC8_index  = gagereader.fields.index(['HUC',        'C',  8, 0]) - 1
        state_index = gagereader.fields.index(['STATE',      'C',  2, 0]) - 1
        site_index  = gagereader.fields.index(['SITE_NO',    'C', 15, 0]) - 1
        nwis_index  = gagereader.fields.index(['NWISWEB',    'C', 75, 0]) - 1
        ave_index   = gagereader.fields.index(['AVE',        'N', 19, 3]) - 1
        name_index  = gagereader.fields.index(['STATION_NM', 'C', 60, 0]) - 1

        # iterate through the records

        for r in gagereader.records():
            
            gage  = r[site_index] 
            day1  = r[day1_index] 
            dayn  = r[dayn_index] 
            drain = r[drain_index]
            state = r[state_index]
            nwis  = r[nwis_index]
            ave   = r[ave_index]  
            name  = r[name_index]
            site  = r[site_index]

            self.gages.append(gage)
            self.day1s.append(day1)
            self.dayns.append(dayn)
            self.drains.append(drain)
            self.states.append(state)
            self.sites.append(site)
            self.nwiss.append(nwis)
            self.aves.append(ave)
            self.names.append(name)

    def download_gagedata(self, 
                          gageid, 
                          start, 
                          end, 
                          output = None, 
                          plot = True
                          ):
        """Downloads the daily flow and water quality data for the given
        period of time for a particular gage in the metadata."""

        if gageid in self.gages:

            i = self.gages.index(gageid)

            gage  = self.gages[i]
            day1  = self.day1s[i]
            dayn  = self.dayns[i]
            drain = self.drains[i]
            state = self.states[i]
            nwis  = self.nwiss[i]
            ave   = self.aves[i]
            name  = self.names[i]

            if not os.path.exists(output):
                
                gagestation = GageStation(gage, name, state, day1, dayn,
                                          drain, ave, nwis)

                # download the daily discharge data if it exists

                try: 

                    gagestation.download_daily_discharge(start, end)

                    # make a plot of the data

                    if plot: gagestation.plot(output)

                    # download the water quality data if it exists

                    try: gagestation.download_water_quality()
                    except: 
                        print('warning, unable to download water quality ' +
                              'data for {}\n'.format(gage))

                    # save for later

                    with open(output, 'wb') as f: pickle.dump(gagestation, f)

                except: 
                    print('warning, unable to download daily flow data ' +
                          'for {}\n'.format(gage))

            else: print('gage data for {} exist\n'.format(gage))

        else:

            print('error: gageid not in metadata\n')
            raise

    def download_all(self, 
                     start, 
                     end, 
                     output = None, 
                     plot = True
                     ):
        """Downloads the daily flow data for the given period of time for 
        all gages in the metadata."""

        if output is None: output = 'gagestations'
        if not os.path.isdir(output): os.mkdir(output)

        # iterate through the records

        iters = zip(self.gages, self.day1s, self.dayns, self.drains, 
                    self.states, self.nwiss, self.aves, self.names)

        for gage, day1, dayn, drain, state, nwis, ave, name in iters:

            # filename to use for data

            filename = '{}/{}'.format(output, gage)

            if not os.path.exists(filename):
                
                gagestation = GageStation(gage, name, state, day1, dayn,
                                          drain, ave, nwis)

                # download the daily discharge data if it exists

                try: 

                    gagestation.download_daily_discharge(start, end)

                    # make a plot of the data

                    if plot: gagestation.plot(filename)

                    # download the water quality data if it exists

                    try: gagestation.download_water_quality()
                    except: 
                        print('warning, unable to download water quality ' +
                              'data for {}\n'.format(gage))

                    # save for later

                    with open(filename, 'wb') as f: pickle.dump(gagestation, f)

                except: 
                    print('warning, unable to download daily flow data ' +
                          'for {}\n'.format(gage))

if __name__ == '__main__':

    NWIS     = '{}/NWIS'.format(os.getcwd())
    HUC8     = '02060006'
    gage     = '01594670'
    output   = 'HSPF_data'
    gagefile = 'gagestations'
    start    = datetime.datetime(1980, 1, 1)
    end      = datetime.datetime(2010, 1, 1)

    # create an instance of the NWIS extractor

    extractor = NWISExtractor(NWIS)

    # extract the gage stations for the HUC8

    extractor.extract_HUC8(HUC8, output = '{}/{}'.format(output, gagefile))

    # download the daily flow and water quality data for each of the stations

    extractor.download_data('{}/{}'.format(output, gagefile), output = output)
