using Convex
using SCS
"""`ar(y, na; λ = 0)`"""
function ar(y::Vector{Float64}, na; λ = 0, normtype=2, verbose=false)
    y_train, A = getARregressor(y,na)
    n_points = size(A,1)
    if normtype == 2
        if λ == 0
            w = A\y_train
            method = :LS
        else
            w = (A'A + λ*eye(size(A,2)))\A'y_train
            method = :LS_reg
        end
    elseif normtype == 1
        w = Variable(size(A,2),1)
        problem = minimize(sum(abs(A*w-y_train )) + λ*norm(w))
        solve!(problem, SCSSolver(verbose=Int(verbose)))
        w = w.value[:]
        method = :L1
    end
    prediction = A*w
    error = y_train - prediction
    model = AR(w,na)
    result = FitResult(y_train,prediction,na, method,λ)
    return model, result
end
ar(iddata::IdData, na; λ = 0) = ar(iddata.y, na; λ = 0)

"""arx(y, u, na, nb; λ = 0)"""
function arx(y::Vector{Float64}, u::VecOrMat{Float64}, na, nb; λ = 0, normtype=2, verbose=false)
    y_train, A = getARXregressor(y,u, na, nb)
    n_points = size(A,1)
    if normtype == 2
        if λ == 0
            w = A\y_train
            method = :LS
        else
            w = (A'A + λ*eye(size(A,2)))\A'y_train
            method = :LS_reg
        end
    elseif normtype == 1
        w = Variable(size(A,2),1)
        problem = minimize(sum(abs(A*w-y_train )) + λ*norm(w))
        solve!(problem, SCSSolver(verbose=Int(verbose)))
        w = w.value[:]
        method = :L1
    end
    prediction = A*w
    error = y_train - prediction

    si = na+1
    b = Polynom[w[si:si+nb[1]-1]]
    si += nb[1]
    for i = 2:length(nb)
        push!(b,w[si:si+nb[i]-1])
        si += nb[i]
    end
    model = ARX(w[1:na],b,na,nb)
    result = FitResult(y_train,prediction,na, method, λ)
    return model, result
end
arx(iddata::IdData, na, nb; λ = 0) = arx(iddata.y, iddata.u, na, nb; λ = 0)


function getARregressor(y::Vector{Float64},na)
    A = toeplitz(y[na+1:end],y[na+1:-1:1])
    y = copy(A[:,1])
    A = A[:,2:end]
    n_points = size(A,1)
    return y,A
end
getARregressor(iddata::IdData, na) = getARXregressor(iddata.y, na)

function getARXregressor(y::Vector{Float64},u::VecOrMat{Float64}, na, nb)
    assert(length(nb) == size(u,2))
    m = max(na+1,maximum(nb))
    n = length(y) - m+1
    offs = m-na-1
    A = toeplitz(y[offs+na+1:n+na+offs],y[offs+na+1:-1:1])
    y = copy(A[:,1])
    A = A[:,2:end]

    for i = 1:length(nb)
        offs = m-nb[i]
        A = [A toeplitz(u[nb[i]+offs:n+nb[i]+offs-1,i],u[nb[i]+offs:-1:1+offs,i])]
    end


    return y,A
end
getARXregressor(iddata::IdData, na, nb) = getARXregressor(iddata.y,iddata.u, na, nb)


"""Plots the RMSE and AIC for model orders up to `n`. Useful for model selection"""
function find_na(y::Vector,n::Int)
    error = zeros(n,2)
    for i = 1:n
        w,e = ar(y,i)
        error[i,1] = rms(e)
        error[i,2] = aic(e,i)
        print(i,", ")
    end
    println("Done")
    scatter(error, show=true)
end

"""
    plotmodel(y,m::AR)

Plots a signal `y` and the output of the model `m`
"""
function plotmodel(y,m::AR)
    na = length(m.a)
    y,A = getARregressor(y,na)
    yh = A*m.a
    error = y-yh
    plot(y,c=:black)
    plot(yh,c=:b)
    plot(error,c=:r, title="Fitresult, AR, n_a: $na, RMSE = $(rms(error)) Fit = $(fit(y,yh))")

end