import random
import requests
import sys
file = open('human-genes.tsv', 'r')
f = open('kozak-learn.fasta', 'w')
g =  open('kozak-test.fasta', 'w')
lines = file.readlines()
random.shuffle(lines[1:])
count = 0
for index, line in enumerate(lines):
    b = line.split("\t")
    if b[4] == "+":
        ext = "https://rest.ensembl.org" +"/sequence/region/human/" + b[0][3:] + ":" + str(int(b[6]) - 6) + ".." + str(int(b[6]) + 6) + ":1?expand_3prime=0;expand_5prime=0"
    elif b[4] == "-":
        ext = "https://rest.ensembl.org" + "/sequence/region/human/" + b[0][3:] + ":" + str(int(b[7]) - 5) + ".." + str(int(b[7]) + 7) + ":-1?expand_3prime=0;expand_5prime=0"
    r = requests.get(ext, headers={"Content-Type": "text/x-fasta"})
    if not r.ok:
        r.raise_for_status()
        sys.exit()
    if r.text[-7:-4] == "ATG" :
        count += 1
        if count <= 40:
            f.write(r.text)
        elif count <= 100:
            g.write(r.text)
    if count == 100:
        break
f.close()
g.close()
lines = []










chr_num = ['1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12', '13', '14', '15', '16', '17', '18', '19', '20', '21', '22', 'X', 'Y']
chr_len = [248956422, 242193529, 198295559, 190214555, 181538259, 170805979, 159345973, 145138636, 138394717, 133797422, 135086622, 133275309, 114364328, 107043718, 101991189, 90338345, 83257441, 80373285, 58617616, 64444167, 46709983, 50818468, 156040895, 57227415]
chr_adit = [248956422, 491149951, 689445510, 879660065, 1061198324, 1232004303, 1391350276, 1536488912, 1674883629, 1808681051, 1943767673, 2077042982, 2191407310, 2298451028, 2400442217, 2490780562, 2574038003, 2654411288, 2713028904, 2777473071, 2824183054, 2875001522, 3031042417, 3088269832]
count = 0
f = open('pseudokozak1.fasta', 'w')
while count < 60:
    a = random.randint(1, chr_adit[-1])
    c = 0
    for j in range(len(chr_adit)):
        if chr_adit[j] >= a:
            c = j
            break
    if chr_len[c] < 1000:
        continue
    ext = "https://rest.ensembl.org/sequence/region/human/" + chr_num[c] + ":" + str(chr_adit[j] - a) + ".." + str(chr_adit[j] - a + 998) + ":1?expand_3prime=0;expand_5prime=0"
    r = requests.get(ext, headers={"Content-Type": "text/x-fasta"})
    if not r.ok:
        r.raise_for_status()
        sys.exit()
    r = r.text[-999:]
    s = ''.join(ch for ch in r if ch.isalnum())
    l = s[7:].find("ATG")
    if l == -1:
        continue
    s = s[l: l + 13]
    f.write(">" + str(count) +"\n")
    f.write(s+"\n")
    count += 1
f.close()









import math
f = open('kozak-learn.fasta', 'r')
d = {"A" : 0, "T" : 1, "G" : 2, "C" : 3}
let = ["A", "T", "G", "C"]
fr = {0 : 0.295, 1 : 0.295, 2 : 0.205, 3 : 0.205}
ep = 0.2
lines = f.readlines()
pnm = [[0]*4 for i in range(13)]
pwm = [[0]*4 for i in range(13)]
ic = [[0]*4 for i in range(13)]
for index, line in enumerate(lines):
    if index % 2 == 0:
        continue
    for i in range(13):
        pnm[i][d[line[i]]] += 1
for i in range(13):
    for j in range(4):
        pwm[i][j] = math.log((pnm[i][j]+ep)/(40+ep)/fr[j])
f.close()
for i in range(13):
    for j in range(4):
        if pnm[i][j] == 0:
            ic[i][j] = 0
        else:
            ic[i][j] = math.log((pnm[i][j]+ep)/(40)/fr[j])/math.log(2)*((pnm[i][j]+ep)/(40))

import statsmodels.stats.weightstats
f = open('kozak-test.fasta', 'r')
lines = f.readlines()
positive = []
negative = []
for index, line in enumerate(lines):
    if index % 2 == 0:
        continue
    score = 0
    for i in range(13):
        score += pwm[i][d[line[i]]]
    positive.append(score)
f.close()
f = open("pseudokozak1.fasta", "r")
lines = f.readlines()
for index, line in enumerate(lines):
    if index % 2 == 0:
        continue
    score = 0
    for i in range(13):
        score += pwm[i][d[line[i]]]
    negative.append(score)
f.close()
icsum = 0
for i in ic:
    icsum += sum(i)
with open("result.txt", "w") as r:
    r.write("positive control mean score: " + str(sum(positive)/len(positive))+"\n")
    r.write("negative control mean score: " + str(sum(negative) / len(negative)) + "\n")
    r.write("p-value: "+ str(statsmodels.stats.weightstats.ttest_ind(positive, negative, alternative='two-sided', usevar='unequal', weights=(None, None), value=0)[1])+"\n")
    r.write("ic:" + str(icsum))
with open("result.csv", "w") as r:
    r.write("letter,1,2,3,4,5,6,7,8,9,10,11,12,13\n")
    for i in range(4):
        r.write(let[i])
        for j in range(13):
            r.write("," + str(pwm[j][i]))
        r.write("\n")
with open("ic.csv", "w") as r:
    r.write("letter,1,2,3,4,5,6,7,8,9,10,11,12,13\n")
    for i in range(4):
        r.write(let[i])
        for j in range(13):
            r.write("," + str(ic[j][i]))
            r.write("\n")