using LinearAlgebra

function flag_poles(poles)
    N = length(poles);
    cindex = zeros(Int, N);
    for m=1:N
        if abs(imag(poles[m])) >= eps(Float64)
            if m==1
                cindex[m] = 1;
            else
                if cindex[m-1]==0 || cindex[m-1]==2
                    cindex[m] = 1;
                    cindex[m+1] = 2;
                else
                    cindex[m] = 2;
                end
            end
        end
    end
    return cindex
end
function vectfit3(f, s, poles, weight; relax=true, stable=true, asymp=2,
                  skip_pole=false, skip_res=false)
                  """
                  ===========================================================
                  =   Fast Relaxed Vector Fitting                           =
                  =   Version 1.0                                           =
                  =   Last revised: 08.08.2008                              =
                  =   Written by: Bjorn Gustavsen                           =
                  =   SINTEF Energy Research, N-7465 Trondheim, NORWAY      =
                  =   bjorn.gustavsen@sintef.no                             =
                  =   http://www.energy.sintef.no/Produkt/VECTFIT/index.asp =
                  =   Note: RESTRICTED to NON-COMMERCIAL use                =
                  ===========================================================
              PURPOSE : Approximate f(s) with a state-space model
                      f(s)=C*(s*I-A)^(-1)*B +D +s*E
                        where f(s) is a singe element or a vector of elements.
                        When f(s) is a vector, all elements become fitted with a common
                        pole set.
              INPUT :
              f(s) : function (vector) to be fitted.
                     dimension : (Nc,Ns)
                                  Nc : number of elements in vector
                                  Ns : number of frequency samples
              s : vector of frequency points [rad/sec]
                     dimension : (1,Ns)
              poles : vector of initial poles [rad/sec]
                      dimension : (1,N)
              weight: the rows in the system matrix are weighted using this array. Can be used
                      for achieving higher accuracy at desired frequency samples.
                      If no weighting is desired, use unitary weights: weight=ones(1,Ns).
                      Two dimensions are allowed:
                        dimension : (1,Ns) --> Common weighting for all vector elements.
                        dimension : (Nc,Ns)--> Individual weighting for vector elements.
              relax: optional, default=true
                 true --> Use relaxed nontriviality constraint
                 false --> Use nontriviality constraint of "standard" vector fitting
              stable: optional, default=true
                  false --> unstable poles are kept unchanged
                  true --> unstable poles are made stable by 'flipping' them
                                into the left half-plane
              asymp: optional, default=2
                  1 --> Fitting with D=0,  E=0
                  2 --> Fitting with D!=0, E=0
                  3 --> Fitting with D!=0, E!=0
              skip_pole: optional, default=false
                  true --> The pole identification part is skipped, i.e (C,D,E)
                            are identified using the initial poles (A) as final poles.
              skip_res: optional, default=false
                  true --> The residue identification part is skipped, i.e. only the
                             poles (A) are identified while C,D,E are returned as zero.
              ========
              OUTPUT :
                  fit(s) = C*(s*I-(A)^(-1)*B +D +s.*E
              poles(1,N)    : new poles
              rmserr(1) : root-mean-square error of approximation for f(s).
                                (0 is returned if skip_res==1)
              fit(Nc,Ns): Rational approximation at samples. (0 is returned if
                          skip_res==1).
              APPROACH:
              The identification is done using the pole relocating method known as Vector Fitting [1],
              with relaxed non-triviality constraint for faster convergence and smaller fitting errors [2],
              and utilization of matrix structure for fast solution of the pole identifion step [3].
             ********************************************************************************
              NOTE: The use of this program is limited to NON-COMMERCIAL usage only.
              If the program code (or a modified version) is used in a scientific work,
              then reference should be made to the following:
              [1] B. Gustavsen and A. Semlyen, "Rational approximation of frequency
                  domain responses by Vector Fitting", IEEE Trans. Power Delivery,
                  vol. 14, no. 3, pp. 1052-1061, July 1999.
              [2] B. Gustavsen, "Improving the pole relocating properties of vector
                  fitting", IEEE Trans. Power Delivery, vol. 21, no. 3, pp. 1587-1592,
                  July 2006.
              [3] D. Deschrijver, M. Mrozowski, T. Dhaene, and D. De Zutter,
                  "Macromodeling of Multiport Systems Using a Fast Implementation of
                  the Vector Fitting Method", IEEE Microwave and Wireless Components
                  Letters, vol. 18, no. 6, pp. 383-385, June 2008.
             ********************************************************************************
              This example script is part of the vector fitting package (VFIT3.zip)
              Last revised: 08.08.2008.
              Created by:   Bjorn Gustavsen.
             """
    TOLlow = 1e-18;
    TOLhigh = 1e18;
    a = 0;
    b = 0;
    try
        a, b = size(poles);
    catch
        a = 1;
        b = size(poles)[1];
    end
    if s[1]==0 && a==1
        if poles[1]==0 && poles[2]!=0
            poles[1] = -1;
        elseif poles[2]==0 && poles[1]!=0
            poles[2] = -1;
        elseif poles[1]==0 && poles[2]==0
            poles[1] = -1+10im;
            poles[2] = -1-10im;
        end
    end
    rmserr = [];
    #try
    #    a, b = size(s);
    #catch
    #    a = 1;
    #    b = size(s)[1];
    #end
    #if a < b
        s = transpose(s);
    #end
    #TODO sanity check the arguments
    LAMBD = diagm(0 => poles);
    Ns = length(s);
    N = size(LAMBD)[1]; #length(LAMBD);
    Nc = length(f[:,1]);
    B = ones(N,1);
    SERA = poles;
    SERC = zeros(Nc,N);
    SERD = zeros(Nc,1);
    SERE = zeros(Nc,1);
    roetter = poles;
    fit = zeros(ComplexF64, Nc, Ns);

    weight = transpose(weight);
    if length(weight[1,:])==1
        common_weight = 1;
    elseif length(weight[1,:])==Nc
        common_weight = 0;
    else
        display("ERROR in vectfit3.m: Invalid size of array weight")
        return
    end

    if asymp==1
        offs = 0;
    elseif asymp==2
        offs = 1;
    else
        offs = 2;
    end
    # =======================================================
    ## POLE IDENTIFICATION:
    # =======================================================
    if !(skip_pole)
        Escale = zeros(1,Nc+1);
        # =======================================================
        # Finding out which starting poles are complex :
        # =======================================================
        cindex = flag_poles(poles);
        # =======================================================
        ## Building system - matrix :
        # =======================================================
        if asymp==1 || asymp==2
            Dk = zeros(ComplexF64, Ns,N+1);
            for i=1:Ns
                Dk[i,N+1] = 1;
            end
        elseif asymp==3
            Dk = zeros(ComplexF64, Ns,N+2);
            for i=1:Ns
                Dk[i,N+1] = 1;
                Dk[i,N+2] = s[i];
            end
        else
            Dk = zeros(ComplexF64, Ns,N);
        end
        for m=1:N
            if cindex[m]==0      #real pole
                Dk[:,m] = 1 ./(s .- LAMBD[m,m]);
            elseif cindex[m]==1  #complex pole, 1st part
                Dk[:,m] = 1 ./(s .- LAMBD[m,m]) + 1 ./(s .- conj(LAMBD[m,m]));
                Dk[:,m+1] = 1im ./(s .- LAMBD[m,m]) - 1im ./(s .- conj(LAMBD[m,m]));
            end
        end

        #Scaling for last row of LS-problem (pole identification)
        scale = 0;
        for m=1:Nc
            if length(weight[1,:])==1
                scale = scale + (norm(weight[:,1].*transpose(f[m,:])))^2;
            else
                scale = scale + (norm(weight[:,m].*transpose(f[m,:])))^2;
            end
        end
        scale = sqrt(scale)/Ns;

        if relax
            #Escale=zeros(1,Nc*(N+offs)+N+1);
            #scale=norm(f);#/Ns; #Scaling for sigma in LS problem
            AA = zeros(ComplexF64, Nc*(N+1), N+1);
            bb = zeros(ComplexF64, Nc*(N+1), 1);
            Escale = zeros(ComplexF64, 1, length(AA[1,:]));

            for n=1:Nc
                A = zeros(ComplexF64, Ns, (N+offs)+N+1); #b=zeros(Ns*Nc+1,1);
                if common_weight==1
                    weig = weight;
                else
                    weig = weight[:,n];
                end
                for m=1:N+offs #left block
                    A[1:Ns,m] = weig .* Dk[1:Ns,m];
                end
                inda = N + offs;
                for m=1:N+1 #right block
                    A[1:Ns,inda+m] = -weig .* Dk[1:Ns,m] .* f[n,1:Ns];
                end
                A = [real(A); imag(A); zeros(1, (N+offs)+N+1)];
                #Integral criterion for sigma:
                offset = (N+offs);
                if n==Nc
                    for mm=1:N+1
                        A[2*Ns+1, offset+mm] = real(scale*sum(Dk[:,mm]));
                    end
                end
                F = qr(A);
                Q = Matrix(F.Q);
                R = Matrix(F.R);
                ind1 = N + offs + 1;
                ind2 = N + offs + N + 1;
                R22 = R[ind1:ind2,ind1:ind2];
                AA[(n-1)*(N+1)+1:n*(N+1),:] = R22;
                if n==Nc
                    bb[(n-1)*(N+1)+1:n*(N+1),1] = (Q[end,N+offs+1:end]')*Ns*scale;
                end
            end #for n=1:Nc
            for col=1:length(AA[1,:])
                Escale[col] = 1/norm(AA[:,col]);
                AA[:,col] = Escale[col] .* AA[:,col];
            end
            x = AA\bb;
            #size(x),size(Escale)
            x = x .* transpose(Escale);
        end #if relax

        #Situation: No relaxation, or produced D of sigma extremely small and large. Solve again, without relaxation
        if !relax | (abs(x[end]) < TOLlow) | (abs(x[end]) > TOLhigh)
            AA = zeros(ComplexF64, Nc*(N),N);
            bb = zeros(ComplexF64, Nc*(N),1);
            if !relax
                Dnew = 1;
            else
                if x[end]==0
                    Dnew = 1;
                elseif abs(x[end])<TOLlow
                    Dnew = sign(x[end])*TOLlow;
                elseif abs(x[end])>TOLhigh
                    Dnew = sign(x[end])*TOLhigh;
                end
            end

            for n=1:Nc
                A=zeros(Ns,(N+offs) +N); #b=zeros(Ns*Nc+1,1);
                Escale=zeros(1,N);
                if common_weight==1
                    weig = weight;
                else
                    weig=weight[:,n];
                end

                for m=1:N+offs #left block
                    A[1:Ns,m] = weig.*Dk[1:Ns,m];
                end
                inda = N+offs;
                for m=1:N #right block
                    A[1:Ns,inda+m] = -weig.*Dk[1:Ns,m].*transpose(f[n,1:Ns]);
                end
                b = Dnew*weig.*transpose(f[n,1:Ns]);
                A = [real(A);imag(A)];
                b = [real(b);imag(b)];
                offset = (N+offs);
                Q, R = qr(A,0);
                ind1 = N+offs+1;
                ind2 = N+offs+N;
                R22 = R[ind1:ind2,ind1:ind2];
                AA[(n-1)*N+1:n*N,:] = R22;
                bb[((n-1)*N+1:n*N,1)] = transpose(Q[:,ind1:ind2])*b;
            end #for n=1:Nc
            for col=1:length(AA[1,:])
                Escale[col] = 1 ./ norm(AA[:,col]);
                AA[:,col] = Escale[col] .* AA[:,col];
            end
            #if use_normal==1
            #  x=AA.'*AA\(AA.'*bb);
            #else
            x = AA\bb;
            #end
            x = x .* transpose(Escale);
            x = [x;Dnew];
        end  #if relax==0 | abs(x[end])<TOLlow | abs(x[end])>TOLhigh
        # ************************************
        C = x[1:end-1];
        D = x[end]; #NEW!!
        #We now change back to make C complex :
        # **************
        for m=1:N
            if cindex[m]==1
                for n=1:1 #Nc+1
                    r1 = C[m];
                    r2 = C[m+1];
                    C[m] = r1 + 1im*r2;
                    C[m+1] = r1 - 1im*r2;
                end
            end
        end
        # **************

        # =============================================================================
        # We now calculate the zeros for sigma :
        # =============================================================================
        #oldLAMBD=LAMBD;oldB=B;oldC=C;7
        m = 0;
        for n=1:N
            m = m+1;
            if m<N
                if ( abs(LAMBD[m,m]) > abs(real(LAMBD[m,m])) ) #complex number?
                    LAMBD[m+1,m] = -imag(LAMBD[m,m]);
                    LAMBD[m,m+1] = imag(LAMBD[m,m]);
                    LAMBD[m,m] = real(LAMBD[m,m]);
                    LAMBD[m+1,m+1] = LAMBD[m,m];
                    B[m,1] = 2;
                    B[m+1,1] = 0;
                    koko = C[m];
                    C[m] = real(koko);
                    C[m+1] = imag(koko);
                    m = m+1;
                end
            end
        end
        ZER = LAMBD - B*transpose(C)/D;
        ZER = Array{Float64,2}(ZER); #FIXME that should not be needed, but gives non-conjugate complex pairs...
        roetter = transpose(eigvals(ZER));
        unstables = map(r -> real(r) > 0, roetter);
        if stable
            roetter[unstables] = roetter[unstables] - 2*real(roetter[unstables]); #Forcing unstable poles to be stable...
        end
        roetter = transpose(roetter);
        sort!(roetter, by = x -> (real(x), imag(x)));
        N = length(roetter);
        # =============================================
        #Sorterer polene s.a. de reelle kommer først:
        for n=1:N
            for m=n+1:N
                if (abs(imag(roetter[m])) <= eps(Float64)
                    && abs(imag(roetter[n])) >= eps(Float64))
                    trans = roetter[n];
                    roetter[n] = roetter[m];
                    roetter[m] = trans;
                end
            end
        end
        N1=0;
        for m=1:N
            if imag(roetter[m])==0
                N1 = m;
            end
        end
        if N1<N
            #roetter[N1+1:N] = sort(roetter[N1+1:N]);
            sort!(roetter[N1+1:N], by = x -> (real(x), imag(x)));
        end       # N1: n.o. real poles
        #N2=N-N1; # N2: n.o. imag.poles
        roetter = roetter - 2im*imag(roetter); #10.11.97 !!!
        SERA = transpose(roetter);
    end #if !skip_pole
    # =========================================================================
    #  RESIDUE IDENTIFICATION:
    # =========================================================================
    if !(skip_res)
        # ============================================================================
        # We now calculate SER for f, using the modified zeros of sigma as new poles :
        # ============================================================================
        LAMBD = diagm(0 => roetter);
        #B=ones(N,1);
        # Finding out which poles are complex :
        cindex = flag_poles(roetter);

        # ===============================================================================
        # We now calculate the SER for f (new fitting), using the above calculated
        # zeros as known poles :
        # ===============================================================================
        if asymp==1
            A = zeros(ComplexF64, 2*Ns,N);
            BB = zeros(ComplexF64, 2*Ns,Nc);
        elseif asymp==2
            A = zeros(ComplexF64, 2*Ns,N+1);
            BB = zeros(ComplexF64, 2*Ns,Nc);
        else
            A = zeros(ComplexF64, 2*Ns,N+2);
            BB = zeros(ComplexF64, 2*Ns,Nc);
        end
        for m=1:N
            if cindex[m]==0      #real pole
                Dk[:,m] = 1 ./(s .- LAMBD[m,m]);
            elseif cindex[m]==1  #complex pole, 1st part
                Dk[:,m] = 1 ./(s .- LAMBD[m,m]) + 1 ./(s .- conj(LAMBD[m,m]));
                Dk[:,m+1] = 1im ./(s .- LAMBD[m,m]) - 1im ./(s .- conj(LAMBD[m,m]));
            end
        end
        if common_weight==1
            weight = transpose(weight);
            #I3=diag(ones(1,Nc));I3(:,Nc)=[];
            Dk = zeros(ComplexF64, Ns,N);
            for m=1:N
                if cindex[m]==0      #real pole
                    Dk[:,m] = weight ./(s .- LAMBD[m,m]);
                elseif cindex[m]==1  #complex pole, 1st part
                    Dk[:,m] = weight ./(s .- LAMBD[m,m]) + weight ./(s .- conj(LAMBD[m,m]));
                    Dk[:,m+1] = 1im .*weight ./(s .- LAMBD[m,m]) - 1im .*weight ./(s .- conj(LAMBD[m,m]));
                end
            end
            if asymp==1
                A[1:Ns,1:N] = Dk;
            elseif asymp==2
                A[1:Ns,1:N] = Dk;
                A[1:Ns,N+1] = weight;
            else
                A[1:Ns,1:N] = Dk;
                A[1:Ns,N+1] = weight;
                A[1:Ns,N+2] = weight.*s;
            end
            for m=1:Nc
                BB[1:Ns,m] = weight.*transpose(f[m,:]);
            end
            A[Ns+1:2*Ns,:] = imag(A[1:Ns,:]);
            A[1:Ns,:] = real(A[1:Ns,:]);
            BB[Ns+1:2*Ns,:] = imag(BB[1:Ns,:]);
            BB[1:Ns,:] = real(BB[1:Ns,:]);

            if asymp==2
                A[1:Ns,N+1] = A[1:Ns,N+1];
            elseif asymp==3
                A[1:Ns,N+1] = A[1:Ns,N+1];
                A[Ns+1:2*Ns,N+2] = A[Ns+1:2*Ns,N+2];
            end
            Escale = zeros(ComplexF64, 1,length(A[1,:]));
            for col=1:length(A[1,:]);
                Escale[col] = norm(A[:,col],2);
                A[:,col] = A[:,col]./Escale[col];
            end
            X = A\BB;
            for n=1:Nc
                X[:,n] = X[:,n]./transpose(Escale);
            end
            X = transpose(X);
            C = X[:,1:N];

            if asymp==2
                SERD = X[:,N+1];
            elseif asymp==3
                SERE = X[:,N+2];
                SERD = X[:,N+1];
            end

        else #if common_weight==1
            SERD = zeros(ComplexF64, Nc,1);
            SERE = zeros(ComplexF64, Nc,1);
            C = zeros(ComplexF64, Nc,N);
            for n=1:Nc
                if asymp==1
                    A[1:Ns,1:N] = Dk;
                elseif asymp==2
                    A[1:Ns,1:N] = Dk;
                    A[1:Ns,N+1] = 1;
                else
                    A[1:Ns,1:N] = Dk;
                    A[1:Ns,N+1] = 1;
                    A[1:Ns,N+2] = s;
                end
                for m=1:length(A[1,:])
                    A[1:Ns,m] = weight[:,n].*A[1:Ns,m];
                end
                BB = weight[:,n].*transpose(f(n,:));
                A[Ns+1:2*Ns,:] = imag(A[1:Ns,:]);
                A[1:Ns,:] = real(A[1:Ns,:]);
                BB[Ns+1:2*Ns] = imag(BB[1:Ns]);
                BB[1:Ns] = real(BB[1:Ns]);

                if asymp==2
                    A[1:Ns,N+1] = A[1:Ns,N+1];
                elseif asymp==3
                    A[1:Ns,N+1] = A[1:Ns,N+1];
                    A[Ns+1:2*Ns,N+2] = A[Ns+1:2*Ns,N+2];
                end
                Escale = zeros(1,length(A[1,:]));
                for col=1:length(A[1,:]);
                    Escale[col] = norm(A[:,col],2);
                    A[:,col] = A[:,col]./Escale[col];
                end
                x = A\BB;
                x = x./transpose(Escale);
                C[n,1:N] = transpose(x[1:N]);
                if asymp==2
                    SERD[n] = x[N+1];
                elseif asymp==3
                    SERE[n] = x[N+2];
                    SERD[n] = x[N+1];
                end
            end #for n=1:Nc
        end #if common_weight==1
        # =========================================================================
        #We now change back to make C complex.
        for m=1:N
            if cindex[m]==1
                for n=1:Nc
                    r1 = C[n,m];
                    r2 = C[n,m+1];
                    C[n,m] = r1 + 1im*r2;
                    C[n,m+1] = r1 - 1im*r2;
                end
            end
        end
        # **************
        B = ones(N,1);
        # ====================================================
        SERA = LAMBD;
        SERB = B;
        SERC = C;
        Dk = zeros(ComplexF64, Ns,N);
        for m=1:N
            Dk[:,m] = 1 ./(s .- SERA[m,m]);
        end
        for n=1:Nc
            fit[n,:] = transpose(Dk*SERC[n,:]);
            if asymp==2
                fit[n,:] = fit[n,:] .+ SERD[n];
            elseif asymp==3
                fit[n,:] = fit[n,:] .+ SERD[n] .+ transpose(s).*SERE[n];
            end
        end
        #fit = transpose(fit);
        #f = transpose(f);
        dif1 = fit - f;
        rmserr = sqrt( sum( sum(map(abs, dif1.^2)) ) )/sqrt(Nc*Ns);
        #fit = transpose(fit);
    end # if !skip_res
    A = SERA;
    poles = A;
    if !(skip_res)
        B = SERB;
        C = SERC;
        D = SERD;
        E = SERE;
    else
        B = ones(N,1);
        C = zeros(Nc,N);
        D = zeros(Nc,Nc);
        E = zeros(Nc,Nc);
        rmserr = 0;
    end;
    #A = sparse(diag(A));    % A is complex, make it diagonal
    #SER.A=A; SER.B=B; SER.C=C; SER.D=D; SER.E=E;
    return SERD[1] , diag(poles)[:], SERC[:]
end


nlp_test = Dict()

nlp_test[1] = function (optimizer;opt...)

    m = SimpleNL.Model(optimizer;opt...)

    x = variable(m;lb=2,ub=4)
    objective(m,(x-1)^2)

    optimize!(m)

    return compare(value(x), 2.)
end

nlp_test[2] = function (optimizer;opt...)

    m = SimpleNL.Model(optimizer;opt...)

    x = variable(m)
    constraint(m,x+1<= -1)
    constraint(m,x+1>= -1)
    constraint(m,-2<= x+1)
    constraint(m,0>=x+1)

    optimize!(m)

    return compare(value(x), -2)
end

nlp_test[3] = function (optimizer;opt...)

    m = SimpleNL.Model(optimizer;opt...)

    x = [variable(m,name="s[$i]",start = .1) for i=1:3]

    c = constraint(m,0 ==x[1]+sin(x[2])-x[3]/2 + 1. )
    objective(m,x[2]^2 + x[2]^2 + 1.)

    optimize!(m)

    return compare(value.(x), [-0.74,0,.52]) && compare(dual(c),0.)
end

nlp_test[4] = function (optimizer;opt...)

    m = SimpleNL.Model(optimizer;opt...)

    x = [variable(m;start=mod(i,2)==1 ? -1.2 : 1.) for i=1:10];
    p = [parameter(m,2) for i=1:10];
    setvalue.(p,2)

    objective(m,sum(100(x[i-1]^2-x[i])^2+(x[i-1]-1)^2 for i=2:10))

    c = [ constraint(m,3x[i+1]^3+p[i]*x[i+2]-5+sin(x[i+1]-x[i+2])sin(x[i+1]+x[i+2])+4x[i+1]-x[i]exp(x[i]-x[i+1])-3 == 0) for i=1:8 ];

    optimize!(m)
    return compare(value.(x), [-0.95055636, 0.91390082, 0.98909052, 0.99855924, 0.99980874, 0.99997459, 0.99999662, 0.99999955, 0.99999994, 0.99999993])
end

m = SimpleNL.Model()

x = variable(m)
y = variable(m)
setvalue(x,3)
set_lower_bound(x,-1)
set_upper_bound(x,1)
@test SimpleNL.index(x) == 1
@test value(x) == 3
@test lower_bound(x) == -1
@test upper_bound(x) == 1

c = constraint(m,x==1.)
constraint(m,x>=1.)
constraint(m,x<=1.)

constraint(m,1. ==y)
constraint(m,1. <=y)
constraint(m,-1. >=y)

constraint(m,x==y)
constraint(m,x-4. <=y)
constraint(m,x+4. >=y)

setvalue(c,2)
set_lower_bound(c,-1)
set_upper_bound(c,1)
@test value(c) == 2
@test lower_bound(c) == -1
@test upper_bound(c) == 1

p = parameter(m,1.)
setvalue(p,4)
@test value(p) == 4

objective(m,sum([p,x,y]))

m[:test] = "test"

@test objective_value(m) == 7.
@test num_variables(m) == 2
@test num_constraints(m) == 9
@test m[:test] == "test"

for (optimizer,opt) in [(Ipopt.Optimizer,[:print_level=>0]),
                        (MadNLP.Optimizer,[:print_level=>MadNLP.ERROR])]
    for f in values(nlp_test)
        @test f(optimizer;opt...)
    end
end


using Stheno, Plots, Random
using Stheno: @model

###########################  Define and inspect our model  ###########################

# Define the vanilla GP-regression generative model.
@model function model()

    # Define a smooth latent process.
    f = 1.5 * GP(EQ())

    # Define a latent noise process.
    noise = GP(Noise(1e-2))

    # Sum them to get the process of which we shall make observations.
    y = f + noise

    # Return all three processes so that we can inspect all of them.
    return f, noise, y
end

# Generate some input locations and sample from the prior.
rng, N = MersenneTwister(123546), 500;
X_prior = sort(rand(rng, N) * 10);
f, noise, y = model();

# Take a look at the mean function and kernel of y.
@show mean(y), kernel(y);

# Take a look at the cross-covariance between f and y.
@show kernel(f, y);

# Inspect the cross-covariance between noise and y.
@show kernel(noise, y);

###########################  Sample jointly from the prior  ###########################

# Sample jointly from the prior distribution over the three processes.
fX, noiseX, yX = rand(rng, [f(X_prior), noise(X_prior), y(X_prior)]);

#######################  Do posterior inference give a few samples  #######################

# Sample a few points from the prior and compute the posterior processes.
N = 15;
X = sort(rand(rng, N) * 10);
ŷ = rand(rng, y(X));
f′, noise′, y′ = (f, noise, y) | (y(X) ← ŷ);

Nplot, S = 500, 100;
Xplot = range(-2.0, stop=12.0, length=Nplot);
f′Xplot_, noise′Xplot_ = f′(Xplot), noise′(Xplot)
f′Xp, noise′Xp = rand(rng, [f′Xplot_, noise′Xplot_], S);

# Get posterior mean and marginals f′ and y′ and write them for plotting.
μ′f, σ′f = marginals(f′(Xplot));
σ′y = marginal_std(y′(Xplot));

####################################  Plot results  ####################################

plotly();

prior_plot = plot(X_prior, fX, label="f");
scatter!(prior_plot, X_prior, noiseX,
    label="noise",
    markershape=:xcross,
    markerstrokewidth=0.0,
    markersize=2);
scatter!(prior_plot, X_prior, yX,
    label="y",
    markershape=:xcross,
    markerstrokewidth=0.0,
    markersize=2);
plot!(prior_plot, X_prior, fX + noiseX - yX,
    label="f + noise - y",
    linecolor=:red);

posterior_plot = plot(Xplot, f′Xp;
    linecolor=:blue,
    linealpha=0.2,
    legend=false);
plot!(posterior_plot, Xplot, μ′f;
    linewidth=2.0,
    linecolor=:blue);
plot!(posterior_plot, Xplot, [μ′f μ′f];
    linewidth=0.0,
    fillrange=[μ′f .- 3 .* σ′y, μ′f .+ 3 * σ′y],
    fillalpha=0.3,
    fillcolor=:red);
plot!(posterior_plot, Xplot, [μ′f μ′f];
    linewidth=0.0,
    fillrange=[μ′f.- 3  .* σ′f μ′f .+ 3 .* σ′f],
    fillalpha=0.5,
    fillcolor=:blue);
scatter!(posterior_plot, X, ŷ;
    markercolor=:red,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=4,
    markeralpha=0.7,
    legend=false);
joint_plot = plot(prior_plot, posterior_plot, layout=(2, 1));
display(joint_plot);


"""
    gumbelfitpwm(...)

Estimate the Gumbel parameters with the probability weighted moments.

# Implementation

Estimation with the probability weighted moments, as described by [Landwehr *et al. (1979)](https://agupubs.onlinelibrary.wiley.com/doi/abs/10.1029/WR015i005p01055),
is only possible in the stationary case.

See also [`gumbelfitpwm`](@ref) for the other methods, [`gevfit`](@ref), [`gevfitbayes`](@ref) and [`BlockMaxima`](@ref).

# Reference

Landwehr, J. M., Matalas, N. C. and Wallis, J. R. (1979). Probability weighted moments compared with
    some traditional techniques in estimating Gumbel parameters and quantiles. *Water Resources Research*,
    15:1055–1064.
"""
function gumbelfitpwm end

"""
    gumbelfitpwm(y::Vector{<:Real})

Estimate the Gumbel parameters with the probability weighted moments.

"""
function gumbelfitpwm(y::Vector{<:Real})::pwmEVA

    model = BlockMaxima(Variable("y", y))

    fittedmodel = gumbelfitpwm(model)

    return fittedmodel

end

"""
    gumbelfitpwm(df::DataFrame, datacol::Symbol)::pwmEVA

Estimate the Gumbel parameters with the probability weighted moments.

Block maxima data are in the column `datacol` of the dataframe `df`.
"""
function gumbelfitpwm(df::DataFrame, datacol::Symbol)::pwmEVA

    model = BlockMaxima(Variable(string(datacol), df[:, datacol]))

    fittedmodel = gumbelfitpwm(model)

    return fittedmodel

end

"""
    gumbelfitpwm(model::BlockMaxima)

Estimate the Gumbel parameters with the probability weighted moments.
"""
function gumbelfitpwm(model::BlockMaxima)::pwmEVA

    model = validatestationarity(model)

    y = model.data.value

    a₀ = pwm(y,1,0,0)
    a₁ = pwm(y,1,0,1)

    σ̂ = (a₀ - 2a₁)/log(2)

    μ̂ = a₀ - Base.MathConstants.eulergamma*σ̂

    ϕ̂ = log(σ̂)

    θ̂ = [μ̂, ϕ̂, 0.0]

    fm = pwmEVA{BlockMaxima, Gumbel}(model,θ̂)

    return fm

end


#####################################################
### Project of IN480 - From Python 2.7 to Julia 0.6
### --- SIMPLEXN - serial unit tests ---
### Authors: Fabio Fatelli, Emanuele Loprevite
#####################################################

include("../src/simplexn-serial.jl")

using Base.Test

@testset "Unit Tests" begin
	@testset "larExtrude1" begin
		VOID = [Int64[]],[[0]] # the empty simplicial model
		pattern1 = [1,-1]
		model1 = larExtrude1(VOID,pattern1) # 1D
		#println(model1)
		@test typeof(model1) == Tuple{Array{Array{Int64,1},1},Array{Array{Int64,1},1}}
		@test length(model1[1]) == length(pattern1)+1 # num of vertices, no rep
		@test length(model1[2]) == length(filter(x->x>0,pattern1)) # num of 1D-simplices

		pattern2 = [1,1,1,-1]
		model2 = larExtrude1(VOID,pattern2) # 1D
		#println(model2)
		@test typeof(model2) == Tuple{Array{Array{Int64,1},1},Array{Array{Int64,1},1}}
		@test length(model2[1]) == length(pattern2)+1 # num of vertices, no rep
		@test length(model2[2]) == length(filter(x->x>0,pattern2)) # num of 1D-simplices

		pattern3 = [1,1,1,1,1,-1]
		model3 = larExtrude1(VOID,pattern3) # 1D
		#println(model3)
		@test typeof(model3) == Tuple{Array{Array{Int64,1},1},Array{Array{Int64,1},1}}
		@test length(model3[1]) == length(pattern3)+1 # num of vertices, no rep
		@test length(model3[2]) == length(filter(x->x>0,pattern3)) # num of 1D-simplices

		model1 = larExtrude1(model1,pattern1) # 2D
		#println(model1)
		@test typeof(model1) == Tuple{Array{Array{Int64,1},1},Array{Array{Int64,1},1}}
		@test length(model1[1]) == (length(pattern1)+1)^2 # num of vertices, no rep
		@test length(model1[2]) == 2*length(filter(x->x>0,pattern1))^2 # num of 2D-simplices

		model2 = larExtrude1(model2,pattern2) # 2D
		#println(model2)
		@test typeof(model2) == Tuple{Array{Array{Int64,1},1},Array{Array{Int64,1},1}}
		@test length(model2[1]) == (length(pattern2)+1)^2 # num of vertices, no rep
		@test length(model2[2]) == 2*length(filter(x->x>0,pattern2))^2 # num of 2D-simplices

		model3 = larExtrude1(model3,pattern3) # 2D
		#println(model3)
		@test typeof(model3) == Tuple{Array{Array{Int64,1},1},Array{Array{Int64,1},1}}
		@test length(model3[1]) == (length(pattern3)+1)^2 # num of vertices, no rep
		@test length(model3[2]) == 2*length(filter(x->x>0,pattern3))^2 # num of 2D-simplices

		model1 = larExtrude1(model1,pattern1) # 3D
		#println(model1)
		@test typeof(model1) == Tuple{Array{Array{Int64,1},1},Array{Array{Int64,1},1}}
		@test length(model1[1]) == (length(pattern1)+1)^3 # num of vertices, no rep
		@test length(model1[2]) == 6*length(filter(x->x>0,pattern1))^3 # num of 3D-simplices

		model2 = larExtrude1(model2,pattern2) # 3D
		#println(model2)
		@test typeof(model2) == Tuple{Array{Array{Int64,1},1},Array{Array{Int64,1},1}}
		@test length(model2[1]) == (length(pattern2)+1)^3 # num of vertices, no rep
		@test length(model2[2]) == 6*length(filter(x->x>0,pattern2))^3 # num of 3D-simplices

		model3 = larExtrude1(model3,pattern3) # 3D
		#println(model3)
		@test typeof(model3) == Tuple{Array{Array{Int64,1},1},Array{Array{Int64,1},1}}
		@test length(model3[1]) == (length(pattern3)+1)^3 # num of vertices, no rep
		@test length(model3[2]) == 6*length(filter(x->x>0,pattern3))^3 # num of 3D-simplices

		model = ([[0.0,0],[0,1]],[[1,0],[1,1]]) # 2D with float
		@test typeof(model[1]) == Array{Array{Float64,1},1}
		model = larExtrude1(model,pattern1)
		#println(model)
		@test typeof(model) == Tuple{Array{Array{Float64,1},1},Array{Array{Int64,1},1}}
	end
	@testset "larSimplexGrid1" begin
		shape = [3] # 1D
		model1 = larSimplexGrid1(shape)
		#println(model1)
		@test typeof(model1) == Tuple{Array{Array{Int64,1},1},Array{Array{Int64,1},1}}
		@test length(model1[1]) == shape[1]+1 # num of vertices, no rep
		@test length(model1[2]) == shape[1] # num of 1D-simplices

		shape = [3,5] # 2D
		model2 = larSimplexGrid1(shape)
		#println(model2)
		@test typeof(model2) == Tuple{Array{Array{Int64,1},1},Array{Array{Int64,1},1}}
		# (shape[1]+1)*(shape[2]+1)
		@test length(model2[1]) == prod(shape)+sum(shape)+1 # num of vertices, no rep
		@test length(model2[2]) == 2*prod(shape) # num of 2D-simplices

		shape = [3,5,7] # 3D
		model3 = larSimplexGrid1(shape)
		#println(model3)
		@test typeof(model3) == Tuple{Array{Array{Int64,1},1},Array{Array{Int64,1},1}}
		# (shape[1]+1)*(shape[2]+1)*(shape[3]+1); num of vertices, no rep
		@test length(model3[1]) == prod(shape)+sum(prod.(collect(combinations(shape,2))))+sum(shape)+1
		@test length(model3[2]) == 6*prod(shape) # num of 3D-simplices
	end
	@testset "larSimplexFacets" begin
		s = larSimplexGrid1([3])[2] # 1D
		#println(s)
		sOut = larSimplexFacets(s)
		#println(sOut)
		@test typeof(sOut) == Array{Array{Int64,1},1}
		d1 = sum([length(s[k]) for k in 1:length(s)])/length(s)
		d2 = sum([length(sOut[k]) for k in 1:length(sOut)])/length(sOut)
		@test d1-1 == d2 # dimension
		# length(s)*binomial(length(s[1]),(length(s[1])-1))
		@test length(sOut) <= length(s)*length(s[1]) # "<=" because no rep

		s = larSimplexGrid1([3,5])[2] # 2D
		#println(s)
		sOut = larSimplexFacets(s)
		#println(sOut)
		@test typeof(sOut) == Array{Array{Int64,1},1}
		d1 = sum([length(s[k]) for k in 1:length(s)])/length(s)
		d2 = sum([length(sOut[k]) for k in 1:length(sOut)])/length(sOut)
		@test d1-1 == d2 # dimension
		# length(s)*binomial(length(s[1]),(length(s[1])-1))
		@test length(sOut) <= length(s)*length(s[1]) # "<=" because no rep

		s = larSimplexGrid1([3,5,7])[2] # 3D
		#println(s)
		sOut = larSimplexFacets(s)
		#println(sOut)
		@test typeof(sOut) == Array{Array{Int64,1},1}
		d1 = sum([length(s[k]) for k in 1:length(s)])/length(s)
		d2 = sum([length(sOut[k]) for k in 1:length(sOut)])/length(sOut)
		@test d1-1 == d2 # dimension
		# length(s)*binomial(length(s[1]),(length(s[1])-1))
		@test length(sOut) <= length(s)*length(s[1]) # "<=" because no rep
	end
	@testset "quads2tria" begin
		modIn = ([[0,0,0],[0,1,0],[1,0,0],[1,1,0]],[[0,1,2,3]]) # 2D
		modOut = quads2tria(modIn)
		#println(modOut)
		@test typeof(modIn) == Tuple{Array{Array{Int64,1},1},Array{Array{Int64,1},1}}
		@test typeof(modOut) == Tuple{Array{Array{Float64,1},1},Array{Array{Int64,1},1}}
		edges = 2*sum([length(modIn[2][k]) for k in 1:length(modIn[2])])
		EulerChar = length(modOut[1])-edges+length(modOut[2]) # V - E + F
		@test EulerChar == 1

		modIn = ([[0,0,0],[0,1,0],[1,0,0],[1,1,0],[0.5,1.5,0]],[[0,1,2,3,4]]) # 2D
		modOut = quads2tria(modIn)
		#println(modOut)
		@test typeof(modIn) == Tuple{Array{Array{Float64,1},1},Array{Array{Int64,1},1}}
		@test typeof(modOut) == Tuple{Array{Array{Float64,1},1},Array{Array{Int64,1},1}}
		edges = 2*sum([length(modIn[2][k]) for k in 1:length(modIn[2])])
		EulerChar = length(modOut[1])-edges+length(modOut[2]) # V - E + F
		@test EulerChar == 1

		modIn = ([[0,0,0],[0,1,0],[1,0,0],[1,1,0],[0,0,1],[0,1,1],[1,0,1],[1,1,1]],[[0,1,2,3,4,5,6,7]]) # 3D
		modOut = quads2tria(modIn)
		#println(modOut)
		@test typeof(modIn) == Tuple{Array{Array{Int64,1},1},Array{Array{Int64,1},1}}
		@test typeof(modOut) == Tuple{Array{Array{Float64,1},1},Array{Array{Int64,1},1}}
		edges = 2*sum([length(modIn[2][k]) for k in 1:length(modIn[2])])
		EulerChar = length(modOut[1])-edges+length(modOut[2]) # V - E + F
		@test EulerChar == 1

		modIn = ([[0,0,0],[0,1,0],[1,0,0],[1,1,0],[0,0,1],[0,1,1],[1,0,1],[1,1,1],[0.5,0.5,1.5]],[[0,1,2,3,4,5,6,7,8]]) # 3D
		modOut = quads2tria(modIn)
		#println(modOut)
		@test typeof(modIn) == Tuple{Array{Array{Float64,1},1},Array{Array{Int64,1},1}}
		@test typeof(modOut) == Tuple{Array{Array{Float64,1},1},Array{Array{Int64,1},1}}
		edges = 2*sum([length(modIn[2][k]) for k in 1:length(modIn[2])])
		EulerChar = length(modOut[1])-edges+length(modOut[2]) # V - E + F
		@test EulerChar == 1
	end
end


# Define the AbstractGP type and its API.

"""
    abstract type AbstractGP end

Supertype for various Gaussian process (GP) types. A common interface is provided for
interacting with each of these objects. See [1] for an overview of GPs.

[1] - C. E. Rasmussen and C. Williams. "Gaussian processes for machine learning".
MIT Press. 2006.
"""
abstract type AbstractGP end

"""
    mean(f::AbstractGP, x::AbstractVector)

Computes the mean vector of the multivariate Normal `f(x)`.
"""
Statistics.mean(::AbstractGP, ::AbstractVector)

"""
    cov(f::AbstractGP, x::AbstractVector)

Compute the `length(x)` by `length(x)` covariance matrix of the multivariate Normal `f(x)`.
"""
Statistics.cov(::AbstractGP, x::AbstractVector)

"""
    cov_diag(f::AbstractGP, x::AbstractVector)

Compute only the diagonal elements of `cov(f(x))`.
"""
cov_diag(::AbstractGP, ::AbstractVector)

"""
    cov(f::AbstractGP, x::AbstractVector, y::AbstractVector)

Compute the `length(x)` by `length(y)` cross-covariance matrix between `f(x)` and `f(y)`.
"""
Statistics.cov(::AbstractGP, x::AbstractVector, y::AbstractVector)

"""
    mean_and_cov(f::AbstractGP, x::AbstractVector)

Compute both `mean(f(x))` and `cov(f(x))`. Sometimes more efficient than separately
computation, particularly for posteriors.
"""
mean_and_cov(f::AbstractGP, x::AbstractVector) = (mean(f, x), cov(f, x))

"""
    mean_and_cov_diag(f::AbstractGP, x::AbstractVector)

Compute both `mean(f(x))` and the diagonal elements of `cov(f(x))`. Sometimes more efficient
than separately computation, particularly for posteriors.
"""
mean_and_cov_diag(f::AbstractGP, x::AbstractVector) = (mean(f, x), cov_diag(f, x))


immutable DRIFT{T} <: Model
    β₀::T
    β₁::T
    β₂::T
    β₃::T
    β₄::T
end

Base.show(io::IO, ::Type{DRIFT}) = print(io, "Heuristic DRIFT")

function Base.show(io::IO, m::DRIFT)
    @printf(
        io,
        "Heuristic DRIFT (β₀ = %.3f, β₁ = %.3f, β₂ = %.3f, β₃ = %.3f, β₄ = %.3f)",
        m.β₀,
        m.β₁,
        m.β₂,
        m.β₃,
        m.β₄,
    )
end

initial_parameters(::Type{DRIFT}) = Float64[0.0, 0.0, 0.0, 0.0, 0.0]

canonical_parameters(m::DRIFT) = (m.β₀, m.β₁, m.β₂, m.β₃, m.β₄)

function DRIFT(parameters::Vector)
    β₀, β₁, β₂, β₃, β₄ = (
        parameters[1],
        parameters[2],
        parameters[3],
        parameters[4],
        parameters[5],
    )
    return DRIFT(β₀, β₁, β₂, β₃, β₄)
end

@inline function predict(m::DRIFT, x2::Real, t2::Real, x1::Real, t1::Real)
    β₀, β₁, β₂, β₃, β₄ = m.β₀, m.β₁, m.β₂, m.β₃, m.β₄
    δx = x2 - x1
    δt = t2 - t1
    z = (
        β₀ +
        β₁ * δx +
        β₂ * (δx / x1) +
        β₃ * ((x2 / x1)^(1 / δt) - 1) +
        β₄ * δt
    )
    p = invlogit(z)
    return p
end

@inline function gradient_component(
    m::DRIFT,
    x2::Real,
    t2::Real,
    x1::Real,
    t1::Real,
)
    δx = x2 - x1
    δt = t2 - t1
    gr = (
        1.0,
        δx,
        (δx / x1),
        ((x2 / x1)^(1 / δt) - 1),
        δt,
    )
    return gr
end


# Based on https://github.com/JuliaGPU/CUDA.jl/blob/master/examples/peakflops.jl

_kernel_fma_nfmas()::Int = 100_000
_kernel_fma_N()::Int = Int((_kernel_fma_nfmas() - 1) ÷ 3)

using CUDA: i32
"Dummy kernel doing `_kernel_fma_nfmas()` many FMAs (default: `100_000`)."
function kernel_fma(a, b, c, out)
    i = (blockIdx().x - 1i32) * blockDim().x + threadIdx().x
    @inbounds if i <= length(out)
        a_val = a[i]
        b_val = b[i]
        c_val = c[i]

        # We suggest LLVM to unroll the loop. For large values, however,
        # there might still be a loop (with a large loop body of ~1000 fmas).
        # Investigate with `@device_code_ptx @cuda kernel_fma(d_a, d_b, d_c, d_out)`.
        @unroll for j in 1:_kernel_fma_N()
            a_val = CUDA.fma(a_val, b_val, c_val)
            b_val = CUDA.fma(a_val, b_val, c_val)
            c_val = CUDA.fma(a_val, b_val, c_val)
        end

        out[i] = CUDA.fma(a_val, b_val, c_val)
    end

    return nothing
end

"""
    peakflops_gpu_fmas(; size::Integer=5_000_000, dtype=Float32, nbench=5, nkernel=5, device=CUDA.device(), verbose=true)
Tries to estimate the peak performance of a GPU in TFLOP/s by measuring the time
it takes to perform `_kernel_fma_nfmas() * size` many FMAs on CUDA cores.

**Keyword arguments:**
* `device` (default: `CUDA.device()`): CUDA device to be used.
* `dtype` (default: `Float32`): element type of the matrices.
* `size` (default: `5_000_000`): length of vectors.
* `nkernel` (default: `5`): number of kernel calls that make up one benchmarking sample.
* `nbench` (default: `5`): number of measurements to be performed the best of which is used for the TFLOP/s computation.
* `verbose` (default: `true`): toggle printing.
* `io` (default: `stdout`): set the stream where the results should be printed.
"""
function peakflops_gpu_fmas(;
    size::Integer=5_000_000,
    dtype=Float32,
    nbench=5,
    nkernel=5,
    device::CuDevice=CUDA.device(),
    verbose=true,
    io::IO=stdout,
)
    device!(device) do
        d_a = CUDA.rand(dtype, size)
        d_b = CUDA.rand(dtype, size)
        d_c = CUDA.rand(dtype, size)
        d_out = CUDA.zeros(dtype, size)

        kernel = @cuda launch = false kernel_fma(d_a, d_b, d_c, d_out)
        config = launch_configuration(kernel.fun)
        threads = min(size, config.threads)
        blocks = cld(size, threads)

        # warm-up
        CUDA.@elapsed kernel(d_a, d_b, d_c, d_out)

        t = Inf
        for _ in 1:nbench
            Δt = CUDA.@elapsed begin
                for _ in 1:nkernel
                    kernel(d_a, d_b, d_c, d_out; threads=threads, blocks=blocks)
                end
            end
            t = min(t, Δt)
        end
        t /= nkernel

        flopcount = 2 * _kernel_fma_nfmas() * size
        flops = (flopcount * 1e-12) / t

        if verbose
            printstyled(io,"Peakflops (TFLOP/s):\n"; bold=true)
            print(io," └ max: ")
            printstyled(io, round(flops; digits=2), "\n"; color=:green, bold=true)
        end
        return flops
    end
end


@testset "modpi" begin
    x=cbrt(41)*sqrt(Double64(pi))

    @test DoubleFloats.mod2pi(x) == Double64(6.111805926475162, -1.667563077572613e-16)
    @test DoubleFloats.mod1pi(x) == Double64(2.9702132728853683, 1.54868222178066e-16)
    # these tests are unusable, sometimes they pass and othertimes they fail
    # @test DoubleFloats.modhalfpi(x) == Double64(1.3994169460904717, 9.362948122802731e-17)
    # @test DoubleFloats.modqrtrpi(x) == Double64(0.6140187826930236, -4.8008991213172176e-17)
end

@testset "rempi" begin
    x=cbrt(41)*sqrt(Double64(pi))

    @test DoubleFloats.rem2pi(x) == Double64(6.111805926475162, -1.667563077572613e-16)
    @test DoubleFloats.rem1pi(x) == Double64(2.9702132728853683, 1.54868222178066e-16)
    @test DoubleFloats.rem2pi(-x) == Double64(-0.1713793807044248, 4.647966647701768e-18)
    @test DoubleFloats.rem1pi(-x) == Double64(2.9702132728853683, 1.54868222178066e-16)
end


#*************************************************************************#
#
# Linear Type
#
#*************************************************************************#
"""
    Linear(ninputs, noutputs)

Initialize `ParametricModel` representing linear operator `z = A*x`.
Where `A` is a dense matrix with dimension `(noutputs, ninputs) `
cointaining the reshaped parameter vector `Θ`.
"""
struct Linear <: ParametricModel
    ninputs::Int
    noutputs::Int
    nparams::Int
end

# Constructor
function Linear(ninputs::Int, noutputs::Int)
    nparams = ninputs*noutputs
    Linear(ninputs, noutputs, nparams)
end

#*************************************************************************#
#
# Generate Buffers
#
#*************************************************************************#
function dΘ_buffer(lin::Linear)
    if lin.noutputs == 1
        return Matrix{Float64}(lin.noutputs, lin.nparams)
    else
        return hcat((speye(lin.noutputs) for i=1:lin.ninputs)...)
    end
end

#*************************************************************************#
#
# Compute Output and its Derivatives
#
#*************************************************************************#
function _evaluate!(lin::Linear,
                    x::AbstractVector{<:Real},
                    Θ::AbstractVector{<:Real},
                    z::AbstractVector{Float64},
                    dx::AbstractMatrix{Float64},
                    dΘ::AbstractMatrix{Float64})
    # Get Matrix
    matrix = reshape(Θ, lin.noutputs, lin.ninputs)

    # Compute z
    z .= matrix*x

    # Compute dz
    if !isempty(dx)
        dx .= matrix
    end

    # Compute dΘ
    if !isempty(dΘ)
        for i = 1:lin.ninputs
            for j = 1:lin.noutputs
                dΘ[j, (i-1)*lin.noutputs+j] = x[i]
            end
        end
    end
end

#*************************************************************************#
#
# Initializer
#
#*************************************************************************#
"""
    ScaledGaussianInitializer()

The `ScaledGaussianInitializer` will initialize the parameter vector
with of random values picked from normal distribution with standard
deviation of `1/sqrt(ninputs)`.

Valid only for Linear static model.
"""
immutable ScaledGaussianInitializer <: Initializer
end

function initial_guess(lin::Linear, init::ScaledGaussianInitializer)
    d = Distributions.Normal(0, 1/sqrt(lin.ninputs))

    return rand(d, lin.nparams)
end

function initial_guess(lin::Linear, init::DefaultInitializer)
    initial_guess(lin, ScaledGaussianInitializer())
end


immutable Truncated{D<:UnivariateDistribution} <: ContinuousUnivariateDistribution
    untruncated::D      # the original distribution (untruncated)
    lower::Float64      # lower bound
    upper::Float64      # upper bound
    lcdf::Float64       # cdf of lower bound
    ucdf::Float64       # cdf of upper bound

    tp::Float64         # the probability of the truncated part, i.e. ucdf - lcdf
    logtp::Float64      # log(tp), i.e. log(ucdf - lcdf)
end

### Constructors

function Truncated(d::ContinuousUnivariateDistribution, l::Float64, u::Float64)
    l < u || error("lower bound should be less than upper bound.")
    lcdf = isinf(l) ? 0.0 : cdf(d, l)
    ucdf = isinf(u) ? 1.0 : cdf(d, u)
    tp = ucdf - lcdf
    Truncated{typeof(d)}(d, l, u, lcdf, ucdf, tp, log(tp))
end

@compat Truncated(d::ContinuousUnivariateDistribution, l::Real, u::Real) = Truncated(d, Float64(l), Float64(u))

### range and support

islowerbounded(d::Truncated) = islowerbounded(d.untruncated) || isfinite(d.lower)
isupperbounded(d::Truncated) = isupperbounded(d.untruncated) || isfinite(d.upper)

minimum(d::Truncated) = max(minimum(d.untruncated), d.lower)
maximum(d::Truncated) = min(maximum(d.untruncated), d.upper)

insupport(d::Truncated, x::Real) =
    d.lower <= x <= d.upper && insupport(d.untruncated, x)

### evaluation

pdf(d::Truncated, x::Float64) = d.lower <= x <= d.upper ? pdf(d.untruncated, x) / d.tp : 0.0

logpdf(d::Truncated, x::Float64) = d.lower <= x <= d.upper ? logpdf(d.untruncated, x) - d.logtp : -Inf

cdf(d::Truncated, x::Float64) = x <= d.lower ? 0.0 :
                             x >= d.upper ? 1.0 :
                             (cdf(d.untruncated, x) - d.lcdf) / d.tp

logcdf(d::Truncated, x::Float64) = x <= d.lower ? -Inf :
                                x >= d.upper ? 0.0 :
                                log(cdf(d.untruncated, x) - d.lcdf) - d.logtp

ccdf(d::Truncated, x::Float64) = x <= d.lower ? 1.0 :
                              x >= d.upper ? 0.0 :
                              (d.ucdf - cdf(d.untruncated, x)) / d.tp

logccdf(d::Truncated, x::Float64) = x <= d.lower ? 0.0 :
                                 x >= d.upper ? -Inf :
                                 log(d.ucdf - cdf(d.untruncated, x)) - d.logtp

quantile(d::Truncated, p::Float64) = quantile(d.untruncated, d.lcdf + p * d.tp)

## random number generation

function rand(d::Truncated)
    d0 = d.untruncated
    if d.tp > 0.25
        while true
            r = rand(d0)
            if d.lower <= r <= d.upper
                return r
            end
        end
    else
        return quantile(d0, d.lcdf + rand() * d.tp)
    end
end

## show

function show(io::IO, d::Truncated)
    print(io, "Truncated(")
    d0 = d.untruncated
    uml, namevals = _use_multline_show(d0)
    uml ? show_multline(io, d0, namevals) :
          show_oneline(io, d0, namevals)
    print(io, ", range=($(d.lower), $(d.upper)))")
    uml && println(io)
end

_use_multline_show(d::Truncated) = _use_multline_show(d.untruncated)

### specialized truncated distributions

include(joinpath("truncated", "normal.jl"))


using TensorOperations
using Base.Test

include("methods.jl")
include("tensor.jl")


@testset "Nodes initialization algorithm entry  dimension >1 output dimension 1" begin
    # Comprobamos que las hipótesis de selección son correctas
    M = 1
    @test RampFunction(M) == 1
    @test RampFunction(-M) == 0

    # Bien definido para tamaño n = 2 y salida de dimensión 1
    f_regression(x,y,z)=x*y-z
    data_set_size =  4
    entry_dimension = 3
    output_dimension = 1
    # Número de neuronas
    n =  data_set_size# Debe de ser mayor que 1 para que no de error
    X_train= rand(Float64, data_set_size, entry_dimension)
    Y_train = map(x->f_regression(x...), eachrow(X_train))

    h = nn_from_data(X_train, Y_train, n, M)

    # veamos que el tamaño de la salida es la adecuada
    @test size(h.W1) == (n,entry_dimension+1)
    @test size(h.W2) == (output_dimension,n)

    # Si ha sido bien construida:
    # Evaluar la red neuronal en los datos con los que se construyó
    # debería de resultar el valor de Y_train respectivo
    evaluar(x)=forward_propagation(h,
     RampFunction,x)

    for (x,y) in zip(eachrow(X_train),Y_train)
        @test evaluar(x) ≈  [y]
    end
end


using JLD2
using Plots
using Printf
using Oceananigans
using Oceananigans.Grids

filename = "oceananigans_unstable_bickley_jet_Nh512_WENO5_fields.jld2"

file = jldopen(filename)

iterations = parse.(Int, keys(file["timeseries/t"]))
grid = file["serialized/grid"]

xc, yc, zc = nodes((Cell, Cell, Cell), grid)
xω, yω, zω = nodes((Face, Face, Cell), grid)

snapshots = []

for iteration in iterations[1:6:54]

    t = file["timeseries/t/$iteration"]
    ω = file["timeseries/ω/$iteration"][:, :, 1]
    c = file["timeseries/c/$iteration"][:, :, 1]

    kwargs = Dict(
                  :aspectratio => 1,
                  :linewidth => 0,
                  :colorbar => false,
                  :xticks => false,
                  :yticks => false,
                  :clims => (-1, 1),
                  :xlims => (-grid.Lx/2, grid.Lx/2),
                  :ylims => (-grid.Ly/2, grid.Ly/2),
                  :title => @sprintf("t = %.1f", t)
                 )

    #snap = heatmap(xc, yc, clamp.(c, -1, 1)'; color = :thermal, kwargs...)
    snap = heatmap(xω, yω, clamp.(ω, -1, 1)'; color = :balance, kwargs...)

    push!(snapshots, snap)
end

close(file)

plt = plot(snapshots..., size = (1800, 1600))

display(plt)

savefig("vorticity.png")
#savefig("dye.png")


function configPotential(y, nodepot, edgepot, es::EdgeStruct)
  nnodes = size(nodepot, 1)
  nedges = size(es.edgeends,1)
  pot = 1
  for n = 1:nnodes
    pot = pot*nodepot[n, y[n]]
  end

  for e = 1:nedges
    i = es.edgeends[e,1]
    j = es.edgeends[e,2]
    pot = pot*edgepot[y[i],y[j],e];
  end
  return pot
end


function Triangle{T}(domain::Triangle{T}, parent::ArbField) where {T<:Union{arb,fmpq}}
    Triangle(domain.angles[1], domain.angles[2], parent)
end

function Base.show(io::IO, domain::Triangle{fmpq})
    angles_string = ["$(numerator(a))π/$(denominator(a))" for a in domain.angles]
    print(
        io,
        "Triangle with $(domain.parent.prec) bits of precision and angles ($(angles_string[1]), $(angles_string[2]), $(angles_string[3]))",
    )
end

function Base.show(io::IO, domain::Triangle{arb})
    angles_string = ["$a" for a in domain.angles]
    print(
        io,
        "Triangle with $(domain.parent.prec) bits of precision and angles ($(angles_string[1]), $(angles_string[2]), $(angles_string[3]))",
    )
end

boundaries(::Triangle) = 1:3

"""
    angle(domain::Triangle, i::Integer)

Return the angle for vertex `i` of the triangle.
"""
function angle(domain::Triangle{fmpq}, i::Integer)
    if i ∈ boundaries(domain)
        return domain.parent(π) * domain.angles[i]
    else
        throw(ArgumentError("attempt to get vertex $i from a $(typeof(domain))"))
    end
end

function angle(domain::Triangle{arb}, i::Integer)
    if i ∈ boundaries(domain)
        return domain.angles[i]
    else
        throw(ArgumentError("attempt to get vertex $i from a $(typeof(domain))"))
    end
end

"""
    angledivπ(domain::Triangle, i::Integer)

Return the angle for vertex `i` of the triangle divided by `π`.
"""
function angledivπ(domain::Triangle{fmpq}, i::Integer)
    if i ∈ boundaries(domain)
        return domain.angles[i]
    else
        throw(ArgumentError("attempt to get vertex $i from a $(typeof(domain))"))
    end
end

function angledivπ(domain::Triangle{arb}, i::Integer)
    if i ∈ boundaries(domain)
        return domain.angles[i] / domain.parent(π)
    else
        throw(ArgumentError("attempt to get vertex $i from a $(typeof(domain))"))
    end
end

"""
    angles(domain::Triangle)

Return the angles of the triangle.
"""
angles(domain::Triangle) = tuple((angle(domain, i) for i in boundaries(domain))...)

"""
    anglesdivπ(domain::Triangle)

Return the angles of the triangle divided by `π`.
"""
anglesdivπ(domain::Triangle) = tuple((angledivπ(domain, i) for i in boundaries(domain))...)

"""
    vertex(domain::Triangle, i::Integer)

Return the Cartesian coordinates of vertex `i` of the triangle.
"""
function vertex(domain::Triangle{T}, i::Integer) where {T}
    if i == 1
        return SVector(domain.parent(0), domain.parent(0))
    elseif i == 2
        return SVector(domain.parent(1), domain.parent(0))
    elseif i == 3
        if T == arb
            x = sin(angle(domain, 2)) / sin(angle(domain, 3))
            s, c = sincos(angle(domain, 1))
        else
            x =
                sinpi(angledivπ(domain, 2), domain.parent) /
                sinpi(angledivπ(domain, 3), domain.parent)
            s, c = sincospi(angledivπ(domain, 1), domain.parent)
        end
        return SVector(c * x, s * x)
    else
        throw(ArgumentError("attempt to get vertex $i from a $(typeof(domain))"))
    end
end

"""
    vertices(domain::Triangle)

Return all three vertices of the triangle.
"""
vertices(domain::Triangle) = tuple((vertex(domain, i) for i in boundaries(domain))...)

area(domain::Triangle) = vertex(domain, 3)[2] / 2

"""
    center(domain::Triangle)

Return the centroid of the triangle.
"""
center(domain::Triangle) = sum(vertices(domain)) / 3

function Base.in(xy, domain::Triangle)
    x, y = xy
    return (
        y >= 0 &&
        atan(y, x) <= angle(domain, 1) &&
        domain.parent(π) - atan(y, x - 1) <= angle(domain, 2) # To the left of the right edge
    )
end

function boundary_parameterization(t, domain::Triangle, i::Integer)
    i ∈ boundaries(domain) ||
        throw(ArgumentError("attempt to get vertex $i from a $(typeof(domain))"))
    v = vertex(domain, mod1(i + 1, 3))
    w = vertex(domain, mod1(i + 2, 3))
    return v .+ t .* (w - v)
end

function boundary_points(
    domain::Triangle,
    i::Integer,
    n::Integer;
    distribution = :chebyshev,
)
    i ∈ boundaries(domain) ||
        throw(ArgumentError("attempt to get vertex $i from a $(typeof(domain))"))

    v = vertex(domain, mod1(i + 1, 3))
    w = vertex(domain, mod1(i + 2, 3))

    if distribution == :linear
        f = j -> domain.parent(j // (n + 1))
    elseif distribution == :chebyshev
        f = j -> (cospi(domain.parent((2j - 1) // 2n)) + 1) / 2
    elseif distribution == :exponential
        m = (n + 1) // 2
        f = j -> begin
            d = exp(-domain.parent(m - j))
            if j < m
                return d / 2
            elseif j > m
                return 1 - inv(d) / 2
            else
                return domain.parent(0.5)
            end
        end
    elseif distribution == :root_exponential
        m = (n + 1) // 2
        f = j -> begin
            d = exp(-domain.parent(m - j) / sqrt(domain.parent(n)))
            if j < m
                return d / 2
            elseif j > m
                return 1 - inv(d) / 2
            else
                return domain.parent(0.5)
            end
        end
    else
        throw(ArgumentError("no distribution named $distribution"))
    end

    points = Vector{SVector{2,arb}}(undef, n)
    for j = 1:n
        t = f(j)
        points[j] = v .+ t .* (w - v)
    end

    return points, fill(i, n)
end

function interior_points(domain::Triangle, n::Integer; rng = MersenneTwister(42))
    v = vertex(domain, 2)
    w = vertex(domain, 3)

    points = Vector{SVector{2,arb}}(undef, n)
    for i = 1:n
        u1, u2 = rand(rng), rand(rng)
        if u1 + u2 > 1
            u1 = 1 - u1
            u2 = 1 - u2
        end
        points[i] = u1 .* v + u2 .* w
    end

    return points
end


∑ = sum
using StatsFuns:logistic
sigmoid = logistic
using CSV:File
using DataFrames
using Printf:@sprintf

# DataTypes
using LinearAlgebra:Transpose
Numeric = Union{Int64,Float64}
NumericV = Union{Array{Int64,1},Array{Float64,1}}
NumericM = Union{Array{Int64,2},Array{Float64,2},Transpose{Float64,Array{Float64,2}},Transpose{Int64,Array{Int64,2}}}
arrayNumericV = Union{Array{Array{Int64,1},1},Array{Array{Float64,1},1}}
TupleNumericM = Tuple{NumericM, NumericM};

function maskBiasUnits(array::NumericV)::NumericV
    mask = ones(size(array));
    mask[1,1] = 0;
    maskedArray = array.*mask;
end

function maskBiasUnits(array::NumericM)::NumericM
    mask = ones(size(array));
    mask[:,1] .= 0;
    maskedArray = array.*mask;
end

function powerElements(array::Union{NumericV,NumericM}, power::Int)::NumericM
    # function defined to be broadcasted to elements of a tuple
    array = array.^power
end

function getLabels(y::NumericV)::NumericV
    labels = collect(BitSet(y));
    sort!(labels)
end

function sortArrayByCol(array::Array{Any,2}, col::Int)
    array = DataFrame(array);
    sort!(array, col)
    array = Matrix(array)
end

function getLayersQ(all_θ)::Int
    θ_matricesQ = length(collect(all_θ));
    layersQ = θ_matricesQ+1;
end

function nnClasificationFeedForward(all_θ, X, y, labels, layersQ::Int=getLayersQ(all_θ))
    labelsQ = length(labels);
    hx_and_indices = reshape([],0,labelsQ+1);
    for lb in labels
        ly = 1;
        lb_idx = findall(==(lb), y)
        a_ly = X[lb_idx,:];
        ly += 1;
        while ly ≤ layersQ
            # add bias units to a_ly_lb
            a_ly = [ones(size(a_ly)[1],1) a_ly];
            a_previous_ly = a_ly;
            θ_previous_ly = all_θ[ly-1];
            z_ly = a_previous_ly*θ_previous_ly';
            a_ly = sigmoid.(z_ly);
            ly += 1;
        end
        ly = 1;
        hx_lb = a_ly;
        hx_and_indices_lb = hcat(hx_lb, lb_idx);
        hx_and_indices = vcat(hx_and_indices, hx_and_indices_lb);
    end
    hx_and_indices = sortArrayByCol(hx_and_indices, labelsQ+1);
    hx = hx_and_indices[:,1:labelsQ];
end

function nnCostClasification(all_θ::TupleNumericM, hx, y::NumericV, labels::NumericV, λ::Numeric)::Numeric
    J = 0;
    for lb in labels
        lb_idx = findall(==(lb), y)
        hx_lb = hx[lb_idx,:];
        y_lb = vec(zeros(size(labels))); y_lb[lb] = 1;
        J += ∑(-((log.(hx_lb))*y_lb)-(log.(1 .- hx_lb)*(1 .- y_lb)));
    end
    casesQ = size(hx)[1];
    J = (1/casesQ)*J;
    # add regularization cost
    J += (λ/(2*casesQ))* ∑(∑.(powerElements.(maskBiasUnits.(all_θ),2)));
end

function calcSigmoidDerivate(z::Numeric)::Numeric
    sigmoid(z)*(1-sigmoid(z));
end

function calcSigmoidDerivate(z::NumericM)::NumericM
    sigmoid.(z).*(1 .- sigmoid.(z));
end

calcSigmoidGradient = calcSigmoidDerivate;

function calcNiceInitϵ(layerInputsQ::Int, layerOutputsQ::Int)::Numeric
    initϵ = √6/√(layerInputsQ+layerOutputsQ)
end

function genRandomInitialθ(layerInputsQ::Int, layerOutputsQ::Int, initϵ::Numeric=calcNiceInitϵ(layerInputsQ,layerOutputsQ))::NumericM
    initϵInt = floor(initϵ*100);
    randInitialθ = rand(-initϵInt:initϵInt,layerOutputsQ,layerInputsQ+1);
    randInitialθ = randInitialθ./100;
end

function substractElements(sub_array::Array{Any,1},vector::NumericV)
    sub_array .- vectors
end

function gen_initial_z(all_θ, labelsQ::Int)
    initial_z_singleLy = [];
    for θ in all_θ
        initial_z_singleLy = [initial_z_singleLy, ones(size(θ))];
    end
    initial_z_singleLy[1]=initial_z_singleLy[1][2];

    initial_z = [];
    for lb in labelsQ
        initial_z = [initial_z, initial_z_singleLy];
    end
    initial_z = initial_z[2]
end

function nnClasificationBackPropagation(hx, all_θ, layersQ::Int, y, labels)
    # calculate the gradients -i.e. partial derivates with respect to θ- of the cost (J)
    nodesQ = labelsQ = length(labels);
    casesQ = size(hx)[1];
    ly = layersQ;
    idx_v = collect(1:labelsQ);
    outputLayer = map(idx -> (vec(hx[:,idx])), idx_v); #Matrix of subarrays where each of them correspond to each row of hx'
    δ_ly =  map(idx -> allCasesOutputLayer[idx] .- y, idx_v);
    ly -= 1;
    while ly ≥ 2
        δ_previous_ly = δ_ly;
        θ_ly = all_θ[ly];
        z_ly =

end

#parameters
# csv files -> DataFrames -> arrays
X =  Matrix{Float64}(File("./X.csv", header=false)
    |> DataFrame);
y = vec(Matrix(File("./y.csv", header=false) |> DataFrame));
θ1 = Matrix{Float64}(File("./theta1.csv", header=false)
    |> DataFrame);
θ2 = Matrix{Float64}(File("./theta2.csv", header=false) |> DataFrame);
λ = 1;

#implementation
all_θ = (θ1, θ2);
labels = getLabels(y);
layersQ = getLayersQ(all_θ);
hx = nnClasificationFeedForward(all_θ, X, y, labels, layersQ);
J = nnCostClasification(all_θ, hx, y, labels, λ);
allCasesδ_ly = nnClasificationBackPropagation(hx, layersQ, y, labels);


## Compute the clustering coefficient ##

maxtrianglecount(g::DirectedGraph, k::Integer) = k * (k - 1)
maxtrianglecount(g::UndirectedGraph, k::Integer) = div(k * (k - 1), 2)

# Adapted from graph-tool with improvements
function triangles(g::Graph, n::Integer, mark::Vector)
    @inbounds for i = neighbors(g, n)
        i == n && continue
        mark[i] = true
    end
    count = 0
    @inbounds for i = neighbors(g, n)
        i == n && continue
        for j = neighbors(g, i)
            j == n && continue
            if mark[j]; count += 1; end
        end
    end
    @inbounds for i = neighbors(g, n)
        mark[i] = false
    end
    return count
end

function dirtriangles(g::DirectedGraph, n::Integer, mark::Vector)
    @inbounds for i = neighbors(g, n)
        i == n && continue
        mark[i] = true
    end
    count = 0
    @inbounds for i = outneighbors(g, n)
        i == n && continue
        for j = outneighbors(g, i)
            j == n && continue
            if mark[j]; count += 1; end
        end
    end
    @inbounds for i = neighbors(g, n)
        mark[i] = false
    end
    return count
end

function _trackcycles!(g::DirectedGraph, n::Integer, mark::Vector)
    @inbounds for i = inneighbors(g, n)
        mark[i] = -1
    end
    @inbounds for i = outneighbors(g, n)
        mark[i] = 1
    end
    count = 0
    @inbounds for i = outneighbors(g, n)
        i == n && continue
        for j = outneighbors(g, i)
            j == n && continue
            if mark[j] == -1; count += 1; end
        end
    end
    @inbounds for i = neighbors(g, n)
        mark[i] = 0
    end
    return count
end

function cycle3count(g::Graph)
    n    = nodecount(g)
    mark = fill(0, n)
    tri, tot = 0, 0
    @inbounds for i = nodes(g)
        d = degree(g, i)
        if d > 1
            tri += _trackcycles!(g, i, mark)
        end
        tot += maxtrianglecount(g, d)
    end
    return tri
end

function clustering(g::Graph)
    n    = nodecount(g)
    mark = fill(false, n)
    tri, tot = 0, 0
    @inbounds for i = nodes(g)
        d = degree(g, i)
        if d > 1
            tri += dirtriangles(g, i, mark)
        end
        tot += maxtrianglecount(g, d)
    end
    return tri * 4. / tot
end


using Base.Test

using CompScienceMeshes
using BEAST
using StaticArrays

T = Float64
P = SVector{3,T}

omega = 1.0;

id = Identity()
nc = NCross()

fn = Pkg.dir("BEAST","test","sphere2.in")
m = readmesh(fn)
rt = raviartthomas(m)
bc = buffachristiansen(m)

G1 = zeros(Complex128, numfunctions(rt), numfunctions(rt))
BEAST.assemble_local_matched!(id, rt, rt, (v,m,n)->(G1[m,n]+=v))

G2 = zeros(Complex128, numfunctions(rt), numfunctions(rt))
BEAST.assemble_local_mixed!(id, rt, rt, (v,m,n)->(G2[m,n]+=v))

# G1 = BEAST.assemble_local_matched(id, rt, rt)
# G2 = BEAST.assemble_local_mixed(id, rt, rt)

@test norm(G1-G2,  Inf) ≈ 0 atol = sqrt(eps())
@test norm(G1-G1', Inf) ≈ 0 atol = sqrt(eps())
@test norm(G2-G2', Inf) ≈ 0 atol = sqrt(eps())

G = zeros(Complex128, numfunctions(rt), numfunctions(bc))
BEAST.assemble_local_mixed!(nc, rt, bc, (v,m,n)->(G[m,n]+=v))
@test cond(G) ≈ 2.919026332637947

G = zeros(Complex128, numfunctions(rt), numfunctions(rt))
BEAST.assemble_local_mixed!(nc, rt, rt, (v,m,n)->(G[m,n]+=v))
@test norm(G+G',  Inf) ≈ 0 atol = sqrt(eps())

v = [
    point(0.0, 0.0, 0.0),
    point(1.0, 0.0, 0.0),
    point(1.0, 1.0, 0.0),
    point(0.0, 1.0, 0.0),]
f = [
    index(1,2,3),
    index(1,3,4)
]
Γ = Mesh(v,f)
X = raviartthomas(Γ)
I = Identity()
G = assemble(I,X,X)
@test size(G) == (1,1)
@test norm(G[1,1] - 1/3) < sqrt(eps())


using Statistics

Statistics._var(A::CuArray, corrected::Bool, mean, dims) =
    sum((A .- something(mean, Statistics.mean(A, dims=dims))).^2, dims=dims)/(prod(size(A)[dims])-corrected)

Statistics._var(A::CuArray, corrected::Bool, mean, dim::Integer) =
    sum((A .- something(mean, Statistics.mean(A, dims=dim))).^2, dims=dim)/(size(A)[dim]-corrected)

Statistics._var(A::CuArray, corrected::Bool, mean, ::Colon) =
    sum((A .- something(mean, Statistics.mean(A))).^2)/(length(A)-corrected)

Statistics._std(A::CuArray, corrected::Bool, mean, dims) =
    Base.sqrt.(Statistics.var(A; corrected=corrected, mean=mean, dims=dims))

Statistics._std(A::CuArray, corrected::Bool, mean, ::Colon) =
    Base.sqrt.(Statistics.var(A; corrected=corrected, mean=mean, dims=:))

# Revert https://github.com/JuliaLang/Statistics.jl/pull/25
Statistics._mean(A::CuArray, ::Colon)    = sum(A) / length(A)
Statistics._mean(f, A::CuArray, ::Colon) = sum(f, A) / length(A)
Statistics._mean(A::CuArray, dims)    = mean!(Base.reducedim_init(t -> t/2, +, A, dims), A)
Statistics._mean(f, A::CuArray, dims) = sum(f, A, dims=dims) / mapreduce(i -> size(A, i), *, unique(dims); init=1)


heat_eq_dir = joinpath(output_root,"HeatEquation")

@testset "∂_t T = K ΔT + 1, T = 0 ∈ ∂Ω, diffusion-explicit" begin
  dd = DomainDecomp(gm=1)
  dss = DomainSubSet(gm=true)
  K, maxiter, Δt = 1.0, 1000, 0.005
  grid = Grid(0.0, 1.0, 10)
  q = StateVec(( (:T, dss), ), grid, dd)
  tmp = StateVec(( (:ΔT, dss), ), grid, dd)
  rhs = deepcopy(q)
  for i in 1:maxiter
    for k in over_elems_real(grid)
      tmp[:ΔT, k] = Δ_z(q[:T, Cut(k)], grid)
      rhs[:T, k] = K*tmp[:ΔT, k] + 1
    end
    for k in over_elems(grid)
      q[:T, k] += Δt*rhs[:T, k]
    end
    apply_Dirichlet!(q, :T, grid, 0.0, Zmax())
    apply_Dirichlet!(q, :T, grid, 0.0, Zmin())
  end
  sol_analtyic = 1/2*grid.zc - grid.zc.^2/2
  sol_error = [abs(q[:T, k] - sol_analtyic[k]) for k in over_elems_real(grid)]
  @test all(sol_error .< grid.Δz)
  export_plots && export_state(q, grid, heat_eq_dir, "T.csv")
  export_plots && plot_state(q, grid, heat_eq_dir, :T, filename="T_unit_source.png")
end

@testset "∂_t T = K ΔT, T={1, 0} ∈ {z_{min}, z_{max}}, diffusion-explicit" begin
  dd = DomainDecomp(gm=1)
  dss = DomainSubSet(gm=true)
  K, maxiter, Δt = 1.0, 1000, 0.002
  grid = Grid(0.0, 1.0, 10)
  q = StateVec(( (:T, dss), ), grid, dd)
  tmp = StateVec(( (:ΔT, dss), ), grid, dd)
  rhs = deepcopy(q)
  for i in 1:maxiter
    for k in over_elems_real(grid)
      tmp[:ΔT, k] = Δ_z(q[:T, Cut(k)], grid)
      rhs[:T, k] = K*tmp[:ΔT, k]
    end
    for k in over_elems(grid)
      q[:T, k] += Δt*rhs[:T, k]
    end
    apply_Dirichlet!(q, :T, grid, 1.0, Zmin())
    apply_Dirichlet!(q, :T, grid, 0.0, Zmax())
  end
  sol_analtyic = 1 .- grid.zc
  sol_error = [abs(q[:T, k] - sol_analtyic[k]) for k in over_elems_real(grid)]
  @test all(sol_error .< grid.Δz)
  export_plots && plot_state(q, grid, heat_eq_dir, :T, filename="T1_diffExp.png")
end

@testset "∂_t T = K ΔT, T = {1, 0} ∈ {z_{min}, z_{max}}, diffusion implicit" begin
  dd = DomainDecomp(gm=1)
  dss = DomainSubSet(gm=true)
  K, maxiter, Δt = 1.0, 1000, 0.1
  grid = Grid(0.0, 1.0, 10)
  k_star1 = first_interior(grid, Zmin())
  k_star2 = first_interior(grid, Zmax())
  q = StateVec(( (:T, dss), (:a, dss), ), grid, dd)
  tmp = StateVec(( (:ΔT, dss), (:a_tau, dss), (:ρ_0, dss), (:K, dss), (:a, dss), ), grid, dd)
  rhs = deepcopy(q)
  for k in over_elems(grid)
    q[:a, k] = 1.0
    tmp[:ρ_0, k] = 1.0
    tmp[:a_tau, k] = 1.0
    tmp[:K, k] = 1.0
  end
  for i in 1:maxiter
    for k in over_elems_real(grid)
      rhs[:T, k] = 0.0
      if k==k_star1
        gv = 2*1 - q[:T, k]
        grad_T = ∇_z_flux([gv, 0], grid)
        bc_surf = ∇_z_flux([grad_T, 0], grid)
        rhs[:T, k] += bc_surf
      end
      if k==k_star2
        gv = - q[:T, k]
        grad_T = ∇_z_flux([0, gv], grid)
        bc_surf = ∇_z_flux([0, grad_T], grid)
        rhs[:T, k] += bc_surf
      end
    end
    solve_tdma!(q, rhs, tmp, :T, :ρ_0, :a, :a_tau, :K, grid, Δt)
    extrap!(q, :T, grid)
  end
  sol_analtyic = 1 .- grid.zc
  sol_numerical = [q[:T, k] for k in over_elems_real(grid)]
  sol_error = [abs(q[:T, k] - sol_analtyic[k]) for k in over_elems_real(grid)]
  @test all(sol_error .< grid.Δz)
  export_plots && plot_state(q, grid, heat_eq_dir, :T, filename= "T1_diffImp.png")
end

@testset "∂_t T = K ΔT, -K∇T = F ∈ z_{min}, T = 0 ∈ z_{max}, explicit Euler" begin
  dd = DomainDecomp(gm=1)
  dss = DomainSubSet(gm=true)
  K, maxiter, Δt = 1.0, 10, 0.001, 10
  grid = Grid(0.0, 1.0, 10)
  k_star1 = first_interior(grid, Zmin())
  k_star2 = first_interior(grid, Zmax())
  q = StateVec(( (:T1, dss), (:T2, dss), ), grid, dd)
  tmp = StateVec(( (:ΔT1, dss), (:ΔT2, dss), ), grid, dd)
  rhs = deepcopy(q)
  for i in 1:maxiter
    q_out = -K*1
    q_in = -q_out
    for k in over_elems_real(grid)
      tmp[:ΔT1, k] = Δ_z(q[:T1, Cut(k)], grid)
      tmp[:ΔT2, k] = Δ_z(q[:T2, Cut(k)], grid)
      rhs[:T1, k] = K*tmp[:ΔT1, k]
      rhs[:T2, k] = K*tmp[:ΔT2, k]
      if k==k_star1 && i>1
        rhs[:T2, k] += -∇_z_flux([q_in, 0.0], grid)
      end
    end
    for k in over_elems(grid)
      q[:T1, k] += Δt*rhs[:T1, k]
      q[:T2, k] += Δt*rhs[:T2, k]
    end
    apply_Neumann!(q   , :T1, grid, q_out, Zmin())
    apply_Dirichlet!(q , :T1, grid, 0.0  , Zmax())
    apply_Neumann!(q   , :T2, grid, 0.0  , Zmin())
    apply_Dirichlet!(q , :T2, grid, 0.0  , Zmax())
    @test all([q[:T1, k] ≈ q[:T2, k] for k in over_elems_real(grid)])
  end
end

@testset "∂_t T = ∇ • (K(z) ∇T) + 1, T = 0 ∈ ∂Ω, explicit Euler" begin
  dd = DomainDecomp(gm=1)
  dss = DomainSubSet(gm=true)
  K, maxiter, Δt = 1.0, 1000, 0.001
  grid = Grid(0.0, 1.0, 10)
  unknowns = ( (:T, dss), )
  vars = ( (:ΔT, dss), (:K_thermal, dss) )
  q = StateVec(unknowns, grid, dd)
  tmp = StateVec(vars, grid, dd)
  rhs = deepcopy(q)
  cond_thermal(z) = z > .5 ? 1 : .1
  for k in over_elems(grid)
    tmp[:K_thermal, k] = cond_thermal(grid.zc[k])
  end
  for i in 1:maxiter
    for k in over_elems_real(grid)
      tmp[:ΔT, k] = Δ_z(q[:T, Cut(k)], grid, tmp[:K_thermal, Cut(k)])
      rhs[:T, k] = K*tmp[:ΔT, k] + 1
    end
    for k in over_elems(grid)
      q[:T, k] += Δt*rhs[:T, k]
    end
    apply_Dirichlet!(q, :T, grid, 0.0, Zmax())
    apply_Dirichlet!(q, :T, grid, 0.0, Zmin())
  end
end

@testset "∂_t T = ∇ • (K(z) ∇T) + 1, T = {0, 0} ∈ {z_{min}, z_{max}}, comparison" begin
  dd = DomainDecomp(gm=1)
  dss = DomainSubSet(gm=true)
  VS, TS, F, K = 1.0, 1.0, 1.0, 1.0
  maxiter, Δt = 5000, 0.005
  grid = Grid(0.0, 1.0, 10)
  k_star1 = first_interior(grid, Zmin())
  k_star2 = first_interior(grid, Zmax())
  q = StateVec(( (:T_explicit_surf, dss), (:T_explicit_vol, dss), (:T_implicit, dss), (:a, dss), ), grid, dd)
  tmp = StateVec(( (:ΔT, dss), (:a_tau, dss), (:ρ_0, dss), (:K_thermal, dss), (:a, dss), ), grid, dd)
  rhs = deepcopy(q)
  for k in over_elems(grid)
    q[:a, k] = 1.0
    tmp[:ρ_0, k] = 1.0
    tmp[:a_tau, k] = 1.0
  end
  cond_thermal(z) = z > .5 ? 1 : .1
  for k in over_elems(grid)
    tmp[:K_thermal, k] = cond_thermal(grid.zc[k])
  end
  for i in 1:maxiter
    for k in over_elems_real(grid)
      tmp[:ΔT, k] = Δ_z(q[:T_explicit_surf, Cut(k)], grid, tmp[:K_thermal, Cut(k)])
      rhs[:T_explicit_surf, k] = tmp[:ΔT, k] + VS
      tmp[:ΔT, k] = Δ_z(q[:T_explicit_vol, Cut(k)], grid, tmp[:K_thermal, Cut(k)])
      rhs[:T_explicit_vol, k] = tmp[:ΔT, k] + VS
      rhs[:T_implicit, k] = VS
      k==k_star1 && (rhs[:T_implicit, k] += bc_source(q, grid, tmp, :T_implicit, :ρ_0, :a, :K_thermal, Zmin(), TS, Dirichlet(), DiffusionAbsorbed()))
      k==k_star2 && (rhs[:T_implicit, k] += bc_source(q, grid, tmp, :T_implicit, :ρ_0, :a, :K_thermal, Zmax(), TS, Dirichlet(), DiffusionAbsorbed()))
      k==k_star1 && (rhs[:T_explicit_vol, k] += bc_source(q, grid, tmp, :T_explicit_vol, :ρ_0, :a, :K_thermal, Zmin(), TS, Dirichlet(), DiffusionAbsorbed()))
      k==k_star2 && (rhs[:T_explicit_vol, k] += bc_source(q, grid, tmp, :T_explicit_vol, :ρ_0, :a, :K_thermal, Zmax(), TS, Dirichlet(), DiffusionAbsorbed()))
    end

    for k in over_elems_real(grid)
      q[:T_explicit_surf, k] += Δt*rhs[:T_explicit_surf, k]
      q[:T_explicit_vol, k] += Δt*rhs[:T_explicit_vol, k]
    end
    solve_tdma!(q, rhs, tmp, :T_implicit, :ρ_0, :a, :a_tau, :K_thermal, grid, Δt)

    apply_Dirichlet!(q, :T_explicit_surf, grid, TS, Zmax())
    apply_Dirichlet!(q, :T_explicit_surf, grid, TS, Zmin())

    assign_ghost!(q, :T_explicit_vol, grid, 0.0, Zmin())
    assign_ghost!(q, :T_explicit_vol, grid, 0.0, Zmax())

    assign_ghost!(q, :T_implicit, grid, 0.0, Zmin())
    assign_ghost!(q, :T_implicit, grid, 0.0, Zmax())
  end
  extrap!(q, :T_explicit_vol, grid, TS, Zmin())
  extrap!(q, :T_explicit_vol, grid, TS, Zmax())
  extrap!(q, :T_implicit, grid, TS, Zmin())
  extrap!(q, :T_implicit, grid, TS, Zmax())

  @test q[:T_explicit_vol, Dual(k_star1)][1] ≈ TS
  @test q[:T_explicit_vol, Dual(k_star2)][2] ≈ TS
  @test all([q[:T_explicit_vol, k] ≈ q[:T_explicit_surf, k] for k in over_elems_real(grid)])
  @test all([q[:T_implicit, k] ≈ q[:T_explicit_surf, k] for k in over_elems_real(grid)])

  export_plots && plot_state(q, grid, heat_eq_dir, :T_explicit_surf, filename="T_varK_explicit_surf.png")
  export_plots && plot_state(q, grid, heat_eq_dir, :T_explicit_vol, filename="T_varK_explicit_vol.png")
  export_plots && plot_state(q, grid, heat_eq_dir, :T_implicit, filename="T_varK_implicit.png")
end

@testset "∂_t T = ∇ • (K(z) ∇T) + VS, -K(z) ∇T = 1 ∈ z_{min}, T = 0 ∈ z_{max}, comparison" begin
  dd = DomainDecomp(gm=1)
  dss = DomainSubSet(gm=true)
  VS, TS, F = 1.0, 0.0, 2.0
  IC = 0.0
  maxiter, Δt = 10000, 0.0001
  grid = Grid(0.0, 1.0, 20)
  k_star1 = first_interior(grid, Zmin())
  k_star2 = first_interior(grid, Zmax())
  q = StateVec(( (:T_explicit_surf, dss), (:T_explicit_vol, dss), (:T_implicit, dss), (:a, dss), ), grid, dd)
  tmp = StateVec(( (:ΔT, dss), (:a_tau, dss), (:ρ_0, dss), (:K_thermal, dss), ), grid, dd)
  rhs = deepcopy(q)
  for k in over_elems(grid)
    q[:a, k] = 1.0
    tmp[:ρ_0, k] = 1.0
    tmp[:a_tau, k] = 1.0
  end
  cond_thermal(z) = 5+3*cos(6*π*z)
  for k in over_elems(grid)
    tmp[:K_thermal, k] = cond_thermal(grid.zc[k])
    q[:T_explicit_surf, k] = IC
    q[:T_explicit_vol, k] = IC
    q[:T_implicit, k] = IC
  end
  K_surf = tmp[:K_thermal, Dual(k_star1)][1]
  apply_Neumann!(q, :T_explicit_surf, grid, F/K_surf, Zmin())
  apply_Dirichlet!(q, :T_explicit_surf, grid, TS, Zmax())
  for i in 1:maxiter
    assign_ghost!(q, :T_explicit_vol, grid, 0.0, Zmin())
    assign_ghost!(q, :T_explicit_vol, grid, 0.0, Zmax())
    for k in over_elems_real(grid)
      tmp[:ΔT, k] = Δ_z_dual(q[:T_explicit_surf, Cut(k)], grid, tmp[:K_thermal, Dual(k)])
      rhs[:T_explicit_surf, k] = tmp[:ΔT, k] + VS

      tmp[:ΔT, k] = Δ_z_dual(q[:T_explicit_vol, Cut(k)], grid, tmp[:K_thermal, Dual(k)])
      rhs[:T_explicit_vol, k] = tmp[:ΔT, k] + VS
      k==k_star1 && (rhs[:T_explicit_vol, k] += bc_source(q, grid, tmp, :T_explicit_vol, :ρ_0, :a, :K_thermal, Zmin(), F, Neumann(), DiffusionAbsorbed()))
      k==k_star2 && (rhs[:T_explicit_vol, k] += bc_source(q, grid, tmp, :T_explicit_vol, :ρ_0, :a, :K_thermal, Zmax(), TS, Dirichlet(), DiffusionAbsorbed()))

      rhs[:T_implicit, k] = VS

      k==k_star1 && (rhs[:T_implicit, k] += bc_source(q, grid, tmp, :T_implicit, :ρ_0, :a, :K_thermal, Zmin(), F, Neumann(), DiffusionAbsorbed()))
      k==k_star2 && (rhs[:T_implicit, k] += bc_source(q, grid, tmp, :T_implicit, :ρ_0, :a, :K_thermal, Zmax(), TS, Dirichlet(), DiffusionAbsorbed()))
    end

    for k in over_elems_real(grid)
      q[:T_explicit_surf, k] += Δt*rhs[:T_explicit_surf, k]
      q[:T_explicit_vol, k] += Δt*rhs[:T_explicit_vol, k]
    end
    solve_tdma!(q, rhs, tmp, :T_implicit, :ρ_0, :a, :a_tau, :K_thermal, grid, Δt)

    apply_Neumann!(q, :T_explicit_surf, grid, F/K_surf, Zmin())
    apply_Dirichlet!(q, :T_explicit_surf, grid, TS, Zmax())

    apply_Neumann!(q, :T_explicit_vol, grid, F/K_surf, Zmin())
    apply_Dirichlet!(q, :T_explicit_vol, grid, TS, Zmax())

    apply_Neumann!(q, :T_implicit, grid, F/K_surf, Zmin())
    apply_Dirichlet!(q, :T_implicit, grid, TS, Zmax())
  end

  apply_Neumann!(q, :T_explicit_vol, grid, F/K_surf, Zmin())
  extrap!(q, :T_explicit_vol, grid, TS, Zmax())

  apply_Neumann!(q, :T_implicit, grid, F/K_surf, Zmin())
  extrap!(q, :T_implicit, grid, TS, Zmax())

  @test q[:T_explicit_vol, Dual(k_star2)][2] ≈ TS
  @test q[:T_implicit, Dual(k_star2)][2] ≈ TS

  export_plots && plot_state(q, grid, heat_eq_dir, :T_explicit_surf, filename="T_cosK_explicit_surf.png")
  export_plots && plot_state(q, grid, heat_eq_dir, :T_explicit_vol, filename="T_cosK_explicit_vol.png")
  export_plots && plot_state(q, grid, heat_eq_dir, :T_implicit, filename="T_cosK_implicit.png")

  @test all([abs(q[:T_explicit_vol, k] - q[:T_explicit_surf, k]) < 10*Δt for k in over_elems(grid)])
  @test all([abs(q[:T_explicit_vol, k] - q[:T_explicit_surf, k]) < 10*Δt for k in over_elems(grid)])
  @test all([abs(q[:T_implicit, k]     - q[:T_explicit_surf, k]) < 10*Δt for k in over_elems(grid)])
end


using OrdinaryDiffEq, ModelingToolkit, Test, SparseArrays

N = 3
xyd_brusselator = range(0,stop=1,length=N)
brusselator_f(x, y, t) = (((x-0.3)^2 + (y-0.6)^2) <= 0.1^2) * (t >= 1.1) * 5.
limit(a, N) = ModelingToolkit.ifelse(a == N+1, 1, ModelingToolkit.ifelse(a == 0, N, a))
function brusselator_2d_loop(du, u, p, t)
  A, B, alpha, dx = p
  alpha = alpha/dx^2
  @inbounds for I in CartesianIndices((N, N))
    i, j = Tuple(I)
    x, y = xyd_brusselator[I[1]], xyd_brusselator[I[2]]
    ip1, im1, jp1, jm1 = limit(i+1, N), limit(i-1, N), limit(j+1, N), limit(j-1, N)
    du[i,j,1] = alpha*(u[im1,j,1] + u[ip1,j,1] + u[i,jp1,1] + u[i,jm1,1] - 4u[i,j,1]) +
                B + u[i,j,1]^2*u[i,j,2] - (A + 1)*u[i,j,1] + brusselator_f(x, y, t)
    du[i,j,2] = alpha*(u[im1,j,2] + u[ip1,j,2] + u[i,jp1,2] + u[i,jm1,2] - 4u[i,j,2]) +
                A*u[i,j,1] - u[i,j,1]^2*u[i,j,2]
    end
end

# Test with tuple parameters
p = (3.4, 1., 10., step(xyd_brusselator))

function init_brusselator_2d(xyd)
  N = length(xyd)
  u = zeros(N, N, 2)
  for I in CartesianIndices((N, N))
    x = xyd[I[1]]
    y = xyd[I[2]]
    u[I,1] = 22*(y*(1-y))^(3/2)
    u[I,2] = 27*(x*(1-x))^(3/2)
  end
  u
end

u0 = init_brusselator_2d(xyd_brusselator)
prob_ode_brusselator_2d = ODEProblem(brusselator_2d_loop,
                                        u0,(0.,11.5),p)
sys = modelingtoolkitize(prob_ode_brusselator_2d)

# test sparse jacobian pattern only.
prob = ODEProblem(sys, u0, (0, 11.5), sparse=true, jac=false)
JP = prob.f.jac_prototype
@test findnz(Symbolics.jacobian_sparsity(map(x->x.rhs, equations(sys)), states(sys)))[1:2] == findnz(JP)[1:2]

# test sparse jacobian
prob = ODEProblem(sys, u0, (0, 11.5), sparse=true, jac=true)
@test findnz(calculate_jacobian(sys, sparse=true))[1:2] == findnz(prob.f.jac_prototype)[1:2]

# test when not sparse
prob = ODEProblem(sys, u0, (0, 11.5), sparse=false, jac=true)
@test prob.f.jac_prototype == nothing

prob = ODEProblem(sys, u0, (0, 11.5), sparse=false, jac=false)
@test prob.f.jac_prototype == nothing

# test when u0 is nothing
f = DiffEqBase.ODEFunction(sys, u0=nothing, sparse=true, jac=true)
@test findnz(f.jac_prototype)[1:2] == findnz(JP)[1:2]
@test eltype(f.jac_prototype) == Float64

f = DiffEqBase.ODEFunction(sys, u0=nothing, sparse=true, jac=false)
@test findnz(f.jac_prototype)[1:2] == findnz(JP)[1:2]
@test eltype(f.jac_prototype) == Float64

# test when u0 is not Float64
u0 = similar(init_brusselator_2d(xyd_brusselator), Float32)
prob_ode_brusselator_2d = ODEProblem(brusselator_2d_loop,
                                        u0,(0.,11.5),p)
sys = modelingtoolkitize(prob_ode_brusselator_2d)

prob = ODEProblem(sys, u0, (0, 11.5), sparse=true, jac=false)
@test eltype(prob.f.jac_prototype) == Float32

prob = ODEProblem(sys, u0, (0, 11.5), sparse=true, jac=true)
@test eltype(prob.f.jac_prototype) == Float32


mutable struct DiagonalPreconditioner{T, S<:AbstractVector{T}} <: AbstractPreconditioner
    D::S
end
function EmptyDiagonalPreconditioner(A::AbstractMatrix)
    D = zeros(eltype(A), size(A,1))
    return DiagonalPreconditioner{eltype(D), typeof(D)}(D)
end
function DiagonalPreconditioner(A::AbstractMatrix)
    D = diag(A, 0)
    return DiagonalPreconditioner{eltype(D), typeof(D)}(D)
end
function diag!(D, A::AbstractMatrix)
    length(D) == size(A,1) == size(A,2) || throw("D and A sizes are not compatible.")
    @inbounds @simd for i in 1:length(D)
        D[i] = A[i,i]
    end
    return
end
function UpdatePreconditioner!(D::DiagonalPreconditioner, K::AbstractMatrix)
    diag!(D.D, K)
    return D
end
function ldiv!(y, C::DiagonalPreconditioner, b)
    @inbounds @simd for i in 1:length(C.D)
        y[i,:] .= view(b, i, :) ./ C.D[i]
    end
    return y
end
function (\)(C::DiagonalPreconditioner, b)
    y = zeros(b)
    @inbounds @simd for i in 1:length(C.D)
        y[i,:] .= view(b, i, :) / C.D[i]
    end
    return y
end


import ImageFiltering: imfilter, KernelFactors.gaussian

"""
    GaussianBlur <: ColorOperation

Description
--------------

Blurs an image using a Gaussian filter.

Usage
--------------

    GaussianBlur(k, [σ])

Arguments
--------------

- **`k`** : `Integer` or `AbstractVector` of `Integer` that denote
    the kernel size. It must be an odd positive number.
- **`σ`** : Optional. `Real` or `AbstractVector` of `Real` that denote the
    standard deviation. It must be a positive number.
    Defaults to `0.3 * ((k - 1) / 2 - 1) + 0.8`.

Examples
--------------

```
using Augmentor
img = testpattern()

# use exactly k=3 and σ=1.0
augment(img, GaussianBlur(3, 1.0))

# pick k and σ randomly from the specified ranges
augment(img, GaussianBlur(3:2:7, 1.0:0.1:2.0))
```
"""
struct GaussianBlur{K <: AbstractVector, S <: AbstractVector} <: ColorOperation
    k::K
    σ::S

    function GaussianBlur(k::K, σ::S) where {K <: AbstractVector{<:Integer},
                                             S <: AbstractVector{<:Real}}
        minimum(k) > 0 || throw(ArgumentError("Kernel size must be positive: $(k)"))
        minimum(σ) > 0 || throw(ArgumentError("σ must be positive: $(σ)"))
        new{K, S}(k, σ)
    end
end

# The default value for σ is taken from Albumentations
GaussianBlur(k) = GaussianBlur(k, 0.3 * ((k - 1) / 2 - 1) + 0.8)
GaussianBlur(k, σ) = GaussianBlur(vectorize(k), vectorize(σ))

randparam(op::GaussianBlur, img) = (safe_rand(op.k), safe_rand(op.σ))

@inline supports_eager(::Type{<:GaussianBlur}) = true

function applyeager(op::GaussianBlur, img::AbstractArray, (k, σ))
    n = ndims(img)
    kernel = gaussian(ntuple(_->σ, n), ntuple(_->k, n))
    return imfilter(img, kernel)
end

function showconstruction(io::IO, op::GaussianBlur)
    print(io, typeof(op).name.name, '(', op.k, ", ", op.σ,')')
end

function Base.show(io::IO, op::GaussianBlur)
    if get(io, :compact, false)
        print(io, "GaussianBlur with k=$(op.k) and σ=$(op.σ)")
    else
        print(io, "Augmentor.")
        showconstruction(io, op)
    end
end


# This file is part of RLinearAlgebra.jl

module ProceduralTestLSVRProjStd

using Test, RLinearAlgebra, LinearAlgebra, Random

@testset "LSVR Projection Standard -- Procedural" begin

    # Supertype and aliases
    @test supertype(LinSysVecRowProjStd) == LinSysVecRowProjection
    @test Kaczmarz == LinSysVecRowProjStd
    @test ART == LinSysVecRowProjStd

    # Verify that residual projection is zero
    Random.seed!(1010)

    A = rand(10,5)
    x = rand(5)
    b = A * x

    rsub = Kaczmarz()

    for i = 1:10
        @test let

            z = rand(5)
            RLinearAlgebra.rsubsolve!(rsub, z, (A[i,:], b[i]), i)

            dot(z, A[i,:]) ≈ b[i]
        end
    end

end

end # End module


function cmp(models::Vector{SampleModel}, type::Symbol)
    mnames = String[]
    lps = Matrix{Float64}[]
    for m in models
        nt = read_samples(m, :namedtuple)
        if :log_lik in keys(nt)
            append!(mnames, String.([m.name]))
            append!(lps, [Matrix(nt.log_lik')])
        else
            @warn "Model $(m.name) does not produce a log_lik matrix."
        end
    end
    (lps, type, mnames)
end

function cmp(m::Vector{Matrix{Float64}}, type::Symbol;
    mnames=String[])

    df = DataFrame()
    loo = Vector{Float64}(undef, length(m))
    loos = Vector{Vector{Float64}}(undef, length(m))
    pk = Vector{Vector{Float64}}(undef, length(m))
    for i in 1:length(m)
        loo[i], loos[i], pk[i] = psisloo(m[i])
    end
    ind = sortperm([-2loo[i][1] for i in 1:length(m)])

    mods = m[ind]
    loo = loo[ind]
    loos = loos[ind]
    pk = pk[ind]

    if length(mnames) > 0
        df.models = String.(mnames[ind])
    end

    df.PSIS = round.([-2loo[i] for i in 1:length(loo)], digits=1)
    df.SE = round.([sqrt(size(m[i], 2)*var2(-2loos[i])) for i in 1:length(m)],
        digits=2)

    dloo = zeros(length(m))
    for i in 2:length(m)
        dloo[i] = df[i, :PSIS] - df[1, :PSIS]
    end
    df.dPSIS = round.(dloo, digits=1)

    dse = zeros(length(m))
    for i in 2:length(m)
        diff = 2(loos[1] .- loos[i])
        dse[i] = √(length(loos[i]) * var2(diff))
    end
    df.dSE = round.(dse, digits=2)

    ps = zeros(length(m))
    lppds = zeros(length(m))
    for j in 1:length(m)
        n_sam, n_obs = size(mods[j])
        pd = zeros(length(m), n_obs)
        pd[j, :] = [var2(mods[j][:,i]) for i in 1:n_obs]
        ps[j] = sum(pd[j, :])
    end
    df.pPSIS = round.(ps, digits=2)

    weights = ones(length(m))
    sumval = sum([exp(-0.5df[i, :PSIS]) for i in 1:length(m)])
    for i in 1:length(m)
        weights[i] = exp(-0.5df[i, :PSIS])/sumval
    end
    df.weight = round.(weights, digits=2)
    df
end

compare([m8_1s, m8_2s], :waic) |> display

lps, type, mnames = cmp([m8_1s, m8_2s], :psis)

cmp(lps, type; mnames) |> display

md"
```

> PSIS(m8.1)
       PSIS    lppd  penalty  std_err
1 -188.6114 94.3057 2.755778 13.35391

> PSIS(m8.2)
       PSIS     lppd  penalty  std_err
1 -252.1136 126.0568 4.272533 15.29773

> compare( m8.1 , m8.2 )
       WAIC   SE dWAIC   dSE pWAIC weight
m8.2 -252.3 15.3   0.0    NA   4.3      1
m8.1 -188.7 13.3  63.5 15.15   2.7      0

> compare( m8.1 , m8.2, func=PSIS )
Some Pareto k values are high (>0.5). Set pointwise=TRUE to inspect individual points.
       PSIS    SE dPSIS   dSE pPSIS weight
m8.2 -252.1 15.28   0.0    NA   4.3      1
m8.1 -188.7 13.40  63.4 15.16   2.7      0
```
"

l, ls, pk = psisloo(m8_1s)
sum(ls)


# White noise
"
White noise process

$FIELDS
"
struct WhiteNoise{T} <: TimeSeriesModel{T}
  "Standard deviation"
  σ::T
end

WhiteNoise() = WhiteNoise(1.0)

function Base.show(io::IO,wn::WhiteNoise)
  σstr = @sprintf "%0.3g" wn.σ
  print(io,"y = ϵ | σ = $σstr")
end

function Base.show(io::IO,::MIME"text/plain",wn::WhiteNoise)
  print(io,"Scalar white noise\n",wn)
end

order(m::WhiteNoise) = 0

Distributions.std(x::WhiteNoise) = x.σ

"
$SIGNATURES

log pdf of white noise model.
If `has_z` is `true`, the first observations is ignored, as is done for SDE(p)
processes in `logpdf` when `z` is provided. In this way, closer agreement
with the log probability of an SDE(p) process is achieved.
"
function Distributions.logpdf(proc::WhiteNoise,data,has_z=false)
  σ2 = proc.σ^2
  istart = has_z ? 2 : 1
  y = [x.y for x in data]
  Nl = length(y)
  lp = -1/2*(sum(y.^2)/σ2+Nl*log(σ2)+Nl*log(2π))
  return lp
end

function Base.rand(rng::AbstractRNG,proc::WhiteNoise,t::AbstractVector)
  data = map(x->Obs(x,proc.σ*randn(rng)),t)
  return data
end


using IndexedTables
using FileIO
using SurrealNumbers
out_dir = "Data/"

# subtraction table
x = convert.(SurrealFinite, [0, 1/2, 1, 2])
A = zeros(SurrealFinite, length(x), length(x))
B = zeros(Int, length(x), length(x))
for i=1:length(x)
    for j=1:length(x)
        A[i,j] = x[i] - x[j]
        B[i,j] = generation(x[i]) + generation(x[j])
     end
end
float.(A)
A
generation.(A)
# A .== canonicalise.(A)
iscanonical.(A)
B
C = dag_stats.(A)
n = [C[i,j].nodes for i=1:size(C)[1], j=1:size(C)[2] ]

# subtraction table 2
ns1 = [4; 4; 4; 4; 4; 4; 1; 2; 3]
ns2 = [0; 1; 2; 3; 4; 5; 1; 2; 3]
n = length(ns1)
m = [10; 10; 10; 10; 10; 10; 10; 10; 10]
x = Array{SurrealFinite,1}(n)
t = Array{Float64,1}(n)
bytes = Array{Float64,1}(n)
gctime = Array{Float64,1}(n)
c_plus = Array{Integer,1}(n)
c_times = Array{Integer,1}(n)
c_created = Array{Integer,1}(n)
c_equals = Array{Integer,1}(n)
c_leq = Array{Integer,1}(n)
i = 0
for i=1:n
    println(ns1[i], " - ", ns2[i])
    t[i] = 0
    b = 0
    g = 0
    mem = 0
    val = 0
    for j=1:m[i]
        clearcache()
        val, time, b, g, mem = @timed convert(SurrealFinite, ns1[i]) - convert(SurrealFinite, ns2[i])
        t[i] += time
    end
    x[i] = val
    bytes[i] = b
    gctime[i] = g

    c_plus[i] = Count['+']
    c_times[i] = Count['*']
    c_created[i] = Count['c']
    c_equals[i] = Count['=']
    c_leq[i] = Count['≦']
end
s = dag_stats.( x )
no = [s[i].nodes for i=1:length(s)]
ed = [s[i].edges for i=1:length(s)]
p = [s[i].paths for i=1:length(s)]
d = (e+1) ./ no
g = [s[i].generation for i=1:length(s)]
t = t./m[1:n]
bytes = bytes

t1 = table(@NT(i=0:n-1,
               n1=ns1,
               n2=ns2,
               nodes=no,
               edges=ed,
               paths=p,
               density=d,
               generation=g,
               time=t,
               bytes=bytes,
               adds=c_plus,
               prods=c_times,
               created=c_created,
               equals=c_equals,
               leq=c_leq);
           pkey = [:i])

save("$(out_dir)subtraction_tables.csv", t1)

# examples with diagrams
a21 = convert(SurrealFinite, 1) - convert(SurrealFinite, 1)
file = "$(out_dir)subtraction_ex_1.dot"
FID = open(file, "w")
surreal2dag(FID, a21)
close(FID)
run(`dot -Tpdf -O $file`)

a22 = convert(SurrealFinite, 2) - convert(SurrealFinite, 2)
file = "$(out_dir)subtraction_ex_2.dot"
FID = open(file, "w")
surreal2dag(FID, a22)
close(FID)
run(`dot -Tpdf -O $file`)

a22 = convert(SurrealFinite, 2) - convert(SurrealFinite, 2)
file = "$(out_dir)subtraction_ex_2.dot"
FID = open(file, "w")
surreal2dag(FID, a22)
close(FID)
run(`dot -Tpdf -O $file`)

a23 = convert(SurrealFinite, 1/2) - convert(SurrealFinite, 1/2)
file = "$(out_dir)subtraction_ex_3.dot"
FID = open(file, "w")
surreal2dag(FID, a23)
close(FID)
run(`dot -Tpdf -O $file`)

a41 = convert(SurrealFinite, 4) - convert(SurrealFinite, 1)
file = "$(out_dir)subtraction_ex_41.dot"
FID = open(file, "w")
surreal2dag(FID, a41)
close(FID)
run(`dot -Tpdf -O $file`)
run(`dot -Tsvg -O $file`)

a42 = convert(SurrealFinite, 4) - convert(SurrealFinite, 2)
file = "$(out_dir)subtraction_ex_42.dot"
FID = open(file, "w")
surreal2dag(FID, a42)
close(FID)
run(`dot -Tpdf -O $file`)
run(`dot -Tsvg -O $file`)

a43 = convert(SurrealFinite, 4) - convert(SurrealFinite, 3)
file = "$(out_dir)subtraction_ex_43.dot"
FID = open(file, "w")
surreal2dag(FID, a43)
close(FID)
run(`dot -Tpdf -O $file`)
run(`dot -Tsvg -O $file`)

a5 = convert(SurrealFinite, 3) - (convert(SurrealFinite, 1)-convert(SurrealFinite, 1))
file = "$(out_dir)subtraction_ex_5.dot"
FID = open(file, "w")
surreal2dag(FID, a5)
close(FID)
run(`dot -Tpdf -O $file`)
run(`dot -Tsvg -O $file`)


using Test
using SumOfSquares
using DynamicPolynomials

function bivariate_quadratic_test(optimizer,
                                  config::MOIT.TestConfig,
                                  cone::SumOfSquares.PolyJuMP.PolynomialSet)
    atol = config.atol
    rtol = config.rtol

    model = _model(optimizer)

    @variable(model, α)

    @polyvar x
    cref = @constraint(model, x^2 + α*x + 1 in cone)

    # See https://github.com/JuliaOpt/MathOptInterface.jl/issues/676
    @objective(model, Max, α + 1)
    optimize!(model)

    @test termination_status(model) == MOI.OPTIMAL
    @test objective_value(model) ≈ 3.0 atol=atol rtol=rtol

    @test primal_status(model) == MOI.FEASIBLE_POINT
    @test value(α) ≈ 2.0 atol=atol rtol=rtol

    p = gram_matrix(cref)
    @test getmat(p) ≈ ones(2, 2) atol=atol rtol=rtol
    @test p.x == [x, 1]

    @test dual_status(model) == MOI.FEASIBLE_POINT
    μ = dual(cref)
    @test μ isa AbstractMeasure{Float64}
    @test length(moments(μ)) == 3
    a = moment_value.(moments(μ))
    @test a[2] ≈ -1.0  atol=atol rtol=rtol
    @test a[1] + a[3] ≈ 2.0  atol=atol rtol=rtol
    @test monomial.(moments(μ)) == [x^2, x, 1]

    ν = moment_matrix(cref)
    @test getmat(ν) ≈ [a[1] a[2]
                       a[2] a[3]] atol=atol rtol=rtol
    @test ν.x == [x, 1]

    S = SumOfSquares.SOSPolynomialSet{
        SumOfSquares.FullSpace, typeof(cone), SumOfSquares.MonomialBasis,
        Monomial{true},MonomialVector{true},Tuple{}
    }
    @test list_of_constraint_types(model) == [(Vector{AffExpr}, S)]
    test_delete_bridge(
        model, cref, 1,
        ((MOI.VectorOfVariables, MOI.Nonnegatives, 0),
         (MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}, 0),
         (MOI.VectorAffineFunction{Float64},
          SumOfSquares.PolyJuMP.ZeroPolynomialSet{
              SumOfSquares.FullSpace, SumOfSquares.MonomialBasis,
              Monomial{true}, MonomialVector{true}},
          0)))
end
sos_bivariate_quadratic_test(optimizer, config)   = bivariate_quadratic_test(optimizer, config, SOSCone())
sdsos_bivariate_quadratic_test(optimizer, config) = bivariate_quadratic_test(optimizer, config, SDSOSCone())
dsos_bivariate_quadratic_test(optimizer, config)  = bivariate_quadratic_test(optimizer, config, DSOSCone())