#Packages 
using LinearAlgebra
using Plots
using FFTW
using Printf
using BenchmarkTools

#%Complex definition
i=sqrt(Complex(-1));

#Definition of individual spin operators
nspins=1
include("basis.jl")
iu,ix,iy,iz,ip,im,ia,ib=basis(nspins)

#General Setup
timestep = 0.0000001
w = zeros(Float64,3)
w = [45.0 * pi/180.0, 2500.0, 1000.0]

eps = 0.00000001
th = w * timestep
ham =  2 * pi * th[1] * ix[:,:,1] +
       2 * pi * th[2] * iy[:,:,1] +
       2 * pi * th[3] * iz[:,:,1] 
uham = exp(- i * ham)

rho = iy[:,:,1]
rho = uham * rho * uham'

function derivative_alpha_xact(w::Vector{Float64},timestep::Float64)
    # mit w wird alpha, theta_xy und theta_z übergeben. w[1] ist alpha.
    dQdalpha    = Vector{Float64}(undef,4)
    theta_xy    = Float64
    z           = Float64
    r           = Float64
    si          = Float64
    co          = Float64
    sinalpha    = Float64
    cosalpha    = Float64
    local theta_xy, z = w[2:3] .* (2 * pi * timestep ) 
    local r = sqrt(theta_xy^2 + z^2) 
    local si,co = sincos(r/2)
    local sinalpha, cosalpha = sincos(w[1])

    local dQdalpha = [theta_xy*sinalpha*si/r, theta_xy*cosalpha*si/r, 0.0, 0.0]

    return dQdalpha
end

function derivative_thetaxy_xact(w::Vector{Float64},timestep::Float64)
    # mit w wird alpha, theta_xy und theta_z übergeben. w[1] ist alpha.
    dQdthetaxy    = Vector{Float64}(undef,4)
    theta_xy    = Float64
    z           = Float64
    r           = Float64
    si          = Float64
    co          = Float64
    sinalpha    = Float64
    cosalpha    = Float64
    local theta_xy, z = w[2:3] .* (2 * pi * timestep ) 
    local r = sqrt(theta_xy^2 + z^2) 
    local si,co = sincos(r/2)
    local sinalpha, cosalpha = sincos(w[1])

    local dQdthetaxy = [cosalpha * ((r^2 - theta_xy^2) / (r^3) * sin(r/2) + theta_xy^2 / (2*r^2) * cos(r/2)),
                        sinalpha * ((r^2 - theta_xy^2) / (r^3) * sin(r/2) + theta_xy^2 / (2*r^2) * cos(r/2)),
                        theta_xy * z * (co / (2*r^2) - si / r^3),
                        - theta_xy * sin(r/2) / (2*r)]

    return dQdthetaxy
end

function derivative_z_xact(w::Vector{Float64},timestep::Float64)
    dQdz = Vector{Float64}(undef,4)
    theta_xy    = Float64
    z           = Float64
    r           = Float64
    si          = Float64
    co          = Float64
    sinalpha    = Float64
    cosalpha    = Float64
    local theta_xy, z = w[2:3] .* (2 * pi * timestep ) 
    local r = sqrt(theta_xy^2 + z^2) 
    local si,co = sincos(r/2)
    local sinalpha, cosalpha = sincos(w[1])

    local dQdz = [theta_xy * z * cosalpha * (co / (2*r^2) - si / r^3),
                    theta_xy * z * sinalpha * (co / (2*r^2) - si / r^3),
                    (r^2 - z^2) / (r^3) * si + z^2 / (2*r^2) * co,
                    - z * sin(r/2) / (2*r)]
    return dQdz
end

function derivative_alpha_Qdiff(w::Vector{Float64},timestep::Float64,eps::Float64)
    Q           = Vector{Float64}(undef,4)
    Qe          = Vector{Float64}(undef,4)
    dQdalpha    = Vector{Float64}(undef,4)
    theta_xy    = Float64
    z           = Float64
    r           = Float64
    alpha       = Float64
    si          = Float64
    co          = Float64
    sinalpha    = Float64
    cosalpha    = Float64
    local theta_xy, z = w[2:3] .* (2 * pi * timestep ) 
    local r = sqrt(theta_xy^2 + z^2) 
    local alpha = w[1]
    local si,co = sincos(r/2)
    local sinalpha, cosalpha = sincos(alpha)
    local Q = [theta_xy * cosalpha * si /r, theta_xy * sinalpha * si /r, z * si /r, co]
    alpha = alpha + eps

    local sinalpha, cosalpha = sincos(alpha)
    local Qe = [theta_xy * cosalpha * si /r, theta_xy * sinalpha * si /r, z * si /r, co]

    dQdalpha = (Qe-Q)/eps
    return dQdalpha
end

function derivative_thetaxy_Qdiff(w::Vector{Float64},timestep::Float64,eps::Float64)
    Q           = Vector{Float64}(undef,4)
    Qe          = Vector{Float64}(undef,4)
    dQdthetaxy    = Vector{Float64}(undef,4)
    theta_xy    = Float64
    z           = Float64
    r           = Float64
    si          = Float64
    co          = Float64
    sinalpha    = Float64
    cosalpha    = Float64
    local theta_xy, z = w[2:3] .* (2 * pi * timestep ) 
    local r = sqrt(theta_xy^2 + z^2) 
    local si,co = sincos(r/2)
    local sinalpha, cosalpha = sincos(w[1])
    local Q = [theta_xy * cosalpha * si /r, theta_xy * sinalpha * si /r, z * si /r, co]
    theta_xy = theta_xy + eps
    local r = sqrt(theta_xy^2 + z^2) 
    local si,co = sincos(r/2)
    local Qe = [theta_xy * cosalpha * si /r, theta_xy * sinalpha * si /r, z * si /r, co]

    dQdthetaxy = (Qe-Q)/eps
    return dQdthetaxy
end

function derivative_z_Qdiff(w::Vector{Float64},timestep::Float64,eps::Float64)
    Q           = Vector{Float64}(undef,4)
    Qe          = Vector{Float64}(undef,4)
    dQdz    = Vector{Float64}(undef,4)
    theta_xy    = Float64
    z           = Float64
    r           = Float64
    si          = Float64
    co          = Float64
    sinalpha    = Float64
    cosalpha    = Float64
    local theta_xy, z = w[2:3] .* (2 * pi * timestep ) 
    local r = sqrt(theta_xy^2 + z^2) 
    local si,co = sincos(r/2)
    local sinalpha, cosalpha = sincos(w[1])
    local Q = [theta_xy * cosalpha * si /r, theta_xy * sinalpha * si /r, z * si /r, co]

    z = z + eps
    local r = sqrt(theta_xy^2 + z^2) 
    local si,co = sincos(r/2)
    local Qe = [theta_xy * cosalpha * si /r, theta_xy * sinalpha * si /r, z * si /r, co]

    dQdz = (Qe-Q)/eps
    return dQdz
end

function printcmat(mat)
    nrow,ncol = size(mat)
    for irow = 1:nrow
        @printf("  ")
        for icol = 1:ncol
            @printf("(%f %f*i)  ",real(mat[irow,icol]), imag(mat[irow,icol]))
        end
        @printf("\n")
    end
end
    

testxact1 = derivative_alpha_xact(w,timestep)
testQdiff1 = derivative_alpha_Qdiff(w,timestep,eps)
testxact2 = derivative_thetaxy_xact(w,timestep)
testQdiff2 = derivative_thetaxy_Qdiff(w,timestep,eps)
testxact3 = derivative_z_xact(w,timestep)
testQdiff3 = derivative_z_Qdiff(w,timestep,eps)

println("Analytische Quaternionen alpha :")
println(testxact1)
println("testdiff (Ableitung aus finiten Differenzen) alpha:")
println(testQdiff1)
println("")
println("Analytische Quaternionen txy :")
println(testxact2)
println("testdiff (Ableitung aus finiten Differenzen) txy :")
println(testQdiff2)
println("")
println("Analytische Quaternionen z :")
println(testxact3)
println("testdiff (Ableitung aus finiten Differenzen) z :")
println(testQdiff3)
println("")


ntest = 1000
wtest = rand(Float64,ntest,3)
@btime begin
    for itest = 1:ntest
        testxact1a = derivative_alpha_xact(wtest[itest,:],timestep)
    end
end

@btime begin
    for itest = 1:ntest
        testxact1a = derivative_thetaxy_xact(wtest[itest,:],timestep)
    end
end

@btime begin
    for itest = 1:ntest
        testxact1a = derivative_z_xact(wtest[itest,:],timestep)
    end
end
    
@btime begin
    for itest = 1:ntest
        testdiff1a = derivative_alpha_Qdiff(wtest[itest,:],timestep,eps)
    end
end

@btime begin
    for itest = 1:ntest
        testdiff1a = derivative_thetaxy_Qdiff(wtest[itest,:],timestep,eps)
    end
end

@btime begin
    for itest = 1:ntest
        testdiff1a = derivative_z_Qdiff(wtest[itest,:],timestep,eps)
    end
end