text:
"""

 █████╗  ██████╗██████╗ ██╗
██╔══██╗██╔════╝██╔══██╗██║
███████║██║     ██████╔╝██║
██╔══██║██║     ██╔══██╗██║
██║  ██║╚██████╗██║  ██║███████╗
╚═╝  ╚═╝ ╚═════╝╚═╝  ╚═╝╚══════╝

File:       PlanarQuadMPC.jl
Author:     Gabriel Barsi Haberfeld, 2020. gbh2@illinois.edu
Function:   This program implemetns an MPC controller for trajectory tracking of
            a planar quadrotor.

Instructions:   Run this file in juno with Julia 1.5.1 or later.
Requirements:   JuMP, Ipopt, Plots, LinearAlgebra, BenchmarkTools.

"""

using JuMP, Ipopt
using Plots, LinearAlgebra
using Polynomials
using SparseArrays
using ControlSystems

function computeTraj()
    m_q = 4                                #mass of a quadrotor
    I_q = diag([0.3 0.3 0.3])              #moment of inertia of a quadrotor
    g = 9.81                               #gravitational acceleration

    #Trajectory
    global m
    n = 4                                  #number of flat outputs (x, y, z, psi)
    t_f = 10                               #final time of the trajectory

    order = 6                              #order of polynomial functions

    time_interval_selection_flag = true    #true : fixed time interval, false : optimal time interval
    if (time_interval_selection_flag)
        t = collect(range(0, stop = t_f, length = m + 1))
    end

    n_intermediate = 5
    corridor_width = 0.05
    corridor_position = [3 4]

    #keyframes, column = index, rows = xyz yaw
    keyframe = [
        0 7.1 14.3 7.9
        0 8.6 0.7 -5.3
        0 1.0 1.3 0.9
        0 0 0 0
    ]
    m = size(keyframe,2)
    c = zeros(4 * (order + 1) * m)
    mu_r = 1
    mu_psi = 1
    k_r = 4
    k_psi = 2
    A = computeCostMat(order, m, mu_r, mu_psi, k_r, k_psi, t)
end

function computeCostMat(order, m, mu_r, mu_psi, k_r, k_psi, t)
    polynomial_r = Polynomial(ones(order + 1))
    for i = 1:k_r
        polynomial_r = derivative(polynomial_r)       #Differentiation up to k
    end

    polynomial_psi = Polynomial(ones(order + 1))
    for i = 1:k_psi
        polynomial_psi = derivative(polynomial_psi)   #Differentiation up to k
    end

    A = []
    for i = 1:m
        A_x = zeros(order + 1, order + 1)
        A_y = zeros(order + 1, order + 1)
        A_z = zeros(order + 1, order + 1)
        A_psi = zeros(order + 1, order + 1)
        for j = 1:order+1
            for k = j:order+1

                #Position
                if (j <= length(polynomial_r) && (k <= length(polynomial_r)))
                    order_t_r = ((order - k_r - j + 1) + (order - k_r - k + 1))
                    if (j == k)
                        A_x[j, k] =
                            polynomial_r[j-1]^2 / (order_t_r + 1) *
                            (t[i+1]^(order_t_r + 1) - t[i]^(order_t_r + 1))
                        A_y[j, k] =
                            polynomial_r[j-1]^2 / (order_t_r + 1) *
                            (t[i+1]^(order_t_r + 1) - t[i]^(order_t_r + 1))
                        A_z[j, k] =
                            polynomial_r[j-1]^2 / (order_t_r + 1) *
                            (t[i+1]^(order_t_r + 1) - t[i]^(order_t_r + 1))
                    else
                        A_x[j, k] =
                            2 * polynomial_r[j-1] * polynomial_r[k-1] / (order_t_r + 1) *
                            (t[i+1]^(order_t_r + 1) - t[i]^(order_t_r + 1))
                        A_y[j, k] =
                            2 * polynomial_r[j-1] * polynomial_r[k-1] / (order_t_r + 1) *
                            (t[i+1]^(order_t_r + 1) - t[i]^(order_t_r + 1))
                        A_z[j, k] =
                            2 * polynomial_r[j-1] * polynomial_r[k-1] / (order_t_r + 1) *
                            (t[i+1]^(order_t_r + 1) - t[i]^(order_t_r + 1))
                    end
                end

                #Yaw
                if (j <= length(polynomial_psi) && (k <= length(polynomial_psi)))
                    order_t_psi = ((order - k_psi - j + 1) + (order - k_psi - k + 1))
                    if (j == k)
                        A_psi[j, k] =
                            polynomial_psi[j-1]^2 / (order_t_psi + 1) *
                            (t[i+1]^(order_t_psi + 1) - t[i]^(order_t_psi + 1))
                    else
                        A_psi[j, k] =
                            2 * polynomial_psi[j-1] * polynomial_psi[k-1] /
                            (order_t_psi + 1) *
                            (t[i+1]^(order_t_psi + 1) - t[i]^(order_t_psi + 1))
                    end
                end

            end
        end
        if i == 1
            blocks = [mu_r * A_x, mu_r * A_y, mu_r * A_z, mu_psi * A_psi]
        else
            blocks = [A, mu_r * A_x, mu_r * A_y, mu_r * A_z, mu_psi * A_psi]
        end
        A = ControlSystems.blockdiag(blocks...)
    end
    A = 0.5 * (A + A') #Make it symmetric
end

function computeConstraint(order, m, k_r, k_psi, t, keyframe, corridor_position, n_intermediate, corridor_width)

n = 4;                              #State number

#Waypoint constraints
C1 = zeros(2*m*n,n*(order+1)*m);
b1 = zeros(2*m*n);
computeMat = diagm(ones(order+1));          #Required for computation of polynomials
for i=1:m
    waypoint = keyframe[:,i];       #Waypoint at t(i)

    if(i==1)                        #Initial and Final Position
        #Initial
        values = zeros(1,order+1);
        for j=1:order+1
            poly = Polynomial(computeMat[j,:])
            values[j] = poly(t[i])
        end

        for k=1:n
            c = zeros(1,n*(order+1)*m);
            c[ ((i-1)*(order+1)*n+(k-1)*(order+1)+1) : ((i-1)*(order+1)*n+(k-1)*(order+1))+order+1 ] = values;
            C1[k,:] = c;
        end
        b1[1:n] = waypoint;

        #Final
        for j=1:order+1
            values(j) = polyval(computeMat(j,:),t(m+1));
        end

        for k=1:n
            c = zeros(1,n*(order+1)*m);
            c( ((m-1)*(order+1)*n+(k-1)*(order+1)+1) : ((m-1)*(order+1)*n+(k-1)*(order+1))+order+1) = values;
            C1(k+n,:) = c;
        end
        b1(n+1:2*n) = waypoint;

    else
        #Elsewhere
        values = zeros(1,order+1);
        for j=1:order+1
            values(j) = polyval(computeMat(j,:),t(i));
        end

        for k=1:n
            c = zeros(1,n*(order+1)*m);
            c( ((i-2)*(order+1)*n+(k-1)*(order+1)+1) : ((i-2)*(order+1)*n+(k-1)*(order+1))+order+1) = values;
            C1(k+2*n*(i-1),:) = c;
        end
        b1(2*n*(i-1)+1:2*n*(i-1)+n) = waypoint;

        for k=1:n
            c = zeros(1,n*(order+1)*m);
            c( ((i-1)*(order+1)*n+(k-1)*(order+1)+1) : ((i-1)*(order+1)*n+(k-1)*(order+1))+order+1) = values;
            C1(k+2*n*(i-1)+n,:) = c;
        end
        b1(2*n*(i-1)+n+1:2*n*(i-1)+2*n) = waypoint;

    end

end

# Derivative constraints

# Position
C2 = zeros(2*m*(n-1)*k_r,n*(order+1)*m);                    #(n-1) : yaw excluded here
b2 = ones(2*m*(n-1)*k_r,1)*eps;
constraintData;
#constraintData_r = zeros(m,k_r,3);
for i=1:m
    for h=1:k_r
        if(i==1)
            #Initial
            values = zeros(1,order+1);
            for j=1:order+1
                tempCoeffs = computeMat(j,:);
                for k=1:h
                    tempCoeffs = polyder(tempCoeffs);
                end
                values(j) = polyval(tempCoeffs,t(i));
            end

            continuity = zeros(1,n-1);
            for k=1:n-1
                if(constraintData_r(i,h,k)==eps)
                    %Continuity
                    continuity(k) = true;
                end

                c = zeros(1,n*(order+1)*m);
                if(continuity(k))
                    c( ((i-1)*(order+1)*n+(k-1)*(order+1)+1) : ((i-1)*(order+1)*n+(k-1)*(order+1))+order+1) = values;
                    c( ((m-1)*(order+1)*n+(k-1)*(order+1)+1) : ((m-1)*(order+1)*n+(k-1)*(order+1))+order+1) = -values;
                    C2(k + (h-1)*(n-1),:) = c;
                    b2(k + (h-1)*(n-1)) = 0;
                else
                    c( ((i-1)*(order+1)*n+(k-1)*(order+1)+1) : ((i-1)*(order+1)*n+(k-1)*(order+1))+order+1) = values;
                    C2(k + (h-1)*(n-1),:) = c;
                    b2(k + (h-1)*(n-1)) = constraintData_r(i,h,k);
                end
            end

            #Final
            values = zeros(1,order+1);
            for j=1:order+1
                tempCoeffs = computeMat(j,:);
                for k=1:h
                    tempCoeffs = polyder(tempCoeffs);
                end
                values(j) = polyval(tempCoeffs,t(m+1));
            end

            for k=1:n-1
                if(constraintData_r(i,h,k)==eps)
                    #Continuity
                end
                c = zeros(1,n*(order+1)*m);
                if(~continuity(k))
                    c( ((m-1)*(order+1)*n+(k-1)*(order+1)+1) : ((m-1)*(order+1)*n+(k-1)*(order+1))+order+1) = values;
                    C2(k + (h-1)*(n-1) + (n-1)*k_r,:) = c;
                    b2(k + (h-1)*(n-1) + (n-1)*k_r) = constraintData_r(i,h,k);
                end
            end

        else

            #Elsewhere
            values = zeros(1,order+1);
            for j=1:order+1
                tempCoeffs = computeMat(j,:);
                for k=1:h
                    tempCoeffs = polyder(tempCoeffs);
                end
                values(j) = polyval(tempCoeffs,t(i));
            end

            continuity = zeros(1,n-1);
            for k=1:n-1
                if(constraintData_r(i,h,k)==eps)
                    #Continuity
                    continuity(k) = true;
                end

                c = zeros(1,n*(order+1)*m);
                if(continuity(k))
                    c( ((i-2)*(order+1)*n+(k-1)*(order+1)+1) : ((i-2)*(order+1)*n+(k-1)*(order+1))+order+1) = values;
                    c( ((i-1)*(order+1)*n+(k-1)*(order+1)+1) : ((i-1)*(order+1)*n+(k-1)*(order+1))+order+1) = -values;
                    C2(k + (h-1)*(n-1) + 2*(i-1)*(n-1)*k_r,:) = c;
                    b2(k + (h-1)*(n-1) + 2*(i-1)*(n-1)*k_r) = 0;
                else
                    c( ((i-2)*(order+1)*n+(k-1)*(order+1)+1) : ((i-2)*(order+1)*n+(k-1)*(order+1))+order+1) = values;
                    C2(k + (h-1)*(n-1) + 2*(i-1)*(n-1)*k_r,:) = c;
                    b2(k + (h-1)*(n-1) + 2*(i-1)*(n-1)*k_r) = constraintData_r(i,h,k);
                end
            end

            continuity = zeros(1,n-1);
            for k=1:n-1
                if(constraintData_r(i,h,k)==eps)
                    #Continuity
                    continuity(k) = true;
                end
                c = zeros(1,n*(order+1)*m);

                if(~continuity(k))
                    c( ((i-1)*(order+1)*n+(k-1)*(order+1)+1) : ((i-1)*(order+1)*n+(k-1)*(order+1))+order+1) = values;
                    C2(k + (h-1)*(n-1) + 2*(i-1)*(n-1)*k_r + (n-1)*k_r,:) = c;
                    b2(k + (h-1)*(n-1) + 2*(i-1)*(n-1)*k_r + (n-1)*k_r) = constraintData_r(i,h,k);
                end

            end

        end
    end
end

#Corridor constraints

C3 = [];

b3 = [];
t_vector = (keyframe(1:3,corridor_position(2)) - keyframe(1:3,corridor_position(1)))...
/norm(keyframe(1:3,corridor_position(2)) - keyframe(1:3,corridor_position(1)));
#unit vector of direction of the corridor

t_intermediate = linspace(t(corridor_position(1)),t(corridor_position(2)),n_intermediate+2);
t_intermediate = t_intermediate(2:end-1);
#intermediate time stamps

computeMat = eye(order+1);          #Required for computation of polynomials
for i = 1:n_intermediate
    values = zeros(1,order+1);
    for j=1:order+1
        values(j) = polyval(computeMat(j,:),t_intermediate(i));
    end

    c = zeros(6, n*(order+1)*m);       #Absolute value constraint : two inequality constraints
    b = zeros(6, 1);

    rix = keyframe(1,corridor_position(1));
    riy = keyframe(2,corridor_position(1));
    riz = keyframe(3,corridor_position(1));
    #x
    c(1,(corridor_position(1)-1)*n*(order+1)+0*(order+1)+1:(corridor_position(1)-1)*n*(order+1)+3*(order+1))...
        = [values zeros(1,2*(order+1))]...
        - t_vector(1)*[t_vector(1)*values t_vector(2)*values t_vector(3)*values];
    b(1) = corridor_width +...
        rix+t_vector(1)*(-rix*t_vector(1) -riy*t_vector(2) -riz*t_vector(3));
    c(2,(corridor_position(1)-1)*n*(order+1)+0*(order+1)+1:(corridor_position(1)-1)*n*(order+1)+3*(order+1))...
        = -c(1,(corridor_position(1)-1)*n*(order+1)+0*(order+1)+1:(corridor_position(1)-1)*n*(order+1)+3*(order+1));
    b(2) = corridor_width +...
        -rix-t_vector(1)*(-rix*t_vector(1) -riy*t_vector(2) -riz*t_vector(3));
    #y
    c(3,(corridor_position(1)-1)*n*(order+1)+0*(order+1)+1:(corridor_position(1)-1)*n*(order+1)+3*(order+1))...
        = [zeros(1,order+1) values zeros(1,order+1)]...
        - t_vector(2)*[t_vector(1)*values t_vector(2)*values t_vector(3)*values];
    b(3) = corridor_width +...
        riy+t_vector(2)*(-rix*t_vector(1) -riy*t_vector(2) -riz*t_vector(3));
    c(4,(corridor_position(1)-1)*n*(order+1)+0*(order+1)+1:(corridor_position(1)-1)*n*(order+1)+3*(order+1))...
        = -c(3,(corridor_position(1)-1)*n*(order+1)+0*(order+1)+1:(corridor_position(1)-1)*n*(order+1)+3*(order+1));
    b(4) = corridor_width +...
        -riy-t_vector(2)*(-rix*t_vector(1) -riy*t_vector(2) -riz*t_vector(3));
    #z
    c(5,(corridor_position(1)-1)*n*(order+1)+0*(order+1)+1:(corridor_position(1)-1)*n*(order+1)+3*(order+1))...
        = [zeros(1,2*(order+1)) values]...
        - t_vector(3)*[t_vector(1)*values t_vector(2)*values t_vector(3)*values];
    b(5) = corridor_width +...
        riz+t_vector(3)*(-rix*t_vector(1) -riy*t_vector(2) -riz*t_vector(3));
    c(6,(corridor_position(1)-1)*n*(order+1)+0*(order+1)+1:(corridor_position(1)-1)*n*(order+1)+3*(order+1))...
        = -c(5,(corridor_position(1)-1)*n*(order+1)+0*(order+1)+1:(corridor_position(1)-1)*n*(order+1)+3*(order+1));
    b(6) = corridor_width +...
        -riz-t_vector(3)*(-rix*t_vector(1) -riy*t_vector(2) -riz*t_vector(3));

    C3 = [C3; c];
    b3 = [b3; b];
end

C = [C1; C2];
b = [b1; b2];
end

max_stars_repo_name: gbarsih/Safe-Optimal-Rendezvous

max_issues_repo_name: gbarsih/Safe-Optimal-Rendezvous

max_forks_repo_name: gbarsih/Safe-Optimal-Rendezvous


text:
function quandlget(id::String; order="des", rows=100, period="daily", transformation="none", format="TimeArray", auth_token="")

    # Create a dictionary with the Query arguments that we pass to get() function
    query_args = {"sort_order" => order, "rows" => rows, "collapse" => period, "transformation" => transformation}

    # Open the auth_token file and add the token (if any) to the Query dictionary
    if auth_token == ""
        auth_token = open(readall, Pkg.dir("Quandl/src/token/auth_token.jl"))
    end

    length(auth_token) < 50 && auth_token != "" ? query_args["auth_token"] = auth_token : nothing

    # Get the response from Quandl's API, using Query arguments (see Response.jl README)
    response = get("http://www.quandl.com/api/v1/datasets/$id.csv", query = query_args)

    # Convert the response to the right DataType
    if format == "TimeArray"
        timearray(response)
    elseif format == "DataFrame"
        dataframe(response)
    else
        error("Invalid $format format. If you want this format implemented, please report an issue or submit a pull request.")
    end
end

# alias quandl/quandlget
quandl(args) = quandlget(args)

function quandlsearch(query::ASCIIString, page="1", results="20", format="Dict")

    # Parsing query argument
    query = replace(query, " ", "+")

    # Create a dictionary with the Query arguments that we pass to get() function
    query_args = {"query" => query, "page" => page, "per_page" => results}

    # Getting response from Quandl and parsing it
    response = get("http://www.quandl.com/api/v1/datasets.json", query = query_args)
    jsondict = JSON.parse(response.data)

    # Printing summary
    print("Returning $results results of $jsondict["total_count"] from page $page")

    # Convert the response to the right DataType
    if format == "Dict"
        return jsondict["docs"]
    elseif format == "DataFrame"
        error("DataFrame format yet to be implemented")
    else
        error("Invalid $format format. If you want this format implemented, please report an issue or submit a pull request.")
    end

function set_auth_token(token::String)

	# TODO: Verify if the token is valid
	token_file = open(Pkg.dir("Quandl/src/token/auth_token.jl"), "w")

	try
		write(token_file, token * "\n") # Put a newline after the token
	finally
	    close(token_file)
	end

	return nothing
end

ext: jl


text:
#
# A Julia to Tiramisu C++ Generator
# Transformation happens in two main steps:
#
#   - First, the Julia AST from ParallelAccelerator gets transformed
#     into a Tiramisu AST of TExpr's, TBuffer's, TComputation's, etc.
#
#   - Then, each Tiramisu AST node is passed into a tiramisu_translate
#     method to generate the proper Tiramisu C++ code.
#
# Optional passes could be added in the future in between the two steps
# to perform other optimizations (eliminate redundant buffers, evaluate
# gotos, etc)
#
# tiramisu_from_root_entry is the entry point where all the magic happens
#
# Originally Created By David Vargas
#

module Tiramisu

import CompilerTools.DebugMsg
using CompilerTools.LambdaHandling
using CompilerTools.Helper
DebugMsg.init()

using ..ParallelAccelerator

importall Base

export tiramisu_from_root_entry, tcanonicalize

primitives = DataType[Int8, UInt8,
                     Int16, UInt16,
                     Int32, UInt32,
                     Int64, UInt64,
                     Float16, Float32,
                     Float64, Bool,
                     Char, Void]

jToC = Dict(
            Int8    =>  "int8_t",
            UInt8   =>  "uint8_t",
            Int16   =>  "int16_t",
            UInt16  =>  "uint16_t",
            Int32   =>  "int32_t",
            UInt32  =>  "uint32_t",
            # Int64   =>  "int64_t",
            # UInt64  =>  "uint64_t",
            Int64   =>  "int32_t",
            UInt64  =>  "uint32_t",
            Float16 =>  "float",
            Float32 =>  "float",
            Float64 =>  "double",
            Bool    =>  "bool",
            Char    =>  "char",
            Void    =>  "void"
    ) #Map from Julia types to C++ types

jToT = Dict(
            Int8    =>  "p_int8",
            UInt8   =>  "p_uint8",
            Int16   =>  "p_int16",
            UInt16  =>  "p_uint16",
            Int32   =>  "p_int32",
            UInt32  =>  "p_uint32",
            Int64   =>  "p_int32",
            UInt64  =>  "p_uint32",
            Float16 =>  "p_float32",
            Float32 =>  "p_float32",
            Float64 =>  "p_float64",
            Bool    =>  "p_boolean"
    ) #Map from Julia types to Tiramisu types

jToH = Dict(
            Int8    =>  "Halide::Int(8)",
            UInt8   =>  "Halide::UInt(8)",
            Int16   =>  "Halide::Int(16)",
            UInt16  =>  "Halide::UInt(16)",
            Int32   =>  "Halide::Int(32)",
            UInt32  =>  "Halide::UInt(32)",
            Int64   =>  "Halide::Int(32)",
            UInt64  =>  "Halide::UInt(32)",
            Float16 =>  "Halide::Float(32)",
            Float32 =>  "Halide::Float(32)",
            Float64 =>  "Halide::Float(64)",
            Bool    =>  "Halide::Boolean()"
)

Primitive = Union{Int8, UInt8, Int16, UInt16, Int32, UInt32, Int64, UInt64,
                  Float16, Float32, Float64, Bool}

# Get type used by arch (Int32 or Int64)
IntArch = typeof(1)

# Generated slot
immutable GenSlot
    id::IntArch
end

Variable = Union{SlotNumber, SSAValue, GenSlot}

# TODO : multiple computations per variable for SlotNumber

#TODO: Edit these constants to match the proper directory of all the following
const HALIDE_INCLUDE_DIR = "/home/malek/tiramisu/Halide/include"
const HALIDE_TOOLS_DIR = "/home/malek/tiramisu/Halide/tools"
const HALIDE_LIB_DIR = "/home/malek/tiramisu/Halide/lib"
const ISL_INCLUDE_DIR = "/home/malek/tiramisu/3rdParty/isl/build/include"
const ISL_LIB_DIR = "/home/malek/tiramisu/3rdParty/isl/build/lib"
const TIRAMISU_INCLUDE_DIR = "/home/malek/tiramisu/include"
const TIRAMISU_BUILD_DIR = "/home/malek/tiramisu/build"

#The following constants are the output files created
const FILENAME = "/home/malek/generated/thello.cpp"
const FILENAME_WRAPPER = "/home/malek/generated/thello_wrapper.cpp"
const GEN = "$TIRAMISU_BUILD_DIR/thello_generator"
const OBJ = "$TIRAMISU_BUILD_DIR/thello.o"
const OBJ_WRAPPER = "$TIRAMISU_BUILD_DIR/thello_wrapper.o"
const LIB = "$TIRAMISU_BUILD_DIR/thello.so"

#Globals
function_name = ""
computationIndex = 0
computationPrefix = "S"

function make_counter()
    counter = -1
    return function()
        counter += 1
        return counter
    end
end

type Constant
    id::Int
end

constant_counter = make_counter()
new_constant() = Constant(constant_counter())

new_computation_id = make_counter()

genslot_counter = make_counter()
new_genslot() = GenSlot(genslot_counter())

@enum TBufferArgument TBOutput TBInput TBTemporary
argNames = Dict(
    TBOutput => "a_output",
    TBInput => "a_input",
    TBTemporary => "a_temporary",
)
@enum TOpStyle TOpFuncCall TOpInfix

abstract TSig

type InputSig <: TSig
    var::Variable
    pos::Int
end

type InputSizeSig <: TSig
    var::Variable
    pos::Int
end

type OutputSig <: TSig
    var::Variable
end

signature = TSig[]

abstract TExpr

type TEmptyExpr <: TExpr
end

type TValue <: TExpr
    #expr e_p0 = expr((int32_t) SIZE0);
    name::String
    val
    typ::DataType
end

#The following types are found on lower levels of a Tiramisu AST to help translate,
#modify, and optimize the Tiramisu AST
type TOp <: TExpr
    name::String
    fcn::String
    style::TOpStyle
    args::Array{TExpr, 1}
    typ::DataType
end

#The following types are what should be found on the top level of a Tiramisu AST
type TBuffer
    #buffer buf0("buf0", 2, {tiramisu::expr(10),tiramisu::expr(10)}, p_uint8, NULL, a_output, &function0);
    name
    dimensions :: Array{TExpr,1}
    typ
    argument::TBufferArgument
    is_scalar::Bool
end

type TAlloc <: TExpr
    dims::Vector{TExpr}
    typ::DataType
end

type TConstant <: TExpr
    #constant N("N", e_N, p_int32, true, NULL, 0, &function0);
    id::Constant
    expr :: TExpr
    typ::DataType
end

indexNames = "ijklmnop"

type ISLCondition
    var::String
    op::String
    lim::String
end

type TAccess
    computation::Vector{TExpr}
    buffer::Vector{TExpr}
end

type TComputation
    #computation S0("[N]->{S0[i,j]: 0<=i<N and 0<=j<N}", e3, true, p_uint8, &function0);
    id::Integer
    condition::TExpr
    dim::Unsigned
    expr::TExpr
    typ::DataType
    access::TAccess
    vars::Vector{Variable}
    is_input::Bool
end

type TSignature
    args
    ret
end

abstract TVar

type JVar <: TVar # Julia Variable
    id::Variable
    typ::DataType
    buffer::Union{TBuffer, Void}
    computation::Union{TComputation, Void} # Might not be there already
end

copy(x::JVar) = JVar(x.id, x.typ, x.buffer, x.computation)

type DVar <: TVar # Dummy variable TODO remove
    id::GenSlot
end

abstract TIndex

type TConstIndex <: TIndex
    val::Integer
end

type TVarIndex <: TIndex
    letter::Unsigned
end

abstract TRead <: TExpr

type TLoopVarRead <: TRead
    var::Variable
end

type TSimpleRead <: TRead
    var::TVar
end

type TIndexedRead <: TRead
    var::JVar # TODO should this be a TVar?
    indices::Vector{TExpr}
end

vars = Dict{Variable, JVar}()
inputSizeGenvars = Vector{JVar}()
readGenvars = Vector{JVar}()
genvars = Vector{JVar}()

idToBufferName(id::SSAValue) = "BSSA" * string(id.id)
idToBufferName(id::SlotNumber) = "BSN" * string(id.id)
idToBufferName(id::GenSlot) = "BGS" * string(id.id)

macro typToBufferGen(types)
    return esc(Expr(:call,
                    :Dict,
                    (:($(typ) => function(id, typ)
                        return TBuffer(idToBufferName(id),
                                       [TValue("", 1, IntArch)],
                                       $(typ),
                                       TBTemporary,
                                       true)
                    end) for typ in eval(types))...))
end

typToTBuffer = @typToBufferGen primitives

function parseNode(node::Expr)
    return exprHandlers[node.head](node)
end

function parseNode(node::Variable)
    # TODO does this always work???
    if !haskey(vars, node)
        throw(UndefVarError(Symbol(node)))
    end
    # TODO: Indexing should not be done here
    if node in loopVars
        return TLoopVarRead(node)
    end
    return TSimpleRead(vars[node])
end

function parseNode(node::Primitive)
    return TValue("", node, typeof(node))
end

function parseNode(expr::Any)
    println("Unsupported: ", typeof(expr))
end

# parsing a computation should be idem potent
function parseComputation(expr::TComputation, typ::DataType)
    println("Parsing a computation, should I be?")
    assert(expr.typ === typ)
    return expr
end

function parseComputation(expr::TExpr, typ::DataType) # only scalar for now
    if typ <: Primitive
        if isempty(condStack)
            var = new_genslot()
            push!(loopVars, var)
            dummy(var, IntArch)
            cond = (wrap(var) == wrap(0))
            access = TAccess([wrap(var)], [wrap(0)])
            indices = Variable[var]
            pop!(loopVars)
        else
            cond = condStack[end]
            access = TAccess(map(wrap, loopVars), [wrap(0)])
            indices = collect(loopVars)
        end
        comp = TComputation(
            new_computation_id(),
            cond,
            length(indices),
            expr,
            typ,
            access,
            indices,
            false # TODO : fix
        )
        schedule(comp)
        return comp
    end
end

@enum SchedTime SchedNext SchedNow SchedNever

const RootDimension = -1

SchedLevel = Int

type DelayedSchedule
    current::SchedLevel
    min::SchedLevel
    time::SchedTime
end

type OnceSchedule
    level::SchedLevel
    time::SchedTime
end

type ScheduleState
    for_level::DelayedSchedule
    fuse_level::OnceSchedule
    previous_comp::Union{TComputation, Void}
    schedule::Vector{Pair{Pair{TComputation, TComputation}, SchedLevel}}
    parallelize_level::Union{Int, Void}
    parallelize::Vector{Pair{TComputation, Int}}
end

readSchedLevel(level::SchedLevel) = level
function readSchedLevel(level::DelayedSchedule)
    result, level.min = level.min, level.current
    return result
    #
    # if level.time == SchedNext
    #     level.time = SchedNow
    #     return level.prev_level
    # elseif level.time == SchedNow
    #     return level.level
    # else
    #     # SchedNever
    #     return RootDimension
    # end
end
function readSchedLevel(level::OnceSchedule)
    if level.time == SchedNext
        level.time = SchedNow
        return RootDimension
    elseif level.time == SchedNow
        level.time = SchedNever
        return level.level
    else
        # SchedNever
        return RootDimension
    end
end

sched = ScheduleState(DelayedSchedule(RootDimension, RootDimension, SchedNever),
                      OnceSchedule(RootDimension, SchedNever),
                      nothing, [], nothing, [])

function schedule(comp::TComputation)
    level = maximum(map(readSchedLevel, [sched.for_level, sched.fuse_level]))
    if sched.previous_comp != nothing
        # Not the first computation
        push!(sched.schedule, (sched.previous_comp => comp) => level)
    end
    sched.previous_comp = comp
    if sched.parallelize_level != nothing
        push!(sched.parallelize, comp => sched.parallelize_level)
        sched.parallelize_level = nothing
    end
end

function typedOp!(op::TOp, typ::DataType)
    op.typ = typ
    return op
end

function opToISL(op::String)
    if op == "=="
        return "="
    elseif op == "&&"
        return "and"
    else
        return op
    end
end

+(args::Vararg{TExpr}) = TOp("", "+", TOpInfix, collect(args), Void)
*(args::Vararg{TExpr}) = TOp("", "*", TOpInfix, collect(args), Void)
-(args::Vararg{TExpr}) = TOp("", "-", TOpInfix, collect(args), Void)
/(args::Vararg{TExpr}) = TOp("", "/", TOpInfix, collect(args), Void)
==(args::Vararg{TExpr}) = TOp("", "==", TOpInfix, collect(args), Bool)
<=(args::Vararg{TExpr}) = TOp("", "<=", TOpInfix, collect(args), Bool)
and(args::Vararg{TExpr}) = TOp("", "&&", TOpInfix, collect(args), Bool)

function dummy(id::Variable, typ::DataType)
    return vars[id] = JVar(
        id,
        typ,
        nothing,
        nothing
    )
end

wrap(val::TExpr) = val

wrap(val::Union{Primitive, Variable}) = parseNode(val)

function wrap(val::TVar)
    if val.id in loopVars
        return TLoopVarRead(val.id)
    end
    return TSimpleRead(val)
end

function wrap(val::Any)
    # Shouldn't have nested calls
    print_with_color(:yellow, "Trying to wrap:\n")
    dump(val)
    return val
end

abstract Clause

immutable ForLoop <: Clause
    id::Int
    from::Union{Int, SlotNumber}
    to::Union{Int, SlotNumber}
    inc::Int
    var::TVar
end

clauseStack = Vector{Clause}()
condStack = Vector{TExpr}()
loopVars = Vector{Variable}()

callHandlers = Dict(
    (Core.Intrinsics, :add_int) => function(args::Vector, typ::DataType)
        return typedOp!(+(map(wrap, args)...), typ)
    end,
    (Core.Intrinsics, :sub_int) => function(args::Vector, typ::DataType)
        return typedOp!(-(map(wrap, args)...), typ)
    end,
    (Core.Intrinsics, :add_float) => function(args::Vector, typ::DataType)
        return typedOp!(+(map(wrap, args)...), typ)
    end,
    (Core.Intrinsics, :mul_float) => function(args::Vector, typ::DataType)
        return typedOp!(*(map(wrap, args)...), typ)
    end,
    (Core.Intrinsics, :div_float) => function(args::Vector, typ::DataType)
        return typedOp!(/(map(wrap, args)...), typ)
    end,
    (Core.Intrinsics, :sub_float) => function(args::Vector, typ::DataType)
        return typedOp!(-(map(wrap, args)...), typ)
    end,
    (Core.Intrinsics, :box) => function(args::Vector, typ::DataType)
        return parseNode(args[2])
    end,
    (Core.Intrinsics, :checked_trunc_sint) => function(args::Vector, typ::DataType)
        res = args[1](args[2]) # throw an error if invalid
        return parseNode(res)
    end,
    (ParallelAccelerator.TiramisuPrepass, :for_loop_start) => function(args::Vector, typ::DataType)
        id = args[1]
        from = args[2]
        to = args[3]
        inc = args[4]
        varid = args[6]
        var = dummy(varid, IntArch)
        push!(clauseStack, ForLoop(id, from, to, inc, var))
        push!(loopVars, varid)
        cond = <=(wrap(from) - wrap(1), wrap(var), wrap(to) - wrap(1))
        if isempty(condStack)
            push!(condStack, cond)
        else
            push!(condStack, and(condStack[end], cond))
        end
    #    sched.for_level = DelayedSchedule(length(loopVars) - 1, sched.for_level.prev_level, SchedNext)
        sched.for_level = DelayedSchedule(length(loopVars) - 1, sched.for_level.min, SchedNext)
    end,
    (ParallelAccelerator.TiramisuPrepass, :for_loop_end) => function(args::Vector, typ::DataType)
        id = args[1]
        assert(id == clauseStack[end].id)
        pop!(clauseStack)
        pop!(loopVars)
        pop!(condStack)
        # if isempty(loopVars)
        #     # sched.for_level = DelayedSchedule(RootDimension, RootDimension, SchedNever)
        #     sched.for_level = DelayedSchedule(RootDimension, RootDimension, SchedNever)
        # else
        #     sched.for_level = DelayedSchedule(length(loopVars) - 1, min(sched.for_level.prev_level, length(loopVars) - 1), SchedNext)
        # end
        current = length(loopVars) - 1
        sched.for_level = DelayedSchedule(current, min(current, sched.for_level.min), SchedNext)
    end,
    (ParallelAccelerator.API, :setindex!) => function(args::Vector, typ::DataType)
        arr = copy(vars[args[1]])
        texpr = parseNode(args[2])
        # TODO handle temporal
        # indices = Variable[]
        # added_vars = Variable[]
        # cond = condStack[end]
        # for arg in args[3:end]
        #     if isa(arg, Primitive)
        #         varid = new_genslot()
        #         var = dummy(varid, Int)
        #         cond = and(cond, wrap(var) == wrap(arg - 1))
        #         push!(indices, varid)
        #         push!(added_vars, varid)
        #         push!(loopVars, varid)
        #     else
        #         push!(indices, arg)
        #     end
        # end
        indices = collect(Variable, args[3:end])
        assert(issubset(indices, loopVars)) # TODO more thourough checking
        access = TAccess(map(wrap, loopVars), [wrap(i) for i in indices])
        comp = TComputation(
            new_computation_id(),
            # cond,
            condStack[end],
            length(loopVars),
            texpr,
            arr.typ,
            access,
            copy(loopVars),
            false # TODO fix
        )
        # for var in added_vars
        #     pop!(loopVars)
        # end
        schedule(comp)
        arr.computation = comp
        push!(used_buffers, arr.buffer.name)
        push!(genvars, arr)
    end,
    (ParallelAccelerator.API, :getindex) => function(args::Vector, typ::DataType)
        arr = vars[args[1]]
        indices = collect(Variable, args[2:end])
        assert(issubset(indices, loopVars)) # TODO more thourough checking
        return TIndexedRead(arr, map(wrap, indices))
    end,
)

function handleReturn(arg::Void, typ::DataType)
end

function handleReturn(arg::Variable, typ::DataType)
    vars[arg].buffer.argument = TBOutput
    push!(signature, OutputSig(arg))
end

function handleReturn(arg, typ::DataType)
    id = new_genslot()
    buffer = typToTBuffer[typ](id, typ)
    buffer.argument = TBOutput
    comp = parseComputation(parseNode(arg), typ)
    vars[id] = JVar(id, typ, buffer, comp)
    push!(genvars, vars[id])
    push!(signature, OutputSig(id))
end

function createEmptyComputation(typ, dims, is_scalar)
    l = length(loopVars)
    indices = Variable[new_genslot() for _ in dims]
    append!(loopVars, indices)
    map(i -> dummy(i, Int), indices)
    if is_scalar
        access = TAccess(map(wrap, indices), [wrap(0) for i in indices])
    else
        access = TAccess(map(wrap, indices), [wrap(i) for i in indices])
    end
    cond = and((<=(wrap(0), wrap(i), wrap(c) - wrap(1))
                for (i, c) in zip(indices, dims))...)
    resize!(loopVars, l)
    return TComputation(
        new_computation_id(),
        cond,
        length(dims),
        TEmptyExpr(),
        typ,
        access,
        indices,
        true,
    )
end

used_buffers = Set{String}()

function createVar(id, typ, rhs)
    parsed_rhs = parseNode(rhs)
    if isa(parsed_rhs, TAlloc)
        vars[id] = JVar(
            id,
            parsed_rhs.typ,
            TBuffer(idToBufferName(id),
                    parsed_rhs.dims,
                    parsed_rhs.typ,
                    TBTemporary,
                    false),
            createEmptyComputation(parsed_rhs.typ, parsed_rhs.dims, false),
        )
        push!(readGenvars, vars[id])
        push!(used_buffers, idToBufferName(id))
    else
        if !haskey(vars, id) # in case of reductions
            vars[id] = JVar(
                id,
                typ,
                typToTBuffer[typ](id, typ),
                createEmptyComputation(typ, TExpr[wrap(1)], true),
            )
            push!(readGenvars, vars[id])
        end
        actualComputation = parseComputation(parsed_rhs, typ)
        newVar = copy(vars[id])
        newVar.computation = actualComputation
        push!(genvars, newVar)
    end
end

function make_const(e::TExpr,
                    typ::DataType)
    C = TConstant(
        new_constant(),
        e,
        typ,
    )
    return C
end

function createInputSizeVar(n::Int, idx)
    indices = Variable[new_genslot()]
    id = new_genslot()
    map(i -> dummy(i, Int), indices)
    push!(loopVars, indices[1])
    access = TAccess(map(wrap, indices), [wrap(indices[1])])
    cond = <=(wrap(0), wrap(indices[1]), wrap(n) - wrap(1))
    b = TBuffer(idToBufferName(id),
                TExpr[wrap(n)],
                Int,
                TBInput,
                false)
    comp = TComputation(
        new_computation_id(),
        cond,
        n,
        TEmptyExpr(),
        Int,
        access,
        indices,
        true,
    )
    pop!(loopVars)

    vars[id] = JVar(
        id,
        Int,
        b,
        comp,
    )
    push!(inputSizeGenvars, vars[id])
    push!(signature, InputSizeSig(id, idx))

    bounds = TExpr[]

    for i in 0:(n-1)
        push!(bounds, TIndexedRead(vars[id], TExpr[wrap(i)]))
    end

    return bounds
end

function createVarInput(id, typ, idx)
    element_type = eltype(typ)
    n = ndims(typ)
    if n == 0
        indices = Variable[new_genslot()]
        push!(loopVars, indices[1])
        map(i -> dummy(i, Int), indices)
        access = TAccess(map(wrap, indices), [wrap(0)])
        b = typToTBuffer[typ](id, typ)
        b.argument = TBInput
        cond = wrap(indices[1]) == wrap(0)
        n = 1
        pop!(loopVars)
    else
        bounds = createInputSizeVar(n, idx)
        l = length(loopVars)
        indices = Variable[new_genslot() for _ in 1:n]
        append!(loopVars, indices)
        map(i -> dummy(i, Int), indices)
        access = TAccess(map(wrap, indices), [wrap(i) for i in indices])
        b = TBuffer(idToBufferName(id),
                    bounds,
                    element_type,
                    TBInput,
                    false)
        cond = and((<=(wrap(0), wrap(i), wrap(c) - wrap(1))
                    for (i, c) in zip(indices, bounds))...)
        resize!(loopVars, l)
    end
    comp = TComputation(
        new_computation_id(),
        cond,
        n,
        TEmptyExpr(),
        element_type,
        access,
        indices,
        true,
    )

    vars[id] = JVar(
        id,
        eltype(typ),
        b,
        comp,
    )

    push!(readGenvars, vars[id])
    push!(signature, InputSig(id, idx))
end

metaHandlers = Dict(
    :parallel => function(args)
        sched.parallelize_level = length(loopVars)
    end,

    :endparallel => function(args)
        sched.parallelize_level = nothing
    end,

    :fuse => function(args)
        # do nothing
    end,

    :endfuse => function(args)
        sched.fuse_level = OnceSchedule(Unsigned(args[1]), SchedNow)
    end,

)

exprHandlers = Dict(
    # In case of call, lookup the function in the list of supported functions
    :call => function(expr)
        called_function_id = (expr.args[1].mod, expr.args[1].name)
        return get(callHandlers,
                   called_function_id,
                   function(args) end)(expr.args[2:end], expr.typ)
    end,

    # In case of assignment, create buffer, computation
    :(=) => function(expr)
        id = expr.args[1]
        typ = expr.typ
        rhs = expr.args[2]

        if isa(rhs, Variable)
            # TODO: does this break things?
            vars[id] = vars[rhs]
        else
            createVar(id, typ, rhs)
        end
    end,

    :alloc => function(expr)
        return TAlloc([parseNode(n) for n in expr.args[2]], expr.args[1])
    end,

    :meta => function(expr)
        metaHandlers[expr.args[1]](expr.args[2:end])
    end,

    # For return, set buffer to output
    :return => function(expr)
        assert(length(expr.args) == 1) # TODO support tuples
        handleReturn(expr.args[1], expr.typ)
    end,
)

#IMPORTANT: This is the main entry point for Julia to Tiramisu Translation
#This functions performs a number of key steps:
#   -First it gets the header string which includes imports and the namespace
#   -Second it sets the name of the function, which declares the method signature,
#    setting Tiramisu default options, and declaring the function to be generated.
#   -Third it translates the body of the program using the steps outlined in the
#    beginning. This is where the majority of the work takes place
#   -Next it writes the combination of this string to a file
#   -Then the file is compiled, run, and converted to a library file
#   -Finally, the library file name is returned to be used by Julia
function tiramisu_from_root_entry(body, linfo, functionName::AbstractString)

    #Form Tiramisu String
    thdr = tiramisu_from_header(linfo)
    twrap = tiramisu_set_function(functionName)
    tbod = tiramisu_analyze_body(body, linfo)
    tc = thdr * twrap* tbod * "}\n"

    tcw = generate_wrapper()

    println(functionName)

    #Write to File
    fil = open(FILENAME,"w")
    write(fil, tc)
    close(fil)
    #println("Wrote to file")

    #Write to File
    fil = open(FILENAME_WRAPPER,"w")
    write(fil, tcw)
    close(fil)
    #println("Wrote to file")

    #Compile the Function Generator, Generate the function and link
    run(tiramisu_get_command())
    #println("Compiled Tiramisu File")
    run(tiramisu_get_gen())
    #println("Generated Function")
    run(tiramisu_get_lib_obj())
    run(tiramisu_get_lib())
    #println("Converted to .so file")
    TIRAMISU_BUILD_DIR, LIB, get_return()
end

# Converts a buffer to the corresponding type to pass to the C function
function sig_to_ccall_type(sig)
    # need to construct it in this obscure way to be equivalent to
    # the symbol equivalent of writing
    return :(Ptr{$(vars[sig.var].buffer.typ)})
end

function buffer_to_ccall_container(buffer)
    return Array(buffer.typ, [dim.val for dim in buffer.dimensions]...)
end

# TODO: handle Tuple return
function decouple_scalar(containers)
    return containers[1][1]
end

# TODO: handle Tuple return
function decouple_buffer(containers)
    return containers[1]
end

function transform_dim(d::TValue, interface, args)
    return d.val
end

function transform_dim(d::TSimpleRead, interface, args)
    sig = [s for s in interface if d.var.id == s.var][1]
    return args[sig.pos]
end

function transform_generator()
    interface = copy(signature)
    n = length([1 for s in interface if isa(s, InputSig)])
    outputs = Dict(s => vars[s.var] for s in interface if isa(s, OutputSig))
    function transform(args)
        assert(length(args) == n)
        transformed = []
        for s in interface
            if isa(s, InputSizeSig)
                push!(transformed, collect(Int32, size(args[s.pos])))
            elseif isa(s, InputSig)
                arg = args[s.pos]
                if ndims(arg) == 0
                    T = eltype(arg)
                    arg = T[arg]
                end
                push!(transformed, arg)
            else
                dims = (transform_dim(d, interface, args) for d in outputs[s].buffer.dimensions)
                push!(transformed, Array{outputs[s].buffer.typ}(dims...))
            end
        end
        return transformed
    end
end

function return_parser()
    interface = copy(signature)
    returns = []
    is_scalar = Dict(s => vars[s.var].buffer.is_scalar
                     for s in interface if isa(s, OutputSig))
    function get_return(args)
        dump(args)
        for (s, arg) in zip(interface, args)
            if isa(s, OutputSig)
                if is_scalar[s]
                    return arg[1]
                else
                    return arg
                end
            end
        end
    end
end

# returns a tuple of:
# - A tuple of types; (Ptr{Int32},)...
# - The corresponding containers in a list
function get_return()
    return (Expr(:tuple, map(sig_to_ccall_type, signature)...),
            transform_generator(),
            return_parser())
end

#Imports for Tiramisu file
function tiramisu_from_header(linfo)
    s = string("#include <isl/set.h>\n",
        "#include <isl/union_map.h>\n",
        "#include <isl/union_set.h>\n",
        "#include <isl/ast_build.h>\n",
        "#include <isl/schedule.h>\n",
        "#include <isl/schedule_node.h>\n\n",

        "#include <$(TIRAMISU_INCLUDE_DIR)/tiramisu/debug.h>\n",
        "#include <$(TIRAMISU_INCLUDE_DIR)/tiramisu/core.h>\n\n",

        "#include <stdint.h>\n",
        "#include <float.h>\n",
        "#include <limits.h>\n",
        "#include <complex>\n",
        "#include <math.h>\n",
        "#include <stdio.h>\n",
        "#include <iostream>\n",
        "#include <sstream>\n",
        "#include <vector>\n",
        "#include <string>\n",
        "#include <Halide.h>\n\n",
        "using namespace tiramisu;\n\n")
    params = [vdef for vdef in linfo.var_defs
                   if vdef.desc == 0 && vdef.name != Symbol("#self#")]
    for (i, def) in enumerate(params)
        createVarInput(SlotNumber(def.id), def.typ, i)
    end
    return s
end

#Method Signature for Tiramisu File
function tiramisu_set_function(f)
    global function_name
    global computationPrefix
    function_name = f
    computationPrefix = f*"_S"
    s = string("int main()\n{\n",
                "\tglobal::set_default_tiramisu_options();\n",
                "\tglobal::set_loop_iterator_default_data_type(tiramisu::p_int32);\n",
                "\tfunction $function_name(\"$function_name\");\n",
                "\ttiramisu::var ",
                join(("$v(\"$v\")" for v in indexNames), ", "),
                ";\n")
end

getIndex(index::TConstIndex) = string(index.val)

getIndex(index::TVarIndex) = indexNames[index.letter]

getIndices(indices::Vector{TIndex}) = join(map(getIndex, indices), ", ")

VarToString = Dict{Variable, String}

createTExpr(nthg::TEmptyExpr, varMap::VarToString) = "expr()"

createTExpr(c::TConstant, varMap::VarToString) = c.name

createTExpr(val::TValue, varMap::VarToString) = "expr(($(jToC[val.typ])) $(val.val))"
createTExpr(val::TValue, varMap::VarToString, _, __) = "expr(($(jToC[val.typ])) $(val.val))"
function createTExpr(var::TSimpleRead, varMap::VarToString)
    return "S$(var.var.computation.id)(0)"
end

function createTExpr(var::TSimpleRead, varMap::VarToString, constSet, used)
    pair = Pair{Variable, Int}(var.var.id, 0)
    name = "C$(constSet[pair].id.id)"
    push!(used, constSet[pair])
    return name
end

function createTExpr(op::TOp, varMap::VarToString)
    if op.style == TOpInfix
        return "(" * join((createTExpr(a, varMap) for a in op.args), " $(op.fcn) ") * ")"
    end
    return op.fcn * "(" * join((createTExpr(a, varMap) for a in op.args), ", ") * ")"
end

function createTExpr(op::TOp, varMap::VarToString, constSet, used)
    if op.style == TOpInfix
        return "(" * join((createTExpr(a, varMap, constSet, used) for a in op.args), " $(op.fcn) ") * ")"
    end
    return op.fcn * "(" * join((createTExpr(a, varMap, constSet, used) for a in op.args), ", ") * ")"
end

function createTExpr(var::TLoopVarRead, varMap::VarToString)
    return varMap[var.var]
end

function createTExpr(var::TIndexedRead, varMap::VarToString)
    # TODO why shouldn't this be reversed
    return ("S$(var.var.computation.id)(" *
            join((createTExpr(v, varMap) for v in var.indices), ", ") *
            ")")
end

function createTExpr(var::TIndexedRead, varMap::VarToString, constSet, used)
    assert(length(var.indices) == 1 && isa(var.indices[1], TValue))
    pair = Pair{Variable, Int}(var.var.id, var.indices[1].val)
    name = "C$(constSet[pair].id.id)"
    push!(used, constSet[pair])
    return name
end

createISLCond(val::TValue, varMap::VarToString, _, __, ___) = string(val.val)

function createISLCond(var::TLoopVarRead, varMap::VarToString, _, __, ___)
    return varMap[var.var]
end

function createISLCond(var::TSimpleRead, varMap::VarToString, constSet, createConst, usedConsts)
    pair = Pair{Variable, Int}(var.var.id, 0)
    if !haskey(constSet, pair)
        constSet[pair] = make_const(var, var.var.typ)
        push!(createConst, constSet[pair])
    end
    name = "C$(constSet[pair].id.id)"
    push!(usedConsts, name)
    return name
end

function createISLCond(var::TIndexedRead, varMap::VarToString, constSet, createConst, usedConsts)
    assert(length(var.indices) == 1 && isa(var.indices[1], TValue))
    pair = Pair{Variable, Int}(var.var.id, var.indices[1].val)
    if !haskey(constSet, pair)
        constSet[pair] = make_const(var, var.var.typ)
        push!(createConst, constSet[pair])
    end
    name = "C$(constSet[pair].id.id)"
    push!(usedConsts, name)
    return name
end

function createISLCond(op::TOp, varMap::VarToString,
                       constSet, createConst, usedConsts)
    createCond = arg -> createISLCond(arg, varMap, constSet, createConst, usedConsts)
    if op.style == TOpInfix
        return "(" * join(map(createCond, op.args), " $(opToISL(op.fcn)) ") * ")"
    end
    return opToISL(op.fcn) * "(" * join(map(createCond, op.args), ", ") * ")"
end

function getIndices(comp::TComputation)
    return join(indexNames[1:length(comp.vars)], ", ")
end

function createISL(comp::TComputation, varMap::VarToString, constSet, createConst)
    for (var, index) in zip(comp.vars, indexNames)
        varMap[var] = string(index)
    end
    usedConsts = Set{String}()
    # dump(comp)
    # dump(comp.vars)
    cond = createISLCond(comp.condition, varMap, constSet, createConst, usedConsts)
    usedConsts = join(usedConsts, ", ")
    indices = getIndices(comp)
    return "[$(usedConsts)] -> {S$(comp.id)[$indices] : $cond}"
end

function createAccess(access::Vector{TExpr}, varMap, constSet, createConst)
    return join((createISLCond(expr, varMap, constSet, createConst, Set{String}()) for expr in access), ",")
end

function createAccess(comp::TComputation, varMap::VarToString, constSet, createConst)
    for (var, index) in zip(comp.vars, indexNames)
        varMap[var] = string(index)
    end
    return (createAccess(comp.access.computation, varMap, constSet, createConst),
            createAccess(reverse(comp.access.buffer), varMap, constSet, createConst))
end

createSchedLevel(level::SchedLevel) = ((level == RootDimension) ?
                                       "tiramisu::computation::root" :
                                       indexNames[level + 1]) # Julia is 1 indexed

function get_arg_type(buffer)
    if buffer.name in used_buffers && buffer.argument == TBInput
        return argNames[TBOutput]
    end
    return argNames[buffer.argument]
end

function tiramisu_analyze_body(ast::Expr, linfo)
    # dump(ast.args)
    # throw(InterruptException())
    for expr in ast.args
        dump(expr)
        parseNode(expr)
    end
    uvars = unique(genvars)
    # computation pass
    lines = Vector{String}()
    constSet = Dict{Pair{Variable, Int}, TConstant}()

    println("Input Computations")
    for v in inputSizeGenvars
        computation = v.computation
        varMap = VarToString()
        createConst = TConstant[]
        push!(lines,
              "\ttiramisu::computation S$(computation.id)(" *
              "\"" * createISL(computation, varMap, constSet, createConst) * "\", " *
              createTExpr(computation.expr, varMap) * ", " *
              string(!computation.is_input) * ", " *
              jToT[computation.typ] * ", " *
              "&$function_name);")
    end
    # for c in consts
    #     varMap = VarToString()
    #     push!(lines,
    #           "\ttiramisu::constant $(c.name)(" *
    #           "\"$(c.name)\", " *
    #           createTExpr(c.expr, varMap) * ", " *
    #           jToT[c.typ] * ", " *
    #           "true, " *
    #           "nullptr, " *
    #           "0, " *
    #           "&$function_name);")
    # end
    println("Read Computation")
    computations = sort([var.computation for var in readGenvars], by=(c -> c.id))
    for computation in computations
        varMap = VarToString()
        createConst = TConstant[]
        # dump(computation)
        result = (
              "\ttiramisu::computation S$(computation.id)(" *
              "\"" * createISL(computation, varMap, constSet, createConst) * "\", " *
              createTExpr(computation.expr, varMap) * ", " *
              "false, " *
              jToT[computation.typ] * ", " *
              "&$function_name);")
        for c in createConst
            varMap = VarToString()
            push!(lines,
                  "\ttiramisu::constant C$(c.id.id)(" *
                  "\"C$(c.id.id)\", " *
                  createTExpr(c.expr, varMap) * ", " *
                  jToT[c.typ] * ", " *
                  "true, " *
                  "nullptr, " *
                  "0, " *
                  "&$function_name);")
        end
        push!(lines, result)
    end
    println("Computation")
    computations = sort([var.computation for var in uvars], by=(c -> c.id))
    for computation in computations
        # dump(computation)
        varMap = VarToString()
        createConst = TConstant[]
        result = (
              "\ttiramisu::computation S$(computation.id)(" *
              "\"" * createISL(computation, varMap, constSet, createConst) * "\", " *
              createTExpr(computation.expr, varMap) * ", " *
              string(!computation.is_input) * ", " *
              jToT[computation.typ] * ", " *
              "&$function_name);")
        for c in createConst
            varMap = VarToString()
            push!(lines,
                  "\ttiramisu::constant C$(c.id.id)(" *
                  "\"C$(c.id.id)\", " *
                  createTExpr(c.expr, varMap) * ", " *
                  jToT[c.typ] * ", " *
                  "true, " *
                  "nullptr, " *
                  "0, " *
                  "&$function_name);")
        end
        push!(lines, result)
    end
    println("Buffers")
    buffers_to_allocate = TBuffer[]
    buffers_to_allocate_lets = Set[]
    for var in vcat(inputSizeGenvars, readGenvars)
        varMap = VarToString()
        createConst = TConstant[]
        cIndices, bIndices = createAccess(var.computation, varMap, constSet, createConst)
        used = Set()
        push!(lines,
              "\ttiramisu::buffer $(var.buffer.name)(" *
              "\"$(var.buffer.name)\", " *
              "{" * join((createTExpr(d, varMap, constSet, used) for d in var.buffer.dimensions), ", ") * "}, " *
              jToT[var.buffer.typ] * ", " *
              get_arg_type(var.buffer) * ", " *
              "&$function_name);")
        push!(lines,
              "\tS$(var.computation.id).set_access(\"" *
              "{S$(var.computation.id)[$cIndices] -> $(var.buffer.name)[$bIndices]}" *
              "\");")
        if !isempty(used) && var.buffer.argument == TBTemporary
            push!(buffers_to_allocate, var.buffer)
            push!(buffers_to_allocate_lets, used)
        end
    end
    for var in uvars
        varMap = VarToString()
        createConst = TConstant[]
        cIndices, bIndices = createAccess(var.computation, varMap, constSet, createConst)
        push!(lines,
              "\tS$(var.computation.id).set_access(\"" *
              "{S$(var.computation.id)[$cIndices] -> $(var.buffer.name)[$bIndices]}" *
              "\");")
    end
    # Scheduling
    println("Scheduling")
    # const_sched_list = collect(values(constSet))
    # for (i, (b, cs)) in enumerate(zip(buffers_to_allocate, buffers_to_allocate_lets))
    #
    #     if i == 1
    #         push!(lines, "\ttiramisu::computation* _allocate_$(b.name) = $(b.name).allocate_at(S$(sched.schedule[1].first.first.id), tiramisu::computation::root_dimension);");
    #     else
    #         push!(lines, "\ttiramisu::computation* _allocate_$(b.name) = $(b.name).allocate_at(*_allocate_$(buffers_to_allocate[i - 1].name), tiramisu::computation::root_dimension);");
    #     end
    #     varMap = VarToString()
    #     for c in cs
    #         push!(lines, "\t_allocate_$(b.name)->add_associated_let_stmt(\"C$(c.id.id)\", $(createTExpr(c.expr, varMap)));");
    #     end
    # end
    # # push!(const_sched_list, sched.schedule[1].first.first)
    # reverse!(buffers_to_allocate)
    # for i in 1:(length(const_sched_list) - 1)
    #     (c1, c2) = (const_sched_list[i], const_sched_list[i + 1])
    #     push!(lines, "\tC$(c2.id.id).after(C$(c1.id.id), tiramisu::computation::root);")
    # end
    # if !isempty(buffers_to_allocate)
    #     cnst = const_sched_list[end]
    #     b = buffers_to_allocate[1]
    #     push!(lines, "\t_allocate_$(b.name)->after(C$(cnst.id.id), tiramisu::computation::root);")
    #     for i in 1:(length(buffers_to_allocate) - 1)
    #         (b1, b2) = (buffers_to_allocate[i], buffers_to_allocate[i + 1])
    #         push!(lines, "\t_allocate_$(b2.name)->after(*_allocate_$(b1.name), tiramisu::computation::root);")
    #     end
    #     if !isempty(sched.schedule)
    #         b = buffers_to_allocate[end]
    #         comp = sched.schedule[1].first.first
    #         push!(lines, "\tS$(comp.id).after(*_allocate_$(b.name), tiramisu::computation::root);")
    #     end
    # elseif !isempty(const_sched_list) && !isempty(sched.schedule)
    #     cnst = const_sched_list[end]
    #     comp = sched.schedule[1].first.first
    #     push!(lines, "\tS$(comp.id).after(C$(cnst.id.id), tiramisu::computation::root);")
    # end
    for ((c1, c2), level) in sched.schedule
        # level = Int(c2.dim) - level - 1
        push!(lines, "\tS$(c2.id).after(S$(c1.id), $(createSchedLevel(level)));")
    end

    for (c, level) in sched.parallelize
        # level = Int(c.dim) - level - 1
        push!(lines, "\tS$(c.id).tag_parallel_level($(createSchedLevel(level)));")
    end

    println("Interface")
    interface = [vars[s.var].buffer.name for s in signature]
    push!(lines, "\t$function_name.set_arguments({" *
                  join(map(s -> "&" * s, interface), ", ")
                  *"});")
    push!(lines, "\t$function_name.gen_time_space_domain();")
    push!(lines, "\t$function_name.gen_isl_ast();")
    push!(lines, "\t$function_name.gen_halide_stmt();")
    push!(lines, "\t$function_name.gen_halide_obj(\"$OBJ\");")
    push!(lines, "\treturn 0;")
    return join(lines, '\n')
end

function buffer_to_arg_prototype(buffer)
    return (jToC[buffer.typ] *
            # repeat("*", length(buffer.dimensions)) *
            "* $(buffer.name)")
end

getdim(dim::TValue) = dim.val
getdim(dim::TSimpleRead) = "$(dim.var.buffer.name)[0]"
function getdim(dim::TIndexedRead)
    name = dim.var.buffer.name
    indices = join(("$(i.val)" for i in dim.indices), ", ")
    return "$name[$indices]"
end

function buffer_to_halide_buffer(buffer)
    lines = String[]
    push!(lines,
          "auto *shape_$(buffer.name) = new halide_dimension_t[$(length(buffer.dimensions))];")
    for (i, dim) in enumerate(reverse(buffer.dimensions))
        push!(lines, "shape_$(buffer.name)[$(i - 1)].min = 0;")
        push!(lines, "shape_$(buffer.name)[$(i - 1)].extent = $(getdim(dim));") # TODO not always TValue
        push!(lines, "shape_$(buffer.name)[$(i - 1)].stride = 1;")
    end
    push!(lines, "Halide::Buffer<> " *
                 "$(buffer.name)_halide($(jToH[buffer.typ]), $(buffer.name), " *
                 "$(length(buffer.dimensions)), shape_$(buffer.name));")
    return join(lines, "\n")
end

function buffer_to_raw_buffer(buffer)
    return "$(buffer.name)_halide.raw_buffer()"
end

function generate_wrapper()
    # computation pass
    lines = Vector{String}()
    push!(lines, "#include \"Halide.h\"")
    push!(lines, "#include <tiramisu/utils.h>")
    push!(lines, "#include <cstdlib>")
    push!(lines, "")
    push!(lines, "int main(int, char**) {")
    push!(lines, "return 0;")
    push!(lines, "}")
    interface = [vars[s.var].buffer for s in signature]
    push!(lines, "extern \"C\" int $function_name(" * join(("halide_buffer_t *" for _ in interface), ", ") * ");")
    push!(lines, "#pragma GCC visibility push(default)")
    push!(lines, "extern \"C\" void $(function_name)_wrapper(" *
                  join(map(buffer_to_arg_prototype, interface), ", ") *
                  ") {")
    push!(lines, "#pragma GCC visibility pop")
    for buffer in interface
        push!(lines, buffer_to_halide_buffer(buffer))
    end
    push!(lines, "$function_name(" *
                  join(map(buffer_to_raw_buffer, interface), ", ") *
                  ");")
    push!(lines, "}")
    return join(lines, '\n')
end

#Returns the command to compile the written Tiramisu file
function tiramisu_get_command()
    s = ["g++","-g","-std=c++11","-O3","-Wall","-Wno-sign-compare","-fno-rtti","-fvisibility=hidden","-fPIC"]
    push!(s,"$FILENAME","-o","$GEN","-I$TIRAMISU_INCLUDE_DIR")
    push!(s,"-I$ISL_INCLUDE_DIR")
    push!(s,"-I/usr/local/include","-I$HALIDE_INCLUDE_DIR")
    push!(s,"-I$HALIDE_TOOLS_DIR","-I$TIRAMISU_BUILD_DIR")
    push!(s,"-L$TIRAMISU_BUILD_DIR", "-ltiramisu")
    push!(s,"-L$ISL_LIB_DIR")
    push!(s,"-L/usr/local/lib","-lisl","-lgmp","-L$HALIDE_LIB_DIR","-lHalide","-ldl")
    push!(s,"-lpthread","-lz","-ljpeg","-ltinfo")#`libpng-config --cflags --ldflags`","-ljpeg")
    Cmd(s)
end

#Returns the command to run the function generator that was compiled by the Tiramisu file
function tiramisu_get_gen()
    env = copy(ENV)
    if haskey(env, "LD_LIBRARY_PATH") && env["LD_LIBRARY_PATH"] != ""
        env["LD_LIBRARY_PATH"]*= ":$HALIDE_INCLUDE_DIR/../lib:/usr/local/lib:$TIRAMISU_BUILD_DIR:$ISL_LIB_DIR"
    else
        env["LD_LIBRARY_PATH"] = "$HALIDE_INCLUDE_DIR/../lib:/usr/local/lib:$TIRAMISU_BUILD_DIR:$ISL_LIB_DIR"
    end
    if haskey(env, "DYLD_LIBRARY_PATH") && env["DYLD_LIBRARY_PATH"] != ""
        env["DYLD_LIBRARY_PATH"]*= ":$HALIDE_INCLUDE_DIR/../lib:$TIRAMISU_BUILD_DIR:$ISL_LIB_DIR"
    else
        env["DYLD_LIBRARY_PATH"] = "$HALIDE_INCLUDE_DIR/../lib:$TIRAMISU_BUILD_DIR:$ISL_LIB_DIR"
    end
    s = Cmd(["$GEN"])
    setenv(s,env)
end

#Returns the command to convert the generated function from an object file to a
#dynamic shared library file. This is necessary because Julia could only use c++
#code stored in libraries as opposed to object files
function tiramisu_get_lib_obj()
    s = ["g++","-c","-std=c++11","-O3","-Wall","-Wno-sign-compare","-fno-rtti","-fvisibility=hidden","-fPIC"]
    push!(s,"-I$TIRAMISU_INCLUDE_DIR")
    push!(s,"-I$ISL_INCLUDE_DIR")
    push!(s,"-I/usr/local/include","-I$HALIDE_INCLUDE_DIR")
    push!(s,"-I$HALIDE_TOOLS_DIR","-I$TIRAMISU_BUILD_DIR")
    push!(s,"-L$ISL_LIB_DIR")
    push!(s,"-L$TIRAMISU_BUILD_DIR", "-ltiramisu")
    push!(s,"-L/usr/local/lib","-lisl","-lgmp","-L$HALIDE_LIB_DIR","-lHalide","-ldl")
    push!(s,"-lpthread","-lz","-ljpeg","-ltinfo")#`libpng-config --cflags --ldflags`","-ljpeg")
    push!(s,"$FILENAME_WRAPPER")
    push!(s,"-shared","-fPIC","-o","$OBJ_WRAPPER")
    Cmd(s)
end

#Returns the command to convert the generated function from an object file to a
#dynamic shared library file. This is necessary because Julia could only use c++
#code stored in libraries as opposed to object files
function tiramisu_get_lib()
    s = ["g++","-g","-std=c++11","-O3","-Wall","-Wno-sign-compare","-fno-rtti","-fvisibility=hidden","-fPIC"]
    push!(s,OBJ)
    push!(s,OBJ_WRAPPER)
    push!(s,"-I$TIRAMISU_INCLUDE_DIR")
    push!(s,"-I$ISL_INCLUDE_DIR")
    push!(s,"-I/usr/local/include","-I$HALIDE_INCLUDE_DIR")
    push!(s,"-I$HALIDE_TOOLS_DIR","-I$TIRAMISU_BUILD_DIR")
    push!(s,"-L$ISL_LIB_DIR")
    push!(s,"-L/usr/local/lib","-lisl","-lgmp","-L$HALIDE_LIB_DIR","-lHalide","-ldl")
    push!(s,"-lpthread","-lz","-ljpeg","-ltinfo")#`libpng-config --cflags --ldflags`","-ljpeg")
    push!(s,"-shared","-fPIC","-o","$LIB")
    Cmd(s)
end

#Canonicalizing method used to remove #'s and other undesirable chars from function name
function tcanonicalize(tok)
    replacedTokens = Set("#")
    scrubbedTokens = Set(",.({}):")
    tokenXlate = Dict(
        '*' => "star",
        '/' => "slash",
        '-' => "minus",
        '!' => "bang",
        '.' => "dot",
        '^' => "hat",
        '|' => "bar",
        '&' => "amp",
        '=' => "eq",
        '\\' => "backslash"
    )
    s = string(tok)
    s = replace(s, scrubbedTokens, "")
    s = replace(s, replacedTokens, "t")
    s = replace(s, "∇", "del")
    for (k,v) in tokenXlate
       s = replace(s, k, v)
    end
    s = replace(s, r"[^a-zA-Z0-9]", "_")
    s
end

end

max_stars_repo_name: malekbr/julia-to-tiramisu

max_issues_repo_name: malekbr/julia-to-tiramisu

max_forks_repo_name: malekbr/julia-to-tiramisu


text:
using MittagLefflerFunctions
import MittagLefflerFunctions.Quadrature.MLQuad
using RationalApproximations
using PyPlot
using Printf

G(t) = (1-t)/(1+t)
big_α_vals = collect(1:5) ./ big(5)
α_vals = collect(1:5) ./ 5
β = big(1.0)
#big_Eα = [ MLQuad(α, β, 20, :hyperbola) for α in big_α_vals ]
Eα = [ MLQuad(α, 1.0, 14, :hyperbola) for α in α_vals ]
fα = Vector{Function}(undef, 5)
for col = 1:5
#    fα[col] = ( t -> convert(Float64, (big_Eα[col](-G(big(t))))) )
    fα[col] = ( t -> Eα[col](-G(t)) )
end
t = range(-1, 1, length=800)
max_m = 10
tol = 1e-10
max_iterations = 5

rows = 10
cols = length(α_vals)
errors = Array{Float64}(undef, rows, cols)
δ = similar(errors)
for col = 1:cols
    α = α_vals[col]
    println("α = ", α)
    F = fα[col].(t)
    r, err = aaa_v2(F, t, max_m, tol)
    M = length(r)
    for row = 1:rows
        if row ≤ M
            m = row
            resid = F - r[m].(t)
            idx = locate_extrema_idx(resid, m)
            if length(idx) ≠ 2m
                errors[row,col] = NaN
                δ[row,col] = NaN
            else
                supp_pt = r[m].supp_pt
                opt_r, zmin, zmax = minimax(fα[col], t[idx], supp_pt,
                                    (-1.0,1.0), max_iterations)
                errors[row,col] = zmax[end]
                δ[row,col] = zmax[end] - zmin[end]
            end
        else
            errors[row,col] = NaN
            δ[row,col] = NaN
        end
    end
end

@printf("%5s ", "m")
for col = 1:cols
    @printf("&    α = %3.1f", α_vals[col])
end
@printf("\n\n")

for row = 1:rows
    @printf("%5d ", row)
    for col = 1:cols
        @printf("& %10.2e", errors[row,col])
    end
    @printf("\\\\\n")
end


text:
suite["gauge"] = BenchmarkGroup(["build_envs", "orthogonalize", "gaugeQR"])

col = 2
next_col = 1

#suite["gauge"]["gaugeQR"] = @benchmarkable PEPS.gaugeQR(B, $col, :left; maxdim=$D, cutoff=0.0) setup=(B=copy($A))
Q, QR_inds, next_col_inds = PEPS.initQs(A, col, next_col; maxdim=D, cutoff=0.0)
dummy_nexts   = [Index(ITensors.dim(next_col_inds[row]), "DM,Site,r$row") for row in 1:Ny]
suite["gauge"]["build_envs"] = @benchmarkable PEPS.buildEnvs(B, $Q, $next_col_inds, $QR_inds, $col) setup=(B=copy($A))
Envs = PEPS.buildEnvs(A, Q, next_col_inds, QR_inds, col)
suite["gauge"]["orthogonalize_q"] = @benchmarkable PEPS.orthogonalize_Q!($Envs, Q_, $A, $QR_inds, $next_col_inds, $dummy_nexts, $col, true; maxdim=D, cutoff=0.0) setup=(Q_=copy($Q))


text:
function download_dep(orig, dest, hash)
    dest_file = joinpath("data", dest)
    if isfile(dest_file)
        dest_hash = open(dest_file, "r") do f
            bytes2hex(sha256(f))
        end
        if dest_hash == hash
            return nothing
        end
    end
    mkpath("data")
    download(orig, dest_file)
    return nothing
end

@testset "ccd2rgb" begin
    download_dep("http://chandra.harvard.edu/photo/2009/casa/fits/casa_0.5-1.5keV.fits", "casa_0.5-1.5keV.fits",
                "5794b9ebced6b991a3e53888d129a38fbf4309250112be530cb6442be812dea6")
    download_dep("http://chandra.harvard.edu/photo/2009/casa/fits/casa_1.5-3.0keV.fits", "casa_1.5-3.0keV.fits",
                "a48b2502ceb979dfad0d05fd5ec19bf3e197ff2d1d9c604c9340992d1bf7eec9")
    download_dep("http://chandra.harvard.edu/photo/2009/casa/fits/casa_4.0-6.0keV.fits", "casa_4.0-6.0keV.fits",
                "15e90a14515c121c2817e97b255c604ad019c9c2340fda4fb6c5c3da55e1b0c2")
    download_dep("https://bintray.com/aquatiko/AstroImages.jl/download_file?file_path=ccd2rgb_rounded.jld","ccd2rgb_rounded.jld",
                "5191e59e527c3667486c680e92c8f77fcdbed1e82d3230317a514d928092107d")

    f(x) = isnan(x) ? RGB.(0.0,0.0,0.0) : x
    r = FITS(joinpath("data","casa_0.5-1.5keV.fits"))[1]
    b = FITS(joinpath("data","casa_1.5-3.0keV.fits"))[1]
    g = FITS(joinpath("data","casa_4.0-6.0keV.fits"))[1]
    linear_res = ccd2rgb(r, b, g, shape_out = (1000,1000))
    asinh_res = ccd2rgb(r, b, g, shape_out = (1000,1000), stretch = asinh)
    linear_res = f.(RGB.(colorview(RGB, round.(red.(linear_res), digits = 8), round.(green.(linear_res), digits = 8),
                                                round.(blue.(linear_res), digits = 8))))
    asinh_res = f.(RGB.(colorview(RGB, round.(red.(asinh_res), digits = 8), round.(green.(asinh_res), digits = 8),
                                                round.(blue.(asinh_res), digits = 8))))

    linear_ans = f.(load(joinpath("data","ccd2rgb_rounded.jld"), "linear"))
    asinh_ans = f.(load(joinpath("data","ccd2rgb_rounded.jld"), "asinh"))

    @test isapprox(red.(linear_res), red.(linear_ans), nans = true, rtol = 3e-5)
    @test isapprox(blue.(linear_res), blue.(linear_ans), nans = true, rtol = 3e-5)
    @test isapprox(green.(linear_res), green.(linear_ans), nans = true, rtol = 3e-5)

    @test isapprox(red.(asinh_res), red.(asinh_ans), nans = true, rtol = 3e-5)
    @test isapprox(blue.(asinh_res), blue.(asinh_ans), nans = true, rtol = 3e-5)
    @test isapprox(green.(asinh_res), green.(asinh_ans), nans = true, rtol = 3e-5)

    @testset "AstroImage using ccd2rgb and properties" begin
        img = AstroImage(RGB, (joinpath("data","casa_0.5-1.5keV.fits"), joinpath("data","casa_1.5-3.0keV.fits"),
                                joinpath("data","casa_4.0-6.0keV.fits")))

        @test RGB.(img.property.rgb_image) isa Array{RGB{Float64},2}
        @test img.property.brightness == 0.0
        @test img.property.contrast == 1.0
        @test img.property.label == []

        test_img = deepcopy(img)
        set_brightness!(img, 1.5)
        @test img.property.brightness == 1.5
        @test isapprox(red.(img.property.rgb_image), red.(test_img.property.rgb_image) .+ 1.5, nans = true)
        @test isapprox(green.(img.property.rgb_image), green.(test_img.property.rgb_image) .+ 1.5, nans = true)
        @test isapprox(blue.(img.property.rgb_image), blue.(test_img.property.rgb_image) .+ 1.5, nans = true)
        set_brightness!(test_img, 1.5)

        set_contrast!(img, 2.5)
        @test img.property.contrast == 2.5
        @test isapprox(red.(img.property.rgb_image), red.(test_img.property.rgb_image) .* 2.5, nans = true)
        @test isapprox(green.(img.property.rgb_image), green.(test_img.property.rgb_image) .* 2.5, nans = true)
        @test isapprox(blue.(img.property.rgb_image), blue.(test_img.property.rgb_image) .* 2.5, nans = true)
        set_contrast!(test_img, 2.5)

        add_label!(img, 12, 13, "This is an important coordinate")
        add_label!(img, 13, 11, "This a random coordinate")
        @test img.property.label[1] == ((12,13), "This is an important coordinate")
        @test img.property.label[2] == ((13,11), "This a random coordinate")

        @info "Running..."
        reset!(img)
        @test img.property.contrast == 1.0
        @test img.property.brightness == 0.0
        @test img.property.label == []
        ans = ccd2rgb(r, b, g)
        @test isapprox(red.(img.property.rgb_image), red.(ans), nans = true)
        @test isapprox(green.(img.property.rgb_image), green.(ans), nans = true)
        @test isapprox(blue.(img.property.rgb_image), blue.(ans), nans = true)

        img = AstroImage(joinpath("data","casa_0.5-1.5keV.fits"))
        @test_throws DomainError set_brightness!(img, 1.2)
        @test_throws DomainError set_contrast!(img, 1.2)
    end
end


text:
# Generate scenarios of the demands for use within the J-SAA and J+-SAA frameworks

function generateJSAAScenarios(outputs::Array{Float64,2},inputs::Array{Float64,2},covariates_obs::Array{Float64,1},regressionMethod::String)

	numDependentVar::Int64 = size(outputs,2)
	numIndependentVar::Int64 = size(inputs,2)
	numSamples::Int64 = size(outputs,1)
	demand_scen_JSAA = zeros(Float64,numSamples,numDependentVar)
	demand_scen_JpSAA = zeros(Float64,numSamples,numDependentVar)

	if(size(outputs,1) != size(inputs,1))
		throw(ErrorException("Number of output samples != number of input samples"))
	end

	if(regressionMethod == "ols")

		if(numSamples < numIndependentVar)
			throw(ErrorException("Expected number of samples > number of covariates in OLS!"))
		end

		XtX_fact = lufact((inputs')*inputs)
		XtY = (inputs')*outputs

		# first, compute the OLS coefficients using all data points
		coeff_full = zeros(Float64,numIndependentVar,numDependentVar)
		for d = 1:numDependentVar
			coeff_full[:,d] = XtX_fact\(XtY[:,d])
		end

		# next, compute the leave-one-out coefficients
		for s = 1:numSamples
			coeff_loo = zeros(Float64,numIndependentVar,numDependentVar)

			if(numSamples == numIndependentVar)
				range_loo = union([1:s-1,s+1:numSamples]...)
				for d = 1:numDependentVar
					coeff_loo[:,d] = (inputs[range_loo,:])\(outputs[range_loo,d])
				end
			else
				rhs_pert = -inputs[s,:]*((outputs[s,:])')
				term1 = XtX_fact\(inputs[s,:])
				for d = 1:numDependentVar
					term2 = XtX_fact\(rhs_pert[:,d])
					term3 = coeff_full[:,d] + term2

					coeff_loo[:,d] = term3 + ((((inputs[s,:])')*term3)/(1.0 - ((inputs[s,:])')*term1))*term1
				end
			end

			demand_scen_JSAA[s,:] = ((covariates_obs')*coeff_full)' + outputs[s,:] - (((inputs[s,:])')*coeff_loo)'
			demand_scen_JpSAA[s,:] = ((covariates_obs')*coeff_loo)' + outputs[s,:] - (((inputs[s,:])')*coeff_loo)'
		end

	elseif(regressionMethod == "lasso")

		# first, compute the OLS coefficients using all data points
		coeff_full = getLassoCoeff(outputs,inputs)

		# next, compute the leave-one-out coefficients
		for s = 1:numSamples
			range_loo = union([1:s-1,s+1:numSamples]...)
			coeff_loo = getLassoCoeff(outputs[range_loo,:],inputs[range_loo,:])

			demand_scen_JSAA[s,:] = ((covariates_obs')*coeff_full)' + outputs[s,:] - (((inputs[s,:])')*coeff_loo)'
			demand_scen_JpSAA[s,:] = ((covariates_obs')*coeff_loo)' + outputs[s,:] - (((inputs[s,:])')*coeff_loo)'
		end

	end

	return demand_scen_JSAA, demand_scen_JpSAA
end

max_stars_repo_name: rohitkannan/DD-SAA

max_issues_repo_name: rohitkannan/DD-SAA

max_forks_repo_name: rohitkannan/DD-SAA


text:
using SpecialFunctions, StatsFuns

@testset "logmvgamma" begin
    @testset "type behavior" for eltya in (Float32, Float64)
        p = rand(1:50)
        a = rand(eltya)
        # add p since loggamma is only define for positive arguments
        @test typeof(logmvgamma(p, a + p)) == eltya
    end

    @testset "consistent with loggamma" for eltya in (Float32, Float64)
        #  Γ₁(a) = Γ(a), Γ₂(a) = √π Γ(a) Γ(a - 0.5), etc
        a = rand(eltya) + 1 # add one since loggamma is only define for positive arguments
        @test logmvgamma(1, a) ≈ loggamma(a)
        @test logmvgamma(2, a) ≈    eltya(0.5logπ) + loggamma(a) + loggamma(a - eltya(0.5))
        @test logmvgamma(3, a) ≈ eltya((3/2)*logπ) + loggamma(a) + loggamma(a - eltya(0.5)) + loggamma(a - one(a))
    end

    @testset "consistent with itself" for eltya in (Float32, Float64)
        #  Γᵢ(a) = (π^{i-1/2}) Γ(a) Γᵢ₋₁(a - 0.5)
        for p in 1:50
            a = rand(eltya) + p # add p since loggamma is only define for positive arguments
            @test logmvgamma(p, a) ≈ eltya((p/2-1/2)*logπ) + loggamma(a) + logmvgamma(p - 1, a - eltya(0.5))
        end
    end
end

@testset "logmvbeta" begin
    @testset "symmetry" for eltya in (Float32, Float64)
                        for eltyb in (Float32, Float64)
            #  Bᵢ(a, b) = Bᵢ(b, a)
            for p in 1:50
                a = rand(eltya) + p
                b = rand(eltyb) + p
                @test logmvbeta(p, a, b) ≈ logmvbeta(p, b, a)
            end
        end
    end

    @testset "consistent with logbeta" for eltya in (Float32, Float64)
                                       for eltyb in (Float32, Float64)
            #  B₁(a, b) = B(a, b)
            a = rand(eltya)
            b = rand(eltyb)
            @test logmvbeta(1, a, b) ≈ logbeta(a, b)
        end
    end

    @testset "type promotion behaves" for eltya in (Float32, Float64)
                                      for eltyb in (Float32, Float64)
            a = rand(eltya)
            b = rand(eltyb)
            T = Base.promote_eltype(eltya, eltyb)
            @test typeof(logmvbeta(1, a, b)) == T
        end
    end
end

# Owen's T function tests
# test values for accurate and precise calculation
hvec = [0.0625, 6.5, 7.0, 4.78125, 2.0, 1.0, 0.0625, 1, 1, 1, 1, 0.5, 0.5, 0.5, 0.5, 0.25, 0.25, 0.25, 0.25, 0.125, 0.125, 0.125, 0.125, 0.0078125
, 0.0078125, 0.0078125, 0.0078125, 0.0078125, 0.0078125, 0.0625, 0.5, 0.9, 2.5, 7.33, 0.6, 1.6, 2.33, 2.33]
avec = [0.25, 0.4375, 0.96875, 0.0625, 0.5, 0.9999975, 0.999999125, 0.5, 1, 2, 3, 0.5, 1, 2, 3, 0.5, 1, 2, 3, 0.5, 1, 2, 3, 0.5, 1, 2, 3, 10, 100
, 0.999999999999999, 0.999999999999999, 0.999999999999999, 0.999999999999999, 0.999999999999999, 0.999999999999999, 0.999999999999999, 0.999999999999999
, 0.99999]
cvec = [big"0.0389119302347013668966224771378", big"2.00057730485083154100907167685e-11", big"6.399062719389853083219914429e-13"
, big"1.06329748046874638058307112826e-7", big"0.00862507798552150713113488319155", big"0.0667418089782285927715589822405"
, big"0.1246894855262192"
, big"0.04306469112078537", big"0.06674188216570097", big"0.0784681869930841", big"0.0792995047488726", big"0.06448860284750375", big"0.1066710629614485"
, big"0.1415806036539784", big"0.1510840430760184", big"0.07134663382271778", big"0.1201285306350883", big"0.1666128410939293", big"0.1847501847929859"
, big"0.07317273327500386", big"0.1237630544953746", big"0.1737438887583106", big"0.1951190307092811", big"0.07378938035365545"
, big"0.1249951430754052", big"0.1761984774738108", big"0.1987772386442824", big"0.2340886964802671", big"0.2479460829231492"
, big"0.1246895548850743676554299881345328280176736760739893903915691894"
, big"0.1066710629614484543187382775527753945264849005582264731161129477"
, big"0.0750909978020473015080760056431386286348318447478899039422181015"
, big"0.0030855526911589942124216949767707430484322201889568086810922629"
, big"5.7538182971139187466647478665179637676531179007295252996453e-14", big"0.0995191725772188724714794470740785702586033387786949658229016920"
, big"0.0258981646643923680014142514442989928165349517076730515952020227"
, big"0.0049025023268168675126146823752680242063832053551244071400100690"
, big"0.0049024988349089450612896251009169062698683918433614542387524648"]

# test values for type stability checking
ht = [0.0625, 0.0, 0.5, 0.5, 0.5]
at = [0.025, 0.5, 0.0, 1.0, +Inf]

@testset "Owen's T" begin

    # check that error for calc is within specification
    @testset "Owen's T value checks" begin

        for i in 1:size(hvec,1)
            h = hvec[i] # h test value
            a = avec[i] # a test value
            c = cvec[i] # the "correct" answer
            t = owent(h,a)

            err = round(log10(abs(t-c)),digits=0)

            @test err ≤ -16.0
        end

    end

    @testset "Owen's T Shortcut Evaluations" begin
        @test owent(-0.0625,0.025) == owent(0.0625,0.025) # if h<0 equals owent(abs(h),a)
        @test owent(0.0,0.025) == atan(0.025)*inv2π  # if h=0, t = atan(a)*inv2π
        @test owent(0.0625,0.0) == 0.0  # when a=0, owent=0
        @test owent(0.0625,-0.025) == -owent(0.0625,0.025) # if a<0, t = -owent(h,abs(a))
        @test owent(0.0625,1.0) == 0.125*erfc(-0.0625*invsqrt2)*erfc(0.0625*invsqrt2) # if a=1, t = (formula)
        @test owent(0.0625,+Inf) == 0.25*erfc(sqrt(0.0625^2)*invsqrt2) # if a=∞, t = (formula)
        @test owent(0.0625,-Inf) == -0.25*erfc(sqrt(0.0625^2)*invsqrt2) # should also work, due to a<0 condition
    end

    @testset "Owen's T type stability" begin

        for T1 in [Float16, Float32, Float64, BigFloat]
            for T2 in [BigFloat, Float64, Float32, Float16]
                for i in 1:size(ht,1)
                    h=T1(ht[i])
                    a=T2(at[i])
                    (p1, p2) = promote(h,a)
                    t=owent(h,a)
                    @test typeof(t) == typeof(p1)
                    #println("T1: ",T1," ",typeof(p1),"\tT2: ",T2," ",typeof(p2),"\t",i,"\th ",h,"\ta ",a,"\tt ",t,"\t",typeof(t)," ",typeof(t)==typeof(p1))
                end
            end
        end

    end # test type stability

end


text:
function to_ndim(T::Type{<: StaticVector{N,ET}}, vec::StaticVector{N2}, fillval) where {N,ET,N2}
    T(ntuple(Val(N)) do i
        i > N2 && return ET(fillval)
        @inbounds return vec[i]
    end)
end

# bs_length(x::NativeFont) = 1 # these are our rules, and for what we do, Vecs are usually scalars
bs_length(x::StaticArray) = 1 # these are our rules, and for what we do, Vecs are usually scalars
bs_length(x::Tuple) = 1 # these are our rules, and for what we do, Vecs are usually scalars
bs_length(x::AbstractArray) = length(x)
bs_length(x::AbstractString) = length(x)
bs_length(x) = 1

# bs_getindex(x::NativeFont, i) = x # these are our rules, and for what we do, Vecs are usually scalars
bs_getindex(x::StaticArray, i) = x # these are our rules, and for what we do, Vecs are usually scalars
bs_getindex(x::Tuple, i) = x # these are our rules, and for what we do, Vecs are usually scalars
bs_getindex(x::AbstractArray, i) = x[i]
bs_getindex(x::AbstractString, i) = x[i]
bs_getindex(x, i) = x

function extract_expr(extract_func, dictlike, args)
    if args.head != :tuple
        error("Usage: args need to be a tuple. Found: $args")
    end
    expr = Expr(:block)
    for elem in args.args
        push!(expr.args, :($(esc(elem)) = $(extract_func)($(esc(dictlike)), $(QuoteNode(elem)))))
    end
    push!(expr.args, esc(args))
    expr
end

"""
usage @exctract scene (a, b, c, d)
"""
macro extract(scene, args)
    extract_expr(getfield, scene, args)
end

"""
Like broadcast but for foreach. Doesn't care about shape and treats Tuples && StaticVectors as scalars.
"""
function broadcast_foreach(f, args...)
    lengths = bs_length.(args)
    maxlen = maximum(lengths)

    # all non scalars should have same length
    if any(x -> !(x in (0, 1, maxlen)), lengths)
        error("All non scalars need same length, Found lengths for each argument: $lengths, $(typeof.(args))")
    end

    # skip if there's a zero length element (like an empty annotations collection, etc)
    # this differs from standard broadcasting logic in which all non-scalar shapes have to match
    0 in lengths && return

    for i in 1:maxlen
        f(bs_getindex.(args, i)...)
    end
    return
end


text:
module MembershipAttack

import ..EvaluateOdds
import Random
using Statistics

function matrixtorowdict(mat)
   matrows = [mat[j,:] for j in 1:size(mat, 1)]
   Dict([(row, count(x -> x == row, matrows)) for row in matrows])
end

function matrixtorowset(mat)
   Set([mat[j,:] for j in 1:size(mat, 1)])
end

const training_dicts = [matrixtorowdict(EvaluateOdds.training_datasets[datasetindex])
      for datasetindex in EvaluateOdds.datasetids]

function get_membership_testdata(datasetindex)
   EvaluateOdds.getdata(datasetindex)[1001:(1000 + EvaluateOdds.nsamples), 1:EvaluateOdds.nvariables]
end

const test_dicts = [matrixtorowdict(get_membership_testdata(datasetindex))
      for datasetindex in EvaluateOdds.datasetids]

function membership_attack(datasetindex, genset::Set, distance)
   trainingdict = training_dicts[datasetindex]
   testdict = test_dicts[datasetindex]
   true_pos = 0
   true_neg = 0
   false_neg = 0
   false_pos = 0

   # nocc: number of occurences of same vector
   for (testsample, nocc) in testdict
      samples_nearby = false
      for gensample in genset
         if sum(abs.(testsample .- gensample)) <= distance
            samples_nearby = true
            break
         end
      end

      if !samples_nearby # attacker assumes: sample is not in training set
         true_neg += nocc
      else # attacker assumes: sample is in the training set
         false_pos += nocc
      end
   end

   for (trainingsample, nocc) in trainingdict
      samples_nearby = false
      for gensample in genset
         if sum(abs.(trainingsample .- gensample)) <= distance
            samples_nearby = true
            break
         end
      end

      if !samples_nearby # attacker assumes: sample is not in training set
         false_neg += nocc
      else # attacker assumes: sample is in the training set
         true_pos += nocc
      end
   end

   (true_pos = true_pos, true_neg = true_neg,
      false_pos = false_pos, false_neg = false_neg,
      distance = distance)
end

const distances = [0, 2, 3, 5, 6, 8, 10]

function train_and_sample(modeltype, trainingdata, epochs)
   ngensamples = size(trainingdata, 1)

   if modeltype == "dbm"
      dbm = EvaluateOdds.my_fitdbm(trainingdata, epochs)
      return EvaluateOdds.BMs.samples(dbm, ngensamples)

   elseif modeltype == "mice"
      ret = []
      try
         model = EvaluateOdds.MICE.fitmice(trainingdata)
         ret = EvaluateOdds.MICE.genmice(model, ngensamples)
      catch ex
      end
      return ret

   elseif modeltype == "im"
      varmeans = mean(trainingdata, dims = 1)
      return float.(rand(Float64, (ngensamples, size(trainingdata,2))) .< varmeans)

   elseif modeltype == "vae"
      vae = EvaluateOdds._initvae()
      EvaluateOdds._trainvae!(vae, trainingdata, epochs)
      return EvaluateOdds.VAEs.samples(vae, ngensamples)

   elseif modeltype == "gan"
      gan = EvaluateOdds._initgan()
      EvaluateOdds._traingan!(gan, trainingdata, epochs)
      return EvaluateOdds.GANs.samples(gan, ngensamples)

   else
      error("Unknown model type $modeltype")
   end
end

function traingenattack(;seed, modeltype, epochs, datasetindex, nsites)
   # train models on parts, sample and combine

   xgens = map(datasetpart -> begin
      Random.seed!(seed)
      xpart = EvaluateOdds.extract_trainingdata(datasetindex, nsites, datasetpart)
      train_and_sample(modeltype, xpart, epochs)
   end, 1:nsites)
   xgens = filter(x -> length(size(x)) == 2, xgens)
   xgen = vcat(xgens...)

   genset = matrixtorowset(xgen)
   [merge((seed = seed, modeltype = modeltype, epochs = epochs, datasetindex = datasetindex, nsites = nsites),
         membership_attack(datasetindex, genset, d)) for d in distances]
end

end # module MembershipAttack

max_stars_repo_name: stefan-m-lenz/dist-gen-comp

max_issues_repo_name: stefan-m-lenz/dist-gen-comp

max_forks_repo_name: stefan-m-lenz/dist-gen-comp