#Downloading sequences:

import numpy as np, pandas as pd, requests
from scipy.stats import mannwhitneyu
genes=pd.read_csv('human-genes.tsv', sep='\t')
with open('seq_train.txt','w') as seq_train, open('seq_pos.txt','w') as seq_pos, open('seq_neg.txt','w') as seq_neg:
    num_train=0
    num_pos=0
    num_neg=0
    req_train=100
    req_pos=100
    req_neg=req_pos
    nonstart=200
    i=0
    while (num_train<req_train) or (num_pos<req_pos) or (num_neg<req_neg):
        chromosome = genes.loc[i, '#chrom']
        start = genes.loc[i, 'thickStart']
        strand = genes.loc[i, 'strand']
        i+=1
        if str(strand) == '+':
            link = f'https://rest.ensembl.org/sequence/region/human/{chromosome}:{start-6}..{start+400}:1?expand_3prime=0;expand_5prime=0'
            request = requests.get(link, headers={ 'Content-Type' : 'text/x-fasta'})
        else:
            continue
        if 'error' in request.text:
            continue
        else:
            seq=''.join(request.text.split('\n')[1:])
            if (num_train<req_train) and (seq[7:10]=='ATG'):
                print(seq[:13],file=seq_train)
                num_train+=1
            elif (num_pos<req_pos) and (seq[7:10]=='ATG'):
                print(seq[:13],file=seq_pos)
                num_pos+=1
            if (num_neg<req_neg):
                ATG=seq[nonstart:-3].find('ATG')
                if ATG!=(-1):
                    print(seq[ATG+nonstart-7:ATG+nonstart+6],file=seq_neg)
                    num_neg+=1
                    
#Defining PWM:
                    
class PWM():
    def __init__(self, e=1, freq=[0.295, 0.295, 0.205, 0.205]):
        self.e = e
        self.positions = {"A": 0, "T": 1, "G": 2, "C": 3}
        self.frequencies = np.array(freq).reshape((4, 1))
        
    def fit(self, seq_list):
        self.seq_arr = np.array(list(map(list, seq_list)), dtype="str")
        self.n_pos = len(seq_list[0])
        self.n_seq = len(seq_list)
        A_n = np.count_nonzero(self.seq_arr == "A", axis=0)
        T_n = np.count_nonzero(self.seq_arr == "T", axis=0)
        G_n = np.count_nonzero(self.seq_arr == "G", axis=0)
        C_n = np.count_nonzero(self.seq_arr == "C", axis=0)
        self.matrix = np.stack([A_n, T_n, G_n, C_n], axis=0).astype("float")
        pseudocount = self.e / 4
        self.matrix += pseudocount
        self.matrix = self.matrix / (self.n_seq+self.e)
        self.pwm_matrix = np.log(self.matrix / self.frequencies)
        
    def test(self, seq):
        weight = 0
        for i, nuc in enumerate(list(seq)):
            weight += self.pwm_matrix[self.positions[nuc], i]
        return weight
    
    def ic_calculate(self):
        self.ic_matrix = self.matrix * np.log2(self.matrix / self.frequencies)
        
    def return_pwm(self):
        return pd.DataFrame(self.pwm_matrix, columns = np.arange(1, self.n_pos + 1), index = ["A", "T", "G", "C"])
    
    def return_ic(self):
        return pd.DataFrame(self.ic_matrix, columns = np.arange(1, self.n_pos + 1), index = ["A", "T", "G", "C"])

#Constructing PWM:
    
pd.set_option('display.max_rows', None)
pd.set_option('display.max_columns', None)
pd.set_option('display.max_colwidth', None)
with open('seq_train.txt','r') as seq_train:
    pwm = PWM()
    pwm.fit([x.strip() for x in seq_train.readlines()])
    pwm.return_pwm().round(2).to_csv(sep='\t', path_or_buf='PWM.tsv')

#Testing PWM:
    
with open('seq_pos.txt','r') as seq_pos, open('seq_neg.txt','r') as seq_neg:
    pos_scores = []
    neg_scores = []
    for pos, neg in zip([x.strip() for x in seq_pos.readlines()], [x.strip() for x in seq_neg.readlines()]):
        pos_scores.append(pwm.test(pos))
        neg_scores.append(pwm.test(neg))

#Analyzing results:

print('Avg. weight for positives: ',np.mean(pos_scores).round(3),'; Avg. weight for negatives: ',np.mean(neg_scores).round(3),'.',sep='')
print(mannwhitneyu(pos_scores, neg_scores))

#Constructing IC matrix:

pwm.ic_calculate()
pwm.return_ic().round(2).to_csv(sep='\t', path_or_buf='IC.tsv')