using DataStructures

abstract type DiscretizeMODL <: DiscretizationAlgorithm end

struct DiscretizeMODL_Optimal <: DiscretizeMODL end
struct DiscretizeMODL_Greedy <: DiscretizeMODL end
struct DiscretizeMODL_PostGreedy <: DiscretizeMODL
    max_bin_count :: Int
    DiscretizeMODL_PostGreedy(max_bin_count::Integer = 0) = new(round(Int, max_bin_count))
end

function binedges(
    alg             :: DiscretizeMODL,
    data_continuous :: AbstractArray{N},
    data_class      :: AbstractArray{I},
    ) where {N<:Real, I<:Integer}

    @assert(length(data_continuous)==length(data_class))
    p = sortperm(data_continuous)
    data_continuous = data_continuous[p]
    data_class = data_class[p]

    if isa(alg, DiscretizeMODL_Optimal)
        return optimal_result(data_continuous,data_class)
    elseif isa(alg, DiscretizeMODL_Greedy)
        return greedy_merge(data_continuous,data_class)
    else
        @assert(isa(alg, DiscretizeMODL_PostGreedy))
        return post_greedy_result(data_continuous,data_class,alg.max_bin_count)
    end
end

function MODL_value2_oneintval(
    continuous_ij  :: AbstractArray{T},
    class_value_ij :: AbstractArray{S},
    class_uniq     :: AbstractArray{S},
    ) where {T<:AbstractFloat, S<:Integer}

    n = length(continuous_ij)
    J = length(class_uniq)

    first_part = lfactorial(n+J-1) - lfactorial(J-1) - lfactorial(n)
    second_part = lfactorial(n)
    # Note(Yi-Chun): lfactorial(n) is log(n!)

    for J_index = 1:J
            n_J = count(x->x==class_uniq[J_index],class_value_ij)
            second_part = second_part - lfactorial(n_J)
    end

    return first_part+second_part
end

function optimal_result(
    continuous      :: AbstractArray{T},
    discrete_target :: AbstractArray{S}
    ) where {T<:AbstractFloat, S<:Integer}

    n = length(continuous)

    Disc_ijk_retval = Array{Array}(undef,n,n)
    Disc_ijk_MODL_value = Array{AbstractFloat}(undef,n,n)
    for j = 1:n
            for k = 1:n
                    if k > j
                            Disc_ijk_MODL_value[j,k] = Inf
                            Disc_ijk_retval[j,k] = [0]
                    end
            end
    end

    class_uniq = unique(discrete_target)

    for k = 1:n
        for j = k:n
            if k == 1
                Disc_ijk_retval[j,k] = [j]
                Disc_ijk_MODL_value[j,k] = MODL_value2_oneintval(
                                      continuous[1:j], discrete_target[1:j],class_uniq)
            else
                MODL_value = Inf
                select_intval = 0

                for i = 1:j-1
                    second_MODL = MODL_value2_oneintval(
                                  continuous[i+1:j],discrete_target[i+1:j],class_uniq)
                    current_MODL = Disc_ijk_MODL_value[i,k-1] + second_MODL
                    if current_MODL < MODL_value
                        MODL_value = current_MODL
                        select_intval = i
                    end
                end

                Disc_ijk_retval[j,k] = append!(copy(Disc_ijk_retval[select_intval,k-1]),[j])
                Disc_ijk_MODL_value[j,k] = MODL_value
            end

        end
    end

    full_length_k_intval = copy(Disc_ijk_MODL_value[n,:])

    for l = 1:n
        full_length_k_intval[l] += lfactorial(n+l-1) - lfactorial(l-1) - lfactorial(n)
    end

    desired_intval_number = argmin(full_length_k_intval)

    bin_edges_index = Disc_ijk_retval[n,desired_intval_number]
    bin_edges = append!([1],bin_edges_index)

    bin_edge_value = Vector{AbstractFloat}(undef,length(bin_edges))

    for index = 1 : length(bin_edge_value)
        if index == 1
                bin_edge_value[index] = continuous[1]
        elseif index == length(bin_edge_value)
                bin_edge_value[index] = continuous[end]
        else
                bin_edge_value[index] = 0.5*(continuous[bin_edges[index]]+
                                             continuous[bin_edges[index]+1])
        end
    end
    return bin_edge_value
end

function merge_adj_intval(
    A_distr :: AbstractArray{S}, # Note(Yi-Chun): Distribution of class values in A
    B_distr :: AbstractArray{S}  # Note(Yi-Chun): Distribution of class values in B
    ) where {S<:Integer}

    @assert(length(A_distr) == length(B_distr))
    J = length(B_distr)

    n_A = sum(A_distr)
    n_B = sum(B_distr)

    Delta = lfactorial(n_A+n_B+J-1) + lfactorial(J-1) - lfactorial(n_A+J-1) - lfactorial(n_B+J-1)

    for j = 1:J
        n_A_J = A_distr[j]
        n_B_J = B_distr[j]
        Delta = Delta - lfactorial(n_A_J+n_B_J) + lfactorial(n_A_J) + lfactorial(n_B_J)
    end
    Delta
end

function greedy_merge_index(
    continuous  :: AbstractArray{T},
    class_value :: AbstractArray{S},
    ) where {T<:AbstractFloat,S<:Integer}

    @assert(length(continuous) == length(class_value))

    n = length(continuous)
    all_class_value = unique(class_value)
    J = length(all_class_value)

    if n==1
        return [1]
    end

    pq = DataStructures.PriorityQueue()
    # Note(Yi-Chun): The PriorQueue sorts differences of MODL after merging of adjacent intervals
    distr_in_intval = Dict()
    # Note(Yi-Chun): This dictionary contains the distribution in each interval
    adj_for_intval = Dict()
    # Note(Yi-Chun): This dictionary contains the two adjacent interval for interval i

    MODL = log(n) + lfactorial(2n-1) - lfactorial(n-1) - lfactorial(n) + n*log(J-1)
    # Note(Yi-Chun): Intial MODL when each attribute is treated as a single bin

    for i = 1:n-1
        A_Distrib_in_intval = zeros(Int64,J)
        A_value_index = something(findfirst(isequal(class_value[i]), all_class_value), 0)

        A_Distrib_in_intval[A_value_index] = 1
        # Note(Yi-Chun): Set up distribution in each bin as [0,0,..0,1,0,..0]
        distr_in_intval[i] = A_Distrib_in_intval
        # Note(Yi-Chun): Here we label intervals by the index of their first boundary.

        # Note(Yi-Chun): For each interval we assign 4 adjacent intervals for it.
        #                Two in the front and two in the behind.
        if i==1
            adj_for_intval[i] = [2,3]
        elseif i==2
            adj_for_intval[i] = [1,3,4]
        elseif i==n-1
            adj_for_intval[i] = [i-2,i-1,i+1]
        else
            adj_for_intval[i] = [i-2,i-1,i+1,i+2]
        end

        B_Distrib_in_intval = zeros(Int64,J)
        B_value_index = something(findfirst(isequal(class_value[i+1]), all_class_value), 0)

        B_Distrib_in_intval[B_value_index] = 1
        if i == (n-1) # Note(Yi-Chun): The adjacents of the last interval are n-2 and n-1
            distr_in_intval[n] = B_Distrib_in_intval
            adj_for_intval[i+1] = [i-1,i]
        end

        X = merge_adj_intval(A_Distrib_in_intval,B_Distrib_in_intval)
        pq[i+1] = X
    end

    n_intval = n
    first_could_remove = 2
    last_could_remove = n

    while n_intval > 4
        min_diff = DataStructures.peek(pq)[2]
        first_term_in_delta = log((n_intval-1)/(n+n_intval-1))

        if (min_diff + first_term_in_delta) < 0
            MODL = MODL + min_diff + first_term_in_delta
            removed = DataStructures.dequeue!(pq)
            if removed == first_could_remove
                # Note(Yi-Chun): Interval labels: y=1,(removed),z,w,s
                @assert(length(adj_for_intval[removed]) == 3)
                @assert(adj_for_intval[removed][1] == 1)
                z = adj_for_intval[removed][2]
                w = adj_for_intval[removed][3]
                first_could_remove = z
                # Note(Yi-Chun): Merge the distribution
                Merge_distr = distr_in_intval[1] + distr_in_intval[removed]
                # Note(Yi-Chun): Update the Delta value for the only one adjacent
                pq[z] = merge_adj_intval(Merge_distr,distr_in_intval[z])
                # Note(Yi-Chun): Update the structure
                distr_in_intval[1] = Merge_distr

                # Note(Yi-Chun): Update the adjacent structure
                adj_for_intval[1][1] = z
                adj_for_intval[1][2] = w
                s = adj_for_intval[z][4]
                adj_for_intval[z] = [1,w,s]
                adj_for_intval[w][1] = 1

                # Note(Yi-Chun): Delete the information of removed one
                delete!(adj_for_intval,removed)
                delete!(distr_in_intval,removed)

            elseif removed == last_could_remove
                @assert(length(adj_for_intval[removed]) == 2)
                x = adj_for_intval[removed][1]
                y = adj_for_intval[removed][2]
                last_could_remove = y
                Merge_distr = distr_in_intval[y] + distr_in_intval[removed]
                pq[y] = merge_adj_intval(distr_in_intval[x],Merge_distr)
                distr_in_intval[y] = Merge_distr
                pop!(adj_for_intval[x])
                pop!(adj_for_intval[y])
                delete!(adj_for_intval,removed)
                delete!(distr_in_intval,removed)

            else
                # Note(Yi-Chun): Interval labels: x,y,(removed),z,w
                # Note(Yi-Chun): Merge the distribution
                y = adj_for_intval[removed][2]
                Merge_distr = distr_in_intval[y] + distr_in_intval[removed]
                # Note(Yi-Chun): Update the Delta value for two adjacents
                x = adj_for_intval[removed][1]
                pq[y] = merge_adj_intval(distr_in_intval[x],Merge_distr)
                z= adj_for_intval[removed][3]
                pq[z] = merge_adj_intval(Merge_distr,distr_in_intval[z])
                # Note(Yi-Chun): Update the structure
                distr_in_intval[y] = Merge_distr

                # Note(Yi-Chun): Update the adjacent structure
                if length(adj_for_intval[removed]) == 3
                    adj_for_intval[y][3] = z
                    pop!(adj_for_intval[y])
                    adj_for_intval[x][end] = z
                    adj_for_intval[z][1] = x
                    adj_for_intval[z][2] = y
                else
                    w = adj_for_intval[removed][4]
                    adj_for_intval[x][end] = z
                    adj_for_intval[y][end-1] = z
                    adj_for_intval[y][end] = w
                    adj_for_intval[z][1] = x
                    adj_for_intval[z][2] = y
                    adj_for_intval[w][1] = y
                end

                # Note(Yi-Chun): Delete the information of removed one
                delete!(adj_for_intval,removed)
                delete!(distr_in_intval,removed)
            end
            n_intval = n_intval - 1
        else
            break
        end
    end
    # Note(Yi-Chun): Here we manage the case that when bin number goes under 4,
    #                since we have to use different label for intervals

    if n_intval > 4
        binedges = sort(collect(keys(distr_in_intval)))
        return binedges
    else
        binedges = sort(collect(keys(distr_in_intval)))
        while ((DataStructures.peek(pq)[2] + log((n_intval-1)/(n+n_intval-1)))<0) && (n_intval>2)
            MODL = MODL + DataStructures.peek(pq)[2] + log((n_intval-1)/(n+n_intval-1))
            n_intval = n_intval-1
            removed = DataStructures.dequeue!(pq)
            prior = binedges[something(findfirst(isequal(removed), binedges), 0)-1]

            Merge_distr = distr_in_intval[removed] + distr_in_intval[prior]
            distr_in_intval[prior] = Merge_distr
            delete!(adj_for_intval,removed)
            delete!(distr_in_intval,removed)
            binedges = sort(collect(keys(distr_in_intval)))

            for i = 1 : n_intval-1
                X = distr_in_intval[binedges[i]]
                Y = distr_in_intval[binedges[i+1]]
                pq[binedges[i+1]] = merge_adj_intval(X,Y)
            end
        end

        if (DataStructures.peek(pq)[2] + log((n_intval-1)/(n+n_intval-1)))<0
            return [1]
        else
            return binedges
        end
    end
end

function greedy_merge(
    continuous  :: AbstractArray{T},
    class_value :: AbstractArray{S},
    ) where {T<:AbstractFloat,S<:Integer}

    binedges = greedy_merge_index(continuous,class_value)
    append!(binedges,[length(continuous)])
    n = length(binedges)
    bin_edges = Vector{Float64}(undef,n)
    for i = 2: n-1
        binedges[i] = binedges[i]-1
    end

    for j = 1: n
        if j == 1
            bin_edges[1] = continuous[binedges[1]]
        elseif j==n
            bin_edges[n] = continuous[binedges[n]]
        else
            bin_edges[j] = 0.5*(continuous[binedges[j]] + continuous[binedges[j]+1])
        end
    end

    return bin_edges
end

# Note(Yi-Chun): The following 4 functions are used in post_greedy methods.

function methods_split(continuous,class_value,distr,i,j,uniq_class,I)

    N = length(continuous)
    continuous_ij = continuous[i:j-1]
    class_ij = class_value[i:j-1]
    distr_ij = distr[i]

    n = length(continuous_ij)
    J = length(uniq_class)
    distr_A = zeros(Int64,J)

    current_MODL = Inf
    current_split = 0
    current_distr_A = zeros(Int64,J)
    current_distr_B = distr_ij

    for n_A_end = 1:n-1
        distr_A[something(findfirst(isequal(class_ij[n_A_end]), uniq_class),0)] += 1
        distr_B = distr_ij - distr_A

        Del = lfactorial(n_A_end+J-1)+lfactorial(n-n_A_end+J-1)-lfactorial(n+J-1)-lfactorial(J-1)

        for j = 1:J
            Del += lfactorial(distr_A[j]+distr_B[j])-lfactorial(distr_A[j])-lfactorial(distr_B[j])
        end

        if Del < current_MODL
            current_MODL = Del
            current_split = n_A_end
            current_distr_A = deepcopy(distr_A)
            current_distr_B = deepcopy(distr_B)
        end
    end

    current_MODL += log((N+I)/(I))
    current_split += i

    return (current_MODL,current_split,current_distr_A,current_distr_B)
end

function methods_merge(
    A_distr   :: AbstractArray{S},
    B_distr   :: AbstractArray{S},
    N,
    I,
    ) where {S<:Integer}

    # Note(Yi-Chun): This function is similar to the previous function merge_adj_intval.
    #                The only different part is the difference of MODL value after merging.
    @assert(length(A_distr) == length(B_distr))
    J = length(B_distr)
    n_A = sum(A_distr)
    n_B = sum(B_distr)
    Delta = lfactorial(n_A+n_B+J-1) + lfactorial(J-1) - lfactorial(n_A+J-1) - lfactorial(n_B+J-1)
    for j = 1:J
        n_A_J = A_distr[j]
        n_B_J = B_distr[j]
        Delta = Delta - lfactorial(n_A_J+n_B_J) + lfactorial(n_A_J) + lfactorial(n_B_J)
    end
    Delta += log((I-1)/(N+I-1))
    return Delta
end

function methods_merge_split(continuous,class_value,distr,i,j,k,uniq_class,I)

    continuous_ij = continuous[i:j-1]
    class_ij = class_value[i:j-1]
    continuous_jk = continuous[j:k-1]
    class_jk = class_value[j:k-1]
    distr_ij = distr[i]
    distr_jk = distr[j]
    N = length(continuous)
    Del_1 = methods_merge(distr_ij,distr_jk,N,I)

    pseudo_dict = Dict()
    pseudo_dict[i] = distr_ij + distr_jk
    split_result = methods_split(continuous,class_value,pseudo_dict,i,k,uniq_class,I-1)

    Del = Del_1 + split_result[1]
    return (Del,split_result[2],split_result[3],split_result[4])
end

function methods_merge2_split(continuous,class_value,distr,i,j,k,l,uniq_class,I)

    N = length(continuous)
    continuous_ij = continuous[i:j-1]
    continuous_jk = continuous[j:k-1]
    continuous_kl = continuous[k:l-1]
    class_ij = class_value[i:j-1]
    class_jk = class_value[j:k-1]
    class_kl = class_value[k:l-1]
    distr_ij = distr[i]
    distr_jk = distr[j]
    distr_kl = distr[k]
    Del_1 = methods_merge(distr_ij,distr_jk,N,I)
    Del_2 = methods_merge(distr_ij + distr_jk,distr_kl,N,I-1)
    pseudo_dict = Dict()
    pseudo_dict[i] = distr_ij + distr_jk + distr_kl
    split_result = methods_split(continuous,class_value,pseudo_dict,i,l,uniq_class,I-2)

    Del = Del_1 + Del_2 + split_result[1]

    return (Del,split_result[2],split_result[3],split_result[4])
end

function PQ_methods_updata(continuous,class_value,uniq_class,binedges,distr,post_pq,effects,index,method)
    # Note(Yi-Chun): (index,method) is the index of the edge that we want to update its method in priorityqueue
    I = length(binedges) - 1
    # Note(Yi-Chun): Since in post_greedy part, we add the index N+1 as ending of bin edges for convenience
    N = length(continuous)

    if method == "M"
        A_distr = distr[binedges[index]]
        B_distr = distr[binedges[index+1]]
        post_pq[(binedges[index],method)] = methods_merge(A_distr,B_distr,N,I)
    elseif method == "S"
        method_result = methods_split(continuous,class_value,distr,binedges[index],binedges[index+1],uniq_class,I)
        post_pq[(binedges[index],method)] = method_result[1]
        effects[(binedges[index],method)] = method_result[2:end]
    elseif method == "MS"
        method_result = methods_merge_split(continuous,class_value,distr,binedges[index],
                                            binedges[index+1],binedges[index+2],uniq_class,I)
        post_pq[(binedges[index],method)] = method_result[1]
        effects[(binedges[index],method)] = method_result[2:end]
    else
        method_result = methods_merge2_split(continuous,class_value,distr,binedges[index],binedges[index+1],
                                             binedges[index+2],binedges[index+3],uniq_class,I)
        post_pq[(binedges[index],method)] = method_result[1]
        effects[(binedges[index],method)] = method_result[2:end]
    end
end

function remove_methods_on_index(index,methods_on_index,post_pq,effects)
    meth = methods_on_index[index]
    for i = 1: length(meth)
        post_pq[(index,meth[i])] = Inf
        delete!(effects,(index,meth[i]))
    end
end

insert_and_dedup!(v::Vector, x) = (splice!(v, searchsorted(v,x), [x]); v)

function uncondi_greedy_merge_index(
    continuous  :: AbstractArray{T},
    class_value :: AbstractArray{S}
    ) where {T<:AbstractFloat,S<:Integer}
    # Note(Yi-Chun): This function processes the unconditional bin merge, Adapted from the function greedy_merge
    @assert(length(continuous) == length(class_value))

    n = length(continuous)
    all_class_value = unique(class_value)
    J = length(all_class_value)
    length_of_Greedy_merge = length(greedy_merge_index(continuous,class_value))

    pq = DataStructures.PriorityQueue()
    distr_in_intval = Dict()
    adj_for_intval = Dict()

    MODL = log(n) + lfactorial(2n-1) - lfactorial(n-1) - lfactorial(n) + n*log(J-1)

    total_distrib = zeros(Int64,J)

    for i = 1:n-1
        A_Distrib_in_intval = zeros(Int64,J)
        A_value_index = something(findfirst(isequal(class_value[i]), all_class_value), 0)
        A_Distrib_in_intval[A_value_index] = 1
        distr_in_intval[i] = A_Distrib_in_intval

        if i==1
            adj_for_intval[i] = [2,3]
        elseif i==2
            adj_for_intval[i] = [1,3,4]
        elseif i==n-1
            adj_for_intval[i] = [i-2,i-1,i+1]
        else
            adj_for_intval[i] = [i-2,i-1,i+1,i+2]
        end

        total_distrib += A_Distrib_in_intval

        B_Distrib_in_intval = zeros(Int64,J)
        B_value_index = something(findfirst(isequal(class_value[i+1]), all_class_value), 0)
        B_Distrib_in_intval[B_value_index] = 1
        if i == (n-1)
            distr_in_intval[n] = B_Distrib_in_intval
            adj_for_intval[i+1] = [i-1,i]
            total_distrib += B_Distrib_in_intval
        end

        X = merge_adj_intval(A_Distrib_in_intval,B_Distrib_in_intval)
        pq[i+1] = X
    end

    one_bin_distrib = Dict()
    one_bin_distrib[1] = total_distrib

    if n==1|J==1
        best_for_far = [1]
        best_for_far_MODL = -Inf
        return (best_for_far,one_bin_distrib)
    end

    n_intval = n
    first_could_remove = 2
    last_could_remove = n

    # Note(Yi-Chun): This best_so_far part will record the best dicretization we meet while merging unconditionally
    best_so_far = collect(1:n)
    best_so_far_MODL = Inf
    best_adj = Dict()
    best_distr = Dict()

    while n_intval > 4
        min_diff = DataStructures.peek(pq)[2]
        first_term_in_delta = log((n_intval-1)/(n+n_intval-1))

        MODL = MODL + min_diff + first_term_in_delta
        removed = DataStructures.dequeue!(pq)
        if removed == first_could_remove
            @assert(length(adj_for_intval[removed]) == 3)
            @assert(adj_for_intval[removed][1] == 1)
            z = adj_for_intval[removed][2]
            w = adj_for_intval[removed][3]
            first_could_remove = z
            # Note(Yi-Chun): Merge the distribution
            Merge_distr = distr_in_intval[1] + distr_in_intval[removed]
            # Note(Yi-Chun): Update the Delta value for the only one adjacent
            pq[z] = merge_adj_intval(Merge_distr,distr_in_intval[z])
            # Note(Yi-Chun): Update the structure
            distr_in_intval[1] = Merge_distr
            # Note(Yi-Chun): Update the pointers
            adj_for_intval[1][1] = z
            adj_for_intval[1][2] = w
            s = adj_for_intval[z][4]
            adj_for_intval[z] = [1,w,s]
            adj_for_intval[w][1] = 1
            # Note(Yi-Chun): Delete the information of removed one
            delete!(adj_for_intval,removed)
            delete!(distr_in_intval,removed)
        elseif removed == last_could_remove
            @assert(length(adj_for_intval[removed]) == 2)
            x = adj_for_intval[removed][1]
            y = adj_for_intval[removed][2]
            last_could_remove = y
            Merge_distr = distr_in_intval[y] + distr_in_intval[removed]
            pq[y] = merge_adj_intval(distr_in_intval[x],Merge_distr)
            distr_in_intval[y] = Merge_distr
            pop!(adj_for_intval[x])
            pop!(adj_for_intval[y])
            delete!(adj_for_intval,removed)
            delete!(distr_in_intval,removed)
        else
            # Note(Yi-Chun): Interval labels: x,y,(removed),z,w
            # Note(Yi-Chun): Merge the distribution
            y = adj_for_intval[removed][2]
            Merge_distr = distr_in_intval[y] + distr_in_intval[removed]
            # Note(Yi-Chun): Update the Delta value for two adjacents
            x = adj_for_intval[removed][1]
            pq[y] = merge_adj_intval(distr_in_intval[x],Merge_distr)
            z= adj_for_intval[removed][3]
            pq[z] = merge_adj_intval(Merge_distr,distr_in_intval[z])
            # Note(Yi-Chun): Update the structure
            distr_in_intval[y] = Merge_distr
            if length(adj_for_intval[removed]) == 3
                adj_for_intval[y][3] = z
                pop!(adj_for_intval[y])
                adj_for_intval[x][end] = z
                adj_for_intval[z][1] = x
                adj_for_intval[z][2] = y
            else
                w = adj_for_intval[removed][4]
                adj_for_intval[x][end] = z
                adj_for_intval[y][end-1] = z
                adj_for_intval[y][end] = w
                adj_for_intval[z][1] = x
                adj_for_intval[z][2] = y
                adj_for_intval[w][1] = y
            end
            # Note(Yi-Chun): Delete the information of removed one
            delete!(adj_for_intval,removed)
            delete!(distr_in_intval,removed)
        end
        n_intval = n_intval - 1

        if (MODL < best_so_far_MODL) && (n_intval<length_of_Greedy_merge+5)
            best_so_far_MODL = MODL
            best_so_far = deepcopy(collect(keys(distr_in_intval)))
            best_distr = deepcopy(distr_in_intval)
        end
    end

    if n_intval > 4
        binedges = sort(collect(keys(distr_in_intval)))
        return binedges
    else
        binedges = sort(collect(keys(distr_in_intval)))
        while (n_intval>2)
            MODL = MODL + DataStructures.peek(pq)[2] + log((n_intval-1)/(n+n_intval-1))
            n_intval = n_intval-1
            removed = DataStructures.dequeue!(pq)
            prior = binedges[something(findfirst(isequal(removed), binedges),0)-1]
            Merge_distr = distr_in_intval[removed] + distr_in_intval[prior]
            distr_in_intval[prior] = Merge_distr
            delete!(adj_for_intval,removed)
            delete!(distr_in_intval,removed)
            binedges = sort(collect(keys(distr_in_intval)))

            for i = 1 : n_intval-1
                X = distr_in_intval[binedges[i]]
                Y = distr_in_intval[binedges[i+1]]
                pq[binedges[i+1]] = merge_adj_intval(X,Y)
            end

            if MODL < best_so_far_MODL
                best_so_far_MODL = MODL
                best_so_far = deepcopy(collect(keys(distr_in_intval)))
                best_distr = deepcopy(distr_in_intval)
            end
        end

        if (DataStructures.peek(pq)[2] + log((n_intval-1)/(n+n_intval-1)))<0
            best_for_far = [1]
            return (best_for_far,one_bin_distrib)
        else
            return (sort(best_so_far),best_distr)
        end
    end
end

function post_greedy_index(continuous,class_value,upper_bound = 0)

    N = length(continuous)
    uniq_class = unique(class_value)
    J = length(uniq_class)

    # Note(Yi-Chun): Check whether we have restriction on bin size
    if upper_bound == 0 # Note(Yi-Chun): the case for no upper bound of bin size
        (bins,distr) = uncondi_greedy_merge_index(continuous,class_value)
        @assert(bins == sort(collect(keys(distr))))
        I = length(bins)
        methods_on_index = Dict()
        for i = 1:I
            methods_on_index[bins[i]] = ["S"]
        end
        append!(bins,[N+1]) # Note(Yi-Chun): for convenience
    else # Note(Yi-Chun): the case for having upper bound of bin size
        I = upper_bound-1
        reminder = N%I
        new_edge = 1
        bins = [1]
        for i = 1 : I-1
            new_edge += div(N,I)
            if reminder > 0
                new_edge += 1
                reminder -= 1
            end
            append!(bins,[new_edge])
        end
        append!(bins,[N+1])
        distr = Dict()
        for bin_index = 1 : I
            distr[bins[bin_index]] = Vector{Int64}(undef,J)
            for j_index = 1 : J
                n_J = count(x->x==uniq_class[j_index],class_value[bins[bin_index]:bins[bin_index+1]-1])
                distr[bins[bin_index]][j_index] = n_J
            end
        end
        methods_on_index = Dict()
        for i = 1:I
            methods_on_index[bins[i]] = ["S"]
        end
    end

        # Note(Yi-Chun): Add methods on each bin into PriorityQueue
        post_pq = DataStructures.PriorityQueue()
        # Note(Yi-Chun): keys in post_pq are represented as: (x,y), where x is the index of bin, y are methods.
        # Note(Yi-Chun): If "S"->split, "M"->merge, "MS"->merge_split, "MMS"->merge_merge_split

        effects = Dict()
        # Note(Yi-Chun): This dictionary record the effects of the methods

        # Note(Yi-Chun): Here we initialize the methods in PriorityQueue
        # Note(Yi-Chun): Methods on bin[1],... ,bin[I-2]. Each bin has 4 methods
        for i = 1 : I-2
            # Note(Yi-Chun): splitting in bin[i]
            method_result = methods_split(continuous,class_value,distr,bins[i],bins[i+1]-1,uniq_class,I)
            post_pq[(bins[i],"S")] = method_result[1]
            effects[(bins[i],"S")] = method_result[2:end]
            # Note(Yi-Chun): merge bin[i] with bin[i+1]
            method_result = methods_merge(distr[bins[i]],distr[bins[2]],N,I)
            post_pq[(bins[i],"M")] = method_result
            append!(methods_on_index[bins[i]],["M"])
            # Note(Yi-Chun): merge_split on bin[i]
            method_result = methods_merge_split(continuous,class_value,distr,
                                                bins[i],bins[i+1],bins[i+2],uniq_class,I)
            post_pq[(bins[i],"MS")] = method_result[1]
            effects[(bins[i],"MS")] = method_result[2:end]
            append!(methods_on_index[bins[i]],["MS"])
            # merge_merge_split on bin[i]
            method_result = methods_merge2_split(continuous,class_value,distr,
                                                 bins[i],bins[i+1],bins[i+2],bins[i+3],uniq_class,I)
            post_pq[(bins[i],"MMS")] = method_result[1]
            effects[(bins[i],"MMS")] = method_result[2:end]
            append!(methods_on_index[bins[i]],["MMS"])
        end

        # Note(Yi-Chun): Methods on bin[I-1]. Only Split, Merge, Merge_Split.
        if I>1
            method_result = methods_split(continuous,class_value,distr,bins[I-1],bins[I]-1,uniq_class,I)
            post_pq[(bins[I-1],"S")] = method_result[1]
            effects[(bins[I-1],"S")] = method_result[2:end]
            method_result = methods_merge(distr[bins[I-1]],distr[bins[I]],N,I)
            post_pq[(bins[I-1],"M")] = method_result
            append!(methods_on_index[bins[I-1]],["M"])
            method_result = methods_merge_split(continuous,class_value,distr,bins[I-1],bins[I],N+1,uniq_class,I)
            post_pq[(bins[I-1],"MS")] = method_result[1]
            effects[(bins[I-1],"MS")] = method_result[2:end]
            append!(methods_on_index[bins[I-1]],["MS"])
        end

        # Note(Yi-Chun): Method on bin[I]. Only Split.
        method_result = methods_split(continuous,class_value,distr,bins[I],N+1,uniq_class,I)
        post_pq[(bins[I],"S")] = method_result[1]
        effects[(bins[I],"S")] = method_result[2:end]

        # Note(Yi-Chun): Here we start the dequeue process
        current_I = I
        time_dequeue = 0
        while (DataStructures.peek(post_pq)[2] < -0.00001 )
            MODL = DataStructures.peek(post_pq)[2]
            time_dequeue += 1

            (removed_index,removed_method) = DataStructures.dequeue!(post_pq)
            removed_posit_bins = searchsorted(bins,removed_index)[1]
            # Note(Yi-Chun): Ittells the position of removed_index in binedges

            # Note(Yi-Chun): This part updates the binedges and distribution of each bin after the dequeue
            if removed_method == "M"
                merged_index = bins[removed_posit_bins+1]
                distr[removed_index] = distr[removed_index] + distr[merged_index]
                delete!(distr,merged_index)
                remove_methods_on_index(merged_index,methods_on_index,post_pq,effects)
                current_I -= 1
                splice!(bins,removed_posit_bins+1)
                methods_on_index[merged_index] = ["S"]

            elseif removed_method == "S"
                effect = effects[(removed_index,removed_method)]
                delete!(effects,(removed_index,removed_method))
                distr[removed_index] = effect[2]
                distr[effect[1]] = effect[3]
                methods_on_index[effect[1]] = ["S"]
                current_I += 1
                insert_and_dedup!(bins, effect[1])

            elseif removed_method == "MS"
                delete!(distr,bins[removed_posit_bins+1])
                effect = effects[(removed_index,removed_method)]
                delete!(effects,(removed_index,removed_method))
                distr[removed_index] = effect[2]
                distr[effect[1]] = effect[3]
                methods_on_index[effect[1]] = ["S"]
                remove_methods_on_index(bins[removed_posit_bins+1],methods_on_index,post_pq,effects)
                methods_on_index[bins[removed_posit_bins+1]] = ["S"]
                splice!(bins,removed_posit_bins+1)
                insert_and_dedup!(bins, effect[1])

            else # Note(Yi-Chun): removed_method == "MMS"
                delete!(distr,bins[removed_posit_bins+1])
                delete!(distr,bins[removed_posit_bins+2])
                remove_methods_on_index(bins[removed_posit_bins+1],methods_on_index,post_pq,effects)
                remove_methods_on_index(bins[removed_posit_bins+2],methods_on_index,post_pq,effects)
                effect = effects[(removed_index,removed_method)]
                delete!(effects,(removed_index,removed_method))
                distr[removed_index] = effect[2]
                distr[effect[1]] = effect[3]
                methods_on_index[effect[1]] = ["S"]
                current_I -= 1
                methods_on_index[bins[removed_posit_bins+1]] = ["S"]
                methods_on_index[bins[removed_posit_bins+2]] = ["S"]
                splice!(bins,removed_posit_bins+1)
                splice!(bins,removed_posit_bins+1)
                insert_and_dedup!(bins, effect[1])
            end

            # Note(Yi-Chun): This part updates the priorityqueue and the effect of each method
            for be_modified = (removed_posit_bins - 2) : (removed_posit_bins + 1)
                 # Note(Yi-Chun): Update the method "Split"
                if (be_modified > 0)
                    PQ_methods_updata(continuous,class_value,uniq_class,bins,distr,
                                        post_pq,effects,be_modified,"S")
                    methods_on_index[bins[be_modified]] = ["S"]
                end
                # Note(Yi-Chun): Update the method "Merge"
                if (be_modified > 0) && (be_modified+1 < current_I+1)
                    PQ_methods_updata(continuous,class_value,uniq_class,bins,distr,
                                        post_pq,effects,be_modified,"M")

                    append!(methods_on_index[bins[be_modified]],["M"])
                end
                # Note(Yi-Chun): Update the method "MergeSplit"
                if (be_modified > 0) && (be_modified+1 < current_I+1)
                    PQ_methods_updata(continuous,class_value,uniq_class,bins,distr,
                                        post_pq,effects,be_modified,"MS")
                    append!(methods_on_index[bins[be_modified]],["MS"])
                end
                # Note(Yi-Chun): Update the method "MergeMergeSplit"
                if (be_modified > 0) && (be_modified+2 < current_I+1)
                    PQ_methods_updata(continuous,class_value,uniq_class,bins,distr,
                                        post_pq,effects,be_modified,"MMS")
                    append!(methods_on_index[bins[be_modified]],["MMS"])
                end
            end

            if upper_bound != 0
                if current_I >= upper_bound
                    for i = 1 : current_I
                        post_pq[(bins[i],"S")] = Inf
                    end
                end
            end
        end
        return bins
end

function post_greedy_result(continuous,class_value,upper_bound=0)

    binedges = post_greedy_index(continuous,class_value,upper_bound)
    bins = Vector{Float64}(undef,length(binedges))
    for i = 1 : length(binedges)
        if i==1
            bins[1] = continuous[1]
        elseif i == length(binedges)
            bins[i] = continuous[end]
        else
            bins[i] = 0.5*(continuous[binedges[i]-1] + continuous[binedges[i]])
        end
    end
    return bins
end


include("../src/HssMatrices.jl")
using .HssMatrices
using LinearAlgebra
using AbstractTrees
using Plots

# generate Cauchy matrix
K(x,y) = (x-y) > 0 ? 0.001/(x-y) : 2.
#K(x,y) = (x-y) != 0 ? 1/(x-y) : 1.
A = [ K(x,y) for x=-1:0.001:1, y=-1:0.001:1];
C = inv(A);

# test the simple implementation of cluster trees
m, n = size(A)
HssMatrices.setopts(leafsize=64)
rcl = bisection_cluster(1:m)
ccl = bisection_cluster(1:n)

# test compression
hssA = compress(A, rcl, ccl);
println("rel. approximation error with direct compression: ", norm(A - full(hssA))/norm(A))
println("abs. approximation error with direct compression: ", norm(A - full(hssA)))
println("hss-rank with direct compression: ", hssrank(hssA))

# test randomized compression
hssB = randcompress_adaptive(A, rcl, ccl);
println("rel. approximation error with randomized compression: ", norm(A - full(hssB))/norm(A))
println("abs. approximation error with randomized compression: ", norm(A - full(hssB)))
println("hss-rank with randomized compression: ", hssrank(hssB))

# test recompression
hssB = recompress!(copy(hssA))
println("rel. approximation error after recompression: ", norm(A - full(hssB))/norm(A))
println("abs. approximation error after recompression: ", norm(A - full(hssB)))
println("hss-rank after recompression: ", hssrank(hssB))

# test mat-vec
x = randn(size(A,2), 3);
println("rel. error in the matrix-vector products: ", norm(A*x - hssA*x)/norm(A*x))
println("abs. error in the matrix-vector products: ", norm(A*x - hssA*x))

# test the ULV based solver
b = randn(size(A,2), 5);
x = ulvfactsolve(hssA, b);
xcor = A\b;
println("rel. error in the solution of Ax = b: ", norm(x-xcor)/norm(xcor))
println("abs. error in the solution of Ax = b: ", norm(x-xcor))

# test left HSS division
Id(i,j) = Matrix{Float64}(i.*ones(length(j))' .== ones(length(i)).*j')
IdOp = LinearMap{Float64}(n, n, (y,_,x) -> x, (y,_,x) -> x, (i,j) -> Id(i,j))
hssI = randcompress(IdOp, ccl, ccl, 0)
hssC = ldiv!(hssA, hssI)
println("rel. error in the solution of AX = I: ", norm(full(hssC) - C) / norm(C) )
println("abs. error in the solution of AX = I: ", norm(full(hssC) - C) )

# test left HSS division
Id(i,j) = Matrix{Float64}(i.*ones(length(j))' .== ones(length(i)).*j')
IdOp = LinearMap{Float64}(n, n, (y,_,x) -> x, (y,_,x) -> x, (i,j) -> Id(i,j))
hssI = randcompress(IdOp, ccl, ccl, 0)
hssC = rdiv!(hssI, hssA)
println("rel. error in the solution of XA = I: ", norm(full(hssC) - C) / norm(C) )
println("abs. error in the solution of XA = I: ", norm(full(hssC) - C) )

# desequilibrate cluster trees and try again
hssA.A11 = prune_leaves!(hssA.A11)
hssI = randcompress(IdOp, ccl, ccl, 0)
hssI.A11 = prune_leaves!(hssI.A11)
hssC = ldiv!(hssA, hssI)
println("rel. error in the solution of XA = I: ", norm(full(hssC) - C) / norm(C) )
println("abs. error in the solution of XA = I: ", norm(full(hssC) - C) )

# # test computation of generators
# U1, V2 = generators(hssA, (1,2))
# A12 = A[1:hssA.m1,hssA.n1+1:end];
# println(norm(A12 - U1*hssA.B12*V2'))
# U2, V1 = generators(hssA, (2,1))
# A21 = A[hssA.m1+1:end,1:hssA.n1];|
# println(norm(A21 - U2*hssA.B21*V1'))

# #println(typeof(hssA))

# test the ULV based solver
# x = ulvfactsolve(hssA, b);

# @time x = ulvfactsolve(hssA, b);

# # test plotting
# plot = plotranks(hssA)


export nnlasso_kernel, lasso_kernel, nnlasso
using SparseArrays
using LinearAlgebra

function nnlasso(A,y,tau; x :: Vector{Float64} = tau*rand(size(A,2))/(size(A,2)*2), l = ones(size(A,2)+1)/(2*size(A,2)))
  K = Array(A'*A);
  b = Array(A'*y);
  return nnlasso_kernel(K,b,tau)
end

function nnlasso_kernel(K,b,tau; x :: Vector{Float64} = tau*rand(length(b))/(length(b)*2), l = ones(length(b)+1)/(2*length(b)))
  K = Array(K)
  b = Array(b)
  n = size(K,2)
  if (n == 0)
    return zeros(0)
  end

  function f_g_h(x)
    0.5*dot(x,K*x) - dot(x,b) ,K*x -b,K
  end
  p = OptimizationProblem(n,n+1,f_g_h,[-I; ones(1,n)], [zeros(n);tau])
  x,l = primalDualSolve(p :: OptimizationProblem, x, l, 10.0, 1e-12 , 1e-12 ,1000, 0.1 , 0.4)
  x[x.<1e-9] .= 0.0
  return x
end

function lasso_kernel(K,b,tau; x :: Vector{Float64} = tau*rand(length(b))/(length(b)*2), l = ones(2*length(b)+1)/(2*length(b)))
  K = Array(K)
  b = Array(b)
  n = size(K,2)
  if ( n == 0)
    return zeros(0)
  end
  function f_g_h(x)
    x = x[1:n]
    f,g,h = 0.5*dot(x,K*x) - dot(x,b) ,K*x -b,K
    f,[g; zeros(n)], [h zeros(n,n); zeros(n,n) zeros(n,n)]
  end
  consrt1 = [eye(n) -eye(n)]
  consrt2 = [-eye(n) -eye(n)]
  p = OptimizationProblem(2*n,2*n + 1,f_g_h,[consrt1;consrt2; [zeros(1,n) ones(1,n)]], [zeros(2*n);tau])
  x,l = primalDualSolve(p :: OptimizationProblem, [x; 1.1*abs(x)], l, 10.0, 1e-10 , 1e-10 ,1000, 0.1 , 0.4)
  x = x[1:n]
  x[abs(x).<1e-9] = 0.0
  return x
end


using CSV, DataFrames, Dates, Test, BusinessDays
using AbnormalReturns

##

df_firm = CSV.File(joinpath("data", "daily_ret.csv")) |> DataFrame
df_mkt = CSV.File(joinpath("data", "mkt_ret.csv")) |> DataFrame
df_res = CSV.File(joinpath("data", "car_results.csv")) |> DataFrame
df_mkt[!, :mkt] = df_mkt.mktrf .+ df_mkt.rf
##

data = MarketData(
    df_mkt,
    df_firm
)

##

rr = quick_reg(data[18428, Date(2019, 4) .. Date(2019, 10)], @formula(ret ~ mktrf + hml))
@test coefnames(rr) == ["(Intercept)", "mktrf", "hml"]
@test responsename(rr) == "ret"
@test nobs(rr) == 126
@test all(isapprox.(coef(rr), [-.00125105, 1.40602071, 1.19924984]; atol=.00001))
@test isapprox(r2(rr), .42340085667)
@test isapprox(adjr2(rr), .41402526)
@test dof_residual(rr) == 123
@test islinear(rr)
@test alpha(rr) == rr.coef[1]
@test beta(rr) == rr.coef[2]

##

@test_throws KeyError quick_reg(data[1, Date(2020) .. Date(2021)], @formula(ret ~ mktrf))

rr = quick_reg(data[18428, Date(2019, 4) .. Date(2019, 10)], @formula(ret ~ mktrf + hml); minobs=1000)# not enough data case

@test ismissing(coefnames(rr))
@test ismissing(coef(rr))
@test nobs(rr) == 126
@test ismissing(r2(rr))
@test ismissing(adjr2(rr))
@test alpha(rr) |> ismissing
@test beta(rr) |> ismissing

##

# the SAS code that I tested this against appears to round results to 3 significant digits

# Since these are over specific periods, I specify those as functions to make this easier

event_start(x; data=data) = advancebdays(data.calendar, tobday(data.calendar, x), -10)
event_end(x; data=data) = advancebdays(data.calendar, tobday(data.calendar, x), 10)
est_end(x; data=data) = advancebdays(data.calendar, event_start(x), -16)
est_start(x; data=data) = advancebdays(data.calendar, est_end(x), -149)

rr_market = quick_reg(
    data[df_res.permno, est_start.(df_res.event_date) .. est_end.(df_res.event_date)],
    @formula(ret ~ mktrf)
)

@test isapprox(round.(alpha.(rr_market), digits=5), df_res.alpha_market_model_)
@test isapprox(round.(beta.(rr_market), digits=3), df_res.beta_market_model)
cars = car(data[df_res.permno, event_start.(df_res.event_date) .. event_end.(df_res.event_date)], rr_market)
@test isapprox(round.(cars, sigdigits=3), df_res.car_mm)
bhars = bhar(data[df_res.permno, event_start.(df_res.event_date) .. event_end.(df_res.event_date)], rr_market)
@test isapprox(round.(bhars, sigdigits=3), df_res.bhar_mm)
stds = std.(rr_market)
vars = var.(rr_market)
@test isapprox(round.(vars, digits=10), df_res.estimation_period_variance_market_model_)

##

rr_ff = quick_reg(
    data[df_res.permno, est_start.(df_res.event_date) .. est_end.(df_res.event_date)],
    @formula(ret ~ mktrf + smb + hml)
)

cars = car(data[df_res.permno, event_start.(df_res.event_date) .. event_end.(df_res.event_date)], rr_ff)
@test isapprox(round.(cars, sigdigits=3), df_res.car_ff)
bhars = bhar(data[df_res.permno, event_start.(df_res.event_date) .. event_end.(df_res.event_date)], rr_ff)
@test isapprox(round.(bhars, sigdigits=3), df_res.bhar_ff)
stds = std.(rr_ff)
vars = var.(rr_ff)
@test isapprox(round.(vars, digits=10), df_res.estimation_period_variance_ff_model_)

##

rr_ffm = quick_reg(
    data[df_res.permno, est_start.(df_res.event_date) .. est_end.(df_res.event_date)],
    @formula(ret ~ mktrf + smb + hml + umd)
)

cars = car(data[df_res.permno, event_start.(df_res.event_date) .. event_end.(df_res.event_date)], rr_ffm)
@test isapprox(round.(cars, sigdigits=3), df_res.car_ffm)
bhars = bhar(data[df_res.permno, event_start.(df_res.event_date) .. event_end.(df_res.event_date)], rr_ffm)
@test isapprox(round.(bhars, sigdigits=3), df_res.bhar_ffm)
stds = std.(rr_ffm)
vars = var.(rr_ffm)
@test isapprox(round.(vars, digits=10), df_res.estimation_period_variance_carhart_model_)

##


import JuMP, Cbc, XLSX

# make sure u set the working directory to current files folder
data = XLSX.readxlsx("../../data/data3_6.xlsx")

#General
General = data["General"];
l = General["B1"]
num_providers = General["B2"]
num_clients = General["B3"]

#Providers
Providers = data["Providers"];
provCostString  = string("A2:A", num_providers+1)
provStartString = string("B2:B", num_providers+1)
f = Providers[provCostString]
w = Providers[provStartString]

#Providers
Clients = data["Clients"];
clientStartString  = string("A2:A", num_clients+1)
clientDurationString  = string("B2:B", num_clients+1)
s_times = Clients[clientStartString]
d = Clients[clientDurationString]

#Distances
Distances = data["Distances"];
t = Distances[Distances.dimension]

# need to convert to floats to use 'sortperm' below
w = convert(Array{Float64,1}, w[:,1])
s_times = convert(Array{Float64,1}, s_times[:,1])
d = convert(Array{Float64,1}, d[:,1])
f = convert(Array{Float64,1}, f[:,1])

client_sorted = sortperm(s_times)  # client start times sorted (in index)

arr = [[f[i], w[i]] for i in 1:length(w)]
provider_sorted = sortperm(f) # providers cost and start time sorted

# heuristic function that computes feasible routes (but likely not optimal, i.e. minimum cost)
function heuristics_result(provider_sorted, client_sorted, w, s_times, t, d)
    past_client = zeros(0)
    num_served = zeros(num_providers, 2)
    num_served[:,1] = provider_sorted

    for m in provider_sorted  # for each index in sorted provider, and then for each client
        current_location = 1  # 1 is at depot(14HS) currently
        for n in client_sorted
            # only proceed if this client has not been assigned to a provider yet.
            if n ∉ past_client
                currentTime_provider = w[m]
                currentTime_client = s_times[n]
                travelTime = t[current_location, n+1] # n+1 is nth client in the t matrix

                #time to travel from provider's current location to the client plus the time when the provider is free.
                time_to_arrive = travelTime + currentTime_provider

                #Time to serve client "n".
                duration = d[n]
                #Time finished serving client "n".
                finish_time = currentTime_client + duration

                #Ensure the provider can arrive in time (at or earlier than service start times)
                # Must return back the origin (14HS) by hour l at latest.
                if currentTime_client >= time_to_arrive && finish_time <= l - t[n+1, 1]
                        println("Provider: ", m, " is assigned to client: ", n)
                        num_served[findfirst(isequal(m), num_served[:,1]), 2] += 1
                        w[m] = finish_time # update the time of provider m
                        current_location = n+1  # n+1 is nth client in the t-matrix
                        append!(past_client, n) # record service provided for client n
                end
            end
        end
        # end the algorithm if we have served all clients. so O(N*<M)?
        if length(past_client) == length(client_sorted)
            break
        end
    end
    println("\n")
    return convert(Array{Int,2}, num_served), convert(Array{Int,1}, past_client)
end

function compute_cost(num_served)
    cost = 0
    for i in 1:length(num_served[:,2])
        if num_served[:,2][i] > 0
            cost += f[num_served[:,1][i]]
        end
    end
    return cost
end

function is_feasible(past_client)
    return length(past_client) == length(client_sorted)
end

function get_heuristics_path(num_served, past_client, k)
    path = zeros(1)
    k_ind = findfirst(isequal(k), num_served[:,1])
    if k_ind == 1
        for i in 1:num_served[1,2]
            append!(path, past_client[i])
        end
    end
    if k_ind != 1
        start_ind = sum(num_served[1:k_ind-1,2]) + 1
        for i in start_ind:start_ind+num_served[k_ind,2]-1
            append!(path, past_client[i])
        end
    end
    return convert(Array{Int,1}, append!(path, zeros(1))) # convert to integer array
end

num_served, past_client = heuristics_result(provider_sorted, client_sorted, copy(w), s_times, t, d)

num_served

past_client

is_feasible(past_client)
compute_cost(num_served)

# k is the provider# you want to get the path for
k=4
path = get_heuristics_path(num_served, past_client, k)
path += convert(Array{Int,1}, ones(length(get_heuristics_path(num_served, past_client, k))))
println(path)


file = open(f->read(f, String), "day_12.txt")

regex = r"([A-Z])(\d+)"

input =[x.captures for x in eachmatch(regex,file)]

#converting to cartezian format
@NamedTuple{x::Int64,y::Int64,angle::Int64}

cartezinaize_comand = Dict()
cartezinaize_comand["N"] = (north(i) = (x=0,y=i,angle=0))
cartezinaize_comand["S"] = (south(i) = (x=0,y=-i,angle=0))
cartezinaize_comand["E"] = (east(i) =  (x=i,y=0,angle=0))
cartezinaize_comand["W"] = (west(i) =  (x=-i,y=0,angle=0))
cartezinaize_comand["L"] = (left(i) =  (x=0,y=0,angle=i))
cartezinaize_comand["R"] = (right(i) = (x=0,y=0,angle=-i))

angle_to_card = Dict([(0,"E"),(90,"N"),(180,"W"),(270,"S")])
function change_foward(letter,angle)
    if (letter != "F")
        return letter
    else

        angle = (angle % 360)#overflow angles
        angle += 360*(angle < 0)#fixing negative angle
        return angle_to_card[angle]
    end
end

#start moving
pos = (x=0,y=0,angle=0)

for (letter,value) in input
    #change F to equivalent command
    letter = change_foward(letter,copy(pos.angle))

    movement = cartezinaize_comand[letter](parse(Int64,value))

    newX = pos.x + movement.x
    newY = pos.y + movement.y
    newAngle = pos.angle + movement.angle
    global pos = (x=newX,y=newY,angle=newAngle)
end

print(abs(pos.x)+abs(pos.y))

##Question 2

matrix_rotation = Dict()
matrix_rotation[0] = [1 0;0 1]
matrix_rotation[90] = [0 -1;1 0]
matrix_rotation[180] = [-1 0;0 -1]
matrix_rotation[270] = [0 1;-1 0]

waypoint = [10;1]
pos = (x=0,y=0,angle=0)
for (letter,value) in input
    value = parse(Int64,value)
    if letter == "F"
        newX = pos.x + (waypoint[1] * value)
        newY = pos.y + (waypoint[2] * value)
        global pos = (x=newX,y=newY,angle=0)
    elseif letter == "N"
        global waypoint += [0;value]
    elseif letter == "S"
        global waypoint += [0;-value]
    elseif letter == "E"
        global waypoint += [value;0]
    elseif letter == "W"
        global waypoint += [-value;0]
    else #angle commands
        angle = if (letter == "L") value else -value end# left or right
        angle = (angle % 360)#overflow angles
        angle += 360*(angle < 0)#fixing negative angle

        #applying angle with matrix of rotation
        global waypoint = matrix_rotation[angle] * waypoint
    end

    print("\n $(pos)")
    print("\n $(waypoint)")
end

print("\n$(abs(pos.x)+abs(pos.y))")


using NearestNeighbors
using StaticArrays
"""
    Fkmnn(xy, r, window, N)

Return empty space function based on Kaplan-Meier estimator.
"""
function Fkmnn(xy, window, r, N)
    if length(xy) == 0
        return zeros(length(r))
    end
    xy = @. SVector(getx(xy), gety(xy))
    (x1,x2), (y1, y2) = window.x, window.y
    tree = KDTree(xy)
    sx = (x2-x1)/N
    Ny = round(Int, N*(y2-y1)/(x2-x1))
    sy = (y2-y1)/Ny
    rx = range(x1 + 0.5sx, x2 - 0.5sx, length=N)
    ry = range(y1 + 0.5sy, y2 - 0.5sy, length=Ny)
    grid = vec([SVector(x, y) for x in rx, y in ry])
    _, D = knn(tree, grid, 1)
    D = first.(D)
    B = [bdist(window, x) for x in grid]
    ind = D .<= B
    T = min.(D, B)
    km(ind, T, r)
end


type FullModel <: GP
	d::Int64
	k::Kernel
	lik::Likelihood
	X::Array{Float64, 2}
	y::Array{Float64, 1}
	y_raw::Array{Float64, 1}
	K::Array{Float64, 2} #Covariance matrix
	Cfull::Array{Float64, 2}
    C::Union(Array{Float64, 2}, Base.LinAlg.Cholesky{Float64}) #Cholesky decomp of C
    CinvY::Array{Float64, 1} # C^-1  * Y
    mu::Float64
    std::Float64
end

function FullModel(k::Kernel, lik::Likelihood)
	FullModel(0, k, lik, Array(Float64, (0, 0)), Array(Float64, (0,)),
		Array(Float64, (0,)),
		Array(Float64, (0,0)), Array(Float64, (0, 0)), Array(Float64, (0,0)),
		Array(Float64, (0,)), 0.0, 1.0)
end

function observe!{T <: FloatingPoint}(m::FullModel, X::Array{T}, y::Array{T})
	if size(m.X) == (0,0)
		m.d = size(X, 2)
		m.X = X
		m.y_raw = y
	else
		@assert m.d == size(X, 2) "Ensure X has the same dimension as the model!"
		m.X = vcat(m.X, X)
		m.y_raw = vcat(m.y_raw, y)
	end
	m.mu = mean(m.y_raw)
	m.std = std(m.y_raw)
	m.y = (m.y_raw .- m.mu) / (m.std + eps())
	refresh_cache!(m)
end

function predict(m::FullModel, Xp)
	Kpd = kernel(m.k, Xp, m.X)
	d = size(Xp, 2)
	z = zeros(1, d)
	c = kernel(m.k, z, z)[1]
	mn = (Kpd * m.CinvY)
	variance = (c .- sum(Kpd .* (m.C \ Kpd')', 2))
	mn * (m.std + eps()) .+ m.mu, variance * (m.std + eps())^2
end

function refresh_cache!(m::FullModel)
	m.K = kernel(m.k, m.X)
	sigma2 = exp(2m.lik.log_sigma)
	m.Cfull = m.K + eye(m.K) * sigma2
	try
		m.C = cholfact(m.Cfull)
	catch e
		m.C = m.Cfull
	end
	m.CinvY = vec(m.C \ m.y)
end


# PeriodicLoss

using LaTeXStrings, LossFunctions, Measures, Latexify

Plots.plot(PeriodicLoss(1.5), leg=:outertopright, c=:olive, linewidth=2, legendtitle= L"L\left( r \right) = 1 - cos(\frac {2r\pi}{c})", size=(850,499))
Plots.plot!(PeriodicLoss(1), c=:orangered4, linewidth=2.5)
Plots.plot!(PeriodicLoss(.5), c=:springgreen4, linewidth=2.5)


using StaticArrays

struct Cylinder{T}
    midpoint::SVector{3,T}
    rot::SMatrix{3,3,T,9}
    r::T
    h::T
end

struct CylinderGenerator{Fr,Fh}
    r::Fr
    h::Fh
end

function (g::CylinderGenerator)(midpoint, rot)
    Cylinder(midpoint, rot, g.r(), g.h())
end

volume(shape::Cylinder) = π * shape.r^2 * h

function set_midpoint(shape::Cylinder, midpoint)
    return Cylinder(midpoint, t.rot, t.r, t.h)
end

function bounding_box(shape::Cylinder)
    lower = (shape.midpoint + shape.rot * SA[0, 0, -shape.h/2]) .- shape.r
    upper = (shape.midpoint + shape.rot * SA[0, 0,  shape.h/2]) .+ shape.r
    return lower, upper
end

function inside(shape::Cylinder, p0)
    p = shape.rot' * (p0 - shape.midpoint)
    return -shape.h/2 < p[3] < shape.h/2 && shape.r^2 < (p[1]^2 + p[2]^2)
end


"""
    SeisMute(d; <keyword arguments>) -> Array{Real,2}

Mute trace noise before the first arrival.
The muting function depends on the offset between source and receiver and the velocity of the first layer as
t_mute = t_0 + offset/v1.

# Arguments
- `d`: Array of input traces.

# Keyword arguments
- `offset=[0.]`: Vector of distances between source and receiver.
- `tmute=0.`: initial muting time
- `vmute=1500.`: Velocity of the first layer.
- `taper=0.1`: Taper of the muting function
- `dt::Real=0.004`: sampling interval in secs.

"""
function SeisMute(in;offset=[0.],tmute=0.,vmute=1500.,taper=0.1,dt=0.004)

	vec= size(in)
	in = reshape(in,vec[1],:)
	nt,nx = size(in)
	out = copy(in)
	for it = 1:nt
		for ix = 1:nx
			t = sqrt(tmute^2 + (offset[ix]/vmute).^2)
			if ((it-1)*dt <= t )
					if ((it-1)*dt > t - taper)
						out[it,ix] *= 1 - (t - (it-1)*dt)/taper
					else
						out[it,ix] = 0.
					end
			end
		end
	end
	out = reshape(out,vec)
	return out
end

"""
    SeisMute(in,out,parameters; <keyword arguments>)

Mutes noise before first arrival of traces in a file. Saves the muted traces to an output file.
The muting function depends on the offset between source and receiver and the velocity of the first layer as
t_mute = t_0 + offset/v1.

# Arguments
- `in::String`: Input file - Seis format.
- `out::String`: Output file - Seis format.
- `parameters` : list of the keyword arguments for the function SeisMute.

# Keyword arguments
- `group="gather"` : Options are all, some or gather
- `key=["imx","imy"]` : Defines type of gather
- `itrace=1` : Initial trace number
- `ntrace=10000` : Total number of traces to process at once

# Example
```
julia> param = Dict(:tmute=>0.0, :vmute=>10000, :taper=>0.05,:dt=>0.01)
julia> SeisMute(filein,fileout, param,group="some")
```

"""
function SeisMute(in::String,out::String,parameters;group="gather",key=["imx","imy"],itrace=1,ntrace=10000)

	tmute = get(parameters,:tmute,0.)
	vmute = get(parameters,:vmute,1500.)
	taper = get(parameters,:taper,0.1)

	if (group=="all")
		headers = SeisMain.SeisReadHeaders(in);
		offset = SeisMain.ExtractHeader(headers,"h")
		dt = headers[1].d1
		parameters = Dict(:offset=>offset,:tmute=>tmute,:vmute=>vmute,:taper=>taper,:dt=>dt)

		SeisProcessFile(in,out,[SeisMute],[parameters];group=group)
	else
		itrace_in = itrace #1
		#itrace_out = 1
		nx = SeisMain.GetNumTraces(in)
		while itrace_in <= nx
			h = SeisMain.SeisReadHeaders(in,group=group,key=key,itrace=itrace_in,ntrace=ntrace)
			offset = SeisMain.ExtractHeader(h,"h")
			dt = h[1].d1
			#num_traces = size(h,1)
			parameters = Dict(:offset=>offset,:tmute=>tmute,:vmute=>vmute,:taper=>taper,:dt=>dt)
			#SeisProcessFile(in,out,[SeisMute],[parameters];group=group,key=key,itrace=itrace_in,ntrace=ntrace)
			d1,h1,e1 = SeisMain.SeisRead(in,group=group,key=key,itrace=itrace_in,ntrace=ntrace)

			d2 = SeisMute(d1;parameters...)
			d1 = copy(d2)

			SeisMain.SeisWrite(out,d1,h1,e1,itrace=itrace_in)

			#println(itrace_in," ",ntrace)

			itrace_in += ntrace

			#itrace_in += num_traces
			#itrace_out += num_traces
		end
	end
end


const scipy_opt = PyNULL()

const igl = PyNULL()
const scipy = PyNULL()

# igl = pyimport("igl")
# scipy = pyimport("scipy")
function igl_gaussian_curvature(verts,faces)
    println("new gaussian curv")
    v = array_of_array_to_matrix(verts)
    f = array_of_array_to_matrix(faces) .-1
    K = igl.gaussian_curvature(v, f)
    m = igl.massmatrix(v, f, igl.MASSMATRIX_TYPE_BARYCENTRIC)
    # m = igl.massmatrix(v, f, igl.MASSMATRIX_TYPE_VORONOI)
    minv = scipy.sparse.diags(1.0 ./ m.diagonal())
    kn = minv.dot(K)
    return kn
end
function igl_mean_curvature(verts,faces)
    v = array_of_array_to_matrix(verts)
    f = array_of_array_to_matrix(faces) .-1
    l = igl.cotmatrix(v, f)
    m = igl.massmatrix(v, f, igl.MASSMATRIX_TYPE_BARYCENTRIC)
    # m = igl.massmatrix(v, f, igl.MASSMATRIX_TYPE_VORONOI)

    minv = scipy.sparse.diags(1.0 ./ m.diagonal())

    hn = -minv.dot(l.dot(v))
    h = scipy.linalg.norm(hn, axis=1)
    return h
end

function igl_mean_curvature_flow(v,f,n_step = 10, delta = 0.001)
    new_v = v .- centroid(v)
    new_v = new_v / sqrt(surface_area(new_v,f))
    l = igl.cotmatrix(new_v, f .- 1)

    for i in 1:n_step
        m = igl.massmatrix(new_v,f .- 1, igl.MASSMATRIX_TYPE_BARYCENTRIC)
        s = (m - delta * l)
        b = m.dot(new_v)
        new_v = scipy.sparse.linalg.spsolve(s, b)
        new_v = new_v .- centroid(new_v)
        new_v = new_v / sqrt(surface_area(new_v,f))

    end
    return new_v
end


# weighted Average

Wavg([1,2,3,4,5,6,7],2)==28.0

Wavg([2,3,5,7,8], [7,3,8,2,9])==5.137931034482759


function polynomial_features(data_in::Array; dim::Int=2,
                             subset::Vector=collect(1:size(data_in, dim)), degree::Int=2,
                             include_bias::Bool=true, interaction_only::Bool=false)

    # Data type of the elements of the input array
    T = eltype(data_in)

    # If features are in columns and examples are in rows
    if dim == 2

        # Size of original input matrix
        num_obs, num_feats_in = size(data_in)

        # Create matrix of the chosen subset of features that will generate the new polynomial features
        data_sub = data_in[:, subset]

        # How many original features are chosen to generate the polynomial features
        num_feats_sub = length(subset)

        # If we're only generating the interaction terms (i.e. only single power cross terms)
        if interaction_only

            # Number of columns needed for the new features
            num_feats_poly = 0
            for deg = 2:degree
                num_feats_poly += binomial(num_feats_sub, deg)
            end

            # Initialise design matrix with:
            # column for bias if wanted i.e. if include_bias = true
            # columns for original features i.e. data_in
            # columns for the polynomial features that will be generated (initialised to zero)
            if include_bias
                data_out = [ones(T, num_obs) data_in zeros(T, num_obs, num_feats_poly)]
            else
                data_out = [data_in zeros(T, num_obs, num_feats_poly)]
            end

            # Generate polynomial features and add to data_out
            ind = (include_bias ? 1 : 0) + num_feats_in + 1        # column number to start placing the new features
            for deg = 2:degree                                     # generate features for each power - 2:degree
                for c in combinations(collect(1:num_feats_sub), deg)      # create combinations iterator
                    p = zeros(Int, num_feats_sub)                  # zeros vector of length = number of features in the subset
                    p[c] = 1                                       # change this combination's elements to 1
                    data_out[:, ind] = prod(data_sub .^ p', 2)     # use the power vector to create the new feature
                    ind += 1                                       # move one column to the right to place next new feature
                end
            end

        # If we're generating all the terms
        else

            # Number of columns needed for the new features
            num_feats_poly = 0
            for deg = 2:degree
                num_feats_poly += binomial(num_feats_sub + deg - 1, num_feats_sub - 1)
            end

            # Initialise design matrix with:
            # column for bias if wanted i.e. if include_bias = true
            # columns for original features i.e. data_in
            # columns for the polynomial features that will be generated (initialised to zero)
            if include_bias
                data_out = [ones(T, num_obs) data_in zeros(T, num_obs, num_feats_poly)]
            else
                data_out = [data_in zeros(T, num_obs, num_feats_poly)]
            end

            # Generate polynomial features and add to data_out
            ind = (include_bias ? 1 : 0) + num_feats_in + 1                                 # column number to start placing the new features
            for deg = 2:degree                                                              # generate features for each power - 2:degree
                tmp = [(num_feats_sub + deg); zeros(Int, num_feats_sub)]                    # temp vector used to convert format of combination
                for c in combinations(collect(1:(num_feats_sub + deg - 1)), num_feats_sub - 1)     # create combinations iterator
                    tmp[num_feats_sub:-1:2] = c                                             # place combination into tmp vector
                    p = -diff(tmp)-1                                                        # vector of powers
                    data_out[:, ind] = prod(data_sub .^ p', 2)                              # use the power vector to create the new feature
                    ind += 1                                                                # move one column to the right to place next new feature
                end
            end

        end

    # If features are in rows and examples are in cols
    elseif dim == 1

        # Size of original input matrix
        num_feats_in, num_obs = size(data_in)

        # Create matrix of the chosen subset of features that will generate the new polynomial features
        data_sub = data_in[subset, :]

        # How many original features are chosen to generate the polynomial features
        num_feats_sub = length(subset)

        # If we're only generating the interaction terms (i.e. only single power cross terms)
        if interaction_only

            # Number of rows needed for the new features
            num_feats_poly = 0
            for deg = 2:degree
                num_feats_poly += binomial(num_feats_sub, deg)
            end

            # Initialise design matrix with:
            # row for bias if wanted i.e. if include_bias = true
            # rows for original features i.e. data_in
            # rows for the polynomial features that will be generated (initialised to zero)
            if include_bias
                data_out = [ones(T, 1, num_obs) ; data_in ; zeros(T, num_feats_poly, num_obs)]
            else
                data_out = [data_in ; zeros(T, num_feats_poly, num_obs)]
            end

            # Generate polynomial features and add to data_out
            ind = (include_bias ? 1 : 0) + num_feats_in + 1       # row number to start placing the new features
            for deg = 2:degree                                    # generate features for each power - 2:degree
                for c in combinations(collect(1:num_feats_sub), deg)     # create combinations iterator
                    p = zeros(Int, num_feats_sub)                 # zeros vector of length = number of features in the subset
                    p[c] = 1                                      # change this combination's elements to 1
                    data_out[ind, :] = prod(data_sub .^ p, 1)     # use the power vector to create the new feature
                    ind += 1                                      # move one column to the right to place next new feature
                end
            end

        # If we're generating all the terms
        else

            # Number of columns needed for the new features
            num_feats_poly = 0
            for deg = 2:degree
                num_feats_poly += binomial(num_feats_sub + deg - 1, num_feats_sub - 1)
            end

            # Initialise design matrix with:
            # row for bias if wanted i.e. if include_bias = true
            # rows for original features i.e. data_in
            # rows for the polynomial features that will be generated (initialised to zero)
            if include_bias
                data_out = [ones(T, 1, num_obs) ; data_in ; zeros(T, num_feats_poly, num_obs)]
            else
                data_out = [data_in ; zeros(T, num_feats_poly, num_obs)]
            end

            # Generate polynomial features and add to data_out
            ind = (include_bias ? 1 : 0) + num_feats_in + 1                                 # row number to start placing the new features
            for deg = 2:degree                                                              # generate features for each power - 2:degree
                tmp = [(num_feats_sub + deg); zeros(Int, num_feats_sub)]                    # temp vector used to convert format of combination
                for c in combinations(collect(1:(num_feats_sub + deg - 1)), num_feats_sub - 1)     # create combinations iterator
                    tmp[num_feats_sub:-1:2] = c                                             # place combination into tmp vector
                    p = -diff(tmp)-1                                                        # vector of powers
                    data_out[ind, :] = prod(data_sub .^ p, 1)                               # use the power vector to create the new feature
                    ind += 1                                                                # move one column to the right to place next new feature
                end
            end

        end

    # If dim entered is not in range i.e. must be 1 or 2
    else throw(DimensionMismatch("Dimension must be 1 (features are rows) or 2, the default, (features are cols)."))

    end

    return data_out
end


module OTSIN

using LinearAlgebra

include("optimaltransport.jl")
include("metrics.jl")
include("fitting.jl")
include("utils.jl")
include("perturbation.jl")
include("posterior.jl")
include("plotting.jl")

export isprobability, marginals
export optimaltransport
export entropy, relative_entropy, KL, utility
export uniform, report_utility, normalize
export fitM, fitMsgd, r, global_loss
export Regularization, L1, L2, ElasticNet, Exp, Entropic
export perturbation, pert_utility, pert_entropy
export parallel_cross_entropy, parallel_cross_entropy_A, parallel_cross_entropy_B,
        parallel_cross_entropy_free
export laplace_posterior
export marginalsheatmap

end # module


using SchumakerSpline
using Test

gridx = collect(1:10)
gridy = collect(1:10)
grid = gridx * gridy'

schum = Schumaker2d(gridx, gridy, grid; extrapolation = (Linear, Linear))

abs(schum(5,5) - 25) < 100 * eps()
abs(schum(5.5,5.5) - 30.25) < 100 * eps()


function Base.real(f::Fun{Laurent{DD,RR}}) where {DD,RR}
    n=ncoefficients(f)
    cfs=f.coefficients

    ret=Array{real(cfstype(f))}(undef, iseven(n) ? n+1 : n)
    ret[1]=real(cfs[1])

    for k=2:2:n
        # exp(1im(k-1)/2*x)=cos((k-1)/2 x) +i sin((k-1)/2 x)
        ret[k]=imag(cfs[k])
        ret[k+1]=real(cfs[k])
    end
    for k=3:2:n
        # exp(1im(k-1)/2*x)=cos((k-1)/2 x) +i sin((k-1)/2 x)
        ret[k]+=real(cfs[k])
        ret[k-1]-=imag(cfs[k])
    end

    Fun(Fourier(domain(f)),ret)
end

function Base.imag(f::Fun{Laurent{DD,RR}}) where {DD,RR}
    n=ncoefficients(f)
    cfs=f.coefficients

    ret=Array{real(cfstype(f))}(undef, iseven(n) ? n+1 : n)
    ret[1]=imag(cfs[1])

    for k=2:2:n
        # exp(1im(k-1)/2*x)=cos((k-1)/2 x) +i sin((k-1)/2 x)
        ret[k]=-real(cfs[k])
        ret[k+1]=imag(cfs[k])
    end
    for k=3:2:n
        # exp(1im(k-1)/2*x)=cos((k-1)/2 x) +i sin((k-1)/2 x)
        ret[k]+=imag(cfs[k])
        ret[k-1]+=real(cfs[k])
    end

    Fun(Fourier(domain(f)),ret)
end


"""
    edm_completion(A::SparseMatrixCSC{Tv, Ti}) where {Tv <: AbstractFloat, Ti <: Integer}

Returns a Euclidean distance matrix completion of `A` if it is EDM-completable
using Algorithm 11.1 in [VA15].

# Reference
[VA15] [Chordal Graphs and Semidefinite Optimization](https://www.seas.ucla.edu/~vandenbe/publications/chordalsdp.pdf)
by Lieven Vandenberghe and Martin Andersen
"""
function edm_completion(A::SparseMatrixCSC{Tv, Ti}) where {Tv <: AbstractFloat, Ti <: Integer}
    !issymmetric(A) && error(ArgumentError("A must be symmetric"))
    n = size(A, 1)

    sp = sparsity_pattern(A)
    !is_chordal(A; peo=1:n) && error(ArgumentError("A must have a chordal sparsity pattern"))
    perm, iperm, L = get_chordal_extension(sp; perm=nothing, verbose=false)
    ct = CliqueTree(L)
    order_snds!(ct)
    iperm = invperm(ct.perm)

    !is_edm_completable(A[ct.perm, ct.perm], ct) && error("A is not EDM completable")

    W = Matrix(A[ct.perm, ct.perm])

    for j in (length(ct.snds)-1):-1:1
        vrep_ind = ct.postordering[j]
        vrep = ct.vreps[vrep_ind]

        ν = ct.snds[vrep_ind]
        α = ct.seps[vrep_ind]
        η = filter(x->(!(x in ν) && !(x in α)), vrep+1:n)

        aj = first(α)

        # See eq (11.4)
        # (I - 1*eₖ') * A * (I - eₖ*1') = A - A[k,k]*1*1' - 1*ãₖ' - aₖ*1'
        Yνν = -0.5*W[ν, ν] .- 0.5*W[aj, aj]
        Yνν .+= 0.5 .* W[ν, aj]
        Yνν .+= 0.5 .* W[aj, ν]'

        Yαα = -0.5*W[α, α] .- 0.5*W[aj, aj]
        Yαα .+= 0.5 .* W[α, aj]
        Yαα .+= 0.5 .* W[aj, α]'

        Yηη = -0.5*W[η, η] .- 0.5*W[aj, aj]
        Yηη .+= 0.5 .* W[η, aj]
        Yηη .+= 0.5 .* W[aj, η]'

        Yηα = -0.5*W[η, α] .- 0.5*W[aj, aj]
        @. Yηα += 0.5 * W[η, aj]
        @. Yηα += 0.5 * W[aj, α]'

        Yαν = -0.5*W[α, ν] .- 0.5*W[aj, aj]
        @. Yαν += 0.5 * W[α, aj]
        @. Yαν += 0.5 * W[aj, ν]'

        d, V = LAPACK.syev!('V', 'L', Yαα)
        λmax = maximum(d)
        Yαα_pinv = V * Diagonal(map(x->(x > λmax*eps() ? 1.0/x : 0.0), d)) * V'
        cache = Yαα_pinv * Yαν
        @views mul!(W[η, ν], Yηα, cache)
        W[η, ν] .*= -2
        W[η, ν] .+= diag(Yηη)
        W[η, ν] .+= diag(Yνν)'

        # W[η, ν] = -2*Yηα*Yαα_pinv*Yαν + ones(length(η))*diag(Yνν)' + diag(Yηη)*ones(length(ν))'
        # BLAS.ger!(1.0, ones(length(η)), diag(Yνν), W[η, ν])
        # BLAS.ger!(1.0, diag(Yηη), ones(length(ν)), W[η, ν])

        # μ = sparsevec([aj], [1.0], n)
        # tmp  = -0.5*(I - ones(n)*μ')*W*(I - μ*ones(n)')
        # rtol = sqrt(eps(real(float(one(eltype(tmp[α, α]))))))
        # W[η, ν] .= -2*tmp[η, α]*pinv(tmp[α, α]; rtol=rtol)*tmp[α, ν] +
        #           ones(length(η))*diag(tmp[ν,ν])' + diag(tmp[η,η])*ones(length(ν))'

        W[ν, η] .= W[η, ν]'
        # display(Matrix(W[iperm, iperm] .* .!(sp .> 0)))
        # break
    end

    return W[iperm, iperm]

end

"""
    is_edm_completable(A::SparseMatrixCSC, ct::CliqueTree)

Checks if `A` is EDM completable. Requires the associated CliqueTree `ct`.
"""
function is_edm_completable(A::SparseMatrixCSC, ct::CliqueTree)
    cliques = get_cliques(ct)
    for c in cliques
        !is_edm(A[c,c]) && return false
    end
    return true
end

"""
    is_edm(M::AbstractMatrix)

Checks if `M` is a Euclidean Distance Matrix. Note: a matrix is a EDM iff it is
negative semidefinite on the subspace on the subspace orthogonal to the all ones
vector.
"""
function is_edm(M::AbstractMatrix)
    (!iszero(diag(M)) || !issymmetric(M)) && return false
    n = size(M, 1)
    Q = spdiagm(n, n-1, 0 => ones(n-1), -1 => -ones(n-1))

    return isposdef(-Symmetric(Q'*M*Q) + 1e-10*I)
end


function SeisFXDecon(d,param=Dict())
	#FX_DECON: SNR enhancement using fx-deconvolution.

	#  [DATA_f] = fx_decon(D1,D2,dt,lf,mu,flow,fhigh);
	# param should be defined, for example
	# param = Dict("dt"=>dt,"filter_length"=>lf,"mu"=>mu,"fmax"=>fhigh)

	#  OUT  de:     filtered data
	#
	#  Reference: Canales, 1984, Random noise reduction, 54.th. Ann. Internat.
	#             Mtg., Soc. Expl. Geophys., Expanded Abstracts, pp. 525-527
	#
	#  Note: Canales' method is modified to use non Toeplitz system of equations
	#        with backward and forward prediction filters
	#

	dt = get(param,"dt",0.002)
	filter_length = get(param,"filter_length",10)
	mu = get(param,"mu",0.01)
	nt = size(d,1)
	nx = size(d,2)
	nf = nt
	dw = 2. *pi/nf/dt
	nw = convert(Int,floor(nf/2)) + 1
	fmax = get(param,"fmax",convert(Int,floor(0.5/dt)))
	if(fmax*dt*nf < nw)
		iw_max = convert(Int,floor(fmax*dt*nf))
	else
		iw_max = convert(Int,floor(0.5/dt))
	end
	D = fft(d,1)
	Df = zeros(ComplexF64,size(D))
	Db = zeros(ComplexF64,size(D))
	for iw=1:iw_max
		y = vec(D[iw,:])
		yf = fxdecon(y,"f",filter_length,mu);
		yb = fxdecon(y,"b",filter_length,mu);
		Df[iw,:] = yf
		Db[iw,:] = yb
	end
	# symmetries
	for iw=nw+1:nf
		Df[iw,:] = conj(Df[nf-iw+2,:])
		Db[iw,:] = conj(Db[nf-iw+2,:])
	end
	dout = ( ifft(Df,1) .+ ifft(Db,1) )/ 2
	out = real(dout[1:nt,:])

	return out
end

function fxdecon(d,direction,filter_length,mu)
	nx = length(d);
	if (direction=="b")
		y  = d[1:nx-filter_length];
		C  = d[2:nx-filter_length+1];
		R  = d[nx-filter_length+1:nx];
		M = hankel(C,R);
	else
		y  = d[filter_length+1:nx,1];
		C  = d[filter_length:nx-1];
		R = reverse(d[1:filter_length]);
		M = toeplitz(C,R);
	end
	B = M'*M;
	beta = B[1,1]*mu/100
	ab = (B + beta*Matrix(I,filter_length,filter_length))\M'*y
	temp = M*ab
	if (direction=="b")
		d = [temp ; zeros(filter_length)]
	else
		d = [zeros(filter_length) ; temp]
	end

	return d
end

function hankel(c, r)

	nc = length(r)
	nr = length(c)
	if (nc > 1)
		c = [ c ; r[2:nc] ]
	end
	m = c[ones(Int,nr) * collect(Int,1:nc)' + collect(Int,0:nr-1)*ones(Int,nc)']
	m = reshape(m,nr,nc)

	return m
end

function toeplitz(c, r)

	nc = length(r)
	nr = length(c)
	m = zeros(ComplexF64,nr, nc);
	for i = 1:min(nc, nr)
		m[i:nr,i] .= c[1:nr-i+1]
	end
	for i = 1:min(nr, nc-1)
		m[i,i+1:nc] .= r[2:nc-i+1];
	end

	return m
end


# hyperbolic_integral_expansion.jl

"""
Note: This is meant to be the full expansion for the hyperbolic integral using the correct arbitrary precision methods, but it currently does not converge to the correct values, so made a typo/error somewhere. It's not massively efficient either, so optimisation will be important!
"""

include("arb_hypgeom.jl")
import .one_F_two_fast

include("hypergeometric_expansion.jl")
import .hypergeom_exp

using ArbNumerics
using QuadGK

"""
Solving the hyperbolic integral with expansions and summations.
"""
function arb_binomial(x, y)
    x = ArbReal("$x")
    y = ArbReal("$y")
    one = ArbReal("1")
    ArbNumerics.gamma(x + one) / (ArbNumerics.gamma(y + one) * ArbNumerics.gamma(x - y + one))
end

function hyperbolic_expansion(Ω, β, α, v, w; prec = 64)

    # Initialise precision of ArbReal to prec.
    setextrabits(0)
    setprecision(ArbReal, prec + 8)

    Ω = ArbReal("$Ω")
    β = ArbReal("$β")
    α = ArbReal("$α")
    v = ArbReal("$v")
    w = ArbReal("$w")

    R = ArbReal((v^2 - w^2) / (w^2 * v))
    a = ArbReal(sqrt(β^2 / 4 + R * β * coth(β * v / 2)))
    b = ArbReal(R * β / sinh(β * v / 2))

    coefficient = ArbReal(α * β^(3/2) * v^3 / (3 * a * √π * w^3 * sinh(β / 2)))

    M_c(z, n) = ArbReal(hypergeometric_expansion.hypergeom_exp(β * z / 2, n, β, a, 0; prec = prec))
    M_s(z, n) = ArbReal(hypergeometric_expansion.hypergeom_exp(β * z / 2, n, β, a, 1; prec = prec))

    c = ArbReal(cosh(Ω * β / 2))
    s = ArbReal(sinh(Ω * β / 2))

    n = ArbReal("0")
    result = ArbReal("0.0")
    err = eps(ArbReal(0, bits = prec))  # Machine accuracy of specified precision prec.

    while true

        previous_result = ArbReal("$result")

        if mod(n, 2) == 0
            even_term = ArbReal(arb_binomial(n, n/2) * M_c(1, n))
            @fastmath @inbounds @simd for z in [ArbReal(Ω + 1), ArbReal(Ω - 1)]
                even_term += ArbReal(arb_binomial(n, n/2) * (s * M_s(z, n) - c * M_c(z, n)) / 2)
            end
        else
            even_term = ArbReal("0.0")
        end

        k_term = ArbReal("0.0")
        @fastmath @inbounds @simd for k in ArbReal("0"):ArbReal(floor(n / 2 - 1 / 2))
            @fastmath @inbounds @simd for z in [ArbReal(1 + v * (n - 2 * k)), ArbReal(1 - v * (n - 2 * k))]
                k_term += ArbReal(arb_binomial(n, k) * M_c(z, n))
            end
            @fastmath @inbounds @simd for z in [ArbReal(Ω + 1 + v * (n - 2 * k)), ArbReal(Ω - 1 + v * (n - 2 * k)), ArbReal(Ω + 1 - v * (n - 2 * k)), ArbReal(Ω - 1 - v * (n - 2 * k))]
                k_term += ArbReal(arb_binomial(n, k) * (s * M_s(z, n) - c * M_c(z, n)) / 2)
            end
        end

        term = ArbReal(arb_binomial(-3/2, n) * (-b / (2 * a))^n * (even_term
         + k_term))
        # Break loop if term smaller than accuracy of result.
        if abs(term) < err
            break
        end
        result += ArbReal(term)
        println("term: n = ", n, "\nterm value: ", coefficient * term, "\ncurrent result: ", coefficient * result, "\n")
        n += ArbReal("1")

        # Double precision if rounding error in result exceeds accuracy specified by prec.
        if ball(result)[2] > err
            setprecision(ArbReal, precision(result) * 2)

            println("Not precise enough. Required error < ", err, ". Increasing precision to ", precision(result) * 2, " bits.\n")

            Ω = ArbReal("$Ω")
            β = ArbReal("$β")
            α = ArbReal("$α")
            v = ArbReal("$v")
            w = ArbReal("$w")

            R = ArbReal((v^2 - w^2) / (w^2 * v))
            a = ArbReal(sqrt(β^2 / 4 + R * β * coth(β * v / 2)))
            b = ArbReal(R * β / sinh(β * v / 2))

            coefficient = ArbReal(α * β^(3/2) * v^3 / (3 * a * √π * w^3 * sinh(β / 2)))

            M_c(z, n) = ArbReal(hypergeometric_expansion.hypergeom_exp(β * z / 2, n, β, a, 0; prec = precision(result)))
            M_s(z, n) = ArbReal(hypergeometric_expansion.hypergeom_exp(β * z / 2, n, β, a, 1; prec = precision(result)))

            c = ArbReal(cosh(Ω * β / 2))
            s = ArbReal(sinh(Ω * β / 2))

            n = ArbReal("$(n - 1)")
            result = ArbReal("$previous_result")
        end
    end
    result *= ArbReal(coefficient)
    println("Frequency: ", ArbReal(Ω, bits = prec + 8), ". Final result: ", ArbReal(result, bits = prec + 8))
    ArbReal(result, bits = prec + 8)
end

"""
Solving the hyperbolic integral with quadgk
"""
function hyperbolic_integral(Ω, β, α, v, w)

    # Initialise constants.
    R = (v^2 - w^2) / (w^2 * v)
    a = sqrt(β^2 / 4 + R * β * coth(β * v / 2))
    b = R * β / sinh(β * v / 2)

    coefficient = 2 * α * β^(3 / 2) * v^3 / (3 * sqrt(π) * sinh(β / 2) * w^3)

    integrand(x) = (1 - cosh(Ω * x)) * cosh(x - β / 2) / (a^2 - β^2 / 4 + x * (β - x) - b * cosh(v * (x - β / 2)))^(3 / 2)

    integral = quadgk(x -> integrand(x), 0, β / 2)
    println(Ω, " ", coefficient * integral[1])
    return coefficient * integral[1]
end

"""
Plot results for comparison.
"""

Ω_range = 7.01
β = 2.0
α = 7.0
v = 5.8
w = 1.6
R_int = [hyperbolic_integral(Ω, β, α, v, w) for Ω in Ω_range]
# R_osc = [oscillatory_integral_expansion(Ω, β, α, v, w; prec = 32) for Ω in Ω_range]
R_exp = [hyperbolic_expansion(Ω, β, α, v, w; prec = 32) for Ω in Ω_range]
# @show(R_exp)
# p = plot(Ω_range, abs.(R_exp), yaxis = :log)
# plot!(Ω_range, abs.(R_int))
# display(p)


using SubsetSelection
using SubsetSelectionCIO
using Test
using StatsBase

n = 500; p = 1000; k = 10;
indices = sort(sample(1:p, StatsBase.Weights(ones(p)/p), k, replace=false));
w = sample(-1:2:1, k);
X = randn(n,p); Y = X[:,indices]*w;

γ = 1/sqrt(size(X,1))

indices0, w0, Δt, status, Gap, cutCount = oa_formulation(SubsetSelection.OLS(), Y, X, k, γ)
for i in 1:k
  @test indices0[i]==indices[i]
end


infeasible = @stochastic_model begin
    @stage 1 begin
        @decision(model, x₁ >= 0)
        @decision(model, x₂ >= 0)
        @objective(model, Min, 3*x₁ + 2*x₂)
    end
    @stage 2 begin
        @uncertain ξ₁ ξ₂
        @recourse(model, 0.8*ξ₁ <= y₁ <= ξ₁)
        @recourse(model, 0.8*ξ₂ <= y₂ <= ξ₂)
        @objective(model, Min, -15*y₁ - 12*y₂)
        @constraint(model, 3*y₁ + 2*y₂ <= x₁)
        @constraint(model, 2*y₁ + 5*y₂ <= x₂)
    end
end

ξ₁ = @scenario ξ₁ = 6. ξ₂ = 8. probability = 0.5
ξ₂ = @scenario ξ₁ = 4. ξ₂ = 4. probability = 0.5

infeasible_res = SPResult([27.2,41.6], Dict(1 => [4.8, 6.4], 2 => [4., 4.]), 36.4, 9.2, 27.2, Inf, 9.2, Inf)
push!(problems, (infeasible, [ξ₁,ξ₂], infeasible_res, "Infeasible"))


# SPDX-License-Identifier: X11
# 2020-11-15
# Collinearity (300pt)

function main()
  N = parse(Int, readline())
  A = Vector{Tuple{Int, Int}}()
  for i ∈ 1:N
    x,y = parse.(Int, (split ∘ readline)())
    push!(A, (x,y))
  end

  for i ∈ 1:N - 2, j ∈ i + 1:N - 1, k ∈ j + 1:N
    if trianglearea(A[i], A[j], A[k]) == 0
      println("Yes")
      return
    end
  end

  println("No")
end

trianglearea((xa,ya)::Tuple{<:Real, <:Real},
             (xb,yb)::Tuple{<:Real, <:Real},
             (xc,yc)::Tuple{<:Real, <:Real}) =
  abs(xa * yb - xa * yc + xb * yc - xb * ya + xc * ya - xc * yb) / 2

main()


using Distributions: MultivariateNormal

function optimaltransport_safe(M::AbstractMatrix, a::AbstractVector, b::AbstractVector;
                        λ::Real=1.0, maxitter::Int=10, ϵ::Real=1e-5)
    Q = exp.(λ * M)
    iter = 0
    while iter < maxitter && maximum(abs.(sum(Q, dims=2)[:] - a)) > ϵ
        iter += 1
        Q = a .* (sum(Q, dims=2) .\ Q)
        Q = Q .* transpose(b) ./ sum(Q, dims=1)
    end
    return Q
end

function laplace_posterior(M, P; γ=1e-2,
                maxitter::Int=20, tol=1e-4)
    a, b = marginals(P)
    Q = M -> Q = optimaltransport_safe(exp.(M), a, b)
    l = M -> relative_entropy(P, Q(M)) + γ * sum(M .* M)
    hess = Zygote.hessian(l, M)
    return MultivariateNormal(vec(exp.(M)), Symmetric(hess+tol*I)^-1)
end

#TODO: test and explore this.