"""Test for the implementation of calc_fields."""

from . import grating
from . import grating_diffuse
import numpy as np
import matplotlib.pyplot as plt
from scipy.special import erf
from genx.gui_logging import iprint


def test_calc_fields_single_interface():
    qz = np.linspace(0.0001, 0.5, 500)
    t, r, k_z = grating.calc_fields(qz/2., 1.54, np.array([9e-5, 9e-5, 0], dtype=np.complex128),
                                    np.array([0.0, 0.0, 0.0]))

    t_theory = 2*k_z[:, -1]/(k_z[:, 0] + k_z[:, -1])
    r_theory = (k_z[:, -1] - k_z[:, 0])/ (k_z[:, 0] + k_z[:, -1])

    iprint((t_theory - t[:, 0]).max())
    iprint((r_theory - r[:, 1]).max())

    if False:
        plt.subplot(211)
        plt.plot(qz, np.abs(t[:, 0])**2)
        plt.plot(qz, np.abs(t_theory) ** 2)
        plt.subplot(212)
        plt.semilogy(qz, np.abs(r[:, -1]) ** 2)
        plt.semilogy(np.abs(r_theory) ** 2)
        plt.show()


def test_calc_fields_single_layer():
    qz = np.linspace(0.0001, 0.5, 500)
    v = np.array([9e-5 - 2e-6, 5e-5 - 8e-6, 0], dtype=np.complex128)
    d = np.array([0.0, 100.0, 0.0])
    t, r, k_z = grating.calc_fields(qz / 2., 1.54, v, d)
    r1 = (k_z[:, 1] - k_z[:, 0])/(k_z[:, 1] + k_z[:, 0])
    r2 = (k_z[:, 2] - k_z[:, 1]) / (k_z[:, 2] + k_z[:, 1])
    t1 = 2*k_z[:, 1]/(k_z[:, 0] + k_z[:, 1])
    t2 = 2 * k_z[:, 2] / (k_z[:, 1] + k_z[:, 2])

    r_theory = ((r1*np.exp(1.0J*d[1]*k_z[:, 1])/(t1*t2) + r2*np.exp(-1.0J*d[1]*k_z[:, 1])/(t1*t2)) /
                (r1*r2*np.exp(1.0J*d[1]*k_z[:, 1])/(t1*t2) + np.exp(-1.0J*d[1]*k_z[:, 1])/(t1*t2)))

    t_lay_theory = t2*np.exp(1.0J*d[1]*k_z[:, 1])/(r1*r2*np.exp(2*1.0J*d[1]*k_z[:, 1]) + 1.)
    r_lay_theory = r1*t2*np.exp(1.0J*d[1]*k_z[:, 1])/(r1*r2*np.exp(2*1.0J*d[1]*k_z[:, 1]) + 1.)

    iprint((t_lay_theory - t[:, -2]).max())
    iprint((r_theory - r[:, -1]).max())
    iprint((r_lay_theory - r[:, -2]).max())

    if False:
        plt.subplot(211)
        plt.plot(qz, np.abs(t[:, -2]), 'b:')
        plt.subplot(212)
        plt.plot(qz, np.abs(r[:, -1]), 'g:')
        plt.plot(qz, np.abs(r[:, -2]), 'm:')

        plt.subplot(211)
        plt.plot(qz, np.abs(t_lay_theory), 'b-.')
        plt.subplot(212)
        plt.semilogy(qz, np.abs(r_theory), 'g-.')
        plt.plot(qz, np.abs(r_lay_theory), 'm-.')
        plt.show()


def test_calc_fields_slicing_single_interface():
    lamda = 1.54
    v_lay = (2 * np.pi / 1.54) ** 2 * (1 - (1 - (1.57 - 0.4J) * 2.82e-5 * lamda ** 2 / 2 / np.pi) ** 2)
    v_sub = (2 * np.pi / 1.54) ** 2 * (1 - (1 - (0.71 - 0.09J) * 2.82e-5 * lamda ** 2 / 2 / np.pi) ** 2)
    v_layers = np.array([v_lay, v_lay, 0.0])
    sigma = np.array([0., 0., 0.])
    z_int = np.array([100., 0., -20.0])
    # NOTE: The first element has to be the substrate => largest z-value!
    z_slice = np.linspace(150, -40, 160)
    h = 0.7

    plt.subplot(211)
    kz_spec = np.arange(0.001, 0.2, 0.001)
    T_1, R_1, kz_1 = grating.calc_fields(kz_spec, lamda, v_layers, np.r_[0, z_int[:-1] - z_int[1:]])
    plt.plot(2 * kz_spec, np.abs(T_1[:, -2]) ** 2, 'b:')

    v_mean, int_layer_indices = grating_diffuse.calc_layer_prop(z_slice * 0 + 1, sigma, v_layers, z_int, z_slice)
    d = np.r_[0, z_slice[:-1] - z_slice[1:]]
    T_1, R_1, kz_1 = grating.calc_fields(kz_spec, lamda, v_mean, d)
    z_bottom_index = np.argmin(np.abs(z_slice - z_int[0]))
    plt.plot(2 * kz_spec, np.abs(T_1[:, z_bottom_index]) ** 2, 'g-.')

    plt.subplot(212)
    plt.plot(z_slice, v_mean.real, 'b', z_slice, v_mean.imag, 'g')
    plt.plot(z_slice, np.ones_like(z_slice) * v_lay.real, 'k', z_slice, np.ones_like(z_slice) * v_lay.imag, 'k')
    plt.show()


def test_calc_fields_slicing_single_layer():
    lamda = 1.54
    v_lay = (2 * np.pi / 1.54) ** 2 * (1 - (1 - (1.57 - 0.4J) * 2.82e-5 * lamda ** 2 / 2 / np.pi) ** 2)
    v_sub = (2 * np.pi / 1.54) ** 2 * (1 - (1 - (0.71 - 0.09J) * 2.82e-5 * lamda ** 2 / 2 / np.pi) ** 2)
    v_layers = np.array([v_sub, v_lay, 0.0])
    sigma = np.array([0., 0., 0.])
    z_int = np.array([100., 0., -20.0])
    z_slice = np.linspace(150, -40, 190*4 + 1)

    kz_spec = np.arange(0.001, 0.20, 0.001)
    T_1_lay, R_1_lay, kz_1 = grating.calc_fields(kz_spec, lamda, v_layers, np.r_[0, z_int[:-1] - z_int[1:]])

    v_mean, int_layer_indices = grating_diffuse.calc_layer_prop(z_slice * 0 + 1, sigma, v_layers, z_int, z_slice)
    d = np.r_[0, z_slice[:-1] - z_slice[1:]]
    T_1_slice, R_1_slice, kz_1 = grating.calc_fields(kz_spec, lamda, v_mean, d)
    z_bottom_index = np.argmin(np.abs(z_slice - z_int[0]))

    # Difficult to get the numbers to agree with high precision due to the slicing.
    iprint((((np.abs(T_1_lay[:, -2]) - np.abs(T_1_slice[:, z_bottom_index]))/np.abs(T_1_lay[:, -2])).max()))
    iprint((((np.abs(R_1_lay[:, -2]) - np.abs(R_1_slice[:, z_bottom_index])) / np.abs(R_1_lay[:, -2])).max()))
    iprint((((np.abs(R_1_lay[:, -1]) - np.abs(R_1_slice[:, -1]))/np.abs(R_1_lay[:, -1])).max()))

    if False:
        plt.subplot(311)
        plt.plot(2 * kz_spec, np.abs(T_1_lay[:, -2]), 'b:')
        plt.subplot(312)
        plt.plot(2 * kz_spec, np.abs(R_1_lay[:, -1]), 'g:')
        plt.plot(2 * kz_spec, np.abs(R_1_lay[:, -2]), 'm:')

        plt.subplot(311)
        plt.plot(2 * kz_spec, np.abs(T_1_slice[:, z_bottom_index]), 'b-.')
        plt.subplot(312)
        plt.semilogy(2 * kz_spec, np.abs(R_1_slice[:, -1]), 'g-.')
        plt.plot(2 * kz_spec, np.abs(R_1_slice[:, z_bottom_index]), 'm-.')

        plt.subplot(313)
        plt.plot(z_slice, v_mean.real, 'b', z_slice, v_mean.imag, 'g')
        plt.plot(z_slice, np.ones_like(z_slice) * v_lay.real, 'k', z_slice, np.ones_like(z_slice) * v_lay.imag, 'k')
        plt.plot(z_slice, np.ones_like(z_slice) * v_sub.real, 'k', z_slice, np.ones_like(z_slice) * v_sub.imag, 'k')
        plt.show()


def test_correlation_function_h05():
    """Test so that the correlation function for h=0.5 matches the analytical solution."""
    eta_r = 1050.
    q_r = np.logspace(-5., 0., 100)
    c_func = grating_diffuse.correlation_function(q_r, 0.5, eta_r, 0.0, 1e-10)
    a = 1/eta_r
    analytical_func = 2.*a / (a**2 + q_r ** 2)

    iprint((np.abs((c_func - analytical_func)/c_func).max()))

    if True:
        plt.loglog(q_r, c_func)
        plt.loglog(q_r, analytical_func)
        plt.show()


def test_marginal_distributions_bvn_int():
    """Test so that the marginal distributions of polynomial_coefficentis_bvn_int is as expected"""
    delta_z_slice_int1 = np.linspace(-2.9, 2.9, 200.0)
    delta_z_slice_int2 = 2.9*np.ones_like(delta_z_slice_int1)
    sigma1 = 2.0
    sigma2 = 1.0
    poly = grating_diffuse.polynomial_cofficents_bvn_int(delta_z_slice_int1, delta_z_slice_int2, sigma1, sigma2)
    values = np.polyval(poly, 0.5*np.ones_like(delta_z_slice_int1))

    ref_values = 0.5*(1 + erf(delta_z_slice_int1/np.sqrt(2)/sigma1))

    iprint((np.abs((values - ref_values)/(ref_values)).max()))

    if True:
        plt.plot(delta_z_slice_int1, values)
        plt.plot(delta_z_slice_int1, ref_values)
        plt.show()

#test_calc_fields()
#test_calc_fields_single_layer()
#test_calc_fields_slicing_single_interface()
#test_calc_fields_slicing_single_layer()
test_correlation_function_h05()
#test_marginal_distributions_bvn_int()
