!pip install -q -U aicrowd-cli

%load_ext aicrowd.magic

!pip install sweetviz
!pip install -U jupyter

import sweetviz as sv

import os

# Please use the absolute for the location of the pip install Shapelydataset.
# Or you can use relative path with `os.getcwd() + "test_data/validation.csv"`
AICROWD_DATASET_PATH = os.getenv("DATASET_PATH", "/ds_shared_drive/validation.csv")
AICROWD_PREDICTIONS_PATH = os.getenv("PREDICTIONS_PATH", "predictions.csv")
AICROWD_ASSETS_DIR = "assets"

#!pip install ipywidgets
#!jupyter nbextension enable --py widgetsnbextension
#!conda install -y jupyterlab_widgets
#!pip install aquirdturtle_collapsible_headings

!pip install numpy pandas
!pip install -U imbalanced-learn
!pip install xgboost
!pip install lightgbm
!pip install catboost
!pip install tensorflow
!pip install shap
!pip install torch torchvision torchaudio

import numpy as np
import pandas as pd
import joblib
import matplotlib.pyplot as plt
from collections import Counter
import torch
from tqdm.notebook import tqdm
%matplotlib inline

target_col = "diagnosis"
key_col = "row_id"
cat_cols = ['intersection_pos_rel_centre']
seed = 2021

target_values = ["normal", "post_alzheimer", "pre_alzheimer"]

scale = 4
translator2d = {1: [4,1], 2 : [5,2], 3: [6,3], 4:[5,4] ,5:[4,5] ,6:[3,6] ,7:[2,5] ,8:[1,4] ,9:[0,3],10:[1,2], 11:[2,1] , 12:[3,0]}
ccw_translate = {i: 12 - i  for i in range(1,13,1)}
ccw_translate[12] = 12
translator2d_ccw = {ccw_translate[k]:v for k,v in translator2d.items()}

import torchvision
import torch, torch.nn as nn
import torchvision.models as models
from torch.autograd import Variable
import torch.nn.functional as F
import torch.optim as optim
from sklearn.model_selection import train_test_split
from sklearn.metrics import f1_score
np.random.seed(0)
torch.manual_seed(0)

z_dim = 1 # image_repr_features.shape[1]
n_classes = 3
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(z_dim, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, n_classes)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.softmax(x)
model = Net()

train = pd.read_csv('/ds_shared_drive/train.csv')

# valid = pd.read_csv('/ds_shared_drive/validation.csv')
# valid_truth = pd.read_csv('/ds_shared_drive/validation_ground_truth.csv')
# valid_all = valid.merge(valid_truth,how='left')
# train = pd.concat([train, valid_all],axis = 0)

train = train[train[target_col].isin(target_values)].copy().reset_index(drop=True)

# Remove Constant Columns
train = train.loc[:, (train != train.iloc[0]).any()]
features = train.columns[1:-1].to_list()

numeric_features = [c for c in features if c not in cat_cols]

for c in numeric_features:
    train[c] = train[c].astype(float)

print(train[target_col].value_counts())
print(train.shape)

df_pos = train[train[target_col].isin(target_values[1:])]
nb_pos = df_pos.shape[0]
nb_neg = nb_pos*2
df_neg = train[train[target_col] == "normal"].sample(n=nb_neg, random_state=seed)
# df_neg = df_normal 
df_samples = pd.concat([df_pos, df_neg]).sample(frac=1).reset_index(drop=True)
# df_samples = train
df_samples.shape

df_samples.shape

print(cat_cols)
for c in cat_cols:
    df_samples[c].fillna("NA", inplace=True)
    
df_dummies = pd.get_dummies(df_samples[cat_cols], columns=cat_cols, dummy_na=True).add_prefix('CAT_')
dummy_cols = df_dummies.columns.to_list()
print(dummy_cols)

df_samples = pd.concat([df_samples, df_dummies], axis=1)
df_samples['cnt_NaN'] = df_samples[numeric_features].isna().sum(axis=1)
# df_samples.fillna(-1, inplace=True)
model_features = df_samples.columns.to_list()
model_features = [c for c in model_features if c not in [key_col, target_col] + cat_cols]
print(len(model_features))
X_train = df_samples[model_features]
y_train_all = df_samples[target_col].map(dict(zip(target_values, list(range(len(target_values))))))

df_samples[target_col].value_counts()

image_repr_features = None
for n,row in tqdm(X_train.iterrows()):
    image_repr = np.zeros((1,7,7))
    centre_repr = np.zeros((1,7,7))
    for i in range(1,13):
        col = f'missing_digit_{i}'
        present = row[col]

        if present:
            ccw_flag = row["sequence_flag_ccw"] == 1
            translator = translator2d
            if ccw_flag:
                translator = translator2d_ccw
            pos = translator[i]
            image_repr[0,pos[1],pos[0]] = 1
    
    image_repr = np.kron(image_repr, np.ones((scale,scale)))
#     rot_angle_z = image_repr * row["final_rotation_angle"]/360
#     centre_dot = row["centre_dot_detect"]
#     if centre_dot == 1:
#         centre_repr[0,3,3] = 1
#     centre_repr = np.kron(centre_repr, np.ones((scale,scale))) 
#     image_repr = np.vstack([image_repr,rot_angle_z,centre_repr])
    image_repr = np.expand_dims(image_repr, axis = 0)
    if n > 0:
        image_repr_features = np.vstack([image_repr_features,image_repr])
    else:
        image_repr_features = image_repr

image_repr_features_no_nan = image_repr_features # np.nan_to_num(image_repr_features)

model = Net()

train_x, val_x, train_y, val_y = train_test_split(image_repr_features_no_nan, y_train_all, test_size = 0.1)
(train_x.shape, train_y.shape), (val_x.shape, val_y.shape)
train_x = torch.from_numpy(train_x).float()
val_x = torch.from_numpy(val_x).float()
train_y = torch.from_numpy(train_y.values).long()
val_y = torch.from_numpy(val_y.values).long()
def train(epoch):
    tr_loss = 0
    # getting the training set
    x_train, y_train = Variable(train_x), Variable(train_y)
    # getting the validation set
    x_val, y_val = Variable(val_x), Variable(val_y)
    # clearing the Gradients of the model parameters
    optimizer.zero_grad()
    
    # prediction for training and validation set
    output_train = model(x_train)
    output_val = model(x_val)

    # computing the training and validation loss
    loss_train = criterion(output_train, y_train)
    loss_val = criterion(output_val, y_val)
    train_losses.append(loss_train)
    val_losses.append(loss_val)

    # computing the updated weights of all the model parameters
    loss_train.backward()
    optimizer.step()
    tr_loss = loss_train.item()
    if epoch%2 == 0:
        # printing the validation loss
        print('Epoch : ',epoch+1, '\t', 'loss :', loss_val)


# defining the optimizer
optimizer = optim.Adam(model.parameters(), lr=0.005)
# defining the loss function
criterion = nn.CrossEntropyLoss()
# defining the number of epochs
n_epochs = 15
# empty list to store training losses
train_losses = []
# empty list to store validation losses
val_losses = []
# training the model
for epoch in range(n_epochs):
    train(epoch)
# prediction for training set
with torch.no_grad():
    output = model(train_x)
    
softmax = output.cpu()
prob = list(softmax.numpy())
predictions = np.argmax(prob, axis=1)
print(predictions.sum())
# f1 score on training set
f1_score(train_y.numpy(), predictions, average='weighted')

# load your data

train_x, val_x, train_y, val_y = train_test_split(image_repr_features_no_nan, y_train_all, test_size = 0.1)
(train_x.shape, train_y.shape), (val_x.shape, val_y.shape)
train_x = torch.from_numpy(train_x).float()
val_x = torch.from_numpy(val_x).float()
train_y = torch.from_numpy(train_y.values).long()
val_y = torch.from_numpy(val_y.values).long()
def train(epoch):
    tr_loss = 0
    # getting the training set
    x_train, y_train = Variable(train_x), Variable(train_y)
    # getting the validation set
    x_val, y_val = Variable(val_x), Variable(val_y)
    # clearing the Gradients of the model parameters
    optimizer.zero_grad()
    
    # prediction for training and validation set
    output_train = model(x_train)
    output_val = model(x_val)

    # computing the training and validation loss
    loss_train = criterion(output_train, y_train)
    loss_val = criterion(output_val, y_val)
    train_losses.append(loss_train)
    val_losses.append(loss_val)

    # computing the updated weights of all the model parameters
    loss_train.backward()
    optimizer.step()
    tr_loss = loss_train.item()
    if epoch%2 == 0:
        # printing the validation loss
        print('Epoch : ',epoch+1, '\t', 'loss :', loss_val)

model = Net()
# defining the optimizer
optimizer = optim.Adam(model.parameters(), lr=0.005)
# defining the loss function
criterion = nn.CrossEntropyLoss()
# defining the number of epochs
n_epochs = 15
# empty list to store training losses
train_losses = []
# empty list to store validation losses
val_losses = []
# training the model
for epoch in range(n_epochs):
    train(epoch)
# prediction for training set
with torch.no_grad():
    output = model(train_x)
    
softmax = output.cpu()
prob = list(softmax.numpy())
predictions = np.argmax(prob, axis=1)
print(predictions.sum())
# f1 score on training set
f1_score(train_y.numpy(), predictions, average='weighted')

filename = f'{AICROWD_ASSETS_DIR}/model_checkpoint'

check_point = {'params': model.state_dict(),
              'optimizer': optimizer.state_dict()}

torch.save(check_point, filename)

file = f'{AICROWD_ASSETS_DIR}/model_checkpoint'
check_point = torch.load(file)
model.load_state_dict(check_point['params'])

test_data = pd.read_csv(AICROWD_DATASET_PATH)
test_data.head()

image_repr_features_test = None
for n,row in tqdm(test_data.iterrows()):
    image_repr = np.zeros((1,7,7))
    centre_repr = np.zeros((1,7,7))
    for i in range(1,13):
        col = f'missing_digit_{i}'
        present = row[col]

        if present:
            ccw_flag = row["sequence_flag_ccw"] == 1
            translator = translator2d
            if ccw_flag:
                translator = translator2d_ccw
            pos = translator[i]
            image_repr[0,pos[1],pos[0]] = 1
    
    image_repr = np.kron(image_repr, np.ones((scale,scale)))
#     rot_angle_z = image_repr * row["final_rotation_angle"]/360
#     centre_dot = row["centre_dot_detect"]
#     if centre_dot == 1:
#         centre_repr[0,3,3] = 1
#     centre_repr = np.kron(centre_repr, np.ones((scale,scale))) 
#     image_repr = np.vstack([image_repr,rot_angle_z,centre_repr])
    image_repr = np.expand_dims(image_repr, axis = 0)
    if n > 0:
        image_repr_features_test = np.vstack([image_repr_features_test,image_repr])
    else:
        image_repr_features_test = image_repr

image_repr_features_test_no_nan = image_repr_features_test # np.nan_to_num(image_repr_features)

# prediction for training set
test_x = torch.from_numpy(image_repr_features_test_no_nan).float()
with torch.no_grad():
    output = model(test_x)
    
softmax = output.cpu()
prob = list(softmax.numpy())
predictions = np.argmax(prob, axis=1)
print(predictions.sum())

predictions = {
    "row_id": test_data["row_id"].values,
    "normal_diagnosis_probability": [x[0] for x in prob],
    "post_alzheimer_diagnosis_probability": [x[1] for x in prob],
    "pre_alzheimer_diagnosis_probability": [x[2] for x in prob],
}

predictions_df = pd.DataFrame.from_dict(predictions)

predictions_df.head()

predictions_df.to_csv(AICROWD_PREDICTIONS_PATH, index=False)

!DATASET_PATH=$AICROWD_DATASET_PATH \
aicrowd notebook submit \
    --assets-dir $AICROWD_ASSETS_DIR \
    --challenge addi-alzheimers-detection-challenge