import numpy as np
import matplotlib.pyplot as plt
import time

def green(r1, r2, k):
    R = r2 - r1
    RR = np.outer(R, R)
    lR = np.sqrt((R**2).sum())
    
    G0 = k**2*np.exp(1j*k*lR) / (4*np.pi*lR)  # I think the k^2 should be there
    
    G = np.zeros((3,3), dtype='complex128')
    I = np.eye(3)
    G += (3/(k**2*lR**2) - 3j/(k*lR) - 1)*RR/lR**2
    G -= (1/(k**2*lR**2) - 1j/(k*lR) - 1)*I
    G *= G0
    return G

def field_sc(X,Y,Z, c, p, k):
    # Calculate field created by a bunch of dipoles.
    # naive implementation
    E = np.zeros((len(X.flatten()), 3), dtype='complex128')
    for i in range(len(X.flat)):
        # Loop over all requested points
        cp = np.array([X.flat[i], Y.flat[i], Z.flat[i]])
        for j in range(c.shape[0]):
            # Loop over all the dipoles
            if (c[j] == cp).all():
                continue
            G = green(c[j], cp, k)
            pthis = p[3*j:3*j+3]
            E[i,:] += G @ pthis
    E = np.reshape(E, X.shape+(3,))
    return E

def scattering_cross_section(c, p, k, Nth=5, Nph=5):
    # Evaluate field on a sphere outside
    sth = np.linspace(0,np.pi,Nth)
    sph = np.linspace(0,2*np.pi,Nph)
    sth = sth[1:-1]
    sph = sph[0:-1]
    delta_sth = sth[1]-sth[0]
    delta_sph = sph[1]-sph[0]
    sth,sph = np.meshgrid(sth,sph,indexing='ij')
    R = 1000  # should be far enough that only far-field components survive
    X = R*np.sin(sth)*np.cos(sph)
    Y = R*np.sin(sth)*np.sin(sph)
    Z = R*np.cos(sth)
    E = field_sc(X,Y,Z, c, p, k)
    # Assuming the incident wave is a plane wave with an amplitude of 1
    area = R**2*np.sin(sth)*delta_sth*delta_sph
    area = np.stack([area, area, area], axis=-1)
    return (abs(E)**2*area).sum()

def dda_dipole_moments(c, alpha, k, Ei):
    # Calculate the electric dipole moments for N dipoles under a prescribed
    # illumination.
    # Arguments:
    #   c: coordinates, N-by-3 matrix, each row holds the (x,y,z) coordinates of
    #      a dipole
    #   alpha: polarizability, scalar
    #   k: wave number
    #   Ei: incident field, vector of length 3N such that Ei[0::3] = Ex for each
    #       dipole, Ei[1::3] = Ey for each dipole etc.
    subMlist = []
    for i1 in range(N):
        subMlist.append([])
        for i2 in range(N):
            if i1 == i2:
                subMlist[i1].append(np.eye(3))
            else:
                subMlist[i1].append(-alpha*green(c[i1,:],c[i2,:],k))
    M = np.block(subMlist)
    p = np.linalg.solve(M, al*Ei)
    return p
    
# Here, length is in microns and so the unit of frequency is "the frequency
# corresponding to 1 micron wavelength radiation".
k = 2*np.pi*310/300
Nl = 9 # Discretization, increase this to make the result better
# Coordinates of the edges of the grid
l = np.linspace(-0.1, 0.1, Nl+1)
deltal = l[1]-l[0]
# Coordinates of the centers of the unit cells
l = 0.5*(l[0:-1] + l[1:])
x,y,z = np.meshgrid(l, l, l)
c = np.stack([x.flatten(), y.flatten(), z.flatten()], axis=-1)
N = c.shape[0]
eps_sc = 12.8
# Polarizability of the dipoles
al = (eps_sc-1)/(eps_sc+2)*3/(1/deltal)**3

Ei = np.zeros(3*N, dtype='complex128')
Ei[0::3] = np.exp(1j*k*c[:,2])

t1 = time.time()
p = dda_dipole_moments(c, al, k, Ei)
t2 = time.time()
print(f'Time for system construction and solution: {t2-t1}')

th = np.linspace(0, 2*np.pi, 128)
R = 1e6
Z = R*np.cos(th)
X = R*np.sin(th)
Y = np.zeros_like(X)
t1 = time.time()
Ef = field_sc(X, Y, Z, c, p, k)
t2 = time.time()
print(f'Time for field computation: {t2-t1}')
t1 = time.time()
sigma = scattering_cross_section(c, p, k, Nth=32, Nph=32)
t2 = time.time()
print(f'Time for scattering cross section: {t2-t1}')
print(f'sigma_sc = {sigma}')
Isc = (abs(Ef)**2).sum(axis=-1)
plt.figure()
plt.polar(th, np.sqrt(Isc))
plt.show()