content: # -*- coding: utf-8 -*-
  """
  Authors: Tim Hessels
           UNESCO-IHE 2016
  Contact: t.hessels@unesco-ihe.org
  Repository: https://github.com/wateraccounting/wa
  Module: Collect/CFSR
  """
  # General modules
  import pandas as pd
  import os
  import numpy as np
  from netCDF4 import Dataset
  import re
  from joblib import Parallel, delayed
  
  # WA+ modules
  from wa.Collect.CFSR.Download_data_CFSR import Download_data
  from wa.General import data_conversions as DC
  
  def CollectData(Dir, Var, Startdate, Enddate, latlim, lonlim, Waitbar, cores, Version):    
      """
      This function collects daily CFSR data in geotiff format
  
      Keyword arguments:
      Dir -- 'C:/file/to/path/'
      Var -- 'dlwsfc','dswsfc','ulwsfc', or 'uswsfc'
      Startdate -- 'yyyy-mm-dd'
      Enddate -- 'yyyy-mm-dd'
      latlim -- [ymin, ymax] (values must be between -50 and 50)
      lonlim -- [xmin, xmax] (values must be between -180 and 180)
      Waitbar -- 1 (Default) will print a wait bar
      cores -- The number of cores used to run the routine.
               It can be 'False' to avoid using parallel computing
  		    routines.
      Version -- 1 or 2 (1 = CFSR, 2 = CFSRv2)	
      """
  
  				
      # Creates an array of the days of which the ET is taken
      Dates = pd.date_range(Startdate,Enddate,freq = 'D') 
  
      # Create Waitbar
      if Waitbar == 1:
          import wa.Functions.Start.WaitbarConsole as WaitbarConsole
          total_amount = len(Dates)
          amount = 0
          WaitbarConsole.printWaitBar(amount, total_amount, prefix = 'Progress:', suffix = 'Complete', length = 50)
  
      
      # For collecting CFSR data				
      if Version == 1:			
          # Check the latitude and longitude and otherwise set lat or lon on greatest extent
          if latlim[0] < -89.9171038899 or latlim[1] > 89.9171038899:
              print 'Latitude above 89.917N or below 89.917S is not possible. Value set to maximum'
              latlim[0] = np.maximum(latlim[0],-89.9171038899)
              latlim[1] = np.minimum(latlim[1],89.9171038899)
          if lonlim[0] < -180 or lonlim[1] > 179.843249782:
              print 'Longitude must be between 179.84E and 179.84W. Now value is set to maximum'
              lonlim[0] = np.maximum(lonlim[0],-180)
              lonlim[1] = np.minimum(lonlim[1],179.843249782)  
  												
          # Make directory for the CFSR data
          output_folder=os.path.join(Dir,'Radiation','CFSR')
          if not os.path.exists(output_folder):
              os.makedirs(output_folder)  										
  												
      # For collecting CFSRv2 data    
      if Version == 2:
              # Check the latitude and longitude and otherwise set lat or lon on greatest extent
          if latlim[0] < -89.9462116040955806 or latlim[1] > 89.9462116040955806:
              print 'Latitude above 89.917N or below 89.946S is not possible. Value set to maximum'
              latlim[0] = np.maximum(latlim[0],-89.9462116040955806)
              latlim[1] = np.minimum(latlim[1],89.9462116040955806)
          if lonlim[0] < -180 or lonlim[1] > 179.8977275:
              print 'Longitude must be between 179.90E and 179.90W. Now value is set to maximum'
              lonlim[0] = np.maximum(lonlim[0],-180)
              lonlim[1] = np.minimum(lonlim[1],179.8977275)  			
  				
          # Make directory for the CFSRv2 data
          output_folder=os.path.join(Dir,'Radiation','CFSRv2')
          if not os.path.exists(output_folder):
              os.makedirs(output_folder)  				
  				
  				
      # Pass variables to parallel function and run
      args = [output_folder, latlim, lonlim, Var, Version]
      if not cores:
          for Date in Dates:
              RetrieveData(Date, args)
              if Waitbar == 1:
                  amount += 1
                  WaitbarConsole.printWaitBar(amount, total_amount, prefix = 'Progress:', suffix = 'Complete', length = 50)
          results = True
      else:
          results = Parallel(n_jobs=cores)(delayed(RetrieveData)(Date, args)
                                           for Date in Dates)
      
      # Remove all .nc and .grb2 files
      for f in os.listdir(output_folder):
          if re.search(".nc", f):
              os.remove(os.path.join(output_folder, f))
      for f in os.listdir(output_folder):
          if re.search(".grb2", f):
              os.remove(os.path.join(output_folder, f))	
      for f in os.listdir(output_folder):
          if re.search(".grib2", f):
              os.remove(os.path.join(output_folder, f))	
  				
      return results
  		
  def RetrieveData(Date, args):
  
      # unpack the arguments
      [output_folder, latlim, lonlim, Var, Version] = args
  
      # Name of the model      
      if Version == 1:
          version_name = 'CFSR'
      if Version == 2:
          version_name = 'CFSRv2'								
  				
      # Name of the outputfile
      if Var == 'dlwsfc':
          Outputname = 'DLWR_%s_W-m2_' %version_name + str(Date.strftime('%Y')) + '.' + str(Date.strftime('%m')) + '.' + str(Date.strftime('%d')) + '.tif'
      if Var == 'dswsfc':
          Outputname = 'DSWR_%s_W-m2_' %version_name + str(Date.strftime('%Y')) + '.' + str(Date.strftime('%m')) + '.' + str(Date.strftime('%d')) + '.tif'
      if Var == 'ulwsfc':
          Outputname = 'ULWR_%s_W-m2_' %version_name + str(Date.strftime('%Y')) + '.' + str(Date.strftime('%m')) + '.' + str(Date.strftime('%d')) + '.tif'
      if Var == 'uswsfc':
          Outputname = 'USWR_%s_W-m2_' %version_name + str(Date.strftime('%Y')) + '.' + str(Date.strftime('%m')) + '.' + str(Date.strftime('%d')) + '.tif'
  
      # Create the total end output name        
      outputnamePath = os.path.join(output_folder, Outputname)
  
      # If the output name not exists than create this output								
      if not os.path.exists(outputnamePath):								
  
         								
          local_filename = Download_data(Date, Version, output_folder, Var)				
  
          # convert grb2 to netcdf (wgrib2 module is needed)    
          for i in range(0,4):
              nameNC = 'Output' + str(Date.strftime('%Y')) + str(Date.strftime('%m')) + str(Date.strftime('%d')) + '-' + str(i+1) + '.nc'
  				
              # Total path of the output 				
              FileNC6hour = os.path.join(output_folder, nameNC)
  				
  	       # Band number of the grib data which is converted in .nc			
              band=(int(Date.strftime('%d')) - 1) * 28 + (i + 1) * 7
  
              # Convert the data
              DC.Convert_grb2_to_nc(local_filename, FileNC6hour, band)
       
          if Version == 1:
  									
              if Date < pd.Timestamp(pd.datetime(2011, 01, 01)):
  	            		   							
                  # Convert the latlim and lonlim into array
                  Xstart = np.floor((lonlim[0] + 180.1562497) / 0.3125)
                  Xend = np.ceil((lonlim[1] + 180.1562497) / 0.3125) + 1
                  Ystart = np.floor((latlim[0] + 89.9171038899) / 0.3122121663)
                  Yend = np.ceil((latlim[1] + 89.9171038899) / 0.3122121663)
              
                  # Create a new dataset   
                  Datatot = np.zeros([576, 1152])	
  
              else:
                  Version = 2																
  												
          if Version == 2:	
  					
              # Convert the latlim and lonlim into array
              Xstart = np.floor((lonlim[0] + 180.102272725) / 0.204545)
              Xend = np.ceil((lonlim[1] + 180.102272725) / 0.204545) + 1
              Ystart = np.floor((latlim[0] + 89.9462116040955806) / 0.204423)
              Yend = np.ceil((latlim[1] + 89.9462116040955806) / 0.204423)
              
              # Create a new dataset   
              Datatot = np.zeros([880, 1760])		
  					
          # Open 4 times 6 hourly dataset 
          for i in range (0, 4):
              nameNC = 'Output' + str(Date.strftime('%Y')) + str(Date.strftime('%m')) + str(Date.strftime('%d')) + '-' + str(i + 1) + '.nc'
              FileNC6hour = os.path.join(output_folder, nameNC)
              f = Dataset(FileNC6hour, mode = 'r')
              Data = f.variables['Band1'][0:int(Datatot.shape[0]), 0:int(Datatot.shape[1])]
              f.close()
              data = np.array(Data)
              Datatot = Datatot + data
             			
          # Calculate the average in W/m^2 over the day   
          DatatotDay = Datatot / 4
          DatatotDayEnd = np.zeros([int(Datatot.shape[0]), int(Datatot.shape[1])])
          DatatotDayEnd[:,0:int(Datatot.shape[0])] = DatatotDay[:, int(Datatot.shape[0]):int(Datatot.shape[1])]
          DatatotDayEnd[:,int(Datatot.shape[0]):int(Datatot.shape[1])] = DatatotDay[:, 0:int(Datatot.shape[0])]
              
          # clip the data to the extent difined by the user
          DatasetEnd = DatatotDayEnd[int(Ystart):int(Yend), int(Xstart):int(Xend)]											
  											
          # save file
          if Version == 1:
  	        pixel_size = 0.3125
          if Version == 2:
  	        pixel_size = 0.204545   
          geo = [lonlim[0],pixel_size,0,latlim[1],0,-pixel_size]  
          DC.Save_as_tiff(data = np.flipud(DatasetEnd), name = outputnamePath, geo = geo, projection = "WGS84") 	
  				
      return()
  