#@title
from IPython.display import HTML,clear_output
from base64 import b64encode
!gdown "https://drive.google.com/uc?id=1PuQU_NZzKYAhXMYLMBU1Ff3VQ02PuOs7"
mp4 = open('render.mp4','rb').read()
clear_output()
data_url = "data:video/mp4;base64," + b64encode(mp4).decode()
HTML("""
<video width=480 controls>
      <source src="%s" type="video/mp4">
</video>
""" % data_url)

!gdown "https://drive.google.com/uc?id=14u7fkARS8WRJUdhvU79kDxg8EKTqg606"
!gdown "https://drive.google.com/uc?id=1--tADAa10l2M1iaSEslGXK-RaBv8UbMf"

!pip install segmentation-models-pytorch==0.1.2  # easy to use some famous model architecture. visit https://github.com/qubvel/segmentation_models.pytorch/
!pip install albumentations               # easy image manipulation for data augmentation

import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import segmentation_models_pytorch as smp
import albumentations as A
import os
from ipywidgets import IntProgress
from IPython.display import display
import time
import cv2
os.environ['CUDA_LAUNCH_BLOCKING'] = "1"

train_data_full = np.load('data_train.npz', allow_pickle=True, mmap_mode='r')['data']
train_label_full = np.load('labels_train.npz', allow_pickle=True, mmap_mode='r')['labels']

print('shape            :',train_data_full.shape)
print('min-max amplitude:',train_data_full.min(),'&',train_data_full.max())

import math
from scipy.signal.windows import triang
from scipy.signal import convolve2d as conv2
def gain(data,dt,parameters):
    nt,nx = data.shape
    dout = np.zeros(data.shape)
    L = parameters/dt+1
    L = np.floor(L/2)
    h = triang(2*L+1)
    shaped_h  = h.reshape(len(h),1)
    for k in range(nx):
        aux = data[:,k]
        e = aux**2
        shaped_e = e.reshape(len(e),1)
        rms = np.sqrt(conv2(shaped_e,shaped_h,"same"))
        epsi = 1e-10*max(rms)
        op = rms/(rms**2+epsi)
        op = op.reshape(len(op),)
        dout[:,k] = data[:,k]*op
    for k in range(nx):
        aux = dout[:,k]
        amax = np.sqrt(sum(aux**2)/nt)
        dout[:,k] = dout[:,k]/amax
    return dout

test_proc=train_data_full[:,0,:]
dat_proc=gain(test_proc,3e-3,0.8)

fig, ax = plt.subplots(1, 2, figsize=(10,8))
ax[0].imshow(test_proc,interpolation='none',cmap='seismic')
ax[1].imshow(dat_proc,interpolation='none',cmap='seismic')

ax[0].set_title("Vanilla")
ax[1].set_title("Gain+RMS")
plt.show()

print('preprocessing the data :')
f = IntProgress(min=0, max=train_data_full.shape[1])
display(f)
for i in range(0,train_data_full.shape[1]):
    #print('reprocess : ',i+1,'of',train_data_full.shape[1])
    train_data_full[:,i,:]=gain(train_data_full[:,i,:],3e-3,0.8)
    f.value += 1

!gdown "https://drive.google.com/uc?id=1JZ5LZz_f2Vfg9BxuGGBY9LliJQAAHi_H"
train_data_full=np.load('data_train_processed.npz')['data']

train_data_full = (train_data_full - train_data_full.min()) / (train_data_full.max() - train_data_full.min())

fig = plt.figure(figsize=(10,5))
labels    = ["1", "2", "3", "4", "5", "6"]
colors = ['hotpink', 'lightskyblue', 'mediumpurple','cornsilk', 'pink', 'lightgrey']
N, bins, patches = plt.hist(train_label_full.flatten(),6,density=True, edgecolor='gray', linewidth=1)
for i in range(6):
    patches[i].set_facecolor(colors[i])
    patches[i].set_label(labels[i])
plt.gca().axes.xaxis.set_ticklabels([])
plt.title('Full Train Data Dist.')
plt.show()

from matplotlib.patches import Rectangle
fig, ax = plt.subplots(figsize=(5,10))
im = ax.imshow(train_label_full[:,-1,:],interpolation='none',cmap='Pastel1')
rect = plt.Rectangle((420,0),300,1500,facecolor='red',alpha=0.5)
ax.add_patch(rect)
plt.show()

split_sample = 0.7
test_data = train_data_full[:,:train_data_full.shape[1]-int(train_data_full.shape[1]*split_sample),:]
train_data = train_data_full[:,train_data_full.shape[1]-int(train_data_full.shape[1]*split_sample):,:]
test_label = train_label_full[:,:train_label_full.shape[1]-int(train_label_full.shape[1]*split_sample),:]
train_label = train_label_full[:,train_label_full.shape[1]-int(train_label_full.shape[1]*split_sample):,:]

print('train cube shape :',train_data.shape)
print('test cube shape :',test_data.shape)

fig = plt.figure(figsize=(10,10))
plt.imshow(train_label[:,-1,:],interpolation='none',cmap='Pastel1')
plt.colorbar()
plt.show()

fig, ax = plt.subplots(1, 2, figsize=(10,5))
labels    = ["1", "2", "3", "4", "5", "6"]
colors = ['hotpink', 'lightskyblue', 'mediumpurple','cornsilk', 'pink', 'lightgrey']
N, bins, patches = ax[0].hist(train_label.flatten(),[1, 2, 3, 4, 5, 6, 7],density=True, edgecolor='gray', linewidth=1)
for i in range(6):
    patches[i].set_facecolor(colors[i])
N2, bins2, patches2 = ax[1].hist(test_label.flatten(),[1, 2, 3, 4, 5, 6, 7],density=True, edgecolor='gray', linewidth=1)
for i in range(6):
    patches2[i].set_facecolor(colors[i])
    patches2[i].set_label(labels[i])
ax[0].get_xaxis().set_visible(False)
ax[1].get_xaxis().set_visible(False)
ax[0].set_title("Training Set Dist.")
ax[1].set_title("Testing Set Dist.")
plt.legend(title="Label")
plt.show()

print('label 5 count of train data : ',np.count_nonzero(train_label.flatten() == 5))
print('label 5 count of test data : ',np.count_nonzero(test_label.flatten() == 5))

from sklearn.metrics import multilabel_confusion_matrix
from sklearn.metrics import f1_score,accuracy_score  #for score metric calculation
def _prf_divide(numerator, denominator, ):
    """Performs division and handles divide-by-zero.
    On zero-division, sets the corresponding result elements equal to
    0 or 1 (according to ``zero_division``). 
    """
    mask = denominator == 0.0
    denominator = denominator.copy()
    denominator[mask] = 1  # avoid infs/nans
    result = numerator / denominator

    return result

def compute_scores(y_true, y_pred, class_weights=[1, 1, 1, 1, 20, 20]):
    """
    Computes the weighted & unweighted f1_score and accuracy
    Using the standard F1-Score and class-wise accuracy computations were quite 
    slow as we were doing a lot of redundant work across all score computations,
    hence we have implemented this from the base principles.
    Please refer to the inline comments.
    """

    # Initial Housekeeping Taks1
    y_true = np.array(y_true).flatten()
    y_pred = np.array(y_pred).flatten()
    class_weights = np.array(class_weights)
    # print(np.max(y_true))
    # print(np.max(y_pred))
    # print(np.min(y_true))
    # print(np.min(y_pred))
    # Computing Multilabel Confusion Matrix
    #print("--------- Computing MCM... ")
    #begin_time = time.time()
    MCM = multilabel_confusion_matrix(y_true, y_pred,labels=[1,2,3,4,5,6])
    #print("MCM computation time  : ", time.time() - begin_time)
    
    """
    Gather True Positives, True Negatives, False Positives, False Negatives
    """
    tp_sum = MCM[:, 1, 1]
    tn_sum = MCM[:, 0, 0]
    fn_sum = MCM[:, 1, 0]
    fp_sum = MCM[:, 0, 1]
    
    #print("--------- Computing per class instances... ")
    per_class_instances = np.bincount(y_true) # Helps keep a track of total number of instances per class
    per_class_instances = per_class_instances[1:] # as the class names in the dataset are NOT zero-indexed
    
    assert class_weights.shape == per_class_instances.shape
    
    #print("--------- Computing precision... ")
    # precision : tp / (tp + fp)
    precision = _prf_divide(
                    tp_sum,
                    (tp_sum + fp_sum)
                )
    #print("--------- Computing recall... ")                        
    # recall : tp / (tp + fn)
    recall = _prf_divide(
                    tp_sum,
                    (tp_sum + fn_sum)
                )

    #print("--------- Computing F1 score... ")
    # f1 : 2 * (recall * precision) / (recall + precision)
    f1_score = _prf_divide(
                    2 * precision * recall,
                    precision + recall
                )
    #print("--------- Computing Accuracy... ")
    # accuracy = tp_sum / instances_per_class
    # NOTE: we are computing the accuracy independently for all the class specific subgroups
    # accuracy = _prf_divide(
    #                 tp_sum,
    #                 per_class_instances
    #             )
    # print(class_weights)
    # print(f1_score)
    f1_score_weighted = np.dot(class_weights, f1_score) / np.sum(class_weights)
    f1_score_unweighted = f1_score.mean()

    # accuracy_weighted = np.dot(class_weights, accuracy) / np.sum(class_weights)
    # accuracy_unweighted = accuracy.mean()

    return f1_score_weighted, f1_score_unweighted#, accuracy_weighted, f1_score_unweighted, accuracy_unweighted

batch_size = 8      
num_epochs = 40      
num_classes = 6       
learning_rate = 0.00085

model = smp.PSPNet(
    encoder_name="efficientnet-b3",        
    encoder_weights=None,     
    in_channels=1,                  
    classes=num_classes,                    
)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
print("")

test = torch.rand(1, 1, 320, 320).cuda()
out = model(test)
out.shape

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(),lr = learning_rate)

class seisdataset_train(Dataset):
    def __init__(self, x_set, y_set):
        self.x, self.y = x_set, y_set
        self.n_sample = self.x.shape[1]+self.x.shape[2]+self.x.shape[1]
        self.aug = A.Compose([
            A.Resize(p=1, height=640, width=320, interpolation=1)
        ]) 
        self.aug2 = A.Compose([
            #A.RandomSizedCrop(p=1.0, min_max_height=(1006, 1006), height=1006, width=256, w2h_ratio=1.0, interpolation=0),
            #A.GridDistortion(p=0.3, num_steps=6, distort_limit=(-0.2, -0.05), border_mode=1),
            A.ElasticTransform(p=0.2,alpha=100, sigma=8, alpha_affine=0, border_mode=1),
            A.ShiftScaleRotate(p=0.5, shift_limit=(0.0, 0.0), scale_limit=(0.01, 0.25), rotate_limit=(-15, 15), interpolation=0, border_mode=1),
            A.RandomCrop(900, 250, p=0.3),
            A.Resize(p=1, height=640, width=320, interpolation=0)
            
        ])  

    def __len__(self):
        return self.n_sample
    
    def __getitem__(self, index):
        if index < self.x.shape[1]:
          idx = index
          batch_x = self.x[:,idx,:]
          batch_y = self.y[:,idx,:]-1
          augmented = self.aug(image=batch_x, mask=batch_y)
        elif self.x.shape[1] <= index < (self.x.shape[1]+self.x.shape[2]):
          idx = index-self.x.shape[1]
          batch_x = self.x[:,:,idx]
          batch_y = self.y[:,:,idx]-1
          augmented = self.aug2(image=batch_x, mask=batch_y)
        elif index >= self.x.shape[1]+self.x.shape[2]:
          idx = index-self.x.shape[1]-self.x.shape[2]
          batch_x = self.x[:,idx,:]
          batch_y = self.y[:,idx,:]-1
          augmented = self.aug2(image=batch_x, mask=batch_y)          
        image, mask = augmented['image'], augmented['mask']
        return image[None,:,:], mask