#
# separable_methods.jl
#
# Provides the state-of-the-art methods Double Ratio
# and Double Difference and separable inverse approach.
#
# State of the arts methods have been implemented from:
# [Tinbergen, 2005]  Tinbergen,  J.  (2005).Astronomical Polarimetry.
# Cambridge  University
#
# Inverse approch :
# [Denneulin, 2020] Approche inverse pour la reconstruction
# des environnements circumstellaires en polarimétrie
# avec l’instrument d’imagerie directe ESO/VLT SPHERE IRDIS..
# Instrumentation et méthodes pour l'astrophysique [astro-ph.IM].
# Université Claude Bernard Lyon 1 (UCBL), 2020. Français. ⟨tel-03200282⟩
#
#
#------------------------------------------------------
#
# This file is part of Rhapsodie
#
# Copyright (c) 2017-2021 Laurence DENNEULIN (see Licence.md)
#TODO : Add a way to return the CRLB

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

"""
    Double_Difference(d,ind) -> X
    Double_Difference(d) -> X #if indices loaded with

where X is:
    - a PolarimetricMap if d is of size (N1,N2,K,2)
    - a PolarimetricPixel if d is of size (K,2).
"""
function Double_Difference(data::Array{Float64,4},ind::Array{Int64,2})
        n1,n2,n3,n4= size(data);
        S=Array{PolarimetricPixel}(undef,n1,n2);
        @inbounds for i2 in 1:n2
            @simd for i1 in 1:n1
                S[i1,i2]=Double_Difference(data[i1,i2,:,:], ind);
            end
        end
        return PolarimetricMap(S)
end

function Double_Difference(data::Array{Float64,2},ind::Array{Int64,2})
		Q=mean((data[ind[1,:],1] .-data[ind[1,:],2]
		          .-(data[ind[2,:],1] .-data[ind[2,:],2]))/2);
		U=mean((data[ind[3,:],1] .-data[ind[3,:],2]
		          .-(data[ind[4,:],1] .-data[ind[4,:],2]))/2);
		Iq=(data[ind[1,:],1] .+data[ind[1,:],2]
		    .+data[ind[2,:],1] .+data[ind[2,:],2])/2;
		Iu=(data[ind[3,:],1] .+data[ind[3,:],2]
		    .+data[ind[4,:],1] .+data[ind[4,:],2])/2;

        return PolarimetricPixel("stokes",
                                 (mean(Iq)+ mean(Iu))/2,
                                 Q,
                                 U)
end

    Double_Difference(data::Array{Float64,4}) = (Double_Difference(data,
                                                                   get_par().indices))

    Double_Difference(data::Array{Float64,2}) = (Double_Difference(data,
                                                                   get_par().indices))

#
#-----------------------------------------------------
#
"""
    Double_Ratio(d,ind) -> X
    Double_Ratio(d) -> X #if indices loaded with

where X is:
    - a PolarimetricMap if d is of size (N1,N2,K,2)
    - a PolarimetricPixel if d is of size (K,2).

"""
function Double_Ratio(data::Array{Float64,4},ind::Array{Int64,2})
        n1,n2,n3,n4= size(data);
        S=Array{PolarimetricPixel}(undef,n1,n2);
        @inbounds for i2 in 1:n2
            @simd for i1 in 1:n1
                S[i1,i2]=Double_Ratio(data[i1,i2,:,:], ind);
            end
        end
        return PolarimetricMap(S)
end

function Double_Ratio(data::Array{Float64,2},ind::Array{Int64,2})
		Rq=(data[ind[1,:],1] ./data[ind[1,:],2]) ./(data[ind[2,:],1] ./data[ind[2,:],2]);
		Ru=(data[ind[3,:],1] ./data[ind[3,:],2]) ./(data[ind[4,:],1] ./data[ind[4,:],2]);
		for i=1:length(Rq)
			Rq[i]=sqrt(Rq[i]*(Rq[i]>0));
			Ru[i]=sqrt(Ru[i]*(Ru[i]>0));
		end
		pq=(Rq .-1)./(Rq .+1);
		pu=(Ru .-1)./(Ru .+1);
		Iq=(data[ind[1,:],1] .+data[ind[1,:],2] .+data[ind[2,:],1] .+data[ind[2,:],2])/2;
		Iu=(data[ind[3,:],1] .+data[ind[3,:],2] .+data[ind[4,:],1] .+data[ind[4,:],2])/2;

        return PolarimetricPixel("stokes",
                                 (mean(Iq)+ mean(Iu))/2,
                                 mean(pq .*Iq),
                                 mean(pu .*Iu))

end

Double_Ratio(data::Array{Float64,4}) = (Double_Ratio(data, get_par().indices))

Double_Ratio(data::Array{Float64,2}) = (Double_Ratio(data, get_par().indices))

#
#-----------------------------------------------------
#
"""
    Linear_Method(d,w) -> X
    Linear_Method(d) -> X #i.e. W=Id

where X is:
    - a PolarimetricMap if d is of size (N1,N2,K,2)
    - a PolarimetricPixel if d is of size (K,2).

"""
function Linear_Method(data::Array{Float64,4}, weight::Array{Float64,4})
        n1,n2,n3,n4= size(data);
        x=Array{PolarimetricPixel}(undef,n1,n2);
        @inbounds for i2 in 1:n2
            @simd for i1 in 1:n1
                if weight[i1,i2] !=0
                    x[i1,i2]=Linear_Method(data[i1,i2,:,:],weight[i1,i2,:,:]);
                else
                    x[i1,i2] = PolarimetricPixel("stokes",0.0, 0.0, 0.0);
                end
            end
        end
        return PolarimetricMap(x)
end

function Linear_Method(data::Array{Float64,2}, weight::Array{Float64,2})
    T = Float64 # floating point type used for computations

    # Local variables to integrate the normal equations and other
    # quantities.  Only the lower triangular part of the left-hand-side
    # matrix is needed.
    K = get_par().dataset_length # index range for 2nd dimension of `ind`
    @assert size(data)[1] == K;
    @assert size(weight)[1] == K;

    local A11::T = 0
    local A21::T = 0
    local A22::T = 0
    local A31::T = 0
    local A32::T = 0
    local A33::T = 0
    local b1::T = 0.0
    local b2::T = 0.0
    local b3::T = 0.0
    local W::T, Wd::T
    for j in 1:2
        W = 0.0
        Wd = 0.0
        for k=1:K
            v1 = get_par().v[k][j][1]
            v2 = get_par().v[k][j][2]
            v3 = get_par().v[k][j][3]
            w, d = weight[k,j], data[k,j]
            W = w
            Wd = w*d

        # Integrate the upper-triangular part of A.
        A11 += v1*W*v1
        A21 += v2*W*v1
        A22 += v2*W*v2
        A31 += v3*W*v1
        A32 += v3*W*v2
        A33 += v3*W*v3

        # Integrate the right-hand side vector b.
        b1 += v1*Wd
        b2 += v2*Wd
        b3 += v3*Wd
        end
    end

    # Solve the normal equations using fast StaticArrays.
    A = SMatrix{3,3,T}((A11, A21, A31,
                        A21, A22, A32,
                        A31, A32, A33))
    b = SVector{3,T}((b1,b2,b3))

    #CLRB=inv(A);
    x = A\b
    #x=CLRB*b;
    return PolarimetricPixel("stokes", x)#, CLRB
end

Linear_Method(data::Array{Float64,4}) = (Linear_Method(data, ones(size(data))))

Linear_Method(data::Array{Float64,2}) = (Linear_Method(data,ones(size(data))))

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

"""
    NonLinear_Method(d,w) -> X
    NonLinear_Method(d) -> X #i.e. W=Id

where X is:
    - a PolarimetricMap if d is of size (N1,N2,K,2)
    - a PolarimetricPixel if d is of size (K,2).

"""

function NonLinear_Method(data::Array{Float64,4}, weight::Array{Float64,4})
        n1,n2,n3,n4= size(data);
        x=Array{PolarimetricPixel}(undef,n1,n2);
        @inbounds for i2 in 1:n2
            @simd for i1 in 1:n1
                if weight[i1,i2] !=0
                    x[i1,i2]=NonLinear_Method(data[i1,i2,:,:],weight[i1,i2,:,:]);
                 else
                    x[i1,i2] = PolarimetricPixel("intensities",0.0, 0.0, 0.0);
                end
            end
        end
        return PolarimetricMap(x)
end

function NonLinear_Method(data::A, weight::A) where {T<:AbstractFloat, A<:AbstractArray{T,2}}
	cost(Θ)=opti_nonlinear(Θ,data,weight)[1]
	(Θ_opt,c)=Bradi.minimize(cost,range(-pi/2,stop=pi/2,length=7))
	f,x=opti_nonlinear(Θ_opt,data,weight)
	return x
end

function opti_nonlinear(Θ::T,data::H, weight::H) where {T<:AbstractFloat, H<:AbstractArray{T,2}}
    K = get_par().dataset_length # index range for 2nd dimension of `ind`
    @assert size(data)[1] == K;
    @assert size(weight)[1] == K;

	local s1 =0.0;
	local s2 =0.0;
	local s3 =0.0;
	local b1 =0.0;
	local b2 =0.0;
	local W::T, Wd::T

	for j=1:2
	    W = 0.0
	    for k=1:K
	        v1 = get_par().v[k][j][1]
            v2 = get_par().v[k][j][2]
            v3 = get_par().v[k][j][3]

            V1 = v1;
            V2 = (v1 + v2*cos(2*Θ) + v3*cos(2*Θ));

            w, d = weight[k,j], data[k,j]
            W = w
            Wd = w*d

	        s1 += V1*W*V1;
            s2 += V1*W*V2;
     		s3 += V2*W*V2;
     		b1 += V1*Wd;
            b2 += V2*Wd;
       end
	end

    A = SMatrix{2,2,T}((s1,s2,s2,s3))
    b = SVector{2,T}((b1,b2))

	Iu, Ip = A\b;

	if (Iu>=0) && (Ip>=0)
	    x=PolarimetricPixel("intensities",Iu,Ip,Θ);
		return f_nonlinear(x,data,weight), x
	elseif (Iu<0) || (Ip<0)
		AIu=A*[1;0];
		Iu0=(AIu\b)[];
		if Iu>0
		    x1=PolarimetricPixel("intensities",Iu,0.0,Θ);
		else
	        x1=PolarimetricPixel("intensities",0.0,0.0,Θ);
	    end

		AIp=A*[0;1];
		Ip=(AIp\b)[];
		if Ip > 0
            x2=PolarimetricPixel("intensities",0.0,Ip,Θ);
	    else
	        x2=PolarimetricPixel("intensities",0.0,0.0,Θ);
	    end

        x3 = PolarimetricPixel("intensities",0.0,0.0,Θ);

		C1=f_nonlinear(x1,data,weight);
		C2=f_nonlinear(x2,data,weight);
		C3=f_nonlinear(x3,data,weight);

        if (C1<=C2)&&(C1<=C3)
            return C1,x1
        elseif (C2<=C1)&&(C2<=C3)
            return C2,x2
        elseif (C3<=C1)&&(C3<=C2)
            return C3,x3
        else
            error("can't find a good configuration")
	    end
	end
end

function f_nonlinear(x::PolarimetricPixel{T}, data::A, weight::A) where {T<:AbstractFloat, A<:AbstractArray{T,2}}
	local f::T =0.0;
	local W::T, Wd::T

	K = get_par().dataset_length # index range
    @assert size(data)[1] == K;
    @assert size(weight)[1] == K;

	for j=1:2
	    W=0.0;
	    Wd =0.0;
	    for k=1:K
            w, d = weight[k,j], data[k,j]
            W = w
            Wd = w*d
	        v1 = get_par().v[k][j][1]
            v2 = get_par().v[k][j][2]
            v3 = get_par().v[k][j][3]

	        m = v1*x.Iu + (v1 + v2*cos(2*x.θ) + v3*cos(2*x.θ))*x.Ip
	        f += w *(m - d)^2
	    end
	end

	return f
end

NonLinear_Method(data::Array{Float64,4}) = (NonLinear_Method(data, ones(size(data))))

NonLinear_Method(data::Array{Float64,2}) = (NonLinear_Method(data,ones(size(data))))

#=
function FDCR_nonlin(RHO,THETA,W,ind)
    (v,v_norm)=V_calc(alpha, psi)
    C=[cos(THETA); sin(THETA)]
    a11=0;
    a22=0;
    a33=0;
    a12=0;
    a13=0;
    a23=0;
    for i=1:4
        for j=1:2
	        prodsca=abs(dot(v[:,(i-1)*2+j],C))^2
	        dprodsca=cos(THETA)*sin(THETA)*(abs(v[2,(i-1)*2+j])^2-abs(v[1,(i-1)*2+j])^2)+cos(2*THETA)*2*real(v[1,(i-1)*2+j]*conj(v[2,(i-1)*2+j]))
	        a11=a11+sum(W[ind[i,:],j]*v_norm[(i-1)*2+j]^2);
	        a22=a22+sum(W[ind[i,:],j]*prodsca^2);
	        a33=a33+sum(W[ind[i,:],j]*dprodsca^2);
	        a12=a12+sum(W[ind[i,:],j]*v_norm[(i-1)*2+j]*prodsca);
	        a13=a13+sum(W[ind[i,:],j]*v_norm[(i-1)*2+j]*dprodsca);
	        a23=a23+sum(W[ind[i,:],j]*prodsca*dprodsca);
	    end
    end
    I=[a11 a12 RHO*a13; a12 a22 RHO*a23; RHO*a13 RHO*a23 (RHO^2)*a33];
    if det(I) !=0
        F=inv(I)
        return F
    else
    return zeros(size(I));
    end
end
=#


"""
    CrossingPhase
Crossing phases for SidewalkPedestrianModel.

For a crossing pedestrian, phases go: -2, -1, 0, 1
For a non-crossing pedestrian that pauses at the crosswalk, phases go: -2, -1, 1
For a non-crossing pedestrian that does not pause at the crosswalk, phases go: -2, 1

Model based on Feliciani et al (2017) - A simulation model for non-signalized pedestrian
crosswalks based on evidence from on field observation.
"""
struct CrossingPhase
    phase::Int # -2, -1, 0, 1
end

const APPROACHING = -2 # Keep walking at constant speed
const APPRAISING = -1 # Decelerate until reach curbside
const CROSSING = 0 # Accelerate until desired speed
const LEAVING = 1 # Keep walking at constant speed

"""
    SidewalkPedestrianModel

Walks along the sidewalk until approaching the crosswalk. Waits for the cars to pass, then crosses.

# Fields
- `timestep::Float64`
- `phase::Int = APPROACHING`
- `ttc_threshold::Float64 = clamp(rand(Normal(4.0, 2.5)), 1.0, Inf)`
- `crosswalk::Lane = Lane()`
- `sw_origin::Lane = Lane()`
- `sw_dest::Lane = Lane()`
- `a::PedestrianLatLonAccel = PedestrianLatLonAccel(0.0, 0.0, sw_origin)` makes you turn, left/right
- `σ::Float64 = NaN` optional stdev on top of the model, set to zero or NaN for deterministic
- `v_des_approach::Float64 = clamp(rand(Normal(1.28, 0.18)), 0.0, Inf)` Based on Feliciani et al. results
- `v_des_appraise::Float64 = clamp(rand(Normal(0.94, 0.21)), 0.0, Inf)`
- `v_des_cross::Float64 = clamp(rand(Normal(1.35, 0.18)), 0.0, Inf)`
- `ped_accel::Float64 = 0.30`
- `ped_decel::Float64 = -0.50`
"""
@with_kw mutable struct SidewalkPedestrianModel <: DriverModel{PedestrianLatLonAccel}
    timestep::Float64
    phase::Int = APPROACHING

    ttc_threshold::Float64 = clamp(rand(Normal(4.0, 2.5)), 1.0, Inf)
    crosswalk::Lane = Lane()
    sw_origin::Lane = Lane()
    sw_dest::Lane = Lane()
    # conflict_lanes::Vector{Lane} = Lane[]

    # Values from Feliciani
    a::PedestrianLatLonAccel = PedestrianLatLonAccel(0.0, 0.0, sw_origin) # makes you turn, left/right
    σ::Float64 = NaN # optional stdev on top of the model, set to zero or NaN for deterministic behavior
    v_des_approach::Float64 = clamp(rand(Normal(1.28, 0.18)), 0.0, Inf) # Based on Feliciani et al. results
    v_des_appraise::Float64 = clamp(rand(Normal(0.94, 0.21)), 0.0, Inf)
    v_des_cross::Float64 = clamp(rand(Normal(1.35, 0.18)), 0.0, Inf)

    ped_accel::Float64 = 0.30
    ped_decel::Float64 = -0.50

    # For graphing
    pos_x = Float64[]
    pos_y = Float64[]
    velocities = Float64[]
    phases = Int[]
end

Base.rand(rng::AbstractRNG, model::SidewalkPedestrianModel) = model.a

function AutomotiveDrivingModels.observe!(model::SidewalkPedestrianModel, scene::Scene{Entity{VehicleState, D, I}}, roadway::Roadway, egoid::I) where {D, I}
    ped = scene[findfirst(egoid, scene)]

    push!(model.pos_x, posg(ped.state).x)
    push!(model.pos_y, posg(ped.state).y)
    push!(model.velocities, vel(ped.state))
    push!(model.phases, model.phase)

    if model.phase == APPROACHING
        update_approaching(ped, model, roadway, model.crosswalk, model.sw_origin)
    elseif model.phase == APPRAISING
        update_appraising(ped, scene, model, roadway, model.crosswalk, model.sw_origin)
    elseif model.phase == CROSSING
        update_crossing(ped, scene, model, roadway, model.crosswalk)
    elseif model.phase == LEAVING
        update_leaving(ped, scene, model, roadway, model.sw_dest)
    else
        # Model phase is out of bounds
    end
end

function update_approaching(ped::Entity{S, D, I},
                            model::SidewalkPedestrianModel,
                            roadway::Roadway,
                            crosswalk::Lane,
                            sw_origin::Lane
                            ) where {S, D, I}
    # Keep walking at constant speed until close to crosswalk
    # If waiting to cross, go to appraising. Else, go to leaving phase
    Δs = get_distance_to_lane(ped, crosswalk, roadway)
    if Δs < 3.0
        model.phase = APPRAISING
    end
end

function update_appraising(ped::Entity{S, D, I},
                           scene::Scene,
                           model::SidewalkPedestrianModel,
                           roadway::Roadway,
                           crosswalk::Lane,
                           sw_origin::Lane,
                           ) where {S, D<:AbstractAgentDefinition, I}

    Δs = get_distance_to_lane(ped, crosswalk, roadway)
    if Δs < 0.1
        set_inst_vel(ped, model, sw_origin, 0.0)

        for veh in scene
            if class(veh.def) == AgentClass.PEDESTRIAN
                continue
            end

            # Check vehicle distance to the crosswalk
            posG = posg(veh.state)
            cw_posF = Frenet(posG, model.crosswalk, roadway)
            lane = get_lane(roadway, veh)
            Δs = get_distance_to_lane(veh, model.crosswalk, roadway)
            ttc = Δs/vel(veh.state)

            if 0.0 < ttc < model.ttc_threshold
                return # Exit if any car is not within threshold
            end
        end
        model.phase = CROSSING
    else
        # Decelerate to appraise speed
        set_accel(ped, model, model.v_des_appraise, sw_origin)
    end
end

function update_crossing(ped::Entity{S, D, I},
                         scene::Scene,
                         model::SidewalkPedestrianModel,
                         roadway::Roadway,
                         crosswalk::Lane
                         ) where {S, D, I}
    # Accelerate until desired speed
    set_accel(ped, model, model.v_des_cross, crosswalk)
    # If ped is at end of crosswalk, transition to leaving phase.
    Δs = get_distance_to_lane(ped, model.sw_dest, roadway)

    if Δs < 0.2
        set_inst_vel(ped, model, model.crosswalk, 0.0)
        model.phase = LEAVING
    end
end

function update_leaving(ped::Entity{S, D, I},
                        scene::Scene,
                        model::SidewalkPedestrianModel,
                        roadway::Roadway,
                        sw_dest::Lane
                        ) where {S, D, I}
    # Keep walking at constant speed
    set_accel(ped, model, model.v_des_approach, model.sw_dest)
end

function set_inst_vel(ped::Entity{VehicleState, D, I}, model::SidewalkPedestrianModel, lane_des::Lane, v_des::Float64) where {D, I}
    model.a = PedestrianLatLonAccel(0.0, calc_a(vel(ped.state), v_des, model.timestep), lane_des)
end

function calc_a(v_init::Float64, v::Float64, Δt::Float64)
    return (v - v_init)/Δt
end

function set_accel(ped::Entity{VehicleState, D, I}, model::SidewalkPedestrianModel, v_des::Float64, lane_des::Lane) where {D, I}
    if vel(ped.state) < v_des
        model.a = PedestrianLatLonAccel(0.0, model.ped_accel, lane_des)
    elseif vel(ped.state) > v_des
        model.a = PedestrianLatLonAccel(0.0, model.ped_decel, lane_des)
    else
        model.a = PedestrianLatLonAccel(0.0, 0.0, lane_des)
    end
end

# Edited from AutomotivePOMDPs.jl - projects collision point onto the lane.
# Should work in any orientation.
function get_distance_to_lane(veh::Entity{VehicleState, D, I}, lane_des::Lane, roadway::Roadway) where {D, I}
    vehlane = get_lane(roadway, veh)
    lane_segment = LineSegment(lane_des.curve[1].pos, lane_des.curve[end].pos)
    collision_point = intersect(Ray(posg(veh.state)), lane_segment)
    lane_des_proj = Frenet(VecSE2(collision_point), vehlane, roadway)
    Δs = lane_des_proj.s - posf(veh.state).s
    return Δs
end


function CompleteObservationFunction(PreObservationFunction::Function)
    numargs = MaximalNumberOfArguments(PreObservationFunction)
    if numargs == 1
        return (u::AbstractArray{<:Number}, t::Real, θ::AbstractVector{<:Number}) -> PreObservationFunction(u)
    elseif numargs == 2
        return (u::AbstractArray{<:Number}, t::Real, θ::AbstractVector{<:Number}) -> PreObservationFunction(u, t)
    elseif numargs == 3
        return (u::AbstractArray{<:Number}, t::Real, θ::AbstractVector{<:Number}) -> PreObservationFunction(u, t, θ)
    else
        throw("Given ObservationFunction should accept either arguments (u) or (u,t) or (u,t,θ). Got function which accepts $numargs arguments.")
    end
end

"""
    InformNames(DS::AbstractDataSet, sys::ODESystem, observables::AbstractVector{<:Int})
Copy the state names saved in `ODESystem` to `DS`.
"""
function InformNames(DS::AbstractDataSet, sys::ModelingToolkit.AbstractSystem, observables::Union{BoolArray,AbstractVector{<:Int}})
    newxnames = xnames(DS) == CreateSymbolNames(xdim(DS),"x") ? [string(ModelingToolkit.get_iv(sys))] : xnames(DS)
    newynames = ynames(DS) == CreateSymbolNames(ydim(DS),"y") ? string.(ModelingToolkit.get_states(sys)[observables]) : ynames(DS)
    InformNames(DS, newxnames, newynames)
end

# No ObservationFunction, therefore try to use sys to infer state names of ODEsys
# Extend for other DEFunctions in the future
function DataModel(DS::AbstractDataSet, sys::Union{ModelingToolkit.AbstractSystem,SciMLBase.AbstractDiffEqFunction}, u0::Union{AbstractArray{<:Number},Function},
                        observables::Union{AbstractVector{<:Int},BoolArray,Function}=collect(1:length(u0)), args...; tol::Real=1e-7, Domain::Union{HyperCube,Nothing}=nothing, kwargs...)
    newDS = (typeof(observables) <: AbstractVector{<:Int} && sys isa ModelingToolkit.AbstractSystem) ? InformNames(DS, sys, observables) : DS
    DataModel(newDS, GetModel(sys, u0, observables; tol=tol, Domain=Domain, kwargs...), args...)
end

# """
# Given `func` is converted to `ODEFunction`. Will probably deprecate this in the future.
# """
# function GetModel(func::Function, u0::Union{AbstractArray{<:Number},Function}, observables::Union{Function,AbstractVector{<:Int},BoolArray}=collect(1:length(u0)); tol::Real=1e-7,
#                     meth::OrdinaryDiffEqAlgorithm=GetMethod(tol), Domain::Union{HyperCube,Nothing}=nothing, inplace::Bool=true)
#     GetModel(ODEFunction{inplace}(func), u0, observables; tol=tol, Domain=Domain, meth=meth, inplace=inplace)
# end

function GetModel(sys::ModelingToolkit.AbstractSystem, u0::Union{AbstractArray{<:Number},Function}, observables::Union{AbstractVector{<:Int},BoolArray,Function}=collect(1:length(u0));
                Domain::Union{HyperCube,Nothing}=nothing, inplace::Bool=true, kwargs...)
    # Is there some optimization that can be applied here? Modellingtoolkitize(sys) or something?
    # sys = Sys isa Catalyst.ReactionSystem ? convert(ODESystem, Sys) : Sys
    Model = if sys isa ModelingToolkit.AbstractODESystem
        GetModel(ODEFunction{inplace}(sys), u0, observables; Domain=Domain, inplace=inplace, kwargs...)
    else
        throw("Not programmed for $(typeof(sys)) yet, please convert to a ModelingToolkit.AbstractODESystem first.")
    end
    if Model isa ModelMap       Model = Model.Map    end
    pnames = ModelingToolkit.get_ps(sys) .|> string
    ylen = if observables isa Function      # ObservationFunction
        # Might still fail if states u are a Matrix.
        argnum = MaximalNumberOfArguments(observables)
        F = if argnum==1  z->observables(z)   elseif argnum==2  z->observables(z,0.1)
            elseif argnum==3 z->observables(z,0.1,GetStartP(length(pnames)))    else throw("Error") end
        num = GetArgLength(F)
        length(F(ones(num)))
    else
        observables isa BoolArray ? sum(observables) : length(observables)
    end
    plen = if Domain isa HyperCube
        length(Domain)
    elseif u0 isa AbstractArray     # Vector / Matrix
        # initial conditions given as array means the parameters are only the ps in sys
        length(pnames)
    else        # SplitterFunction
        # May well fail depending on how splitter function is implemented
        GetArgLength(u0)
    end
    xyp = (1, ylen, plen)
    Domain = isa(Domain, Bool) ? FullDomain(xyp[3], 1e5) : Domain

    pnames = plen - length(pnames) > 0 ? vcat(CreateSymbolNames(plen - length(pnames), "u"), pnames) : pnames
    # new(Map, InDomain, Domain, xyp, pnames, StaticOutput, inplace, CustomEmbedding)
    ModelMap(Model, nothing, Domain, xyp, pnames, Val(false), Val(false), Val(true))
end

# Although specialized methods for constant initial condition and specification of observed components in terms of arrays is faster than using functions, also use general method here.
function GetModel(func::SciMLBase.AbstractDiffEqFunction{T}, Initial, Observables; kwargs...) where T
    SplitterFunction = if Initial isa Function
        Initial
    else
        u0 = PromoteStatic(Initial, T)
        θ -> (u0, θ)
    end
    GetModel(func, SplitterFunction, (Observables isa Function ? Observables : u -> u[Observables]); kwargs...)
end

function GetModel(func::SciMLBase.AbstractDiffEqFunction, SplitterFunction::Function, ObservationFunction::Function; kwargs...)
    throw("If you see this error, it is most likely because no specialized method has been implemented for $(typeof(func)) yet.")
end

# Vanilla version with constant array of initial conditions and vector of observables.
function GetModel(func::AbstractODEFunction{T}, u0::AbstractArray{<:Number}, observables::Union{AbstractVector{<:Int},BoolArray}=collect(1:length(u0)); tol::Real=1e-7,
                    meth::OrdinaryDiffEqAlgorithm=GetMethod(tol), Domain::Union{HyperCube,Nothing}=nothing, inplace::Bool=true) where T
    @assert T == inplace
    u0 = PromoteStatic(u0, inplace)

    function GetSol(θ::AbstractVector{<:Number}, u0::AbstractArray{<:Number}; tol::Real=tol, max_t::Number=10., meth::OrdinaryDiffEqAlgorithm=meth, kwargs...)
        odeprob = ODEProblem(func, u0, (0., max_t), θ)
        solve(odeprob, meth; reltol=tol, abstol=tol, kwargs...)
    end
    function ODEmodel(t::Number, θ::AbstractVector{<:Number}; observables::Union{AbstractVector{<:Int},BoolArray}=observables, u0::AbstractArray{<:Number}=u0,
                                                        tol::Real=tol, max_t::Number=t, meth::OrdinaryDiffEqAlgorithm=meth, FullSol::Bool=false, kwargs...)
        FullSol && return GetSol(θ, u0; tol=tol, max_t=max_t, meth=meth, kwargs...)
        sol = GetSol(θ, u0; tol=tol, max_t=max_t, meth=meth, save_everystep=false, save_start=false, save_end=true, kwargs...)
        sol.u[end][observables]
    end
    function ODEmodel(ts::AbstractVector{<:Number}, θ::AbstractVector{<:Number}; observables::Union{AbstractVector{<:Int},BoolArray}=observables, u0::AbstractArray{<:Number}=u0,
                                                tol::Real=tol, max_t::Number=maximum(ts), meth::OrdinaryDiffEqAlgorithm=meth, FullSol::Bool=false, kwargs...)
        FullSol && return GetSol(θ, u0; tol=tol, max_t=max_t, meth=meth, tstops=ts, kwargs...)
        sol = GetSol(θ, u0; tol=tol, max_t=max_t, meth=meth, saveat=ts, kwargs...)
        length(sol.u) != length(ts) && throw("ODE integration failed, maybe try using a lower tolerance value. θ=$θ.")
        [sol.u[i][observables] for i in 1:length(ts)] |> Reduction
    end
    MakeCustom(ODEmodel, Domain)
end

"""
    GetModel(func::ODEFunction, u0::AbstractArray, ObservationFunction::Function; tol::Real=1e-7, meth::OrdinaryDiffEqAlgorithm=Tsit5(), Domain::Union{HyperCube,Nothing}=nothing, inplace::Bool=true)
Returns a `ModelMap` which evolves the given system of ODEs from the initial configuration `u0` and afterwards applies the `ObservationFunction` to produce its predictions.

`ObservationFunction` should either be of the form `F(u) -> Vector` or `F(u,t) -> Vector` or `F(u,t,θ) -> Vector`.
Internally, the `ObservationFunction` is automatically wrapped as `F(u,t,θ)` if it is not already defined to accept three arguments.

A `Domain` can be supplied to constrain the parameters of the model to particular ranges which can be helpful in the fitting process.
"""
function GetModel(func::AbstractODEFunction{T}, u0::AbstractArray{<:Number}, PreObservationFunction::Function; tol::Real=1e-7,
                    meth::OrdinaryDiffEqAlgorithm=GetMethod(tol), Domain::Union{HyperCube,Nothing}=nothing, inplace::Bool=true) where T
    @assert T == inplace
    u0 = PromoteStatic(u0, inplace)
    ObservationFunction = CompleteObservationFunction(PreObservationFunction)

    function GetSol(θ::AbstractVector{<:Number}, u0::AbstractArray{<:Number}; tol::Real=tol, max_t::Number=10., meth::OrdinaryDiffEqAlgorithm=meth, kwargs...)
        odeprob = ODEProblem(func, u0, (0., max_t), θ)
        solve(odeprob, meth; reltol=tol, abstol=tol, kwargs...)
    end
    function ODEmodel(t::Number, θ::AbstractVector{<:Number}; ObservationFunction::Function=ObservationFunction, u0::AbstractArray{<:Number}=u0,
                                                tol::Real=tol, max_t::Number=t, meth::OrdinaryDiffEqAlgorithm=meth, FullSol::Bool=false, kwargs...)
        FullSol && return GetSol(θ, u0; tol=tol, max_t=max_t, meth=meth, kwargs...)
        sol = GetSol(θ, u0; tol=tol, max_t=max_t, meth=meth, save_everystep=false, save_start=false, save_end=true, kwargs...)
        ObservationFunction(sol.u[end], t, θ)
    end
    function ODEmodel(ts::AbstractVector{<:Number}, θ::AbstractVector{<:Number}; ObservationFunction::Function=ObservationFunction, u0::AbstractArray{<:Number}=u0,
                                            tol::Real=tol, max_t::Number=maximum(ts), meth::OrdinaryDiffEqAlgorithm=meth, FullSol::Bool=false, kwargs...)
        FullSol && return GetSol(θ, u0; tol=tol, max_t=max_t, meth=meth, tstops=ts, kwargs...)
        sol = GetSol(θ, u0; tol=tol, max_t=max_t, meth=meth, saveat=ts, kwargs...)
        length(sol.u) != length(ts) && throw("ODE integration failed, maybe try using a lower tolerance value. θ=$θ.")
        [ObservationFunction(sol.u[i], sol.t[i], θ) for i in 1:length(ts)] |> Reduction
    end
    MakeCustom(ODEmodel, Domain)
end

"""
    GetModel(func::ODEFunction, SplitterFunction::Function, observables::Union{AbstractVector{<:Int},BoolArray}=collect(1:length(u0)); tol::Real=1e-7, meth::OrdinaryDiffEqAlgorithm=Tsit5(), Domain::Union{HyperCube,Nothing}=nothing, inplace::Bool=true)
Returns a `ModelMap` which evolves the given system of ODEs and returns `u[observables]` to produce its predictions.
Here, the initial conditions for the ODEs are produced from the parameters `θ` using the `SplitterFunction` which for instance allows one to estimate them from data.

`SplitterFunction` should be of the form `F(θ) -> (u0, p)`, i.e. the output is a `Tuple` whose first entry is the initial condition for the ODE model and the second entry constitutes the parameters which go on to enter the `ODEFunction`.
Typically, a fair bit of performance can be gained from ensuring that `SplitterFunction` outputs the initial condition `u0` as type `MVector` or `MArray`, if it has less than ~100 components.

A `Domain` can be supplied to constrain the parameters of the model to particular ranges which can be helpful in the fitting process.
"""
function GetModel(func::AbstractODEFunction{T}, SplitterFunction::Function, observables::Union{AbstractVector{<:Int},BoolArray}=collect(1:length(u0)); tol::Real=1e-7,
                    meth::OrdinaryDiffEqAlgorithm=GetMethod(tol), Domain::Union{HyperCube,Nothing}=nothing, inplace::Bool=true) where T
    @assert T == inplace

    function GetSol(θ::AbstractVector{<:Number}, SplitterFunction::Function; tol::Real=tol, max_t::Number=10., meth::OrdinaryDiffEqAlgorithm=meth, kwargs...)
        u0, p = SplitterFunction(θ);        odeprob = ODEProblem(func, u0, (0., max_t), p)
        solve(odeprob, meth; reltol=tol, abstol=tol, kwargs...)
    end
    function ODEmodel(t::Number, θ::AbstractVector{<:Number}; observables::Union{AbstractVector{<:Int},BoolArray}=observables, SplitterFunction::Function=SplitterFunction,
                                tol::Real=tol, max_t::Number=t, meth::OrdinaryDiffEqAlgorithm=meth, FullSol::Bool=false, kwargs...)
        FullSol && return GetSol(θ, SplitterFunction; tol=tol, max_t=max_t, meth=meth, kwargs...)
        sol = GetSol(θ, SplitterFunction; tol=tol, max_t=max_t, meth=meth, save_everystep=false, save_start=false, save_end=true, kwargs...)
        sol.u[end][observables]
    end
    function ODEmodel(ts::AbstractVector{<:Number}, θ::AbstractVector{<:Number}; observables::Union{AbstractVector{<:Int},BoolArray}=observables, SplitterFunction::Function=SplitterFunction,
                                tol::Real=tol, max_t::Number=maximum(ts), meth::OrdinaryDiffEqAlgorithm=meth, FullSol::Bool=false, kwargs...)
        FullSol && return GetSol(θ, SplitterFunction; tol=tol, max_t=max_t, meth=meth, tstops=ts, kwargs...)
        sol = GetSol(θ, SplitterFunction; tol=tol, max_t=max_t, meth=meth, saveat=ts, kwargs...)
        length(sol.u) != length(ts) && throw("ODE integration failed, maybe try using a lower tolerance value. θ=$θ.")
        [sol.u[i][observables] for i in 1:length(ts)] |> Reduction
    end
    MakeCustom(ODEmodel, Domain)
end

"""
    GetModel(func::AbstractODEFunction{T}, SplitterFunction::Function, PreObservationFunction::Function; tol::Real=1e-7, meth::OrdinaryDiffEqAlgorithm=Tsit5(), Domain::Union{HyperCube,Nothing}=nothing, inplace::Bool=true)
Returns a `ModelMap` which evolves the given system of ODEs and afterwards applies the `ObservationFunction` to produce its predictions.
Here, the initial conditions for the ODEs are produced from the parameters `θ` using the `SplitterFunction` which for instance allows one to estimate them from data.

`SplitterFunction` should be of the form `F(θ) -> (u0, p)`, i.e. the output is a `Tuple` whose first entry is the initial condition for the ODE model and the second entry constitutes the parameters which go on to enter the `ODEFunction`.
Typically, a fair bit of additional performance can be gained from ensuring that `SplitterFunction` outputs the initial condition `u0` as type `MVector` or `MArray`, if it has less than ~100 components.

`ObservationFunction` should either be of the form `F(u) -> Vector` or `F(u,t) -> Vector` or `F(u,t,θ) -> Vector`.
Internally, the `ObservationFunction` is automatically wrapped as `F(u,t,θ)` if it is not already defined to accept three arguments.

!!! note
    The vector `θ` passed to `ObservationFunction` is the same `θ` that is passed to `SplitterFunction`, i.e. before splitting.
    This is because `ObservationFunction` might also depend on the initial conditions in general.

A `Domain` can be supplied to constrain the parameters of the model to particular ranges which can be helpful in the fitting process.
"""
function GetModel(func::AbstractODEFunction{T}, SplitterFunction::Function, PreObservationFunction::Function; tol::Real=1e-7,
                    meth::OrdinaryDiffEqAlgorithm=GetMethod(tol), Domain::Union{HyperCube,Nothing}=nothing, inplace::Bool=true) where T
    @assert T == inplace
    ObservationFunction = CompleteObservationFunction(PreObservationFunction)

    function GetSol(θ::AbstractVector{<:Number}, SplitterFunction::Function; tol::Real=tol, max_t::Number=10., meth::OrdinaryDiffEqAlgorithm=meth, kwargs...)
        u0, p = SplitterFunction(θ);        odeprob = ODEProblem(func, u0, (0., max_t), p)
        solve(odeprob, meth; reltol=tol, abstol=tol, kwargs...)
    end
    function ODEmodel(t::Number, θ::AbstractVector{<:Number}; ObservationFunction::Function=ObservationFunction, SplitterFunction::Function=SplitterFunction,
                                                    tol::Real=tol, max_t::Number=t, meth::OrdinaryDiffEqAlgorithm=meth, FullSol::Bool=false, kwargs...)
        FullSol && return GetSol(θ, SplitterFunction; tol=tol, max_t=max_t, meth=meth, kwargs...)
        sol = GetSol(θ, SplitterFunction; tol=tol, max_t=max_t, meth=meth, save_everystep=false, save_start=false, save_end=true, kwargs...)
        ObservationFunction(sol.u[end], sol.t[end], θ)
    end
    function ODEmodel(ts::AbstractVector{<:Number}, θ::AbstractVector{<:Number}; ObservationFunction::Function=ObservationFunction, SplitterFunction::Function=SplitterFunction,
                                            tol::Real=tol, max_t::Number=maximum(ts), meth::OrdinaryDiffEqAlgorithm=meth, FullSol::Bool=false, kwargs...)
        FullSol && return GetSol(θ, SplitterFunction; tol=tol, max_t=max_t, meth=meth, tstops=ts, kwargs...)
        sol = GetSol(θ, SplitterFunction; tol=tol, max_t=max_t, meth=meth, saveat=ts, kwargs...)
        length(sol.u) != length(ts) && throw("ODE integration failed, maybe try using a lower tolerance value. θ=$θ.")
        [ObservationFunction(sol.u[i], sol.t[i], θ) for i in 1:length(ts)] |> Reduction
    end
    MakeCustom(ODEmodel, Domain)
end

"""
    ModifyODEmodel(DM::AbstractDataModel, NewObservationFunc::Function) -> ModelMap
Constructs a new `ModelMap` with new observation function `f(u,t,θ)` from a given ODE-based `DataModel`.
"""
ModifyODEmodel(DM::AbstractDataModel, NewObservationFunc::Function) = ModifyODEmodel(DM, Predictor(DM), NewObservationFunc)
function ModifyODEmodel(DM::AbstractDataModel, Model::ModelMap, NewObservationFunc::Function)
    # Model.Map isa DEModel && return ModifyDEModel(Model, NewObservationFunc)
    Eval = try
        Model(WoundX(DM)[1], MLE(DM); FullSol=true).u[1]
    catch;
        throw("It appears as though the given model is not an ODEmodel.")
    end
    F = CompleteObservationFunction(NewObservationFunc)
    out = F(Eval, WoundX(DM)[1], MLE(DM))
    function NewODEmodel(x::AbstractVector{<:Number}, θ::AbstractVector{<:Number}; FullSol=false, kwargs...)
        FullSol && return Model.Map(x, θ; FullSol=true, kwargs...)
        sol = Model.Map(x, θ; FullSol=true, saveat=x, kwargs...)
        length(sol.u) != length(x) && throw("ODE integration failed, maybe try using a lower tolerance value. θ=$θ.")
        [F(sol.u[i], sol.t[i], θ) for i in 1:length(x)] |> Reduction
    end
    function NewODEmodel(x::Number, θ::AbstractVector{<:Number}; FullSol=false, kwargs...)
        FullSol && return Model.Map(x, θ; FullSol=true, kwargs...)
        sol = Model.Map(x, θ; FullSol=true, save_everystep=false, save_start=false, save_end=true, kwargs...)
        F(sol.u[end], sol.t[end], θ)
    end
    ModelMap(NewODEmodel, Model.InDomain, Model.Domain, (Model.xyp[1], length(out), Model.xyp[3]), Model.pnames, Val(out isa SVector), Model.inplace, Model.CustomEmbedding)
end


module CrossValidators

using Statistics: mean, std

using MLBase: StratifiedKfold, LOOCV,
              RandomSub, StratifiedRandomSub

# standard included modules
using Random
using AutoMLPipeline
using AutoMLPipeline.Utils
using DataFrames: DataFrame

export crossvalidate, holdoutcrossvalidate

function holdoutcrossvalidate(pl::Machine,X::DataFrame,Y::Vector,metric::Function,ntimes=10)
  @assert size(X)[1] == length(Y)
end

"""
    crossvalidate(pl::Machine,X::DataFrame,Y::Vector,pfunc::Function,kfolds=10)

Run K-fold crossvalidation where:
- `pfunc` is a performance metric
- `X` and `Y` are input and target
"""
function crossvalidate(pl::Machine,X::DataFrame,Y::Vector,
		       pfunc::Function,nfolds::Int=10,verbose::Bool=true)
  ## flatten arrays
  @assert size(X)[1] == length(Y)
  rowsize = size(X)[1]
  folds = StratifiedKfold(Y,nfolds) |> collect
  pacc = Float64[]
  fold = 0
  error = 0
  for trndx in folds
    ppl = deepcopy(pl)
    rlist = collect(1:rowsize) # nasty bug if not collected and placed inside setdiff
    tsndx = setdiff(rlist,trndx)
    trX = X[trndx,:]
    trY = Y[trndx] |> collect
    tstX = X[tsndx,:]
    tstY = Y[tsndx] |> collect
    res = 0.0
    try
      res = pipe_accuracy(ppl,pfunc,trX,trY,tstX,tstY)
      push!(pacc,res)
      fold += 1
      if verbose == true
	println("fold: ",fold,", ",res)
      end
    catch e
      #println(e)
      error += 1
    end
  end
  if verbose == true
    println("errors: ",error)
  end
  (mean=mean(pacc),std=std(pacc),folds=nfolds,errors=error)
end

function pipe_accuracy(plearner,perf::Function,trX::DataFrame,
		       trY::Vector,tstX::DataFrame,tstY::Vector)
  learner = deepcopy(plearner)
  fit!(learner,trX,trY)
  pred = transform!(learner,tstX)
  res = perf(pred,tstY)
  return res
end

end


include("vpg.jl")
include("A2C.jl")
include("ppo.jl")
include("A2CGAE.jl")
include("MAC.jl")
include("ddpg.jl")
include("td3.jl")
include("sac.jl")


Base.@kwdef mutable struct NFFTParams{T}
  m::Int = 4
  σ::T = 2.0
  reltol::T = 1e-9
  window::Symbol = :kaiser_bessel
  LUTSize::Int64 = 0
  precompute::PrecomputeFlags = POLYNOMIAL
  sortNodes::Bool = false
  storeApodizationIdx::Bool = false
  blocking::Bool = true
end

mutable struct NFFTPlan{T,D,R} <: AbstractNFFTPlan{T,D,R}
    N::NTuple{D,Int64}
    NOut::NTuple{R,Int64}
    M::Int64
    x::Matrix{T}
    n::NTuple{D,Int64}
    dims::UnitRange{Int64}
    params::NFFTParams{T}
    forwardFFT::FFTW.cFFTWPlan{Complex{T},-1,true,D,UnitRange{Int64}}
    backwardFFT::FFTW.cFFTWPlan{Complex{T},1,true,D,UnitRange{Int64}}
    tmpVec::Array{Complex{T},D}
    tmpVecHat::Array{Complex{T},D}
    # Caches for deconvolve
    deconvolveIdx::Array{Int64,1}
    windowHatInvLUT::Vector{Vector{T}}
    # Cache for precompute = LUT
    windowLinInterp::Vector{T}
    # Cache for precompute = POLYNOMIAL
    windowPolyInterp::Matrix{T}
    # Caches for blocking
    blocks::Array{Array{Complex{T},D},D}
    nodesInBlock::Array{Vector{Int64},D}
    blockOffsets::Array{NTuple{D,Int64},D}
    idxInBlock::Array{Matrix{Tuple{Int,T}},D}
    windowTensor::Array{Array{T,3},D}
    # Cache for precompute = FULL
    B::SparseMatrixCSC{T,Int64}
end

function Base.copy(p::NFFTPlan{T,D,R}) where {T,D,R}
    tmpVec = similar(p.tmpVec)
    tmpVecHat = similar(p.tmpVecHat)
    deconvolveIdx = copy(p.deconvolveIdx)
    windowLinInterp = copy(p.windowLinInterp)
    windowPolyInterp = copy(p.windowPolyInterp)
    windowHatInvLUT = copy(p.windowHatInvLUT)
    B = copy(p.B)
    blocks = deepcopy(p.blocks)
    nodesInBlock = deepcopy(p.nodesInBlock)
    blockOffsets = copy(p.blockOffsets)
    idxInBlock = copy(p.idxInBlock)
    windowTensor = copy(p.windowTensor)
    x = p.x

    FP = plan_fft!(tmpVec, p.dims; flags = p.forwardFFT.flags)
    BP = plan_bfft!(tmpVec, p.dims; flags = p.backwardFFT.flags)

    return NFFTPlan{T,D,R}(p.N, p.NOut, p.M, x, p.n, p.dims, p.params, FP, BP, tmpVec,
                           tmpVecHat, deconvolveIdx, windowHatInvLUT, windowLinInterp, windowPolyInterp,
                           blocks, nodesInBlock, blockOffsets, idxInBlock, windowTensor, B)
end

################
# constructor
################

function NFFTPlan(x::Matrix{T}, N::NTuple{D,Int}; dims::Union{Integer,UnitRange{Int64}}=1:D,
                 fftflags=nothing, kwargs...) where {T,D}

    checkNodes(x)

    params, N, NOut, M, n, dims_ = initParams(x, N, dims; kwargs...)

    if length(NOut) > 1
      params.precompute = LINEAR
    end

    tmpVec = Array{Complex{T},D}(undef, n)

    fftflags_ = (fftflags != nothing) ? (flags=fftflags,) : NamedTuple()
    FP = plan_fft!(tmpVec, dims_; fftflags_...)
    BP = plan_bfft!(tmpVec, dims_; fftflags_...)

    calcBlocks = (params.precompute == LINEAR ||
                  params.precompute == TENSOR ||
                  params.precompute == POLYNOMIAL ) &&
                     params.blocking && length(dims_) == D

    blocks, nodesInBlocks, blockOffsets, idxInBlock, windowTensor = precomputeBlocks(x, n, params, calcBlocks)

    windowLinInterp, windowPolyInterp, windowHatInvLUT, deconvolveIdx, B =
            precomputation(x, N[dims_], n[dims_], params)

    U = params.storeApodizationIdx ? N : ntuple(d->0,D)
    tmpVecHat = Array{Complex{T},D}(undef, U)

    NFFTPlan(N, NOut, M, x, n, dims_, params, FP, BP, tmpVec, tmpVecHat,
                       deconvolveIdx, windowHatInvLUT, windowLinInterp, windowPolyInterp,
                       blocks, nodesInBlocks, blockOffsets, idxInBlock, windowTensor, B)
end

function AbstractNFFTs.nodes!(p::NFFTPlan{T}, x::Matrix{T}) where {T}
    checkNodes(x)

    # Sort nodes in lexicographic way
    if p.params.sortNodes
        x .= sortslices(x, dims=2)
    end

    calcBlocks = (p.params.precompute == LINEAR ||
                  p.params.precompute == TENSOR ||
                  p.params.precompute == POLYNOMIAL ) &&
                     p.params.blocking && length(p.dims) == length(p.N)

    blocks, nodesInBlocks, blockOffsets, idxInBlock, windowTensor = precomputeBlocks(x, p.n, p.params, calcBlocks)

    windowLinInterp, windowPolyInterp, windowHatInvLUT, deconvolveIdx, B =
       precomputation(x, p.N, p.n, p.params)

    p.blocks = blocks
    p.nodesInBlock = nodesInBlocks
    p.blockOffsets = blockOffsets
    p.idxInBlock = idxInBlock
    p.windowTensor = windowTensor

    p.M = size(x, 2)
    p.windowLinInterp = windowLinInterp
    p.windowPolyInterp = windowPolyInterp
    p.windowHatInvLUT = windowHatInvLUT
    p.deconvolveIdx = deconvolveIdx
    p.B = B
    p.x = x

    return p
end

function Base.show(io::IO, p::NFFTPlan{T,D,R}) where {T,D,R}
    print(io, "NFFTPlan with ", p.M, " sampling points for an input array of size",
           p.N, " and an output array of size", p.NOut, " with dims ", p.dims)
end

AbstractNFFTs.size_in(p::NFFTPlan) = p.N
AbstractNFFTs.size_out(p::NFFTPlan) = p.NOut

################
# nfft functions
################

function LinearAlgebra.mul!(fHat::StridedArray, p::NFFTPlan{T,D,R}, f::AbstractArray;
               verbose=false, timing::Union{Nothing,TimingStats} = nothing) where {T,D,R}
    consistencyCheck(p, f, fHat)

    fill!(p.tmpVec, zero(Complex{T}))
    t1 = @elapsed @inbounds deconvolve!(p, f, p.tmpVec)
    t2 = @elapsed p.forwardFFT * p.tmpVec
    t3 = @elapsed @inbounds convolve!(p, p.tmpVec, fHat)
    if verbose
        @info "Timing: apod=$t1 fft=$t2 conv=$t3"
    end
    if timing != nothing
      timing.conv = t3
      timing.fft = t2
      timing.apod = t1
    end
    return fHat
end

function LinearAlgebra.mul!(f::StridedArray, pl::Adjoint{Complex{T},<:NFFTPlan{T}}, fHat::AbstractArray;
                       verbose=false, timing::Union{Nothing,TimingStats} = nothing) where {T}
    p = pl.parent
    consistencyCheck(p, f, fHat)

    t1 = @elapsed @inbounds convolve_transpose!(p, fHat, p.tmpVec)
    t2 = @elapsed p.backwardFFT * p.tmpVec
    t3 = @elapsed @inbounds deconvolve_transpose!(p, p.tmpVec, f)
    if verbose
        @info "Timing: conv=$t1 fft=$t2 apod=$t3"
    end
    if timing != nothing
      timing.conv_adjoint = t1
      timing.fft_adjoint = t2
      timing.apod_adjoint = t3
    end
    return f
end


@testset "UC With DC - PF" begin
    ps_model = PSI._canonical_model_init(bus_numbers, nothing, PM.AbstractPowerFormulation, time_range; parameters = false)
    PSI.construct_device!(ps_model, PSY.ThermalGen, PSI.ThermalUnitCommitment, PM.DCPlosslessForm, sys5b_uc, time_range; parameters = false);
    @test !(:params in keys(ps_model.JuMPmodel.ext))
    @test JuMP.num_variables(ps_model.JuMPmodel) == 480
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.LessThan{Float64}) == 504
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.GreaterThan{Float64}) == 120
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.EqualTo{Float64}) == 96
end

@testset "UC With AC - PF" begin
    ps_model = PSI._canonical_model_init(bus_numbers, nothing, PM.AbstractPowerFormulation, time_range; parameters = false)
    PSI.construct_device!(ps_model, PSY.ThermalGen, PSI.ThermalUnitCommitment, PM.StandardACPForm, sys5b_uc, time_range; parameters = false);
    @test !(:params in keys(ps_model.JuMPmodel.ext))
    @test JuMP.num_variables(ps_model.JuMPmodel) == 600
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.LessThan{Float64}) == 624
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.GreaterThan{Float64}) == 240
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.EqualTo{Float64}) == 96
end

@testset "Dispatch With DC - PF" begin
    ps_model = PSI._canonical_model_init(bus_numbers, nothing, PM.AbstractPowerFormulation, time_range)
    PSI.construct_device!(ps_model, PSY.ThermalGen, PSI.ThermalDispatch, PM.DCPlosslessForm, sys5b_uc, time_range);
    @test JuMP.num_variables(ps_model.JuMPmodel) == 120
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.Interval{Float64}) == 120
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.LessThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.GreaterThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.EqualTo{Float64}) == 0

    ps_model = PSI._canonical_model_init(bus_numbers, nothing, PM.AbstractPowerFormulation, time_range; parameters = false)
    PSI.construct_device!(ps_model, PSY.ThermalGen, PSI.ThermalDispatch, PM.DCPlosslessForm, sys5b_uc, time_range; parameters = false)
    @test !(:params in keys(ps_model.JuMPmodel.ext))
    @test JuMP.num_variables(ps_model.JuMPmodel) == 120
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.Interval{Float64}) == 120
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.LessThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.GreaterThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.EqualTo{Float64}) == 0
    end

@testset "Dispatch With AC - PF" begin
    ps_model = PSI._canonical_model_init(bus_numbers, nothing, PM.AbstractPowerFormulation, time_range)
    PSI.construct_device!(ps_model, PSY.ThermalGen, PSI.ThermalDispatch, PM.StandardACPForm, sys5b_uc, time_range);
    @test JuMP.num_variables(ps_model.JuMPmodel) == 240
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.Interval{Float64}) == 240
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.LessThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.GreaterThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.EqualTo{Float64}) == 0

    ps_model = PSI._canonical_model_init(bus_numbers, nothing, PM.AbstractPowerFormulation, time_range; parameters = false)
    PSI.construct_device!(ps_model, PSY.ThermalGen, PSI.ThermalDispatch, PM.StandardACPForm, sys5b_uc, time_range; parameters = false)
    @test !(:params in keys(ps_model.JuMPmodel.ext))
    @test JuMP.num_variables(ps_model.JuMPmodel) == 240
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.Interval{Float64}) == 240
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.LessThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.GreaterThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.EqualTo{Float64}) == 0
    end

@testset "Dispatch No-Minimum With DC - PF" begin
    ps_model = PSI._canonical_model_init(bus_numbers, nothing, PM.AbstractPowerFormulation, time_range)
    PSI.construct_device!(ps_model, PSY.ThermalGen, PSI.ThermalDispatchNoMin, PM.DCPlosslessForm, sys5b_uc, time_range);
    @test JuMP.num_variables(ps_model.JuMPmodel) == 120
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.Interval{Float64}) == 120
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.LessThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.GreaterThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.EqualTo{Float64}) == 0

    ps_model = PSI._canonical_model_init(bus_numbers, nothing, PM.AbstractPowerFormulation, time_range; parameters = false)
    PSI.construct_device!(ps_model, PSY.ThermalGen, PSI.ThermalDispatchNoMin, PM.DCPlosslessForm, sys5b_uc, time_range; parameters = false)
    @test !(:params in keys(ps_model.JuMPmodel.ext))
    @test JuMP.num_variables(ps_model.JuMPmodel) == 120
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.Interval{Float64}) == 120
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.LessThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.GreaterThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.EqualTo{Float64}) == 0
end

@testset "Dispatch No-Minimum With AC - PF" begin
    ps_model = PSI._canonical_model_init(bus_numbers, nothing, PM.AbstractPowerFormulation, time_range)
    PSI.construct_device!(ps_model, PSY.ThermalGen, PSI.ThermalDispatchNoMin, PM.StandardACPForm, sys5b_uc, time_range);
    @test JuMP.num_variables(ps_model.JuMPmodel) == 240
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.Interval{Float64}) == 240
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.LessThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.GreaterThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.EqualTo{Float64}) == 0

    ps_model = PSI._canonical_model_init(bus_numbers, nothing, PM.AbstractPowerFormulation, time_range; parameters = false)
    PSI.construct_device!(ps_model, PSY.ThermalGen, PSI.ThermalDispatchNoMin, PM.StandardACPForm, sys5b_uc, time_range; parameters = false);
    @test !(:params in keys(ps_model.JuMPmodel.ext))
    @test JuMP.num_variables(ps_model.JuMPmodel) == 240
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.Interval{Float64}) == 240
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.LessThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.GreaterThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.EqualTo{Float64}) == 0
end

@testset "Dispatch No-Minimum With DC - PF" begin
    ps_model = PSI._canonical_model_init(bus_numbers, nothing, PM.AbstractPowerFormulation, time_range)
    PSI.construct_device!(ps_model, PSY.ThermalGen, PSI.ThermalDispatchNoMin, PM.DCPlosslessForm, sys5b_uc, time_range);
    @test JuMP.num_variables(ps_model.JuMPmodel) == 120
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.Interval{Float64}) == 120
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.LessThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.GreaterThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.EqualTo{Float64}) == 0

    ps_model = PSI._canonical_model_init(bus_numbers, nothing, PM.AbstractPowerFormulation, time_range; parameters = false)
    PSI.construct_device!(ps_model, PSY.ThermalGen, PSI.ThermalDispatchNoMin, PM.DCPlosslessForm, sys5b_uc, time_range; parameters = false)
    @test !(:params in keys(ps_model.JuMPmodel.ext))
    @test JuMP.num_variables(ps_model.JuMPmodel) == 120
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.Interval{Float64}) == 120
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.LessThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.GreaterThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.EqualTo{Float64}) == 0
end

@testset "Ramp Limited Dispatch With AC - PF" begin
    ps_model = PSI._canonical_model_init(bus_numbers, nothing, PM.AbstractPowerFormulation, time_range)
    PSI.construct_device!(ps_model, PSY.ThermalGen, PSI.ThermalRampLimited, PM.StandardACPForm, sys5b_uc, time_range);
    @test JuMP.num_variables(ps_model.JuMPmodel) == 240
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.Interval{Float64}) == 240
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.LessThan{Float64}) == 192
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.GreaterThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.EqualTo{Float64}) == 0

    ps_model = PSI._canonical_model_init(bus_numbers, nothing, PM.AbstractPowerFormulation, time_range; parameters = false)
    PSI.construct_device!(ps_model, PSY.ThermalGen, PSI.ThermalRampLimited, PM.StandardACPForm, sys5b_uc, time_range; parameters = false)
    @test !(:params in keys(ps_model.JuMPmodel.ext))
    @test JuMP.num_variables(ps_model.JuMPmodel) == 240
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.Interval{Float64}) == 240
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.LessThan{Float64}) == 192
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.GreaterThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.EqualTo{Float64}) == 0
end

@testset "Ramp Limited Dispatch With DC - PF" begin
    ps_model = PSI._canonical_model_init(bus_numbers, nothing, PM.AbstractPowerFormulation, time_range)
    PSI.construct_device!(ps_model, PSY.ThermalGen, PSI.ThermalRampLimited, PM.DCPlosslessForm, sys5b_uc, time_range);
    @test JuMP.num_variables(ps_model.JuMPmodel) == 120
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.Interval{Float64}) == 120
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.LessThan{Float64}) == 192
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.GreaterThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.EqualTo{Float64}) == 0

    ps_model = PSI._canonical_model_init(bus_numbers, nothing, PM.AbstractPowerFormulation, time_range; parameters = false)
    PSI.construct_device!(ps_model, PSY.ThermalGen, PSI.ThermalRampLimited, PM.DCPlosslessForm, sys5b_uc, time_range; parameters = false)
    @test !(:params in keys(ps_model.JuMPmodel.ext))
    @test JuMP.num_variables(ps_model.JuMPmodel) == 120
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.Interval{Float64}) == 120
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.LessThan{Float64}) == 192
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.GreaterThan{Float64}) == 0
    @test JuMP.num_constraints(ps_model.JuMPmodel,GenericAffExpr{Float64,VariableRef},MOI.EqualTo{Float64}) == 0
end


"""
    FromState

a basic thermodynamic model, used to extract thermodynamic specifications from a state without using a Equation of State.
"""
struct FromState <: ThermoModel end
const URVec = AbstractVector{T} where T<:Number

moles2(a::ThermodynamicState,mw) = moles3(amount_type(a),a,mw)
mass2(a::ThermodynamicState,mw) =  mass3(amount_type(a),a,mw)
kg_per_mol2(a::ThermodynamicState,mw) =  kg_per_mol3(amount_type(a),a,mw)

function mw_error(::Nothing)
    throw(error("molecular weight information is required for this combination of properties."))
end

function mw_error(mw)
    return nothing
end
##one mol cases

#those cases should rarely be called
function moles3(::Tuple{T,OneMol},st,mw::Nothing) where T
    return one(numtype(st))
end

function moles3(::Tuple{T,OneMol},st,mw::T2) where {T,T2 <:Number}
    return one(promote_type(T2,numtype(st)))
end

#error?
function moles3(::Tuple{T,OneMol},st,mw::T2) where T where T2<:AbstractVector{T3} where T3 <: Number
    return one(promote_type(T3,numtype(st)))
end

function moles3(::Tuple{SingleComponent,MaterialAmount{MOLAR}},st,mw::T) where T <:Number
    return normalize_units(value(get_spec(MaterialAmount{MOLAR}(),st)))*_one(st)
end

function moles3(::Tuple{SingleComponent,MaterialAmount{MASS}},st,mw::T) where T <:Number
    # g/(g/mol) = mol
    mw_error(mw)
    x = mw_div(value(get_spec(MaterialAmount{MASS}(),st)),mw)
    return normalize_units(x)*_one(st)
end

## total ammounts:
function moles3(::Tuple{T,T},st,mw) where T<:MaterialCompounds{MOLAR,TOTAL_AMOUNT}
    return normalize_units(sum(value(get_spec(T(),st))))*_one(st)
end

function moles3(::Tuple{T,T},st,mw::T2) where T<:MaterialCompounds{MASS,TOTAL_AMOUNT} where T2<: URVec
    ###### needs mol
    ### exists mass numbers
    #mw_i * xmi = gc* (molc/gc)= molc, molmix = sum(molc)
    mw_error(mw)
    compounds::Spec = get_spec(T(),st)
    v = compounds.val
    sum_n_mw = zero(numtype(st))
    for i in 1:length(mw)
        sum_n_mw += mw_div(v[i],mw[i])
    end
    #sum_n_mw = mapreduce(mw_div,+,compounds,mw)
    return normalize_units(sum_n_mw)
end

function moles3(::Tuple{T,MaterialAmount{MOLAR}},st,mw) where T
    return normalize_units(value(get_spec(MaterialAmount{MOLAR}(),st)))*_one(st)
end

#inverse operations with fractions
function moles3(::Tuple{MaterialCompounds{MOLAR,FRACTION},MaterialAmount{MASS}},st,mw::T) where T<:URVec
    #mol fraction and mass
    #xni * mwi = mc/mmix * (gc/molc)= gc/molmix (sum)-> gmix/molmix
    #molmix = gmix / (gmix/molmix)
    mw_error(mw)
    compounds::Spec = get_spec(MaterialCompounds{MOLAR,FRACTION}(),st)
    amount::Spec = get_spec(MaterialAmount{MASS}(),st)
    v = compounds.val
    #sum_x_mw = mapreduce(mw_mul,+,value(compounds),mw)
    sum_x_mw = zero(numtype(st))
    for i in 1:length(mw)
        sum_x_mw += mw_mul(v[i],mw[i])
    end
    return normalize_units(value(amount)/sum_x_mw)
end

function moles3(::Tuple{MaterialCompounds{MASS,FRACTION},MaterialAmount{MASS}},st,mw::T) where T<:URVec
    #mass fractions and mass
    #mw_i * xmi = gc/gmix / (molc/gc)= molc/gmix (sum)-> molmix/gmix
    # molmix = molmix/gmix * (gmix)
    mw_error(mw)
    compounds::Spec = get_spec(MaterialCompounds{MASS,FRACTION}(),st)
    amount::Spec  = get_spec(MaterialAmount{MASS}(),st)
    v = compounds.val
    sum_x_mw = zero(numtype(st))
    for i in 1:length(mw)
        sum_x_mw += mw_div(v[i],mw[i])
    end
    #sum_x_mw = mapreduce(mw_div,+,value(compounds),mw)
    return normalize_units(value(amount)*sum_x_mw)
end

function mass3(::Tuple{SingleComponent,OneMol},st,mw::T2) where T2 <:Number
    mw_error(mw)
    return normalize_units(mw_mul(one(T2),mw))*_one(st)
end

#mass defined
function mass3(::Tuple{T,MaterialAmount{MASS}},st,mw) where T
    return normalize_units(value(get_spec(MaterialAmount{MASS}(),st)))*_one(st)
end

function mass3(::Tuple{SingleComponent,MaterialAmount{MOLAR}},st,mw) where T
    mw_error(mw)
    mol = value(get_spec(MaterialAmount{MOLAR}(),st))
    return normalize_units(mw_mul(mol,mw))*_one(st)
end

## total ammounts:
function mass3(::Tuple{T,T},st,mw) where T<:MaterialCompounds{MASS,TOTAL_AMOUNT}
    return normalize_units(sum(value(get_spec(T(),st))))*_one(st)
end

function mass3(::Tuple{T,T},st::S,mw::T2) where {S<:ThermodynamicState,T<:MaterialCompounds{MOLAR,TOTAL_AMOUNT},T2<:URVec}
    ###### needs mass
    ### exists mol numbers
    #mw_i * ni = (gc/molc) * molc = gc, sum(gc) = gmix
    mw_error(mw)
    compounds = get_spec(T(),st)::Spec
    v = compounds.val
    sum_n_mw = zero(numtype(st))
    for i in 1:length(mw)
        sum_n_mw += mw_mul(v[i],mw[i])
    end
    return normalize_units(sum_n_mw)
end

#inverse operations with fractions
function mass3(::Tuple{MaterialCompounds{MOLAR,FRACTION},MaterialAmount{MOLAR}},st,mw::T2) where T2<:URVec
    #moles, mol fractions
    #mw_i * xi = (gc/molc) * molc/molmix = gc/molmix (*molmix)
    mw_error(mw)
    amount = get_spec(MaterialAmount{MOLAR}(),st)
    compounds::Spec = get_spec(MaterialCompounds{MOLAR,FRACTION}(),st)
    sum_x_mw = zero(numtype(st))
    v = compounds.val
    for i in 1:length(mw)
        sum_x_mw += mw_mul(v[i],mw[i])
    end
    return normalize_units(value(amount)*sum_x_mw)
end

function mass3(::Tuple{MaterialCompounds{MASS,FRACTION},MaterialAmount{MOLAR}},st,mw::T2) where T2<:URVec
    #mass fractions and moles
    #sum_x_mw = xmi/mwi = mc/mmix / mc/molc = mc/mmix * molc/mc = molc/mmix -> sum -> molmix/mmix
    #a = molmix/mmix
    #inv(a) = mmix/molmix
    # molmix = molmix/gmix * (gmix)
    mw_error(mw)
    compounds::Spec = get_spec(MaterialCompounds{MASS,FRACTION}(),st)
    v = compounds.val
    amount = get_spec(MaterialAmount{MOLAR}(),st)
    sum_x_mw = zero(numtype(st))
    for i in 1:length(mw)
        sum_x_mw += mw_div(v[i],mw[i])
    end
    return normalize_units(value(amount)/sum_x_mw)
end

#single component cases, all the same
function kg_per_mol3(::Tuple{SingleComponent,OneMol},st,mw::T2) where T2 <:Number
    return normalize_units(mw_mul(one(T2),mw))*_one(st)
end

function kg_per_mol3(::Tuple{SingleComponent,MaterialAmount},st,mw::T2) where {T2 <:Number}
    return normalize_units(mw_mul(one(T2),mw))*_one(st)
end

#total amounts:
function kg_per_mol3(::Tuple{T,T},st,mw) where T<:MaterialCompounds{MASS,TOTAL_AMOUNT}
    #mw from mass numbers
    #gc / (gc/molc) = molc sum-> molmix
    mw_error(mw)
    compounds::Spec = get_spec(T(),st)
    v = compounds.val
    molmix = zero(numtype(st))
    for i in 1:length(mw)
        molmix += mw_div(v[i],mw[i])
    end
    gmix = sum(v)
    return normalize_units(gmix/molmix)
end

function kg_per_mol3(::Tuple{T,T},st,mw) where T<:MaterialCompounds{MOLAR,TOTAL_AMOUNT}
    #mw from mol numbers
    #molc * gc/molc = gc sum-> gmix
    mw_error(mw)
    compounds::Spec = get_spec(T(),st)
    v = compounds.val
    gmix = zero(numtype(st))
    for i in 1:length(mw)
        gmix += mw_mul(v[i],mw[i])
    end
    #gmix = mapreduce(mw_mul,+,value(compounds),mw)
    molmix = sum(value(compounds))
    return normalize_units(gmix/molmix)*_one(st)
end

#fraction amounts
function kg_per_mol3(::Tuple{MaterialCompounds{MOLAR,FRACTION},MaterialAmount{T}},st,mw) where T
    #mol fraction and mass
    #xni * mwi = mc/mmix * (gc/molc)= gc/molmix (sum)-> gmix/molmix
    mw_error(mw)
    compounds::Spec = get_spec(MaterialCompounds{MOLAR,FRACTION}(),st)
    #sum_x_mw = mapreduce(mw_mul,+,value(compounds),mw)*_one(st)
    v = compounds.val
    sum_x_mw = zero(numtype(st))
    for i in 1:length(mw)
        sum_x_mw += mw_mul(v[i],mw[i])
    end
    return normalize_units(one(sum_x_mw)/sum_x_mw)
end

function kg_per_mol3(::Tuple{MaterialCompounds{MASS,FRACTION},MaterialAmount{T}},st,mw) where T
    mw_error(mw)
    #mass fractions and mass
    compounds::Spec = get_spec(MaterialCompounds{MASS,FRACTION}(),st)
    #xmi * mwi = gc/gmix / gc/molc = molc/gmix (sum)-> molmix/gmix
    #sum_x_mw = mapreduce(mw_div,+,value(compounds),mw)*_one(st)
    v = compounds.val
    sum_x_mw = zero(numtype(st))
    for i in 1:length(mw)
        sum_x_mw += mw_div(v[i],mw[i])
    end
    return normalize_units(one(sum_x_mw)/sum_x_mw)
end

#invariant
function to_spec(st,sp::Spec{SP},mw,::MOLAR) where {SP<:AbstractIntensiveSpec{MOLAR}}
    return normalize_units(value(sp))*_one(st)
end

function to_spec(st,sp::Spec{SP},mw,any_value) where {SP<:Union{Pressure,Temperature}}
    return normalize_units(value(sp))*_one(st)
end

function to_spec(st,sp::Spec{SP},mw,::MASS) where {SP<:AbstractIntensiveSpec{MASS}}
    return normalize_units(value(sp))*_one(st)
end

function to_spec(st,sp::Spec{SP},mw,::TOTAL) where {SP<:AbstractIntensiveSpec{TOTAL}}
    return normalize_units(value(sp))*_one(st)
end

#to mass

#from molar to mass
function to_spec(st,sp::Spec{SP},mw,::MASS) where {SP<:AbstractIntensiveSpec{MOLAR}}
    return normalize_units(value(sp))/kg_per_mol2(st,mw)*_one(st)
end

function to_spec(st,sp::Spec{SP},mw,::MASS) where {SP<:AbstractIntensiveSpec{TOTAL}}
    return normalize_units(value(sp))/mass2(st,mw)*_one(st)
end

#to mol
function to_spec(st,sp::Spec{SP},mw,::MOLAR) where {SP<:AbstractIntensiveSpec{MASS}}
    return normalize_units(value(sp))*kg_per_mol2(st,mw)
end

function to_spec(st,sp::Spec{SP},mw,::MOLAR) where {SP<:AbstractIntensiveSpec{TOTAL}}
    return normalize_units(value(sp))/moles2(st,mw)
end

#to total
function to_spec(st,sp::Spec{SP},mw,::TOTAL) where {SP<:AbstractIntensiveSpec{MOLAR}}
    return normalize_units(value(sp))*moles2(st,mw)
end

function to_spec(st,sp::Spec{SP},mw,::TOTAL) where {SP<:AbstractIntensiveSpec{MASS}}
    return normalize_units(value(sp))*mass2(st,mw)
end

#volume and density
#same value
function to_spec_vol(st,sp::Spec{T},mw,::T) where {T<:VolumeAmount}
    return normalize_units(value(sp))
end

#inversion

function to_spec_vol(st,sp::Spec{VolumeAmount{T,VOLUME}},mw,::VolumeAmount{T,DENSITY}) where T
    val = value(sp)
    return normalize_units(one(val)/val)
end

function to_spec_vol(st,sp::Spec{VolumeAmount{T,DENSITY}},mw,::VolumeAmount{T,VOLUME}) where T
    val = value(sp)
    return normalize_units(one(val)/val)
end

#same type, volume
function to_spec_vol(st,sp::Spec{VolumeAmount{MOLAR,VOLUME}},mw,::VolumeAmount{MASS,VOLUME})
    val = value(sp)
    return normalize_units(val/kg_per_mol2(st,mw))
end

function to_spec_vol(st,sp::Spec{VolumeAmount{MASS,VOLUME}},mw,::VolumeAmount{MOLAR,VOLUME})
    val = value(sp)
    return normalize_units(val*kg_per_mol2(st,mw))
end

function to_spec_vol(st,sp::Spec{VolumeAmount{MOLAR,VOLUME}},mw,::VolumeAmount{TOTAL,VOLUME})
    val = value(sp)
    return normalize_units(val*moles2(st,mw))
end

function to_spec_vol(st,sp::Spec{VolumeAmount{TOTAL,VOLUME}},mw,::VolumeAmount{MOLAR,VOLUME})
    val = value(sp)
    return normalize_units(val/moles2(st,mw))
end

function to_spec_vol(st,sp::Spec{VolumeAmount{MASS,VOLUME}},mw,::VolumeAmount{TOTAL,VOLUME})
    val = value(sp)
    return normalize_units(val*mass2(st,mw))
end

function to_spec_vol(st,sp::Spec{VolumeAmount{TOTAL,VOLUME}},mw,::VolumeAmount{MASS,VOLUME})
    val = value(sp)
    return normalize_units(val/mass2(st,mw))
end

#same type, density
function to_spec_vol(st,sp::Spec{VolumeAmount{MOLAR,DENSITY}},mw,::VolumeAmount{MASS,DENSITY})
    val = value(sp)
    return normalize_units(val*kg_per_mol2(st,mw))
end

function to_spec_vol(st,sp::Spec{VolumeAmount{MASS,DENSITY}},mw,::VolumeAmount{MOLAR,DENSITY})
    val = value(sp)
    return normalize_units(val/kg_per_mol2(st,mw))
end

function to_spec_vol(st,sp::Spec{VolumeAmount{MOLAR,DENSITY}},mw,::VolumeAmount{TOTAL,DENSITY})
    val = value(sp)
    return normalize_units(val/moles2(st,mw))
end

function to_spec_vol(st,sp::Spec{VolumeAmount{TOTAL,DENSITY}},mw,::VolumeAmount{MOLAR,DENSITY})
    val = value(sp)
    return  (val*moles2(st,mw))
end

function to_spec_vol(st,sp::Spec{VolumeAmount{MASS,DENSITY}},mw,::VolumeAmount{TOTAL,DENSITY})
    val = value(sp)
    return normalize_units(val/mass2(st,mw))
end

function to_spec_vol(st,sp::Spec{VolumeAmount{TOTAL,DENSITY}},mw,::VolumeAmount{MASS,DENSITY})
    val = value(sp)
    return normalize_units(val*mass2(st,mw))
end

#all different combinations
function to_spec_vol(st,sp::Spec{VolumeAmount{MOLAR,VOLUME}},mw,::VolumeAmount{TOTAL, DENSITY})
    val = value(sp)
    totv = val*moles2(st,mw)
    return normalize_units(one(totv)/totv)
  end

function to_spec_vol(st,sp::Spec{VolumeAmount{MOLAR,VOLUME}},mw,::VolumeAmount{MASS, DENSITY})
    val = value(sp)
    val2 = val/kg_per_mol2(st,mw)
    return normalize_units(one(val2)/val2)
end

function to_spec_vol(st,sp::Spec{VolumeAmount{TOTAL,VOLUME}},mw,::VolumeAmount{MOLAR, DENSITY})
    val = value(sp)
    val2 = val/moles2(st,mw)
    return normalize_units(one(val2)/val2)
end
function to_spec_vol(st,sp::Spec{VolumeAmount{TOTAL,VOLUME}},mw,::VolumeAmount{MASS, DENSITY})
    val = value(sp)
    val2 = val/mass2(st,mw)
    return normalize_units(one(val2)/val2)
end

function to_spec_vol(st,sp::Spec{VolumeAmount{MASS,VOLUME}},mw,::VolumeAmount{MOLAR, DENSITY})
    val = value(sp)
    val2 = val*kg_per_mol2(st,mw)
    return normalize_units(one(val2)/val2)
end
function to_spec_vol(st,sp::Spec{VolumeAmount{MASS,VOLUME}},mw,::VolumeAmount{TOTAL, DENSITY})
    val = value(sp)
    val2 = val*mass2(st,mw)
    return normalize_units(one(val2)/val2)
end

function to_spec_vol(st,sp::Spec{VolumeAmount{TOTAL,DENSITY}},mw,::VolumeAmount{MOLAR, VOLUME})
    val = value(sp)
    val2 = val*moles2(st,mw)
    return normalize_units(one(val2)/val2)
end
function to_spec_vol(st,sp::Spec{VolumeAmount{MASS,DENSITY}},mw,::VolumeAmount{MOLAR, VOLUME})
    val = value(sp)
    val2 = val/kg_per_mol2(st,mw)
    return normalize_units(one(val2)/val2)
end

function to_spec_vol(st,sp::Spec{VolumeAmount{MOLAR,DENSITY}},mw,::VolumeAmount{TOTAL, VOLUME})
    val = value(sp)
    val2 = val/moles2(st,mw)
    return normalize_units(one(val2)/val2)
end
function to_spec_vol(st,sp::Spec{VolumeAmount{MASS,DENSITY}},mw,::VolumeAmount{TOTAL, VOLUME})
    val = value(sp)
    val2 = val/mass2(st,mw)
    return normalize_units(one(val2)/val2)
end

function to_spec_vol(st,sp::Spec{VolumeAmount{MOLAR,DENSITY }},mw,::VolumeAmount{MASS, VOLUME})
    val = value(sp)
    val2 = val*kg_per_mol2(st,mw)
    return normalize_units(one(val2)/val2)
end
function to_spec_vol(st,sp::Spec{VolumeAmount{TOTAL, DENSITY }},mw,::VolumeAmount{MASS, VOLUME})
    val = value(sp)
    val2 = val*mass2(st,mw)
    return normalize_units(one(val2)/val2)
end

function to_spec(st,sp::Spec{SP1},mw,x::SP2) where {SP1<:VolumeAmount,SP2<:VolumeAmount}
    return to_spec_vol(st,sp,mw,x)
end

#material compounds part

#fraction <-> total part, same base
function to_spec_mat(st,sp::Spec{MaterialCompounds{MOLAR,FRACTION}},mw,::MaterialCompounds{MOLAR,TOTAL_AMOUNT})
    return moles2(st,mw) .* normalize_units(value(sp))
end

function to_spec_mat(st,sp::Spec{MaterialCompounds{MASS,FRACTION}},mw,::MaterialCompounds{MASS,TOTAL_AMOUNT})
    return mass2(st,mw) .* normalize_units(value(sp))
end

function to_spec_mat(st,sp::Spec{MaterialCompounds{MOLAR,TOTAL_AMOUNT}},mw,::MaterialCompounds{MOLAR,FRACTION})
    return normalize_units(value(sp)) ./ moles2(st,mw)
end

function to_spec_mat(st,sp::Spec{MaterialCompounds{MASS,TOTAL_AMOUNT}},mw,::MaterialCompounds{MASS,FRACTION})
    return normalize_units(value(sp)) ./ mass2(st,mw)
end

#molar total <-> mass total
function to_spec_mat(st,sp::Spec{MaterialCompounds{MOLAR,TOTAL_AMOUNT}},mw,::MaterialCompounds{MASS,TOTAL_AMOUNT})
    return normalize_units(map(mw_mul,value(sp),mw))
end

function to_spec_mat(st,sp::Spec{MaterialCompounds{MASS,TOTAL_AMOUNT}},mw,::MaterialCompounds{MOLAR,TOTAL_AMOUNT})
    return  normalize_units(map(mw_div,value(sp),mw))
end

#molar fraction <-> mass fraction
function to_spec_mat(st,sp::Spec{MaterialCompounds{MOLAR,FRACTION}},mw,::MaterialCompounds{MASS,FRACTION})
    n =  normalize_units(value(sp)) .* moles2(st,mw)
    m = normalize_units(map(mw_mul,n,mw))
    return m ./ sum(m)
end

function to_spec_mat(st,sp::Spec{MaterialCompounds{MASS,FRACTION}},mw,::MaterialCompounds{MOLAR,FRACTION})
    m =  normalize_units(value(sp)) .* mass2(st,mw)
    n = map(mw_div,m,mw)
    return n ./ sum(n)
end

#crossed properties:

function to_spec_mat(st,sp::Spec{MaterialCompounds{MOLAR,FRACTION}},mw,::MaterialCompounds{MASS,TOTAL_AMOUNT})
    n =  value(sp).* moles2(st,mw)
    m = normalize_units(map(mw_mul,n,mw))
end

function to_spec_mat(st,sp::Spec{MaterialCompounds{MASS,FRACTION}},mw,::MaterialCompounds{MOLAR,TOTAL_AMOUNT})
    m =  value(sp) .* mass2(st,mw)
    n = normalize_units(map(mw_div,m,mw))
end

function to_spec_mat(st,sp::Spec{MaterialCompounds{MOLAR,TOTAL_AMOUNT}},mw,::MaterialCompounds{MASS,FRACTION})
    mw_error(mw)
    m =  map(mw_mul,value(sp),mw)
    return normalize_units(m ./ sum(m))
end

function to_spec_mat(st,sp::Spec{MaterialCompounds{MASS,TOTAL_AMOUNT}},mw,::MaterialCompounds{MOLAR,FRACTION})
    mw_error(mw)
    n =  map(mw_div,value(sp),mw)
    return normalize_units(n ./ sum(n))
end

function to_spec_compounds(st,mw,x::SP2) where {SP2<:MaterialCompounds}
    return to_spec_compounds(amount_type(st),st,mw,x)
end

function to_spec_compounds(amount_type::Tuple{SingleComponent,T},st,mw,x) where {T}
    return single_component_to_spec(st,mw,x)
end

function to_spec_compounds(amount_type::Tuple{T1,T2},st,mw,x::T1) where {T1<:MaterialCompounds,T2}
    sp = get_spec(T1,st)
    return normalize_units(value(sp))
end

function to_spec_compounds(amount_type::Tuple{T1,T2},st,mw,x) where {T1<:MaterialCompounds,T2}
    sp = get_spec(T1,st)
    return to_spec_mat(st,sp,mw,x)
end

function single_component_to_spec(st,mw,x::MaterialCompounds{MOLAR,FRACTION})
    return [_one(st)]
end

function single_component_to_spec(st,mw,x::MaterialCompounds{MASS,FRACTION})
    return [_one(st)]
end

function single_component_to_spec(st,mw,x::MaterialCompounds{MOLAR,TOTAL_AMOUNT})
    return [moles2(st,mw)]
end

function single_component_to_spec(st,mw,x::MaterialCompounds{MASS,TOTAL_AMOUNT})
    return [mass2(st,mw)]
end


module InfiniteOpt

# Import the necessary packages.
import JuMP
import MathOptInterface
import Distributions
import DataStructures
import Polynomials

# Make useful aliases
const MOI = MathOptInterface
const MOIU = MOI.Utilities
const JuMPC = JuMP.Containers
const MOIUC = MOIU.CleverDicts

# Import the Collections module
include("Collections/Collections.jl")
using .Collections: VectorTuple, same_structure

# Import methods
include("datatypes.jl")
include("infinite_sets.jl")
include("scalar_parameters.jl")
include("array_parameters.jl")
include("variable_basics.jl")
include("infinite_variables.jl")
include("reduced_variables.jl")
include("point_variables.jl")
include("hold_variables.jl")
include("expressions.jl")
include("macro_utilities.jl")
include("measures.jl")

# Import MeasureToolbox
include("MeasureToolbox/MeasureToolbox.jl")
using .MeasureToolbox: Automatic, UniTrapezoid, UniMCSampling,
UniIndepMCSampling, Quadrature, GaussHermite, GaussLegendre, GaussLaguerre,
MultiMCSampling, MultiIndepMCSampling, uni_integral_defaults,
set_uni_integral_defaults, integral, multi_integral_defaults,
set_multi_integral_defaults, expect, support_sum, @integral, @expect, @support_sum,
generate_integral_data, 𝔼, ∫, @𝔼, @∫

# Import more methods
include("derivatives.jl")
include("constraints.jl")
include("macros.jl")
include("objective.jl")
include("measure_expansions.jl")
include("derivative_evaluations.jl")
include("optimize.jl")
include("results.jl")
include("show.jl")
include("utilities.jl")
include("general_variables.jl")

# Import TranscriptionOpt
include("TranscriptionOpt/TranscriptionOpt.jl")
using .TranscriptionOpt

# Export core datatypes
export AbstractDataObject, InfiniteModel

# Export macros
export @independent_parameter, @dependent_parameters, @infinite_parameter,
@finite_parameter, @infinite_variable, @point_variable, @hold_variable,
@infinite_parameter, @BDconstraint, @set_parameter_bounds, @add_parameter_bounds,
@measure, @integral, @expect, @support_sum, @deriv, @derivative_variable, @∂, @𝔼,
@∫

# Export variable types
export InfOptVariableType, Infinite, Hold, Point, Deriv

# Export index types
export AbstractInfOptIndex, ObjectIndex, IndependentParameterIndex,
DependentParametersIndex, DependentParameterIndex, FiniteParameterIndex,
InfiniteVariableIndex, PointVariableIndex, ReducedVariableIndex,
HoldVariableIndex, MeasureIndex, ConstraintIndex, InfiniteParameterIndex,
FiniteVariableIndex, DerivativeIndex, ParameterFunctionIndex

# Export infinite set types
export AbstractInfiniteSet, InfiniteScalarSet, InfiniteArraySet, IntervalSet,
UniDistributionSet, MultiDistributionSet, CollectionSet, AbstractSupportLabel,
PublicLabel, InternalLabel, SampleLabel, UserDefined, MCSample, UniformGrid,
Mixture, All, WeightedSample, UniqueMeasure, OrthogonalCollocationNode

# Export infinite set methods
export collection_sets, supports_in_set, generate_support_values

# Export parameter datatypes
export InfOptParameter, ScalarParameter, IndependentParameter, FiniteParameter,
DependentParameters, ScalarParameterData, MultiParameterData,
IndependentParameterRef, DependentParameterRef, FiniteParameterRef,
ScalarParameterRef, InfiniteParameter, AbstractDerivativeMethod,
GenerativeDerivativeMethod, NonGenerativeDerivativeMethod, OrthogonalCollocation,
FiniteDifference, OCTechnique, Lobatto, FDTechnique, Forward, Central,
Backward

# Export parameter methods
export build_parameter, add_parameter, add_parameters,
infinite_set, set_infinite_set, parameter_by_name, num_parameters,
all_parameters, num_supports, has_supports,
set_supports, add_supports, delete_supports, supports, used_by_constraint,
used_by_measure, used_by_infinite_variable, is_used, generate_and_add_supports!,
significant_digits, parameter_value, used_by_derivative, derivative_method,
set_derivative_method, set_all_derivative_methods, has_internal_supports,
has_derivative_supports, has_derivative_constraints, used_by_parameter_function

# Export the paramter function datatypes and methods
export InfiniteParameterFunction, ParameterFunctionData, ParameterFunctionRef,
build_parameter_function, add_parameter_function, raw_function, parameter_function

# Export variable datatypes
export InfOptVariable, InfiniteVariable, ReducedVariable,
PointVariable, HoldVariable, ParameterBounds, VariableData, GeneralVariableRef,
DispatchVariableRef, InfiniteVariableRef, ReducedVariableRef,
MeasureFiniteVariableRef, FiniteVariableRef, PointVariableRef, HoldVariableRef,
DecisionVariableRef, UserDecisionVariableRef

# Export variable methods
export used_by_objective, infinite_variable_ref, parameter_refs,
used_by_point_variable, parameter_values,
eval_supports, used_by_reduced_variable, has_parameter_bounds, parameter_bounds,
set_parameter_bounds, add_parameter_bounds, delete_parameter_bounds,
parameter_list, raw_parameter_refs, raw_parameter_values,
intervals, dispatch_variable_ref, internal_reduced_variable, start_value_function,
set_start_value_function, reset_start_value_function

# Export derivative datatypes
export Derivative, DerivativeRef

# Export derivative methods
export build_derivative, add_derivative, derivative_argument, operator_parameter,
deriv, evaluate_derivative, evaluate, evaluate_all_derivatives!, num_derivatives,
all_derivatives, derivative_constraints, delete_derivative_constraints, ∂

# Export measure datatypes
export AbstractMeasureData, DiscreteMeasureData, FunctionalDiscreteMeasureData,
Measure, MeasureData, MeasureRef

# Export measure methods
export default_weight, support_label, coefficient_function, coefficients,
weight_function, build_measure, min_num_supports, add_measure, measure_function,
measure_data, is_analytic, measure, num_measures, all_measures,
add_supports_to_parameters, measure_data_in_hold_bounds, make_point_variable_ref,
make_reduced_variable_ref, add_measure_variable, delete_internal_reduced_variable,
delete_reduced_variable, expand, expand_all_measures!, expand_measure

# Export the measure toolbox functions and datatypes
export Automatic, UniTrapezoid, UniMCSampling, UniIndepMCSampling, Quadrature,
GaussHermite, GaussLegendre, GaussLaguerre, MultiMCSampling,
MultiIndepMCSampling, uni_integral_defaults, set_uni_integral_defaults,
integral, multi_integral_defaults, set_multi_integral_defaults, expect,
support_sum, generate_integral_data, 𝔼, ∫

# Export objective methods
export objective_has_measures

# Export constraint datatypes
export BoundedScalarConstraint, ConstraintData, InfOptConstraintRef,
InfiniteConstraintRef, FiniteConstraintRef

# Export constraint methods
export original_parameter_bounds

# Export transcription datatypes
export TranscriptionData, TranscriptionModel

# Export transcription methods
export is_transcription_model, transcription_data, transcription_variable,
transcription_constraint, transcription_model, transcription_expression

# Export optimize methods
export optimizer_model, set_optimizer_model, optimizer_model_ready,
set_optimizer_model_ready, build_optimizer_model!, optimizer_model_key,
optimizer_model_variable, optimizer_model_constraint, constraint_parameter_refs,
clear_optimizer_model_build!, optimizer_model_expression

# Export support generation and fill-in methods
export fill_in_supports!, generate_supports

end # module


struct PosteriorLTISDE{Tprior<:LTISDE, Tdata} <: AbstractGP
    prior::Tprior
    data::Tdata
end

# Avoids method ambiguity.
posterior(fx::FiniteLTISDE, y::AbstractVector) = _posterior(fx, y)
posterior(fx::FiniteLTISDE, y::AbstractVector{<:Real}) = _posterior(fx, y)

_posterior(fx, y) = PosteriorLTISDE(fx.f, (y=y, x=fx.x, Σy=fx.Σy))

const FinitePosteriorLTISDE = FiniteGP{<:PosteriorLTISDE}

function AbstractGPs.cov(fx::FinitePosteriorLTISDE)
    @error "Intentionally not implemented. Please don't try to explicitly compute this cov. matrix."
end

function AbstractGPs.marginals(fx::FinitePosteriorLTISDE)
    if fx.x != fx.f.data.x
        x, Σys, ys, _, pr_indices = build_inference_data(fx.f, fx.x)

        model = build_lgssm(fx.f.prior, x, Σys)
        Σys_new = noise_var_to_time_form(fx.x, fx.Σy)
        σ²s_pr_full = build_prediction_obs_vars(pr_indices, x, Σys_new)
        model_post = replace_observation_noise_cov(posterior(model, ys), σ²s_pr_full)
        return destructure(x, map(marginals, marginals(model_post))[pr_indices])
    else
        f = Zygote.literal_getfield(fx, Val(:f))
        prior = Zygote.literal_getfield(f, Val(:prior))
        x = Zygote.literal_getfield(fx, Val(:x))
        data = Zygote.literal_getfield(f, Val(:data))
        Σy = Zygote.literal_getfield(data, Val(:Σy))
        Σy_diag = Zygote.literal_getfield(Σy, Val(:diag))
        y = Zygote.literal_getfield(data, Val(:y))

        Σy_new = Zygote.literal_getfield(fx, Val(:Σy))

        model = build_lgssm(AbstractGPs.FiniteGP(prior, x, Σy))
        Σys_new = noise_var_to_time_form(x, Σy_new)
        ys = observations_to_time_form(x, y)
        model_post = replace_observation_noise_cov(posterior(model, ys), Σys_new)
        return destructure(x, map(marginals, marginals(model_post)))
    end
end

function AbstractGPs.mean_and_var(fx::FinitePosteriorLTISDE)
    ms = marginals(fx)
    return map(mean, ms), map(var, ms)
end

AbstractGPs.mean(fx::FinitePosteriorLTISDE) = mean_and_var(fx)[1]

AbstractGPs.var(fx::FinitePosteriorLTISDE) = mean_and_var(fx)[2]

function AbstractGPs.rand(rng::AbstractRNG, fx::FinitePosteriorLTISDE)
    x, Σys, ys, _, pr_indices = build_inference_data(fx.f, fx.x)

    # Get time-form observation vars and observations.
    Σys_pr = noise_var_to_time_form(fx.x, fx.Σy)

    model = build_lgssm(fx.f.prior, x, Σys)
    σ²s_pr_full = build_prediction_obs_vars(pr_indices, x, Σys_pr)
    model_post = replace_observation_noise_cov(posterior(model, ys), σ²s_pr_full)
    return destructure(x, rand(rng, model_post)[pr_indices])
end

AbstractGPs.rand(fx::FinitePosteriorLTISDE) = rand(Random.GLOBAL_RNG, fx)

function AbstractGPs.logpdf(fx::FinitePosteriorLTISDE, y_pr::AbstractVector{<:Real})

    x, Σys, ys, tr_indices, pr_indices = build_inference_data(
        fx.f, fx.x, fx.Σy, fill(missing, length(y_pr)),
    )

    # Get time-form observation vars and observations.
    Σys_pr = noise_var_to_time_form(fx.x, fx.Σy)
    ys_pr = observations_to_time_form(fx.x, y_pr)

    Σys_pr_full = build_prediction_obs_vars(pr_indices, x, Σys_pr)
    ys_pr_full = build_prediction_obs(tr_indices, pr_indices, x, ys_pr)

    model = build_lgssm(fx.f.prior, x, Σys)
    model_post = replace_observation_noise_cov(posterior(model, ys), Σys_pr_full)
    return logpdf(model_post, ys_pr_full)
end

# Join the dataset used to construct `f` and the one specified by `x_pred`, `σ²s_pred`, and
# `y_pred`. This is used in all of the inference procedures. Also provide a collection of
# indices that can be used to obtain the requested prediction locations from the joint data.

# This is the most naive way of going about this.
# The present implementation assumes that there are no overlapping data, which will lead to
# numerical issues if violated.
function build_inference_data(
    f::PosteriorLTISDE,
    x_pred::AbstractVector,
    Σy_pred::AbstractMatrix{<:Real},
    y_pred::AbstractVector{<:Union{Missing, Real}},
)
    d = f.data
    return merge_datasets((d.x, x_pred), (d.Σy, Σy_pred), (d.y, y_pred))
end

function merge_datasets(
    (x1, x2)::Tuple{AbstractVector, AbstractVector},
    (Σy1, Σy2)::Tuple{AbstractMatrix{<:Real}, AbstractMatrix{<:Real}},
    (y1, y2)::Tuple{
        AbstractVector{<:Union{Missing, Real}}, AbstractVector{<:Union{Missing, Real}},
    },
)
    # Merge and sort the inputs temporally.
    x_raw = merge_inputs(x1, x2)
    sort_indices, x = sort_in_time(x_raw)

    # Merge and sort the observation covariances temporally.
    Σys_raw = vcat(noise_var_to_time_form(x1, Σy1), noise_var_to_time_form(x2, Σy2))
    Σys = Σys_raw[sort_indices]

    # Merge and sort the observations temporally.
    ys_1 = observations_to_time_form(x1, y1)
    ys_2 = observations_to_time_form(x2, y2)
    ys_raw = vcat(ys_1, ys_2)
    ys = ys_raw[sort_indices]

    # The last length(x_pred) indices belong to the predictions.
    x1_indices = sortperm(sort_indices)[1:length(ys_1)]
    x2_indices = sortperm(sort_indices)[end-length(ys_2) + 1:end]

    return x, Σys, ys, x1_indices, x2_indices
end

# If no observations or variances are provided, make the observations arbitrary and the
# variances very large to simulate missing data.
function build_inference_data(f::PosteriorLTISDE, x_pred::AbstractVector)
    Σy_pred = Diagonal(fill(_large_var_const(), length(x_pred)))
    y_pred = fill(missing, length(x_pred))
    return build_inference_data(f, x_pred, Σy_pred, y_pred)
end

# Functions that make predictions at new locations require missings to be placed at the
# locations of the training data.
function build_prediction_obs_vars(
    pr_indices::AbstractVector{<:Integer},
    x_full::AbstractVector,
    Σys_pr::AbstractVector,
)
    σ²s_pr_full = get_zeros(x_full)
    σ²s_pr_full[pr_indices] .= Σys_pr
    return σ²s_pr_full
end

get_zeros(x::AbstractVector{T}) where {T<:Real} = zeros(T, length(x))

function build_prediction_obs(
    tr_indices::AbstractVector{<:Integer},
    pr_indices::AbstractVector{<:Integer},
    x_full::AbstractVector,
    y_pr::AbstractVector{T},
) where {T}
    y_pr_full = Vector{Union{Missing, T}}(undef, length(get_times(x_full)))
    y_pr_full[tr_indices] .= missing
    y_pr_full[pr_indices] .= y_pr
    return y_pr_full
end