#!/usr/bin/env python
# coding: utf-8

# # Problem 7
# Let $A$ be the $20000 \times 20000$ matrix whose entries are 0 everywhere except for the primes 2,3,5,7,...,224737 along the main diagonal and the number 1 in all the positions $ a_{ij} $ with $ |i-j| = 1,2,4,8,\cdots,16384$. What is the (1,1) entry of $A^{-1}$

# ### Some introductory code 

# In[1]:


def is_prime(N):
    if N < 2: return False
    if N == 2: return True
    for i in range(2, int(N**0.5)+1):
        if N%i == 0: return False
    return True
def is_prime(N):
    return N == 2 or any([N%i == 0 for i in range()])
entries = []
primes = [i for i in range(2,250000) if is_prime(i)][:20000]
entries.extend([(i,i,primes[i]) for i in range(len(primes))])
for diff in [2**j for j in range(15)]:
    for i in range(20000):
        j1, j2 = i+diff, i-diff
        if 0<=j1<20000: entries.append((i,j1,1))
        if 0<=j2<20000: entries.append((i,j2,1))
print(len(entries))


# ### Attempt 1: Using standard libraries
# 
# We avoid calculating $A^{-1}$ explicitly. Instead, note that if we have $\hat{v}$ be the first column of $ A^{-1} $, then we get that $ A \hat{v} = (1,0,\cdots, 0)^{T}$. Let's solve this system. 

# In[1]:


import numpy as np
from scipy.sparse import coo_matrix
from scipy.sparse.linalg import spsolve

row = np.array([i[0] for i in entries])
col = np.array([i[1] for i in entries])
data = np.array([i[2] for i in entries])

A = coo_matrix((data, (row, col))).tocsc()
b = np.zeros(A.shape[0])
b[0] = 1
v = spsolve(A,b)
# Timeout


# The above code times out - so instead, let's try to use an iterative method to solve this linear system.

# In[11]:


import numpy as np
from scipy.sparse import coo_matrix
from scipy.sparse.linalg import bicg
row = np.array([i[0] for i in entries])
col = np.array([i[1] for i in entries])
data = np.array([i[2] for i in entries])

A = coo_matrix((data, (row, col))).tocsc()
b = np.zeros(A.shape[0])
b[0] = 1

iterations = 0
def callback(xk):
    global iterations
    iterations += 1
    if iterations % 10 == 0: print("Iteration {}: current val = {}".format(iterations, xk[0]))
x, exit_code = bicg(A,b,tol=1e-50,callback=callback)
print(x[0])


# In[ ]: