#  conserve.jl -- an energy conserving IRK integrator

using Printf, LinearAlgebra

logYs=zeros(10); loghs=zeros(10)

function f(t,y)
    r=[y[2]*y[3]*sin(t)-y[1]*y[2]*y[3],
       -y[1]*y[3]*sin(t)+1/20*y[1]*y[3],
       y[1]^2*y[2]-1/20*y[1]*y[2]]
    return r
end

function euler(t,y,h)
    return y+h*f(t,y)
end

function rk2(t,y,h)
    k1=f(t,y)
    k2=f(t,y+h*2/3*k1)
    return y+h*(1/4*k1+3/4*k2)
end

# adding const here makes it run 10 times faster
const s3o6=sqrt(3)/6
const a=[1/4 1/4-s3o6; 1/4+s3o6 1/4]
const b=[1/2,1/2]
const c=[1/2-s3o6,1/2+s3o6]

function fG(k,t,h,y)
    G=zeros(6)
G[1] = k[1] - sin(t + h*c[1])*(h*(a[1, 1]*k[2] + a[1, 2]*k[5]) + y[2])*(h*(a[1, 1]*k[3] + a[1, 2]*k[6]) + y[3]) + (h*(a[1, 1]*k[1] + a[1, 2]*k[4]) + y[1])*(h*(a[1, 1]*k[2] + a[1, 2]*k[5]) + y[2])*(h*(a[1, 1]*k[3] + a[1, 2]*k[6]) + y[3])
G[2] = k[2] - sin(t + h*c[1])*(-(h*(a[1, 1]*k[1] + a[1, 2]*k[4])) - y[1])*(h*(a[1, 1]*k[3] + a[1, 2]*k[6]) + y[3]) - ((h*(a[1, 1]*k[1] + a[1, 2]*k[4]) + y[1])*(h*(a[1, 1]*k[3] + a[1, 2]*k[6]) + y[3]))/20
G[3] = k[3] + ((h*(a[1, 1]*k[1] + a[1, 2]*k[4]) + y[1])*(h*(a[1, 1]*k[2] + a[1, 2]*k[5]) + y[2]))/20 - (h*(a[1, 1]*k[1] + a[1, 2]*k[4]) + y[1])^2*(h*(a[1, 1]*k[2] + a[1, 2]*k[5]) + y[2])
G[4] = k[4] - sin(t + h*c[2])*(h*(a[2, 1]*k[2] + a[2, 2]*k[5]) + y[2])*(h*(a[2, 1]*k[3] + a[2, 2]*k[6]) + y[3]) + (h*(a[2, 1]*k[1] + a[2, 2]*k[4]) + y[1])*(h*(a[2, 1]*k[2] + a[2, 2]*k[5]) + y[2])*(h*(a[2, 1]*k[3] + a[2, 2]*k[6]) + y[3])
G[5] = k[5] - sin(t + h*c[2])*(-(h*(a[2, 1]*k[1] + a[2, 2]*k[4])) - y[1])*(h*(a[2, 1]*k[3] + a[2, 2]*k[6]) + y[3]) - ((h*(a[2, 1]*k[1] + a[2, 2]*k[4]) + y[1])*(h*(a[2, 1]*k[3] + a[2, 2]*k[6]) + y[3]))/20
G[6] = k[6] + ((h*(a[2, 1]*k[1] + a[2, 2]*k[4]) + y[1])*(h*(a[2, 1]*k[2] + a[2, 2]*k[5]) + y[2]))/20 - (h*(a[2, 1]*k[1] + a[2, 2]*k[4]) + y[1])^2*(h*(a[2, 1]*k[2] + a[2, 2]*k[5]) + y[2])
    return G
end

function fDG(k,t,h,y)
    DG=zeros(6,6)

DG[1, 1] = 1 + h*a[1, 1]*(h*(a[1, 1]*k[2] + a[1, 2]*k[5]) + y[2])*(h*(a[1, 1]*k[3] + a[1, 2]*k[6]) + y[3])
DG[1, 2] = -(sin(t + h*c[1])*h*a[1, 1]*(h*(a[1, 1]*k[3] + a[1, 2]*k[6]) + y[3])) + h*a[1, 1]*(h*(a[1, 1]*k[1] + a[1, 2]*k[4]) + y[1])*(h*(a[1, 1]*k[3] + a[1, 2]*k[6]) + y[3])
DG[1, 3] = -(sin(t + h*c[1])*h*a[1, 1]*(h*(a[1, 1]*k[2] + a[1, 2]*k[5]) + y[2])) + h*a[1, 1]*(h*(a[1, 1]*k[1] + a[1, 2]*k[4]) + y[1])*(h*(a[1, 1]*k[2] + a[1, 2]*k[5]) + y[2])
DG[1, 4] = h*a[1, 2]*(h*(a[1, 1]*k[2] + a[1, 2]*k[5]) + y[2])*(h*(a[1, 1]*k[3] + a[1, 2]*k[6]) + y[3])
DG[1, 5] = -(sin(t + h*c[1])*h*a[1, 2]*(h*(a[1, 1]*k[3] + a[1, 2]*k[6]) + y[3])) + h*a[1, 2]*(h*(a[1, 1]*k[1] + a[1, 2]*k[4]) + y[1])*(h*(a[1, 1]*k[3] + a[1, 2]*k[6]) + y[3])
DG[1, 6] = -(sin(t + h*c[1])*h*a[1, 2]*(h*(a[1, 1]*k[2] + a[1, 2]*k[5]) + y[2])) + h*a[1, 2]*(h*(a[1, 1]*k[1] + a[1, 2]*k[4]) + y[1])*(h*(a[1, 1]*k[2] + a[1, 2]*k[5]) + y[2])
DG[2, 1] = -(h*a[1, 1]*(h*(a[1, 1]*k[3] + a[1, 2]*k[6]) + y[3]))/20 + sin(t + h*c[1])*h*a[1, 1]*(h*(a[1, 1]*k[3] + a[1, 2]*k[6]) + y[3])
DG[2, 2] = 1
DG[2, 3] = -(sin(t + h*c[1])*h*a[1, 1]*(-(h*(a[1, 1]*k[1] + a[1, 2]*k[4])) - y[1])) - (h*a[1, 1]*(h*(a[1, 1]*k[1] + a[1, 2]*k[4]) + y[1]))/20
DG[2, 4] = -(h*a[1, 2]*(h*(a[1, 1]*k[3] + a[1, 2]*k[6]) + y[3]))/20 + sin(t + h*c[1])*h*a[1, 2]*(h*(a[1, 1]*k[3] + a[1, 2]*k[6]) + y[3])
DG[2, 5] = 0
DG[2, 6] = -(sin(t + h*c[1])*h*a[1, 2]*(-(h*(a[1, 1]*k[1] + a[1, 2]*k[4])) - y[1])) - (h*a[1, 2]*(h*(a[1, 1]*k[1] + a[1, 2]*k[4]) + y[1]))/20
DG[3, 1] = (h*a[1, 1]*(h*(a[1, 1]*k[2] + a[1, 2]*k[5]) + y[2]))/20 - 2*h*a[1, 1]*(h*(a[1, 1]*k[1] + a[1, 2]*k[4]) + y[1])*(h*(a[1, 1]*k[2] + a[1, 2]*k[5]) + y[2])
DG[3, 2] = (h*a[1, 1]*(h*(a[1, 1]*k[1] + a[1, 2]*k[4]) + y[1]))/20 - h*a[1, 1]*(h*(a[1, 1]*k[1] + a[1, 2]*k[4]) + y[1])^2
DG[3, 3] = 1
DG[3, 4] = (h*a[1, 2]*(h*(a[1, 1]*k[2] + a[1, 2]*k[5]) + y[2]))/20 - 2*h*a[1, 2]*(h*(a[1, 1]*k[1] + a[1, 2]*k[4]) + y[1])*(h*(a[1, 1]*k[2] + a[1, 2]*k[5]) + y[2])
DG[3, 5] = (h*a[1, 2]*(h*(a[1, 1]*k[1] + a[1, 2]*k[4]) + y[1]))/20 - h*a[1, 2]*(h*(a[1, 1]*k[1] + a[1, 2]*k[4]) + y[1])^2
DG[3, 6] = 0
DG[4, 1] = h*a[2, 1]*(h*(a[2, 1]*k[2] + a[2, 2]*k[5]) + y[2])*(h*(a[2, 1]*k[3] + a[2, 2]*k[6]) + y[3])
DG[4, 2] = -(sin(t + h*c[2])*h*a[2, 1]*(h*(a[2, 1]*k[3] + a[2, 2]*k[6]) + y[3])) + h*a[2, 1]*(h*(a[2, 1]*k[1] + a[2, 2]*k[4]) + y[1])*(h*(a[2, 1]*k[3] + a[2, 2]*k[6]) + y[3])
DG[4, 3] = -(sin(t + h*c[2])*h*a[2, 1]*(h*(a[2, 1]*k[2] + a[2, 2]*k[5]) + y[2])) + h*a[2, 1]*(h*(a[2, 1]*k[1] + a[2, 2]*k[4]) + y[1])*(h*(a[2, 1]*k[2] + a[2, 2]*k[5]) + y[2])
DG[4, 4] = 1 + h*a[2, 2]*(h*(a[2, 1]*k[2] + a[2, 2]*k[5]) + y[2])*(h*(a[2, 1]*k[3] + a[2, 2]*k[6]) + y[3])
DG[4, 5] = -(sin(t + h*c[2])*h*a[2, 2]*(h*(a[2, 1]*k[3] + a[2, 2]*k[6]) + y[3])) + h*a[2, 2]*(h*(a[2, 1]*k[1] + a[2, 2]*k[4]) + y[1])*(h*(a[2, 1]*k[3] + a[2, 2]*k[6]) + y[3])
DG[4, 6] = -(sin(t + h*c[2])*h*a[2, 2]*(h*(a[2, 1]*k[2] + a[2, 2]*k[5]) + y[2])) + h*a[2, 2]*(h*(a[2, 1]*k[1] + a[2, 2]*k[4]) + y[1])*(h*(a[2, 1]*k[2] + a[2, 2]*k[5]) + y[2])
DG[5, 1] = -(h*a[2, 1]*(h*(a[2, 1]*k[3] + a[2, 2]*k[6]) + y[3]))/20 + sin(t + h*c[2])*h*a[2, 1]*(h*(a[2, 1]*k[3] + a[2, 2]*k[6]) + y[3])
DG[5, 2] = 0
DG[5, 3] = -(sin(t + h*c[2])*h*a[2, 1]*(-(h*(a[2, 1]*k[1] + a[2, 2]*k[4])) - y[1])) - (h*a[2, 1]*(h*(a[2, 1]*k[1] + a[2, 2]*k[4]) + y[1]))/20
DG[5, 4] = -(h*a[2, 2]*(h*(a[2, 1]*k[3] + a[2, 2]*k[6]) + y[3]))/20 + sin(t + h*c[2])*h*a[2, 2]*(h*(a[2, 1]*k[3] + a[2, 2]*k[6]) + y[3])
DG[5, 5] = 1
DG[5, 6] = -(sin(t + h*c[2])*h*a[2, 2]*(-(h*(a[2, 1]*k[1] + a[2, 2]*k[4])) - y[1])) - (h*a[2, 2]*(h*(a[2, 1]*k[1] + a[2, 2]*k[4]) + y[1]))/20
DG[6, 1] = (h*a[2, 1]*(h*(a[2, 1]*k[2] + a[2, 2]*k[5]) + y[2]))/20 - 2*h*a[2, 1]*(h*(a[2, 1]*k[1] + a[2, 2]*k[4]) + y[1])*(h*(a[2, 1]*k[2] + a[2, 2]*k[5]) + y[2])
DG[6, 2] = (h*a[2, 1]*(h*(a[2, 1]*k[1] + a[2, 2]*k[4]) + y[1]))/20 - h*a[2, 1]*(h*(a[2, 1]*k[1] + a[2, 2]*k[4]) + y[1])^2
DG[6, 3] = 0
DG[6, 4] = (h*a[2, 2]*(h*(a[2, 1]*k[2] + a[2, 2]*k[5]) + y[2]))/20 - 2*h*a[2, 2]*(h*(a[2, 1]*k[1] + a[2, 2]*k[4]) + y[1])*(h*(a[2, 1]*k[2] + a[2, 2]*k[5]) + y[2])
DG[6, 5] = (h*a[2, 2]*(h*(a[2, 1]*k[1] + a[2, 2]*k[4]) + y[1]))/20 - h*a[2, 2]*(h*(a[2, 1]*k[1] + a[2, 2]*k[4]) + y[1])^2
DG[6, 6] = 1

    return DG
end

function irk(t,y,h)
# use explicit method to make the guesses
    k1=f(t+h*c[1],rk2(t,y,h*c[1]))
    k2=f(t+h*c[2],rk2(t,y,h*c[2]))
    K=[k1;k2]
# newton's iteration
    for i=1:5
        Gi=fG(K,t,h,y)
        DGi=fDG(K,t,h,y)
        DK=-DGi\Gi
        K=K+DK
    end
# time step
    k1=K[1:3]
    k2=K[4:6]
    return y+h*(b[1]*k1+b[2]*k2)
end

function solve(t0,y0,h,n)
    yn=copy(y0)
    for j=1:n
        tn=t0+(j-1)*h
        yn=irk(tn,yn,h)
    end
    return yn
end

function main()
    t0=0; T=1
    y0=[1.0,-1,1]
    yold=zeros(3)
    @printf("%5s %7s %14s    %14s %14s %14s\n",
        "p","n","h","E(1)","dE(1)","error")
    for p=6:16
        n=2^p
        h=T/n
        yn=solve(t0,y0,h,n)
        @printf("%5d %7d %14.7e    %14.7e %14.7e %14.7e\n",
            p,n,h,yn'*yn,yn'*yn-y0'*y0,norm(yold-yn))
        if p>6
            t=log(norm(yold-yn))
            global logYs[p-6]=t
            global loghs[p-6]=log(2*h)
        end
        yold=copy(yn)
    end
end

main()