import sympy as sp

class HighLevel():

    def __init__(self, j):
        self.j = j

    def initial(self):
        j = self.j

        #init parameters
        h, b, n, u = 0.1, None, None, None
        x, y, z, px, py, pz = None, 0.1, 0.001, 0.01, None, 0.0001
        xa, ya, za, pxa, pya, pza = None, None, None, None, None, None

        if j == 1:
            b = 5.0 / 4
            n = b / (1 + b) ** 2
            u = 1.0 / (1.0 / b + b + 2.0)
            x = 10.0
            py = 0.5
        elif j == 2:
            b = 3.0 / 4
            n = b / (1 + b) ** 2
            u = 1.0 / (1.0 / b + b + 2.0)
            x = 8.3
            py = 0.6
        elif j == 3:
            b = 3.0 / 2
            x = 12.0
            py = 0.4
        elif j == 4:
            b = 7.0 / 4
            n = b / (1 + b) ** 2
            u = 1.0 / (1.0 / b + b + 2.0)
            x = 15.0
            py = 0.35
        elif j == 5:
            b = 1.0
            n = b / (1 + b) ** 2
            u = 1.0 / (1.0 / b + b + 2.0)
            x = 18.0
            py = 0.3
        elif j == 6:
            b = 3.0 / 5
            n = b / (1 + b) ** 2
            u = 1.0 / (1.0 / b + b + 2.0)
            x = 20.0
            py = 0.25
        elif j == 7:
            b = 5.0 / 7
            n = b / (1 + b) ** 2
            u = 1.0 / (1.0 / b + b + 2.0)
            x = 22.0
            py = 0.22
        elif j == 8:
            b = 2.0
            x = 26.0
            py = 0.2
        elif j == 9:
            b = 0.5
            n = b / (1 + b) ** 2
            u = 1.0 / (1.0 / b + b + 2.0)
            x = 30.0
            y = 0.5
            z = 0.1
            pz = 0.01
        elif j == 10:
            b = 5.0
            n = b / (1 + b) ** 2
            u = 1.0 / (1.0 / b + b + 2.0)
            x = 35.0
            y = 2.0
            z = 0.1
            pz = 0.03
            py = 0.15

        xa, ya, za, pxa, pya, pza = x, y, z, px, py, pz
        return j, h, b, n, x, y, z, xa, ya, za, px, py, pz, pxa, pya, pza

    def f(self, x, y, z, px, py, pz, b):
        x_val, y_val, z_val, px_val, py_val, pz_val, b_val = x, y, z, px, py, pz, b
        x, y, z, px, py, pz, b = sp.symbols('x y z px py pz b')

        c = 1.0

        u = 1 / (1 / b + b + 2)
        ht = px ** 2 / 2 + py ** 2 / 2 + pz ** 2 / 2
        hv = -1 / (x ** 2 + y ** 2 + z ** 2) ** (1 / 2)
        h1pn = 1 / (2 * x ** 2 + 2 * y ** 2 + 2 * z ** 2) - (((u + 3) * (px ** 2 + py ** 2 + pz ** 2)) / 2 + (u * (
                    (px * x) / (x ** 2 + y ** 2 + z ** 2) ** (1 / 2) + (py * y) / (x ** 2 + y ** 2 + z ** 2) ** (
                        1 / 2) + (pz * z) / (x ** 2 + y ** 2 + z ** 2) ** (1 / 2)) ** 2) / 2) / (
                           x ** 2 + y ** 2 + z ** 2) ** (1 / 2) + ((3 * u) / 8 - 1 / 8) * (
                           px ** 2 + py ** 2 + pz ** 2) ** 2

        e = ht + hv + h1pn

        de_dx = sp.diff(e, x)
        de_dy = sp.diff(e, y)
        de_dz = sp.diff(e, z)
        de_dpx = sp.diff(e, px)
        de_dpy = sp.diff(e, py)
        de_dpz = sp.diff(e, pz)

        de_dx_val = de_dx.subs({x: x_val, y: y_val, z: z_val, px: px_val, py: py_val, pz: pz_val, b: b_val})
        de_dy_val = de_dy.subs({x: x_val, y: y_val, z: z_val, px: px_val, py: py_val, pz: pz_val, b: b_val})
        de_dz_val = de_dz.subs({x: x_val, y: y_val, z: z_val, px: px_val, py: py_val, pz: pz_val, b: b_val})
        de_dpx_val = de_dpx.subs({x: x_val, y: y_val, z: z_val, px: px_val, py: py_val, pz: pz_val, b: b_val})
        de_dpy_val = de_dpy.subs({x: x_val, y: y_val, z: z_val, px: px_val, py: py_val, pz: pz_val, b: b_val})
        de_dpz_val = de_dpz.subs({x: x_val, y: y_val, z: z_val, px: px_val, py: py_val, pz: pz_val, b: b_val})

        e_val = e.subs({x: x_val, y: y_val, z: z_val, px: px_val, py: py_val, pz: pz_val, b: b_val})

        return de_dx_val, de_dy_val, de_dz_val, de_dpx_val, de_dpy_val, de_dpz_val, e_val

    def rejust(self, x, y, z, px, py, pz, xa, ya, za, pxa, pya, pza):

        x = (x + xa) / 2
        y = (y + ya) / 2
        z = (z + za) / 2

        px = (px + pxa) / 2
        py = (py + pya) / 2
        pz = (pz + pza) / 2
        xa = x
        ya = y
        za = z
        pxa = px
        pya = py
        pza = pz

        return x, y, z, px, py, pz, xa, ya, za, pxa, pya, pza

    def symplectic(self, h, x, y, z, px, py, pz, xa, ya, za, pxa, pya, pza, b):

        vxa, vya, vza, vpx, vpy, vpz, e = self.f(xa, ya, za, px, py, pz, b)
        x = x + h / 2 * vpx
        y = y + h / 2 * vpy
        z = z + h / 2 * vpz
        pxa = pxa - h / 2 * vxa
        pya = pya - h / 2 * vya
        pza = pza - h / 2 * vza

        vx, vy, vz, vpxa, vpya, vpza, e = self.f(x, y, z, pxa, pya, pza, b)
        xa = xa + h * vpxa
        ya = ya + h * vpya
        za = za + h * vpza
        px = px - h * vx
        py = py - h * vy
        pz = pz - h * vz

        vxa, vya, vza, vpx, vpy, vpz, e = self.f(xa, ya, za, px, py, pz, b)
        x = x + h / 2 * vpx
        y = y + h / 2 * vpy
        z = z + h / 2 * vpz
        pxa = pxa - h / 2 * vxa
        pya = pya - h / 2 * vya
        pza = pza - h / 2 * vza

        return x, y, z, px, py, pz, xa, ya, za, pxa, pya, pza