import os
import pandas as pd
import numpy as np
import keras.backend as K
import keras
import cv2

from keras.models import load_model

def grad_cam(model, category_index, layer_name):
    print(layer_name)

    inp = model.input
    y_c = model.output.op.inputs[0][0, category_index]
    A_k = model.get_layer(layer_name).output
    print(y_c, A_k)

    get_output = K.function(
        [inp], [A_k, K.gradients(y_c, A_k)[0], model.output])

    return get_output


def apply_grad_cam(get_output, input):
    input_image = input[0]

    [conv_output, grad_val, model_output] = get_output(input)

    conv_output = conv_output[0]
    grad_val = grad_val[0]

    weights = np.mean(grad_val, axis=(0, 1))

    grad_cam = np.zeros(dtype=np.float32, shape=conv_output.shape[0:2])
    for k, w in enumerate(weights):
        grad_cam += w * conv_output[:, :, k]

    grad_cam = np.maximum(grad_cam, 0)
    heatmap = grad_cam/np.max(grad_cam)

    image = input_image.squeeze()
    image -= np.min(image)
    image = 255 * image / np.max(image)
    image1 = cv2.resize(cv2.cvtColor(image.astype(np.uint8), cv2.COLOR_GRAY2RGB), (240, 1460),
                        interpolation=cv2.INTER_NEAREST)

    print(np.percentile(heatmap/np.max(heatmap), [0, 25, 50, 75, 100]))
    heatmap1 = cv2.resize(heatmap/np.max(heatmap),
                          (240, 1460), interpolation=cv2.INTER_LINEAR)

    grad_cam = cv2.applyColorMap(np.uint8(255*heatmap1), cv2.COLORMAP_JET)
    grad_cam = np.float32(grad_cam) + np.float32(image1)
    grad_cam = 255 * grad_cam / np.max(grad_cam)

    return np.uint8(grad_cam), heatmap


pca_exp = pd.read_excel("data/PCA_EXP.xlsx", header=None)
pca_cnv = pd.read_excel("data/PCA_CNV.xlsx", header=None)
pca_mt = pd.read_excel("data/PCA_MT.xlsx", header=None)
clinical = pd.read_excel("data/Clinical.xlsx")

# data size: 146 pathways, 5 PCs
n = len(pca_exp)  # sample size: number of Pts
path_n = 146  # number of pathways
pc = 5  # number of PCs

# data creation-EXP
pca_exp = pca_exp.to_numpy()
print(pca_exp.shape)
exp_data = np.zeros((n, path_n, pc))
for i in range(n):
    for j in range(path_n):
        exp_data[i, j, :] = pca_exp[i, j * pc:(j + 1) * pc]

# data creation-CNV
pca_cnv = pca_cnv.to_numpy()
cnv_data = np.zeros((n, path_n, pc))
for i in range(n):
    for j in range(path_n):
        cnv_data[i, j, :] = pca_cnv[i, j * pc:(j + 1) * pc]

# data creation-MT
pca_mt = pca_mt.to_numpy()
mt_data = np.zeros((n, path_n, pc))
for i in range(n):
    for j in range(path_n):
        mt_data[i, j, :] = pca_mt[i, j * pc:(j + 1) * pc]

# data merge: mRNA expression, CNV, and MT with a specific number of PCs
no_pc = 2  # use the first 2 PCs among 5 PCs
all_data = np.zeros((n, path_n, no_pc * 3))
for i in range(n):
    all_data[i, :, :] = np.concatenate((exp_data[i, :, 0:no_pc], cnv_data[i, :, 0:no_pc], mt_data[i, :, 0:no_pc]),
                                       axis=1)

index = np.arange(0, pca_exp.shape[1], pc)

clinical = clinical.to_numpy()
survival = clinical[:, 5]
os_months = clinical[:, 6]
idx0 = np.where((survival == 1) & (os_months <= 24))  # class0
idx1 = np.where(os_months > 24)                       # class1

all_data = all_data[:, :, :]

data_0 = all_data[idx0, :, :]
data_0 = data_0[0, :, :, :]
data_1 = all_data[idx1, :, :]
data_1 = data_1[0, :, :, :]

outcomes_0 = np.zeros(len(idx0[0]))
outcomes_1 = np.ones(len(idx1[0]))

# data merge
data = np.concatenate((data_0, data_1))
outcomes = np.concatenate((outcomes_0, outcomes_1))

print('test')
event_rate = np.array(
    [outcomes_0.shape[0], outcomes_1.shape[0]])/outcomes_0.shape[0]
print(event_rate)

vmin = -10
vmax = 10
m = 1

model_file = 'PathCNN_model.h5'

model = load_model(model_file)
layers = model.layers
print(model.summary())

orig_weights = model.get_weights().copy()

output_layer = model.get_layer('dense_2')
output_weights = output_layer.get_weights()
print(output_weights[1])

col = ['outcome', 'prediction', 'target_class',
       'probability']+list(range(1, no_pc*path_n*3+1))
print(col)

layer = layers[4]  # the last layer before flatten
print(layer.name)
weights = layer.get_weights()

get_output = [None, None]
for target_class in range(2):
  get_output[target_class] = grad_cam(model, target_class, layer.name)

os.makedirs('output/'+layer.name, exist_ok=True)
for i in range(0, data.shape[0], 1):
    oimg = data[i, :, :]

    print(f'\ncase {i:03d}_{outcomes[i]:.0f}')
    print(oimg.shape)

    img = ((np.clip(oimg, vmin, vmax) - vmin) /
           (vmax - vmin) * 255).astype('uint8')

    timg = oimg.reshape(1, oimg.shape[0], oimg.shape[1], 1)

    preprocessed_input = [timg]

    model.set_weights(orig_weights)
    predictions = model.predict(preprocessed_input)
    print('Predicted class:')
    print(predictions)
    predicted_class = np.argmax(predictions)
    print(predicted_class)
    
    for target_class in range(2):
        gradcam, heatmap = apply_grad_cam(get_output[target_class], preprocessed_input)
        filename = f'output/{layer.name}/gradcam{i:03d}_{outcomes[i]:.0f}_{predicted_class:d}_{target_class:d}_{predictions[0][target_class]:0.2f}.png'
        cv2.imwrite(filename, gradcam)

        heatmap1 = cv2.resize(heatmap, (no_pc*3, path_n))
        #print(heatmap.shape)
        data_row = [[outcomes[i], predicted_class, target_class,
                     predictions[0][target_class]]+heatmap1.flatten().tolist()]
        gradcam_pd = pd.DataFrame(data_row, columns=col, index=[i])
        if i == 0 and target_class == 0:
            gradcam_pd.to_csv(
                f'output/{layer.name}/pathcnn_gradcam.csv', mode='w')
        else:
            gradcam_pd.to_csv(f'output/{layer.name}/pathcnn_gradcam.csv',
                              mode='a', header=False)