'''
Sun, 2019/5/20
Cubist回归进行多光谱水深反演,利用rpy2调用R库Cubist
* 输入参数[n_samples, n_features],即一个shape为（样点数，波段数）的矩阵
* predice对象也要有相同的维数
'''
import os
import rpy2
print(rpy2.__version__)
import pandas as pd
import numpy as np
from osgeo import gdal
import math
os.environ['R_USER'] = 'C:/ProgramData/Anaconda3/Lib/site-packages/rpy2'
import rpy2.robjects as robjects
from rpy2.robjects.packages import importr
from rpy2.robjects.vectors import FloatVector
import rpy2.robjects.conversion as conversion
from rpy2.robjects import pandas2ri
import rpy2.robjects.numpy2ri
rpy2.robjects.pandas2ri.activate()
rpy2.robjects.numpy2ri.activate()
Cubist = importr('Cubist')
lattice = importr('lattice')
r = robjects.r
import gc
gc.get_threshold()


# prepare the training point
dt = r['read.csv']            ('H:/GF6_yantai/Bathymetry/area/yangdian/data_control.csv',sep=',',header=True)
Z = FloatVector(dt[3])
X = FloatVector(dt[4])
X1 = FloatVector(dt[5])
X2 = FloatVector(dt[11])
X3 = FloatVector(dt[16])
X4 = FloatVector(dt[17])
X5 = FloatVector(dt[20])
X6 = FloatVector(dt[24])
T = r['cbind'](X=X,X1=X1,X2=X2,X3=X3,X4=X4,X5=X5,X6=X6)
regr = r['cubist'](x=T,y=Z,committees=10)
print(regr)
del X,X1,X2,X3,X4,X5,X6
gc.collect()
# mapping water depth
# rasters to be computed
dt = gdal.Open("H:/bathmetry/GF6/cart/layer_stac")
blue = dt.GetRasterBand(1).ReadAsArray(0,0,dt.RasterXSize,         dt.RasterYSize)
green = dt.GetRasterBand(2).ReadAsArray(0,0,dt.RasterXSize, dt.RasterYSize)
yellow = dt.GetRasterBand(3).ReadAsArray(0,0,dt.RasterXSize,     dt.RasterYSize)
blue_violet = dt.GetRasterBand(4).ReadAsArray(0,0,dt.RasterXSize, dt.RasterYSize)
blue_green = dt.GetRasterBand(5).ReadAsArray(0,0,dt.RasterXSize, dt.RasterYSize)
green_violet = dt.GetRasterBand(6).ReadAsArray(0,0,dt.RasterXSize, dt.RasterYSize)
yellow_violet = dt.GetRasterBand(7).ReadAsArray(0,0,dt.RasterXSize, dt.RasterYSize)

Blue = np.array(blue)
B = Blue.ravel()
Green = np.array(green)
G = Green.ravel()
Yellow = np.array(yellow)
Y = Yellow.ravel()
Blue_violet = np.array(blue_violet)
BV = Blue_violet.ravel()
Blue_green = np.array(blue_green)
BG = Blue_green.ravel()
Green_violet = np.array(green_violet)
GV = Green_violet.ravel()
Yellow_violet = np.array(yellow_violet)
YV = Yellow_violet.ravel()

X_test = np.vstack((B,G,Y,BV,BG,GV,YV))
print(X_test.shape)

del blue,green,yellow,blue_violet,blue_green,green_violet,yellow_violet
del Blue,Green,Yellow,Blue_violet,Blue_green,Green_violet,Yellow_violet
del B,G,Y,BV,BG,GV,YV
gc.collect()

# Turn nan to 0
X_test2 = []
for i in range(X_test.shape[0]):
    for j in range(X_test.shape[1]):
        if math.isnan(X_test[i,j]) is True:
            X_test2.append(0)
        else:
            X_test2.append(X_test[i,j])
X_test1 = np.array(X_test2).reshape(7,46623690).transpose()

del X_test
del X_test2
gc.collect()

X_test3 = r['as.matrix'](X_test1)
Y_pred = r['predict'](regr,X_test3)
print(Y_pred)
z = np.array(Y_pred).reshape(5898,7905)

# 导出预测结果
driver = gdal.GetDriverByName('GTiff')
outDataSet = driver.Create('H:/bathmetry/GF6/Cubist1.tif',            dt.RasterXSize, dt.RasterYSize, 1)
outBand = outDataSet.GetRasterBand(1)
outBand.WriteArray(z,0,0)
outDataSet.SetProjection(dt.GetProjection())
outDataSet.SetGeoTransform(dt.GetGeoTransform()) #写入仿射变换参数