text:
module Bridge_format
    import ..Gaugefields:GaugeFields,SU3GaugeFields,SU2GaugeFields,set_wing!
#=
Bridge++ Text file format
U(x,y,z,t;mu)_{ab} is like
(outer) [mu][t][z][y][x][a][b][re/im] (inner)
re/im;
  0: Real
  1: Imaginary
mu: (x,y,z,t)=(0,1,2,3)

Re of U(0,0,0,0;mu=0)_00  # mu=0, site (x,y,z,t)=(0,0,0,0)
Im of U(0,0,0,0;mu=)_00
Re of U(0,0,0,0;mu=0)_01
Im of U(0,0,0,0;mu=0)_01
...
Re of U(1,0,0,0;0)_00     # mu=0, site (x,y,z,t)=(1,0,0,0)
Im of U(1,0,0,0;0)_00
...
Re of U(0,0,0,0;mu=1)_00  # mu=1, site (x,y,z,t)=(0,0,0,0)
Im of U(0,0,0,0;mu=1)_00

=#

    function save_textdata(U,filename)

        NX = U[1].NX
        NY = U[1].NY
        NZ = U[1].NZ
        NT = U[1].NT
        NC = U[1].NC

        #li = LIME_header((NX,NY,NZ,NT),"su3gauge",1,64)
        #print(li.doc)
        #write("test.xml", li.doc)

        fp = open(filename,"w")
        #fp = open("testbin.dat","w")
        i = 0
        i = 0
        #for μ=1:4
            for it=1:NT
                for iz=1:NZ
                    for iy=1:NY
                        for ix=1:NX
                            for μ=1:4
                            for a = 1:NC
                                for b = 1:NC
                                    i+= 1
                                    #rvalue = read(fp, floattype)
                                    rvalue = real(U[μ][a,b,ix,iy,iz,it])
                                    println(fp,rvalue)
                                    ivalue = imag(U[μ][a,b,ix,iy,iz,it])
                                    println(fp,ivalue)
                                end
                            end
                        end
                    end
                end
            end
        end
        close(fp)

        return

    end

    function load_BridgeText!(initial,U,L,NC)
        NX = L[1]
        NY = L[2]
        NZ = L[3]
        NT = L[4]
        @assert U[1].NX == NX "NX mismatch"
        @assert U[1].NY == NY "NY mismatch"
        @assert U[1].NZ == NZ "NZ mismatch"
        @assert U[1].NT == NT "NT mismatch"
        @assert U[1].NC == NC "NC mismatch"
        fp = open(initial,"r")
        numdata = countlines(initial)
        @assert numdata == 4*NX*NY*NT*NZ*NC*NC*2 "data shape is wrong"

        #for μ=1:4
            for it=1:NT
                for iz=1:NZ
                    for iy=1:NY
                        for ix=1:NX
                            for μ=1:4
                            for a=1:NC
                                for b=1:NC
                                    u = split(readline(fp))
                                    #println(u)
                                    rvalue = parse(Float64,u[1])
                                    u = split(readline(fp))
                                    ivalue = parse(Float64,u[1])
                                    U[μ][a,b,ix,iy,iz,it] = rvalue + im*ivalue
                                    #U[μ][b,a,ix,iy,iz,it] = rvalue + im*ivalue
                                end
                            end
                            #u = U[μ][:,:,ix,iy,iz,it]
                            #println(u'*u)
                        end
                    end
                end
            end
        end

        set_wing!(U)

        close(fp)

    end
end


text:
flbl = ARGS[1]
fout = ARGS[2]

using EMIRT

lbl = readseg(flbl)

add_lbl_boundary!( lbl )

saveseg(fout, lbl, "main")

ext: jl


text:
using WAV
using Dierckx

const A4 = 440
const C4 = 440/2^(9/12)
const C5 = C4*2
const C6 = C4*4
const C3 = C4/2
const C2 = C4/4
const C1 = C4/8

fs = 48000
p = Parameters(0,2,5,15)

maxval = 1-1/32768

square_i = Instruction(inst_pulse_exp2, IContext(0, 0))
triangle_i = Instruction(inst_triangle, IContext(0,-1))

f_env = 60
#f_env = 480
#f_env = 1000
n_env = div(fs, f_env)
h_env = 1/f_env

range_clamp(x) = clamp.(x, -1, maxval)

function f_to_freq_exp(f)
	freq = 2*f/fs
	e = 0
	if freq > 0
		while (freq > 1)
			freq /= 2
			e += 1
		end
		while (freq < 0.5)
			freq *= 2
			e -= 1
		end
	end
	return (Fix(freq), e)
end

function fm_program(fc, fm; gain_m_exp4=0, m_inst = inst_sin, c_inst = inst_sin)
	freq_c, exp_c = f_to_freq_exp(fc)
	freq_m, exp_m = f_to_freq_exp(fm)

	gain_m = gain_c = Fix(0.0)

	phase_c_addr, phase_m_addr = 4, 2
	gain_c_addr, gain_m_addr = 5, 3
	mix_addr = 6
	program = [
		(Instruction(inst_acc_imm_to_mem, IContext(phase_c_addr, exp_c)), freq_c),
		(Instruction(inst_acc_imm_to_mem, IContext(phase_m_addr, exp_m)), freq_m),
		(m_inst, Fix(0)), # phase_m
		(Instruction(inst_wadd, IContext(phase_c_addr, gain_m_exp4)), gain_m),
		(c_inst, Fix(0)), # phase_c
		(Instruction(inst_macc_to_mem, IContext(mix_addr)), gain_c),

		(inst_output, Fix(0)),
		(inst_dither_update, Fix(2^-15))
	]

	d = Dict(:gain_c => gain_c_addr, :gain_m => gain_m_addr)
	return (program, d)
end

function fm_program_lp(fc, fm, lp_exp; gain_m_exp4=0, m_inst = inst_sin, c_inst = inst_sin)
	freq_c, exp_c = f_to_freq_exp(fc)
	freq_m, exp_m = f_to_freq_exp(fm)

	gain_m = gain_c = Fix(0.0)

	phase_c_addr, phase_m_addr = 6, 4
	gain_c_addr, gain_m_addr = 0, 1
	filt_gain_c_addr, filt_gain_m_addr = 7, 5
	mix_addr = 8
	program = [
		(Instruction(inst_lpfilter, IContext(filt_gain_c_addr, lp_exp)), gain_c), # gain_c parameter (unfiltered)
		(Instruction(inst_lpfilter, IContext(filt_gain_m_addr, lp_exp)), gain_m), # gain_m parameter (unfiltered)
		(Instruction(inst_acc_imm_to_mem, IContext(phase_c_addr, exp_c)), freq_c),
		(Instruction(inst_acc_imm_to_mem, IContext(phase_m_addr, exp_m)), freq_m),
		(m_inst, Fix(0)), # phase_m
		(Instruction(inst_wadd, IContext(phase_c_addr, gain_m_exp4)), gain_m),
		(c_inst, Fix(0)), # phase_c
		(Instruction(inst_macc_to_mem, IContext(mix_addr)), gain_c),

		(inst_output, Fix(0)),
		(inst_dither_update, Fix(2^-15))
	]

	d = Dict(:gain_c => gain_c_addr, :gain_m => gain_m_addr)
	return (program, d)
end

function fm_and_noise_program(fc, fm, fn; gain_m_exp4=0, m_inst = inst_sin, c_inst = inst_sin)
	freq_c, exp_c = f_to_freq_exp(fc)
	freq_m, exp_m = f_to_freq_exp(fm)
	freq_n, exp_n = f_to_freq_exp(fn)

	gain_m = gain_c = gain_n = Fix(0.0)

	phase_c_addr, phase_m_addr, phase_n_addr = 4, 2, 7
	gain_c_addr, gain_m_addr, gain_n_addr = 5, 3, 9

	mix_addr = 10
	program = [
		(Instruction(inst_acc_imm_to_mem, IContext(phase_c_addr, exp_c)), freq_c),
		(Instruction(inst_acc_imm_to_mem, IContext(phase_m_addr, exp_m)), freq_m),
		(m_inst, Fix(0)), # phase_m
		(Instruction(inst_wadd, IContext(phase_c_addr, gain_m_exp4)), gain_m),
		(c_inst, Fix(0)), # phase_c
		(Instruction(inst_macc_to_mem, IContext(mix_addr)), gain_c),

		(Instruction(inst_acc_imm_to_mem, IContext(phase_n_addr, exp_n)), freq_n),
		(inst_nop, 0), # phase_n
		(inst_lfsr, -1),
		(Instruction(inst_macc_to_mem, IContext(mix_addr)), gain_n),

		(inst_output, Fix(0)),
		(inst_dither_update, Fix(2^-15))
	]

	d = Dict(:gain_c => gain_c_addr, :gain_m => gain_m_addr, :gain_n => gain_n_addr)
	return (program, d)
end

function synthesize!(filename, program, t, env)
	dsp = create_dsp(p, program)
	s = Fix[]
	for k=1:length(t)
		for (addr, data) in env
			dsp.data[addr+1] = Fix(range_clamp(data[k]))
		end
		run_loops!(dsp, s, n_env)
	end
	wavwrite(float.(s), filename; Fs=fs)
	return s
end

phi = (1+sqrt(5))/2

function bell()
	T = 4
	lp_exp = -8

	fc = C5; Imax = 4.5
	#fc = C6; Imax = 3
	fm = phi*fc
	#fm = 4/3*fc+2
	#fm = 5/3*fc+2; Imax = 4

	gain_m_exp4 = 1
	gain_m_factor = Imax/pi/(4^gain_m_exp4)

	t = 0:h_env:T
	fa = 1 .- exp.(-t*25)
	gain_c = fa.*exp.(-0.5*t)
	gain_m = gain_m_factor*fa.*exp.(-0.85*t)

	#program, d = fm_program(fc, fm; gain_m_exp4=gain_m_exp4)
	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4)
	synthesize!("../wav/bell.wav", program, t, [d[:gain_m] => gain_m, d[:gain_c] => gain_c])

	fc *= 3/4
	fm = 4/3*fc+2

	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4)
	synthesize!("../wav/bell_4th.wav", program, t, [d[:gain_m] => gain_m, d[:gain_c] => gain_c])
end

function wood_drum()
	fc = C3
	fm = fc/phi
	Imax = 25/2
	gain_m_exp4 = 2
	T = 2
	lp_exp = -4

	t = 0:h_env:T

	I = Imax*max.(0, 1 .- t/0.02)
	tVmax = 0.04
	V0 = 0.8
	Te = 0.02
	V = min.(t/tVmax*(1-V0) .+ 0.8, exp.(-(t .- tVmax)/Te))

	#program, d = fm_program(fc, fm; gain_m_exp4=gain_m_exp4)
	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4)
	synthesize!("../wav/wood_drum.wav", program, t, [d[:gain_m] => I/pi, d[:gain_c] => V])
end

function brass()
	fc = C4
	fm = fc + 1
	T = 2
	lp_exp = -8
	#Imax = 3; gain_m_exp4 = 0
	Imax = 3.5; gain_m_exp4 = 0
	v = [0,1,0.75,0.63,0]
	#dtv = [0, 85, 86, 256, 86]; Tenv = 0.6
	dtv = [0, 85, 86, 256+512, 86]; Tenv = 1

	tv = cumsum(dtv)
	tv = tv*(Tenv/maximum(tv))

	t = 0:h_env:T

	V = Spline1D(tv, v; k=1)(t)
	I = Imax*V

	#program, d = fm_program(fc, fm; gain_m_exp4=gain_m_exp4)
	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4)
	synthesize!("../wav/brass.wav", program, t, [d[:gain_m] => I/(pi*4^gain_m_exp4), d[:gain_c] => V])
end

function clarinet()
	fm = C5
	fc = 3*fm + 0.1
	Imax = 2.0; gain_m_exp4 = 0
	Tenv = 1
	T = 2
	lp_exp = -8

	i = [0,1,1]
	dti = [0, 50, 463+512]
	v = [0,1,1,0]
	dtv = [0, 50, 443+512, 20]
	ti = cumsum(dti)
	tv = cumsum(dtv)
	ti = ti*(Tenv/maximum(tv))
	tv = tv*(Tenv/maximum(tv))

	t = 0:h_env:T

	Ienv = Spline1D(ti, i; k=1)(t)
	V = Spline1D(tv, v; k=1)(t)

	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4)
	synthesize!("../wav/clarinet.wav", program, t, [d[:gain_m] => Imax*Ienv/(pi*4^gain_m_exp4), d[:gain_c] => V])
	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4, m_inst = inst_sina2)
	synthesize!("../wav/clarinet2.wav", program, t, [d[:gain_m] => Imax*Ienv/(pi*4^gain_m_exp4), d[:gain_c] => V])
	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4, c_inst = inst_sina2, m_inst = inst_sina2)
	synthesize!("../wav/clarinet3.wav", program, t, [d[:gain_m] => Imax*Ienv/(pi*4^gain_m_exp4), d[:gain_c] => V])

	fm = C4
	fc = 3*fm + 0.1
	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4, c_inst = inst_sina2)
	Imax = 0.2
	synthesize!("../wav/pipe.wav", program, t, [d[:gain_m] => Imax*Ienv/(pi*4^gain_m_exp4), d[:gain_c] => V])

	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4, m_inst = triangle_i)
	Imax = 0.25
	synthesize!("../wav/pipe2.wav", program, t, [d[:gain_m] => Imax*Ienv/(pi*4^gain_m_exp4), d[:gain_c] => V])

	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4, m_inst = triangle_i, c_inst = inst_sina2)
	Imax = 0.25
	synthesize!("../wav/pipe3.wav", program, t, [d[:gain_m] => Imax*Ienv/(pi*4^gain_m_exp4), d[:gain_c] => V])
end

function bassoon()
	fm = C2
	fc = 5*fm + 2
	Imax = 2.0; gain_m_exp4 = 0
	Tenv = 1
	T = 2
	lp_exp = -8

	i = [0,1,1]
	dti = [0, 100, 463+512]
	v = [0,1,1,0]
	dtv = [0, 100, 443+512, 100]
	ti = cumsum(dti)
	tv = cumsum(dtv)
	ti = ti*(Tenv/maximum(tv))
	tv = tv*(Tenv/maximum(tv))

	t = 0:h_env:T

	Ienv = Spline1D(ti, i; k=1)(t)
	V = Spline1D(tv, v; k=1)(t)

	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4)
	synthesize!("../wav/bassoon.wav", program, t, [d[:gain_m] => Imax*Ienv/(pi*4^gain_m_exp4), d[:gain_c] => V])
	#program, d = fm_program_lp(fc*2, fm*2, lp_exp; gain_m_exp4=gain_m_exp4)
	#synthesize!("../wav/bassoon-h.wav", program, t, [d[:gain_m] => Imax*Ienv/(pi*4^gain_m_exp4), d[:gain_c] => V])

	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4, c_inst = inst_sina2)
	synthesize!("../wav/bassoon1.wav", program, t, [d[:gain_m] => Imax*Ienv/(pi*4^gain_m_exp4), d[:gain_c] => V])
	#program, d = fm_program_lp(fc*2, fm*2, lp_exp; gain_m_exp4=gain_m_exp4, c_inst = inst_sina2)
	#synthesize!("../wav/bassoon1-h.wav", program, t, [d[:gain_m] => Imax*Ienv/(pi*4^gain_m_exp4), d[:gain_c] => V])

	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4, m_inst = triangle_i)
	synthesize!("../wav/bassoon2.wav", program, t, [d[:gain_m] => Imax*Ienv/(pi*4^gain_m_exp4), d[:gain_c] => V])
	#program, d = fm_program_lp(fc*2, fm*2, lp_exp; gain_m_exp4=gain_m_exp4, m_inst = triangle_i)
	#synthesize!("../wav/bassoon2-h.wav", program, t, [d[:gain_m] => Imax*Ienv/(pi*4^gain_m_exp4), d[:gain_c] => V])

	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4, m_inst = triangle_i, c_inst = inst_sina2)
	synthesize!("../wav/bassoon3.wav", program, t, [d[:gain_m] => Imax*Ienv/(pi*4^gain_m_exp4), d[:gain_c] => V])

	Imax = 3.5; gain_m_exp4 = 1
	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4)
	synthesize!("../wav/bassooni.wav", program, t, [d[:gain_m] => Imax*Ienv/(pi*4^gain_m_exp4), d[:gain_c] => V])
	#program, d = fm_program_lp(fc*2, fm*2, lp_exp; gain_m_exp4=gain_m_exp4)
	#synthesize!("../wav/bassooni-h.wav", program, t, [d[:gain_m] => Imax*Ienv/(pi*4^gain_m_exp4), d[:gain_c] => V])

	#program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4, m_inst = triangle_i)
	#synthesize!("../wav/bassooni2.wav", program, t, [d[:gain_m] => Imax*Ienv/(pi*4^gain_m_exp4), d[:gain_c] => V])
end

function bass()
	fc = C2
	fm = fc-1
	Imax = 5; gain_m_exp4 = 1
	Te = 0.5
	T = 2
	lp_exp = -8

	t = 0:h_env:T
	fa = 1 .- exp.(-t*100)
	V = fa.*exp.(-t/Te)
	I = Imax*fa.*exp.(-0.5*t/Te)

	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4)
	synthesize!("../wav/bass.wav", program, t, [d[:gain_m] => I/(pi*4^gain_m_exp4), d[:gain_c] => V])
	#program, d = fm_program_lp(fc*2, fm*2, lp_exp; gain_m_exp4=gain_m_exp4)
	#synthesize!("../wav/bass-h.wav", program, t, [d[:gain_m] => I/(pi*4^gain_m_exp4), d[:gain_c] => V])

	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4, c_inst = inst_sina2)
	synthesize!("../wav/bass-cs2.wav", program, t, [d[:gain_m] => I/(pi*4^gain_m_exp4), d[:gain_c] => V])
	#program, d = fm_program_lp(fc*2, fm*2, lp_exp; gain_m_exp4=gain_m_exp4, c_inst = inst_sina2)
	#synthesize!("../wav/bass-cs2-h.wav", program, t, [d[:gain_m] => I/(pi*4^gain_m_exp4), d[:gain_c] => V])

	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4, m_inst=inst_sina2)
	synthesize!("../wav/bass-ms2.wav", program, t, [d[:gain_m] => I/(pi*4^gain_m_exp4), d[:gain_c] => V])
	#program, d = fm_program_lp(fc*2, fm*2, lp_exp; gain_m_exp4=gain_m_exp4, m_inst=inst_sina2)
	#synthesize!("../wav/bass-ms2-h.wav", program, t, [d[:gain_m] => I/(pi*4^gain_m_exp4), d[:gain_c] => V])

	#program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4, m_inst=triangle_i)
	#synthesize!("../wav/bass-mt.wav", program, t, [d[:gain_m] => I/(pi*4^gain_m_exp4), d[:gain_c] => V])
	#program, d = fm_program_lp(fc*2, fm*2, lp_exp; gain_m_exp4=gain_m_exp4, m_inst=triangle_i)
	#synthesize!("../wav/bass-mt-h.wav", program, t, [d[:gain_m] => I/(pi*4^gain_m_exp4), d[:gain_c] => V])
end

function drum()
	T = 2

	fc = C2
	Imax = 25; gain_m_exp4 = 2
	fm = fc/phi

	t = 0:h_env:T

	#fa = 1 .- exp.(-t*200); lp_exp = 0
	fa = 1; lp_exp = -8

	V = fa.*exp.(-3*t)
	I = Imax*fa.*exp.(-85*t)

	#program, d = fm_program(fc, fm; gain_m_exp4=gain_m_exp4)
	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4)
	synthesize!("../wav/drum.wav", program, t, [d[:gain_m] => I/(pi*4^gain_m_exp4), d[:gain_c] => V])
end

function snare()
	T = 2
	t = 0:h_env:T

	fn = 5000

	fa = 1 .- exp.(-t*200)
	Vn = fa.*exp.(-10*t)
	#Vn = max.(0, 1 .- 5*t)

	fc = C2*2^(7/12)
	Imax = 18; gain_m_exp4 = 2
	fm = fc/phi

	#V = fa.*exp.(-7*t)
	V = max.(0, 1 .- 5*t)
	I = Imax*fa.*exp.(-40*t) .+ 1.5

	program, d = fm_and_noise_program(fc, fm, fn; gain_m_exp4=gain_m_exp4)
	synthesize!("../wav/snare.wav", program, t, [d[:gain_m] => I/(pi*4^gain_m_exp4), d[:gain_c] => V, d[:gain_n] => Vn])
end

function hihat()
	T = 1
	t = 0:h_env:T

	fn = fs

	#fa = 1 .- exp.(-t*200)
	Vn = max.(0, 1 .- 20*t)

	gain_m_exp4 = 0
	fc = C2*2^(7/12)
	fm = fc/phi

	V = 0*t
	I = 0*t

	program, d = fm_and_noise_program(fc, fm, fn; gain_m_exp4=gain_m_exp4)
	synthesize!("../wav/hihat.wav", program, t, [d[:gain_m] => I/(pi*4^gain_m_exp4), d[:gain_c] => V, d[:gain_n] => Vn])
end

function keyboard()
	fc = C3
	fm = fc*1.5 + 1
	lp_exp = -8

	Imax = 2.5; gain_m_exp4 = 0

	T = 2
	t = 0:h_env:T

	fa = 1 .- exp.(-t*75)
	V = fa.*exp.(-3*t)
	I = Imax.*fa.*exp.(-1.8*t)

	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4, m_inst = inst_sina2)
	synthesize!("../wav/keyboard0.wav", program, t, [d[:gain_m] => I/(pi*4^gain_m_exp4), d[:gain_c] => V])
	program, d = fm_program_lp(fc*2, fm*2, lp_exp; gain_m_exp4=gain_m_exp4, m_inst = inst_sina2)
	synthesize!("../wav/keyboard0-h.wav", program, t, [d[:gain_m] => I/(pi*4^gain_m_exp4), d[:gain_c] => V])
	#program, d = fm_program_lp(fc*2, fm*2, lp_exp; gain_m_exp4=gain_m_exp4)
	#synthesize!("../wav/keyboard00-h.wav", program, t, [d[:gain_m] => I/(pi*4^gain_m_exp4), d[:gain_c] => V])

	fm = fc + 1
	Imax = 4.8; gain_m_exp4 = 1
	I = Imax.*fa.*exp.(-1.8*t)
	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4, m_inst = inst_sina2)
	synthesize!("../wav/keyboard1.wav", program, t, [d[:gain_m] => I/(pi*4^gain_m_exp4), d[:gain_c] => V])
	program, d = fm_program_lp(fc*2, fm*2, lp_exp; gain_m_exp4=gain_m_exp4, m_inst = inst_sina2)
	synthesize!("../wav/keyboard1-h.wav", program, t, [d[:gain_m] => I/(pi*4^gain_m_exp4), d[:gain_c] => V])

	fm = fc*2 + 1
	Imax = 2.0; gain_m_exp4 = 0
	I = Imax.*fa.*exp.(-1.5*t)
	program, d = fm_program_lp(fc, fm, lp_exp; gain_m_exp4=gain_m_exp4, m_inst = inst_sina2)
	synthesize!("../wav/keyboard2.wav", program, t, [d[:gain_m] => I/(pi*4^gain_m_exp4), d[:gain_c] => V])
	program, d = fm_program_lp(fc*2, fm*2, lp_exp; gain_m_exp4=gain_m_exp4, m_inst = inst_sina2)
	synthesize!("../wav/keyboard2-h.wav", program, t, [d[:gain_m] => I/(pi*4^gain_m_exp4), d[:gain_c] => V])
end

function sinewave()
	fc = A4
	T = 4
	t = 0:h_env:T

	freq_c, exp_c = f_to_freq_exp(fc)

	gain_c = Fix(1.0)

	phase_c_addr = 1
	program = [
		(Instruction(inst_acc_imm_to_mem, IContext(phase_c_addr, exp_c)), freq_c),
		(inst_sin, Fix(0)),
		(inst_output_a, Fix(0)),
		(inst_dither_update, Fix(2^-15))
	]

	synthesize!("../wav/sine-approx.wav", program, t, [])

	t = (0:fs*T)/fs
	phase = 2*fc*t
	phase = round.(phase*32768)/32768
	s = sin.(pi*phase)
	wavwrite(s, "../wav/sinewave.wav"; Fs=fs)
end

bell()
#wood_drum()
brass()
clarinet()
bassoon()
bass()
drum()
snare()
hihat()
keyboard()
sinewave()

nothing

max_stars_repo_name: toivoh/fmicrosynth

max_issues_repo_name: toivoh/fmicrosynth

max_forks_repo_name: toivoh/fmicrosynth


text:
#=

Author: Shunsuke Hori

=#

module CES

using LS

function baseline(;
    beta = 0.9,
    sigma = 2.0,
    gamma = 2.0,
    Pi = 0.5 *ones(2,2),
    G = [0.1, 0.2],
    Theta = ones(Float64, 2),
    transfers = false
    )
    function U(c,n)
        if sigma == 1.0
            U = log(c)
        else
            U = (c.^(1.0-sigma)-1.0)/(1.0-sigma)
        end
        return U - n.^(1+gamma)/(1+gamma)
    end
    #derivatives of utiltiy function
    Uc(c,n) =  c.^(-sigma)
    Ucc(c,n) = -sigma*c.^(-sigma-1.0)
    Un(c,n) = -n.^gamma
    Unn(c,n) = -gamma*n.^(gamma-1.0)
    n_less_than_one = false
    return LS.Para(beta, Pi, G, Theta, transfers,
                    U, Uc, Ucc, Un, Unn, n_less_than_one)
end

#Model 1
M1 = baseline()

#Model 2

M2 = baseline(G=[0.15], Pi=ones(1,1), Theta=[1.0])

#Model 3 with time varying

M_time_example = baseline(G=[0.1, 0.1, 0.1, 0.2, 0.1, 0.1],
                            Theta = ones(6)# Theta can in principle be random
                            )

M_time_example.Pi = [0.0 1.0 0.0 0.0 0.0 0.0;
                     0.0 0.0 1.0 0.0 0.0 0.0;
                     0.0 0.0 0.0 0.5 0.5 0.0;
                     0.0 0.0 0.0 0.0 0.0 1.0;
                     0.0 0.0 0.0 0.0 0.0 1.0;
                     0.0 0.0 0.0 0.0 0.0 1.0]
end


text:
global SANITIZE_NAMES = false

function set_sanitize_names_option(value::Bool = false)
    global SANITIZE_NAMES
    SANITIZE_NAMES = value
    return nothing
end

# Find columns in a `TS` object corresponding to the given indexing indicator
findcols(c::C, x::TS) where {C<:Symbol} = x.fields .== c
findcols(c::C, x::TS) where {C<:AbstractVector{<:Symbol}} = map(sym->sym in c, x.fields)
findcols(c::C, x::TS) where {C<:Int} = [c]
findcols(c::C, x::TS) where {C<:AbstractVector{<:Integer}} = c

# Find rows in a `TS` object corresponding to the given indexing indicator
findrows(t::T, x::TS) where {T<:AbstractVector{<:TimeType}} = map(ti->ti in t, x.index)
findrows(t::T, x::TS) where {T<:TimeType} = x.index .== t
findrows(r::R, x::TS) where {R<:Int} = [r]
findrows(r::R, x::TS) where {R<:AbstractVector{<:Integer}} = r
findrows(s::S, x::TS) where {S<:AbstractString} = dtidx(s, x.index)

# Find the alphanumeric characters be able to generate sanitized column names (`fields`)
function findalphanum(s::AbstractString, drop_underscores::Bool=false)
   if drop_underscores
       return sort(union(find(isalpha,s), find(isnumber,s)))
   else
       return sort(union(union(find(isalpha,s), find(isnumber,s)), find(c->c=='_', s)))
   end
end

# Find the alphanumeric characters be able to generate sanitized column names (`fields`)
findalphanum(s::String)::Vector{Int} = find(isalpha.(split(s,"")).+isnumber.(split(s,"")))

# Change a String or Symbol to an alphanumeric-only version
namefix(s::AbstractString)::AbstractString = s[findalphanum(s)]

# Change a String or Symbol to an alphanumeric-only version
namefix(s::Symbol)::Symbol = Symbol(namefix(string(s)))

# Convert all column names (`field`s) of a TS object to alphanumeric-only `Symbol`s
namefix!(x::TS)::Nothing = x.fields = namefix.(x.fields)

# Generate automatic column names following the Excel spreadsheet naming convention (A,B,C,..,X,Y,Z,AA,AB,AC,...)
function autocol(col::Int)
    @assert col >= 1 "Invalid column number $col - cannot generate column name"
    if col <= 26
        return Symbol(Char(64 + col))
    end
    colname = ""
    modulo = 0
    dividend = col
    while dividend > 0
        modulo = (dividend - 1) % 26
        colname = string(Char(65 + modulo)) * colname
        dividend = Int(round((dividend - modulo) / 26))
    end
    return Symbol(colname)
end
autocol(cols::AbstractArray{Int,1}) = autocol.(cols)

# Automatically generate an `index` vector of `Date`/`DateTime` values from today's date back through $n$ days
autoidx(n::Int; dt::Period=Day(1), from::Date=today()-(n-1)*dt, thru::Date=from+(n-1)*dt) = collect(from:dt:thru)

# function overlaps(x::AbstractArray, y::AbstractArray, n::Int=1)::Vector{Bool}
#     @assert n == 1 || n == 2
#     if n == 1
#         xx = falses(x)
#         @inbounds for i = 1:size(x,1), j = 1:size(y,1)
#             if x[i] == y[j]
#                 xx[i] = true
#             end
#         end
#         return xx
#     elseif n == 2
#         yy = falses(y)
#         @inbounds for i = 1:size(x,1), j = 1:size(y,1)
#             if x[i] == y[j]
#                 yy[i] = true
#             end
#         end
#         return yy
#     end
# end


text: _1 = 4

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:
function diversificationCore(wgts::Array{Float64, 2})
    ## calculates several diversification measures
    ## Note: low-level function not to be called directly!
    ##
    ## Inputs:
    ## 	wgts 	nObs x nAss matrix of portfolio weights
    ##
    ## Outputs:
    ## 	DataFrame containing information on number of assets
    ## 	divIndicators table containing information on number of
    ## 	assets, number of significant assets, maximum weight,
    ## 	standard deviation of weights, number of negative
    ## 	weights and gini coefficient (herfindahl) for each day

    ## tolerance level for significant assets
    sigLev = 0.0001;

    ## transpose if column vector
    if size(wgts, 2) == 1
        wgts = wgts'
    end

    ## basic statistics
    nAss = sum(wgts .!= 0, 2);
    nSignAss = sum(abs(wgts .- 0) .> sigLev, 2)
    maxWgt = maximum(wgts, 2)
    stdWgt = std(wgts, 2)
    nShort = sum(wgts .< 0, 2)

    ## gini coefficient with negative weights
    absExposure = sum(abs(wgts), 2);
    relExposures = wgts ./ repmat(absExposure, 1, size(wgts, 2))
    gini = sum(relExposures.^2, 2);

    divIndicators = DataFrame()
    divIndicators[:nAss] = nAss[:]
    divIndicators[:nSignAss] = nSignAss[:]
    divIndicators[:maxWgt] = maxWgt[:]
    divIndicators[:stdWgt] = stdWgt[:]
    divIndicators[:nShort] = nShort[:]
    divIndicators[:Gini] = gini[:]

    return divIndicators
end

function diversification(invs::Investments)
    ## calculates several diversification measures
    ##
    ## Inputs:
    ## 	wgts 	nObs x nAss matrix of portfolio weights
    ##
    ## Outputs:
    ## 	DataFrame containing information on number of assets
    ## 	divIndicators table containing information on number of
    ## 	assets, number of significant assets, maximum weight,
    ## 	standard deviation of weights, number of negative
    ## 	weights and gini coefficient (herfindahl) for each day

    wgts = AssetMgmt.core(invs)

    divIndicators = diversificationCore(wgts)
    return Timematr(divIndicators, AssetMgmt.idx(invs))
end


text:
"""
	PosEvo(fp::FluPop, i::Int64;
			ambiguous = false,
			threshold = 0.05)
"""
function PosEvo(fp::FluPop, i::Int64;
			ambiguous = false,
			threshold = 0.05)
	A = eltype(first(fp.strains)[2].seq)
	out = PosEvo(A,i)
	#
	if isempty(fp.datebin)
		@error "`fp.datebin` is empty"
	end
	#
	datebins = sort(collect(keys(fp.datebin)))
	for (j,d) in enumerate(datebins)
		S = fp.datebin[d]
		if !isempty(S) || j == 1
			# Frequencies for that site and timebin
			f = SiteFrequency(A, i, 0)
			for s in S
				if ambiguous || !isambiguous(s.seq[i])
					f[s.seq[i]] = get(f, s.seq[i], 0.) + 1.
					f.M += 1
				end
			end
			for a in BioTools.alphabet(f)
				f[a] /= f.M
				if !in(a, out.alphabet)
					push!(out.alphabet, a)
				end
			end
			#
			out.data[d] = f
		else j > 1 # Use frequencies of previous datebin (i>1 here)
			out.data[d] = deepcopy(out.data[datebins[j-1]])
		end
	end
	remove_rare_symbols!(out, threshold)
	return out
end
function PosEvo(fp::FluPop;
				ambiguous=false,
				threshold=0.05)
	A = eltype(first(fp.strains)[2].seq)
	ph = PosEvo{A}[]
	for i in 1:length(first(fp.strains)[2].seq)
	    print("$i       \r")
	    push!(ph, PosEvo(fp, i, ambiguous=ambiguous, threshold=threshold))
	end
	return ph
end
# Regional weights could be implemented at the PosEvo level
# For instance, a function `PosEvo(fp, i, weights)`
# It's the only place where we compute frequencies, and we also know the strains since we have `fp`

# Why is this useful? I'll leave it empty for now...
function remove_rare_symbols!(ph::PosEvo, threshold)

end

"""
	frequency_series(ph::PosEvo)

Return arrays `X`, `Y` and `Z`, with `X` containing dates, `Y` a matrix with columns containing frequencies of symbols, and `Z` a vector containing population sizes frequencies are based on.
"""
function frequency_series(ph::PosEvo; freq_threshold = 0.05)
	X = Array{Date,1}(undef, 0)
	Y = zeros(Float64, length(ph.data), length(ph.alphabet))
	pop = Array{Int64,1}(undef, 0)
	for (i, (d,f)) in enumerate(sort(OrderedDict(ph.data)))
		push!(pop, f.M)
		push!(X, datebin_to_date(d))
		for (k,a) in enumerate(ph.alphabet)
			Y[i,k] = f[a]
		end
	end
	# Delete positions with frequencies always smaller than `freq_threshold`
	if freq_threshold > 0.
		idx = Int64[]
		for i in 1:size(Y,2)
			for t in 1:size(Y,1)
				if Y[t,i] > freq_threshold
					push!(idx, i)
					break
				end
			end
		end
	else
		idx = 1:length(Y,2)
	end
	return (X,Y[:,idx],pop,ph.alphabet[idx])
end

"""
	entropy(Z::PosEvo)

Return a dictionary `Date => Float` giving the entropy as a function of time for position history `Z`.
"""
function entropy(Z::PosEvo; freq_threshold = 0.05)
	X,Y = frequency_series(Z, freq_threshold=freq_threshold)[1:2]
	S = zeros(Float64, length(X))
	for i in 1:length(S)
		f = Y[i,:] ./ sum(Y[i,:])
		S[i] = isnan(StatsBase.entropy(f)) ? 0. : StatsBase.entropy(f)
	end
	return S,X
end

max_stars_repo_name: PierreBarrat/FluPredictibility

max_issues_repo_name: PierreBarrat/FluPredictibility

max_forks_repo_name: PierreBarrat/FluPredictibility


text:
# +
import Revise
import GLRenderer
import PoseComposition
import Rotations
import Images
I = Images

R = Rotations
P = PoseComposition
GL = GLRenderer
# -

obj_path = joinpath(@__DIR__, "035_power_drill/textured_simple.obj")
texture_path = joinpath(@__DIR__, "035_power_drill/texture_map.png")

# +
camera_intrinsics = GL.CameraIntrinsics(
    640, 480,
    1000.0, 1000.0,
    320.0, 240.0,
    0.01, 5.0
)

renderer = GL.setup_renderer(camera_intrinsics, GL.TextureMixedMode())

mesh_data = GL.get_mesh_data_from_obj_file(obj_path; tex_path=texture_path)
GL.load_object!(renderer, mesh_data)

rgb_image, depth_image = GL.gl_render(
    renderer, [1], [P.Pose([0.0, 0.0, 1.0], R.RotXYZ(0.1, 0.4, 0.9))],
    P.IDENTITY_POSE
)
img = I.colorview(I.RGBA, permutedims(rgb_image,(3,1,2)))

# +
camera_intrinsics = GL.CameraIntrinsics(
    640, 480,
    1000.0, 1000.0,
    320.0, 240.0,
    0.01, 30.0
)

renderer = GL.setup_renderer(camera_intrinsics, GL.TextureMixedMode())

mesh_data = GL.get_mesh_data_from_obj_file(obj_path; tex_path=texture_path)
GL.load_object!(renderer, mesh_data)

box_mesh = GL.box_mesh_from_dims([10.0, 0.5, 5.0])
GL.load_object!(renderer, box_mesh)

rgb_image, depth_image = GL.gl_render(
    renderer, [1,2], [P.Pose([0.0, 0.0, 1.0], R.RotXYZ(0.1, 0.4, 0.9)), P.Pose([0.0, 2.0, 10.0],P.IDENTITY_ORN)],
    P.IDENTITY_POSE; colors = [I.colorant"red" for _ in 1:2]
)

img = I.colorview(I.RGBA, permutedims(rgb_image,(3,1,2)))
I.save(joinpath(@__DIR__, "imgs/texture_mixed_rendering_rgb_image.png"), img)

img = I.colorview(I.Gray, depth_image ./ maximum(depth_image))
I.save(joinpath(@__DIR__, "imgs/texture_mixed_rendering_depth_image.png"), img)


text:
function inex03()
#
#  test polynomial suggested by Goedecker
#
    inex(5)
end


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

# ----------------
# IMPLEMENTATIONS
# ----------------
include("solutions/estimation.jl")
include("solutions/simulation.jl")


text:
@info "Hello, World!"
@info ARGS


text:
module SIIP_Tutorial

export precompile_tutorial

using PowerSystems
using PowerSystemCaseBuilder
using PowerSimulations
using PowerGraphics
using TimeSeries
using DataFrames
using Plots
using HiGHS
using Logging

function precompile_tutorial()
    system = build_system(SIIPExampleSystems, "5_bus_matpower_DA"; force_build = true, skip_serialization = true)
    template = template_unit_commitment()
    optimizer = optimizer_with_attributes(HiGHS.Optimizer, "mip_rel_gap" => 0.05)
    model = DecisionModel(
        template,
        system;
        optimizer = optimizer,
        horizon = 1
    )
    build!(model, output_dir = mktempdir())
    solve!(model)
    ProblemResults(model)
    nothing
end

end # module SIIP_Tutorial

max_stars_repo_name: tlienart/SIIP-Tutorial

max_issues_repo_name: tlienart/SIIP-Tutorial

max_forks_repo_name: tlienart/SIIP-Tutorial


text:
## Transforms take values at Chebyshev points of the first and second kinds and produce Chebyshev coefficients

abstract type ChebyshevPlan{T} <: Plan{T} end

size(P::ChebyshevPlan) = isdefined(P, :plan) ? size(P.plan) : (0,)
length(P::ChebyshevPlan) = isdefined(P, :plan) ? length(P.plan) : 0

# Check whether a ChebyshevPlan is applicable to a given input array, and
# throw an informative error if not:
function assert_applicable(p::ChebyshevPlan{T}, X::StridedArray{T}) where T
    if size(X) != size(p)
        throw(ArgumentError("Chebyshev plan applied to wrong-size array"))
    end
end

struct ChebyshevTransformPlan{T,kind,inplace,P} <: ChebyshevPlan{T}
    plan::FFTW.r2rFFTWPlan{T,P,true,1,UnitRange{Int}}
    ChebyshevTransformPlan{T,kind,inplace,P}(plan) where {T,kind,inplace,P} = new{T,kind,inplace,P}(plan)
    ChebyshevTransformPlan{T,kind,inplace,P}() where {T,kind,inplace,P} = new{T,kind,inplace,P}()
end

ChebyshevTransformPlan{T,kind,inplace}(plan::FFTW.r2rFFTWPlan{T,P}) where {T,kind,inplace,P} =
    ChebyshevTransformPlan{T,kind,inplace,P}(plan)

ChebyshevTransformPlan{T,kind,inplace}(plan::ChebyshevTransformPlan{T,kind,inp,P}) where {T,kind,inplace,inp,P} =
    ChebyshevTransformPlan{T,kind,inplace,P}(plan.plan)

function plan_chebyshevtransform!(x::AbstractVector{T}, ::Val{1}) where T<:fftwNumber
    if isempty(x)
        ChebyshevTransformPlan{T,1,true,(5,)}()
    else
        ChebyshevTransformPlan{T,1,true,(5,)}(FFTW.plan_r2r!(x, FFTW.REDFT10))
    end
end
function plan_chebyshevtransform!(x::AbstractVector{T}, ::Val{2}) where T<:fftwNumber
    length(x) ≤ 1 && throw(ArgumentError("Vector must contain at least 2 entries"))
    ChebyshevTransformPlan{T,2,true,(3,)}(FFTW.plan_r2r!(x, FFTW.REDFT00))
end

function plan_chebyshevtransform(x::AbstractVector{T}, ::Val{1}) where T<:fftwNumber
    if isempty(x)
        ChebyshevTransformPlan{T,1,false,(5,)}()
    else
        ChebyshevTransformPlan{T,1,false,(5,)}(FFTW.plan_r2r!(x, FFTW.REDFT10))
    end
end
function plan_chebyshevtransform(x::AbstractVector{T}, ::Val{2}) where T<:fftwNumber
    length(x) ≤ 1 && throw(ArgumentError("Vector must contain at least 2 entries"))
    ChebyshevTransformPlan{T,2,false,(3,)}(FFTW.plan_r2r!(x, FFTW.REDFT00))
end

plan_chebyshevtransform!(x::AbstractVector) = plan_chebyshevtransform!(x, Val(1))
plan_chebyshevtransform(x::AbstractVector) = plan_chebyshevtransform(x, Val(1))

function *(P::ChebyshevTransformPlan{T,1,true}, x::AbstractVector{T}) where T
    n = length(x)
    assert_applicable(P, x)
    n == 0 && return x

    y = P.plan*x # will be  === x if in-place
    y[1] /= 2
    lmul!(inv(convert(T,n)), y)
end

function *(P::ChebyshevTransformPlan{T,2,true}, x::AbstractVector{T}) where T
    n = length(x)
    y = P.plan*x # will be  === x if in-place
    y[1] /= 2; y[end] /= 2
    lmul!(inv(convert(T,n-1)),y)
end

*(P::ChebyshevTransformPlan{T,k,false}, x::AbstractVector{T}) where {T,k} =
    ChebyshevTransformPlan{T,k,true}(P)*copy(x)

chebyshevtransform!(x::AbstractVector{T}, kind=Val(1)) where T =
    plan_chebyshevtransform!(x, kind)*x

"""
    chebyshevtransform(x, kind=Val(1))

transforms from values on a Chebyshev grid of the first or second kind to Chebyshev
coefficients.
"""
chebyshevtransform(x, kind=Val(1)) = chebyshevtransform!(copy(x), kind)

## Inverse transforms take Chebyshev coefficients and produce values at Chebyshev points of the first and second kinds

struct IChebyshevTransformPlan{T,kind,inplace,P} <: ChebyshevPlan{T}
    plan::FFTW.r2rFFTWPlan{T,P,true,1,UnitRange{Int}}
    IChebyshevTransformPlan{T,kind,inplace,P}(plan) where {T,kind,inplace,P} = new{T,kind,inplace,P}(plan)
    IChebyshevTransformPlan{T,kind,inplace,P}() where {T,kind,inplace,P} = new{T,kind,inplace,P}()
end

IChebyshevTransformPlan{T,kind,inplace}(F::FFTW.r2rFFTWPlan{T,P}) where {T,kind,inplace,P} =
    IChebyshevTransformPlan{T,kind,inplace,P}(F)

IChebyshevTransformPlan{T,kind,true}(F::IChebyshevTransformPlan{T,kind,false,P}) where {T,kind,P} =
    IChebyshevTransformPlan{T,kind,true,P}(F.plan)

size(P::IChebyshevTransformPlan) = isdefined(P, :plan) ? size(P.plan) : (0,)
length(P::IChebyshevTransformPlan) = isdefined(P, :plan) ? length(P.plan) : 0

# second kind Chebyshev transforms share a plan with their inverse
# so we support this via inv
inv(P::ChebyshevTransformPlan{T,2,inp}) where {T,inp} = IChebyshevTransformPlan{T,2,inp}(P.plan)
inv(P::IChebyshevTransformPlan{T,2,inp}) where {T,inp} = ChebyshevTransformPlan{T,2,inp}(P.plan)

\(P::ChebyshevTransformPlan, x::AbstractArray) = inv(P) * x
\(P::IChebyshevTransformPlan, x::AbstractArray) = inv(P) * x

function plan_ichebyshevtransform!(x::AbstractVector{T}, ::Val{1}) where T<:fftwNumber
    if isempty(x)
        IChebyshevTransformPlan{T,1,true,(4,)}()
    else
        IChebyshevTransformPlan{T,1,true,(4,)}(FFTW.plan_r2r!(x, FFTW.REDFT01))
    end
end

function plan_ichebyshevtransform!(x::AbstractVector{T}, ::Val{2}) where T<:fftwNumber
    inv(plan_chebyshevtransform!(x, Val(2)))
end

function plan_ichebyshevtransform(x::AbstractVector{T}, ::Val{1}) where T<:fftwNumber
    if isempty(x)
        IChebyshevTransformPlan{T,1,false,(4,)}()
    else
        IChebyshevTransformPlan{T,1,false,(4,)}(FFTW.plan_r2r!(x, FFTW.REDFT01))
    end
end

function plan_ichebyshevtransform(x::AbstractVector{T}, ::Val{2}) where T<:fftwNumber
    inv(plan_chebyshevtransform(x, Val(2)))
end

plan_ichebyshevtransform!(x::AbstractVector) = plan_ichebyshevtransform!(x, Val(1))
plan_ichebyshevtransform(x::AbstractVector) = plan_ichebyshevtransform(x, Val(1))

function *(P::IChebyshevTransformPlan{T,1,true}, x::AbstractVector{T}) where T<:fftwNumber
    n = length(x)
    assert_applicable(P, x)
    n == 0 && return x

    x[1] *= 2
    x = lmul!(convert(T,0.5), P.plan*x)
    x
end
function *(P::IChebyshevTransformPlan{T,2, true}, x::AbstractVector{T}) where T<:fftwNumber
    n = length(x)
    assert_applicable(P, x)

    x[1] *= 2; x[end] *= 2
    x = ChebyshevTransformPlan{T,2,true}(P.plan)*x
    x[1] *= 2; x[end] *= 2
    lmul!(convert(T,0.5(n-1)),x)
end

*(P::IChebyshevTransformPlan{T,k,false},x::AbstractVector{T}) where {T,k} =
    IChebyshevTransformPlan{T,k,true}(P)*copy(x)

ichebyshevtransform!(x::AbstractVector{T}, kind=Val(1)) where T =
    plan_ichebyshevtransform!(x, kind)*x

ichebyshevtransform(x, kind=Val(1)) = ichebyshevtransform!(copy(x), kind)

# Matrix inputs
#
#
function chebyshevtransform!(X::AbstractMatrix{T}, ::Val{1}) where T<:fftwNumber
    if size(X) == (1,1)
        X
    else
        X=FFTW.r2r!(X,FFTW.REDFT10)
        X[:,1]/=2;X[1,:]/=2;
        lmul!(1/(size(X,1)*size(X,2)),X)
    end
end

function chebyshevtransform!(X::AbstractMatrix{T}, ::Val{2}) where T<:fftwNumber
    if size(X,1) < 2 || size(X,2) < 2
        throw(ArgumentError("Chebyshev plan applied to wrong-size array"))
    else
        X=FFTW.r2r!(X,FFTW.REDFT00)
        lmul!(1/((size(X,1)-1)*(size(X,2)-1)),X)
        X[:,1]/=2;X[:,end]/=2
        X[1,:]/=2;X[end,:]/=2
        X
    end
end
#
function ichebyshevtransform!(X::AbstractMatrix{T}, ::Val{1}) where T<:fftwNumber
    if size(X) == (1,1)
        X
    else
        X[1,:]*=2;X[:,1]*=2
        X = FFTW.r2r(X,FFTW.REDFT01)
        lmul!(0.25, X)
    end
end
function ichebyshevtransform!(X::AbstractMatrix{T}, ::Val{2}) where T<:fftwNumber
    if size(X,1) < 2 || size(X,2) < 2
        throw(ArgumentError("Chebyshev plan applied to wrong-size array"))
    else
        X[1,:]*=2;X[end,:]*=2;X[:,1]*=2;X[:,end]*=2
        X=chebyshevtransform!(X, Val(2))
        X[1,:]*=2;X[end,:]*=2;X[:,1]*=2;X[:,end]*=2
        lmul!((size(X,1)-1)*(size(X,2)-1)/4,X)
    end
end
#

## Chebyshev U

struct ChebyshevUTransformPlan{T,kind,inplace,P} <: ChebyshevPlan{T}
    plan::FFTW.r2rFFTWPlan{T,P,true,1,UnitRange{Int}}
    ChebyshevUTransformPlan{T,kind,inplace,P}(plan) where {T,kind,inplace,P} = new{T,kind,inplace,P}(plan)
    ChebyshevUTransformPlan{T,kind,inplace,P}() where {T,kind,inplace,P} = new{T,kind,inplace,P}()
end

ChebyshevUTransformPlan{T,kind,inplace}(plan::FFTW.r2rFFTWPlan{T,P}) where {T,kind,inplace,P} =
    ChebyshevUTransformPlan{T,kind,inplace,P}(plan)

ChebyshevUTransformPlan{T,kind,inplace}(plan::ChebyshevUTransformPlan{T,kind,inp,P}) where {T,kind,inplace,inp,P} =
    ChebyshevUTransformPlan{T,kind,inplace,P}(plan.plan)

function plan_chebyshevutransform!(x::AbstractVector{T}, ::Val{1}) where T<:fftwNumber
    if isempty(x)
        ChebyshevUTransformPlan{T,1,true,(9,)}()
    else
        ChebyshevUTransformPlan{T,1,true,(9,)}(FFTW.plan_r2r!(x, FFTW.RODFT10))
    end
end
function plan_chebyshevutransform!(x::AbstractVector{T}, ::Val{2}) where T<:fftwNumber
    length(x) ≤ 1 && throw(ArgumentError("Vector must contain at least 2 entries"))
    ChebyshevUTransformPlan{T,2,true,(7,)}(FFTW.plan_r2r!(x, FFTW.RODFT00))
end

function plan_chebyshevutransform(x::AbstractVector{T}, ::Val{1}) where T<:fftwNumber
    if isempty(x)
        ChebyshevUTransformPlan{T,1,false,(9,)}()
    else
        ChebyshevUTransformPlan{T,1,false,(9,)}(FFTW.plan_r2r!(x, FFTW.RODFT10))
    end
end
function plan_chebyshevutransform(x::AbstractVector{T}, ::Val{2}) where T<:fftwNumber
    length(x) ≤ 1 && throw(ArgumentError("Vector must contain at least 2 entries"))
    ChebyshevUTransformPlan{T,2,false,(7,)}(FFTW.plan_r2r!(x, FFTW.RODFT00))
end

plan_chebyshevutransform!(x::AbstractVector) = plan_chebyshevutransform!(x, Val(1))
plan_chebyshevutransform(x::AbstractVector) = plan_chebyshevutransform(x, Val(1))

function *(P::ChebyshevUTransformPlan{T,1,true},x::AbstractVector{T}) where T
    n = length(x)
    assert_applicable(P, x)
    n ≤ 1 && return x

    for k=1:n # sqrt(1-x_j^2) weight
        x[k] *= sinpi(one(T)/(2n) + (k-one(T))/n)/n
    end
    P.plan * x
end

function *(P::ChebyshevUTransformPlan{T,2,true}, x::AbstractVector{T}) where T
    n = length(x)
    assert_applicable(P, x)
    n ≤ 1 && return x

    c = one(T)/ (n+1)
    for k=1:n # sqrt(1-x_j^2) weight
        x[k] *= sinpi(k*c)
    end
    lmul!(c, P.plan * x)
end

chebyshevutransform!(x::AbstractVector{T}, kind=Val(1)) where {T<:fftwNumber} =
    plan_chebyshevutransform!(x, kind)*x

"""
    chebyshevutransform(x, ::Val{kind}=Val(1))

transforms from values on a Chebyshev grid of the first or second kind to Chebyshev
coefficients of the 2nd kind (Chebyshev U expansion).
"""
chebyshevutransform(x, kind=Val(1)) = chebyshevutransform!(copy(x), kind)

*(P::ChebyshevUTransformPlan{T,k,false}, x::AbstractVector{T}) where {T,k} = ChebyshevUTransformPlan{T,k,true}(P)*copy(x)

## Inverse transforms take ChebyshevU coefficients and produce values at ChebyshevU points of the first and second kinds

struct IChebyshevUTransformPlan{T,kind,inplace,P} <: ChebyshevPlan{T}
    plan::FFTW.r2rFFTWPlan{T,P,true,1,UnitRange{Int}}
    IChebyshevUTransformPlan{T,kind,inplace,P}(plan) where {T,kind,inplace,P} = new{T,kind,inplace,P}(plan)
    IChebyshevUTransformPlan{T,kind,inplace,P}() where {T,kind,inplace,P} = new{T,kind,inplace,P}()
end

IChebyshevUTransformPlan{T,kind,inplace}(F::FFTW.r2rFFTWPlan{T,P}) where {T,kind,inplace,P} =
    IChebyshevUTransformPlan{T,kind,inplace,P}(F)

IChebyshevUTransformPlan{T,kind,true}(F::IChebyshevUTransformPlan{T,kind,false,P}) where {T,kind,P} =
    IChebyshevUTransformPlan{T,kind,true,P}(F.plan)

function plan_ichebyshevutransform!(x::AbstractVector{T}, ::Val{1}) where T<:fftwNumber
    if isempty(x)
        IChebyshevUTransformPlan{T,1,true,(8,)}()
    else
        IChebyshevUTransformPlan{T,1,true,(8,)}(FFTW.plan_r2r!(x, FFTW.RODFT01))
    end
end
function plan_ichebyshevutransform!(x::AbstractVector{T}, ::Val{2}) where T<:fftwNumber
    length(x) ≤ 1 && throw(ArgumentError("Vector must contain at least 2 entries"))
    IChebyshevUTransformPlan{T,2,true,(7,)}(FFTW.plan_r2r!(x, FFTW.RODFT00))
end

function plan_ichebyshevutransform(x::AbstractVector{T}, ::Val{1}) where T<:fftwNumber
    if isempty(x)
        IChebyshevUTransformPlan{T,1,false,(8,)}()
    else
        IChebyshevUTransformPlan{T,1,false,(8,)}(FFTW.plan_r2r!(x, FFTW.RODFT01))
    end
end
function plan_ichebyshevutransform(x::AbstractVector{T}, ::Val{2}) where T<:fftwNumber
    length(x) ≤ 1 && throw(ArgumentError("Vector must contain at least 2 entries"))
    IChebyshevUTransformPlan{T,2,false,(7,)}(FFTW.plan_r2r!(x, FFTW.RODFT00))
end

plan_ichebyshevutransform!(x::AbstractVector) = plan_ichebyshevutransform!(x, Val(1))
plan_ichebyshevutransform(x::AbstractVector) = plan_ichebyshevutransform(x, Val(1))

function *(P::IChebyshevUTransformPlan{T,1,true}, x::AbstractVector{T}) where T<:fftwNumber
    n = length(x)
    assert_applicable(P, x)
    n ≤ 1 && return x

    x = P.plan * x
    for k=1:n # sqrt(1-x_j^2) weight
        x[k] /= 2sinpi(one(T)/(2n) + (k-one(T))/n)
    end
    x
end

function *(P::IChebyshevUTransformPlan{T,2,true}, x::AbstractVector{T}) where T<:fftwNumber
    n = length(x)
    assert_applicable(P, x)
    n ≤ 1 && return x

    c = one(T)/ (n+1)
    lmul!((n+1)/(2n+2*one(T)), x)
    x = P.plan * x
    for k=1:n # sqrt(1-x_j^2) weight
        x[k] /= sinpi(k*c)
    end
    x
end

ichebyshevutransform!(x::AbstractVector{T}, kind=Val(1)) where {T<:fftwNumber} =
    plan_ichebyshevutransform!(x, kind)*x

ichebyshevutransform(x, kind=Val(1)) = ichebyshevutransform!(copy(x), kind)

*(P::IChebyshevUTransformPlan{T,k,false},x::AbstractVector{T}) where {T,k} =
    IChebyshevUTransformPlan{T,k,true}(P)*copy(x)

## Code generation for integer inputs

for func in (:chebyshevtransform,:ichebyshevtransform,:chebyshevutransform,:ichebyshevutransform)
    @eval $func(x::AbstractVector{T}, kind=Val(1)) where {T<:Integer} = $func(convert(AbstractVector{Float64},x), kind)
end

## points

struct ChebyshevGrid{kind,T} <: AbstractVector{T}
    n::Int
    function ChebyshevGrid{1,T}(n::Int) where T
        n ≥ 0 || throw(ArgumentError("Number of points must be nonnehative"))
        new{1,T}(n)
    end
    function ChebyshevGrid{2,T}(n::Int) where T
        n ≥ 2 || throw(ArgumentError("Number of points must be greater than 2"))
        new{2,T}(n)
    end
end

ChebyshevGrid{kind}(n::Integer) where kind = ChebyshevGrid{kind,Float64}(n)

size(g::ChebyshevGrid) = (g.n,)
getindex(g::ChebyshevGrid{1,T}, k::Integer) where T =
    sinpi(convert(T,g.n-2k+1)/(2g.n))

getindex(g::ChebyshevGrid{2,T}, k::Integer) where T =
    sinpi(convert(T,g.n-2k+1)/(2g.n-2))

chebyshevpoints(::Type{T}, n::Integer, ::Val{kind}) where {T<:Number,kind} = ChebyshevGrid{kind,T}(n)
chebyshevpoints(::Type{T}, n::Integer) where T = chebyshevpoints(T, n, Val(1))
chebyshevpoints(n::Integer, kind=Val(1)) = chebyshevpoints(Float64, n, kind)

# sin(nθ) coefficients to values at Clenshaw-Curtis nodes except ±1
#
# struct DSTPlan{T,kind,inplace,P} <: Plan{T}
#     plan::P
# end
#
# DSTPlan{k,inp}(plan) where {k,inp} =
#     DSTPlan{eltype(plan),k,inp,typeof(plan)}(plan)
#
#
# plan_DSTI!(x) = length(x) > 0 ? DSTPlan{1,true}(FFTW.FFTW.plan_r2r!(x, FFTW.FFTW.RODFT00)) :
#                                 fill(one(T),1,length(x))
#
# function *(P::DSTPlan{T,1}, x::AbstractArray) where {T}
#     x = P.plan*x
#     rmul!(x,half(T))
# end

###
# BigFloat
# Use `Nothing` and fall back too FFT
###

plan_chebyshevtransform(x::AbstractVector{T}, ::Val{kind}) where {T,kind} =
    ChebyshevTransformPlan{T,kind,false,Nothing}()
plan_ichebyshevtransform(x::AbstractVector{T}, ::Val{kind}) where {T,kind} =
    IChebyshevTransformPlan{T,kind,false,Nothing}()

plan_chebyshevtransform!(x::AbstractVector{T}, ::Val{kind}) where {T,kind} =
    ChebyshevTransformPlan{T,kind,true,Nothing}()
plan_ichebyshevtransform!(x::AbstractVector{T}, ::Val{kind}) where {T,kind} =
    IChebyshevTransformPlan{T,kind,true,Nothing}()

#following Chebfun's @Chebtech1/vals2coeffs.m and @Chebtech2/vals2coeffs.m
function *(P::ChebyshevTransformPlan{T,1,false,Nothing}, x::AbstractVector{T}) where T
    n = length(x)
    if n == 1
        x
    else
        w = [2exp(im*convert(T,π)*k/2n) for k=0:n-1]
        ret = w.*ifft([x;reverse(x)])[1:n]
        ret = T<:Real ? real(ret) : ret
        ret[1] /= 2
        ret
    end
end

# function *(P::ChebyshevTransformPlan{T,2,false,Nothing}, x::AbstractVector{T}) where T
#     n = length(x)
#     if n == 1
#         x
#     else
#         ret = ifft([x;x[end:-1:2]])[1:n]
#         ret = T<:Real ? real(ret) : ret
#         ret[2:n-1] *= 2
#         ret
#     end
# end

*(P::ChebyshevTransformPlan{T,1,true,Nothing}, x::AbstractVector{T}) where T =
    copyto!(x, ChebyshevTransformPlan{T,1,false,Nothing}() * x)
# *(P::ChebyshevTransformPlan{T,2,true,Nothing}, x::AbstractVector{T}) where T =
#     copyto!(x, ChebyshevTransformPlan{T,2,false,Nothing}() * x)

#following Chebfun's @Chebtech1/vals2coeffs.m and @Chebtech2/vals2coeffs.m
function *(P::IChebyshevTransformPlan{T,1,false,Nothing}, x::AbstractVector{T}) where T
    n = length(x)
    if n == 1
        x
    else
        w = [exp(-im*convert(T,π)*k/2n)/2 for k=0:2n-1]
        w[1] *= 2;w[n+1] *= 0;w[n+2:end] *= -1
        ret = fft(w.*[x;one(T);x[end:-1:2]])
        ret = T<:Real ? real(ret) : ret
        ret[1:n]
    end
end

# function *(P::IChebyshevTransformPlan{T,2,true,Nothing}, x::AbstractVector{T}) where T
#     n = length(x)
#     if n == 1
#         x
#     else
#         x[1] *= 2; x[end] *= 2
#         chebyshevtransform!(x, Val(2))
#         x[1] *= 2; x[end] *= 2
#         lmul!(convert(T,n-1)/2, x)
#         x
#     end
# end

*(P::IChebyshevTransformPlan{T,1,true,Nothing}, x::AbstractVector{T}) where T =
    copyto!(x, IChebyshevTransformPlan{T,1,false,Nothing}() * x)
# *(P::IChebyshevTransformPlan{T,2,false,Nothing}, x::AbstractVector{T}) where T =
#     IChebyshevTransformPlan{T,2,true,Nothing}() * copy(x)


text:
#!/usr/bin/env julia
#
# Jerzy Wroczyński (nr. indeksu 250075)
# Obliczenia Naukowe, Laboratorium
# Lista 3, Zadanie 5.
#

include("iterative-methods.jl")
using .IterativeMethods

fl = IterativeMethods.FL

r1, fr1, k1, err1 =  IterativeMethods.mbisekcji(
    x -> 3x - exp(x),
    fl(0),
    fl(1),
    fl(10^(-4)),
    fl(10^(-4))
)

println(
    "first root:\n",
    "  r = ", r1, "\n",
    "  f(r) = ", fr1, "\n",
    "  k = ", k1, "\n",
    err1 == 0 ? "" : string("  error code: ", err1)
)

r2, fr2, k2, err2 =  IterativeMethods.mbisekcji(
    x -> 3x - exp(x),
    fl(1),
    fl(2),
    fl(10^(-4)),
    fl(10^(-4))
)

println(
    "\nsecond root:\n",
    "  r = ", r2, "\n",
    "  f(r) = ", fr2, "\n",
    "  k = ", k2, "\n",
    err2 == 0 ? "" : string("  error code: ", err2)
)

max_stars_repo_name: jerry-sky/academic-notebook

max_issues_repo_name: jerry-sky/academic-notebook

max_forks_repo_name: jerry-sky/academic-notebook


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

using JuliaFEM
using JuliaFEM.Testing

@testset "geometry missing" begin
    el = Element(Quad4, [1, 2, 3, 4])
    pr = Problem(Elasticity, "problem", 2)
    # this throws KeyError: geometry not found.
    # it's descriptive enough to give hint to user
    # what went wrong
    @test_throws KeyError assemble!(pr, el)
end

@testset "connectivity information missing" begin
    el = Element(Quad4)
    nodes = Vector{Float64}[[0,0],[1,0],[1,1],[0,1]]
    update!(el, "geometry", nodes)
    pr = Problem(Elasticity, "problem", 2)
    @test_throws Exception assemble!(pr, el)
end


text:
@doc raw"""
    GCNConv(in => out, σ=identity; [bias, init, add_self_loops, use_edge_weight])

Graph convolutional layer from paper [Semi-supervised Classification with Graph Convolutional Networks](https://arxiv.org/abs/1609.02907).

Performs the operation
```math
\mathbf{x}'_i = \sum_{j\in N(i)} a_{ij} W \mathbf{x}_j
```
where ``a_{ij} = 1 / \sqrt{|N(i)||N(j)|}`` is a normalization factor computed from the node degrees.

If the input graph has weighted edges and `use_edge_weight=true`, than ``a_{ij}`` will be computed as
```math
a_{ij} = \frac{e_{j\to i}}{\sqrt{\sum_{j \in N(i)}  e_{j\to i}} \sqrt{\sum_{i \in N(j)}  e_{i\to j}}}
```

The input to the layer is a node feature array `X` of size `(num_features, num_nodes)`
and optionally an edge weight vector.

# Arguments

- `in`: Number of input features.
- `out`: Number of output features.
- `σ`: Activation function. Default `identity`.
- `bias`: Add learnable bias. Default `true`.
- `init`: Weights' initializer. Default `glorot_uniform`.
- `add_self_loops`: Add self loops to the graph before performing the convolution. Default `false`.
- `use_edge_weight`: If `true`, consider the edge weights in the input graph (if available).
                     If `add_self_loops=true` the new weights will be set to 1. Default `false`.

# Examples

```julia
# create data
s = [1,1,2,3]
t = [2,3,1,1]
g = GNNGraph(s, t)
x = randn(3, g.num_nodes)

# create layer
l = GCNConv(3 => 5)

# forward pass
y = l(g, x)       # size:  5 × num_nodes

# convolution with edge weights
w = [1.1, 0.1, 2.3, 0.5]
y = l(g, x, w)

# Edge weights can also be embedded in the graph.
g = GNNGraph(s, t, w)
l = GCNConv(3 => 5, use_edge_weight=true)
y = l(g, x) # same as l(g, x, w)
```
"""
struct GCNConv{A<:AbstractMatrix, B, F} <: GNNLayer
    weight::A
    bias::B
    σ::F
    add_self_loops::Bool
    use_edge_weight::Bool
end

@functor GCNConv

function GCNConv(ch::Pair{Int,Int}, σ=identity;
                 init=glorot_uniform,
                 bias::Bool=true,
                 add_self_loops=true,
                 use_edge_weight=false)
    in, out = ch
    W = init(out, in)
    b = bias ? Flux.create_bias(W, true, out) : false
    GCNConv(W, b, σ, add_self_loops, use_edge_weight)
end

function (l::GCNConv)(g::GNNGraph{<:COO_T}, x::AbstractMatrix)
    # Extract edge_weight from g if available and l.edge_weight == true,
    # otherwise return nothing.
    edge_weight = GNNGraphs._get_edge_weight(g, l.use_edge_weight) # vector or nothing
    return l(g, x, edge_weight)
end

function (l::GCNConv)(g::GNNGraph{<:COO_T}, x::AbstractMatrix{T}, edge_weight::EW) where
    {T, EW<:Union{Nothing,AbstractVector}}

    if l.add_self_loops
        g = add_self_loops(g)
        if edge_weight !== nothing
            # TODO for ADJMAT_T the new edges are not generally at the end
            edge_weight = [edge_weight; fill!(similar(edge_weight, g.num_nodes), 1)]
            @assert length(edge_weight) == g.num_edges
        end
    end
    Dout, Din = size(l.weight)
    if Dout < Din
        # multiply before convolution if it is more convenient, otherwise multiply after
        x = l.weight * x
    end
    d = degree(g, T; dir=:in, edge_weight)
    c = 1 ./ sqrt.(d)
    x = x .* c'
    if edge_weight === nothing
        x = propagate(copy_xj, g, +, xj=x)
    else
        x = propagate(e_mul_xj, g, +, xj=x, e=edge_weight)
    end
    x = x .* c'
    if Dout >= Din
        x = l.weight * x
    end
    return l.σ.(x .+ l.bias)
end

# TODO merge the ADJMAT_T and COO_T methods
# The main problem is handling the weighted case for both.
function (l::GCNConv)(g::GNNGraph{<:ADJMAT_T}, x::AbstractMatrix{T}) where T
    if l.add_self_loops
        g = add_self_loops(g)
    end
    Dout, Din = size(l.weight)
    if Dout < Din
        # multiply before convolution if it is more convenient, otherwise multiply after
        x = l.weight * x
    end
    d = degree(g, T; dir=:in, edge_weight=l.use_edge_weight)
    c = 1 ./ sqrt.(d)
    x = x .* c'
    A = adjacency_matrix(g, weighted=l.use_edge_weight)
    x = x * A
    x = x .* c'
    if Dout >= Din
        x = l.weight * x
    end
    return l.σ.(x .+ l.bias)
end

function (l::GCNConv)(g::GNNGraph{<:ADJMAT_T}, x::AbstractMatrix, edge_weight::AbstractVector)
    g = GNNGraph(edge_index(g)...; g.num_nodes)  # convert to COO
    return l(g, x, edge_weight)
end

function Base.show(io::IO, l::GCNConv)
    out, in = size(l.weight)
    print(io, "GCNConv($in => $out")
    l.σ == identity || print(io, ", ", l.σ)
    print(io, ")")
end

@doc raw"""
    ChebConv(in => out, k; bias=true, init=glorot_uniform)

Chebyshev spectral graph convolutional layer from
paper [Convolutional Neural Networks on Graphs with Fast Localized Spectral Filtering](https://arxiv.org/abs/1606.09375).

Implements

```math
X' = \sum^{K-1}_{k=0}  W^{(k)} Z^{(k)}
```

where ``Z^{(k)}`` is the ``k``-th term of Chebyshev polynomials, and can be calculated by the following recursive form:

```math
Z^{(0)} = X \\
Z^{(1)} = \hat{L} X \\
Z^{(k)} = 2 \hat{L} Z^{(k-1)} - Z^{(k-2)}
```

with ``\hat{L}`` the [`scaled_laplacian`](@ref).

# Arguments

- `in`: The dimension of input features.
- `out`: The dimension of output features.
- `k`: The order of Chebyshev polynomial.
- `bias`: Add learnable bias.
- `init`: Weights' initializer.
"""
struct ChebConv{A<:AbstractArray{<:Number,3}, B} <: GNNLayer
    weight::A
    bias::B
    k::Int
end

function ChebConv(ch::Pair{Int,Int}, k::Int;
                  init=glorot_uniform, bias::Bool=true)
    in, out = ch
    W = init(out, in, k)
    b = bias ? Flux.create_bias(W, true, out) : false
    ChebConv(W, b, k)
end

@functor ChebConv

function (c::ChebConv)(g::GNNGraph, X::AbstractMatrix{T}) where T
    check_num_nodes(g, X)
    @assert size(X, 1) == size(c.weight, 2) "Input feature size must match input channel size."

    L̃ = scaled_laplacian(g, eltype(X))

    Z_prev = X
    Z = X * L̃
    Y = view(c.weight,:,:,1) * Z_prev
    Y += view(c.weight,:,:,2) * Z
    for k = 3:c.k
        Z, Z_prev = 2*Z*L̃ - Z_prev, Z
        Y += view(c.weight,:,:,k) * Z
    end
    return Y .+ c.bias
end

function Base.show(io::IO, l::ChebConv)
    out, in, k = size(l.weight)
    print(io, "ChebConv(", in, " => ", out)
    print(io, ", k=", k)
    print(io, ")")
end

@doc raw"""
    GraphConv(in => out, σ=identity; aggr=+, bias=true, init=glorot_uniform)

Graph convolution layer from Reference: [Weisfeiler and Leman Go Neural: Higher-order Graph Neural Networks](https://arxiv.org/abs/1810.02244).

Performs:
```math
\mathbf{x}_i' = W_1 \mathbf{x}_i + \square_{j \in \mathcal{N}(i)} W_2 \mathbf{x}_j
```

where the aggregation type is selected by `aggr`.

# Arguments

- `in`: The dimension of input features.
- `out`: The dimension of output features.
- `σ`: Activation function.
- `aggr`: Aggregation operator for the incoming messages (e.g. `+`, `*`, `max`, `min`, and `mean`).
- `bias`: Add learnable bias.
- `init`: Weights' initializer.
"""
struct GraphConv{A<:AbstractMatrix, B} <: GNNLayer
    weight1::A
    weight2::A
    bias::B
    σ
    aggr
end

@functor GraphConv

function GraphConv(ch::Pair{Int,Int}, σ=identity; aggr=+,
                   init=glorot_uniform, bias::Bool=true)
    in, out = ch
    W1 = init(out, in)
    W2 = init(out, in)
    b = bias ? Flux.create_bias(W1, true, out) : false
    GraphConv(W1, W2, b, σ, aggr)
end

function (l::GraphConv)(g::GNNGraph, x::AbstractMatrix)
    check_num_nodes(g, x)
    m = propagate(copy_xj, g, l.aggr, xj=x)
    x = l.σ.(l.weight1 * x .+ l.weight2 * m .+ l.bias)
    return x
end

function Base.show(io::IO, l::GraphConv)
    in_channel = size(l.weight1, ndims(l.weight1))
    out_channel = size(l.weight1, ndims(l.weight1)-1)
    print(io, "GraphConv(", in_channel, " => ", out_channel)
    l.σ == identity || print(io, ", ", l.σ)
    print(io, ", aggr=", l.aggr)
    print(io, ")")
end

@doc raw"""
    GATConv(in => out, σ=identity;
            heads=1,
            concat=true,
            init=glorot_uniform
            bias=true,
            negative_slope=0.2f0)

Graph attentional layer from the paper [Graph Attention Networks](https://arxiv.org/abs/1710.10903).

Implements the operation
```math
\mathbf{x}_i' = \sum_{j \in N(i) \cup \{i\}} \alpha_{ij} W \mathbf{x}_j
```
where the attention coefficients ``\alpha_{ij}`` are given by
```math
\alpha_{ij} = \frac{1}{z_i} \exp(LeakyReLU(\mathbf{a}^T [W \mathbf{x}_i; W \mathbf{x}_j]))
```
with ``z_i`` a normalization factor.

# Arguments

- `in`: The dimension of input features.
- `out`: The dimension of output features.
- `bias`: Learn the additive bias if true.
- `heads`: Number attention heads.
- `concat`: Concatenate layer output or not. If not, layer output is averaged over the heads.
- `negative_slope`: The parameter of LeakyReLU.
"""
struct GATConv{T, A<:AbstractMatrix, B} <: GNNLayer
    weight::A
    bias::B
    a::A
    σ
    negative_slope::T
    channel::Pair{Int, Int}
    heads::Int
    concat::Bool
end

@functor GATConv
Flux.trainable(l::GATConv) = (l.weight, l.bias, l.a)

function GATConv(ch::Pair{Int,Int}, σ=identity;
                 heads::Int=1, concat::Bool=true, negative_slope=0.2,
                 init=glorot_uniform, bias::Bool=true)
    in, out = ch
    W = init(out*heads, in)
    if concat
        b = bias ? Flux.create_bias(W, true, out*heads) : false
    else
        b = bias ? Flux.create_bias(W, true, out) : false
    end
    a = init(2*out, heads)
    negative_slope = convert(eltype(W), negative_slope)
    GATConv(W, b, a, σ, negative_slope, ch, heads, concat)
end

function (l::GATConv)(g::GNNGraph, x::AbstractMatrix)
    check_num_nodes(g, x)
    g = add_self_loops(g)
    chin, chout = l.channel
    heads = l.heads

    Wx = l.weight * x
    Wx = reshape(Wx, chout, heads, :)                   # chout × nheads × nnodes

    function message(Wxi, Wxj, e)
        aWW = sum(l.a .* vcat(Wxi, Wxj), dims=1)   # 1 × nheads × nedges
        α = exp.(leakyrelu.(aWW, l.negative_slope))
        return (α = α, β = α .* Wxj)
    end

    m = propagate(message, g, +; xi=Wx, xj=Wx)                 ## chout × nheads × nnodes
    x = m.β ./ m.α

    if !l.concat
        x = mean(x, dims=2)
    end
    x = reshape(x, :, size(x, 3))  # return a matrix
    x = l.σ.(x .+ l.bias)

    return x
end

function Base.show(io::IO, l::GATConv)
    out_channel, in_channel = size(l.weight)
    print(io, "GATConv(", in_channel, "=>", out_channel ÷ l.heads)
    print(io, ", LeakyReLU(λ=", l.negative_slope)
    print(io, "))")
end

@doc raw"""
    GATv2Conv(in => out, σ=identity;
            heads=1,
            concat=true,
            init=glorot_uniform
            bias=true,
            negative_slope=0.2f0)

GATv2 attentional layer from the paper [How Attentive are Graph Attention Networks?](https://arxiv.org/abs/2105.14491).

Implements the operation
```math
\mathbf{x}_i' = \sum_{j \in N(i) \cup \{i\}} \alpha_{ij} W_1 \mathbf{x}_j
```
where the attention coefficients ``\alpha_{ij}`` are given by
```math
\alpha_{ij} = \frac{1}{z_i} \exp(\mathbf{a}^T LeakyReLU([W_2 \mathbf{x}_i; W_1 \mathbf{x}_j]))
```
with ``z_i`` a normalization factor.

# Arguments

- `in`: The dimension of input features.
- `out`: The dimension of output features.
- `bias`: Learn the additive bias if true.
- `heads`: Number attention heads.
- `concat`: Concatenate layer output or not. If not, layer output is averaged over the heads.
- `negative_slope`: The parameter of LeakyReLU.
"""
struct GATv2Conv{T, A1, A2, B, C<:AbstractMatrix} <: GNNLayer
    dense_i::A1
    dense_j::A2
    bias::B
    a::C
    σ
    negative_slope::T
    channel::Pair{Int, Int}
    heads::Int
    concat::Bool
end

@functor GATv2Conv
Flux.trainable(l::GATv2Conv) = (l.dense_i, l.dense_j, l.bias, l.a)

function GATv2Conv(
    channel::Pair{Int,Int},
    σ=identity;
    heads::Int=1,
    concat::Bool=true,
    negative_slope=0.2,
    init=glorot_uniform,
    bias::Bool=true,
)
    in, out = channel
    dense_i = Dense(in, out*heads; bias=bias, init=init)
    dense_j = Dense(in, out*heads; bias=false, init=init)
    if concat
        b = bias ? Flux.create_bias(dense_i.weight, bias, out*heads) : false
    else
        b = bias ? Flux.create_bias(dense_i.weight, bias, out) : false
    end
    a = init(out, heads)

    negative_slope = convert(eltype(dense_i.weight), negative_slope)
    GATv2Conv(dense_i, dense_j, b, a, σ, negative_slope, channel, heads, concat)
end

function (l::GATv2Conv)(g::GNNGraph, x::AbstractMatrix)
    check_num_nodes(g, x)
    g = add_self_loops(g)
    in, out = l.channel
    heads = l.heads

    Wix = reshape(l.dense_i(x), out, heads, :)                                  # out × heads × nnodes
    Wjx = reshape(l.dense_j(x), out, heads, :)                                  # out × heads × nnodes

    function message(Wix, Wjx, e)
        eij = sum(l.a .* leakyrelu.(Wix + Wjx, l.negative_slope), dims=1)   # 1 × heads × nedges
        α = exp.(eij)
        return (α = α, β = α .* Wjx)
    end

    m = propagate(message, g, +; xi=Wix, xj=Wjx)                            # out × heads × nnodes
    x = m.β ./ m.α

    if !l.concat
        x = mean(x, dims=2)
    end
    x = reshape(x, :, size(x, 3))
    x = l.σ.(x .+ l.bias)
    return x
end

function Base.show(io::IO, l::GATv2Conv)
    out, in = size(l.weight_i)
    print(io, "GATv2Conv(", in, "=>", out ÷ l.heads)
    print(io, ", LeakyReLU(λ=", l.negative_slope)
    print(io, "))")
end

@doc raw"""
    GatedGraphConv(out, num_layers; aggr=+, init=glorot_uniform)

Gated graph convolution layer from [Gated Graph Sequence Neural Networks](https://arxiv.org/abs/1511.05493).

Implements the recursion
```math
\mathbf{h}^{(0)}_i = [\mathbf{x}_i || \mathbf{0}] \\
\mathbf{h}^{(l)}_i = GRU(\mathbf{h}^{(l-1)}_i, \square_{j \in N(i)} W \mathbf{h}^{(l-1)}_j)
```

 where ``\mathbf{h}^{(l)}_i`` denotes the ``l``-th hidden variables passing through GRU. The dimension of input ``\mathbf{x}_i`` needs to be less or equal to `out`.

# Arguments

- `out`: The dimension of output features.
- `num_layers`: The number of gated recurrent unit.
- `aggr`: Aggregation operator for the incoming messages (e.g. `+`, `*`, `max`, `min`, and `mean`).
- `init`: Weight initialization function.
"""
struct GatedGraphConv{A<:AbstractArray{<:Number,3}, R} <: GNNLayer
    weight::A
    gru::R
    out_ch::Int
    num_layers::Int
    aggr
end

@functor GatedGraphConv

function GatedGraphConv(out_ch::Int, num_layers::Int;
                        aggr=+, init=glorot_uniform)
    w = init(out_ch, out_ch, num_layers)
    gru = GRUCell(out_ch, out_ch)
    GatedGraphConv(w, gru, out_ch, num_layers, aggr)
end

# remove after https://github.com/JuliaDiff/ChainRules.jl/pull/521
@non_differentiable fill!(x...)

function (l::GatedGraphConv)(g::GNNGraph, H::AbstractMatrix{S}) where {S<:Real}
    check_num_nodes(g, H)
    m, n = size(H)
    @assert (m <= l.out_ch) "number of input features must less or equals to output features."
    if m < l.out_ch
        Hpad = similar(H, S, l.out_ch - m, n)
        H = vcat(H, fill!(Hpad, 0))
    end
    for i = 1:l.num_layers
        M = view(l.weight, :, :, i) * H
        M = propagate(copy_xj, g, l.aggr; xj=M)
        H, _ = l.gru(H, M)
    end
    H
end

function Base.show(io::IO, l::GatedGraphConv)
    print(io, "GatedGraphConv(($(l.out_ch) => $(l.out_ch))^$(l.num_layers)")
    print(io, ", aggr=", l.aggr)
    print(io, ")")
end

@doc raw"""
    EdgeConv(nn; aggr=max)

Edge convolutional layer from paper [Dynamic Graph CNN for Learning on Point Clouds](https://arxiv.org/abs/1801.07829).

Performs the operation
```math
\mathbf{x}_i' = \square_{j \in N(i)} nn(\mathbf{x}_i || \mathbf{x}_j - \mathbf{x}_i)
```

where `nn` generally denotes a learnable function, e.g. a linear layer or a multi-layer perceptron.

# Arguments

- `nn`: A (possibly learnable) function acting on edge features.
- `aggr`: Aggregation operator for the incoming messages (e.g. `+`, `*`, `max`, `min`, and `mean`).
"""
struct EdgeConv <: GNNLayer
    nn
    aggr
end

@functor EdgeConv

EdgeConv(nn; aggr=max) = EdgeConv(nn, aggr)

function (l::EdgeConv)(g::GNNGraph, x::AbstractMatrix)
    check_num_nodes(g, x)
    message(xi, xj, e) = l.nn(vcat(xi, xj .- xi))
    x = propagate(message, g, l.aggr, xi=x, xj=x)
    return x
end

function Base.show(io::IO, l::EdgeConv)
    print(io, "EdgeConv(", l.nn)
    print(io, ", aggr=", l.aggr)
    print(io, ")")
end

@doc raw"""
    GINConv(f, ϵ; aggr=+)

Graph Isomorphism convolutional layer from paper [How Powerful are Graph Neural Networks?](https://arxiv.org/pdf/1810.00826.pdf).

Implements the graph convolution
```math
\mathbf{x}_i' = f_\Theta\left((1 + \epsilon) \mathbf{x}_i + \sum_{j \in N(i)} \mathbf{x}_j \right)
```
where ``f_\Theta`` typically denotes a learnable function, e.g. a linear layer or a multi-layer perceptron.

# Arguments

- `f`: A (possibly learnable) function acting on node features.
- `ϵ`: Weighting factor.
"""
struct GINConv{R<:Real} <: GNNLayer
    nn
    ϵ::R
    aggr
end

@functor GINConv
Flux.trainable(l::GINConv) = (l.nn,)

GINConv(nn, ϵ; aggr=+) = GINConv(nn, ϵ, aggr)

function (l::GINConv)(g::GNNGraph, x::AbstractMatrix)
    check_num_nodes(g, x)
    m = propagate(copy_xj, g, l.aggr, xj=x)
    l.nn((1 + ofeltype(x, l.ϵ)) * x + m)
end

function Base.show(io::IO, l::GINConv)
    print(io, "GINConv($(l.nn)")
    print(io, ", $(l.ϵ)")
    print(io, ")")
end

@doc raw"""
    NNConv(in => out, f, σ=identity; aggr=+, bias=true, init=glorot_uniform)

The continuous kernel-based convolutional operator from the
[Neural Message Passing for Quantum Chemistry](https://arxiv.org/abs/1704.01212) paper.
This convolution is also known as the edge-conditioned convolution from the
[Dynamic Edge-Conditioned Filters in Convolutional Neural Networks on Graphs](https://arxiv.org/abs/1704.02901) paper.

Performs the operation

```math
\mathbf{x}_i' = W \mathbf{x}_i + \square_{j \in N(i)} f_\Theta(\mathbf{e}_{j\to i})\,\mathbf{x}_j
```

where ``f_\Theta``  denotes a learnable function (e.g. a linear layer or a multi-layer perceptron).
Given an input of batched edge features `e` of size `(num_edge_features, num_edges)`,
the function `f` will return an batched matrices array whose size is `(out, in, num_edges)`.
For convenience, also functions returning a single `(out*in, num_edges)` matrix are allowed.

# Arguments

- `in`: The dimension of input features.
- `out`: The dimension of output features.
- `f`: A (possibly learnable) function acting on edge features.
- `aggr`: Aggregation operator for the incoming messages (e.g. `+`, `*`, `max`, `min`, and `mean`).
- `σ`: Activation function.
- `bias`: Add learnable bias.
- `init`: Weights' initializer.
"""
struct NNConv <: GNNLayer
    weight
    bias
    nn
    σ
    aggr
end

@functor NNConv

function NNConv(ch::Pair{Int,Int}, nn, σ=identity; aggr=+, bias=true, init=glorot_uniform)
    in, out = ch
    W = init(out, in)
    b = bias ? Flux.create_bias(W, true, out) : false
    return NNConv(W, b, nn, σ, aggr)
end

function (l::NNConv)(g::GNNGraph, x::AbstractMatrix, e)
    check_num_nodes(g, x)

    function message(xi, xj, e)
        nin, nedges = size(xj)
        W = reshape(l.nn(e), (:, nin, nedges))
        xj = reshape(xj, (nin, 1, nedges)) # needed by batched_mul
        m = NNlib.batched_mul(W, xj)
        return reshape(m, :, nedges)
    end

    m = propagate(message, g, l.aggr, xj=x, e=e)
    return l.σ.(l.weight*x .+ m .+ l.bias)
end

(l::NNConv)(g::GNNGraph) = GNNGraph(g, ndata=l(g, node_features(g), edge_features(g)))

function Base.show(io::IO, l::NNConv)
    out, in = size(l.weight)
    print(io, "NNConv( $in => $out")
    print(io, ", aggr=", l.aggr)
    print(io, ")")
end

@doc raw"""
    SAGEConv(in => out, σ=identity; aggr=mean, bias=true, init=glorot_uniform)

GraphSAGE convolution layer from paper [Inductive Representation Learning on Large Graphs](https://arxiv.org/pdf/1706.02216.pdf).

Performs:
```math
\mathbf{x}_i' = W \cdot [\mathbf{x}_i; \square_{j \in \mathcal{N}(i)} \mathbf{x}_j]
```

where the aggregation type is selected by `aggr`.

# Arguments

- `in`: The dimension of input features.
- `out`: The dimension of output features.
- `σ`: Activation function.
- `aggr`: Aggregation operator for the incoming messages (e.g. `+`, `*`, `max`, `min`, and `mean`).
- `bias`: Add learnable bias.
- `init`: Weights' initializer.
"""
struct SAGEConv{A<:AbstractMatrix, B} <: GNNLayer
    weight::A
    bias::B
    σ
    aggr
end

@functor SAGEConv

function SAGEConv(ch::Pair{Int,Int}, σ=identity; aggr=mean,
                   init=glorot_uniform, bias::Bool=true)
    in, out = ch
    W = init(out, 2*in)
    b = bias ? Flux.create_bias(W, true, out) : false
    SAGEConv(W, b, σ, aggr)
end

function (l::SAGEConv)(g::GNNGraph, x::AbstractMatrix)
    check_num_nodes(g, x)
    m = propagate(copy_xj, g, l.aggr, xj=x)
    x = l.σ.(l.weight * vcat(x, m) .+ l.bias)
    return x
end

function Base.show(io::IO, l::SAGEConv)
    out_channel, in_channel = size(l.weight)
    print(io, "SAGEConv(", in_channel ÷ 2, " => ", out_channel)
    l.σ == identity || print(io, ", ", l.σ)
    print(io, ", aggr=", l.aggr)
    print(io, ")")
end

@doc raw"""
    ResGatedGraphConv(in => out, act=identity; init=glorot_uniform, bias=true)

The residual gated graph convolutional operator from the [Residual Gated Graph ConvNets](https://arxiv.org/abs/1711.07553) paper.

The layer's forward pass is given by

```math
\mathbf{x}_i' = act\big(U\mathbf{x}_i + \sum_{j \in N(i)} \eta_{ij} V \mathbf{x}_j\big),
```
where the edge gates ``\eta_{ij}`` are given by

```math
\eta_{ij} = sigmoid(A\mathbf{x}_i + B\mathbf{x}_j).
```

# Arguments

- `in`: The dimension of input features.
- `out`: The dimension of output features.
- `act`: Activation function.
- `init`: Weight matrices' initializing function.
- `bias`: Learn an additive bias if true.
"""
struct ResGatedGraphConv <: GNNLayer
    A
    B
    U
    V
    bias
    σ
end

@functor ResGatedGraphConv

function ResGatedGraphConv(ch::Pair{Int,Int}, σ=identity;
                      init=glorot_uniform, bias::Bool=true)
    in, out = ch
    A = init(out, in)
    B = init(out, in)
    U = init(out, in)
    V = init(out, in)
    b = bias ? Flux.create_bias(A, true, out) : false
    return ResGatedGraphConv(A, B, U, V, b, σ)
end

function (l::ResGatedGraphConv)(g::GNNGraph, x::AbstractMatrix)
    check_num_nodes(g, x)

    message(xi, xj, e) = sigmoid.(xi.Ax .+ xj.Bx) .* xj.Vx

    Ax = l.A * x
    Bx = l.B * x
    Vx = l.V * x

    m = propagate(message, g, +, xi=(; Ax), xj=(; Bx, Vx))

    return l.σ.(l.U*x .+ m .+ l.bias)
end

function Base.show(io::IO, l::ResGatedGraphConv)
    out_channel, in_channel = size(l.A)
    print(io, "ResGatedGraphConv(", in_channel, "=>", out_channel)
    l.σ == identity || print(io, ", ", l.σ)
    print(io, ")")
end

@doc raw"""
    CGConv((nin, ein) => nout, f, act=identity; bias=true, init=glorot_uniform, residual=false)

The crystal graph convolutional layer from the paper
[Crystal Graph Convolutional Neural Networks for an Accurate and
Interpretable Prediction of Material Properties](https://arxiv.org/pdf/1710.10324.pdf).
Performs the operation

```math
\mathbf{x}_i' = \mathbf{x}_i + \sum_{j\in N(i)}\sigma(W_f \mathbf{z}_{ij} + \mathbf{b}_f)\, act(W_s \mathbf{z}_{ij} + \mathbf{b}_s)
```

where ``\mathbf{z}_{ij}``  is the node and edge features concatenation
``[\mathbf{x}_i; \mathbf{x}_j; \mathbf{e}_{j\to i}]``
and ``\sigma`` is the sigmoid function.
The residual ``\mathbf{x}_i`` is added only if `residual=true` and the output size is the same
as the input size.

# Arguments

- `nin`: The dimension of input node features.
- `nout`: The dimension of input edge features.
- `out`: The dimension of output node features.
- `act`: Activation function.
- `bias`: Add learnable bias.
- `init`: Weights' initializer.
- `residual`: Add a residual connection.

# Examples

```julia
x = rand(Float32, 2, g.num_nodes)
e = rand(Float32, 3, g.num_edges)

l = CGConv((2,3) => 4, tanh)

y = l(g, x, e)    # size: (4, num_nodes)
```
"""
struct CGConv <: GNNLayer
    ch
    dense_f::Dense
    dense_s::Dense
    residual::Bool
end

@functor CGConv

CGConv(nin::Int, ein::Int, out::Int, args...; kws...) = CGConv((nin, ein) => out, args...; kws...)

function CGConv(ch::Pair{NTuple{2,Int},Int}, act=identity; residual=false, bias=true, init=glorot_uniform)
    (nin, ein), out = ch
    dense_f = Dense(2nin+ein, out, sigmoid; bias, init)
    dense_s = Dense(2nin+ein, out, act; bias, init)
    return CGConv(ch, dense_f, dense_s, residual)
end

function (l::CGConv)(g::GNNGraph, x::AbstractMatrix, e::AbstractMatrix)
    check_num_nodes(g, x)
    check_num_edges(g, e)

    function message(xi, xj, e)
        z = vcat(xi, xj, e)
        return l.dense_f(z) .* l.dense_s(z)
    end

    m = propagate(message, g, +, xi=x, xj=x, e=e)
    if l.residual
        if size(x, 1) == size(m, 1)
            m += x
        else
            @warn "number of output features different from number of input features, residual not applyed."
        end
    end
    return m
end

(l::CGConv)(g::GNNGraph) = GNNGraph(g, ndata=l(g, node_features(g), edge_features(g)))

function Base.show(io::IO, l::CGConv)
    print(io, "CGConv($(l.ch)")
    l.dense_s.σ == identity || print(io, ", ", l.dense_s.σ)
    print(io, ", residual=$(l.residual)")
    print(io, ")")
end

@doc raw"""
    AGNNConv(init_beta=1f0)

Attention-based Graph Neural Network layer from paper [Attention-based
Graph Neural Network for Semi-Supervised Learning](https://arxiv.org/abs/1803.03735).

THe forward pass is given by
```math
\mathbf{x}_i' = \sum_{j \in {N(i) \cup \{i\}} \alpha_{ij} W \mathbf{x}_j
```
where the attention coefficients ``\alpha_{ij}`` are given by
```math
\alpha_{ij} =\frac{e^{\beta \cos(\mathbf{x}_i, \mathbf{x}_j)}}
                  {\sum_{j'}e^{\beta \cos(\mathbf{x}_i, \mathbf{x}_j'}}
```
with the cosine distance defined by
```math
\cos(\mathbf{x}_i, \mathbf{x}_j) =
  \mathbf{x}_i \cdot \mathbf{x}_j / \lVert\mathbf{x}_i\rVert \lVert\mathbf{x}_j\rVert``
```
and ``\beta`` a trainable parameter.

# Arguments

- `init_beta`: The initial value of ``\beta``.
"""
struct AGNNConv{A<:AbstractVector} <: GNNLayer
    β::A
end

@functor AGNNConv

function AGNNConv(init_beta = 1f0)
    AGNNConv([init_beta])
end

function (l::AGNNConv)(g::GNNGraph, x::AbstractMatrix)
    check_num_nodes(g, x)
    g = add_self_loops(g)

    xn = x ./ sqrt.(sum(x.^2, dims=1))
    cos_dist = apply_edges(xi_dot_xj, g, xi=xn, xj=xn)
    α = softmax_edge_neighbors(g, l.β .* cos_dist)

    x = propagate(g, +; xj=x, e=α) do xi, xj, α
            α .* xj
        end

    return x
end

@doc raw"""
    MEGNetConv(ϕe, ϕv; aggr=mean)
    MEGNetConv(in => out; aggr=mean)

Convolution from [Graph Networks as a Universal Machine Learning Framework for Molecules and Crystals](https://arxiv.org/pdf/1812.05055.pdf)
paper. In the forward pass, takes as inputs node features `x` and edge features `e` and returns
updated features `x'` and `e'` according to

```math
\mathbf{e}_{i\to j}'  = \phi_e([\mathbf{x}_i;\,  \mathbf{x}_j;\,  \mathbf{e}_{i\to j}]),\\
\mathbf{x}_{i}'  = \phi_v([\mathbf{x}_i;\, \square_{j\in \mathcal{N}(i)}\,\mathbf{e}_{j\to i}']).
```

`aggr` defines the aggregation to be performed.

If the neural networks `ϕe` and  `ϕv` are not provided, they will be constructed from
the `in` and `out` arguments instead as multi-layer perceptron with one hidden layer and `relu`
activations.

# Examples

```julia
g = rand_graph(10, 30)
x = randn(3, 10)
e = randn(3, 30)
m = MEGNetConv(3 => 3)
x′, e′ = m(g, x, e)
```
"""
struct MEGNetConv <: GNNLayer
    ϕe
    ϕv
    aggr
end

@functor MEGNetConv

MEGNetConv(ϕe, ϕv; aggr=mean) = MEGNetConv(ϕe, ϕv, aggr)

function MEGNetConv(ch::Pair{Int,Int}; aggr=mean)
    nin, nout = ch
    ϕe = Chain(Dense(3nin, nout, relu),
               Dense(nout, nout))

    ϕv = Chain(Dense(nin + nout, nout, relu),
               Dense(nout, nout))

    MEGNetConv(ϕe, ϕv; aggr)
end

function (l::MEGNetConv)(g::GNNGraph)
    x, e = l(g, node_features(g), edge_features(g))
    g = GNNGraph(g, ndata=x, edata=e)
end

function (l::MEGNetConv)(g::GNNGraph, x::AbstractMatrix, e::AbstractMatrix)
    check_num_nodes(g, x)

    ē = apply_edges(g, xi=x, xj=x, e=e) do xi, xj, e
        l.ϕe(vcat(xi, xj, e))
    end

    xᵉ = aggregate_neighbors(g, l.aggr, ē)

    x̄ = l.ϕv(vcat(x, xᵉ))

    return x̄, ē
end


text:
using Test, Statistics, CategoricalArrays, Random, StableRNGs
#using StableRNGs
#rng = StableRNG(123)
using BetaML
#import BetaML.Utils

TESTRNG = FIXEDRNG # This could change...
#TESTRNG = StableRNG(123)

println("*** Testing individual utility functions (module `Utils`)...")

# ==================================
# TEST 1: oneHotEncoder
println("Going through Test1 (oneHotEncoder)...")

a = [[1,3,4],[1,4,2,2,3],[2,3]]
b = [2,1,5,2,1]
c = 2
ae = oneHotEncoder(a;d=5,count=true)
be = oneHotEncoder(b;d=6,count=true)
ce = oneHotEncoder(c;d=6)
@test sum(ae*be*ce') == 4

a = [["aa","ee","dd","aa"],["aa","dd","bb","bb","dd"],["bb","ee"]]
b = ["a","d","b","b","d"]
c = "dd"
ae = oneHotEncoder(a,factors=["aa","bb","cc","dd","ee","ff"],count=true)
be = oneHotEncoder(b,factors=["a","b","c","d","e","f"],count=true)
ce = oneHotEncoder(c,factors=["aa","bb","cc","dd","ee","ff"],count=false)
@test sum(ae * be') + ce[1,2] == 15

ae = oneHotEncoder(a,count=true)
be = oneHotEncoder(b,count=false)
ce = oneHotEncoder(c,count=true)
@test sum(ae)+sum(size(be))+sum(ce) == 20

@test BetaML.Utils.singleUnique([[1,2,3],[7,8],[5]]) == [1,2,3,7,8,5]
@test BetaML.Utils.singleUnique([1,4,5,4]) == [1,4,5]
@test BetaML.Utils.singleUnique("aaa") == ["aaa"]

@test BetaML.Utils.oneHotEncoderRow(4,d=5) == [0,0,0,1,0]
@test BetaML.Utils.oneHotEncoderRow("b",factors=["a","b","c"]) == [0,1,0]
@test BetaML.Utils.oneHotEncoderRow(["b","c","c"],factors=["a","b","c"],count=true) == [0,1,2]

# ==================================
# NEW TEST
println("Testing findFirst/ findall / integerEncoder / integerDecoder...")

a = ["aa","cc","cc","bb","cc","aa"]
A = [3 2 1; 2 3 1; 1 2 3]
@test findfirst("cc",a)  == 2
@test findall("cc",a)    == [2,3,5]
@test findfirst(2,A)     == (2,1)
@test findall(2,A)       == [(2,1),(1,2),(3,2)]
@test findfirst(2,A;returnTuple=false) == CartesianIndex(2,1)

factors =  ["aa","bb","cc"]
encoded  = integerEncoder(a,factors=factors)
@test encoded == [1,3,3,2,3,1]
decoded  = integerDecoder(encoded,factors)
@test a == decoded

# ==================================
# TEST 2: softMax
println("** Going through Test2 (softmax and other activation functions)...")
@test isapprox(softmax([2,3,4],β=0.1),[0.3006096053557272,0.3322249935333472,0.36716540111092544])

@test didentity(-1) == 1 && relu(-1) == 0 && drelu(-1) == 0 && celu(-1,α=0.1) == -0.09999546000702375 && dcelu(-1) == 0.36787944117144233 &&
      elu(-1,α=0.1) == -0.06321205588285576 && delu(-1) == 0.36787944117144233 && plu(-1) == -1 && dplu(-1) == 1

@test didentity(1) == 1 && relu(1) == 1 && drelu(1) == 1 && celu(1) == 1 && dcelu(1) == 1 &&
    elu(1) == 1 && delu(1) == 1 && plu(1) == 1 && dplu(1) == 1

@test dtanh(1) == 0.41997434161402614 && sigmoid(1) == 0.7310585786300049 == dsoftplus(1) && dsigmoid(1) == 0.19661193324148188 &&
      softplus(1) == 1.3132616875182228 && mish(1) == 0.8650983882673103 && dmish(1) == 1.0490362200997922

# ==================================
# TEST 3: autoJacobian
println("** Going through Test3 (autoJacobian)...")
@test isapprox(softmax([2,3,4],β=0.1),[0.3006096053557272,0.3322249935333472,0.36716540111092544])

#import BetaML.Utils: autoJacobian
@test autoJacobian(x -> (x[1]*2,x[2]*x[3]),[1,2,3]) == [2.0 0.0 0.0; 0.0 3.0 2.0]

b = softmax([2,3,4],β=1/2)
c = softmax([2,3.0000001,4],β=1/2)
# Skipping this test as gives problems in CI for Julia < 1.6
#autoGrad = autoJacobian(x->softmax(x,β=1/2),[2,3,4])
realG2 = [(c[1]-b[1])*10000000,(c[2]-b[2])*10000000,(c[3]-b[3])*10000000]
#@test isapprox(autoGrad[:,2],realG2,atol=0.000001)
manualGrad = dsoftmax([2,3,4],β=1/2)
@test isapprox(manualGrad[:,2],realG2,atol=0.000001)

# Manual way is hundred of times faster
#@benchmark autoJacobian(softMax2,[2,3,4])
#@benchmark dSoftMax([2,3,4],β=1/2)

# ==================================
# New test
println("** Testing cross-entropy...")
or = crossEntropy([0.8,0.001,0.001],[1.0,0,0],weight = [2,1,1])
@test or ≈ 0.4462871026284194
d = dCrossEntropy([0.8,0.001,0.001],[1.0,0,0],weight = [2,1,1])
δ = 0.001
dest = crossEntropy([0.8+δ,0.101,0.001],[1.0,0,0],weight = [2,1,1])
@test isapprox(dest-or, d[1]*δ,atol=0.0001)

# ==================================
# New test
println("** Testing permutations...")
y = ["a","a","a","b","b","c","c","c"]
yp = getPermutations(y,keepStructure=true)
ypExpected = [
  ["a", "a", "a", "b", "b", "c", "c", "c"],
  ["a", "a", "a", "c", "c", "b", "b", "b"],
  ["b", "b", "b", "a", "a", "c", "c", "c"],
  ["b", "b", "b", "c", "c", "a", "a", "a"],
  ["c", "c", "c", "a", "a", "b", "b", "b"],
  ["c", "c", "c", "b", "b", "a", "a", "a"],
]
@test yp == ypExpected

# ==================================
# New test
println("** Going through testing accuracy...")

y = ["a","a","a","b","b","c","c","c"]
ŷ = ["b","b","a","c","c","a","a","c"]

accuracyConsideringClassLabels = accuracy(ŷ,y) # 2 out of 8
accuracyConsideringAnyClassLabel = accuracy(ŷ,y,ignoreLabels=true) # 6 out of 8

@test accuracyConsideringClassLabels == 2/8
@test accuracyConsideringAnyClassLabel == 6/8

# with categorical arrays..
y = CategoricalArray(y)
ŷ = CategoricalArray(ŷ)

accuracyConsideringClassLabels = accuracy(ŷ,y) # 2 out of 8
accuracyConsideringAnyClassLabel = accuracy(ŷ,y,ignoreLabels=true) # 6 out of 8

@test accuracyConsideringClassLabels == 2/8
@test accuracyConsideringAnyClassLabel == 6/8

x = [0.01 0.02 0.1 0.05 0.2 0.1  0.05 0.27  0.2;
     0.05 0.01 0.2 0.02 0.1 0.27 0.1  0.05  0.2]
y = [3,3]
@test [accuracy(x,y,tol=i) for i in 1:10] == [0.0, 0.5, 0.5, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
x = [0.3 0.2 0.5; 0.5 0.25 0.25; 0.1 0.1 0.9]
y = [1, 3, 1]
@test accuracy(x,y) == 0.0
@test accuracy(x,y, ignoreLabels=true) == 1.0

yest = [0 0.12; 0 0.9]
y    = [1,2]
accuracy(yest,y)
yest = [0.0  0.0  0.0;0.0  0.0  0.0;0.0  0.0  0.706138]
y    = [2,2,1]
accuracy(yest,y)

# ==================================
# New test
println("** Going through testing scaling...")

skip = [1,4,5]
x = [1.1 4.1 8.1 3 8; 2 4 9 7 2; 7 2 9 3 1]
scaleFactors = getScaleFactors(x,skip=skip)
y  = scale(x,scaleFactors)
y2 = scale(x)
x2 = copy(x)
scale!(x2)
@test y2 == x2
@test all((sum(mean(y,dims=1)), sum(var(y,corrected=false,dims=1)) ) .≈ (11.366666666666667, 21.846666666666668))
x3 = scale(y,scaleFactors,rev=true)
@test x3 == x

# ==================================
# New test
println("** Testing batch()...")

@test size.(batch(10,3,rng=copy(TESTRNG)),1) == [3,3,3]
@test size.(batch(10,12,rng=copy(TESTRNG)),1) == [10]

# ==================================
# New test
println("** Testing meanRelError()...")

ŷ = [22 142 328; 3 9 31; 5 10 32; 3 10 36]
y = [20 140 330; 1 11 33; 3 8 30; 5 12 38]
p=2
(n,d) = size(y)

# case 1 - average of the relative error (records and dimensions normalised)
avgϵRel = sum(abs.((ŷ-y)./ y).^p)^(1/p)/(n*d)
#avgϵRel = (norm((ŷ-y)./ y,p)/(n*d))
meanRelError(ŷ,y,normDim=true,normRec=true,p=p) == avgϵRel
# case 2 - normalised by dimensions (i.e.  all dimensions play the same)
avgϵRel_byDim = (sum(abs.(ŷ-y) .^ (1/p),dims=1).^(1/p) ./ n) ./   (sum(abs.(y) .^ (1/p) ,dims=1) ./n)
avgϵRel = mean(avgϵRel_byDim)
@test meanRelError(ŷ,y,normDim=true,normRec=false,p=p) == avgϵRel
# case 3
avgϵRel_byRec = (sum(abs.(ŷ-y) .^ (1/p),dims=2).^(1/p) ./ d) ./   (sum(abs.(y) .^ (1/p) ,dims=2) ./d)
avgϵRel = mean(avgϵRel_byRec)
@test meanRelError(ŷ,y,normDim=false,normRec=true,p=p) == avgϵRel
# case 4 - average error relativized
avgϵRel = (sum(abs.(ŷ-y).^p)^(1/p) / (n*d)) / (sum( abs.(y) .^p)^(1/p) / (n*d))
#avgϵRel = (norm((ŷ-y),p)/(n*d)) / (norm(y,p) / (n*d))
@test meanRelError(ŷ,y,normDim=false,normRec=false,p=p) == avgϵRel

# ==================================
# New test
println("** Testing pca()...")

X = [1 10 100; 1.1 15 120; 0.95 23 90; 0.99 17 120; 1.05 8 90; 1.1 12 95]
out = pca(X,error=0.05)
@test out.error ≈ 1.0556269747774571e-5
@test sum(out.X) ≈ 662.3492034128955
#X2 = out.X*out.P'
@test out.explVarByDim ≈ [0.873992272007021,0.9999894437302522,1.0]
X = [1 8; 4.5 5.5; 9.5 0.5]
out = pca(X;K=2)
@test isapprox(out.X,[-4.58465   6.63182;-0.308999  7.09961; 6.75092   6.70262],atol=0.00001)
@test isapprox(out.P,[0.745691  0.666292;-0.666292  0.745691],atol=0.00001)

# ==================================
# New test
println("** Testing accuracy() on probs given in a dictionary...")

ŷ = Dict("Lemon" => 0.33, "Apple" => 0.2, "Grape" => 0.47)
y = "Lemon"
@test accuracy(ŷ,y) == 0
@test accuracy(ŷ,y,tol=2) == 1
y = "Grape"
@test accuracy(ŷ,y) == 1
y = "Something else"
@test accuracy(ŷ,y) == 0

ŷ1 = Dict("Lemon" => 0.6, "Apple" => 0.4)
ŷ2 = Dict("Lemon" => 0.33, "Apple" => 0.2, "Grape" => 0.47)
ŷ3 = Dict("Lemon" => 0.2, "Apple" => 0.5, "Grape" => 0.3)
ŷ4 = Dict("Apple" => 0.2, "Grape" => 0.8)
ŷ = [ŷ1,ŷ2,ŷ3,ŷ4]
y = ["Lemon","Lemon","Apple","Lemon"]
@test accuracy(ŷ,y) == 0.5
@test accuracy(ŷ,y,tol=2) == 0.75

# ==================================
# New test
println("** Testing mode(dicts)...")
ŷ1 = Dict("Lemon" => 0.6, "Apple" => 0.4)
ŷ2 = Dict("Lemon" => 0.33, "Apple" => 0.2, "Grape" => 0.47)
ŷ3 = Dict("Lemon" => 0.2, "Apple" => 0.5, "Grape" => 0.3)
ŷ4 = Dict("Apple" => 0.2, "Grape" => 0.8)
ŷ5 = Dict("Lemon" => 0.4, "Grape" => 0.4, "Apple" => 0.2)
ŷ = [ŷ1,ŷ2,ŷ3,ŷ4,ŷ5]
@test mode(ŷ,rng=copy(TESTRNG)) == ["Lemon","Grape","Apple","Grape","Lemon"]

y1 = [1,4,2,5]
y2 = [2,6,6,4]
y = [y1,y2]
@test mode(y,rng=copy(TESTRNG)) == [4,3]
y = vcat(y1',y2')
mode(y,rng=copy(TESTRNG)) == [4,3]

# ==================================
# New test
println("** Testing ConfusionMatrix()...")
cm = ConfusionMatrix(ŷ,["Lemon","Lemon","Apple","Grape","Lemon"],rng=copy(FIXEDRNG))
@test cm.scores == [2 0 1; 0 1 0; 0 0 1]
@test cm.tp == [2,1,1] && cm.tn == [2,4,3] && cm.fp == [0,0,1] && cm.fn == [1, 0, 0]

# Example from https://scikit-learn.org/stable/modules/model_evaluation.html#classification-report
y = [0,1,2,2,0]
ŷ = [0,0,2,1,0]
labels = ["Class 0", "Class 1 with an extra long name super long", "Class 2"]
#y = integerDecoder(y .+ 1,labels)
#ŷ = integerDecoder(ŷ .+ 1,labels)
cm = ConfusionMatrix(ŷ,y,labels=labels)

@test cm.precision ≈ [0.6666666666666666, 0.0, 1.0]
@test cm.recall ≈ [1.0, 0.0, 0.5]
@test cm.specificity ≈ [0.6666666666666666, 0.75, 1.0]
@test cm.f1Score ≈ [0.8, 0.0, 0.6666666666666666]
@test cm.meanPrecision == (0.5555555555555555, 0.6666666666666666)
@test cm.meanRecall == (0.5, 0.6)
@test cm.meanSpecificity == (0.8055555555555555, 0.8166666666666667)
@test cm.meanF1Score == (0.48888888888888893, 0.5866666666666667)
@test cm.accuracy == 0.6
@test cm.misclassification == 0.4

@test BetaML.Utils.print(cm,["report"]) == nothing
@test BetaML.Utils.println(cm) == nothing

# ==================================
# New test
println("** Testing classCounts()...")

@test classCountsWithLabels(["a","b","a","c","d"]) == Dict("a"=>2,"b"=>1,"c"=>1,"d"=>1)
@test classCountsWithLabels(['a' 'b'; 'a' 'c';'a' 'b']) == Dict(['a', 'b'] => 2,['a', 'c'] => 1)
@test classCounts(["a","b","a","c","d"],classes=["a","b","c","d","e"]) == [2,1,1,1,0]
@test collect(classCounts(['a' 'b'; 'a' 'c';'a' 'b'])) in [[2,1],[1,2]] # Order doesn't matter

# ==================================
# New test
println("** Testing gini()...")
@test gini(['a','b','c','c']) == 0.625 # (1/4) * (3/4) + (2/4) * (2/4) + (1/4)*(3/4)
@test gini([1 10; 2 20; 3 30; 2 20]) == 0.625

#a = -0.01*log2(0.01)
#b = -0.99*log2(0.99)
#c = a+b

#A = -0.49*log2(0.49)
#B = -0.51*log2(0.51)
#C = A+B

# ==================================
# New test
println("** Testing entropy()...")
@test isapprox(entropy([1,2,3]), 1.584962500721156) #-(1/3)*log2(1/3)-(1/3)*log2(1/3)-(1/3)*log2(1/3)
@test isapprox(entropy([1 10; 2 20; 3 30]), 1.584962500721156)
#par = entropy([1,1,1,1,1,0,0,0,0,0,0,0,0,0])
#k1 = entropy([1,1,1,0,0,0])
#k2 = entropy([1,1,0,0,0,0,0,0])
#kidsEntropy = k1 *(6/14) + k2*(8/14)
#gain = par - kidsEntropy
#entropy([0,1,2,3,4,5,6,7])

# ==================================
# New test
println("** Testing meanDicts()...")

a = Dict('a'=> 0.2,'b'=>0.3,'c'=>0.5)
b = Dict('a'=> 0.3,'b'=>0.1,'d'=>0.6)
c = Dict('b'=>0.6,'e'=>0.4)
d = Dict('a'=>1)
dicts = [a,b,c,d]
@test meanDicts(dicts) == Dict('a' => 0.375,'c' => 0.125,'d' => 0.15,'e' => 0.1,'b' => 0.25)
@test meanDicts(dicts,weights=[1,1,97,1]) == Dict('a' => 0.015, 'b' => 0.586, 'c' => 0.005, 'd' =>0.006, 'e' =>0.388)

a = Dict(1=> 0.1,2=>0.4,3=>0.5)
b = Dict(4=> 0.3,1=>0.1,2=>0.6)
c = Dict(5=>0.6,4=>0.4)
d = Dict(2=>1)
dicts = [a,b,c,d]
@test meanDicts(dicts)  == Dict(4 => 0.175, 2 => 0.5, 3 => 0.125, 5 => 0.15, 1 => 0.05)

# ==================================
# New test
println("** Testing partition()...")

m1 = [1:10 11:20 31:40]
m2 = convert(Array{Float64,2},[41:50 51:60])
m3 = makeMatrix(collect(61:70))
m4 = collect(71:80)
parts = [0.33,0.27,0.4]
out = partition([m1,m2,m3,m4],parts,shuffle=true,rng=copy(TESTRNG))
@test size(out,1) == 4 && size(out[1][3]) == (4,3)
x = [1:10 11:20]
y = collect(31:40)
((xtrain,xtest),(ytrain,ytest)) = partition([x,y],[0.7,0.3], shuffle=false,rng=copy(TESTRNG))
@test xtest[2,2] == 19 && ytest[2] == 39

m1 = [1:2 11:12 31:32 41:42 51:52 61:62]
out = partition(m1,[0.7,0.3],dims=2,rng=copy(TESTRNG))
@test out == [[31 1 51 61; 32 2 52 62],[11 41; 12 42]]

# Testing not numeric matrices
ms = [[11:16 string.([21:26;])],[31:36;]]
out = partition(ms,[0.7,0.3],dims=1,rng=copy(TESTRNG))
@test out[1][2] == [12 "22"; 14 "24"] && out[2][2]== [32,34]

# ==================================
# New test
println("** Testing KFold sampler with a single X matrix...")
data           = [11:13 21:23 31:33 41:43 51:53 61:63]
sampleIterator = SamplerWithData(KFold(nSplits=3,nRepeats=1,shuffle=true,rng=copy(TESTRNG)),data,2)
for (i,d) in enumerate(sampleIterator)
    if i == 1
        @test d[1][1] == [51 61 21 41; 52 62 22 42; 53 63 23 43] && d[2][1] == [31 11; 32 12; 33 13]
    elseif i == 2
        @test d[1][1] == [31 11 21 41; 32 12 22 42; 33 13 23 43] && d[2][1] == [51 61; 52 62; 53 63]
    elseif i ==3
        @test d[1][1] == [31 11 51 61; 32 12 52 62; 33 13 53 63] && d[2][1] == [21 41; 22 42; 23 43]
    else
        @error "There shoudn't be more than 3 iterations for this iterator!"
    end
end

println("** Testing KFold sampler with multiple matrices...")
data           = [[11:16 string.([21:26;])],[31:36;]]
sampleIterator = SamplerWithData(KFold(nSplits=3,nRepeats=2,shuffle=false,rng=copy(TESTRNG)),data,1)
for (i,d) in enumerate(sampleIterator)
    local xtrain, ytrain, xval, yval
    (xtrain,ytrain),(xval,yval) = d
    if i == 1
        @test xtrain == [13 "23"; 14 "24"; 15 "25"; 16 "26"] && ytrain == [33, 34, 35, 36] && xval == [11 "21"; 12 "22"] && yval == [31, 32]
    elseif i == 5
        @test xtrain == [11 "21"; 12 "22"; 15 "25"; 16 "26"] && ytrain == [31, 32, 35, 36] && xval == [13 "23"; 14 "24"]  && yval == [33, 34]
    elseif i > 6
        @error "There shoudn't be more than 6 iterations for this iterator!"
    end
end

println("** Testing crossValidation...")

X = [11:19 21:29 31:39 41:49 51:59 61:69]
Y = [1:9;]
sampler = KFold(nSplits=3,nRepeats=1,shuffle=true,rng=copy(TESTRNG))
(μ,σ) = crossValidation([X,Y],sampler) do trainData,valData,rng
        (xtrain,ytrain) = trainData; (xval,yval) = valData
        trainedModel = buildForest(xtrain,ytrain,30,rng=rng)
        predictions = predict(trainedModel,xval,rng=rng)
        ϵ = meanRelError(predictions,yval,normRec=false)
        return ϵ
    end

@test (μ,σ) == (0.3202242202242202, 0.04307662219315022)

# ==================================
# New test
println("** Testing shuffle()...")

a = [1 2 3; 10 20 30; 100 200 300; 1000 2000 4000; 10000 20000 40000]; b = [4,40,400,4000,40000]
out = shuffle([a,b],rng=copy(FIXEDRNG))
@test out[1] ==  [100 200 300; 1 2 3;10000 20000 40000; 10 20 30; 1000 2000 4000] && out[2] == [400,4,40000,40,4000]
out2 = shuffle(copy(FIXEDRNG),[a,b])
@test out2 == out

a = [1 2 3 4 5; 10 20 30 40 50]; b = [100 200 300 400 500]
out = shuffle([a,b],rng=copy(FIXEDRNG),dims=2)
@test out[1] == [3 1 5 2 4; 30 10 50 20 40] && out[2] == [300 100 500 200 400]

a = [1 2 30; 10 20 30]; b = [100 200 300];
(aShuffled, bShuffled) = shuffle([a,b],dims=2)

a = [1 2 3; 10 20 30]; b = [4,40]
out = shuffle([a,b],rng=copy(FIXEDRNG))

data = [a,b]
dims = 1
rng=Random.GLOBAL_RNG
Ns = [size(m,dims) for m in data]
length(Set(Ns)) == 1 || @error "In `shuffle(arrays)` all individual arrays need to have the same size on the dimension specified"
N    = Ns[1]
ridx = Random.shuffle(rng, 1:N)
out = similar(data)
for (i,a) in enumerate(data)
i = 1
a = data[1]
aidx = convert(Vector{Union{UnitRange{Int64},Vector{Int64}}},[1:i for i in size(a)])
aidx[dims] = ridx
out[i] = a[aidx...]
Vector{Union{UnitRange{Int64},Vector{Int64}}}
end
return out

# ==================================
# New test
println("** Testing generateParallelRngs()...")
x = rand(copy(TESTRNG),100)

function innerFunction(bootstrappedx; rng=Random.GLOBAL_RNG)
     sum(bootstrappedx .* rand(rng) ./ 0.5)
end
function outerFunction(x;rng = Random.GLOBAL_RNG)
    masterSeed = rand(rng,100:9999999999999) # important: with some RNG it is important to do this before the generateParallelRngs to guarantee independance from number of threads
    rngs       = generateParallelRngs(rng,Threads.nthreads()) # make new copy instances
    results    = Array{Float64,1}(undef,30)
    Threads.@threads for i in 1:30
        tsrng = rngs[Threads.threadid()]    # Thread safe random number generator: one RNG per thread
        Random.seed!(tsrng,masterSeed+i*10) # But the seeding depends on the i of the loop not the thread: we get same results indipendently of the number of threads
        toSample = rand(tsrng, 1:100,100)
        bootstrappedx = x[toSample]
        innerResult = innerFunction(bootstrappedx, rng=tsrng)
        results[i] = innerResult
    end
    overallResult = mean(results)
    return overallResult
end

# Different sequences..
@test outerFunction(x) != outerFunction(x)

# Different values, but same sequence
mainRng = copy(TESTRNG)
a = outerFunction(x, rng=mainRng)
b = outerFunction(x, rng=mainRng)

mainRng = copy(TESTRNG)
A = outerFunction(x, rng=mainRng)
B = outerFunction(x, rng=mainRng)

@test a != b && a == A && b == B

# Same value at each call
a = outerFunction(x,rng=copy(TESTRNG))
b = outerFunction(x,rng=copy(TESTRNG))
@test a == b

# ==================================
# New test
println("** Testing pool1d()...")
x = [1,2,3,4,5,6]
poolSize = 3

out = pool1d(x,poolSize)
@test out == [2.0,3.0,4.0,5.0]
out = pool1d(x,poolSize;f=maximum)
@test out == [3,4,5,6]

#=
using Random, StableRNGs
rDiff(rngFunction,seedBase,seedDiff,repetitions) = norm(rand(rngFunction(seedBase),repetitions) .- rand(rngFunction(seedBase+seedDiff),repetitions))/repetitions

# Seed base 1000: ok
rDiff(StableRNG,1000,1,100000)          # 0.00129
rDiff(StableRNG,1000,10,100000)         # 0.00129
rDiff(StableRNG,1000,1000,100000)       # 0.00129

rDiff(MersenneTwister,1000,1,100000)    # 0.00129
rDiff(MersenneTwister,1000,10,100000)   # 0.00129
rDiff(MersenneTwister,1000,1000,100000) # 0.00129

# Seed base 10: Still ok
rDiff(StableRNG,10,1,100000)            # 0.00129
rDiff(StableRNG,10,10,100000)           # 0.00129
rDiff(StableRNG,10,1000,100000)         # 0.00129

rDiff(MersenneTwister,10,1,100000)      # 0.00129
rDiff(MersenneTwister,10,10,100000)     # 0.00129
rDiff(MersenneTwister,10,1000,100000)   # 0.00129

# Seed base 1: We start seeing problems for StableRNG..
rDiff(StableRNG,1,1,100000)             # 0.00125  <--
rDiff(StableRNG,1,10,100000)            # 0.00129
rDiff(StableRNG,1,1000,100000)          # 0.00129

rDiff(MersenneTwister,1,1,100000)       # 0.00129
rDiff(MersenneTwister,1,10,100000)      # 0.00129
rDiff(MersenneTwister,1,1000,100000)    # 0.00129

# Seed base 0: Unexpected results for for StableRNG..
rDiff(StableRNG,0,1,100000)             # 0.00105 <----------
rDiff(StableRNG,0,2,100000)             # 0.00116 <-----
rDiff(StableRNG,0,5,100000)             # 0.00123 <---
rDiff(StableRNG,0,10,100000)            # 0.00126 <--
rDiff(StableRNG,0,1000,100000)          # 0.00129

rDiff(MersenneTwister,0,1,100000)       # 0.00130 <-
rDiff(MersenneTwister,0,2,100000)       # 0.00129
rDiff(MersenneTwister,0,5,100000)       # 0.00129
rDiff(MersenneTwister,0,10,100000)      # 0.00129
rDiff(MersenneTwister,0,1000,100000)    # 0.00129
=#