import matplotlib as mpl
mpl.rcParams.update({
    'font.family': 'serif',
})

from mpl_toolkits.axes_grid1.anchored_artists import AnchoredSizeBar
import matplotlib.font_manager as fm

from functools import partial

import jax
import numpy as np
from jax import numpy as jnp
from jax import random

key = random.PRNGKey(42)

import matplotlib
matplotlib.rcParams.update(matplotlib.rcParamsDefault)

# Defining geometry
from jwave import FourierSeries
from jwave.geometry import Domain, Medium

N = (320, 512)  # Grid size
dx = (1e-4, 1e-4)  # Spatial resolution
omega = 1.7e6*2*jnp.pi  # Wavefield omega = 2*pi*f
target = [160, 360]  # Target location

# Making geometry
domain = Domain(N, dx)

# Constructing medium physical properties
sound_speed = jnp.ones(N)
sound_speed = sound_speed.at[30:80, 50:80].set(1.6)
sound_speed = sound_speed.at[80:140, 50:100].set(1.4)
sound_speed = sound_speed.at[140:220, 45:130].set(1.2)
sound_speed = jnp.expand_dims(sound_speed.at[220:280, 70:100].set(1.5), -1)*1480
sound_speed = FourierSeries(sound_speed, domain)

medium = Medium(domain=domain, sound_speed=sound_speed, pml_size=25)

# Build the vector that holds the parameters of the apodization an the
# functions required to transform it into a source wavefield
transmit_phase = jnp.concatenate([jnp.ones((32,)), jnp.ones((32,))])
position = list(range(32, 32 + (8 * 32), 8))


def phase_to_apod(phases):
    dim = len(phases) // 2
    return jnp.exp(1j * phases[dim:]) / (1 + (phases[:dim]) ** 2)


def phases_to_field(phases, domain):
    phases = phase_to_apod(phases)
    src_field = jnp.zeros(domain.N).astype(jnp.complex64)
    src_field = src_field.at[position, 25].set(phases)
    return FourierSeries(jnp.expand_dims(src_field, -1), domain)


linear_phase = phases_to_field(transmit_phase, domain)

from jax import jit

from jwave.acoustics.operators import helmholtz
from jwave.acoustics.time_harmonic import helmholtz_solver, helmholtz_solver_verbose

op_params = helmholtz.default_params(linear_phase, medium, omega=1.0)

from jax import value_and_grad, vmap

def loss(field):
    field = field.on_grid
    return -jnp.sum(jnp.abs(field[target[0], target[1]]))

def get_field(transmit_phase, tol, guess, op_params):
    transmit_field = phases_to_field(transmit_phase, domain)
    return fixed_medium_solver(transmit_field, op_params, guess, tol)

def full_loss(transmit_phase, tol, guess, op_params):
    field = get_field(transmit_phase, tol, guess, op_params)
    return loss(field), field

loss_with_grad = value_and_grad(full_loss, has_aux=True)

from jax import jit
from jax.example_libraries import optimizers
from tqdm import tqdm

losshistory = []

init_fun, update_fun, get_params = optimizers.adam(0.1, b1=0.9, b2=0.9)
opt_state = init_fun(transmit_phase)

@partial(jit, static_argnums=(1,))
def update(opt_state, tol, guess, op_params):
    loss_and_field, gradient = loss_with_grad(
        get_params(opt_state), tol, guess, op_params
    )
    lossval = loss_and_field[0]
    field = loss_and_field[1]
    return lossval, field, update_fun(k, gradient, opt_state)


pbar = tqdm(range(100))
tol = 1e-3
guess = None
for k in pbar:
    lossval, new_field, opt_state = update(opt_state, tol, guess, op_params)

    # For logging
    pbar.set_description("Ampl: {:01.4f}".format(-lossval))
    losshistory.append(lossval)

transmit_phase = get_params(opt_state)

from jax import random

from jwave.signal_processing import smooth

key = random.PRNGKey(12)

target = [60, 360]  # Target location

# Constructing medium physical properties
def get_sos(segments, start_point=30, height=4, width=30):
    sos = jnp.ones(N)
    for k in range(len(segments)):
        sos = sos.at[
            start_point + k * height : start_point + (k + 1) * height, 50 : 50 + width
        ].add(jax.nn.sigmoid(segments[k]))
    return FourierSeries(jnp.expand_dims(sos, -1), domain)

key, _ = random.split(key)
sos_control_points = random.normal(key, shape=(65,))
sos = get_sos(sos_control_points)

from jax import value_and_grad

from jwave.acoustics.time_harmonic import helmholtz_solver


def loss(field):
    field = field.on_grid
    return -jnp.sum(jnp.abs(field[target[0], target[1]]))


def get_field(params, tol, field):
    medium = Medium(domain, sound_speed=get_sos(params))
    return helmholtz_solver(
        medium, 1.0, linear_phase, guess=field, tol=tol, checkpoint=False
    )


def full_loss(params, tol, field):
    field = get_field(params, tol, field)
    return loss(field), field


loss_with_grad = value_and_grad(full_loss, has_aux=True)

opt_sos_vector = get_params(opt_state)