text:
using Pkg
 Pkg.activate(".")

verbose = true
if verbose   println("# Loading RvSpecML")    end
 using RvSpectML
 if verbose   println("# Loading other packages")    end
 using DataFrames, Query, Statistics
 using Dates

# USER:  The default paths that specify where datafiles can be entered here or overridden in examples/data_paths.jl
target_subdir = "101501"   # USER: Replace with directory of your choice
 fits_target_str = "101501"
 output_dir = "examples/output"
 default_paths_to_search = [pwd(),"examples",joinpath(pkgdir(RvSpectML),"examples"),"/gpfs/group/ebf11/default/ebf11/expres/inputs"]

 make_plots = true
 write_ccf_to_csv = true
 write_rvs_to_csv = true
 write_order_ccf_to_csv = true
 write_template_to_csv = true
 write_spectral_grid_to_jld2 = true
 write_dcpca_to_csv = true
 write_lines_to_csv = true

has_loaded_data = false
 has_computed_ccfs = false
 has_computed_rvs = false
 has_computed_order_ccfs = false
 has_computed_template = false
 has_computed_dcpca = false
 has_found_lines = false
 has_computed_ccfs2 = false
 has_computed_rvs2 = false
 has_computed_template2 = false
 has_computed_dcpca2 = false

if !has_loaded_data
  if verbose println("# Finding what data files are avaliable.")  end
  df_files = make_manifest(target_subdir, EXPRES )

  if verbose println("# Reading in customized parameters.")  end
  eval(code_to_include_param_jl())

  if verbose println("# Reading in FITS files.")  end
  @time expres_data = map(EXPRES.read_data,eachrow(df_files_use))
  has_loaded_data = true

  if verbose println("# Extracting order list timeseries from spectra.")  end
  order_list_timeseries = RvSpectML.make_order_list_timeseries(expres_data)
  order_list_timeseries = RvSpectML.filter_bad_chunks(order_list_timeseries,verbose=true)
  RvSpectML.normalize_spectra!(order_list_timeseries,expres_data);
end

if !has_computed_ccfs || true
if verbose println("# Reading line list for CCF: ", linelist_for_ccf_filename, ".")  end
 lambda_range_with_good_data = get_λ_range(expres_data)
 espresso_filename = joinpath(pkgdir(RvSpectML),"data","masks",linelist_for_ccf_filename)
 espresso_df = RvSpectML.read_linelist_espresso(espresso_filename)
 line_list_df = EXPRES.filter_line_list(espresso_df,first(expres_data).inst)

if verbose println("# Removing lines with telluric contamination.")  end    # Currently only works for EXPRES data
 Δv_to_avoid_tellurics = 14000.0
 line_list_to_search_for_tellurics = copy(line_list_df)
 line_list_to_search_for_tellurics.lambda_lo = line_list_to_search_for_tellurics.lambda./calc_doppler_factor(Δv_to_avoid_tellurics)
 line_list_to_search_for_tellurics.lambda_hi = line_list_to_search_for_tellurics.lambda.*calc_doppler_factor(Δv_to_avoid_tellurics)
 chunk_list_timeseries = RvSpectML.make_chunk_list_timeseries(expres_data,line_list_to_search_for_tellurics)
 line_list_to_search_for_tellurics.min_telluric_model_all_obs = RvSpectML.find_worst_telluric_in_each_chunk( chunk_list_timeseries, expres_data)
 line_list_no_tellurics_df = line_list_to_search_for_tellurics |> @filter(_.min_telluric_model_all_obs == 1.0) |> DataFrame
 #=
 Δv_to_avoid_tellurics = RvSpectMLBase.max_bc # default_Δv_to_avoid_tellurics
 line_list_no_tellurics_df.lambda_lo = line_list_no_tellurics_df.lambda./calc_doppler_factor(Δv_to_avoid_tellurics)
 line_list_no_tellurics_df.lambda_hi = line_list_no_tellurics_df.lambda.*calc_doppler_factor(Δv_to_avoid_tellurics)
 #find_overlapping_chunks(line_list_no_tellurics_df)
 #chunk_list_no_tellurics_merged_df = chunk_list_no_tellurics_merged_df |> @take(1) |> DataFrame
 #chunk_list_no_tellurics_merged_df = RvSpectML.merge_chunks(line_list_no_tellurics_df)
 #@assert find_overlapping_chunks(chunk_list_no_tellurics_merged_df) == nothing
 #chunk_list_timeseries = RvSpectML.make_chunk_list_timeseries(expres_data,chunk_list_no_tellurics_merged_df)
 chunk_list_timeseries = RvSpectML.make_chunk_list_timeseries(expres_data,line_list_no_tellurics_df)
 #chunk_list_timeseries = RvSpectML.make_chunk_list_timeseries(expres_data,chunk_list_no_tellurics_merged_df)
=#

 # Compute CCF's & measure RVs
if verbose println("# Computing CCF.")  end
 mask_shape = RvSpectML.CCF.TopHatCCFMask(order_list_timeseries.inst, scale_factor=tophap_ccf_mask_scale_factor)
 line_list = RvSpectML.CCF.BasicLineList(line_list_df.lambda, line_list_df.weight)
 line_list = RvSpectML.CCF.BasicLineList(line_list_no_tellurics_df.lambda, line_list_no_tellurics_df.weight)
 #line_list = RvSpectML.CCF.BasicLineList(chunk_list_no_tellurics_merged_df.lambda, chunk_list_no_tellurics_merged_df.weight)
 ccf_plan = RvSpectML.CCF.BasicCCFPlan(mask_shape = mask_shape, line_list=line_list, midpoint=ccf_mid_velocity)
 v_grid = RvSpectML.CCF.calc_ccf_v_grid(ccf_plan)
 @time ccfs = RvSpectML.CCF.calc_ccf_chunklist_timeseries(order_list_timeseries, ccf_plan)
 #@time ccfs = RvSpectML.CCF.calc_ccf_chunklist_timeseries(chunk_list_timeseries, ccf_plan)
 # Write CCFs to file
 if write_ccf_to_csv
    using CSV
    CSV.write(joinpath(output_dir,target_subdir * "_ccfs.csv"),Tables.table(ccfs',header=Symbol.(v_grid)))
 end
 has_computed_ccfs = true
end

#make_plots = true
if make_plots
   using Plots
   t_idx = 1:size(ccfs,2)
   using Plots
   #plot(v_grid,ccfs[:,t_idx]./maximum(ccfs[:,t_idx],dims=1),label=:none)
   plot(v_grid,ccfs[:,t_idx],label=:none)
   xlabel!("v (m/s)")
   ylabel!("CCF")
 end

if !has_computed_rvs
 if verbose println("# Measuring RVs from CCF.")  end
 rvs_ccf_gauss = RvSpectML.RVFromCCF.measure_rv_from_ccf(v_grid,ccfs,fit_type = "gaussian")
 # Store estimated RVs in metadata
 map(i->order_list_timeseries.metadata[i][:rv_est] = rvs_ccf_gauss[i]-mean(rvs_ccf_gauss), 1:length(order_list_timeseries) )
 if write_rvs_to_csv
    using CSV
    CSV.write(joinpath(output_dir,target_subdir * "_rvs_ccf.csv"),DataFrame("Time [MJD]"=>order_list_timeseries.times,"CCF RV [m/s]"=>rvs_ccf_gauss))
 end
 has_computed_rvs = true
end

if !has_computed_order_ccfs  # Compute order CCF's & measure RVs
tstart = now()    # Compute CCFs for each order
 # TODO:  Need to mask tellurics somehow
 RvSpectML.CCF.calc_order_ccf_chunklist_timeseries(order_list_timeseries, ccf_plan)
 #=
 orders_to_use = RvSpectML.orders_to_use_default(chunk_list_timeseries.inst)
 num_obs = length(chunk_list_timeseries)
 order_ccfs = Array{Float64,3}(undef,length(v_grid), length(orders_to_use), num_obs)
 for (i,ord) in enumerate(orders_to_use)
     chunk_list_order_timeseries = RvSpectML.extract_chunk_list_timeseries_for_order(order_list_timeseries, ord)
     if isnothing(chunk_list_order_timeseries)
        println("# Warning: No usable lines in order = ", ord)
        order_ccfs[:,i,:] .= 0.0
     else
        order_ccfs[:,i,:] .= RvSpectML.CCF.calc_ccf_chunklist_timeseries(order_list_order_timeseries, ccf_plan)

     end
 end
 =#
 println("# Order CCFs runtime: ", now()-tstart)

 if write_order_ccf_to_csv
    using CSV
    inst = EXPRES.EXPRES2D()
    for (i, order) in orders_to_use
      if !(sum(order_ccfs[:,i,:]) > 0)   continue    end
      local t = Tables.table( order_ccfs[:,i,:]', header=Symbol.(v_grid) )
      CSV.write(joinpath(output_dir,target_subdir * "_ccf_order=" * string(order) * ".csv"),t)
    end
 end
 has_computed_order_ccfs = true
end

if !has_computed_template
   if verbose println("# Making template spectra.")  end
   @time ( spectral_orders_matrix, f_mean, var_mean, deriv, deriv2, order_grids )  = RvSpectML.make_template_spectra(order_list_timeseries)

   if write_template_to_csv
      using CSV
      CSV.write(joinpath(output_dir,target_subdir * "_template.csv"),DataFrame("λ"=>spectral_orders_matrix.λ,"flux_template"=>f_mean,"var"=>var_mean, "dfluxdlnλ_template"=>deriv,"d²fluxdlnλ²_template"=>deriv2))
   end
   if write_spectral_grid_to_jld2
      using JLD2, FileIO
      save(joinpath(output_dir,target_subdir * "_matrix.jld2"), Dict("λ"=>spectral_orders_matrix.λ,"spectra"=>spectral_orders_matrix.flux,"var_spectra"=>spectral_orders_matrix.var,"flux_template"=>f_mean,"var"=>var_mean, "dfluxdlnλ_template"=>deriv,"d²fluxdlnλ²_template"=>deriv2))
   end
   has_computed_template = true
end

if !has_computed_dcpca
   if verbose println("# Performing Doppler constrained PCA analysis.")  end
   using MultivariateStats
   dcpca_out, M = RvSpectML.DCPCA.doppler_constrained_pca(spectral_orders_matrix.flux, deriv, rvs_ccf_gauss)
   frac_var_explained = 1.0.-cumsum(principalvars(M))./tvar(M)
   println("# Fraction of variance explained = ", frac_var_explained[1:min(5,length(frac_var_explained))])
   if write_dcpca_to_csv
      using CSV
      CSV.write(joinpath(output_dir,target_subdir * "_dcpca_basis.csv"),  Tables.table(M.proj) )
      CSV.write(joinpath(output_dir,target_subdir * "_dcpca_scores.csv"), Tables.table(dcpca_out) )
   end
   has_computed_dcpca = true
end

if make_plots  # Ploting results from DCPCA
  using Plots
  # Set parameters for plotting analysis
  plt_order = 42
  plt_order_pix = 4500:5000
  plt = scatter(frac_var_explained, xlabel="Number of PCs", ylabel="Frac Variance Unexplained")
  display(plt)
end

if make_plots
  plt_order = 13
  RvSpectML.plot_basis_vectors(order_grids, f_mean, deriv, M.proj, idx_plt = spectral_orders_matrix.chunk_map[plt_order], num_basis=3 )
  #xlims!(5761.5,5766)
end

if make_plots
  RvSpectML.plot_basis_scores(order_list_timeseries.times, rvs_ccf_gauss, dcpca_out, num_basis=3 )
end

if make_plots
  RvSpectML.plot_basis_scores_cor( rvs_ccf_gauss, dcpca_out, num_basis=3)
end

has_found_lines = false
write_lines_to_csv= false
if !has_found_lines
   if verbose println("# Performing fresh search for lines in template spectra.")  end
   cl = ChunkList(map(grid->ChuckOfSpectrum(spectral_orders_matrix.λ,f_mean, var_mean, grid), spectral_orders_matrix.chunk_map))
   spectral_orders_matrix = nothing # We're done with this, so can release memory
   GC.gc()
   lines_in_template = RvSpectML.LineFinder.find_lines_in_chunklist(cl)
   if verbose println("# Finding above lines in all spectra.")  end
   @time fits_to_lines = RvSpectML.LineFinder.fit_all_lines_in_chunklist_timeseries(order_list_timeseries, lines_in_template )

   if verbose println("# Rejecting lines that have telluric contamination at any time.")  end
   telluric_info = RvSpectML.LineFinder.find_worst_telluric_in_each_line_fit!(fits_to_lines, order_list_timeseries, expres_data )

   # Look at distribution of standard deviations for line properties
    fit_distrib = fits_to_lines |> @groupby(_.line_id) |>
            @map( { median_a=median(_.fit_a), median_b=median(_.fit_b), median_depth=median(_.fit_depth), median_σ²=median(_.fit_σ²), median_λc=median(_.fit_λc),
                   std_a=std(_.fit_a), std_b=std(_.fit_b), std_depth=std(_.fit_depth), std_σ²=std(_.fit_σ²), std_λc=std(_.fit_λc), min_telluric_model_all_obs=minimum(_.min_telluric_model_this_obs), line_id=_.line_id } ) |>
            @filter(_.min_telluric_model_all_obs == 1.0 ) |> # No telluric in any observations
            @filter(_.std_b < 0.257) |>     # ~90th quantile
            @filter( 0.001 < _.median_σ² ) |>  # ~99th quantile
            @filter( _.std_σ² < 0.00844) |>  # ~99th quantile
            @filter( _.median_depth > 0.05 ) |>  # ~99th quantile
            @filter( _.std_depth < 0.11 ) |>  # ~95th quantile
            #@filter( -0.25 < _.median_b < 0.25) |>  #~5th to 95th percentiles
            DataFrame # ~75th quantile
      good_lines = fit_distrib |> @filter(_.std_λc < 0.0316 ) |> DataFrame
      bad_lines = fit_distrib |> @filter(_.std_λc > 0.0316 ) |> DataFrame
      #scatter(good_lines.median_σ², good_lines.median_depth )
      #scatter!(bad_lines.median_σ², bad_lines.median_depth )
      lines_to_try = lines_in_template[first.(good_lines[!,:line_id]),:]

      if write_lines_to_csv
         using CSV
         CSV.write(joinpath(output_dir,target_subdir * "_linefinder_lines.csv"), lines_in_template )
         CSV.write(joinpath(output_dir,target_subdir * "_linefinder_line_fits.csv"), fits_to_lines )
         CSV.write(joinpath(output_dir,target_subdir * "_linefinder_line_fits_clean.csv"), lines_to_try )
         fits_to_lines = nothing
         telluric_info = nothing
         fit_distrib = nothing
         good_lines = nothing
         bad_lines = nothing
         GC.gc()
      end

   has_found_lines = true
end

if has_computed_ccfs2
   if verbose println("# Computing CCFs with new line list.")  end
   #mask_shape = RvSpectML.CCF.TopHatCCFMask(order_list_timeseries.inst, scale_factor=tophap_ccf_mask_scale_factor)
   perm = sortperm(lines_to_try.fit_λc)
   line_list2 = RvSpectML.CCF.BasicLineList(lines_to_try.fit_λc[perm], lines_to_try.fit_depth[perm].*lines_to_try.fit_a[perm] )
   ccf_plan2 = RvSpectML.CCF.BasicCCFPlan(mask_shape = mask_shape, line_list=line_list2, midpoint=0.0)
   v_grid2 = RvSpectML.CCF.calc_ccf_v_grid(ccf_plan2)
   @time ccfs2 = RvSpectML.CCF.calc_ccf_chunklist_timeseries(order_list_timeseries, ccf_plan2)
   # Write CCFs to file
   if write_ccf_to_csv
      using CSV
      CSV.write(joinpath(output_dir,target_subdir * "_ccfs2.csv"),Tables.table(ccfs2',header=Symbol.(v_grid)))
   end
   has_computed_ccfs2 = true
end

if !has_computed_rvs2
 if verbose println("# Measuring RVs from CCF.")  end
 #rvs_ccf_gauss = [ RvSpectML.RVFromCCF.measure_rv_from_ccf(v_grid,ccfs[:,i],fit_type = "gaussian") for i in 1:length(order_list_timeseries) ]
 rvs_ccf_gauss2 = RvSpectML.RVFromCCF.measure_rv_from_ccf(v_grid2,ccfs2,fit_type = "gaussian")
 # Store estimated RVs in metadata
 map(i->order_list_timeseries.metadata[i][:rv_est] = rvs_ccf_gauss2[i]-mean(rvs_ccf_gauss2), 1:length(order_list_timeseries) )
 if write_rvs_to_csv
    using CSV
    CSV.write(joinpath(output_dir,target_subdir * "_rvs_ccf2.csv"),DataFrame("Time [MJD]"=>order_list_timeseries.times,"CCF RV [m/s]"=>rvs_ccf_gauss))
 end
 has_computed_rvs2 = true
end

if !has_computed_template2
   # Do we need to expand region around lines to include more of wings for DCPCA to work well?
   # If so, should probably avoid tellurics encroaching on exapnded wavelength region.  Sigh.
   chunk_list_df2 = lines_to_try |> @select(:fit_min_λ,:fit_max_λ) |> @rename(:fit_min_λ=>:lambda_lo, :fit_max_λ=>:lambda_hi) |> DataFrame
   chunk_list_timeseries2 = RvSpectML.make_chunk_list_timeseries(expres_data,chunk_list_df2)

   # Check that no NaN's included
   (chunk_list_timeseries2, chunk_list_df2) = RvSpectML.filter_bad_chunks(chunk_list_timeseries2,chunk_list_df2)
   println(size(chunk_list_df2), " vs ", num_chunks(chunk_list_timeseries2) )

   if verbose println("# Making template spectra.")  end
   @time ( spectral_orders_matrix2, f_mean2, var_mean2, deriv_2, deriv2_2, order_grids2 )  = RvSpectML.make_template_spectra(chunk_list_timeseries2)

   if write_template_to_csv
      using CSV
      CSV.write(joinpath(output_dir,target_subdir * "_template2.csv"),DataFrame("λ"=>spectral_orders_matrix2.λ,"flux_template"=>f_mean2,"var"=>var_mean2, "dfluxdlnλ_template"=>deriv_2,"d²fluxdlnλ²_template"=>deriv2_2))
   end
   if write_spectral_grid_to_jld2
      using JLD2, FileIO
      save(joinpath(output_dir,target_subdir * "_matrix2.jld2"), Dict("λ"=>spectral_orders_matrix2.λ,"spectra"=>spectral_orders_matrix2.flux,"var_spectra"=>spectral_orders_matrix2.var,"flux_template"=>f_mean2,"var"=>var_mean2, "dfluxdlnλ_template"=>deriv_2,"d²fluxdlnλ²_template"=>deriv2_2))
   end
   has_computed_template2 = true
end

has_computed_dcpca2 = false
if !has_computed_dcpca2
   if verbose println("# Performing Doppler constrained PCA analysis.")  end
   using MultivariateStats
   dcpca2_out, M2 = RvSpectML.DCPCA.doppler_constrained_pca(spectral_orders_matrix2.flux, deriv_2, rvs_ccf_gauss2)
   frac_var_explained2 = 1.0.-cumsum(principalvars(M2))./tvar(M2)
   println("# Fraction of variance explained = ", frac_var_explained2[1:min(5,length(frac_var_explained2))])
   if write_dcpca_to_csv
      using CSV
      CSV.write(joinpath(output_dir,target_subdir * "_dcpca_basis2.csv"),  Tables.table(M.proj) )
      CSV.write(joinpath(output_dir,target_subdir * "_dcpca_scores2.csv"), Tables.table(dcpca_out) )
   end
   has_computed_dcpca2 = true
end

if make_plots  # Ploting results from DCPCA
  using Plots
  # Set parameters for plotting analysis
  plt_order = 42
  plt_order_pix = 4500:5000
  plt = scatter(frac_var_explained, xlabel="Number of PCs", ylabel="Frac Variance Unexplained")
  display(plt)
end

if make_plots
  plt_line = 170
  RvSpectML.plot_basis_vectors(order_grids, f_mean, deriv, M.proj, idx_plt = spectral_orders_matrix.chunk_map[plt_line], num_basis=3 )
  #xlims!(5761.5,5766)
end

if make_plots
  RvSpectML.plot_basis_scores(order_list_timeseries.times, rvs_ccf_gauss, dcpca_out, num_basis=3 )
end

if make_plots
  RvSpectML.plot_basis_scores_cor( rvs_ccf_gauss, dcpca_out, num_basis=3)
end

has_loaded_data = false
 has_computed_ccfs = false
 has_computed_rvs = false
 has_computed_template = false
 has_computed_dcpca = false
 has_found_lines = false


text:
function cooperative_astar!(
    solution::Solution,
    agents::AbstractVector{Int},
    mapf::MAPF,
    edge_weights::AbstractVector=mapf.edge_weights;
    conflict_price=Inf,
)
    reservation = compute_reservation(solution, mapf)
    @showprogress for a in agents
        s, d, t0 = mapf.sources[a], mapf.destinations[a], mapf.starting_times[a]
        dist = mapf.distances_to_destinations[d]
        heuristic(v) = dist[v]
        path = temporal_astar(
            mapf.graph,
            s,
            d,
            t0;
            edge_indices=mapf.edge_indices,
            edge_weights=edge_weights,
            heuristic=heuristic,
            reservation=reservation,
            conflict_price=conflict_price,
        )
        solution[a] = path
        update_reservation!(reservation, path, mapf)
    end
end

function cooperative_astar(
    mapf::MAPF,
    permutation::AbstractVector{Int}=1:nb_agents(mapf),
    edge_weights::AbstractVector=mapf.edge_weights;
    conflict_price=Inf,
)
    solution = [Path() for a in 1:nb_agents(mapf)]
    cooperative_astar!(
        solution, permutation, mapf, edge_weights; conflict_price=conflict_price
    )
    return solution
end

ext: jl


text:
using CurveFit
using Polynomials
using Statistics

abstract type AbstractCalibr end
abstract type AbstractCalibr1d end

struct CalibrCurve{Fit} <: AbstractCalibr1d
    "Anemometer calibration output"
    E::Vector{Float64}
    "Calibration Velocity"
    U::Vector{Float64}
    "Polynomial fitting the calibration curve"
    fit::Fit
    "Reference temperature of the calibration"
    T0::Float64
end
Base.broadcastable(cal::CalibrCurve) = Ref(cal)

"""
    cal = calibr_curve(sensor, V, E, temp, N)

Create a calibration for a thermal anemometer

Parameters:
 * `sensor` an anemometer object (<:AbstractThermalAnemometer)
 * `V` a vector with calibration velocities
 * `E` a vector with calibration voltages
 * `temp` a vector with calibration temperatures or a single calibration temperature
 * `N` degree of polynomial fit

This method will correct all voltages to the mean calibration temperature. Then it will fit the nonzero calibration velocities using a polynomial fit of degree N.

The speeds between zero and the first nonzero velocity are then interpolated using King's law.

# Examples
```jldoctest
julia> U = [0.0, 2.0, 3.0, 4.0, 6.0, 8.0, 10.0, 12.0, 18.0];

julia> E = [0.979, 1.333, 1.385, 1.423, 1.48, 1.521, 1.551, 1.575, 1.624];

julia> temp = [20.02, 19.75, 19.79, 19.73, 19.67, 19.75, 19.85, 19.85, 20.12] .+ 273.15;

julia> Rw = Thermistor(5e3, 3500, 25.0+273.15);

julia> cta = CTASensor(Rw, 100.0, 20/120);

julia> cal = CalibrCurve(cta, U, E, temp, 5)
CalibrCurve{CTASensor{Thermistor}}(CTASensor{Thermistor}(Thermistor(5000.0, 3500.0, 298.15), 100.0, 0.16666666666666666, 1.0), Polynomial(2526.816258773263 - 3916.1946161935334*x - 1547.7951632370953*x^2 + 6109.915344234665*x^3 - 3944.5331575193045*x^4 + 814.5086726564189*x^5), 19.836666666666666, (0.9897386382393427, 0.5558472575990494, 0.5), 1.332625512625792)

julia> cal.(E, temp) - U  # Check the calibration
9-element Vector{Float64}:
  0.0
  0.00033274230130109217
 -0.002686962091977829
  0.005428070340446567
 -0.007450571094204328
  0.006889739326171096
 -0.0019912770775185606
 -0.0007380705280706223
  0.00021632881207622745

julia>

```
"""
function calibr_curve(sensor::CTA, V, E, temp, N=5; extrapolate=true) where{CTA<:Hotwire.AbstractThermalAnemometer}

    nz = V .> 0  # Points with velocity above zero
    Tm = mean(temp)

    # Correct output to reference temperature
    Ec = tempcorr.(Ref(sensor), E, temp, Tm)
    Uc = [u for u in V]

    V1 = V[nz]
    E1 = Ec[nz]

    # Fit the nonzero velocities to a polynomial of degree N
    fit = curve_fit(Polynomial, E1, V1, N)

    if extrapolate
        # Now if the speed is lower than te minimum calibration speed,
        # attempt to extrapolate using King's law: E^2 = A + B + U^n
        Ea = minimum(Ec[.!nz])
        Eb = minimum(E1)
        efit = ExtrapolateFit(fit, Ea, Eb)
        return CalibrCurve(Ec, Uc, efit, Tm)
    else
        return CalibrCurve(Ec, Uc, fit, Tm)
    end

end

struct ExtrapolateFit{Fit}
    "A curve fit object obtained during calibration"
    fit::Fit
    "Minimum output WITH velocity"
    E0::Float64
    "King's law A coefficient"
    A::Float64
    "King's law B coefficient"
    B::Float64
    "King's law n coefficient"
    n::Float64
end

"""
    `efit = ExtrapolateFit(fit, Ea, Eb)`

For very low speeds, the behavior of thermal anemometers change
as natural convection begins to dominate the heat transfer from the probe.
Most calibrations reach has a nonzero minimum velocity. This function
extrapolates the calibration curve to zero.

It uses the calibration point with lowest speed and its derivative.
It also uses the output at 0 velocity as a data point.

With these information, a King's law is fitted. Notice that the zero
velocity point is unstable since any air current whatsoever will enhance
heat transfer.

The following curve will be used to approximate the low speeds:

``
E² = A + B⋅Uⁿ
``

Arguments:

 - `fit` curve fit obtained from calibration
 - `Ea`  Anemometer output for 0 velocity
 - `Eb` Anemometer output for minimum calibration velocity

**Returns**: An `ExtrapolateFit` object

"""
function ExtrapolateFit(fit::Fit, Ea, Eb) where {Fit}

    Ua = 0.0
    Ub = fit(Eb)
    f = eps(Eb) * 1e7
    Eb1 = Eb + f
    Eb0 = Eb - f
    Ub1 = fit(Eb1)
    Ub0 = fit(Eb0)
    dE² = ( (Eb1 - Eb0) / (Ub1 - Ub0) ) * (Eb1 + Eb0)
    A = Ea*Ea
    n = Ub * dE² / (Eb*Eb - A)

    B = (Eb*Eb-A)/Ub^n

    return ExtrapolateFit(fit, Eb, A, B, n)
end
Base.broadcastable(fit::ExtrapolateFit) = Ref(cal)

"""
    `fit(E)`

Uses an `ExtrapolateFit` object as a function call.

If E > fit.E0, use the original fit. Otherwise use the extrapolation.
"""
function (fit::ExtrapolateFit)(E)
    if E < fit.E0
        E2 = E*E
        if E2 < fit.A
            return zero(E)
        else
            return ((E2 - fit.A) / fit.B)^(1/fit.n)
        end
    else
        return fit.fit(E)
    end
end

"""
    `BlendFit(fit₁, fit₂, E₁, E₂)`

Blend two curve fits together.
The approximation `fit₁` is used for ``E < E₁`` and
`fit₂` is used for ``E > E₂``. For intermediary values,
a linear blending is used

"""
struct BlendFit{Fit₁,Fit₂}
    "Fit for lower speeds"
    fit₁::Fit₁
    "Fit for higher speeds"
    fit₂::Fit₂
    "Lower threshold"
    E₁::Float64
    "Higher threshold"
    E₂::Float64
end
Base.broadcastable(fit::BlendFit) = Ref(cal)

"""
    `fit(E)`

Calculates the blended fit
"""
function (fit::BlendFit)(E)
    if E < fit.E₁
        return fit.fit₁(E)
    elseif E > fit.E₂
        return fit.fit₂(E)
    else
        u₁ =  fit.fit₁(E)
        u₂ = fit.fit₂(E)
        ΔE = fit.E₂ - fit.E₁
        return u₁ * (fit.E₂-E) / ΔE + u₂*(E - fit.E₁)/ΔE
    end
end

struct HWKingFit
    A::Float64
    B::Float64
    n::Float64
end
Base.broadcastable(fit::HWKingFit) = Ref(cal)

function HWKingFit(E::AbstractVector, U::AbstractVector, ϵ=1e-8, maxiter=200)

    A, B, n = king_fit(E, U, ϵ, maxiter)

    return HWKingFit(A, B, n)
end

function (fit::HWKingFit)(E)

    E² = E*E
    if E² < fit.A
        return zero(E)
    else
        return (  (E²-fit.A) / fit.B )^(1/fit.n)
    end
end

"""
    `U = cal(sensor, E [,temp])`

Apply the calibration to obtain wind speed.

Parameters:
 * `sensor` Structure representing the anemometer
 * `E` Output voltage
 * `temp` fluid temperature. If not provided assume calibration temperature
 * Returns velocity

# Examples
```jldoctest
julia> cal(w, 1.5)
6.925627034695708

julia> cal(w, 1.5, 15)
5.902143445069172

julia> cal(w, 1.5, 20)
6.9642311432613395

julia> cal(w, 1.5, 25)
8.309847269677324

```

"""
function (cal::CalibrCurve)(sensor, E, temp)
    Ec = tempcorr(sensor, E, temp, cal.T0)

    return cal.fit(Ec)
end

(cal::CalibrCurve)(sensor, E) = cal(sensor, E, cal.T0)


text:
using MeshCat, MeshCatMechanisms, RigidBodyDynamics
using FileIO, MeshIO, GeometryBasics, CoordinateTransformations, Rotations, Meshing
using Plots

include_ddp()

contact_control_path = "/home/taylor/Research/ContactControl.jl/src"

using Parameters
# Utilities
include(joinpath(contact_control_path, "utils.jl"))

# Solver
include(joinpath(contact_control_path, "solver/cones.jl"))
include(joinpath(contact_control_path, "solver/interior_point.jl"))
include(joinpath(contact_control_path, "solver/lu.jl"))

# Environment
include(joinpath(contact_control_path, "simulator/environment.jl"))

# Dynamics
include(joinpath(contact_control_path, "dynamics/model.jl"))

# Simulator
include(joinpath(contact_control_path, "simulation/contact_methods.jl"))
include(joinpath(contact_control_path, "simulation/simulation.jl"))
include(joinpath(contact_control_path, "simulator/trajectory.jl"))

include(joinpath(contact_control_path, "dynamics/code_gen_dynamics.jl"))
include(joinpath(contact_control_path, "dynamics/fast_methods_dynamics.jl"))

# Models
include(joinpath(contact_control_path, "dynamics/quaternions.jl"))
include(joinpath(contact_control_path, "dynamics/mrp.jl"))
include(joinpath(contact_control_path, "dynamics/euler.jl"))

# Simulation
include(joinpath(contact_control_path, "simulation/environments/flat.jl"))

include(joinpath(contact_control_path, "simulation/residual_approx.jl"))
include(joinpath(contact_control_path, "simulation/code_gen_simulation.jl"))

# Visuals
using MeshCatMechanisms
include(joinpath(contact_control_path, "dynamics/visuals.jl"))
include(joinpath(contact_control_path, "dynamics/visual_utils.jl"))

include(joinpath(pwd(),"models/kuka/kuka_utils.jl"))

urdf_path = joinpath(pwd(), "models/kuka/temp/kuka.urdf")

kuka = MeshCatMechanisms.parse_urdf(urdf_path, remove_fixed_tree_joints = true)
kuka_visuals = MeshCatMechanisms.URDFVisuals(urdf_path)

state = MechanismState(kuka)
state_cache = StateCache(kuka)
result = DynamicsResult(kuka)
result_cache = DynamicsResultCache(kuka)

vis = Visualizer()
mvis = MechanismVisualizer(kuka, kuka_visuals, vis[:base])
render(vis)
default_background!(vis)

nq = num_positions(kuka)
q_init = zeros(nq)
q_init[1] = 0
q_init[2] = pi/4

q_init[4] = -pi/2
q_init[5] = 0
q_init[6] = -pi/4
q_init[7] = 0
set_configuration!(state, q_init)
set_configuration!(mvis, q_init)
ee = findbody(kuka, "iiwa_link_7")
ee_point = Point3D(default_frame(ee), 0.0, 0.0, 0.0)
setelement!(mvis, ee_point, 0.05)
ee_jacobian_frame = ee_point.frame
ee_jacobian_path = path(kuka, root_body(kuka), ee)

world = root_frame(kuka)
ee_in_world = transform(state, ee_point, world).v
desired = Point3D(world, 0.5, 0.0, 0.0)
q_res1 = jacobian_transpose_ik!(state, ee, ee_point, desired,
    visualize = true, mvis = mvis)
q_res1 = Array(q_res1)
set_configuration!(mvis, q_res1)
set_configuration!(state, q_res1)
ee_in_world = transform(state, ee_point, world).v

desired = Point3D(world, 0.75, 0.0, 0.0)
q_res2 = jacobian_transpose_ik!(state, ee, ee_point, desired,
    visualize = true, mvis = mvis)
q_res2 = Array(q_res2)
set_configuration!(mvis, q_res2)
set_configuration!(state, q_res2)
ee_in_world = transform(state, ee_point, world).v

# Kuka iiwa arm parsed from URDF using RigidBodyDynamics.jl
# + 3D particle
struct KukaParticle <: ContactModel
	dim::Dimensions

	# particle
	mp

    state_cache1
    state_cache2
    state_cache3

    result_cache1
    result_cache2
    result_cache3

    world

	ee
    ee_point
    ee_jacobian_frame
    ee_jacobian_path
end

# Dimensions
nq = 7 + 3 # configuration dim
nu = 7 # control dim
nc = 1 # number of contact points
nf = 0 # number of faces for friction cone
nb = 0
ns = 0

n = 2 * nq
m_contact = 0 #nc + nb + nc + nb + 1
m = nu #nu + m_contact
d = 0

mp = 0.01

results_cache1 = DynamicsResultCache(kuka)
results_cache2 = DynamicsResultCache(kuka)
results_cache3 = DynamicsResultCache(kuka)

state_cache1 = StateCache(kuka)
state_cache2 = StateCache(kuka)
state_cache3 = StateCache(kuka)

function kuka_q(q)
	view(q, 1:7)
end

function particle_q(q)
	view(q, 8:10)
end

# Methods
function M_func(m::KukaParticle, q::AbstractVector{T}) where T
    # return Diagonal(m.mp * ones(3))
    state = m.state_cache3[T]
    result = m.result_cache3[T]
    set_configuration!(state, kuka_q(q))
    mass_matrix!(result.massmatrix, state)
    return SMatrix{10, 10}(cat(result.massmatrix, Diagonal(m.mp * ones(3)), dims = (1, 2)))
end

function C_func(m::KukaParticle, q, q̇)
    # return SVector{3}([0.0, 0.0, m.mp * 9.81])
    if eltype(q) <: ForwardDiff.Dual
		return SVector{10}([dynamics_bias_q(m, kuka_q(q), kuka_q(q̇))..., 0.0, 0.0, m.mp * 9.81])
	elseif eltype(q̇) <: ForwardDiff.Dual
		return SVector{10}([dynamics_bias_q̇(m, kuka_q(q), kuka_q(q̇))..., 0.0 ,0.0, m.mp * 9.81])
	else
		return SVector{10}([_dynamics_bias(m, kuka_q(q), kuka_q(q̇))..., 0.0, 0.0, m.mp * 9.81])
	end
end

function _dynamics_bias(m::KukaParticle, q, q̇)
	T = eltype(m.mp)
	state = m.state_cache1[T]
	result = m.result_cache1[T]
	set_configuration!(state, q)
	set_velocity!(state, q̇)
	dynamics_bias!(result, state)
    return result.dynamicsbias
end

function dynamics_bias_q(m::KukaParticle, q::AbstractVector{T}, q̇) where T
    state = m.state_cache2[T]
    result = m.result_cache2[T]
    set_configuration!(state, q)
    set_velocity!(state,q̇)
    dynamics_bias!(result, state)
    return result.dynamicsbias
end

function dynamics_bias_q̇(m::KukaParticle, q, q̇::AbstractVector{T}) where T
    state = m.state_cache3[T]
    result = m.result_cache3[T]
    set_configuration!(state, q)
    set_velocity!(state, q̇)
    dynamics_bias!(result, state)
    return result.dynamicsbias
end

function B_func(m::KukaParticle, q)
    # return zeros(3, 3)
    SMatrix{10, 7}([Diagonal(ones(7)); zeros(3, 7)])
end

function kinematics_ee(m::KukaParticle, q::AbstractVector{T}) where T
    state = m.state_cache3[T]
    set_configuration!(state, kuka_q(q))
    return transform(state, m.ee_point, m.world).v
end

function ϕ_func(m::KukaParticle, q::AbstractVector{T}) where T
	p_ee = kinematics_ee(m, q)
	p_p = q[8:10]
	diff = (p_ee - p_p)
	d_ee_p = sqrt(diff' * diff)
    SVector{1}([0.5 - d_ee_p])
end

function P_func(m::KukaParticle, q::AbstractVector{T}) where T
    state = m.state_cache3[T]

    set_configuration!(state, kuka_q(q))

    pj1 = PointJacobian(transform(state, m.ee_point, m.world).frame,
		zeros(T, 3, 7))

    ee_in_world = transform(state, m.ee_point, m.world)

    point_jacobian!(pj1, state, m.ee_jacobian_path,ee_in_world) #TODO confirm this is correct

	p_ee = ee_in_world.v
	p_p = q[8:10]
	diff = (p_ee - p_p)

    return Array(-1.0 * 2.0 * diff' * [pj1.J[1:3,:] -1.0 * Diagonal(ones(3))])
end

ϕ(z) = ϕ_func(_model, z)
norm(ForwardDiff.jacobian(ϕ, q0) - P_func(_model, q1))

nq = 7 + 3
nu = 7
nw = 0
nc = 1

h = 0.1

_model = KukaParticle(
	Dimensions(nq, nu, nw, nc),
	mp,
    state_cache1, state_cache2, state_cache3,
    results_cache1, results_cache2, results_cache3,
    world,
	ee,
	ee_point,
	ee_jacobian_frame,
	ee_jacobian_path)

q0 = zeros(_model.dim.q)
q0[1] = 0
q0[3] = 0
q0[4] = -pi/2
q0[5] = 0.0
p_pos0 = Array(kinematics_ee(_model, q0))
p_pos0[1] += 0.0
p_pos0[3] -= 0.5
q0[8:10] = p_pos0

# visualize!(mvis, _model, [q0], Δt = h)

q1 = copy(q0)

qN = zeros(_model.dim.q)
qN[1] = 0
qN[3] = 0
qN[4] = -pi/2
qN[5] = 0.

ee_posN = Array(kinematics_ee(_model, qN))
p_posN = ee_posN
p_posN[1] += 0.0
p_posN[3] += 0.1
qN[8:10] = p_posN

qD = zeros(_model.dim.q)
qD[1] = 0
qD[3] = 0
qD[4] = -pi/2
qD[5] = 0.

ee_posD = Array(kinematics_ee(_model, qD))
p_posD = ee_posD
# p_posD[1] += 0.5
qD[8:10] = p_posD

visualize!(mvis, _model, [qD], Δt = h)

# Visualization
function visualize!(mvis, model::KukaParticle, q;
		verbose = false, r_ball = 0.035, Δt = 0.1)

	setobject!(vis["ball"], Sphere(Point3f0(0),
				convert(Float32, r_ball)),
				MeshPhongMaterial(color = RGBA(1.0, 0.0, 0.0,1.0)))

	settransform!(vis["ball"], compose(Translation(0.66,3.0,0.0)))

	state = model.state_cache1[Float64]

	anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt)))

	T = length(q)
    for t = 1:T
        q_kuka = kuka_q(q[t])
		q_particle = particle_q(q[t])
		set_configuration!(state, kuka_q(q[t]))

        MeshCat.atframe(anim,t) do
			set_configuration!(mvis,q_kuka)
            settransform!(vis["ball"], compose(Translation(q_particle), LinearMap(RotZ(0))))
		end
    end

    MeshCat.setanimation!(vis, anim)
end

include(joinpath(contact_control_path, "simulator/environment.jl"))
include(joinpath(contact_control_path, "simulation/environments/flat.jl"))

env = flat_3D_lc

function lagrangian_derivatives(model::KukaParticle, q, v)
	D1L = -1.0 * C_func(model, q, v)
    D2L = M_func(model, q) * v
	return D1L, D2L
end

function gravity_compensation(model::KukaParticle, h, q0, q1)
    q2 = q1
	# evalutate at midpoint
	qm1 = 0.5 * (q0 + q1)
    vm1 = (q1 - q0) / h[1]
    qm2 = 0.5 * (q1 + q2)
    vm2 = (q2 - q1) / h[1]

	D1L1, D2L1 = lagrangian_derivatives(model, qm1, vm1)
	D1L2, D2L2 = lagrangian_derivatives(model, qm2, vm2)

	# return 0.0
	return (0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2)[1:7]
end

function dynamics(model::KukaParticle, h, q0, q1, u1, λ1, q2)
	# evalutate at midpoint
	qm1 = 0.5 * (q0 + q1)
    vm1 = (q1 - q0) / h[1]
    qm2 = 0.5 * (q1 + q2)
    vm2 = (q2 - q1) / h[1]

	D1L1, D2L1 = lagrangian_derivatives(model, qm1, vm1)
	D1L2, D2L2 = lagrangian_derivatives(model, qm2, vm2)

    u_gravity = gravity_compensation(model, h, q0, q1)

	# return 0.0
	return (0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2
		+ B_func(model, qm2) * (u1 - u_gravity)
        + transpose(P_func(model, q2)) * λ1)
        # - h[1] * 0.5 .* vm2)
end

function residual(model::KukaParticle, env::Environment{<:World,LinearizedCone}, z, θ, κ)
	# nc = model.dim.c
	# nb = nc * friction_dim(env)
	# nf = Int(nb / nc)
	# np = dim(env)

    q0 = θ[1:nq]
    q1 = θ[nq .+ (1:nq)]
    u1 = θ[2nq .+ (1:nu)]
    h = θ[2nq + nu .+ (1:1)]

    q2 = z[1:nq]
    λ1 = z[nq .+ (1:nc)]
    s1 = z[nq + nc .+ (1:nc)]

	# q0, q1, u1, w1, μ, h = unpack_θ(model, θ)
	# q2, γ1, b1, ψ1, η1, s1, s2 = unpack_z(model, env, z)

	# ϕ = ϕ_func(model, env, q2)
    #
	# k = kinematics(model, q2)
	# λ1 = contact_forces(model, env, γ1, b1, q2, k)
	# Λ1 = transpose(J_func(model, env, q2)) * λ1 #@@@@ maybe need to use J_fast
	# vT_stack = velocity_stack(model, env, q1, q2, k, h)
	# ψ_stack = transpose(E_func(model, env)) * ψ1
    # return dynamics(model, h, q0, q1, u1, zeros(nc), q2)

    [
     dynamics(model, h, q0, q1, u1, λ1, q2);
     s1 .- ϕ_func(model, q2);
     λ1 .* s1 .- κ;
    ]
	# [
    #  dynamics(model, h, q0, q1, u1, q2);
	#  # s1 - ϕ;
	#  # vT_stack + ψ_stack - η1;
	#  # s2 .- (μ[1] * γ1 .- E_func(model, env) * b1);
	#  # γ1 .* s1 .- κ;
	#  # b1 .* η1 .- κ;
	#  # ψ1 .* s2 .- κ
    #  ]
end

nz = nq + nc + nc
nθ = nq + nq + nu + 1

u0 = zeros(_model.dim.u)
z0 = copy([q1; 0.1 * ones(2 * nc)])
θ0 = copy([q0; q1; u0; h])

residual(_model, env, z0, θ0, [1.0])

function r_func(r, z, θ, κ)
    r .= residual(_model, env, z, θ, κ)
end

function rz_func(rz, z, θ)
    r(a) = residual(_model, env, a, θ, 0.0)
    rz .= ForwardDiff.jacobian(r, z)
end

function rθ_func(rθ, z, θ)
    r(a) = residual(_model, env, z, a, 0.0)
    rθ .= ForwardDiff.jacobian(r, θ)
end

r0 = zeros(nz)
rz0 = zeros(nz, nz)
rθ0 = zeros(nz, nθ)
r_func(r0, z0, θ0, [1.0])
rz_func(rz0, z0, θ0)
rθ_func(rθ0, z0, θ0)

# options
opts = InteriorPointOptions(
   κ_init = 0.1,
   κ_tol = 1.0e-4,
   r_tol = 1.0e-8,
   diff_sol = true)

idx_ineq = collect([nq .+ (1:(nc + nc))]...)

# solver
ip = interior_point(z0, θ0,
   r! = r_func, rz! = rz_func,
   rz = similar(rz0, Float64),
   rθ! = rθ_func,
   rθ = similar(rθ0, Float64),
   idx_ineq = idx_ineq,
   opts = opts)

# simulate
T = 26
q_hist = [q0, q0]
u_hist = []
h = 0.1
for t = 1:T-1
    u_gravity = zeros(_model.dim.u)#-1.0 * gravity_compensation(_model, h, q_hist[end-1], q_hist[end])
    push!(u_hist, 0.0 * u_gravity)
    # if t == 5
    #     u_gravity[6] = -1.0
    # end
    # u_gravity = -1.0 * gravity_compensation(_model, h, q0, q1)

    ip.z .= copy([q_hist[end]; 0.1 * ones(nc + nc)])
    ip.θ .= copy([q_hist[end-1]; q_hist[end];
        # u0;
        u_hist[end];
        h])
    status = interior_point_solve!(ip)

    if status
        push!(q_hist, ip.z[1:nq])
    else
        println("dynamics failure t = $t")
        println("res norm = $(norm(ip.r, Inf))")
        break
    end
end

vz_particle = [(((q_hist[t+2] - q_hist[t+1]) / h - (q_hist[t+1] - q_hist[t]) / h) / h)[10] for t = 1:T-2]

visualize!(mvis, _model, q_hist, Δt = h)

s = Simulation(_model, env)
s.res.r! = r_func
s.res.rz! = rz_func
s.res.rθ! = rθ_func
s.rz = ip.rz
s.rθ = ip.rθ

struct Dynamics{T}
	s::Simulation
	ip_dyn::InteriorPoint
	ip_jac::InteriorPoint
	h::T
end

function gen_dynamics(s::Simulation, h;
		dyn_opts =  InteriorPointOptions{Float64}(
						r_tol = 1.0e-8,
						κ_tol = 1.0e-4,
						κ_init = 0.1,
						diff_sol = true),
		jac_opts =  InteriorPointOptions{Float64}(
						r_tol = 1.0e-8,
						κ_tol = 1.0e-4,
						κ_init = 0.1,
						diff_sol = true))

	z = zeros(nz)
	θ = zeros(nθ)

	ip_dyn = interior_point(z, θ,
		idx_ineq = idx_ineq,
		r! = s.res.r!,
		rz! = s.res.rz!,
		rθ! = s.res.rθ!,
		rz = s.rz,
		rθ = s.rθ,
		opts = dyn_opts)

	ip_dyn.opts.diff_sol = false

	ip_jac = interior_point(z, θ,
		idx_ineq = idx_ineq,
		r! = s.res.r!,
		rz! = s.res.rz!,
		rθ! = s.res.rθ!,
		rz = s.rz,
		rθ = s.rθ,
		opts = jac_opts)

	ip_jac.opts.diff_sol = true

	Dynamics(s, ip_dyn, ip_jac, h)
end

d = gen_dynamics(s, h,
	dyn_opts = InteriorPointOptions{Float64}(κ_tol = 1.0e-4, κ_init = 0.1),
	jac_opts = InteriorPointOptions{Float64}(κ_tol = 1.0e-2, κ_init = 0.1))

function f!(d::Dynamics, q0, q1, u1, mode = :dynamics)
	s = d.s
	ip = (mode == :dynamics ? d.ip_dyn : d.ip_jac)
	h = d.h

    ip.z .= copy([q1; 0.1 * ones(nc + nc)])
	ip.θ .= copy([q0; q1; u1; h])

	status = interior_point_solve!(ip)

	!status && (@warn "dynamics failure")
end

function f(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :dynamics)
	return copy(d.ip_dyn.z[1:d.s.model.dim.q])
end

f(d, q0, q1, zeros(m))

function fq0(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :jacobian)
	return copy(d.ip_jac.δz[1:d.s.model.dim.q, 1:d.s.model.dim.q])
end

fq0(d, q0, q1, zeros(m))

function fq1(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :jacobian)
	return copy(d.ip_jac.δz[1:d.s.model.dim.q, d.s.model.dim.q .+ (1:d.s.model.dim.q)])
end

fq1(d, q0, q1, zeros(m))

function fx1(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :jacobian)
	return copy(d.ip_jac.δz[1:d.s.model.dim.q, 1:(2 * d.s.model.dim.q)])
end

fx1(d, q0, q1, zeros(m))

function fu1(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :jacobian)
	return copy(d.ip_jac.δz[1:d.s.model.dim.q, 2 * d.s.model.dim.q .+ (1:d.s.model.dim.u)])
end

fu1(d, q0, q1, zeros(m))

struct KukaParticleI{I, T} <: Model{I, T}
    n::Int
    m::Int
    d::Int
	dynamics::Dynamics
end

model = KukaParticleI{Midpoint, FixedTime}(2 * s.model.dim.q, s.model.dim.u, 0, d)

function fd(model::KukaParticleI{Midpoint, FixedTime}, x, u, w, h, t)
	nq = model.dynamics.s.model.dim.q
	q0 = x[1:nq]
	q1 = x[nq .+ (1:nq)]

	q2 = f(model.dynamics, q0, q1, u)

	return [q1; q2]
end

fd(model, [q0; q1], zeros(model.m), zeros(model.d), h, 1)

function fdx(model::KukaParticleI{Midpoint, FixedTime}, x, u, w, h, t)
	nq = model.dynamics.s.model.dim.q
	q0 = x[1:nq]
	q1 = x[nq .+ (1:nq)]
	dq2dx1 = fx1(model.dynamics, q0, q1, u)

	return [zeros(nq, nq) I; dq2dx1]
end

fdx(model, [q0; q1], zeros(model.m), zeros(model.d), h, 1)

function fdu(model::KukaParticleI{Midpoint, FixedTime}, x, u, w, h, t)
	nq = model.dynamics.s.model.dim.q
	q0 = x[1:nq]
	q1 = x[nq .+ (1:nq)]
	dq2du1 = fu1(model.dynamics, q0, q1, u)
	return [zeros(nq, model.m); dq2du1]
end

fdu(model, [q0; q1], zeros(model.m), zeros(model.d), h, 1)

T = 21
x1 = [q1; q1]
xT = [qN; qN]
ū = [0.1 * randn(_model.dim.u) for t = 1:T-1]#[t == 1 ? [0.0, 10.0, 0.0, 0.0, 0.0, -0.0, 0.0] : t == 2 ? [0.0, -10.0, 0.0, -0.0, 0.0, 0.0, 0.0] : [0.0, 0.0, 0.0, -0.0, 0.0, 0.0, 0.0] for t = 1:T-1] #u_hist #[-1.0 * gravity_compensation(_model, h, q1, q1) for t = 1:T-1]
w = [zeros(model.d) for t = 1:T-1]

# Rollout
x̄ = rollout(model, x1, ū, w, h, T)
q̄ = state_to_configuration(x̄)
visualize!(mvis, _model, q̄, Δt = h)

# visualize!(mvis, _model, [x_track_extend[1:nq]], Δt = h)

# Objective
V = Diagonal(ones(s.model.dim.q))
Q_velocity = 1.0e-3 * [V -V; -V V] ./ h^2.0
Q_track = 1.0 * Diagonal(ones(model.n))
Q_track_extend = copy(Q_track)
Q_track_extend[8, 8] = 10.0
Q_track_extend[nq + 8, nq + 8] = 10.0
x_track = [qD; qD]
x_track_extend = copy(x_track)
x_track_extend[8] += 0.5
x_track_extend[nq + 8] += 0.5
Q = [t < T ? Q_velocity + (t == 11 ? Q_track_extend : 1.0 * Q_track) : Q_velocity + 1.0 * Q_track for t = 1:T]
q = [t < T ? -2.0 * (t == 11 ? Q_track_extend * x_track_extend : 1.0 * Q_track * x_track) : -2.0 * 1.0 * Q_track * x_track for t = 1:T]
R = [Diagonal(1.0e-3 * ones(model.m)) for t = 1:T-1]
r = [zeros(model.m) for t = 1:T-1]

obj = StageCosts([QuadraticCost(Q[t], q[t],
	t < T ? R[t] : nothing, t < T ? r[t] : nothing) for t = 1:T], T)

function g(obj::StageCosts, x, u, t)
	T = obj.T
    if t < T
		Q = obj.cost[t].Q
		q = obj.cost[t].q
	    R = obj.cost[t].R
		r = obj.cost[t].r
        return x' * Q * x + q' * x + u' * R * u + r' * u
    elseif t == T
		Q = obj.cost[T].Q
		q = obj.cost[T].q
        return x' * Q * x + q' * x
    else
        return 0.0
    end
end

# Constraints
ul = -100.0 * ones(model.m)
uu = 100.0 * ones(model.m)
p = [t < T ? (t == 11 ? 2 * m + 0 * nq : 2 * m) : n for t = 1:T]
info_t = Dict(:ul => ul, :uu => uu, :inequality => (1:2 * m))
info_m = Dict(:qD => x_track_extend[1:nq], :ul => ul, :uu => uu, :inequality => (1:2 * m))
info_T = Dict(:xT => xT)
con_set = [StageConstraint(p[t], t < T ? (t == 11 ? info_m : info_t) : info_T) for t = 1:T]

function c!(c, cons::StageConstraints, x, u, t)
	T = cons.T
	p = cons.con[t].p

	if t < T
		ul = cons.con[t].info[:ul]
		uu = cons.con[t].info[:uu]
		c[1:2 * m] .= [ul - u; u - uu]
        # if t == 11
        #     c[2 * m .+ (1:nq)] = x[1:nq] - cons.con[t].info[:qD]
        # end
	else
		c[1:n] .= x - cons.con[T].info[:xT]
	end
end

prob = problem_data(model, obj, con_set, copy(x̄), copy(ū), w, h, T,
	analytical_dynamics_derivatives = true)

ū = [t == 1 ? [0.0, 6.5, 0.0, 0.0, 0.0, -0.0, 0.0] : t == 2 ? [0.0, -6.5, 0.0, -0.0, 0.0, 0.0, 0.0] : [0.0, 0.0, 0.0, -0.0, 0.0, 0.0, 0.0] for t = 1:T-1] #u_hist #[-1.0 * gravity_compensation(_model, h, q1, q1) for t = 1:T-1]
w = [zeros(model.d) for t = 1:T-1]

# Rollout
x̄ = rollout(model, x1, ū, w, h, T)
q̄ = state_to_configuration(x̄)
visualize!(mvis, _model, q̄, Δt = h)
open(mvis)
# Solve
@time constrained_ddp_solve!(prob,
	max_iter = 1000, max_al_iter = 10,
	ρ_init = 1.0, ρ_scale = 10.0,
	con_tol = 0.005)

x, u = current_trajectory(prob)
x̄, ū = nominal_trajectory(prob)

q̄ = state_to_configuration(x̄)
visualize!(mvis, _model, q̄, Δt = h)

max_issues_repo_name: thowell/DirectMotionPlanning


text:
using ValueOrientedRiskManagementInsurance

using Distributions
using DataFrames
include("SSTLife_Input.jl")
include("SSTLife.jl")

println("Start SSTLife test")

@test_approx_eq(round(rtk_start, 2), 158.58)
@test_approx_eq(round(tc, 2), 147.60)
@test_approx_eq(round(sst_ratio, 4), 1.0744)

println("End SSTLife test")


text:
"""
    overwrite_netcdf2d(fname,var,varname)

Overwrite variable `varname` with field `var` in netcdf file `fname`
"""
function overwrite_netcdf2d(fname,var,varname)
ncin= NetCDF.open(fname)
ncwrite(var,fname,varname)
ncclose(fname)
end


text:
using DataStructures: nlargest

struct GlobalPool{A}
    aggr::Symbol
    cluster::A
    function GlobalPool(aggr::Symbol, dim...)
        cluster = ones(Int64, dim)
        new{typeof(cluster)}(aggr, cluster)
    end
end

(g::GlobalPool)(X::AbstractArray) = pool(g.aggr, g.cluster, X)

struct LocalPool{A}
    aggr::Symbol
    cluster::A
end

function LocalPool(aggr::Symbol, cluster::AbstractArray)
    LocalPool{typeof(cluster)}(aggr, cluster)
end

(l::LocalPool)(X::AbstractArray) = pool(l.aggr, l.cluster, X)

struct TopKPool{T,S}
    A::AbstractMatrix{T}
    k::Int
    p::AbstractVector{S}
    Ã::AbstractMatrix{T}
end

function TopKPool(adj::AbstractMatrix, k::Int, in_channel::Integer; init=glorot_uniform)
    TopKPool(adj, k, init(in_channel), similar(adj, k, k))
end

function (t::TopKPool)(X::AbstractArray)
    y = t.p' * X / norm(t.p)
    idx = topk_index(y, t.k)
    t.Ã .= view(t.A, idx, idx)
    X_ = view(X, :, idx) .* σ.(view(y, idx)')
    return X_
end

function sumpool(cluster::AbstractArray{Int}, X::AbstractArray{T},
                 c::Integer=length(Set(cluster))) where {T<:Real}
    dims = Dims(cluster, X)
    Y = zeros(T, dims.us_dims[1], c)
    scatter_add!(Y, X, cluster)
    Y
end

function subpool(cluster::AbstractArray{Int}, X::AbstractArray{T},
                 c::Integer=length(Set(cluster))) where {T<:Real}
    dims = Dims(cluster, X)
    Y = zeros(T, dims.us_dims[1], c)
    scatter_sub!(Y, X, cluster)
    Y
end

function prodpool(cluster::AbstractArray{Int}, X::AbstractArray{T},
                  c::Integer=length(Set(cluster))) where {T<:Real}
    dims = Dims(cluster, X)
    Y = ones(T, dims.us_dims[1], c)
    scatter_mul!(Y, X, cluster)
    Y
end

function divpool(cluster::AbstractArray{Int}, X::AbstractArray{T},
                 c::Integer=length(Set(cluster))) where {T<:Real}
    dims = Dims(cluster, X)
    FT = floattype(T)
    Y = ones(FT, dims.us_dims[1], c)
    scatter_div!(Y, FT.(X), cluster)
    Y
end

function maxpool(cluster::AbstractArray{Int}, X::AbstractArray{T},
                 c::Integer=length(Set(cluster))) where {T<:Real}
    dims = Dims(cluster, X)
    Y = fill(typemin(T), dims.us_dims[1], c)
    scatter_max!(Y, X, cluster)
    Y
end

function minpool(cluster::AbstractArray{Int}, X::AbstractArray{T},
                 c::Integer=length(Set(cluster))) where {T<:Real}
    dims = Dims(cluster, X)
    Y = fill(typemax(T), dims.us_dims[1], c)
    scatter_min!(Y, X, cluster)
    Y
end

function meanpool(cluster::AbstractArray{Int}, X::AbstractArray{T},
                  c::Integer=length(Set(cluster))) where {T<:Real}
    dims = Dims(cluster, X)
    FT = floattype(T)
    Y = zeros(FT, dims.us_dims[1], c)
    scatter_mean!(Y, FT.(X), cluster)
    Y
end

struct Dims
    xs_dims::Tuple
    us_dims::Tuple

    function Dims(xs_dims::Tuple, us_dims::Tuple)
        @assert xs_dims == us_dims[2:end] "X must have the same latter dimension with cluster."
        new(xs_dims, us_dims)
    end
end

Dims(xs::AbstractArray{Int}, us::AbstractArray) = Dims(size(xs), size(us))

@adjoint sumpool(cluster::AbstractArray{Int}, X::AbstractArray{T}) where {T<:Real} =
    sumpool(cluster, X), Δ -> (nothing, gather(zero(Δ)+Δ, cluster))
@adjoint subpool(cluster::AbstractArray{Int}, X::AbstractArray{T}) where {T<:Real} =
    subpool(cluster, X), Δ -> (nothing, -gather(zero(Δ)+Δ, cluster))

@adjoint function prodpool(cluster::Array{Int}, X::Array{T}) where {T<:Real}
    prodpool(cluster, X), function (Δ)
        rev_cluster = gather_indices(cluster)
        ∇X = gather(zero(Δ)+Δ, cluster)
        @inbounds for ind = CartesianIndices(cluster)
            inds = filter(x -> x != ind, rev_cluster[cluster[ind]])
            for i = 1:size(X, 1)
                ∇X[i, ind] *= prod(j -> X[i, j], inds)
            end
        end
        (nothing, ∇X)
    end
end

@adjoint function divpool(cluster::Array{Int}, X::Array{T}) where {T<:Real}
    divpool(cluster, X), function (Δ)
        rev_cluster = gather_indices(cluster)
        ∇X = -gather(zero(Δ)+Δ, cluster) ./ X.^2
        @inbounds for ind = CartesianIndices(cluster)
            inds = filter(x -> x != ind, rev_cluster[cluster[ind]])
            for i = 1:size(X, 1)
                ∇X[i, ind] /= prod(j -> X[i, j], inds)
            end
        end
        (nothing, ∇X)
    end
end

@adjoint function maxpool(cluster::Array{Int}, X::Array{T}) where {T<:Real}
    max = maxpool(cluster, X)
    max, function (Δ)
       Δu = (X .== gather(max, cluster)) .* gather(zero(Δ)+Δ, cluster)
       (nothing, Δu)
    end
end

@adjoint function minpool(cluster::Array{Int}, X::Array{T}) where {T<:Real}
    min = minpool(cluster, X)
    min, function (Δ)
       Δu = (X .== gather(min, cluster)) .* gather(zero(Δ)+Δ, cluster)
       (nothing, Δu)
    end
end

@adjoint function meanpool(cluster::AbstractArray{Int}, X::AbstractArray{T}) where {T<:Real}
    m = meanpool(cluster, X)
    m, function (Δ)
        ΔX = gather(zero(Δ)+Δ, cluster)
        counts = zero.(cluster)
        @inbounds for i = 1:size(m, 2)
            counts += sum(cluster.==i) * (cluster.==i)
        end
        @inbounds for ind = CartesianIndices(counts)
            ΔX[:, ind] ./= counts[ind]
        end
        (nothing, ΔX)
    end
end

function pool(op::Symbol, cluster::AbstractArray, X::AbstractArray)
    if op == :add
        return sumpool(cluster, X)
    elseif op == :sub
        return subpool(cluster, X)
    elseif op == :mul
        return prodpool(cluster, X)
    elseif op == :div
        return divpool(cluster, X)
    elseif op == :max
        return maxpool(cluster, X)
    elseif op == :min
        return minpool(cluster, X)
    elseif op == :mean
        return meanpool(cluster, X)
    end
end


text:
@option struct JuliaProject
    version::VersionNumber = v"0.1.0"
    uuid::UUID = uuid1()
    deps::Vector{String} = String[]
end

function compile(bp::JuliaProject, ctx::Context)
    project_name = ctx.kwargs.name
    project_path = ctx.kwargs.path
    if ispath(project_path)
        ctx.kwargs.force == true &&
            rm(project_path; force=true, recursive=true)
    end
    mkpath(project_path)

    # Project.toml
    d = Dict{String, Any}()
    d["name"] = project_name
    d["uuid"] = string(bp.uuid)
    d["version"] = string(bp.version)
    isempty(ctx.kwargs.authors) || (d["authors"] = ctx.kwargs.authors)
    d["deps"] = Dict{String, String}()
    d["compat"] = Dict{String, String}(
        "julia" => string(VERSION.major, ".", VERSION.minor)
    )

    file_path = joinpath(project_path, "Project.toml")
    open(file_path, "w+") do f
        write_project_toml(file_path, d)
    end

    isempty(bp.deps) || with_project(project_path) do
        Pkg.add(bp.deps)
    end
    return
end

function write_project_toml(filepath::String, d::Dict{String, Any})
    open(filepath, "w+") do f # following whatever Pkg does
        TOML.print(f, d; sorted=true, by=key -> (Pkg.Types.project_key_order(key), key))
    end
end


text:
using Test
using BenchmarkTools
using RecipesPipeline

function _extract_group_attributes_old_slow_known_good_implementation(v, args...; legend_entry = string)
    group_labels = collect(unique(sort(v)))
    n = length(group_labels)
    #if n > 100
    #    @warn("You created n=$n groups... Is that intended?")
    #end
    group_indices = Vector{Int}[filter(i -> v[i] == glab, eachindex(v)) for glab in group_labels]
    RecipesPipeline.GroupBy(map(legend_entry, group_labels), group_indices)
end

sc = [ "C","C","C","A", "A", "A","B","B","D"]
mc = rand([ "xx"*"$(i%6)" for i in 1:6],300)
mp = rand([ "xx"*"$(i%73)" for i in 1:73],1000)
lp = [ "xx"*"$(i%599)" for i in 1:2000]

@testset "All" begin
    @testset "Correctness" begin
        res1 = _extract_group_attributes_old_slow_known_good_implementation(sc)
        res2 = RecipesPipeline._extract_group_attributes(sc)
        @test res1.group_labels == res2.group_labels
        @test res1.group_indices == res2.group_indices
    end
    @testset "Correctness (medium)" begin
        res1 = _extract_group_attributes_old_slow_known_good_implementation(mc)
        res2 = RecipesPipeline._extract_group_attributes(mc)
        @test res1.group_labels == res2.group_labels
        @test res1.group_indices == res2.group_indices
    end
    @testset "Performance (medium)" begin
        t1 = @benchmark res1 = _extract_group_attributes_old_slow_known_good_implementation(mp)
        t2 = @benchmark res2 = RecipesPipeline._extract_group_attributes(mp)
        @test !BenchmarkTools.isregression(judge(median(t2),median(t1)))
    end
    @testset "Performance (large ish)" begin
        t1 = @benchmark res1 = _extract_group_attributes_old_slow_known_good_implementation(lp)
        t2 = @benchmark res2 = RecipesPipeline._extract_group_attributes(lp)
        @test BenchmarkTools.isimprovement(judge(median(t2),median(t1)))
    end
end


text:
"""
Rules that involved the interaction between
two grids
"""
abstract type Interaction{Keys} <: Rule end

Base.keys(::Interaction{Keys}) where Keys = Keys

# Provide a constructor for generic rule reconstruction in Flatten.jl and Setfield.jl
ConstructionBase.constructorof(::Type{T}) where T<:Interaction{Keys} where Keys =
    T{Keys}

abstract type CellInteraction{Keys} <: Interaction{Keys} end

abstract type PartialInteraction{Keys} <: Interaction{Keys} end

"""
Interactions that use a neighborhood and write to the current cell.
"""
abstract type NeighborhoodInteraction{Keys} <: Interaction{Keys} end

neighborhood(interaction::NeighborhoodInteraction) = interaction.neighborhood

"""
Interactions that write to a neighborhood.
"""
abstract type PartialNeighborhoodInteraction{Keys} <: PartialInteraction{Keys} end

neighborhood(interaction::PartialInteraction) = interaction.neighborhood


text:
"""
  module JCInput
The module required for reading in and processing the selected input file.
Import this module into the script when you need to process an input file
(which will be every single calculation).
"""
module JCInput

include("xyz_to_molecule.jl")

using MPI
using JSON
using Base.Threads

"""
  run(args::String)
Perform the operations necessary to read in, process, and extract data from the
selected input file.

One input variable is required:
1. args = The name of the input file.

Two variables are output:
1. input_info = Information gathered from the input file.
2. basis = The basis set shells, determined from the input file.

Thus, proper use of the Input.run() function would look like this:

```
input_info, basis = Input.run(args)
```
"""
function run(args; output=0)
  comm=MPI.COMM_WORLD

  if MPI.Comm_rank(comm) == 0 && output >= 2
    println("--------------------------------------------------------------------------------")
    println("                       ========================================                 ")
    println("                                READING INPUT DATA FILE                         ")
    println("                       ========================================                 ")
    println(" ")
  end

  #== output parallelization information ==#
  directory = pwd()
  #println("Input file: ", directory*"/"*input_file)
  if MPI.Comm_rank(comm) == 0 && output >= 2
    println(" ")
    println("Number of worker processes: ", MPI.Comm_size(comm))
    println("Number of threads per process: ", Threads.nthreads())
    println("Number of threads in total: ",
    MPI.Comm_size(comm)*Threads.nthreads())
  end

  #== read in input file ==#
  input_file = open(args)
    input_string = read(input_file,String)
  close(input_file)

  #== initialize variables ==#
  molecule = Dict([])
  driver = ""
  model = Dict([])
  keywords = Dict([])

  #== do extraction ==#
  json_parse = JSON.parse(input_string)

  merge!(molecule,Dict("geometry" => json_parse["molecule"]["geometry"]))
  merge!(molecule,Dict("symbols" => json_parse["molecule"]["symbols"]))
  merge!(molecule,Dict("molecular_charge" => json_parse["molecule"]["molecular_charge"]))

  #if (MPI.Comm_rank(comm) == 0) && (Threads.threadid() == 1)
  #  jldopen("tei_all.jld", "w") do file
  #    eri_array::Vector{Float64} = json_parse["molecule"]["tei"]
  #    write(file, "Integrals/All",eri_array)
  #  end
  #end

  driver = json_parse["driver"]
  merge!(model,json_parse["model"])
  merge!(keywords,json_parse["keywords"])

  if MPI.Comm_rank(comm) == 0 && output >= 2
    println(" ")
    println("                       ========================================                 ")
    println("                                       END INPUT                                ")
    println("                       ========================================                 ")
    println("--------------------------------------------------------------------------------")
  end

  return molecule, driver, model, keywords
end
export run

end


text:
# LoadError

struct L
    head
    tail
end

function oups(l)
    print(l.head)
end

oups(nothing)

max_stars_repo_name: gabriel-doriath-dohler/pjulia

max_issues_repo_name: gabriel-doriath-dohler/pjulia

max_forks_repo_name: gabriel-doriath-dohler/pjulia


text:
using Test

@testset "CombinatorialMaps" begin
  include("CombinatorialMaps.jl")
end

@testset "SimplicialSets" begin
  include("SimplicialSets.jl")
end

@testset "ExteriorCalculus" begin
  include("ExteriorCalculus.jl")
  include("DiscreteExteriorCalculus.jl")
end

@testset "Meshes" begin
  include("MeshInterop.jl")
  include("MeshGraphics.jl")
end


text: isleapyear(y) = (y % 4 == 0) && ((y % 100 != 0) || (y % 400 == 0))


text:
using Documenter, FinEtools, FinEtoolsDeforLinear, FinEtoolsFlexBeams, FinEtoolsFlexBeamsTutorials

makedocs(
	modules = [FinEtoolsFlexBeamsTutorials],
	doctest = false, clean = true,
	format = Documenter.HTML(prettyurls = false),
	authors = "Petr Krysl",
	sitename = "FinEtoolsFlexBeamsTutorials.jl",
	pages = Any[
			"Home" => "index.md",
			"Tutorials" => "tutorials/tutorials.md",
		],
	)

deploydocs(
    repo = "github.com/PetrKryslUCSD/FinEtoolsFlexBeamsTutorials.jl.git",
)


text:
struct Sketch{T} <: Surface where T <: Real
    rect::Rectangle
    width::Integer
    height::Integer
    img::Array{RGBA{T}, 2}

    function Sketch(p1::AbstractVec, p2::AbstractVec, p3::AbstractVec, img::Array{RGBA{T}, 2}) where T <: Real
        rect::Rectangle = Rectangle(p1, p2, p3)
        image_size::Tuple{Integer, Integer} = size(img)
        new{T}(rect, image_size[2], image_size[1], img)
    end
end

function hit(sketch::Sketch, ray::Ray, range::NamedTuple{(:min, :max), Tuple{T, T}} where T <: AbstractFloat)::Tuple{Bool, AbstractFloat, Union{Nothing, HitPoint}}
    succeeded::Bool, t::AbstractFloat, hp::Union{Nothing, HitPoint} = hit(sketch.rect, ray, range)
    if succeeded
        if solid(sketch, hp.u, hp.v)
            return true, t, hp
        end
    end
    false, NaN, nothing
end

function solid(sketch::Sketch, u::AbstractFloat, v::AbstractFloat)::Bool
    i::Integer = convert(Integer, floor(u * sketch.width)) + 1
    j::Integer = convert(Integer, floor((1.0 - v) * sketch.height)) + 1
    i = max(1, i)
    j = max(1, j)
    i = min(sketch.width, i)
    j = min(sketch.height, j)
    pixel::RGBA = sketch.img[j, i]
    pixel.alpha > 0
end

max_stars_repo_name: goooxu/JuliaRayTracing

max_issues_repo_name: goooxu/JuliaRayTracing

max_forks_repo_name: goooxu/JuliaRayTracing


text:
# Authors: Ali Siahkoohi, alisk@gatech.edu
# Date: March 2022

using DrWatson
@quickactivate :ConditionalNFs4Imaging

using InvertibleNetworks
using HDF5
using Random
using Statistics
using ProgressMeter
using PyPlot
using Seaborn
using LinearAlgebra
using Flux
using PyCall: @py_str

font_prop, sfmt = sef_plot_configs()
args = read_config("test_amortized_imaging.json")
args = parse_input_args(args)

py"""
from scipy.signal import hilbert
import numpy as np

def normalize_std(mu, sigma):
    analytic_mu = hilbert(mu, axis=1)
    return sigma*np.abs(analytic_mu)/(np.abs(analytic_mu)**2 + 5e+2), analytic_mu
"""
normalize_std(mu, sigma) = py"normalize_std"(mu, sigma)

max_epoch = args["max_epoch"]
lr = args["lr"]
lr_step = args["lr_step"]
batchsize = args["batchsize"]
n_hidden = args["n_hidden"]
depth = args["depth"]
sim_name = args["sim_name"]
if args["epoch"] == -1
    args["epoch"] = args["max_epoch"]
end
epoch = args["epoch"]

# Define raw data directory
mkpath(datadir("training-data"))
data_path = datadir("training-data", "training-pairs.h5")

# Download the dataset into the data directory if it does not exist
if isfile(data_path) == false
    run(`wget https://www.dropbox.com/s/53u8ckb9aje8xv4/'
        'training-pairs.h5 -q -O $data_path`)
end

# Load seismic images and create training and testing data
file = h5open(data_path, "r")
X_train = file["dm"][:, :, :, :]
Y_train = file["rtm"][:, :, :, :]

nx, ny, nc, nsamples = size(X_train)

AN_x = ActNorm(nsamples)
AN_y = ActNorm(nsamples)
X_train = AN_x.forward(X_train)
Y_train = AN_y.forward(Y_train)

# Dimensions after wavelet squeeze to increase no. of channels
nx = Int(nx / 2)
ny = Int(ny / 2)
n_in = Int(nc * 4)

# Create network
CH = NetworkConditionalHINT(n_in, n_hidden, depth)

# Loading the experiment—only network weights and training loss
loaded_keys = load_experiment(args, ["Params", "fval", "fval_eval", "train_idx"])
Params = loaded_keys["Params"]
fval = loaded_keys["fval"]
fval_eval = loaded_keys["fval_eval"]
train_idx = loaded_keys["train_idx"]
put_params!(CH, convert(Array{Any,1}, Params))

# test data pairs
idx = shuffle(setdiff(1:nsamples, train_idx))[1]
X_fixed = wavelet_squeeze(X_train[:, :, :, idx:idx])
Y_fixed = wavelet_squeeze(Y_train[:, :, :, idx:idx])

# Now select single fixed sample from all Ys
Zy_fixed = CH.forward_Y(Y_fixed)

# Draw new Zx, while keeping Zy fixed
n_samples = 1000
X_post = zeros(Float32, nx, ny, n_in, n_samples)
CH = CH |> gpu

test_batchsize = 4
test_loader = Flux.DataLoader(
    (randn(Float32, nx, ny, n_in, n_samples), repeat(Zy_fixed, 1, 1, 1, n_samples)),
    batchsize = test_batchsize,
    shuffle = false,
)

p = Progress(length(test_loader))

for (itr, (X, Y)) in enumerate(test_loader)
    Base.flush(Base.stdout)
    counter = (itr - 1) * test_batchsize + 1

    X = X |> gpu
    Y = Y |> gpu

    X_post[:, :, :, counter:(counter+size(X)[4]-1)] = (CH.inverse(X, Y)[1] |> cpu)
    ProgressMeter.next!(p)
end

X_post = wavelet_unsqueeze(X_post)
X_post = AN_x.inverse(X_post)
X_post[1:10, :, :, :] .= 0.0f0

# Some stats
X_post_mean = mean(X_post; dims = 4)
X_post_std = std(X_post; dims = 4)

X_fixed = wavelet_unsqueeze(X_fixed)
Y_fixed = wavelet_unsqueeze(Y_fixed)
Zy_fixed = wavelet_unsqueeze(Zy_fixed)

X_fixed = AN_x.inverse(X_fixed)
Y_fixed = AN_y.inverse(Y_fixed)
Y_fixed[1:10, :, :, :] .= 0.0f0

save_dict = @strdict max_epoch epoch lr lr_step batchsize n_hidden depth sim_name
save_path = plotsdir(sim_name, savename(save_dict; digits = 6))

spacing = [20.0, 12.5]
extent = [0.0, size(X_fixed, 1) * spacing[1], size(X_fixed, 2) * spacing[2], 0.0] / 1e3

signal_to_noise(xhat, x) = -20.0 * log(norm(x - xhat) / norm(x)) / log(10.0)

snr_list = []
for j = 1:n_samples
    push!(snr_list, signal_to_noise(X_post[:, :, :, j], X_fixed[:, :, :, 1]))
end
X_post_mean_snr = signal_to_noise(X_post_mean[:, :, :, 1], X_fixed[:, :, :, 1])
Y_fixed_snr = signal_to_noise(Y_fixed[:, :, :, 1], X_fixed[:, :, :, 1])

# Training loss
fig = figure("training logs", figsize = (7, 2.5))
if epoch == args["max_epoch"]
    plot(
        range(0, epoch, length = length(fval_eval)),
        fval_eval,
        color = "#addbc2",
        label = "validation loss",
    )
else
    plot(
        range(0, epoch, length = length(fval_eval[1:findfirst(fval_eval .== 0.0f0)-1])),
        fval_eval[1:findfirst(fval_eval .== 0.0f0)-1],
        color = "#addbc2",
        label = "validation loss",
    )
end
ticklabel_format(axis = "y", style = "sci", useMathText = true)
title("Negative log-likelihood")
ylabel("Validation objective")
xlabel("Epochs")
wsave(joinpath(save_path, "log.png"), fig)
close(fig)

# Plot the true model
fig = figure("x", figsize = (7.68, 4.8))
imshow(
    X_fixed[:, :, 1, 1],
    vmin = -1.5e3,
    vmax = 1.5e3,
    aspect = 1,
    cmap = "Greys",
    resample = true,
    interpolation = "lanczos",
    filterrad = 1,
    extent = extent,
)
title(L"High-fidelity image, $\mathbf{x}$")
colorbar(fraction = 0.03, pad = 0.01, format = sfmt)
grid(false)
xlabel("Horizontal distance (km)")
ylabel("Depth (km)")
safesave(joinpath(save_path, string(idx), "true_model.png"), fig)
close(fig)

# Plot the observed data
fig = figure("y", figsize = (7.68, 4.8))
imshow(
    Y_fixed[:, :, 1, 1],
    vmin = -1.5e6,
    vmax = 1.5e6,
    aspect = 1,
    cmap = "Greys",
    resample = true,
    interpolation = "lanczos",
    filterrad = 1,
    extent = extent,
)
title("RTM image")
colorbar(fraction = 0.03, pad = 0.01, format = sfmt)
grid(false)
xlabel("Horizontal distance (km)")
ylabel("Depth (km)")
safesave(joinpath(save_path, string(idx), "observed_data.png"), fig)
close(fig)

# Plot the conditional mean estimate
fig = figure("x_cm", figsize = (7.68, 4.8))
imshow(
    X_post_mean[:, :, 1, 1],
    vmin = -1.5e3,
    vmax = 1.5e3,
    aspect = 1,
    cmap = "Greys",
    resample = true,
    interpolation = "lanczos",
    filterrad = 1,
    extent = extent,
)
title("Conditional mean estimate")
colorbar(fraction = 0.03, pad = 0.01, format = sfmt)
grid(false)
xlabel("Horizontal distance (km)")
ylabel("Depth (km)")
safesave(joinpath(save_path, string(idx), "conditional_mean.png"), fig)
close(fig)

# Plot the pointwise standard deviation
fig = figure("x_std", figsize = (7.68, 4.8))
imshow(
    X_post_std[:, :, 1, 1],
    vmin = 40,
    vmax = 250.0,
    aspect = 1,
    cmap = "OrRd",
    resample = true,
    interpolation = "lanczos",
    filterrad = 1,
    extent = extent,
    norm = matplotlib.colors.LogNorm(),
)
title(L"Pointwise standard deviation, $\sigma$")
cp = colorbar(fraction = 0.03, pad = 0.01)
grid(false)
xlabel("Horizontal distance (km)")
ylabel("Depth (km)")
safesave(joinpath(save_path, string(idx), "pointwise_std.png"), fig)
close(fig)

for ns = 1:10
    fig = figure("x_cm", figsize = (7.68, 4.8))
    imshow(
        X_post[:, :, 1, ns],
        vmin = -1.5e3,
        vmax = 1.5e3,
        aspect = 1,
        cmap = "Greys",
        resample = true,
        interpolation = "lanczos",
        filterrad = 1,
        extent = extent,
    )
    title(L"Posterior sample, $\mathbf{x} \sim p(\mathbf{x} \mid \mathbf{y})$")
    colorbar(fraction = 0.03, pad = 0.01, format = sfmt)
    grid(false)
    xlabel("Horizontal distance (km)")
    ylabel("Depth (km)")
    safesave(joinpath(save_path, string(idx), "sample.png"), fig)
    close(fig)
end

# Plot the pointwise standard deviation
normalized_std, analytic_mu = normalize_std(X_post_mean[:, :, 1, 1], X_post_std[:, :, 1, 1])

fig = figure("x_std", figsize = (7.68, 4.8))
imshow(
    normalized_std,
    vmin = 3.0f-2,
    vmax = 1.2f0,
    aspect = 1,
    cmap = "OrRd",
    resample = true,
    interpolation = "hermite",
    filterrad = 1,
    extent = extent,
    norm = matplotlib.colors.LogNorm(),
)
title("Normalized pointwise standard deviation")
colorbar(fraction = 0.03, pad = 0.01, format = sfmt)
grid(false)
xlabel("Horizontal distance (km)")
ylabel("Depth (km)")
wsave(joinpath(save_path, string(idx), "normalized_pointwise_std.png"), fig)
close(fig)

println("SNR of conditional mean: ", X_post_mean_snr)
println("SNR of RTM image: ", Y_fixed_snr)


text:
#!/usr/bin/env julia

fb = open("/dev/fb0", "w")
width = 1280
height = 1024
nframes = 0
nsec = 10

frame = zeros(UInt32, width, height)
t0 = time()
while time() - t0 < nsec
    frame[:] = 0x00000000 + nframes
    seekstart(fb)
    write(fb, frame)
    #sleep(0.1)
    nframes += 1
end
fps = nframes / nsec
@show fps

max_stars_repo_name: vtjnash/orange-winner

max_issues_repo_name: vtjnash/orange-winner

max_forks_repo_name: vtjnash/orange-winner


text:
### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 4c148672-8b79-11eb-0133-6de14d906224
using HeroIcons, PlutoUI

# ╔═╡ 39d69132-8b7a-11eb-06cf-bfa0618705d8
begin
	opts = map(x -> replace(x, ".svg" => ""), readdir(
		joinpath(@__DIR__, "..", "deps", "heroicons-0.4.2", "optimized", "outline")
	))

	md"""
	# Choose an Icon to View
	$(@bind icon Select(opts))
	"""
end

# ╔═╡ 4c30529a-8b7d-11eb-314d-1ba4c5db4103
HeroIcon(icon, css="height:3em; width:3em")

# ╔═╡ 555361a0-8b7d-11eb-189e-31e44fd4fd2d
md"# In HTML"

# ╔═╡ a3b68f84-8b7d-11eb-0562-1911d449f75d
let
	s = repr("text/html", HeroIcon(icon, css="height:1em; width:1em; display:inline-block;"))
	HTML("""<h4>$s This is inside an H4 Tag</h4>""")
end

# ╔═╡ 9355cd84-8b7e-11eb-0278-05736b4c3983
md"# In Markdown"

# ╔═╡ 9792e224-8b7e-11eb-2816-f767eac59cf4
md"""
$(HeroIcon(icon, css="height: 2em; width:2em"))
"""

# ╔═╡ Cell order:
# ╠═4c148672-8b79-11eb-0133-6de14d906224
# ╟─39d69132-8b7a-11eb-06cf-bfa0618705d8
# ╟─4c30529a-8b7d-11eb-314d-1ba4c5db4103
# ╟─555361a0-8b7d-11eb-189e-31e44fd4fd2d
# ╠═a3b68f84-8b7d-11eb-0562-1911d449f75d
# ╟─9355cd84-8b7e-11eb-0278-05736b4c3983
# ╠═9792e224-8b7e-11eb-2816-f767eac59cf4


text:
module TestExamples3D

using Test
using Trixi

include("test_trixi.jl")

# pathof(Trixi) returns /path/to/Trixi/src/Trixi.jl, dirname gives the parent directory
EXAMPLES_DIR = joinpath(pathof(Trixi) |> dirname |> dirname, "examples", "3d")

# Start with a clean environment: remove Trixi output directory if it exists
outdir = "out"
isdir(outdir) && rm(outdir, recursive=true)

@testset "3D" begin

# Run basic tests
@testset "Examples 3D" begin
  # Linear scalar advection
  include("test_examples_3d_advection.jl")

  # Hyperbolic diffusion
  include("test_examples_3d_hypdiff.jl")

  # Compressible Euler
  include("test_examples_3d_euler.jl")

  # MHD
  include("test_examples_3d_mhd.jl")

  # Compressible Euler with self-gravity
  include("test_examples_3d_eulergravity.jl")

  # Lattice-Boltzmann
  include("test_examples_3d_lbm.jl")
end

@testset "Displaying components 3D" begin
  @test_nowarn include(joinpath(EXAMPLES_DIR, "elixir_advection_amr.jl"))

  # test both short and long printing formats
  @test_nowarn show(mesh); println()
  @test_nowarn println(mesh)
  @test_nowarn display(mesh)

  @test_nowarn show(equations); println()
  @test_nowarn println(equations)
  @test_nowarn display(equations)

  @test_nowarn show(solver); println()
  @test_nowarn println(solver)
  @test_nowarn display(solver)

  @test_nowarn show(solver.basis); println()
  @test_nowarn println(solver.basis)
  @test_nowarn display(solver.basis)

  @test_nowarn show(solver.mortar); println()
  @test_nowarn println(solver.mortar)
  @test_nowarn display(solver.mortar)

  @test_nowarn show(semi); println()
  @test_nowarn println(semi)
  @test_nowarn display(semi)

  @test_nowarn show(summary_callback); println()
  @test_nowarn println(summary_callback)
  @test_nowarn display(summary_callback)

  @test_nowarn show(amr_controller); println()
  @test_nowarn println(amr_controller)
  @test_nowarn display(amr_controller)

  @test_nowarn show(amr_callback); println()
  @test_nowarn println(amr_callback)
  @test_nowarn display(amr_callback)

  @test_nowarn show(stepsize_callback); println()
  @test_nowarn println(stepsize_callback)
  @test_nowarn display(stepsize_callback)

  @test_nowarn show(save_solution); println()
  @test_nowarn println(save_solution)
  @test_nowarn display(save_solution)

  @test_nowarn show(analysis_callback); println()
  @test_nowarn println(analysis_callback)
  @test_nowarn display(analysis_callback)

  @test_nowarn show(alive_callback); println()
  @test_nowarn println(alive_callback)
  @test_nowarn display(alive_callback)

  @test_nowarn println(callbacks)
end

@testset "Additional tests in 3D" begin
  @testset "compressible Euler" begin
    eqn = CompressibleEulerEquations3D(1.4)

    @test isapprox(energy_total([1.0, 2.0, 3.0, 4.0, 20.0], eqn), 20.0)
    @test isapprox(energy_kinetic([1.0, 2.0, 3.0, 4.0, 20], eqn), 14.5)
    @test isapprox(energy_internal([1.0, 2.0, 3.0, 4.0, 20], eqn), 5.5)
  end

  @testset "hyperbolic diffusion" begin
    @test_nowarn HyperbolicDiffusionEquations3D(nu=1.0)
    eqn = HyperbolicDiffusionEquations3D(nu=1.0)

    @test isapprox(initial_condition_sedov_self_gravity(collect(1:3), 4.5, eqn), zeros(4))
    @test isapprox(boundary_condition_sedov_self_gravity(collect(1:4), 1, 1, collect(11:13), 2.3, flux_central, eqn), [-1.0, -19.739208802178712, 0.0, 0.0])
    @test isapprox(boundary_condition_sedov_self_gravity(collect(1:4), 2, 2, collect(11:13), 4.5, flux_central, eqn), [-1.5, 0.0, -19.739208802178712, 0.0])
  end

  @testset "ideal GLM MHD" begin
    eqn = IdealGlmMhdEquations3D(1.4)
    u = [1.0, 2.0, 3.0, 4.0, 20.0, 0.1, 0.2, 0.3, 1.5]

    @test isapprox(density(u, eqn), 1.0)
    @test isapprox(pressure(u, eqn), 1.7219999999999995)
    @test isapprox(density_pressure(u, eqn), 1.7219999999999995)

    @test isapprox(Trixi.entropy_thermodynamic(u, eqn), 0.5434864060055388)
    @test isapprox(Trixi.entropy_math(u, eqn), -1.3587160150138473)
    @test isapprox(Trixi.entropy(u, eqn), -1.3587160150138473)

    @test isapprox(energy_total(u, eqn), 20.0)
    @test isapprox(energy_kinetic(u, eqn), 14.5)
    @test isapprox(energy_magnetic(u, eqn), 0.07)
    @test isapprox(energy_internal(u, eqn), 4.305)

    @test isapprox(cross_helicity(u, eqn), 2.0)
  end
end

# Clean up afterwards: delete Trixi output directory
@test_nowarn rm(outdir, recursive=true)

end # 3D

end #module


text:
# it has unique keys
# leftChild has keys which are less than the node
# rightChild has keys which are greater than the node
# height stores the height of the subtree.
mutable struct AVLTreeNode{K}
    height::Int8
    leftChild::Union{AVLTreeNode{K}, Nothing}
    rightChild::Union{AVLTreeNode{K}, Nothing}
    subsize::Int32
    data::K

    AVLTreeNode{K}(d::K) where K = new{K}(1, nothing, nothing, 1, d)
end

AVLTreeNode(d) = AVLTreeNode{Any}(d)

AVLTreeNode_or_null{T} = Union{AVLTreeNode{T}, Nothing}

mutable struct AVLTree{T}
    root::AVLTreeNode_or_null{T}
    count::Int

    AVLTree{T}() where T = new{T}(nothing, 0)
end

AVLTree() = AVLTree{Any}()

Base.length(tree::AVLTree) = tree.count

get_height(node::Union{AVLTreeNode, Nothing}) = (node == nothing) ? 0 : node.height

# balance is the difference of height between leftChild and rightChild of a node.
function get_balance(node::Union{AVLTreeNode, Nothing})
    if node == nothing
        return 0
    else
        return get_height(node.leftChild) - get_height(node.rightChild)
    end
end

# computes the height of the subtree, which basically is
# one added the maximum of the height of the left subtree and right subtree
compute_height(node::AVLTreeNode) = 1 + max(get_height(node.leftChild), get_height(node.rightChild))

get_subsize(node::AVLTreeNode_or_null) = (node == nothing) ? 0 : node.subsize

# compute the subtree size
function compute_subtree_size(node::AVLTreeNode_or_null)
    if node == nothing
        return 0
    else
        L = get_subsize(node.leftChild)
        R = get_subsize(node.rightChild)
        return (L + R + 1)
    end
end

"""
    left_rotate(node_x::AVLTreeNode)

Performs a left-rotation on `node_x`, updates height of the nodes, and returns the rotated node.
"""
function left_rotate(z::AVLTreeNode)
    y = z.rightChild
    α = y.leftChild
    y.leftChild = z
    z.rightChild = α
    z.height = compute_height(z)
    y.height = compute_height(y)
    z.subsize = compute_subtree_size(z)
    y.subsize = compute_subtree_size(y)
    return y
end

"""
    right_rotate(node_x::AVLTreeNode)

Performs a right-rotation on `node_x`, updates height of the nodes, and returns the rotated node.
"""
function right_rotate(z::AVLTreeNode)
    y = z.leftChild
    α = y.rightChild
    y.rightChild = z
    z.leftChild = α
    z.height = compute_height(z)
    y.height = compute_height(y)
    z.subsize = compute_subtree_size(z)
    y.subsize = compute_subtree_size(y)
    return y
end

"""
   minimum_node(tree::AVLTree, node::AVLTreeNode)

Returns the AVLTreeNode with minimum value in subtree of `node`.
"""
function minimum_node(node::Union{AVLTreeNode, Nothing})
    while node != nothing && node.leftChild != nothing
        node = node.leftChild
    end
    return node
end

function search_node(tree::AVLTree{K}, d::K) where K
    prev = nothing
    node = tree.root
    while node != nothing && node.data != nothing && node.data != d

        prev = node
        if d < node.data
            node = node.leftChild
        else
            node = node.rightChild
        end
    end

    return (node == nothing) ? prev : node
end

function Base.haskey(tree::AVLTree{K}, d::K) where K
    (tree.root == nothing) && return false
    node = search_node(tree, d)
    return (node.data == d)
end

Base.in(key, tree::AVLTree) = haskey(tree, key)

function Base.insert!(tree::AVLTree{K}, d::K) where K

    function insert_node(node::Union{AVLTreeNode, Nothing}, key)
        if node == nothing
            return AVLTreeNode{K}(key)
        elseif key < node.data
            node.leftChild = insert_node(node.leftChild, key)
        else
            node.rightChild = insert_node(node.rightChild, key)
        end

        node.subsize = compute_subtree_size(node)
        node.height = compute_height(node)
        balance = get_balance(node)

        if balance > 1
            if key < node.leftChild.data
                return right_rotate(node)
            else
                node.leftChild = left_rotate(node.leftChild)
                return right_rotate(node)
            end
        end

        if balance < -1
            if key > node.rightChild.data
                return left_rotate(node)
            else
                node.rightChild = right_rotate(node.rightChild)
                return left_rotate(node)
            end
        end

        return node
    end

    haskey(tree, d) && return tree

    tree.root = insert_node(tree.root, d)
    tree.count += 1
    return tree
end

function Base.push!(tree::AVLTree{K}, key0) where K
    key = convert(K, key0)
    insert!(tree, key)
end

function Base.delete!(tree::AVLTree{K}, d::K) where K

    function delete_node!(node::Union{AVLTreeNode, Nothing}, key)
        if key < node.data
            node.leftChild = delete_node!(node.leftChild, key)
        elseif key > node.data
            node.rightChild = delete_node!(node.rightChild, key)
        else
            if node.leftChild == nothing
                result = node.rightChild
                return result
            elseif node.rightChild == nothing
                result = node.leftChild
                return result
            else
                result = minimum_node(node.rightChild)
                node.data = result.data
                node.rightChild = delete_node!(node.rightChild, result.data)
            end
        end

        node.subsize = compute_subtree_size(node)
        node.height = compute_height(node)
        balance = get_balance(node)

        if balance > 1
            if get_balance(node.leftChild) >= 0
                return right_rotate(node)
            else
                node.leftChild = left_rotate(node.leftChild)
                return right_rotate(node)
            end
        end

        if balance < -1
            if get_balance(node.rightChild) <= 0
                return left_rotate(node)
            else
                node.rightChild = right_rotate(node.rightChild)
                return left_rotate(node)
            end
        end

        return node
    end

    # if the key is not in the tree, do nothing and return the tree
    !haskey(tree, d) && return tree

    # if the key is present, delete it from the tree
    tree.root = delete_node!(tree.root, d)
    tree.count -= 1
    return tree
end

"""
    sorted_rank(tree::AVLTree, key)

Returns the rank of `key` present in the `tree`, if it present. A KeyError is thrown if `key` is not present.
"""
function sorted_rank(tree::AVLTree{K}, key::K) where K
    !haskey(tree, key) && throw(KeyError(key))
    node = tree.root
    rank = 0
    while node.data != key
        if (node.data < key)
            rank += (1 + get_subsize(node.leftChild))
            node = node.rightChild
        else
            node = node.leftChild
        end
    end
    rank += (1 + get_subsize(node.leftChild))
    return rank
end

function Base.getindex(tree::AVLTree{K}, ind::Integer) where K
    @boundscheck (1 <= ind <= tree.count) || throw(BoundsError("$ind should be in between 1 and $(tree.count)"))
    function traverse_tree(node::AVLTreeNode_or_null, idx)
        if (node != nothing)
            L = get_subsize(node.leftChild)
            if idx <= L
                return traverse_tree(node.leftChild, idx)
            elseif idx == L + 1
                return node.data
            else
                return traverse_tree(node.rightChild, idx - L - 1)
            end
        end
    end
    value = traverse_tree(tree.root, ind)
    return value
end

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
using MPIMeasurements

scanner = MPIScanner("MPS.toml")
daq = getDAQ(scanner)

params = toDict(daq.params)
#<<<<<<< HEAD

params["studyName"]="TestTobi"
params["studyDescription"]="A very cool measurement"
params["scannerOperator"]="Tobi"
params["dfStrength"]=[0.02]
params["acqNumAverages"]=1000
params["calibFieldToVolt"]=[19.4]#dfstrength*calibF2V=Uout1Rp
params["calibRefToField"]=[0.01639344262295082]#[0.013242]calibR2F*Uref=Field

measurementCont(daq, params, controlPhase=true)
#=======
params["studyName"]="TestTobi"
params["studyDescription"]="A very cool measurement"
params["scannerOperator"]="Tobi"
params["dfStrength"]=[0.001]
params["acqNumAverages"]=1000
#iterativ setting of calibFieldToVolt
#start with dfstrength=1 and calibFieldToVolt=0.98
params["calibFieldToVolt"]=[12.91]
params["calibRefToField"]=[0.012195]

#<<<<<<< HEAD
measurementCont(daq, params, controlPhase=false)
=======
#
measurementCont(daq, params, controlPhase=false)
#>>>>>>> 5c0f102c515252749e73e93636483ed0aa25267d
#>>>>>>> 9375e758066b195c64f90ab5757ce438a2fdddd8
=#


text:
export CachedMinkowskiSumArray,
       forget_sets!

# =============================================================
# Minkowski sum of an array of sets with a support vector cache
# =============================================================

"""
    CachedPair{N}

A mutable pair of an index and a vector.

### Fields

- `idx` -- index
- `vec` -- vector
"""
mutable struct CachedPair{N}
    idx::Int
    vec::AbstractVector{N}
end

function getindex(cp::CachedPair, idx::Int)
    if idx == 1
        return cp.idx
    elseif idx == 2
        return cp.vec
    end
    error("invalid index $idx, can only access a pair at index 1 or 2")
end

"""
    CachedMinkowskiSumArray{N<:Real, S<:LazySet{N}} <: LazySet{N}

Type that represents the Minkowski sum of a finite number of convex sets.
Support vector queries are cached.

### Fields

- `array` -- array of convex sets
- `cache` -- cache of support vector query results

### Notes

This type assumes that the dimensions of all elements match.

The `ZeroSet` is the neutral element and the `EmptySet` is the absorbing element
for `CachedMinkowskiSumArray`.

The cache (field `cache`) is implemented as dictionary whose keys are directions
and whose values are pairs `(k, s)` where `k` is the number of elements in the
array `array` when the support vector was evaluated last time, and `s` is the
support vector that was obtained.
Thus this type assumes that `array` is not modified except by adding new sets at
the end.

Constructors:

- `CachedMinkowskiSumArray(array::Vector{<:LazySet})` -- default constructor

- `CachedMinkowskiSumArray([n]::Int=0, [N]::Type=Float64)`
  -- constructor for an empty sum with optional size hint and numeric type
"""
struct CachedMinkowskiSumArray{N<:Real, S<:LazySet{N}} <: LazySet{N}
    array::Vector{S}
    cache::Dict{AbstractVector{N}, CachedPair{N}}

    # default constructor that initializes cache
    CachedMinkowskiSumArray(arr::Vector{S}) where {N<:Real, S<:LazySet{N}} =
        new{N, S}(arr, Dict{AbstractVector{N}, CachedPair{N}}())
end

isoperationtype(::Type{<:CachedMinkowskiSumArray}) = true
isconvextype(::Type{CachedMinkowskiSumArray{N, S}}) where {N, S} = isconvextype(S)

# constructor for an empty sum with optional size hint and numeric type
function CachedMinkowskiSumArray(n::Int=0, N::Type=Float64)
    arr = Vector{LazySet{N}}()
    sizehint!(arr, n)
    return CachedMinkowskiSumArray(arr)
end

# ZeroSet is the neutral element for CachedMinkowskiSumArray
@neutral(CachedMinkowskiSumArray, ZeroSet)

# EmptySet and Universe are the absorbing element for CachedMinkowskiSumArray
@absorbing(CachedMinkowskiSumArray, EmptySet)
# @absorbing(CachedMinkowskiSumArray, Universe)  # TODO problematic

"""
    array(cms::CachedMinkowskiSumArray{N, S}) where {N<:Real, S<:LazySet{N}}

Return the array of a caching Minkowski sum.

### Input

- `cms` -- caching Minkowski sum

### Output

The array of a caching Minkowski sum.
"""
function array(cms::CachedMinkowskiSumArray{N, S}
              ) where {N<:Real, S<:LazySet{N}}
    return cms.array
end

"""
    dim(cms::CachedMinkowskiSumArray)

Return the dimension of a caching Minkowski sum.

### Input

- `cms` -- caching Minkowski sum

### Output

The ambient dimension of the caching Minkowski sum.
"""
function dim(cms::CachedMinkowskiSumArray)
    return length(cms.array) == 0 ? 0 : dim(cms.array[1])
end

"""
    σ(d::AbstractVector{N}, cms::CachedMinkowskiSumArray{N}) where {N<:Real}

Return the support vector of a caching Minkowski sum in a given direction.

### Input

- `d`   -- direction
- `cms` -- caching Minkowski sum

### Output

The support vector in the given direction.
If the direction has norm zero, the result depends on the summand sets.

### Notes

The result is cached, i.e., any further query with the same direction runs in
constant time.
When sets are added to the caching Minkowski sum, the query is only performed
for the new sets.
"""
function σ(d::AbstractVector{N}, cms::CachedMinkowskiSumArray{N}) where {N<:Real}
    arr = array(cms)
    l = length(arr)
    cache = cms.cache
    if haskey(cache, d)
        pair = cache[d]
        k = pair.idx
        svec1 = pair.vec
        if k == l
            # has already stored the full support vector
            return svec1
        else
            # has only stored the support vector of the first k sets
            @assert k < l "invalid cache index"
            svec = svec1 + σ_helper(d, @view arr[k+1:l])
        end
    else
        # first-time computation of support vector
        svec = σ_helper(d, arr)
    end
    # NOTE: make a copy of the direction vector (can be modified outside)
    cache[copy(d)] = CachedPair(l, svec)
    return svec
end

"""
	isbounded(cms::CachedMinkowskiSumArray)

Determine whether a caching Minkowski sum is bounded.

### Input

- `cms` -- caching Minkowski sum

### Output

`true` iff all wrapped sets are bounded.
"""
function isbounded(cms::CachedMinkowskiSumArray)
    return all(isbounded, cms.array)
end

"""
    isempty(cms::CachedMinkowskiSumArray)

Return if a caching Minkowski sum array is empty or not.

### Input

- `cms` -- caching Minkowski sum

### Output

`true` iff any of the wrapped sets are empty.

### Notes

Forgotten sets cannot be checked anymore.
Usually they have been empty because otherwise the support vector query should
have crashed before.
In that case, the caching Minkowski sum should not be used further.
"""
function isempty(cms::CachedMinkowskiSumArray)
    return any(isempty, array(cms))
end

"""
    forget_sets!(cms::CachedMinkowskiSumArray)

Tell a caching Minkowski sum to forget the stored sets (but not the support
vectors).
Only those sets are forgotten such that for each cached direction the support
vector has been computed before.

### Input

- `cms` -- caching Minkowski sum

### Output

The number of sets that have been forgotten.

### Notes

This function should only be used under the assertion that no new directions are
queried in the future; otherwise such support vector results will be incorrect.

This implementation is optimistic and first tries to remove all sets.
However, it also checks that for all cached directions the support vector has
been computed before.
If it finds that this is not the case, the implementation identifies the
biggest index ``k`` such that the above holds for the ``k`` oldest sets, and
then it only removes these.
See the example below.

### Examples

```jldoctest
julia> x1 = BallInf(ones(3), 3.); x2 = Ball1(ones(3), 5.);

julia> cms1 = CachedMinkowskiSumArray(2); cms2 = CachedMinkowskiSumArray(2);

julia> d = ones(3);

julia> a1 = array(cms1); a2 = array(cms2);

julia> push!(a1, x1); push!(a2, x1);

julia> σ(d, cms1); σ(d, cms2);

julia> push!(a1, x2); push!(a2, x2);

julia> σ(d, cms1);

julia> idx1 = forget_sets!(cms1) # support vector was computed for both sets
2

julia> idx1 = forget_sets!(cms2) # support vector was only computed for first set
1
```
"""
function forget_sets!(cms::CachedMinkowskiSumArray)
    len = length(cms.array)
    sets_to_remove = len
    for pair in values(cms.cache)
        if pair.idx != len
            @assert pair.idx < len "invalid cache index"
            sets_to_remove = min(sets_to_remove, pair.idx)
            pair.idx -= len
        else
            pair.idx = 0
        end
    end

    if sets_to_remove != len
        # not all sets shall be forgotten
        offset = len - sets_to_remove
        for pair in values(cms.cache)
            # fix the wrong index assignments from above
            pair.idx += offset
        end
        # forget only the first k sets
        splice!(cms.array, 1:sets_to_remove)
    else
        # forget all sets
        empty!(cms.array)
    end
    return sets_to_remove
end


text:
# Unit tests
include("unit/test_io.jl")

# Functional tests
include("functional/test_main.jl")


text:
"""
    NAryState{T <: Number, BT <:Unsigned} <: DualValuedState
"""
mutable struct NAryState{T <: Number, N} <: FiniteBasisState
    n::Int                  # number of elements
    σ::Vector{T}            # the states as 0/1
    i_σ::Int64              # the number corresponding to σ_binary

    space_dim::Int64        # N^Nsites
end
NAryState{T,N}(n, i_σ=0) where {T, N} = NAryState(T, N, n, i_σ)

NAryState(N::Int, n, args...) = NAryState(Float32, N, n, args...)
function NAryState(T::Type, N::Int, n, i_σ=0)
    @assert N>= 0
    @assert i_σ >= 0

    NAryState{T,N}(n,
              IntegerToState!(Vector{T}(undef, n), N, i_σ),
              i_σ,
              N^n)
end

export NAryState

# custom accessors
@inline config(state::NAryState) = state.σ

## Property Accessors
@inline spacedimension(st::NAryState) = st.space_dim
@inline nsites(state::NAryState) = state.n
@inline local_dimension(state::Type{NAryState{T,N}}) where {T,N} = N
@inline local_dimension(state::NAryState{T,N}) where {T,N} = local_dimension(typeof(state))
@inline eltype(state::NAryState{T,N}) where {T,N} = T

@inline toint(state::NAryState) = state.i_σ
@inline index(state::NAryState) = toint(state)+1
@inline index_to_int(state::NAryState, id) = id - 1

# checks
same_basis(s1::NAryState, s2::NAryState)  =
    nsites(s1)==nsites(s2) && local_dimension(s1) == local_dimension(s2) && eltype(s1) == eltype(s2)

# Operations on tuples
@inline _toint(left::NAryState{T,N}, right::NAryState{T,N}) where {T,N} =
    toint(left) * spacedimension(right) + toint(right)
#@inline index(row::NAryState, col::NAryState) = toint(row, col) + 1
#=function index_to_int(id, σrow::NAryState{T,N}, σcol::NAryState{T,N}) where {T,N}
    i = id - 1
    row = div(i, spacedimension(σrow)) #div(i, N^nsites(σrow))
    col = i - row * spacedimension(σrow)
    return (row, col)
end=#

# Operations
"""
    flipat!(rng, state, site) -> (old_val, new_val)

Flips `state[site]` to a random new state. The old value is returned together
with the new value.
"""
function flipat!(rng::AbstractRNG, state::NAryState{T, N}, i::Int) where {T, N}
    old_val = state.σ[i]

    # Generate new value
    new_val = T(rand(rng, 0:(N-1)))
    while new_val == old_val
        new_val = T(rand(rng, 0:(N-1)))
    end

    # Modify
    state.i_σ += (Int(new_val)-Int(old_val))*N^(i-1)
    state.σ[i] = new_val

    return (old_val, new_val)
end

# For Nary states the fast method is equivalent to the standard one
flipat_fast!(rng::AbstractRNG, state::NAryState, i::Int) =
    flipat!(rng, state, i)

"""
    setat!(state, site, value) -> old_val

Sets `state[site] = value`. Returns the old value of `state[site]`
"""
function setat!(state::NAryState{T, N}, i::Int, val::T) where {T, N}
    old_val = state.σ[i]

    state.i_σ += (Int(val)-Int(old_val))*N^(i-1)
    state.σ[i] = val
    return old_val
end

set_index!(state::NAryState, val::Integer) = set!(state, index_to_int(state, val))
function set!(state::NAryState{T, N}, val::Integer) where {T, N}
    state.i_σ = val
    for i=1:state.n
        val, state.σ[i] = divrem(val, N)
    end
    state
end

function add!(state::NAryState, val::Integer)
    set!(state, val+toint(state))
    state
end

function rand!(rng::AbstractRNG, state::NAryState)
    val = rand(rng, 0:(spacedimension(state)-1))
    set!(state, val)
end

flipped(a::NAryState{T,N}, b::NAryState{T,N}) where {T,N} =
    throw(ErrorException("Flipped for Nary states not implemented"))

IntegerToState(n_sites, loc_dim, val, T::Type{<:Number}) =
    IntegerToState!(Vector{T}(undef, n_sites), loc_dim, val)

@inline function IntegerToState!(arr, loc_dim, val)
    for i=1:length(arr)
        val, arr[i] = divrem(val, loc_dim)
    end
    arr
end

# --- Indexing function for local spaces used in operators
local_index(s::NAryState, i::T) where {T<:Integer} =
    T(s.σ[i])+1

function local_index(s::NAryState{T,Nb}, is::AbstractVector{T2}) where {T,Nb,T2<:Integer}
    idx = 1
    for (i,j)=enumerate(is)
        idx += T2(s.σ[j]) * Nb^(i-1)
    end
    return idx
end
# -- end

function apply!(state::NAryState, changes::StateChanges)
    for (id, val)=changes
        setat!(state, id, val)
    end
    return state
end

# -------------- Base.show extension for nice printing -------------- #
Base.show(io::IO, ::MIME"text/plain", bs::NAryState) = print(io, "NAryState(",bs.n,") : ", String(bs.σ,false),
                                           " = ", bs.i_σ)
Base.show(io::IO, bs::NAryState) = print(io, bs.i_σ, String(bs.σ,false))

function String(bv::Vector, toInt=true)
    str="["
    for el in bv
        str *= Char(Int(el) + '0')
    end
    str *= ']'
    str
end

Base.show(io::IO, ::MIME"text/plain", bs::NAryState{T,2}) where T = print(io, "NAryState(",bs.n,") : ", StringToSpin(bs.σ,false),
                                           " = ", bs.i_σ)
Base.show(io::IO, bs::NAryState{T,2}) where T = print(io, bs.i_σ, StringToSpin(bs.σ,false))
function StringToSpin(bv::Vector, toInt=true)
    str="["
    for el in bv
        str *= el > 0.0 ? Char('↑') : Char('↓')
    end
    str *= ']'
    str
end


text:
@testset "isless() for BibInternal.Date" begin
    @test BibInternal.Date("", "", "2000") < BibInternal.Date("", "", "2001")
    @test BibInternal.Date("", "10", "2000") < BibInternal.Date("", "11", "2000")
    @test BibInternal.Date("5", "01", "2000") < BibInternal.Date("15", "01", "2000")

    @test BibInternal.Date("", "", "1970") > BibInternal.Date("", "", "1969")
    @test BibInternal.Date("", "7", "1970") > BibInternal.Date("", "6", "1970")
    @test BibInternal.Date("7", "7", "1970") > BibInternal.Date("6", "7", "1970")

    @test BibInternal.Date("", "", "1805") == BibInternal.Date("", "", "1805")
    @test BibInternal.Date("", "4", "1805") == BibInternal.Date("", "4", "1805")
    @test BibInternal.Date("3", "4", "1805") == BibInternal.Date("3", "4", "1805")

    @test BibInternal.Date("", "", "1805") !== BibInternal.Date("", "", "1905")
    @test BibInternal.Date("", "4", "1805") !== BibInternal.Date("", "5", "1805")
    @test BibInternal.Date("3", "4", "1805") !== BibInternal.Date("4", "4", "1805")

    @test_skip BibInternal.Date("", "May", "1805") == BibInternal.Date("", "5", "1805")
    @test_skip BibInternal.Date("1th", "5", "1805") !== BibInternal.Date("1", "5", "1805")

    @test_throws ArgumentError BibInternal.Date("", "", "1") <
                               BibInternal.Date("", "", "2000a")
    @test_throws ArgumentError BibInternal.Date("", "", "1") < BibInternal.Date("", "", "")
    @test_throws ArgumentError BibInternal.Date("", "", "1") <
                               BibInternal.Date("", "", "90ies")
    @test_throws ArgumentError BibInternal.Date("", "", "40k") <
                               BibInternal.Date("", "", "40000")
end

@testset "isless() for BibInternal.Name" begin
    # TODO: consider testing for alphabetizing rules

    @test BibInternal.Name("", "Cow", "", "John", "") <
          BibInternal.Name("", "Doe", "", "John", "")
    @test BibInternal.Name("", "Doe", "", "John", "A.") <
          BibInternal.Name("", "Doe", "", "John", "B.")
    @test BibInternal.Name("", "Doe", "", "Bronn", "") <
          BibInternal.Name("", "Doe", "", "John", "")
    @test BibInternal.Name("", "Bronn", "", "Would", "") <
          BibInternal.Name("", "Bronn", "", "Would", "Not")

    @test BibInternal.Name("", "Doe", "", "John", "") ==
          BibInternal.Name("", "Doe", "", "John", "")
    @test BibInternal.Name("", "Doe", "jun.", "John", "") ==
          BibInternal.Name("", "Doe", "jun.", "John", "")
    @test BibInternal.Name("", "Doe", "", "John", "E.") ==
          BibInternal.Name("", "Doe", "", "John", "E.")

    @test BibInternal.Name("", "Doe", "", "Bronn", "") !==
          BibInternal.Name("", "Doe", "", "John", "")
    @test BibInternal.Name("", "Doe", "jun.", "John", "") !==
          BibInternal.Name("", "Doe", "sen.", "John", "")
    @test BibInternal.Name("", "Doe", "", "John", "E.") !==
          BibInternal.Name("", "Doe", "", "John", "A.")
end


text:
using ReinLearn
using Test

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


text:
using Markdown
using GraphRecipes
using Plots

cd(@__DIR__)

readme = read("../README.md", String) |> Markdown.parse
content = readme.content

code_blocks = []
for paragraph in content
    if paragraph isa Markdown.Code
        push!(code_blocks, paragraph.code)
    end
end

# Parse the code examples on the README into expressions. Ignore the first one, which is
# the installation instructions.
readme_exprs = [Meta.parse("begin $(code_blocks[i]) end") for i in 2:length(code_blocks)]

julia_logo_pun() = eval(readme_exprs[1])