import numpy as np


def first_order_matrix(z, L, theta, k1, gamma_ref):
    """
    Calculate the first order matrix for the transfer map.
    This function is an adaptation of the one found in the particle tracking
    code Ocelot (see https://github.com/ocelot-collab/ocelot) written by
    S. Tomin.

    Parameters
    ----------
    z : float
        Longitudinal position in which to calculate the transfer matrix
    L : float
        Total length of the beamline element
    theta : float
        Bending angle of the beamline element
    k1 : float
        Quadrupole gradient of the beamline element in units of 1/m^2. A
        positive value implies focusing on the 'x' plane, while a negative
        gradient corresponds to focusing on 'y'
    k2 : float
        Sextupole gradient of the beamline element in units of 1/m^3. A
        positive value implies focusing on the 'x' plane, while a negative
        gradient corresponds to focusing on 'y'
    gamma_ref : float
        Reference energy with respect to which the particle momentum dp is
        calculated.

    """
    gamma = gamma_ref
    hx = theta / L
    kx2 = k1 + hx * hx
    ky2 = -k1
    kx = np.sqrt(kx2 + 0.0j)
    ky = np.sqrt(ky2 + 0.0j)
    cx = np.cos(z * kx).real
    cy = np.cos(z * ky).real
    sy = (np.sin(ky * z) / ky).real if ky != 0 else z
    igamma2 = 0.0
    if gamma != 0:
        igamma2 = 1.0 / (gamma * gamma)
    beta = np.sqrt(1.0 - igamma2)
    if kx != 0:
        sx = (np.sin(kx * z) / kx).real
        dx = hx / kx2 * (1.0 - cx)
        r56 = hx * hx * (z - sx) / kx2 / beta**2
    else:
        sx = z
        dx = z * z * hx / 2.0
        r56 = hx * hx * z**3 / 6.0 / beta**2
    r56 -= z / (beta * beta) * igamma2
    u_matrix = np.array(
        [
            [cx, sx, 0.0, 0.0, 0.0, dx / beta],
            [-kx2 * sx, cx, 0.0, 0.0, 0.0, sx * hx / beta],
            [0.0, 0.0, cy, sy, 0.0, 0.0],
            [0.0, 0.0, -ky2 * sy, cy, 0.0, 0.0],
            [hx * sx / beta, dx / beta, 0.0, 0.0, 1.0, r56],
            [0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
        ]
    )
    return u_matrix


def second_order_matrix(z, L, theta, k1, k2, gamma_ref):
    """
    Calculate the second order matrix for the transfer map.
    This function is an adaptation of the one found in the particle tracking
    code Ocelot (see https://github.com/ocelot-collab/ocelot) written by
    S. Tomin.

    Parameters
    ----------
    beam_matrix : array
        6 x N matrix, where N is the number of particles, containing the
        phase-space information of the bunch as (x, x', y, y', xi, dp) in
        units of (m, rad, m, rad, m, -). dp is defined as
        dp = (g-g_ref)/g_ref, while x' = px/p_kin and y' = py/p_kin, where
        p_kin is the reference kinetic momentum.
    z : float
        Longitudinal position in which to obtain the bunch distribution
    L : float
        Total length of the beamline element
    theta : float
        Bending angle of the beamline element
    k1 : float
        Quadrupole gradient of the beamline element in units of 1/m^2. A
        positive value implies focusing on the 'x' plane, while a negative
        gradient corresponds to focusing on 'y'
    k2 : float
        Sextupole gradient of the beamline element in units of 1/m^3. A
        positive value implies focusing on the 'x' plane, while a negative
        gradient corresponds to focusing on 'y'
    gamma_ref : float
        Reference energy with respect to which the particle momentum dp is
        calculated.

    """
    igamma2 = 0.0
    if gamma_ref != 0:
        gamma = gamma_ref
        gamma2 = gamma * gamma
        igamma2 = 1.0 / gamma2

    beta = np.sqrt(1.0 - igamma2)
    beta2 = beta * beta
    beta3 = beta2 * beta
    h = theta / L
    L = z  # Easy fix for the calculation below, where L should actually be z.
    h2 = h * h
    h3 = h2 * h
    kx2 = k1 + h * h
    ky2 = -k1
    kx4 = kx2 * kx2
    kx = np.sqrt(kx2 + 0.0j)
    ky = np.sqrt(-k1 + 0.0j)
    cx = np.cos(kx * L).real
    sx = (np.sin(kx * L) / kx).real if kx != 0 else L
    cy = np.cos(ky * L).real

    sy = (np.sin(ky * L) / ky).real if ky != 0 else L

    sx2 = sx * sx
    sy2 = sy * sy
    L2 = L * L
    L3 = L2 * L
    L4 = L3 * L
    L5 = L4 * L
    dx = (1.0 - cx) / kx2 if kx != 0.0 else L * L / 2.0
    dx2 = dx * dx
    K1 = k1
    K2 = k2

    d2y = 0.5 * sy * sy
    s2y = sy * cy
    c2y = np.cos(2 * ky * L).real
    fx = (L - sx) / kx2 if kx2 != 0 else L3 / 6.0
    f2y = (L - s2y) / ky2 if ky2 != 0 else L3 / 6
    J1 = (L - sx) / kx2 if kx2 != 0 else L3 / 6.0
    J2 = (3.0 * L - 4.0 * sx + sx * cx) / (2 * kx4) if kx != 0 else L**5 / 20.0
    J3 = (
        (
            (15 * L - 22.5 * sx + 9 * sx * cx - 1.5 * sx * cx * cx + kx2 * sx2 * sx)
            / (6 * kx4 * kx2)
        )
        if kx != 0
        else L**7 / 56.0
    )
    # J3 = (15 * L - 22 * sx + 9 * sx * cx - 2 * sx * cx * cx ) / (
    #        6 * kx4 * kx2) if kx != 0 else L ** 7 / 56.
    J_denom = kx2 - 4 * ky2
    Jc = (c2y - cx) / J_denom if J_denom != 0 else 0.5 * L2
    Js = (cy * sy - sx) / J_denom if J_denom != 0 else L3 / 6
    Jd = (d2y - dx) / J_denom if J_denom != 0 else L4 / 24
    Jf = (f2y - fx) / J_denom if J_denom != 0 else L5 / 120

    khk = K2 + 2 * h * K1
    T = np.zeros((6, 6, 6))

    T111 = -1 / 6 * khk * (sx2 + dx) - 0.5 * h * kx2 * sx2
    T112 = -1 / 6 * khk * sx * dx + 0.5 * h * sx * cx
    T122 = -1 / 6 * khk * dx2 + 0.5 * h * dx * cx
    T116 = (
        -h / 12 / beta * khk * (3 * sx * J1 - dx2)
        + 0.5 * h2 / beta * sx2
        + 0.25 / beta * K1 * L * sx
    )
    T126 = (
        -h / 12 / beta * khk * (sx * dx2 - 2 * cx * J2)
        + 0.25 * h2 / beta * (sx * dx + cx * J1)
        - 0.25 / beta * (sx + L * cx)
    )
    T166 = (
        -h2 / 6 / beta2 * khk * (dx2 * dx - 2 * sx * J2)
        + 0.5 * h3 / beta2 * sx * J1
        - 0.5 * h / beta2 * L * sx
        - 0.5 * h / (beta2) * igamma2 * dx
    )
    T133 = K1 * K2 * Jd + 0.5 * (K2 + h * K1) * dx
    T134 = 0.5 * K2 * Js
    T144 = K2 * Jd - 0.5 * h * dx

    T211 = -1 / 6 * khk * sx * (1 + 2 * cx)
    T212 = -1 / 6 * khk * dx * (1 + 2 * cx)
    T222 = -1 / 3 * khk * sx * dx - 0.5 * h * sx
    T216 = -h / 12 / beta * khk * (3 * cx * J1 + sx * dx) - 0.25 / beta * K1 * (
        sx - L * cx
    )
    T226 = -h / 12 / beta * khk * (3 * sx * J1 + dx2) + 0.25 / beta * K1 * L * sx
    T266 = (
        -h2 / 6 / beta2 * khk * (sx * dx2 - 2 * cx * J2)
        - 0.5 * h / beta2 * K1 * (cx * J1 - sx * dx)
        - 0.5 * h / beta2 * igamma2 * sx
    )
    T233 = K1 * K2 * Js + 0.5 * (K2 + h * K1) * sx
    T234 = 0.5 * K2 * Jc
    T244 = K2 * Js - 0.5 * h * sx

    T313 = 0.5 * K2 * (cy * Jc - 2 * K1 * sy * Js) + 0.5 * h * K1 * sx * sy
    T314 = 0.5 * K2 * (sy * Jc - 2 * cy * Js) + 0.5 * h * sx * cy
    T323 = 0.5 * K2 * (cy * Js - 2 * K1 * sy * Jd) + 0.5 * h * K1 * dx * sy
    T324 = 0.5 * K2 * (sy * Js - 2 * cy * Jd) + 0.5 * h * dx * cy
    T336 = (
        0.5 * h / beta * K2 * (cy * Jd - 2 * K1 * sy * Jf)
        + 0.5 * h2 / beta * K1 * J1 * sy
        - 0.25 / beta * K1 * L * sy
    )
    T346 = (
        0.5 * h / beta * K2 * (sy * Jd - 2 * cy * Jf)
        + 0.5 * h2 / beta * J1 * cy
        - 0.25 / beta * (sy + L * cy)
    )

    T413 = 0.5 * K1 * K2 * (2 * cy * Js - sy * Jc) + 0.5 * (K2 + h * K1) * sx * cy
    T414 = 0.5 * K2 * (2 * K1 * sy * Js - cy * Jc) + 0.5 * (K2 + h * K1) * sx * sy
    T423 = 0.5 * K1 * K2 * (2 * cy * Jd - sy * Js) + 0.5 * (K2 + h * K1) * dx * cy
    T424 = 0.5 * K2 * (2 * K1 * sy * Jd - cy * Js) + 0.5 * (K2 + h * K1) * dx * sy
    T436 = (
        0.5 * h / beta * K1 * K2 * (2 * cy * Jf - sy * Jd)
        + 0.5 * h / beta * (K2 + h * K1) * J1 * cy
        + 0.25 / beta * K1 * (sy - L * cy)
    )
    T446 = (
        0.5 * h / beta * K2 * (2 * K1 * sy * Jf - cy * Jd)
        + 0.5 * h / beta * (K2 + h * K1) * J1 * sy
        - 0.25 / beta * K1 * L * sy
    )

    T511 = h / 12 / beta * khk * (sx * dx + 3 * J1) - 0.25 / beta * K1 * (L - sx * cx)
    T512 = h / 12 / beta * khk * dx2 + 0.25 / beta * K1 * sx2
    T522 = h / 6 / beta * khk * J2 - 0.5 / beta * sx - 0.25 / beta * K1 * (J1 - sx * dx)

    T516 = (
        h2 / 12 / beta2 * khk * (3 * dx * J1 - 4 * J2)
        + 0.25 * h / beta2 * K1 * J1 * (1 + cx)
        + 0.5 * h / beta2 * igamma2 * sx
    )

    T526 = (
        h2 / 12 / beta2 * khk * (dx * dx2 - 2 * sx * J2)
        + 0.25 * h / beta2 * K1 * sx * J1
        + 0.5 * h / beta2 * igamma2 * dx
    )
    T566 = (
        h3 / 6 / beta3 * khk * (3 * J3 - 2 * dx * J2)
        + h2 / 6 / beta3 * K1 * (sx * dx2 - J2 * (1 + 2 * cx))
        + 1.5 / beta3 * igamma2 * (h2 * J1 - L)
    )
    T533 = (
        -h / beta * K1 * K2 * Jf
        - 0.5 * h / beta * (K2 + h * K1) * J1
        + 0.25 / beta * K1 * (L - cy * sy)
    )
    T534 = -0.5 * h / beta * K2 * Jd - 0.25 / beta * K1 * sy2
    T544 = -h / beta * K2 * Jf + 0.5 * h2 / beta * J1 - 0.25 / beta * (L + cy * sy)

    T[0, 0, 0] = T111
    T[0, 0, 1] = T112 * 2
    T[0, 1, 1] = T122
    T[0, 0, 5] = T116 * 2
    T[0, 1, 5] = T126 * 2
    T[0, 5, 5] = T166
    T[0, 2, 2] = T133
    T[0, 2, 3] = T134 * 2
    T[0, 3, 3] = T144

    T[1, 0, 0] = T211
    T[1, 0, 1] = T212 * 2
    T[1, 1, 1] = T222
    T[1, 0, 5] = T216 * 2
    T[1, 1, 5] = T226 * 2
    T[1, 5, 5] = T266
    T[1, 2, 2] = T233
    T[1, 2, 3] = T234 * 2
    T[1, 3, 3] = T244

    T[2, 0, 2] = T313 * 2
    T[2, 0, 3] = T314 * 2
    T[2, 1, 2] = T323 * 2
    T[2, 1, 3] = T324 * 2
    T[2, 2, 5] = T336 * 2
    T[2, 3, 5] = T346 * 2

    T[3, 0, 2] = T413 * 2
    T[3, 0, 3] = T414 * 2
    T[3, 1, 2] = T423 * 2
    T[3, 1, 3] = T424 * 2
    T[3, 2, 5] = T436 * 2
    T[3, 3, 5] = T446 * 2

    T[4, 0, 0] = -T511
    T[4, 0, 1] = -T512 * 2
    T[4, 1, 1] = -T522
    T[4, 0, 5] = -T516 * 2
    T[4, 1, 5] = -T526 * 2
    T[4, 5, 5] = -T566
    T[4, 2, 2] = -T533
    T[4, 2, 3] = -T534 * 2
    T[4, 3, 3] = -T544
    return T
