%matplotlib inline
%load_ext autoreload
%autoreload 2

import numpy as np
import matplotlib.pyplot as plt
import fipy as fp

def fipy_response(phi0, dt, N):
    from fipy.solvers.scipy.linearLUSolver import LinearLUSolver
    dx = 0.005
    a = np.sqrt(200.)
    epsilon = 0.1
    nx = ny = N
    mesh = fp.PeriodicGrid2D(nx=nx, ny=ny, dx=dx, dy=dx)
    phi = fp.CellVariable(name=r"$\phi$", mesh=mesh, value=phi0.copy())
    PHI = phi.arithmeticFaceValue
    D = 1.
    eq = (fp.TransientTerm()
      == fp.DiffusionTerm(coeff=D * a**2 * (1 - 6 * PHI * (1 - PHI)))
      - fp.DiffusionTerm(coeff=(D, epsilon**2)))
    
    eq.solve(phi, dt=dt, solver=LinearLUSolver())
    
    return (np.array(phi) - phi0) / dt
    


from IPython.display import clear_output
import time

N = 21
np.random.seed(0)
phi0 = np.random.random(N * N)
dt = 1e-8

fig = plt.figure()

for i in range(30):
    response = fipy_response(phi0, dt=dt, N=N)
    #Euler backward in this case since implicit
    phi0 = response * dt + phi0
    #print phi0
    plt.contourf(phi0.reshape((N,N)))
    time.sleep(1)
    clear_output()
    display(fig)

N = 10
Nbin = 6
Nsample = 5

np.random.seed(1)
microstructures = np.random.random((Nsample, N**2))
responses = np.array([fipy_response(M, dt=dt, N=N) for M in microstructures])
print microstructures.shape
print responses.shape

from pymks import bin
?bin

from pymks import draw_microstructure_discretization
for s in range(3):
    plt.figure()
    draw_microstructure_discretization(microstructures, a=0, s=s)

binnedMicrostructures = np.array([bin(M, Nbin) for M in microstructures])
print binnedMicrostructures.shape

X = np.linspace(0, 1, Nbin)
reconstructedMicrostructure = np.sum(binnedMicrostructures * X[np.newaxis, np.newaxis, :], axis=-1)
print np.allclose(reconstructedMicrostructure, microstructures)

print X

??bin

def rollMatrix(m, N, Nbin):
    matrix = np.zeros((N**2, N**2 * Nbin))
    for i in range(N**2):
        matrix[i] = np.roll(m, -i, axis=0).swapaxes(0,1).flatten()
    return matrix

microstructureMatrix = np.concatenate([rollMatrix(m, N, Nbin) for m in binnedMicrostructures])

print microstructureMatrix.shape


responses = responses.flatten()
coefficients = np.linalg.lstsq(microstructureMatrix, responses)[0]
print coefficients.shape


print np.allclose(np.dot(microstructureMatrix, coefficients), responses)

N = 20
Nbin = 11
Nsample = 160

np.random.seed(101)
microstructures = np.array([np.random.random(N**2) for i in range(Nsample)])
responses = np.array([fipy_response(m, dt=dt, N=N) for m in microstructures])
binnedMicrostructures = np.array([bin(m, Nbin) for m in microstructures])
print microstructures.shape
print responses.shape

microstructuresRS = binnedMicrostructures.reshape((Nsample, N, N, Nbin))
responsesRS = responses.reshape((Nsample, N, N))
Fmicrostructures = np.fft.fft2(microstructuresRS, axes=(1, 2))
Fresponses = np.fft.fft2(responsesRS, axes=(1, 2))

Fcoeff = np.zeros((N, N, Nbin), dtype=np.complex)

for ki in range(N):
    for kj in range(N):
        Fcoeff[ki,kj,:] = np.linalg.lstsq(Fmicrostructures[:,ki,kj,:], Fresponses[:,ki,kj] )[0]
        

reconstructedResponses = np.sum(Fmicrostructures * Fcoeff[np.newaxis], axis=-1)

from sklearn import metrics
mse = metrics.mean_squared_error
responses = np.fft.ifft2(reconstructedResponses, axes=(1, 2)).real
MSE = mse(responsesRS, np.fft.ifft2(reconstructedResponses, axes=(1, 2)).real)
print 'Mean square error: {0:1.3e}'.format(MSE)
print 'Typical values'
print responsesRS[0,0,:10]
print responses[0,0,:10]

np.random.seed(103)
test_microstructure = np.random.random(N**2)
test_response = fipy_response(test_microstructure, dt=dt, N=N)

binned_test_microstructure = bin(test_microstructure, Nbin).reshape((N, N, Nbin))
Fm = np.fft.fft2(binned_test_microstructure, axes=(0, 1))
Fr = np.sum(Fm * Fcoeff, axis=-1)
            
calc_response = np.fft.ifft2(Fr, axes=(0, 1)).real.flatten()


print test_response[:20]
print calc_response[:20]

coeff = np.fft.ifft2(Fcoeff, axes=(0, 1)).real

plt.figure()
plt.title(r'$\beta_s^0$', fontsize=18)
plt.contourf(coeff[:,:,0])
plt.colorbar()
plt.figure()
plt.title(r'$\beta_s^1$', fontsize=18)
plt.contourf(coeff[:,:,1])
c = plt.colorbar()