import sys
import os.path
import re
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import h5py
import sympy

## These are optional configurations
#mpl.rcParams['figure.figsize'] = 8,6  # Set default figure size (in inches)
#%config InlineBackend.figure_format='svg'  # If plotting inline, use SVG or PNG
sympy.init_printing()

## These import some of the custom code; they must be installed properly first
import GWFrames
import GWFrames.plot
import Quaternions

## These are a few convenient symbols we use below
v, nu, delta, Sigma_n, Sigma_lambda = sympy.symbols('v, nu, delta, Sigma_n, Sigma_lambda', real=True)

## Set up the paths correctly
PathToSXSCatalogLinks = os.path.expanduser('~')+'/Research/Data/SimulationAnnex/CatalogLinks/'
PathToSXSCatalog = os.path.expanduser('~')+'/Research/Data/SimulationAnnex/Catalog/'
MyFavoriteSystem = PathToSXSCatalog + 'q1.0/SBBH/d19.0_q1.0_s0.5_0_0_s0_0_0/Lev5/'
MyFavoriteWaveform = 'rhOverM_Asymptotic_GeometricUnits.h5/Extrapolated_N2.dir'
MyFavoriteData = MyFavoriteSystem + MyFavoriteWaveform

## Read in the waveform and adjust it
W_NR = GWFrames.ReadFromNRAR(MyFavoriteData);
t0 = -W_NR.MaxNormTime()
t1 = 200
W_NR.SetT(W_NR.T()+t0)
print('The initial time has been moved back to {0}'.format(t0))

## Get the coordinate-position data from the h5 file
with h5py.File(MyFavoriteSystem+'Horizons.h5', 'r') as f:
    tA = f['AhA.dir/CoordCenterInertial.dat'][:,0]+t0
    xA = f['AhA.dir/CoordCenterInertial.dat'][:,1:]
    mA = f['AhA.dir/ArealMass.dat'][:,1]
    chiA = f['AhA.dir/chiInertial.dat'][:,1:]
    tB = f['AhB.dir/CoordCenterInertial.dat'][:,0]+t0
    xB = f['AhB.dir/CoordCenterInertial.dat'][:,1:]
    mB = f['AhB.dir/ArealMass.dat'][:,1]
    chiB = f['AhB.dir/chiInertial.dat'][:,1:]

## Do some calculations with the coordinates
d = xA-xB
nHat = array([[q[1],q[2],q[3]] for q in Quaternions.normalized(Quaternions.QuaternionArray(d))])
dnHatdt = array([[q[1],q[2],q[3]] for q in Quaternions.QuaternionDerivative(Quaternions.QuaternionArray(nHat), tA)])
lambdaHat = array([[q[1],q[2],q[3]] for q in Quaternions.normalized(Quaternions.QuaternionArray(dnHatdt))])
Omega_orb = cross(nHat, dnHatdt)
v = array([(o[0]**2+o[1]**2+o[2]**2)**(1./6.) for o in Omega_orb])
COM = ((xA.transpose()*mA+xB.transpose()*mB)/(mA+mB)).transpose()

## Set the average initial velocity from the inital-data ADM quantities
initial_ADM_energy = 0.0
initial_ADM_linear_momentum = []
number_regex = r'[-+eE\.\d]*'
with open(MyFavoriteSystem+'metadata.txt','r') as f:
    for line in f:
        if line.startswith('initial-ADM-energy'):
            initial_ADM_energy = float(re.search('.* = *({0})'.format(number_regex),line).group(1))
        if line.startswith('initial-ADM-linear-momentum'):
            initial_ADM_linear_momentum = array([float(x)
                                                 for x in re.search('.* = *({0}), *({0}), *({0})'.format(number_regex),line).groups()])
        if (len(initial_ADM_linear_momentum)!=0 and initial_ADM_energy!=0.0):
            break
v_avg = initial_ADM_linear_momentum/initial_ADM_energy
print("The speed of the drift we are removing is {0}".format(sqrt(v_avg.dot(v_avg))))

## Translate the waveform to remove average velocity drift
W_NR = W_NR.Translate([v_avg*(t-W_NR.T(0)) for t in W_NR.T()])

## Make a copy of the waveform, and transform it to the co-rotating frame
UseEllModes = [2,3,4] # Which ell modes to use to calculate the angular velocity and hence corotating frame
t_alignment = t0+2000. # Time at which to align the modes to the frame
W_NR_corot = GWFrames.Waveform(W_NR)
W_NR_corot.TransformToCorotatingFrame(UseEllModes)
W_NR_corot.AlignDecompositionFrameToModes(t_alignment, Quaternions.Quaternion(nHat[argmin(abs(tA-t_alignment))]), UseEllModes);

## Just briefly read in the non-precessing waveforms
W_NR_nonpreces = GWFrames.ReadFromNRAR(PathToSXSCatalogLinks + 'SXS:BBH:0002/Lev6/' + MyFavoriteWaveform);
W_NR_nonpreces.SetT(W_NR_nonpreces.T()-W_NR_nonpreces.MaxNormTime());
W_NR_nonpreces_corot = GWFrames.Waveform(W_NR_nonpreces);
W_NR_nonpreces_corot.TransformToCorotatingFrame(UseEllModes);
W_NR_nonpreces_corot.AlignDecompositionFrameToModes(t_alignment, Quaternions.Quaternion(nHat[argmin(abs(tA-t_alignment))]), UseEllModes);
W_NR_eccentric = GWFrames.ReadFromNRAR(PathToSXSCatalogLinks + 'SXS:BBH:0091/Lev5/' + MyFavoriteWaveform);
W_NR_eccentric.SetT(W_NR_eccentric.T()-W_NR_eccentric.MaxNormTime());
W_NR_eccentric_corot = GWFrames.Waveform(W_NR_eccentric);
W_NR_eccentric_corot.TransformToCorotatingFrame(UseEllModes);
W_NR_eccentric_corot.AlignDecompositionFrameToModes(t_alignment, Quaternions.Quaternion(nHat[argmin(abs(tA-t_alignment))]), UseEllModes);

## Plot the amplitudes of the (2,2) and (2,-2) modes
W_NR_corot.plot('Abs', [[2,2],[2,-2]])
plt.gca().set_color_cycle(None)
W_NR_nonpreces_corot.plot('Abs', [[2,2],[2,-2]], ls='dashed')
plt.gca().set_color_cycle(None)
W_NR_eccentric_corot.plot('Abs', [[2,2],[2,-2]], ls='dotted')
ylim(7e-2,0.15)
xlim(t0,t1)
title(r'Comparing the $(\ell,m) = (2,\pm 2)$ mode pair amplitudes')
legend(loc='upper left');

## Save the data to file, for direct incorporation into the main LaTeX file via `pgfplot`
numpy.savetxt('data/ModeAmplitudes_NR_2_2.dat',
              array([W_NR_corot.T(),
                     W_NR_corot.Abs(W_NR_corot.FindModeIndex(2,2)),
                     W_NR_corot.Abs(W_NR_corot.FindModeIndex(2,-2))]).transpose()[::60],
              fmt='%.5e')
numpy.savetxt('data/ModeAmplitudes_Nonspinning_2_2.dat',
              array([W_NR_nonpreces_corot.T(),
                     W_NR_nonpreces_corot.Abs(W_NR_nonpreces_corot.FindModeIndex(2,2)),
                     W_NR_nonpreces_corot.Abs(W_NR_nonpreces_corot.FindModeIndex(2,-2))]).transpose()[::100],
              fmt='%.5e')
numpy.savetxt('data/ModeAmplitudes_Eccentric_2_2.dat',
              array([W_NR_eccentric_corot.T(),
                     W_NR_eccentric_corot.Abs(W_NR_eccentric_corot.FindModeIndex(2,2)),
                     W_NR_eccentric_corot.Abs(W_NR_eccentric_corot.FindModeIndex(2,-2))]).transpose()[::100],
              fmt='%.5e')

## Plot the phases
W_NR_corot.plot('Arg', [[2,2],[2,-2]])
plt.gca().set_color_cycle(None)
W_NR_nonpreces_corot.plot('Arg', [[2,2],[2,-2]], ls='dashed')
plt.gca().set_color_cycle(None)
W_NR_eccentric_corot.plot('Arg', [[2,2],[2,-2]], ls='dotted')
ylim(-0.06,0.06)
xlim(t0,t1)
title(r'Comparing the $(\ell,m) = (2,\pm 2)$ mode pair phases')
legend(loc='lower left');

## Save the data to file, for direct incorporation into the main LaTeX file via `pgfplot`
numpy.savetxt('data/ModePhases_NR_2_2.dat',
              array([W_NR_corot.T(),
                     W_NR_corot.Arg(W_NR_corot.FindModeIndex(2,2)),
                     W_NR_corot.Arg(W_NR_corot.FindModeIndex(2,-2))]).transpose()[::60],
              fmt='%.5e')
numpy.savetxt('data/ModePhases_Nonspinning_2_2.dat',
              array([W_NR_nonpreces_corot.T(),
                     W_NR_nonpreces_corot.Arg(W_NR_nonpreces_corot.FindModeIndex(2,2)),
                     W_NR_nonpreces_corot.Arg(W_NR_nonpreces_corot.FindModeIndex(2,-2))]).transpose()[::100],
              fmt='%.5e')
numpy.savetxt('data/ModePhases_Eccentric_2_2.dat',
              array([W_NR_eccentric_corot.T(),
                     W_NR_eccentric_corot.Arg(W_NR_eccentric_corot.FindModeIndex(2,2)),
                     W_NR_eccentric_corot.Arg(W_NR_eccentric_corot.FindModeIndex(2,-2))]).transpose()[::100],
              fmt='%.5e')

v,M = sympy.symbols('v,M', real=True)
h2p2 = 1 - v**2*(Sigma_lambda + sympy.I*Sigma_n)/(2*M**2)
h2m2 = 1 + v**2*(Sigma_lambda - sympy.I*Sigma_n)/(2*M**2)

sympy.series((sympy.sqrt(h2p2*h2p2.conjugate()) - sympy.sqrt(h2m2*h2m2.conjugate()))
             /(sympy.sqrt(h2p2*h2p2.conjugate()) + sympy.sqrt(h2m2*h2m2.conjugate())), x=v, x0=0, n=3)

def RelativeAmplitudeDifference(TotalMassInSolarMasses, FrequencyInHertz):
    "Scale of relative amplitude differences as a percentage"
    from scipy.constants import c, G
    from numpy import pi
    SolarMass = 1.9891e30 # kilograms
    return 50*(G*TotalMassInSolarMasses*SolarMass*pi*FrequencyInHertz/c**3)**(2./3.)

RelativeAmplitudeDifference(18.275, 10.), RelativeAmplitudeDifference(18.275, 100.), RelativeAmplitudeDifference(18.275, 1000.)

RelativeAmplitudeDifference(10., 10.), RelativeAmplitudeDifference(10., 100.), RelativeAmplitudeDifference(10., 1000.)

%%time
W_NR_corot_short = W_NR_corot.Interpolate(W_NR_corot.T()[::10])
MinimalParityViolation_NR = W_NR_corot_short.MinimalParityViolation()
Antisymmetry_NR = W_NR_corot_short.NormalizedAntisymmetry()

# Plot the two data sets
semilogy(W_NR_corot_short.T(), MinimalParityViolation_NR, label=r'$p$')
semilogy(W_NR_corot_short.T(), Antisymmetry_NR, label=r'$a$')
ylim((2.2e-3,1))
xlim((t0,t1))
xlabel(r'$(t-t_{\mathrm{merger}})/M$')
ylabel(r'Rotationally invariant antisymmetries')
legend(loc='upper left');

# Save the data to file, for direct incorporation into the main LaTeX file via `pgfplot`
numpy.savetxt('data/Asymmetry_NR.dat', array([W_NR_corot_short.T(), Antisymmetry_NR, MinimalParityViolation_NR]).transpose(),
              fmt='%.5e')

# Relative antisymmetries near merger, as percentages
i_merger = W_NR_corot_short.MaxNormIndex()
Antisymmetry_NR[i_merger:i_merger+10]*100

# Just plot between these times, so it's not too messy
t_1 = W_NR.T(0)+375.
t_2 = -3000.
i_1 = argmin(abs(W_NR.T()-t_1))
i_2 = argmin(abs(W_NR.T()-t_2))

# Get the waveform vectors
AV_NRHat = array([o/np.sqrt(o.dot(o)) for o in W_NR.AngularVelocityVector()[i_1:i_2,:]])
Vh_NRHat = array([-o/np.sqrt(o.dot(o)) for o in W_NR.LLDominantEigenvector()[i_1:i_2,:]])
Ldt_NRHat = array([-o/np.sqrt(o.dot(o)) for o in W_NR.LdtVector()[i_1:i_2,:]])

# Now get the vectors for the orbital motions
i_3 = argmin(abs(tA-t_1))
i_4 = argmin(abs(tA-t_2))
Omega_orbHat = array([o/np.sqrt(o.dot(o)) for o in Omega_orb])[i_3:i_4,:]
R_orb = Quaternions.UnflipRotors(
            Quaternions.FrameFromXY(
                Quaternions.QuaternionArray(nHat), Quaternions.QuaternionArray(lambdaHat)
            )
        )
Omega_tot = Quaternions.vec(Quaternions.FrameAngularVelocity(R_orb, tA))
Omega_tot = array(Omega_tot)
Omega_totHat = array([o/np.sqrt(o.dot(o)) for o in Omega_tot])[i_3:i_4,:]

# Plot the vector directions
plot(Omega_orbHat[:,0], Omega_orbHat[:,1], label=r"$\hat{\Omega}_{\mathrm{orb}}=\hat{\ell}$")
plot(Omega_totHat[:,0], Omega_totHat[:,1], label=r"$\hat{\Omega}_{\mathrm{tot}}$")
plot(Vh_NRHat[:,0], Vh_NRHat[:,1], label=r"$\hat{V}_h$")
plot(Ldt_NRHat[:,0], Ldt_NRHat[:,1], label=r"$\langle L \partial_t \rangle$")
plot(AV_NRHat[:,0], AV_NRHat[:,1], label=r"$\hat{\omega}$")

# Plot the starting points as dots
plot([Omega_orbHat[0,0],], [Omega_orbHat[0,1],], 'o')
plot([Omega_totHat[0,0],], [Omega_totHat[0,1],], 'o')
plot([Vh_NRHat[0,0],], [Vh_NRHat[0,1],], 'o')
plot([Ldt_NRHat[0,0],], [Ldt_NRHat[0,1],], 'o')
plot([AV_NRHat[0,0],], [AV_NRHat[0,1],], 'o')
xlabel(r'$x$')
ylabel(r'$y$')
xlim((-0.1, 0.3))
ylim((-0.125,0.125))
legend();

# Save the data to file, for direct incorporation into the main LaTeX file via `pgfplot`
numpy.savetxt('data/WaveformVectors_NR.dat',
              array([AV_NRHat[:,0], AV_NRHat[:,1], Vh_NRHat[:,0], Vh_NRHat[:,1],
                     Ldt_NRHat[:,0], Ldt_NRHat[:,1],]).transpose()[::20],
              fmt='%.5e')
numpy.savetxt('data/WaveformVectors_Coordinates_NR.dat',
              array([Omega_orbHat[:,0], Omega_orbHat[:,1], Omega_totHat[:,0], Omega_totHat[:,1]]).transpose()[::10],
              fmt='%.5e')

# Get the initial conditions for the PN data
t_alignment = t0+800.
i_i = argmin(abs(tA-t_alignment))
t_i = tA[i_i]
approximant = 'TaylorT1'
delta = float( (mA[i_i]-mB[i_i])/(mA[i_i]+mB[i_i]) )
chiA_i = [float(c) for c in chiA[i_i]]
chiB_i = [float(c) for c in chiB[i_i]]
Omega_orb_i = float( np.sqrt(sum([o**2 for o in Omega_orb[i_i]])) )
R_frame_i = Quaternions.FrameFromXY([Quaternions.Quaternion(nHat[i_i]),], [Quaternions.Quaternion(lambdaHat[i_i]),])[0]

# Evolve the PN system
W_PN_coorb = GWFrames.PNWaveform(approximant, delta, chiA_i, chiB_i, Omega_orb_i, 0.9*Omega_orb_i, R_frame_i)
W_PN_coorb.SetT(W_PN_coorb.T()+t0)

# Make copies in the inertial and corotating frames
W_PN = GWFrames.PNWaveform(W_PN_coorb)
W_PN.TransformToInertialFrame();
W_PN_corot = GWFrames.PNWaveform(W_PN)
W_PN_corot.TransformToCorotatingFrame();

# Align PN to NR
GWFrames.AlignWaveforms(W_NR_corot, W_PN_corot, t_alignment+250., t_alignment+2000., 100);
W_PN = GWFrames.PNWaveform(W_PN_corot)
W_PN.TransformToInertialFrame();

# Plot the amplitudes of the (2,2) and (2,-2) modes
W_PN_corot.plot('Abs', [[2,2],[2,-2]])
plt.gca().set_color_cycle(None)
W_NR_corot.plot('Abs', [[2,2],[2,-2]], ls='dotted')
ylim((7e-2,0.15))
xlim((t0,t1))
title(r'Comparing the $(\ell,m) = (2,\pm 2)$ mode pair amplitudes')
legend(loc='upper left');

# Save the data to file, for direct incorporation into the main LaTeX file via `pgfplot`
numpy.savetxt('data/ModeAmplitudes_PN_2_2.dat', array([W_PN_corot.T(), W_PN_corot.Abs(W_PN_corot.FindModeIndex(2,2)),
                                                       W_PN_corot.Abs(W_PN_corot.FindModeIndex(2,-2))]).transpose()[::180],
              fmt='%.5e')


# Plot the phases
figure()
W_PN_corot.plot('Arg', [[2,2],[2,-2]])
plt.gca().set_color_cycle(None)
W_NR_corot.plot('Arg', [[2,2],[2,-2]], ls='dotted')
ylim((-0.06,0.06))
xlim((t0,t1))
title(r'Comparing the $(\ell,m) = (2,\pm 2)$ mode pair phases')
legend(loc='lower left');

# Save the data to file, for direct incorporation into the main LaTeX file via `pgfplot`
numpy.savetxt('data/ModePhases_PN_2_2.dat', array([W_PN_corot.T(), W_PN_corot.Arg(W_PN_corot.FindModeIndex(2,2)),
                                                   W_PN_corot.Arg(W_PN_corot.FindModeIndex(2,-2))]).transpose()[::180],
              fmt='%.5e')

%%time
PN_short_times = W_NR_corot_short.T()[W_NR_corot_short.T()<W_PN_corot.T(W_PN_corot.NTimes()-1)]
W_PN_corot_short = W_PN_corot.Interpolate(PN_short_times)
MinimalParityViolation_PN = W_PN_corot_short.MinimalParityViolation()
Antisymmetry_PN = W_PN_corot_short.NormalizedAntisymmetry()

# Plot the two data sets
missed_minimizations = [176,]
semilogy(np.delete(W_PN_corot_short.T(),missed_minimizations), np.delete(MinimalParityViolation_PN, missed_minimizations), label=r'$p$')
semilogy(W_PN_corot_short.T(), Antisymmetry_PN, label=r'$a$')
ylim((2.2e-3,1))
xlim((t0,t1))
xlabel(r'$(t-t_{\mathrm{merger}})/M$')
ylabel(r'Rotationally invariant antisymmetries')
legend(loc='upper left');

# Save the data to file, for direct incorporation into the main LaTeX file via `pgfplot`
numpy.savetxt('data/Asymmetry_PN.dat',
              np.delete(array([W_PN_corot_short.T(), Antisymmetry_PN, MinimalParityViolation_PN]),
                        missed_minimizations, axis=1).transpose(),
              fmt='%.5e')

# Just plot between these times, so it's not too messy
W_PN_vectors = GWFrames.PNWaveform(W_PN)
W_PN_vectors.InterpolateInPlace(W_NR.T()[i_1:i_2])

i_5 = argmin(abs(W_PN_corot.T()-t_1))
i_6 = argmin(abs(W_PN_corot.T()-t_2))

# Get the waveform vectors
AV_PNHat = array([o/np.sqrt(o.dot(o)) for o in W_PN_vectors.AngularVelocityVector()])
Vh_PNHat = array([o/np.sqrt(o.dot(o)) for o in W_PN_vectors.LLDominantEigenvector()])
Ldt_PNHat = array([-o/np.sqrt(o.dot(o)) for o in W_PN_vectors.LdtVector()])

# Now get the vectors for the orbital motions
Omega_orbHat_PN = array([o/np.sqrt(o.dot(o)) for o in W_PN_coorb.Omega_orb()[i_5:i_6]])
Omega_totHat_PN = array([o/np.sqrt(o.dot(o)) for o in W_PN_coorb.Omega_tot()[i_5:i_6]])

# Plot the vector directions
plot(Omega_orbHat_PN[:,0], Omega_orbHat_PN[:,1], label=r"$\hat{\Omega}_{\mathrm{orb}}=\hat{\ell}$")
plot(Omega_totHat_PN[:,0], Omega_totHat_PN[:,1], label=r"$\hat{\Omega}_{\mathrm{tot}}$")
plot(Vh_PNHat[:,0], Vh_PNHat[:,1], label=r"$\hat{V}_h$")
plot(Ldt_PNHat[:,0], Ldt_PNHat[:,1], label=r"$\langle L \partial_t \rangle$")
plot(AV_PNHat[:,0], AV_PNHat[:,1], label=r"$\hat{\omega}$")

# Plot the starting points as dots
plot([Omega_orbHat_PN[0,0],], [Omega_orbHat_PN[0,1],], 'o')
plot([Omega_totHat_PN[0,0],], [Omega_totHat_PN[0,1],], 'o')
plot([Vh_PNHat[0,0],], [Vh_PNHat[0,1],], 'o')
plot([Ldt_PNHat[0,0],], [Ldt_PNHat[0,1],], 'o')
plot([AV_PNHat[0,0],], [AV_PNHat[0,1],], 'o')

xlabel(r'$x$')
ylabel(r'$y$')
xlim((-0.1, 0.3))
ylim((-0.125,0.125))
legend();

# Save the data to file, for direct incorporation into the main LaTeX file via `pgfplot`
numpy.savetxt('data/WaveformVectors_PN.dat',
              array([AV_PNHat[:,0], AV_PNHat[:,1], Vh_PNHat[:,0], Vh_PNHat[:,1],
                     Ldt_PNHat[:,0], Ldt_PNHat[:,1],]).transpose()[::50],
              fmt='%.5e')
numpy.savetxt('data/WaveformVectors_Coordinates_PN.dat',
              array([Omega_orbHat_PN[:,0], Omega_orbHat_PN[:,1], Omega_totHat_PN[:,0], Omega_totHat_PN[:,1]]).transpose()[::50],
              fmt='%.5e')



ell = sympy.Symbol('ell', integer=True)
ell_Max = sympy.Symbol(r'\ell_{\mathrm{max}}', integer=True)

# At each time step, we would need these many elements of the Wigner D matrices:
WignerDEvaluations = sympy.summation((2*ell+1)**2, (ell,2,ell_Max)).simplify()
display(WignerDEvaluations)
for lmax in range(2,8+1):
    print('Including up to ell={0} would require {1} Wigner D evaluations.'.format(lmax,WignerDEvaluations.subs(ell_Max,lmax)))
print('\n')

# We would also need these many evaluations of the SWSHs, but just once:
SWSHEvaluations = sympy.summation((2*ell+1), (ell,2,ell_Max)).simplify()
display(SWSHEvaluations)
for lmax in range(2,8+1):
    print('Including up to ell={0} would require {1} SWSH evaluations.'.format(lmax,SWSHEvaluations.subs(ell_Max,lmax)))

# Import the basics for symbolic calculations with quaternions
from sympy.galgebra.ga import *
vartheta, varphi, gamma = sympy.symbols('varthetabreve, varphibreve, gammabreve')
X = (x,y,z) = sympy.symbols('x y z')
(ex,ey,ez,grad) = MV.setup('e_x e_y e_z',metric='[1,1,1]',coords=X)
xHat = ey^ez
yHat = -ex^ez
zHat = ex^ey
def detan(e):
    "This is just a utility function to simplify atan(tan(x)) to x."
    from sympy import Wild, atan, tan
    a = Wild('a')
    return e.replace( atan(tan(a)), a )
def decos(e):
    "This is just a utility function to simplify acos(cos(x)) to x."
    from sympy import Wild, acos, cos
    a = Wild('a')
    return e.replace( acos(cos(a)), a )

R_tot = zHat.exp(varphi/2)*yHat.exp(vartheta/2)*zHat.exp(gamma/2)

r_0 = R_tot.scalar().simplify()
r_1 = R_tot.coef(xHat).simplify()
r_2 = (R_tot.coef(ex^ez)).simplify()
r_3 = R_tot.coef(zHat).simplify()

display(r_0, r_1, r_2, r_3)

r_a = detan(sympy.atan((-r_1/r_2).simplify()))
r_b = detan(sympy.atan((r_3/r_0).simplify()))

decos(2*sympy.acos(sympy.sqrt(r_0**2 + r_3**2).simplify())), sympy.simplify(r_b+r_a), sympy.simplify(r_b-r_a)

angles = np.array([[2*np.pi*a[0]-np.pi, np.pi*a[1], 2*np.pi*a[2]-np.pi] for a in rand(1000000,3)])
rotors = np.array([Quaternions.Quaternion(*a) for a in angles])
deduce = np.array([[np.arctan2(r[3],r[0])+np.arctan2(-r[1],r[2]), # varphi
                    2*np.arccos(np.sqrt(r[0]**2 + r[3]**2)), # vartheta
                    np.arctan2(r[3],r[0])-np.arctan2(-r[1],r[2])] # gamma
                   for r in rotors])
errors = np.array([np.linalg.norm(a-d) for a,d in zip(angles,deduce)])
max(errors)