text: # -*- coding: utf-8 -*-
  ################################################################################
  # Copyright 2014, Distributed Meta-Analysis System
  ################################################################################
  
  """
  This file provides methods for extracting data from impact bundles (.nc4 files)
  """
  
  __copyright__ = "Copyright 2014, Distributed Meta-Analysis System"
  
  __author__ = "James Rising"
  __credits__ = ["James Rising"]
  __maintainer__ = "James Rising"
  __email__ = "jar2234@columbia.edu"
  
  __status__ = "Production"
  __version__ = "$Revision$"
  # $Source$
  
  import os, csv
  import numpy as np
  from netCDF4 import Dataset
  import configs
  
  deltamethod_vcv = None
  
  
  def read_region(config, *args, **kwargs):
      """Snip-out target regions from nc4 file
  
      Quick and dirty hax to reduce the size of data read in from netCDF files.
      Keeps a memory leak in the module from blowing up the script. Not
      the best way to handle this.
  
      Parameters
      ----------
      config : dict
          Run configuration dictionary. Used to parse out target regions.
      *args :
          Passed on to read().
      **kwargs :
          Passed on to read().
  
      Returns
      -------
      years : array-like
      regions : array-like
      data : array-like
      """
      years, regions, data = read(*args, **kwargs)
  
      if configs.is_allregions(config):
          regions_msk = np.ones(regions.shape, dtype='bool')
      else:
          target_regions = configs.get_regions(config, regions)
          regions_msk = np.isin(regions, target_regions)
  
      return years, regions[regions_msk], data[..., regions_msk]
  
  
  def read(filepath, column='rebased', deltamethod=False):
      """If deltamethod is True, treat as a deltamethod file."""
      global deltamethod_vcv
  
      try:
          rootgrp = Dataset(filepath, 'r', format='NETCDF4')
      except:
          print "Error: Cannot read %s" % filepath
          exit()
  
      years = rootgrp.variables['year'][:]
      regions = rootgrp.variables['regions'][:]
  
      if deltamethod is None:
          # Infer from the file
          deltamethod = 'vcv' in rootgrp.variables
          
      if deltamethod:
          data = rootgrp.variables[column + '_bcde'][:, :, :]
          if deltamethod_vcv is None:
              deltamethod_vcv = rootgrp.variables['vcv'][:, :]
          else:
              assert(np.all(deltamethod_vcv == rootgrp.variables['vcv'][:, :]))
      else:
          data = rootgrp.variables[column][:, :]
  
      rootgrp.close()
  
      # Correct bad regions in costs
      if filepath[-10:] == '-costs.nc4' and not isinstance(regions[0], str) and not isinstance(regions[0], unicode) and np.isnan(regions[0]):
          rootgrp = Dataset(filepath.replace('-costs.nc4', '.nc4'), 'r', format='NETCDF4')
          regions = rootgrp.variables['regions'][:]
          rootgrp.close()
  
      return years, regions, data
  
  def iterate_regions(filepath, column, config={}):
      global deltamethod_vcv
  
      do_deltamethod = False if configs.is_parallel_deltamethod(config) else config.get('deltamethod', None)
      if column is not None or 'costs' not in filepath:
          years, regions, data = read_region(config, filepath, column if column is not None else 'rebased', do_deltamethod)
      else:
          years, regions, data1 = read_region(config, filepath, 'costs_lb', do_deltamethod)
          years, regions, data2 = read_region(config, filepath, 'costs_ub', do_deltamethod)
          data = data2 / 1e5
  
      if deltamethod_vcv is not None and not config.get('deltamethod', False):
          ## Inferred that these were deltamethod files
          config['deltamethod'] = True
  
      if config.get('multiimpact_vcv', None) is not None and deltamethod_vcv is not None:
          assert isinstance(config['multiimpact_vcv'], np.ndarray)
          # Extend data to conform to multiimpact_vcv
          foundindex = None
          for ii in range(config['multiimpact_vcv'].shape[0] - deltamethod_vcv.shape[0] + 1):
              if np.allclose(deltamethod_vcv, config['multiimpact_vcv'][ii:(ii+deltamethod_vcv.shape[0]), ii:(ii+deltamethod_vcv.shape[1])]):
                  foundindex = ii
                  break
          if foundindex is None:
              print np.sum(np.abs(deltamethod_vcv - config['multiimpact_vcv'][:deltamethod_vcv.shape[0], :deltamethod_vcv.shape[1]]))
              print np.sum(np.abs(deltamethod_vcv - config['multiimpact_vcv'][deltamethod_vcv.shape[0]:, deltamethod_vcv.shape[1]:]))
          assert foundindex is not None, "Cannot find the VCV for " + filepath + " within the master VCV."
          newdata = np.zeros(tuple([config['multiimpact_vcv'].shape[0]] + list(data.shape[1:])))
          if len(data.shape) == 2:
              newdata[foundindex:(foundindex + deltamethod_vcv.shape[0]),:] = data
          else:
              newdata[foundindex:(foundindex + deltamethod_vcv.shape[0]),:,:] = data
          data = newdata
  
          deltamethod_vcv = None # reset for next file
  
      config['regionorder'] = list(regions)
  
      if configs.is_allregions(config):
          yield 'all', years, data
          return
  
      regions = list(regions)
      for region in configs.get_regions(config, regions):
          ii = regions.index(region)
          if config.get('deltamethod', False) and not configs.is_parallel_deltamethod(config):
              yield regions[ii], years, data[:, :, ii]
          else:
              yield regions[ii], years, data[:, ii]
  
  def iterate_values(years, values, config={}):
      """
      Config options: yearsets, years
      """
  
      if 'yearsets' in config and config['yearsets']:
          yearsets = config['yearsets']
          if yearsets == True:
              yearsets = [(2000, 2019), (2020, 2039), (2040, 2059), (2080, 2099)]
  
          for yearset in yearsets:
              if isinstance(yearset, list):
                  yearset = tuple(yearset)
              if config.get('deltamethod', False):
                  if values.ndim == 1:
                      yield "%d-%d" % yearset, np.mean(values[:, np.logical_and(years >= yearset[0], years < yearset[1])], axis=1)
                  else: # multiple regions included
                      yield "%d-%d" % yearset, np.mean(values[:, np.logical_and(years >= yearset[0], years < yearset[1]), :], axis=1)
              else:
                  if values.ndim == 1:
                      yield "%d-%d" % yearset, np.mean(values[np.logical_and(years >= yearset[0], years < yearset[1])])
                  else: # multiple regions included
                      yield "%d-%d" % yearset, np.mean(values[np.logical_and(years >= yearset[0], years < yearset[1]), :], axis=0)
          return
  
      years = list(years)
      for year in configs.get_years(config, years):
          if config.get('deltamethod', False) and not configs.is_parallel_deltamethod(config):
              if values.ndim == 2:
                  yield year, values[:, years.index(year)]
              else:
                  yield year, values[:, years.index(year), :]
          else:
              if values.ndim == 1:
                  yield year, values[years.index(year)]
              else:
                  yield year, values[years.index(year), :]
  