import matplotlib.pyplot as plt


f = open('Test.txt', 'r')
s = f.readlines()
a = [0]*len(s)
Sample_Mutations = {}
for i in range (len(s) - 1):
    a[i] = s[i].split('\t')
    if a[i][32] in Sample_Mutations:
        Sample_Mutations[a[i][32]] += 1
    else:
        Sample_Mutations[a[i][32]] = 1

p = open('!ДАТАСЕТ.txt', 'r')
t = p.readlines()
b = [0]*len(t)
Sample_Matches = {}
for i in range (len(t) - 1):
    b[i] = t[i].split('\t')

for key in Sample_Mutations:
    Sample_Matches[key] = 0

for i in range (len(a) - 1):
    for j in range (len(b) - 1):
        if a[i][4] == b[j][0] and (int(b[j][1]) <= int(a[i][5]) <= int(b[j][2])):
            Sample_Matches[a[i][32]] += 1


w = open('Out.txt', 'w')
w.write('Tumor_Sample_UUID' + '\t' + 'Mutations' + '\t' +'Matches' + '\n')
for key in Sample_Mutations:
    w.write(key + '\t' + str(Sample_Mutations[key]) + '\t' + str(Sample_Matches[key])+ '\n')
w.close()


if __name__ == '__main__':
    x = []
    y = []
    for key in Sample_Mutations:
        if Sample_Mutations[key] != 405:
            x.append(Sample_Mutations[key])
            y.append(Sample_Matches[key])

    plt.title("Correlation")
    plt.xlabel("Matches")
    plt.ylabel("Mutations")
    plt.scatter(y, x, edgecolors='r', s=10)
    plt.grid(True, linestyle='-', color='0.75')
    plt.show()

from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score
M = [[0] * len(x) for i in range(len(y))]
lr2 = LinearRegression()
xs = x
lr2.fit([[x] for x in xs], y)
y2 = lr2.predict([[x] for x in xs])
#plt.figure()
#plt.plot(x, y, color='b', linewidth=1, label='Linear')
plt.scatter(y, x, edgecolors='r', s=10)
plt.plot(y2, x, color='b', linewidth=1, label='Linear')

plt.grid(True, linestyle='-', color='0.75')
plt.title("Linear Regression")
plt.ylabel('Mutations')
plt.xlabel('Matches')
plt.show()

print('Coefficients: \n', 'a = ', lr2.coef_,
      'b = ', lr2.intercept_, '\n',
      'R^2 = ', lr2.score([[x] for x in xs], y),)