#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 = [3000.0, 7500.0, 1000.0]
eps = 0.000000001
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_x_expm(w::Vector{Float64},timestep::Float64)
    local ham = Array{ComplexF64,2}(undef,2,2)
    local ctrl= Array{ComplexF64,2}(undef,2,2)
    local res = Array{ComplexF64,2}(undef,4,4)
    local i   = sqrt(ComplexF64(-1));

    local ham =  2 * pi * w[1] * ix[:,:,1] +
                 2 * pi * w[2] * iy[:,:,1] +
                 2 * pi * w[3] * iz[:,:,1]
    local ctrl = - 2 * pi * ix[:,:,1] ./(-timestep*2*pi)
    local mnull = zeros(2,2)
    res = exp(-i * timestep * [ham  ctrl; mnull  ham])
    return res
end

function derivative_y_expm(w::Vector{Float64},timestep::Float64)
    local ham = Array{ComplexF64,2}(undef,2,2)
    local ctrl= Array{ComplexF64,2}(undef,2,2)
    local res = Array{ComplexF64,2}(undef,4,4)
    local i   = sqrt(ComplexF64(-1));

    local ham =  2 * pi * w[1] * ix[:,:,1] +
                 2 * pi * w[2] * iy[:,:,1] +
                 2 * pi * w[3] * iz[:,:,1]
    local ctrl = - 2 * pi * iy[:,:,1] ./(-timestep*2*pi)
    local mnull = zeros(2,2)
    res = exp(-i * timestep * [ham  ctrl; mnull  ham])
    return res
end

function derivative_z_expm(w::Vector{Float64},timestep::Float64)
    local ham = Array{ComplexF64,2}(undef,2,2)
    local ctrl= Array{ComplexF64,2}(undef,2,2)
    local res = Array{ComplexF64,2}(undef,4,4)
    local i   = sqrt(ComplexF64(-1));

    local ham =  2 * pi * w[1] * ix[:,:,1] +
                 2 * pi * w[2] * iy[:,:,1] +
                 2 * pi * w[3] * iz[:,:,1]
    local ctrl = - 2 * pi * iz[:,:,1] ./(-timestep*2*pi)
    local mnull = zeros(2,2)
    res = exp(-i * timestep * [ham  ctrl; mnull  ham])
    return res
end

function derivative_x_diff(w::Vector{Float64},timestep::Float64,eps::Float64)
    local ham = Array{ComplexF64,2}(undef,2,2)
    local hame= Array{ComplexF64,2}(undef,2,2)
    local res = Array{ComplexF64,2}(undef,2,2)
    local i   = sqrt(ComplexF64(-1));


    local ham =  2 * pi * w[1] * ix[:,:,1] * timestep +
                 2 * pi * w[2] * iy[:,:,1] * timestep +
                 2 * pi * w[3] * iz[:,:,1] * timestep
    local hame=  ((2 * pi * w[1] * timestep)+eps) * ix[:,:,1] +
                 2 * pi * w[2] * iy[:,:,1] * timestep +
                 2 * pi * w[3] * iz[:,:,1] * timestep
    res = (exp(-i*hame) - exp(-i*ham)) ./ eps
    return res
end

function derivative_y_diff(w::Vector{Float64},timestep::Float64,eps::Float64)
    local ham = Array{ComplexF64,2}(undef,2,2)
    local hame= Array{ComplexF64,2}(undef,2,2)
    local res = Array{ComplexF64,2}(undef,2,2)
    local i   = sqrt(ComplexF64(-1));


    local ham =  2 * pi * w[1] * ix[:,:,1] * timestep +
                 2 * pi * w[2] * iy[:,:,1] * timestep +
                 2 * pi * w[3] * iz[:,:,1] * timestep
    local hame=  2 * pi * w[1] * ix[:,:,1] * timestep +
                ((2 * pi * w[2] * timestep)+eps) * iy[:,:,1] +
                 2 * pi * w[3] * iz[:,:,1] * timestep
    res = (exp(-i*hame) - exp(-i*ham)) ./ eps
    return res
end

function derivative_z_diff(w::Vector{Float64},timestep::Float64,eps::Float64)
    local ham = Array{ComplexF64,2}(undef,2,2)
    local hame= Array{ComplexF64,2}(undef,2,2)
    local res = Array{ComplexF64,2}(undef,2,2)
    local i   = sqrt(ComplexF64(-1));


    local ham =  2 * pi * w[1] * ix[:,:,1] * timestep +
                 2 * pi * w[2] * iy[:,:,1] * timestep +
                 2 * pi * w[3] * iz[:,:,1] * timestep
    local hame=  2 * pi * w[1] * ix[:,:,1] * timestep +
                 2 * pi * w[2] * iy[:,:,1] * timestep +
                 ((2 * pi * w[3] * timestep)+eps) * iz[:,:,1]
    res = (exp(-i*hame) - exp(-i*ham)) ./ eps
    return res
end

function derivative_x_xact(w::Vector{Float64},timestep::Float64)
    dQdx = Vector{Float64}(undef,4)
    x    = Float64
    y    = Float64
    z    = Float64
    r    = Float64
    si   = Float64
    co   = Float64
    local x,y,z = w[1:3] .* (2 * pi * timestep ) 
    local r = sqrt(x^2 + y^2 + z^2) 
    local si,co = sincos(r/2)
    local dQdx = [((r^2 - x^2)*si) / (r^3) + (x^2)*co / (2*r*r),
                  x*y*(co/(2*r*r) - si/(r^3)), 
                  x*z*(co/(2*r*r) - si/(r^3)), 
                  - x*si / (2*r)]
    return dQdx
end

function derivative_y_xact(w::Vector{Float64},timestep::Float64)
    dQdy = Vector{Float64}(undef,4)
    x    = Float64
    y    = Float64
    z    = Float64
    r    = Float64
    si   = Float64
    co   = Float64
    local x,y,z = w[1:3] .* (2 * pi * timestep ) 
    local r = sqrt(x^2 + y^2 + z^2) 
    local si,co = sincos(r/2)
    local dQdy = [x * y * (co / (2*r^2) - si / r^3),
                    ((r^2 - y^2) / (r^3)) * si + (y^2 / (2*r^2)) * co,
                    y * z * (co / (2*r^2) - si / r^3),
                    -y * si / (2*r)]
                  
    return dQdy
end

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

function derivative_x_Qdiff(w,timestep,eps)
    local x,y,z = w .* (2 * pi * timestep ) 
    local r = sqrt(x^2 + y^2 + z^2) 
    local si,co = sincos(r/2)
    A = x * si /r
    B = y * si /r
    C = z * si /r
    D = co
    x = x + eps
    local r = sqrt(x^2 + y^2 + z^2) 
    local si,co = sincos(r/2)
    Ae = x * si /r
    Be = y * si /r
    Ce = z * si /r
    De = co
    dAdx = (Ae-A)/eps
    dBdx = (Be-B)/eps
    dCdx = (Ce-C)/eps
    dDdx = (De-D)/eps
    return [dAdx,dBdx,dCdx,dDdx]
end

function derivative_y_Qdiff(w,timestep,eps)
    local x,y,z = w .* (2 * pi * timestep ) 
    local r = sqrt(x^2 + y^2 + z^2) 
    local si,co = sincos(r/2)
    A = x * si /r
    B = y * si /r
    C = z * si /r
    D = co
    y = y + eps
    local r = sqrt(x^2 + y^2 + z^2) 
    local si,co = sincos(r/2)
    Ae = x * si /r
    Be = y * si /r
    Ce = z * si /r
    De = co
    dAdy = (Ae-A)/eps
    dBdy = (Be-B)/eps
    dCdy = (Ce-C)/eps
    dDdy = (De-D)/eps
    return [dAdy,dBdy,dCdy,dDdy]
end

function derivative_z_Qdiff(w,timestep,eps)
    local x,y,z = w .* (2 * pi * timestep ) 
    local r = sqrt(x^2 + y^2 + z^2) 
    local si,co = sincos(r/2)
    A = x * si /r
    B = y * si /r
    C = z * si /r
    D = co
    z = z + eps
    local r = sqrt(x^2 + y^2 + z^2) 
    local si,co = sincos(r/2)
    Ae = x * si /r
    Be = y * si /r
    Ce = z * si /r
    De = co
    dAdz = (Ae-A)/eps
    dBdz = (Be-B)/eps
    dCdz = (Ce-C)/eps
    dDdz = (De-D)/eps
    return [dAdz, dBdz, dCdz, dDdz]
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_x_xact(w,timestep)
testxact2 = derivative_y_xact(w,timestep)
testxact3 = derivative_z_xact(w,timestep)

testQdiff1 = derivative_x_Qdiff(w,timestep,eps)
testQdiff2 = derivative_y_Qdiff(w,timestep,eps)
testQdiff3 = derivative_z_Qdiff(w,timestep,eps)

println("Analytische Quaternionen dABCDdx:")
println(testxact1)
println("testdiff (Ableitung aus finiten Differenzen) x:")
println(testQdiff1)

println("Analytische Quaternionen dABCDDdy:")
println(testxact2)
println("testdiff (Ableitung aus finiten Differenzen) y:")
println(testQdiff2)

println("Analytische Quaternionen dABCDdz:")
println(testxact3)
println("testdiff (Ableitung aus finiten Differenzen) z:")
println(testQdiff3)


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

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

@btime begin
    for itest = 1:ntest
        testxact1 = derivative_z_xact(wtest[itest,:],timestep)
    end
end
    
    
@btime begin
    for itest = 1:ntest
        testexpm1 = derivative_x_expm(wtest[itest,:],timestep)
    end
end

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

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

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

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

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