text:
# test projections onto intersection for julia for simultaneous image inpainting&deblurring
# as set theoretical image recovery problem (a feasibility problem)
# We will use a VERY simple learning appraoch to obtain 'good' constraints. This
# learing works with just a few even <10 training examples.
# South America dataset
using Distributed
using LinearAlgebra
using SparseArrays
using Random
using Statistics
using StatsBase

@everywhere using SetIntersectionProjection
using MAT

@everywhere mutable struct compgrid
  d :: Tuple
  n :: Tuple
end

#select working precision
FL=32
if     FL==64
  TF = Float64
  TI = Int64
elseif FL==32
  TF = Float32
  TI = Int32
end

#load a very small data set (12 images only) (Mablab files for compatibility with matlab only solvers for comparison...)
file = matopen(joinpath(dirname(pathof(SetIntersectionProjection)), "../examples/Data/SA_patches.mat"))
mtrue = read(file, "SA_patches")
mtrue = convert(Array{TF,3},mtrue)

#split data tensor in a training and evaluation data
#patches have been randomized already in matlab
m_train      = mtrue[1:35,:,:]
m_evaluation = mtrue[36:39,:,:]
m_est        = zeros(TF,size(m_evaluation))

#computational grid for the training images (all images assumed to be on the same grid here)
comp_grid = compgrid((1, 1),(size(m_evaluation,2), size(m_evaluation,3)))
#convert model and computational grid parameters
comp_grid=compgrid( ( convert(TF,comp_grid.d[1]),convert(TF,comp_grid.d[2]) ), comp_grid.n )

#create true observed data by blurring, setting pixels to zero, and adding random zero mean noise

#first, motion blur as a sparse matrix
bkl   = 25; #blurring kernel length
d_obs = zeros(TF,size(m_evaluation,1),comp_grid.n[1]-bkl,comp_grid.n[2])
#temp  = zeros(TF,comp_grid.n[1]-bkl,comp_grid.n[2])

n1 = comp_grid.n[1]
Bx = SparseMatrixCSC{TF}(LinearAlgebra.I,n1,n1)./bkl
for i=2:bkl
  temp =  spdiagm(i => ones(TF,n1)) ./ bkl #spdiagm(ones(n1)./bkl,i)
  temp =  temp[1:n1,1:n1];
  global Bx   += temp;
end
Bx = Bx[1:end-bkl,:]
Iz = SparseMatrixCSC{TF}(LinearAlgebra.I,comp_grid.n[2],comp_grid.n[2])# speye(TF,comp_grid.n[2]);
BF = kron(Iz,Bx);
#BF = convert(SparseMatrixCSC{TF,TI},BF);

#second, subsample
(e1,e2,e3) = size(d_obs)
mask       = ones(TF,e2*e3)
s          = randperm(e3*e2)
zero_ind   = s[1:Int(8.0 .* round((e2*e3)/10))]
mask[zero_ind] .= TF(0.0)
mask       = spdiagm(0 => mask)#spdiagm(mask,0)
FWD_OP     = mask*BF#convert(SparseMatrixCSC{TF,TI},mask*BF)

#blur and subsample to create observed data
for i=1:size(d_obs,1)
  temp = FWD_OP*vec(m_evaluation[i,:,:])

  #add noise (integers between -10 and 10 for each nonzero observation point)
  noise         = rand((-10:10),count(!iszero, temp))#countnz(temp))
  nz_ind        = findall(x->x!=0, temp)
  temp[nz_ind] .= temp[nz_ind].+noise
  d_obs[i,:,:]  = reshape(temp,comp_grid.n[1]-bkl,comp_grid.n[2])
end
println("finished creating observed data")

#"train" by observing constraints on the data images in training data set
observations = constraint_learning_by_obseration(comp_grid,m_train)
println("finished learning by observation")

#define a few constraints
constraint = Vector{SetIntersectionProjection.set_definitions}()

#bounds:
m_min     = 0.0
m_max     = 255.0
set_type  = "bounds"
TD_OP     = "identity"
app_mode  = ("matrix","")
custom_TD_OP = ([],false)
push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP))

#relaxed histogram constraint:
m_min     = observations["hist_min"]
m_max     = observations["hist_max"]
set_type  = "histogram"
TD_OP     = "identity"
app_mode  = ("matrix","")
custom_TD_OP = ([],false)
push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP))

#relaxed histogram constraint on discrete derivative of the image:
# m_min     = observations["hist_TV_min"]
# m_max     = observations["hist_TV_max"]
# set_type  = "histogram"
# TD_OP     = "TV"
# app_mode  = ("matrix","")
# custom_TD_OP = ([],false)
# push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP))

# #rank that preserves 95% of the training images:
# m_min     = 0
# observations["rank_095"]=sort(vec(observations["rank_095"]))
# m_max     = convert(TI,round(quantile(observations["rank_095"],0.33)))
# set_type  = "rank"
# TD_OP     = "identity"
# app_mode  = ("matrix","")
# custom_TD_OP = ([],false)
# push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP))

#nuclear norm constraint
m_min     = 0
m_max     = convert(TF,quantile(vec(observations["nuclear_norm"]),0.33))
set_type  = "nuclear"
TD_OP     = "identity"
app_mode  = ("matrix","")
custom_TD_OP = ([],false)
push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP))

#nuclear norm constraint on the x-derivative of the image
m_min     = 0
m_max     = convert(TF,quantile(vec(observations["nuclear_Dx"]),0.33))
set_type  = "nuclear"
TD_OP     = "D_x"
app_mode  = ("matrix","")
custom_TD_OP = ([],false)
push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP))

#nuclear norm constraint on the z-derivative of the image
m_min     = 0
m_max     = convert(TF,quantile(vec(observations["nuclear_Dz"]),0.33))
set_type  = "nuclear"
TD_OP     = "D_z"
app_mode  = ("matrix","")
custom_TD_OP = ([],false)
push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP))

#anisotropic total-variation constraint:
m_min     = 0
m_max     = convert(TF,quantile(vec(observations["TV"]),0.33))
set_type  = "l1"
TD_OP     = "TV"
app_mode  = ("matrix","")
custom_TD_OP = ([],false)
push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP))

#anisotropic total-variation constraint:
m_min     = 0
m_max     = convert(TF,quantile(vec(observations["wavelet_l1"]),0.33))
set_type  = "l1"
TD_OP     = "wavelet"
app_mode  = ("matrix","")
custom_TD_OP = ([],false)
push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP))

#l2 constraint on discrete derivatives of the image:
m_min     = 0
m_max     = convert(TF,quantile(vec(observations["D_l2"]),0.33))
set_type  = "l2"
TD_OP     = "TV"
app_mode  = ("matrix","")
custom_TD_OP = ([],false)
push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP))

# constraint["use_TD_l1_2"]=false
# constraint["TD_l1_operator_2"]="curvelet"
# constraint["TD_l1_sigma_2"] = 0.5f0*convert(TF,quantile(vec(observations["curvelet_l1"]),0.33))
#
# constraint["use_TD_l1_3"]=false
# constraint["TD_l1_operator_3"]="DFT"
# constraint["TD_l1_sigma_3"] = convert(TF,quantile(vec(observations["DFT_l1"]),0.33))

#bound constraints on x-derivative
m_min     = convert(TF,quantile(vec(observations["D_x_min"]),0.15))
m_max     = convert(TF,quantile(vec(observations["D_x_max"]),0.85))
set_type  = "bounds"
TD_OP     = "D_x"
app_mode  = ("matrix","")
custom_TD_OP = ([],false)
push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP))

#bound constraints on z-derivative
m_min     = convert(TF,quantile(vec(observations["D_z_min"]),0.15))
m_max     = convert(TF,quantile(vec(observations["D_z_max"]),0.85))
set_type  = "bounds"
TD_OP     = "D_z"
app_mode  = ("matrix","")
custom_TD_OP = ([],false)
push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP))

#curvelet cardinality constraint
# m_min     = 0
# m_max     = convert(TI,round(quantile(vec(observations["curvelet_card_095"]),0.85)))
# set_type  = "cardinality"
# TD_OP     = "curvelet"
# app_mode  = ("matrix","")
# custom_TD_OP = ([],false)
# push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP))

#annulus constraint
m_min     = quantile(observations["annulus"],0.15)
m_max     = quantile(observations["annulus"],0.85)
set_type  = "annulus"
TD_OP     = "identity"
app_mode  = ("matrix","")
custom_TD_OP = ([],false)
push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP))

#annulus constraint on discrete gradients
m_min     = quantile(observations["TV_annulus"],0.15)
m_max     = quantile(observations["TV_annulus"],0.85)
set_type  = "annulus"
TD_OP     = "TV"
app_mode  = ("matrix","")
custom_TD_OP = ([],false)
push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP))

#PARSDMM options:
options=PARSDMM_options()
options=default_PARSDMM_options(options,options.FL)
options.evol_rel_tol = 1f-6
options.feas_tol     = 0.001f0
options.obj_tol      = 0.0002f0
options.adjust_gamma           = true
options.adjust_rho             = true
options.adjust_feasibility_rho = true
options.Blas_active            = true
options.maxit                  = 300
set_zero_subnormals(true)

options.feasibility_only     = true #this is important to set
options.parallel             = true
options.zero_ini_guess       = true
BLAS.set_num_threads(2)
#FFTW.set_num_threads(2) (this is not supported in Julia 0.7)

(P_sub,TD_OP,set_Prop) = setup_constraints(constraint,comp_grid,options.FL)

#add the mask*blurring filer sparse matrix as a transform domain matrix
push!(TD_OP,FWD_OP)
push!(set_Prop.AtA_offsets,[0])#convert(Vector{TI},0:bkl)) #these are dummy values, actual ofsetts are automatically detected
push!(set_Prop.ncvx,false)
push!(set_Prop.banded,true)
push!(set_Prop.AtA_diag,true)
push!(set_Prop.dense,false)
push!(set_Prop.tag,("bounds","subsampled x-motion-blur","matrix",""))

#also add a projector onto the data constraint: i.e. ||A*x-m||=< sigma, or l<=(A*x-m)<=u
data = vec(d_obs[1,:,:])
LBD=data.-15.0;  LBD=convert(Vector{TF},LBD)
UBD=data.+15.0;  UBD=convert(Vector{TF},UBD)
push!(P_sub,x -> project_bounds!(x,LBD,UBD))
println("finished setting up constraints")

dummy=zeros(TF,size(BF,2))
(TD_OP,AtA,l,y) = PARSDMM_precompute_distribute(TD_OP,set_Prop,comp_grid,options)
println("finished precomputing and distributing")

options.rho_ini      = ones(TF,length(TD_OP))*TF(1000.0)
for i=1:length(options.rho_ini)
  if set_Prop.ncvx[i]==true
    options.rho_ini[i] = TF(10.0)
  end
end

for i=1:size(d_obs,1)
  SNR(in1,in2)=20*log10(norm(in1)/norm(in1-in2))
  @time (x,log_PARSDMM) = PARSDMM(dummy,AtA,TD_OP,set_Prop,P_sub,comp_grid,options)
  m_est[i,:,:]=reshape(x,comp_grid.n)
  println("SNR:", round(SNR(vec(m_evaluation[i,(bkl*2):end-(bkl*2),:]),vec(m_est[i,(bkl*2):end-(bkl*2),:])),digits=2))
  println("PSNR:", round(psnr(vec(m_evaluation[i,(bkl*2):end-(bkl*2),:]),vec(m_est[i,(bkl*2):end-(bkl*2),:]),maximum(m_evaluation[i,:,:])),digits=2))
  if i+1<=size(d_obs,1)
    data = vec(d_obs[i+1,:,:])
    LBD=data.-15.0;  LBD=convert(Vector{TF},LBD);
    UBD=data.+15.0;  UBD=convert(Vector{TF},UBD);
    P_sub[end] = x -> project_bounds!(x,LBD,UBD)
  end
end

# for i=1:35
#   figure();title(string("training image", i), fontsize=10)
#   imshow(m_train[i,:,:],cmap="gray",vmin=0.0,vmax=255.0);axis("off") #title("training image", fontsize=10)
#   savefig(string("training_data_", i,".eps"),bbox_inches="tight",dpi=600)
#   savefig(string("training_data_", i,".eps"),bbox_inches="tight")
#   close()
# end
# close("all")

ENV["MPLBACKEND"]="qt5agg"
using PyPlot

#plot first 8 training in 1 figure
figure();
for i=1:8
  subplot(2,4,i);imshow(m_train[i,:,:],cmap="gray",vmin=0.0,vmax=255.0);axis("off") #title("training image", fontsize=10)
end
savefig("training_data_first8.eps",bbox_inches="tight",dpi=300)
close()

SNR(in1,in2)=20*log10(norm(in1)/norm(in1-in2))

file = matopen("m_est.mat", "w")
write(file, "m_est", convert(Array{Float64,3},m_est))
close(file)

file = matopen("m_evaluation.mat", "w")
write(file, "m_evaluation", convert(Array{Float64,3},m_evaluation))
close(file)

file = matopen("m_train.mat", "w")
write(file, "m_train", convert(Array{Float64,3},m_train))
close(file)

file = matopen("d_obs.mat", "w")
write(file, "d_obs", convert(Array{Float64,3},d_obs))
close(file)

file = matopen("FWD_OP.mat", "w")
write(file, "FWD_OP", convert(SparseMatrixCSC{Float64,Int64},FWD_OP))
close(file)

file = matopen("x_SPGL1_wavelet_save_SA.mat")
x_SPGL_wavelet_save=read(file, "x_SPGL1_wavelet_save_SA")
x_SPGL_wavelet_save = convert(Array{TF,3},x_SPGL_wavelet_save)

figure(figsize=(4.8, 4.8))
FS  = 5
LS  = 3
TML = 1
PD  = 1
for i=1:size(m_est,1)
  subplot(4,4,i);imshow(d_obs[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)],cmap="gray",vmin=0.0,vmax=255.0); title("observed",FontSize=FS,verticalalignment="bottom");tick_params(labelsize=LS,length=TML,pad=PD)
end
for i=1:size(m_est,1)
  subplot(4,4,i+8);imshow(m_est[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)],cmap="gray",vmin=0.0,vmax=255.0); title(FontSize=FS,string("PARSDMM, PSNR=", round(psnr(vec(m_evaluation[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)]),vec(m_est[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)]),maximum(m_evaluation[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)])),digits=2)));tick_params(labelsize=LS,length=TML,pad=PD)
end
for i=1:size(m_est,1)
  subplot(4,4,i+4);imshow(m_evaluation[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)],cmap="gray",vmin=0.0,vmax=255.0); title(FontSize=FS,"True");tick_params(labelsize=LS,length=TML,pad=PD)
end
for i=1:size(m_est,1)
  subplot(4,4,i+12);imshow(x_SPGL_wavelet_save[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)],cmap="gray",vmin=0.0,vmax=255.0); title(FontSize=FS,string("BPDN-wavelet, PSNR=", round(psnr(vec(m_evaluation[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)]),vec(x_SPGL_wavelet_save[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)]),maximum(m_evaluation[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)])),digits=2)));tick_params(labelsize=LS,length=TML,pad=PD)
end
tight_layout()
PyPlot.subplots_adjust(wspace=-0.05, hspace=0.25)
savefig("deblurring_inpainting_results.png",bbox_inches="tight",dpi=600)

#First 8 in 1 figure
figure(figsize=(6.4, 3.2))
for i=1:8
  subplot(2,4,i);imshow(m_train[i,:,:],cmap="gray",vmin=0.0,vmax=255.0);axis("off") #title("training image", fontsize=10)
end
tight_layout()
PyPlot.subplots_adjust(wspace=-0.15, hspace=0.1)
savefig("training_data_first8.png",bbox_inches="tight",dpi=600)
close()

for i=1:size(m_est,1)
    figure();imshow(d_obs[i,(bkl*2):end-(bkl*2),:],cmap="gray",vmin=0.0,vmax=255.0); title("observed");
    savefig(string("deblurring_inpainting_observed",i,".eps"),bbox_inches="tight",dpi=600)
    figure();imshow(m_est[i,(bkl*2):end-(bkl*2),:],cmap="gray",vmin=0.0,vmax=255.0); title(string("PARSDMM, SNR=", round(SNR(vec(m_evaluation[i,(bkl*2):end-(bkl*2),:]),vec(m_est[i,(bkl*2):end-(bkl*2),:])),digits=2)))
    savefig(string("PARSDMM_deblurring_inpainting",i,".eps"),bbox_inches="tight",dpi=600)
    figure();imshow(m_evaluation[i,(bkl*2):end-(bkl*2),:],cmap="gray",vmin=0.0,vmax=255.0); title("True")
    savefig(string("deblurring_inpainting_evaluation",i,".eps"),bbox_inches="tight",dpi=600)
    close("all")
end

#all results in one figure (PARSDMM + SPGL1)
# figure()
# for i=1:size(m_est,1)
#   subplot(4,4,i);imshow(d_obs[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)],cmap="gray",vmin=0.0,vmax=255.0); title("observed");
# end
# for i=1:size(m_est,1)
#   subplot(4,4,i+8);imshow(m_est[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)],cmap="gray",vmin=0.0,vmax=255.0); title(string("PARSDMM, SNR=", round(SNR(vec(m_evaluation[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)]),vec(m_est[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)])),digits=2)))
# end
# for i=1:size(m_est,1)
#   subplot(4,4,i+4);imshow(m_evaluation[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)],cmap="gray",vmin=0.0,vmax=255.0); title("True")
# end
# for i=1:size(m_est,1)
#   subplot(4,4,i+12);imshow(x_SPGL_wavelet_save[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)],cmap="gray",vmin=0.0,vmax=255.0); title(string("BPDN-wavelet, SNR=", round(SNR(vec(m_evaluation[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)]),vec(x_SPGL_wavelet_save[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)])),digits=2)))
# end

#plot training images
figure();title("training image", fontsize=10)
for i=1:16
  subplot(4,4,i);imshow(m_train[i,:,:],cmap="gray",vmin=0.0,vmax=255.0);axis("off") #title("training image", fontsize=10)
end
savefig("training_data_all.eps",bbox_inches="tight",dpi=600)
close()

SNR(in1,in2)=20*log10(norm(in1)/norm(in1-in2))
for i=1:size(m_evaluation,1)
  figure()
  imshow(x_SPGL_wavelet_save[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)],cmap="gray",vmin=0.0,vmax=255.0); title(string("SPGL1 BPDN-wavelet, SNR=", round(SNR(vec(m_evaluation[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)]),vec(x_SPGL_wavelet_save[i,(bkl*2):end-(bkl*2),(bkl*2):end-(bkl*2)])),digits=2)))
  savefig(string("SPGL1_wavelet_inpainting",i,".eps"),bbox_inches="tight",dpi=600)
  close()
end

max_stars_repo_name: slimgroup/SetIntersectionProjection

max_issues_repo_name: slimgroup/SetIntersectionProjection

max_forks_repo_name: slimgroup/SetIntersectionProjection


text:
# *****************************************************************************
# Written by Ritchie Lee, ritchie.lee@sv.cmu.edu
# *****************************************************************************
# Copyright ã ``2015, United States Government, as represented by the
# Administrator of the National Aeronautics and Space Administration. All
# rights reserved.  The Reinforcement Learning Encounter Simulator (RLES)
# platform is licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. You
# may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0. Unless required by applicable
# law or agreed to in writing, software distributed under the License is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the specific language
# governing permissions and limitations under the License.
# _____________________________________________________________________________
# Reinforcement Learning Encounter Simulator (RLES) includes the following
# third party software. The SISLES.jl package is licensed under the MIT Expat
# License: Copyright (c) 2014: Youngjun Kim.
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED
# "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
# NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
# *****************************************************************************

"""
Depth aware Probabilistic Context-Free Grammar
Wraps a PCFG and adds depth information
"""
module DepthAwarePCFGs

export DepthAwarePCFG, fit_mle!, weighted_sum!

using GrammaticalEvolution
using ExprSearch, LinearDerivTrees, PCFGs, MinDepths
using RLESUtils, StatUtils

import PCFGs: fit_mle!, weighted_sum!
import Base: rand, rand!, normalize!, fill!, copy

typealias CFG Grammar
typealias ProbDict Dict{Symbol,Vector{Float64}}

type DepthAwarePCFG
    pcfg::PCFG
    depths_by_rule::MinDepthByRule
    depths_by_action::MinDepthByAction
end

"""
DepthAwarePCFG constructor
"""
function DepthAwarePCFG(cfg::CFG)
    DepthAwarePCFG(
        PCFG(cfg),
        min_depth_rule(cfg),
        min_depth_actions(cfg)
        )
end

"""
Draw a random sample from depth-aware pcfg.  The output is in-place
into tree.
"""
function rand!(tree::LinearDerivTree, dpcfg::DepthAwarePCFG, maxdepth::Int64)
    rand!(Base.GLOBAL_RNG, tree, dpcfg, maxdepth)
end
function rand!(rng::AbstractRNG, tree::LinearDerivTree, dpcfg::DepthAwarePCFG,
    maxdepth::Int64)
    probs = dpcfg.pcfg.probs
    initialize!(tree)
    while !isdone(tree)
        actions = actionspace(tree)
        node = current_opennode(tree)
        depths = dpcfg.depths_by_action[Symbol(node.rule.name)]
        w = probs[get_sym(tree)]
        ids = find(x->x <= maxdepth-node.depth, depths)
        @assert !isempty(ids) #shouldn't be empty if started sampling from root
        aid = weighted_rand(rng, ids, w[ids])
        step!(tree, actions[aid])
    end
    iscomplete(tree)
end

"""
Draw N random samples from depth-aware pcfg.  Derivation trees are limited to depth of maxdepth.
Defaults to unlimited number of steps (non-terminal expansions)
Outputs a vector of LinearDerivTrees.
"""
function rand(dpcfg::DepthAwarePCFG, N::Int64, maxdepth::Int64)
    rand(Base.GLOBAL_RNG, dpcfg, N, maxdepth)
end
function rand(rng::AbstractRNG, dpcfg::DepthAwarePCFG, N::Int64, maxdepth::Int64;
    maxsteps::Int64=typemax(Int64))
    params = LDTParams(dpcfg.pcfg.cfg, maxsteps)
    samples = Array(LinearDerivTree, N)
    for i = 1:N
        s = LinearDerivTree(params)
        rand!(rng, s, dpcfg, maxdepth)
        samples[i] = s
    end
    samples
end

"""
Draw random samples according to pcfg and output them in-place into samples.
"""
function rand!(samples::Vector{LinearDerivTree}, dpcfg::DepthAwarePCFG, maxdepth::Int64)
    rand!(Base.GLOBAL_RNG, samples, dpcfg, maxdepth)
end
function rand!(rng::AbstractRNG, samples::Vector{LinearDerivTree}, dpcfg::DepthAwarePCFG,
    maxdepth::Int64)
    for i = 1:length(samples)
        rand!(rng, samples[i], dpcfg, maxdepth)
    end
end

"""
Compute the maximum likelihood estimator from all the transitions in the
trees in samples and store the output in-place into pcfg.
"""
function fit_mle!(dpcfg::DepthAwarePCFG, samples::Vector{LinearDerivTree})
    #passthrough
    fit_mle!(dpcfg.pcfg, samples)
end

"""
Computes the weighted sum of two pcfgs and returns the output in-place
pcfg = w1 * pcfg + w2 * pcfg2
"""
function weighted_sum!(dpcfg::DepthAwarePCFG, w1::Float64, dpcfg2::DepthAwarePCFG, w2::Float64)
    #passthrough
    weighted_sum!(dpcfg.pcfg, w1, dpcfg2.pcfg, w2)
end

"""
Fill all probabilities in pcfg to value x
"""
function fill!(dpcfg::DepthAwarePCFG, x)
    #passthrough
    fill!(dpcfg.pcfg, x)
end

"""
Normalize the sum of each probability vector in pcfg
"""
normalize!(dpcfg::DepthAwarePCFG) = normalize!(dpcfg.pcfg)

function copy(dpcfg::DepthAwarePCFG)
    DepthAwarePCFG(copy(dpcfg.pcfg), dpcfg.depths_by_rule, dpcfg.depths_by_action)
end

end #module

ext: jl

max_stars_repo_name: sisl/ExprSearch

max_issues_repo_name: sisl/ExprSearch


text:
using JLD2, PyPlot, JUDI.TimeModeling

PyPlot.rc("font", family="serif")
PyPlot.rc("xtick", labelsize=20)
PyPlot.rc("ytick", labelsize=20)

@load "./data/bg_tti_g.jld"

n = model_true.n
d = model_true.d
extent = [0, (n[1] -1) * d[1],  (n[2] -2) * d[2], 0]

### Plot model
figure(figsize=(30, 30))
subplot(221)
imshow(model_true.m', vmin=1/4.5^2, vmax=.44, cmap="jet", extent=extent, aspect="auto")
ax1 = gca()
ax1.set_xticks([])
ylabel("Depth (m)", fontsize=20)
title("m", fontsize=20)
colorbar(ax=ax1)
subplot(222)
imshow(model_true.epsilon', vmin=0.0, vmax=.3, cmap="jet", extent=extent, aspect="auto")
ax1 = gca()
ax1.set_xticks([])
ax1.set_yticks([])
title("ε", fontsize=20)
colorbar(ax=ax1)
subplot(223)
imshow(model_true.delta', vmin=0.0, vmax=.3, cmap="jet", extent=extent, aspect="auto")
ax1 = gca()
xlabel("X (m)", fontsize=20)
ylabel("Depth (m)", fontsize=20)
title("δ", fontsize=20)
colorbar(ax=ax1)
subplot(224)
imshow(model_true.theta', vmin=-.5, vmax=.5, cmap="jet", extent=extent, aspect="auto")
xlabel("X (m)", fontsize=20)
ax1 = gca()
ax1.set_yticks([])
title("θ", fontsize=20)
colorbar(ax=ax1)

savefig("./data/bg_model.png", bbox_inches="tight")

# Plot model perturbations
# True water layer
@load "./data/bg_tti_g.jld"
plt.figure(figsize=(30, 30))
subplot(221)
imshow(model0_tti.m', vmin=1/4.5^2, vmax=.44, cmap="jet", extent=extent, aspect="auto")
ylabel("Depth (m)", fontsize=20)
ax1 = gca()
ax1.set_xticks([])
title("m0", fontsize=20)
ax1 = gca()
colorbar(ax=ax1)
subplot(223)
imshow(model0_tti.m' - model_true.m', vmin=-.05, vmax=.05, cmap="jet", extent=extent, aspect="auto")
xlabel("X (m)", fontsize=20)
ylabel("Depth (m)", fontsize=20)
ax1 = gca()
title("δ m", fontsize=20)
colorbar(ax=ax1)
# Wrong water layer
@load "./data/bg_tti_g_w.jld"
subplot(222)
imshow(model0_tti.m', vmin=1/4.5^2, vmax=.44, cmap="jet", extent=extent, aspect="auto")
ax1 = gca()
ax1.set_yticks([])
ax1.set_xticks([])
title("m0", fontsize=20)
colorbar(ax=ax1)
subplot(224)
imshow(model0_tti.m' - model_true.m', vmin=-.05, vmax=.05, cmap="jet", extent=extent, aspect="auto")
xlabel("X (m)", fontsize=20)
ax1 = gca()
ax1.set_yticks([])
title("δ m", fontsize=20)
colorbar(ax=ax1)

savefig("./data/bg_dm.png", bbox_inches="tight")

### gradients with true water
@load "./data/bg_tti_g.jld"

# true tti params
figure(figsize=(30, 30))
subplot(321)
imshow(g_wri_tti', vmin=-.1, vmax=.1, cmap="jet", extent=extent, aspect="auto")
ylabel("Depth (m)", fontsize=20)
ax1 = gca()
ax1.set_xticks([])
title("WRI true anisotropy", fontsize=20)
colorbar(ax=ax1)
subplot(322)
imshow(g_fwi_tti', vmin=-.1, vmax=.1, cmap="jet", extent=extent, aspect="auto")
ax1 = gca()
ax1.set_xticks([])
ax1.set_yticks([])
title("FWI true anisotropy", fontsize=20)
colorbar(ax=ax1)

# Error in tti params
@load "./data/bg_tti_g_w.jld"
subplot(323)
imshow(g_wri_tti_w', vmin=-.05, vmax=.05, cmap="jet", extent=extent, aspect="auto")
ylabel("Depth (m)", fontsize=20)
ax1 = gca()
ax1.set_xticks([])
title("WRI inaccurate anisotropy", fontsize=20)
colorbar(ax=ax1)
subplot(324)
imshow(g_fwi_tti_w', vmin=-.1, vmax=.1, cmap="jet", extent=extent, aspect="auto")
ax1 = gca()
ax1.set_xticks([])
ax1.set_yticks([])
title("FWI inaccurate anisotropy", fontsize=20)
colorbar(ax=ax1)

# Acoustic
@load "./data/bg_acou_g.jld"
subplot(325)
imshow(g_wri_a', vmin=-.01, vmax=.01, cmap="jet", extent=extent, aspect="auto")
xlabel("X (m)", fontsize=20)
ylabel("Depth (m)", fontsize=20)
ax1 = gca()
title("WRI acoustic", fontsize=20)
colorbar(ax=ax1)
subplot(326)
imshow(g_fwi_a', vmin=-.1, vmax=.1, cmap="jet", extent=extent, aspect="auto")
xlabel("X (m)", fontsize=20)
ax1 = gca()
ax1.set_yticks([])
title("FWI acoustic", fontsize=20)
colorbar(ax=ax1)

savefig("./data/bg_true_water.png", bbox_inches="tight")

# Wrong water layer
@load "./data/bg_tti_g_ww.jld"

# true tti params
figure(figsize=(30, 30))
subplot(321)
imshow(g_wri_tti_ww', vmin=-.1, vmax=.1, cmap="jet", extent=extent, aspect="auto")
ylabel("Depth (m)", fontsize=20)
ax1 = gca()
ax1.set_xticks([])
title("WRI true anisotropy", fontsize=20)
colorbar(ax=ax1)
subplot(322)
imshow(g_fwi_tti_ww', vmin=-.1, vmax=.1, cmap="jet", extent=extent, aspect="auto")
ax1 = gca()
ax1.set_xticks([])
ax1.set_yticks([])
title("FWI true anisotropy", fontsize=20)
colorbar(ax=ax1)

# Error in tti params
@load "./data/bg_tti_g_ww_w.jld"
subplot(323)
imshow(g_wri_tti_ww_w', vmin=-.05, vmax=.05, cmap="jet", extent=extent, aspect="auto")
ylabel("Depth (m)", fontsize=20)
ax1 = gca()
ax1.set_xticks([])
title("WRI inaccurate anisotropy", fontsize=20)
colorbar(ax=ax1)
subplot(324)
imshow(g_fwi_tti_ww_w', vmin=-.1, vmax=.1, cmap="jet", extent=extent, aspect="auto")
ax1 = gca()
ax1.set_xticks([])
ax1.set_yticks([])
title("FWI inaccurate anisotropy", fontsize=20)
colorbar(ax=ax1)

# Acoustic
@load "./data/bg_acou_g_ww.jld"
subplot(325)
imshow(g_wri_a_ww', vmin=-.01, vmax=.01, cmap="jet", extent=extent, aspect="auto")
xlabel("X (m)", fontsize=20)
ylabel("Depth (m)", fontsize=20)
ax1 = gca()
title("WRI acoustic", fontsize=20)
colorbar(ax=ax1)
subplot(326)
imshow(g_fwi_a_ww', vmin=-.1, vmax=.1, cmap="jet", extent=extent, aspect="auto")
xlabel("X (m)", fontsize=20)
ax1 = gca()
ax1.set_yticks([])
title("FWI acoustic", fontsize=20)
colorbar(ax=ax1)

savefig("./data/bg_wrong_water.png", bbox_inches="tight")


text:
using DataStructures

randint(μ, σ) = max(σ, round(Int, randn()σ + μ))

"""
Generate testdata that NIST sclite can read.  Run sclite like:
```
sclite -r ref.txt -h hyp.txt -i rm -o pralign
```
"""
function gentestdata(N::Integer; nw=100, nsub=10, nins=10, ndel=10, V=25, σ=10)
	srand(N)
	reffd = open("test/ref.txt", "w")
	hypfd = open("test/hyp.txt", "w")
	res = DefaultOrderedDict{AbstractString, Vector}(Vector)
	for i in 1:N
		ref = rand(1:V, randint(nw, σ))
		hyp = ref[:]
		for j in 1:randint(nsub, σ)
			hyp[rand(1:length(hyp))] = rand(1:V)
		end
		for j in 1:randint(nins, σ)
			insert!(hyp, rand(1:length(hyp)), rand(1:V))
		end
		for j in 1:randint(ndel, σ)
			deleteat!(hyp, rand(1:length(hyp)))
		end
		id = "a-$i"
		println(reffd, join([string(x) for x in ref], " "), " ($id)")
		println(hypfd, join([string(x) for x in hyp], " "), " ($id)")
		push!(res[id], [string(x) for x in ref])
		push!(res[id], [string(x) for x in hyp])
	end
	close(reffd)
	close(hypfd)
	return res
end

function readtestdata()
	res = DefaultOrderedDict{AbstractString, Vector}(Vector)
	for file in ["ref.txt", "hyp.txt"]
		open(file) do fd
			for line in eachline(fd)
				words = line |> chomp |> split
				id = strip(pop!(words), ['(', ')'])
				push!(res[id], words)
			end
		end
	end
	return res
end

function readsclitepralign(file="hyp.txt.pra")
	res = OrderedDict()
	open(file) do fd
		while !eof(fd)
			m = match(r"id: \(([^()]+)\)", readline(fd))
			if m != nothing
				id = m[1]
				res[id] = join([readline(fd) for i in 1:4], "\n")
			end
		end
	end
	return res
end


text:
using DelayDiffEq, DiffEqProblemLibrary, Test
using DiffEqProblemLibrary.DDEProblemLibrary: importddeproblems; importddeproblems()
import DiffEqProblemLibrary.DDEProblemLibrary: prob_dde_1delay, prob_dde_1delay_notinplace,
       prob_dde_1delay_scalar_notinplace,
       prob_dde_2delays, prob_dde_2delays_notinplace, prob_dde_2delays_scalar_notinplace,
       prob_dde_1delay_long, prob_dde_1delay_long_notinplace,
       prob_dde_1delay_long_scalar_notinplace, prob_dde_2delays_long,
       prob_dde_2delays_long_notinplace, prob_dde_2delays_long_scalar_notinplace,
       prob_dde_mackey


text:
module DerivativeTestUtils

using Celeste: Model, DeterministicVI, SensitiveFloats
using Base.Test
import ForwardDiff

export forward_diff_model_params, unwrap_vp_vector, wrap_vp_vector,
       test_with_autodiff, trim_tiles!

""""
Return an ElboArgs with the corresponding type.
"""
function forward_diff_model_params{T<:Number}(::Type{T}, ea0::ElboArgs{Float64})
    P = length(ea0.vp[1])
    vp = Vector{T}[zeros(T, P) for s=1:ea0.S]
    # Set the values (but not gradient numbers) for parameters other
    # than the galaxy parameters.
    for s=1:ea0.S, i=1:length(ids)
        vp[s][i] = ea0.vp[s][i]
    end

    ElboArgs(ea0.images,
             vp,
             ea0.tile_source_map,
             ea0.patches,
             ea0.active_sources)
end

"""
Wrap a vector of canonical parameters for active sources into a ElboArgs
object of the appropriate type. Used for testing with forward
autodifferentiation.
"""
function unwrap_vp_vector{NumType <: Number}(
            vp_vec::Vector{NumType}, ea::ElboArgs)
    vp_array = reshape(vp_vec, length(CanonicalParams), length(ea.active_sources))
    ea_local = forward_diff_model_params(NumType, ea)
    for sa = 1:length(ea.active_sources)
        ea_local.vp[ea.active_sources[sa]] = vp_array[:, sa]
    end
    ea_local
end

"""
Convert the variational params into a vector for autodiff.
"""
function wrap_vp_vector(ea::ElboArgs, use_active_sources::Bool)
    P = length(CanonicalParams)
    S = use_active_sources ? length(ea.active_sources) : ea.S
    x_mat = zeros(Float64, P, S)
    for s in 1:S
        ind = use_active_sources ? ea.active_sources[s] : s
        x_mat[:, s] = ea.vp[ind]
    end
    x_mat[:]
end

"""
Use ForwardDiff to test that fun(x) = sf (to abuse some notation)
"""
function test_with_autodiff(fun::Function, x::Vector{Float64}, sf::SensitiveFloat)

    const test_grad = false
    const test_hess = false

    @test_approx_eq fun(x) sf.v

    if test_grad
        ad_grad = ForwardDiff.gradient(fun, x)
        @test_approx_eq ad_grad sf.d[:]
    end

    if test_hess
        ad_hess = ForwardDiff.hessian(fun, x)
        @test_approx_eq ad_hess sf.h
    end
end

"""
Set all but a few pixels to NaN to speed up autodiff Hessian testing.
"""
function trim_tiles!(tiled_blob::Vector{TiledImage}, keep_pixels)
    for b = 1:length(tiled_blob)
	    pixels1 = tiled_blob[b].tiles[1,1].pixels
        h_width, w_width = size(pixels1)
	    pixels1[setdiff(1:h_width, keep_pixels), :] = NaN
        pixels1[:, setdiff(1:w_width, keep_pixels)] = NaN
	end
end

end # end module


text:
module CoolPkg

export add2

"""
    add2(a, b)

Adds two numbers together...

* `a`: A number
* `b`: Another number
"""
add2(a, b) = a + b

end


text:
"""
	mutable struct ArrayEntry <: JSONEntry
		items
		l::Dict{Int,Int}
		updated::Int
	end

keeps statistics about an array entry in JSON.
- `items` is typeof `Entry` or nothing and keeps statistics about the elements of the array
- `l` keeps histogram of message length
- `updated` counts how many times the struct was updated.
"""
mutable struct ArrayEntry <: JSONEntry
	items
	l::Dict{Int,Int}
	updated::Int
end

ArrayEntry(items) = ArrayEntry(items,Dict{Int,Int}(),0)
Base.isempty(e::ArrayEntry) = e.items isa ArrayEntry ? isempty(e.items) : isnothing(e.items)

function update!(a::ArrayEntry, b::Vector; path = "")
	n = length(b)
	a.updated += 1
	a.l[n] = get(a.l,n,0) + 1
	if isnothing(a.items)
		 a.items = newentry(b).items
	end

	for (i, v) in enumerate(b)
		a.items = safe_update!(a.items,v,path="$path[$i]")
	end
	return true
end

# todo:
#   logiku jestli to je missing nebo ne přesunout z sample_synthetic do suggestextractor
#   pokud vrátím správné typy u full samplu, ověřit, jestli potřebuju empty sample nebo se to odvodí z plného
function suggestextractor(node::ArrayEntry, settings = NamedTuple(); path = "", child_less_than_parent = false)
	if isempty(node)
		@warn "$(path) is an empty array, therefore I can not suggest extractor."
		return nothing
	end

	if length(node.l) == 1 && typeof(node.items) <: Entry && promote_type(unique(typeof.(keys(node.items.counts)))...) <: Number
		@info "$(path) is an array of numbers with of same length, therefore it will be treated as a vector."
		return ExtractVector(only(collect(keys(node.l))))
	end
	e = suggestextractor(node.items, settings, path = path, child_less_than_parent = child_less_than_parent)
	isnothing(e) ? e : ExtractArray(e)
end

function merge(es::ArrayEntry...)
	updates_merged = sum(updated.(es))
	l_merged = merge(+, l.(es)...)
	nonempty_items = items.(filter(!isempty, collect(es)))
	items_merged = isempty(nonempty_items) ? nothing : merge(merge, nonempty_items...)
	ArrayEntry(items_merged, l_merged, updates_merged)
end

l(s::T) where {T<:ArrayEntry} = s.l
items(s::T) where {T<:ArrayEntry} = s.items
Base.hash(e::ArrayEntry, h::UInt) = hash((e.items, e.l, e.updated), h)
Base.:(==)(e1::ArrayEntry, e2::ArrayEntry) = e1.updated === e2.updated && e1.l == e2.l && e1.items == e2.items
sample_synthetic(e::ArrayEntry) = isnothing(e.items) ? [] : repeat([sample_synthetic(e.items)], 2)


text:
function maximize_elbo(rng, logp, dists, ndraws)
    elbo_ϕ_logqϕ = map(dists) do dist
        elbo_and_samples(rng, logp, dist, ndraws)
    end
    lopt = argmax(first.(elbo_ϕ_logqϕ))
    return (lopt, elbo_ϕ_logqϕ[lopt]...)
end

function elbo_and_samples(rng, logp, dist, ndraws)
    ϕ, logqϕ = rand_and_logpdf(rng, dist, ndraws)
    logpϕ = logp.(eachcol(ϕ))
    elbo = elbo_from_logpdfs(logpϕ, logqϕ)
    return elbo, ϕ, logqϕ
end

elbo_from_logpdfs(logpϕ, logqϕ) = Statistics.mean(logpϕ) - Statistics.mean(logqϕ)


text:
module FrameFunInterfaceExtension

using ..FrameFunInterface, ..Platforms, BasisFunctions, DomainSets, ..ExtensionFrames
import ..FrameFunInterface: oversampling_grid, azdual_dict, correct_sampling_parameter,
    approximationproblem, deduce_samplingparameter, solver, discretemeasure
using ..ApproximationProblems: setsamplingparam!
using ..FrameFunInterface: deduce_oversampling_parameter, SamplingStrategy, DefaultSamplingStrategy
using BasisFunctions: AbstractMeasure
approximationproblem(dict::Dictionary, domain::Domain) =
    approximationproblem(promote_type(coefficienttype(dict),prectype(eltype(domain))), dict, domain)

approximationproblem(::Type{T}, dict::Dictionary) where {T} =
    approximationproblem(BasisFunctions.ensure_coefficienttype(T, dict))

# If a dictionary and a domain is specified, we make an extension frame.
function approximationproblem(::Type{T}, dict::Dictionary, domain::Domain) where {T}
    if domain == support(dict)
        approximationproblem(T, dict)
    else
        approximationproblem(T, extensionframe(domain, BasisFunctions.ensure_coefficienttype(T, dict)))
    end
end

include("tridiagonalsolver.jl")
solver(::TridiagonalProlateStyle, ap, A; scaling = Zt_scaling_factor(dictionary(ap), A)
    , options...) =
    FE_TridiagonalSolver(A, scaling; options...)
# They have to do with a normalization of the
# sampling operators.
Zt_scaling_factor(S::Dictionary, A) = length(supergrid(grid(dest(A))))
Zt_scaling_factor(S::DerivedDict, A) = Zt_scaling_factor(superdict(S), A)
Zt_scaling_factor(S::ChebyshevT, A) = length(supergrid(grid(dest(A))))/2

using ..ExtensionFrames, ..ExtensionFramePlatforms
oversampling_grid(dict::ExtensionFrame, L) = subgrid(oversampling_grid(superdict(dict),L), support(dict))

discretemeasure(ss::SamplingStyle, platform::ExtensionFramePlatform, param, ap_old; options...) =
    restrict(discretemeasure(ss, platform.basisplatform, param,
        (ap=approximationproblem(platform.basisplatform, param);setsamplingparam!(ap,samplingparameter(ap_old));ap)
            ; options...), platform.domain)

azdual_dict(sstyle::SamplingStyle, platform::ExtensionFramePlatform, param, L, measure::AbstractMeasure; options...) =
   extensionframe(azdual_dict(sstyle, platform.basisplatform, param, L, supermeasure(measure); options...), platform.domain)

correct_sampling_parameter(samplingstrategy::SamplingStrategy, platform::ExtensionFramePlatform, param, L_trial; options...) =
   correct_sampling_parameter(samplingstrategy, platform.basisplatform, param, L_trial; options...)

using ..AugmentationPlatforms
azdual_dict(sstyle::SamplingStyle, platform::AugmentationPlatform, param, L, measure::AbstractMeasure; options...) =
   MultiDict([azdual_dict(sstyle, platform.basis, param, L, measure; options...), platform.functions])

oversampling_grid(samplingstyle::SamplingStyle, platform::AugmentationPlatform, param, L; options...) =
   oversampling_grid(samplingstyle, platform.basis, param, L; options...)

correct_sampling_parameter(strategy::SamplingStrategy, platform::AugmentationPlatform, param, L; options...) =
    correct_sampling_parameter(strategy, platform.basis, param, L; options...)

using ..WeightedSumPlatforms
function azdual_dict(sstyle::SamplingStyle, platform::WeightedSumPlatform, param, L, measure::AbstractMeasure; options...)
    denom = (x...)->sum(map(w->abs(w(x...))^2, platform.weights))
    # TODO: discuss, what is the relation between param, L of a platform and platform.P
    MultiDict([((x...)->(platform.weights[j](x...)/denom(x...))) * azdual_dict(sstyle, platform.P, param, L, measure; options...) for j=1:length(platform.weights)])
end

oversampling_grid(samplingstyle::SamplingStyle, platform::WeightedSumPlatform, param::NTuple, L; options...) =
  oversampling_grid(samplingstyle, platform.P, param[1], L; options...)

azdual_dict(sstyle::SamplingStyle, dict::ExtensionFrame, L, measure::AbstractMeasure; options...) =
   extensionframe(support(dict), azdual_dict(sstyle, superdict(dict), L, supermeasure(measure); options...),)

function azdual_dict(dict::MultiDict, measure::AbstractMeasure; options...)
    dictionary = dict.dicts[1].superdict
    weights = map(weightfunction, elements(dict))
    denom = (x...)->sum(map(w->abs(w(x...))^2, weights))
    MultiDict([((x...)->(weights[j](x...)/denom(x...))) * azdual_dict(dictionary, measure; options...) for j=1:length(weights)])
end

function deduce_samplingparameter(ss::OversamplingStyle, strategy::SamplingStrategy, platform::WeightedSumPlatform, param::NTuple; options...)
    mix_samplingparameters(strategy, platform, map(parami->FrameFunInterface.deduce_oversampling_parameter(ss, strategy, platform.P, parami; options...), param))
end

mix_samplingparameters(::DefaultSamplingStrategy, platform::WeightedSumPlatform, Ls::NTuple) =
    round.(Int,(sum(prod.(Ls))/prod(Ls[1]))^(1/length(Ls[1]))     .*Ls[1])

correct_sampling_parameter(s::SamplingStrategy, platform::WeightedSumPlatform, param, L; options...) =
    correct_sampling_parameter(s, platform.P, param, L; options...)

end


text:
using Qaintessent
using LinearAlgebra: I
using SparseArrays: sparse
using Memoize

# Pauli X and Y matrices
X = [0 1; 1 0]
Y = [0 -im; im 0]

# Utility function for a XY mixer
# Implements X_a X_{a+1} + Y_a Y_{a+1}, or more generally (⊗_{i ∈ xy_indices} X_i) + (⊗_{i ∈ xy_indices} Y_i)
function XY_sum(xy_indices::Vector{Int64}, d::Int64)::Matrix{ComplexF64}
    # passing generator into kron via varargs syntax
    return kron((i ∈ xy_indices ? X : I(2) for i ∈ 1:d)...
        ) + kron((i ∈ xy_indices ? Y : I(2) for i ∈ 1:d)...)
end

"""
    r-nearby values single-qudit mixer gate, which acts on a single qudit (but implemented
    here not for qudits, but the one-hot encoding)

``U_{r\\text{-NV}}(\\beta) = e^{-i \\beta H_{r\\text{-NV}}}``
``H_{r\\text{-NV}} = \\sum_{i=1}^r \\left(\\breve{X}^i + \\left(\\breve{X}^\\dagger\\right)^i\\right)``

Reference:\n
    Stuart Hadfield, Zhihui Wang, Bryan O'Gorman, Eleanor G. Rieffel,Davide Venturelli and Rupak Biswas\n
    From the Quantum Approximate Optimization Algorithm to a Quantum Alternating Operator Ansatz\n
    Algorithms 12.2 (2019), p.34
"""
struct RNearbyValuesMixerGate <: AbstractGate
    # use a reference type (array with 1 entry) for compatibility with Flux
    β::Vector{Float64}
    r::Integer # this is the r in r-Nearby-values
    d::Integer # d (= κ) = number of colors

    function RNearbyValuesMixerGate(β::Float64, r::Integer, d::Integer)
        (r > 0 && d > 0) || throw(ArgumentError("Parameters r and d must be positive integers."))
        (r <= d-1) || throw(ArgumentError("Parameter d must be between 1 and d-1."))
        new([β], r, d)
    end
end

# Compute the r-NV mixer Hamiltonian
@memoize function r_nearby_values_hamiltonian_onehot(g::RNearbyValuesMixerGate)
    H_rNV = sum(XY_sum([a, ((a + j - 1) % g.d) + 1], g.d)
        for a in 1:g.d for j in 1:g.r)

    return H_rNV
end

# Compute the r-NV mixer gate matrix: Implementation of Eq. (6)
function Qaintessent.matrix(g::RNearbyValuesMixerGate)
    H_rNV = r_nearby_values_hamiltonian_onehot(g)
    U_rNV = exp(-im * g.β[] * H_rNV)

    return U_rNV
end

# Adjoint of r-NV = r-NV gate with negated β parameter
Qaintessent.adjoint(g::RNearbyValuesMixerGate) = RNearbyValuesMixerGate(-g.β[], g.r, g.d)

Qaintessent.sparse_matrix(g::RNearbyValuesMixerGate) = sparse(matrix(g))

# Number of wires (= g.d since we use the one-hot encoding canonically)
Qaintessent.num_wires(g::RNearbyValuesMixerGate)::Int = g.d

"""
    Parity single-qudit ring mixer gate, which acts on a single qudit (but implemented
    here not for qudits, but the one-hot encoding)

``U_{\\text{parity}}(\\beta) = U_{\\text{last}}(\\beta) U_{\\text{even}}(\\beta) U_{\\text{odd}}(\\beta)``
``U_{\\text{odd}}(\\beta) = \\prod_{a~\\text{odd}, a \\neq d} e^{-i \\beta (X_a X_{a+1} + Y_a Y_{a+1})}``
``U_{\\text{even}}(\\beta) = \\prod_{a~\\text{even}} e^{-i \\beta (X_a X_{a+1} + Y_a Y_{a+1})}``
``U_{\\text{last}}(\\beta) = e^{-i \\beta (X_d X_1 + Y_d Y_1)} ~\\text{if}~ d ~\\text{is odd,}~ I ~\\text{otherwise.}``

The formulas require some interpretation. Assumptions for this implementation are:
- the indices a start at 1 (not at zero like elsewhere in the paper)
- X_{a+1} and Y_{a+1} actually means X_1 and Y_1 if a = d

Reference:\n
    Stuart Hadfield, Zhihui Wang, Bryan O'Gorman, Eleanor G. Rieffel, Davide Venturelli and Rupak Biswas\n
    From the Quantum Approximate Optimization Algorithm to a Quantum Alternating Operator Ansatz\n
    Algorithms 12.2 (2019), equations (7) - (10), p. 11
"""
struct ParityRingMixerGate <: AbstractGate
    β::Vector{Float64}
    d::Integer # d (= κ) = number of colors
    is_adjoint::Bool # if yes, the order in the product is reversed

    function ParityRingMixerGate(β::Float64, d::Integer; is_adjoint::Bool = false)
        d > 0 || throw(ArgumentError("Parameter d must be a positive integer."))
        new([β], d, is_adjoint)
    end
end

# Compute the parity ring mixer gate matrix: Implements Eq. (8)
function Qaintessent.matrix(g::ParityRingMixerGate)
    # assumption: by a ≠ n, the paper actually means a ≠ d.
    # assumption: by X_a for a = d+1, the paper means X_1.
    U_odd = prod([exp(-im * g.β[] * XY_sum([a, a+1], g.d)) for a ∈ 1:2:(g.d - 1)], init=I)
    U_even = prod([exp(-im * g.β[] * XY_sum([a, a < g.d ? (a+1) : 1], g.d)) for a ∈ 2:2:g.d], init=I)

    # Implements Eq. (9)
    U_last =  isodd(g.d) ? exp(-im * g.β[] * XY_sum([g.d, 1], g.d)) : I

    # Implements Eq. (7) (U_parity)
    if !g.is_adjoint
        return U_last * U_even * U_odd
    else
        return U_odd * U_even * U_last
    end
end

# Adjoint of parity ring mixer -> negated β parameter, reversed order in the product
Qaintessent.adjoint(g::ParityRingMixerGate) = ParityRingMixerGate(-g.β[], g.d, is_adjoint = !g.is_adjoint)

Qaintessent.sparse_matrix(g::ParityRingMixerGate) = sparse(matrix(g))

# Number of wires (= g.d since we use the one-hot encoding canonically)
Qaintessent.num_wires(g::ParityRingMixerGate)::Int = g.d

"""
    Partition single-qudit mixer gate (implemented not for qudits, but the one-hot encoding)

``U_{\\mathcal{P}-r-\\text{NV}}(\\beta) = U_{P_p-\\text{XY}}(\\beta) \\dots U_{P_1-\\text{XY}}(\\beta)``
``U_{P-\\text{XY}}(\\beta) = \\prod_{\\{a, b\\} \\in P} e^{-i \\beta (\\ket{a}\\bra{b} + \\ket{b}\\bra{a})}``

Note: the time evolution term is effectively implemented with an additional factor of two in the exponent:
``e^{-i \\beta 2 (\\ket{a}\\bra{b} + \\ket{b}\\bra{a})}``
to be consistent with the XY gates used for the other mixers.

Reference:\n
    Stuart Hadfield, Zhihui Wang, Bryan O'Gorman, Eleanor G. Rieffel, Davide Venturelli and Rupak Biswas\n
    From the Quantum Approximate Optimization Algorithm to a Quantum Alternating Operator Ansatz\n
    Algorithms 12.2 (2019), equations (11) - (12), p. 12
"""
struct PartitionMixerGate <: AbstractGate
    β::Vector{Float64}
    d::Int64
    partition::Vector{Vector{Tuple{Int, Int}}} # the Tuple stops Flux.@functor from misinterpreting these as params

    function PartitionMixerGate(β::Float64, d::Int64, partition::Vector{Vector{Tuple{Int, Int}}})
        d > 0 || throw(ArgumentError("Parameter d must be a positive integer."))

        # check that no duplicate indices occur in a part (s.t. the XY mixers within one part commute)
        for partition_part ∈ partition
            part_indices = union(reduce(vcat, collect.(partition_part)))
            part_indices ⊆ 1:d || throw("Indices in partition must be between 1 and d.")
            length(part_indices) == 2 * length(partition_part) ||
                throw("No index must occur more than once within each partition part in `partition`.")
        end

        new([β], d, partition)
    end
end

# Compute the partition mixer Hamiltonians
@memoize function partition_mixer_hamiltonians(g::PartitionMixerGate)::Vector{Matrix{ComplexF64}}
    hamiltonians = Matrix{ComplexF64}[]

    # Implements parts of Eqs. (11), (12)
    # iterate through partition
    for partition_part ∈ g.partition
        # we can represent each part by one Hamiltonian, because the individual XY gates commute
        # and therefore exp(-iβ ∑H_{a,b}) = ∏exp(-iβ H_{a,b})
        H_part = sum(XY_sum([a, b], g.d) for (a, b) ∈ partition_part)
        push!(hamiltonians, H_part)
    end

    return hamiltonians
end

# Compute the partition mixer gate matrix, but allowing to pass a different β
# than is stored in the gate struct (useful for backward pass)
function partition_mixer_gate_matrix(g::PartitionMixerGate, β::Float64)
    hamiltonians = partition_mixer_hamiltonians(g)

    # Implements parts of Eqs. (11), (12)
    # reverse terms in product to have matrix application from right to left
    Us = exp.(-im * β * hamiltonians)
    return prod(reverse(Us))
end

Qaintessent.matrix(g::PartitionMixerGate) = partition_mixer_gate_matrix(g, g.β[])

# Adjoint of partition ring mixer -> negated β parameter, reversed order in the product
Qaintessent.adjoint(g::PartitionMixerGate) = PartitionMixerGate(-g.β[], g.d, reverse(g.partition))

Qaintessent.sparse_matrix(g::PartitionMixerGate) = sparse(matrix(g))

# Number of wires (= g.d since we use the one-hot encoding canonically)
Qaintessent.num_wires(g::PartitionMixerGate)::Int = g.d


text:
# Integrator specifically for SSA
# Built to have 0-overhead stepping

struct SSAStepper <: DiffEqBase.DEAlgorithm end

mutable struct SSAIntegrator{F,uType,tType,P,S,CB,SA} <: DiffEqBase.DEIntegrator{SSAStepper,Nothing,uType,tType}
    f::F
    u::uType
    t::tType
    p::P
    sol::S
    i::Int
    tstop::tType
    cb::CB
    saveat::SA
    save_everystep::Bool
    save_end::Bool
    cur_saveat::Int
end

(integrator::SSAIntegrator)(t) = copy(integrator.u)
(integrator::SSAIntegrator)(out,t) = (out .= integrator.u)

function DiffEqBase.solve(jump_prob::JumpProblem,
                         alg::SSAStepper;
                         kwargs...)
    integrator = init(jump_prob,alg;kwargs...)
    solve!(integrator)
    integrator.sol
end

function DiffEqBase.solve!(integrator)
    end_time = integrator.sol.prob.tspan[2]
    while integrator.t < integrator.tstop # It stops before adding a tstop over
        step!(integrator)
    end

    integrator.t = end_time

    if integrator.saveat !== nothing && !isempty(integrator.saveat)
        # Split to help prediction
        while integrator.cur_saveat < length(integrator.saveat) &&
           integrator.saveat[integrator.cur_saveat] < integrator.t

            push!(integrator.sol.t,integrator.saveat[integrator.cur_saveat])
            push!(integrator.sol.u,copy(integrator.u))
            integrator.cur_saveat += 1

        end
    end

    if integrator.save_end && integrator.sol.t[end] != end_time
        push!(integrator.sol.t,end_time)
        push!(integrator.sol.u,copy(integrator.u))
    end
end

function DiffEqBase.init(jump_prob::JumpProblem,
                         alg::SSAStepper;
                         save_start = true,
                         save_end = true,
                         saveat = nothing)
    if !(jump_prob.prob isa DiscreteProblem)
        error("SSAStepper only supports DiscreteProblems.")
    end
    @assert isempty(jump_prob.jump_callback.continuous_callbacks)
    @assert length(jump_prob.jump_callback.discrete_callbacks) == 1
    cb = jump_prob.jump_callback.discrete_callbacks[1]
    prob = jump_prob.prob
    if save_start
        t = [prob.tspan[1]]
        u = [copy(prob.u0)]
    else
        t = typeof(prob.tspan[1])[]
        u = typeof(prob.u0)[]
    end
    sol = DiffEqBase.build_solution(prob,alg,t,u,dense=false,
                         calculate_error = false,
                         interp = DiffEqBase.ConstantInterpolation(t,u))
    save_everystep = any(cb.save_positions)

    if typeof(saveat) <: Number
        _saveat = prob.tspan[1]:saveat:prob.tspan[2]
    else
        _saveat = saveat
    end

   if _saveat !== nothing && !isempty(_saveat) && _saveat[1] == prob.tspan[1]
       cur_saveat = 2
   else
       cur_saveat = 1
   end

   if _saveat !== nothing && !isempty(_saveat)
     sizehint!(u,length(_saveat)+1)
     sizehint!(t,length(_saveat)+1)
   elseif save_everystep
     sizehint!(u,10000)
     sizehint!(t,10000)
   else
     sizehint!(u,2)
     sizehint!(t,2)
   end

    integrator = SSAIntegrator(prob.f,copy(prob.u0),prob.tspan[1],prob.p,
                               sol,1,prob.tspan[1],
                               cb,_saveat,save_everystep,save_end,cur_saveat)
    cb.initialize(cb,u[1],prob.tspan[1],integrator)

    integrator
end

DiffEqBase.add_tstop!(integrator::SSAIntegrator,tstop) = integrator.tstop = tstop

function DiffEqBase.step!(integrator::SSAIntegrator)
    integrator.t = integrator.tstop
    integrator.cb.affect!(integrator)
    if integrator.save_everystep
        push!(integrator.sol.t,integrator.t)
        push!(integrator.sol.u,copy(integrator.u))
    end
    @inbounds if integrator.saveat !== nothing && !isempty(integrator.saveat)
        # Split to help prediction
        while integrator.cur_saveat < length(integrator.saveat) &&
           integrator.saveat[integrator.cur_saveat] < integrator.t

            push!(integrator.sol.t,integrator.saveat[integrator.cur_saveat])
            push!(integrator.sol.u,copy(integrator.u))
            integrator.cur_saveat += 1

        end
    end
    nothing
end

export SSAStepper


text:
# # DCP analysis
using Convex
x = Variable();
y = Variable();
expr = quadoverlin(x - y, 1 - max(x, y));
println("expression curvature = ", vexity(expr));
println("expression sign = ", sign(expr));


text:
module CPUSummary

using Static, Hwloc
using Static: Zero, One, gt, lt
using IfElse: ifelse

export cache_size, cache_linesize, cache_associativity, cache_type,
  cache_inclusive, num_cache, num_cores, num_threads

include("topology.jl")
const BASELINE_CORES = Int(num_cores()) * ((Sys.ARCH === :aarch64) && Sys.isapple() ? 2 : 1)
function __init__()
  Sys.isapple() && Sys.ARCH === :aarch64 && return # detect M1
  ccall(:jl_generating_output, Cint, ()) == 1 && return
  safe_topology_load!()
  if count_attr(:Core) ≢ BASELINE_CORES
    redefine_attr_count()
    foreach(redefine_cache, 1:4)
  end
  redefine_num_threads()
  return nothing
end

end

max_stars_repo_name: JuliaSIMD/CPUSummary

max_issues_repo_name: JuliaSIMD/CPUSummary

max_forks_repo_name: JuliaSIMD/CPUSummary


text:
# https://www.sciencedirect.com/science/article/pii/S0895717710005935

function boxMuller(u,v)
    r = sqrt(-2 * log(u))
    θ = 2π*v
    (r * cos(θ), r * sin(θ))
end

normals() = ()
normals(x::Vector) = normals(x...)
normals(x::Tuple) = normals(x...)
normals(u,v,args...) = (boxMuller(u,v)..., normals(args...)...)
normals(u) = norminvcdf(u)

export StdNormals
export Normals

StdNormals(N) = Quasi(N,normals)

Normals(N,μ,σ) = map(x -> σ*x + μ, StdNormals(N))

# MvNormals(X)


text:
"""
    PCA{T<:AbstractFloat}

Principal Components Analysis

## Fields

* `covcorr` covariance or correlation matrix
* `sv` singular value decomposition
* `rnames` rownames of the original matrix
* `corr` is this a correlation matrix?
"""
struct PCA{T<:AbstractFloat}
    covcor::Symmetric{T,<:AbstractMatrix{T}}
    sv::SVD{T,T,<:AbstractMatrix{T}}
    rnames::Union{Vector{String},Missing}
    corr::Bool
end

"""
    PCA(::AbstractMatrix; corr::Bool=true)
    PCA(::ReMat; corr::Bool=true)
    PCA(::LinearMixedModel; corr::Bool=true)

Constructs a [`MixedModels.PCA`](@ref]) object from a covariance matrix.

For `LinearMixedModel`, a named tuple of PCA on each of the random-effects terms is returned.

If `corr=true`, then the covariance is first standardized to the correlation scale.
"""

function PCA(covfac::AbstractMatrix, rnames=missing; corr::Bool=true)
    covf = corr ? rownormalize(covfac) : covfac
    return PCA(Symmetric(covf * covf', :L), svd(covf), rnames, corr)
end

function Base.getproperty(pca::PCA, s::Symbol)
    if s == :cumvar
        cumvv = cumsum(abs2.(pca.sv.S))
        cumvv ./ last(cumvv)
    elseif s == :loadings
        pca.sv.U
    else
        getfield(pca, s)
    end
end

function Base.propertynames(pca::PCA, private::Bool=false)
    return (
        :covcor,
        :sv,
        :corr,
        :cumvar,
        :loadings,
        #    :rotation,
    )
end

Base.show(io::IO, pca::PCA; kwargs...) = Base.show(io, MIME"text/plain"(), pca; kwargs...)

function Base.show(
    io::IO,
    ::MIME"text/plain",
    pca::PCA;
    ndigitsmat=2,
    ndigitsvec=2,
    ndigitscum=4,
    covcor=true,
    loadings=true,
    variances=false,
    stddevs=false,
)
    println(io)
    if covcor
        println(
            io,
            "Principal components based on ",
            pca.corr ? "correlation" : "(relative) covariance",
            " matrix",
        )
        # only display the lower triangle of symmetric matrix
        if pca.rnames !== missing
            n = length(pca.rnames)
            cv = string.(round.(pca.covcor, digits=ndigitsmat))
            dotpad = lpad(".", div(maximum(length, cv), 2))
            for i in 1:n, j in (i + 1):n
                cv[i, j] = dotpad
            end
            neg = startswith.(cv, "-")
            if any(neg)
                cv[.!neg] .= " " .* cv[.!neg]
            end
            # this hurts type stability,
            # but this show method shouldn't be a bottleneck
            printmat = Text.([pca.rnames cv])
        else
            # if there are no names, then we cheat and use the print method
            # for LowerTriangular, which automatically covers the . in the
            # upper triangle
            printmat = round.(LowerTriangular(pca.covcor), digits=ndigitsmat)
        end

        Base.print_matrix(io, printmat)
        println(io)
    end
    if stddevs
        println(io, "\nStandard deviations:")
        sv = pca.sv
        show(io, round.(sv.S, digits=ndigitsvec))
        println(io)
    end
    if variances
        println(io, "\nVariances:")
        vv = abs2.(sv.S)
        show(io, round.(vv, digits=ndigitsvec))
        println(io)
    end
    println(io, "\nNormalized cumulative variances:")
    show(io, round.(pca.cumvar, digits=ndigitscum))
    println(io)
    if loadings
        println(io, "\nComponent loadings")
        printmat = round.(pca.loadings, digits=ndigitsmat)
        if pca.rnames !== missing
            pclabs = [Text(""); Text.("PC$i" for i in 1:length(pca.rnames))]
            pclabs = reshape(pclabs, 1, :)
            # this hurts type stability,
            # but this show method shouldn't be a bottleneck
            printmat = [pclabs; Text.(pca.rnames) Matrix(printmat)]
        end

        Base.print_matrix(io, printmat)
    end

    return nothing
end


text:
# -*- coding: utf-8 -*-
# + {}
module RelOlogModels
using Catlab
using Catlab.Syntax
using Catlab.WiringDiagrams
using Catlab.Doctrines
import Catlab.Doctrines.⊗

using SemanticModels.ModelTools
using SemanticModels.ModelTools.CategoryTheory
import SemanticModels.ModelTools.CategoryTheory: ⊔, FinSetMorph
import SemanticModels.ModelTools: model

export RelOlogModel, model, ⊔

struct RelOlogModel <: AbstractModel
  model::Presentation
end

function model(::Type{RelOlogModel}, r::Presentation)
    return RelOlogModel(r)
end

function ⊔(gModel::RelOlogModel, hModel::RelOlogModel)
    g = gModel.model
    h = hModel.model
    dict = Dict(i => any(isequal.(i, generators(g, FreeBicategoryRelations.Ob))) ? Ob(FreeBicategoryRelations, Symbol(i.args[1], "_2")) : i for i in generators(h, FreeBicategoryRelations.Ob))
    temp = Presentation()
    map(x->begin if !has_generator(temp, x.args[1]) add_generator!(temp, x) end end, map(x->functor((FreeBicategoryRelations.Ob, FreeBicategoryRelations.Hom), x; generators=dict), generators(h)))
    map(z->add_equation!(temp, z[1], z[2]), map(y->map(x->functor((FreeBicategoryRelations.Ob, FreeBicategoryRelations.Hom), x; generators=dict), y), equations(h)))

    out = deepcopy(g)
    merge_presentation!(out, temp)
    return RelOlogModel(out)
end

function (f::FinSetMorph)(gModel::G) where G <: RelOlogModel
    g = gModel.model
    states = generators(g, FreeBicategoryRelations.Ob)
    # check if length of states = codomain
    length(dom(f)) == length(states) || throw(DomainError(states))
    ϕ = func(f)
    # build renaming dictionary
    outS = Array{FreeBicategoryRelations.Ob}(undef, length(Set(f.fun)))
    for i in dom(f)
        if !isassigned(outS, ϕ(i)) outS[ϕ(i)] = states[i] end
    end
    dict = Dict(states[i] => outS[ϕ(i)] for i in dom(f))
    # map renaming functor to relations
    g′ = Presentation()

    # get the states in the correct order
    add_generators!(g′, outS)
    map(x->begin if !has_generator(g′, x.args[1]) add_generator!(g′, x) end end, map(x->functor((FreeBicategoryRelations.Ob, FreeBicategoryRelations.Hom), x; generators=dict), generators(g)))
    map(z->add_equation!(g′, z[1], z[2]), map(y->map(x->functor((FreeBicategoryRelations.Ob, FreeBicategoryRelations.Hom), x; generators=dict), y), equations(g)))
    return RelOlogModel(g′)
end

end


text:
#    Copyright (c) 2021 J.A. Duffek
#    Copyright (c) 2000 D.M. Spink
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program; if not, see <http://www.gnu.org/licenses/>.

using Plots;

"""
    democirc()

Demonstration of a circle arcs in the x-y plane.

# Examples:
```julia
julia> democirc()
```
"""
function democirc()
plot(title = "NURBS construction of a 2D circle and arc.",
     framestyle=:box,
     legend = false);
for r in 1.0:9.0
  crv = nrbcirc(r,[0.0;0.0;0.0],deg2rad(45),deg2rad(315));
  nrbplot!(crv,50);
end # for r
current();
end # democirc


text:
struct NotNamelistError <: Exception
    filename::String
end
Base.showerror(io::IO, e::NotNamelistError) = print("Namelist format not valid : ", e.filename)

const OptionHeader = Symbol
const OptionFileName = String

const OptionBody = OrderedDict{Symbol, Any}
const FpOptions = Vector{OptionBody}

const OptionsGroup = Dict{OptionHeader, FpOptions}

const FileOptions = Dict{OptionFileName, OptionsGroup}

struct FlexpartOptions
    dir::FlexpartDir
    options::FileOptions
end
const OPTION_FILE_NAMES = ["COMMAND", "RELEASES", "OUTGRID", "OUTGRID_NEST"]

function to_fpoption(fpdir::FlexpartDir, name::OptionFileName)
    # dir = FlexpartDir(path)
    name = name |> uppercase
    namelist2dict(joinpath(getdir(fpdir, :options), name))
    # OptionsGroup(Symbol(key) => to_fpoption(dr[key], key) for key in keys(dr))
    # Tuple(fpoption)
end

# function to_fpoption(dict_option, option::OptionName)
#     option = option |> String |> uppercase |> Symbol
#     FpOption(k |> String |> lowercase |> Symbol => v for (k, v) in dict_option)
# end

FlexpartOptions(path::FlexpartPath) = FlexpartOptions(FlexpartDir(path))
# function FlexpartOptions(fpdir::FlexpartDir)
#     FlexpartOptions(
#         fpdir,
#         FileOptions(filename => to_fpoption(fpdir, filename) for filename in OPTION_FILE_NAMES)
#     )
# end
function FlexpartOptions(fpdir::FlexpartDir)
    FlexpartOptions(
        fpdir,
        getnamelists(getdir(fpdir, :options))
    )
end
Base.getindex(fp::FlexpartOptions, name::OptionFileName) = fp.options[name]
function Base.setindex!(fp::FlexpartOptions, value, name::OptionFileName)
    fp.options[name] = value
end
Base.keys(fpoptions::FlexpartOptions) = Base.keys(fpoptions.options)

function add(fpoptions::FlexpartOptions, name::OptionFileName, header::OptionHeader, value)
    optionbody = convert(OptionBody, value)
    group = OptionsGroup(header => [optionbody])
    fpoptions[name] = group
end

function getnamelists(path::String)
    fileoptions = FileOptions()
    for (root, dirs, files) in walkdir(path)
        for file in files
            rel = basename(relpath(path, root))
            name = root == path ? file : joinpath(relpath(root, path), file)
            absfile = joinpath(root, file)
            dict = try
                namelist2dict(absfile)
            catch e
                if e isa NotNamelistError
                    nothing
                else
                    throw(e)
                end
            end
            !isnothing(dict) && push!(fileoptions, name => dict)
        end
    end
    fileoptions
end

# function getnamelists(pathinit::String, path::String)
#     files = readdir(path, join=true)
#     for file in files
#         if isdir(file)
#             getnamelists(pathinit, file)
#         else

#         end
#     end
# end

# function to_fpoption(dict_option, option::String)
#     option = option |> uppercase
#     fpoption = STR_TO_TYPE[option]
#     d = Dict(k |> String |> lowercase |> Symbol => v for (k, v) in dict_option)
#     fn = fieldnames(fpoption)
#     fd = Dict(key => try
#             if stype <: Vector
#                 parse.(Float64, split(d[key], ","))
#             else
#                 parse(stype, d[key])
#             end
#         catch
#             d[key]
#         end for (key, stype) in zip(fn, fpoption.types)
#     )
#     nt = NamedTuple{fn}([fd[key] for key in fn])
#     fpoption(nt...)
# end

# function FlexpartOptions(dir::FlexpartDir)
#     FlexpartOptions(
#         to_fpoption(dir, "COMMAND")...,
#         Releases(
#             to_fpoption(dir, "RELEASES")...
#         ),
#         to_fpoption(dir, "OUTGRID")...,
#         to_fpoption(dir, "OUTGRID_NEST")...
#     )
# end
# FlexpartOptions(path::String) = FlexpartOptions(FlexpartDir(path))

# function write(file::IOStream, option::FpOption)
#     for line in format(option) write(file, line*"\n") end
# end
# function write(file::IOStream, options::Vector{<:FpOption})
#     for option in options write(file, option) end
# end

# function write_options(options::Vector{FpOption}, filepath::String, dest::String="")
#     (tmppath, tmpio) = mktemp()

#     write(tmpio, options)

#     close(tmpio)
#     newf = dest=="" ? mv(tmppath, filepath, force=true) : mv(tmppath, joinpath(dest, basename(filepath)), force=true)
#     chmod(newf, stat(filepath).mode)
#     newf
# end
# write_options(option::FpOption, filepath::String, dest::String="") = write_options([option], filepath, dest)

function format(options::OptionsGroup)::Vector{String}
    str = []
    for (header, fpoptions) in options
        for line in option2lines(header, fpoptions)
            push!(str, line)
        end
    end
    str
end

function option2lines(header::Symbol, options::FpOptions) :: Vector{String}
    head = header |> string |> uppercase
    str = String[]
    for option in options
        push!(str, "&$(head)")
        for line in body2lines(option)
            push!(str, line)
        end
    end
    str
end
function body2lines(body::OptionBody) :: Vector{String}
    str = String[]
    for (k, v) in body
        key = uppercase(String(k))
        # val = v |> typeof <: Vector ? join(field, ",") : field
        push!(str, " $key = $v,")
    end
    push!(str, " /")
    str
end

function area2outgrid(area::Vector{<:Real}, gridres=0.01; nested=false)
    # mult = 1/gridres
    # area = round_area(area, mult)
    outlon0 = area[2]
    outlat0 = area[3]
    Δlon = area[4] - outlon0
    Δlat = area[1] - outlat0
    (numxgrid, numygrid) = try
        convert(Int, Δlon/gridres), convert(Int, Δlat/gridres)
    catch
        error("gridres must divide area")
    end
    dxout = gridres
    dyout = gridres

    # var = [:outlon0, :outlat0, :numxgrid, :numygrid, :dxout, :dyout]
    # values = eval.(var)
    # Dict(k=> v for (k, v) in zip(var, values))
    res = OrderedDict(
        :OUTLON0 => outlon0, :OUTLAT0 => outlat0, :NUMXGRID => numxgrid, :NUMYGRID => numygrid, :DXOUT => dxout, :DYOUT => dyout,
    )
    nested ? Dict(
        String(k)*'N' |> Symbol => v for (k, v) in res
    ) : res
end

function area2outgrid(fpdir::FlexpartDir, gridres::Real; nested=false)
    firstinput = readdir(getdir(fpdir, :input), join=true)[1]
    area = grib_area(firstinput)

    area2outgrid(area, gridres; nested)
end

function set!(option::OptionBody, newv)
    newv = newv isa Pair ? Dict(newv) : newv
    merge!(option, newv)
end
function set(option::OptionBody, newv)
    newv = newv isa Pair ? Dict(newv) : newv
    merge(option, newv)
end

function setfromdates!(fpoptions::FlexpartOptions, start::DateTime, finish::DateTime)
    toset = OrderedDict(
        :IBDATE => Dates.format(start, "yyyymmdd"),
        :IEDATE => Dates.format(finish, "yyyymmdd"),
        :IBTIME => Dates.format(start, "HHMMSS"),
        :IETIME => Dates.format(finish, "HHMMSS"),
    )
    set!(fpoptions["COMMAND"][:command][1], toset)
end

function setrelease!(fpoptions::FlexpartOptions, start::DateTime, finish::DateTime)
    set!(fpoptions["RELEASE"][:command][1], toset)
end
# function broadcast(set, option::OptionBody, newvs::Array{Dict{Symbol, <:Any}})
#     [set(option, newv) for newv in newvs]
# end

function write(flexpartoption::FlexpartOptions, newpath::String = "")
    options_dir = newpath == "" ? getdir(flexpartoption.dir, :options) : joinpath(newpath, OPTIONS_DIR)
    try
        mkdir(options_dir)
    catch
    end

    for (name, options) in flexpartoption.options
        filepath = joinpath(options_dir, name)
        write(options, filepath)
    end
end
function write(options::OptionsGroup, path::String)
    (tmppath, tmpio) = mktemp()

    for line in format(options) Base.write(tmpio, line*"\n") end

    close(tmpio)
    dest = path

    mv(tmppath, dest, force=true)
end

# function namelist2dict(filepath)
#     options = FpOption[]
#     headers = OptionName[]
#     count = 0
#     f = open(filepath, "r")
#     for line in eachline(f)
#         if !((m = match(r"\s*(.*?)\s*=\s*(\".*?\"|[^\s,]*)\s*,", line)) |> isnothing) #captures the field name in group 1 and the value in group 2
#             push!(options[count], m.captures[1] |> Symbol => m.captures[2])
#         elseif !((m = match(r"\&(\w*)", line)) |> isnothing) #captures the name of the header in group 1
#             count = count + 1
#             push!(options, Dict{Symbol, Any}())
#             push!(headers, m.captures[1] |> lowercase |> Symbol)
#         end
#     end
#     close(f)
#     OptionsGroup(k => v for (k, v) in zip(headers, options))
# end

diffkeys(dict1, dict2) = [k for k in keys(dict1) if dict1[k] != get(dict2, k, nothing)]

function compare(opt1::OptionsGroup, opt2::OptionsGroup)
    d = Dict()
    for (header, fpoption) in opt1
        push!(d, header => [diffkeys(d1, opt2[header][i]) for (i, d1) in enumerate(fpoption)])
    end
    println("HEADER \t NAME \t FILE1 \t FILE2")
    for (header, v) in d
        println(header)
        for (i, opt) in enumerate(v)
            for name in opt
                println("      \t $name \t $(opt1[header][i][name]) \t $(opt2[header][i][name])")
            end
        end
    end
    d
end

function compare(file1::String, file2::String)
    opt1 = namelist2dict(file1)
    opt2 = namelist2dict(file2)
    compare(opt1, opt2)
    # diffs
end

function compare(fpdir1::FlexpartDir, fpdir2::FlexpartDir, filename1::String; filename2::String = "", which = :output)
    filename2 = filename2 |> isempty ? filename1 : filename2
    file1 = joinpath(getdir(fpdir1, which), filename1)
    file2 = joinpath(getdir(fpdir2, which), filename2)
    compare(file1, file2)
end

function namelist2dict(filepath)
    lines = readlines(filepath)
    reg_header = r"\&(\w*)"
    # match only '/' that are not between brackets.
    # src: https://stackoverflow.com/questions/11502598/how-to-match-something-with-regex-that-is-not-between-two-special-characters
    reg_endofbody = r"\/(?=(?:[^\"]*\"[^\"]*\")*[^\"]*\Z)"
    cleared = [split(line, '!')[1] for line in lines] # remove the comments
    iheader = findall(line -> !isnothing(match(reg_header, line)), cleared) # find headers
    iend = findall(line -> !isnothing(match(reg_endofbody, line)), cleared) # find end of body
    (length(iheader) == 0 || (length(iheader) !== length(iend))) && throw(NotNamelistError(filepath))
    bodys = [lines[i:j] for (i, j) in zip(iheader, iend)]
    headers = lines[iheader]
    headers = [match(reg_header, line).captures[1] |> lowercase |> Symbol for line in headers]
    optgroup = OptionsGroup(k => OptionBody[] for k in unique(headers))
    # optgroup = Dict(k => Dict{Symbol, Any}[] for k in unique(headers))
    for (n, header) in enumerate(headers)
        push!(optgroup[header], lines2option(bodys[n]))
    end
    optgroup
end

# function isnamelist(file)
#     try
#         namelist2dict(file)
#     catch
#     end
# end

function lines2option(lines::Vector{String})
    opt = OptionBody()
    for line in lines
        if !((m = match(r"\s*(.*?)\s*=\s*(\".*?\"|[^\s,]*)\s*,", line)) |> isnothing) #captures the field name in group 1 and the value in group 2
            push!(opt, m.captures[1] |> Symbol => m.captures[2])
        end
    end
    opt
end

# Base.convert(::Type{OutgridN}, outgrid::Outgrid) = OutgridN(getfield.(Ref(outgrid), filter(x -> x != :outheights, fieldnames(typeof(outgrid))))...)

# function set_heights(outgrid::Outgrid, heights::Vector{Float64})
#     outgrid.outheights = heights
#     outgrid
# end


text:
using DynamicPolynomials
using MultivariatePolynomials
using Test
using LinearAlgebra

# TODO move to MP
@testset "Issue #70" begin
    @ncpolyvar y0 y1 x0 x1
    p = x1 * x0 * x1
    @test subs(p, x0 => y0, x1 => y1) == y1 * y0 * y1
    @test subs(p, x0 => 1) == x1^2
    @test p(x0 => y0, x1 => y1) == y1 * y0 * y1
end

include("mono.jl")
include("poly.jl")
include("comp.jl")
include("mutable_arithmetics.jl")

module newmodule
    using Test
    import DynamicPolynomials
    @testset "Polyvar macro hygiene" begin
        # Verify that the @polyvar macro works when the package has been activated
        # with `import` instead of `using`.
        DynamicPolynomials.@polyvar x y
        @test isa(x, DynamicPolynomials.PolyVar)
        @test isa(y, DynamicPolynomials.PolyVar)
    end
end

include("mvp.jl")


text:
mutable struct ActionNode
    r::Float64      # reward for taking this action
    pind::Int       # index of parent belief node in Graph.belief_nodes
    ai::Int         # which action does this correspond to

    L::Float64
    U::Float64

    children::UnitRange{Int64}
end

mutable struct BeliefNode{B}
    b::B            # belief associated with this node
    ind::Int        # index of this node in Graph.belief_nodes
    pind::Int       # index of parent action node in Graph.action_nodes
    aind::Int       # index of best child action node in action_nodes

    oi::Int         # index of observation corresponding with this belief
    po::Float64     # probability of seeing that observation
    poc::Float64    # prob of all observations leading here (cumulative)

    L::Float64
    U::Float64
    d::Int          # depth
    gd::Float64     # discount^depth, stored to improve computation

    children::UnitRange{Int64}
end
# constructor for root node
function BeliefNode(b, L::Float64, U::Float64)
    return BeliefNode(b, 1, 0, 0, 0, 1.0, 1.0, L, U, 0, 1.0, 0:0)
end
function BeliefNode(b,
                    ind::Int,
                    pind::Int,
                    oi::Int,
                    po::Float64,
                    poc::Float64,
                    L::Float64,
                    U::Float64,
                    d::Int,
                    gd::Float64
                   )

    return BeliefNode(b, ind, pind, 0, oi, po, poc, L, U, d, gd, 0:0)
end

mutable struct Graph{B <: BeliefNode}
    action_nodes::Vector{ActionNode}
    belief_nodes::Vector{B}

    na::Int     # number of action nodes
    nb::Int     # number of belief nodes

    root_ind::Int

    df::Float64     # discount factor

    # constructor
    function Graph{B}(df::Real) where B
        new(ActionNode[], BeliefNode{B}[], 0, 0, 1, df)
    end
end

function clear_graph!(G::Graph)
    G.action_nodes = ActionNode[]
    G.belief_nodes = BeliefNode[]
    G.na = 0
    G.nb = 0
    G.root_ind = 1
end

isroot(G::Graph, bn::BeliefNode) = G.root_ind == bn.ind

get_root(G::Graph) = G.belief_nodes[G.root_ind]

# assumes all beliefs can have children
# even if it is a terminal belief, I suppose action could do nothing
isfringe(bn::BeliefNode) = iszero(bn.children)

function add_node(G::Graph, an::ActionNode)
    push!(G.action_nodes, an)
    G.na += 1
end
function add_node(G::Graph, bn::BeliefNode)
    push!(G.belief_nodes, bn)   # add new node to list of belief nodes
    G.nb += 1                   # length of G.belief_nodes increases by 1
end

parent_node(G::Graph, bn::BeliefNode) = G.action_nodes[bn.pind]
parent_node(G::Graph, an::ActionNode) = G.belief_nodes[an.pind]