# using Pkg; Pkg.add("ControlSystemIdentification"); Pkg.add("Optim"); Pkg.add("Plots")

using ControlSystemIdentification, ControlSystems, Optim, Plots
using Random, LinearAlgebra
using ControlSystemIdentification: newpem
default(size=(500,280))

function generate_system(nx,ny,nu)
    A   = randn(nx,nx)
    A   = A - A'
    A  -= 0.1I
    A   = exp(A)
    B   = randn(nx,nu)
    C   = randn(ny,nx)
    sys = ss(A,B,C,0,1)
end

Random.seed!(1)
T   = 300                       # Number of time steps
nx  = 3                         # Number of poles in the true system
nu  = 1                         # Number of control inputs
ny  = 1                         # Number of outputs
x0  = randn(nx)                 # Initial state
sim(sys,u,x0=x0) = lsim(sys, u, 1:T, x0=x0)[1] # Helper function
sys = generate_system(nx,nu,ny)
u   = randn(nu,T)               # Generate random input
y   = sim(sys, u, x0)           # Simulate system
d   = iddata(y,u,1)

sysh,opt = newpem(d, nx) # Estimate model

yh = predict(sysh, d)      # Predict using estimated model
plot([y; yh]', lab=["y" "ŷ"])   # Plot prediction and true output, see also predplot

# Both noises
σu = 0.2
σy = 0.2

sysn = generate_system(nx,ny,ny)

u  = randn(nu,T)
un = u + sim(sysn, σu*randn(size(u)),0*x0)
y  = sim(sys, un, x0)
yn = y + sim(sysn, σy*randn(size(u)),0*x0)
dn = iddata(yn, un, 1)
sysh,opt = newpem(dn,nx, iterations=400)
f1 = bodeplot([sys,sysh.sys], lab=["True system" "estimated system"])
f2 = predplot(sysh,dn)
f3 = simplot(sysh,dn)
plot(f1,f2,f3, layout=(1,3), size=(900,300))

u  = randn(nu,T)
un = u + sim(sysn, σu*randn(size(u)),0*x0)
y  = sim(sys, un, x0)
yn = y + sim(sysn, σy*randn(size(u)),0*x0)
d = iddata(y, u, 1)
dn = iddata(yn, un, 1)
sysh,opt = newpem(dn,nx, metric=abs, regularizer=(p, P)->(0.1/T)*norm(p), iterations=3000)
plot(predplot(sysh,d),simplot(sysh,d))

σy = 0.5
sim(sys,u) = lsim(sys, u, 1:T)[1]
sys = tf(1,[1,2*0.1,0.1])
sysn = tf(σy,[1,2*0.1,0.3])

# Training data
u  = randn(nu,T)
un = u + 0.1randn(size(u))
y  = sim(sys, u)
yn = y + sim(sysn, σy*randn(size(u)))
dn = iddata(yn, un, 1)

# Validation data
uv  = randn(nu,T)
yv  = sim(sys, uv)
ynv = yv + sim(sysn, σy*randn(size(uv)))
dnv = iddata(ynv, uv, 1)

# fit 4 different models
res = [newpem(dn,nx, iterations=50) for nx = 1:4]

ω = exp10.(range(-2, stop=log10(pi), length=150))
fig = plot(layout=4, size=(1000,600))
for i in eachindex(res)
    (sysh,opt) = res[i]
    predplot!(sysh,dnv; subplot=1, ploty=i==1)
    simplot!(sysh,dnv; subplot=2, ploty=i==1)
end
bodeplot!(getindex.(res,1), ω, plotphase=false, subplot=3, title="Process", linewidth=2*[4 3 2 1])
bodeplot!([noise_model(r[1]) for r in res], ω, plotphase=false, subplot=4, linewidth=2*[4 3 2 1])
bodeplot!(sys, ω, plotphase=false, subplot=3, lab="True", linecolor=:blue, l=:dash, legend = :bottomleft, title="System model")
bodeplot!(ControlSystems.innovation_form(ss(sys),syse=σy^2*ss(sysn), R2=I), ω, plotphase=false, subplot=4, lab="True", linecolor=:blue, l=:dash, ylims=(0.1, 100), legend = :bottomleft, title="Noise model")