# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.metrics import *
import seaborn as sns
sns.set()
# %matplotlib inline

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from scripts.common import utils
from scripts.common import metrics 
from scripts.common.dataset import SimpleSeqsDataset
from scripts.common.training import *

print(torch.__version__)
os.chdir('./Data')

os.listdir()[:3]
# -

chromosomes = list(range(1,23)) + ['M','X','Y']
data_full = utils.merge_chrom_data('chr', chromosomes)
print(data_full.shape)
print('not longer than 500bp:', sum(data_full.len <= 500) / data_full.shape[0])
pd.DataFrame(data_full.len.describe())

ax = data_full[data_full.len<1000].len.hist(bins=np.arange(data_full.len.min(),1000,30),figsize=(12,6))
ax.set_ylabel('Количество G-квадруплексов')
ax.set_xlabel('Длина участка')
ax.set_title('Распределение длин G-квадруплексов')
ax.axvline(500,color='r', linestyle='--')
plt.savefig('../g_quadruplexes.png')

print('Координаты G-квадруплексов после удаления слишком длинных и центрирования:')
utils.center_seq(500).head()

chromosomes = list(range(1,23)) + ['M','X','Y']
df = utils.merge_chrom_data('chr', chromosomes, '_centered')
df.head()

# Load shuffled data 
chromosomes = list(range(1,23)) + ['M','X','Y']
shuf = utils.merge_chrom_data('shuffled_chr', chromosomes)
shuf.head()

# +
# Split to test & train
test_chromosomes_list = [3,6,9,12,15,18,21]
test_chromosomes = list(map(lambda n:'chr'+str(n),test_chromosomes_list))
specific_chromosomes = ['chrX','chrM','chrY']

shuf['label'] = 0
test_shuf = shuf[shuf.chr.isin(test_chromosomes)]
train_shuf = shuf[~shuf.chr.isin(test_chromosomes + specific_chromosomes)]

df['label'] = 1
test_data = df[df.chr.isin(test_chromosomes)]
train_data = df[~df.chr.isin(test_chromosomes + specific_chromosomes)]

train_full = pd.concat([train_data,train_shuf]).reset_index(drop=True)
test_full = pd.concat([test_data,test_shuf]).reset_index(drop=True)
# -

# Get X & y
train_X = utils.encode_seq(train_full.seq)
test_X = utils.encode_seq(test_full.seq)
train_y = train_full.label
test_y = test_full.label
np.array(test_y).shape

# +
# Get train & test datasets
device = utils.get_device()      
BATCH_SIZE = 32
batch_num = len(train_X)//BATCH_SIZE
     
train_dataset = SimpleSeqsDataset(train_X, train_y)
test_dataset = SimpleSeqsDataset(test_X, test_y)
# -

train_dataset.__getitem__(3)[0]


# +
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv1d(in_channels=4,out_channels=16, kernel_size=5)
        self.pool = nn.MaxPool1d(5)
        self.fc1 = nn.Linear(16*99, 128)
        self.fc2 = nn.Linear(128, 1)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool(x)
        x = x.view(-1, 16*99)
        x = F.relu(self.fc1(x))
        x = torch.sigmoid(self.fc2(x))
        return x
    
net = Net()
print(net)
# -

all_losses = []
all_accs = []

net = Net()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
all_losses, all_accs = train(net, 13, optimizer, train = train_dataset, batch_size = BATCH_SIZE,
                             eval_epochs = 0, use_wandb = False, all_losses = all_losses, all_accs = all_accs)

net = Net()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
all_losses, all_accs = train(net, 13, optimizer, train = train_dataset, batch_size = BATCH_SIZE,
                             eval_epochs = 0, use_wandb = False, all_losses = all_losses, all_accs = all_accs)


def plot_metric(metric_values):
    plt.plot(np.arange(len(metric_values)), metric_values)
plt.subplots(1,2, figsize=(14,6))
plt.subplot(1,2,1)
plot_metric(all_losses[-1])
plt.subplot(1,2,2)
plot_metric(all_accs[-1])

y_labels_list, y_pred_list = evaluation(net, test_dataset, batch_size = 1)
metrics.plot_conf_matrix(y_labels_list, y_pred_list)

roc_auc_score(y_labels_list, y_pred_list)



StarcoderdataJupyterScriptsDedupFiltered