text:
export
    PDB,
    PDBXML,
    MMCIF,
    MMTF,
    pdbextension,
    PDBParseError,
    pdbentrylist,
    pdbstatuslist,
    pdbrecentchanges,
    pdbobsoletelist,
    downloadpdb,
    downloadentirepdb,
    updatelocalpdb,
    downloadallobsoletepdb,
    retrievepdb,
    readpdb,
    spaceatomname,
    pdbline,
    writepdb

# PDB file formats

"Protein Data Bank (PDB) file format."
struct PDB <: BioCore.IO.FileFormat end

"Protein Data Bank (PDB) XML file format."
struct PDBXML <: BioCore.IO.FileFormat end

"Protein Data Bank (PDB) mmCIF file format."
struct MMCIF <: BioCore.IO.FileFormat end

"Protein Data Bank (PDB) MMTF file format."
struct MMTF <: BioCore.IO.FileFormat end

"Mapping of Protein Data Bank (PDB) formats to their file extensions."
const pdbextension = Dict{Type, String}(PDB=> ".pdb", PDBXML=> ".xml",
                                        MMCIF=> ".cif", MMTF=> ".mmtf")

"Error arising from parsing a Protein Data Bank (PDB) file."
struct PDBParseError <: Exception
    message::String
    line_number::Int
    line::String
end

function Base.showerror(io::IO, e::PDBParseError)
    return print(io,
            e.message,
            " at line ",
            e.line_number,
            " of file:\n",
            e.line)
end

"""
    pdbentrylist()

Obtain the list of all Protein Data Bank (PDB) entries from the RCSB server.

Requires an internet connection.
"""
function pdbentrylist()
    pdbidlist = String[]
    @info "Fetching the list of all PDB entries from the RCSB server"
    tempfilepath = tempname()
    try
        download("ftp://ftp.wwpdb.org/pub/pdb/derived_data/index/entries.idx", tempfilepath)
        open(tempfilepath) do input
            # Skips the first two lines as it contains headers
            linecount = 1
            for line in eachline(input)
                if linecount > 2
                    # The first 4 characters in the line is the PDB ID
                    pdbid = uppercase(line[1:4])
                    # Check PDB ID is 4 characters long and only consits of alphanumeric characters
                    if !occursin(r"^[a-zA-Z0-9]{4}$", pdbid)
                        throw(ArgumentError("Not a valid PDB ID: \"$pdbid\""))
                    end
                    push!(pdbidlist, pdbid)
                end
                linecount +=1
            end
        end
    finally
        rm(tempfilepath, force=true)
    end
    return pdbidlist
end

"""
    pdbstatuslist(url::AbstractString)

Obtain the list of Protein Data Bank (PDB) entries from a RCSB weekly status
file by specifying its URL.

An example URL is ftp://ftp.wwpdb.org/pub/pdb/data/status/latest/added.pdb.
Requires an internet connection.
"""
function pdbstatuslist(url::AbstractString)
    statuslist = String[]
    filename = split(url, "/")[end]
    @info "Fetching weekly status file $filename from the RCSB server"
    tempfilepath = tempname()
    try
        download(url, tempfilepath)
        # Some operating systems don't create a file if the download is empty
        if isfile(tempfilepath)
            open(tempfilepath) do input
                for line in eachline(input)
                    # The first 4 characters in the line is the PDB ID
                    pdbid = uppercase(line[1:4])
                    # Check PDB ID is 4 characters long and only consits of alphanumeric characters
                    if !occursin(r"^[a-zA-Z0-9]{4}$", pdbid)
                        throw(ArgumentError("Not a valid PDB ID: \"$pdbid\""))
                    end
                    push!(statuslist, pdbid)
                end
            end
        end
    finally
        rm(tempfilepath, force=true)
    end
    return statuslist
end

"""
    pdbrecentchanges()

Obtain three lists giving the added, modified and obsolete Protein Data Bank
(PDB) entries from the recent RCSB weekly status files.

Requires an internet connection.
"""
function pdbrecentchanges()
    addedlist = pdbstatuslist("ftp://ftp.wwpdb.org/pub/pdb/data/status/latest/added.pdb")
    modifiedlist = pdbstatuslist("ftp://ftp.wwpdb.org/pub/pdb/data/status/latest/modified.pdb")
    obsoletelist = pdbstatuslist("ftp://ftp.wwpdb.org/pub/pdb/data/status/latest/obsolete.pdb")
    return addedlist, modifiedlist, obsoletelist
end

"""
    pdbobsoletelist()

Obtain the list of all obsolete Protein Data Bank (PDB) entries from the RCSB
server.

Requires an internet connection.
"""
function pdbobsoletelist()
    obsoletelist = String[]
    @info "Fetching the list of all obsolete PDB entries from the RCSB server"
    tempfilepath = tempname()
    try
        download("ftp://ftp.wwpdb.org/pub/pdb/data/status/obsolete.dat", tempfilepath)
        open(tempfilepath) do input
            for line in eachline(input)
                # Check if its an obsolete pdb entry and not headers
                if line[1:6] == "OBSLTE"
                    # The 21st to 24th characters in obsolete pdb entry has the pdb id
                    pdbid = uppercase(line[21:24])
                    # Check PDB ID is 4 characters long and only consits of alphanumeric characters
                    if !occursin(r"^[a-zA-Z0-9]{4}$", pdbid)
                        throw(ArgumentError("Not a valid PDB ID: \"$pdbid\""))
                    end
                    push!(obsoletelist, pdbid)
                end
            end
        end
    finally
        rm(tempfilepath, force=true)
    end
    return obsoletelist
end

"""
    downloadpdb(pdbid::AbstractString; <keyword arguments>)
    downloadpdb(pdbid::Array{String, 1}; <keyword arguments>)
    downloadpdb(f::Function, args...)

Download files from the Protein Data Bank (PDB) via RCSB.

When given an `AbstractString`, e.g. `"1AKE"`, downloads the PDB file and
returns the path to the file.
When given an `Array{String, 1}`, downloads the PDB files in the array and
returns an array of the paths to the files.
When given a function as the first argument, runs the function with the
downloaded filepath(s) as an argument then removes the file(s).
Requires an internet connection.

# Arguments
- `pdb_dir::AbstractString=pwd()`: the directory to which the PDB file is
    downloaded; defaults to the current working directory.
- `file_format::Type=PDB`: the format of the PDB file; options are PDB, PDBXML,
    MMCIF and MMTF.
- `obsolete::Bool=false`: if set `true`, the PDB file is downloaded in the
    auto-generated "obsolete" directory inside the specified `pdb_dir`.
- `overwrite::Bool=false`: if set `true`, overwrites the PDB file if it exists
    in `pdb_dir`; by default skips downloading the PDB file if it exists.
- `ba_number::Integer=0`: if set > 0 downloads the respective biological
    assembly; by default downloads the PDB file.
"""
function downloadpdb(pdbid::AbstractString;
                pdb_dir::AbstractString=pwd(),
                file_format::Type=PDB,
                obsolete::Bool=false,
                overwrite::Bool=false,
                ba_number::Integer=0)
    pdbid = uppercase(pdbid)
    # Check PDB ID is 4 characters long and only consits of alphanumeric characters
    if !occursin(r"^[a-zA-Z0-9]{4}$", pdbid)
        throw(ArgumentError("Not a valid PDB ID: \"$pdbid\""))
    end
    # Check if PDB file format is valid
    if !haskey(pdbextension, file_format)
        throw(ArgumentError("Invalid PDB file format!"))
    end
    # Check if the PDB file is marked as obsolete
    if obsolete
        # Set the download path to obsolete directory inside the "pdb_dir"
        pdb_dir = joinpath(pdb_dir, "obsolete")
    end
    # Check and create directory if it does not exists in filesystem
    if !isdir(pdb_dir)
        @info "Creating directory: $pdb_dir"
        mkpath(pdb_dir)
    end
    # Standard file name format for PDB and biological assembly
    if ba_number==0
        pdbpath = joinpath(pdb_dir, "$pdbid$(pdbextension[file_format])")
    else
        pdbpath = joinpath(pdb_dir, "$(pdbid)_ba$ba_number$(pdbextension[file_format])")
    end
    # Download the PDB file only if it does not exist in the "pdb_dir" and when "overwrite" is true
    if isfile(pdbpath) && !overwrite
        @info "PDB exists: $pdbid"
    else
        # Temporary location to download compressed PDB file.
        archivefilepath = tempname()
        try
            # Download the compressed PDB file to the temporary location
            @info "Downloading PDB: $pdbid"
            if ba_number == 0
                if file_format == PDB || file_format == PDBXML || file_format == MMCIF
                    download("http://files.rcsb.org/download/$pdbid$(pdbextension[file_format]).gz", archivefilepath)
                else
                    download("http://mmtf.rcsb.org/v1.0/full/$pdbid.mmtf.gz", archivefilepath)
                end
            else
                if file_format == PDB
                    download("http://files.rcsb.org/download/$pdbid$(pdbextension[file_format])$ba_number.gz", archivefilepath)
                elseif file_format == MMCIF
                    download("http://files.rcsb.org/download/$pdbid-assembly$ba_number$(pdbextension[file_format]).gz", archivefilepath)
                else
                    throw(ArgumentError("Biological assemblies are available in the PDB and mmCIF formats only"))
                end
            end
            # Verify if the compressed file is downloaded properly and extract it
            if isfile(archivefilepath) && filesize(archivefilepath) > 0
                stream = GzipDecompressorStream(open(archivefilepath))
                open(pdbpath, "w") do output
                    write(output, stream)
                end
                close(stream)
            end
            # Verify if the PDB file is downloaded and extracted without any error
            if !isfile(pdbpath) || filesize(pdbpath)==0
                throw(ErrorException("Error downloading PDB: $pdbid"))
            end
        finally
            # Remove the temporary compressd PDB file downloaded to clear up space
            rm(archivefilepath, force=true)
        end
    end

    return pdbpath
end

function downloadpdb(pdbidlist::Array{String, 1}; kwargs...)
    pdbpaths = String[]
    failedlist = String[]
    for pdbid in pdbidlist
        try
            pdbpath = downloadpdb(pdbid; kwargs...)
            push!(pdbpaths, pdbpath)
        catch
            @warn "Error downloading PDB: $pdbid"
            push!(failedlist, pdbid)
        end
    end
    if length(failedlist) > 0
        @warn "$(length(failedlist)) PDB file(s) failed to download: $failedlist"
    end
    return pdbpaths
end

function downloadpdb(f::Function, args...; kwargs...)
    pdbpath = downloadpdb(args...; kwargs...)
    try
        f(pdbpath)
    finally
        isa(pdbpath, Array) ? rm.(pdbpath) : rm(pdbpath)
    end
end

"""
    downloadentirepdb(; <keyword arguments>)

Download the entire Protein Data Bank (PDB) from the RCSB server.

Returns the list of PDB IDs downloaded.
Ensure you have enough disk space and time before running.
The function can be stopped any time and called again to resume downloading.
Requires an internet connection.

# Arguments
- `pdb_dir::AbstractString=pwd()`: the directory to which the PDB files are
    downloaded; defaults to the current working directory.
- `file_format::Type=PDB`: the format of the PDB file; options are PDB, PDBXML,
    MMCIF and MMTF.
- `overwrite::Bool=false`: if set `true`, overwrites the PDB file if it exists
    in `pdb_dir`; by default skips downloading the PDB file if it exists.
"""
function downloadentirepdb(; pdb_dir::AbstractString=pwd(), file_format::Type=PDB, overwrite::Bool=false)
    pdblist = pdbentrylist()
    @info "About to download $(length(pdblist)) PDB files, make sure you have enough disk space and time"
    @info "The function can be stopped any time and called again to resume downloading"
    downloadpdb(pdblist, pdb_dir=pdb_dir, overwrite=overwrite, file_format=file_format)
end

"""
    updatelocalpdb(; pdb_dir::AbstractString=pwd(), file_format::Type=PDB)

Update a local copy of the Protein Data Bank (PDB).

Obtains the recent weekly lists of new, modified and obsolete PDB entries and
automatically updates the PDB files of the given `file_format` inside the local
`pdb_dir` directory.
Requires an internet connection.
"""
function updatelocalpdb(; pdb_dir::AbstractString=pwd(), file_format::Type=PDB)
    addedlist, modifiedlist, obsoletelist = pdbrecentchanges()
    # Download the newly added and modified pdb files
    downloadpdb(vcat(addedlist, modifiedlist), pdb_dir=pdb_dir, overwrite=true, file_format=file_format)
    # Set the obsolete directory to be inside pdb_dir
    obsolete_dir=joinpath(pdb_dir, "obsolete")
    for pdbid in obsoletelist
        oldfile = joinpath(pdb_dir, "$pdbid$(pdbextension[file_format])")
        newfile = joinpath(obsolete_dir, "$pdbid$(pdbextension[file_format])")
        # if obsolete pdb is in the "pdb_dir", move it to "obsolete" directory inside "pdb_dir"
        if isfile(oldfile)
            if !isdir(obsolete_dir)
                mkpath(obsolete_dir)
            end
            mv(oldfile, newfile)
        # If obsolete pdb is already in the obsolete directory, inform the user and skip
        elseif isfile(newfile)
            @info "PDB $pdbid is already moved to the obsolete directory"
        # If obsolete pdb not available in both pdb_dir and obsolete, inform the user and skip
        else
            @info "Obsolete PDB $pdbid is missing"
        end
    end
end

"""
    downloadallobsoletepdb(; <keyword arguments>)

Download all obsolete Protein Data Bank (PDB) files from the RCSB server.

Returns the list of PDB IDs downloaded.
Requires an internet connection.

# Arguments
- `obsolete_dir::AbstractString=pwd()`: the directory where the PDB files are
    downloaded; defaults to the current working directory.
- `file_format::Type=PDB`: the format of the PDB file; options are PDB, PDBXML,
    MMCIF and MMTF.
- `overwrite::Bool=false`: if set `true`, overwrites the PDB file if it exists
    in `pdb_dir`; by default skips downloading the PDB file if it exists.
"""
function downloadallobsoletepdb(; obsolete_dir::AbstractString=pwd(), file_format::Type=PDB, overwrite::Bool=false)
    obsoletelist = pdbobsoletelist()
    downloadpdb(obsoletelist, pdb_dir=obsolete_dir, file_format=file_format, overwrite=overwrite)
end

"""
    retrievepdb(pdbid::AbstractString; <keyword arguments>)

Download and read a Protein Data Bank (PDB) file or biological assembly from the
RCSB server, returning a `ProteinStructure`.

Requires an internet connection.

# Arguments
- `pdbid::AbstractString`: the PDB ID to be downloaded and read.
- `pdb_dir::AbstractString=pwd()`: the directory to which the PDB file is
    downloaded; defaults to the current working directory.
- `obsolete::Bool=false`: if set `true`, the PDB file is downloaded in the
    auto-generated "obsolete" directory inside the specified `pdb_dir`.
- `overwrite::Bool=false`: if set `true`, overwrites the PDB file if it exists
    in `pdb_dir`; by default skips downloading the PDB file if it exists.
- `ba_number::Integer=0`: if set > 0 downloads the respective biological
    assembly; by default downloads the PDB file.
- `structure_name::AbstractString="\$pdbid.pdb"`: the name given to the returned
    `ProteinStructure`; defaults to the PDB ID.
- `remove_disorder::Bool=false`: whether to remove atoms with alt loc ID not ' '
    or 'A'.
- `read_std_atoms::Bool=true`: whether to read standard ATOM records.
- `read_het_atoms::Bool=true`: whether to read HETATOM records.
"""
function retrievepdb(pdbid::AbstractString;
            pdb_dir::AbstractString=pwd(),
            obsolete::Bool=false,
            overwrite::Bool=false,
            ba_number::Integer=0,
            structure_name::AbstractString="$(uppercase(pdbid)).pdb",
            kwargs...)
    downloadpdb(pdbid, pdb_dir=pdb_dir, obsolete=obsolete, overwrite=overwrite, ba_number=ba_number)
    if obsolete
        # If obsolete is set true, the PDB file is present in the obsolete directory inside "pdb_dir"
        pdb_dir = joinpath(pdb_dir, "obsolete")
    end
    readpdb(pdbid; pdb_dir=pdb_dir, ba_number=ba_number, structure_name=structure_name, kwargs...)
end

"""
    readpdb(pdbid::AbstractString; <keyword arguments>)

Read a Protein Data Bank (PDB) file and return a `ProteinStructure`.

This is an alternative to `read("\$pdb_dir/\$pdbid.pdb", PDB)` but is
effectively the same.

# Arguments
- `pdbid::AbstractString`: the PDB ID to be read.
- `pdb_dir::AbstractString=pwd()`: the directory to which the PDB file is
    downloaded; defaults to the current working directory.
- `ba_number::Integer=0`: if set > 0 downloads the respective biological
    assembly; by default downloads the PDB file.
- `structure_name::AbstractString="\$pdbid.pdb"`: the name given to the returned
    `ProteinStructure`; defaults to the PDB ID.
- `remove_disorder::Bool=false`: whether to remove atoms with alt loc ID not ' '
    or 'A'.
- `read_std_atoms::Bool=true`: whether to read standard ATOM records.
- `read_het_atoms::Bool=true`: whether to read HETATOM records.
"""
function readpdb(pdbid::AbstractString;
            pdb_dir::AbstractString=pwd(),
            ba_number::Integer=0,
            structure_name::AbstractString="$pdbid.pdb",
            kwargs...)
    pdbid = uppercase(pdbid)
    # Standard file name format for PDB and biological assembly
    if ba_number==0
        pdbpath = joinpath(pdb_dir, "$pdbid.pdb")
    else
        pdbpath = joinpath(pdb_dir, "$(pdbid)_ba$ba_number.pdb")
    end
    read(pdbpath, PDB; structure_name=structure_name, kwargs...)
end

function Base.read(input::IO,
            ::Type{PDB};
            structure_name::AbstractString="",
            remove_disorder::Bool=false,
            read_std_atoms::Bool=true,
            read_het_atoms::Bool=true)
    # Define ProteinStructure and add to it incrementally
    struc = ProteinStructure(structure_name)
    struc[1] = Model(1, struc)
    # Entries outside of a MODEL/ENDMDL block are added to model 1
    curr_model = 1
    line_n = 0
    for line in eachline(input)
        line_n += 1
        # Read ATOM and HETATM records as required
        if (read_std_atoms && startswith(line, "ATOM  ")) ||
                (read_het_atoms && startswith(line, "HETATM"))
            unsafe_addatomtomodel!(
                    struc[curr_model],
                    AtomRecord(line, line_n), remove_disorder=remove_disorder)
        # Read MODEL record
        elseif startswith(line, "MODEL ")
            try
                curr_model = parse(Int, line[11:14])
            catch
                throw(PDBParseError(
                    "Could not read model serial number", line_n, line))
            end
            # Create model if required
            if !haskey(models(struc), curr_model)
                struc[curr_model] = Model(curr_model, struc)
            end
        # Read ENDMDL record
        elseif startswith(line, "ENDMDL")
            curr_model = 1
        end
    end
    # Remove any models that were not added to
    for model in struc
        if countchains(model) == 0
            delete!(models(struc), modelnumber(model))
        end
    end
    # Generate lists for iteration
    fixlists!(struc)
    return struc
end

function Base.read(filepath::AbstractString,
            t::Type{<:Union{PDB, MMCIF}};
            structure_name::AbstractString=splitdir(filepath)[2],
            kwargs...)
    open(filepath) do input
        read(input, t; structure_name=structure_name, kwargs...)
    end
end

# Constructor from PDB ATOM/HETATM line
AtomRecord(pdb_line::String, line_n::Integer=1) = AtomRecord(
    pdb_line[1] == 'H', # This assumes the line has already been checked as an ATOM/HETATM record
    parseserial(pdb_line, line_n),
    parseatomname(pdb_line, line_n),
    parsealtloc(pdb_line, line_n),
    parseresname(pdb_line, line_n),
    parsechainid(pdb_line, line_n),
    parseresnumber(pdb_line, line_n),
    parseinscode(pdb_line, line_n),
    [
        parsecoordx(pdb_line, line_n),
        parsecoordy(pdb_line, line_n),
        parsecoordz(pdb_line, line_n)
    ],
    parseoccupancy(pdb_line),
    parsetempfac(pdb_line),
    parseelement(pdb_line),
    parsecharge(pdb_line)
)

function parseserial(line::String, line_n::Integer=1)
    try
        return parse(Int, line[7:11])
    catch
        throw(PDBParseError("Could not read atom serial number", line_n, line))
    end
end

function parseatomname(line::String, line_n::Integer=1)
    try
        return line[13:16]
    catch
        throw(PDBParseError("Could not read atom name", line_n, line))
    end
end

function parsealtloc(line::String, line_n::Integer=1)
    try
        return line[17]
    catch
        throw(PDBParseError("Could not read alt loc identifier", line_n, line))
    end
end

function parseresname(line::String, line_n::Integer=1)
    try
        return line[18:20]
    catch
        throw(PDBParseError("Could not read residue name", line_n, line))
    end
end

function parsechainid(line::String, line_n::Integer=1)
    try
        return string(line[22])
    catch
        throw(PDBParseError("Could not read chain ID", line_n, line))
    end
end

function parseresnumber(line::String, line_n::Integer=1)
    try
        return parse(Int, line[23:26])
    catch
        throw(PDBParseError("Could not read residue number", line_n, line))
    end
end

function parseinscode(line::String, line_n::Integer=1)
    try
        return line[27]
    catch
        throw(PDBParseError("Could not read insertion code", line_n, line))
    end
end

function parsecoordx(line::String, line_n::Integer=1)
    try
        return parse(Float64, line[31:38])
    catch
        throw(PDBParseError("Could not read x coordinate", line_n, line))
    end
end

function parsecoordy(line::String, line_n::Integer=1)
    try
        return parse(Float64, line[39:46])
    catch
        throw(PDBParseError("Could not read y coordinate", line_n, line))
    end
end

function parsecoordz(line::String, line_n::Integer=1)
    try
        return parse(Float64, line[47:54])
    catch
        throw(PDBParseError("Could not read z coordinate", line_n, line))
    end
end

function parseoccupancy(line::String)
    try
        return parse(Float64, line[55:60])
    catch
        return 1.0
    end
end

function parsetempfac(line::String)
    try
        return parse(Float64, line[61:66])
    catch
        return 0.0
    end
end

function parseelement(line::String)
    try
        return line[77:78]
    catch
        return "  "
    end
end

function parsecharge(line::String)
    try
        return line[79:80]
    catch
        return "  "
    end
end

# Form a string of a certain length from a value by adding spaces to the left
# Throws an error if the value is too long
function spacestring(val_in, new_length::Integer)
    string_out = string(val_in)
    if length(string_out) > new_length
        throw(ArgumentError("Cannot fit value \"$string_out\" into $new_length space(s)"))
    end
    return lpad(string_out, new_length)
end

"""
    spaceatomname(at::Atom)

Space an `Atom` name such that the last element letter (generally) appears in
the second column.

If the `element` property of the `Atom` is set it is used to get the element,
otherwise the name starts from the second column where possible.
This function is generally not required as spacing is recorded when atom names
are read in from a Protein Data Bank (PDB) file.
However this spacing can be important, for example distinguising between C-alpha
and calcium atoms.
"""
function spaceatomname(at::Atom)
    at_name = atomname(at, strip=false)
    chars = length(at_name)
    if chars == 4
        return at_name
    end
    strip_el = element(at, strip=true)
    if chars > 4
        throw(ArgumentError("Atom name is greater than four characters: \"$at_name\""))
    end
    # In the absence of the element, the first index goes in column two
    if strip_el == "" || findfirst(isequal(strip_el[1]), at_name) == nothing
        cent_ind = 1
    # The last letter of the element goes in column two where possible
    else
        cent_ind = something(findfirst(isequal(strip_el[1]), at_name), 0) + length(strip_el) - 1
    end
    if cent_ind > 2
        throw(ArgumentError("Atom name is too long to space correctly: \"$at_name\""))
    end
    if cent_ind == 1 && chars < 4
        out_string = " $at_name"
    else
        out_string = "$at_name"
    end
    return rpad(out_string, 4)
end

# Decimal places to format output to
const coordspec = FormatSpec(".3f")
const floatspec = FormatSpec(".2f")

"""
    pdbline(at::Atom)
    pdbline(at::AtomRecord)

Form a Protein Data Bank (PDB) format ATOM or HETATM record as a `String` from
an `Atom` or `AtomRecord`.

This will throw an `ArgumentError` if a value cannot fit into the allocated
space, e.g. the chain ID is longer than one character or the atom serial is
greater than 99999.
In this case consider using `writemmcif` to write a mmCIF file.
"""
function pdbline(at::Atom)
    return (ishetero(at) ? "HETATM" : "ATOM  ") *
            # This will throw an error for serial numbers over 99999
            spacestring(serial(at), 5) *
            " " *
            spaceatomname(at) *
            string(altlocid(at)) *
            spacestring(resname(at), 3) *
            " " *
            chainid(at) * # This is checked during writing for being one character
            spacestring(resnumber(at), 4) *
            string(inscode(at)) *
            "   " *
            # This will throw an error for large coordinate values, e.g. -1000.123
            spacestring(pyfmt(coordspec, round(x(at), digits=3)), 8) *
            spacestring(pyfmt(coordspec, round(y(at), digits=3)), 8) *
            spacestring(pyfmt(coordspec, round(z(at), digits=3)), 8) *
            spacestring(pyfmt(floatspec, round(occupancy(at), digits=2)), 6) *
            # This will throw an error for large temp facs, e.g. 1000.12
            spacestring(pyfmt(floatspec, round(tempfactor(at), digits=2)), 6) *
            "          " *
            spacestring(element(at), 2) *
            spacestring(charge(at), 2)
end

function pdbline(at_rec::AtomRecord)
    return (at_rec.het_atom ? "HETATM" : "ATOM  ") *
            # This will throw an error for serial numbers over 99999
            spacestring(at_rec.serial, 5) *
            " " *
            spacestring(at_rec.atom_name, 4) *
            string(at_rec.alt_loc_id) *
            spacestring(at_rec.res_name, 3) *
            " " *
            at_rec.chain_id * # This is checked during writing for being one character
            spacestring(at_rec.res_number, 4) *
            string(at_rec.ins_code) *
            "   " *
            # This will throw an error for large coordinate values, e.g. -1000.123
            spacestring(pyfmt(coordspec, round(at_rec.coords[1], digits=3)), 8) *
            spacestring(pyfmt(coordspec, round(at_rec.coords[2], digits=3)), 8) *
            spacestring(pyfmt(coordspec, round(at_rec.coords[3], digits=3)), 8) *
            spacestring(pyfmt(floatspec, round(at_rec.occupancy, digits=2)), 6) *
            # This will throw an error for large temp facs, e.g. 1000.12
            spacestring(pyfmt(floatspec, round(at_rec.temp_factor, digits=2)), 6) *
            "          " *
            spacestring(at_rec.element, 2) *
            spacestring(at_rec.charge, 2)
end

function checkchainerror(el::Union{Chain, AbstractResidue, AbstractAtom})
    if length(chainid(el)) != 1
        throw(ArgumentError("Cannot write valid PDB file with non-single character chain ID: \"$(chainid(el))\""))
    end
end

"""
    writepdb(output, element, atom_selectors...)

Write a `StructuralElementOrList` to a Protein Data Bank (PDB) format file or
output stream.

Only ATOM, HETATM, MODEL and ENDMDL records are written - there is no header and
there are no TER records.
Atom selector functions can be given as additional arguments - only atoms that
return `true` from all the functions are retained.
"""
function writepdb(output::IO,
                el::Union{ProteinStructure, Vector{Model}},
                atom_selectors::Function...)
    # If there are multiple models, write out MODEL/ENDMDL lines
    if length(el) > 1
        for mod in sort(collect(el))
            println(output, "MODEL     ", spacestring(modelnumber(mod), 4), repeat(" ", 66))
            writepdb(output, mod, atom_selectors...)
            println(output, "ENDMDL$(repeat(" ", 74))")
        end
    # If there is only one model, do not write out MODEL/ENDMDL lines
    elseif isa(el, ProteinStructure)
        writepdb(output, defaultmodel(el), atom_selectors...)
    else
        writepdb(output, el[1], atom_selectors...)
    end
end

function writepdb(output::IO,
                el::Union{Model, Chain, AbstractResidue, Vector{Chain},
                    Vector{AbstractResidue}, Vector{Residue}, Vector{DisorderedResidue}},
                atom_selectors::Function...)
    # Collect residues then expand out disordered residues
    for res in collectresidues(el)
        checkchainerror(res)
        if isa(res, Residue)
            for at in collectatoms(res, atom_selectors...)
                writepdb(output, at)
            end
        else
            for res_name in resnames(res)
                for at in collectatoms(disorderedres(res, res_name), atom_selectors...)
                    writepdb(output, at)
                end
            end
        end
    end
end

function writepdb(output::IO, at::AbstractAtom, atom_selectors::Function...)
    checkchainerror(at)
    for atom_record in at
        println(output, pdbline(atom_record))
    end
end

function writepdb(output::IO,
                ats::Vector{<:AbstractAtom},
                atom_selectors::Function...)
    for at in collectatoms(ats, atom_selectors...)
        writepdb(output, at)
    end
end

function writepdb(filepath::AbstractString,
                el::StructuralElementOrList,
                atom_selectors::Function...)
    open(filepath, "w") do output
        writepdb(output, el, atom_selectors...)
    end
end


text:
"""
struct BagModel{T <: MillModel, U <: MillModel} <: MillModel
im::T
a::Aggregation
bm::U
end

use a `im` model on data in `BagNode`, the uses function `a` to aggregate individual bags,
and finally it uses `bm` model on the output
"""
struct BagModel{T <: MillModel, A, U <: MillModel} <: MillModel
    im::T
    a::A
    bm::U
end

Flux.@functor BagModel

BagModel(im::MillFunction, a, bm::MillFunction) = BagModel(ArrayModel(im), a, ArrayModel(bm))
BagModel(im::MillModel, a, bm::MillFunction) = BagModel(im, a, ArrayModel(bm))
BagModel(im::MillFunction, a, bm::MillModel) = BagModel(ArrayModel(im), a, bm)
BagModel(im::MillFunction, a) = BagModel(im, a, identity)
BagModel(im::MillModel, a) = BagModel(im, a, ArrayModel(identity))

(m::BagModel)(x::WeightedBagNode{<: AbstractNode}) = m.bm(m.a(m.im(x.data), x.bags, x.weights))
# if the data is missing, we do not use the mapping
(m::BagModel)(x::WeightedBagNode{<: Missing}) = m.bm(ArrayNode(m.a(x.data, x.bags, x.weights)))

(m::BagModel)(x::BagNode{<: AbstractNode}) = m.bm(m.a(m.im(x.data), x.bags))
# if the data is missing, we do not use the mapping
(m::BagModel)(x::BagNode{<: Missing}) = m.bm(ArrayNode(m.a(x.data, x.bags)))

function modelprint(io::IO, m::BagModel{T, A}; pad=[], s="", tr=false) where {T, A}
    c = COLORS[(length(pad)%length(COLORS))+1]
    paddedprint(io, "BagModel$(tr_repr(s, tr))\n", color=c)
    paddedprint(io, "  ├── ", color=c, pad=pad)
    modelprint(io, m.im, pad=[pad; (c, "  │   ")], s=s * encode(1, 1), tr=tr)
    println(io);
    paddedprint(io, "  ├── ", color=c, pad=pad)
    modelprint(io, m.a, pad=[pad; (c, "  │   ")])
    # A == Aggregation || paddedprint(io, '\n', color=c)
    println(io);
    paddedprint(io, "  └── ", color=c, pad=pad)
    modelprint(io, m.bm, pad=[pad; (c, "  │   ")])
end

function HiddenLayerModel(m::BagModel, x::BagNode, k::Int)
    im, o = HiddenLayerModel(m.im, x.data, k)
    a = SegmentedMax(k)
    b = m.a(o, x.bags)
    bm, o = HiddenLayerModel(m.bm, b, k)
    BagModel(im, a, bm), o
end

function mapactivations(hm::BagModel, x::BagNode{<:AbstractNode, B, C}, m::BagModel) where {B, C}
    hmi, mi = mapactivations(hm.im, x.data, m.im)
    ai = m.a(mi, x.bags)
    hai = hm.a(hmi, x.bags)
    hbo, bo = mapactivations(hm.bm, ai, m.bm)
    (ArrayNode(hbo.data + hai.data), bo)
end

function mapactivations(hm::BagModel, x::BagNode{<: Missing, B,C}, m::BagModel) where {B,C}
    ai = m.a(missing, x.bags)
    hai = hm.a(missing, x.bags)
    hbo, bo = mapactivations(hm.bm, ArrayNode(ai), m.bm)
    (ArrayNode(hbo.data + hai), bo)
end

ext: jl


text:
# using ArviZ
using CSV
using DataFrames
using Distributions
using NPZ
using StatsPlots
using Random
using Transits
using Turing
using NestedSamplers

rootdir(args...) = joinpath(@__DIR__, args...)
datadir(args...) = rootdir("data", args...)

data = npzread(datadir("generated_data_100ppm.npz"))
t = data["t"]
flux = data["flux"]
data = npzread(datadir("generated_data.npz"))
true_mod = data["flux"]

ground_truth = Dict(
    "period" => 3.5,
    "t0" => 1.3,
    "us" => [0.5, 0.2],
    "r" => 0.03,
    "a" => 10.0,
    "yerr" => 1e-4
)

phases = @. (t - ground_truth["t0"] + ground_truth["period"] / 2) % ground_truth["period"] - ground_truth["period"] / 2
# phases, flux_phase, mod_phase = Transits.phaseup(phases, flux, true_mod; period=ground_truth["period"], t0=ground_truth["t0"])
maxphase = 0.1
inds = @. abs(phases) < maxphase

# show data
scatter(phases[inds], flux[inds], yerr=ground_truth["yerr"], c=:black, label="data")
plot!(phases[inds], true_mod[inds], label="truth")

orbit = KeplerianOrbit(ground_truth["a"], ground_truth["period"], 0, ground_truth["t0"], 0)
law = PolynomialLimbDark(ground_truth["us"])
f = @. law(orbit, t, ground_truth["r"])
plot(phases[inds], f[inds], label="transits")
##
# using Zygote
# Turing.setadbackend(:zygote)
rng = Random.seed!(8462852)

@model function transit_model(t, flux)
    logP ~ Normal(log(ground_truth["period"]), 0.1)
    period = exp(logP)

    t0 ~ Normal(ground_truth["t0"], 10)

    us ~ Kipping13()
    r ~ Uniform(0.01, 0.1)
    loga ~ Normal(log(ground_truth["a"]), 0.1)
    a = exp(loga)

    # define orbit
    orbit = KeplerianOrbit(a, period, 0, t0, 0)
    law = PolynomialLimbDark(us)

    yerr ~ Truncated(Cauchy(0, 0.1), 0, Inf)
    for i in eachindex(flux)
        flux[i] ~ Normal(law(orbit, t[i], r), yerr)
    end
end

model = transit_model(t, flux)

chain = sample(rng, model, NUTS(5000, 0.9), MCMCThreads(), 5000, 4)

max_stars_repo_name: mileslucas/yeehaw

max_issues_repo_name: mileslucas/yeehaw

max_forks_repo_name: mileslucas/yeehaw


text:
using Cairo

function common_header()
    ## header to provide surface and context
    c = CairoRGBSurface(256,256);
    cr = CairoContext(c);

    save(cr);
    set_source_rgb(cr,0.8,0.8,0.8);    # light gray
    rectangle(cr,0.0,0.0,256.0,256.0); # background
    fill(cr);
    restore(cr);

    return c,cr
end

function write_out_picture(filename::AbstractString,c::CairoSurface)
    ## mark picture with current date
    move_to(cr,0.0,12.0);
    set_source_rgb(cr, 0,0,0);
    show_text(cr,Libc.strftime(time()));
    ##
    write_to_png(c,filename);
    nothing
end

function example_copy_path(cr)

    save(cr);
    # single (large) character set with text_path+stroke
    select_font_face(cr, "Sans", Cairo.FONT_SLANT_NORMAL,
                        Cairo.FONT_WEIGHT_BOLD);
    set_font_size(cr, 100.0);
    translate(cr, 10.0, 100.0);
    text_path(cr, "J");
    stroke_preserve(cr);

    # copied and converted
    opath = Cairo.convert_cairo_path_data(Cairo.copy_path(cr));

    restore(cr);

    l = 2 # something like a line counter

    # loop over all entries and just print out command and available points
    for x in opath
        if x.element_type == Cairo.CAIRO_PATH_MOVE_TO
            s0 = "moveto"
        elseif x.element_type == Cairo.CAIRO_PATH_LINE_TO
            s0 = "lineto"
        elseif x.element_type == Cairo.CAIRO_PATH_CURVE_TO
            s0 = "curveto"
        elseif x.element_type == Cairo.CAIRO_PATH_CLOSE_PATH
            s0 = "closepath"
        end
        s1 = repr(x.points)
        move_to(cr,100.0,14.0*l)
        l += 1
        show_text(cr,s0*s1)
    end
    nothing
end

c,cr = common_header();
save(cr);
example_copy_path(cr);
restore(cr);
write_out_picture("sample_copy_path.png",c);


text:
"""
    normalize_and_add_constraint(
        model::MOI.ModelLike,
        func::MOI.AbstractScalarFunction,
        set::MOI.AbstractScalarSet;
        allow_modify_function::Bool = false,
    )

Adds the scalar constraint obtained by moving the constant term in `func` to the
set in `model`. If `allow_modify_function` is `true` then the function `func`
can be modified.
"""
function normalize_and_add_constraint(
    model::MOI.ModelLike,
    func::MOI.AbstractScalarFunction,
    set::MOI.AbstractScalarSet;
    allow_modify_function::Bool = false,
) where {T}
    return MOI.add_constraint(
        model,
        normalize_constant(
            func,
            set;
            allow_modify_function = allow_modify_function,
        )...,
    )
end

"""
    normalize_constant(
        func::MOI.AbstractScalarFunction,
        set::MOI.AbstractScalarSet;
        allow_modify_function::Bool = false,
    )

Return the `func`-in-`set` constraint in normalized form. That is, if `func` is
[`MOI.ScalarQuadraticFunction`](@ref) or [`MOI.ScalarAffineFunction`](@ref), the
constant is moved to the set. If `allow_modify_function` is `true` then the
function `func` can be modified.
"""
function normalize_constant(
    func::MOI.AbstractFunction,
    set::MOI.AbstractSet;
    allow_modify_function::Bool = false,
)
    return func, set
end

function normalize_constant(
    func::Union{MOI.ScalarAffineFunction{T},MOI.ScalarQuadraticFunction{T}},
    set::MOI.AbstractScalarSet;
    allow_modify_function::Bool = false,
) where {T}
    if supports_shift_constant(typeof(set))
        set = shift_constant(set, -func.constant)
        if !allow_modify_function
            func = copy(func)
        end
        func.constant = zero(T)
    end
    return func, set
end


text:
C2hs = [Symmetry.Symel("E", [1 0 0; 0 1 0; 0 0 1]),
    Symmetry.Symel("sigmah", [1 0 0; 0 1 0; 0 0 -1]),
    Symmetry.Symel("i", [-1.0 0.0 0.0; 0.0 -1.0 0.0; 0.0 0.0 -1.0]),
    Symmetry.Symel("C_2^1", [-1.0 -1.2246467991473532e-16 0.0; 1.2246467991473532e-16 -1.0 0.0; 0.0 0.0 1.0])]

C3hs = [Symmetry.Symel("E", [1 0 0; 0 1 0; 0 0 1]),
    Symmetry.Symel("sigmah", [1 0 0; 0 1 0; 0 0 -1]),
    Symmetry.Symel("C_3^1", [-0.4999999999999998 -0.8660254037844387 0.0; 0.8660254037844387 -0.4999999999999998 0.0; 0.0 0.0 1.0]),
    Symmetry.Symel("C_3^2", [-0.5000000000000003 0.8660254037844384 0.0; -0.8660254037844384 -0.5000000000000003 0.0; 0.0 0.0 1.0]),
    Symmetry.Symel("S_3^1", [-0.4999999999999998 -0.8660254037844387 -0.0; 0.8660254037844387 -0.4999999999999998 0.0; 0.0 0.0 -1.0]),
    Symmetry.Symel("S_3^5", [-0.5000000000000003 0.8660254037844384 0.0; -0.8660254037844384 -0.5000000000000003 -0.0; 0.0 0.0 -1.0])]

C4hs = [Symmetry.Symel("E", [1 0 0; 0 1 0; 0 0 1]),
    Symmetry.Symel("sigmah", [1 0 0; 0 1 0; 0 0 -1]),
    Symmetry.Symel("i", [-1.0 0.0 0.0; 0.0 -1.0 0.0; 0.0 0.0 -1.0]),
    Symmetry.Symel("C_4^1", [6.123233995736766e-17 -1.0 0.0; 1.0 6.123233995736766e-17 0.0; 0.0 0.0 1.0]),
    Symmetry.Symel("C_2^1", [-1.0 -1.2246467991473532e-16 0.0; 1.2246467991473532e-16 -1.0 0.0; 0.0 0.0 1.0]),
    Symmetry.Symel("C_4^3", [-1.8369701987210297e-16 1.0 0.0; -1.0 -1.8369701987210297e-16 0.0; 0.0 0.0 1.0]),
    Symmetry.Symel("S_4^1", [6.123233995736766e-17 -1.0 0.0; 1.0 6.123233995736766e-17 0.0; 0.0 0.0 -1.0]),
    Symmetry.Symel("S_4^3", [-1.8369701987210297e-16 1.0 0.0; -1.0 -1.8369701987210297e-16 -0.0; 0.0 0.0 -1.0])]

C5hs = [Symmetry.Symel("E", [1 0 0; 0 1 0; 0 0 1]),
    Symmetry.Symel("sigmah", [1 0 0; 0 1 0; 0 0 -1]),
    Symmetry.Symel("C_5^1", [0.30901699437494745 -0.9510565162951535 0.0; 0.9510565162951535 0.30901699437494745 0.0; 0.0 0.0 1.0]),
    Symmetry.Symel("C_5^2", [-0.8090169943749473 -0.5877852522924731 0.0; 0.5877852522924731 -0.8090169943749473 0.0; 0.0 0.0 1.0]),
    Symmetry.Symel("C_5^3", [-0.8090169943749475 0.587785252292473 0.0; -0.587785252292473 -0.8090169943749475 0.0; 0.0 0.0 1.0]),
    Symmetry.Symel("C_5^4", [0.30901699437494734 0.9510565162951535 0.0; -0.9510565162951535 0.30901699437494734 0.0; 0.0 0.0 1.0]),
    Symmetry.Symel("S_5^1", [0.30901699437494745 -0.9510565162951535 0.0; 0.9510565162951535 0.30901699437494745 0.0; 0.0 0.0 -1.0]),
    Symmetry.Symel("S_5^7", [-0.8090169943749473 -0.5877852522924731 -0.0; 0.5877852522924731 -0.8090169943749473 0.0; 0.0 0.0 -1.0]),
    Symmetry.Symel("S_5^3", [-0.8090169943749475 0.587785252292473 0.0; -0.587785252292473 -0.8090169943749475 -0.0; 0.0 0.0 -1.0]),
    Symmetry.Symel("S_5^9", [0.30901699437494734 0.9510565162951535 0.0; -0.9510565162951535 0.30901699437494734 0.0; 0.0 0.0 -1.0])]

C6hs = [Symmetry.Symel("E", [1 0 0; 0 1 0; 0 0 1]),
    Symmetry.Symel("sigmah", [1 0 0; 0 1 0; 0 0 -1]),
    Symmetry.Symel("i", [-1.0 0.0 0.0; 0.0 -1.0 0.0; 0.0 0.0 -1.0]),
    Symmetry.Symel("C_6^1", [0.5000000000000001 -0.8660254037844386 0.0; 0.8660254037844386 0.5000000000000001 0.0; 0.0 0.0 1.0]),
    Symmetry.Symel("C_3^1", [-0.4999999999999998 -0.8660254037844388 0.0; 0.8660254037844388 -0.4999999999999998 0.0; 0.0 0.0 1.0]),
    Symmetry.Symel("C_2^1", [-1.0 -3.885780586188048e-16 0.0; 3.885780586188048e-16 -1.0 0.0; 0.0 0.0 1.0]),
    Symmetry.Symel("C_3^2", [-0.5000000000000006 0.8660254037844385 0.0; -0.8660254037844385 -0.5000000000000006 0.0; 0.0 0.0 1.0]),
    Symmetry.Symel("C_6^5", [0.49999999999999944 0.8660254037844392 0.0; -0.8660254037844392 0.49999999999999944 0.0; 0.0 0.0 1.0]),
    Symmetry.Symel("S_6^1", [0.5000000000000001 -0.8660254037844386 0.0; 0.8660254037844386 0.5000000000000001 0.0; 0.0 0.0 -1.0]),
    Symmetry.Symel("S_3^1", [-0.4999999999999998 -0.8660254037844388 -0.0; 0.8660254037844388 -0.4999999999999998 0.0; 0.0 0.0 -1.0]),
    Symmetry.Symel("S_3^8", [-0.5000000000000006 0.8660254037844385 0.0; -0.8660254037844385 -0.5000000000000006 -0.0; 0.0 0.0 -1.0]),
    Symmetry.Symel("S_6^5", [0.49999999999999944 0.8660254037844392 0.0; -0.8660254037844392 0.49999999999999944 0.0; 0.0 0.0 -1.0])]

C2hcn = ["Ag", "Bg", "Au", "Bu"]
C2hct = [1.0 1.0 1.0 1.0; 1.0 -1.0 1.0 -1.0; 1.0 1.0 -1.0 -1.0; 1.0 -1.0 -1.0 1.0]
C3hcn = ["A′", "E′", "A″", "E″"]
C3hct = [1.0 1.0 1.0 1.0 1.0 1.0; 2.0 -1.0 -1.0 2.0 -1.0 -1.0; 1.0 1.0 1.0 -1.0 -1.0 -1.0; 2.0 -1.0 -1.0 -2.0 1.0 1.0]
C4hcn = ["Ag", "Bg", "Eg", "Au", "Bu", "Eu"]
C4hct = [1.0   1.0   1.0   1.0   1.0   1.0   1.0   1.0;
         1.0  -1.0   1.0  -1.0   1.0  -1.0   1.0  -1.0;
         2.0   0.0  -2.0   0.0   2.0   0.0  -2.0   0.0;
         1.0   1.0   1.0   1.0  -1.0  -1.0  -1.0  -1.0;
         1.0  -1.0   1.0  -1.0  -1.0   1.0  -1.0   1.0;
         2.0   0.0  -2.0   0.0  -2.0   0.0   2.0   0.0]
C5hcn = ["A′", "E1′", "E2′", "A″", "E1″", "E2″"]
C5hct = [1.0   1.0   1.0   1.0   1.0   1.0   1.0   1.0   1.0   1.0;
         2.0  0.6180339887498949 -1.6180339887498947 -1.618033988749895   0.6180339887498945  2.0  0.6180339887498949 -1.6180339887498947 -1.618033988749895   0.6180339887498945;
         2.0 -1.6180339887498947  0.6180339887498945  0.6180339887498955 -1.6180339887498953  2.0 -1.6180339887498947  0.6180339887498945  0.6180339887498955 -1.6180339887498953;
         1.0   1.0   1.0   1.0   1.0  -1.0  -1.0  -1.0  -1.0  -1.0;
         2.0  0.6180339887498949 -1.6180339887498947 -1.618033988749895  0.6180339887498945 -2.0 -0.6180339887498949  1.6180339887498947  1.618033988749895 -0.6180339887498945;
         2.0 -1.6180339887498947  0.6180339887498945  0.6180339887498955 -1.6180339887498953 -2.0  1.6180339887498947 -0.6180339887498945 -0.6180339887498955  1.6180339887498953]

C6hcn = ["Ag", "Bg", "E1g", "E2g", "Au", "Bu", "E1u", "E2u"]
C6hct = [1.0   1.0   1.0   1.0   1.0   1.0   1.0   1.0   1.0   1.0   1.0   1.0;
         1.0  -1.0   1.0  -1.0   1.0  -1.0   1.0  -1.0   1.0  -1.0   1.0  -1.0;
         2.0   1.0  -1.0  -2.0  -1.0   1.0   2.0   1.0  -1.0  -2.0  -1.0   1.0;
         2.0  -1.0  -1.0   2.0  -1.0  -1.0   2.0  -1.0  -1.0   2.0  -1.0  -1.0;
         1.0   1.0   1.0   1.0   1.0   1.0  -1.0  -1.0  -1.0  -1.0  -1.0  -1.0;
         1.0  -1.0   1.0  -1.0   1.0  -1.0  -1.0   1.0  -1.0   1.0  -1.0   1.0;
         2.0   1.0  -1.0  -2.0  -1.0   1.0  -2.0  -1.0   1.0   2.0   1.0  -1.0;
         2.0  -1.0  -1.0   2.0  -1.0  -1.0  -2.0   1.0   1.0  -2.0   1.0   1.0]


text:
module RuiExamples

using BenchmarkTools, IntervalCensoredEventTimes, Ipopt, KNITRO, MAT, Test

@info "load n=200, p=1000 example data"
datafile = "/Users/huazhou/Desktop/interval_censoring_examples/n200p1000_Identity.mat"
file = matopen(datafile)
Z = read(file, "Z")
βtrue = read(file, "beta_tr")
L = vec(read(file, "L"))
R = vec(read(file, "R"))
close(file)

# form a data set with with p=30 for benchmark
Zbm = [ones(size(Z, 1)) Z[:, [1, 50, 100, 150, 200]] Z[:, 101:124]]
βbm = [0; βtrue[[1, 50, 100, 150, 200]]; zeros(24)]
@show βbm
# form model object
icm = IntervalCensoredModel(Zbm, L, R)
display(icm); println()
# fit model
# KNITRO options:
# algorithm (1, IP/Direct; 2, IP/CG; 3, Active Set; 4, SQP (fastest); 5, try all)
# hessopt (2, BFGS; 6, L-BFGS)
# derivcheck (1, gradient; 2, hessian; 3, both)
# derivcheck_tol
# derivcheck_type (1, forward; 2, central)
# maxit
# solver = KNITRO.KnitroSolver(outlev = 3, derivcheck = 3, derivcheck_tol = 1e-4, maxit = 10)
solver = KNITRO.KnitroSolver(outlev = 3)
@time fit!(icm, solver, init=initialize_uniform!(icm))
display(icm); println()
# IPOPT options:
# hessian_approximation = "limited-memory"
# watchdog_shortened_iter_trigger = 3
# solver = Ipopt.IpoptSolver(print_level = 0, tol = 1e-4)
# @time fit!(icm, solver, init=initialize_uniform!(icm))
# display(icm); println()

# @info "Benchmark loglikelihood!"
# initialize_uniform!(icm)
# bm = @benchmark loglikelihood!($icm, false, false)
# display(bm); println()
# bm = @benchmark loglikelihood!($icm, true, false)
# display(bm); println()
# bm = @benchmark loglikelihood!($icm, true, true)
# display(bm); println()
end


text:
# Types used for rational expectations

type Blanchard_Kahn_Form{T<:AbstractFloat,S<:Integer}
  # E_tY[t+1] = A*Y[t] + C*V[t+1]
  nx::S                               # Number of predetermined variables
  ny::S                               # Number of nonpredetermined variables
  a::Array{T,2}                       # Companion matrix
  c::Union{Array{T,2},Array{T,1}}     # Innovation loading matrix
  sigma::Union{Array{T,2},Array{T,1}} # Innovation variance-covariance matrix
end

type Blanchard_Kahn_Soln{T<:AbstractFloat,S<:Integer}
  p::Union{Array{T,2},Array{T,1}}     # Transition matrix for predetermined variables
  k::Union{Array{T,2},Array{T,1}}     # Innovation loading matrix
  h::Union{Array{T,2},Array{T,1}}     # Decision rule matrix linking nonpredetermined variables to predetermined variables
  sigma::Union{Array{T,2},Array{T,1}} # Innovation variance-covariance matrix
  grc::S                              # Number of eigenvalues greater than cutoff
  soln_type::AbstractString                   # "determinate", "indeterminate", or "explosive"
end

type Klein_Form{T<:AbstractFloat,S<:Integer}
  # B*E_tY[t+1] = A*Y[t] + C*V[t+1]
  nx::S                               # Number of predetermined variables
  ny::S                               # Number of nonpredetermined variables
  a::Array{T,2}                       # Companion matrix; this is B in Klein's paper
  b::Array{T,2}                       # Lead matrix; this is A in Klein's paper
  c::Union{Array{T,2},Array{T,1}}     # Innovation loading matrix
  sigma::Union{Array{T,2},Array{T,1}} # Innovation variance-covariance matrix
end

type Klein_Soln{T<:AbstractFloat,S<:Integer}
  p::Union{Array{T,2},Array{T,1}}     # Transition matrix for predetermined variables
  k::Union{Array{T,2},Array{T,1}}     # Innovation loading matrix
  h::Union{Array{T,2},Array{T,1}}     # Decision rule matrix linking nonpredetermined variables to predetermined variables
  sigma::Union{Array{T,2},Array{T,1}} # Innovation variance-covariance matrix
  grc::S                              # Number of eigenvalues greater than cutoff
  soln_type::AbstractString                   # "determinate", "indeterminate", or "explosive"
end

type Binder_Pesaran_Form{T<:AbstractFloat}
  # A*Y[t] = A1*Y[t-1]+ B*E_tY[t+1] + C*V[t]
  a::Array{T,2}                       # Contemporaneous matrix
  a1::Array{T,2}                      # Lag matrix matrix
  b::Array{T,2}                       # Lead matrix
  c::Union{Array{T,2},Array{T,1}}     # Innovation matrix
  sigma::Union{Array{T,2},Array{T,1}} # Innovation variance-covariance matrix
end

type Binder_Pesaran_Soln{T<:AbstractFloat,S<:Integer}
  p::Union{Array{T,2},Array{T,1}}     # Transition matrix for predetermined and nonpredetermined variables
  k::Union{Array{T,2},Array{T,1}}     # Innovation loading matrix
  sigma::Union{Array{T,2},Array{T,1}} # Innovation variance-covariance matrix
  grc::S                              # Number of eigenvalues greater than cutoff
  soln_type::AbstractString                   # "determinate", "indeterminate", or "explosive"
end

type Sims_Form{T<:AbstractFloat}
  # gamma0*Y[t] = gamma1*Y[t-1] + C + psi*V[t] + pi*Eta[t]
  gamma0::Array{T,2}                  # Lead matrix
  gamma1::Array{T,2}                  # Companion matrix
  c::Array{T,1}                       # Constant vector
  psi::Array{T,2}                     # Innovation loading matrix
  pi::Union{Array{T,2},Array{T,1}}    # Expectational-errors loading matrix
  sigma::Union{Array{T,2},Array{T,1}} # Innovation variance-covariance matrix
end

type Sims_Soln{T<:AbstractFloat,S<:Integer}
  # y(t) = g1*y(t-1) + c + impact*z(t) + ywt*inv(I-fmat*inv(L))*fwt*z(t+1)
  g1::Union{Array{T,2},Array{T,1}}     # Transition matrix
  c::Array{T,1}                        # Constant vector
  impact::Union{Array{T,2},Array{T,1}} # Constant vector
  ywt::Union{Array{T,2},Array{T,1}}    # Innovation loading matrix
  fmat::Union{Array{T,2},Array{T,1}}   # Expectational-errors loading matrix
  fwt::Union{Array{T,2},Array{T,1}}    # Expectational-errors loading matrix
  sigma::Union{Array{T,2},Array{T,1}}  # Innovation variance-covariance matrix
  grc::S                               # Number of eigenvalues greater than cutoff
  iid_exist::Bool                      # true or false
  general_exist::Bool                  # true or false
  uniqueness::Bool                     # true or false
end

type Gomme_Klein_Form{T<:AbstractFloat,S<:Integer}
  nx::S                                # Number of predetermined variables
  ny::S                                # Number of nonpredetermined variables
  derivs::Array{T,2}                   # Matrix of first derivatives (n*2n, where n = nx+ny)
  hessians::Array{T,2}                 # Matrix of stacked second derivatives ((2n)^2*2n, where n = nx+ny)
  eta::Union{Array{T,2},Array{T,1}}    # Innovation loading-matrix in predetermined equation-block
  sigma::Union{Array{T,2},Array{T,1}} # Innovation variance-covariance matrix
end

type Gomme_Klein_Soln{T<:AbstractFloat,S<:Integer}
  # x(t+1) = ssh + hx*x(t) + [kron(I,x(t))]'hxx*[kron(I,x(t))] + v(t+1)
  #   y(t) = ssg + gx*x(t) + [kron(I,x(t))]'gxx*[kron(I,x(t))]
  ssh::Array{T,1}                     # Intercepts in predetermined block
  hx::Union{Array{T,2},Array{T,1}}    # Linear component in predetermined block
  hxx::Array{T,2}                     # Quadratic component in predetermined block
  ssg::Array{T,1}                     # Intercepts in predetermined block
  gx::Union{Array{T,2},Array{T,1}}    # Linear component in predetermined block
  gxx::Array{T,2}                     # Quadratic component in predetermined block
  eta::Union{Array{T,2},Array{T,1}}   # Innovation loading-matrix in predetermined equation-block
  sigma::Union{Array{T,2},Array{T,1}} # Innovation variance-covariance matrix
  grc::S                              # Number of eigenvalues greater than cutoff
  soln_type::AbstractString                   # "determinate", "indeterminate", or "explosive"
end

type Lombardo_Sutherland_Form{T<:AbstractFloat,S<:Integer}
  nx::S                                # Number of predetermined variables
  ny::S                                # Number of nonpredetermined variables
  derivs::Array{T,2}                   # Matrix of first derivatives (n*2n, where n = nx+ny)
  hessians::Array{T,2}                 # Matrix of stacked second derivatives ((2n)^2*2n, where n = nx+ny)
  eta::Union{Array{T,2},Array{T,1}}    # Innovation loading-matrix in predetermined equation-block
  sigma::Union{Array{T,2},Array{T,1}}  # Innovation variance-covariance matrix
end

type Lombardo_Sutherland_Soln{T<:AbstractFloat,S<:Integer}
  ssh::Union{Array{T,2},Array{T,1}}   # Intercepts in predetermined block
  hx::Union{Array{T,2},Array{T,1}}    # Linear component in predetermined block
  hxx::Array{T,2}                     # Quadratic component in predetermined block
  ssg::Union{Array{T,2},Array{T,1}}   # Intercepts in predetermined block
  gx::Union{Array{T,2},Array{T,1}}    # Linear component in predetermined block
  gxx::Array{T,2}                     # Quadratic component in predetermined block
  eta::Union{Array{T,2},Array{T,1}}   # Innovation loading-matrix in predetermined equation-block
  phi::Array{T,2}
  gamma::Array{T,2}
  psi::Array{T,2}
  sigma::Union{Array{T,2},Array{T,1}} # Innovation variance-covariance matrix
  grc::S                              # Number of eigenvalues greater than cutoff
  soln_type::AbstractString                   # "determinate", "indeterminate", or "explosive"
end

# Types used for optimal policy

type State_Space_Objective{T<:AbstractFloat}
  beta::T
  q::Array{T,2}
  u::Array{T,2}
  r::Array{T,2}
end

type State_Space_Form{T<:AbstractFloat,S<:Integer}
  # E_tY[t+1] = A*Y[t] + B*U[t] + C*V[t+1]
  nx::S             # Number of predetermined variables
  ny::S             # Number of nonpredetermined variables
  a::Array{T,2}     # Companion matrix
  b::Array{T,2}     # Policy loadings
  c::Array{T,2}     # Innovation loading matrix
  sigma::Array{T,2} # Innovation variance-covariance matrix
end

type Generalized_State_Space_Form{T<:AbstractFloat,S<:Integer}
  # E_tY[t+1] = A*Y[t] + B*U[t] + C*V[t+1]
  nx::S             # Number of predetermined variables
  ny::S             # Number of nonpredetermined variables
  a0::Array{T,2}    # Companion matrix
  a::Array{T,2}     # Companion matrix
  b::Array{T,2}     # Policy loadings
  c::Array{T,2}     # Innovation loading matrix
  sigma::Array{T,2} # Innovation variance-covariance matrix
end

type State_Space_Soln{T<:AbstractFloat}
  p::Array{T,2}     # Transition matrix for predetermined variables
  k::Array{T,2}     # Innovation loading matrix
  h::Array{T,2}
  f::Array{T,2}     # Decision rule matrix linking nonpredetermined variables to predetermined variables
  v::Array{T,2}
  sigma::Array{T,2} # Innovation variance-covariance matrix
  converged::Bool
end

type Structural_Objective{T<:AbstractFloat}
  beta::T
  q::Array{T,2}
  r::Array{T,2}
end

type Structural_Form{T<:AbstractFloat}
  # E_tY[t+1] = A*Y[t] + B*U[t] + C*V[t+1]
  a0::Array{T,2}     # Companion matrix
  a1::Array{T,2}     # Policy loadings
  a2::Array{T,2}     # Innovation loading matrix
  a3::Array{T,2}     # Policy loadings
  a5::Array{T,2}     # Innovation loading matrix
  sigma::Array{T,2}  # Innovation variance-covariance matrix
end

type Generalized_Structural_Form{T<:AbstractFloat}
  # E_tY[t+1] = A*Y[t] + B*U[t] + C*V[t+1]
  a0::Array{T,2}     # Companion matrix
  a1::Array{T,2}     # Policy loadings
  a2::Array{T,2}     # Innovation loading matrix
  a3::Array{T,2}     # Policy loadings
  a4::Array{T,2}     # Policy loadings
  a5::Array{T,2}     # Innovation loading matrix
  sigma::Array{T,2}  # Innovation variance-covariance matrix
end

type Structural_Soln{T<:AbstractFloat}
  p::Array{T,2}     # Transition matrix for predetermined variables
  k::Array{T,2}     # Innovation loading matrix
  v::Array{T,2}
  sigma::Array{T,2} # Innovation variance-covariance matrix
  converged::Bool
end

# Types used for impulse response functions

Perturbable_Soln = Union{Blanchard_Kahn_Soln,Klein_Soln,Binder_Pesaran_Soln,Gomme_Klein_Soln,Lombardo_Sutherland_Soln,State_Space_Soln,Structural_Soln}
Second_Order_State_Space_Soln = Union{Gomme_Klein_Soln,Lombardo_Sutherland_Soln}


text:
# This file is a modified version of the script at
# https://github.com/JuliaPy/pyjulia/blob/master/src/julia/compile.jl
# from the pyjulia projects, which ships under the MIT license

compiler_env, script, output, morbit_repo_url, morbit_env = ARGS

include("prepare_path.jl")

const Pkg =
    Base.require(Base.PkgId(Base.UUID("44cfe95a-1eb2-52ea-b672-e2afdf69b78f"), "Pkg"))

Pkg.activate(compiler_env)
@info "Loading PackageCompiler..."
using PackageCompiler

Pkg.activate(prepare_path(morbit_env))
# install morbit if need be
if !(haskey( Pkg.dependencies(), Base.UUID("88936782-c8cd-4a0f-b259-ffb12bfd2869") ) )
    Pkg.add(; url = morbit_repo_url )
end

using Morbit

# install test dependencies
test_toml = joinpath(pkgdir(Morbit), "test", "Project.toml");
target_toml = joinpath(compiler_env, "test_env", "Project.toml")
mkpath(joinpath(compiler_env,"test_env"))
cp(test_toml, target_toml; force = true)
chmod(target_toml,0o777)
Pkg.activate(target_toml)
Pkg.instantiate()

@info "Compiling Morbit system image..."
Pkg.activate(prepare_path(morbit_env))
create_sysimage(
    [:Morbit],
    sysimage_path = output,
    precompile_execution_file = script,
)

@info "Morbit image is created at $output"

max_stars_repo_name: manuelbb-upb/MorbitWrapper

max_issues_repo_name: manuelbb-upb/MorbitWrapper

max_forks_repo_name: manuelbb-upb/MorbitWrapper


text:
"""
    aerodynamic_conductance!(df;
      Gb_model = Thom1972(), Ram_model = ResistanceWindZr(),
      zr=nothing, zh=nothing, d = isnothing(zh) ? nothing : 0.7*zh,
      ...
      )

Bulk aerodynamic conductance, including options for the boundary layer conductance
formulation and stability correction functions.

# Arguments
- `df`: DataFrame with columns
  - `ustar`     : Friction velocity (m s-1)
  - `wind`      : Wind speed at sensor height (m s-1)
- `Gb_model`  : model for computing boundary layer conductance (see [`compute_Gb!`](@ref))
- `Ram_model` : model for computing aerodynamic resistance (see [`compute_Ram`](@ref))
- `zh`        : canopy height (m)
- `zr`        : Instrument (reference) height (m)

Further required columns of `df` and keyword argument depend on `Gb_model`
(see [`compute_Gb!`](@ref)) and `Ram_model` (see [`compute_Ram`](@ref)).

If only columns `ustar` and `wind` are available, use default models
(`Thom1972()` and `ResistanceWindZr()`).

# Details

Aerodynamic conductance for heat (Ga_h) is calculated as:

``Ga_h = 1 / (Ra_m + Rb_h)``

where ``Ra_m`` is the aerodynamic resistance for momentum and ``Rb_h = 1/Gb_h`` the
(quasi-laminar) canopy boundary layer resistance ('excess resistance') for heat.

`Ra_m` is computed and described with [`compute_Ram`](@ref) using model `Ram_model`.

`Rb_h` is computed and described with 1/[`compute_Gb!`](@ref) using a given `Gb_model`.

# Value
combined results of [`compute_Gb!`](@ref) and
- `Ra_m`: Aerodynamic resistance for momentum transfer (s m-1)
- `Ga_m`: Aerodynamic conductance for momentum transfer (m s-1)
- `Ga_h`: Aerodynamic conductance for heat transfer (m s-1)
- `Ra_h`: Aerodynamic resistance for heat transfer (s m-1)
- `Ga_CO2`: Aerodynamic conductance for CO2 transfer (m s-1)
- `z0h`: roughness length for heat (m)

# Note
The roughness length for water and heat (z0h) is computed by [`roughness_length_heat`](@ref).

TODO check
Input variables such as LAI, Dl, or zh can be either constants, or
vary with time, i.e. are vectors of the same length as `df`.

Note that boundary layer conductance to water vapor transfer (`Gb_w`) is often
assumed to equal `Gb_h`. This assumption is also made in `Bigleaf.jl`, for
example in the function [`surface_conductance`](@ref).

If the roughness length for momentum (`z0m`) is not provided as input, it is estimated
using [`roughness_parameters`](@ref), which estimates a single `z0m`
value for the entire time period. If a varying `z0m` value
(e.g. across seasons or years) is required, `z0m` should be provided as input argument.

# Examples
```jldoctest; output = false
using DataFrames
df = DataFrame(Tair=25.0,pressure=100.0,wind=[3.0,4,5],
  ustar=[0.5,0.6,0.65],H=[200.0,230,250])
# simple calculation of Ga
aerodynamic_conductance!(df;Gb_model=Thom1972())
# calculation of Ram using a model derived from the logarithmic wind profile
aerodynamic_conductance!(df;Gb_model=Thom1972(),Ram_model = ResistanceWindProfile(),
  zr=40,zh=25,d=17.5,z0m=2)
# simple calculation of Ga, but a physically based canopy boundary layer model
aerodynamic_conductance!(df,Gb_model=Su2001(),
  zr=40,zh=25,d=17.5,Dl=0.05,N=2,fc=0.8)
all(isfinite.(df.psi_h))
# output
true
```
"""
function aerodynamic_conductance!(df; Gb_model = Thom1972(), Ram_model = ResistanceWindZr(),
  zr=nothing,zh=nothing, d = isnothing(zh) ? nothing : 0.7*zh ,
  z0m=nothing,Dl=nothing,N=2,fc=nothing,LAI=nothing,Cd=0.2,hs=0.01,
  leafwidth=nothing,
  stab_formulation=Dyer1970(),
  kB_h=nothing,constants=BigleafConstants()
  )
  # add zeta
  if !isnothing(zr) && !isnothing(d) && !(stab_formulation isa NoStabilityCorrection)
    stability_parameter!(df::AbstractDataFrame; z=zr, d, constants)
  else
    df[!,:zeta] .= missing
  end
  # adds columns psi_m and psi_h, (:no_stability_correction: just add 0s without using zeta)
  stability_correction!(df; zeta=df.zeta, stab_formulation, constants)
  # pre-estimate z0m to use it in both Gb and Ga
  needs_windprofile =
    Gb_model isa Union{Choudhury1988, Su2001} ||
    Ram_model isa ResistanceWindProfile
  if needs_windprofile
    if isnothing(z0m)
      z0m = roughness_parameters(Roughness_wind_profile(), df; zh, zr, psi_m = df.psi_m).z0m
    end
    wind_zh = wind_profile(zh, df, d, z0m; stab_formulation, constants)
  end
  #
  # calculate canopy boundary layer conductance (Gb)
  Gb_model isa Thom1972 && compute_Gb!(df, Gb_model; constants)
  Gb_model isa ConstantKB1 && compute_Gb!(df, Gb_model; kB_h, constants)
  Gb_model isa Choudhury1988 && compute_Gb!(
    df, Gb_model; leafwidth, LAI, wind_zh, constants)
  Gb_model isa Su2001 && compute_Gb!(
    df, Gb_model; wind_zh, Dl, fc, N, Cd, hs, LAI, constants)
  compute_Gb_quantities!(df)
  #
  # calculate aerodynamic risistance for momentum (Ra_m)
  Ram_model isa ResistanceWindProfile && compute_Ram!(df, Ram_model; zr, d, z0m, constants)
  Ram_model isa ResistanceWindZr && compute_Ram!(df, Ram_model)
  fr = (Ra_m, Gb_h, Gb_CO2) -> begin
    Ga_m = 1/Ra_m
    Ra_h = Ra_m + 1/Gb_h
    Ga_h = 1/Ra_h
    Ga_CO2 =  1/(Ra_m + 1/Gb_CO2)
    (;Ga_m, Ga_h, Ra_h, Ga_CO2)
  end
  transform!(df, [:Ra_m, :Gb_h, :Gb_CO2] => ByRow(fr) => AsTable)
  if !(isnothing(z0m))
    df[!,:z0h] .= roughness_length_heat.(z0m, df.kB_h)
  else
    df[!,:z0h] .= missing
  end
  df
end

"""
    roughness_length_heat(z0m, kB_h)

# Arguments
- `z0m` : Roughness length for momentum (m). Can be calculated
          by [`roughness_parameters`](@ref).
- `kB_h` : kB^(-1) parameter of heat transfer, Output of [`aerodynamic_conductance!`](@ref)

# Details
The roughness length for water and heat (z0h) is calculated from the
relationship (e.g. Verma 1989):

``{k_B}_h = ln(z_{0m}/z_{0h})``

it follows:

``z_{0h} = z_{0m} / e^{k_{B_h}}``

# References

- Verma, S., 1989: Aerodynamic resistances to transfers of heat, mass and momentum.
  In: Estimation of areal evapotranspiration, IAHS Pub, 177, 13-20.
- Rigden, A., Li, D., Salvucci, G., 2018: Dependence of thermal roughness length on friction
  velocity across land cover types: A synthesis analysis using AmeriFlux data. Agricultural
  and Forest Meteorology 249, 512-519.
"""
roughness_length_heat(z0m, kB_h) = z0m / exp(kB_h)

abstract type ResistanceMethod end
struct ResistanceWindZr <: ResistanceMethod end
struct ResistanceWindProfile <: ResistanceMethod end

"""
    compute_Ram(::ResistanceWindProfile(), ustar;
      zr, d, z0m, psi_h, constants=BigleafConstants())
    compute_Ram!(df, method::ResistanceWindProfile();
      zr, d, z0m, psi_h = df.psi_h, kwargs...)

    compute_Ram(::ResistanceWindZr(), ustar, wind)
    compute_Ram!(df, method::ResistanceWindZr(); kwargs...)

Estimate bulk aerodynamic conductance.

# Arguments
- `ustar`          : Friction velocity (m s-1)
- `df`             : DataFrame with above columns
- `zr`             : Instrument (reference) height (m)
- `d`              : Zero-plane displacement height (-), can be estimated using
                      [`roughness_parameters`](@ref)
- `z0m`            : Roughness length for momentum (m). Can be estimated using
                     from [`roughness_parameters`](@ref)
- `psi_h`          : the value of the stability function for heat and water vapor (-)
                     see [`stability_correction`](@ref)

# Details

The aerodynamic resistance for momentum ``R_{a_m}`` is given by (`Ram_method = ResistanceWindZr()`):

``R_{a_m} = u/{u^*}^2``

Where u is the horizontal wind velocity.
Note that this formulation accounts for changes in atmospheric stability, and does not
require an additional stability correction function.

An alternative method to calculate ``Ra_m`` is provided (`Ram_method = ResistanceWindProfile()`):

``R_{a_m} = (ln((z_r - d)/z_{0m}) - \\psi_h) / (k \\, u^*)``

If the roughness parameters `z0m` and `d` are unknown, they can be estimated using
[`roughness_parameters`](@ref). The argument `stab_formulation` determines the stability
correction function used to account for the effect of atmospheric stability on `Ra_m`
(`Ra_m` is lower for unstable and higher for stable stratification). Stratification is based
on a stability parameter `zeta` ``\\zeta=(z-d/L)``, where `z` is the height,
`d` the zero-plane displacement height, and `L` the Monin-Obukhov length, calculated with
[`Monin_Obukhov_length`](@ref)
The stability correction function is chosen by the argument `stab_formulation`.
Options are `Dyer1970()` and `Businger1971()` and `NoStabilityCorrection()`.

# Note
For adding aerodynamic conductance for other species see [`add_Ga!`](@ref).

# Value
Aerodynamic resistance for momentum transfer (s m-1) (``Ra_m``)

# References
- Verma, S., 1989: Aerodynamic resistances to transfers of heat, mass and momentum.
  In: Estimation of areal evapotranspiration, IAHS Pub, 177, 13-20.
- Verhoef, A., De Bruin, H., Van Den Hurk, B., 1997: Some practical notes on the parameter kB^(-1)
  for sparse vegetation. Journal of Applied Meteorology, 36, 560-572.
- Hicks, B_B., Baldocchi, D_D., Meyers, T_P., Hosker, J_R., Matt, D_R., 1987:
  A preliminary multiple resistance routine for deriving dry deposition velocities
  from measured quantities. Water, Air, and Soil Pollution 36, 311-330.
- Monteith, J_L., Unsworth, M_H., 2008: Principles of environmental physics.
  Third Edition. Elsevier Academic Press, Burlington, USA.

# See also
[`aerodynamic_conductance!`](@ref), [`add_Ga!`](@ref)
"""
function compute_Ram(::ResistanceWindProfile, ustar::Union{Missing,Number};
  zr, d, z0m, psi_h, constants=BigleafConstants()
  )
  Ra_m = max((log((zr - d)/z0m) - psi_h),0) / (oftype(ustar,constants.k)*ustar)
end
function compute_Ram!(df, method::ResistanceWindProfile;
  zr, d, z0m, psi_h = df.psi_h, kwargs...
  )
  # put keyword arguments to positional arguments for proper broadcast
  ftpos(ustar, zr, d, z0m, psi_h) =
    compute_Ram(method, ustar; zr, d, z0m, psi_h, kwargs...)
  ft = (ustar) -> ftpos.(ustar, zr, d, z0m, psi_h)
  transform!(df, :ustar => ft => :Ra_m)
end

function compute_Ram(::ResistanceWindZr, ustar::Union{Missing,Number}, wind)
  Ra_m = wind / ustar^2
end
function compute_Ram!(df, method::ResistanceWindZr; kwargs...)
  fr = (ustar, wind) -> compute_Ram(method, ustar, wind; kwargs...)
  transform!(df, SA[:ustar, :wind] => ByRow(fr) => :Ra_m)
end

"""
    add_Ga(Gb_h, Ga_m, Sc::Vararg{Pair,N}; constants)
    add_Ga!(df::AbstractDataFrame, Sc; Gb_h = df.Gb_h, Ga_m = df.Ga.m, kwargs...)

compute additional aerodynamic conductance quantities for given Schmidt-numbers

# Arguments
- `Gb_h`     : Boundary layer conductance for heat transfer (m s-1)
- `Ga_m`     : Aerodynamic conductance for momentum (m s-1)
- `Sc`       : several `Pair{Symbol,Number}` Output name and Schmidt number of
               additional conductances to be calculated
- `df`       : DataFrame to add output columns
optional
- `constants=`[`BigleafConstants`](@ref)`()`: Dictionary with entries
  - `Pr` - Prandtl number

# Details
Aerodynamic conductance is calculated as

``G_{a_x} = 1/(1/G_{a_m} + 1/G_{b_x})``

where `Gb_x` is the Boundary layer conductance for other quantities x is calculated
based on boundary layer for heat transfer, Schmidt-Number, and  Prantl number,
as documented in [`add_Gb!`](@ref).

# Value
a NameTuple or `df` with keys `Ga_x` where `x` are the keys in `Sc` and
corresponding aerodynamic conductances (m s-1).

# Examples
```jldoctest; output=false
using DataFrames
df = DataFrame(Gb_h=[0.02, missing, 0.055], Ga_m = 1 ./ [0.03, 0.03, 0.03])
add_Ga!(df, :O2 => 0.84, :CH4 => 0.99)
propertynames(df)[3:4] == [:Ga_O2, :Ga_CH4]
# output
true
```
"""
function add_Ga!(df::AbstractDataFrame, Sc::Vararg{Pair,N};
  Gb_h = df.Gb_h, Ga_m = df.Ga_m, kwargs...) where N
  N == 0 && return(df)
  Scn, Scv = get_names_and_values("Ga_", Sc...)
  ft = () -> add_Ga_.(Gb_h, Ga_m, Ref(Scn), Ref(Scv); kwargs...)
  transform!(df, [] => ft => AsTable)
end
function add_Ga(Gb_h::Union{Missing,Number}, Ga_m::Union{Missing,Number},
  Sc::Vararg{Pair,N}; kwargs...) where N
  Scn, Scv = get_names_and_values("Ga_", Sc...)
  add_Ga_(Gb_h, Ga_m, Scn, Scv; kwargs...)
end
function add_Ga_(Gb_h::Union{Missing,Number}, Ga_m::Union{Missing,Number},
  Scn::NTuple{N,Symbol}, Scv::NTuple{N};  constants=BigleafConstants()) where N
  Gbx = add_Gb_(Gb_h, Scn, Scv; constants)
  Gaxv = ntuple(i -> 1/(1/Ga_m + 1/Gbx[i]), length(Gbx))
  Gax = NamedTuple{Scn}(Gaxv)
end


text:
module Calphad
using Symbolics
import LinearAlgebra
using OrderedCollections

export CompSet, PhaseRecord
export R, P, T

include("constants.jl")
include("structs.jl")
include("solver.jl")

end


text:
using AssemblyCalculus: PartialArea, generate_synapses

@testset "Constructors" begin
    pna = PartialArea(Float32)
    @test length(pna) == 0
    @test num_neurons(pna) == 0
    pna = PartialArea(10, 3, 0.01f0, 0.2f0)
    @test length(pna) == 10
    @test num_neurons(pna) == 10
end

@testset "Utilities" begin
    T = Float32
    area1 = PartialArea(T)
    neuron1 = Neuron(area1, 1)
    push!(area1.neurons, neuron1)
    area2 = PartialArea(T)
    neuron2 = Neuron(area2, 1)
    push!(area2.neurons, neuron2)
    neuron1[area2, neuron2] = 1.5f0
    neuron2[area1, neuron1] = 1.7f0
    @test num_synapses(area1) == 1
    @test num_synapses(area2) == 1
    for n in area1
        @test num_synapses(n) == 1
    end
    @test length(area1.firing) == 0
    area1.assembly_size = 1
    random_firing!(area1)
    @test length(area1.firing) == 1
    area1.assembly_size = 2
    @test_throws ErrorException random_firing!(area1)
    # k-winners
    area1.assembly_size = 3
    @test_throws ErrorException winners!(area1, rand(Float32, 3))
    for i in 2:9
        push!(area1.neurons, Neuron(area1, i))
    end
    currents = zeros(Float32, 10)
    @test_throws ErrorException winners!(area1, currents)
    push!(area1.neurons, Neuron(area1, 10))
    winners!(area1, currents)
    @test length(area1.firing) == 0
    currents[2:4] .= 1.0f0
    winners!(area1, currents)
    @test all(area1.firing .== [2, 3, 4])
    currents = rand(Float32, 10)
    perm = sortperm(currents, rev=true)
    winners!(area1, currents)
    @test all(area1.firing .== perm[1:3])
    # Prev assembly comparison
    area1.firing = [1, 2, 3]
    area1.firing_prev = [1, 2, 3]
    @test change_in_assembly(area1) == 0
    area1.firing = [1, 2, 4]
    @test change_in_assembly(area1) == 1
end

@testset "Hebbian Updates" begin
    # Make an area
    num_neurons = 10
    T = Float32
    area1 = PartialArea(num_neurons, 3, 1.f0, 0f0)
    area2 = PartialArea(num_neurons, 3, 1.f0, 0f0)
    current = ones(T, 10)
    ag = neuron_attrib_graph(T, PartialArea{T})
    # Make synapses
    for i in 1:num_neurons
        source = area1[i]
        target = area2[i]
        source[area2, target] = 1.f0
        if i > (num_neurons / 2)
            contrib = get!(ag, target)
            push!(contrib, source)
        end
    end

    # Ion current array update
    hebb_update!(area1, current)
    @test all(current .== 1.f0)
    random_firing!(area1)
    hebb_update!(area1, current)
    @test length(findall(current .> 1f0)) == 3
    area2.firing =[1, 9, 10]
    # Synapse Update
    hebb_update!(area2, ag)
    syn = map(x -> collect(values(x[area2]))[1], area1)
    @test all(syn[1:8] .== 1f0)
    @test all(syn[9:10] .== 2f0)

    init_currents = zeros(T, 10)
    init_currents[1:5] .= rand(T, 5)
    currents = zeros(T, 10)
    pows = [1., 2, 4, 1, 8]
    currents[1:5] .= init_currents[1:5] .* (2f0 .^ pows)
    ag = neuron_attrib_graph(T, PartialArea{T})
    ic_currents = zeros(T, 10)
    ic_currents[1:5] .= init_currents[1:5]
    ic = Stimulus(area2, ic_currents)
    hebb_update!(ag, ic, currents, init_currents)
    @test all(ic.currents .== currents)

    ic_currents = zeros(T, 10)
    ic_currents[1:3] .= init_currents[1:3]
    ic = Stimulus(area2, ic_currents)
    for i in 1:3
        contrib = get!(ag, area2[i])
        push!(contrib, area1[i])
    end
    hebb_update!(ag, ic, currents, init_currents)

    @test all(ic.currents[1:3] .== currents[1:3])
    @test all(ic.currents[4:5] .!= currents[4:5])
    syn = map(x -> collect(values(x[area2]))[1], area1)
    @test all(syn[1:3] .== (2f0 .^ pows[1:3]))
    @test all(syn[4:5] .== 1f0)
end

@testset "Firing" begin
    T = Float32
    num_neurons = 10
    areas = [PartialArea(num_neurons, 3, 1.f0, 0f0) for i in 1:3]
    a1, a2, a3 = areas
    # Make synapses
    for i in 1:num_neurons
        source = a1[i]
        target = a2[i]
        target2 = a3[i]
        source[a2, target] = 1.1f0
        source[a3, target2] = 1.2f0
        target[a3, target2] = 1.3f0
    end

    # Assembly firing
    currents = zeros(Float32, num_neurons)
    random_firing!.(areas)
    assem = Assembly(a2, Stimulus{T}[], Assembly{T}[])
    ag = neuron_attrib_graph(Float32, typeof(a2))
    fire!(assem, a1, currents, ag)
    @test all(currents .== 0f0)
    @test length(ag) == 0
    fire!(assem, a3, currents, ag)
    @test all(currents[assem.neurons] .== 1.3f0)
    currents[assem.neurons] .= 0f0
    @test all(currents .== 0f0)
    @test length(ag) .== assem.area.assembly_size
    @test all([length(v) for v in values(ag)] .== 1)

    # Active Area firing
    active_areas = [a1, a3]
    area_currents  = [zeros(T, num_neurons) for i in 1:2]
    ag = fire!(active_areas, area_currents)
    @test all(area_currents[1] .== 0f0)
    @test all(area_currents[2][a1.firing] .== 1.2f0)
    @test length(ag) == a1.assembly_size
    @test all([length(v) for v in values(ag)] .== 1)
    @test length(intersect(collect(keys(ag)), a3.neurons)) == a1.assembly_size
end

@testset "Simulation" begin
    k = 100
    n = k ^ 2
    β = 1f0
    mean_degree = 100f0
    p = mean_degree / n

    ba = Brain(num_areas=1, n=n, k=k, β=β, p=p)
    random_firing!(ba[1])
    init_firing = ba[1].firing
    stims = [rand_stim(ba[1])]
    init_currents = deepcopy(stims[1].currents)
    assems = Assembly{Float32}[]
    new_assems, spikes, dists = simulate!(
        stims,
        assems,
        tol=0,
        fire_past_convergence=10,
        random_initial=false
    )
    @test length(new_assems) == 1
    @test all([length(a.neurons) for a in assems] .== n)

    adj = adjacency_matrix(ba)
    adj[adj .!= 0f0] .= 1.f0
    ba2 = Brain(adj, [n], [k], [β])
    ba2[1].firing = init_firing
    stims2 = [Stimulus(ba2[1], init_currents)]
    new_assems2, spikes2, dists2 = simulate!(
        stims2,
        assems,
        tol=0,
        fire_past_convergence=10,
        random_initial=false
    )
    @test all(dists2 .== dists)

end


text:
@info "---------- Modal test ----------"
st=Structure()
lcset=LoadCaseSet()

add_uniaxial_metal!(st,"steel",2e11,0.2,7849.0474)
add_general_section!(st,"frame",4.26e-3,3.301e-6,6.572e-5,9.651e-8,1e-3,1e-3,0,0)

add_node!(st,1,0,0,0)
add_node!(st,2,18,18,0)
add_beam!(st,1,1,2,"steel","frame")

set_nodal_restraint!(st,1,true,true,true,true,true,true)

add_modal_case!(lcset,"MODAL")
set_modal_params!(lcset,"MODAL",modal_type="eigen")

@info "---------- Time history test ----------"

t=0:1/64:20
f=sin.(t)

add_static_case!(lcset,"STATIC")
add_nodal_force!(lcset,"STATIC",2,0,0,-1,0,0,0)

add_time_history_case!(lcset,"DIFF",t,f)
set_time_history_params!(lcset,"DIFF",0)
add_nodal_force!(lcset,"DIFF",2,0,0,-1,0,0,0)

add_time_history_case!(lcset,"NEWMARK",t,f)
add_nodal_force!(lcset,"NEWMARK",2,0,0,-1,0,0,0)
add_nodal_force!(lcset,"NEWMARK",2,0,0,-1,0,0,0)

add_time_history_case!(lcset,"WILSON",t,f)
set_time_history_params!(lcset,"WILSON",2)
add_nodal_force!(lcset,"WILSON",2,0,0,-1,0,0,0)

add_time_history_case!(lcset,"MODALDECOMP",t,f)
set_time_history_params!(lcset,"MODALDECOMP",3,modal_case="MODAL")
add_nodal_force!(lcset,"MODALDECOMP",2,0,0,-1,0,0,0)

assembly=assemble!(st,lcset,path=PATH)
solve(assembly)

T=result_modal_period(assembly,"MODAL")
@show T
# u=result_nodal_time_history(assembly,"NEWMARK",2,0,3)
r=result_nodal_time_history(assembly,"MODALDECOMP",2,0,3)