import numpy as N
from functions import solarang,time_zone,residuals
from constantes import *
from genutil import statistics
from scipy.optimize import leastsq
#-------------------------------------------------------------------------------------------
# function gap_fill_func gapfills the meteorlogical data
# argument1 : the weather dataset (weather)
# argument2 : the climatology dataset (clim)
# argument3 : the ratio between clim and weather time steps (diff_clim_weather)

# returns a vector that contains the weather dataset gapfilled
def gap_fill_func(weather,clim,diff_clim_weather,weather_period,climato_period,julian,year_length,lon,lat,gapmax,avg,climatoshift,timeshift):
   weather_clim_period=[]
   weather_clim_period_test=[]

   weather_clim_period_nogap=[]
   clim_nogap=[]
   weather_clim_period_gapfill=[]
   weather_gapfill=[]

   if(timeshift==-9999):
      timezone=time_zone(0,lon)
      # east of Greenwich => timeshift>0
      # west of Greenwich => timeshift<0
      if(timezone[1]<13):
         timeshift=timezone[1]
      else:
         timeshift=timezone[1]-24

   stat_vec=[]

   print 'shift to UTC time =',timeshift,' hours'


   for k in range(len(weather)):


      weather_clim_period.append(N.zeros(len(weather[k])/diff_clim_weather, N.float32, 0))
      if(k==id_precip):
         freq_precip=N.zeros(len(weather[k])/diff_clim_weather, N.float32, 0)
      weather_clim_period_test.append(N.zeros(len(weather[k])/diff_clim_weather, N.float32, 0))


      # climatoshift indicates to which time step or time period
      # unit is fraction of a time period (between 2 consecutive time steps)
      # a climatic field value corresponds to
      # for field that is a mean value (avg=1)
      #    climatoshift=0 when the mean value is calculated from one time step to the next one
      #    climatoshift=-0.5 when the mean value is centered on one time step
      # for field that is an instantaneous value (avg=0)
      #    climatoshift=0 when the value correspond to the current time step
      #    climatoshift=1 when the value corresponds to the next time step
      totalshift=timeshift+climatoshift[k]*climato_period
      print '\tVar=',label_fig[k]
      for t in range(len(weather[k])):
         tshift=t-(float)(totalshift)/weather_period
         if((tshift>=0)and(tshift<len(weather[k]))):
            # in case of a mean value calculation, we sum all weather element within each element of weather_clim_period
            # if one weather element equals -9999, the related weather_clim_period is equal to -9999
            if (avg[k]==1):
               if((weather[k][t] != -9999) and (weather_clim_period[k][tshift/diff_clim_weather] != -9999)):
                  weather_clim_period[k][tshift/diff_clim_weather]+=weather[k][t]/diff_clim_weather
               else:
                  weather_clim_period[k][tshift/diff_clim_weather]=-9999
               weather_clim_period_test[k][tshift/diff_clim_weather]=1

               if(k==id_precip):
                  if((weather[k][t] != -9999) and (freq_precip[tshift/diff_clim_weather] != -9999)):
                     if(weather[k][t]>0):
                        freq_precip[tshift/diff_clim_weather]+=1./diff_clim_weather
                  else:
                     freq_precip[tshift/diff_clim_weather]=-9999
               
            # in case of a instantaneaous calculation, each weather_clim_period element corresponds
            # to the first weather element associated to this weather_clim_period element
            else:
               if(weather_clim_period_test[k][tshift/diff_clim_weather] == 0):
                  if((weather[k][t] != -9999)):
                     weather_clim_period[k][tshift/diff_clim_weather]=weather[k][t]
                  else:
                     weather_clim_period[k][tshift/diff_clim_weather]=-9999
                  weather_clim_period_test[k][tshift/diff_clim_weather]=1
      # in case of a mean value calculcation (avg==1)
      # elements of weather_clim_period that have been partly filled (at the beginning or at the end)
      # are set to -9999
      if (avg[k]==1):
         if(totalshift%climato_period !=0):
            if(totalshift>0):
               weather_clim_period[k][(len(weather[k])-1-(float)(totalshift)/weather_period)/diff_clim_weather]=-9999
            else:
               if(totalshift<0):
                  weather_clim_period[k][(-(float)(totalshift)/weather_period)/diff_clim_weather]=-9999
      # elements of weather_clim_period that have NOT been filled (at the beginning or at the end)
      # are set to -9999          
      weather_clim_period[k]=N.where(weather_clim_period_test[k]==1,weather_clim_period[k],-9999)
               

      weather_clim_period_nogap.append([])
      clim_nogap.append([])
      for t in range(len(weather_clim_period[k])):
         if(weather_clim_period[k][t] != -9999):
            weather_clim_period_nogap[k].append(weather_clim_period[k][t])
            clim_nogap[k].append(clim[k][t])


      weather_clim_period_nogap[k]=N.array(weather_clim_period_nogap[k],N.float)
      clim_nogap[k]=N.array(clim_nogap[k],N.float)



      # Evaluate the correlation
      # between clim_nogap and weather_clim_period_nogap
      # intercept and slope of the relation
      # will be used for correcting clim data when filling gaps
      a=clim_nogap[k]
      b=weather_clim_period_nogap[k]

      if(k==id_swdown):
         stat=leastsq(residuals,1.,args=(b,a))
         stat=N.array(stat)
         stat[1]=0
      elif(k==id_precip):
         stat=N.zeros(2)
         stat[0]=sum(b)/sum(a)
         stat[1]=0
      else:
         stat=statistics.linearregression(b,x=a)
         
      if(str(stat[0])!='nan'):
         slope=stat[0]
         intercept=stat[1]
      else:
         slope=1
         intercept=0

      stat_vec.append([slope,intercept])
      print '\t\tslope of the linear relation in-situ VS reanalysis=',slope
      print '\t\tintercept of the linear relation in-situ VS reanalysis=',intercept
      if(str(stat[0])!='nan'):
         print '\t\tRMSE without correction=',statistics.rms(a,b)
         print '\t\tRMSE with correction=',statistics.rms(a*slope+intercept,b)
         

      weather_clim_period_gapfill.append([])
      weather_clim_period_gapfill[k]=N.where(weather_clim_period[k]==-9999, slope*clim[k]+intercept, weather_clim_period[k])

      weather_gapfill.append([])
      mean_csang=0.
      if(k==id_precip):
         freq_precip_nogap=N.ma.masked_values(freq_precip,-9999)
         freq_precip_nogap_nonull=N.ma.masked_values(freq_precip_nogap,0.)
         freq_precip_scalar=freq_precip_nogap_nonull.mean()
         print '\t\tfreq_precip_scalar=',freq_precip_scalar
         number_precip_per_diff_clim_weather=round(1./freq_precip_scalar)
       
      for t in range(len(weather[k])):
         tshift=t-(float)(totalshift)/weather_period
         
         if(k==id_swdown):
            if((t-(float)(timeshift)/weather_period)<(len(weather[k])-(diff_clim_weather-1))):
               if(tshift%diff_clim_weather==0):
                  mean_csang=0.
                  for l in range(diff_clim_weather):
                     mean_csang+=solarang(julian[(t-(float)(timeshift)/weather_period)+l],0,lon,lat,year_length[tshift])/diff_clim_weather
               if(mean_csang == 0.):
                  ratio=0.
               else:
                  ratio=solarang(julian[(t-(float)(timeshift)/weather_period)],0,lon,lat,year_length[(t-(float)(timeshift)/weather_period)])/mean_csang
            else:
               ratio=1
         if(k==id_lwdown):
            ratio=1
         if(k==id_precip):
            if(t%number_precip_per_diff_clim_weather==0):
               ratio=number_precip_per_diff_clim_weather
            else:
               ratio=0.

         # if the current value is -9999, we have to fill the gap
         if(weather[k][t] == -9999):
            # if gapmax is > 0
            # for short gap, we will try to interpolate
            # between the last and the next
            # defined values in the weather dataset
            active_linear=1
            if(gapmax>0):
               tlow=t-1
               while((tlow>=0)and(weather[k][tlow]==-9999)and((t-tlow)<=gapmax)):
                  tlow=tlow-1
               tup=t+1
               while((tup<len(weather[k]))and(weather[k][tup]==-9999)and((tup-t)<=gapmax)):
                  tup=tup+1
               # if the last and the next defined values are not too far
               # we linearly interpolate with the weather dataset
               if(tup-tlow<=gapmax+1):
                  if(tlow<0):
                     weather_gapfill[k].append(weather[k][tup])
                  elif(tup>=len(weather[k])):
                     weather_gapfill[k].append(weather[k][tlow])
                  else:
                     step=(float)(t-tlow)/(tup-tlow)
                     weather_gapfill[k].append(weather[k][tlow]*(1-step)+weather[k][tup]*step)
               else:
                  active_linear=0
            # if gapmax == 0 or
            # if the last and the next defined values are too far each other
            # we use the clim dataset for filling the gap
            if((gapmax==0)or(not active_linear)):   
               if (avg[k]==0):
                  step=(float)(tshift%diff_clim_weather)/diff_clim_weather
                  if(tshift<0.):
                     weather_gapfill[k].append(weather_clim_period_gapfill[k][0])
                  elif((tshift/diff_clim_weather+1)>=(len(weather_clim_period_gapfill[k]))):
                     weather_gapfill[k].append(weather_clim_period_gapfill[k][len(weather_clim_period_gapfill[k])-1])
                  else:
                     weather_gapfill[k].append(weather_clim_period_gapfill[k][tshift/diff_clim_weather]*(1-step)+weather_clim_period_gapfill[k][tshift/diff_clim_weather+1]*step)
               else:
                  if(tshift<0.):
                     weather_gapfill[k].append(weather_clim_period_gapfill[k][0]*ratio)
                  elif((tshift/diff_clim_weather)>=(len(weather_clim_period_gapfill[k]))):
                     weather_gapfill[k].append(weather_clim_period_gapfill[k][len(weather_clim_period_gapfill[k])-1]*ratio)
                  else:
                     weather_gapfill[k].append(weather_clim_period_gapfill[k][tshift/diff_clim_weather]*ratio)
         else:
            weather_gapfill[k].append(weather[k][t])

   return weather_gapfill,weather_clim_period_gapfill,weather_clim_period_nogap,weather_clim_period,clim_nogap,stat_vec
# end function
#-------------------------------------------------------------------------------------------