"""
# Final project: Fault Segmentation 
### This assignment is to implement convolutional neural networks for fault data segmentation
### AI in geosciences, ESS1502
"""
"""
#### Task: Train and validate your networks/models on the provide dataset. You are free to use Tensorflow, Keras, or Pytorch to implement your networks. The test result need to upload to the kaggle platform to get the final score 
"""
from modelarts.session import Session
session = Session()
session.download_data(bucket_path="20ai-project/data.npy", path="/home/ma-user/work/data.npy")
session.download_data(bucket_path="20ai-project/label.npy", path="/home/ma-user/work/label.npy")
session.download_data(bucket_path="20ai-project/new_test_data.npy", path="/home/ma-user/work/new_test_data.npy")
# session.download_data(bucket_path="20ai-project/new_test_label.npy", path="/home/ma-user/work/new_test_label.npy")
"""
## 1. Import the requirements
"""
import os
import random
import numpy as np
import skimage
import pandas as pd
import matplotlib.pyplot as plt
from numpy.random import seed
from keras import backend as K
from keras.models import *
from keras.layers import *
from keras.optimizers import *
from keras.callbacks import ModelCheckpoint, LearningRateScheduler, ReduceLROnPlateau, TensorBoard
from keras import backend as keras
from keras.models import load_model
from keras.utils import to_categorical
# os.environ["CUDA_VISIBLE_DEVICES"] = "0"
seed(12345)
os.mkdir('fig')
os.mkdir('check1channel')
"""
## 2. Load the train dataset
"""
# select your own datapath
data = np.load('data.npy')
label = np.load('label.npy')
"""
### 2.1 visualize the train dataset
"""
j = 1500
plt.figure(figsize=(8,4),dpi=200)
plt.subplot(1,2,1)
plt.imshow(np.squeeze(data[j,:,:]),cmap=plt.cm.gray)
plt.subplot(1,2,2)
plt.imshow(np.squeeze(label[j,:,:]),interpolation="bilinear",vmin=0.2,vmax=1.0,cmap=plt.cm.gray)
figname = './fig/data_' +str(j)+'.jpg'
plt.savefig(figname)
def Z_ScoreNormalization(x): #灰度值归一化
    for i in range(len(x)):
        mu = np.average(x[i,:])
        sigma = np.std(x[i,:])
        x[i,:] = (x[i,:] - mu) / sigma;
    return x;
# normalize data
data = Z_ScoreNormalization(data)
data = np.reshape(data,(2000,128,128,1))
label = np.reshape(label,(2000,128,128,1))
# split the train_data and val_data
train_data = data[0:1600]
val_data = data[1600:2000]
train_label = label[0:1600]
val_label = label[1600:2000]
print(train_data.shape)
print(val_data.shape)
print(train_label.shape)
print(val_label.shape)
"""
## 3. Define the neural network
"""
#given network
def unet(pretrained_weights=None, input_size1=(None, None, 1)):
    input = Input(input_size1, name='input')
    conv1 = Conv2D(16, (3,3), activation='relu', padding='same')(input)
    conv1 = Conv2D(16, (3,3), activation='relu', padding='same')(conv1)
    pool1 = MaxPooling2D(pool_size=(2,2))(conv1)

    conv2 = Conv2D(32, (3,3), activation='relu', padding='same')(pool1)
    conv2 = Conv2D(32, (3,3), activation='relu', padding='same')(conv2)
    pool2 = MaxPooling2D(pool_size=(2,2))(conv2)

    conv3 = Conv2D(64, (3,3), activation='relu', padding='same')(pool2)
    conv3 = Conv2D(64, (3,3), activation='relu', padding='same')(conv3)
    pool3 = MaxPooling2D(pool_size=(2,2))(conv3)

    conv4 = Conv2D(128, (3,3), activation='relu', padding='same')(pool3)
    conv4 = Conv2D(128, (3,3), activation='relu', padding='same')(conv4)

    up5 = concatenate([UpSampling2D(size=(2,2))(conv4), conv3], axis=-1)
    conv5 = Conv2D(64, (3,3), activation='relu', padding='same')(up5)
    conv5 = Conv2D(64, (3,3), activation='relu', padding='same')(conv5)

    up6 = concatenate([UpSampling2D(size=(2,2))(conv5), conv2], axis=-1)
    conv6 = Conv2D(32, (3,3), activation='relu', padding='same')(up6)
    conv6 = Conv2D(32, (3,3), activation='relu', padding='same')(conv6)

    up7 = concatenate([UpSampling2D(size=(2,2))(conv6), conv1], axis=-1)
    conv7 = Conv2D(16, (3,3), activation='relu', padding='same')(up7)
    conv7 = Conv2D(16, (3,3), activation='relu', padding='same')(conv7)

    conv8 = Conv2D(1, (1,1), activation='sigmoid')(conv7)
    model = Model(inputs=[input], outputs=[conv8])
    model.summary()
    return model
#one more layer
#each layer one more convlution with a (1,1) kernel
def unet(pretrained_weights=None, input_size1=(None, None, 1)):
    input = Input(input_size1, name='input')
    conv1 = Conv2D(16, (3,3), activation='relu', padding='same')(input)
    conv1 = Conv2D(16, (3,3), activation='relu', padding='same')(conv1)
    conv1 = Conv2D(16, (3,3), activation='relu', padding='same')(conv1)
    conv1 = Conv2D(16, (3,3), activation='relu', padding='same')(conv1)
    conv1 = Conv2D(16, (3,3), activation='relu', padding='same')(conv1)
    conv1 = Conv2D(16, (1,1), activation='relu', padding='same')(conv1)
    pool1 = MaxPooling2D(pool_size=(2,2))(conv1)

    conv2 = Conv2D(32, (3,3), activation='relu', padding='same')(pool1)
    conv2 = Conv2D(32, (3,3), activation='relu', padding='same')(conv2)
    conv2 = Conv2D(32, (3,3), activation='relu', padding='same')(conv2)
    conv2 = Conv2D(32, (3,3), activation='relu', padding='same')(conv2)
    conv2 = Conv2D(32, (3,3), activation='relu', padding='same')(conv2)
    conv2 = Conv2D(32, (1,1), activation='relu', padding='same')(conv2)
    pool2 = MaxPooling2D(pool_size=(2,2))(conv2)

    conv3 = Conv2D(64, (3,3), activation='relu', padding='same')(pool2)
    conv3 = Conv2D(64, (3,3), activation='relu', padding='same')(conv3)
    conv3 = Conv2D(64, (3,3), activation='relu', padding='same')(conv3)
    conv3 = Conv2D(64, (3,3), activation='relu', padding='same')(conv3)
    conv3 = Conv2D(64, (3,3), activation='relu', padding='same')(conv3)
    conv3 = Conv2D(64, (1,1), activation='relu', padding='same')(conv3)
    pool3 = MaxPooling2D(pool_size=(2,2))(conv3)

    conv4 = Conv2D(128, (3,3), activation='relu', padding='same')(pool3)
    conv4 = Conv2D(128, (3,3), activation='relu', padding='same')(conv4)
    conv4 = Conv2D(128, (3,3), activation='relu', padding='same')(conv4)
    conv4 = Conv2D(128, (3,3), activation='relu', padding='same')(conv4)
    conv4 = Conv2D(128, (3,3), activation='relu', padding='same')(conv4)
    conv4 = Conv2D(128, (1,1), activation='relu', padding='same')(conv4)
    pool4 = MaxPooling2D(pool_size=(2,2))(conv4)
    
    conv5 = Conv2D(256, (3,3), activation='relu', padding='same')(pool4)
    conv5 = Conv2D(256, (3,3), activation='relu', padding='same')(conv5)
    conv5 = Conv2D(256, (3,3), activation='relu', padding='same')(conv5)
    conv5 = Conv2D(256, (3,3), activation='relu', padding='same')(conv5)
    conv5 = Conv2D(256, (3,3), activation='relu', padding='same')(conv5)
    conv5 = Conv2D(256, (1,1), activation='relu', padding='same')(conv5)

    up6 = concatenate([UpSampling2D(size=(2,2))(conv5), conv4], axis=-1)
    conv6 = Conv2D(128, (3,3), activation='relu', padding='same')(up6)
    conv6 = Conv2D(128, (3,3), activation='relu', padding='same')(conv6)
    conv6 = Conv2D(128, (3,3), activation='relu', padding='same')(conv6)
    conv6 = Conv2D(128, (3,3), activation='relu', padding='same')(conv6)
    conv6 = Conv2D(128, (3,3), activation='relu', padding='same')(conv6)
    conv6 = Conv2D(128, (1,1), activation='relu', padding='same')(conv6)

    up7 = concatenate([UpSampling2D(size=(2,2))(conv6), conv3], axis=-1)
    conv7 = Conv2D(64, (3,3), activation='relu', padding='same')(up7)
    conv7 = Conv2D(64, (3,3), activation='relu', padding='same')(conv7)
    conv7 = Conv2D(64, (3,3), activation='relu', padding='same')(conv7)
    conv7 = Conv2D(64, (3,3), activation='relu', padding='same')(conv7)
    conv7 = Conv2D(64, (3,3), activation='relu', padding='same')(conv7)
    conv7 = Conv2D(64, (1,1), activation='relu', padding='same')(conv7)

    up8 = concatenate([UpSampling2D(size=(2,2))(conv7), conv2], axis=-1)
    conv8 = Conv2D(32, (3,3), activation='relu', padding='same')(up8)
    conv8 = Conv2D(32, (3,3), activation='relu', padding='same')(conv8)
    conv8 = Conv2D(32, (3,3), activation='relu', padding='same')(conv8)
    conv8 = Conv2D(32, (3,3), activation='relu', padding='same')(conv8)
    conv8 = Conv2D(32, (3,3), activation='relu', padding='same')(conv8)
    conv8 = Conv2D(32, (1,1), activation='relu', padding='same')(conv8)
    
    up9 = concatenate([UpSampling2D(size=(2,2))(conv8), conv1], axis=-1)
    conv9 = Conv2D(16, (3,3), activation='relu', padding='same')(up9)
    conv9 = Conv2D(16, (3,3), activation='relu', padding='same')(conv9)
    conv9 = Conv2D(16, (3,3), activation='relu', padding='same')(conv9)
    conv9 = Conv2D(16, (3,3), activation='relu', padding='same')(conv9)
    conv9 = Conv2D(16, (3,3), activation='relu', padding='same')(conv9)
    conv9 = Conv2D(16, (1,1), activation='relu', padding='same')(conv9)

    conv10 = Conv2D(1, (1,1), activation='sigmoid')(conv9)
    model = Model(inputs=[input], outputs=[conv10])
    model.summary()
    return model
#one more layer
#begin with 32
#each layer one more convlution with a (1,1) kernel
def unet(pretrained_weights=None, input_size1=(None, None, 1)):
    input = Input(input_size1, name='input')
    conv1 = Conv2D(32, (3,3), activation='relu', padding='same')(input)
    conv1 = Conv2D(32, (3,3), activation='relu', padding='same')(conv1)
    conv1 = Conv2D(32, (3,3), activation='relu', padding='same')(conv1)
    conv1 = Conv2D(32, (3,3), activation='relu', padding='same')(conv1)
    conv1 = Conv2D(32, (3,3), activation='relu', padding='same')(conv1)
    conv1 = Conv2D(32, (1,1), activation='relu', padding='same')(conv1)
    pool1 = MaxPooling2D(pool_size=(2,2))(conv1)

    conv2 = Conv2D(64, (3,3), activation='relu', padding='same')(pool1)
    conv2 = Conv2D(64, (3,3), activation='relu', padding='same')(conv2)
    conv2 = Conv2D(64, (3,3), activation='relu', padding='same')(conv2)
    conv2 = Conv2D(64, (3,3), activation='relu', padding='same')(conv2)
    conv2 = Conv2D(64, (3,3), activation='relu', padding='same')(conv2)
    conv2 = Conv2D(64, (1,1), activation='relu', padding='same')(conv2)
    pool2 = MaxPooling2D(pool_size=(2,2))(conv2)

    conv3 = Conv2D(128, (3,3), activation='relu', padding='same')(pool2)
    conv3 = Conv2D(128, (3,3), activation='relu', padding='same')(conv3)
    conv3 = Conv2D(128, (3,3), activation='relu', padding='same')(conv3)
    conv3 = Conv2D(128, (3,3), activation='relu', padding='same')(conv3)
    conv3 = Conv2D(128, (3,3), activation='relu', padding='same')(conv3)
    conv3 = Conv2D(128, (1,1), activation='relu', padding='same')(conv3)
    pool3 = MaxPooling2D(pool_size=(2,2))(conv3)

    conv4 = Conv2D(256, (3,3), activation='relu', padding='same')(pool3)
    conv4 = Conv2D(256, (3,3), activation='relu', padding='same')(conv4)
    conv4 = Conv2D(256, (3,3), activation='relu', padding='same')(conv4)
    conv4 = Conv2D(256, (3,3), activation='relu', padding='same')(conv4)
    conv4 = Conv2D(256, (3,3), activation='relu', padding='same')(conv4)
    conv4 = Conv2D(256, (1,1), activation='relu', padding='same')(conv4)
    pool4 = MaxPooling2D(pool_size=(2,2))(conv4)
    
    conv5 = Conv2D(512, (3,3), activation='relu', padding='same')(pool4)
    conv5 = Conv2D(512, (3,3), activation='relu', padding='same')(pool4)
    conv5 = Conv2D(512, (3,3), activation='relu', padding='same')(conv5)
    conv5 = Conv2D(512, (3,3), activation='relu', padding='same')(conv5)
    conv5 = Conv2D(512, (3,3), activation='relu', padding='same')(conv5)
    conv5 = Conv2D(512, (1,1), activation='relu', padding='same')(conv5)

    up6 = concatenate([UpSampling2D(size=(2,2))(conv5), conv4], axis=-1)
    conv6 = Conv2D(256, (3,3), activation='relu', padding='same')(up6)
    conv6 = Conv2D(256, (3,3), activation='relu', padding='same')(conv6)
    conv6 = Conv2D(256, (3,3), activation='relu', padding='same')(conv6)
    conv6 = Conv2D(256, (3,3), activation='relu', padding='same')(conv6)
    conv6 = Conv2D(256, (3,3), activation='relu', padding='same')(conv6)
    conv6 = Conv2D(256, (1,1), activation='relu', padding='same')(conv6)

    up7 = concatenate([UpSampling2D(size=(2,2))(conv6), conv3], axis=-1)
    conv7 = Conv2D(128, (3,3), activation='relu', padding='same')(up7)
    conv7 = Conv2D(128, (3,3), activation='relu', padding='same')(conv7)
    conv7 = Conv2D(128, (3,3), activation='relu', padding='same')(conv7)
    conv7 = Conv2D(128, (3,3), activation='relu', padding='same')(conv7)
    conv7 = Conv2D(128, (3,3), activation='relu', padding='same')(conv7)
    conv7 = Conv2D(128, (1,1), activation='relu', padding='same')(conv7)

    up8 = concatenate([UpSampling2D(size=(2,2))(conv7), conv2], axis=-1)
    conv8 = Conv2D(64, (3,3), activation='relu', padding='same')(up8)
    conv8 = Conv2D(64, (3,3), activation='relu', padding='same')(conv8)
    conv8 = Conv2D(64, (3,3), activation='relu', padding='same')(conv8)
    conv8 = Conv2D(64, (3,3), activation='relu', padding='same')(conv8)
    conv8 = Conv2D(64, (3,3), activation='relu', padding='same')(conv8)
    conv8 = Conv2D(64, (1,1), activation='relu', padding='same')(conv8)
    
    up9 = concatenate([UpSampling2D(size=(2,2))(conv8), conv1], axis=-1)
    conv9 = Conv2D(32, (3,3), activation='relu', padding='same')(up9)
    conv9 = Conv2D(32, (3,3), activation='relu', padding='same')(conv9)
    conv9 = Conv2D(32, (3,3), activation='relu', padding='same')(conv9)
    conv9 = Conv2D(32, (3,3), activation='relu', padding='same')(conv9)
    conv9 = Conv2D(32, (3,3), activation='relu', padding='same')(conv9)
    conv9 = Conv2D(32, (1,1), activation='relu', padding='same')(conv9)

    conv10 = Conv2D(1, (1,1), activation='sigmoid')(conv9)
    model = Model(inputs=[input], outputs=[conv10])
    model.summary()
    return model
model = unet()
model.compile(optimizer=Adam(lr=5e-4), loss='binary_crossentropy',metrics=['accuracy'])
# checkpoint
filepath="check1channel/fault-{epoch:02d}.hdf5"
checkpoint = ModelCheckpoint(filepath, monitor='val_accuracy',verbose=1, save_best_only=False, mode='max')
callbacks_list = [checkpoint]
print("data prepared, ready to train!")
# Fit the model
history = model.fit(train_data, train_label,
#                    validation_split=0.2,
                    validation_data=(val_data,val_label),
                    epochs=40,
                    batch_size=8,
                    shuffle=True,
                    verbose=1,callbacks=callbacks_list)

"""
## 4. Plot the loss and acc
"""
def showHistory(history):
  # list all data in history
  print(history.history.keys())
  fig = plt.figure(figsize=(10,6))

  # summarize history for accuracy
  plt.plot(history.history['acc'])
  plt.plot(history.history['val_acc'])
  plt.title('Model accuracy',fontsize=20)
  plt.ylabel('Accuracy',fontsize=20)
  plt.xlabel('Epoch',fontsize=20)
  plt.legend(['train', 'test'], loc='center right',fontsize=20)
  plt.tick_params(axis='both', which='major', labelsize=18)
  plt.tick_params(axis='both', which='minor', labelsize=18)
  plt.show()

  # summarize history for loss
  fig = plt.figure(figsize=(10,6))
  plt.plot(history.history['loss'])
  plt.plot(history.history['val_loss'])
  plt.title('Model loss',fontsize=20)
  plt.ylabel('Loss',fontsize=20)
  plt.xlabel('Epoch',fontsize=20)
  plt.legend(['train', 'test'], loc='center right',fontsize=20)
  plt.tick_params(axis='both', which='major', labelsize=18)
  plt.tick_params(axis='both', which='minor', labelsize=18)
  plt.show()
showHistory(history)
"""
## 5. Validate the network using val_data
"""
model = load_model('./check1channel/fault-28.hdf5')
# need a threshold to get the true predict result 
threshold = 0.4
result = model.predict(val_data)
result[result>threshold]=1
result[result<threshold]=0
# np.save('./Data/data/result.npy',result)
j=200
plt.figure(figsize=(12,4),dpi=150)
plt.subplot(1,3,1)
plt.imshow(np.squeeze(val_data[j,:,:,:]),cmap=plt.cm.gray)
plt.subplot(1,3,2)
plt.imshow(np.squeeze(result[j,:,:,:]),interpolation="bilinear",vmin=0.2,vmax=1.0,cmap=plt.cm.gray)
plt.subplot(1,3,3)
plt.imshow(np.squeeze(val_label[j,:,:]),interpolation="bilinear",vmin=0.2,vmax=1.0,cmap=plt.cm.gray)
# savefig
figname = './fig/data_pred_nothreshold_' +str(j)+'.jpg'
plt.savefig(figname)
"""
## 6. Predictions on the testing datasets and visualize the prediction results
"""
new_test_data = np.load('new_test_data.npy')
# new_test_label = np.load('new_test_label.npy')
real_test_data = np.reshape(new_test_data,(100,256,640,1))
# real_test_label = np.reshape(new_test_label,(100,256,640,1))
temp = 0.25
real_result = model.predict(real_test_data)
real_result[real_result>temp]=1
real_result[real_result<temp]=0
result = np.squeeze(real_result)
# np.save('./Data/data/real_result.npy',real_result)
j=50
plt.figure(figsize=(6,6),dpi=150)
plt.subplot(2,1,1)
plt.imshow(np.squeeze(real_test_data[j,:,:,:]),cmap=plt.cm.gray)
plt.subplot(2,1,2)
plt.imshow(np.squeeze(real_result[j,:,:]),interpolation="bilinear",vmin=0.2,vmax=1.0,cmap=plt.cm.gray)
"""
## 7. Output the results as csv and upload it to the kaggle
"""
result_csv = np.reshape(result,(-1))
print(result_csv.shape)
df = pd.DataFrame(result_csv)
df=df.astype(int)
df=df.astype(str)
df['index'] = range(len(df))
new_col = ['value', 'index']
df.columns = new_col
order = ['index','value']
df=df[order]
df=df.astype(str)
print(df)
#choose your own name
df.to_csv('2_13_5.csv',index=False)
import moxing as mox
mox.file.copy('2_13_5.csv','obs://geo-ai-dataset/2_13_5.csv')
print("done")