#
# grad_tools.jl
#
# Provide tools to load the instrumental parameters and the
# calibrated data, and for the calculus of the RHAPSODIE data fidelity term.
#
# ------------------------------------------------
#
# This file is part of Rhapsodie
#
#
# Copyright (c) 2017-2021 Laurence Denneulin (see LICENCE.md)
#

#------------------------------------------------

const ImageInterpolator{T<:AbstractFloat, K<:Kernel{T}} = TwoDimensionalTransformInterpolator{T,K,K}
const MyKer = LinearInterpolators.CatmullRomSpline(Float64, LinearInterpolators.Flat)
#const MyKer = LinearInterpolators.RectangularSpline(Float64, LinearInterpolators.Flat)
#const MyKer = LinearInterpolators.LinearSpline(Float64, LinearInterpolators.Flat)

struct parameters_table{T<: AbstractFloat}
     cols::NTuple{3,Int64} #Imput size (reconstruction)
     rows::NTuple{2,Int64} #Output size (data)
     dataset_length::Int64
     Nframe::Int64
     Nrot::Int64
     Nangle::Int64
     v::Vector{NTuple{2, NTuple{3, T}}}
     indices::Array{Int64,2}
     center::Array{T,1}
     psf_center::NTuple{2,Array{T,1}}
     epsilon::Vector{NTuple{2,Array{T,1}}}
     derotang::Vector{T}
end

struct FieldTransformOperator{T<:AbstractFloat, L<:Mapping, R<:Mapping} <: LinearMapping
    cols::NTuple{3,Int}
    rows::NTuple{2,Int}
    v_l::NTuple{3,T}
    v_r::NTuple{3,T}
    H_l::L
    H_r::R
end

struct data_table{T<:AbstractFloat,A<:FieldTransformOperator{T}} #autres types
    data::Array{T, 2};
    weights::Array{T, 2};
    H::A
end

const Trans_Table = Vector{NTuple{2,AffineTransform2D}}(); #Contains all the affine transform used
const Parameters = parameters_table[];
get_par()::parameters_table = Parameters[1];
const dataset = data_table[];

const PSF_save = Vector{Array{Float64,2}}();
get_PSF()=PSF_save[1];
const EPSILON_save = Array{Float64,1}(undef,1);
function set_epsilon(epsilon)
    EPSILON_save[1]=epsilon
end
get_epsilon()::Float64=EPSILON_save[1];
const PRECOND_SAVE= Vector{Any}();
U()=PRECOND_SAVE[1];

const MASK_save = Vector{Array{Float64,2}}();
get_MASK()=MASK_save[1];

function load_parameters(size_data::NTuple{3,Int64},
                         Nframe::Int64,
                         Nrot::Int64,
                         Nangle::Int64,
                         center::Array{Float64,1},
                         psf_center::NTuple{2,Array{Float64,1}},
                         epsilon::Vector{NTuple{2,Array{Float64,1}}})
    sets_indices=Indices(size_data[3],Nrot,Nframe)
    sets_v=Set_Vi(sets_indices)
    bbox_output=(0,0)
    Id = AffineTransform2D{Float64}()
    for k=1:size_data[3]
        A_left=translate( epsilon[k][1][1], epsilon[k][1][2], Id)
        A_right=translate( epsilon[k][2][1], epsilon[k][2][2], Id)

        out_left=bbox_size((size_data[1], Int64(size_data[2]/2)), A_left)[1]
        out_right=bbox_size((size_data[1], Int64(size_data[2]/2)), A_left)[1]

        xmax=max(bbox_output[1],out_left[1],out_right[1]);
        ymax=max(bbox_output[2],out_left[2],out_right[2]);

        bbox_output =(xmax,ymax);
    end
    bbox_output = bbox_output .+ 4;
    push!(Parameters, parameters_table((bbox_output[1],bbox_output[2],3),
                                       (size_data[1], size_data[2]),
                                       size_data[3],
                                       Nframe,
                                       Nrot,
                                       Nangle,
                                       sets_v,
                                       sets_indices,
                                       center,
                                       psf_center,
                                       epsilon,
                                       Float64[]));

        newcenter= (bbox_output .+1)./2;
        centerdiff=newcenter .- (center[1], center[2])
        for k=1:size_data[3]
        A_left=translate( epsilon[k][1][1] + centerdiff[1], epsilon[k][1][2] + centerdiff[2], Id)
        A_right=translate( epsilon[k][2][1] + centerdiff[1], epsilon[k][2][2] + centerdiff[2], Id)
        push!(Trans_Table, (A_left, A_right))
    end
end

function load_parameters(size_data::NTuple{3,Int64},
                         Nframe::Int64,
                         Nrot::Int64,
                         Nangle::Int64,
                         center::Array{Float64,1},
                         psf_center::NTuple{2,Array{Float64,1}},
                         epsilon::Vector{NTuple{2,Array{Float64,1}}},
                         derotang::Vector{Float64})
    sets_indices=Indices(size_data[3],Nrot,Nframe)
    sets_v=Set_Vi(sets_indices)
    bbox_output=(0,0)
    Id = AffineTransform2D{Float64}()
    for k=1:size_data[3]
        A_left=TransRotate(Id, epsilon[k][1], derotang[k], center, center)
        A_right=TransRotate(Id, epsilon[k][2], derotang[k], center, center)

        out_left=bbox_size((size_data[1], Int64(size_data[2]/2)), A_left)[1]
        out_right=bbox_size((size_data[1], Int64(size_data[2]/2)), A_left)[1]

        xmax=max(bbox_output[1],out_left[1],out_right[1]);
        ymax=max(bbox_output[2],out_left[2],out_right[2]);

        bbox_output =(xmax,ymax);
    end
    bbox_output = bbox_output .+ 4;
    push!(Parameters, parameters_table((bbox_output[1],bbox_output[2],3),
                                       (size_data[1], size_data[2]),
                                       size_data[3],
                                       Nframe,
                                       Nrot,
                                       Nangle,
                                       sets_v,
                                       sets_indices,
                                       center,
                                       psf_center,
                                       epsilon,
                                       derotang));

    newcenter= (bbox_output .+1)./2
    for k=1:size_data[3]
        A_left=inv(TransRotate(Id, epsilon[k][1], derotang[k], center, newcenter))
        A_right=inv(TransRotate(Id, epsilon[k][2], derotang[k], center, newcenter))
        push!(Trans_Table, (A_left, A_right))
    end
end

function load_parameters(size_data::NTuple{3,Int64},
                         Nframe::Int64,
                         Nrot::Int64,
                         Nangle::Int64,
                         center::Array{Float64,1},
                         psf_center::NTuple{2,Array{Float64,1}},
                         epsilon::Vector{NTuple{2,Array{Float64,1}}},
                         v::Array{Float64,2})
    sets_indices=Indices(size_data[3],Nrot,Nframe)
    sets_v=Set_Vi(Nframe, size_data[3], v)
    bbox_output=(0,0)
    Id = AffineTransform2D{Float64}()
    for k=1:size_data[3]
        A_left=translate( epsilon[k][1][1], epsilon[k][1][2], Id)
        A_right=translate( epsilon[k][2][1], epsilon[k][2][2], Id)

        out_left=bbox_size((size_data[1], Int64(size_data[2]/2)), A_left)[1]
        out_right=bbox_size((size_data[1], Int64(size_data[2]/2)), A_left)[1]

        xmax=max(bbox_output[1],out_left[1],out_right[1]);
        ymax=max(bbox_output[2],out_left[2],out_right[2]);

        bbox_output =(xmax,ymax);
    end
    bbox_output = bbox_output .+ 4;
    push!(Parameters, parameters_table((bbox_output[1],bbox_output[2],3),
                                       (size_data[1], size_data[2]),
                                       size_data[3],
                                       Nframe,
                                       Nrot,
                                       Nangle,
                                       sets_v,
                                       sets_indices,
                                       center,
                                       psf_center,
                                       epsilon,
                                       Float64[]));

        newcenter= (bbox_output .+1)./2;
        centerdiff=newcenter .- (center[1], center[2])
        for k=1:size_data[3]
        A_left=translate( epsilon[k][1][1] + centerdiff[1], epsilon[k][1][2] + centerdiff[2], Id)
        A_right=translate( epsilon[k][2][1] + centerdiff[1], epsilon[k][2][2] + centerdiff[2], Id)
        push!(Trans_Table, (A_left, A_right))
    end
end

function load_parameters(size_data::NTuple{3,Int64},
                         Nframe::Int64,
                         Nrot::Int64,
                         Nangle::Int64,
                         center::Array{Float64,1},
                         psf_center::NTuple{2,Array{Float64,1}},
                         epsilon::Vector{NTuple{2,Array{Float64,1}}},
                         derotang::Vector{Float64},
                         v::Array{Float64,2})
    sets_indices=Indices(size_data[3],Nrot,Nframe)
    sets_v=Set_Vi(Nframe, size_data[3], v)
    bbox_output=(0,0)
    Id = AffineTransform2D{Float64}()
    for k=1:size_data[3]
        A_left=TransRotate(Id, epsilon[k][1], derotang[k], center, center)
        A_right=TransRotate(Id, epsilon[k][2], derotang[k], center, center)

        out_left=bbox_size((size_data[1], Int64(size_data[2]/2)), A_left)[1]
        out_right=bbox_size((size_data[1], Int64(size_data[2]/2)), A_left)[1]

        xmax=max(bbox_output[1],out_left[1],out_right[1]);
        ymax=max(bbox_output[2],out_left[2],out_right[2]);

        bbox_output =(xmax,ymax);
    end
    bbox_output = bbox_output .+ 4;
    push!(Parameters, parameters_table((bbox_output[1],bbox_output[2],3),
                                       (size_data[1], size_data[2]),
                                       size_data[3],
                                       Nframe,
                                       Nrot,
                                       Nangle,
                                       sets_v,
                                       sets_indices,
                                       center,
                                       psf_center,
                                       epsilon,
                                       derotang));

    newcenter= (bbox_output .+1)./2
    for k=1:size_data[3]
        A_left=inv(TransRotate(Id, epsilon[k][1], derotang[k], center, newcenter))
        A_right=inv(TransRotate(Id, epsilon[k][2], derotang[k], center, newcenter))
        push!(Trans_Table, (A_left, A_right))
    end
end

function Load_Data(name_data, name_weight, name_psf)
	data=read(FitsArray, name_data);
	weight=read(FitsArray, name_weight);
	psf=read(FitsArray, name_psf);
	A_l=set_fft_op((psf[1:end÷2,:]'[:,:]), get_par().psf_center[1]);
	A_r=set_fft_op((psf[end÷2+1:end,:]'[:,:]), get_par().psf_center[1]);
	ker= MyKer;
	SetCropOperator()
    for k=1:size(data)[3]
        output_size=(get_par().rows[1], Int64(get_par().rows[2]/2));
        input_size= get_par().cols[1:2];
    	T1=TwoDimensionalTransformInterpolator(output_size, input_size, ker, ker, Trans_Table[k][1])
    	T2=TwoDimensionalTransformInterpolator(output_size, input_size, ker, ker, Trans_Table[k][2])
	    push!(dataset, data_table((data[:,:,k]')[:,:],
	                              (weight[:,:,k]')[:,:],
	                              FieldTransformOperator(get_par().cols,
	                                                     get_par().rows,
	                                                     get_par().v[k][1],
	                                                     get_par().v[k][2],
	                                                     T1,
	                                                     T2,
	                                                     A_l)));#,
	                                                     #A_r)));
	end
end

function Load_Data(name_data, name_weight)
	data=read(FitsArray, name_data);
	weight=read(FitsArray, name_weight);
	ker= MyKer;
	SetCropOperator()
    for k=1:size(data)[3]
        output_size=(get_par().rows[1], Int64(get_par().rows[2]/2));
        input_size= get_par().cols[1:2];
    	T1=TwoDimensionalTransformInterpolator(output_size, input_size, ker, ker, Trans_Table[k][1])
    	T2=TwoDimensionalTransformInterpolator(output_size, input_size, ker, ker, Trans_Table[k][2])
    	promote
	    push!(dataset, data_table((data[:,:,k]')[:,:],
	                              (weight[:,:,k]')[:,:],
	                              FieldTransformOperator(get_par().cols,
	                                                     get_par().rows,
	                                                     get_par().v[k][1],
	                                                     get_par().v[k][2],
	                                                     T1,
	                                                     T2,
	                                                     LazyAlgebra.Id)));
	end
end

function Indices(S::Int64,Nrot::Int64,Nframe::Int64)
#S = le nombre total de frame gauche ou droite
#Nrot : Nombre de positions du derotateur
#Nframe : Nombre de frame par positions de la lame demi-onde

	Nangle=Int32.(S/(Nrot*Nframe*4)) #Nombre de rotations de la lame demi-onde
	INDICES=zeros(4,Nframe*Nangle*Nrot)
	for i=1:4
		ind=repeat(range(0,stop=4*Nframe*(Nrot*Nangle-1),length=Nrot*Nangle), inner=Nframe)+(Nframe*i .-mod.(range(1,stop=Nframe*Nangle*Nrot,length=Nframe*Nangle*Nrot),Nframe))
		INDICES[i,:]=ind
	end
	INDICES=round.(Int64, INDICES);
	return INDICES
end

function Set_Vi(Indices::AbstractArray{Int64,2}; alpha=[0, pi/4, pi/8, 3*pi/8], psi=[0,pi/2])
	J(a)=[cos(2*a) sin(2*a); sin(2*a) -cos(2*a)]
	P(a)=[cos(a), -sin(a)]
    v=Vector{NTuple{2, NTuple{3, Float64}}}(undef, length(Indices));
	for i=1:4
		v1=J(alpha[i])*P(psi[1])
		vn1=norm(v1)^2;
		v2=J(alpha[i])*P(psi[2])
		vn2=norm(v2)^2;
		for k=1:length(Indices[i,:])
		v[Indices[i,k]] =((round(vn1/2, digits=4), round((abs(v1[1])^2-abs(v1[2])^2)/2, digits=4), round(real(v1[1]*v1[2]), digits=4)),
		                  (round(vn2/2, digits=4), round((abs(v2[1])^2-abs(v2[2])^2)/2, digits=4), round(real(v2[1]*v2[2]), digits=4)));
        end
	end
	return v
end

function Set_Vi(Nframe::Int, dataset_length::Int, mueller_instru::AbstractArray{Float64,2})
    @assert dataset_length == Nframe * size(mueller_instru)[1]

    v=Vector{NTuple{2, NTuple{3, Float64}}}(undef, dataset_length);
    for k=1:size(mueller_instru)[1]
        for l=1:Nframe
            v[Nframe*(k-1) + l] =((mueller_instru[k,1], mueller_instru[k,2], mueller_instru[k,3]),
                                  (mueller_instru[k,4], mueller_instru[k,5], mueller_instru[k,6]));
        end
    end
	return v
end

function reset_instrument(V::Vector{NTuple{2, NTuple{3, T}}})  where {T <: AbstractFloat}
    push!(Parameters, parameters_table(get_par().cols,
                                       get_par().rows,
                                       get_par().dataset_length,
                                       get_par().Nframe,
                                       get_par().Nrot,
                                       get_par().Nangle,
                                       V,
                                       get_par().indices,
                                       get_par().center,
                                       get_par().psf_center,
                                       get_par().epsilon,
                                       get_par().derotang));

    popfirst!(Parameters);
end

function TransRotate(A::AffineTransform2D{Float64}, EPSILON, ANGLE, CENTER, NEWCENTER)
	return translate(rotate(translate(CENTER[1]-EPSILON[1],CENTER[2]-EPSILON[2], A), ANGLE),-NEWCENTER[1], -NEWCENTER[2])
end

function bbox_size(inp_dims::NTuple{2,Integer},
                  A::AffineTransform2D{Float64})

	xmin=typemax(Float64);
	ymin=typemax(Float64);
	xmax=typemin(Float64);
	ymax=typemin(Float64);

	xmin_int=typemin(Int64);
	ymin_int=typemin(Int64);
	xmax_int=typemin(Int64);
	ymax_int=typemin(Int64);

	width=typemin(Float64);
	height=typemin(Float64);

	Ind=[repeat(1:inp_dims[1], inner=inp_dims[2]) repeat(1:inp_dims[2], outer=inp_dims[1])]
	for it=1:(inp_dims[1]*inp_dims[2])
		(x,y)=A(Ind[it,1], Ind[it,2]);
		xmin=min(x, xmin);
		ymin=min(y, ymin);
		xmax=max(x, xmax);
		ymax=max(y, ymax);
	end
	xmin_int=floor(Int64,xmin);
	ymin_int=floor(Int64,ymin);
	xmax_int=ceil(Int64,xmax);
	ymax_int=ceil(Int64,ymax);

	width=xmax_int-xmin_int+1+50;
	height=ymax_int-ymin_int+1+50;
    out_dims=(width, height);

	return(out_dims, (xmin_int, xmax_int, ymin_int, ymax_int))
end

function set_fft_op(PSF::AbstractArray{T,2}, PSFCenter::AbstractArray{T,1}) where {T <: AbstractFloat}
 	MapSize=get_par().cols[1:2];
	MapCenter=floor.(MapSize./2).+1
	MAP=zeros(MapSize);
	ker = LinearInterpolators.CatmullRomSpline(Float64, LinearInterpolators.Flat)

	Id = AffineTransform2D{Float64}()
	Recentering=translate(-(MapCenter[1]-PSFCenter[1]), -(MapCenter[2]-PSFCenter[2]), Id)

	LazyAlgebra.apply!(MAP, ker, Recentering, PSF);
	MAP./=sum(MAP);
	push!(PSF_save,MAP);
	F=FFTOperator(MAP)
	FFT=F\Diag(F*ifftshift(MAP)) .*F;

	return FFT
end

function FieldTransformOperator(cols::NTuple{3,Int64},
                                rows::NTuple{2,Int64},
                                v_l::NTuple{3,T},
                                v_r::NTuple{3,T},
                                T_l::TwoDimensionalTransformInterpolator{T},
                                T_r::TwoDimensionalTransformInterpolator{T},
                                A::M) where {T <: AbstractFloat, M <: Mapping}
    FieldTransformOperator(cols, rows, v_l, v_r, T_l*A, T_r*A)
end

function vcreate(::Type{LazyAlgebra.Direct}, A::FieldTransformOperator{T},
                 x::AbstractArray{T,3}, scratch::Bool = false) where {T <: AbstractFloat}
    @assert !Base.has_offset_axes(x)
    @assert size(x) == A.cols
    Array{T,2}(undef, A.rows)
end

function vcreate(::Type{LazyAlgebra.Adjoint}, A::FieldTransformOperator{T},
                 x::AbstractArray{T,2}, scratch::Bool = false) where {T <: AbstractFloat}
    @assert !Base.has_offset_axes(x)
    @assert size(x) == A.rows
    Array{T,3}(undef, A.cols)
end

function apply!(α::Real,
                ::Type{LazyAlgebra.Direct},
                R::FieldTransformOperator{T},
                src::AbstractArray{T,3},
                scratch::Bool,
                β::Real,
                dst::AbstractArray{T,2}) where {T<:AbstractFloat}
    @assert β==0 && α==1
    @assert size(src) == R.cols
    @assert size(dst) == R.rows
    n = R.rows[2]
    @assert iseven(n)

    z_l = zeros(R.cols[1:2]);
    z_r = zeros(R.cols[1:2]);
    for i=1:length(R.v_l)
        # FIXME: use Generalized Matrix-Vector Multiplication of LazyAlgebra here.
        x_i = view(src,:,:,i)
	z_l .+= R.v_l[i]*x_i;
	z_r .+= R.v_r[i]*x_i;
    end
    dst[:, 1:(n÷2)  ] = R.H_l*z_l;
    dst[:, (n÷2)+1:n] = R.H_r*z_r;
    return dst
end

function apply!(α::Real,
                ::Type{LazyAlgebra.Adjoint},
                R::FieldTransformOperator{T},
                src::AbstractArray{T,2},
                scratch::Bool,
                β::Real,
                dst::AbstractArray{T,3}) where {T<:AbstractFloat}
    @assert β==0 && α==1
    @assert size(src) == R.rows
    @assert size(dst) == R.cols
    n = R.rows[2]
    @assert iseven(n)
    y_l = R.H_l'*view(src, :, 1:(n÷2))
    y_r = R.H_r'*view(src, :, (n÷2)+1:n)
    for i=1:length(R.v_l)
        #dst[:,:,i] = R.v_l[i]*y_l + R.v_r[i]*y_r;
        # FIXME: The following should be equivalent and much faster:
         vcombine!(view(dst,:,:,i), R.v_l[i], y_l, R.v_r[i], y_r)
    end
    return dst;
end

#=
function apply!(α::Real,
                ::Type{LazyAlgebra.Direct},
                R::FieldTransformOperator{T},
                src::AbstractArray{T,3},
                scratch::Bool,
                β::Real,
                dst::AbstractArray{T,2}) where {T<:AbstractFloat}

    @assert size(src) == R.cols
    @assert size(dst) == R.rows
    @assert β==0 && α==1
    vfill!(dst,0);

    dst_left=zeros(R.cols[1:2]);
    dst_right=zeros(R.cols[1:2]);

    for i=1:length(R.v_l)
        dst_left .+=(R.v_l[i]*src[:,:,i]);
        dst_right .+=(R.v_r[i]*src[:,:,i]);
    end
    dst .+= [R.T_l*R.A_l*dst_left R.T_r*R.A_r*dst_right];
    return dst
end

function apply!(α::Real,
                ::Type{LazyAlgebra.Adjoint},
                R::FieldTransformOperator{T},
                src::AbstractArray{T,2},
                scratch::Bool,
                β::Real,
                dst::AbstractArray{T,3}) where {T<:AbstractFloat}
    @assert β==0 && α==1
    @assert size(src) == R.rows
    @assert size(dst) == R.cols
    for i=1:length(R.v_l)
        dst_left=R.A_l'*R.T_l'*(R.v_l[i]*src[:,1:end÷2]);
        dst_right=R.A_r'*R.T_r'*(R.v_r[i]*src[:,end÷2+1:end]);
        dst[:,:,i] = dst_left + dst_right;
    end
    return dst;
end
=#

function mydot(x::AbstractArray{Tx,N},
               y::AbstractArray{Ty,N}) where {Tx<:AbstractFloat,Ty<:AbstractFloat,N}
    @assert axes(x) == axes(y)
    s = zero(promote_type(Tx,Ty))
    @inbounds @simd for i in eachindex(x,y)
       s += x[i]*y[i]
    end
    return s
end

function fg!(x::AbstractArray{T,3},g::AbstractArray{T,3}) where {T<:AbstractFloat}
    local f::Float64 = 0.0;
    vfill!(g,0.0)
    for k=1:length(dataset)
        if sum(dataset[k].weights) !=0
            f += fg!(x, g, dataset[k])
        end
    end
    #crop!(g)
    return f
end

function fg!(x::AbstractArray{T,3},g::AbstractArray{T,3}, d::data_table{T}) where {T<:AbstractFloat}
    r = d.H*x - d.data;
    wr = d.weights .* r;
    g .+= d.H'*wr;
    #f = vdot(Float64, r,wr); #FIXME : mettre une issue !
    f = mydot(r,wr);
    #crop!(g)
    return Float64(f)/2
end

#--------------------------------------------UTILS------------------------------------------------

function SetCropOperator()
    DSIZE=get_par().rows[1]
    MASK=ones(get_par().cols[1:2]);
    for k=1:length(Trans_Table)
    X1=Trans_Table[k][1](1,1);
    X2=Trans_Table[k][1](1,DSIZE);
    X3=Trans_Table[k][1](DSIZE, DSIZE);
    X4=Trans_Table[k][1](DSIZE, 1);
    Mask1=zeros(get_par().cols[1:2]);

    for i=1:get_par().cols[1]
	    for j=1:get_par().cols[2]
		    if ((i-X1[1])*(X2[1]-X1[1]) +(j-X1[2])*(X2[2]-X1[2]) >0)&&((i-X2[1])*(X3[1]-X2[1]) +(j-X2[2])*(X3[2]-X2[2]) >0)&&((i-X3[1])*(X4[1]-X3[1]) +(j-X3[2])*(X4[2]-X3[2]) >0) &&((i-X4[1])*(X1[1]-X4[1]) +(j-X4[2])*(X1[2]-X4[2]) >0)
		    Mask1[i,j,:] .=1;
		    end
	    end
    end
    X1=Trans_Table[k][2](1,1);
    X2=Trans_Table[k][2](1,DSIZE);
    X3=Trans_Table[k][2](DSIZE, DSIZE);
    X4=Trans_Table[k][2](DSIZE, 1);
    Mask2=zeros(get_par().cols[1:2]);

    for i=1:get_par().cols[1]
	    for j=1:get_par().cols[2]
		    if ((i-X1[1])*(X2[1]-X1[1]) +(j-X1[2])*(X2[2]-X1[2]) >0)&&((i-X2[1])*(X3[1]-X2[1]) +(j-X2[2])*(X3[2]-X2[2]) >0)&&((i-X3[1])*(X4[1]-X3[1]) +(j-X3[2])*(X4[2]-X3[2]) >0) &&((i-X4[1])*(X1[1]-X4[1]) +(j-X4[2])*(X1[2]-X4[2]) >0)
		    Mask2[i,j,:] .=1;
		    end
	    end
    end
    MASK .*= Mask1.*Mask2
    end
    #write(FitsFile, "MASK.fits", MASK, overwrite=true)
    push!(MASK_save, MASK)
end

function crop(X::M)  where {T<:AbstractFloat, M<:AbstractArray{T,2}}
    #@assert size(X) .==   get_par().cols
    xymin_1=findfirst(get_MASK() .!= 0.)
    xymin_2=findfirst(get_MASK()' .!= 0.)
    xymax_1=findlast(get_MASK() .!= 0.)
    xymax_2=findlast(get_MASK()' .!= 0.)

    xmin=min(xymin_1[1], xymin_2[2]);
    ymin=min(xymin_1[2], xymin_2[1]);
    xmax=max(xymax_1[1], xymax_2[2]);
    ymax=max(xymax_1[2], xymax_2[1]);

    Y=X[xmin:xmax, ymin:ymax];
    Y[.!isfinite.(Y)].=0
    return Y
end

function crop!(X::M)  where {T<:AbstractFloat, M<:AbstractArray{T,2}}
    #@assert size(X) .==   get_par().cols
    X[.!isfinite.(X)].=0
    X .*= get_MASK()
end

function crop(X::M)  where {T<:AbstractFloat, M<:AbstractArray{T,3}}
    #@assert size(X) .==   get_par().cols
    Y=copy(X);
    return crop!(Y)
end

function crop!(X::M)  where {T<:AbstractFloat, M<:AbstractArray{T,3}}
    #@assert size(X) .==   get_par().cols
    for k in size(X)[3]
        crop!(view(X,:,:,k))
    end
end

function crop(X::PolarimetricMap{T}) where {T<:AbstractFloat}
    #@assert size(X) .==   get_par().cols
    return PolarimetricMap(X.parameter_type,
                           crop(view(X.I,:,:)),
                           crop(view(X.Q,:,:)),
                           crop(view(X.U,:,:)),
                           crop(view(X.Iu,:,:)),
                           crop(view(X.Ip,:,:)),
                           crop(view(X.θ,:,:)))
end

function crop!(X::PolarimetricMap{T})  where {T<:AbstractFloat}
        crop!(view(X.I,:,:));
        crop!(view(X.Q,:,:));
        crop!(view(X.U,:,:));
        crop!(view(X.Iu,:,:));
        crop!(view(X.Ip,:,:));
        crop!(view(X.θ,:,:));
end


#####
##### AlphaZero.jl
##### Jonathan Laurent, Carnegie Mellon University (2019-2020)
#####

module AlphaZero

  import Distributed
  using Formatting
  using Base: @kwdef
  using DataStructures: CircularBuffer
  using Distributions: Categorical, Dirichlet
  using Statistics: mean

  # When running on a CPU, having multiple threads does not play
  # well with BLAS multithreading
  import LinearAlgebra
  LinearAlgebra.BLAS.set_num_threads(1)

  # Internal helper functions
  include("util.jl")
  using .Util
  export Util
  export apply_temperature

  # A generic interface for single-player or zero-sum two-players games.
  include("game.jl")
  using .GameInterface
  const GI = GameInterface
  export GameInterface, GI
  export AbstractGameEnv
  export AbstractGameSpec

  # A standalone, generic MCTS implementation
  include("mcts.jl")
  using .MCTS
  export MCTS

  # A generic network interface
  include("networks/network.jl")
  using .Network
  export Network
  export AbstractNetwork
  export OptimiserSpec
  export CyclicNesterov, Adam

  # Utilities to batch oracle calls
  include("batchifier.jl")
  using .Batchifier
  export Batchifier

  # Schedules
  include("schedule.jl")
  export AbstractSchedule
  export ConstSchedule, PLSchedule, StepSchedule, CyclicSchedule

  # Training params
  include("params.jl")
  export Params
  export MctsParams
  export SimParams
  export SelfPlayParams
  export LearningParams
  export ArenaParams
  export MemAnalysisParams
  export SamplesWeighingPolicy, CONSTANT_WEIGHT, LOG_WEIGHT, LINEAR_WEIGHT

  # Stats about training
  include("report.jl")
  export Report

  # Game traces
  include("trace.jl")
  export Trace
  export total_reward

  # Players and games
  include("play.jl")
  export AbstractPlayer, think, select_move, reset_player!, player_temperature
  export MctsPlayer
  export RandomPlayer
  export NetworkPlayer
  export PlayerWithTemperature
  export EpsilonGreedyPlayer
  export TwoPlayers
  export Human, interactive!
  export play_game

  # Utilities for distributed games simulation
  include("simulations.jl")
  export Simulator, simulate, simulate_distributed
  export record_trace
  export rewards_and_redundancy

  # Memory buffer to hold samples generated during self-play
  include("memory.jl")
  export MemoryBuffer, get_experience

  # Utilities to train the neural network based on collected samples
  include("learning.jl")

  # Main training algorithm
  include("training.jl")
  export Env, train!, initial_report
  export Handlers
  export AlphaZeroPlayer

  # A minmax player to be used as a baseline
  include("minmax.jl")
  using .MinMax
  export MinMax

  # Utilities to write benchmarks
  include("benchmark.jl")
  using .Benchmark
  export Benchmark

  # We provide a library of standard network, both in Knet and Flux.
  # Which backend is used to implement this library is determined during precompilation
  # based on the value of the ALPHAZERO_DEFAULT_DL_FRAMEWORK environment variable.
  const DEFAULT_DL_FRAMEWORK = get(ENV, "ALPHAZERO_DEFAULT_DL_FRAMEWORK", "FLUX")

  # Only Flux works with CUDA,NNlib#master right now.

  if DEFAULT_DL_FRAMEWORK == "FLUX"
    # @info "Using the Flux implementation of AlphaZero.NetLib."
    KNET_ADVICE =
    """
    For optimal performances, we recommend that you configure AlphaZero.jl
    to use Knet as a backend by setting ALPHAZERO_DEFAULT_DL_FRAMEWORK to KNET.
    """
    # Knet does not work with the latest CUDA.jl so we disable this advice for now.
    # @info KNET_ADVICE
    @eval begin
      include("networks/flux.jl")
      const NetLib = FluxLib
    end
  elseif DEFAULT_DL_FRAMEWORK == "KNET"
    # @info "Using the Knet implementation of AlphaZero.NetLib."
    @eval begin
      include("networks/knet.jl")
      const NetLib = KnetLib
    end
  else
    error("Unknown DL framework: $(DEFAULT_DL_FRAMEWORK)")
  end

  using .NetLib
  export NetLib
  export SimpleNet, SimpleNetHP, ResNet, ResNetHP

  # A structure that contains the information necessary to replicate a training session
  include("experiments.jl")
  using .Experiments
  export Experiments
  export Experiment

  # The default user interface is included here for convenience but it could be
  # replaced or separated from the main AlphaZero.jl package (which would also
  # enable dropping some dependencies such as Crayons or JSON3).
  include("ui/ui.jl")
  using .UserInterface
  const UI = UserInterface
  export UserInterface, UI
  export Session, resume!, save
  export explore

  # Bridge with CommonRLInterface.jl
  include("common_rl_intf.jl")
  export CommonRLInterfaceWrapper

  # A small library of standard examples
  include("examples.jl")
  export Examples

  # Scripts
  include("scripts/scripts.jl")
  export Scripts

end


using seisimu

function extrap!{T1,T2<:Real}(model::elmod3d, threadim::Array{Int64}, CPU_mem::T1, GPU_mem::T2; num_stream=1)

    if length(threadim) != 3
        error("Thread dimension should be 3. Use [threadim1, threadim2, threadim3].")
    else
        chk_group_full, chk_group_half, chk_stream_full, chk_stream_half = gpuchunk("elastic",
        model.medium.BDnDZ, model.medium.BDnHX,
        model.medium.BDnHY,"double",GPU_mem, CPU_mem, num_stream)

        Num_group = length(chk_group_half[2:end])
        Max_group_dim = maximum(chk_group_half)
        Max_stream_dim = maximum(chk_stream_half[3:end,:])
        Max_num_stream = maximum(chk_stream_half[1,:])

        blockdim = [Int64(ceil(model.medium.BDnHX/threadim[1])),
        Int64(ceil(Max_stream_dim/threadim[2])),
        Int64(ceil(model.medium.BDnDZ/threadim[3]))]

        path="/home/lzh/Dropbox/Zhenhua/Ongoing/Seisimu/deps/src/cump/el3d_cump.so"
        ccall((:el3d_cump,path),
        Void,
        (Ptr{Cdouble}, Cint, Cint, Cint, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble},
         Ptr{Cdouble}, Cint, Cint, Cint, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble},
         Ptr{Cdouble}, Cint, Cint, Cint, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble},
         Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble},
         Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble},
         Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble},
         Cint, Cint, Cint,
         Ptr{Cdouble}, Cint, Cint, Cint, Ptr{Cdouble}, Ptr{Cdouble},
         Ptr{Cdouble}, Cint, Cint, Cint, Ptr{Cdouble}, Ptr{Cdouble},
         Ptr{Cdouble}, Cint, Cint, Cint, Ptr{Cdouble}, Ptr{Cdouble},
         Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble},
         Cdouble, Cdouble, Cdouble, Cdouble, Cint,
         Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble},
         Ptr{Cint}, Ptr{Cint}, Cint, Cint,
         Ptr{Cint}, Ptr{Cint}, Cint, Cint,
         Ptr{Cint}, Ptr{Cint}),
         model.wf.vx, model.nwf.BDnvx[2], model.nwf.BDnvx[3], model.nwf.BDnvx[1],
         model.pml.PVxBTxx, model.pml.PVxBTxy, model.pml.PVxBTxz,
         model.wf.vy, model.nwf.BDnvy[2], model.nwf.BDnvy[3], model.nwf.BDnvy[1],
         model.pml.PVyBTxy, model.pml.PVyBTyy, model.pml.PVyBTyz,
         model.wf.vz, model.nwf.BDnvz[2], model.nwf.BDnvz[3], model.nwf.BDnvz[1],
         model.pml.PVzBTxz, model.pml.PVzBTyz, model.pml.PVzBTzz,
         model.wf.txx, model.pml.PTxxBVx, model.pml.PTxxBVy, model.pml.PTxxBVz,
         model.wf.tyy, model.pml.PTyyBVx, model.pml.PTyyBVy, model.pml.PTyyBVz,
         model.wf.tzz, model.pml.PTzzBVx, model.pml.PTzzBVy, model.pml.PTzzBVz,
         model.nwf.BDntpp[2], model.nwf.BDntpp[3], model.nwf.BDntpp[1],
         model.wf.txy, model.nwf.BDntxy[2], model.nwf.BDntxy[3], model.nwf.BDntxy[1],
         model.pml.PTxyBVx, model.pml.PTxyBVy,
         model.wf.tyz, model.nwf.BDntyz[2], model.nwf.BDntyz[3], model.nwf.BDntyz[1],
         model.pml.PTyzBVy, model.pml.PTyzBVz,
         model.wf.txz, model.nwf.BDntxz[2], model.nwf.BDntxz[3], model.nwf.BDntxz[1],
         model.pml.PTxzBVx, model.pml.PTxzBVz,
         model.medium.rho, model.medium.lambda, model.medium.mu, model.fdc,
         model.medium.dt, model.medium.dx, model.medium.dy, model.medium.dz, model.medium.ext,
         model.pml.bhalf, model.pml.ahalf, model.pml.bfull, model.pml.afull,
         chk_group_full, chk_group_half, Num_group, Max_group_dim,
         chk_stream_full, chk_stream_half, Max_num_stream, Max_stream_dim,
         threadim, blockdim)
    end
end

function extrap!(model::elmod3d)

        path="/home/lzh/Dropbox/Zhenhua/Ongoing/Seisimu/deps/src/cump/el3d_openmp.so"

        ccall((:el3d_openmp,path),
        Void,
        (Ptr{Cdouble}, Cint, Cint, Cint, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble},
         Ptr{Cdouble}, Cint, Cint, Cint, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble},
         Ptr{Cdouble}, Cint, Cint, Cint, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble},
         Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble},
         Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble},
         Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble},
         Cint, Cint, Cint,
         Ptr{Cdouble}, Cint, Cint, Cint, Ptr{Cdouble}, Ptr{Cdouble},
         Ptr{Cdouble}, Cint, Cint, Cint, Ptr{Cdouble}, Ptr{Cdouble},
         Ptr{Cdouble}, Cint, Cint, Cint, Ptr{Cdouble}, Ptr{Cdouble},
         Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble},
         Cdouble, Cdouble, Cdouble, Cdouble, Cint,
         Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}),
         model.wf.vx, model.nwf.BDnvx[2], model.nwf.BDnvx[3], model.nwf.BDnvx[1],
         model.pml.PVxBTxx, model.pml.PVxBTxy, model.pml.PVxBTxz,
         model.wf.vy, model.nwf.BDnvy[2], model.nwf.BDnvy[3], model.nwf.BDnvy[1],
         model.pml.PVyBTxy, model.pml.PVyBTyy, model.pml.PVyBTyz,
         model.wf.vz, model.nwf.BDnvz[2], model.nwf.BDnvz[3], model.nwf.BDnvz[1],
         model.pml.PVzBTxz, model.pml.PVzBTyz, model.pml.PVzBTzz,
         model.wf.txx, model.pml.PTxxBVx, model.pml.PTxxBVy, model.pml.PTxxBVz,
         model.wf.tyy, model.pml.PTyyBVx, model.pml.PTyyBVy, model.pml.PTyyBVz,
         model.wf.tzz, model.pml.PTzzBVx, model.pml.PTzzBVy, model.pml.PTzzBVz,
         model.nwf.BDntpp[2], model.nwf.BDntpp[3], model.nwf.BDntpp[1],
         model.wf.txy, model.nwf.BDntxy[2], model.nwf.BDntxy[3], model.nwf.BDntxy[1],
         model.pml.PTxyBVx, model.pml.PTxyBVy,
         model.wf.tyz, model.nwf.BDntyz[2], model.nwf.BDntyz[3], model.nwf.BDntyz[1],
         model.pml.PTyzBVy, model.pml.PTyzBVz,
         model.wf.txz, model.nwf.BDntxz[2], model.nwf.BDntxz[3], model.nwf.BDntxz[1],
         model.pml.PTxzBVx, model.pml.PTxzBVz,
         model.medium.rho, model.medium.lambda, model.medium.mu, model.fdc,
         model.medium.dt, model.medium.dx, model.medium.dy, model.medium.dz, model.medium.ext,
         model.pml.bhalf, model.pml.ahalf, model.pml.bfull, model.pml.afull)
    end


#!/usr/bin/env julia

using Dtree
if VERSION > v"0.5.0-dev"
    using Base.Threads
else
    # Pre-Julia 0.5 there are no threads
    nthreads() = 1
    threadid() = 1
    macro threads(x)
        x
    end
    SpinLock() = 1
    lock(l) = ()
    unlock(l) = ()
end

const cpu_hz = 2.3e9

const libdtree = joinpath(Pkg.dir("Dtree"), "deps", "Dtree",
        "libdtree.$(Libdl.dlext)")

@inline rdtsc() = ccall((:rdtsc, libdtree), Culonglong, ())
@inline cpu_pause() = ccall((:cpu_pause, libdtree), Void, ())
@inline secs2cpuhz(s) = s * cpu_hz::Float64

@inline function ntputs(nid,tid,s)
    ccall(:puts, Cint, (Cstring,), string("[", nid, "]<", tid, "> ", s))
    return
end

function threadfun(dt, ni, ci, li, ilock, rundt, dura)
    tid = threadid()
    if rundt && tid == 1
        ntputs(dt_nodeid, tid, "running tree")
        while runtree(dt)
            cpu_pause()
        end
    else
        ntputs(dt_nodeid, tid, string("begin, ", ni[], " items, ", length(dura), " available delays"))
        while ni[] > 0
            lock(ilock)
            if li[] == 0
                ntputs(dt_nodeid, tid, string("out of work"))
                unlock(ilock)
                break
            end
            if ci[] == li[]
                ntputs(dt_nodeid, tid, string("work consumed (last was ", li[], "); requesting more"))
                ni[], (ci[], li[]) = getwork(dt)
                ntputs(dt_nodeid, tid, string("got ", ni[], " work items (", ci[], " to ", li[], ")"))
                unlock(ilock)
                continue
            end
            item = ci[]
            ci[] = ci[] + 1
            unlock(ilock)

            # wait dura[item] seconds
            ticks = secs2cpuhz(dura[item])
            #ntputs(dt_nodeid, tid, string("item ", item, ", ", ticks, " ticks"))
            startts = rdtsc()
            while rdtsc() - startts < ticks
                cpu_pause()
            end
        end
    end
end

function bench(nwi, meani, stddevi, first_distrib, rest_distrib, min_distrib, fan_out = 1024)
    # create the tree
    dt, is_parent = DtreeScheduler(fan_out, nwi, true, 1.0, first_distrib, rest_distrib, min_distrib)

    # ---
    if dt_nodeid == 1
        println("dtreebench -- ", dt_nnodes, " nodes")
        println("  system clock speed is ", cpu_hz, " GHz")
    end

    # roughly how many work items will each node will handle?
    each, r = divrem(nwi, dt_nnodes)
    if r > 0
        each = each + 1
    end

    # ---
    if dt_nodeid == 1
        println("  ", nwi, " work items, ~", each, " per node")
    end

    # generate random numbers for work item durations
    dura = Float64[]
    mn = repmat([meani-0.5*stddevi, meani+0.5*stddevi], ceil(Int, dt_nnodes/2))
    mt = MersenneTwister(7777777)
    for i = 1:dt_nnodes
        r = randn(mt)*stddevi*0.25+mn[i]
        append!(dura, max(randn(mt, each)*stddevi+r, zero(Float64)))
    end

    # ---
    if dt_nodeid == 1
        println("  initializing...")
    end

    # get the initial allocation
    ilock = SpinLock()
    ni, (ci, li) = initwork(dt)

    # ---
    if dt_nodeid == 1
        println("  ...done.")
    end

    # start threads and run, or run single-threaded
    if VERSION > v"0.5.0-dev"
        tfargs = Core.svec(threadfun, dt, Ref(ni), Ref(ci), Ref(li), ilock, runtree(dt), dura)
        ccall(:jl_threading_run, Void, (Any,), tfargs)
    else
        threadfun(dt, Ref(ni), Ref(ci), Ref(li), ilock, runtree(dt), dura)
    end

    # ---
    if dt_nodeid == 1
        println("complete")
    end
    tic()
    finalize(dt)
    wait_done = toq()
    ntputs(dt_nodeid, 1, "wait for done: $wait_done secs")
end

#bench(80, 0.5, 0.125, 0.2, 0.5, nthreads())
bench(100, 0.5, 0.125, 0.5, 0.5, nthreads())


"""
    JuMP.build_variable(::Function, info::JuMP.VariableInfo, set::T) where T <: MOI.AbstractScalarSet

DOCSTRING

# Arguments:
- ``: DESCRIPTION
- `info`: DESCRIPTION
- `set`: DESCRIPTION
"""
function JuMP.build_variable(
    ::Function,
    info::JuMP.VariableInfo,
    set::T,
) where {T<:MOI.AbstractScalarSet}
    return JuMP.VariableConstrainedOnCreation(JuMP.ScalarVariable(info), set)
end

"""
    Optimizer <: MOI.AbstractOptimizer

DOCSTRING

# Arguments:
- `solver::Solver`: DESCRIPTION
- `status::MOI.TerminationStatusCode`: DESCRIPTION
- `options::Options`: DESCRIPTION
"""
mutable struct Optimizer <: MOI.AbstractOptimizer
    solver::LS.MainSolver
    int_vars::Set{Int}
    compare_vars::Set{Int}
end

"""
    Optimizer(model = Model(); options = Options())

DOCSTRING
"""
function Optimizer(model = model(); options = Options())
    return Optimizer(
        solver(model, options = options),
        Set{Int}(),
        Set{Int}()
        )
end

# forward functions from Solver
@forward Optimizer.solver LS.variable!, LS._set_domain!, LS.constraint!, LS.solution
@forward Optimizer.solver LS.max_domains_size, LS.objective!, Base.empty!, LS._inc_cons!
@forward Optimizer.solver LS._best_bound, LS.best_value, LS.is_sat, LS.get_value
@forward Optimizer.solver LS.domain_size, LS.best_values, LS._max_cons, LS.update_domain!
@forward Optimizer.solver LS.get_variable, LS.has_solution, LS.sense, LS.sense!
@forward Optimizer.solver LS.time_info, LS.status

# forward functions from Solver (from Options)
@forward Optimizer.solver LS._verbose, LS.set_option!, LS.get_option

"""
    MOI.get(::Optimizer, ::MOI.SolverName) = begin

DOCSTRING
"""
MOI.get(::Optimizer, ::MOI.SolverName) = "LocalSearchSolvers"

"""
    MOI.set(::Optimizer, ::MOI.Silent, bool = true) = begin

DOCSTRING

# Arguments:
- ``: DESCRIPTION
- ``: DESCRIPTION
- `bool`: DESCRIPTION
"""
MOI.set(::Optimizer, ::MOI.Silent, bool = true) = @debug "TODO: Silent"

"""
    MOI.is_empty(model::Optimizer) = begin

DOCSTRING
"""
MOI.is_empty(model::Optimizer) = LS._is_empty(model.solver)

"""
Copy constructor for the optimizer
"""
MOIU.supports_default_copy_to(::Optimizer, copy_names::Bool) = false
function MOI.copy_to(model::Optimizer, src::MOI.ModelLike; kws...)
    return MOIU.automatic_copy_to(model, src; kws...)
end

"""
    MOI.optimize!(model::Optimizer)
"""
MOI.optimize!(optimizer::Optimizer) = solve!(optimizer.solver)

"""
    DiscreteSet(values)
"""
struct DiscreteSet{T <: Number} <: MOI.AbstractScalarSet
    values::Vector{T}
end
DiscreteSet(values) = DiscreteSet(collect(values))
DiscreteSet(values::T...) where {T<:Number} = DiscreteSet(collect(values))

"""
    Base.copy(set::DiscreteSet) = begin

DOCSTRING
"""
Base.copy(set::DiscreteSet) = DiscreteSet(copy(set.values))

"""
    MOI.empty!(opt) = begin

DOCSTRING
"""
MOI.empty!(opt) = empty!(opt)

function MOI.is_valid(optimizer::Optimizer, index::CI{SVF, MOI.Integer})
    return index.value ∈ optimizer.int_vars
end


module SpeechAugment

using WAV: wavread
using FFTW: fft, ifft

function veclen(a::Array{T}) where T
    TYPE = typeof(a)
    if TYPE <: Array{T,1}
        return length(a)
    end
    if TYPE <: Array{T,2}
        @assert 1 in size(a)
        return length(a)
    end
end

include("./addecho.jl")
include("./addnoise.jl")
include("./distort.jl")
include("./drop.jl")
include("./farfield.jl")
include("./speed.jl")
include("./augment.jl")

end


# This file stores data structures used in imputation and phasing.
# For JLSO compression, the relevant data structures are in compress.jl

"""
Data structure for recording haplotype mosaic of one strand:
`start[i]` to `start[i+1]` has haplotype `haplotypelabel[i]`
in `window[i]`. The haplotype label is the column index of
`CW_typed[window[i]].uniqueH`. `start[end]` to`length` has
haplotype `haplotypelabel[end]` in `window[end]`.
"""
struct HaplotypeMosaic
    length::Int64
    start::Vector{Int64}
    window::Vector{Int32}
    haplotypelabel::Vector{Int32}
end
HaplotypeMosaic(len) = HaplotypeMosaic(len, Int64[], Int32[], Int32[])

function push_Mosaic!(x::HaplotypeMosaic, y::Tuple{Int64, T}) where T <: Integer
    newstart, newlabel = y[1], y[2]
    xlen = length(x.start)
    if xlen != 0
        # check if start occurs before previous start position, since searching
        # breakpoints for 2 consecutive windows can cause "overlaps"
        xstart = last(x.start)
        if newstart ≤ xstart
            deleteat!(x.start, xlen)
            deleteat!(x.haplotypelabel, xlen)
        end
    end
    push!(x.start, newstart)
    push!(x.haplotypelabel, newlabel)
end

function push_Mosaic!(x::HaplotypeMosaic, y::Tuple{Int64, T, T}) where T <: Integer
    newstart, newlabel, newwindow = y[1], y[2], y[3]
    xlen = length(x.start)
    if xlen != 0
        # check if start occurs before previous start position, since searching
        # breakpoints for 2 consecutive windows can cause "overlaps"
        xstart = last(x.start)
        if newstart ≤ xstart
            deleteat!(x.start, xlen)
            deleteat!(x.haplotypelabel, xlen)
            deleteat!(x.window, xlen)
        end
    end
    push!(x.start, newstart)
    push!(x.haplotypelabel, newlabel)
    push!(x.window, newwindow)
end

"""
data structure for recording haplotype mosaic of two strands
"""
struct HaplotypeMosaicPair
    strand1::HaplotypeMosaic
    strand2::HaplotypeMosaic
end
HaplotypeMosaicPair(len) = HaplotypeMosaicPair(HaplotypeMosaic(len),
    HaplotypeMosaic(len))

# reset function for dynamic programming phasing
function initialize!(x::Vector{Vector{Vector{Tuple{Int32, Int32}}}})
    n = length(x)
    win = length(x[1])
    @inbounds for i in 1:n, w in 1:win
        empty!(x[i][w])
    end
end

# resize function for dynamic programming phasing
function resize!(x::Vector{Vector{Vector{Tuple{Int32, Int32}}}}, windows::Int)
    n = length(x)
    @inbounds for i in 1:n
        Base.resize!(x[i], windows)
        sizehint!(x[i], windows)
    end
end


## Deprecated

## old initialization algorithms
#
# deprecated at 0.2.7
#

function randseed_initialize!{T<:AbstractFloat}(X::Matrix{T}, centers::Matrix{T})
    Base.depwarn("randseed_initialize! is deprecated. Please use initseeds! instead.",
                 :randseed_initialize!)

    K = size(centers, 2)
    iseeds = initseeds(RandSeedAlg(), X, K)
    copyseeds!(centers, X, iseeds)
end

function kmeanspp_initialize!{T<:AbstractFloat}(X::Matrix{T}, centers::Matrix{T})
    Base.depwarn("kmeanspp_initialize! is deprecated. Please use initseeds! instead.",
                 :kmeanspp_initialize!)

    K = size(centers, 2)
    iseeds = initseeds(KmppAlg(), X, K)
    copyseeds!(centers, X, iseeds)
end

function initial_medoids{T<:AbstractFloat}(costs::Matrix{T}, k::Int)
    Base.depwarn("initial_medoids is deprecated. Please use initseeds_by_costs instead.",
                 :initial_medoids)

    initseeds_by_costs(KmCentralityAlg(), costs, k)
end

function affinity_propagation(x::Matrix; opts...)
    Base.depwarn("affinity_propagation is deprecated. Please use affinityprop instead.",
                 :affinity_propagation)
    affinityprop(x; opts...)
end


module DataCreationDriver

    using CSVFiles
    using Base.Filesystem

    using CareData:
        DataSettings,
        load_all,
        create_base_datasets,
        default_data_settings
    using GlobalDecls

    export createdata

    #settings.reaching_18s_source = DFE
    #settings.num_iterations = 200
    # push!(settings.targets, "E10000030" )
    #settings.agglevel = national

    function create_data( settings :: DataSettings )
        run_data_dir = DATADIR*"/populations/$(settings.datayear)/$(settings.dataset)/"
        mkpath( run_data_dir )
        alldata = load_all( settings.datayear )
        for k in 1:settings.num_iterations
            println( "creating iteration $k")
            created = create_base_datasets(
                alldata.ofdata,
                settings );
            CSVFiles.save( run_data_dir*"/yp_data_$k.csv", created.yp_data )
            CSVFiles.save( run_data_dir*"/carer_data_$k.csv", created.carer_data )
        end
    end # create data

end # module


# test creation of RawData, FFTData, CorrData

# create RawData from SeisData
fs = 100.
nx = Int(86400 * fs)
S = SeisData(2)
S.id[1] = "N1.SS1.00.HHZ"
S.name[1] = "New Station"
S.loc[1] = GeoLoc()
S.fs[1] = 100.0
S.gain[1] = 2.0
S.resp[1] = PZResp([0.0+0.0*im 1.0+0.767*im; 0.0+0.0*im 1.0-0.767*im])
S.units[1] = ""
S.src[1] = ""
S.notes[1] = ["asdfasdf","dfhdfgh"]
S.misc[1] = Dict{String,Any}("P" => 2.0)
S.t[1] = [1 round(Int, d2u(DateTime(Date(now())))*1e6); nx 0]
S.x[1] = rand(Float32,nx)

S.id[2] = "N2.SS2.00.HHZ"
S.name[2] = "Old Station"
S.loc[2] = GeoLoc()
S.fs[2] = fs
S.gain[2] = 22.0
S.resp[2] = PZResp64(z = [0.0+0.0*im, 0.0+0.0*im], p = [1.0+1.0*im, 1.0-1.0*im])
S.units[2] = "ms/2"
S.src[2] = "file"
S.notes[2] = ["0913840183","klnelgng"]
S.misc[2] = Dict{String,Any}("S" => 6.5)
S.t[2] = [1 round(Int, d2u(DateTime(Date(now())))*1e6); nx 0]
S.x[2] = rand(Float32,nx)

# test empty allocator
@testset "Empty RawData" begin
@test !isa(try RawData() catch ex ex end, Exception)    # Passes
R = RawData()
@test isempty(R)
end

# check SeisData to RawData with similar cc_len, cc_step
@testset "SeisData to RawData cc_len == cc_step" begin
cc_len = cc_step = 1800.
R = RawData(S,cc_len,cc_step)
@test !isa(try RawData(S,cc_len,cc_step) catch ex ex end, Exception)    # Passes
@test length(R.t) == size(R.x,2)        # check number of windows
@test size(R.x,1) == R.cc_len * R.fs    # check length of window
@test R.freqmin == 1 / R.cc_len         # check freqmin
@test R.freqmax == R.fs / 2             # check freqmax
@test ndims(R.x) == 2                   # check dimensions
@test size(R.x,2) == 48                 # check number of windows
@test R == R                            # test equality
@test u2d(R.t[1]) == u2d(S.t[1][1,2] * 1e-6) # test starttimes equal
@test eltype(R.x) == eltype(S.x[1])
end

# check uncommon cc_len & cc_step
@testset "SeisData odd cc_len/cc_step" begin
cc_len = π
cc_step = 1800.
@test_throws DomainError RawData(S,cc_len,cc_step)
cc_len = 1800.
cc_step = π
@test_throws DomainError RawData(S,cc_len,cc_step)

# check huge cc_len
cc_len = 86401.
cc_step = 1800.
@test_throws DomainError RawData(S,cc_len,cc_step)
end

# check SeisChannel to RawData with similar cc_len, cc_step
@testset "SeisChannel to RawData cc_len == cc_step" begin
cc_len = cc_step = 1800.
Ch = S[1]
R = RawData(Ch,cc_len,cc_step)
@test !isa(try RawData(Ch,cc_len,cc_step) catch ex ex end, Exception)    # Passes
@test length(R.t) == size(R.x,2)        # check number of windows
@test size(R.x,1) == R.cc_len * R.fs    # check length of window
@test R.freqmin == 1 / R.cc_len         # check freqmin
@test R.freqmax == R.fs / 2             # check freqmax
@test ndims(R.x) == 2                   # check dimensions
@test size(R.x,2) == 48                 # check number of windows
@test R == R                            # test equality
@test u2d(R.t[1]) == u2d(Ch.t[1,2] * 1e-6) # test starttimes equal
@test eltype(R.x) == eltype(Ch.x)
end

# check SeisChannel to RawData with different cc_len, cc_step
@testset "SeisChannel to RawData cc_len != cc_step" begin
cc_len = 1800.
cc_step = 450.
Ch = S[1]
R = RawData(Ch,cc_len,cc_step)
@test !isa(try RawData(Ch,cc_len,cc_step) catch ex ex end, Exception)    # Passes
@test length(R.t) == size(R.x,2)        # check number of windows
@test size(R.x,1) == R.cc_len * R.fs    # check length of window
@test R.freqmin ==  1 / R.cc_len        # check freqmin
@test R.freqmax == R.fs / 2             # check freqmax
@test ndims(R.x) == 2                   # check dimensions
@test size(R.x,2) == 189                # check number of windows
@test R == R                            # test equality
@test u2d(R.t[1]) == u2d(Ch.t[1,2] * 1e-6) # test starttimes equal
@test eltype(R.x) == eltype(Ch.x)
end

# check uncommon cc_len & cc_step
@testset "SeisChannel odd cc_len/cc_step" begin
Ch = S[1]
cc_len = π
cc_step = 1800.
@test_throws DomainError RawData(Ch,cc_len,cc_step)
cc_len = 1800.
cc_step = π
@test_throws DomainError RawData(Ch,cc_len,cc_step)

# check huge cc_len
cc_len = 86401.
cc_step = 1800.
@test_throws DomainError RawData(Ch,cc_len,cc_step)
end

## test when starttime is not aligned with cc_len
# this SeisData starts at 00:16:40
@testset "SeisChannel to RawData starttime offset" begin
S.t[1] = [1 round(Int, (d2u(DateTime(Date(now()))) + 1e3) *1e6 ); nx 0]
cc_len = cc_step = 1800.
Ch = S[1]
R = RawData(S,cc_len,cc_step)
@test !isa(try RawData(Ch,cc_len,cc_step) catch ex ex end, Exception)    # Passes
@test length(R.t) == size(R.x,2)        # check number of windows
@test size(R.x,1) == R.cc_len * R.fs    # check length of window
@test R.freqmin == 1 / R.cc_len         # check freqmin
@test R.freqmax == R.fs / 2             # check freqmax
@test ndims(R.x) == 2                   # check dimensions
@test size(R.x,2) == 47                 # check number of windows
@test eltype(R.x) == eltype(Ch.x)
end

# test CPU/ GPU
@testset "RawData CPU/GPU" begin
cc_len = cc_step = 1800.
Ch = S[1]
R = RawData(S,cc_len,cc_step)
@test R == cpu(R)
@test cpu(R) == cpu(R)
@test R == cpu(gpu(R))
end

# test in
@testset "RawData appends" begin
cc_len = cc_step = 1800.
Ch = S[1]
R = RawData(S,cc_len,cc_step)
@test in(R.id,R)

# test addition
R = RawData(S[1],cc_len,cc_step)
R1 = RawData()
@test append!(R1,R) == R
@test append!(R1,R) == R + R
end

## test FFTData

# test empty allocator
@testset "FFTData creation" begin
@test !isa(try FFTData() catch ex ex end, Exception)    # Passes
F = FFTData()
@test isempty(F)

# test FFTData creation
diveven(x) = Int(ceil(x / 2)) + iseven(Int(x)) # function to check rfft size
S.t[1] = [1 round(Int, d2u(DateTime(Date(now())))*1e6); nx 0]
cc_len = cc_step = 1800.
R = RawData(S,cc_len,cc_step)
FFT = rfft(R.x,1)

@test !isa(try FFTData(R.name, R.id, R.loc, R.fs, R.gain, R.freqmin,
                       R.freqmax, R.cc_len, R.cc_step, false, "", R.resp,
                       R.misc, R.notes, R.t, FFT) catch ex ex end, Exception)    # Passes
F = FFTData(R.name, R.id, R.loc, R.fs, R.gain, R.freqmin,
                       R.freqmax, R.cc_len, R.cc_step, false, "", R.resp,
                       R.misc, R.notes, R.t, FFT)
@test length(F.t) == size(F.fft,2)              # check number of windows
@test size(F.fft,1) == diveven(F.cc_len * F.fs) # check length of window
@test F.freqmin ==  1 / F.cc_len                # check freqmin
@test F.freqmax == F.fs / 2                     # check freqmax
@test size(F.fft,2) == 48                       # check number of windows
@test ndims(F.fft) == 2                         # check dimensions
@test F == F                                    # test equality
@test u2d(F.t[1]) == u2d(S.t[1][1,2] * 1e-6)    # test starttimes equal
@test eltype(F.fft) <: Complex

# test CPU/ GPU

@test F == cpu(F)
@test cpu(F) == cpu(F)
@test F == cpu(gpu(F))

# test in
@test in(F.id,F)

# test addition
F1 = FFTData()
@test append!(F1,F) == F
@test append!(F1,F) == F + F
@test size(F1.fft,2) == 2 * size(F.fft,2)
end
## test CorrData
# test empty allocator
@testset "CorrData creation" begin
@test !isa(try CorrData() catch ex ex end, Exception)    # Passes
C = CorrData()
@test isempty(C)

# test CorrData creation
maxlag = 100.
cc_len = cc_step = 1800.
R1 = RawData(S[1],cc_len,cc_step)
R2 = RawData(S[2],cc_len,cc_step)
F1 = rfft(R1)
F2 = rfft(R2)
comp = F1.name[end] * F2.name[end]
corr = correlate(F1.fft,F2.fft,Int(F1.cc_len * F1.fs),Int(maxlag*F1.fs))

@test !isa(try CorrData(F1.name, F1.id, F1.loc, comp, false, "CC", F1.fs, F1.gain,
  F1.freqmin, F1.freqmax, F1.cc_len, F1.cc_step, F1.whitened, F1.time_norm,
  F1.resp, F1.misc, F1.notes, 0.,0., 0., maxlag, F1.t, corr) catch ex ex end, Exception)    # Passes
C = CorrData(F1.name, F1.id, F1.loc, comp, false, "CC", F1.fs, F1.gain,
  F1.freqmin, F1.freqmax, F1.cc_len, F1.cc_step, F1.whitened, F1.time_norm,
  F1.resp, F1.misc, F1.notes, 0.,0., 0., maxlag, F1.t, corr)
@test length(C.t) == size(C.corr,2)             # check number of windows
@test size(C.corr,1) == Int(2 * C.maxlag * C.fs + 1) # check length of window
@test C.freqmin ==  1 / C.cc_len                # check freqmin
@test C.freqmax == C.fs / 2                     # check freqmax
@test size(C.corr,2) == 48                      # check number of windows
@test ndims(C.corr) == 2                        # check dimensions
@test C == C                                    # test equality
@test u2d(C.t[1]) == u2d(S.t[1][1,2] * 1e-6)    # test starttimes equal
@test eltype(C.corr) <: Real

# test CPU/ GPU
@test C == cpu(C)
@test cpu(C) == cpu(C)
@test C == cpu(gpu(C))

# test in
@test in(C.id,C)

# test addition
C1 = CorrData()
@test append!(C1,C) == C
@test append!(C1,C) == C + C
@test size(C1.corr,2) == 2 * size(C.corr,2)
end

## test indexing
@testset "Indexing" begin
    cc_len = cc_step = 1800.
    R1 = RawData(S[1],cc_len,cc_step)

    # test length
    @test length(R1) == length(R1.t)

    # test firstindex
    @test firstindex(R1) == 1

    # test lastindex
    @test lastindex(R1) == length(R1.t)

    # test indexing
    ind = [1,2,4,8,16,32]
    R2 = R1[ind]
    @test length(R2) == length(ind)
    @test size(R2.x,2) == length(ind)
    @test all(R2.x .== R1.x[:,ind])
    @test all(R2.t .== R1.t[ind])

    # test indexing with collection
    ind = 1:2:10
    R2 = R1[ind]
    @test length(R2) == length(ind)
    @test size(R2.x,2) == length(ind)
    @test all(R2.x .== R1.x[:,ind])
    @test all(R2.t .== R1.t[ind])

    # test indexing with DateTime
    ind = u2d(R1.t[1]):Hour(1):u2d(R1.t[end])
    R2 = R1[ind]
    @test length(R2) == length(ind)
    @test size(R2.x,2) == length(ind)

    # test setindex
    ind = 1:2:48
    R3 = deepcopy(R1)
    setindex!(R3,R2,ind)
    # test
    @test R3[ind] == R2

    # test sorting
    R4 = sort(R1[1:2:48] + R1[2:2:48])
    @test R4 == R1
end


using Distributed
@everywhere using Pkg
@everywhere Pkg.activate(dirname(dirname(@__DIR__)))

@everywhere begin
    using CalibrateEDMF
    using CalibrateEDMF.HelperFuncs
    using CalibrateEDMF.ReferenceStats
    using CalibrateEDMF.LESUtils
    using CalibrateEDMF.ReferenceModels
    import CalibrateEDMF.ModelTypes: LES
    import CalibrateEDMF.Pipeline: get_ref_model_kwargs, get_ref_stats_kwargs
    import LinearAlgebra: dot
    import Statistics: mean
    import JSON
end

using Combinatorics
import NCDatasets
const NC = NCDatasets
using ArgParse
using Glob

@everywhere begin
    "Get index of a case, defined local to each case name"
    get_case_ind(cases, case_i) = length(cases[1:case_i][(cases .== cases[case_i])[1:case_i]])

    """
        compute_loss(t) = compute_loss(t...)
        compute_loss(value1::Number, value2::Number, case_ind::I, ens_ind::I, nt::NamedTuple) where {I <: Integer}

    Compute loss function from forward model output data.

    Arguments:
    value1      :: a parameter value (in config > grid_search > parameters > param_name_1)
    value2      :: a parameter value (in config > grid_search > parameters > param_name_2)
    case_ind    :: index of a case name (in config > reference/validation > case_name)
    ens_ind     :: index of the ensemble member
    nt          :: Named tuple that contains information that is constant across forward model data.
        Has the entries `sim_dir`, `group_name`, `config`, `sim_type`, `ref_stats`, `ref_models`, denoting
        the root path to forward model data, the parameter pair name, the config file, the simulation type
        (reference or validation), the reference statistics, and the reference models, respectively.

    """
    compute_loss(t) = compute_loss(t...)
    function compute_loss(value1::Number, value2::Number, case_ind::I, ens_ind::I, nt::NamedTuple) where {I <: Integer}
        # fetch from config
        get_from_config(x) = nt.config[nt.sim_type][x][case_ind]
        (case_name, loss_names, y_dir) = get_from_config.(["case_name", "y_names", "y_dir"])
        y_ref_type = nt.config[nt.sim_type]["y_reference_type"]
        # Get the case index for this case (if there are duplicate cases)
        cases = nt.config[nt.sim_type]["case_name"]
        case_j = get_case_ind(cases, case_ind)

        # path to .nc simulation data
        param_path = joinpath(nt.sim_dir, nt.group_name, "$(value1)_$(value2)")
        output_dir = joinpath(param_path, "$case_name.$case_j/Output.$case_name.$ens_ind")
        scm_file = joinpath(output_dir, "stats/Stats.$case_name.nc")
        if !isfile(scm_file)
            @warn("No NetCDF file found on path $scm_file")
            return NaN  # if case data does not exist, return NaN
        end
        # compute loss
        y_loss_names = if (y_ref_type isa LES)
            get_les_names(loss_names, get_stats_path(y_dir))
        else
            loss_names
        end

        RS = nt.ref_stats
        m = nt.ref_models[case_ind]

        # Check that the simulation interval fully covers the averaging interval (i.e. simulation completed)
        t = nc_fetch(scm_file, "t")
        ti, tf = get_t_start(m), get_t_end(m)
        dt = (length(t) > 1) * mean(diff(t))
        if (t[end] < ti) || (t[end] < tf - dt)
            # check is simulation completed to t_max in the namelist
            namelist_path = joinpath(output_dir, "namelist_$case_name.in")
            namelist = open(namelist_path, "r") do io
                JSON.parse(io; dicttype = Dict, inttype = Int64)
            end
            t_max = namelist["time_stepping"]["t_max"]
            if t[end] < t_max
                @warn "Simulation did not finish during grid search deleting output directory: \n $output_dir"
                rm(output_dir, force = true, recursive = true)
            end

            @warn string(
                "The requested averaging interval: ($ti s, $tf s) is not in the simulation interval ($(t[1]) s, $(t[end]) s), ",
                "which means the simulation stopped before reaching the requested t_end. Returning NaN. ",
                "\n Simulation file: $scm_file",
            )
            return NaN
        end

        # case PCA and covariance matrix
        pca_vec = RS.pca_vec[case_ind]'
        pca_case_inds = pca_inds(RS, case_ind)
        Γ = RS.Γ[pca_case_inds, pca_case_inds]

        # Reference data
        y_full_case = get_profile(m, get_stats_path(y_dir), y_loss_names, z_scm = get_z_obs(m))
        y_norm = normalize_profile(y_full_case, num_vars(m), RS.norm_vec[case_ind])

        # SCM data
        g_full_case_i = get_profile(m, scm_file, z_scm = get_z_obs(m))
        g_norm = normalize_profile(g_full_case_i, num_vars(m), RS.norm_vec[case_ind])

        # PCA
        yg_diff_pca = pca_vec * (y_norm - g_norm)
        # loss
        sim_loss = dot(yg_diff_pca, Γ \ yg_diff_pca)
        return sim_loss
    end
end  # end @everywhere begin

"""
    compute_loss_map(config::Dict, sim_dir::AbstractString)

Compute loss map for a set of simulations and store the output in a netcdf4 file

Parameters:
config  :: Dictionary of model calibration and grid search configuration options
sim_dir :: Directory of grid search simulation output
"""
function compute_loss_map(config::Dict, sim_dir::AbstractString)

    # construct reference models and reference statistics
    sim_type = get_entry(config["grid_search"], "sim_type", "reference")
    @assert sim_type in ("reference", "validation")
    ref_config = config[sim_type]  # or validation
    kwargs_ref_model = get_ref_model_kwargs(ref_config)
    reg_config = config["regularization"]
    kwargs_ref_stats = get_ref_stats_kwargs(ref_config, reg_config)

    n_ens = get_entry(config["grid_search"], "ensemble_size", nothing)
    cases = get_entry(config[sim_type], "case_name", nothing)
    case_names_unique = ["$case.$(get_case_ind(cases, i))" for (i, case) in enumerate(cases)]
    n_cases = length(cases)

    parameters = get_entry(config["grid_search"], "parameters", nothing)
    param_names = collect(keys(parameters))
    NC.Dataset(joinpath(sim_dir, "loss.nc"), "c") do ds
        for (param_name_1, param_name_2) in combinations(param_names, 2)
            group_name = "$param_name_1.$param_name_2"

            # iterate all folders, e.g. "0.1_0.2", ...
            # compute loss, store in matrix:
            # (p1, p2, case, ens_i)
            value1 = parameters[param_name_1]
            value2 = parameters[param_name_2]
            n_value1 = length(value1)
            n_value2 = length(value2)

            loss_2D_sec = zeros((n_value1, n_value2, n_cases, n_ens))

            # Construct `ReferenceModel`s and `ReferenceStatistics`
            ref_models = construct_reference_models(kwargs_ref_model)
            ref_stats = ReferenceStatistics(ref_models; kwargs_ref_stats...)

            # configurations for `compute_loss`
            nt = (; sim_dir, group_name, config, sim_type, ref_stats, ref_models)
            loss_configs = vec(collect(Iterators.product(value1, value2, 1:n_cases, 1:n_ens, [nt])))

            # compute loss
            sim_loss = pmap(compute_loss, loss_configs)
            loss_2D_sec = reshape(sim_loss, size(loss_2D_sec))

            # save output
            NC.defGroup(ds, group_name, attrib = [])
            ensemble_member = 1:n_ens
            group_root = ds.group[group_name]  # group is e.g. 'sorting_power.entrainment_factor'
            # Define dimensions: param_name_1, param_name_2, case, ensemble_member
            NC.defDim(group_root, "param_name_1", n_value1)
            NC.defDim(group_root, "param_name_2", n_value2)
            NC.defDim(group_root, "case", n_cases)
            NC.defDim(group_root, "ensemble_member", n_ens)
            # Define variables:
            ncvar = NC.defVar(group_root, "case", case_names_unique, ("case",))
            ncvar[:] = case_names_unique  # "Bomex", "DYCOMS"
            ncvar = NC.defVar(group_root, "ensemble_member", ensemble_member, ("ensemble_member",))
            ncvar[:] = ensemble_member  # 1, 2, ..., n_ens
            # param_name_1 and param_name_2 values
            ncvar = NC.defVar(group_root, param_name_1, value1, ("param_name_1",))
            ncvar[:] = value1
            ncvar = NC.defVar(group_root, param_name_2, value2, ("param_name_2",))
            ncvar[:] = value2
            ncvar = NC.defVar(
                group_root,
                "loss_data",
                loss_2D_sec,
                ("param_name_1", "param_name_2", "case", "ensemble_member"),
            )
            ncvar[:, :, :, :] = loss_2D_sec
        end  # end for param combinations loop
    end  # NC.Dataset do-block
end

function parse_commandline_lm()

    s = ArgParseSettings(; description = "Run data path input")

    @add_arg_table s begin
        "--sim_dir"
        help = "Grid search simulations directory"
        arg_type = String
    end

    return ArgParse.parse_args(s)  # parse_args(ARGS, s)
end

if abspath(PROGRAM_FILE) == @__FILE__
    args = parse_commandline_lm()
    sim_dir = args["sim_dir"]
    include(joinpath(sim_dir, "config.jl"))
    config = get_config()
    compute_loss_map(config, sim_dir)
end


struct BlockBasisCoo <: Field end

function evaluate_field!(cache,b::BlockBasisCoo,x)
  msg =
  """
  Not implemented! We not support iteration of vectors of block bases before
  evaluation. This is unlikely to be be needed in the future.
  """
  @notimplemented msg
end

struct VectorOfBlockBasisCoo <: AbstractVector{BlockBasisCoo}
  blocks::Tuple
  blockids
  axes
end

Base.size(v::VectorOfBlockBasisCoo) = (length(first(v.blocks)) ,)
Base.getindex(v::VectorOfBlockBasisCoo,i::Integer) = BlockBasisCoo()

function evaluate_field_array(v::VectorOfBlockBasisCoo,x::AbstractArray)
  blocks = evaluate_field_arrays(v.blocks,x)
  axs = apply(_new_axes,x,v.axes)
  blockids = map(i->(1,i...),v.blockids)
  VectorOfBlockArrayCoo(blocks,blockids,axs)
end

function field_array_gradient(v::VectorOfBlockBasisCoo)
  blocks = map(field_array_gradient,v.blocks)
  VectorOfBlockBasisCoo(blocks,v.blockids,v.axes)
end

function reindex(v::VectorOfBlockBasisCoo,j_to_i::AbstractArray)
  blocks = map(b->reindex(b,j_to_i),v.blocks)
  axs = reindex(v.axes,j_to_i)
  VectorOfBlockBasisCoo(blocks,v.blockids,axs)
end

function compose_field_arrays(v::VectorOfBlockBasisCoo,f)
  blocks = map(b->compose_field_arrays(b,f),v.blocks)
  VectorOfBlockBasisCoo(blocks,v.blockids,v.axes)
end

function insert_array_of_bases_in_block(i::Integer,a,ax1,ax2...)
  blocks = (a,)
  blockids = _compute_blockids(eltype(ax1),i)
  axs = create_array_of_blocked_axes(ax1,ax2...)
  VectorOfBlockBasisCoo(blocks,blockids,axs)
end

#TODO perhaps this should go into Arrays/BlockArraysCoo.jl
function create_array_of_blocked_axes(axs...)
  axs = apply(_cat_axes,axs...)
end

function _compute_blockids(::Type{<:NTuple{1}},i)
  [(i,)]
end

function _compute_blockids(::Type{<:NTuple{2}},i)
  [(1,i)]
end

function _cat_axes(a::NTuple{1},b::NTuple{1})
  (_blockedrange([a[1],b[1]]),)
end

function _cat_axes(a::NTuple{1}...)
  (_blockedrange([map(i->i[1],a)...]),)
end

function _cat_axes(a::NTuple{2},b::NTuple{2})
  @assert length(a[1]) == 1
  @assert length(b[1]) == 1
  ran = (_blockedrange([a[2],b[2]]),)
  _add_singleton_block(ran)
end

function _cat_axes(a::NTuple{2}...)
  ran = (_blockedrange([map(i->i[2],a)...]),)
  _add_singleton_block(ran)
end

function _new_axes(x,ran::NTuple{N,<:BlockedUnitRange} where N)
  np = length(x)
  (blockedrange([np]),ran...)
end

function _new_axes(x,ran::NTuple{N,<:TwoLevelBlockedUnitRange} where N)
  np = length(x)
  r = blockedrange([np])
  (blockedrange([r]),ran...)
end

function _add_singleton_block(ran::NTuple{N,<:Base.OneTo} where N)
  (Base.OneTo(1),ran...)
end

function _add_singleton_block(ran::NTuple{N,<:BlockedUnitRange} where N)
  (blockedrange([1]),ran...)
end

function _add_singleton_block(ran::NTuple{N,<:TwoLevelBlockedUnitRange} where N)
  r = blockedrange([1])
  (blockedrange([r]),ran...)
end

function _blockedrange(a::Vector{<:Base.OneTo})
  blockedrange(map(length,a))
end

function _blockedrange(a::Vector{<:BlockedUnitRange})
  blockedrange(a)
end


module YaoTargetQASM

using MLStyle
using OpenQASM
using OpenQASM.Types
using OpenQASM.Tools
using YaoCompiler
using YaoLocations
using Expronicon
using Configurations
using CompilerPluginTools
using RBNF: Token
using YaoCompiler.Intrinsics
using YaoCompiler: YaoCompileTarget
using Core: CodeInfo
using YaoLocations: plain
using OpenQASM.Types: UGate

export OpenQASMTarget, RegInfo, code_qasm, QASM, @qasm_str

@option struct OpenQASMTarget <: YaoCompileTarget
    version::VersionNumber = v"2.0"
    toplevel::Bool = true
    inline_intrinsics::Bool = false
    include_qelib1::Bool = false
end

module QASM

Base.@pure @noinline sin(x) = Base.Math.sin(x)
Base.@pure @noinline cos(x) = Base.Math.cos(x)
Base.@pure @noinline tan(x) = Base.Math.tan(x)
Base.@pure @noinline exp(x) = Base.Math.exp(x)
Base.@pure @noinline log(x) = Base.Math.log(x)
Base.@pure @noinline sqrt(x) = Base.Math.sqrt(x)

end

function code_qasm(op::Type{<:Operation}; kw...)
    target = OpenQASMTarget(;kw...)
    return YaoCompiler.compile(target, Intrinsics.apply, Tuple{AnyReg, op})
end

include("validate.jl")
include("intrinsic.jl")
include("reginfo.jl")
include("codegen.jl")
include("frontend.jl")

end


using MeshCat, MeshCatMechanisms, RigidBodyDynamics

urdf_path = joinpath(pwd(), "models/a1_quadruped/temp/a1.urdf")
a1 = MeshCatMechanisms.parse_urdf(urdf_path, remove_fixed_tree_joints = true)
a1_visuals = MeshCatMechanisms.URDFVisuals(joinpath(pwd(), "models/a1_quadruped/temp/a1.urdf"))

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

vis = Visualizer()
mvis = MechanismVisualizer(a1, a1_visuals, vis[:base])
render(vis)


# MadNLP.jl
# Created by Sungho Shin (sungho.shin@wisc.edu)

module MadNLPMa77

import ..MadNLPHSL:
    @kwdef, Logger, @debug, @warn, @error, libhsl,
    SparseMatrixCSC, SparseMatrixCSC, SubVector, StrideOneVector,
    get_tril_to_full, transfer!,
    AbstractOptions, AbstractLinearSolver, set_options!,
    SymbolicException,FactorizationException,SolveException,InertiaException,
    introduce, factorize!, solve!, improve!, is_inertia, inertia
import ..MadNLPHSL: Mc68

const INPUT_MATRIX_TYPE = :csc

@enum(Ordering::Int,AMD = 1, METIS = 3)

@kwdef mutable struct Options <: AbstractOptions
    ma77_buffer_lpage::Int = 4096
    ma77_buffer_npage::Int = 1600
    ma77_file_size::Int = 2097152
    ma77_maxstore::Int = 0
    ma77_nemin::Int = 8
    ma77_order::Ordering = METIS
    ma77_print_level::Int = -1
    ma77_small::Float64 = 1e-20
    ma77_static::Float64 = 0.
    ma77_u::Float64 = 1e-8
    ma77_umax::Float64 = 1e-4
end

@kwdef mutable struct Control
    f_arrays::Cint = 0
    print_level::Cint = 0
    unit_diagnostics::Cint = 0
    unit_error::Cint = 0
    unit_warning::Cint = 0
    bits::Cint = 0

    buffer_lpage_1::Cint = 0
    buffer_lpage_2::Cint = 0
    buffer_npage_1::Cint = 0
    buffer_npage_2::Cint = 0

    file_size::Clong = 0
    maxstore::Clong = 0

    storage_1::Clong = 0
    storage_2::Clong = 0
    storage_3::Clong = 0

    nemin::Cint = 0
    maxit::Cint = 0
    infnorm::Cint = 0
    thresh::Cdouble = 0.
    nb54::Cint = 0
    action::Cint = 0
    multiplier::Cdouble = 0.
    nb64::Cint = 0
    nbi::Cint = 0
    small::Cdouble = 0.
    static_::Cdouble = 0.
    storage_indef::Clong = 0
    u::Cdouble = 0.
    umin::Cdouble = 0.
    consist_tol::Cdouble = 0.

    ispare_1::Cint = 0
    ispare_2::Cint = 0
    ispare_3::Cint = 0
    ispare_4::Cint = 0
    ispare_5::Cint = 0

    lspare_1::Clong = 0
    lspare_2::Clong = 0
    lspare_3::Clong = 0
    lspare_4::Clong = 0
    lspare_5::Clong = 0

    rspare_1::Cdouble = 0.
    rspare_2::Cdouble = 0.
    rspare_3::Cdouble = 0.
    rspare_4::Cdouble = 0.
    rspare_5::Cdouble = 0.
end

@kwdef mutable struct Info
    detlog::Cdouble = 0.
    detsign::Cint = 0
    flag::Cint = 0
    iostat::Cint = 0
    matrix_dup::Cint = 0
    matrix_rank::Cint = 0
    matrix_outrange::Cint = 0
    maxdepth::Cint = 0
    maxfront::Cint = 0
    minstore::Clong = 0
    ndelay::Cint = 0
    nfactor::Clong = 0
    nflops::Clong = 0
    niter::Cint = 0
    nsup::Cint = 0
    num_neg::Cint = 0
    num_nothresh::Cint = 0
    num_perturbed::Cint = 0
    ntwo::Cint = 0

    stat_1::Cint = 0
    stat_2::Cint = 0
    stat_3::Cint = 0
    stat_4::Cint = 0

    nio_read_1::Clong = 0 # 2
    nio_read_2::Clong = 0 # 2

    nio_write_1::Clong = 0 # 2
    nio_write_2::Clong = 0 # 2

    nwd_read_1::Clong = 0 # 2
    nwd_read_2::Clong = 0 # 2

    nwd_write_1::Clong = 0 # 2
    nwd_write_2::Clong = 0 # 2

    num_file_1::Cint = 0 # 4
    num_file_2::Cint = 0 # 4
    num_file_3::Cint = 0 # 4
    num_file_4::Cint = 0 # 4

    storage_1::Clong = 0 # 4
    storage_2::Clong = 0 # 4
    storage_3::Clong = 0 # 4
    storage_4::Clong = 0 # 4

    tree_nodes::Cint = 0
    unit_restart::Cint = 0
    unused::Cint = 0
    usmall::Cdouble = 0.

    ispare_1::Cint = 0
    ispare_2::Cint = 0
    ispare_3::Cint = 0
    ispare_4::Cint = 0
    ispare_5::Cint = 0

    lspare_1::Clong = 0
    lspare_2::Clong = 0
    lspare_3::Clong = 0
    lspare_4::Clong = 0
    lspare_5::Clong = 0

    rspare_1::Cdouble = 0.
    rspare_2::Cdouble = 0.
    rspare_3::Cdouble = 0.
    rspare_4::Cdouble = 0.
    rspare_5::Cdouble = 0.
end

mutable struct Solver <: AbstractLinearSolver
    tril::SparseMatrixCSC{Float64,Int32}
    full::SparseMatrixCSC{Float64,Int32}
    tril_to_full_view::SubVector{Float64}

    control::Control
    info::Info

    mc68_control::Mc68.Control
    mc68_info::Mc68.Info

    order::Vector{Int32}
    keep::Vector{Ptr{Nothing}}

    opt::Options
    logger::Logger
end

ma77_default_control_d(control::Control) = ccall(
    (:ma77_default_control_d,libhsl),
    Cvoid,
    (Ref{Control},),
    control)
ma77_open_d(n::Cint,fname1::String,fname2::String,fname3::String,fname4::String,
            keep::Vector{Ptr{Cvoid}},control::Control,info::Info) = ccall(
                (:ma77_open_d,libhsl),
                Cvoid,
                (Cint,Ptr{Cchar},Ptr{Cchar},Ptr{Cchar},Ptr{Cchar},
                 Ptr{Ptr{Cvoid}},Ref{Control},Ref{Info}),
                n,fname1,fname2,fname3,fname4,keep,control,info)
ma77_input_vars_d(idx::Cint,nvar::Cint,list::StrideOneVector{Cint},
                  keep::Vector{Ptr{Cvoid}},control::Control,info::Info) = ccall(
                      (:ma77_input_vars_d,libhsl),
                      Cvoid,
                      (Cint,Cint,Ptr{Cint},
                       Ptr{Ptr{Cvoid}},Ref{Control},Ref{Info}),
                      idx,nvar,list,keep,control,info)
ma77_input_reals_d(idx::Cint,length::Cint,reals::StrideOneVector{Cdouble},
                   keep::Vector{Ptr{Cvoid}},control::Control,info::Info) = ccall(
                       (:ma77_input_reals_d,libhsl),
                       Cvoid,
                       (Cint,Cint,Ptr{Cdouble},
                        Ptr{Ptr{Cvoid}},Ref{Control},Ref{Info}),
                       idx,length,reals,keep,control,info)
ma77_analyse_d(order::Vector{Cint},
               keep::Vector{Ptr{Cvoid}},control::Control,info::Info) = ccall(
                   (:ma77_analyse_d,libhsl),
                   Cvoid,
                   (Ptr{Cint},Ptr{Ptr{Cvoid}},Ref{Control},Ref{Info}),
                   order,keep,control,info)
ma77_factor_d(posdef::Cint,keep::Vector{Ptr{Cvoid}},control::Control,info::Info,
              scale::Ptr{Nothing}) = ccall(
                  (:ma77_factor_d,libhsl),
                  Cvoid,
                  (Cint,Ptr{Ptr{Cvoid}},Ref{Control},Ref{Info},Ptr{Nothing}),
                  posdef,keep,control,info,scale)
ma77_factor_solve_d(posdef::Cint,keep::Vector{Ptr{Cvoid}},control::Control,info::Info,
                    scale::Vector{Cdouble},nrhs::Cint,lx::Cint,rhs::Vector{Cdouble}) = ccall(
                        (:ma77_factor_solve_d,libhsl),
                        Cvoid,
                        (Cint,Ptr{Ptr{Cvoid}},Ref{Control},Ref{Info},
                         Ptr{Cdouble},Cint,Cint,Ptr{Cdouble}),
                        posdef,keep,control,info,scale,nrhs,lx,rhs);
ma77_solve_d(job::Cint,nrhs::Cint,lx::Cint,x::Vector{Cdouble},
             keep::Vector{Ptr{Cvoid}},control::Control,info::Info,
             scale::Ptr{Nothing})=ccall(
                 (:ma77_solve_d,libhsl),
                 Cvoid,
                 (Cint,Cint,Cint,Ptr{Float64},
                  Ptr{Ptr{Cvoid}},Ref{Control},Ref{Info},Ptr{Nothing}),
                 job,nrhs,lx,x,keep,control,info,scale);
ma77_finalize_d(keep::Vector{Ptr{Cvoid}},control::Control,info::Info) = ccall(
    (:ma77_finalise_d,libhsl),
    Cvoid,
    (Ptr{Ptr{Cvoid}},Ref{Control},Ref{Info}),
    keep,control,info)

function Solver(csc::SparseMatrixCSC{Float64,Int32};
                option_dict::Dict{Symbol,Any}=Dict{Symbol,Any}(),
                opt=Options(),logger=Logger(),
                kwargs...)

    set_options!(opt,option_dict,kwargs)

    full,tril_to_full_view = get_tril_to_full(csc)
    order = Vector{Int32}(undef,csc.n)

    mc68_info = Mc68.Info()
    mc68_control = Mc68.get_mc68_default_control()

    keep = [C_NULL]

    mc68_control.f_array_in=1
    mc68_control.f_array_out=1
    Mc68.mc68_order_i(Int32(opt.ma77_order),Int32(csc.n),csc.colptr,csc.rowval,order,mc68_control,mc68_info)

    info=Info()
    control=Control()
    ma77_default_control_d(control)
    control.f_arrays = 1
    control.bits = 32
    control.file_size = opt.ma77_file_size
    control.maxstore = opt.ma77_maxstore
    control.print_level = -1

    control.buffer_lpage_1=opt.ma77_buffer_lpage
    control.buffer_lpage_2=opt.ma77_buffer_lpage
    control.buffer_npage_1=opt.ma77_buffer_npage
    control.buffer_npage_2=opt.ma77_buffer_npage

    control.nemin = opt.ma77_nemin
    control.small = opt.ma77_small
    control.static_ = opt.ma77_static
    control.u = opt.ma77_u

    isfile(".ma77_int")   && rm(".ma77_int")
    isfile(".ma77_real")  && rm(".ma77_real")
    isfile(".ma77_work")  && rm(".ma77_work")
    isfile(".ma77_delay") && rm(".ma77_delay")

    ma77_open_d(Int32(full.n),".ma77_int", ".ma77_real", ".ma77_work", ".ma77_delay",
                keep,control,info)

    info.flag < 0 && throw(SymbolicException())

    for i=1:full.n
        ma77_input_vars_d(Int32(i),full.colptr[i+1]-full.colptr[i],
                          view(full.rowval,full.colptr[i]:full.colptr[i+1]-1),
                          keep,control,info);
        info.flag < 0 && throw(LinearSymbolicException())
    end

    ma77_analyse_d(order,keep,control,info)
    info.flag<0 && throw(SymbolicException())

    M = Solver(csc,full,tril_to_full_view,
               control,info,mc68_control,mc68_info,order,keep,opt,logger)
    finalizer(finalize,M)
    return M
end

function factorize!(M::Solver)
    M.full.nzval.=M.tril_to_full_view
    for i=1:M.full.n
        ma77_input_reals_d(Int32(i),M.full.colptr[i+1]-M.full.colptr[i],
                           view(M.full.nzval,M.full.colptr[i]:M.full.colptr[i+1]-1),
                           M.keep,M.control,M.info)
        M.info.flag < 0 && throw(FactorizationException())
    end
    ma77_factor_d(Int32(0),M.keep,M.control,M.info,C_NULL)
    M.info.flag < 0 && throw(FactorizationException())
    return M
end
function solve!(M::Solver,rhs::StrideOneVector{Float64})
    ma77_solve_d(Int32(0),Int32(1),Int32(M.full.n),rhs,M.keep,M.control,M.info,C_NULL);
    M.info.flag < 0 && throw(SolveException())
    return rhs
end

is_inertia(::Solver) = true
function inertia(M::Solver)
    return (M.info.matrix_rank-M.info.num_neg,M.full.n-M.info.matrix_rank,M.info.num_neg)
end

finalize(M::Solver) = ma77_finalize_d(M.keep,M.control,M.info)

function improve!(M::Solver)
    if M.control.u == M.opt.ma77_umax
        @debug(M.logger,"improve quality failed.")
        return false
    end
    M.control.u = min(M.opt.ma77_umax,M.control.u^.75)
    @debug(M.logger,"improved quality: pivtol = $(M.control.u)")
    return true
end

introduce(::Solver)="ma77"

end