text:
# # How-to Guide
using AeroMDAO # hide
using Plots # hide
gr(dpi = 300) # hide
using LaTeXStrings # hide

# ## Airfoil Geometry
# How to work with airfoil geometry.

# ### Import Coordinates File
# You can specify the path consisting of the foil's coordinates to the `read_foil` function. The format for the coordinates file requires a header for the name in the first line of the file. A custom name can be provided by setting the optional `name` variable.

## Airfoil coordinates file path
foilpath = string(@__DIR__, "/misc/s1223.dat")

## Read coordinates file
my_foil = read_foil(foilpath;
                    name   = "S1223" # To overwrite name in header
                   )

#
plot(my_foil.x, my_foil.y,
     xlabel = L"x", ylabel = L"y", aspect_ratio = 1, label = "$(my_foil.name)")

# ### Interpolate and Process Coordinates

# A basic cosine interpolation functionality is provided for foils.
## Cosine spacing with approx. 51 points on upper and lower surfaces each
cos_foil = cosine_interpolation(my_foil, 51)

# The upper and lower surfaces can be obtained by the following variety of functions.

## Upper surface
upper = upper_surface(my_foil)

## Lower surface
lower = lower_surface(my_foil)

## Upper and lower surfaces
upper, lower = split_surface(my_foil)

# The camber-thickness distribution can be obtained as follows. It internally performs a cosine interpolation to standardize the $x$--coordinates for both surfaces; hence the number of points for the interpolation can be specified.
xcamthick = camber_thickness(my_foil, 60)

# You can also do the inverse transformation.
coords = camber_thickness_to_coordinates(xcamthick[:,1], xcamthick[:,2], xcamthick[:,3])

# ### Control Surfaces

# You can (somewhat) mimic the behaviour of a control surface by specifying a deflection angle $\delta$ (in degrees, clockwise-positive convention) with the specification of the hinge location's $x$-coordinate normalized in $[0,1]$ to the chord length.
con_foil = control_surface(cos_foil; angle = 10., hinge = 0.75)

plot!(con_foil.x, con_foil.y, label = "$(my_foil.name) Deflected")

# ## Doublet-Source Aerodynamic Analyses
# The `solve_case` method runs the analysis given a `Foil` containing the airfoil coordinates, a `Uniform2D` defining the boundary conditions, and an optional named specification for the number of panels. It returns a system which can be used to obtain the aerodynamic quantities of interest and post-processing.

## Define freestream boundary conditions
uniform = Uniform2D(1.0, 4.0)

## Solve system
system  = solve_case(my_foil, uniform;
                     num_panels = 80)

# The following functions compute the quantities of interest, such as the inviscid edge velocities, lift coefficient, and the sectional lift, moment, and pressure coefficients.
cls, cms, cps = surface_coefficients(system);
u_es   = surface_velocities(system)
cl     = lift_coefficient(system)

# AeroMDAO provides more helper functions for the panel geometry.
panels   = system.surface_panels
pts      = collocation_point.(panels) # Collocation points
tangents = tangent_vector.(panels)     # Tangents
normals  = normal_vector.(panels)      # Normals
locs     = panel_location.(panels);   # Upper or lower surface

# ## Wing Geometry
# How to work with wing geometry.
#
# To define one side of a wing, AeroMDAO provides a `HalfWing` constructor.
airfoil    = naca4((2,4,1,2))
wing_right = HalfWing(foils     = [ airfoil for i in 1:3 ],
                      chords    = [0.4, 0.2, 0.1],
                      twists    = [0., 2., 5.],
                      spans     = [1.0, 0.1],
                      dihedrals = [0., 60.],
                      sweeps    = [0., 30.],
                      w_sweep   = 0.25)

# The `Wing` constructor takes left and right `HalfWing`s to define a full wing. For example, the following generates a symmetric wing.
wing = Wing(wing_right, wing_right)

# The following "getter" functions provide quantities of interest such as chord lengths, spans, twist, dihedral, and sweep angles.

# ```@repl howto
# foils(wing)
# chords(wing)
# spans(wing)
# twists(wing)
# dihedrals(wing)
# sweeps(wing)
# ```

# There is also a convenient function for pretty-printing information, in which the first argument takes the `Wing` type and the second takes a name (as a `String` or `Symbol`).
print_info(wing, "Wing")

#md # !!! tip
#md #     You can use [Setfield.jl](https://github.com/jw3126/Setfield.jl) to conveniently copy and modify properties of an existing object.

## Import Setfield
using Setfield

## Set only chords with other properties remaining identical.
wing_left = @set wing_right.chords = [0.4, 0.1, 0.05]

# To create an asymmetric wing, feed the left and right halves to `Wing` in the particular order.
wing = Wing(wing_left, wing_right);

print_info(wing, "My Wing")

# ## Vortex Lattice Aerodynamic Analyses
#
# How to run a generic aerodynamic analysis on a conventional aircraft configuration.
#
# ### Geometry
# First we define the lifting surfaces. These can be a combination of `Wing` or `HalfWing` types constructed using the various methods available.
# !!! warning "Alert"
#     Support for fuselages and control surfaces will be added soon.

## Wing
wing  = WingSection(span       = 8.0,
                    dihedral   = 5.0,
                    sweep      = 15.0,
                    taper      = 0.4,
                    root_chord = 2.0,
                    root_twist = 0.0,
                    tip_twist  = -2.0,
                    root_foil  = naca4(2,4,1,2),
                    tip_foil   = naca4(2,4,1,2),
                    position   = [0., 0., 0.])

## Horizontal tail
htail = WingSection(span       = 2.0,
                    dihedral   = 0.0,
                    sweep      = 15.0,
                    taper      = 0.6,
                    root_chord = 0.8,
                    root_twist = 0.0,
                    tip_twist  = 0.0,
                    root_foil  = naca4(0,0,1,2),
                    tip_foil   = naca4(0,0,0,9),
                    position   = [5., 0., -0.1],
                    angle      = 0.,
                    axis       = [0., 1., 0.]);

## Vertical tail
vtail = HalfWingSection(span       = 0.8,
                        dihedral   = 0.0,
                        sweep      = 8.0,
                        taper      = 0.6,
                        root_chord = 0.8,
                        root_twist = 0.0,
                        tip_twist  = 0.,
                        root_foil  = naca4(0,0,0,9),
                        tip_foil   = naca4(0,0,0,9),
                        position   = [5., 0., 0.],
                        angle      = 90.,
                        axis       = [1., 0., 0.])

# ### Meshing

# The `WingMesh` type takes a `HalfWing` or `Wing` type with a vector of integers consisting of the spanwise panel distribution corresponding to the number of sections, and an integer for the chordwise distribution.

## Wing meshes
wing_mesh  = WingMesh(wing, [20], 10)

# Optionally the type of spanwise spacing can be specified by the keyword `span_spacing` and providing types `Sine(), Cosine(), Uniform()` or a vector of the combination with length corresponding to the number of sections.

## Horizontal tail
htail_mesh = WingMesh(htail, [10], 5;
                      span_spacing = Cosine()) # Uniform(), Sine()

## Vertical tail
vtail_mesh = WingMesh(vtail, [10], 5); # Vertical tail

# ### Inviscid Analysis

# The inviscid 3D analysis uses a vortex lattice method. The `WingMesh` type allows you to generate horseshoe elements easily.
wing_horsies = make_horseshoes(wing_mesh)

# For multiple lifting surfaces, it is most convenient to define a single vector consisting of all the components' horseshoes using [ComponentArrays.jl](https://github.com/jonniedie/ComponentArrays.jl).
aircraft = ComponentVector(
                           wing  = wing_horsies,
                           htail = make_horseshoes(htail_mesh),
                           vtail = make_horseshoes(vtail_mesh)
                          );

# To define boundary conditions, use the following `Freestream` type, which takes named arguments for angles of attack and sideslip (in degrees), and a quasi-steady rotation vector.

## Define freestream conditions
fs = Freestream(alpha = 1.0,
                beta  = 0.0,
                omega = [0.,0.,0.])

# To define reference values, use the following `References` type.

## Define reference values
refs = References(
           speed     = 10.0,
           density   = 1.225,
           viscosity = 1.5e-5,
           area      = projected_area(wing),
           span      = span(wing),
           chord     = mean_aerodynamic_chord(wing),
           location  = mean_aerodynamic_center(wing)
          )

# The vortex lattice analysis can be executed with the horseshoes, freestream condition, and reference values defined.

## Solve system
system = solve_case(
             aircraft, fs, refs;
             print            = true, # Prints the results for only the aircraft
             print_components = true, # Prints the results for all components
            )

# If needed, you can access the relevant component influence matrix values and boundary conditions with the following attributes, e.g.
system.influence_matrix[:wing] # Or :htail, :vtail
system.influence_matrix[:wing, :htail];

system.boundary_vector[:wing]

# ### Dynamics
#
# The analysis computes the aerodynamic forces by surface pressure integration for nearfield forces. You can specify the axis system for the nearfield forces, with choices of geometry, body, wind, or stability axes. The wind axes are used by default.

## Compute dynamics
ax = Wind() # Body(), Stability(), Geometry()

## Compute non-dimensional coefficients
CFs, CMs = surface_coefficients(system; axes = ax)

# You can access the corresponding values of the components' by the name provided in the `ComponentVector`.
CFs.wing

# Special functions are provided for directly retrieving the dimensionalized forces and moments.
Fs, Ms   = surface_dynamics(system; axes = ax)

# A Trefftz plane integration is performed to compute farfield forces.
#
# !!! note
#     The farfield forces are usually more accurate compared to nearfield forces, as the components do not interact as in the evaluation of the Biot-Savart integrals for the latter.
#
# To obtain the nearfield and farfield coefficients of the components (in wind axes by definition):
nfs = nearfield_coefficients(system)
ffs = farfield_coefficients(system)

# You can similarly access the components by name.
@show (nfs.wing, ffs.wing)

# To obtain the total nearfield and farfield force coefficients:
nf = nearfield(system)
ff = farfield(system)

# You can also print the relevant information as a pretty table, if necessary.
print_coefficients(nfs.wing, ffs.wing, :wing)
print_coefficients(nf, ff, :aircraft)

# ## Aerodynamic Stability Analyses
# The derivatives of the aerodynamic coefficients with respect to the freestream values is obtained by automatic differentiation enabled by [ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl). To compute the values, simply replace `solve_case` with `solve_case_derivatives`. You can also optionally provide the axes for the reference frame of the coefficients.
dv_data = solve_case_derivatives(
             aircraft, fs, refs;
             axes             = Wind(),
             name             = :aircraft,
             print            = true,    # Prints the results for only the aircraft
             print_components = true,    # Prints the results for all components
            );

# ## Euler-Bernoulli Beam Structural Analysis
# The tubular beam's relevant properties, viz. the Young's (elastic) modulus $E$, shear modulus $G$, and torsional moment of inertia $J$ must be specified to define the stiffness matrix for its discretization with $n$ elements.

## Material properties
E = 1. # Elastic modulus
G = 1. # Shear modulus
J = 2. # Torsional moment of inertia
n = 2  # Number of sections

## Stiffness matrix
K = bending_stiffness_matrix(
                             fill(E, 2),
                             fill(G, 2),
                             fill(J, 2),
                             :z         # Direction of deflection
                            )

# Fixed, hinged beam subjected to force and moment at the center.

## Stiffness matrix
A = K[[3,4,6],[3,4,6]]  # v2, φ2, φ3

## Load vector
b = [-1000, 1000, 0]    # F2, M2, M3

## Solution
x = A \ b

## Forces
F1 = K * [ 0.; 0.; x[1:2]; 0.; x[3] ]

# Propped cantilever beam subjected to force at free end.

## Stiffness matrix
A = K[[1,2,4],[1,2,4]] # v1, φ1, φ2

## Load vector
b = [10, 0, 0]

## Solution
x = A \ b

## Forces
F2 = K * [ x[1:2]; 0.; x[3]; 0.; 0. ]


text:
# Return a substring of input that is of equal or greater indentation than the
# line starting at i.
function get_indented_block(input::String, i::Integer=1)
    output = IOBuffer()
    n = length(input)
    block_indent = 0
    while i < n && (input[i] == ' ' || input[i] == '\t')
        block_indent += 1
        write(output, input[i])
        i = nextind(input, i)
    end

    j = search(input, '\n', i)
    if j == 0
        j = n
    end

    write(output, input[i:j])
    i = nextind(input, j)

    while i < n
        j = i
        line_indent = 0
        while j < n && (input[j] == ' ' || input[j] == '\t')
            line_indent += 1
            j = nextind(input, j)
        end

        k = search(input, '\n', j)
        if k == 0
            k = n
        end

        if line_indent < block_indent && j < k
            break
        end

        write(output, input[i:k])
        i = nextind(input, k)
    end
    takebuf_string(output)
end

# pattern used to find fileds in a function declaration comment
const func_field_pat = r"^\h*(Args|Returns|Modifies|Throws)\s*:\h*\r?\n"im

# Parse a function declaration comment
#
# Args:
#   input: Comment preceeding a funciton declaration, stripped of the leading
#          '#' characters.
#
# Returns:
#   A four-tuple of the form (description, args, returns, modifies, throws)
#
function parse_function_comment(input::String)
    mat = match(func_field_pat, input)
    if mat == nothing
        return (strip(input), nothing, nothing, nothing, nothing, nothing)
    end

    description = input[1:mat.offset-1]
    args = nothing
    returns = nothing
    modifies = nothing
    throws = nothing

    while mat != nothing
        typ = mat.captures[1]
        if typ == "Args"
            args = get_indented_block(input, mat.offset + length(mat.match))
            offset = mat.offset + length(args)
            args = parse_function_comment_args(args)
        elseif typ == "Returns"
            returns = strip(get_indented_block(input, mat.offset + length(mat.match)))
            offset = mat.offset + length(returns)
        elseif typ == "Modifies"
            modifies = strip(get_indented_block(input, mat.offset + length(mat.match)))
            offset = mat.offset + length(modifies)
        elseif typ == "Throws"
            throws = strip(get_indented_block(input, mat.offset + length(mat.match)))
            offset = mat.offset + length(throws)
        end

        mat = match(func_field_pat, input, offset)
    end

    (strip(description), args, returns, modifies, throws)
end

# pattern to match argument names/descriptions.
const arg_desc_pat = r"^(\h*)([\w_][\w\d_\!]*)\h*:\h*(.*)\r?"m

# Parse the "Args" section of a function declaration comment.
#
# Args:
#   input: Everything under the "Args" block.
#
# Returns:
#   A dictionary mapping arugment names to their descriptions.
#
function parse_function_comment_args(input::String)
    args = Dict()
    mat = match(arg_desc_pat, input)
    while mat != nothing
        indent = length(mat.captures[1])
        name = mat.captures[2]
        desc = IOBuffer()
        write(desc, mat.captures[3])

        offset = mat.offset + length(mat.match)
        if offset < length(input) && input[offset] == '\n'
            offset = nextind(input, offset)
            next_line_indent = 0
            i = offset
            while i < length(input) && (input[i] == ' ' || input[i] == '\t')
                next_line_indent += 1
                i = nextind(input, i)
            end
            println(STDERR, (next_line_indent, indent))
            if next_line_indent > indent
                desc_rest = get_indented_block(input, offset)
                write(desc, desc_rest)
                offset += length(desc_rest)
            end
        end

        args[name] = takebuf_string(desc)
        mat = match(arg_desc_pat, input, offset)
    end

    args
end

ext: jl


text:
using Test
using Unfold
using StatsModels
using DataFrames
using Statistics
using Random
using CategoricalArrays
using StatsBase

using MixedModels
using SparseArrays

include("test_utilities.jl")

max_stars_repo_name: unfoldtoolbox/jlunfold

max_issues_repo_name: unfoldtoolbox/jlunfold

max_forks_repo_name: unfoldtoolbox/jlunfold


text:
# -*- coding: utf-8 -*-
#
# To do force field transform using FFT
#
# Aaron LI
# 2015/07/16
#

function forcefieldtransform_fft(img)
    rows, cols = size(img)
    pic = zeros(3*rows, 3*cols)
    pic[1:rows, 1:cols] = img
    # unit force field
    unit_ff = complex(zeros(3*rows, 3*cols))
    for r = 1:(2*rows-1)
        for c = 1:(2*cols)
            d = (rows+cols*im) - (r+c*im)
            if (r, c) == (rows, cols)
                unit_ff[r, c] = 0 + 0im
            else
                unit_ff[r, c] = d / abs(d)^3
            end
        end
    end
    # FIXME matrix sizes
    ff = sqrt(rows*cols) * ifft(fft(pic) .* fft(unit_ff))
    #ff_crop = ff[rows:2*rows, cols:2*cols]
end

max_stars_repo_name: liweitianux/atoolbox

max_issues_repo_name: liweitianux/atoolbox

max_forks_repo_name: liweitianux/atoolbox


text:
"""
    macro bu_str(unit)

String macro to easily recall Batsrus units located in the `UnitfulBatsrus` package.
Although all unit symbols in that package are suffixed with `_bu`, the suffix should not be
used when using this macro.
Note that what goes inside must be parsable as a valid Julia expression.
Examples:
```
julia> 1.0bu"u"
1.0 u
julia> 1.0bu"amucc"
1.0 amu/cc
```
"""
macro bu_str(unit)
   ex = Meta.parse(unit)
   esc(replace_value(ex))
end

const allowed_funcs = [:*, :/, :^, :sqrt, :√, :+, :-, ://]
function replace_value(ex::Expr)
   if ex.head == :call
      ex.args[1] in allowed_funcs ||
         error("""$(ex.args[1]) is not a valid function call when parsing a unit.
          Only the following functions are allowed: $allowed_funcs""")
      for i=2:length(ex.args)
         if typeof(ex.args[i])==Symbol || typeof(ex.args[i])==Expr
            ex.args[i] = replace_value(ex.args[i])
         end
      end
      return Core.eval(@__MODULE__, ex)
   elseif ex.head == :tuple
      for i=1:length(ex.args)
         if typeof(ex.args[i]) == Symbol
            ex.args[i] = replace_value(ex.args[i])
         else
            error("only use symbols inside the tuple.")
         end
      end
      return Core.eval(@__MODULE__, ex)
   else
      error("Expr head $(ex.head) must equal :call or :tuple")
   end
end

dottify(s, t, u...) = dottify(Expr(:(.), s, QuoteNode(t)), u...)
dottify(s) = s

function replace_value(sym::Symbol)
   s = Symbol(sym, :_bu)
   if !(isdefined(UnitfulBatsrus, s) && typecheck_bool(getfield(UnitfulBatsrus, s)))
      error("Symbol $s could not be found in UnitfulBatsrus.")
   end

   return getfield(UnitfulBatsrus, s)
end

replace_value(literal::Number) = literal

typecheck_bool(x::Unitful.Number) = true
typecheck_bool(x::Unitful.Unitlike) = true
typecheck_bool(x::Unitful.Quantity) = true
typecheck_bool(x::Unitful.MixedUnits) = true
typecheck_bool(x) = false


text:
#=
[Intermediate] rearranged magic squares

Description

An NxN magic square is an NxN grid of the numbers 1 through N^2 such that each row,
column, and major diagonal adds up to M = N(N^2+1)/2. See this week's Easy problem
for an example.
You will be given an NxN grid that is not a magic square, but whose rows can be
rearranged to form a magic square. In this case, rearranging the rows means to
put the rows (horizontal lines of numbers) in a different order, but within each
row the numbers stay the same. So for instance, the top row can be swapped with
the second row, but the numbers within each row cannot be moved to a different
position horizontally, and the numbers that are on the same row as each other
to begin with must remain on the same row as each other.
Write a function to find a magic square formed by rearranging the rows of the
given grid.
There is more than one correct solution. Format your grid however you like. You
can parse the program's input to get the grid, but you don't have to.

Example

15 14  1  4        12  6  9  7
12  6  9  7   =>    2 11  8 13
 2 11  8 13        15 14  1  4
 5  3 16 10         5  3 16 10

Inputs

Challenge inputs

Any technique is going to eventually run too slowly when the grid size gets too
large, but you should be able to handle 8x8 in a reasonable amount of time (less
than a few minutes). If you want more of a challenge, see how many of the example
inputs you can solve.
I've had pretty good success with just randomly rearranging the rows and checking
the result. Of course, you can use a "smarter" technique if you want, as long as
it works!

Optional bonus

(Warning: hard for 12x12 or larger!) Given a grid whose rows can be rearranged to
form a magic square, give the number of different ways this can be done. That is,
how many of the N! orderings of the rows will result in a magic square?
If you take on this challenge, include the result you get for as many of the
challenge input grids as you can, along with your code.

( source: https://redd.it/4dmm44 )
=#

#   exit at EOF ( in Linux: Ctrl + D and Windows: Ctrl + Z )
read_matrix( data::DataType=Any, input::IO=STDIN, spacing::Char=' ' ) = readdlm(
input, spacing, data, header=false, skipblanks=true, ignore_invalid_chars=true,
quotes=false, comments=false )

magic_N( N::Int ) = convert( Int, N*( N^2+1 )/2 )

function swap_grid_line!( grid::Matrix{ Int }, i::Integer, j::Integer )
    #   @inbounds allows the operation without bounding checking -- great time
    #   saving
    @inbounds for k in 1:size( grid, 2 )
        grid[ i, k ], grid[ j, k ] = grid[ j, k ], grid[ i, k ]
    end
end

#=
    This function basically does the interpolation of a determinated line and see
    all the combinations, if the combination is valid or not
=#
function __magic_square!( grid::Matrix{ Int }, index::Int, magic::Int )
    for i in 1:size( grid, 1 ), j in 1:size( grid, 1 )
        if i != index && j != index
            swap_grid_line!( grid, i, j )
            trace( grid ) == magic ? break : continue
        end
    end
end

#=
    This function call __magic_square! each time with each line of the grid to
    check up all the combinations -- when the solutions is found returns the
    the grid
=#
function magic_square!( grid::Matrix{ Int } )
    local magic::Int = magic_N( size( grid, 1 ) )

    for i in 1:size( grid, 1 )
        trace( grid ) == magic ? break : __magic_square!( grid, i, magic )
    end

    return grid
end

function main( )
    local matrix::Matrix{ Int } = read_matrix( Int )
    println( @time magic_square!( matrix ) )
    println( "The magic number is ", magic_N( size( matrix, 1 ) ), " and your diagonal adds up to ", trace( matrix ) )
end

main( )

max_stars_repo_name: farmfirst/dailyprogrammer

max_issues_repo_name: farmfirst/dailyprogrammer

max_forks_repo_name: farmfirst/dailyprogrammer


text:
using PkgSRA, Test, Random
Random.seed!(13)

# include("../utils/regression_utils.jl")
# include("../utils/sparse_regression_functions.jl")

# Generate test data: Lorenz
include("../examples/example_lorenz.jl")
dat = Array(solve_lorenz_system())
numerical_grad = numerical_derivative(dat, ts)

sindy_library = Dict("cross_terms"=>2,"constant"=>nothing);

make_model(x) = sindyc(slstNumber(2, x), dat, numerical_grad)

#####
##### Test Sparse Regression function
#####
thresh = 0.1
m1 = slstHard(2, thresh)
A = sparse_regression(m1, dat, numerical_grad)
@testset "Hard threshold" begin
    @test all(abs.(A).<thresh)
end

#####
##### Test sparse models: all equations same terms
#####
num_terms5 = [5]
terms5 = get_nonzero_terms(make_model(num_terms5))
num_terms3 = [3]
terms3 = get_nonzero_terms(make_model(num_terms3))
@testset "Number of terms" begin
    @test length(terms5) == size(dat,1)*num_terms5
    @test length(terms3) == size(dat,1)*num_terms5
end

#####
##### Test sparse models: all equations different terms
#####
@testset "Number of terms2" begin
    num_terms = [1, 2, 3]
    terms = get_nonzero_terms(make_model(num_terms))
    @test length(terms) == sum(num_terms)

    num_terms = [2, 3, 4]
    terms = get_nonzero_terms(make_model(num_terms))
    @test length(terms) == sum(num_terms)
end

#####
##### Test ensemble models
#####
make_ensemble(x) = sindyc_ensemble(dat, numerical_grad,
                    sindy_library,
                    x,
                    selection_criterion=my_aicc,
                    sparsification_mode="num_terms")

# Each ensemble has all same
val_list = [1, 2]
(best_model_raw,best_criterion,all_criteria,all_models) =
    make_ensemble(val_list)

@testset "Ensemble models" begin
    i = 1
    @test length(get_nonzero_terms(all_models[i])) == size(dat,1)*val_list[i]
    i = 2
    @test length(get_nonzero_terms(all_models[i])) == size(dat,1)*val_list[i]
end

# Permutations
# val_list = calc_permutations(5,3)
# val_list = combinations(1:5, 3)
val_list = Iterators.product(1:3,1:3,1:3) # DEBUG
(best_model,best_criterion,all_criteria,all_models) =
    make_ensemble(val_list)

@testset "Permutations" begin
    for (i, v) in enumerate(val_list)
        @test length(get_nonzero_terms(all_models[i])) == sum(v)
    end
end

# The best model should be almost the real one
@testset "Accuracy" begin
    for (real_val, model_val) in zip(vec(core_dyn_true.A), vec(best_model.A))
        @test isapprox(Float64(real_val), model_val, atol=1e-1)
    end
end


text:
module SpikeUnits

using DataFrames, Query
import CSV
import MAT

export SpikeUnit, SpikeMarker, import_spike_train_data, marker_validity

"""
Storing information and data for each channel.

You wouldn't normally create a SpikeUnit struct yourself, this is done
for you when retreiving a mat file from your data directory.

```julia
SpikeUnit
    session     ::String
    mouse_id    ::String
    channel     ::String
    spike_train ::Array{Float64,1}
```
"""
struct SpikeUnit
    session     ::String
    mouse_id    ::String
    channel     ::String
    spike_train ::Array{Float64,1}
end

"""
Storing and parsing marker information of each session.

You wouldn't normally create a SpikeMarker object yourself, this is done
for you when retreiving a csv file from your data directory.

Specially, the _raw_table shall be a DataFrame, with coloumns("time", "marker");
if you have a different layout, you should write
your own ```_chunker(_raw_table, _raw_train)``` function,
which returns Dict{String, Array{Float64,1}}

```julia
SpikeMarker
    session     ::String
    mouse_id    ::String
    _raw_table  ::DataFrames.DataFrame
    _raw_train  ::Array{Float64, 1}
    chunked_marker::Dict{String, Array{Float64,1}}
```
"""
struct SpikeMarker
    session     ::String
    mouse_id    ::String
    _raw_table  ::DataFrames.DataFrame
    _raw_train  ::Array{Float64, 1}
    chunked_marker::Dict{String, Array{Float64,1}}
end

"""
Check .csv marker and .mat marker are valid or not.

#### Method:
the time shift between the first and the last marker of .csv and .mat
file shall be less than thresh.

#### Args
- table:  the .csv marker in DataFrame
- train:  the .mat marker in 1d Array
- boundary: the threshold, as seconds [optional, default: 1.0]

#### Returns:
- shift:  the shift value. i.e. how many markers in .csv is not recorded in .mat markers.

#### Raises:
- ErrorException("marker value not match!")
- ErrorException("electrode markers exceed stimulus markers!") when shift < 0
"""
function marker_validity(table, train, boundary=1.0)
    shift = length(table[1]) - length(train)

    shift<0 && throw(ErrorException("electrode markers exceed stimulus markers!"))

    _start_d = table[:time][1] - train[1]
    _end_d = table[:time][end-shift] - train[end]

    if abs(_start_d - _end_d) <= boundary
        return shift
    else
        throw(ErrorException("marker value not match!: $(_start_d - _end_d)"))
    end
end

"""
If you have a different .csv data layout, you should write
your own ```_chunker``` function.

Tips:
- you may need the ```marker_validity``` to get the shift value
  between csv and mat markers.

#### Args:
- _raw_table: the raw csv data
- _raw_train: the raw marker data

#### Returns:
- _chunked: chunked data in Dict{String, Array{Float64,1}}
"""
function _default_marker_chunker(_raw_table, _raw_train)
    _valids = @from item in _raw_table begin
        @where !(item.marker in ["START", "QUIT", "gray"])
        @select item
        @collect DataFrame
        end

    _shift = marker_validity(_valids, _raw_train)
    _neo = DataFrame(time=_raw_train, marker=_valids[:marker][1:end-_shift])
    _chunked = Dict{String,Array{Float64,1}}()

    for stim in unique(_valids[:marker])
        _temp = @from item in _valids begin
            @where item.marker == stim
            @select item.time
            @collect
        end
        _chunked[stim] = _temp
    end

    _chunked
end

"""
Import .mat and .csv data.

#### Args:
- session:    session name in string.
- mouse_id:   mouse name in string.

#### Keyword Args:
- _mat: .mat data file name. [optional, default: ""]
- _csv: .csv data file name. [optional, default: ""]
- _dir: data directory path. [optional, default: data/SpikeTrain]
- _marker_channel: the marker channel name in .mat file.[optional, default: "DIG01"]
- _chunker: custom csv chunker function [optional, default: _default_marker_chunker]

#### Returns:
- _spike_units:   Dict{String, SpikeUnit}
- _spike_marker:  SpikeMarker
"""
function import_spike_train_data(session::String, mouse_id::String;
                                 _mat::String="", _csv::String="",
                                 _dir::String="data/SpikeTrain", _marker_channel::String="DIG01",
                                 _chunker=_default_marker_chunker)
    _mat==""&&(_mat="$session.mat")
    _csv==""&&(_csv="$session.csv")

    _mat_file = joinpath(_dir, _mat)
    _csv_file = joinpath(_dir, _csv)

    _marker_train = Array{Float64, 1}()
    _marker_table = CSV.read(_csv_file, nullable=false)
    _spike_units = Dict{String, SpikeUnit}()

    _mat_data = MAT.matread(_mat_file)
    for (name, raw_train) in _mat_data

        _raw_train = raw_train["times"][:,1]
        if name != _marker_channel
            _spike_units[name] = SpikeUnit(session, mouse_id, name, _raw_train)
        else
            _marker_train = _raw_train
        end
    end

    _spike_marker = SpikeMarker(session, mouse_id, _marker_table,
                                _marker_train, _chunker(_marker_table, _marker_train))

    (_spike_units, _spike_marker)
end

end # SpikeUnit


text:
f(x) = x + 1
@show f(1)


text:
using Test, AxisKeys, Statistics
using NamedDims: unname

M = wrapdims(rand(Int8, 3,4), r='a':'c', c=2:5)
MN = NamedDimsArray(M.data.data, r='a':'c', c=2:5)
V = wrapdims(rand(1:99, 10), v=10:10:100)
VN = NamedDimsArray(V.data.data, v=10:10:100)
A3 = wrapdims(rand(Int8, 3,4,2), r='a':'c', c=2:5, p=[10.0, 20.0])

@testset "dims" begin

    # reductions
    M1 = sum(M, dims=1)
    @test axiskeys(M1, 1) === Base.OneTo(1)

    M2 = prod(M, dims=:c)
    @test axiskeys(M2) == ('a':'c', Base.OneTo(1))

    @test_throws ArgumentError sum(M, dims=:nope)

    M4 = mean(M, dims=:c)
    @test axiskeys(M4) === ('a':'c', Base.OneTo(1))

    # dropdims
    @test axiskeys(dropdims(M1, dims=1)) == (2:5,)
    @test axiskeys(dropdims(M2, dims=:c)) == ('a':'c',)

    M3 = dropdims(M2, dims=:c)
    @test dimnames(M3) == (:r,)

    # permutedims
    @test size(permutedims(M)) == (4,3)
    @test dimnames(transpose(M2)) == (:c, :r)
    @test axiskeys(permutedims(M, (2,1))) == (2:5, 'a':'c')
    @test dimnames(M3') == (:_, :r)
    @test axiskeys(transpose(transpose(V))) == axiskeys(V) # fixed by NamedDims.jl/pull/105
    @test axiskeys(permutedims(transpose(V))) == (axiskeys(V,1), Base.OneTo(1))

    V2 = wrapdims(rand(3), along=[:a, :b, :c])
    axiskeys(permutedims(V2), 2)[1] = :nope
    @test axiskeys(V2,1) == [:a, :b, :c]
    axiskeys(V2', 2)[1] = :zed
    @test axiskeys(V2,1) == [:zed, :b, :c]

    # eachslice
    if VERSION >= v"1.1"
        @test axiskeys(first(eachslice(M, dims=:r))) === (2:5,)
    end

    @test axiskeys(selectdim(M, :r, [true, false, true])) == (['a', 'c'], 2:5)

    # mapslices
    @test axiskeys(mapslices(identity, M, dims=1)) === (Base.OneTo(3), 2:5)
    @test axiskeys(mapslices(sum, M, dims=1)) === (Base.OneTo(1), 2:5)
    @test axiskeys(mapslices(v -> rand(10), M, dims=2)) === ('a':'c', Base.OneTo(10))

    # reshape
    @test reshape(M, 4,3) isa Array
    @test reshape(M, 2,:) isa Array
    if AxisKeys.nameouter() == false  # reshape(keyless(M), (2,:)) # is an error
        @test reshape(M, (4,3)) isa Array
        @test reshape(M, (2,:)) isa Array
    end

end
@testset "mapreduce & keywords, issue 43" begin
    # https://github.com/mcabbott/AxisKeys.jl/issues/43

    z = KeyedArray(zeros(3); foo=[:a, :b, :c])
    @test mapreduce(identity, +, z; init=10) == 10

    m = KeyedArray([1, 2, 3]; foo=[:a, :b, :c])
    @test std(m; corrected=false) == std([1, 2, 3]; corrected=false)

end
@testset "sort" begin

    @test sort(V)(20) == V(20)

    @test axiskeys(sort(M, dims=:c), :c) isa Base.OneTo
    @test axiskeys(sort(M, dims=:c), :r) == 'a':'c'

    # sortslices
    @test sortslices(M, dims=:c) == sortslices(AxisKeys.keyless(M), dims=:c)
    @test axiskeys(sortslices(M, dims=:r), :r) isa Vector # not a steprange, it got sorted

    T = wrapdims(rand(Int8, 5,7,3), a=1:5, b=0:6.0, c='a':'c')
    T′ = sortslices(T, dims=:b)
    @test issorted(T′[a=1, c=1])
    @test T′ == sortslices(AxisKeys.keyless(T), dims=:b, by=vec)

    A = wrapdims([1,0], vec=[:a, :b]) # https://github.com/JuliaArrays/AxisArrays.jl/issues/172
    sort!(A)
    @test A(:a) == 1
    @test axiskeys(A, 1) == [:b, :a]
    @test_throws Exception sort!(KeyedArray([1,0], 'a':'b'))

    # sortkeys
    B = wrapdims(rand(Int8,3,7), 🚣=rand(Int8,3), 🏛=rand(Int8,7))
    p = sortperm(B.🏛; rev=true)
    B′ = sortkeys(B, dims=:🏛, rev=true)
    @test B′ == B[:,p]
    @test issorted(reverse(B′.🏛))

end
@testset "map & collect" begin

    mapM =  map(exp, M)
    @test axiskeys(mapM) == ('a':'c', 2:5)
    @test dimnames(mapM) == (:r, :c)

    @test axiskeys(map(+, M, M, M)) == ('a':'c', 2:5)
    @test axiskeys(map(+, M, parent(M))) == ('a':'c', 2:5)

    @test axiskeys(map(sqrt, V)) == (10:10:100,)

    V2 = wrapdims(rand(1:99, 10), v=2:11) # different keys
    V3 = wrapdims(rand(1:99, 10), w=10:10:100) # different name
    @test_throws Exception map(+, V, V2)
    @test_throws Exception map(+, V, V3) # should throw an error ???

    genM =  [exp(x) for x in M]
    @test axiskeys(genM) == ('a':'c', 2:5)
    @test dimnames(genM) == (:r, :c)

    @test axiskeys([exp(x) for x in V]) == (10:10:100,)

    gen3 = [x+y for x in M, y in V];
    @test axiskeys(gen3) == ('a':'c', 2:5, 10:10:100)
    @test dimnames(gen3) == (:r, :c, :v)

    gen1 = [x^i for (i,x) in enumerate(V)]
    @test axiskeys(gen1) == (10:10:100,)
    @test dimnames(gen1) == (:v,)

    @test axiskeys(filter(isodd, V2),1) isa Vector{Int}
    @test dimnames(filter(isodd, V2)) == (:v,)

    @test filter(isodd, M) isa Array

end
@testset "cat" begin

    # concatenation
    @test axiskeys(hcat(M,M)) == ('a':'c', [2, 3, 4, 5, 2, 3, 4, 5])
    @test axiskeys(vcat(M,M)) == (['a', 'b', 'c', 'a', 'b', 'c'], 2:5)
    @test axiskeys(hcat(MN,MN)) == ('a':'c', [2, 3, 4, 5, 2, 3, 4, 5])
    @test axiskeys(hcat(M,MN)) == ('a':'c', [2, 3, 4, 5, 2, 3, 4, 5])

    V = wrapdims(rand(1:99, 3), r=['a', 'b', 'c'])
    VN = NamedDimsArray(V.data.data, r=['a', 'b', 'c'])
    @test axiskeys(hcat(M,V)) == ('a':'c', [2, 3, 4, 5, 1])
    @test axiskeys(hcat(M,VN)) == ('a':'c', [2, 3, 4, 5, 1])
    @test axiskeys(hcat(MN,V)) == ('a':'c', [2, 3, 4, 5, 1])
    @test axiskeys(hcat(V,V),2) === Base.OneTo(2)
    @test axiskeys(hcat(V,VN),2) === Base.OneTo(2)

    @test axiskeys(vcat(V,V),1) == ['a', 'b', 'c', 'a', 'b', 'c']
    @test axiskeys(vcat(V,VN),1) == ['a', 'b', 'c', 'a', 'b', 'c']
    @test axiskeys(vcat(V', V')) == (1:2, ['a', 'b', 'c'])
    @test axiskeys(vcat(V', VN')) == (1:2, ['a', 'b', 'c'])

    @test hcat(M, ones(3)) == hcat(M.data, ones(3))
    @test axiskeys(hcat(M, ones(3))) == ('a':1:'c', [2, 3, 4, 5, 1])

    @test axiskeys(cat(M,M, dims=3)) == ('a':1:'c', 2:5, Base.OneTo(2))
    @test axiskeys(cat(M.data,M,M, dims=3)) == ('a':1:'c', 2:5, Base.OneTo(3))
    @test axiskeys(cat(M,M, dims=(1,2))) == (['a','b','c', 'a','b','c'], [2,3,4,5, 2,3,4,5])

    @test axiskeys(cat(MN,MN, dims=3)) == ('a':1:'c', 2:5, Base.OneTo(2))
    @test axiskeys(cat(M,MN, dims=3)) == ('a':1:'c', 2:5, Base.OneTo(2))

    @test axiskeys(cat(M,M, dims=:z)) == ('a':1:'c', 2:5, Base.OneTo(2))
    @test dimnames(cat(M,M, dims=:z)) == (:r, :c, :z)

    @test axiskeys(cat(A3, A3, dims=:p)) == ('a':1:'c', 2:5, [10.0, 20.0, 10.0, 20.0])
    @test axiskeys(hcat(A3, A3)) == ('a':1:'c', vcat(2:5,2:5), [10.0, 20.0])

    # reduce(*cat, a_of_as)
    vs = [KeyedArray(i.*[1,2,3], foo=[:a,:b,:c]) for i in 1:4];
    @test axiskeys(reduce(hcat, vs)) == ([:a,:b,:c], Base.OneTo(4))
    @test axiskeys(reduce(hcat, vs')) == (Base.OneTo(1), [:a,:b,:c,  :a,:b,:c,  :a,:b,:c,  :a,:b,:c])
    @test axiskeys(reduce(vcat, vs)) == ([:a,:b,:c,  :a,:b,:c,  :a,:b,:c,  :a,:b,:c],)

    kvs = wrapdims(vs, bar='a':'d') # add an outer name & keys
    @test axiskeys(reduce(hcat, kvs)) == ([:a,:b,:c], 'a':'d')
    @test axiskeys(reduce(vcat, unname(kvs))) == ([:a,:b,:c,  :a,:b,:c,  :a,:b,:c,  :a,:b,:c],)

    twomats = [wrapdims(rand(2,2), x=10:11, y='a':'b'), wrapdims(rand(2,2), x=10:11, y='c':'d')]
    @test axiskeys(reduce(hcat, twomats)) == (10:11, ['a', 'b', 'c', 'd'])
    @test axiskeys(reduce(vcat, twomats')) == (['a', 'b', 'c', 'd'], 10:11)
    @test_throws ArgumentError axiskeys(reduce(hcat, twomats'))
    @test_throws ArgumentError axiskeys(reduce(vcat, twomats))

end
@testset "matmul" begin

    # two matrices
    @test axiskeys(M * M') === ('a':'c', 'a':'c')
    @test axiskeys(M * rand(4,5)) === ('a':'c', Base.OneTo(5))
    @test axiskeys(rand(2,3) * M) === (Base.OneTo(2), 2:5)
    @test axiskeys(MN * MN') === ('a':'c', 'a':'c')
    @test axiskeys(M * MN') === ('a':'c', 'a':'c')

    # two vectors
    @test (V' * V) isa Int
    @test (V' * VN) isa Int
    @test (V' * rand(Int, 10)) isa Int
    @test (rand(Int, 10)' * V) isa Int
    @test axiskeys(V * V') === (10:10:100, 10:10:100)
    @test dimnames(V * V') === (:v, :v)
    @test axiskeys(V * VN') === (10:10:100, 10:10:100)
    @test dimnames(V * VN') === (:v, :v)
    @test axiskeys(V * rand(1,10)) === (10:10:100, Base.OneTo(10))
    @test dimnames(V * rand(1,10)) === (:v, :_)

    # matrix * vector
    @test axiskeys(M * M('a')) === ('a':'c',)
    @test dimnames(M * M('a')) === (:r,)
    @test axiskeys(M(5)' * M) === (Base.OneTo(1), 2:5)
    @test dimnames(M(5)' * M) === (:_, :c)

end
@testset "div" begin # doesn't work for names yet

    A = wrapdims(rand(Int8,3,4), 'a':'c', 10:10:40)
    C = wrapdims(rand(Int8,3), ['a', 'b', 'c'])
    D = wrapdims(rand(Int8,4), [10, 20, 30, 40])

    @test axiskeys(A \ A) == (10:10:40, 10:10:40)
    @test axiskeys(A \ C) == (10:10:40,)
    @test axiskeys(A' \ D) == ('a':'c',)

end
@testset "copy etc" begin

    # copy, similar, etc
    @test axiskeys(copy(M)) == ('a':'c', 2:5)
    @test zero(M)('a',2) == 0

    @test axiskeys(similar(M, Int)) == axiskeys(M)
    @test AxisKeys.haskeys(similar(M, Int, 3,3)) == false
    @test dimnames(similar(M, 3,3)) == (:r, :c)
    @test AxisKeys.hasnames(similar(M, 2,2,2)) == false

end
@testset "equality" begin

    data = parent(parent(M))
    M2 = wrapdims(data, r='a':'c', c=[2,3,4,5]) # same values
    M3 = wrapdims(data, 'a':'c', 2:5) # no names but same keys
    M4 = wrapdims(data, r='a':'c', c=nothing) # missing keys
    @test M == M2 == M3 == M4
    @test isequal(M, M2) && isequal(M, M3) && isequal(M, M4)
    @test M ≈ M2 ≈ M3 ≈ M4

    M5 = wrapdims(data, r='a':'c', c=4:7) # wrong keys
    M6 = wrapdims(data, r='a':'c', nope=2:5) # wrong name
    M7 = wrapdims(2 .* data, r='a':'c', c=2:5) # wrong data
    @test M != M5
    @test M != M6
    @test M != M7
    @test !isapprox(M, M5) && !isapprox(M, M7)

    @test M == MN # order of wrappers
    @test isequal(M, MN)
    @test M ≈ MN

    # Arrays of different sizes
    V = wrapdims(vec(data), r='a':'l')  # vectorised
    A = wrapdims(reshape(data, 2, 2, 3), r=['a','b'], c=[1, 2], l=[:x, :y, :z])  # cube

    @test M != V
    @test M != A
    @test V != A
    @test !isequal(M, V)
    @test !isequal(M, A)
    @test !isequal(V, A)
    @test !isapprox(M, V)
    @test !isapprox(M, A)
    @test !isapprox(V, A)

end

@testset "linalg" begin
    using LinearAlgebra

    # more wrappers -- https://github.com/mcabbott/AxisKeys.jl/issues/69
    @testset "wrappers" begin
        D3 = Diagonal(randn(3))
        U3 = UpperTriangular(rand(3,3))
        U4 = UpperTriangular(rand(4,4) .+ im)
        A3 = rand(10, 3)'
        A10 = rand(3, 10)'
        T3 = transpose(rand(10, 3))
        T10 = transpose(rand(3, 10))

        @test axiskeys(D3 * unname(M), 2) == 2:5
        @test_skip U3 * M # *(::UpperTriangular, ::NamedDimsArray) is ambiguous
        @test axiskeys(U4 * unname(M'), 2) == 'a':'c'
        @test axiskeys(A3 * V, 1) == Base.OneTo(3)
        @test axiskeys(T3 * V, 1) == Base.OneTo(3)

        @test axiskeys(D3 \ M, 2) == 2:5
        @test axiskeys(U3 \ M, 2) == 2:5
        @test_skip U4 / M  # no method matching ldiv!(::NamedDims.NamedFactorization, ::Matrix{ComplexF64})
        @test axiskeys(A10 \ V, 1) == 1:3
        @test axiskeys(T10 \ V, 1) == 1:3
    end

    @testset "cholesky" begin
        A = rand(10, 3)
        ka = KeyedArray(A, (0:9, ["a", "b", "c"]))
        kanda = KeyedArray(A, time=0:9, id=["a", "b", "c"])

        # These should all be equivalent because we aren't wrapping the U and L matrices in
        # the Factorization
        @test cholesky(cor(A)) == cholesky(cor(ka)) == cholesky(cor(kanda))
        @test cholesky(Hermitian(cor(A))) == cholesky(Hermitian(cor(ka))) == cholesky(Hermitian(cor(kanda)))
        @test all(isposdef, (cor(A), cor(ka), cor(kanda)))
    end
end


text:
export Adam
using LinearAlgebra: lmul!, BLAS

mutable struct Adam
    lr::AbstractFloat
    gclip::AbstractFloat
    beta1::AbstractFloat
    beta2::AbstractFloat
    eps::AbstractFloat
    t::Int
    fstm
    scndm
end

Adam(; lr=0.001, gclip=0, beta1=0.9, beta2=0.999, eps=1e-8)=Adam(lr, gclip, beta1, beta2, eps, 0, nothing, nothing)

function update!(w, g, p::Adam)
    gclip!(g, p.gclip)
    if p.fstm===nothing; p.fstm=zero(w); p.scndm=zero(w); end
    p.t += 1
    lmul!(p.beta1, p.fstm)
    BLAS.axpy!(1-p.beta1, g, p.fstm)
    lmul!(p.beta2, p.scndm)
    BLAS.axpy!(1-p.beta2, g .* g, p.scndm)
    fstm_corrected = p.fstm / (1 - p.beta1 ^ p.t)
    scndm_corrected = p.scndm / (1 - p.beta2 ^ p.t)
    BLAS.axpy!(-p.lr, @.(fstm_corrected / (sqrt(scndm_corrected) + p.eps)), w)
end

function gclip!(g, gclip)
    if gclip == 0
        g
    else
        gnorm = vecnorm(g)
        if gnorm <= gclip
            g
        else
            BLAS.scale!(gclip/gnorm, g)
        end
    end
end


text:
## TRAITS FOR MEASURES

is_measure(::Any) = false

const MEASURE_TRAITS =
    [:name, :target_scitype, :supports_weights, :prediction_type, :orientation,
     :reports_each_observation, :is_feature_dependent]

# already defined in model_traits.jl:
# name              - fallback for non-MLJType is string(M) where M is arg
# target_scitype    - fallback value = Unknown
# supports_weights  - fallback value = false
# prediction_type   - fallback value = :unknown (also: :deterministic,
#                                           :probabilistic, :interval)

# specfic to measures:
orientation(::Type) = :loss  # other options are :score, :other
reports_each_observation(::Type) = false
is_feature_dependent(::Type) = false

# extend to instances:
orientation(m) = orientation(typeof(m))
reports_each_observation(m) = reports_each_observation(typeof(m))
is_feature_dependent(m) = is_feature_dependent(typeof(m))

# specific to probabilistic measures:
distribution_type(::Type) = missing

## DISPATCH FOR EVALUATION

# yhat - predictions (point or probabilisitic)
# X - features
# y - target observations
# w - per-observation weights

value(measure, yhat, X, y, w) = value(measure, yhat, X, y, w,
                                      Val(is_feature_dependent(measure)),
                                      Val(supports_weights(measure)))

## DEFAULT EVALUATION INTERFACE

#  is feature independent, weights not supported:
value(measure, yhat, X, y, w, ::Val{false}, ::Val{false}) = measure(yhat, y)

#  is feature dependent:, weights not supported:
value(measure, yhat, X, y, w, ::Val{true}, ::Val{false}) = measure(yhat, X, y)

#  is feature independent, weights supported:
value(measure, yhat, X, y, w, ::Val{false}, ::Val{true}) = measure(yhat, y, w)
value(measure, yhat, X, y, ::Nothing, ::Val{false}, ::Val{true}) = measure(yhat, y)

#  is feature dependent, weights supported:
value(measure, yhat, X, y, w, ::Val{true}, ::Val{true}) = measure(yhat, X, y, w)
value(measure, yhat, X, y, ::Nothing, ::Val{true}, ::Val{true}) = measure(yhat, X, y)

## HELPERS

"""
    check_dimension(ŷ, y)

Check that two vectors have compatible dimensions
"""
function check_dimensions(ŷ::AbstractVector, y::AbstractVector)
    length(y) == length(ŷ) ||
        throw(DimensionMismatch("Differing numbers of observations and "*
                                "predictions. "))
    return nothing
end

function check_pools(ŷ, y)
    levels(y) == levels(ŷ[1]) ||
        error("Conflicting categorical pools found "*
              "in observations and predictions. ")
    return nothing
end

## FOR BUILT-IN MEASURES

abstract type Measure <: MLJType end
is_measure(::Measure) = true

Base.show(stream::IO, ::MIME"text/plain", m::Measure) = print(stream, "$(name(m)) (callable Measure)")
Base.show(stream::IO, m::Measure) = print(stream, name(m))

MLJBase.info(measure, ::Val{:measure}) =
    (name=name(measure),
     target_scitype=target_scitype(measure),
     prediction_type=prediction_type(measure),
     orientation=orientation(measure),
     reports_each_observation=reports_each_observation(measure),
     is_feature_dependent=is_feature_dependent(measure),
     supports_weights=supports_weights(measure))

## REGRESSOR METRICS (FOR DETERMINISTIC PREDICTIONS)

mutable struct MAV<: Measure end
"""
    mav(ŷ, y)
    mav(ŷ, y, w)

Mean absolute error (also known as MAE).

``\\text{MAV} =  n^{-1}∑ᵢ|yᵢ-ŷᵢ|`` or ``\\text{MAV} =  ∑ᵢwᵢ|yᵢ-ŷᵢ|/∑ᵢwᵢ``

For more information, run `info(mav)`.

"""
mav = MAV()
name(::Type{<:MAV}) = "mav"

target_scitype(::Type{<:MAV}) = Union{AbstractVector{Continuous},AbstractVector{Count}}
prediction_type(::Type{<:MAV}) = :deterministic
orientation(::Type{<:MAV}) = :loss
reports_each_observation(::Type{<:MAV}) = false
is_feature_dependent(::Type{<:MAV}) = false
supports_weights(::Type{<:MAV}) = true

function (::MAV)(ŷ::AbstractVector{<:Real}, y::AbstractVector{<:Real})
    check_dimensions(ŷ, y)
    ret = 0.0
    for i in eachindex(y)
        dev = y[i] - ŷ[i]
        ret += abs(dev)
    end
    return ret / length(y)
end

function (::MAV)(ŷ::AbstractVector{<:Real}, y::AbstractVector{<:Real}, w::AbstractVector{<:Real})
    check_dimensions(ŷ, y)
    check_dimensions(y, w)
    ret = 0.0
    for i in eachindex(y)
        dev = w[i]*(y[i] - ŷ[i])
        ret += abs(dev)
    end
    return ret / sum(w)
end

# synonym
"""
mae(ŷ, y)

See also [`mav`](@ref).
"""
const mae = mav

struct RMS <: Measure end
"""
    rms(ŷ, y)
    rms(ŷ, y, w)

Root mean squared error:

``\\text{RMS} = \\sqrt{n^{-1}∑ᵢ|yᵢ-ŷᵢ|^2}`` or ``\\text{RMS} = \\sqrt{\\frac{∑ᵢwᵢ|yᵢ-ŷᵢ|^2}{∑ᵢwᵢ}}``

For more information, run `info(rms)`.

"""
rms = RMS()
name(::Type{<:RMS}) = "rms"
target_scitype(::Type{<:RMS}) = Union{AbstractVector{Continuous},AbstractVector{Count}}
prediction_type(::Type{<:RMS}) = :deterministic
orientation(::Type{<:RMS}) = :loss
reports_each_observation(::Type{<:RMS}) = false
is_feature_dependent(::Type{<:RMS}) = false
supports_weights(::Type{<:RMS}) = true

function (::RMS)(ŷ::AbstractVector{<:Real}, y::AbstractVector{<:Real})
    check_dimensions(ŷ, y)
    ret = 0.0
    for i in eachindex(y)
        dev = y[i] - ŷ[i]
        ret += dev * dev
    end
    return sqrt(ret / length(y))
end

function (::RMS)(ŷ::AbstractVector{<:Real}, y::AbstractVector{<:Real}, w::AbstractVector{<:Real})
    check_dimensions(ŷ, y)
    ret = 0.0
    for i in eachindex(y)
        dev = y[i] - ŷ[i]
        ret += w[i]*dev*dev
    end
    return sqrt(ret / sum(w))
end

struct L2 <: Measure end
"""
    l2(ŷ, y)
    l2(ŷ, y, w)

L2 per-observation loss.

For more information, run `info(l2)`.

"""
l2 = L2()
name(::Type{<:L2}) = "l2"
target_scitype(::Type{<:L2}) = Union{AbstractVector{Continuous},AbstractVector{Count}}
prediction_type(::Type{<:L2}) = :deterministic
orientation(::Type{<:L2}) = :loss
reports_each_observation(::Type{<:L2}) = true
is_feature_dependent(::Type{<:L2}) = false
supports_weights(::Type{<:L2}) = true

function (::L2)(ŷ::AbstractVector{<:Real}, y::AbstractVector{<:Real})
    (check_dimensions(ŷ, y); (y - ŷ).^2)
end

function (::L2)(ŷ::AbstractVector{<:Real}, y::AbstractVector{<:Real}, w::AbstractVector{<:Real})
    check_dimensions(ŷ, y)
    check_dimensions(w, y)
    return (y - ŷ).^2 .* w ./ (sum(w)/length(y))
end

struct L1 <: Measure end
"""
    l1(ŷ, y)
    l1(ŷ, y, w)

L1 per-observation loss.

For more information, run `info(l1)`.

"""
l1 = L1()
name(::Type{<:L1}) = "l1"
target_scitype(::Type{<:L1}) = Union{AbstractVector{Continuous},AbstractVector{Count}}
prediction_type(::Type{<:L1}) = :deterministic
orientation(::Type{<:L1}) = :loss
reports_each_observation(::Type{<:L1}) = true
is_feature_dependent(::Type{<:L1}) = false
supports_weights(::Type{<:L1}) = true

function (::L1)(ŷ::AbstractVector{<:Real}, y::AbstractVector{<:Real})
    (check_dimensions(ŷ, y); abs.(y - ŷ))
end

function (::L1)(ŷ::AbstractVector{<:Real}, y::AbstractVector{<:Real}, w::AbstractVector{<:Real})
    check_dimensions(ŷ, y)
    check_dimensions(w, y)
    return abs.(y - ŷ) .* w ./ (sum(w)/length(y))
end

struct RMSL <: Measure end
"""
    rmsl(ŷ, y)

Root mean squared logarithmic error:

``\\text{RMSL} = n^{-1}∑ᵢ\\log\\left({yᵢ \\over ŷᵢ}\\right)``

For more information, run `info(rmsl)`.

See also [`rmslp1`](@ref).

"""
rmsl = RMSL()
name(::Type{<:RMSL}) = "rmsl"
target_scitype(::Type{<:RMSL}) = Union{AbstractVector{Continuous},AbstractVector{Count}}
prediction_type(::Type{<:RMSL}) = :deterministic
orientation(::Type{<:RMSL}) = :loss
reports_each_observation(::Type{<:RMSL}) = false
is_feature_dependent(::Type{<:RMSL}) = false
supports_weights(::Type{<:RMSL}) = false

function (::RMSL)(ŷ::AbstractVector{<:Real}, y::AbstractVector{<:Real})
    check_dimensions(ŷ, y)
    ret = 0.0
    for i in eachindex(y)
        dev = log(y[i]) - log(ŷ[i])
        ret += dev * dev
    end
    return sqrt(ret / length(y))
end

struct RMSLP1 <: Measure end
"""
    rmslp1(ŷ, y)

Root mean squared logarithmic error with an offset of 1:

``\\text{RMSLP1} = n^{-1}∑ᵢ\\log\\left({yᵢ + 1 \\over ŷᵢ + 1}\\right)``

For more information, run `info(rmslp1)`.

See also [`rmsl`](@ref).
"""
rmslp1 = RMSLP1()
name(::Type{<:RMSLP1}) = "rmslp1"
target_scitype(::Type{<:RMSLP1}) = Union{AbstractVector{Continuous},AbstractVector{Count}}
prediction_type(::Type{<:RMSLP1}) = :deterministic
orientation(::Type{<:RMSLP1}) = :loss
reports_each_observation(::Type{<:RMSLP1}) = false
is_feature_dependent(::Type{<:RMSLP1}) = false
supports_weights(::Type{<:RMSLP1}) = false

function (::RMSLP1)(ŷ::AbstractVector{<:Real}, y::AbstractVector{<:Real})
    check_dimensions(ŷ, y)
    ret = 0.0
    for i in eachindex(y)
        dev = log(y[i] + 1) - log(ŷ[i] + 1)
        ret += dev * dev
    end
    return sqrt(ret / length(y))
end

struct RMSP <: Measure end
"""
    rmsp(ŷ, y)

Root mean squared percentage loss:

``\\text{RMSP} = m^{-1}∑ᵢ \\left({yᵢ-ŷᵢ \\over yᵢ}\\right)^2``

where the sum is over indices such that `yᵢ≂̸0` and `m` is the number
of such indices.

For more information, run `info(rmsp)`.

"""
rmsp = RMSP()
name(::Type{<:RMSP}) = "rmsp"
target_scitype(::Type{<:RMSP}) = Union{AbstractVector{Continuous},AbstractVector{Count}}
prediction_type(::Type{<:RMSP}) = :deterministic
orientation(::Type{<:RMSP}) = :loss
reports_each_observation(::Type{<:RMSP}) = false
is_feature_dependent(::Type{<:RMSP}) = false
supports_weights(::Type{<:RMSP}) = false

function (::RMSP)(ŷ::AbstractVector{<:Real}, y::AbstractVector{<:Real})
    check_dimensions(ŷ, y)
    ret = 0.0
    count = 0
    for i in eachindex(y)
        if y[i] != 0.0
            dev = (y[i] - ŷ[i])/y[i]
            ret += dev * dev
            count += 1
        end
    end
    return sqrt(ret/count)
end

## CLASSIFICATION METRICS (FOR DETERMINISTIC PREDICTIONS)

struct MisclassificationRate <: Measure end

"""
    misclassification_rate(ŷ, y)
    misclassification_rate(ŷ, y, w)

Returns the rate of misclassification of the (point) predictions `ŷ`,
given true observations `y`, optionally weighted by the weights
`w`. All three arguments must be abstract vectors of the same length.

For more information, run `info(misclassification_rate)`.

"""
misclassification_rate = MisclassificationRate()
name(::Type{<:MisclassificationRate}) = "misclassification_rate"
target_scitype(::Type{<:MisclassificationRate}) = AbstractVector{<:Finite}
prediction_type(::Type{<:MisclassificationRate}) = :deterministic
orientation(::Type{<:MisclassificationRate}) = :loss
reports_each_observation(::Type{<:MisclassificationRate}) = false
is_feature_dependent(::Type{<:MisclassificationRate}) = false
supports_weights(::Type{<:MisclassificationRate}) = true

(::MisclassificationRate)(ŷ::AbstractVector{<:CategoricalElement},
                          y::AbstractVector{<:CategoricalElement}) =
                              mean(y .!= ŷ)
(::MisclassificationRate)(ŷ::AbstractVector{<:CategoricalElement},
                          y::AbstractVector{<:CategoricalElement},
                          w::AbstractVector{<:Real}) =
                              sum((y .!= ŷ) .*w) / sum(w)

## CLASSIFICATION METRICS (FOR PROBABILISTIC PREDICTIONS)

struct CrossEntropy <: Measure end

"""
    cross_entropy(ŷ, y::AbstractVector{<:Finite})

Given an abstract vector of `UnivariateFinite` distributions `ŷ` (ie,
probabilistic predictions) and an abstract vector of true observations
`y`, return the negative log-probability that each observation would
occur, according to the corresponding probabilistic prediction.

For more information, run `info(cross_entropy)`.

"""
cross_entropy = CrossEntropy()
name(::Type{<:CrossEntropy}) = "cross_entropy"
target_scitype(::Type{<:CrossEntropy}) = AbstractVector{<:Finite}
prediction_type(::Type{<:CrossEntropy}) = :probabilistic
orientation(::Type{<:CrossEntropy}) = :loss
reports_each_observation(::Type{<:CrossEntropy}) = true
is_feature_dependent(::Type{<:CrossEntropy}) = false
supports_weights(::Type{<:CrossEntropy}) = false

# for single observation:
_cross_entropy(d, y) = -log(pdf(d, y))

function (::CrossEntropy)(ŷ::AbstractVector{<:UnivariateFinite},
                          y::AbstractVector{<:CategoricalElement})
    check_dimensions(ŷ, y)
    check_pools(ŷ, y)
    return broadcast(_cross_entropy, Any[ŷ...], y)
end

# TODO: support many distributions/samplers D below:
struct BrierScore{D} <: Measure end

"""
    brier = BrierScore(; distribution=UnivariateFinite)
    brier(ŷ, y)

Given an abstract vector of distributions `ŷ` and an abstract vector
of true observations `y`, return the corresponding Brier (aka
quadratic) scores.

Currently only `distribution=UnivariateFinite` is supported, which is
applicable to superivised models with `Finite` target scitype. In this
case, if `p(y)` is the predicted probability for a *single*
observation `y`, and `C` all possible classes, then the corresponding
Brier score for that observation is given by

``2p(y) - \\left(\\sum_{η ∈ C} p(η)^2\\right) - 1``

For more information, run `info(brier_score)`.

"""
function BrierScore(; distribution=UnivariateFinite)
    distribution == UnivariateFinite ||
        error("Only `UnivariateFinite` Brier scores currently supported. ")
    return BrierScore{distribution}()
end

name(::Type{<:BrierScore{D}}) where D = "BrierScore{$(string(D))}"
target_scitype(::Type{<:BrierScore{D}}) where D = AbstractVector{<:Finite}
prediction_type(::Type{<:BrierScore}) = :probabilistic
orientation(::Type{<:BrierScore}) = :score
reports_each_observation(::Type{<:BrierScore}) = true
is_feature_dependent(::Type{<:BrierScore}) = false
supports_weights(::Type{<:BrierScore}) = true

# For single observations (no checks):

# UnivariateFinite:
function brier_score(d::UnivariateFinite, y)
    levels = classes(d)
    pvec = broadcast(pdf, d, levels)
    offset = 1 + sum(pvec.^2)
    return 2*pdf(d, y) - offset
end

# For multiple observations:

# UnivariateFinite:
function (::BrierScore{<:UnivariateFinite})(
    ŷ::AbstractVector{<:UnivariateFinite},
    y::AbstractVector{<:CategoricalElement})
    check_dimensions(ŷ, y)
    check_pools(ŷ, y)
    return broadcast(brier_score, Any[ŷ...], y)
end

function (score::BrierScore{<:UnivariateFinite})(
    ŷ, y, w::AbstractVector{<:Real})
    check_dimensions(y, w)
    return w .* score(ŷ, y) ./ (sum(w)/length(y))
end

## DEFAULT MEASURES

default_measure(model::M) where M<:Supervised =
    default_measure(model, target_scitype(M))
default_measure(model, ::Any) = nothing
default_measure(model::Deterministic,
                ::Type{<:Union{AbstractVector{Continuous},
                               AbstractVector{Count}}}) = rms
# default_measure(model::Probabilistic,
#                 ::Type{<:Union{AbstractVector{Continuous},
#                                AbstractVector{Count}}}) = rms
default_measure(model::Deterministic,
                ::Type{<:AbstractVector{<:Finite}}) =
                    misclassification_rate
default_measure(model::Probabilistic,
                ::Type{<:AbstractVector{<:Finite}}) =
                    cross_entropy


text:
#works for general tensors
expectation_value(state::Union{InfiniteMPS,MPSComoving,FiniteMPS},opp::AbstractTensorMap) = expectation_value(state,fill(opp,length(state)))
function expectation_value(state::Union{InfiniteMPS,MPSComoving,FiniteMPS},opps::AbstractArray{<:AbstractTensorMap})
    #todo : gauge gets moved all over the place for finite and comoving states
    #this will invalidate possible caches
    #we should probably not be moving the gauge
    map(zip(state.AC,opps)) do (t,opp)
        tr(t'*permute(opp*permute(t,TensorKit.allind(t)[2:end-1],(1,TensorKit.numind(t))),(TensorKit.numind(t)-1,TensorKit.allind(t)[1:end-2]...),(TensorKit.numind(t),)))
    end
end

"""
calculates the expectation value of op, where op is a plain tensormap where the first index works on site at
"""
function expectation_value(state::Union{FiniteMPS{T},MPSComoving{T},InfiniteMPS{T}},op::AbstractTensorMap,at::Int) where T <: MPSTensor
    expectation_value(state,decompose_localmpo(add_util_leg(op)),at);
end

"""
calculates the expectation value of op = op1*op2*op3*... (ie an N site operator) starting at site at
"""
function expectation_value(state::Union{FiniteMPS{T},MPSComoving{T},InfiniteMPS{T}},op::AbstractArray{<:AbstractTensorMap}, at::Int) where T <: MPSTensor
    @tensor tmp[-1  -2; -3 -4] := state.AC[at+0][1,2,-4]*op[1][-1,3,-3,2]*conj(state.AC[at+0][1,3,-2])

    for index in 2:length(op)
        @tensor tmp[-1, -2, -3, -4] := tmp[-1,1,2,4]*state.AR[at+index-1][4,5,-4]*op[index][2,3,-3,5]*conj(state.AR[at+index-1][1,3,-2])
    end

    return( @tensor tmp[1,2,1,2] )
end

"""
    calculates the expectation value for the given operator/hamiltonian
"""
expectation_value(state,ham::MPOHamiltonian,envs=environments(state,ham)) = expectation_value(state,envs);

function expectation_value(state::MPSComoving,envs::FinEnv)
    ham = envs.opp;

    vals = expectation_value_fimpl(state,envs);

    tot = 0.0+0im;
    for i in 1:ham.odim
        for j in 1:ham.odim

            tot+= @tensor  leftenv(envs,length(state),state)[i][1,2,3]*
                            state.AC[end][3,4,5]*
                            rightenv(envs,length(state),state)[j][5,6,7]*
                            ham[length(state),i,j][2,8,6,4]*
                            conj(state.AC[end][1,8,7])

        end
    end

    return vals,tot/(norm(state.AC[end])^2);
end

expectation_value(state::FiniteMPS,envs::FinEnv) = expectation_value_fimpl(state,envs)
function expectation_value_fimpl(state::Union{MPSComoving,FiniteMPS},envs::FinEnv)
    ham = envs.opp;

    ens=zeros(eltype(eltype(state)),length(state))
    for i=1:length(state)
        for (j,k) in keys(ham,i)

            if !((j == 1 && k!= 1) || (k == ham.odim && j!=ham.odim))
                continue
            end

            cur = @tensor leftenv(envs,i,state)[j][1,2,3]*state.AC[i][3,7,5]*rightenv(envs,i,state)[k][5,8,6]*conj(state.AC[i][1,4,6])*ham[i,j,k][2,4,8,7]
            if !(j==1 && k == ham.odim)
                cur/=2
            end

            ens[i]+=cur
        end
    end

    n = norm(state.AC[end])^2
    return ens./n;
end

function expectation_value(st::InfiniteMPS,prevca::MPOHamInfEnv);
    ham = prevca.opp;
    #calculate energy density
    len = length(st);
    ens = PeriodicArray(zeros(eltype(st.AR[1]),len));
    for i=1:len
        util = Tensor(ones,space(prevca.lw[i+1,ham.odim],2))
        for j=ham.odim:-1:1
            apl = transfer_left(leftenv(prevca,i,st)[j],ham[i,j,ham.odim],st.AL[i],st.AL[i]);
            ens[i] += @tensor apl[1,2,3]*r_LL(st,i)[3,1]*conj(util[2])
        end
    end
    return ens
end

#the mpo hamiltonian over n sites has energy f+n*edens, which is what we calculate here. f can then be found as this - n*edens
expectation_value(st::InfiniteMPS,ham::MPOHamiltonian,size::Int,prevca=environments(st,ham)) = expectation_value(st,prevca,size);
function expectation_value(st::InfiniteMPS,prevca::MPOHamInfEnv,size::Int)
    ham = prevca.opp;

    len=length(st)
    start=leftenv(prevca,1,st)
    start=[@tensor x[-1 -2;-3]:=y[1,-2,3]*st.CR[0][3,-3]*conj(st.CR[0][1,-1]) for y in start]

    for i in 1:size
        start=transfer_left(start,ham[i],st.AR[i],st.AR[i])
    end

    tot=0.0+0im
    for i=1:ham.odim
        tot+=@tensor start[i][1,2,3]*rightenv(prevca,size,st)[i][3,2,1]
    end

    return tot
end

expectation_value(st::InfiniteMPS,ca::PerMPOInfEnv) = expectation_value(convert(MPSMultiline,st),ca);
function expectation_value(st::MPSMultiline,ca::PerMPOInfEnv)
    opp = ca.opp;
    retval = PeriodicArray{eltype(st.AC[1,1]),2}(undef,size(st,1),size(st,2));
    for (i,j) in Iterators.product(1:size(st,1),1:size(st,2))
        retval[i,j] = @tensor   leftenv(ca,i,j,st)[1,2,3]*
                                opp[i,j][2,4,5,6]*
                                st.AC[i,j][3,6,7]*
                                rightenv(ca,i,j,st)[7,5,8]*
                                conj(st.AC[i+1,j][1,4,8])
    end
    return retval
end

expectation_value(state::FiniteQP,opp) = expectation_value(convert(FiniteMPS,state),opp)


text:
function equals(obj::FractionField, arg0::Object)
    return jcall(obj, "equals", jboolean, (Object,), arg0)
end

function get_instance(::Type{FractionField})
    return jcall(FractionField, "getInstance", FractionField, ())
end

function get_one(obj::FractionField)
    return jcall(obj, "getOne", Fraction, ())
end

function get_runtime_class(obj::FractionField)
    return jcall(obj, "getRuntimeClass", Class, ())
end

function get_zero(obj::FractionField)
    return jcall(obj, "getZero", Fraction, ())
end

function hash_code(obj::FractionField)
    return jcall(obj, "hashCode", jint, ())
end


text:
"""
     One()
The Differential which is the multiplicative identity.
Basically, this represents `1`.
"""
struct One <: AbstractDifferential end

extern(x::One) = true  # true is a strong 1.

Base.Broadcast.broadcastable(::One) = Ref(One())
Base.Broadcast.broadcasted(::Type{One}) = One()

Base.iterate(x::One) = (x, nothing)
Base.iterate(::One, ::Any) = nothing


text:
"""
This maps (v, c, log_α) ↦ u = log.(p) - log.(Kb) = -log.(Kb./p) = log.(p./Kb)
"""
function sinkhorn_step_logsep!(u::AbstractArray{T}, v, c, εinv, log_α, tempv, tempm) where T
    for i₁ in eachindex(tempv), j₂ in eachindex(tempv)     # sum over j₁
        tempv .= @view c[1][i₁,:]
        rmul!(tempv, -εinv)
        tempv .+= @view v[:,j₂]     # using c = cᵀ
        tempm[i₁,j₂] = logsumexp(tempv)
    end
    for i₁ in eachindex(tempv), i₂ in eachindex(tempv)     # sum over j₂
        tempv .= @view c[2][i₂,:]
        rmul!(tempv, -εinv)
        tempv .-= log_α[i₁,i₂]
        tempv .+= @view tempm[i₁,:]
        u[i₁,i₂] = -logsumexp(tempv)
    end
    return u
end

"""
wasserstein_distance_logseparated(log_α, log_β, u₀, v₀, log_d₁₀, log_d₂₀, c, SP, caches)
log_α: right marginal weights as nxn matrix in log-scale
log_β: left marginal weights as nxn matrix in log-scale
c: Cost between grid points as dx(nxn) vector of matrices

Only differentiable by log_α

TODO: higher dimensions than d=2, number of grid points different in dimensions/marginals
"""

function sinkhorn_dvg_logseparated( log_α::AbstractArray{T}, log_β::AbstractArray{V},
                                    u₀, v₀, log_d₁₀, log_d₂₀,
                                    c, SP, caches,
                                    ) where {T, V}

    MC = caches.MC
    VC = caches.VC
    εinv = 1/SP.ε

    MC[:log_α,V] .= ForwardDiff.value.(log_α) # Value of input weights to speed up the iterations

    MC[:u,V] .= u₀;   MC[:v,V] .= v₀
    if SP.debias
        MC[:log_d₁,V] .= log_d₁₀; MC[:log_d₂,V] .= log_d₂₀
    end

    for l in 1:SP.L
        if SP.averaged_updates
            sinkhorn_step_logsep!(MC[:u₊,V], MC[:v,V], c, εinv, MC[:log_α,V], VC[:t1,V], MC[:t1,V]) # u = log.(p./Kb)
            sinkhorn_step_logsep!(MC[:v₊,V], MC[:u,V], c, εinv, log_β, VC[:t1,V], MC[:t1,V]) # v = log.(q./Ka) assuming K = Kᵀ

            MC[:u,V] .= 0.5 * (MC[:u₊,V] + MC[:u,V])
            MC[:v,V] .= 0.5 * (MC[:v₊,V] + MC[:v,V])
        else
            sinkhorn_step_logsep!(MC[:u,V], MC[:v,V], c, εinv, MC[:log_α,V], VC[:t1,V], MC[:t1,V]) # u = log.(p./Kb)
            sinkhorn_step_logsep!(MC[:v,V], MC[:u,V], c, εinv, log_β, VC[:t1,V], MC[:t1,V]) # v = log.(q./Ka) assuming K = Kᵀ
        end

        if SP.debias # && (l%2==0 || l==SP.L)
            if SP.averaged_updates
                MC[:t3,V] .= MC[:log_d₁,V]
                MC[:t2,V] .= MC[:log_d₁,V] .+ MC[:log_α,V]
                sinkhorn_step_logsep!(MC[:log_d₁₊,V], MC[:t3,V], c, εinv,
                                        MC[:t2,V], VC[:t1,V], MC[:t1,V]) .*= 0.5

                MC[:t3,V] .= MC[:log_d₂,V]
                MC[:t2,V] .= MC[:log_d₂,V] .+ log_β
                sinkhorn_step_logsep!(MC[:log_d₂₊,V], MC[:t3,V], c, εinv,
                                    MC[:t2,V], VC[:t1,V], MC[:t1,V]) .*= 0.5

                MC[:log_d₂,V] .= 0.5 * (MC[:log_d₂₊,V] + MC[:log_d₂,V])
                MC[:log_d₁,V] .= 0.5 * (MC[:log_d₁₊,V] + MC[:log_d₁,V])
            else
                MC[:t3,V] .= MC[:log_d₁,V]
                MC[:t2,V] .= MC[:log_d₁,V] .+ MC[:log_α,V]
                sinkhorn_step_logsep!(MC[:log_d₁,V], MC[:t3,V], c, εinv,
                                        MC[:t2,V], VC[:t1,V], MC[:t1,V]) .*= 0.5

                MC[:t3,V] .= MC[:log_d₂,V]
                MC[:t2,V] .= MC[:log_d₂,V] .+ log_β
                sinkhorn_step_logsep!(MC[:log_d₂,V], MC[:t3,V], c, εinv,
                                    MC[:t2,V], VC[:t1,V], MC[:t1,V]) .*= 0.5
            end
        end
    end

    if SP.debias
        for i in eachindex(MC[:t2,T])
            MC[:t2,T][i] =  exp(log_α[i]) * ( MC[:u,V][i] - MC[:log_d₁,V][i] ) +
                            exp(log_β[i]) * ( MC[:v,T][i] - MC[:log_d₂,T][i] )
        end
    else
        for i in eachindex(MC[:t2,T])
            MC[:t2,T][i] =  exp(log_α[i]) * MC[:u,V][i] + exp(log_β[i]) * MC[:v,V][i]
        end
    end
    S_ε = SP.ε * sum( MC[:t2,T] )

    if SP.update_potentials
        u₀ .= MC[:u,V]
        v₀ .= MC[:v,V]
        if SP.debias
            log_d₁₀ .= MC[:log_d₁,V]
            log_d₂₀ .= MC[:log_d₂,V]
        end
    end

    return S_ε
end

function sinkhorn_dvg_logseparated( log_α::AbstractArray{T}, log_β::AbstractArray{V},
                                            c, SP, caches,
                                            ) where {T, V}
            sinkhorn_dvg_logseparated(  log_α, log_β,
                                        zeros(V, size(log_α)), zeros(V, size(log_β)),
                                        zeros(V, size(log_α)), zeros(V, size(log_β)),
                                        c, SP, caches)
end

function sinkhorn_barycenter_logseparated(λ::AbstractVector{T}, log_α::AbstractVector{AT},
                                            v₀, log_d₀,
                                            c, SP, caches,
                                            ) where {T, V, AT <: AbstractArray{V}}

    MC = caches.MC
    VC = caches.VC
    VMC = caches.VMC
    εinv = 1/SP.ε

    for s in 1:VMC.S
        VMC[:v,T][s] .= v₀[s]
    end
    if SP.debias
        MC[:log_d,T] .= log_d₀
    end

    #=
    if SP.update_warmstart
        for s in 1:VMC.S
            VMC[:vL,T][s] .= VMC[:v,T][s]
        end
        if SP.debias
            MC[:log_dL,T] .= MC[:log_d,T]
        end
        SP.update_warmstart = false
    end

    if SP.warmstart == "ones"
        for s in 1:VMC.S
            VMC[:v,T][s] .= 0
        end
        if SP.debias
            MC[:log_d,T] .= 0
        end
    elseif SP.warmstart == "last_update"
        for s in 1:VMC.S
            VMC[:v,T][s] .= VMC[:vL,T][s]
        end
        if SP.debias
            VMC[:log_d,T] .= VMC[:log_dL,T]
        end
    elseif SP.warmstart == "latest"
        nothing
    else
        error("No valid warmstart")
    end
    =#

    for l in 1:SP.L # Sinkhorn loop

        MC[:t3,T] .= 0

       for s in 1:VMC.S
            sinkhorn_step_logsep!(MC[:t2,T], VMC[:v,T][s], c, εinv, log_α[s], VC[:t1,T], MC[:t1,T])   # this is uˡₛ
            sinkhorn_step_logsep!(VMC[:log_φ,T][s], MC[:t2,T], c, εinv, MC[:t3,T], VC[:t1,T], MC[:t1,T]) .*= -1   # log φˡₛ = log Kᵀ aˡₛ
        end

        SP.debias ? MC[:log_μ,T] .= MC[:log_d,T] : MC[:log_μ,T] .= 0

        for s in 1:MC.S
            MC[:log_μ,T] .+= λ[s] .* VMC[:log_φ,T][s]
        end

        for s in 1:MC.S
            VMC[:v,T][s] .= MC[:log_μ,T] .- VMC[:log_φ,T][s]
        end

        if SP.debias #&& ( l%2==0 || l==SP.L )
            MC[:t3,T] .= MC[:log_d,T]
            MC[:t2,T] .= MC[:log_d,T] .+ MC[:log_μ,T]
            sinkhorn_step_logsep!(MC[:log_d,T], MC[:t3,T], c, εinv, MC[:t2,T], VC[:t1,T], MC[:t1,T]) .*= 0.5
        end
    end

    if SP.update_potentials
        for s in 1:VMC.S
            v₀[s] .= ForwardDiff.value.( VMC[:v,T][s] )
        end
        if SP.debias
            log_d₀ .= ForwardDiff.value.( MC[:log_d,T] )
        end
    end

    return copy(MC[:log_μ,T])
end

function sinkhorn_barycenter_logseparated(λ::AbstractVector{T}, log_α::AbstractVector{AT},
                                            c, SP, caches,
                                            ) where {T, V, AT <: AbstractArray{V}}
    sinkhorn_barycenter_logseparated(λ, log_α,
                                        [zeros(V, size(log_α[s])) for s in eachindex(log_α)],
                                        zeros(V, size(log_α[1])),
                                        c, SP, caches)
end

max_stars_repo_name: JuliaRCM/WassersteinDictionaries

max_issues_repo_name: JuliaRCM/WassersteinDictionaries

max_forks_repo_name: JuliaRCM/WassersteinDictionaries


text:
using ADCME
using PyCall
using LinearAlgebra
using PyPlot
using Random
Random.seed!(233)

function compute_sin(input)
    compute_sin_ = load_op_and_grad("./build/libComputeSin","compute_sin")
    input = convert_to_tensor(Any[input], [Float64]); input = input[1]
    compute_sin_(input)
end

# TODO: specify your input parameters
input = rand(10)
u = compute_sin(input)

sess = Session(CPU=1); init(sess)
@show run(sess, u)


text:
export make_graphene
export link_relatives!
export make_gatom
export make_gbond
export make_gbond!
export make_bondmatrix
export get_gbond
export make_gpolygon
export make_gpolygon!
export link_bond_polygons!
export make_signature
export make_signature!
export common_members
export common_relatives
export find_defect
export find_stonewales
export find_5775
export find_divacancy
export find_divacancy14
export find_flower
export find_butterfly
export filter_relatives_by_type
export filter_members_by_type
export stepout

function make_graphene(atom_xy; image_sampling=1, max_bondlength=10.0, frame=1, dataset="dataset")

    indexed_atoms = make_atoms(atom_xy)
    gatom_count = length(indexed_atoms)

    gatoms = make_gatom.(collect(indexed_atoms))

    atom_groups = collect_atom_groups(atom_xy; max_bondlength=max_bondlength)
    bonds, paths = collect_bonds_paths(atom_groups, indexed_atoms)
    if length(bonds) == 0
        graphene = gatoms
        return graphene
    end
    bondmatrix = make_bondmatrix(bonds)
    gbond_count = maximum(bondmatrix)

    gbonds = map(x -> make_gbond!(x, bondmatrix[get_id.(x)...] + gatom_count), [gatoms[collect(y)] for y in Set(sort.(collect.(bonds)))])
    sort!(gbonds, by=get_id)

    polygons = make_polygons(paths)
    polygon_atoms = first.(polygons)
    polygon_bonds = last.(polygons)
    gpolygon_count = length(polygons)
    polygon_gatoms = [gatoms[collect(x)] for x in polygon_atoms]
    polygon_gbonds = [gbonds[x] for x in map(x -> get_gbond(bondmatrix, x), polygon_bonds)]
    polygon_ids = collect(range(1 + gatom_count + gbond_count, length=gpolygon_count))

    gpolygons = map(make_gpolygon!, polygon_gatoms, polygon_gbonds, polygon_ids)

    gpolygon_dict = Dict([Pair(x._id, x) for x in gpolygons])
    map(x -> link_bond_polygons!(gpolygon_dict, x), gbonds)

    graphene = vcat(gatoms, gbonds, gpolygons)

    noa_dict = Dict([Pair(x._id, get_noa(x)) for x in graphene])
    map(x -> make_signature!(x, noa_dict), graphene)

    for i in 1:length(graphene)
        graphene[i]._frame = frame
        graphene[i]._dataset = dataset
    end
    return graphene
end

function link_relatives!(a, b)
    push!(a._relatives, b._id)
    push!(b._relatives, a._id)
end

function link_relatives!(a, b_vector::Vector)
    link_relatives!.(tuple(a), b_vector)
end

function make_gatom(indexed_atom)
    id = first(indexed_atom)
    atom = last(indexed_atom)
    GAtom(id, getx(atom), gety(atom))
end

function make_gbond(bond_gatoms, gbond_id)
    GBond(gbond_id, mean(get_x.(bond_gatoms)), mean(get_y.(bond_gatoms)), Set(get_id.(bond_gatoms)))
end

function make_gbond!(bond_gatoms::Vector{GAtom}, gbond_id)
    gbond = GBond(gbond_id, mean(get_x.(bond_gatoms)), mean(get_y.(bond_gatoms)))
    link_relatives!(gbond, bond_gatoms)
    return gbond
end

function make_bondmatrix(bonds)
    bonds_directionless = unique(map(x -> sort(collect(x)), bonds))
    index1 = vcat(first.(bonds_directionless), last.(bonds_directionless))
    index2 = vcat(last.(bonds_directionless), first.(bonds_directionless))
    bondids = collect(1:length(bonds_directionless))
    bondmatrix = sparse(index1, index2, repeat(bondids,2))
end

function get_gbond(bondmatrix, bond_index::Tuple{Int, Int})
    bondmatrix[bond_index...]
end

function get_gbond(bondmatrix, bond_index_vector::Vector)
    [get_gbond(bondmatrix, x) for x in bond_index_vector]
end

function make_gpolygon(polygon_gatoms::Vector{GAtom}, polygon_gbonds::Vector{GBond}, gpolygon_id)
    gpolygon = GPolygon(gpolygon_id, mean(get_x.(polygon_gatoms)), mean(get_y.(polygon_gatoms)), length(polygon_gatoms))
    return gpolygon
end

function make_gpolygon!(polygon_gatoms::Vector{GAtom}, polygon_gbonds::Vector{GBond}, gpolygon_id)
    gpolygon = GPolygon(gpolygon_id, mean(get_x.(polygon_gatoms)), mean(get_y.(polygon_gatoms)), length(polygon_gatoms))
    link_relatives!(gpolygon, polygon_gatoms)
    link_relatives!(gpolygon, polygon_gbonds)
    return gpolygon
end

function link_bond_polygons!(gpolygon_dict, gbond::GBond)
    gpolygon_neighbours = Tuple([gpolygon_dict[x] for x in gbond._relatives if x > gbond._id])
    if length(gpolygon_neighbours) == 2
        link_relatives!(gpolygon_neighbours...)
    end
end

function make_signature(g, noa_dict)
    if isempty(g._relatives)
        signature = "lone atom"
    else
        relatives = [noa_dict[x] for x in g._relatives]
        count_vector = counts(relatives, maximum(relatives))
        # signature = join([string(i, "-", count_vector[i], "|") for i in 3:length(count_vector) if !iszero(count_vector[i])])
        if length(count_vector) >= 2
            count_vector[1:2] = [0 0]
        else
            count_vector = zero(count_vector)
        end
        signature = join(map(signature_string, findnz(sparse(count_vector))...))
    end
    return signature
end

signature_string(a, b) = string(a, "-", b, "|")

function make_signature!(g, noa_dict)
    g._signature = make_signature(g, noa_dict)
end

function find_stonewales(graphene, area)
    defects = GDefect[]
    for g in area
        if get_type(g) == "Bond" && get_signature(g) == "7-2|"
            d = stepout(graphene, g, ["Atom", "Polygon"])
            if sort(get_noa.(filter(ispolygon, graphene[collect(get_members(d))]))) == [5, 5, 7, 7]
                if Set(get_signature.(filter(ispolygon, graphene[collect(get_members(d))]))) == Set(["5-2|6-4|7-1|", "6-3|7-2|"])
                    push!(defects, GDefect(0, get_x(d), get_y(d), get_relatives(d), "", get_frame(d), get_dataset(d), get_noa(d), get_members(d), "Stone-Wales"))
                end
            end
        end
    end
    return defects
end
find_stonewales(graphene) = find_stonewales(graphene, graphene)

function find_5775(graphene, area)
    defects = GDefect[]
    for g in area
        if get_type(g) == "Bond" && get_signature(g) == "7-2|"
            gpolygon_members = filter_relatives_by_type(graphene, g, "Polygon")
            if unique(get_signature.(gpolygon_members)) == ["5-1|6-5|7-1|"]
                key_hexagons = filter(x -> get_noa(x) == 6, graphene[collect(get_members(stepout(graphene, g, ["Atom", "Polygon"])))])
                if length(key_hexagons) == 2 && unique(get_signature.(key_hexagons)) == ["5-1|6-3|7-2|"]
                    key_bonds = filter(x -> get_signature(x) == "5-1|7-1|", filter_relatives_by_type(graphene, gpolygon_members, "Bond"))
                    d = stepout(graphene, key_bonds, "Polygon")
                    push!(defects, GDefect(0, get_x(d), get_y(d), get_relatives(d), "", get_frame(d), get_dataset(d), get_noa(d), get_members(d), "5775"))
                end
            end
        end
    end
    return defects
end
find_5775(graphene) = find_5775(graphene, graphene)

function find_divacancy(graphene, area)
    defects = GDefect[]
    for g in area
        if get_type(g) == "Polygon" && get_signature(g) == "5-2|6-6|"
            hexagon_type1 = filter(x -> get_signature(x) == "6-5|8-1|", filter_relatives_by_type(graphene, g, "Polygon"))
            hexagon_type2 = filter(x -> get_signature(x) == "5-1|6-4|8-1|", filter_relatives_by_type(graphene, g, "Polygon"))
            if length(hexagon_type1) == 2 && length(hexagon_type2) == 4
                atoms_type2 = vcat([filter_relatives_by_type(graphene, x, "Atom") for x in hexagon_type2]...)
                if length(atoms_type2) == length(unique(atoms_type2))
                    key_bonds = filter(x -> get_signature(x) == "5-1|8-1|", filter_relatives_by_type(graphene, g, "Bond"))
                    d = stepout(graphene, key_bonds, "Polygon")
                    push!(defects, GDefect(0, get_x(d), get_y(d), get_relatives(d), "", get_frame(d), get_dataset(d), get_noa(d), get_members(d), "Divacancy"))
                end
            end
        end
    end
    return defects
end
find_divacancy(graphene) = find_divacancy(graphene, graphene)

function find_divacancy14(graphene, area)
    defects = GDefect[]
    for g in area
        if get_type(g) == "Polygon" && get_noa(g) == 14 && get_signature(g) == "6-10|"
            relatives = Set(get_id.(filter_relatives_by_type(graphene, g, "Polygon")))
            members = Set(get_id.([g; filter_relatives_by_type(graphene, g, "Atom", "Bond")]))
            push!(defects, GDefect(0, get_x(g), get_y(g), relatives, "", get_frame(g), get_dataset(g), get_noa(g), members, "Divacancy-14"))
        end
    end
    return defects
end
find_divacancy14(graphene) = find_divacancy14(graphene, graphene)

function find_flower(graphene, area)
    defects = GDefect[]
    for g in area
        if get_type(g) == "Atom" && get_signature(g) == "7-3|"
            gpolygon_members = filter_relatives_by_type(graphene, g, "Polygon")
            if unique(get_noa.(gpolygon_members)) == [7] && unique(get_signature.(gpolygon_members)) == ["5-2|6-3|7-2|"]
                d = stepout(graphene, g, ["Bond", "Atom", "Polygon"])
                push!(defects, GDefect(0, get_x(d), get_y(d), get_relatives(d), "", get_frame(d), get_dataset(d), get_noa(d), get_members(d), "Flower"))
            end
        end
    end
    return defects
end
find_flower(graphene) = find_flower(graphene, graphene)

function common_members(graphene, g1, g2, types...)
    m = graphene[collect(intersect(get_members(g1), get_members(g2)))]
    if !isempty(types)
        m = filter(x -> istype(x, types), m)
    end
    return m
end

function common_relatives(graphene, g1, g2, types...)
    r = graphene[collect(intersect(get_relatives(g1), get_relatives(g2)))]
    if !isempty([types...])
        r = filter(x -> istype(x, types...), r)
    end
    return r
end

function find_butterfly(graphene, area)
    defects = GDefect[]
    for g in area
        if get_type(g) == "Polygon" && get_signature(g) == "5-2|7-4|"
            gpolygon_members = filter_relatives_by_type(graphene, g, "Polygon")
            key_polygons = filter(x -> get_noa(x) == 7 && get_signature(x) == "5-2|6-4|7-1|", gpolygon_members)
            if length(key_polygons) == 4
                temp = vcat([common_relatives(graphene, i..., "Atom") for i in subsets(key_polygons, Val(2))]...)
                key_atoms = unique([filter_relatives_by_type(graphene, g, "Atom"); temp])
                d = stepout(graphene, key_atoms, "Polygon")
                push!(defects, GDefect(0, get_x(d), get_y(d), get_relatives(d), "", get_frame(d), get_dataset(d), get_noa(d), get_members(d), "Butterfly"))
            end
        end
    end
    return defects
end
find_butterfly(graphene) = find_butterfly(graphene, graphene)

function find_defect(graphene, type::String)
    if isempty(graphene)
        id_offset = 0
    else
        id_offset = maximum(get_id.(graphene))
    end
    defects = GDefect[]
    if type == "Stone-Wales"
        defects = [defects; find_stonewales(graphene)]
    elseif type == "5775"
        defects = [defects; find_5775(graphene)]
    elseif type == "Divacancy"
        defects = [defects; find_divacancy(graphene)]
        defects = [defects; find_divacancy14(graphene)]
    elseif type == "V2(585)"
        defects = [defects; find_divacancy(graphene)]
    elseif type == "V2(14)"
        defects = [defects; find_divacancy14(graphene)]
    elseif type == "Flower" || type == "V2(555-777)"
        defects = [defects; find_flower(graphene)]
    elseif type == "Butterfly" || type == "V2(5555-6-7777)"
        defects = [defects; find_butterfly(graphene)]
    end
    for i in 1:length(defects)
        defects[i]._id = id_offset+i
    end
    return defects
end

function find_defect(graphene, types...)
    if isempty(graphene)
        id_offset = 0
    else
        id_offset = maximum(get_id.(graphene))
    end
    defects = vcat(map(t -> find_defect(graphene, t), types)...)
    for i in 1:length(defects)
        defects[i]._id = id_offset+i
    end
    return defects
end

function filter_relatives_by_type(graphene, g, type)
    relatives = graphene[collect(get_relatives(g))]
    filter(x -> get_type(x) == type, relatives)
end
filter_relatives_by_type(graphene, g, types...) = vcat(map(t -> filter_relatives_by_type(graphene, g, t), types)...)
filter_relatives_by_type(graphene, g_vector::Vector, types...) = unique(vcat(map(g -> filter_relatives_by_type(graphene, g, types...), g_vector)...))

function filter_members_by_type(graphene, g, type)
    members = graphene[collect(get_members(g))]
    filter(x -> get_type(x) == type, members)
end
filter_members_by_type(graphene, g, types...) = vcat(map(t -> filter_members_by_type(graphene, g, t), types)...)
filter_members_by_type(graphene, g_vector::Vector, types...) = unique(vcat(map(g -> filter_members_by_type(graphene, g, types...), g_vector)...))

function stepout(graphene, g, n::Int, steps=String[]; id=0)
    gmembers = AbstractGPrimitive[]
    if typeof(g) <: Vector
        members = union(get_members.(g)...)
        relatives = union(get_relatives.(g)...)
        frame = get_frame(first(g))
        dataset = get_dataset(first(g))
    else
        members = get_members(g)
        relatives = get_relatives(g)
        frame = get_frame(g)
        dataset = get_dataset(g)
    end
    if isempty(steps)
        itr = Stateful(cycle(["Atom", "Bond", "Polygon"]))
        if typeof(g) <: AbstractGPrimitive
            for x in itr; x != get_type(g) || break; end
        end
        steps = collect(take(itr, n))
    end

    seed = g
    for t in steps
        seed = filter_relatives_by_type(graphene, seed, t)
        for s in seed
            push!(gmembers, s)
            relatives = union(relatives, get_relatives(s))
        end
    end
    x = mean(get_x.(gmembers))
    y = mean(get_y.(gmembers))
    signature = ""
    if steps[end] != "Atom"
        gmembers = [gmembers; filter(x -> isatom(x)||isbond(x), graphene[collect(relatives)])]
    end
    members = union([get_members.(gmembers); members]...)
    gmembers = [gmembers; filter(x -> isatom(x)||isbond(x), graphene[collect(members)])]
    noa = length(unique(filter(isatom, gmembers)))
    members = Set(get_id.(gmembers))
    relatives = setdiff!(relatives, members)

    g_new = GEntry(id, x, y, relatives, signature, frame, dataset, noa, members)
end
stepout(graphene, g, steps::Vector{String}; id=0) = stepout(graphene, g, 0, steps; id=id)
stepout(graphene, g, steps::String; id=0) = stepout(graphene, g, 0, [steps]; id=id)
stepout(graphene, g; id=0) = stepout(graphene, g, 1; id=id)


text:
using ApproxFun,SingularIntegralEquations,Base.Test
    import SingularIntegralEquations:stieltjesmoment,stieltjesjacobimoment

c = [0.9731840665678853,0.11644664868790366,0.8961305368364185,0.30663942299763747,0.4564158422422153,0.7262396331533589,0.6448725310785985,0.05623500638232981,0.8452582204677404,0.25385833878392283]

for z in [1.+1.0im,0.1+0.1im]
    o=Fun(one)
    x=Fun(identity)

    @test_approx_eq stieltjes(o)(z) sum(1/(z-x))
    @test_approx_eq stieltjes(x)(z) sum(x/(z-x))

    for w in (sqrt(1-x),sqrt(1+x))
        @test_approx_eq stieltjes(w*o)(z) sum(w/(z-x))
        @test_approx_eq stieltjes(w*x)(z) sum(w*x/(z-x))
    end

    for a in -0.5:0.5:0.5, b in -0.5:0.5:0.5
        f = Fun(c,WeightedJacobi(a,b))
        @test_approx_eq stieltjes(f)(z) sum(f/(z-x))
    end

    f = Fun(c,WeightedJacobi(0.123,0.456))
    @test_approx_eq stieltjes(f)(z) sum(f/(z-x))

    f = Fun(identity,[-1.,0.,1.])
    @test_approx_eq cauchy(sqrt(Fun(one,space(f))-f^2))(z) cauchy(sqrt(1-Fun()^2),z)
end

x=Fun(identity,[im,0.,1.])
w=2/(sqrt(1-x)*sqrt(1+im*x))

for x in (0.9im,0.4im,0.4,0.9)
    #@test_approx_eq cauchy(w,x,true)-cauchy(w,x,false) w(x)
end


text:
module Utils

using DataFrames, Dates, Statistics

    function parseDates(df; format="y-m-d",dateCol=:date)
        dfmat = DateFormat(format)
        tCol = df[!,dateCol]
        len = length(tCol)
        dts = Array{Date}(undef,len)
        for i in 1:len
            dts[i] = Date(tCol[i],dfmat)
        end
        return dts
    end

    function renameCols(df,newNames)
        @assert length(names(df)) == length(newNames) "List of new names does not equal to columns"
        names!(df,[Symbol(i) for i in newNames])
        return df
    end

    function nse(sim,obs)
        numerator = sum((obs-sim).^2)
        denominator = sum((obs.-mean(obs)).^2)
        return 1 - (numerator/denominator)
    end

end

max_stars_repo_name: jsta/hymod

max_issues_repo_name: jsta/hymod

max_forks_repo_name: jsta/hymod


text:
# ---
# title: 1545. Find Kth Bit in Nth Binary String
# id: problem1545
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: String
# link: <https://leetcode.com/problems/find-kth-bit-in-nth-binary-string/description/>
# hidden: true
# ---
#
# Given two positive integers `n` and `k`, the binary string  `Sn` is formed as
# follows:
#
#   * `S1 = "0"`
#   * `Si = Si-1 + "1" + reverse(invert(Si-1))` for `i > 1`
#
# Where `+` denotes the concatenation operation, `reverse(x)` returns the
# reversed string x, and `invert(x)` inverts all the bits in x (0 changes to 1
# and 1 changes to 0).
#
# For example, the first 4 strings in the above sequence are:
#
#   * `S1 = "0"`
#   * `S2 = "0 **1** 1"`
#   * `S3 = "011 **1** 001"`
#   * `S4 = "0111001 **1** 0110001"`
#
# Return _the_ `kth` _bit_ _in_  `Sn`. It is guaranteed that `k` is valid for
# the given `n`.
#
#
#
# **Example 1:**
#
#
#
#     Input: n = 3, k = 1
#     Output: "0"
#     Explanation: S3 is " ** _0_** 111001". The first bit is "0".
#
#
# **Example 2:**
#
#
#
#     Input: n = 4, k = 11
#     Output: "1"
#     Explanation: S4 is "0111001101 ** _1_** 0001". The 11th bit is "1".
#
#
# **Example 3:**
#
#
#
#     Input: n = 1, k = 1
#     Output: "0"
#
#
# **Example 4:**
#
#
#
#     Input: n = 2, k = 3
#     Output: "1"
#
#
#
#
# **Constraints:**
#
#   * `1 <= n <= 20`
#   * `1 <= k <= 2n - 1`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
module IterativelyHypergraphEmbedding

using SimpleHypergraphs
using Distributions
using LinearAlgebra
include("./algorithm.jl")

export iteratively_hypergraph_embedding

end


text:
# TODO: Plots.jl seems slow for these kind of plots (many segments). Use ParaView instead?

@recipe function plot(t::TrackGenerator{T}) where {T}
    @unpack tracks_by_uid, n_total_tracks = t

    x = Matrix{T}(undef, 2, n_total_tracks)
    y = Matrix{T}(undef, 2, n_total_tracks)

    for j in 1:n_total_tracks
        track = tracks_by_uid[j]
        x[1, j] = track.p[1]
        x[2, j] = track.q[1]
        y[1, j] = track.p[2]
        y[2, j] = track.q[2]
    end

    seriestype  :=  :path
    linewidth   --> 0.20
    legend      --> false
    framestyle  :=  :none

    return (x, y)
end

@recipe function plot(segments::Vector{Segment{T}}) where {T}

    n = length(segments)
    x = Matrix{T}(undef, 2, n) # TODO: compute Float64
    y = Matrix{T}(undef, 2, n)
    z = Matrix{T}(undef, 2, n)

    for (i, segment) in enumerate(segments)
        x[1, i] = segment.p[1]
        x[2, i] = segment.q[1]
        y[1, i] = segment.p[2]
        y[2, i] = segment.q[2]
        z[1, i] = segment.element
        z[2, i] = segment.element
    end

    seriestype  :=  :path
    linewidth   --> 0.25
    legend      --> false
    framestyle  :=  :none
    line_z      :=  z

    return (x, y)
end

@recipe function plot(tracks::Vector{Track})

    # m = length(tracks)
    # n = length(segments)
    x = Matrix{Float64}(undef, 2, 0)
    y = Matrix{Float64}(undef, 2, 0)
    z = Matrix{Float64}(undef, 2, 0)

    for track in tracks
        for segment in track.segments
            x = hcat(x, [segment.p[1], segment.q[1]])
            y = hcat(y, [segment.p[2], segment.q[2]])
            z = hcat(z, [segment.element, segment.element])
        end
    end

    seriestype  :=  :path
    linewidth   --> 0.25
    legend      --> false
    framestyle  :=  :none
    line_z      :=  z

    return (x, y)
end

@recipe function plot(mesh::Mesh)
    @unpack cell_nodes, model = mesh

    grid = get_grid(model)
    nodes = get_node_coordinates(grid)

    # number of cells and number of nodes per cell
    nc = length(cell_nodes)
    nn = all(l -> isequal(l, length(cell_nodes[1])), length.(cell_nodes)) ? length(cell_nodes[1]) : error("error")

    x = Matrix{Float64}(undef, nn + 1, nc) # TODO: compute Float64
    y = Matrix{Float64}(undef, nn + 1, nc)

    for (i, nodes_ids) in enumerate(cell_nodes)

        for (j, node_id) in enumerate(nodes_ids)
            x[j, i] = nodes[node_id][1]
            y[j, i] = nodes[node_id][2]
        end
        # TODO: improve
        x[nn + 1, i] = nodes[nodes_ids[1]][1]
        y[nn + 1, i] = nodes[nodes_ids[1]][2]
    end

    linecolor   --> :black
    seriestype  :=  :path
    linewidth   --> 0.2
    legend      --> false
    framestyle  :=  :none

    return (x, y)
end


text:
#=
Copyright (c) 2015, Intel Corporation
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
- Redistributions of source code must retain the above copyright notice,
  this list of conditions and the following disclaimer.
- Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
THE POSSIBILITY OF SUCH DAMAGE.
=#

using DocOpt

const spot = 100.0
const strike = 110.0
const vol = 0.25
const maturity = 2.0

function init(numPaths::Int, numSteps::Int)
    # Pre-compute some values
    dt = maturity / numSteps
    volSqrtdt = vol * sqrt(dt)
    fwdFactor = vol^2 * dt * -0.5
    vsqrtdt_log2e = volSqrtdt * log2(exp(1.0))
    fwdFactor_log2e = fwdFactor * log2(exp(1.0))

    # Storage
    # Compute all per-path asset values for each time step
    # NOTE: hand-hoisting obscures the basic formula.
    # Can the compiler do this optimization?
    asset = Array(Array{Float64,1}, numSteps+1)
    asset[1] = Array(Float64, numPaths)
    fill!(asset[1], spot)
    for s=2:numSteps+1
        asset[s] = asset[s-1] .* exp2(fwdFactor_log2e .+ vsqrtdt_log2e .* randn(numPaths))
    end
    return asset
end

function model(strike, numPaths, numSteps, asset)
    avgPathFactor = 1.0/numPaths
    putpayoff = max(strike .- asset[numSteps+1], 0.0)
    cashFlowPut = putpayoff * avgPathFactor
    # Now go back in time using regression
    for s = numSteps:-1:2
        curasset = asset[s]
        vcond = curasset .<= strike
        valPut0 = curasset .* 0.0 # zeros(numPaths)
        valPut0[vcond] = 1.0
        valPut1 = valPut0 .* curasset # assetval
        valPut2 = valPut1 .* valPut1  # assetval ^ 2
        valPut3 = valPut2 .* valPut1  # assetval ^ 3
        valPut4 = valPut2 .* valPut2  # assetval ^ 4
        # compute the regression
        ridgecorrect = 0.01
        sum0 = sum(valPut0)
        sum1 = sum(valPut1)
        sum2 = sum(valPut2)
        sum3 = sum(valPut3)
        sum4 = sum(valPut4)
        sum0 += ridgecorrect
        sum1 += ridgecorrect
        sum2 += ridgecorrect
        sum3 += ridgecorrect
        sum4 += ridgecorrect
        sqrtest = Float64[sum0 sum1 sum2;
                          sum1 sum2 sum3;
                          sum2 sum3 sum4]
        invsqr = inv(sqrtest)
        vp0 = dot(cashFlowPut, valPut0)
        vp1 = dot(cashFlowPut, valPut1)
        vp2 = dot(cashFlowPut, valPut2)
        (betaPut0, betaPut1, betaPut2) = [vp0 vp1 vp2] * invsqr
        regImpliedCashFlow = valPut0 .* betaPut0 .+ valPut1 .* betaPut1 .+ valPut2 .* betaPut2
        payoff = valPut0 .* (strike .- curasset)
        pcond = payoff .> regImpliedCashFlow
        cashFlowPut[pcond] = payoff[pcond]
    end

    amerpayoffput = sum(cashFlowPut)
    finalputpayoff = sum(putpayoff)

    return amerpayoffput/numPaths, finalputpayoff/numPaths
end

function main()
    doc = """quant.jl

Quantitative option pricing model.

Usage:
  quant.jl -h | --help
  quant.jl [--assets=<assets>] [--iterations=<iterations>]

Options:
  -h --help                  Show this screen.
  --assets=<assets>          Specify the number of assets to simulate [default: 524288].
  --iterations=<iterations>  Specify a number of iterations [default: 256].
"""
    arguments = docopt(doc)

    paths = parse(Int, arguments["--assets"])
    steps = parse(Int, arguments["--iterations"])

    srand(0)

    println("assets = ", paths)
    println("iterations = ", steps)

    asset = init(paths, steps)
    tic()
    model(strike, 1, 1, asset)
    compiletime = toq()
    println("SELFPRIMED ", compiletime)

    tic()
    amerpayoffput, finalputpayoff = model(strike, paths, steps, asset)
    selftimed = toq()
    println("European Put Option Price: ", finalputpayoff)
    println("American Put Option Price: ", amerpayoffput)
    println("SELFTIMED ", selftimed)

end

main()


text:
using Perturbo
using Test

@testset "Perturbo.jl" begin
    # Write your own tests here.
end


text:
using LightGraphs
using GraphPlot
using Random, RandomNumbers.Xorshifts

@testset "graph_problem.jl" begin
    NV = 5
    NE = 8
    NP = 20
    seed = 4

    pp1 = PMDPs.graph_pp(;NV=NV, NE=NE, NP=NP, seed=seed)

    g = SimpleDiGraph(NV, NE, seed=seed)
    # display(gplot(g, nodelabel=1:nv(g)))
    pp2 = PMDPs.graph_pp(g, NP)

    # @test pp1==pp2 #likely an issue with not seeding rng somewhere
end

# g = SimpleDiGraph(8, 20, seed=38)
# connected_components(g)
# display(gplot(g, nodelabel=1:nv(g)))