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

#%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.00001
w = StaticVector{3,Float64}
w = [2500.0, 170.0, 30.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_expr(w::Vector{Float64},timestep::Float64)
    local ham = Array{Float64,2}(undef,3,3)
    local ctrl= Array{Float64,2}(undef,3,3)
    local res = Array{Float64,2}(undef,6,6)
    local mnull= zeros(Float64,3,3)

    local th = (2 * pi * timestep) .* w
    local ham = [0      -th[3]  th[2];
                 th[3]  0      -th[1];
                -th[2]  th[1]  0      ]
    local ctrl= [0     0     0 ;
                 0     0    -1 ;
                 0     1     0 ]
    res = exp([ham  ctrl; mnull  ham])
    return res
end

function derivative_y_expr(w::Vector{Float64},timestep::Float64)
    local ham = Array{Float64,2}(undef,3,3)
    local ctrl= Array{Float64,2}(undef,3,3)
    local res = Array{Float64,2}(undef,6,6)
    local mnull= zeros(Float64,3,3)

    local th = (2 * pi * timestep) .* w
    local ham = [0      -th[3]  th[2];
                 th[3]  0      -th[1];
                -th[2]  th[1]  0      ]
    local ctrl= [0 0 1; 
                 0 0 0; 
                 -1 0 0]
    res = exp([ham  ctrl; mnull  ham])
    return res
end

function derivative_z_expr(w::Vector{Float64},timestep::Float64)
    local ham = Array{Float64,2}(undef,3,3)
    local ctrl= Array{Float64,2}(undef,3,3)
    local res = Array{Float64,2}(undef,6,6)
    local mnull= zeros(Float64,3,3)

    local th = (2 * pi * timestep) .* w
    local ham = [0      -th[3]  th[2];
                 th[3]  0      -th[1];
                -th[2]  th[1]  0      ]
    local ctrl= [0 -1 0; 
                 1 0 0; 
                 0 0 0]
    res = exp([ham  ctrl; mnull  ham])
    return res
end

function derivative_x_diffr(w::Vector{Float64},timestep::Float64,eps::Float64)
    local rot = Array{Float64,2}(undef,3,3)
    local rote= Array{Float64,2}(undef,3,3)
    local res = Array{Float64,2}(undef,3,3)

    local th = (2 * pi * timestep) .* w
    local rot  = RotationVec(th[1],th[2],th[3])
    th[1] = th[1] + eps    
    local rote = RotationVec(th[1],th[2],th[3])

    res = (rote - rot) ./ eps
    return res
end

function derivative_y_diffr(w::Vector{Float64},timestep::Float64,eps::Float64)
    local rot = Array{Float64,2}(undef,3,3)
    local rote= Array{Float64,2}(undef,3,3)
    local res = Array{Float64,2}(undef,3,3)

    local th = (2 * pi * timestep) .* w
    local rot  = RotationVec(th[1],th[2],th[3])
    th[2] = th[2] + eps    
    local rote = RotationVec(th[1],th[2],th[3])

    res = (rote - rot) ./ eps
    return res
end

function derivative_z_diffr(w::Vector{Float64},timestep::Float64,eps::Float64)
    local rot = Array{Float64,2}(undef,3,3)
    local rote= Array{Float64,2}(undef,3,3)
    local res = Array{Float64,2}(undef,3,3)

    local th = (2 * pi * timestep) .* w
    local rot  = RotationVec(th[1],th[2],th[3])
    th[3] = th[3] + eps    
    local rote = RotationVec(th[1],th[2],th[3])

    res = (rote - rot) ./ eps
    return res
end


function derivative_x_xact2(w::Vector{Float64},timestep::Float64)
    dRdx = StaticScalar{Float64}
    x    = StaticScalar{Float64}
    y    = StaticScalar{Float64}
    z    = StaticScalar{Float64}
    r    = StaticScalar{Float64}
    si   = StaticScalar{Float64}
    co   = StaticScalar{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)
    local dRdx = zeros(3, 3)

    dRdx[1, 1] = (x^3 / r^3 - x / r) * si - (x / r^2 - x^3 / r^4) * 2 * (co - 1)
    dRdx[2, 2] = (x * y^2 / r^3 - x / r) * si + (x * y^2 / r^4) * 2 * (co - 1)
    dRdx[3, 3] = (x * z^2 / r^3 - x / r) * si + (x * z^2 / r^4) * 2 * (co - 1) 
    dRdx[1, 2] = (x^2 * y + x * z) / r^3 * si + (2 * x^2 * y / r^4 - y / r^2) * (co - 1) - (x * z / r^2) * co
    dRdx[2, 1] = (x^2 * y - x * z) / r^3 * si + (2 * x^2 * y / r^4 - y / r^2) * (co - 1) + (x * z / r^2) * co  
    dRdx[1, 3] = (x^2 * y - x * y) / r^3 * si + (2 * x^2 * z / r^4 - y / r^2) * (co - 1) + (x * y / r^2) * co
    dRdx[3, 1] = (x^2 * y + x * y) / r^3 * si + (2 * x^2 * z / r^4 - y / r^2) * (co - 1) - (x * y / r^2) * co   
    dRdx[2, 3] = (x * y * z - z^2 - y^2) / r^3 * si + 2 * x * y * z / r^4 * (co - 1) - x^2 / r^2 * co
    dRdx[3, 2] = (x * y * z + z^2 + y^2) / r^3 * si + 2 * x * y * z / r^4 * (co - 1) + x^2 / r^2 * co

    return dRdx
end

function derivative_x_xact(w::Vector{Float64},timestep::Float64)
    dRdx = Array{Float64,2}(undef,3,3)
    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)
    local dRdx = zeros(3, 3)

    dRdx[1, 1] = (x^3 / r^3 - x / r) * si - (x / r^2 - x^3 / r^4) * 2 * (co - 1)
    dRdx[2, 2] = (x * y^2 / r^3 - x / r) * si + (x * y^2 / r^4) * 2 * (co - 1)
    dRdx[3, 3] = (x * z^2 / r^3 - x / r) * si + (x * z^2 / r^4) * 2 * (co - 1) 
    dRdx[1, 2] = (x^2 * y + x * z) / r^3 * si + (2 * x^2 * y / r^4 - y / r^2) * (co - 1) - (x * z / r^2) * co
    dRdx[2, 1] = (x^2 * y - x * z) / r^3 * si + (2 * x^2 * y / r^4 - y / r^2) * (co - 1) + (x * z / r^2) * co  
    dRdx[1, 3] = (x^2 * z - x * y) / r^3 * si + (2 * x^2 * z / r^4 - z / r^2) * (co - 1) + (x * y / r^2) * co
    dRdx[3, 1] = (x^2 * z + x * y) / r^3 * si + (2 * x^2 * z / r^4 - z / r^2) * (co - 1) - (x * y / r^2) * co   
    dRdx[2, 3] = (x * y * z - z^2 - y^2) / r^3 * si + 2 * x * y * z / r^4 * (co - 1) - x^2 / r^2 * co
    dRdx[3, 2] = (x * y * z + z^2 + y^2) / r^3 * si + 2 * x * y * z / r^4 * (co - 1) + x^2 / r^2 * co

    return dRdx
end

function derivative_y_xact(w::Vector{Float64},timestep::Float64)
    dRdx = Array{Float64,2}(undef,3,3)
    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)
    local dRdy = zeros(3, 3)


    dRdy[1, 1] = (x^2 * y / r^3 - y / r) * si + (x^2 * y / r^4) * 2 * (co - 1)
    dRdy[2, 2] = (y^3 / r^3 - y / r) * si - (y / r^2 - y^3 / r^4) * 2 * (co - 1)
    dRdy[3, 3] = (y * z^2 / r^3 - y / r) * si + (y * z^2 / r^4) * 2 * (co - 1)  
    dRdy[1, 2] = (y^2 * x + y * z) / r^3 * si + (2 * y^2 * x / r^4 - x / r^2) * (co - 1) - (y * z / r^2) * co
    dRdy[2, 1] = (y^2 * x - y * z) / r^3 * si + (2 * y^2 * x / r^4 - x / r^2) * (co - 1) + (y * z / r^2) * co  
    dRdy[3, 2] = (y^2 * z - x * y) / r^3 * si + (2 * y^2 * z / r^4 - z / r^2) * (co - 1) + (x * y / r^2) * co
    dRdy[2, 3] = (y^2 * z + x * y) / r^3 * si + (2 * y^2 * z / r^4 - z / r^2) * (co - 1) - (x * y / r^2) * co  
    dRdy[3, 1] = (x * y * z - x^2 - z^2) / r^3 * si + 2 * x * y * z / r^4 * (co - 1) - y^2 / r^2 * co
    dRdy[1, 3] = (x * y * z + x^2 + z^2) / r^3 * si + 2 * x * y * z / r^4 * (co - 1) + y^2 / r^2 * co

    return dRdy
end

function derivative_z_xact(w::Vector{Float64},timestep::Float64)
    dRdx = Array{Float64,2}(undef,3,3)
    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)
    local dRdz = zeros(3, 3)

    dRdz[1, 1] = (x^2 * z / r^3 - (z / r)) * si + (x^2 * z / r^4) * 2 * (co - 1)
    dRdz[1, 2] = (x * y * z - x^2 - y^2) / r^3 * si + 2 * x * y * z / r^4 * (co - 1) - z^2 / r^2 * co
    dRdz[1, 3] = (z^2 * x - z * y) / r^3 * si + (2 * z^2 * x / r^4 - x / r^2) * (co - 1) + z * y / r^2 * co
    dRdz[2, 1] = (x * y * z + x^2 + y^2) / r^3 * si + 2 * x * y * z / r^4 * (co - 1) + z^2 / r^2 * co
    dRdz[2, 2] = (z * y^2 / r^3 - z / r) * si + (z * y^2 / r^4) * 2 * (co - 1)
    dRdz[2, 3] = (z^2 * y + x * z)/ r^3 * si + (2 * z^2 * y / r^4 - y / r^2) * (co - 1) - x * z / r^2 * co
    dRdz[3, 1] = (z^2 * x + z * y) / r^3 * si + (2 * z^2 * x / r^4 - x / r^2) * (co - 1) - z * y / r^2 * co
    dRdz[3, 2] = (z^2 * y - x * z) / r^3 * si + (2 * z^2 * y / r^4 - y / r^2) * (co - 1) + x * z / r^2 * co
    dRdz[3, 3] = (z^3 / r^3 - z / r) * si + 2 * (z/r^2- z^3 / r^4) * (1 - co)
    
    return dRdz
end

function derivative_x_diff(w::Vector{Float64},timestep::Float64,eps::Float64)
    R = Array{Float64,2}(undef,3,3)
    Re = Array{Float64,2}(undef,3,3)
    dRdx = Array{Float64,2}(undef,3,3)
    local x,y,z = w .* (2 * pi * timestep ) 
    local r = sqrt(x^2 + y^2 + z^2) 
    local si,co = sincos(r)
    local nx = x/r
    local ny = y/r
    local nz = z/r
    local R = zeros(3,3)
    R[1,1] = co + nx^2 * (1 - co)    
    R[1,2] = - nz * si + nx * ny * (1 - co)  
    R[1,3] = ny * si + nx * nz * (1 - co)
    R[2,1] = nz * si + nx * ny * (1 - co)    
    R[2,2] = co + ny^2 * (1 - co)    
    R[2,3] = - nx * si + ny * nz * (1 - co)
    R[3,1] = -ny * si + nx * nz * (1 - co)   
    R[3,2] = nx * si + ny * nz * (1 - co)    
    R[3,3] = co + nz^2 * (1 - co)

    x = x + eps
    local r = sqrt(x^2 + y^2 + z^2) 
    local si,co = sincos(r)
    local nx = x/r
    local ny = y/r
    local nz = z/r

    local Re = zeros(3,3)
    Re[1,1] = co + nx^2 * (1 - co)    
    Re[1,2] = - nz * si + nx * ny * (1 - co)  
    Re[1,3] = ny * si + nx * nz * (1 - co)
    Re[2,1] = nz * si + nx * ny * (1 - co)    
    Re[2,2] = co + ny^2 * (1 - co)    
    Re[2,3] = - nx * si + ny * nz * (1 - co)
    Re[3,1] = -ny * si + nx * nz * (1 - co)   
    Re[3,2] = nx * si + ny * nz * (1 - co)    
    Re[3,3] = co + nz^2 * (1 - co)

    local dRdx = (Re-R)/eps
    return dRdx
end

function derivative_y_diff(w::Vector{Float64},timestep::Float64,eps::Float64)
    R = Array{Float64,2}(undef,3,3)
    Re = Array{Float64,2}(undef,3,3)
    dRdy = Array{Float64,2}(undef,3,3)
    local x,y,z = w .* (2 * pi * timestep ) 
    local r = sqrt(x^2 + y^2 + z^2) 
    local si,co = sincos(r)
    local nx = x/r
    local ny = y/r
    local nz = z/r
    local R = zeros(3,3)
    R[1,1] = co + nx^2 * (1 - co)    
    R[1,2] = - nz * si + nx * ny * (1 - co)  
    R[1,3] = ny * si + nx * nz * (1 - co)
    R[2,1] = nz * si + nx * ny * (1 - co)    
    R[2,2] = co + ny^2 * (1 - co)    
    R[2,3] = - nx * si + ny * nz * (1 - co)
    R[3,1] = -ny * si + nx * nz * (1 - co)   
    R[3,2] = nx * si + ny * nz * (1 - co)    
    R[3,3] = co + nz^2 * (1 - co)

    y = y + eps
    local r = sqrt(x^2 + y^2 + z^2) 
    local si,co = sincos(r)
    local nx = x/r
    local ny = y/r
    local nz = z/r

    local Re = zeros(3,3)
    Re[1,1] = co + nx^2 * (1 - co)    
    Re[1,2] = - nz * si + nx * ny * (1 - co)  
    Re[1,3] = ny * si + nx * nz * (1 - co)
    Re[2,1] = nz * si + nx * ny * (1 - co)    
    Re[2,2] = co + ny^2 * (1 - co)    
    Re[2,3] = - nx * si + ny * nz * (1 - co)
    Re[3,1] = -ny * si + nx * nz * (1 - co)   
    Re[3,2] = nx * si + ny * nz * (1 - co)    
    Re[3,3] = co + nz^2 * (1 - co)

    local dRdy = (Re-R)/eps
    return dRdy
end

function derivative_z_diff(w::Vector{Float64},timestep::Float64,eps::Float64)
    R = Array{Float64,2}(undef,3,3)
    Re = Array{Float64,2}(undef,3,3)
    dRdz = Array{Float64,2}(undef,3,3)
    local x,y,z = w .* (2 * pi * timestep ) 
    local r = sqrt(x^2 + y^2 + z^2) 
    local si,co = sincos(r)
    local nx = x/r
    local ny = y/r
    local nz = z/r
    local R = zeros(3,3)
    R[1,1] = co + nx^2 * (1 - co)    
    R[1,2] = - nz * si + nx * ny * (1 - co)  
    R[1,3] = ny * si + nx * nz * (1 - co)
    R[2,1] = nz * si + nx * ny * (1 - co)    
    R[2,2] = co + ny^2 * (1 - co)    
    R[2,3] = - nx * si + ny * nz * (1 - co)
    R[3,1] = - ny * si + nx * nz * (1 - co)   
    R[3,2] = nx * si + ny * nz * (1 - co)    
    R[3,3] = co + nz^2 * (1 - co)

    z = z + eps
    local r = sqrt(x^2 + y^2 + z^2) 
    local si,co = sincos(r)
    local nx = x/r
    local ny = y/r
    local nz = z/r

    local Re = zeros(3,3)
    Re[1,1] = co + nx^2 * (1 - co)    
    Re[1,2] = - nz * si + nx * ny * (1 - co)  
    Re[1,3] = ny * si + nx * nz * (1 - co)
    Re[2,1] = nz * si + nx * ny * (1 - co)    
    Re[2,2] = co + ny^2 * (1 - co)    
    Re[2,3] = - nx * si + ny * nz * (1 - co)
    Re[3,1] = -ny * si + nx * nz * (1 - co)   
    Re[3,2] = nx * si + ny * nz * (1 - co)    
    Re[3,3] = co + nz^2 * (1 - co)

    local dRdz = (Re-R)/eps
    return dRdz
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
 
testexpr1 = derivative_x_expr(w,timestep)
testdiffr1 = derivative_x_diffr(w,timestep,eps)
testxact1 = derivative_x_xact(w,timestep)
testdiff1 = derivative_x_diff(w,timestep,eps)

testexpr2 = derivative_y_expr(w,timestep)
testdiffr2 = derivative_y_diffr(w,timestep,eps)
testxact2 = derivative_y_xact(w,timestep)
testdiff2 = derivative_y_diff(w,timestep,eps)

testexpr3 = derivative_z_expr(w,timestep)
testdiffr3 = derivative_z_diffr(w,timestep,eps)
testxact3 = derivative_z_xact(w,timestep)
testdiff3 = derivative_z_diff(w,timestep,eps)

println("derivative_x_diff")
printcmat(testdiff1)
println("derivative_x_xact")
printcmat(testxact1)
#println("")
# printcmat(testdiffr1[1:3,1:3])
# println("")
# printcmat(testexpr1[1:3,4:6])
println("")
println("derivative_y_diff")
printcmat(testdiff2)
println("derivative_y_xact")
printcmat(testxact2)
println("")
println("derivative_z_diff")
printcmat(testdiff3)
println("derivative_z_xact")
printcmat(testxact3)
# println("")
# printcmat(testdiffr2[1:3,1:3])
# println("")
# printcmat(testexpr2[1:3,4:6])
# println("")
# printcmat(testdiff3)
# println("")
# printcmat(testxact3)

ntest = 1000
wtest = rand(Float64,ntest,3)
println("btime runs for xact solutions x,y,z")
@btime begin
    for itest = 1:ntest
        testxact1 = derivative_x_xact(wtest[itest,:],timestep)
    end
end

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

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

println("btime runs for rotation matrix exponentiation x,y,z")
@btime begin
    for itest = 1:ntest
        testexpm1 = derivative_x_expr(wtest[itest,:],timestep)
    end
end

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

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

println("btime runs for finite difference calculations rx,ry,rz,x,y,z")
@btime begin
    for itest = 1:ntest
        testxact1 = derivative_x_diffr(wtest[itest,:],timestep,eps)
    end
end
@btime begin
    for itest = 1:ntest
        testxact1 = derivative_y_diffr(wtest[itest,:],timestep,eps)
    end
end
@btime begin
    for itest = 1:ntest
        testxact1 = derivative_z_diffr(wtest[itest,:],timestep,eps)
    end
end
@btime begin
    for itest = 1:ntest
        testxact1 = derivative_x_diff(wtest[itest,:],timestep,eps)
    end
end
@btime begin
    for itest = 1:ntest
        testxact1 = derivative_y_diff(wtest[itest,:],timestep,eps)
    end
end
@btime begin
    for itest = 1:ntest
        testxact1 = derivative_z_diff(wtest[itest,:],timestep,eps)
    end
end