'''
Vacuum beam propagation
no plasma, no limiter nor reflections;
1D beam parameters;
general astigmatic beam
'''

from .. import TestCase, GrayTest
from .. import options, load_unit

import numpy as np

# SI unit converions
cm = 1e-2
mm = 1e-3
deg = np.pi/180
c = 299792458


def vectorize(*arg, **kwargs):
    '''
    Helper to use np.vectorize with keyword arguments as a decorator
    '''
    return lambda f: np.vectorize(f, *arg, **kwargs)


def rotate(angle: float):
    '''
    Builds a 2D rotation matrix
    '''
    c, s = np.cos(angle), np.sin(angle)
    return np.array([[c, -s], [s, c]])


@vectorize(excluded=['k0', 'W', 'K', 'φ_w', 'φ_k'],
           signature='(),(2,2),(2,2),(),(),(),()->()')
def eikonal_true(k0: float, W: np.array, K: np.array,
                 φ_w: float, φ_k: float, s: float,
                 ray: dict[str, float]):
    '''
    Computes the complex eikonal from the initial
    beam parameters and the final ray position.
    '''
    r = np.array([ray['xt'], ray['yt']]) * cm
    # Note: z = s₀ + z̅ where s₀ is the arclength of the
    # central ray and z̅ the local beam coordinate.
    z = (s + ray['zt']) * cm

    # compute the complex curvature tensor
    K = rotate(φ_k) @ K @ rotate(φ_k).T
    W = rotate(φ_w) @ W @ rotate(φ_w).T
    Q = K - 1j*W

    # propagate the beam from 0→z
    Q = Q @ np.linalg.inv(np.eye(2) + z*Q)

    # Compute the eikonal
    # Note: the Gouy phase is not treated by gray
    S = k0 * (z + 1/2 * r @ Q @ r)

    return S


@vectorize(excluded=['k0', 'W', 'K', 'φ_w', 'φ_k'],
           signature='(),(2,2),(2,2),(),(),()->()')
def eikonal_gray(k0: float, W: np.array, K: np.array,
                 φ_w: float, φ_k: float, ray: dict[str, float]):
    '''
    Extracts the eikonal from the gray results

    Notes:
      1. the coordinates r,θ in the beam cross section
         are defined by:
           { ξ_w = w_ξ r cosθ
           { η_w = w_η r cosθ,

      2. for the (j,k)-th ray we have:
           { r = (j-1)/(n_r - 1)
           { θ = (k-1)/n_θ

        so S_I = -(ξ_w²/w_ξ² + η_w²/w_η²) = -r²
        and this should be independent of s.

      3. S_R is given by k₀s, when integrating in
         the phase (idst=2);
    '''
    r_max = 1.0
    n_r = 6
    δr = r_max / (n_r - 1)
    r = (ray['j'] - 1) * δr

    S_R = k0*ray['sst'] * cm
    S_I = -r**2

    return S_R + 1j * S_I


class Test(GrayTest, TestCase):
    def test_eikonal(self):
        '''
        Check the correctness of beam propagation
        by comparing the values of eikonal on each ray
        at the final step.
        '''
        # Read the initial beam parameters
        w_ξ, w_η, k_ξ, k_η, φ_w, φ_k = np.loadtxt(
            str(self.inputs / 'beam.txt'),
            usecols=[6, 7, 8, 9, 10, 11], skiprows=2,
            unpack=True)

        # Convert to SI
        w_ξ, w_η = w_ξ * mm, w_η * mm
        k_ξ, k_η = k_ξ / mm, k_η / mm
        φ_w, φ_k = φ_w * deg, φ_k * deg

        # Vacuum wavenumber
        k0 = 2*np.pi/c * 170e9

        # Beam width/curvature
        W = 2/k0 * np.diag([1/w_ξ, 1/w_η])**2
        K = np.diag([k_ξ, k_η])

        # Load the final rays data
        data = load_unit(self.candidate / 'fort.9')

        # Complex eikonal: analytical and gray values
        s0 = data['sst'][0]
        S_true = eikonal_true(k0, W, K, φ_w, φ_k, s0, data)
        S_gray = eikonal_gray(k0, W, K, φ_w, φ_k, data)

        if options.visual:
            import matplotlib.pyplot as plt

            plt.suptitle(self.__module__ + '.test_eikonal')
            plt.subplot(121, aspect='equal')
            plt.title('$|S_R - S_{R,gray}|$', loc='left')
            plt.tricontour(data['xt'], data['yt'],
                           abs(S_true.real - S_gray.real),
                           levels=8, linewidths=0.5, colors='k')
            plt.tricontourf(data['xt'], data['yt'],
                            abs(S_true.real - S_gray.real),
                            levels=8, cmap='summer')
            plt.colorbar()
            plt.scatter(data['xt'], data['yt'], c='k', s=5)

            plt.subplot(122, aspect='equal')
            plt.title('$|S_I - S_{I,gray}|$', loc='left')
            plt.tricontour(data['xt'], data['yt'],
                           abs(S_true.imag - S_gray.imag),
                           levels=8, linewidths=0.5, colors='k')
            plt.tricontourf(data['xt'], data['yt'],
                            abs(S_true.imag - S_gray.imag),
                            levels=8, cmap='summer')
            plt.colorbar()
            plt.scatter(data['xt'], data['yt'], c='k', s=5)
            plt.show()

        # Compare the eikonal of each ray
        for jk, ray in enumerate(data):
            with self.subTest(ray=(int(ray['j']), int(ray['k']))):
                self.assertAlmostEqual(S_true[jk].imag, S_gray[jk].imag, 2)
                self.assertAlmostEqual(S_true[jk].real, S_gray[jk].real, 1)