# https://ieeexplore.ieee.org/document/411098
struct DCDCConverter{T} <: AbstractDynamicalSystem{T}
    C₁::T
    C₂::T
    L₃::T
    period::T
    μ::T  # Duty ratio
end

DCDCConverter{T}() where {T} = DCDCConverter{T}(0.1, 0.2, 0.5, 3.0, 0.5)

function (system::DCDCConverter)(du, u, p, t)
    (; C₁, C₂, L₃, period, μ) = system
    V₁, V₂, I₃ = u
    s = (t % period) / period < μ ? 0 : 1

    du[1] = ((1 - s) * I₃) / C₁
    du[2] = (s * I₃) / C₂
    du[3] = (-(1 - s) * V₁ - s * V₂) / L₃

    return nothing
end

function initial_conditions(::DCDCConverter{T}) where {T}
    return rand(T, 3)
end

function get_trajectories(
    system::DCDCConverter{T},
    experiment_version,
    seconds,
    dt,
    transient_seconds,
    solver,
    reltol,
    abstol,
    N,
    steps,
    stabilization_param,
    θ,
    restructure,
) where {T}
    f = NeuralVectorField(system, experiment_version, restructure)
    F = ConstraintsPseudoinverse(system, experiment_version)
    γ = T(stabilization_param)

    systemBF = DCDCConverter{BigFloat}()
    trajectories = []
    for _ = 1:N
        time_series = generate_data(
            systemBF;
            seconds,
            dt,
            transient_seconds,
            solver,
            reltol,
            abstol,
            u0 = initial_conditions(systemBF),
            NF = T,
        )
        u0 = time_series.trajectory[:, 1]
        t0 = time_series.times[1]

        # Set up the SNDE
        if γ == 0
            rhs = f
        else
            g = ConstraintsFunction(system, experiment_version, u0, t0)
            rhs = StabilizedNDE(f, g, F, γ)
        end
        prob =
            ODEProblem{false,SciMLBase.FullSpecialize}(rhs, zeros(T), (zero(T), one(T)), θ)
        data_ms = multiple_shooting(prob, time_series; steps)
        push!(trajectories, data_ms)
    end
    return vcat(trajectories...)
end

function get_mlp(
    hidden_layers,
    hidden_width,
    activation,
    ::DCDCConverter{T},
    ::Val{1},
) where {T}
    return get_mlp(4 => 3, hidden_layers, hidden_width, activation, T)
end

function constraints(u, t, system::DCDCConverter{T}, ::Val{1}) where {T}
    (; C₁, C₂, L₃) = system
    V₁, V₂, I₃ = u
    return [T(0.5) * (C₁ * V₁^2 + C₂ * V₂^2 + L₃ * I₃^2)]
end

function constraints_jacobian(u, t, system::DCDCConverter{T}, ::Val{1}) where {T}
    (; C₁, C₂, L₃) = system
    V₁, V₂, I₃ = u
    return [C₁ * V₁ C₂ * V₂ L₃ * I₃]
end

function rhs_neural(u, θ, t, re::Optimisers.Restructure, system::DCDCConverter, ::Val{1})
    (; period, μ) = system
    s = (t % period) / period < μ ? 0 : 1  # Switch ∈ {0,1}
    return re(θ)(vcat(u, s))
end
