import os
import numpy as np
from keras.layers import Input, Conv3D, LeakyReLU, BatchNormalization, Add, Activation
from keras.optimizers import Adam
from keras.callbacks import ReduceLROnPlateau, EarlyStopping
from keras.models import Model
from osgeo import gdal, gdalconst
import math
import datetime


class GRID:

    # 读图像文件
    def read_img(self, filename):
        dataset = gdal.Open(filename)
        # 打开文件
        im_width = dataset.RasterXSize  # 栅格矩阵的列数
        im_height = dataset.RasterYSize  # 栅格矩阵的行数

        im_geotrans = dataset.GetGeoTransform()  # 仿射矩阵
        im_proj = dataset.GetProjection()  # 地图投影信息
        im_data = dataset.ReadAsArray(0, 0, im_width, im_height)  # 将数据写成数组，对应栅格矩阵
        im_data = np.array(im_data)
        sp = im_data.shape
        if im_data.ndim == 2:
            im_data2 = im_data[:, :, np.newaxis]
        else:
            im_data2 = np.zeros((im_height, im_width, sp[0]))
            for bands in range(0, sp[0]):
                im_data2[:, :, bands] = im_data[bands, :, :]

        del dataset
        return im_geotrans, im_proj, im_data2

    def read_group(self, dir, patchsize, pref, jg, scale, cutsp):
        lenstr = len(pref)
        filenames = os.listdir(dir)
        filenames.sort()
        print(filenames)
        filenums = len(filenames)
        h5f = []
        for fnum in range(0, filenums):
            ptclfile = filenames[fnum]
            if (ptclfile[-4:] != '.hdr') and (ptclfile[0:lenstr] == pref) and (ptclfile[-4:] != '.enp') and (
                    ptclfile[-1] != 'F'):
                print(ptclfile)
                if scale != 1:
                    im_geotrans, im_proj, img = self.reprojectRaster(dir + r'/' + ptclfile, 'temp.tif', scale)
                else:
                    im_geotrans, im_proj, img = self.read_img(dir + r'/' + ptclfile)
                if cutsp != []:
                    img = img[0:cutsp[0], 0:cutsp[1], 0:cutsp[2]]
                img = img / 1000
                sp = img.shape
                img_arr, orisp = self.prepare4d_over(img, patchsize, jg)
                del img
                img_arr = img_arr[:, np.newaxis, :, :, :]
                if h5f == []:
                    h5f = img_arr
                else:
                    h5f = np.concatenate((h5f, img_arr), axis=1)
                del img_arr
        print('shape: ' + str(h5f.shape))
        return im_geotrans, im_proj, h5f, sp

    def reprojectRaster(self, filename, dst_filename, scale):
        src = gdal.Open(filename)
        src_proj = src.GetProjection()
        src_geotrans = src.GetGeoTransform()
        wide = src.RasterXSize
        high = src.RasterYSize
        bands = src.RasterCount  # 波段数
        assert (np.sum(np.array(src.ReadAsArray(0, 0, wide, high) != 0)))
        driver = gdal.GetDriverByName('GTiff')
        dst = driver.Create(dst_filename, int(wide * scale), int(high * scale), int(bands), gdalconst.GDT_Float32)
        dst.SetGeoTransform(
            [src_geotrans[0], src_geotrans[1] / scale, src_geotrans[2], src_geotrans[3], src_geotrans[4],
             src_geotrans[5] / scale, ])
        dst.SetProjection(src_proj)

        gdal.ReprojectImage(src, dst, src_proj, src_proj, gdalconst.GRA_Cubic)

        im_width = dst.RasterXSize  # 栅格矩阵的列数
        im_height = dst.RasterYSize  # 栅格矩阵的行数

        im_geotrans = dst.GetGeoTransform()  # 仿射矩阵
        im_proj = dst.GetProjection()  # 地图投影信息
        im_data = dst.ReadAsArray(0, 0, im_width, im_height)  # 将数据写成数组，对应栅格矩阵
        im_data = np.array(im_data)
        sp = im_data.shape
        if im_data.ndim == 2:
            im_data2 = im_data[:, :, np.newaxis]
        else:
            im_data2 = np.zeros((im_height, im_width, sp[0]))
            for bands in range(0, sp[0]):
                im_data2[:, :, bands] = im_data[bands, :, :]
        print(im_data2.shape)
        del dst  # Flush
        assert (np.sum(im_data2 != 0))
        return im_geotrans, im_proj, im_data2

    def write_group(self, dir, h5f, im_proj, im_geotrans, prefixx, patchsize, jg, sp):

        for fnum in range(0, h5f.shape[1]):
            print(h5f[:, fnum, :, :, :].shape)
            img = self.restore4d_over(np.squeeze(h5f[:, fnum, :, :, :]), patchsize, sp, jg) * 1000
            self.write_img(dir + r'/' + prefixx + str((fnum + 1) * 2) + '.dat', im_proj, im_geotrans, img)
            del img

    # 写文件，以写成tif为例
    def write_img(self, filename, im_proj, im_geotrans, im_data2):
        # gdal数据类型包括
        # gdal.GDT_Byte,
        # gdal .GDT_UInt16, gdal.GDT_Int16, gdal.GDT_UInt32, gdal.GDT_Int32,
        # gdal.GDT_Float32, gdal.GDT_Float64
        sp = im_data2.shape
        if len(sp) > 2:
            im_data = np.zeros((sp[2], sp[0], sp[1]))
            for bands in range(0, sp[2]):
                im_data[bands, :, :] = im_data2[:, :, bands]
            print(im_data.shape)
        else:
            im_data = im_data2

        # 判断栅格数据的数据类型
        if 'int8' in im_data.dtype.name:
            datatype = gdal.GDT_Byte
        elif 'int16' in im_data.dtype.name:
            datatype = gdal.GDT_UInt16
        else:
            datatype = gdal.GDT_Float32

        # 判读数组维数
        if len(im_data.shape) == 3:
            im_bands, im_height, im_width = im_data.shape
        else:
            im_bands, (im_height, im_width) = 1, im_data.shape

        # 创建文件
        driver = gdal.GetDriverByName("ENVI")  # 数据类型必须有，因为要计算需要多大内存空间
        dataset = driver.Create(filename, im_width, im_height, im_bands, datatype)

        dataset.SetGeoTransform(im_geotrans)  # 写入仿射变换参数
        dataset.SetProjection(im_proj)  # 写入投影

        if im_bands == 1:
            dataset.GetRasterBand(1).WriteArray(im_data)  # 写入数组数据
        else:
            for i in range(im_bands):
                dataset.GetRasterBand(i + 1).WriteArray(im_data[i])

        del dataset

    def restore4d_over(self, data4d, patchsize, sp, jg):
        spp = data4d.shape
        sm = math.floor((sp[1] - 2 * jg) / (patchsize - 2 * jg))
        rex = np.zeros((math.floor(spp[0] / sm) * (patchsize - 2 * jg), sm * (patchsize - 2 * jg), sp[2]))
        for ii in range(0, spp[0]):
            mm = math.floor(ii / sm)
            nn = ii % sm

            rex[mm * (patchsize - 2 * jg):(mm + 1) * (patchsize - 2 * jg),
            nn * (patchsize - 2 * jg):(nn + 1) * (patchsize - 2 * jg), :] = data4d[ii, jg:patchsize - jg,
                                                                            jg:patchsize - jg, :]
        return rex

    def prepare4d_over(self, img, patchsize, jg):
        sp = img.shape
        ind2dscd = [-1]
        mnum = math.floor((sp[0] - 2 * jg) / (patchsize - 2 * jg))
        nnum = math.floor((sp[1] - 2 * jg) / (patchsize - 2 * jg))
        result = np.zeros((mnum * nnum, patchsize, patchsize, sp[2]))
        count = -1
        store = [[0, patchsize, 0, patchsize]]
        for i in range(0, mnum):
            for j in range(0, nnum):
                count += 1
                result[count, :, :, :] = img[(patchsize - 2 * jg) * i:(patchsize - 2 * jg) * i + patchsize,
                                         (patchsize - 2 * jg) * j: (patchsize - 2 * jg) * j + patchsize, :]
        return result, sp

    def restore2d(self, data3d, patchsize, sp):
        spp = data3d.shape
        sm = math.floor(sp[1] / patchsize)
        rex = np.zeros((sp[0], sp[1]))
        for ii in range(0, spp[2]):
            mm = math.floor(ii / sm)
            nn = ii % sm
            rex[mm * patchsize:(mm + 1) * patchsize, nn * patchsize:(nn + 1) * patchsize] = data3d[:, :, ii]
        return rex


class Model3dd():
    def __init__(self, dir_m, dir_l, tardir, patchsize1, patchsize2, epochs, jg1, jg2, dates, pairs):
        self.p1_model = None
        self.p2_model = None
        self.f_model = None
        self.grid = GRID()
        self.patchsize1 = patchsize1
        self.patchsize2 = patchsize2
        self.epochs = epochs
        self.mdir = dir_m
        self.ldir = dir_l

        self.bandsm = 6
        self.bandsl = 6
        self.bandsh = 15
        self.tardir = tardir
        self.jg1 = jg1
        self.jg2 = jg2
        self.dates = dates
        self.pairs = pairs
        '''shutil.rmtree(self.tardir) # 能删除该文件夹和文件夹下所有文件
        os.mkdir(self.tardir)
        self.hd5m = None
        self.hd5dl = None
        self.hd5l = None
        self.hd5ld = None
        self.hd5h = None'''

    def dtct_zero(self, h5data):
        ind4zero = [-1]
        for cnt in range(0, h5data.shape[0]):
            for time in range(0, h5data.shape[1]):
                if np.sum(h5data[cnt, time, :, :, :]) == 0:
                    ind4zero.append(cnt)
        return ind4zero[1:]

    def read_pdata(self):
        self.im_geotrans, self.im_proj, self.hd5mu, self.sp = self.grid.read_group(self.mdir, self.patchsize2, 'M',
                                                                                   self.jg2, 1, [])
        self.im_geotrans, self.im_proj, self.hd5l, self.sp = self.grid.read_group(self.ldir, self.patchsize2, 'L',
                                                                                  self.jg2, 1, [])
        id4zr_m = self.dtct_zero(self.hd5mu)
        id4zr_l = self.dtct_zero(self.hd5l)
        id2rmn = list(set(np.arange(0, self.hd5mu.shape[0])).difference(set(id4zr_m).union(set(id4zr_l))))
        print('dscd: ' + str(list(set(id4zr_l))))
        self.mtrain = self.hd5mu[id2rmn, :, :, :, :]
        self.ltrain = self.hd5l[id2rmn, :, :, :, :]
        print(self.hd5l.shape)
        self.bandsm = self.hd5l.shape[4]
        self.bandsl = self.hd5l.shape[4]

    def _convolutional_block(self, X, filters, stage, block):
        conv_name_base = 'res' + stage + block + '_branch'
        Conv_name_base = 'resde' + stage + block + '_branch'
        bn_name_base = 'bn' + stage + block + '_branch'
        lrelu_name_base = 'lrelu' + stage + block + '_branch'
        # Retrieve Filters
        F1, F2, F3 = filters
        # Save the input value
        X_shortcut = X
        # ##### MAIN PATH #####
        # First component of main path
        X = Conv3D(filters=F1, kernel_size=3, strides=1, name=conv_name_base + '2a', padding='same',
                   data_format='channels_last',
                   kernel_initializer='glorot_uniform')(X)
        X = BatchNormalization(axis=-1, name=bn_name_base + '2a')(X)
        X = LeakyReLU(alpha=0.3, name=lrelu_name_base + '1')(X)
        ### START CODE HERE ##
        # Second component of main path (≈3 lines)
        X = Conv3D(F2, 3, strides=1, name=conv_name_base + '2b', padding='same', data_format='channels_last',
                   kernel_initializer='glorot_uniform')(X)
        X = BatchNormalization(axis=-1, name=bn_name_base + '2b')(X)
        X = LeakyReLU(alpha=0.3, name=lrelu_name_base + '2')(X)
        # Third component of main path (≈2 lines)
        X = Conv3D(F1, 3, strides=1, name=conv_name_base + '2c', padding='same', data_format='channels_last',
                   kernel_initializer='glorot_uniform')(X)
        X = BatchNormalization(axis=-1, name=bn_name_base + '2c')(X)
        X = LeakyReLU(alpha=0.3, name=lrelu_name_base + '3')(X)
        # ha
        ##### SHORTCUT PATH #### (≈2 lines)
        '''X_shortcut = Conv2D(F1, 3, strides=(s, s), name=conv_name_base + '1', padding='same',
                            data_format='channels_last',
                            kernel_initializer='glorot_uniform')(X_shortcut)'''
        # X_shortcut = BatchNormalization(axis=-1, name=bn_name_base + '1')(X_shortcut)

        # Final step: Add shortcut value to main path, and pass it through a RELU activation (≈2 lines)
        X = Add()([X, X_shortcut])
        X = Activation('relu')(X)
        return X

    def create_model(self):
        m = Input(shape=(None, self.patchsize2, self.patchsize2, self.bandsm), name='l_like')
        x = Conv3D(filters=64, kernel_size=3, strides=1, padding='same',
                   data_format='channels_last',
                   kernel_initializer='glorot_uniform')(m)

        x = LeakyReLU(alpha=0.3, name='lrelu_p2_f1')(x)
        x = Conv3D(filters=32, kernel_size=3, strides=1, padding='same',
                   data_format='channels_last',
                   kernel_initializer='glorot_uniform')(x)
        x = LeakyReLU(alpha=0.3, name='lrelu_p2_f2')(x)
        x = Conv3D(filters=self.bandsl, kernel_size=3, strides=1, padding='same',
                   data_format='channels_last',
                   kernel_initializer='glorot_uniform')(x)
        x = LeakyReLU(alpha=0.3, name='lrelu_p2_f3')(x)
        self.st_model = Model(outputs=x, inputs=m)


    def arrangedates(self, data, jg, indd):
        for i in range(0, self.pairs):
            if i == 0:
                trainn = data[:, jg * (i + 1), :, :, :] - data[:, jg * i, :, :, :]
                trainn = trainn[:, np.newaxis, :, :, :]
            else:
                tempp = data[:, jg * (i + 1), :, :, :] - data[:, jg * i, :, :, :]
                tempp = tempp[:, np.newaxis, :, :, :]
                trainn = np.concatenate((trainn, tempp), axis=1)
        if indd == []:
            indd = np.arange(0, trainn.shape[0])
            np.random.shuffle(indd)
        trainn = trainn[indd, :, :, :, :]
        return trainn, indd

    def train_model(self):
        self.read_pdata()
        trainx, indd = self.arrangedates(self.hd5mu, 2, [])
        trainy, indd = self.arrangedates(self.hd5l, 1, indd)
        print('trainx: ' + str(trainx.shape))
        optim = Adam(lr=0.001)
        self.st_model.compile(loss='mse', optimizer=optim)
        if os.path.isfile('p_3ls_cl_dif_model' + str(patchsize2) + '.h5'):
            self.st_model.load_weights('p_3ls_cl_dif_model' + str(patchsize2) + '.h5')
        else:

            callbacks1 = EarlyStopping(monitor='loss', min_delta=0.001, patience=15, verbose=0, mode='auto',
                                       baseline=None, restore_best_weights=True)
            reduce_lr = ReduceLROnPlateau(monitor='loss', factor=0.2,
                                          patience=5, min_lr=0.00001)
            # history = self.st_model.fit(x=self.hd5mu[:,np.linspace(0,8,5,dtype='int8'),:,:,:], y=self.hd5l, batch_size=32, callbacks = [callbacks1,reduce_lr],epochs=self.epochs, validation_split=0.1)
            history = self.st_model.fit(x=trainx, y=trainy, batch_size=32, callbacks=[callbacks1, reduce_lr],
                                       epochs=self.epochs, validation_split=0.1)
            self.st_model.save_weights('p_3ls_cl_dif_model' + str(patchsize2) + '.h5')
            loss_history = history.history["loss"]
            val_loss_history = history.history['val_loss']
            numpy_loss_history = np.array(loss_history)
            numpy_val_loss = np.array(val_loss_history)
            np.savetxt('p_3lsdif_cl_loss_' + str(patchsize2) + '.txt', numpy_loss_history, delimiter=",")
            np.savetxt('p_3lsdif_cl_val_loss_' + str(patchsize2) + '.txt', numpy_val_loss, delimiter=",")
        #前时相预测
        self.predn1 = np.zeros(
            (self.hd5l.shape[0], self.pairs, self.hd5l.shape[2], self.hd5l.shape[3], self.hd5l.shape[4]))
        self.predn1 = self.hd5l[:, :-1, :, :, :] + self.st_model.predict(
            self.hd5mu[:, np.linspace(1, self.pairs * 2 - 1, self.pairs, dtype='int8'), :, :, :] -
            self.hd5mu[:, np.linspace(0, self.pairs * 2 - 2, self.pairs, dtype='int8'), :, :, :], verbose=0)
        self.predn2 = np.zeros(
            (self.hd5l.shape[0], self.pairs, self.hd5l.shape[2], self.hd5l.shape[3], self.hd5l.shape[4]))
        self.predn2 = self.hd5l[:, 1:, :, :, :] - self.st_model.predict(
            self.hd5mu[:, np.linspace(2, self.pairs * 2, self.pairs, dtype='int8'), :, :, :] - self.hd5mu[:,
                                                                                               np.linspace(1,
                                                                                                           self.pairs - 1,
                                                                                                           self.pairs,
                                                                                                           dtype='int8'),
                                                                                               :, :, :], verbose=0)
        # 前后时相共同预测
        '''self.predn = np.zeros((self.hd5l.shape[0], self.pairs, self.hd5l.shape[2], self.hd5l.shape[3],
                               self.hd5l.shape[4]))
        for i in range(0, self.pairs):
            w1 = (self.dates[2 * i + 2] - self.dates[2 * i]) / (self.dates[2 * i + 1] - self.dates[i * 2])
            w2 = (self.dates[2 * i + 2] - self.dates[2 * i]) / (self.dates[2 * i + 2] - self.dates[2 * i + 1])
            temp = (self.predn1[:, i, :, :, :] * w1 + self.predn2[:, i, :, :, :] * w2) / (w1 + w2)
            self.predn[:, i, :, :, :] = temp

        print('final: ' + str(self.predn.shape))
        self.grid.write_group(self.tardir, self.predn, self.im_proj, self.im_geotrans, 'res_vilo_', self.patchsize2,
                              self.jg2, self.sp)
        self.predn = (self.predn1 + self.predn2) / 2
        self.grid.write_group(self.tardir, self.predn, self.im_proj, self.im_geotrans, 'res_avg_', self.patchsize2,
                              self.jg2, self.sp)'''
        self.grid.write_group(self.tardir, self.predn1, self.im_proj, self.im_geotrans, 'res_sig1_', self.patchsize2,
                              self.jg2, self.sp)
        self.grid.write_group(self.tardir, self.predn2, self.im_proj, self.im_geotrans, 'res_sig2_', self.patchsize2,
                              self.jg2, self.sp)


if __name__ == '__main__':
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
    print(os.getcwd())
    dirall = r''
    os.chdir(dirall)
    # set parameters
    # 将mdir/lpdir/tardir放到dirall文件夹下
    mdir = r'modisf'#modis数据文件夹
    lpdir = r'landsatf'#landsat数据文件夹
    tardir = r'result'#结果存放文件夹

    patchsize1 = 16
    patchsize2 = 64
    epochs = 1000
    jg1 = 5
    jg2 = 5
    pairs = 8
    dates = []

    # cia
    # dates = [281, 290, 306, 313, 329, 338, 365 + 5, 365 + 12, 365 + 44, 365 + 53,
    #         365 + 69, 365 + 76, 365 + 92, 365 + 101, 365 + 108, 365 + 117, 365 + 124]
    # lgc
    # dates = [107, 123, 187, 219, 235, 299, 331, 347, 362, 379, 395, 411, 427, 459]

    starttime = datetime.datetime.now()
    model3d = Model3dd(mdir,lpdir,tardir,patchsize1,patchsize2,epochs,jg1,jg2,dates,pairs)
    print('data all read.')
    print('train p')
    model3d.create_model()
    model3d.train_model()
    endtime = datetime.datetime.now()
    print('total time is ' + str((endtime - starttime).seconds) + 's')
    np.savetxt('time_cst_' + str(patchsize2) + '.txt', np.array([(endtime - starttime).seconds]), delimiter=",")
    print('Finish!')