text:
using PastaQ
using ITensors
using Test
using LinearAlgebra
using Random

@testset "hilbertspace" begin
  N = 5
  ψ = productstate(N)
  ρ = MPO(productstate(ψ))
  Λ = randomstate(ψ; mixed=true)

  @test PastaQ.hilbertspace(ψ) == siteinds(ψ)
  @test PastaQ.hilbertspace(ψ) == PastaQ.hilbertspace(ρ)
  @test PastaQ.hilbertspace(ψ) == PastaQ.hilbertspace(Λ)
end

@testset "choi tags and MPO/MPS conversion" begin
  N = 4
  circuit = randomcircuit(4; depth =  4)

  U = runcircuit(circuit; process=true)
  ρ = PastaQ.choimatrix(PastaQ.hilbertspace(U), circuit; noise=("DEP", (p=0.01,)))
  Λ = randomprocess(4; mixed=true)

  @test PastaQ.ischoi(ρ) == true
  @test PastaQ.ischoi(U) == false
  @test PastaQ.ischoi(Λ) == true
  @test PastaQ.haschoitags(U) == false
  @test PastaQ.haschoitags(ρ) == true
  @test PastaQ.haschoitags(Λ) == true

  Ψ = PastaQ.choitags(U)
  @test hastags(inds(Ψ[1]), "Input") == true
  @test hastags(inds(Ψ[1]), "Output") == true
  @test plev(firstind(Ψ[1]; tags="Input")) == 0
  @test plev(firstind(Ψ[1]; tags="Output")) == 0

  V = PastaQ.mpotags(Ψ)
  @test hastags(inds(V[1]), "Input") == false
  @test hastags(inds(V[1]), "Output") == false
  @test plev(inds(V[1]; tags="Qubit")[1]) == 1
  @test plev(inds(V[1]; tags="Qubit")[2]) == 0

  Ψ = PastaQ.unitary_mpo_to_choi_mps(U)
  @test Ψ isa MPS
  @test hastags(inds(Ψ[1]), "Input") == true
  @test hastags(inds(Ψ[1]), "Output") == true
  @test plev(firstind(Ψ[1]; tags="Input")) == 0
  @test plev(firstind(Ψ[1]; tags="Output")) == 0

  V = PastaQ.choi_mps_to_unitary_mpo(Ψ)
  @test V isa MPO
  @test hastags(inds(V[1]), "Input") == false
  @test hastags(inds(V[1]), "Output") == false
  @test plev(inds(V[1]; tags="Qubit")[1]) == 1
  @test plev(inds(V[1]; tags="Qubit")[2]) == 0
end


text:
module CuArrays

import Base: sum, prod, minimum, maximum, getindex, convert
using CUDA: CuArray, CuPtr
using Knet.LibKnet8: @knet8, @knet8r, reduction_ops
using Knet.KnetArrays: checkbetween

include("convert.jl")
include("getindex.jl")
include("reduction.jl")
include("cubytes.jl"); export cuarrays, cubytes
include("jld2.jl")

end

ext: jl


text:
using Test
using Qecsim:QecsimError

# taken from https://github.com/JuliaLang/julia/blob/v1.6.1/test/errorshow.jl
macro except_str(expr, err_type)
    return quote
        let err = nothing
            try
                $(esc(expr))
            catch err
            end
            err === nothing && error("expected failure, but no exception thrown")
            @test typeof(err) === $(esc(err_type))
            buf = IOBuffer()
            showerror(buf, err)
            String(take!(buf))
        end
    end
end

@testset "QecsimError" begin
    @test_throws QecsimError throw(QecsimError("something went wrong!"))
    err_str = @except_str throw(QecsimError("something went wrong!")) QecsimError
    @test err_str == "QecsimError: something went wrong!"
end


text:
using BSON
using Serialization
using ReinforcementLearningBase
"""
Makes a single input from a pair of state and action, for dynamics network
"""
function make_state_action(state::Array{Float32,3}, action::Int, conf::Config)::Array{Float32,3}
	action /= length(conf.action_space)
	action = action .* ones(Float32, (conf.observation_shape[1], conf.observation_shape[2]))
	# Scale the gradient by half at the start of the dynamics function (See paper appendix Training)
	state .*= 2.0f0
	state_action = cat(state, action, dims=3)
	return state_action
end
"""
# Holds the Min-Max values of the tree
Simply brings variables('min' and 'max') and functions('update' and 'normalize')
to the scope.
'update_tree!(min,max,value::Float32)'
'value=normalize_tree_value(min,max,value::Float32)'
"""
mutable struct MinMaxStats
    min::Float32
    max::Float32
end

function update_tree!(treeminmax::MinMaxStats, value::Float32)::Nothing
    treeminmax.min = treeminmax.min < value ? treeminmax.min : value
    treeminmax.max = treeminmax.max > value ? treeminmax.max : value
    return nothing
end

function normalize_tree_value(treeminmax::MinMaxStats, value::Float32)::Float32
    if treeminmax.max > treeminmax.min
        return (value - treeminmax.min) / (treeminmax.max - treeminmax.min)
    else
        return value
    end
end

"""
Parameter to alter the visit count distribution to ensure that the action selection becomes greedier as training progresses.
The smaller it is, the more likely the best action (ie with the highest visit count) is chosen.

Returns:
Positive float.
"""
function visit_softmax_temperature_fn(trained_steps::Int)::Float32
    if trained_steps < 500e3
        return 1.0
    elseif trained_steps < 750e3
        return 0.5
    else
        return 0.25
    end
end

using Distributions:Dirichlet, Categorical
using Parameters:@with_kw
using Flux:softmax

@with_kw mutable struct Node
    visit_count::Int = 0
    to_play::Int = 1
    prior::Float32
    value_sum::Float32 = 0.0
    children::Union{Dict{Int,Node},Nothing} = nothing
    hidden_state::Union{Array{Float32,3},Nothing} = nothing
    reward::Union{Float32,Int} = 0
end

function expanded(node::Node)::Bool
    return !isnothing(node.children)
end

function node_value(node::Node)::Float32
    if node.visit_count == 0
        return 0
    else
        return node.value_sum / node.visit_count
    end
end

"""
We expand a node using the value, reward and policy prediction obtained from the
neural network.
"""
function expand_node!(node::Node, actions, to_play::Int, reward::Float32, policy_logits::Vector{Float32}, hidden_state::Array{Float32,3})::Nothing
    policy_values = softmax([policy_logits[a] for a in actions])
    policy = Dict([(a, policy_values[i]) for (i, a) in enumerate(actions)])
    node.children = Dict([(action, Node(prior=prob)) for (action, prob) in policy])
    node.to_play = to_play
    node.reward = reward
    node.hidden_state = hidden_state
	return nothing
end

"""
At the start of each search, we add dirichlet noise to the prior of the root to
encourage the search to explore new actions.
"""
function add_exploration_noise!(node::Node, dirichlet_α::Float32, exploration_ϵ::Float32)::Nothing
    actions = collect(keys(node.children))
    noise = rand(Dirichlet(length(actions), dirichlet_α))
    for (a, n) in zip(actions, noise)
        node.children[a].prior = node.children[a].prior * (1 - exploration_ϵ) + n * exploration_ϵ
    end
	return nothing
end

"""
Stores MCTS stats like `child_visits` and `root_value` to GameHistory
"""
function store_search_stats!(history::GameHistory, root::Node, action_space::Array{Int})
	children=collect(values(root.children))
	children=filter!(x->!isnothing(x), children)
	sum_visits = sum([child.visit_count for child in children])
	history.child_visits = hcat(history.child_visits, [haskey(root.children, a) ? root.children[a].visit_count / sum_visits : 0.0f0 for a in action_space])
	value=node_value(root)
	append!(history.root_values, value)
end

"""
Generate a new observation with the observation at the index position
and 'num_stacked_observations' past observations and actions stacked.
"""
function get_stacked_observations(history::GameHistory, index::Int, num_stacked_observations::Int)::Array{Float32,3}

    # stacked_observations = Array{Float32}(undef, conf.observation_shape[1:2]...,0)
	stacked_observations = copy(history.observation_history[:,:,:,index])
    for past_observation_index = index-1:-1:index - num_stacked_observations
        if 1 <= past_observation_index
            previous_observation = cat(
				ones(Float32, conf.observation_shape[1:2]) .* history.action_history[past_observation_index],
				history.observation_history[:,:,:,past_observation_index],
				dims=3
				)
		else
            previous_observation = cat(
				zeros(Float32, conf.observation_shape[1:2]),
				zeros(Float32, conf.observation_shape),
				dims=3
				)
			end
        stacked_observations = cat(stacked_observations, previous_observation, dims=3)
    end
    return stacked_observations
end

using Random
global rng = MersenneTwister(1234)

"""
Select the child with the highest UCB score.
"""
function select_child(node::Node, treeminmax::MinMaxStats)::Tuple{Int,Node}
	actions= collect(keys(node.children))
	children= collect(values(node.children))
	ucb_scores=[ucb_score(node, child, treeminmax) for child in children]
    max_ucb = maximum(ucb_scores)
	max_ucbs=findall(x-> x==max_ucb, ucb_scores)
	i = 0
    i = rand(max_ucbs)
	return actions[i], children[i]
end

"""
The score for a node is based on its value, plus an exploration bonus based on the prior.
"""
function ucb_score(parent_node::Node, child::Node, treeminmax::MinMaxStats)::Float32
    pb_c = (log2((parent_node.visit_count + conf.pb_c_base + 1) / conf.pb_c_base)
            + conf.pb_c_init)
    pb_c *= sqrt(parent_node.visit_count) / (child.visit_count + 1)
    prior_score = pb_c * child.prior

    if child.visit_count > 0
        # Mean value Q
        value_score = normalize_tree_value(treeminmax, child.reward + (conf.discount * (length(conf.players) == 1 ? node_value(child) : -node_value(child))))
    else
        value_score = 0
    end
    return prior_score + value_score
end

"""
At the end of a simulation, we propagate the evaluation all the way up the tree
to the root.
"""
function backpropagate!(search_path::Vector{Node}, value::Float32, to_play::Int, treeminmax::MinMaxStats)::Nothing
    if length(conf.players) == 1
        for node in reverse(search_path)
            node.value_sum += value
            node.visit_count += 1
            update_tree!(treeminmax, node.reward + conf.discount * node_value(node))
            value = node.reward + conf.discount * value
        end
    elseif length(conf.players) == 2
        for node in reverse(search_path)
            if node.to_play == to_play
                node.value_sum += value
            else
                node.value_sum -= value
            end
            node.visit_count += 1
            update_tree!(treeminmax, node.reward + conf.discount * node_value(node))
            if node.to_play == to_play
                value = -node.reward
            else
                value = node.reward + conf.discount * value
            end
        end
    else
        ErrorException("backpropagate for more than 2 players is not implemented")
	end
	return nothing
end

"""
# Game Independent Monte Carlo Tree Search algorithm.
To decide on an action, we run N simulations, always starting at the root of
the search tree and traversing the tree according to the UCB formula until we
reach a leaf node.

At the root of the search tree we use the representation function to obtain a
hidden state given the current observation.
We then run a Monte Carlo Tree Search using only action sequences and the network
learned by the network.
"""
function run_mcts(observation::Array{Float32,3}, legal_actions::Vector{Int}, to_play::Int, exploration::Bool, NNs::NamedTuple{(:representation, :prediction, :dynamics), Tuple{Any, Any, Any}})::Node

	root = Node(prior=0.)
	observation = unsqueeze(observation, 4)
	hidden_state = NNs.representation(observation)
	if ndims(hidden_state)==2
		hidden_state=reshape(hidden_state, (conf.observation_shape...,1))
	end
	# hidden_state = unsqueeze(hidden_state, 4)
	root_predicted_value, policy_logits = NNs.prediction(hidden_state)
	hidden_state, root_predicted_value, policy_logits = squeeze.([hidden_state, root_predicted_value, policy_logits])
	reward = 0.0f0
	# policy_logits = unsqueeze(policy_logits, 2)
	@assert !isempty(legal_actions) "Legal actions should not be an empty array. Got $(legal_actions)"
	@assert Set(legal_actions) ⊆ Set(conf.action_space) "Legal actions should be a subset of the action space."
	expand_node!(root, legal_actions, to_play, reward, policy_logits, hidden_state)

	if exploration
        add_exploration_noise!(root, conf.dirichlet_α, conf.exploration_ϵ)
    end

    treeminmax = MinMaxStats(Inf, -Inf)

    max_tree_depth = 0
    for iter = 1:conf.num_iters
		# @info "Running MCTS: currently at iteration $(iter)"
        node = root
		virtual_to_play = to_play
        search_path = [node]
        current_tree_depth = 0
		action=0
        while expanded(node)
            current_tree_depth += 1
            action, node = select_child(node, treeminmax)
            push!(search_path, node)

			# Players play turn by turn
			virtual_to_play = mod1(virtual_to_play + 1, length(conf.players))
        end

        parent = search_path[end-1]
		value, policy_logits = NNs.prediction(unsqueeze(parent.hidden_state, 4))
		# policy_logits = unsqueeze(policy_logits, 2)
		state_action = make_state_action(parent.hidden_state, action, conf)
		state_action = unsqueeze(state_action, 4)
		next_hidden_state, reward = NNs.dynamics(state_action)
		if ndims(next_hidden_state)==2
			next_hidden_state=reshape(next_hidden_state, (conf.observation_shape...,1))
		end
		value, policy_logits, next_hidden_state, reward = squeeze.([value, policy_logits, next_hidden_state, reward])
        expand_node!(node, legal_actions, virtual_to_play, reward[1], policy_logits, next_hidden_state)
        backpropagate!(search_path, value[1], virtual_to_play, treeminmax)
        max_tree_depth = maximum([max_tree_depth, current_tree_depth])
    end
    return root
end

"""
Select action according to the visit count distribution and the temperature.
The temperature is changed dynamically with the visit_softmax_temperature function
in the config.
"""
function select_action(node::Node, temperature::Float32)::Int
    visit_counts = Int32[child.visit_count for child in values(node.children)]
    actions = [action for action in keys(node.children)]
    if temperature == 0.0f0
        action = actions[argmax(visit_counts)]
    elseif temperature == Inf
        action = rand(rng, actions)
    else
        visit_count_distribution = visit_counts.^(1 / temperature)
        visit_count_distribution = visit_count_distribution ./ sum(visit_count_distribution)
        action = actions[rand(rng, Categorical(visit_count_distribution))]
    end
    return action
end

"""
Select opponent action for evaluating MuZero level.
"""
function select_opponent_action(env::AbstractEnv, opponent::String, stacked_observations::Array{Float32,3})::Int
    if opponent == "human"
		p = current_player(env)
		las = legal_action_space(env, p)
        return human_input()
    elseif opponent == "expert"
        return expert_agent()
    elseif opponent == "random"
        @assert las "Legal actions should not be an empty array. Got $(legal_actions)"
        @assert Set(las) ⊆ Set(conf.action_space) "Legal actions should be a subset of the action space."
        return rand(rng, las)
    else
        error("Wrong argument: opponent argument should be self, human, expert or random")
    end
end

"""
Play one game with actions based on the Monte Carlo tree search at each moves.
"""
function play_game(env::AbstractEnv, temperature, render::Bool, opponent::String, muzero_player::Int, NNs::NamedTuple{(:representation, :prediction, :dynamics), Tuple{Any, Any, Any}})::GameHistory
    history = GameHistory(Array{Float32}(undef, conf.observation_shape..., 0), Vector{Int}(), Vector{Float32}(), Vector{Int}(), Matrix{Float32}(undef, length(conf.action_space), 0), Vector{Float32}(), nothing, nothing, nothing)

    done = false

    if render
        render_game(env)
    end

	observation=0
	root=0
	action=0

    while !done && length(history.action_history) <= conf.max_moves
		if !isnothing(conf.temperature_threshold) && length(history.action_history) ≥ conf.temperature_threshold
			temperature = 0.0f0
		end

		if length(history.action_history)==0
			observation = reset!(env)
		end
		p = current_player(env)
    	history.observation_history = cat(history.observation_history, observation, dims=4)
		# @info "Self-Playing one game: $(length(history.action_history)-1) moves played"
        @assert ndims(observation) == 3 "Observation should be 3 dimensional instead of $(ndims(observation)) dimensionnal. Got observation of shape: $(size(observation))"
        stacked_observations = get_stacked_observations(history, lastindex(history.observation_history, 4), conf.stacked_observations)

        # Choose the action
        if opponent == "self" || muzero_player == p
            root = run_mcts(stacked_observations, legal_action_space(env, p), p, true, NNs)
            action = select_action(root, temperature)
        else
            action = select_opponent_action(env, opponent, stacked_observations)
        end

        # observation, reward, done = execute_step(action)
        observation = env(action) # execute step and return observation, changes the current player
        reward = convert(Float32, RLBase.reward(env, p))
        done = is_terminated(env) #TODO instead of initializing state info everytime

        if render
        	println("Played action: $(action)")
            render_game(env)
        end

        store_search_stats!(history, root, conf.action_space)

        append!(history.action_history, action)
        append!(history.reward_history, reward)
        push!(history.to_play_history, p)
    end
    return history
end

function self_play!(env,
training_step,
num_played_games,
num_played_steps,
total_samples,
remote_NNs,
remote_buffer::RemoteChannel{BufferChannel})::Bool

	NNs = fetch(remote_NNs)
	training_step_ = 0
    while training_step_ ≤ conf.training_steps

		training_step_ = take!(training_step)
		temperature = visit_softmax_temperature_fn(training_step_)

		if training_step_ % conf.checkpoint_interval == 0 && training_step_ > 1
			NNs = take!(remote_NNs)
			@info "Networks loaded during SelfPlay" training_step_ #size(fetch(remote_buffer))
		end

		# Explore moves during training mode
		history = play_game(
			env,
			temperature,
			false,
			"self",
			conf.muzero_player,
			NNs
		)
		# @info "One episode of Self-Play finished"
		save_game(history, remote_buffer, num_played_games,
			num_played_steps,
			total_samples,)
    end
	return true
end

function competitive_play!(;buffer_to_disk=false, NNs)::Nothing
# Take the best action (no exploration) in competition mode
history = play_game(
	env,
	0.0f0,
	true,
	length(conf.players) == 1 ? "self" : conf.opponent,
	conf.muzero_player,
	NNs
)
# if buffer_to_disk
	# save_game(history, buffer)
# end
return nothing
end


text:
export PostStackWEM,
ShotProfileWEM,
ShotProfileEWEM,
ShotProfileLSWEM,
ShotProfileLSEWEM,
WaveSep,
ComputeAngles,
SeisPSTM
include("PostStackWEM.jl")
include("ShotProfileWEM.jl")
include("ShotProfileEWEM.jl")
include("ShotProfileLSWEM.jl")
include("ShotProfileLSEWEM.jl")
include("WaveSep.jl")
include("ComputeAngles.jl")
include("SeisPSTM.jl")

max_stars_repo_name: fercarozzi/myseismicjulia

max_issues_repo_name: fercarozzi/myseismicjulia

max_forks_repo_name: fercarozzi/myseismicjulia


text:
import Base.Iterators
import JSON

abstract type ColumnImporter end

mutable struct CatImporter <: ColumnImporter
    colname::AbstractString
    filepath::String
    length::Int64
    levelmap::Dict{String, Int64}
    levelindexmap::Dict{Int64, String}
    levelfreq::Dict{String, Int64}
    isdropped::Bool
    isnumeric::Function
    nas::Set{String}
end

mutable struct NumImporter <: ColumnImporter
    colname::AbstractString
    filepath::String
    length::Int64
    nancount::Int64
    zerocount::Int64
end

function CatImporter(colname::AbstractString, outfolder::String, isnumeric::Function, nas::Set{String}, isdropped)
    filepath = joinpath(outfolder, "$(randstring(10)).dat")
    CatImporter(colname, filepath, 0, Dict{String, Int64}(MISSINGLEVEL => 0), Dict{Int64, String}(0 => MISSINGLEVEL), Dict{String, Int64}(), isdropped, isnumeric, nas)
end

struct DataColumnInfo
    name::AbstractString
    length::Int64
    filename::String
    datatype::String
    levels::Vector{String}
end

function DataColumnInfo(catimporter::CatImporter)
    levelcount = length(catimporter.levelmap)
    datatype = levelcount <= typemax(UInt8) + 1 ? "UInt8" : (levelcount <= typemax(UInt16) + 1 ? "UInt16" : UInt32)
    levels = Vector{String}(levelcount - 1)
    for (k, v) in catimporter.levelmap
        if v > 0
            levels[v] = k
        end
    end
    DataColumnInfo(catimporter.colname, catimporter.length, basename(catimporter.filepath), datatype, levels)
end

function DataColumnInfo(numimporter::NumImporter)
    DataColumnInfo(numimporter.colname, numimporter.length, basename(numimporter.filepath), "Float32", Vector{String}())
end

struct DataHeader
    datacolumns::Vector{DataColumnInfo}
end

function widencatcolumn(frompath::String, topath::String, length::Int64)
    buffer = Array{UInt8}(1)
    open(frompath) do fromfile
        open(topath, "a") do tofile
            for i in 1:length
                readbytes!(fromfile, buffer)
                v::UInt16 = UInt16(buffer[1])
                write(tofile, v)
            end
        end
    end
end

function uint8tonumcolumn(levelindexmap::Dict{Int64, String}, frompath::String, topath::String, length::Int64)
    buffer = Array{UInt8}(1)
    open(frompath) do fromfile
        open(topath, "a") do tofile
            for i in 1:length
                readbytes!(fromfile, buffer)
                level::String = levelindexmap[buffer[1]]
                v = tryparse(Float32, level)
                if isnull(v)
                    write(tofile, NaN32)
                else
                    write(tofile, get(v))
                end
            end
        end
    end
end

function uint16tonumcolumn(levelindexmap::Dict{Int64, String}, frompath::String, topath::String, length::Int64)
    buffer = Array{UInt8}(2)
    open(frompath) do fromfile
        open(topath, "a") do tofile
            for i in 1:length
                readbytes!(fromfile, buffer)
                uint16Arr = reinterpret(UInt16, buffer)
                level::String = levelindexmap[uint16Arr[1]]
                v = tryparse(Float32, level)
                if isnull(v)
                    write(tofile, NaN32)
                else
                    write(tofile, get(v))
                end
            end
        end
    end
end

function importlevels(colimporter::NumImporter, datalines::Vector{Vector{SubString{String}}}, colindex::Integer)
    iostream = open(colimporter.filepath, "a")
    collength = colimporter.length
    for line in datalines
        level = line[colindex]
        v = tryparse(Float32, level)
        if isnull(v)
            write(iostream, NaN32)
        else
            write(iostream, get(v))
        end
        collength += 1
    end
    close(iostream)
    colimporter.length = collength
    colimporter
end

function importlevels(colimporter::CatImporter, datalines::Vector{Vector{SubString{String}}}, colindex::Integer)
    if !colimporter.isdropped

        levelmap = colimporter.levelmap
        levelindexmap = colimporter.levelindexmap
        levelfreq = colimporter.levelfreq
        nas = colimporter.nas
        iostream = open(colimporter.filepath, "a")
        collength = colimporter.length

        for line in datalines
            level = strip(strip(line[colindex]), ['"'])
            if level in nas
                level = MISSINGLEVEL
            end
            levelcount = length(levelmap)
            levelindex = get(levelmap, level, levelcount)
            freq = get(levelfreq, level, 0)
            levelmap[level] = levelindex
            levelindexmap[levelindex] = level
            levelfreq[level] = freq + 1
            if freq == 0 && levelcount == typemax(UInt8) + 1
                newpath = joinpath(dirname(colimporter.filepath), "$(randstring(10)).dat")
                close(iostream)
                widencatcolumn(colimporter.filepath, newpath, collength)
                oldpath = colimporter.filepath
                colimporter.filepath = newpath
                iostream = open(newpath, "a")
                rm(oldpath)
            end
            levelcount = length(colimporter.levelmap)
            if levelcount > typemax(UInt16) + 1
                    colimporter.isdropped = true
                    println("dropped: $(colimporter.colname)")
                    break
            elseif levelcount > typemax(UInt8) + 1
                write(iostream, convert(UInt16, levelindex))
            else
                write(iostream, convert(UInt8, levelindex))
            end
            collength += 1
        end
        colimporter.length = collength
        close(iostream)
        if colimporter.isdropped
            rm(colimporter.filepath)
        end
        if !colimporter.isdropped && colimporter.isnumeric(colimporter.colname, levelfreq)
            newpath = joinpath(dirname(colimporter.filepath), "$(randstring(10)).dat")
            if length(colimporter.levelmap) > typemax(UInt8) + 1
                uint16tonumcolumn(levelindexmap, colimporter.filepath, newpath, collength)
            else
                uint8tonumcolumn(levelindexmap, colimporter.filepath, newpath, collength)
            end
            rm(colimporter.filepath)
            colimporter = NumImporter(colimporter.colname, newpath, collength, 0, 0)
        end
    end
    colimporter
end

function importdata(colimporters::Vector{ColumnImporter}, datalines::Vector{Vector{SubString{String}}})
    colcount = length(colimporters)
    Threads.@threads for i = 1:colcount
        colimporters[i] = importlevels(colimporters[i], datalines, i)
    end
    colimporters
end

function isanylevelnumeric(colname::AbstractString, levelfreq::Dict{String, Int64})
    any([!isnull(tryparse(Float32, strip(k))) for k in keys(levelfreq)])
end

function importcsv(path::String; maxobs::Integer = -1, chunksize::Integer = SLICELENGTH, nas::Vector{String} = Vector{String}(),
                   isnumeric::Function = isanylevelnumeric, drop::Vector{String} = Vector{String}())
    nas = Set{String}(nas)
    push!(nas, "")
    path = abspath(path)
    outfolder = splitext(path)[1]
    outtempfolder = joinpath(dirname(path), randstring(10))
    mkpath(outtempfolder)
    iostream = open(path)
    lineseq = maxobs > -1 ? Iterators.take(Seq(String, iostream, nextline), maxobs + 1) : Seq(String, iostream, nextline)
    lines = map((line -> split(line, ",")), lineseq, Vector{SubString{String}})
    colnames, datalines = lines |> tryread
    if !isnull(colnames)
        colnames = map((c -> strip(strip(c), ['"'])), get(colnames))
        colimporters::Vector{ColumnImporter} = map((colname -> CatImporter(colname, outtempfolder, isnumeric, nas, colname in drop)), colnames)
        colimporters = fold(importdata, colimporters, chunkbysize(datalines, chunksize))
    end
    colimporters = filter(colimporters) do colimp
        isa(colimp, NumImporter) || !colimp.isdropped
    end
    header = DataHeader([DataColumnInfo(colimp) for colimp in colimporters])
    headerjson = JSON.json(header)
    headerpath = joinpath(outtempfolder, "header.txt")
    open(headerpath, "a") do f
        write(f, headerjson)
    end
    mv(outtempfolder, outfolder; remove_destination = true)
end

max_stars_repo_name: vikashdwd/JuML

max_issues_repo_name: vikashdwd/JuML

max_forks_repo_name: vikashdwd/JuML


text:
mutable struct UncertaintyTrapEcon{TF<:AbstractFloat, TI<:Integer}
    a::TF          # Risk aversion
    gx::TF         # Production shock precision
    rho::TF        # Correlation coefficient for theta
    sig_theta::TF  # Std dev of theta shock
    num_firms::TI      # Number of firms
    sig_F::TF      # Std dev of fixed costs
    c::TF          # External opportunity cost
    mu::TF         # Initial value for mu
    gamma::TF      # Initial value for gamma
    theta::TF      # Initial value for theta
    sd_x::TF       # standard deviation of shock
end

function UncertaintyTrapEcon(;a::AbstractFloat=1.5, gx::AbstractFloat=0.5,
                             rho::AbstractFloat=0.99, sig_theta::AbstractFloat=0.5,
                             num_firms::Integer=100, sig_F::AbstractFloat=1.5,
                             c::AbstractFloat=-420.0, mu_init::AbstractFloat=0.0,
                             gamma_init::AbstractFloat=4.0,
                             theta_init::AbstractFloat=0.0)
    sd_x = sqrt(a/gx)
    UncertaintyTrapEcon(a, gx, rho, sig_theta, num_firms, sig_F, c, mu_init,
                        gamma_init, theta_init, sd_x)

end

function psi(uc::UncertaintyTrapEcon, F::Real)
    temp1 = -uc.a * (uc.mu - F)
    temp2 = 0.5 * uc.a^2 * (1/uc.gamma + 1/uc.gx)
    return (1/uc.a) * (1 - exp(temp1 + temp2)) - uc.c
end

"""
Update beliefs (mu, gamma) based on aggregates X and M.
"""
function update_beliefs!(uc::UncertaintyTrapEcon, X::Real, M::Real)
    # Simplify names
    gx, rho, sig_theta = uc.gx, uc.rho, uc.sig_theta

    # Update mu
    temp1 = rho * (uc.gamma*uc.mu + M*gx*X)
    temp2 = uc.gamma + M*gx
    uc.mu =  temp1 / temp2

    # Update gamma
    uc.gamma = 1 / (rho^2 / (uc.gamma + M * gx) + sig_theta^2)
end

update_theta!(uc::UncertaintyTrapEcon, w::Real) =
    (uc.theta = uc.rho*uc.theta + uc.sig_theta*w)

"""
Generate aggregates based on current beliefs (mu, gamma).  This
is a simulation step that depends on the draws for F.
"""
function gen_aggregates(uc::UncertaintyTrapEcon)
    F_vals = uc.sig_F * randn(uc.num_firms)

    M = sum(psi.(uc, F_vals) .> 0)  # Counts number of active firms
    if M > 0
        x_vals = uc.theta + uc.sd_x * randn(M)
        X = mean(x_vals)
    else
        X = 0.0
    end
    return X, M
end


text:
function nelder_mead(fcn::Function,
                     x0::Array,
                     args...;
                     parameters            = Optim.AdaptiveParameters(),
                     initial_simplex       = Optim.AffineSimplexer(),
                     iterations::Int       = 1000,
                     store_trace::Bool     = false,
                     show_trace::Bool      = false,
                     extended_trace::Bool  = false,
                     kwargs...)

    Optim.optimize(fcn, x0,
                   method = Optim.NelderMead(parameters = parameters, initial_simplex = initial_simplex),
                   iterations = iterations, store_trace = store_trace, show_trace = show_trace,
                   extended_trace = extended_trace)
end

type MatlabSimplexer <: Optim.Simplexer
    a::Float64
    b::Float64
end
MatlabSimplexer(;a = 0.00025, b = 0.05) = MatlabSimplexer(a, b)

function Optim.simplexer{T, N}(A::MatlabSimplexer, initial_x::Array{T, N})
    n = length(initial_x)
    initial_simplex = Array{T, N}[initial_x for i = 1:n+1]
    for j = 1:n
        initial_simplex[j+1][j] += initial_simplex[j+1][j] == zero(T) ? A.b * initial_simplex[j+1][j] : A.a
    end
    initial_simplex
end


text:
"""
    plot_graphviz(g, node_label= true, edge_label=false; path = zeros(Int, nv(mat))

Render graph `g` in iJulia using `Graphviz` engines.

#### Arguments
- `g::AbstractSimpleWeightedGraph`: a graph representation to export
- (optional) `edge_label::Bool`: if true all edges are labeled with their weights (default = false)
- (optional) `path = []`: Int-Array of nodes. Nodes and their edges are drawn in red color (i.e. shortest path)
- (optional) `colors = zeros(Int, nv(mat))`: Color nodes using Brewer Color Scheme (max 9 colors).
- (optional) `scale = 3.0`: Scale your plot
- (optional) `landscape = false`: if true > set `rankdir` to `LR`
"""
function plot_graphviz(g::AbstractSimpleWeightedGraph;
    edge_label::Bool = false,
    colors = zeros(Int, SimpleWeightedGraphs.nv(g)),
    path = [],
    scale = 3.0,
    landscape = false)

    attrs = GraphvizAttributes(g; node_label = true, edge_label = edge_label)
    set!(attrs.graph_options, "size", string(scale))

    (edge_label) ? set!(attrs.graph_options, "forcelabels", "true") : nothing
    (landscape) ? set!(attrs.graph_options, "rankdir", "LR") : nothing

    gv_dot = string_dot(g, attrs, path, colors)
    plot_graphviz(gv_dot)
end

function plot_graphviz(tup::Tuple{SimpleWeightedDiGraph{Int64,Float64},GraphvizAttributes};
    edge_label::Bool = false,
    colors = zeros(Int, SimpleWeightedGraphs.nv(tup[1])),
    path = [],
    scale = 3.0,
    landscape = false)

    g = tup[1]
    attrs = tup[2]

    (edge_label) ? set!(attrs.graph_options, "forcelabels", "true") : set!(attrs.graph_options, "forcelabels", "false")
    (landscape) ? set!(attrs.graph_options, "rankdir", "LR") : set!(attrs.graph_options, "rankdir", "TB")

    plot_graphviz(g, attrs; path, colors, scale)

end

"""
    plot_graphviz(g, attributes; path = zeros(Int, nv(mat)))

Render graph `g` in **iJulia** using `Graphviz` engines.

#### Arguments
- `g::AbstractSimpleWeightedGraph`: a graph representation to export
- `attributes::AttributeDict`: Render with own set of plotting attributes (see http://www.graphviz.org/ for details)
- (optional) `path = []`: Int-Array of nodes. Nodes and their edges are drawn in red color (i.e. shortest path)
- (optional) `colors = zeros(Int, nv(mat))`: Color nodes using Brewer Color Scheme (max 9 colors).
- (optional) `scale = 3.0`: Scale your plot
- (optional) `landscape = false`: if true > set `rankdir` to `LR`
"""
function plot_graphviz(g::AbstractSimpleWeightedGraph, attributes::GraphvizAttributes;
    path = [],
    colors = zeros(Int, nv(g)),
    scale = 3.0
)
    if !isempty(val(attributes.plot_options, "weights"))
        (val(attributes.plot_options, "weights") == "true") ? set!(attributes.graph_options, "forcelabels", "true") : nothing
    end
    set!(attributes.graph_options, "size", string(scale))
    gv_dot = string_dot(g, attributes, path, colors)
    plot_graphviz(gv_dot)
end

# call ShowGraphviz
function plot_graphviz(str::AbstractString)
    ShowGraphviz.CONFIG.dot_option = `-q` # do not warn in iJulia!
    ShowGraphviz.DOT(str)
end

max_stars_repo_name: tragisch/PlotGraphviz

max_issues_repo_name: tragisch/PlotGraphviz


text:
using ImageReconstruction
using Test

@testset "radon - centered impulse response" begin
    pixels = 129
    views = 200
    I = zeros(pixels, pixels)
    I[pixels÷2, pixels÷2] = 1
    θ = range(0, step = 2π / views, length = views)
    t = -100:100

    P = radon(I, θ, t)

    @test all(P[101, :] .== 1)
end

using Images: shepp_logan, sad
@testset "iradon - shepp logan" begin
    Igt = shepp_logan(128)
    views = 200
    θ = range(0, 2π, length=views)
    t = -150:150
    P = radon(Igt, θ, t)
    I = iradon(P, θ, t)
	@test sad(I, Igt) < 500
end

using Images: shepp_logan
@testset "sinogram - shepp logan" begin
    Igt = shepp_logan(128)
    sinogram(Igt,10)
end


text:
using MyExample
using Test

@testset "my_f tests" begin
    @test my_f(2, 1) == 7;
    @test my_f(2, 3) == 13;
    @test my_f(1, 3) == 11;
end

@testset "Derivative of my_f test" begin
    @test derivative_of_my_f(2, 0) == 2;
    @test derivative_of_my_f(1, 1) == 2;
end

max_stars_repo_name: DomagojGalic/MyExample

max_issues_repo_name: DomagojGalic/MyExample

max_forks_repo_name: DomagojGalic/MyExample


text:
@reexport using Plots
@reexport using ColorSchemes

get_colors(scheme, v::AbstractArray) = get.((scheme, ), 1:length(v), ((1, length(v)),))

@export function plot_energy_barriers(E_barriers, Hs, offset=15; plot_kwargs...)
    o = 0.0
    cs = get_colors(ColorSchemes.rainbow, E_barriers)
    # p1 = plot(xlabel = "State", ylabel="Energy", yticks=([],[]), legend_title="H", legend=false)
    p1 = plot(xlabel = "State", yticks=([],[]), legend_title="", colorbar_title="H")
    ytickvals = []
    yticklabels = []
    horizontal_energy = E_barriers[1][1]
    for (i, (E, h)) in enumerate(zip(E_barriers, Hs))
        o += offset
        n = length(E)
        # E_ = [E[1:30]; E[11:20]]
        E_ = E
        horizontal_energy = E_[1][1]
        # plot!(p1, 1:n, E_ .+ o, color=cs[i], xticks=([1, 11, 21, 31, 40],["1", "2", "3", "4", "1"]), colorbar_entry=true, marker_z = round(h, digits=2))
        plot!(p1, 1:n, E_ .+ o, seriescolor=cs[i], xticks=([1, 11, 21, 31, 40],["1", "2", "3", "4", "1"]), marker_z = round(h, digits=2), markercolor=:rainbow,label="")
        push!(ytickvals, o + horizontal_energy)
        push!(yticklabels, "H=$(round(h, digits=2))")
        plot!(p1, 1:n, fill(o + horizontal_energy, n), color=cs[i], label="")
    end

    # return plot(p1, yticks=(ytickvals,yticklabels), title="Nudged Elastic Band")
    return plot(p1; plot_kwargs...)
end

@export function plot_fields(ledgers, Gd_id1, Gd_id2, L_id; L_H_arrow_scale=0.05, dip_arrow_scale=0.5, cs = get_colors(ColorSchemes.rainbow, ledgers))
    fields_on_Gds = map(field_on_Gds, ledgers)
    fields_on_Ls = map(field_on_L, ledgers)
    p1 = scatter(map(x->x[1], spins(ledgers, Gd_id1)),
                map(x->x[2], spins(ledgers, Gd_id1)), color=cs, title="Gd$(Gd_id1)", xlabel="spin x", ylabel="spin y")

    quiver!(p1, map(x->x[1], spins(ledgers, Gd_id1)),
            map(x->x[2], spins(ledgers, Gd_id1)),
            quiver=(map(x->x.dipole[Gd_id1][1], fields_on_Gds) .* dip_arrow_scale, map(x->x.dipole[Gd_id1][2], fields_on_Gds) .* dip_arrow_scale ),color=:red)

    quiver!(p1, map(x->x[1], spins(ledgers, Gd_id1)),
            map(x->x[2], spins(ledgers, Gd_id1)),
            quiver=(map(x->x.L[Gd_id1][1], fields_on_Gds) .* L_H_arrow_scale, map(x->x.L[Gd_id1][2], fields_on_Gds) .*L_H_arrow_scale ),color=:green)

    quiver!(p1, map(x->x[1], spins(ledgers, Gd_id1)),
            map(x->x[2], spins(ledgers, Gd_id1)),
            quiver=(map(x->x[1], H_vecs(ledgers)) .* L_H_arrow_scale, map(x->x[2], H_vecs(ledgers)) .* L_H_arrow_scale ),color=:purple)

    quiver!(p1, map(x->x[1], spins(ledgers, Gd_id1)),
            map(x->x[2], spins(ledgers, Gd_id1)),
            quiver=(map(x->x.sum[Gd_id1][1], fields_on_Gds) .* L_H_arrow_scale, map(x->x.sum[Gd_id1][2], fields_on_Gds) .* L_H_arrow_scale),color=:cyan)

    p2 = scatter(map(x->x[1], spins(ledgers, Gd_id2)),
                map(x->x[2], spins(ledgers, Gd_id2)), color=cs, title="Gd$(Gd_id2)", xlabel="spin x", ylabel="spin y")

    quiver!(p2, map(x->x[1], spins(ledgers, Gd_id2)),
            map(x->x[2], spins(ledgers, Gd_id2)),
            quiver=(map(x->x.dipole[Gd_id2][1], fields_on_Gds) .*  dip_arrow_scale, map(x->x.dipole[Gd_id2][2], fields_on_Gds) .* dip_arrow_scale ),color=:red)

    quiver!(p2, map(x->x[1], spins(ledgers, Gd_id2)),
            map(x->x[2], spins(ledgers, Gd_id2)),
            quiver=(map(x->x.L[Gd_id2][1], fields_on_Gds) .* L_H_arrow_scale, map(x->x.L[Gd_id2][2], fields_on_Gds) .*L_H_arrow_scale),color=:green)

    quiver!(p2, map(x->x[1], spins(ledgers, Gd_id2)),
            map(x->x[2], spins(ledgers, Gd_id2)),
            quiver=(map(x->x.sum[Gd_id2][1], fields_on_Gds) .* L_H_arrow_scale, map(x->x.sum[Gd_id2][2], fields_on_Gds) .* L_H_arrow_scale),color=:cyan)

    quiver!(p2, map(x->x[1], spins(ledgers, Gd_id2)),
            map(x->x[2], spins(ledgers, Gd_id2)),
            quiver=(map(x->x[1], H_vecs(ledgers)) .* L_H_arrow_scale, map(x->x[2], H_vecs(ledgers)) .* L_H_arrow_scale),color=:purple)

    p4 = scatter(map(x->x[1], L_vecs(ledgers, L_id)),
                 map(x->x[2], L_vecs(ledgers, L_id)), color=cs, title="L$L_id", xlabel="spin x", ylabel="spin y")

    quiver!(p4, map(x->x[1], L_vecs(ledgers, L_id)),
                map(x->x[2], L_vecs(ledgers, L_id)),
            quiver=(map(x->x.L[L_id][1], fields_on_Ls), map(x->x.L[L_id][2], fields_on_Ls)),color=:red)

    quiver!(p4, map(x->x[1], L_vecs(ledgers, L_id)),
                map(x->x[2], L_vecs(ledgers, L_id)),
            quiver=(map(x->x.Gd[L_id][1], fields_on_Ls) .* L_H_arrow_scale, map(x->x.Gd[L_id][2], fields_on_Ls) .* L_H_arrow_scale),color=:green)

    quiver!(p4, map(x->x[1], L_vecs(ledgers, L_id)),
                map(x->x[2], L_vecs(ledgers, L_id)),
            quiver=(map(x->x.easy[L_id][1], fields_on_Ls) .* L_H_arrow_scale, map(x->x.easy[L_id][2], fields_on_Ls).*L_H_arrow_scale),color=:purple)

    p1,p2,p4
end


text:
isdefined(Base, :__precompile__) && __precompile__()

module Grid

using Compat
if VERSION > v"0.4.0-dev+3066"
    using WoodburyMatrices
end

const mp = module_parent(Grid)
if isdefined(mp, :Images) && isdefined(mp.Images, :restrict)
    import ..Images.restrict
end

if !isdefined(:range)
    range = Range
end

include("counter.jl")
include("boundaryconditions.jl")
include("interpflags.jl")
include("interp.jl")
include("restrict_prolong.jl")
include("utilities.jl")
include("coord.jl")

export
# Types
    BoundaryCondition,
    BCnil,
    BCnan,
    BCna,
    BCreflect,
    BCperiodic,
    BCnearest,
    BCfill,
    Counter,
    AbstractInterpGrid,
    InterpGrid,
    InterpGridCoefs,
    InterpIrregular,
    InterpType,
    InterpForward,
    InterpBackward,
    InterpNearest,
    InterpLinear,
    InterpQuadratic,
    InterpCubic,
    CoordInterpGrid,
# Functions
    filledges!,
    interp,
    interp_invert!,
    npoints,
    pad1,
    prolong,
    prolongb,
    prolong_size,
    restrict,
    restrictb,
    restrict_size,
    restrict_extrap,
    set_gradient_coordinate,
    set_hessian_coordinate,
    set_position,
    set_size,
    valgrad,
    valgradhess,
    wrap

end # module


text:
module LightPropagationRunTests

using Test

@time @testset "DAsemiinf" begin include("DAsemiinfTests/runtests.jl") end

@time @testset "DAslab" begin include("DAslabTests/runtests.jl") end

@time @testset "DAparalpip" begin include("DAparalpipTests/runtests.jl") end

@time @testset "DA_Nlayer" begin include("DA_NlayerTests/runtests.jl") end

@time @testset "DAinf" begin include("DAinfTests/runtests.jl") end

@time @testset "DAdcs" begin include("DA_DCSTests/runtests.jl") end

end

max_issues_repo_name: heltonmc/LightPropagation


text:
# This file is a part of SimilaritySearch.jl

using Test, SimilaritySearch, LinearAlgebra

@testset "allknn" begin
    dist = CosineDistance()
    X = rand(Float32, 4, 1000)
    db = VectorDatabase(Vector{Float32}[])
    ϵ = 0.1
    D = neardup(SearchGraph(; db, dist), MatrixDatabase(X), ϵ)
    @test all(x -> x <= ϵ, D.dist)
    @test sum(D.dist) > 0
    @test D.map == sort(unique(D.nn))
end

max_stars_repo_name: sadit/NNS

max_issues_repo_name: sadit/NNS

max_forks_repo_name: sadit/NNS