# Plot transducer field in plane

#
# Nonlinear field generated in by a bowl-shaped HIFU transducer
# ==========================================================
#
# This demo illustrates how to:
#
# * Compute the nonlinear time-harmonic field in a homogeneous medium
# * Use incident field routines to generate the field from a HIFU transducer
# * Make a nice plot of the solution in the domain
#
#
# We consider the field generated by the Sonic Concepts H101 transducer:
# https://sonicconcepts.com/transducer-selection-guide/
# This transducer operates at 1.1 MHz, has a 63.2 mm radius of curvature and a 
# diameter of 64 mm. It has no central aperture.
# The medium of propagation we consider is water.

import os
import sys
# FIXME: figure out how to avoid this sys.path stuff
sys.path.append(os.path.join(os.path.dirname(__file__), '../../'))
import numpy as np
from vines.geometry.geometry import shape
from vines.fields.plane_wave import PlaneWave
from vines.operators.acoustic_operators import volume_potential, volume_potential_cylindrical
from vines.precondition.threeD import circulant_embed_fftw
from vines.operators.acoustic_matvecs import mvp_volume_potential, mvp_vec_fftw
from scipy.sparse.linalg import LinearOperator, gmres
from vines.mie_series_function import mie_function
from matplotlib import pyplot as plt
from vines.geometry.geometry import generatedomain2d, generatedomain
from vines.fields.transducers import bowl_transducer, normalise_power
import time
import matplotlib
from matplotlib import pyplot as plt
import itertools

'''                        Define medium parameters                         '''
# * speed of sound (c)
# * medium density (\rho)
# * the attenuation power law info (\alpha_0, \eta)
# * nonlinearity parameter (\beta)
c = 1487.0
rho = 998.0
alpha0 = 0.217
eta = 2
beta = 3.5e0

# material = 'liver'
# c = 1590.0
# rho = 1040
# alpha0 = 90.0
# eta = 1.1
# beta = 4.4


def attenuation(f, alpha0, eta):
    'Attenuation function'
    alpha = alpha0 * (f * 1e-6)**eta
    alpha = alpha / 8.686
    return alpha


'''                      Define transducer parameters                       '''
# * operating/fundamental frequency f1
# * radius of curvature, focal length (roc)
# * inner diameter (inner_D)
# * outer diameter (outer_D)
# * total acoustic power (power)
# f1 = 1.1e6
# roc = 0.0632
# inner_D = 0.0
# outer_D = 0.064
# power = 50

f1 = 1.1e6
transducername = 'H131'
roc = 0.035
inner_D = 0.0
outer_D = 0.033
power = 100

# f1 = 1.1e6
# transducername = 'H151'
# roc = 0.1
# inner_D = 0.0
# outer_D = 0.064
# power = 100


# FIXME: don't need to define focus location but perhaps handy for clarity?
focus = [roc, 0., 0.]
# FIXME: need source pressure as input

# Mesh resolution (number of voxels per fundamental wavelength)
nPerLam = 12

# Compute useful quantities: wavelength (lam), wavenumber (k0),
# angular frequency (omega)
lam = c / f1
k1 = 2 * np.pi * f1 / c + 1j * attenuation(f1, alpha0, eta)
omega = 2 * np.pi * f1

# Create voxel mesh
dx = lam / nPerLam

# Dimension of computation domain
# x_start needs to be close to the transducer
# x_end can be just beyond the focus
# the width in the y,z directions should be around the width of outer_D,
# but you can shrink this to speed up computations if required
# x_start = 0.001
x_start_0 = roc - 0.99 * np.sqrt(roc**2 - (outer_D/2)**2)
x_start = 0
x_end = roc + 0.01
wx = x_end - x_start_0
wy = outer_D*0.72
wz = wy/1000
# wz = wy
# from IPython import embed;embed()

start = time.time()
# r, L, M = generatedomain2d(dx, wx, wy)
r, L, M, N = generatedomain(dx, wx, wy, wz)

# Adjust r by shifting x locations
r[:, :, :, 0] = r[:, :, :, 0] - r[0, 0, 0, 0] + x_start
# r[:, :, :, 1] = r[:, :, :, 1] - r[0, 0, 0, 1]
end = time.time()
print('Mesh generation time:', end-start)
points = r.reshape(L*M*N, 3, order='F')

from IPython import embed; embed()
exit()

print('Number of voxels = ', L*M*N)

start = time.time()
n_elements = 2**12
x, y, z, p = bowl_transducer(k1, roc, focus, outer_D / 2, n_elements,
                             inner_D / 2, points.T, 'x')
end = time.time()
print('Incident field evaluation time (s):', end-start)
dist_from_focus = np.sqrt((points[:, 0]-focus[0])**2 + points[:, 1]**2 +
                           points[:,2]**2)
idx_near = np.abs(dist_from_focus - roc) < 1e-5
p[idx_near] = 1e-14

# Normalise incident field to achieve desired total acoutic power
p0 = normalise_power(power, rho, c, outer_D/2, k1, roc,
                     focus, n_elements, inner_D/2)

p *= p0

P = p.reshape(L, M, N, order='F')

# Create a pretty plot of the first harmonic in the domain
matplotlib.rcParams.update({'font.size': 22})
plt.rc('font', family='serif')
plt.rc('text', usetex=True)
xmin, xmax = r[0, 0, 0, 0] * 100, r[-1, 0, 0, 0] * 100
ymin, ymax = r[0, 0, 0, 1] * 100, r[0, -1, 0, 1] * 100
fig = plt.figure(figsize=(9, 7))
ax = fig.gca()
# plt.imshow(np.log10(np.abs(P[:, :, np.int(np.floor(N/2))].T / 1e6)),
#            extent=[xmin, xmax, ymin, ymax],
#            cmap=plt.cm.get_cmap('viridis'), interpolation='spline16')
plt.imshow((np.abs(P[:, :, np.int(np.floor(N/2))].T / 1e6)),
           extent=[xmin, xmax, ymin, ymax],
           cmap=plt.cm.get_cmap('viridis'), interpolation='spline16')
plt.xlabel(r'$x$ (cm)')
plt.ylabel(r'$y/z$ (cm)')
# circle = plt.Circle((roc*100, 0), roc*100,fill=False,color='white')
# ax.add_artist(circle)

rect = plt.Rectangle((x_start_0*100, -100*outer_D/2), 100*(roc+0.01-x_start_0), 100*outer_D, 
                     color='white', linestyle='dashed',fill=False)
ax.add_artist(rect)

rect = plt.Rectangle((0.44, -1.18), 100*(roc+0.01-x_start_0), 2*1.18, 
                     color='white', linestyle='dashed',fill=False)

ax.add_artist(rect)

rect = plt.Rectangle((1.82, -0.524), 100*(roc+0.01-0.0182), 2*0.524, 
                     color='white', linestyle='dashed',fill=False)

ax.add_artist(rect)

rect = plt.Rectangle((2.96, -0.166), 100*(roc+0.01-0.0296), 2*0.166, 
                     color='white', linestyle='dashed',fill=False)

ax.add_artist(rect)

rect = plt.Rectangle((3.14, -0.0642), 100*(roc+0.01-0.0314), 2*0.0642, 
                     color='white', linestyle='dashed',fill=False)

ax.add_artist(rect)


theta = np.linspace(np.pi-0.49, np.pi+0.49, 100)
radius = roc*100
a = radius*np.cos(theta)
b = radius*np.sin(theta)
ax.plot(a+roc*100, b,color='white')

# cbar = plt.colorbar()
# cbar.ax.set_ylabel('Pressure (MPa)')
# plt.clim(-4,0.5)

plt.text(2.1, 1.35, 'reference',
         {'color': 'w', 'fontsize': 20})

plt.text(1.1, -0.05, r'$p_2$',
         {'color': 'w', 'fontsize': 20})

plt.text(2.2, -0.05, r'$p_3$',
         {'color': 'w', 'fontsize': 20})

plt.text(2.95, -0.05, r'$p_4$',
         {'color': 'w', 'fontsize': 20})

plt.text(4.1, -0.05, r'$p_5$',
         {'color': 'w', 'fontsize': 20})

# fig = plt.figure()
# ax1 = fig.add_subplot(111)
# ax1.plot(your_data)

# Now adding the colorbar
# cbaxes = fig.add_axes([0.8, 0.1, 0.03, 0.8]) 
# cb = plt.colorbar(ax, cax = cbaxes)  

fig.savefig('results/H131_abs_subdomains.png')
plt.close()