text:
#==============================================================================#
# Health.jl
#
# This file is generated from:
# https://github.com/aws/aws-sdk-js/blob/master/apis/health-2016-08-04.normal.json
#==============================================================================#

__precompile__()

module Health

using AWSCore

"""
    using AWSSDK.Health.describe_affected_entities
    describe_affected_entities([::AWSConfig], arguments::Dict)
    describe_affected_entities([::AWSConfig]; filter=, <keyword arguments>)

    using AWSCore.Services.health
    health([::AWSConfig], "DescribeAffectedEntities", arguments::Dict)
    health([::AWSConfig], "DescribeAffectedEntities", filter=, <keyword arguments>)

# DescribeAffectedEntities Operation

Returns a list of entities that have been affected by the specified events, based on the specified filter criteria. Entities can refer to individual customer resources, groups of customer resources, or any other construct, depending on the AWS service. Events that have impact beyond that of the affected entities, or where the extent of impact is unknown, include at least one entity indicating this.

At least one event ARN is required. Results are sorted by the `lastUpdatedTime` of the entity, starting with the most recent.

# Arguments

## `filter = [ ... ]` -- *Required*
Values to narrow the results returned. At least one event ARN is required.
```
 filter = [
        "eventArns" => <required> [::String, ...],
        "entityArns" =>  [::String, ...],
        "entityValues" =>  [::String, ...],
        "lastUpdatedTimes" =>  [[
            "from" =>  timestamp,
            "to" =>  timestamp
        ], ...],
        "tags" =>  [::Dict{String,String}, ...],
        "statusCodes" =>  ["IMPAIRED", "UNIMPAIRED" or "UNKNOWN", ...]
    ]
```

## `locale = ::String`
The locale (language) to return information in. English (en) is the default and the only supported value at this time.

## `nextToken = ::String`
If the results of a search are large, only a portion of the results are returned, and a `nextToken` pagination token is returned in the response. To retrieve the next batch of results, reissue the search request and include the returned token. When all results have been returned, the response does not contain a pagination token value.

## `maxResults = ::Int`
The maximum number of items to return in one batch, between 10 and 100, inclusive.

# Returns

`DescribeAffectedEntitiesResponse`

# Exceptions

`InvalidPaginationToken` or `UnsupportedLocale`.

See also: [AWS API Documentation](https://docs.aws.amazon.com/goto/WebAPI/health-2016-08-04/DescribeAffectedEntities)
"""
@inline describe_affected_entities(aws::AWSConfig=default_aws_config(); args...) = describe_affected_entities(aws, args)

@inline describe_affected_entities(aws::AWSConfig, args) = AWSCore.Services.health(aws, "DescribeAffectedEntities", args)

@inline describe_affected_entities(args) = describe_affected_entities(default_aws_config(), args)

"""
    using AWSSDK.Health.describe_entity_aggregates
    describe_entity_aggregates([::AWSConfig], arguments::Dict)
    describe_entity_aggregates([::AWSConfig]; <keyword arguments>)

    using AWSCore.Services.health
    health([::AWSConfig], "DescribeEntityAggregates", arguments::Dict)
    health([::AWSConfig], "DescribeEntityAggregates", <keyword arguments>)

# DescribeEntityAggregates Operation

Returns the number of entities that are affected by each of the specified events. If no events are specified, the counts of all affected entities are returned.

# Arguments

## `eventArns = [::String, ...]`
A list of event ARNs (unique identifiers). For example: `"arn:aws:health:us-east-1::event/AWS_EC2_MAINTENANCE_5331", "arn:aws:health:us-west-1::event/AWS_EBS_LOST_VOLUME_xyz"`

# Returns

`DescribeEntityAggregatesResponse`

See also: [AWS API Documentation](https://docs.aws.amazon.com/goto/WebAPI/health-2016-08-04/DescribeEntityAggregates)
"""
@inline describe_entity_aggregates(aws::AWSConfig=default_aws_config(); args...) = describe_entity_aggregates(aws, args)

@inline describe_entity_aggregates(aws::AWSConfig, args) = AWSCore.Services.health(aws, "DescribeEntityAggregates", args)

@inline describe_entity_aggregates(args) = describe_entity_aggregates(default_aws_config(), args)

"""
    using AWSSDK.Health.describe_event_aggregates
    describe_event_aggregates([::AWSConfig], arguments::Dict)
    describe_event_aggregates([::AWSConfig]; aggregateField=, <keyword arguments>)

    using AWSCore.Services.health
    health([::AWSConfig], "DescribeEventAggregates", arguments::Dict)
    health([::AWSConfig], "DescribeEventAggregates", aggregateField=, <keyword arguments>)

# DescribeEventAggregates Operation

Returns the number of events of each event type (issue, scheduled change, and account notification). If no filter is specified, the counts of all events in each category are returned.

# Arguments

## `filter = [ ... ]`
Values to narrow the results returned.
```
 filter = [
        "eventArns" =>  [::String, ...],
        "eventTypeCodes" =>  [::String, ...],
        "services" =>  [::String, ...],
        "regions" =>  [::String, ...],
        "availabilityZones" =>  [::String, ...],
        "startTimes" =>  [[
            "from" =>  timestamp,
            "to" =>  timestamp
        ], ...],
        "endTimes" =>  [[
            "from" =>  timestamp,
            "to" =>  timestamp
        ], ...],
        "lastUpdatedTimes" =>  [[
            "from" =>  timestamp,
            "to" =>  timestamp
        ], ...],
        "entityArns" =>  [::String, ...],
        "entityValues" =>  [::String, ...],
        "eventTypeCategories" =>  ["issue", "accountNotification" or "scheduledChange", ...],
        "tags" =>  [::Dict{String,String}, ...],
        "eventStatusCodes" =>  ["open", "closed" or "upcoming", ...]
    ]
```

## `aggregateField = "eventTypeCategory"` -- *Required*
The only currently supported value is `eventTypeCategory`.

## `maxResults = ::Int`
The maximum number of items to return in one batch, between 10 and 100, inclusive.

## `nextToken = ::String`
If the results of a search are large, only a portion of the results are returned, and a `nextToken` pagination token is returned in the response. To retrieve the next batch of results, reissue the search request and include the returned token. When all results have been returned, the response does not contain a pagination token value.

# Returns

`DescribeEventAggregatesResponse`

# Exceptions

`InvalidPaginationToken`.

See also: [AWS API Documentation](https://docs.aws.amazon.com/goto/WebAPI/health-2016-08-04/DescribeEventAggregates)
"""
@inline describe_event_aggregates(aws::AWSConfig=default_aws_config(); args...) = describe_event_aggregates(aws, args)

@inline describe_event_aggregates(aws::AWSConfig, args) = AWSCore.Services.health(aws, "DescribeEventAggregates", args)

@inline describe_event_aggregates(args) = describe_event_aggregates(default_aws_config(), args)

"""
    using AWSSDK.Health.describe_event_details
    describe_event_details([::AWSConfig], arguments::Dict)
    describe_event_details([::AWSConfig]; eventArns=, <keyword arguments>)

    using AWSCore.Services.health
    health([::AWSConfig], "DescribeEventDetails", arguments::Dict)
    health([::AWSConfig], "DescribeEventDetails", eventArns=, <keyword arguments>)

# DescribeEventDetails Operation

Returns detailed information about one or more specified events. Information includes standard event data (region, service, etc., as returned by [DescribeEvents](@ref)), a detailed event description, and possible additional metadata that depends upon the nature of the event. Affected entities are not included; to retrieve those, use the [DescribeAffectedEntities](@ref) operation.

If a specified event cannot be retrieved, an error message is returned for that event.

# Arguments

## `eventArns = [::String, ...]` -- *Required*
A list of event ARNs (unique identifiers). For example: `"arn:aws:health:us-east-1::event/AWS_EC2_MAINTENANCE_5331", "arn:aws:health:us-west-1::event/AWS_EBS_LOST_VOLUME_xyz"`

## `locale = ::String`
The locale (language) to return information in. English (en) is the default and the only supported value at this time.

# Returns

`DescribeEventDetailsResponse`

# Exceptions

`UnsupportedLocale`.

See also: [AWS API Documentation](https://docs.aws.amazon.com/goto/WebAPI/health-2016-08-04/DescribeEventDetails)
"""
@inline describe_event_details(aws::AWSConfig=default_aws_config(); args...) = describe_event_details(aws, args)

@inline describe_event_details(aws::AWSConfig, args) = AWSCore.Services.health(aws, "DescribeEventDetails", args)

@inline describe_event_details(args) = describe_event_details(default_aws_config(), args)

"""
    using AWSSDK.Health.describe_event_types
    describe_event_types([::AWSConfig], arguments::Dict)
    describe_event_types([::AWSConfig]; <keyword arguments>)

    using AWSCore.Services.health
    health([::AWSConfig], "DescribeEventTypes", arguments::Dict)
    health([::AWSConfig], "DescribeEventTypes", <keyword arguments>)

# DescribeEventTypes Operation

Returns the event types that meet the specified filter criteria. If no filter criteria are specified, all event types are returned, in no particular order.

# Arguments

## `filter = [ ... ]`
Values to narrow the results returned.
```
 filter = [
        "eventTypeCodes" =>  [::String, ...],
        "services" =>  [::String, ...],
        "eventTypeCategories" =>  ["issue", "accountNotification" or "scheduledChange", ...]
    ]
```

## `locale = ::String`
The locale (language) to return information in. English (en) is the default and the only supported value at this time.

## `nextToken = ::String`
If the results of a search are large, only a portion of the results are returned, and a `nextToken` pagination token is returned in the response. To retrieve the next batch of results, reissue the search request and include the returned token. When all results have been returned, the response does not contain a pagination token value.

## `maxResults = ::Int`
The maximum number of items to return in one batch, between 10 and 100, inclusive.

# Returns

`DescribeEventTypesResponse`

# Exceptions

`InvalidPaginationToken` or `UnsupportedLocale`.

See also: [AWS API Documentation](https://docs.aws.amazon.com/goto/WebAPI/health-2016-08-04/DescribeEventTypes)
"""
@inline describe_event_types(aws::AWSConfig=default_aws_config(); args...) = describe_event_types(aws, args)

@inline describe_event_types(aws::AWSConfig, args) = AWSCore.Services.health(aws, "DescribeEventTypes", args)

@inline describe_event_types(args) = describe_event_types(default_aws_config(), args)

"""
    using AWSSDK.Health.describe_events
    describe_events([::AWSConfig], arguments::Dict)
    describe_events([::AWSConfig]; <keyword arguments>)

    using AWSCore.Services.health
    health([::AWSConfig], "DescribeEvents", arguments::Dict)
    health([::AWSConfig], "DescribeEvents", <keyword arguments>)

# DescribeEvents Operation

Returns information about events that meet the specified filter criteria. Events are returned in a summary form and do not include the detailed description, any additional metadata that depends on the event type, or any affected resources. To retrieve that information, use the [DescribeEventDetails](@ref) and [DescribeAffectedEntities](@ref) operations.

If no filter criteria are specified, all events are returned. Results are sorted by `lastModifiedTime`, starting with the most recent.

# Arguments

## `filter = [ ... ]`
Values to narrow the results returned.
```
 filter = [
        "eventArns" =>  [::String, ...],
        "eventTypeCodes" =>  [::String, ...],
        "services" =>  [::String, ...],
        "regions" =>  [::String, ...],
        "availabilityZones" =>  [::String, ...],
        "startTimes" =>  [[
            "from" =>  timestamp,
            "to" =>  timestamp
        ], ...],
        "endTimes" =>  [[
            "from" =>  timestamp,
            "to" =>  timestamp
        ], ...],
        "lastUpdatedTimes" =>  [[
            "from" =>  timestamp,
            "to" =>  timestamp
        ], ...],
        "entityArns" =>  [::String, ...],
        "entityValues" =>  [::String, ...],
        "eventTypeCategories" =>  ["issue", "accountNotification" or "scheduledChange", ...],
        "tags" =>  [::Dict{String,String}, ...],
        "eventStatusCodes" =>  ["open", "closed" or "upcoming", ...]
    ]
```

## `nextToken = ::String`
If the results of a search are large, only a portion of the results are returned, and a `nextToken` pagination token is returned in the response. To retrieve the next batch of results, reissue the search request and include the returned token. When all results have been returned, the response does not contain a pagination token value.

## `maxResults = ::Int`
The maximum number of items to return in one batch, between 10 and 100, inclusive.

## `locale = ::String`
The locale (language) to return information in. English (en) is the default and the only supported value at this time.

# Returns

`DescribeEventsResponse`

# Exceptions

`InvalidPaginationToken` or `UnsupportedLocale`.

See also: [AWS API Documentation](https://docs.aws.amazon.com/goto/WebAPI/health-2016-08-04/DescribeEvents)
"""
@inline describe_events(aws::AWSConfig=default_aws_config(); args...) = describe_events(aws, args)

@inline describe_events(aws::AWSConfig, args) = AWSCore.Services.health(aws, "DescribeEvents", args)

@inline describe_events(args) = describe_events(default_aws_config(), args)

end # module Health

#==============================================================================#
# End of file
#==============================================================================#


text:
@info "Compiling packages for running model..."
#meta
using TickTock
using ConfParser
using Logging
using Profile
#sim
using Dates
using Random
using StaticArrays
using Distributed
using OrdinaryDiffEq
using JLD2
using Plots
@info "Packages compiled, running model"

#######################
## Constants n stuff ##
#######################

const Re   = 6370e3;        # Earth radius, f64
const c    = 3e8;           # speedo lite, f64
const Beq  = 3.e-5;         # B field at equator (T), f64

# Parsing the conf file
conffile = "setup.conf";
conf = ConfParse(conffile)
parse_conf!(conf)
const basename = retrieve(conf, "basename");
const directoryname = retrieve(conf, "directoryname");
numParticles = parse(Int64, retrieve(conf, "numberOfParticles"));
startTime = parse(Float64, retrieve(conf, "startTime"));
endTime = parse(Float64, retrieve(conf, "endTime"));
tspan = (startTime, endTime); # integration time
const lossConeAngle = parse(Float64, retrieve(conf, "lossConeAngle"));
const saveDecimation = parse(Float64, retrieve(conf, "saveDecimation"));
const L = parse(Float64, retrieve(conf, "L"));
const omegam = parse(Float64, retrieve(conf, "omegam"));
const Omegape = parse(Float64, retrieve(conf, "Omegape"));
const z0 = parse(Float64, retrieve(conf, "z0"));
const λ0 = parse(Float64, retrieve(conf, "lambda0"));
const dλ1 = parse(Float64, retrieve(conf, "dLambda1"));
const dλ2 = parse(Float64, retrieve(conf, "dLambda2"));
const a = parse(Float64, retrieve(conf, "a"));
const dPhi = parse(Float64, retrieve(conf, "dPhi"));
const waveAmplitudeModifier = parse(Float64, retrieve(conf, "waveAmplitudeModifier"));
const ELo = parse(Float64, retrieve(conf, "ELo"));
const EHi = parse(Float64, retrieve(conf, "EHi"));
const Esteps = parse(Float64, retrieve(conf, "Esteps"));
const PALo = parse(Float64, retrieve(conf, "PALo"));
const PAHi = parse(Float64, retrieve(conf, "PAHi"));
const PAsteps = parse(Float64, retrieve(conf, "PAsteps"));
ICrange = [ELo, EHi, Esteps, PALo, PAHi, PAsteps];
const batches = parse(Int64, retrieve(conf, "batches"));
const numThreads = parse(Int64, retrieve(conf, "numberOfThreads"))
const B_eq_measured = parse(Float64, retrieve(conf, "B_eq_measured")); # in nT
u(lambda) = tanh((deg2rad(lambda)/(deg2rad(2)))) * (exp(-(deg2rad(lambda)/(deg2rad(dλ2)))^2));
const B_w_normalizer = maximum(u.(1:.01:90))^-1
@info "Parsed Config file: $conffile:"

###################
## Setup Logging ##
###################

function setupDirectories(directoryname::String)
    #=
    directory name
     └ jld2_yymmdd_HH
      | setupasrun.conf
      | yymmdd_HHMMSS.log
      └ basename_numparticles_numbatch.jld2
    =#
    mkpath(directoryname)
    outputFileDirectory = string(directoryname, "/jld2_", Dates.format(now(), DateFormat("yymmdd_HH")))
    mkpath(outputFileDirectory)
    outputFileBaseName = outputFileDirectory*"/"*basename*"_$numParticles";
    loggingFileName = string(outputFileDirectory,"/",Dates.format(now(), DateFormat("yymmdd_HHMMSS")),".log")
    asrunConfFileName = outputFileDirectory*"/setupasrun.conf"
    run(`cp setup.conf $asrunConfFileName`) # copy setup over to as run file
    run(`cp setup.conf $loggingFileName`) # copy setup over to log file

    @info "View logs here: "*loggingFileName
    io = open(loggingFileName, "a")
    global_logger(ConsoleLogger(io))
    @info "^^Used this setup configuration file^^^"
    flush(io)

    return outputFileBaseName, io
end

####################
## Initialization ##
####################

function generateFlatParticleDistribution(numParticles::Int64, ICrange, z0=0::Float64, λ0=0::Float64)
    ELo, EHi, Esteps, PALo, PAHi, PAsteps = ICrange
    @info "Generating a flat particle distribution with"
    @info "$Esteps steps of energy from $ELo KeV to $EHi KeV"
    @info "$PAsteps steps of pitch angles from $PALo deg to $PAHi deg"
    flush(io)

    nBins = PAsteps*Esteps
    N = numParticles ÷ nBins # num of particles per bin
    E_bins = logrange(ELo,EHi, Int64(Esteps))
    PA_bins = range(PALo, PAHi, length = Int64(PAsteps))
    @info "Flat distribution with $N particles/bin in $nBins bins"
    flush(io)

    if numParticles%nBins != 0
        @warn "Truncating $(numParticles%nBins) particles for an even distribution"
        flush(io)
    end
    if iszero(N)
        N = 1;
        @warn "Use higher number of particles next time. Simulating 1 trajectory/bin."
        @warn "Minimum number of particles to simulate is 1 particles/bin."
        flush(io)
    end

    @views f0 = [[(E+511.)/511. deg2rad(PA)] for PA in PA_bins for E in E_bins for i in 1:N] # creates a 2xN array with initial PA and Energy

    #####       [[z0 pz0                          ζ0               mu0                          λ0 Φ0              ]             ]
    @views h0 = [[z0 sqrt(IC[1]^2 - 1)*cos(IC[2]) rand()*2*pi*dPhi .5*(IC[1]^2-1)*sin(IC[2])^2  λ0 rand()*2*pi*dPhi] for IC in f0] # creates a 5xN array with inital h0 terms
    f0 = vcat(f0...) # convert Array{Array{Float64,2},1} to Array{Float64,2}
    h0 = vcat(h0...) # since i used list comprehension it is now a nested list

    # Other ICs that are important
    # Define basic ICs and parameters
    # B0          = Beq*sqrt(1. +3. *sin(λ0)^2.)/L^3.;     # starting B field at eq
    B0          = B_eq_measured*1e-9                        # measured equatorial field strength
    Omegace0    = (1.6e-19*B0)/(9.11e-31);                    # electron gyrofreq @ the equator
    η           = Omegace0*L*Re/c;              # should be like 10^3
    ε           = waveAmplitudeModifier/η;    # normalized wave large amplitude, .1 for small, 15 for large
    resolution  = .1/η;                       # determines max step size of the integrator
    @info "Min integration step of $resolution"
    @info "Created Initial Conditions for $(length(h0[:,1])) particles"
    flush(io)

    return h0, f0, η, ε, resolution;
end

@everywhere function generateModifiableFunction(batches)
    #=
    Takes in the initial condition and splits them into batches.
    This way, we can feed each one in and get a percent completeness during sim.
    =#
    probGeneratorList = [];
    nPerBatch = numParticles÷batches;
    for j in 0:batches-1
        truncatedIC =  h0[nPerBatch*j+1:(nPerBatch*j+nPerBatch),:]
        push!(probGeneratorList, ((prob,i,repeat) -> remake(prob, u0 = truncatedIC[i,:], p = @SVector [η, ε, Omegape, omegam, a, dPhi, dλ1, dλ2, B_w_normalizer])))
    end
    percentage = (round(100/batches))
    @info "Each batch will simulate $nPerBatch particles for $(endTime-startTime) dt and correspond with $percentage%"
    flush(io)
    return probGeneratorList, nPerBatch, percentage
end

##################
## Math n stuff ##
##################

function eom!(dH,H,p::SVector{9, Float64},t::Float64)
    # These equations define the motion.

    # z, pz, zeta, mu, lambda, phi = H
    # eta, epsilon, Omegape, omegam, a, dPhi, dLambda1, dLambda2 = p
    sinλ = sin(H[5]);
    cosλ = cos(H[5]);
    g = exp(-p[5] * (cos(H[6]/(2*π*p[6]))^2)) +  exp(-p[5] * (sin(H[6]/(2*π*p[6]))^2))
    sinζ = g*sin(H[3]);
    cosζ = g*cos(H[3]);

    # double sided wave, grows to max at dLambda1 deg, dissipates by dLambda2 deg
    u = p[9]*tanh((H[5]/(deg2rad(p[7])))) * (exp(-(H[5]/(deg2rad(p[8])))^2));

    # helper variables
    b = sqrt(1+3*sinλ^2)/(cosλ^6);
    db = (3*(27*sinλ-5*sin(3*H[5])))/(cosλ^8*(4+12*sinλ^2));
    γ = sqrt(1 + H[2]^2 + 2*H[4]*b);
    if H[5] < 0
        K = -1 * (p[3] * (cosλ^(-5/2)))/sqrt(b/p[4] - 1);
    else
        K = (p[3] * (cosλ^(-5/2)))/sqrt(b/p[4] - 1);
    end
    psi = p[1]*p[2]*u*sqrt(2*H[4]*b)/γ;

    # actual integration vars
    dH1 = H[2]/γ;
    dH2 = -(H[4]*db)/γ - (psi*cosζ);
    dH3 = p[1]*(K*dH1 - p[4] + b/γ) + (psi*sinζ)/(2*H[4]*K);
    dH4 = -(psi*cosζ)/K;
    dH5 = H[2]/(γ*cosλ*sqrt(1+3*sinλ^2));
    dH6 = p[1]*(K*dH1 - p[4]);

    dH .= SizedVector{6}([ dH1, dH2, dH3, dH4, dH5, dH6 ]);
end

function palostcondition(H,t,integrator)
    # condition: if particle enters loss cone
    b = sqrt(1+3*sin(H[5])^2)/(cos(H[5])^6);
    γ = sqrt(1 + H[2]^2 + 2*H[4]*b);
    return (rad2deg(asin(sqrt((2*H[4])/(γ^2 -1))))) < (lossConeAngle)
end

function ixlostcondition(H,t,integrator)
    # condition: if I_x approaches 0
    b = sqrt(1+3*sin(H[5])^2)/(cos(H[5])^6);
    return 2*H[4]*b < 10*resolution
end

affect!(integrator) = terminate!(integrator); # terminate if condition reached
cb1 = DiscreteCallback(palostcondition,affect!);
cb2 = DiscreteCallback(ixlostcondition,affect!);

function ensemble()
    for i in 1:batches
        ensemble_prob = EnsembleProblem(prob::ODEProblem,prob_func=probGeneratorList[i])
        tick()
        flush(io)
        sol = solve(ensemble_prob, Tsit5(), EnsembleThreads(), save_everystep=false;
                            callback=CallbackSet(cb1, cb2), trajectories=nPerBatch,
                            dtmax=resolution, linear_solver=:LapackDense, maxiters=1e8,
                            saveat = saveDecimation*resolution)
        @save outputFileBaseName*"_$(i).jld2" sol
        @info "$nPerBatch particles simulated in:"
        tock()
        @info "$(i*percentage)% complete..."
        flush(io)
    end
end

# Simple calcs
calcb(b::Vector{Float64},λ::Vector{Float64}) = @. b = sqrt(1+3*sin(λ)^2)/(cos(λ)^6)
calcdb(db::Vector{Float64},λ::Vector{Float64}) = @. db = (3*(27*sin(λ)-5*sin(3*λ)))/(cos(λ)^8*(4+12*sin(λ)^2))
calcGamma(γ::Vector{Float64},pz::Vector{Float64},μ::Vector{Float64},b::Vector{Float64}) = @. γ = sqrt(1 + pz^2 + 2*μ*b)
calcK(K::Vector{Float64},b::Vector{Float64},λ::Vector{Float64}) = @. K = (Omegape * (cos(λ)^-(5/2)))/sqrt(b/omegam - 1)
calcAlpha(α::Vector{Float64},μ::Vector{Float64}, γ::Vector{Float64}) = @. α = rad2deg(asin(sqrt((2*μ)/(γ^2 - 1))))

# Useful helpers
logrange(x1, x2, n::Int64) = [10^y for y in range(log10(x1), log10(x2), length=n)]

# ```
# Useful links
# Community stuff: https://discourse.julialang.org/c/domain/models/21
# Docs for ensemble stuff: https://diffeq.sciml.ai/dev/features/ensemble/
# Marginal Histograms: http://docs.juliaplots.org/latest/recipes/#Marginal-Histograms
# Dumb colorbar ticks issue: https://github.com/JuliaPlots/Plots.jl/issues/2308
# ```

ext: jl

max_stars_repo_name: ethantsai/nlwhistlers

max_issues_repo_name: ethantsai/nlwhistlers

max_forks_repo_name: ethantsai/nlwhistlers


text:
module Transforms
export  dwt, idwt, dwt!, idwt!,
        wpt, iwpt, wpt!, iwpt!,
        modwt, imodwt, cwt
using ..Util, ..WT
using Compat: AbstractRange, copyto!, undef
using FFTW

# TODO Use StridedArray instead of AbstractArray where writing to array.
# TODO change integer dependent wavelets to parametric types (see "Value types", https://docs.julialang.org/en/v1/manual/types/index.html#%22Value-types%22-1)
const DWTArray = AbstractArray
const WPTArray = AbstractVector
const ValueType = Union{AbstractFloat, Complex}

# DWT

"""
    dwt(x, wt[, L=maxtransformlevels(x)])

Perform a discrete wavelet transform of the array `x`.
The wavelet type `wt` determines the transform type
(filter or lifting) and the wavelet class, see `wavelet`.

The number of transformation levels `L` can be any non-negative
integer such that the size of `x` is divisible by `L`.
Performs the identity transformation if `L==0`.

# Examples
```julia
dwt(x, wavelet(WT.coif6))
```

**See also:** `idwt`, `dwt!`, `wavelet`
"""
function dwt end

"""
    idwt(x, wt[, L=maxtransformlevels(x)])

Perform an inverse discrete wavelet transform of the array `x`,
the inverse of `dwt(x, wt, L)`.

**See also:** `dwt`, `idwt!`
"""
function idwt end

"""
    dwt!(y, x, wt::OrthoFilter[, L=maxtransformlevels(x)])

    dwt!(y, wt::GLS[, L=maxtransformlevels(x)])

Same as `dwt` but without array allocation.
Perform "out of place" transform with a filter, or
a inplace transform with a lifting scheme. The difference
between the filter and lifting methods is due to the
structure of the transform algorithms.

**See also:** `idwt!`
"""
function dwt! end

"""
    idwt!(y, x, wt::OrthoFilter[, L=maxtransformlevels(x)])

    idwt!(y, wt::GLS[, L=maxtransformlevels(x)])

The inverse of `dwt!`.

**See also:** `dwt!`
"""
function idwt! end

# WPT

"""
    wpt

Perform a discrete wavelet packet transform of the array `x`.
**See also:** `dwt`, `wavelet`
"""
function wpt end

"""
    iwpt

Perform an inverse discrete wavelet packet transform of the array `x`.
**See also:** `idwt`, `wavelet`
"""
function iwpt end

"""
    wpt!

Same as `wpt` but without array allocation.
"""
function wpt! end

"""
    iwpt!

Same as `iwpt` but without array allocation.
"""
function iwpt! end

# DWT (discrete wavelet transform)

for (Xwt, Xwt!, _Xwt!, fw) in ((:dwt, :dwt!, :_dwt!, true),
                                (:idwt, :idwt!, :_dwt!, false))
@eval begin
    # filter
    function ($Xwt)(x::DWTArray{T}, filter::OrthoFilter,
                    L::Integer=maxtransformlevels(x)) where T<:ValueType
        y = Array{T}(undef, size(x))
        return ($_Xwt!)(y, x, filter, L, $fw)
    end
    function ($Xwt!)(y::DWTArray{T}, x::DWTArray{T}, filter::OrthoFilter,
                    L::Integer=maxtransformlevels(x)) where T<:ValueType
        return ($_Xwt!)(y, x, filter, L, $fw)
    end
    # lifting
    function ($Xwt)(x::DWTArray{T}, scheme::GLS,
                    L::Integer=maxtransformlevels(x)) where T<:ValueType
        y = Array{T}(undef, size(x))
        copyto!(y, x)
        return ($_Xwt!)(y, scheme, L, $fw)
    end
    function ($Xwt!)(y::DWTArray{T}, scheme::GLS,
                    L::Integer=maxtransformlevels(y)) where T<:ValueType
        return ($_Xwt!)(y, scheme, L, $fw)
    end
end # begin
end # for

# CWT (continuous wavelet transform)
# cwt(Y::AbstractVector, ::ContinuousWavelet)

"""
wave = cwt(Y::AbstractArray{T}, c::CFW{W}; J1::Int64=-1, dt::S=NaN, s0::V=NaN) where {T<:Real, S<:Real, V<:Real, W<:WT.WaveletBoundary}

return the continuous wavelet transform wave, which is (nscales)×(signalLength), of type c of Y. J1 is the total number of scales; default (when J1=NaN, or is negative) is just under the maximum possible number, i.e. the log base 2 of the length of the signal, times the number of wavelets per octave. If you have sampling information, you will need to scale wave by δt^(1/2).
"""
function cwt(Y::AbstractArray{T}, c::CFW{W}; J1::Int64=-1, dt::S=NaN, s0::V=NaN) where {T<:Real, S<:Real, V<:Real, W<:WT.WaveletBoundary}
    n1 = length(Y);
    # don't alter scaling with sampling information if it doesn't exists
    fλ = (4*π) / (c.σ[1] + sqrt(2 + c.σ[1]^2))
    if isnan(dt) || (dt<0)
        dt = 1
    end
    # smallest resolvable scale
    if isnan(s0) || (s0<0)
        s0 = 2 * dt / fλ
    end
    # J1 is the total number of scales
    if J1<0
        J1 = Int(round(log2(n1 * dt / s0) * c.scalingFactor))
    end

    # scales from Mallat 1999
    sj = s0 * 2.0.^(collect(0:J1)./c.scalingFactor)
    # Fourier equivalent frequencies
    freqs = 1 ./ (fλ .* sj)

    #....construct time series to analyze, pad if necessary
    if eltypes(c) == WT.padded
        base2 = round(Int,log(n1)/log(2));   # power of 2 nearest to N
        x = [Y; zeros(2^(base2+1)-n1)];
    elseif eltypes(c) == WT.DEFAULT_BOUNDARY
        x = [Y; flipdim(Y,1)]
    else
        x= Y
    end

    n = length(x);

    #....construct wavenumber array used in transform [Eqn(5)]
    ω = fftfreq(n, 1/dt)*2π

    #....compute FFT of the (padded) time series
    x̂ = fft(x);    # [Eqn(3)]

    wave = zeros(Complex{T}, J1+1, n);  # define the wavelet array
    # iterate over scales
    for scale=1:J1+1
        # define daughter wavelet (in frequency domain) for each scale
        daughter = sqrt.(sj[scale] .* ω[2] .* n) .* conj.(π^(-1/4).*exp.(-0.5 * ((sj[scale] .* ω) .- c.σ[1]).^2))
        # convolve in frequency domain, perform ifft
        wave[scale, :] = ifft(x̂ .* daughter, 1)
    end

    # Determines the cone-of-influence. Note that it is returned as a function
    # of time in Fourier periods. Uses triangualr Bartlett window with
    # non-zero end-points.
    coi = (n1 / 2 .- abs.(collect(0:n1-1) .- (n1 - 1) ./ 2))
    coi = (fλ * dt / sqrt(2)).*coi

    return wave, sj, freqs, coi
end
"""
period,scale, coi = caveats(Y::AbstractArray{T}, c::CFW{W}; J1::S=NaN) where {T<:Real, S<:Real, W<:WT.WaveletBoundary}

returns the period, the scales, and the cone of influence for the given wavelet transform. If you have sampling information, you will need to scale the vector scale appropriately by 1/δt, and the actual transform by δt^(1/2).
"""
function caveats(Y::AbstractArray{T}, c::CFW{W}; J1::S=NaN) where {T<:Real, S<:Real, W<:WT.WaveletBoundary}
    n1 = length(Y);
    # J1 is the total number of elements
    if isnan(J1) || (J1<0)
        J1=floor(Int,(log2(n1))*c.scalingFactor);
    end
    println("$(J1 == NaN)")
    #....construct time series to analyze, pad if necessary
    if eltypes(c) == WT.ZPBoundary
        base2 = round(Int,log(n1)/log(2));   # power of 2 nearest to N
        n = length(Y)+2^(base2+1)-n1
    elseif eltypes(c) == WT.PerBoundary
        n = length(Y)*2
    end
    ω = [0:floor(Int, n/2); -floor(Int,n/2)+1:-1]*2π
    period = c.fourierFactor*2 .^((0:J1)/c.scalingFactor)
    scale = [1E-5; 1:((n1+1)/2-1); reverse((1:(n1/2-1)),dims=1); 1E-5]
    coi = c.coi*scale  # COI [Sec.3g]
    return period, scale, coi
end
cwt(Y::AbstractArray{T}, w::WT.ContinuousWaveletClass; J1::Int64=-1, dt::S=NaN, s0::V=NaN) where {T<:Real, S<:Real, V<:Real} = cwt(Y,CFW(w),J1=J1,dt=dt,s0=s0)
caveats(Y::AbstractArray{T}, w::WT.ContinuousWaveletClass; J1::S=NaN) where {T<: Real, S<: Real} = caveats(Y,CFW(w),J1=J1)
cwt(Y::AbstractArray{T}) where T<:Real = cwt(Y,WT.Morlet())
caveats(Y::AbstractArray{T}) where T<:Real = caveats(Y,WT.Morlet())

"""
icwt(WT::AbstractArray{T}, c::CFW{W}, sj::AbstractArray; dt::S=NaN, dj::V=1/12) where {T<:Complex{Real}, S<:Real, V<:Real, W<:WT.WaveletBoundary}

return the inverse continuous wavelet transform
"""
function icwt(WT::AbstractArray, c::CFW{W}, sj::AbstractArray; dt::S=NaN, dj::V=1/12) where {S<:Real, V<:Real, W<:WT.WaveletBoundary}
    # Torrence and Compo (1998), eq. (11)
    iW = (dj * sqrt(dt) / 0.776 * psi(c, 0)) .* sum((real.(WT) ./ sqrt.(sj)), dims=1)

    return iW
end
cwt(Y::AbstractArray{T}, w::WT.ContinuousWaveletClass, sj::AbstractArray; dt::S=NaN, dj::V=NaN) where {T<:Real, S<:Real, V<:Real} = cwt(Y,CFW(w),sj,dt=dt,dj=dj)

function psi(c::CFW{W}, t::Int64) where W<:WT.WaveletBoundary
    return π^(-0.25) * exp.(im*c.σ[1]*t - t^2 / 2)
end

# WPT (wavelet packet transform)

for (Xwt, Xwt!, _Xwt!, fw) in ((:wpt, :wpt!, :_wpt!, true),
                                (:iwpt, :iwpt!, :_wpt!, false))
@eval begin
    function ($Xwt)(x::WPTArray{T}, wt::DiscreteWavelet,
                    L::Integer=maxtransformlevels(x)) where T<:ValueType
        return ($Xwt)(x, wt, maketree(length(x), L, :full))
    end
    # filter
    function ($Xwt)(x::WPTArray{T}, filter::OrthoFilter,
                    tree::BitVector=maketree(x, :full)) where T<:ValueType
        y = Array{T}(undef, size(x))
        return ($_Xwt!)(y, x, filter, tree, $fw)
    end
    function ($Xwt!)(y::WPTArray{T}, x::WPTArray{T},
                    filter::OrthoFilter) where T<:ValueType
        return ($Xwt!)(y, x, filter, maketree(x, :full))
    end
    function ($Xwt!)(y::WPTArray{T}, x::WPTArray{T},
                    filter::OrthoFilter, L::Integer) where T<:ValueType
        return ($Xwt!)(y, x, filter, maketree(length(x), L, :full))
    end
    function ($Xwt!)(y::WPTArray{T}, x::WPTArray{T},
                    filter::OrthoFilter, tree::BitVector) where T<:ValueType
        return ($_Xwt!)(y, x, filter, tree, $fw)
    end
    # lifting
    function ($Xwt)(x::WPTArray{T}, scheme::GLS,
                    tree::BitVector=maketree(x, :full)) where T<:ValueType
        y = Array{T}(undef, size(x))
        copyto!(y, x)
        return ($_Xwt!)(y, scheme, tree, $fw)
    end
    function ($Xwt!)(y::WPTArray{T}, scheme::GLS) where T<:ValueType
        return ($Xwt!)(y, scheme, maketree(y, :full))
    end
    function ($Xwt!)(y::WPTArray{T}, scheme::GLS, L::Integer) where T<:ValueType
        return ($Xwt!)(y, scheme, maketree(length(x), L, :full))
    end
    function ($Xwt!)(y::WPTArray{T}, scheme::GLS, tree::BitVector) where T<:ValueType
        return ($_Xwt!)(y, scheme, tree, $fw)
    end
end # begin
end # for

# DWTC (column-wise discrete wavelet transform)
#dwtc(::AbstractArray, ::DiscreteWavelet)
#idwtc(::AbstractArray, ::DiscreteWavelet)

# SWT (stationary wavelet transform)
#swt(::AbstractVector, ::DiscreteWavelet)
#iswt(::AbstractVector, ::DiscreteWavelet)

# CWT (continuous wavelet transform directly) TODO: direct if sufficiently small

# TODO: continuous inverse, when defined
#icwt(::AbstractVector, ::ContinuousWavelet)

# Int -> Float
for Xwt in (:dwt, :idwt, :dwtc, :idwtc, :wpt, :iwpt)
    @eval $Xwt(x::AbstractArray{<:Integer}, args...) = $Xwt(float(x), args...)
end

# non-exported "out of place" functions
for (Xwt_oop!, Xwt!) in ((:dwt_oop!, :dwt!), (:idwt_oop!, :idwt!))
@eval begin
    # filter
    function ($Xwt_oop!)(y::DWTArray{T}, x::DWTArray{T}, filter::OrthoFilter,
                        L::Integer=maxtransformlevels(x)) where T<:ValueType
        return ($Xwt!)(y, x, filter, L)
    end
    # lifting
    function ($Xwt_oop!)(y::DWTArray{T}, x::DWTArray{T}, scheme::GLS,
                        L::Integer=maxtransformlevels(x)) where T<:ValueType
        copyto!(y, x)
        return ($Xwt!)(y, scheme, L)
    end
end # begin
end # for

# Array with shared memory
function unsafe_vectorslice(A::Array{T}, i::Int, n::Int) where T
    return unsafe_wrap(Array, pointer(A, i), n)::Vector{T}
end

# linear indices of start of rows/cols/planes
# 2-D, size(A) = (m,n)
row_idx(i, m) = i
col_idx(i, m) = 1 + (i-1)*m
# 3-D, size(A) = (m,n,d)
row_idx(i, j, m, n=m) = row_idx(i, n) + (j-1)*n*m
col_idx(i, j, m, n=m) = col_idx(i, m) + (j-1)*n*m
plane_idx(i, j, m) = i + (j-1)*m

# filter transforms
include("transforms_filter.jl")

# lifting transforms
include("transforms_lifting.jl")

include("transforms_maximal_overlap.jl")

end # module


text:
"""
    nsol(F!, x0, FS, FPS, J!=diffjac!; rtol=1.e-6, atol=1.e-12,
               maxit=20, solver="newton", sham=5, armmax=10, resdec=.1,
               dx = 1.e-7, armfix=false,
               pdata = nothing, jfact = klfact,
               printerr = true, keepsolhist = false, stagnationok=false)
)

C. T. Kelley, 2020

Julia versions of the nonlinear solvers from my SIAM books.
Herewith: nsol

You must allocate storage for the function and Jacobian in advance
--> in the calling program <-- ie. in FS and FPS

Inputs:\n
- F!: function evaluation, the ! indicates that F! overwrites FS, your
    preallocated storage for the function.\n
    So FS=F!(FS,x) or FS=F!(FS,x,pdata) returns FS=F(x)

- x0: initial iterate\n

- FS: Preallocated storage for function. It is an N x 1 column vector\n
  You may store it as (n,) or (n,1), depending on what F! likes to see.

- FPS: preallocated storage for Jacobian. It is an N x N matrix\n

- J!: Jacobian evaluation, the ! indicates that J! overwrites FPS, your
    preallocated storage for the Jacobian. If you leave this out the
    default is a finite difference Jacobian.\n
    So, FP=J!(FP,FS,x) or FP=J!(FP,FS,x,pdata) returns FP=F'(x). \n
    (FP,FS, x) must be the argument list, even if FP does not need FS.
    One reason for this is that the finite-difference Jacobian
    does and that is the default in the solver.

- Precision: Lemme tell ya 'bout precision. I designed this code for
    full precision functions and linear algebra in any precision you want.
    You can declare
    FPS as Float64, Float32, or Float16 and nsol will do the right thing if
    YOU do not destroy the declaration in your J! function. I'm amazed
    that this works so easily. If the Jacobian is reasonably well
    conditioned, you can cut the cost of Jacobian factorization and
    storage in half with no loss. For large dense Jacobians and inexpensive
    functions, this is a good deal.\n
    BUT ... There is very limited support for direct sparse solvers in
    anything other than Float64. I recommend that you only use Float64
    with direct sparse solvers unless you really know what you're doing. I
    have a couple examples in the notebook, but watch out.

----------------------

Keyword Arguments (kwargs):\n

rtol and atol: relative and absolute error tolerances\n

maxit: limit on nonlinear iterations\n

solver: default = "newton"\n
Your choices are "newton" or "chord". However,
you have sham at your disposal only if you chose newton. "chord"
will keep using the initial derivative until the iterate converges,
uses the iteration budget, or the line search fails. It is not the
same as sham=Inf, which is smarter.\n

sham: default = 5 (ie Newton)\n
This is the Shamanskii method. If sham=1, you have Newton.
The iteration updates the derivative every sham iterations.
The convergence rate has local q-order sham+1 if you only count
iterations where you update the derivative. You need not
provide your own derivative function to use this option. sham=Inf
is chord only if chord is converging well.\n

I made sham=1 the default for scalar equations. For systems I'm
more aggressive and want to invest as little energy in linear algebra
as possible. So the default is sham=5.

armmax: upper bound on step size reductions in line search\n

resdec: default = .1\n
This is the target value for residual reduction.
The default value is .1. In the old MATLAB codes it was .5.
I only turn Shamanskii on if the residuals are decreasing
rapidly, at least a factor of resdec, and the line search is quiescent.
If you want to eliminate resdec from the method ( you don't ) then set
resdec = 1.0 and you will never hear from it again.

dx: default = 1.e-7\n
difference increment in finite-difference derivatives
      h=dx*norm(x,Inf)+1.e-8

armfix: default = false\n
The default is a parabolic line search (ie false). Set to true and
the step size will be fixed at .5. Don't do this unless you are doing
experiments for research.\n

pdata:\n
precomputed data for the function/Jacobian.
Things will go better if you use this rather than hide the data
in global variables within the module for your function/Jacobian

If you use pdata in either of F! or J!, you must use in in the
calling sequence of both.

jfact: default = klfact (tries to figure out best choice) \n
If your Jacobian has any special structure, please set jfact
to the correct choice for a factorization.

I use jfact when I call PrepareJac! to evaluate the
Jacobian (using your J!) and factor it. The default is to use
klfact (an internal function) to do something reasonable.
For general matrices, klfact picks lu! to compute an LU factorization
and share storage with the Jacobian.  You may change LU to something else by,
for example, setting jfact = cholseky! if your Jacobian is spd.

klfact knows about banded matrices and picks qr. You should,
however RTFM, allocate the extra two upper bands, and use jfact=qr!
to override klfact.

If you give me something that klfact does not know how to dispatch on,
then nothing happens. I just return the original Jacobian matrix and
nsol will use backslash to compute the Newton step.
I know that this is probably not optimal in your situation, so it is
good to pick something else, like jfact = lu.

If you want to manage your own factorization within your Jacobian
evaluation function, then set\n
jfact = nofact\n
and nsol will not attempt to factor your Jacobian. That is also what
happens when klfact does not know what to do. Your Jacobian is sent
directly to Julia's \\  operation

Please do not mess with the line that calls PrepareJac!.

        FPF = PrepareJac!(FPS, FS, x, ItRules)

FPF is not the same as FPS (the storage you allocate for the Jacobian)
for a reason. FPF and FPS do not have the same type, even though they
share storage. So, FPS=PrepareJac!(FPS, FS, ...) will break things.

printerr: default = true\n
I print a helpful message when the solver fails. To suppress that
message set printerr to false.

keepsolhist: default = false\n
Set this to true to get the history of the iteration in the output
tuple. This is on by default for scalar equations and off for systems.
Only turn it on if you have use for the data, which can get REALLY LARGE.

stagnationok: default = false\n
Set this to true if you want to disable the line search and either
observe divergence or stagnation. This is only useful for research
or writing a book.

Output:\n
- A named tuple (solution, functionval, history, stats, idid,
               errcode, solhist)
where

   -- solution = converged result

   -- functionval = F(solution)

   -- history = the vector of residual norms (||F(x)||) for the iteration

   -- stats = named tuple of the history of (ifun, ijac, iarm), the number
of functions/derivatives/steplength reductions at each iteration.

I do not count the function values for a finite-difference derivative
because they count toward a Jacobian evaluation.

   -- idid=true if the iteration succeeded and false if not.

   -- errcode = 0 if if the iteration succeeded

        = -1 if the initial iterate satisfies the termination criteria

        = 10 if no convergence after maxit iterations

        = 1  if the line search failed

   -- solhist:\n
This is the entire history of the iteration if you've set
keepsolhist=true\n

solhist is an N x K array where N is the length of x and K is the number of
iteration + 1. So, for scalar equations, it's a row vector.

------------------------

### Examples from the docstrings for nsol

#### World's easiest problem example. Test 64 and 32 bit Jacobians. No meaningful difference in the residual histories or the converged solutions.

```jldoctest
 julia> function f!(fv,x)
       fv[1]=x[1] + sin(x[2])
       fv[2]=cos(x[1]+x[2])
       end
f (generic function with 1 method)

julia> x=ones(2,); fv=zeros(2,); jv=zeros(2,2); jv32=zeros(Float32,2,2);
julia> nout=nsol(f!,x,fv,jv; sham=1);
julia> nout32=nsol(f!,x,fv,jv32; sham=1);
julia> [nout.history nout32.history]
5×2 Array{Float64,2}:
 1.88791e+00  1.88791e+00
 2.43119e-01  2.43120e-01
 1.19231e-02  1.19231e-02
 1.03266e-05  1.03265e-05
 1.46416e-11  1.50402e-11

julia> [nout.solution nout.solution-nout32.solution]
2×2 Array{Float64,2}:
 -7.39085e-01  5.69433e-13
  2.30988e+00  2.89546e-13
```

#### H-equation example. I'm taking the sham=5 default here, so the convergence is not quadratic. The good news is that we evaluate the Jacobian only once.

```jldoctest
julia> n=16; x0=ones(n,); FV=ones(n,); JV=ones(n,n);
julia> hdata=heqinit(x0, .5);
julia> hout=nsol(heqf!,x0,FV,JV;pdata=hdata);
julia> hout.history
4-element Array{Float64,1}:
 6.17376e-01
 3.17810e-03
 2.75227e-05
 2.35817e-07
```
"""
function nsol(
    F!,
    x0,
    FS,
    FPS,
    J! = diffjac!;
    rtol = 1.e-6,
    atol = 1.e-12,
    maxit = 20,
    solver = "newton",
    sham = 5,
    armmax = 10,
    resdec = 0.1,
    dx = 1.e-7,
    armfix = false,
    pdata = nothing,
    jfact = klfact,
    printerr = true,
    keepsolhist = false,
    stagnationok = false,
)

    itc = 0
    idid = true
    iline = false
    #
    #   If I'm letting the iteration stagnate and turning off the
    #   linesearch, then the line search cannot fail.
    #
    stagflag = stagnationok && (armmax == 0)
    #=
    First evaluation of the function. I evaluate the derivative when
    Shamanskii tells me to, at the first iteration (duh!), and when
    the rate of residual reduction is below the target value of resdec.
    =#
    (ItRules, x, n, solhist) = Newtoninit(
        x0,
        dx,
        F!,
        J!,
        solver,
        sham,
        armmax,
        armfix,
        resdec,
        maxit,
        printerr,
        pdata,
        jfact,
        keepsolhist,
    )
    #    keepsolhist ? (solhist = solhistinit(n, maxit, x)) : (solhist = [])
    #
    # First Evaluation of the function. Initialize the iteration stats.
    # Fix the tolerances for convergence and define the derivative FPF
    # outside of the main loop for scoping.
    #
    FS = EvalF!(F!, FS, x, pdata)
    resnorm = norm(FS)
    tol = rtol * resnorm + atol
    FPF = []
    ItData = ItStats(resnorm)
    newiarm = -1
    newfun = 0
    newjac = 0
    derivative_is_old = false
    residratio = 1.0
    armstop = true
    #
    # Preallocate a few vectors for the step, trial step, trial function
    #
    step = copy(x)
    xt = copy(x)
    FT = copy(x)
    #
    # If the initial iterate satisfies the termination criteria, tell me.
    #
    toosoon = (resnorm <= tol)
    #
    # The main loop stops on convergence, too many iterations, or a
    # line search failure after a derivative evaluation.
    #
    T = eltype(FPS)
    while resnorm > tol && itc < maxit && (armstop || stagnationok)
        #
        # Evaluate and factor the Jacobian.
        #
        newfun = 0
        newjac = 0
        #
        # Evaluate and factor the Jacobian if (1) you are using the chord
        # method and it's the intial iterate, or
        # (2) it's Newton and you are on the right part of the Shamaskii loop,
        # or the line search failed with a stale deriviative, or the residual
        # reduction ratio is too large. This leads to a tedious barrage
        # of conditionals that I have parked in a function.
        #
        # I'm storing the factorization in FPF unless you have asked for
        # jfact = nofact or nsol can't figure out how to factor the Jacobian.
        # In those case you just get the Jacobian back and I use \\.
        #
        evaljac = test_evaljac(ItRules, itc, newiarm, residratio)
        if evaljac
            FPF = PrepareJac!(FPS, FS, x, ItRules)
            newfun += solver == "secant"
            newjac += ~(solver == "secant")
        end
        derivative_is_old = (newjac == 0) && (solver == "newton")
        if n > 1
            # If the Jacobian precision is worse than Float32, you'll have to
            # do some scaling to avoid underflow in the terminal phase of
            # the nonlinear iteratinon. So, I do it for anything worse that
            # Float64 to make the logic simple.
            T == Float64 ? (step .= -(FPF \ FS)) :
            (ns = norm(FS, Inf); step .= -ns * (FPF \ T.(FS / ns)))
        else
            step = -FS / FPF
        end
        #
        # Compute the trial point, evaluate F and the residual norm.
        #
        AOUT = armijosc(xt, x, FT, FS, step, resnorm, ItRules, derivative_is_old)
        #
        # update solution/function value
        #
        if n > 1
            x .= AOUT.ax
            FS .= AOUT.afc
        else
            x = AOUT.ax
            FS = AOUT.afc
        end
        #
        # If the line search fails and the derivative is current,
        # stop the iteration. Print an error message unless
        # stagnationok == true and armmax=0
        #
        armstop = AOUT.idid || derivative_is_old
        iline = ~armstop && ~stagflag
        #
        # Keep the books.
        #
        residm = resnorm
        resnorm = AOUT.resnorm
        residratio = resnorm / residm
        updateStats!(ItData, newfun, newjac, AOUT)
        newiarm = AOUT.aiarm
        itc += 1
        ~keepsolhist || (@views solhist[:, itc+1] .= x)
    end
    (idid, errcode) = NewtonOK(resnorm, iline, tol, toosoon, itc, ItRules)
    newtonout = CloseIteration(x, FS, ItData, idid, errcode, keepsolhist, solhist)
    return newtonout
end


text:
# Uncomment below if the package is not already installed
#using Pkg; Pkg.add("NBInclude")

# Import code from previous section
using NBInclude
@nbinclude("Graph_Code.ipynb");

function dfs(g::Graph, start)
    visited = falses(length(g.vertices))
    function visit(ivertex)
        visited[ivertex] = true
        println("Visiting vertex #$ivertex")
        for nb in g.vertices[ivertex].neighbors
            if !visited[nb]
                visit(nb)
            end
        end
    end
    visit(start)
    return nothing
end

plot_graph(g)

dfs(g,1)

function find_path_dfs(g::Graph, start, finish)
    visited = falses(length(g.vertices))
    path = Int64[]
    function visit(ivertex)
        visited[ivertex] = true
        if ivertex == finish
            pushfirst!(path, ivertex)
            return true
        end
        for nb in g.vertices[ivertex].neighbors
            if !visited[nb]
                if visit(nb)
                    pushfirst!(path, ivertex)
                    return true
                end
            end
        end
        return false
    end
    visit(start)
    return path
end

function plot_path(g, path)
    xypath = [v.coordinates for v in g.vertices][path]
    plot(first.(xypath), last.(xypath), color="red", linewidth=2)
end

plot_graph(g)
path = find_path_dfs(g, 10, 5)
plot_path(g, path)
println("Path from 10 to 5: $path")

function bfs(g::Graph, start)
    visited = falses(length(g.vertices))
    S = [start]
    visited[start] = true
    while !isempty(S)
        ivertex = popfirst!(S)
        println("Visiting vertex #$ivertex")
        for nb in g.vertices[ivertex].neighbors
            if !visited[nb]
                visited[nb] = true
                push!(S, nb)
            end
        end
    end
end

bfs(g,1)

function shortest_path_bfs(g::Graph, start, finish)
    parent = zeros(Int64, length(g.vertices))
    S = [start]
    parent[start] = start
    while !isempty(S)
        ivertex = popfirst!(S)
        if ivertex == finish
            break
        end
        for nb in g.vertices[ivertex].neighbors
            if parent[nb] == 0 # Not visited yet
                parent[nb] = ivertex
                push!(S, nb)
            end
        end
    end
    # Build path
    path = Int64[]
    iv = finish
    while true
        pushfirst!(path, iv)
        if iv == start
            break
        end
        iv = parent[iv]
    end
    return path
end

path = shortest_path_bfs(g, 10, 5)
plot_graph(g)
plot_path(g, path)
println("Shortest path from 10 to 5: $path")

field = """
....O....
O.O.O.O..
..O...O.O
.OOOOOO.O
....OO..O
.O...OOOO
.O.......
.OOOOOO..
......O..
......O..
"""
n = 10

F = Int.(reverse(reshape(collect(field), n, n) .== 'O', dims=2))

g1 = Graph([])
nbr = 0
for j = 1:n
    for i = 1:n
        nbr += 1
        nb = Int64[]
        for dj = -1:1
            for di = -1:1
                if 1 ≤ i+di ≤ n && 1 ≤ j+dj ≤ n
                    if F[i,j] == 0 && F[i+di,j+dj] == 0
                        push!(nb, nbr + di + n*dj)
                    end
                end
            end
        end
        v = Vertex(nb, coordinates=[i,j])
        push!(g1.vertices, v)
    end
end

path = shortest_path_bfs(g1, 6, 57)
xypath = [v.coordinates for v in g1.vertices][path]
plot_graph(g1, scale=0.5)
plot(first.(xypath), last.(xypath), color="red", linewidth=2)


text:
### Functions for plotting stellar cluster results
###

#### aux. function to plot text in a subplot
## plot array of text in a box
function show_text(posx,posy, text, ywidth = 1.)
    dct= Dict("color"=> "black", "fontsize"=> 11)
    dy = ywidth / length(text)
    for t in text
        PyPlot.plt.text(posx,posy,t, dct)
        posy += dy
    end
end

### plot the stats of the DBSCAN parameters from the MCMC
function plot_dbscan_mcmc(plotdir, voname, mc::GaiaClustering.mc , showplot = true)
    PyPlot.plt[:figure](figsize=(10.0,12.0))

    nbins = 50
    PyPlot.plt.subplot(3, 2, 1  )
    h = PyPlot.plt.hist(mc.eps,nbins, alpha=0.75)
    PyPlot.plt.xlabel("ϵ")
    PyPlot.plt.grid(true)

    nbins = 20
    PyPlot.plt.subplot(3, 2, 2 )
    h = PyPlot.plt.hist(mc.mne,nbins, alpha=0.75)
    PyPlot.plt.xlabel("min_neighbor")
    PyPlot.plt.grid(true)

    nbins = 20
    PyPlot.plt.subplot(3, 2, 3 )
    h = PyPlot.plt.hist(mc.mcl,nbins, alpha=0.75)
    PyPlot.plt.xlabel("min_cluster")
    PyPlot.plt.grid(true)

    ### text

    PyPlot.plt.subplot(3, 2, 4)
    PyPlot.plt.axis("off")
    ## text to display
        text =[]
        push!(text, "VOT: "*voname)
        vtext  = mean(mc.eps)
        v2text = std(mc.eps)
        txt = "ϵ : $vtext +/- $v2text "
        push!(text,txt)
        vtext  = mean(mc.mne)
        v2text = std(mc.mne)
        txt = "min_neigh : $vtext +/- $v2text "
        push!(text,txt)
        vtext  = mean(mc.mcl)
        v2text = std(mc.mcl)
        txt = "min_clus : $vtext +/- $v2text "
        push!(text,txt)
        vtext  = mean(mc.qn)
        v2text = std(mc.qn)
        txt = "Qn : $vtext +/- $v2text "
        push!(text,txt)
        vtext  = mean(mc.qc)
        v2text = std(mc.qc)
        txt = "Qc : $vtext +/- $v2text "
        push!(text,txt)
        show_text(-0.01,0.0, text)

    #######
    nbins = 50
    PyPlot.plt.subplot(3, 2, 6 )
    PyPlot.plt.axis("on")
    PyPlot.plt.axis("on")
    h = PyPlot.plt.hist(mc.qn,nbins, color = "g", alpha=0.75)
    PyPlot.plt.xlabel("Qn")
    PyPlot.plt.grid(true)

    nbins = 50
    PyPlot.plt.subplot(3, 2, 5 )
    h = PyPlot.plt.hist(mc.qc,nbins, color="g", alpha=0.75)
    PyPlot.plt.grid(true)

    PyPlot.plt.xlabel("Qc")
    figname = plotdir*"/"*voname*".mcmc.png"
    PyPlot.plt.savefig(figname)
    if showplot PyPlot.plt.show() end
end

### plot the stats of the DBSCAN parameters from the MCMC
function plot_dbscanfull_mcmc(plotdir, voname, mc::mcfull , showplot = true)
    patch= pyimport("matplotlib.patches")
    PyPlot.plt.figure(figsize=(12.0,13.0))

    nbins = 50
    PyPlot.plt.subplot(3, 3, 1  )
    h = PyPlot.plt.hist(mc.eps,nbins, alpha=0.75)
    PyPlot.plt.xlabel("ϵ")
    PyPlot.plt.grid(true)

    nbins = 20
    PyPlot.plt.subplot(3, 3, 2 )
    h = PyPlot.plt.hist(mc.mne,nbins, alpha=0.75)
    PyPlot.plt.xlabel("min_neighbor")
    PyPlot.plt.grid(true)

    nbins = 20
    PyPlot.plt.subplot(3, 3, 3 )
    h = PyPlot.plt.hist(mc.mcl,nbins, alpha=0.75)
    PyPlot.plt.xlabel("min_cluster")
    PyPlot.plt.grid(true)

    nbins = 50
    PyPlot.plt.subplot(3, 3, 4 )
    h = PyPlot.plt.hist(mc.w3d,nbins, alpha=0.75)
    PyPlot.plt.xlabel("W3d")
    PyPlot.plt.grid(true)

    nbins = 50
    PyPlot.plt.subplot(3, 3, 5 )
    h = PyPlot.plt.hist(mc.wvel,nbins, alpha=0.75)
    PyPlot.plt.xlabel("Wvel")
    PyPlot.plt.grid(true)

    nbins = 50
    PyPlot.plt.subplot(3, 3, 6 )
    h = PyPlot.plt.hist(mc.whrd,nbins, alpha=0.75)
    PyPlot.plt.xlabel("Whrd")
    PyPlot.plt.grid(true)

    #######
    nbins = 50
    PyPlot.plt.subplot(3, 3, 7 )
    PyPlot.plt.axis("on")
    PyPlot.plt.axis("on")
    h = PyPlot.plt.hist(mc.qn,nbins, color = "g", alpha=0.75)
    PyPlot.plt.xlabel("Qn")
    PyPlot.plt.grid(true)

    nbins = 50
    PyPlot.plt.subplot(3, 3, 8 )
    h = PyPlot.plt.hist(mc.qc,nbins, color="g", alpha=0.75)
    PyPlot.plt.xlabel("Qc")
    PyPlot.plt.grid(true)

        ### text

    axt= PyPlot.plt.subplot(3, 3, 9)
    PyPlot.plt.axis("off")
    ## text to display
        dct= Dict( :color=> "black", :fontsize=> 10)
        text =[]
        push!(text, "votable: "*voname)
        vtext  =  fmt("3.3f", median(mc.eps))
        v2text =  fmt("3.3f", std(mc.eps))
        txt = "ϵ : $vtext +/- $v2text "
        push!(text,txt)
        vtext  =  fmt("3.3f", median(mc.mne))
        v2text =  fmt("3.3f", std(mc.mne))
        txt = "min_neigh : $vtext +/- $v2text "
        push!(text,txt)
        vtext  =  fmt("3.3f", median(mc.mcl))
        v2text =  fmt("3.3f", std(mc.mcl))
        txt = "min_clus : $vtext +/- $v2text "
        push!(text,txt)
        vtext  =  fmt("3.3f", median(mc.w3d))
        v2text =  fmt("3.3f", std(mc.w3d))
        txt = "W3d : $vtext +/- $v2text "
        push!(text,txt)
        vtext  =  fmt("3.3f", median(mc.wvel))
        v2text =  fmt("3.3f", std(mc.wvel))
        txt = "Wvel : $vtext +/- $v2text "
        push!(text,txt)
        vtext  =  fmt("3.3f", median(mc.whrd))
        v2text =  fmt("3.3f", std(mc.whrd))
        txt = "Whrd : $vtext +/- $v2text "
        push!(text,txt)
        vtext  =  fmt("3.3f", median(mc.qn))
        v2text =  fmt("3.3f", std(mc.qn))
        txt = "Qn : $vtext +/- $v2text "
        push!(text,txt)
        vtext  =  fmt("3.3f", median(mc.qc))
        v2text =  fmt("3.3f", std(mc.qc))
        txt = "Qc : $vtext +/- $v2text "
        push!(text,txt)
        show_text(-0.01,0.0, text)

        rec= patch.Rectangle((-0.07, -0.05), 1.2 , 1.05, color="mediumaquamarine",
            alpha= 0.4, clip_on=false)
        axt.add_artist(rec)

    PyPlot.plt.xlabel("Qc")
    figname = plotdir*"/"*voname*".mcmc.png"
    PyPlot.plt.savefig(figname)
    if showplot PyPlot.plt.show() end
end

function plot_cluster(plotdir, voname, indx, sc::GaiaClustering.SCproperties, df::GaiaClustering.Df, showplot = true , cmap = "gist_stern")

    println("### Cluster plot is centered in Y,Z...")

    PyPlot.plt.figure(figsize=(13.0,12.0))
    PyPlot.plt.subplot(3, 3, 1 , xlim = [-20,20] , ylim = [-20,20])

    xx = df.data[2,indx] .- mean(df.data[2,indx])
    yy = df.data[3,indx] .- mean(df.data[3,indx])

    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("Y (pc)")
    PyPlot.plt.ylabel("Z (pc)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 2 , ylim = [-20,20])
    xx = df.data[1,indx]
    yy = df.data[3,indx] .- mean(df.data[3,indx])
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("X (pc)")
    PyPlot.plt.ylabel("Z (pc)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 4 , xlim = [-20,20])
    xx = df.data[2,indx] .- mean(df.data[2,indx])
    yy = df.data[1,indx]
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("Y (pc)")
    PyPlot.plt.ylabel("X (pc)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 3 )
    xx = df.data[1,indx]
    yy = df.raw[13,indx]
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("X(pc)")
    PyPlot.plt.ylabel("Vrad (km/s)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 5)
    PyPlot.plt.axis("off")
    ## text to display
    text =[]
    v = sc.nstars ; txt = "N stars   : $v" ; push!(text,txt)
    v = fmt("3.1f",sc.distance) ; txt = "Distance  : $v (pc)" ; push!(text,txt)
    v = fmt("3.3f",sc.l) ; txt = "l         : $v (degree)" ; push!(text,txt)
    v = fmt("3.3f",sc.b) ; txt = "b         : $v (degree)" ; push!(text,txt)
    v = fmt("3.2f", sc.vl) ; txt = "Vl      : $v (km/s)" ; push!(text,txt)
    v = fmt("3.2f",sc.vb) ; txt = "Vb      : $v (km/s)"; push!(text,txt)
    v = fmt("3.2f",sc.vrad) ; txt = "Vradial   : $v (km/s)"; push!(text,txt)
    v = fmt("3.2f",sc.xdisp) ; txt = "X disp.   : $v (pc)" ; push!(text,txt)
    v = fmt("3.2f",sc.ydisp) ; txt = "Y disp.   : $v (pc)" ; push!(text,txt)
    v = fmt("3.2f",sc.zdisp) ; txt = "Z disp.   : $v (pc)" ; push!(text,txt)
    v = fmt("3.2f",sc.vldisp) ; txt = "Vl disp. : $v (km/s)" ; push!(text,txt)
    v = fmt("3.2f",sc.vbdisp) ; txt = "Vb disp.: $v (km/s)" ; push!(text,txt)
    v = fmt("3.2f",sc.vraddisp) ; txt = "Vradial disp.: $v (km/s)" ; push!(text,txt)
    show_text(-0.01,0.0, text , 1.0)

    PyPlot.plt.subplot(3, 3, 7 )
    PyPlot.plt.axis("on")
    xx = df.data[7,indx]
    yy = -df.data[6,indx]
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("G-Rp")
    PyPlot.plt.ylabel("G")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 8 )
    xx = df.data[4,indx]
    yy = df.data[5,indx]
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("Vl (km/s)")
    PyPlot.plt.ylabel("Vb (km/s)")
    PyPlot.plt.grid(true)

    figname = plotdir*"/"*voname*".cluster.png"
    PyPlot.plt.savefig(figname)
    if showplot PyPlot.plt.show() end
end

function plot_cluster2(plotdir, voname, indx, sc::GaiaClustering.SCproperties2, df::GaiaClustering.Df,
    showplot = true , extra= 0, cmap = "gist_stern")
    patch= pyimport("matplotlib.patches")

    println("### Cluster plot is centered in Y,Z...")
    PyPlot.plt.rcParams["font.size"]= 25

    PyPlot.plt.figure(figsize=(13.0,12.0))
    PyPlot.plt.subplot(3, 3, 1 , xlim = [-20,20] , ylim = [-20,20])

    xx = df.data[2,indx] .- mean(df.data[2,indx])
    yy = df.data[3,indx] .- mean(df.data[3,indx])

    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("Y (pc)")
    PyPlot.plt.ylabel("Z (pc)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 2 , ylim = [-20,20])
    xx = df.data[1,indx]
    yy = df.data[3,indx] .- mean(df.data[3,indx])
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("X (pc)")
    PyPlot.plt.ylabel("Z (pc)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 4 , xlim = [-20,20])
    xx = df.data[2,indx] .- mean(df.data[2,indx])
    yy = df.data[1,indx]
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("Y (pc)")
    PyPlot.plt.ylabel("X (pc)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 3 )
    xx = df.data[1,indx]
    yy = df.raw[13,indx]
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("X(pc)")
    PyPlot.plt.ylabel("Vrad (km/s)")
    PyPlot.plt.grid(true)

    axt= PyPlot.plt.subplot(3, 3, 5)
    PyPlot.plt.axis("off")

    ## text to display
    dct= Dict("color"=> "black", "fontsize"=> 10)
    text =[]
    v = sc.nstars ; txt = "N stars : $v" ; push!(text,txt)
    v = fmt("3.1f",sc.distance) ; txt = "Distance : $v (pc)" ; push!(text,txt)
    v1 = fmt("3.3f",sc.l) ; v2 = fmt("3.3f",sc.b) ;
    txt = "l , b : $v1  ,  $v2  (degree)" ; push!(text,txt)
    v1 = fmt("3.3f",sc.ra) ; v2 = fmt("3.3f",sc.dec) ;
    txt = "RA , Dec : $v1  ,  $v2  (degree)" ; push!(text,txt)

    v1 = fmt("3.3f",sc.vl) ; v2 = fmt("3.3f",sc.vb) ;
    txt = "vl , vb : $v1  ,  $v2  (km/s)" ; push!(text,txt)
    v = fmt("3.2f",sc.vrad) ; txt  = "Vradial : $v (km/s)"; push!(text,txt)
    v = fmt("3.2f",sc.xdisp) ; txt = "X disp. : $v (pc)" ; push!(text,txt)
    v = fmt("3.2f",sc.ydisp) ; txt = "Y disp. : $v (pc)" ; push!(text,txt)
    v = fmt("3.2f",sc.zdisp) ; txt = "Z disp. : $v (pc)" ; push!(text,txt)
    v = fmt("3.2f",sc.vldisp) ; txt = "Vl disp. : $v (km/s)" ; push!(text,txt)
    v = fmt("3.2f",sc.vbdisp) ; txt = "Vb disp. : $v (km/s)" ; push!(text,txt)
    v = fmt("3.2f",sc.vraddisp) ; txt = "Vradial disp. : $v (km/s)" ; push!(text,txt)
    show_text(-0.01,-0.1, text , 1.1 )

    if extra != 0
        text =[]
        v1= "$(extra.votname[1])"
        txt = "Votable : $v1" ; push!(text,txt)
        v1= "$(extra.cycle[1])"
        txt = "Cycle : $v1" ; push!(text,txt)
        v1= fmt("3.3f",extra.qc[1])
        txt = "Qc : $v1" ; push!(text,txt)
        v1= fmt("3.3f",extra.score_cycle[1])
        txt = "Score : $v1" ; push!(text,txt)
        v1= fmt("3.3f",sc.offdeg)
        txt = "Offset : $v1 (degree)" ; push!(text,txt)
        v1= fmt("3.3f",sc.edgratg)
        txt = "Edge ratio(g) : $v1 " ; push!(text,txt)
        v1= fmt("3.3f",sc.edgratm)
        txt = "Edge ratio(m) : $v1 " ; push!(text,txt)
        v = @sprintf("XG : %6.1f (pc)", extra.xg[1]) ; push!(text,v)
        v = @sprintf("YG : %6.1f (pc)", extra.yg[1]) ; push!(text,v)
        v = @sprintf("ZG : %6.1f (pc)", extra.zg[1]) ; push!(text,v)
        show_text(1.2,-0.1, text , 0.95 )

        rec= patch.Rectangle((-0.07, -0.15), 2.2, 1.15, color="salmon", alpha= 0.4, clip_on=false)
        axt.add_artist(rec)
    end

    PyPlot.plt.subplot(3, 3, 7 )
    PyPlot.plt.axis("on")
    xx = df.data[7,indx]
    yy = -df.data[6,indx]
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("G-Rp")
    PyPlot.plt.ylabel("G")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 8 )
    xx = df.data[4,indx]
    yy = df.data[5,indx]
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("Vl (km/s)")
    PyPlot.plt.ylabel("Vb (km/s)")
    PyPlot.plt.grid(true)

    figname = plotdir*"/"*voname*".cluster.png"
    PyPlot.plt.savefig(figname)
    if showplot PyPlot.plt.show() end
end

## Raw data of the field with Principal Component plots of the extracted cluster
function plot_rawdata(plotdir, voname, indx, sc::GaiaClustering.SCproperties2, df::GaiaClustering.Df, pc, ijump=100,
    showplot = true , extra= 0, cmap = "gist_stern")
    patch= pyimport("matplotlib.patches")

    nsize= size(df.data)
    nstar= nsize[2]
    iter= 1:ijump:nstar

    PyPlot.plt.rcParams["font.size"]= 25

    PyPlot.plt.figure(figsize=(13.0,12.0))
    PyPlot.plt.subplot(3, 3, 1)

    xx = df.data[2,1iter]
    yy = df.data[3,iter]
    PyPlot.plt.scatter(xx, yy , s = 0.1 )
    xx = df.data[2,indx]
    yy = df.data[3,indx]
    PyPlot.plt.scatter(xx, yy , s = 1, c="r", alpha=0.5 )
    PyPlot.plt.xlabel("Y (pc)")
    PyPlot.plt.ylabel("Z (pc)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 2 )
    xx = df.data[1,iter]
    yy = df.data[3,iter]
    PyPlot.plt.scatter(xx, yy , s = 0.1 )
    xx = df.data[1,indx]
    yy = df.data[3,indx]
    PyPlot.plt.scatter(xx, yy , s = 1, c="r", alpha=0.5 )
    PyPlot.plt.xlabel("X (pc)")
    PyPlot.plt.ylabel("Z (pc)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 4 )
    xx = df.data[2,iter]
    yy = df.data[1,iter]
    PyPlot.plt.scatter(xx, yy , s = 0.1 )
    xx = df.data[2,indx]
    yy = df.data[1,indx]
    PyPlot.plt.scatter(xx, yy , s = 1, c="r", alpha=0.5 )
    PyPlot.plt.xlabel("Y (pc)")
    PyPlot.plt.ylabel("X (pc)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 3 )
    xx = df.data[1,iter]
    yy = df.raw[13,iter]
    PyPlot.plt.scatter(xx, yy , s = 0.1 )
    xx = df.data[1,indx]
    yy = df.raw[13,indx]
    PyPlot.plt.scatter(xx, yy , s = 1, c="r", alpha=0.5 )
    PyPlot.plt.xlabel("X(pc)")
    PyPlot.plt.ylabel("Vrad (km/s)")
    PyPlot.plt.grid(true)

    ## text to display
    axt= PyPlot.plt.subplot(3, 3, 5)
    PyPlot.plt.axis("off")
    ## text to display

    if extra != 0
        text =[]
        v1= "$(extra.votname[1])"
        txt = "Votable : $v1" ; push!(text,txt)
        v = nstar ; txt = "N stars in the field  : $v" ; push!(text,txt)
        v1= "$(extra.cycle[1])"
        txt = "Cycle : $v1" ; push!(text,txt)
        v1= @sprintf("%3.3f",extra.qc[1])
        txt = "Qc : $v1" ; push!(text,txt)
        v1= @sprintf("%3.3f",extra.score_cycle[1])
        txt = "Score : $v1" ; push!(text,txt)
        v1= @sprintf("%3.3f",sc.offdeg)
        txt = "Offset : $v1 (degree)" ; push!(text,txt)
        v1= @sprintf("%3.3f",sc.edgratg)
        txt = "Edge ratio(g) : $v1 " ; push!(text,txt)
        v1= @sprintf("%3.3f",sc.edgratm)
        txt = "Edge ratio(m) : $v1 " ; push!(text,txt)
        txt= @sprintf("PC1: %3.1f , PC2: %3.1f , PC3: %3.1f", extra.pc1[1], extra.pc2[1], extra.pc3[1])
        push!(text,txt)

        show_text(-0.01,-0.1, text , 1.1 )
        rec= patch.Rectangle((-0.07, -0.15), 1.2, 1.15, color="skyblue", alpha= 0.4, clip_on=false)
        axt.add_artist(rec)
    end

    PyPlot.plt.subplot(3, 3, 7 )
    PyPlot.plt.axis("on")
    xx = df.data[7,iter]
    yy = -df.data[6,iter]
    PyPlot.plt.scatter(xx, yy , s = 0.1 )
    xx = df.data[7,indx]
    yy = -df.data[6,indx]
    PyPlot.plt.scatter(xx, yy , s = 1, c="r", alpha=0.5 )
    PyPlot.plt.xlabel("G-Rp")
    PyPlot.plt.ylabel("G")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 8 )
    xx = df.data[4,iter]
    yy = df.data[5,iter]
    PyPlot.plt.scatter(xx, yy , s = 0.1 )
    xx = df.data[4,indx]
    yy = df.data[5,indx]
    PyPlot.plt.scatter(xx, yy , s = 1, c="r", alpha=0.5 )
    PyPlot.plt.xlabel("Vl (km/s)")
    PyPlot.plt.ylabel("Vb (km/s)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 9 )
    xx = pc[1,:]
    yy = pc[2,:]
    PyPlot.plt.scatter(xx, yy , s = 1, c="r" )
    PyPlot.plt.xlabel("PC1")
    PyPlot.plt.ylabel("PC2")
    PyPlot.plt.grid(true)

    figname = plotdir*"/"*voname*".raw.png"
    PyPlot.plt.savefig(figname)

    if showplot PyPlot.plt.show() end
end

max_stars_repo_name: bosscha/gaia-shock

max_issues_repo_name: bosscha/gaia-shock

max_forks_repo_name: bosscha/gaia-shock


text:
using IncrementalInference
using DistributedFactorGraphs
using JSON

##

@testset "Test packing of mixture of distributions, 1498" begin
##

# Start with an empty factor graph
fg = initfg()

# add the first node
addVariable!(fg, :x0, ContinuousScalar)
addVariable!(fg, :x1, ContinuousScalar)
mmo = Mixture(LinearRelative,
              (hypo1=Rayleigh(3), hypo2=Uniform(30,55)),
              [0.4; 0.6])
addFactor!(fg, [:x0, :x1], mmo)

##

pf = packFactor(fg, getFactor(fg, :x0x1f1))

##

pf_ = JSON.json(pf)

##

saveDFG("/tmp/caesar/test_mixture.tar.gz", fg)
fg_ = loadDFG("/tmp/caesar/test_mixture.tar.gz")

Base.rm("/tmp/caesar/test_mixture.tar.gz")

##
end


text:
"""
    DOM: flux_kfvs!(ff, fL, fR, u, dt, sfL, sfR)
    1D1F1V: flux_kfvs!(fw, ff, fL, fR, u, ω, dt, sfL, sfR)
    1D1F3V: flux_kfvs!(fw, ff, fL, fR, u, v, w, ω, dt, sfL, sfR)
    1D2F1V: flux_kfvs!(fw, fh, fb, hL, bL, hR, bR, u, ω, dt, shL, sbL, shR, sbR)
    1D4F1V: flux_kfvs!(fw, fh0, fh1, fh2, fh3, h0L, h1L, h2L, h3L, h0R, h1R, h2R, h3R, u, ω, dt, sh0L, sh1L, sh2L, sh3L, sh0R, sh1R, sh2R, sh3R)
    2D1F2V: flux_kfvs!(fw, ff, fL, fR, u, v, ω, dt, len, sfL, sfR)
    2D2F2V: flux_kfvs!(fw, fh, fb, hL, bL, hR, bR, u, v, ω, dt, len, shL, sbL, shR, sbR)

Kinetic flux vector splitting (KFVS) flux

- @args: particle distribution functions and their left/right slopes
- @args: particle velocity quadrature points and weights
- @args: time step and cell size

"""
function flux_kfvs!(
    ff::X,
    fL::Y,
    fR::Y,
    u::Z,
    dt,
    sfL = zero(fL)::Y,
    sfR = zero(fR)::Y,
) where {
    X<:AbstractArray{<:AbstractFloat,1},
    Y<:AbstractArray{<:AbstractFloat,1},
    Z<:AbstractArray{<:AbstractFloat,1},
} # 1F1V flux for pure DOM

    # upwind reconstruction
    δ = heaviside.(u)

    f = @. fL * δ + fR * (1.0 - δ)
    sf = @. sfL * δ + sfR * (1.0 - δ)

    # calculate flux
    @. ff = dt * u * f - 0.5 * dt^2 * u^2 * sf

    return nothing

end

# ------------------------------------------------------------
# 1F1V flux
# ------------------------------------------------------------
function flux_kfvs!(
    fw::X,
    ff::Y,
    fL::Z,
    fR::Z,
    u::A,
    ω::A,
    dt,
    sfL = zero(fL)::Z,
    sfR = zero(fR)::Z,
) where {
    X<:AbstractArray{<:AbstractFloat,1},
    Y<:AbstractArray{<:AbstractFloat,1},
    Z<:AbstractArray{<:AbstractFloat,1},
    A<:AbstractArray{<:AbstractFloat,1},
}

    # upwind reconstruction
    δ = heaviside.(u)

    f = @. fL * δ + fR * (1.0 - δ)
    sf = @. sfL * δ + sfR * (1.0 - δ)

    # calculate fluxes
    fw[1] = dt * sum(ω .* u .* f) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sf)
    fw[2] = dt * sum(ω .* u .^ 2 .* f) - 0.5 * dt^2 * sum(ω .* u .^ 3 .* sf)
    fw[3] = dt * 0.5 * sum(ω .* u .^ 3 .* f) - 0.5 * dt^2 * 0.5 * sum(ω .* u .^ 4 .* sf)

    @. ff = dt * u * f - 0.5 * dt^2 * u^2 * sf

    return nothing

end

#--- mixture ---#
function flux_kfvs!(
    fw::X,
    ff::Y,
    fL::Z,
    fR::Z,
    u::A,
    ω::A,
    dt,
    sfL = zero(fL)::Z,
    sfR = zero(fR)::Z,
) where {
    X<:AbstractArray{<:AbstractFloat,2},
    Y<:AbstractArray{<:AbstractFloat,2},
    Z<:AbstractArray{<:AbstractFloat,2},
    A<:AbstractArray{<:AbstractFloat,2},
}

    for j in axes(fw, 2)
        _fw = @view fw[:, j]
        _ff = @view ff[:, j]
        _fL = @view fL[:, j]
        _fR = @view fR[:, j]
        _u = @view u[:, j]
        _ω = @view ω[:, j]
        _sfL = @view sfL[:, j]
        _sfR = @view sfR[:, j]

        flux_kfvs!(_fw, _ff, _fL, _fR, _u, _ω, dt, _sfL, _sfR)
    end

    return nothing

end

# ------------------------------------------------------------
# 2F1V flux
# ------------------------------------------------------------
function flux_kfvs!(
    fw::X,
    fh::Y,
    fb::Y,
    hL::Z,
    bL::Z,
    hR::Z,
    bR::Z,
    u::A,
    ω::A,
    dt,
    shL = zero(hL)::Z,
    sbL = zero(bL)::Z,
    shR = zero(hR)::Z,
    sbR = zero(bR)::Z,
) where {
    X<:AbstractArray{<:AbstractFloat,1},
    Y<:AbstractArray{<:AbstractFloat,1},
    Z<:AbstractArray{<:AbstractFloat,1},
    A<:AbstractArray{<:AbstractFloat,1},
}

    # upwind reconstruction
    δ = heaviside.(u)

    h = @. hL * δ + hR * (1.0 - δ)
    b = @. bL * δ + bR * (1.0 - δ)

    sh = @. shL * δ + shR * (1.0 - δ)
    sb = @. sbL * δ + sbR * (1.0 - δ)

    # calculate fluxes
    fw[1] = dt * sum(ω .* u .* h) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sh)
    fw[2] = dt * sum(ω .* u .^ 2 .* h) - 0.5 * dt^2 * sum(ω .* u .^ 3 .* sh)
    fw[3] =
        dt * 0.5 * (sum(ω .* u .^ 3 .* h) + sum(ω .* u .* b)) -
        0.5 * dt^2 * 0.5 * (sum(ω .* u .^ 4 .* sh) + sum(ω .* u .^ 2 .* sb))

    @. fh = dt * u * h - 0.5 * dt^2 * u^2 * sh
    @. fb = dt * u * b - 0.5 * dt^2 * u^2 * sb

    return nothing

end

#--- mixture ---#
function flux_kfvs!(
    fw::X,
    fh::Y,
    fb::Y,
    hL::Z,
    bL::Z,
    hR::Z,
    bR::Z,
    u::A,
    ω::A,
    dt,
    shL = zero(hL)::Z,
    sbL = zero(bL)::Z,
    shR = zero(hR)::Z,
    sbR = zero(bR)::Z,
) where {
    X<:AbstractArray{<:AbstractFloat,2},
    Y<:AbstractArray{<:AbstractFloat,2},
    Z<:AbstractArray{<:AbstractFloat,2},
    A<:AbstractArray{<:AbstractFloat,2},
}

    for j in axes(fw, 2)
        _fw = @view fw[:, j]
        _fh = @view fh[:, j]
        _fb = @view fb[:, j]
        _hL = @view hL[:, j]
        _bL = @view bL[:, j]
        _hR = @view hR[:, j]
        _bR = @view bR[:, j]
        _u = @view u[:, j]
        _ω = @view ω[:, j]
        _shL = @view shL[:, j]
        _sbL = @view sbL[:, j]
        _shR = @view shR[:, j]
        _sbR = @view sbR[:, j]

        flux_kfvs!(_fw, _fh, _fb, _hL, _bL, _hR, _bR, _u, _ω, dt, _shL, _sbL, _shR, _sbR)
    end

    return nothing

end

# ------------------------------------------------------------
# 3F1V flux (Rykov)
# ------------------------------------------------------------
function flux_kfvs!(
    fw::X,
    fh::Y,
    fb::Y,
    fr::Y,
    hL::Z,
    bL::Z,
    rL::Z,
    hR::Z,
    bR::Z,
    rR::Z,
    u::A,
    ω::A,
    dt,
    shL = zero(hL)::Z,
    sbL = zero(bL)::Z,
    srL = zero(rL)::Z,
    shR = zero(hR)::Z,
    sbR = zero(bR)::Z,
    srR = zero(rR)::Z,
) where {
    X<:AbstractArray{<:AbstractFloat,1},
    Y<:AbstractArray{<:AbstractFloat,1},
    Z<:AbstractArray{<:AbstractFloat,1},
    A<:AbstractArray{<:AbstractFloat,1},
}

    # upwind reconstruction
    δ = heaviside.(u)

    h = @. hL * δ + hR * (1.0 - δ)
    b = @. bL * δ + bR * (1.0 - δ)
    r = @. rL * δ + rR * (1.0 - δ)

    sh = @. shL * δ + shR * (1.0 - δ)
    sb = @. sbL * δ + sbR * (1.0 - δ)
    sr = @. srL * δ + srR * (1.0 - δ)

    # macro fluxes
    fw[1] = dt * sum(ω .* u .* h) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sh)
    fw[2] = dt * sum(ω .* u .^ 2 .* h) - 0.5 * dt^2 * sum(ω .* u .^ 3 .* sh)
    fw[3] =
        dt * 0.5 * (sum(ω .* u .^ 3 .* h) + sum(ω .* u .* b)) -
        0.5 * dt^2 * 0.5 * (sum(ω .* u .^ 4 .* sh) + sum(ω .* u .^ 2 .* sb))
    fw[4] = dt * 0.5 * sum(ω .* u .* r) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sr)

    # micro fluxes
    @. fh = dt * u * h - 0.5 * dt^2 * u^2 * sh
    @. fb = dt * u * b - 0.5 * dt^2 * u^2 * sb
    @. fr = dt * u * r - 0.5 * dt^2 * u^2 * sr

    return nothing

end

# ------------------------------------------------------------
# 1F3V flux
# ------------------------------------------------------------
function flux_kfvs!(
    fw::X,
    ff::Y,
    fL::Z,
    fR::Z,
    u::A,
    v::A,
    w::A,
    ω::A,
    dt::Real,
    sfL = zero(fL)::Z,
    sfR = zero(fR)::Z,
) where {
    X<:AbstractArray{<:AbstractFloat,1},
    Y<:AbstractArray{<:AbstractFloat,3},
    Z<:AbstractArray{<:AbstractFloat,3},
    A<:AbstractArray{<:AbstractFloat,3},
}

    # upwind reconstruction
    δ = heaviside.(u)

    f = @. fL * δ + fR * (1.0 - δ)
    sf = @. sfL * δ + sfR * (1.0 - δ)

    # calculate fluxes
    fw[1] = dt * sum(ω .* u .* f) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sf)
    fw[2] = dt * sum(ω .* u .^ 2 .* f) - 0.5 * dt^2 * sum(ω .* u .^ 3 .* sf)
    fw[3] = dt * sum(ω .* u .* v .* f) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* v .* sf)
    fw[4] = dt * sum(ω .* u .* w .* f) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* w .* sf)
    fw[5] =
        dt * 0.5 * sum(ω .* u .* (u .^ 2 .+ v .^ 2 .+ w .^ 2) .* f) -
        0.5 * dt^2 * 0.5 * sum(ω .* u .^ 2 .* (u .^ 2 .+ v .^ 2 .+ w .^ 2) .* sf)

    @. ff = dt * u * f - 0.5 * dt^2 * u^2 * sf

    return nothing

end

function flux_kfvs!(
    fw::X,
    ff::Y,
    fL::Z,
    fR::Z,
    u::A,
    v::A,
    w::A,
    ω::A,
    dt::Real,
    len::Real,
    sfL = zero(fL)::Z,
    sfR = zero(fR)::Z,
) where {
    X<:AbstractArray{<:AbstractFloat,1},
    Y<:AbstractArray{<:AbstractFloat,3},
    Z<:AbstractArray{<:AbstractFloat,3},
    A<:AbstractArray{<:AbstractFloat,3},
}

    # upwind reconstruction
    δ = heaviside.(u)

    f = @. fL * δ + fR * (1.0 - δ)
    sf = @. sfL * δ + sfR * (1.0 - δ)

    # calculate fluxes
    fw[1] = dt * sum(ω .* u .* f) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sf)
    fw[2] = dt * sum(ω .* u .^ 2 .* f) - 0.5 * dt^2 * sum(ω .* u .^ 3 .* sf)
    fw[3] = dt * sum(ω .* u .* v .* f) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* v .* sf)
    fw[4] = dt * sum(ω .* u .* w .* f) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* w .* sf)
    fw[5] =
        dt * 0.5 * sum(ω .* u .* (u .^ 2 .+ v .^ 2 .+ w .^ 2) .* f) -
        0.5 * dt^2 * 0.5 * sum(ω .* u .^ 2 .* (u .^ 2 .+ v .^ 2 .+ w .^ 2) .* sf)

    @. fw *= len
    @. ff = (dt * u * f - 0.5 * dt^2 * u^2 * sf) * len

    return nothing

end

# ------------------------------------------------------------
# 4F1V flux
# ------------------------------------------------------------
function flux_kfvs!(
    fw::X,
    fh0::Y,
    fh1::Y,
    fh2::Y,
    fh3::Y,
    h0L::Z,
    h1L::Z,
    h2L::Z,
    h3L::Z,
    h0R::Z,
    h1R::Z,
    h2R::Z,
    h3R::Z,
    u::A,
    ω::A,
    dt,
    sh0L = zero(h0L)::Z,
    sh1L = zero(h1L)::Z,
    sh2L = zero(h2L)::Z,
    sh3L = zero(h3L)::Z,
    sh0R = zero(h0R)::Z,
    sh1R = zero(h1R)::Z,
    sh2R = zero(h2R)::Z,
    sh3R = zero(h3R)::Z,
) where {
    X<:AbstractArray{<:AbstractFloat,1},
    Y<:AbstractArray{<:AbstractFloat,1},
    Z<:AbstractArray{<:AbstractFloat,1},
    A<:AbstractArray{<:AbstractFloat,1},
}

    # upwind reconstruction
    δ = heaviside.(u)

    h0 = @. h0L * δ + h0R * (1.0 - δ)
    h1 = @. h1L * δ + h1R * (1.0 - δ)
    h2 = @. h2L * δ + h2R * (1.0 - δ)
    h3 = @. h3L * δ + h3R * (1.0 - δ)

    sh0 = @. sh0L * δ + sh0R * (1.0 - δ)
    sh1 = @. sh1L * δ + sh1R * (1.0 - δ)
    sh2 = @. sh2L * δ + sh2R * (1.0 - δ)
    sh3 = @. sh3L * δ + sh3R * (1.0 - δ)

    # calculate fluxes
    fw[1] = dt * sum(ω .* u .* h0) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sh0)
    fw[2] = dt * sum(ω .* u .^ 2 .* h0) - 0.5 * dt^2 * sum(ω .* u .^ 3 .* sh0)
    fw[3] = dt * sum(ω .* u .* h1) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sh1)
    fw[4] = dt * sum(ω .* u .* h2) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sh2)
    fw[5] =
        dt * 0.5 * (sum(ω .* u .^ 3 .* h0) + sum(ω .* u .* h3)) -
        0.5 * dt^2 * 0.5 * (sum(ω .* u .^ 4 .* sh0) + sum(ω .* u .^ 2 .* sh3))

    @. fh0 = dt * u * h0 - 0.5 * dt^2 * u^2 * sh0
    @. fh1 = dt * u * h1 - 0.5 * dt^2 * u^2 * sh1
    @. fh2 = dt * u * h2 - 0.5 * dt^2 * u^2 * sh2
    @. fh3 = dt * u * h3 - 0.5 * dt^2 * u^2 * sh3

    return nothing

end

#--- mixture ---#
function flux_kfvs!(
    fw::X,
    fh0::Y,
    fh1::Y,
    fh2::Y,
    fh3::Y,
    h0L::Z,
    h1L::Z,
    h2L::Z,
    h3L::Z,
    h0R::Z,
    h1R::Z,
    h2R::Z,
    h3R::Z,
    u::A,
    ω::A,
    dt,
    sh0L = zero(h0L)::Z,
    sh1L = zero(h1L)::Z,
    sh2L = zero(h2L)::Z,
    sh3L = zero(h3L)::Z,
    sh0R = zero(h0R)::Z,
    sh1R = zero(h1R)::Z,
    sh2R = zero(h2R)::Z,
    sh3R = zero(h3R)::Z,
) where {
    X<:AbstractArray{<:AbstractFloat,2},
    Y<:AbstractArray{<:AbstractFloat,2},
    Z<:AbstractArray{<:AbstractFloat,2},
    A<:AbstractArray{<:AbstractFloat,2},
}

    for j in axes(fw, 2)
        _fw = @view fw[:, j]
        _fh0 = @view fh0[:, j]
        _fh1 = @view fh1[:, j]
        _fh2 = @view fh2[:, j]
        _fh3 = @view fh3[:, j]

        flux_kfvs!(
            _fw,
            _fh0,
            _fh1,
            _fh2,
            _fh3,
            h0L[:, j],
            h1L[:, j],
            h2L[:, j],
            h3L[:, j],
            h0R[:, j],
            h1R[:, j],
            h2R[:, j],
            h3R[:, j],
            u[:, j],
            ω[:, j],
            dt,
            sh0L[:, j],
            sh1L[:, j],
            sh2L[:, j],
            sh3L[:, j],
            sh0R[:, j],
            sh1R[:, j],
            sh2R[:, j],
            sh3R[:, j],
        )
    end

    return nothing

end

# ------------------------------------------------------------
# 1F2V flux
# ------------------------------------------------------------
function flux_kfvs!(
    fw::X,
    ff::Y,
    fL::Z,
    fR::Z,
    u::A,
    v::A,
    ω::A,
    dt,
    len,
    sfL = zero(fL)::Z,
    sfR = zero(fR)::Z,
) where {
    X<:AbstractVector{<:AbstractFloat},
    Y<:Union{AbstractVector{<:AbstractFloat},AbstractMatrix{<:AbstractFloat}},
    Z<:Union{AbstractVector{<:AbstractFloat},AbstractMatrix{<:AbstractFloat}},
    A<:Union{AbstractVector{<:AbstractFloat},AbstractMatrix{<:AbstractFloat}},
}

    # --- upwind reconstruction ---#
    δ = heaviside.(u)

    f = @. fL * δ + fR * (1.0 - δ)
    sf = @. sfL * δ + sfR * (1.0 - δ)

    # --- calculate fluxes ---#
    fw[1] = dt * sum(ω .* u .* f) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sf)
    fw[2] = dt * sum(ω .* u .^ 2 .* f) - 0.5 * dt^2 * sum(ω .* u .^ 3 .* sf)
    fw[3] = dt * sum(ω .* v .* u .* f) - 0.5 * dt^2 * sum(ω .* v .* u .^ 2 .* sf)
    fw[4] =
        dt * 0.5 * sum(ω .* u .* (u .^ 2 .+ v .^ 2) .* f) -
        0.5 * dt^2 * 0.5 * sum(ω .* u .^ 2 .* (u .^ 2 .+ v .^ 2) .* sf)
    fw .*= len

    @. ff = (dt * u * f - 0.5 * dt^2 * u^2 * sf) * len

    return nothing

end

# ------------------------------------------------------------
# 2F2V flux
# ------------------------------------------------------------
function flux_kfvs!(
    fw::X,
    fh::Y,
    fb::Y,
    hL::Z,
    bL::Z,
    hR::Z,
    bR::Z,
    u::A,
    v::A,
    ω::A,
    dt,
    len,
    shL = zero(hL)::Z,
    sbL = zero(bL)::Z,
    shR = zero(hR)::Z,
    sbR = zero(bR)::Z,
) where {
    X<:AbstractVector{<:AbstractFloat},
    Y<:Union{AbstractVector{<:AbstractFloat},AbstractMatrix{<:AbstractFloat}},
    Z<:Union{AbstractVector{<:AbstractFloat},AbstractMatrix{<:AbstractFloat}},
    A<:Union{AbstractVector{<:AbstractFloat},AbstractMatrix{<:AbstractFloat}},
}

    # --- upwind reconstruction ---#
    δ = heaviside.(u)

    h = @. hL * δ + hR * (1.0 - δ)
    b = @. bL * δ + bR * (1.0 - δ)
    sh = @. shL * δ + shR * (1.0 - δ)
    sb = @. sbL * δ + sbR * (1.0 - δ)

    # --- calculate fluxes ---#
    fw[1] = dt * sum(ω .* u .* h) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sh)
    fw[2] = dt * sum(ω .* u .^ 2 .* h) - 0.5 * dt^2 * sum(ω .* u .^ 3 .* sh)
    fw[3] = dt * sum(ω .* v .* u .* h) - 0.5 * dt^2 * sum(ω .* v .* u .^ 2 .* sh)
    fw[4] =
        dt * 0.5 * (sum(ω .* u .* (u .^ 2 .+ v .^ 2) .* h) + sum(ω .* u .* b)) -
        0.5 *
        dt^2 *
        0.5 *
        (sum(ω .* u .^ 2 .* (u .^ 2 .+ v .^ 2) .* sh) + sum(ω .* u .^ 2 .* sb))
    fw .*= len

    @. fh = (dt * u * h - 0.5 * dt^2 * u^2 * sh) * len
    @. fb = (dt * u * b - 0.5 * dt^2 * u^2 * sb) * len

    return nothing

end

# ------------------------------------------------------------
# 3F2V flux
# ------------------------------------------------------------
function flux_kfvs!(
    fw::X,
    fh0::Y,
    fh1::Y,
    fh2::Y,
    h0L::Z,
    h1L::Z,
    h2L::Z,
    h0R::Z,
    h1R::Z,
    h2R::Z,
    u::A,
    v::A,
    ω::A,
    dt,
    len,
    sh0L = zero(h0L)::Z,
    sh1L = zero(h1L)::Z,
    sh2L = zero(h2L)::Z,
    sh0R = zero(h0R)::Z,
    sh1R = zero(h1R)::Z,
    sh2R = zero(h2R)::Z,
) where {
    X<:AbstractArray{<:AbstractFloat,1},
    Y<:AbstractArray{<:AbstractFloat,2},
    Z<:AbstractArray{<:AbstractFloat,2},
    A<:AbstractArray{<:AbstractFloat,2},
}

    #--- upwind reconstruction ---#
    δ = heaviside.(u)

    h0 = @. h0L * δ + h0R * (1.0 - δ)
    h1 = @. h1L * δ + h1R * (1.0 - δ)
    h2 = @. h2L * δ + h2R * (1.0 - δ)
    sh0 = @. sh0L * δ + sh0R * (1.0 - δ)
    sh1 = @. sh1L * δ + sh1R * (1.0 - δ)
    sh2 = @. sh2L * δ + sh2R * (1.0 - δ)

    #--- calculate fluxes ---#
    fw[1] = dt * sum(ω .* u .* h0) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sh0)
    fw[2] = dt * sum(ω .* u .^ 2 .* h0) - 0.5 * dt^2 * sum(ω .* u .^ 3 .* sh0)
    fw[3] = dt * sum(ω .* u .* v .* h0) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* v .* sh0)
    fw[4] = dt * sum(ω .* u .* h1) - 0.5 * dt^2 * sum(ω .* u .^ 2 .* sh1)
    fw[5] =
        dt * 0.5 * (sum(ω .* u .* (u .^ 2 .+ v .^ 2) .* h0) + sum(ω .* u .* h2)) -
        0.5 *
        dt^2 *
        0.5 *
        (sum(ω .* u .^ 2 .* (u .^ 2 .+ v .^ 2) .* sh0) + sum(ω .* u .^ 2 .* sh2))

    fw .*= len
    @. fh0 = (dt * u * h0 - 0.5 * dt^2 * u^2 * sh0) * len
    @. fh1 = (dt * u * h1 - 0.5 * dt^2 * u^2 * sh1) * len
    @. fh2 = (dt * u * h2 - 0.5 * dt^2 * u^2 * sh2) * len

    return nothing

end

#--- mixture ---#
function flux_kfvs!(
    fw::X,
    fh0::Y,
    fh1::Y,
    fh2::Y,
    h0L::Z,
    h1L::Z,
    h2L::Z,
    h0R::Z,
    h1R::Z,
    h2R::Z,
    u::A,
    v::A,
    ω::A,
    dt,
    len,
    sh0L = zero(h0L)::Z,
    sh1L = zero(h1L)::Z,
    sh2L = zero(h2L)::Z,
    sh0R = zero(h0R)::Z,
    sh1R = zero(h1R)::Z,
    sh2R = zero(h2R)::Z,
) where {
    X<:AbstractArray{<:AbstractFloat,2},
    Y<:AbstractArray{<:AbstractFloat,3},
    Z<:AbstractArray{<:AbstractFloat,3},
    A<:AbstractArray{<:AbstractFloat,3},
}

    #--- reconstruct initial distribution ---#
    δ = heaviside.(u)

    h0 = @. h0L * δ + h0R * (1.0 - δ)
    h1 = @. h1L * δ + h1R * (1.0 - δ)
    h2 = @. h2L * δ + h2R * (1.0 - δ)

    sh0 = @. sh0L * δ + sh0R * (1.0 - δ)
    sh1 = @. sh1L * δ + sh1R * (1.0 - δ)
    sh2 = @. sh2L * δ + sh2R * (1.0 - δ)

    for j = 1:2
        fw[1, j] =
            dt * sum(ω[:, :, j] .* u[:, :, j] .* h0[:, :, j]) -
            0.5 * dt^2 * sum(ω[:, :, j] .* u[:, :, j] .^ 2 .* sh0[:, :, j])
        fw[2, j] =
            dt * sum(ω[:, :, j] .* u[:, :, j] .^ 2 .* h0[:, :, j]) -
            0.5 * dt^2 * sum(ω[:, :, j] .* u[:, :, j] .^ 3 .* sh0[:, :, j])
        fw[3, j] =
            dt * sum(ω[:, :, j] .* v[:, :, j] .* u[:, :, j] .* h0[:, :, j]) -
            0.5 * dt^2 * sum(ω[:, :, j] .* u[:, :, j] .^ 2 .* v[:, :, j] .* sh0[:, :, j])
        fw[4, j] =
            dt * sum(ω[:, :, j] .* u[:, :, j] .* h1[:, :, j]) -
            0.5 * dt^2 * sum(ω[:, :, j] .* u[:, :, j] .^ 2 .* sh1[:, :, j])
        fw[5, j] =
            dt *
            0.5 *
            (
                sum(
                    ω[:, :, j] .* u[:, :, j] .* (u[:, :, j] .^ 2 .+ v[:, :, j] .^ 2) .*
                    h0[:, :, j],
                ) + sum(ω[:, :, j] .* u[:, :, j] .* h2[:, :, j])
            ) -
            0.5 *
            dt^2 *
            0.5 *
            (
                sum(
                    ω[:, :, j] .* u[:, :, j] .^ 2 .* (u[:, :, j] .^ 2 .+ v[:, :, j] .^ 2) .*
                    sh0[:, :, j],
                ) + sum(ω[:, :, j] .* u[:, :, j] .^ 2 .* sh2[:, :, j])
            )

        @. fh0[:, :, j] =
            dt * u[:, :, j] * h0[:, :, j] - 0.5 * dt^2 * u[:, :, j]^2 * sh0[:, :, j]
        @. fh1[:, :, j] =
            dt * u[:, :, j] * h1[:, :, j] - 0.5 * dt^2 * u[:, :, j]^2 * sh1[:, :, j]
        @. fh2[:, :, j] =
            dt * u[:, :, j] * h2[:, :, j] - 0.5 * dt^2 * u[:, :, j]^2 * sh2[:, :, j]
    end

    @. fw *= len
    @. fh0 *= len
    @. fh1 *= len
    @. fh2 *= len

    return nothing

end


text:
__precompile__()
module MScThesis
    info("MSc Thesis Implementation by Ekrem Emre Yurdakul")

    using Knet, AutoGrad, JLD, Images, Colors

    dataDIR = joinpath(dirname(dirname(@__FILE__)),"data")
    pretrainedPath(name) = joinpath(dataDIR,"pretrained",name)
    outputPath(name) = joinpath(dataDIR,"output",name)
    worlds=["0001","0002","0006","0018","0020"]

    include("base.jl");
    include("baseData.jl");

    include("DAVIS/Conv.jl");
    include("DAVIS/LSTM.jl");
    include("DAVIS/Test.jl");
    include("DAVIS/Train.jl");

    include("Robot/Conv_D.jl");
    include("Robot/Conv_RGB.jl");
    include("Robot/Conv_RGBD.jl");
    include("Robot/LSTM_D.jl");
    include("Robot/LSTM_RGB.jl");
    include("Robot/LSTM_RGBD.jl");
    include("Robot/Test.jl");
    include("Robot/Train.jl");

    include("vKITTI/Conv_D.jl");
    include("vKITTI/Conv_RGB.jl");
    include("vKITTI/Conv_RGBD.jl");
    include("vKITTI/GRU_D.jl");
    include("vKITTI/GRU_RGB.jl");
    include("vKITTI/GRU_RGBD.jl");
    include("vKITTI/LSTM_D.jl");
    include("vKITTI/LSTM_RGB.jl");
    include("vKITTI/LSTM_RGBD.jl");
    include("vKITTI/RNN_D.jl");
    include("vKITTI/RNN_RGB.jl");
    include("vKITTI/RNN_RGBD.jl");
    include("vKITTI/Test.jl");
    include("vKITTI/Train.jl");

    include("timing.jl");
end

max_stars_repo_name: eeyrdkl/MScThesis

max_issues_repo_name: eeyrdkl/MScThesis

max_forks_repo_name: eeyrdkl/MScThesis


text:
"""
    HalfWing(foils      :: Vector{Foil},
             chords     :: Vector{Real},
             twists     :: Vector{Real},
             spans      :: Vector{Real},
             dihedrals  :: Vector{Real},
             sweeps     :: Vector{Real})

Definition for a `HalfWing` consisting of ``N+1`` airfoils and their associated chord lengths ``c``, twist angles ``\\iota``, for ``N`` sections with span lengths ``b``, dihedrals ``\\delta`` and sweep angles ``\\Lambda``, with all angles in degrees.
"""
struct HalfWing{T <: Real} <: Aircraft
    foils     :: Vector{Foil{T}}
    chords    :: Vector{T}
    twists    :: Vector{T}
    spans     :: Vector{T}
    dihedrals :: Vector{T}
    sweeps    :: Vector{T}
    function HalfWing(foils :: Vector{Foil{T}}, chords, twists, spans, dihedrals, sweeps) where T <: Real
        # Error handling
        check_wing(foils, chords, twists, spans, dihedrals, sweeps)
        # Convert angles to radians, with adjusting twists to leading edge, and generate HalfWing
        new{T}(foils, chords, -deg2rad.(twists), spans, deg2rad.(dihedrals), deg2rad.(sweeps))
    end
end

function check_wing(foils, chords, twists, spans, dihedrals, sweeps)
    # Check if number of sections match up with number of edges
    @assert (length ∘ zip)(foils, chords, twists) == (length ∘ zip)(spans, dihedrals, sweeps) + 1 "N+1 foil sections, chords and twists are required for N section(s)."
    # Check if lengths are positive
    @assert any(x -> x >= 0., chords) || any(x -> x >= 0., spans) "Chord and span lengths must be positive."
    # Check if dihedrals and sweeps are within bounds
    @assert any(x -> x >= -90. || x <= 90., dihedrals) || any(x -> x >= -90. || x <= 90., sweeps) "Dihedrals and sweep angles must not exceed ±90ᵒ."
end

# Named arguments version for ease, with default NACA-4 0012 airfoil shape
HalfWing(; chords :: Vector{T}, twists :: Vector{T}, spans :: Vector{T}, dihedrals :: Vector{T}, sweep_LEs :: Vector{T}, foils :: Vector{Foil{T}} = fill(Foil(naca4((0,0,1,2)), "NACA 0012"), length(chords))) where T <: Real = HalfWing(foils, chords, twists, spans, dihedrals, sweep_LEs)

HalfWingSection(; span = 1., dihedral = 0., sweep_LE = 0., taper = 1., root_chord = 1., root_twist = 0., tip_twist = 0., root_foil = naca4((0,0,1,2)), tip_foil = naca4((0,0,1,2))) = HalfWing([ Foil(root_foil, "Root"), Foil(tip_foil, "Tip") ], [root_chord, taper * root_chord], [root_twist, tip_twist], [span], [dihedral], [sweep_LE])

# Getters
foils(wing     :: HalfWing) = wing.foils
chords(wing    :: HalfWing) = wing.chords
twists(wing    :: HalfWing) = wing.twists
spans(wing     :: HalfWing) = wing.spans
dihedrals(wing :: HalfWing) = wing.dihedrals
sweeps(wing    :: HalfWing) = wing.sweeps

"""
    span(half_wing :: HalfWing)

Compute the planform span of a `HalfWing`.
"""
span(wing :: HalfWing) = sum(wing.spans)

taper_ratio(wing :: HalfWing) = last(wing.chords) / first(wing.chords)

"""
    projected_area(half_wing :: HalfWing)

Compute the projected area of a `HalfWing` onto the ``x``-``y`` plane.
"""
function projected_area(wing :: HalfWing)
    mean_chords = fwdsum(wing.chords) / 2 # Mean chord lengths of sections.
    mean_twists = fwdsum(wing.twists) / 2 # Mean twist angles of sections
    sum(@. wing.spans * mean_chords * cos(mean_twists))
end

section_macs(wing :: HalfWing) = @views mean_aerodynamic_chord.(wing.chords[1:end-1], fwddiv(wing.chords))

section_projected_areas(wing :: HalfWing) = wing.spans .* fwdsum(wing.chords) / 2

"""
    mean_aerodynamic_chord(half_wing :: HalfWing)

Compute the mean aerodynamic chord of a `HalfWing`.
"""
function mean_aerodynamic_chord(wing :: HalfWing)
    areas = section_projected_areas(wing)
    macs  = section_macs(wing)
    sum(macs .* areas) / sum(areas)
end

function mean_aerodynamic_center(wing :: HalfWing, factor = 0.25)
    # Compute mean aerodynamic chords and projected areas
    macs 		= section_macs(wing)
    areas 		= section_projected_areas(wing)

    # Computing leading edge coordinates
    wing_LE 	= leading_edge(wing)
    x_LEs 		= getindex.(wing_LE, 1)
    y_LEs		= getindex.(wing_LE, 2)

    # Compute x-y locations of MACs
    x_mac_LEs	= @views @. y_mac.(x_LEs[1:end-1], 2 * x_LEs[2:end], wing.chords[2:end] / wing.chords[1:end-1])
    y_macs		= @views @. y_mac.(y_LEs[1:end-1], wing.spans, wing.chords[2:end] / wing.chords[1:end-1])

    mac_coords 	= @. SVector(x_mac_LEs + factor * macs, y_macs, 0.)

    sum(mac_coords .* areas) / sum(areas)
end

camber_thickness(wing :: HalfWing, num) = camber_thickness.(wing.foils, num)

max_thickness_to_chord_ratio_location(wing :: HalfWing, num) = max_thickness_to_chord_ratio_location.(camber_thickness(wing, num))

function max_tbyc_sweeps(wing :: HalfWing, num)
    xs_max_tbyc = max_thickness_to_chord_ratio_location(wing, num)
    max_tbyc 	= last.(xs_max_tbyc)
    xs_temp 	= first.(xs_max_tbyc)
    xs 			= first.(leading_edge(wing)) .+ wing.chords .* first.(xs_temp)
    ds 			= fwddiff(xs)
    widths 		= @. wing.spans / cos(wing.dihedrals)

    sweeps 		= @. atan(ds, widths)
    xs 	  		= fwdsum(xs_temp) / 2
    tbycs 		= fwdsum(max_tbyc) / 2

    xs, tbycs, sweeps
end

"""
    leading_edge(wing :: HalfWing, flip = false)

Compute the leading edge coordinates of a `HalfWing`, with an option to flip the signs of the ``y``-coordinates.
"""
function leading_edge(wing :: HalfWing, flip = false)
    spans, dihedrals, sweeps = wing.spans, wing.dihedrals, wing.sweeps

    sweeped_spans 	 = [ 0; cumsum(@. spans * tan(sweeps)) ]
    dihedraled_spans = [ 0; cumsum(@. spans * tan(dihedrals)) ]
    cum_spans 		 = [ 0; cumsum(spans) ]

    leading 		 = SVector.(sweeped_spans, ifelse(flip, -cum_spans, cum_spans), dihedraled_spans)

    ifelse(flip, leading[end:-1:1], leading)
end

"""
    wing_bounds(wing :: HalfWing, flip = false)

Compute the leading and trailing edge coordinates of a `HalfWing`, with an option to flip the signs of the ``y``-coordinates.
"""
function wing_bounds(wing :: HalfWing, flip = false)
    chords 	 		= wing.chords
    twisted_chords 	= @. chords * sin(wing.twists)

    leading  		= leading_edge(wing, flip)
    trailing 		= SVector.(chords, (zeros ∘ length)(chords), twisted_chords)

    shifted_trailing = ifelse(flip, trailing[end:-1:1], trailing) .+ leading

    leading, shifted_trailing
end

trailing_edge(wing :: HalfWing, flip = false) = wing_bounds(wing, flip)[2]

max_forks_repo_name: GodotMisogi/AeroMDAO