from scipy import stats
from scipy.spatial.distance import cdist
from scipy.spatial import cKDTree as KDTree
import statsmodels.api as sm
import numpy as np
import pandas as pd
import gdal
import os
from __future__ import division

st_wmo = [os.path.join(root, name)
           for root, dirs, files in os.walk(r'I:\Documents\Klusjes\GIDS Interpolation Air Temperature\GSOD2009\selection_nl')
             for name in files
             if name.endswith('.op')]

date_all = [20120501]#, 20120502, 20120503, 20120504, 20120505, 20120506, 20120507, 20120508, 20120509, 20120510, 20120511, 20120512, 20120513, 20120514, 20120515, 20120516, 20120517, 20120518, 20120519, 20120520, 20120521, 20120522, 20120523, 20120524, 20120525, 20120526, 20120527, 20120528, 20120529, 20120530]

file_raster = r'I:\Documents\Klusjes\GIDS Interpolation Air Temperature\GSOD//AALBERS_25KM_DEM.tif'
raster = gdal.Open(file_raster, gdal.GA_ReadOnly)
band = raster.GetRasterBand(1)
dem = band.ReadAsArray()
#print band.GetNoDataValue()
#print 'Grid to evaluate\n', dem

def GIDS(x,y):    
    x_y = [(x,y)] # for kd_tree that starts counting at bottom left with 0,0
    
    x_ = (x*extent[1]+extent[0]+extent[1]/2) # longitude aalbers projection (meters)
    y_ = (y*extent[5]+extent[3]+extent[5]/2) # latitude aalbers projection (meters)
    long_lat = np.array([[x_,y_]])
    
    #print long_lat
    
    dem_1 = dem[y,x] # elevation x_y coordinate
    dist_tree, ix_tree = tree.query(x_y, k=10, eps=0, p=1) # returns distance and index
    df_selection = df.ix[ix_tree.ravel()]
    
    #print 'elevation from x_y -', dem_1
    #print '\n10 nearest neighbours\n', df_selection
    
    Longi = df_selection.ix[:,10]
    Lati = df_selection.ix[:,11]
    hi = df_selection.ix[:,5]
    ti = df_selection.ix[:,16]
    ti = (ti-32)*(5/9)
    
    pr_var = zip(Longi,Lati,hi) # combines predictor variables as tuples
    y = zip(ti) # dependent variable
    X = sm.add_constant(pr_var) # multiple linear regression
    
    #fit the model
    mlr = sm.OLS(y,X).fit()
    b1,b2,b3,b0 = mlr.params
    #print mlr.summary()
    
    long_lat_stations = df_selection.as_matrix(columns=['POINT_X','POINT_Y']) 
                                 
    di = cdist(long_lat_stations, long_lat, 'euclidean') # Returns Eucleadian distance in meters between grid cell and selected weather-stations
    #print di
    
    # prepare datasets as flattened numpy array or as single values
    Hi = df_selection.as_matrix(columns=['Elev_scaled']).flatten()
    Ti = df_selection.as_matrix(columns=['TEMP_y']).flatten()
    Ti = (Ti-32)*(5/9)
    di_ = di.flatten()
    long_lat_ = long_lat.flatten()
    Longi_ = df_selection.as_matrix(columns=['POINT_X',]).flatten()
    Lati_ = df_selection.as_matrix(columns=['POINT_Y',]).flatten()
    
    top =    sum( (1/di_)**2 )**-1
    long_f = b1*(long_lat_[0]-Longi_)
    lat_f =  b2*(long_lat_[1]-Lati_)
    h_f =    b3*(dem_1-Hi)
    middle = Ti + long_f + lat_f + h_f
    end = (1/di_)**2 
    comb = top * sum( middle * end )
    #print comb
    return comb

for date in date_all: 
    df = pd.read_csv(r'I:\Documents\Klusjes\GIDS Interpolation Air Temperature\GSOD2009//20120502.csv')
    df_select = None
    df_select = pd.DataFrame()
    for st in st_wmo:
        with open(st) as f:    
            colspecs_data = [(0,6), (14,22), (24,30), (78,83), (118,123)]
            df_date = pd.read_fwf(f, colspecs=colspecs_data)
            for moda in df_date.YEARMODA:
                #print moda
                if moda == date:
                    df_select = df_select.append(df_date[df_date.YEARMODA==date])
                    break
    df = pd.merge(df, df_select, how='inner', left_on='USAF', right_on='STN---')
    #print df.head()
    
    Longscaled = df.ix[:,12]
    Latscaled = df.ix[:,13]
    tree = KDTree(zip(Longscaled,Latscaled), leafsize=11)
    #tree.data
    
    file_raster = r'I:\Documents\Klusjes\GIDS Interpolation Air Temperature\GSOD//AALBERS_25KM_DEM.tif'
    raster = gdal.Open(file_raster, gdal.GA_ReadOnly)
    band = raster.GetRasterBand(1)
    dem = band.ReadAsArray()
    #print band.GetNoDataValue()
    #print 'Grid to evaluate\n', dem
    
    extent = raster.GetGeoTransform()
    #print extent
    #print raster.GetProjection() # fyi
    
    tp = np.zeros([dem.shape[0],dem.shape[1]])
    tp.shape
    
    #Double for-loop
    for x in range(tp.shape[1]-1):
        for y in range(tp.shape[0]-1):
            tp[x][y] = GIDS(x,y)
    #print 'T predicted\n', tp
    
    outFilename = r'I:\Documents\Klusjes\GIDS Interpolation Air Temperature\GSOD2009\outTifs2//'+str(date)+'.tif'
    #outFilename
    
    # Set Driver
    format = "GTiff"
    driver = gdal.GetDriverByName( format )
    driver.Register()
    
    # Set Metadata for Raster output
    cols = raster.RasterXSize
    rows = raster.RasterYSize
    bands = raster.RasterCount
    datatype = band.DataType
    
    # Set Projection for Raster
    outDataset = driver.Create(outFilename, cols, rows, bands, datatype)
    geoTransform = raster.GetGeoTransform()
    outDataset.SetGeoTransform(geoTransform)
    proj = raster.GetProjection()
    outDataset.SetProjection(proj)
    
    # Write output to band 1 of new Raster
    outBand = outDataset.GetRasterBand(1)
    outBand.WriteArray(tp,0,0)
    
    # Close and finalise newly created Raster
    #F_M01 = None
    outBand = None
    proj = None
    geoTransform = None
    outDataset = None
    driver = None
    datatype = None
    bands = None
    rows = None
    cols = None
    river = None
    tp = None