text:
### A Pluto.jl notebook ###
# v0.12.20

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ db4c1f10-7c37-4513-887a-2467ce673458
begin
	using Pkg
	Pkg.add.(["CSV", "DataFrames", "PlutoUI", "Shapefile", "ZipFile", "LsqFit", "Plots", PackageSpec(name="GR", version="0.53"), "StatsPlots"])
	using CSV
	using DataFrames
	using PlutoUI
	using Shapefile
	using ZipFile
	using LsqFit
	using Plots
	using StatsPlots
end

# ╔═╡ cbd9c1aa-fc37-11ea-29d9-e3361406796f
using Dates

# ╔═╡ 0f87cec6-fc31-11ea-23d2-395e61f38b6f
md"# Module 2: Epidemic propagation"

# ╔═╡ 19f4da16-fc31-11ea-0de9-1dbe668b862d
md"We are starting a new module on modelling epidemic propagation.

Let's start off by analysing some of the data that is now available on the current COVID-19 pandemic.
"

# ╔═╡ d3398953-afee-4989-932c-995c3ffc0c40
md"""
## Exploring COVID-19 data
"""

# ╔═╡ efa281da-cef9-41bc-923e-625140ce5a07
md"""
In this notebook we will explore and analyse data on the COVID-19 pandemic. The aim is to use Julia tools to analyse and visualise the data in different ways.

By the end of the notebook we will produce the following visualisation using Julia and Pluto:
"""

# ╔═╡ 7617d970-fce4-11ea-08ba-c7eba3e17f62
@bind day Clock(0.5)

# ╔═╡ e0493940-8aa7-4733-af72-cd6bc0e37d92
md"""
## Download and load data
"""

# ╔═╡ 64d9bcea-7c85-421d-8f1e-17ea8ee694da
url = "https://raw.githubusercontent.com/CSSEGISandData/COVID-19/master/csse_covid_19_data/csse_covid_19_time_series/time_series_covid19_confirmed_global.csv";

# ╔═╡ c460b0c3-6d3b-439b-8cc7-1c58d6547f51
download(url, "covid_data.csv");

# ╔═╡ a7369222-fc20-11ea-314d-4d6b0f0f72eb
md"We will need a couple of new packages. The data is in CSV format, i.e. *C*omma-*S*eparated *V*alues. This is a common data format in which observations, i.e. data points, are separated on different lines. Within each line the different data for that observation are separated by commas or other punctuation (possibly spaces and tabs)."

# ╔═╡ 9073f070-6656-11eb-3ee5-234f5a2461d2
GR.version()

# ╔═╡ 1ca02510-6656-11eb-09a4-ed8597cef105
Plots.GR.version()

# ╔═╡ 1620aa9d-7dcd-4686-b7e4-a72cebe315ed
md"""
We can load the data from a CSV using the `File` function from the `CSV.jl` package, and then convert it to a `DataFrame`:
"""

# ╔═╡ 38344160-fc27-11ea-220e-95aa00e4b083
begin
	csv_data = CSV.File("covid_data.csv");
	data = DataFrame(csv_data)   # it is common to use `df` as a variable name
end

# ╔═╡ ad43cea2-fc28-11ea-2bc3-a9d81e3766f4
md"A `DataFrame` is a standard way of storing **heterogeneous data** in Julia, i.e. a table consisting of columns with different types. As you can see from the display of the `DataFrame` object above, each column has an associated type, but different columns have different types, reflecting the type of the data in that column.

In our case, country names are stored as `String`s, their latitude and longitude as `Float64`s and the (cumulative) case counts for each day as `Int64`s.
."

# ╔═╡ fab64d86-fc28-11ea-0ae1-3ba1b9a14759
md"## Using the data"

# ╔═╡ 3519cf96-fc26-11ea-3386-d97c61ea1b85
md"""Since we need to manipulate the columns, let's rename them to something shorter. We can do this either **in place**, i.e. modifying the original `DataFrame`, or **out of place**, creating a new `DataFrame`. The convention in Julia is that functions that modify their argument have a name ending with `!` (often pronounced "bang").

We can use the `first` function to see only the first few lines of the data.
"""

# ╔═╡ a054e048-4fea-487c-9d06-463723c7151c
begin
	data_2 = rename(data, 1 => "province", 2 => "country", 3 => "latitude", 4 => "longitude")
	first(data_2, 5)
end

# ╔═╡ e9ad97b6-fdef-4f48-bd32-634cfd2ce0e6
begin
	rename!(data, 1 => "province", 2 => "country", 3 => "latitude", 4 => "longitude")
	first(data)
end

# ╔═╡ aaa7c012-fc1f-11ea-3c6c-89630affb1db
md"## Extracting useful information"

# ╔═╡ b0eb3918-fc1f-11ea-238b-7f5d23e424bb
md"How can we extract the list of all the countries? The country names are in the second column.

For some purposes we can think of a `DataFrame`.as a matrix and use similar syntax. For example, we can extract the second column:
"

# ╔═╡ 68f76d3b-b398-459d-bf39-20bf300dcaa2
all_countries = data[:, "country"]

# ╔═╡ 20e144f2-fcfb-11ea-010c-97e21eb0d231
all_countries2 = data[:, :country]

# ╔═╡ 2ec98a16-fcfb-11ea-21ad-15f2f5e68248
all_countries3 = data[:, 2]

# ╔═╡ 382cfc62-fcfb-11ea-26aa-2984d0449dcc
data[5:8, 2]

# ╔═╡ 34440afc-fc2e-11ea-0484-5b47af235bad
md"It turns out that some countries are divided into provinces, so there are repetitions in the `country` column that we can eliminate with the `unique` function:"

# ╔═╡ 79ba0433-2a31-475a-87c9-14103ebbff16
countries = unique(all_countries)

# ╔═╡ 5c1ec9ae-fc2e-11ea-397d-937c7ab1edb2
@bind i Slider(1:length(countries), show_value=true)

# ╔═╡ a39589ee-20e3-4f22-bf81-167fd815f6f9
md"$(Text(countries[i]))"

# ╔═╡ 9484ea9e-fc2e-11ea-137c-6da8212da5bd
md"[Here we used **string interpolation** with `$` to put the text into a Markdown string.]"

# ╔═╡ bcc95a8a-fc2e-11ea-2ccd-3bece42a08e6
md"You can also use `Select` to get a dropdown instead:"

# ╔═╡ ada3ceb4-fc2e-11ea-2cbf-399430fa18b5
@bind country Select(countries)

# ╔═╡ 1633abe8-fc2f-11ea-2c7e-21b3348a3569
md"""How can we extract the data for a particular country? First we need to know the exact name of the country. E.g. is the US written as "USA", or "United States"?

We could scroll through to find out, or **filter** the data to only look at a sample of it, for example those countries that begin with the letter "U".

One way to do this is with an array comprehension:"""

# ╔═╡ ed383524-e0c0-4da2-9a98-ca75aadd2c9e
md"""
Array comprehension:
"""

# ╔═╡ 90810d7e-fcfb-11ea-396a-35543dcc1e06
startswith("david", "d")

# ╔═╡ 977e1a2c-fcfb-11ea-08e9-cd656a631778
startswith("hello", "d")

# ╔═╡ 9ee79840-30ff-4c92-97f4-e178caceceaf
U_countries = [startswith(country, "U") for country in all_countries]

# ╔═╡ 99d5a138-fc30-11ea-2977-71732ca3aead
length(U_countries)

# ╔═╡ 450b4902-fc30-11ea-321d-29faf6188ff5
md"Note that this returns an array of booleans of the same length as the vector `all_countries`. We can now use this to index into the `DataFrame`:"

# ╔═╡ 4f423a75-43da-486f-ac2a-7220032dac9f
data[U_countries, :]

# ╔═╡ a8b2db96-fc30-11ea-2eea-b938a3a430fb
md"""We see that the correct spelling is `"US"`. (And note how the different provinces of the UK are separated.)"""

# ╔═╡ c400ce4e-fc30-11ea-13b1-b54cf8f5630e
md"Now we would like to extract the data for the US alone. How can we access the correct row of the table? We can again filter on the country name. A nicer way to do this is to use the `filter` function.

This is a **higher-order function**: its first argument is itself a function, which must return `true` or `false`.  `filter` will return all the rows of the `DataFrame` that satisfy that **predicate**:
"

# ╔═╡ 7b2496b0-fc35-11ea-0e78-473e5e8eac44
filter(x -> x.country == "United Kingdom", data)

# ╔═╡ 8990f13a-fc35-11ea-338f-0955eeb23c3c
md"Here we have used an **anonymous function** with the syntax `x -> ⋯`. This is a function which takes the argument `x` and returns whatever is on the right of the arrow (`->`)."

# ╔═╡ a772eadc-fc35-11ea-3d38-4b121f88f1d7
md"To extract a single row we need the **index** of the row (i.e. which number row it is in the `DataFrame`). The `findfirst` function finds the first row that satisfies the given predicate:"

# ╔═╡ ef6bd9b0-6658-11eb-3eb9-c797ee94e8d5
==("US")

# ╔═╡ 16a79308-fc36-11ea-16e5-e1087d7ebbda
US_row = findfirst(==("US"), all_countries)

# ╔═╡ a41db8ea-f0e0-461f-a298-bdcea42a67f3
data[US_row, :]

# ╔═╡ f75e1992-fcfb-11ea-1123-b59bf888eac3
data[US_row:US_row, :]

# ╔═╡ 67eebb7e-fc36-11ea-03ef-bd6966487bb5
md"Now we can extract the data into a standard Julia `Vector`:"

# ╔═╡ 7b5db0f4-fc36-11ea-09a5-49def64f4c79
US_data = Vector(data[US_row, 5:end])

# ╔═╡ a26b8742-6a16-445a-ae77-25a4189c0f14

# ╔═╡ f099424c-0e22-42fb-894c-d8c2a65715fb
scatter(US_data, m=:o, alpha=0.5, ms=3, xlabel="day", ylabel="cumulative cases", leg=false)

# ╔═╡ 07c43b90-665b-11eb-37d9-a1ffcbe430c5
data[US_row, 5:end]

# ╔═╡ 7e7d14a2-fc37-11ea-3f1a-870ca98c4b75
md"Note that we are only passing a single vector to the `scatter` function, so the $x$ coordinates are taken as the natural numbers $1$, $2$, etc.

Also note that the $y$-axis in this plot gives the *cumulative* case numbers, i.e. the *total* number of confirmed cases since the start of the epidemic up to the given date.
"

# ╔═╡ 75d2dc66-fc47-11ea-0e35-05f9cf38e901
md"This is an example of a **time series**, i.e. a single quantity that changes over time."

# ╔═╡ b3880f40-fc36-11ea-074a-edc51adeb6f0
md"## Using dates"

# ╔═╡ 6de0800c-fc37-11ea-0d94-2b6f8f66964d
md"We would like to use actual dates instead of just the number of days since the start of the recorded data. The dates are given in the column names of the `DataFrame`:
"

# ╔═╡ bb6316b7-23fb-44a3-b64a-dfb71a7df011
column_names = names(data)

# ╔═╡ 0c098923-b016-4c65-9a37-6b7b56b13a0c
date_strings = names(data)[5:end]  # apply String function to each element

# ╔═╡ 546a40eb-7897-485d-a1b5-c4dfae0a4861
md"""
Now we need to **parse** the date strings, i.e. convert from a string representation into an actual Julia type provided by the `Dates.jl` standard library package:
"""

# ╔═╡ 9e23b0e2-ac13-4d19-a3f9-4a655a1e9f14
date_strings[1]

# ╔═╡ 25c79620-14f4-45a7-b120-05ec72cb77e9
date_format = dateformat"m/d/Y"

# ╔═╡ 31dc4e46-4839-4f01-b383-1a1189aeb0e6
parse(Date, date_strings[1], date_format)

# ╔═╡ ee27bd98-fc37-11ea-163c-1365e194fc2e
md"Since the year was not correctly represented in the original data, we need to manually fix it:"

# ╔═╡ f5c29f0d-937f-4731-8f87-0405ebc966f5
dates = parse.(Date, date_strings, date_format) .+ Year(2000)

# ╔═╡ b0e7f1c6-fce3-11ea-10e5-9101d0f861a2
dates[day]

# ╔═╡ 36c37b4d-eb23-4deb-a593-e511eccd9204
begin
	plot(dates, US_data, xrotation=45, leg=:topleft,
	    label="US data", m=:o, ms=3, alpha=0.5)

	xlabel!("date")
	ylabel!("cumulative US cases")
	title!("US cumulative confirmed COVID-19 cases")
end

# ╔═╡ 0604b46e-665f-11eb-3411-3f635ac8960e
data_long = stack(data, 5:382, variable_name = :date, value_name = :cases)

# ╔═╡ 8df98b80-665f-11eb-0be9-91b0d8467573
data_long.date = parse.(Date, data_long.date, date_format) .+ Year(2000)

# ╔═╡ f7824200-665e-11eb-132d-0d458a20601b
my_data_long = filter(x -> x.province in ["Netherlands", "US", "United Kingdom"],
	data_long)

# ╔═╡ 10c20890-665f-11eb-3b51-1338b0a62986
@df my_data_long scatter(:date, :cases, group=:country)

# ╔═╡ 511eb51e-fc38-11ea-0492-19532da809de
md"## Exploratory data analysis"

# ╔═╡ d228e232-fc39-11ea-1569-a31b817118c4
md"
Working with *cumulative* data is often less intuitive. Let's look at the actual number of daily cases. Julia has a `diff` function to calculate the difference between successive entries of a vector:
"

# ╔═╡ dbaacbb6-fc3b-11ea-0a42-a9792e8a6c4c
begin
	daily_cases = diff(US_data)
	plot(dates[2:end], daily_cases, m=:o, leg=false, xlabel="days", ylabel="daily US cases", alpha=0.5)   # use "o"-shaped markers
end

# ╔═╡ 19bdf146-fc3c-11ea-3c60-bf7823c43a1d
begin
	using Statistics
	running_mean = [mean(daily_cases[i-6:i]) for i in 7:length(daily_cases)]
end

# ╔═╡ 12900562-fc3a-11ea-25e1-f7c91a6940e5
md"Note that discrete data should *always* be plotted with points. The lines are just to guide the eye.

Cumulating data corresponds to taking the integral of a function and is a *smoothing* operation. Note that the cumulative data is indeed visually smoother than the daily data.

The oscillations in the daily data seem to be due to a lower incidence of reporting at weekends. We could try to smooth this out by taking a **moving average**, say over the past week:
"

# ╔═╡ be868a52-fc3b-11ea-0b60-7fea05ffe8e9
begin
	plot(daily_cases, label="raw daily cases")
	plot!(running_mean, m=:o, label="running weakly mean", leg=:topleft)
end

# ╔═╡ 0b01120c-fc3d-11ea-1381-8bab939e6214
md"## Exponential growth

Simple models of epidemic spread often predict a period with **exponential growth**. Do the data corroborate this?
"

# ╔═╡ 252eff18-fc3d-11ea-0c18-7b130ada882e
md"""A visual check for this is to plot the data with a **logarithmic scale** on the $y$ axis (but a standard scale on the $x$ axis).

If we observe a straight line on such a semi-logarithmic plot, then we know that

$$\log(y) \sim \alpha x + \beta,$$

where we are using $\sim$ to denote approximate equality.

Taking exponentials of both sides gives

$$y \sim \exp(\alpha x + \beta),$$

i.e.

$$y \sim c \, \mathrm{e}^{\alpha x},$$

where $c$ is a constant (sometimes called a "pre-factor") and $\alpha$ is the exponential growth rate, found from the slope of the straight line on the semi-log plot.
"""

# ╔═╡ 9626d74a-fc3d-11ea-2ab3-978dc46c0f1f
md"""Since the data contains some zeros, we need to replace those with `NaN`s ("Not a Number"), which `Plots.jl` interprets as a signal to break the line"""

# ╔═╡ 4358c348-91aa-4c76-a443-0a9cefce0e83
begin
	plot(replace(daily_cases, 0 => NaN),
		yscale=:log10,
		leg=false, m=:o)

	xlabel!("day")
	ylabel!("confirmed cases in US")
	title!("US confirmed COVID-19 cases")
end

# ╔═╡ 687409a2-fc43-11ea-03e0-d9a7a48165a8
md"Let's zoom in on the region of the graph where the growth looks linear on this semi-log plot:"

# ╔═╡ 4f23c8fc-fc43-11ea-0e73-e5f89d14155c
xlims!(0, 100)

# ╔═╡ 91f99062-fc43-11ea-1b0e-afe8aa8a1c3d
exp_period = 38:60

# ╔═╡ 07282688-fc3e-11ea-2f9e-5b0581061e65
md"We see that there is a period lasting from around day $(first(exp_period)) to around day $(last(exp_period)) when the curve looks straight on the semi-log plot.
This corresponds to the following date range:"

# ╔═╡ 210cee94-fc3e-11ea-1a6e-7f88270354e1
dates[exp_period]

# ╔═╡ 2f254a9e-fc3e-11ea-2c02-75ed59f41903
md"i.e. the first 3 weeks of March. Fortunately the imposition of lockdown during the last 10 days of March (on different days in different US states) significantly reduced transmission."

# ╔═╡ 84f5c776-fce0-11ea-2d52-39c51d4ab6b5
md"## Data fitting"

# ╔═╡ 539c951c-fc48-11ea-2293-457b7717ea4d
md"""Let's try to fit an exponential function to our data in the relevant region. We will use the Julia package `LsqFit.jl` ("least-squares fit").

This package allows us to specify a model function that takes a vector of data and a vector of parameters, and it finds the best fit to the data.
"""

# ╔═╡ b33e97f2-fce0-11ea-2b4d-ffd7ed7000f8
model(x, (c, α)) = c .* exp.(α .* x)

# ╔═╡ d52fc8fe-fce0-11ea-0a04-b146ee2dbe80
begin
	p0 = [0.5, 0.5]  # initial guess for parameters

	x_data = exp_period
	y_data = daily_cases[exp_period]

	fit = curve_fit(model, x_data, y_data, p0)
end;

# ╔═╡ c50b5e42-fce1-11ea-1667-91c56ea80dcc
md"We are interested in the coefficients of the best-fitting model:"

# ╔═╡ 3060bfa8-fce1-11ea-1047-db0dc06485a2
parameters = coef(fit)

# ╔═╡ 62bdc04a-fce1-11ea-1724-bfc4bc4789d1
md"Now let's add this to the plot:"

# ╔═╡ 6bc8cc20-fce1-11ea-2180-0fa69e86741f
begin
	plot(replace(daily_cases, 0 => NaN),
		yscale=:log10,
		leg=false, m=:o,
		xlims=(1, 100), alpha=0.5)

	line_range = 30:70
	plot!(line_range, model(line_range, parameters), lw=3, ls=:dash, alpha=0.7)

	xlabel!("day")
	ylabel!("confirmed cases in US")
	title!("US confirmed COVID-19 cases")
end

# ╔═╡ 287f0fa8-fc44-11ea-2788-9f3ac4ee6d2b
md"## Geographical data"

# ╔═╡ 3edd2a22-fc4a-11ea-07e5-55ca6d7639e8
md"Our data set contains more information: the geographical locations (latitude and longitude) of each country (or, rather, of a particular point that was chosen as being representative of that country)."

# ╔═╡ c5ad4d40-fc57-11ea-23cb-e55487bc6f7a
filter(x -> startswith(x.country, "A"), data)

# ╔═╡ 57a9bb06-fc4a-11ea-2665-7f97026981dc
md"Let's extract and plot the geographical information. To reduce the visual noise a bit we will only use those "

# ╔═╡ 80138b30-fc4a-11ea-0e15-b54cf6b402df
province = data.province

# ╔═╡ 8709f208-fc4a-11ea-0203-e13eae5f0d93
md"If the `province` is missing we should use the country name instead:"

# ╔═╡ a29c8ad0-fc4a-11ea-14c7-71435769b73e
begin
	indices = ismissing.(province)
	province[indices] .= all_countries[indices]
end

# ╔═╡ 4e4cca22-fc4c-11ea-12ae-2b51545799ec
begin

	scatter(data.longitude, data.latitude, leg=false, alpha=0.5, ms=2)

	for i in 1:length(province)
		annotate!(data.longitude[i], data.latitude[i], text(province[i], :center, 5, color=RGBA{Float64}(0.0,0.0,0.0,0.3)))
	end

	plot!(axis=false)
end

# ╔═╡ 16981da0-fc4d-11ea-37a2-535aa014a298
data.latitude

# ╔═╡ a9c39dbe-fc4d-11ea-2e86-4992896e2abb
md"## Adding maps"

# ╔═╡ b93b88b0-fc4d-11ea-0c45-8f64983f8b5c
md"We would also like to see the outlines of each country. For this we can use, for example, the data from [Natural Earth](https://www.naturalearthdata.com/downloads/110m-cultural-vectors/110m-admin-0-countries), which comes in the form of **shape files**, giving the outlines in terms of latitude and longitude coordinates.

These may be read in using the `Shapefile.jl` package.

The data is provided in a `.zip` file, so after downloading it we first need to decompress it.
"

# ╔═╡ 7ec28cd0-fc87-11ea-2de5-1959ea5dc37c
begin
	zipfile = download("https://www.naturalearthdata.com/http//www.naturalearthdata.com/download/110m/cultural/ne_110m_admin_0_countries.zip")

	r = ZipFile.Reader(zipfile);
	for f in r.files
	    println("Filename: $(f.name)")
		open(f.name, "w") do io
	    	write(io, read(f))
		end
    end
	close(r)
end

# ╔═╡ ada44a56-fc56-11ea-2ab7-fb649be7e066
shp_countries = Shapefile.shapes(Shapefile.Table("./ne_110m_admin_0_countries.shp"))

# ╔═╡ d911edb6-fc87-11ea-2258-d34d61c02245

# ╔═╡ b3e1ebf8-fc56-11ea-05b8-ed0b9e50503d
plot!(shp_countries, alpha=0.2)

# ╔═╡ f8e754ee-fc73-11ea-0c7f-cdc760ab3e94
md"Now we would like to combine the geographical and temporal (time) aspects. One way to do so is to animate time:"

# ╔═╡ 39982810-fc76-11ea-01c3-3987cfc2fd3c
daily = max.(1, diff(Array(data[:, 5:end]), dims=2));

# ╔═╡ 0f329ece-fc74-11ea-1e02-bdbddf551ef3
@bind day2 Slider(1:size(daily, 2), show_value=true)

# @bind day Clock(0.5)

# ╔═╡ b406eec8-fc77-11ea-1a98-d36d6d3e2393
log10(maximum(daily[:, day]))

# ╔═╡ 1f30a1ac-fc74-11ea-2abf-abf437006bab
dates[day2]

# ╔═╡ 24934438-fc74-11ea-12e4-7f7e50f54029
world_plot = begin
	plot(shp_countries, alpha=0.2)
	scatter!(data.longitude, data.latitude, leg=false, ms=2*log10.(daily[:, day2]), alpha=0.7)
	xlabel!("latitude")
	ylabel!("longitude")
	title!("daily cases per country")
end

# ╔═╡ f7a37706-fcdf-11ea-048a-236b8ed0f1f3
world_plot

# ╔═╡ 251c06e4-fc77-11ea-1a0f-73139ba11e83
md"However, we should always be wary about visualisations such as these. Perhaps we should be plotting cases per capita instead of absolute numbers of cases. Or should we divide by the area of the country? Some countries, such as China and Canada, are divided into states or regions in the original data set -- but others, such as the US, are not. You should always check exactly what is being plotted!

Unfortunately, published visualisations often hide some of  this information. This emphasises the need to be able to get our hands on the data, create our own visualisations and draw our own conclusions."

# ╔═╡ Cell order:
# ╟─0f87cec6-fc31-11ea-23d2-395e61f38b6f
# ╟─19f4da16-fc31-11ea-0de9-1dbe668b862d
# ╟─d3398953-afee-4989-932c-995c3ffc0c40
# ╟─efa281da-cef9-41bc-923e-625140ce5a07
# ╠═7617d970-fce4-11ea-08ba-c7eba3e17f62
# ╠═b0e7f1c6-fce3-11ea-10e5-9101d0f861a2
# ╠═f7a37706-fcdf-11ea-048a-236b8ed0f1f3
# ╟─e0493940-8aa7-4733-af72-cd6bc0e37d92
# ╠═64d9bcea-7c85-421d-8f1e-17ea8ee694da
# ╠═c460b0c3-6d3b-439b-8cc7-1c58d6547f51
# ╟─a7369222-fc20-11ea-314d-4d6b0f0f72eb
# ╠═db4c1f10-7c37-4513-887a-2467ce673458
# ╠═9073f070-6656-11eb-3ee5-234f5a2461d2
# ╠═1ca02510-6656-11eb-09a4-ed8597cef105
# ╟─1620aa9d-7dcd-4686-b7e4-a72cebe315ed
# ╠═38344160-fc27-11ea-220e-95aa00e4b083
# ╟─ad43cea2-fc28-11ea-2bc3-a9d81e3766f4
# ╟─fab64d86-fc28-11ea-0ae1-3ba1b9a14759
# ╟─3519cf96-fc26-11ea-3386-d97c61ea1b85
# ╠═a054e048-4fea-487c-9d06-463723c7151c
# ╠═e9ad97b6-fdef-4f48-bd32-634cfd2ce0e6
# ╟─aaa7c012-fc1f-11ea-3c6c-89630affb1db
# ╟─b0eb3918-fc1f-11ea-238b-7f5d23e424bb
# ╠═68f76d3b-b398-459d-bf39-20bf300dcaa2
# ╠═20e144f2-fcfb-11ea-010c-97e21eb0d231
# ╠═2ec98a16-fcfb-11ea-21ad-15f2f5e68248
# ╠═382cfc62-fcfb-11ea-26aa-2984d0449dcc
# ╟─34440afc-fc2e-11ea-0484-5b47af235bad
# ╠═79ba0433-2a31-475a-87c9-14103ebbff16
# ╠═5c1ec9ae-fc2e-11ea-397d-937c7ab1edb2
# ╟─a39589ee-20e3-4f22-bf81-167fd815f6f9
# ╟─9484ea9e-fc2e-11ea-137c-6da8212da5bd
# ╟─bcc95a8a-fc2e-11ea-2ccd-3bece42a08e6
# ╠═ada3ceb4-fc2e-11ea-2cbf-399430fa18b5
# ╟─1633abe8-fc2f-11ea-2c7e-21b3348a3569
# ╟─ed383524-e0c0-4da2-9a98-ca75aadd2c9e
# ╠═90810d7e-fcfb-11ea-396a-35543dcc1e06
# ╠═977e1a2c-fcfb-11ea-08e9-cd656a631778
# ╠═9ee79840-30ff-4c92-97f4-e178caceceaf
# ╠═99d5a138-fc30-11ea-2977-71732ca3aead
# ╟─450b4902-fc30-11ea-321d-29faf6188ff5
# ╠═4f423a75-43da-486f-ac2a-7220032dac9f
# ╟─a8b2db96-fc30-11ea-2eea-b938a3a430fb
# ╟─c400ce4e-fc30-11ea-13b1-b54cf8f5630e
# ╠═7b2496b0-fc35-11ea-0e78-473e5e8eac44
# ╟─8990f13a-fc35-11ea-338f-0955eeb23c3c
# ╟─a772eadc-fc35-11ea-3d38-4b121f88f1d7
# ╠═ef6bd9b0-6658-11eb-3eb9-c797ee94e8d5
# ╠═16a79308-fc36-11ea-16e5-e1087d7ebbda
# ╠═a41db8ea-f0e0-461f-a298-bdcea42a67f3
# ╠═f75e1992-fcfb-11ea-1123-b59bf888eac3
# ╟─67eebb7e-fc36-11ea-03ef-bd6966487bb5
# ╠═7b5db0f4-fc36-11ea-09a5-49def64f4c79
# ╠═a26b8742-6a16-445a-ae77-25a4189c0f14
# ╠═f099424c-0e22-42fb-894c-d8c2a65715fb
# ╠═07c43b90-665b-11eb-37d9-a1ffcbe430c5
# ╟─7e7d14a2-fc37-11ea-3f1a-870ca98c4b75
# ╟─75d2dc66-fc47-11ea-0e35-05f9cf38e901
# ╟─b3880f40-fc36-11ea-074a-edc51adeb6f0
# ╟─6de0800c-fc37-11ea-0d94-2b6f8f66964d
# ╠═bb6316b7-23fb-44a3-b64a-dfb71a7df011
# ╠═0c098923-b016-4c65-9a37-6b7b56b13a0c
# ╟─546a40eb-7897-485d-a1b5-c4dfae0a4861
# ╠═cbd9c1aa-fc37-11ea-29d9-e3361406796f
# ╠═9e23b0e2-ac13-4d19-a3f9-4a655a1e9f14
# ╠═25c79620-14f4-45a7-b120-05ec72cb77e9
# ╠═31dc4e46-4839-4f01-b383-1a1189aeb0e6
# ╟─ee27bd98-fc37-11ea-163c-1365e194fc2e
# ╠═f5c29f0d-937f-4731-8f87-0405ebc966f5
# ╠═36c37b4d-eb23-4deb-a593-e511eccd9204
# ╠═0604b46e-665f-11eb-3411-3f635ac8960e
# ╠═8df98b80-665f-11eb-0be9-91b0d8467573
# ╠═f7824200-665e-11eb-132d-0d458a20601b
# ╠═10c20890-665f-11eb-3b51-1338b0a62986
# ╟─511eb51e-fc38-11ea-0492-19532da809de
# ╟─d228e232-fc39-11ea-1569-a31b817118c4
# ╠═dbaacbb6-fc3b-11ea-0a42-a9792e8a6c4c
# ╟─12900562-fc3a-11ea-25e1-f7c91a6940e5
# ╠═19bdf146-fc3c-11ea-3c60-bf7823c43a1d
# ╠═be868a52-fc3b-11ea-0b60-7fea05ffe8e9
# ╟─0b01120c-fc3d-11ea-1381-8bab939e6214
# ╟─252eff18-fc3d-11ea-0c18-7b130ada882e
# ╟─9626d74a-fc3d-11ea-2ab3-978dc46c0f1f
# ╠═4358c348-91aa-4c76-a443-0a9cefce0e83
# ╟─687409a2-fc43-11ea-03e0-d9a7a48165a8
# ╠═4f23c8fc-fc43-11ea-0e73-e5f89d14155c
# ╟─07282688-fc3e-11ea-2f9e-5b0581061e65
# ╠═91f99062-fc43-11ea-1b0e-afe8aa8a1c3d
# ╠═210cee94-fc3e-11ea-1a6e-7f88270354e1
# ╟─2f254a9e-fc3e-11ea-2c02-75ed59f41903
# ╟─84f5c776-fce0-11ea-2d52-39c51d4ab6b5
# ╟─539c951c-fc48-11ea-2293-457b7717ea4d
# ╠═b33e97f2-fce0-11ea-2b4d-ffd7ed7000f8
# ╠═d52fc8fe-fce0-11ea-0a04-b146ee2dbe80
# ╟─c50b5e42-fce1-11ea-1667-91c56ea80dcc
# ╟─3060bfa8-fce1-11ea-1047-db0dc06485a2
# ╟─62bdc04a-fce1-11ea-1724-bfc4bc4789d1
# ╠═6bc8cc20-fce1-11ea-2180-0fa69e86741f
# ╟─287f0fa8-fc44-11ea-2788-9f3ac4ee6d2b
# ╟─3edd2a22-fc4a-11ea-07e5-55ca6d7639e8
# ╠═c5ad4d40-fc57-11ea-23cb-e55487bc6f7a
# ╟─57a9bb06-fc4a-11ea-2665-7f97026981dc
# ╠═80138b30-fc4a-11ea-0e15-b54cf6b402df
# ╟─8709f208-fc4a-11ea-0203-e13eae5f0d93
# ╠═a29c8ad0-fc4a-11ea-14c7-71435769b73e
# ╠═4e4cca22-fc4c-11ea-12ae-2b51545799ec
# ╠═16981da0-fc4d-11ea-37a2-535aa014a298
# ╟─a9c39dbe-fc4d-11ea-2e86-4992896e2abb
# ╟─b93b88b0-fc4d-11ea-0c45-8f64983f8b5c
# ╠═7ec28cd0-fc87-11ea-2de5-1959ea5dc37c
# ╟─ada44a56-fc56-11ea-2ab7-fb649be7e066
# ╠═d911edb6-fc87-11ea-2258-d34d61c02245
# ╠═b3e1ebf8-fc56-11ea-05b8-ed0b9e50503d
# ╟─f8e754ee-fc73-11ea-0c7f-cdc760ab3e94
# ╠═39982810-fc76-11ea-01c3-3987cfc2fd3c
# ╠═0f329ece-fc74-11ea-1e02-bdbddf551ef3
# ╠═b406eec8-fc77-11ea-1a98-d36d6d3e2393
# ╠═1f30a1ac-fc74-11ea-2abf-abf437006bab
# ╠═24934438-fc74-11ea-12e4-7f7e50f54029
# ╟─251c06e4-fc77-11ea-1a0f-73139ba11e83

max_stars_repo_name: mathijsvdv/ComputationalThinking

max_issues_repo_name: mathijsvdv/ComputationalThinking

max_forks_repo_name: mathijsvdv/ComputationalThinking


text:
module TestReactions

using Base.Test
using SoilTracers.PhysChem: hydrolysis_cos

@time @testset "Reaction functions" begin
    temp_0C = 273.15
    temp_25C = 298.15

    # test hydrolysis_cos

    k_hyd_cos_0C = hydrolysis_cos(temp_0C)
    k_hyd_cos_25C = hydrolysis_cos(temp_25C)

    @test typeof(k_hyd_cos_0C) <: Real
    @test typeof(k_hyd_cos_25C) <: Real

    @test k_hyd_cos_0C < k_hyd_cos_25C
    @test isapprox(k_hyd_cos_25C, 2.26e-5, rtol=1e-3)
end

end

ext: jl


text:
using ForneyLab
using LinearAlgebra
import ForneyLab: SoftFactor, @ensureVariables, generateId, addNode!, associate!,
                  averageEnergy, Interface, Variable, slug, ProbabilityDistribution,
                  differentialEntropy, unsafeLogMean, unsafeMean, unsafeCov, unsafePrecision, unsafeMeanCov,
                  collectAverageEnergyInbounds, localEdgeToRegion, ultimatePartner, region, Region, isClamped, currentInferenceAlgorithm
import SpecialFunctions: digamma
export NAutoRegressiveMovingAverageX, NARMAX

"""
Description:

    A Nonlinear AutoRegressive model with moving average and exogeneous input (NARMAX)

    y_k = fθ(u_k, …, u_k-M1, y_k-1, …, y_k-M2, , e_k, …, e_k-M3)

    where M1 is the number of input delays, M2 the number of output delays
    and M3 the number of error delays. These histories are stored as the following
    vectors:
    - x_k-1 = [u_k-1, …, u_k-M1]
    - z_k-1 = [y_k-1, …, y_k-M2]
    - r_k-1 = [e_k-1, …, e_k-M3]

    Assume u_k, x_k-1, z_k-1 and r_k-1 are observed and e_k ~ N(0, τ^-1).

    !! Currently, fθ is assumed to be a linear product of coefficients θ
    and a basis expansion ϕ of inputs, outputs and errors:

        fθ(...) = θ'*ϕ(u_k, …, u_k-M1, y_k-1, …, y_k-M2, , e_k, …, e_k-M3)

Interfaces:

    1. y (output)
    2. θ (function coefficients)
    3. u (current input)
    4. x (previous inputs)
    5. z (previous outputs)
    6. r (previous errors)
    7. τ (precision)

Construction:

    NAutoRegressiveMovingAverageX(y, θ, u, x, z, r, τ, g=ϕ, id=:some_id)
"""

mutable struct NAutoRegressiveMovingAverageX <: SoftFactor
    id::Symbol
    interfaces::Vector{Interface}
    i::Dict{Symbol,Interface}

    ϕ::Function # Scalar function between autoregression coefficients and state variable

    function NAutoRegressiveMovingAverageX(y, θ, u, x, z, r, τ; ϕ::Function, id=generateId(NAutoRegressiveMovingAverageX))
        @ensureVariables(y, θ, x, u, z, r, τ)
        self = new(id, Array{Interface}(undef, 7), Dict{Symbol,Interface}(), ϕ)
        addNode!(currentGraph(), self)
        self.i[:y] = self.interfaces[1] = associate!(Interface(self), y)
        self.i[:θ] = self.interfaces[2] = associate!(Interface(self), θ)
        self.i[:u] = self.interfaces[3] = associate!(Interface(self), u)
        self.i[:x] = self.interfaces[4] = associate!(Interface(self), x)
        self.i[:z] = self.interfaces[5] = associate!(Interface(self), z)
        self.i[:r] = self.interfaces[6] = associate!(Interface(self), r)
        self.i[:τ] = self.interfaces[7] = associate!(Interface(self), τ)
        return self
    end
end

slug(::Type{NAutoRegressiveMovingAverageX}) = "NARMAX"

function averageEnergy(::Type{NAutoRegressiveMovingAverageX},
                       ϕ::Function,
                       marg_y::ProbabilityDistribution{Univariate},
                       marg_θ::ProbabilityDistribution{Multivariate},
                       marg_u::ProbabilityDistribution{Univariate},
                       marg_x::ProbabilityDistribution{Multivariate},
                       marg_z::ProbabilityDistribution{Multivariate},
                       marg_r::ProbabilityDistribution{Multivariate},
                       marg_τ::ProbabilityDistribution{Univariate})

    # Extract moments of beliefs
    mθ,Vθ = unsafeMeanCov(marg_θ)
    my = unsafeMean(marg_y)
    mu = unsafeMean(marg_u)
    mx = unsafeMean(marg_x)
    mz = unsafeMean(marg_z)
    mr = unsafeMean(marg_r)
    mτ = unsafeMean(marg_τ)
    aτ = marg_τ.params[:a]
    bτ = marg_τ.params[:b]

    # Basis expansion applied to moments
    ϕm = ϕ([mu; mx; mz; mr])

    # Return E_q -log p(y_k | u_k, .., θ, τ)
    return 1/2*log(2*pi) - 1/2*(digamma(aτ)-log(bτ)) + 1/2*mτ*(my^2 - 2*my*(mθ'*ϕm) + (mθ'*ϕm)^2 + ϕm'*Vθ*ϕm)
end

function collectAverageEnergyInbounds(node::NAutoRegressiveMovingAverageX)
    inbounds = Any[]

    # Push function to calling signature (ϕ needs to be defined in user scope)
    push!(inbounds, Dict{Symbol, Any}(:ϕ => node.ϕ, :keyword => false))

    local_edge_to_region = localEdgeToRegion(node)

    encountered_regions = Region[] # Keep track of encountered regions
    for node_interface in node.interfaces
        inbound_interface = ultimatePartner(node_interface)
        current_region = region(node_interface.node, node_interface.edge)

        if isClamped(inbound_interface)
            # Hard-code marginal of constant node in schedule
            push!(inbounds, assembleClamp!(copy(inbound_interface.node), ProbabilityDistribution)) # Copy Clamp before assembly to prevent overwriting dist_or_msg field
        elseif !(current_region in encountered_regions)
            # Collect marginal entry from marginal dictionary (if marginal entry is not already accepted)
            target = local_edge_to_region[node_interface.edge]
            current_inference_algorithm = currentInferenceAlgorithm()
            push!(inbounds, current_inference_algorithm.target_to_marginal_entry[target])
        end

        push!(encountered_regions, current_region)
    end

    return inbounds
end

max_stars_repo_name: wmkouw/NARMAX

max_issues_repo_name: wmkouw/NARMAX

max_forks_repo_name: wmkouw/NARMAX


text:
using Revise, Distributions, StatsFuns, Plots, ProgressMeter, AdvancedVI, ForwardDiff, DiffResults, LinearAlgebra, Random, Flux, StableRNGs, Zygote, ChainRulesCore, Nonconvex
Nonconvex.@load Ipopt

function mixturef(means, stds = ones(length(means)))
    return MixtureModel(Normal.(means, stds))
end

function kldivergence(p::MvNormal, q::MvNormal)
    # This is the generic implementation for AbstractMvNormal, you might need to specialize for your type
    length(p) == length(q) ||
        throw(DimensionMismatch("Distributions p and q have different dimensions $(length(p)) and $(length(q))"))
    # logdetcov is used separately from _cov for any potential optimization done there
    div = (tr(cov(q) \ cov(p)) + sqmahal(q, mean(p)) - length(p) + logdetcov(q) - logdetcov(p)) / 2
    return div
end
getq(θ) = MvNormal([θ[1]], Diagonal([exp(θ[2])]))

rng = StableRNG(123)
Random.seed!(rng, 123)
means = 30 * randn(rng, 20)
initθ = [0.0, 10.0, 3.0]
nsolves = 10
power = 5.0

struct FDFunc <: Function
    f
end
(f::FDFunc)(x) = f.f(x)
function ChainRulesCore.rrule(f::FDFunc, x::AbstractVector)
    ∇ = ForwardDiff.gradient(f, x)
    return f(x), Δ -> (NoTangent(), Δ * ∇)
end

function run_exp_ipopt(means, initθ, rng; power = 3.0, shift = 1.0, nsamples = 1000, max_iter = 100, nsolves = length(means))
    mixture = mixturef(means)
    logπ = x -> logpdf(mixture, x[1])
    plot(logπ, minimum(means) - 5, maximum(means) + 5)
    variational_objective = AdvancedVI.ELBO()

    solutions = []
    function deflation(x, y)
        if length(solutions) > 0
            d = zero(eltype(x))
            for sol in solutions
                d += (1/kldivergence(getq(x), getq(sol)))^power + shift
            end
            @info "d = $d, y = $y"
            return d - y
        else
            return return -one(eltype(x))
        end
    end

    for _ in 1:nsolves
        θ = copy(initθ)
        function f(θ_)
            return -variational_objective(
                rng, ADVI(nsamples, 1000000), getq(θ_[1:end-1]), logπ,
                nsamples
            )
        end
        g(θ_) = deflation(θ_[1:end-1], exp(θ_[end]))

        oalg = IpoptAlg()
        options = IpoptOptions(; max_iter)
        model = Model(FDFunc(f))
        addvar!(model, fill(-Inf, 3), fill(Inf, 3))
        add_ineq_constraint!(model, FDFunc(g))
        r = Nonconvex.optimize(model, oalg, initθ; options)
        push!(solutions, r.minimizer[1:end-1])
    end

    # Optimal distributions
    dists = getq.(solutions)

    # PDF plots
    plt = plot(x -> exp(logπ(x)), minimum(means) - 5, maximum(means) + 5, legend = true)
    map(dists) do dist
        plot!(plt, x -> exp(logpdf(dist, [x])), minimum(means) - 5, maximum(means) + 5)
        display(plt)
    end

    # Minimization objective values
    vals = map(dists) do dist
        -variational_objective(rng, ADVI(nsamples, 1000000), dist, logπ, 1000)
    end
    return vals
end

rng = StableRNG(123)
Random.seed!(rng, 123)
means = 20 * randn(rng, 20)
initθ = [0.0, 5.0, 6.0]
nsolves = 10
power = 3.0
_nsamples = 1000
max_iter = 50

vals = run_exp_ipopt(means, initθ, rng; nsamples = _nsamples, nsolves, max_iter)


text:
# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

@with_kw mutable struct ObjectiveThresholdParameters{T <: AbstractFloat}
    reference::T = 0.0
    τ::T = 1e-6
end

"""
    ObjectiveThreshold

Functor object for using an objective threshold as termination criterion in a quasigradient algorithm. Create by supplying a [`AtObjectiveThreshold`](@ref) object through `terminate` to `QuasiGradient.Optimizer` or by setting the [`Termination`](@ref) attribute.

...
# Parameters
- `reference::AbstractFloat = 0.0`: Reference objective value
- `τ::AbstractFloat = 1e-6`: Relative tolerance
...
"""
struct ObjectiveThreshold{T <: AbstractFloat} <: AbstractTerminationCriterion
    parameters::ObjectiveThresholdParameters{T}
end

function Progress(termination::ObjectiveThreshold, str::AbstractString)
    @unpack τ = termination.parameters
    return ProgressThresh(τ, 0.0, str)
end

function progress_value(termination::ObjectiveThreshold, k::Integer, f::AbstractFloat, ∇f_norm::AbstractFloat)
    @unpack reference = termination.parameters
    return abs(f - reference) / abs(reference + 1e-10)
end

function terminate(termination::ObjectiveThreshold, k::Integer, f::Float64, x::AbstractVector, ∇f::AbstractVector)
    @unpack reference, τ = termination.parameters
    return abs(f - reference) / abs(reference + 1e-10) <= τ
end

# API
# ------------------------------------------------------------
"""
    AtObjectiveThreshold

Factory object for [`ObjectiveThreshold`](@ref). Pass to `terminate` in `Quasigradient.Optimizer` or set the [`Termination`](@ref) attribute. See ?ObjectiveThreshold for parameter descriptions.

"""
struct AtObjectiveThreshold <: AbstractTermination
    parameters::ObjectiveThresholdParameters{Float64}
end
AtObjectiveThreshold(reference::AbstractFloat, τ::AbstractFloat) = AtObjectiveThreshold(ObjectiveThresholdParameters{Float64}(; reference = Float64(reference), τ = Float64(τ)))
AtObjectiveThreshold(; kw...) = AtObjectiveThreshold(ObjectiveThresholdParameters{Float64}(; kw...))

function (criteria::AtObjectiveThreshold)(::Type{T}) where T <: AbstractFloat
    return ObjectiveThreshold(ObjectiveThresholdParameters{T}(; type2dict(criteria.parameters)...))
end


text:
abstract type BoundaryCondition end

###
struct Dirichlet <: BoundaryCondition end

const dir = Dirichlet()

###
struct Neumann <: BoundaryCondition end

const neu = Neumann()

###
struct Mixed{T} <: BoundaryCondition
    info::T
end


text:
using Documenter, ProxSDP

makedocs(
    modules = [ProxSDP],
    doctest  = false,
    clean    = true,
    format   = Documenter.HTML(),
    assets = ["assets/favicon.ico"],
    sitename = " ",
    authors = "Mario Souto, Joaquim D. Garcia and contributors.",
    pages = [
        "Home" => "index.md",
        "manual.md"
    ]
)

deploydocs(
    repo = "github.com/mariohsouto/ProxSDP.jl.git",
)


text:
## UTF-32 in the native byte order, i.e. plain old character arrays ##

immutable UTF32String <: DirectIndexString
    data::Array{Char,1} # includes 32-bit NULL termination after string chars

    function UTF32String(a::Array{Char,1})
        if length(a) < 1 || a[end] != 0
            throw(ArgumentError("UTF32String data must be NULL-terminated"))
        end
        new(a)
    end
end

next(s::UTF32String, i::Int) = (s.data[i], i+1)
endof(s::UTF32String) = length(s.data) - 1
length(s::UTF32String) = length(s.data) - 1

function utf32(c::Integer...)
    a = Array(Char, length(c) + 1)
    for i = 1:length(c)
        a[i] = c[i]
    end
    a[end] = 0
    UTF32String(a)
end

utf32(x) = convert(UTF32String, x)
convert(::Type{UTF32String}, s::UTF32String) = s

function convert(::Type{UTF32String}, s::String)
    a = Array(Char, length(s) + 1)
    i = 0
    for c in s
        a[i += 1] = c
    end
    a[end] = 0 # NULL terminate
    UTF32String(a)
end

function convert(::Type{UTF32String}, data::AbstractVector{Char})
    len = length(data)
    d = Array(Char, len + 1)
    d[end] = 0 # NULL terminate
    UTF32String(copy!(d,1, data,1, len))
end

convert{T<:Union(Int32,Uint32)}(::Type{UTF32String}, data::AbstractVector{T}) =
    convert(UTF32String, reinterpret(Char, data))

convert{T<:String}(::Type{T}, v::AbstractVector{Char}) = convert(T, utf32(v))

# specialize for performance reasons:
function convert{T<:ByteString}(::Type{T}, data::AbstractVector{Char})
    s = IOBuffer(Array(Uint8,length(data)), true, true)
    truncate(s,0)
    for x in data
        print(s, x)
    end
    convert(T, takebuf_string(s))
end

convert(::Type{Array{Char,1}}, s::UTF32String) = s.data
convert(::Type{Array{Char}}, s::UTF32String) = s.data

reverse(s::UTF32String) = UTF32String(reverse!(copy(s.data), 1, length(s)))

sizeof(s::UTF32String) = sizeof(s.data) - sizeof(Char)
convert{T<:Union(Int32,Uint32,Char)}(::Type{Ptr{T}}, s::UTF32String) =
    convert(Ptr{T}, pointer(s))

function convert(T::Type{UTF32String}, bytes::AbstractArray{Uint8})
    isempty(bytes) && return UTF32String(Char[0])
    length(bytes) & 3 != 0 && throw(ArgumentError("need multiple of 4 bytes"))
    data = reinterpret(Char, bytes)
    # check for byte-order mark (BOM):
    if data[1] == 0x0000feff        # native byte order
        d = Array(Char, length(data))
        copy!(d,1, data,2, length(data)-1)
    elseif data[1] == 0xfffe0000    # byte-swapped
        d = Array(Char, length(data))
        for i = 2:length(data)
            d[i-1] = bswap(data[i])
        end
    else
        d = Array(Char, length(data) + 1)
        copy!(d,1, data,1, length(data)) # assume native byte order
    end
    d[end] = 0 # NULL terminate
    UTF32String(d)
end

utf32(p::Ptr{Char}, len::Integer) = utf32(pointer_to_array(p, len))
utf32(p::Union(Ptr{Uint32}, Ptr{Int32}), len::Integer) = utf32(convert(Ptr{Char}, p), len)
function utf32(p::Union(Ptr{Char}, Ptr{Uint32}, Ptr{Int32}))
    len = 0
    while unsafe_load(p, len+1) != 0; len += 1; end
    utf32(p, len)
end

function map(f::Function, s::UTF32String)
    d = s.data
    out = similar(d)
    out[end] = 0

    for i = 1:(length(d)-1)
        c2 = f(d[i])
        if !isa(c2,Char)
            error("map(f,s::String) requires f to return Char; try map(f,collect(s)) or a comprehension instead")
        end
        out[i] = (c2::Char)
    end
    UTF32String(out)
end

max_stars_repo_name: Aerlinger/julia

max_issues_repo_name: Aerlinger/julia

max_forks_repo_name: Aerlinger/julia


text:
"""
    AbstractPowerRepresentation

An abstract representation type of points and tangent vectors on a power manifold.
"""
abstract type AbstractPowerRepresentation end

"""
    NestedPowerRepresentation

Representation of points and tangent vectors on a power manifold using arrays
of size equal to `TSize` of a [`PowerManifold`](@ref).
Each element of such array stores a single point or tangent vector.

For modifying operations, each element of the outer array is modified in-place, differently
than in [`NestedReplacingPowerRepresentation`](@ref).
"""
struct NestedPowerRepresentation <: AbstractPowerRepresentation end

"""
    NestedReplacingPowerRepresentation

Representation of points and tangent vectors on a power manifold using arrays
of size equal to `TSize` of a [`PowerManifold`](@ref).
Each element of such array stores a single point or tangent vector.

For modifying operations, each element of the outer array is replaced using non-modifying
operations, differently than for [`NestedReplacingPowerRepresentation`](@ref).
"""
struct NestedReplacingPowerRepresentation <: AbstractPowerRepresentation end

@doc raw"""
    AbstractPowerManifold{𝔽,M,TPR} <: AbstractManifold{𝔽}

An abstract [`AbstractManifold`](@ref) to represent manifolds that are build as powers
of another [`AbstractManifold`](@ref) `M` with representation type `TPR`, a subtype of
[`AbstractPowerRepresentation`](@ref).
"""
abstract type AbstractPowerManifold{
    𝔽,
    M<:AbstractManifold{𝔽},
    TPR<:AbstractPowerRepresentation,
} <: AbstractManifold{𝔽} end

@doc raw"""
    PowerManifold{𝔽,TM<:AbstractManifold,TSize<:Tuple,TPR<:AbstractPowerRepresentation} <: AbstractPowerManifold{𝔽,TM}

The power manifold $\mathcal M^{n_1× n_2 × … × n_d}$ with power geometry
 `TSize` statically defines the number of elements along each axis.

For example, a manifold-valued time series would be represented by a power manifold with
$d$ equal to 1 and $n_1$ equal to the number of samples. A manifold-valued image
(for example in diffusion tensor imaging) would be represented by a two-axis power
manifold ($d=2$) with $n_1$ and $n_2$ equal to width and height of the image.

While the size of the manifold is static, points on the power manifold
would not be represented by statically-sized arrays.

# Constructor

    PowerManifold(M::PowerManifold, N_1, N_2, ..., N_d)
    PowerManifold(M::AbstractManifold, NestedPowerRepresentation(), N_1, N_2, ..., N_d)
    M^(N_1, N_2, ..., N_d)

Generate the power manifold $M^{N_1 × N_2 × … × N_d}$.
By default, a [`PowerManifold`](@ref} is expanded further, i.e. for `M=PowerManifold(N,3)`
`PowerManifold(M,2)` is equivalend to `PowerManifold(N,3,2)`. Points are then 3×2 matrices
of points on `N`.
Providing a [`NestedPowerRepresentation`](@ref) as the second argument to the constructor
can be used to nest manifold, i.e. `PowerManifold(M,NestedPowerRepresentation(),2)`
represents vectors of length 2 whose elements are vectors of length 3 of points on N
in a nested array representation.

Since there is no default [`AbstractPowerRepresentation`](@ref) within this interface, the
`^` operator is only available for `PowerManifold`s and concatenates dimensions.
"""
struct PowerManifold{𝔽,TM<:AbstractManifold{𝔽},TSize,TPR<:AbstractPowerRepresentation} <:
       AbstractPowerManifold{𝔽,TM,TPR}
    manifold::TM
end

function PowerManifold(
    M::AbstractManifold{𝔽},
    ::TPR,
    size::Integer...,
) where {𝔽,TPR<:AbstractPowerRepresentation}
    return PowerManifold{𝔽,typeof(M),Tuple{size...},TPR}(M)
end
function PowerManifold(
    M::PowerManifold{𝔽,TM,TSize,TPR},
    size::Integer...,
) where {𝔽,TM<:AbstractManifold{𝔽},TSize,TPR<:AbstractPowerRepresentation}
    return PowerManifold{𝔽,TM,Tuple{TSize.parameters...,size...},TPR}(M.manifold)
end
function PowerManifold(
    M::PowerManifold{𝔽,TM,TSize},
    ::TPR,
    size::Integer...,
) where {𝔽,TM<:AbstractManifold{𝔽},TSize,TPR<:AbstractPowerRepresentation}
    return PowerManifold{𝔽,TM,Tuple{TSize.parameters...,size...},TPR}(M.manifold)
end
function PowerManifold(
    M::PowerManifold{𝔽,TM,TSize},
    ::TPR,
    size::Integer...,
) where {
    𝔽,
    TM<:AbstractManifold{𝔽},
    TSize,
    TPR<:Union{NestedPowerRepresentation,NestedReplacingPowerRepresentation},
}
    return PowerManifold{𝔽,PowerManifold{𝔽,TM,TSize},Tuple{size...},TPR}(M)
end

"""
    PowerBasisData{TB<:AbstractArray}

Data storage for an array of basis data.
"""
struct PowerBasisData{TB<:AbstractArray}
    bases::TB
end

const PowerManifoldNested =
    AbstractPowerManifold{𝔽,<:AbstractManifold{𝔽},NestedPowerRepresentation} where {𝔽}

const PowerManifoldNestedReplacing = AbstractPowerManifold{
    𝔽,
    <:AbstractManifold{𝔽},
    NestedReplacingPowerRepresentation,
} where {𝔽}

_access_nested(x, i::Int) = x[i]
_access_nested(x, i::Tuple) = x[i...]

function Base.:^(
    M::PowerManifold{
        𝔽,
        TM,
        TSize,
        <:Union{NestedPowerRepresentation,NestedReplacingPowerRepresentation},
    },
    size::Integer...,
) where {𝔽,TM<:AbstractManifold{𝔽},TSize}
    return PowerManifold(M, size...)
end

function allocate_result(M::PowerManifoldNested, f, x...)
    if representation_size(M.manifold) === () && length(x) > 0
        return allocate(x[1])
    else
        return [
            allocate_result(M.manifold, f, map(y -> _access_nested(y, i), x)...) for
            i in get_iterator(M)
        ]
    end
end
# avoid ambituities - though usually not used
function allocate_result(
    M::PowerManifoldNested,
    f::typeof(get_coordinates),
    p,
    X,
    B::AbstractBasis,
)
    return invoke(
        allocate_result,
        Tuple{AbstractManifold,typeof(get_coordinates),Any,Any,AbstractBasis},
        M,
        f,
        p,
        X,
        B,
    )
end
function allocate_result(M::PowerManifoldNestedReplacing, f, x...)
    if length(x) == 0
        return [allocate_result(M.manifold, f) for _ in get_iterator(M)]
    else
        return copy(x[1])
    end
end
# the following is not used but necessary to avoid ambiguities
function allocate_result(
    M::PowerManifoldNestedReplacing,
    f::typeof(get_coordinates),
    p,
    X,
    B::AbstractBasis,
)
    return invoke(
        allocate_result,
        Tuple{AbstractManifold,typeof(get_coordinates),Any,Any,AbstractBasis},
        M,
        f,
        p,
        X,
        B,
    )
end
function allocate_result(M::PowerManifoldNested, f::typeof(get_vector), p, X)
    return [allocate_result(M.manifold, f, _access_nested(p, i)) for i in get_iterator(M)]
end
function allocate_result(::PowerManifoldNestedReplacing, ::typeof(get_vector), p, X)
    return copy(p)
end
function allocation_promotion_function(M::AbstractPowerManifold, f, args::Tuple)
    return allocation_promotion_function(M.manifold, f, args)
end

"""
    check_point(M::AbstractPowerManifold, p; kwargs...)

Check whether `p` is a valid point on an [`AbstractPowerManifold`](@ref) `M`,
i.e. each element of `p` has to be a valid point on the base manifold.
If `p` is not a point on `M` a [`CompositeManifoldError`](@ref) consisting of all error messages of the
components, for which the tests fail is returned.

The tolerance for the last test can be set using the `kwargs...`.
"""
function check_point(M::AbstractPowerManifold, p; kwargs...)
    rep_size = representation_size(M.manifold)
    e = [
        (i, check_point(M.manifold, _read(M, rep_size, p, i); kwargs...)) for
        i in get_iterator(M)
    ]
    errors = filter((x) -> !(x[2] === nothing), e)
    cerr = [ComponentManifoldError(er...) for er in errors]
    (length(errors) > 1) && return CompositeManifoldError(cerr)
    (length(errors) == 1) && return cerr[1]
    return nothing
end

"""
    check_power_size(M, p)
    check_power_size(M, p, X)

Check whether p hase the right size to represent points on M generically, i.e. just cheking the overall sizes, not the individual ones per manifold
"""
function check_power_size(M::AbstractPowerManifold, p)
    d = prod(representation_size(M.manifold)) * prod(power_dimensions(M))
    (d != length(p)) && return DomainError(
        length(p),
        "The point $p can not be a point on $M, since its number of elements does not match the required overall representation size ($d)",
    )
    return nothing
end
function check_power_size(M::Union{PowerManifoldNested,PowerManifoldNestedReplacing}, p)
    d = prod(power_dimensions(M))
    (d != length(p)) && return DomainError(
        length(p),
        "The point $p can not be a point on $M, since its number of elements does not match the power dimensions ($d)",
    )
    return nothing
end

function check_power_size(M::AbstractPowerManifold, p, X)
    d = prod(representation_size(M.manifold)) * prod(power_dimensions(M))
    (d != length(X)) && return DomainError(
        length(X),
        "The tangent vector $X can not belong to a trangent space at on $M, since its number of elements does not match the required overall representation size ($d)",
    )
    return nothing
end
function check_power_size(M::Union{PowerManifoldNested,PowerManifoldNestedReplacing}, p, X)
    d = prod(power_dimensions(M))
    (d != length(X)) && return DomainError(
        length(X),
        "The point $p can not be a point on $M, since its number of elements does not match the power dimensions ($d)",
    )
    return nothing
end

function check_size(M::AbstractPowerManifold, p)
    cps = check_power_size(M, p)
    (cps === nothing) || return cps
    rep_size = representation_size(M.manifold)
    e = [(i, check_size(M.manifold, _read(M, rep_size, p, i))) for i in get_iterator(M)]
    errors = filter((x) -> !(x[2] === nothing), e)
    cerr = [ComponentManifoldError(er...) for er in errors]
    (length(errors) > 1) && return CompositeManifoldError(cerr)
    (length(errors) == 1) && return cerr[1]
    return nothing
end

function check_size(M::AbstractPowerManifold, p, X)
    cps = check_power_size(M, p, X)
    (cps === nothing) || return cps
    rep_size = representation_size(M.manifold)
    e = [
        (i, check_size(M.manifold, _read(M, rep_size, p, i), _read(M, rep_size, X, i);)) for i in get_iterator(M)
    ]
    errors = filter((x) -> !(x[2] === nothing), e)
    cerr = [ComponentManifoldError(er...) for er in errors]
    (length(errors) > 1) && return CompositeManifoldError(cerr)
    (length(errors) == 1) && return cerr[1]
    return nothing
end

"""
    check_vector(M::AbstractPowerManifold, p, X; kwargs... )

Check whether `X` is a tangent vector to `p` an the [`AbstractPowerManifold`](@ref)
`M`, i.e. atfer [`check_point`](@ref)`(M, p)`, and all projections to
base manifolds must be respective tangent vectors.
If `X` is not a tangent vector to `p` on `M` a [`CompositeManifoldError`](@ref) consisting of all error
messages of the components, for which the tests fail is returned.

The tolerance for the last test can be set using the `kwargs...`.
"""
function check_vector(M::AbstractPowerManifold, p, X; kwargs...)
    rep_size = representation_size(M.manifold)
    e = [
        (
            i,
            check_vector(
                M.manifold,
                _read(M, rep_size, p, i),
                _read(M, rep_size, X, i);
                kwargs...,
            ),
        ) for i in get_iterator(M)
    ]
    errors = filter((x) -> !(x[2] === nothing), e)
    cerr = [ComponentManifoldError(er...) for er in errors]
    (length(errors) > 1) && return CompositeManifoldError(cerr)
    (length(errors) == 1) && return cerr[1]
    return nothing
end

@doc raw"""
    copyto!(M::PowerManifoldNested, q, p)

Copy the values elementwise, i.e. call `copyto!(M.manifold, b, a)` for all elements `a` and
`b` of `p` and `q`, respectively.
"""
function copyto!(M::PowerManifoldNested, q, p)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        copyto!(M.manifold, _write(M, rep_size, q, i), _read(M, rep_size, p, i))
    end
    return q
end

@doc raw"""
    copyto!(M::PowerManifoldNested, Y, p, X)

Copy the values elementwise, i.e. call `copyto!(M.manifold, B, a, A)` for all elements
`A`, `a` and `B` of `X`, `p`, and `Y`, respectively.
"""
function copyto!(M::PowerManifoldNested, Y, p, X)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        copyto!(
            M.manifold,
            _write(M, rep_size, Y, i),
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
        )
    end
    return Y
end

@doc raw"""
    distance(M::AbstractPowerManifold, p, q)

Compute the distance between `q` and `p` on an [`AbstractPowerManifold`](@ref),
i.e. from the element wise distances the Forbenius norm is computed.
"""
function distance(M::AbstractPowerManifold, p, q)
    sum_squares = zero(number_eltype(p))
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        sum_squares +=
            distance(M.manifold, _read(M, rep_size, p, i), _read(M, rep_size, q, i))^2
    end
    return sqrt(sum_squares)
end

@doc raw"""
    exp(M::AbstractPowerManifold, p, X)

Compute the exponential map from `p` in direction `X` on the [`AbstractPowerManifold`](@ref) `M`,
which can be computed using the base manifolds exponential map elementwise.
"""
exp(::AbstractPowerManifold, ::Any...)

function exp!(M::AbstractPowerManifold, q, p, X)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        exp!(
            M.manifold,
            _write(M, rep_size, q, i),
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
        )
    end
    return q
end
function exp!(M::PowerManifoldNestedReplacing, q, p, X)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        q[i...] = exp(M.manifold, _read(M, rep_size, p, i), _read(M, rep_size, X, i))
    end
    return q
end

function get_basis(M::AbstractPowerManifold, p, B::AbstractBasis)
    rep_size = representation_size(M.manifold)
    vs = [get_basis(M.manifold, _read(M, rep_size, p, i), B) for i in get_iterator(M)]
    return CachedBasis(B, PowerBasisData(vs))
end
function get_basis(M::AbstractPowerManifold, p, B::DiagonalizingOrthonormalBasis)
    rep_size = representation_size(M.manifold)
    vs = [
        get_basis(
            M.manifold,
            _read(M, rep_size, p, i),
            DiagonalizingOrthonormalBasis(_read(M, rep_size, B.frame_direction, i)),
        ) for i in get_iterator(M)
    ]
    return CachedBasis(B, PowerBasisData(vs))
end

"""
    get_component(M::AbstractPowerManifold, p, idx...)

Get the component of a point `p` on an [`AbstractPowerManifold`](@ref) `M` at index `idx`.
"""
function get_component(M::AbstractPowerManifold, p, idx...)
    rep_size = representation_size(M.manifold)
    return _read(M, rep_size, p, idx)
end

function get_coordinates(M::AbstractPowerManifold, p, X, B::AbstractBasis)
    rep_size = representation_size(M.manifold)
    vs = [
        get_coordinates(M.manifold, _read(M, rep_size, p, i), _read(M, rep_size, X, i), B) for i in get_iterator(M)
    ]
    return reduce(vcat, reshape(vs, length(vs)))
end
function get_coordinates(
    M::AbstractPowerManifold,
    p,
    X,
    B::CachedBasis{𝔽,<:AbstractBasis,<:PowerBasisData},
) where {𝔽}
    rep_size = representation_size(M.manifold)
    vs = [
        get_coordinates(
            M.manifold,
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            _access_nested(B.data.bases, i),
        ) for i in get_iterator(M)
    ]
    return reduce(vcat, reshape(vs, length(vs)))
end

function get_coordinates!(M::AbstractPowerManifold, c, p, X, B::AbstractBasis)
    rep_size = representation_size(M.manifold)
    dim = manifold_dimension(M.manifold)
    v_iter = 1
    for i in get_iterator(M)
        # TODO: this view is really suboptimal when `dim` can be statically determined
        get_coordinates!(
            M.manifold,
            view(c, v_iter:(v_iter + dim - 1)),
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            B,
        )
        v_iter += dim
    end
    return c
end
function get_coordinates!(
    M::AbstractPowerManifold,
    c,
    p,
    X,
    B::CachedBasis{𝔽,<:AbstractBasis,<:PowerBasisData},
) where {𝔽}
    rep_size = representation_size(M.manifold)
    dim = manifold_dimension(M.manifold)
    v_iter = 1
    for i in get_iterator(M)
        # TODO: this view is really suboptimal when `dim` can be statically determined
        get_coordinates!(
            M.manifold,
            view(c, v_iter:(v_iter + dim - 1)),
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            _access_nested(B.data.bases, i),
        )
        v_iter += dim
    end
    return c
end

get_iterator(::PowerManifold{𝔽,<:AbstractManifold{𝔽},Tuple{N}}) where {𝔽,N} = Base.OneTo(N)
@generated function get_iterator(
    ::PowerManifold{𝔽,<:AbstractManifold{𝔽},SizeTuple},
) where {𝔽,SizeTuple}
    size_tuple = size_to_tuple(SizeTuple)
    return Base.product(map(Base.OneTo, size_tuple)...)
end

function get_vector(
    M::AbstractPowerManifold,
    p,
    c,
    B::CachedBasis{𝔽,<:AbstractBasis{𝔽},<:PowerBasisData},
) where {𝔽}
    Y = allocate_result(M, get_vector, p, c)
    return get_vector!(M, Y, p, c, B)
end
function get_vector!(
    M::AbstractPowerManifold,
    Y,
    p,
    c,
    B::CachedBasis{𝔽,<:AbstractBasis{𝔽},<:PowerBasisData},
) where {𝔽}
    dim = manifold_dimension(M.manifold)
    rep_size = representation_size(M.manifold)
    v_iter = 1
    for i in get_iterator(M)
        get_vector!(
            M.manifold,
            _write(M, rep_size, Y, i),
            _read(M, rep_size, p, i),
            c[v_iter:(v_iter + dim - 1)],
            _access_nested(B.data.bases, i),
        )
        v_iter += dim
    end
    return Y
end
function get_vector!(
    M::PowerManifoldNestedReplacing,
    Y,
    p,
    c,
    B::CachedBasis{𝔽,<:AbstractBasis{𝔽},<:PowerBasisData},
) where {𝔽}
    dim = manifold_dimension(M.manifold)
    rep_size = representation_size(M.manifold)
    v_iter = 1
    for i in get_iterator(M)
        Y[i...] = get_vector(
            M.manifold,
            _read(M, rep_size, p, i),
            c[v_iter:(v_iter + dim - 1)],
            _access_nested(B.data.bases, i),
        )
        v_iter += dim
    end
    return Y
end
function get_vector(M::AbstractPowerManifold, p, c, B::AbstractBasis)
    Y = allocate_result(M, get_vector, p, c)
    return get_vector!(M, Y, p, c, B)
end
function get_vector!(M::AbstractPowerManifold, Y, p, c, B::AbstractBasis)
    dim = manifold_dimension(M.manifold)
    rep_size = representation_size(M.manifold)
    v_iter = 1
    for i in get_iterator(M)
        get_vector!(
            M.manifold,
            _write(M, rep_size, Y, i),
            _read(M, rep_size, p, i),
            c[v_iter:(v_iter + dim - 1)],
            B,
        )
        v_iter += dim
    end
    return Y
end
function get_vector!(M::PowerManifoldNestedReplacing, Y, p, c, B::AbstractBasis)
    dim = manifold_dimension(M.manifold)
    rep_size = representation_size(M.manifold)
    v_iter = 1
    for i in get_iterator(M)
        Y[i...] = get_vector(
            M.manifold,
            _read(M, rep_size, p, i),
            c[v_iter:(v_iter + dim - 1)],
            B,
        )
        v_iter += dim
    end
    return Y
end

function _get_vectors(
    M::PowerManifoldNested,
    p,
    B::CachedBasis{𝔽,<:AbstractBasis{𝔽},<:PowerBasisData},
) where {𝔽}
    zero_tv = zero_vector(M, p)
    rep_size = representation_size(M.manifold)
    vs = typeof(zero_tv)[]
    for i in get_iterator(M)
        b_i = _access_nested(B.data.bases, i)
        p_i = _read(M, rep_size, p, i)
        for v in b_i.data
            new_v = copy(M, p, zero_tv)
            copyto!(M.manifold, _write(M, rep_size, new_v, i), p_i, v)
            push!(vs, new_v)
        end
    end
    return vs
end
function _get_vectors(
    M::PowerManifoldNestedReplacing,
    p,
    B::CachedBasis{𝔽,<:AbstractBasis{𝔽},<:PowerBasisData},
) where {𝔽}
    zero_tv = zero_vector(M, p)
    vs = typeof(zero_tv)[]
    for i in get_iterator(M)
        b_i = _access_nested(B.data.bases, i)
        for v in b_i.data
            new_v = copy(M, p, zero_tv)
            new_v[i...] = v
            push!(vs, new_v)
        end
    end
    return vs
end

"""
    getindex(p, M::AbstractPowerManifold, i::Union{Integer,Colon,AbstractVector}...)
    p[M::AbstractPowerManifold, i...]

Access the element(s) at index `[i...]` of a point `p` on an [`AbstractPowerManifold`](@ref)
`M` by linear or multidimensional indexing.
See also [Array Indexing](https://docs.julialang.org/en/v1/manual/arrays/#man-array-indexing-1) in Julia.
"""
Base.@propagate_inbounds function Base.getindex(
    p::AbstractArray,
    M::AbstractPowerManifold,
    I::Union{Integer,Colon,AbstractVector}...,
)
    return get_component(M, p, I...)
end

@doc raw"""
    injectivity_radius(M::AbstractPowerManifold[, p])

the injectivity radius on an [`AbstractPowerManifold`](@ref) is for the global case
equal to the one of its base manifold. For a given point `p` it's equal to the
minimum of all radii in the array entries.
"""
function injectivity_radius(M::AbstractPowerManifold, p)
    radius = 0.0
    initialized = false
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        cur_rad = injectivity_radius(M.manifold, _read(M, rep_size, p, i))
        if initialized
            radius = min(cur_rad, radius)
        else
            radius = cur_rad
            initialized = true
        end
    end
    return radius
end
injectivity_radius(M::AbstractPowerManifold) = injectivity_radius(M.manifold)
function injectivity_radius(M::AbstractPowerManifold, ::AbstractRetractionMethod)
    return injectivity_radius(M)
end

@doc raw"""
    inner(M::AbstractPowerManifold, p, X, Y)

Compute the inner product of `X` and `Y` from the tangent space at `p` on an
[`AbstractPowerManifold`](@ref) `M`, i.e. for each arrays entry the tangent
vector entries from `X` and `Y` are in the tangent space of the corresponding
element from `p`.
The inner product is then the sum of the elementwise inner products.
"""
function inner(M::AbstractPowerManifold, p, X, Y)
    result = zero(number_eltype(X))
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        result += inner(
            M.manifold,
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            _read(M, rep_size, Y, i),
        )
    end
    return result
end

function Base.isapprox(M::AbstractPowerManifold, p, q; kwargs...)
    result = true
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        result &= isapprox(
            M.manifold,
            _read(M, rep_size, p, i),
            _read(M, rep_size, q, i);
            kwargs...,
        )
    end
    return result
end
function Base.isapprox(M::AbstractPowerManifold, p, X, Y; kwargs...)
    result = true
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        result &= isapprox(
            M.manifold,
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            _read(M, rep_size, Y, i);
            kwargs...,
        )
    end
    return result
end

@doc raw"""
    inverse_retract(M::AbstractPowerManifold, p, q, m::AbstractInverseRetractionMethod)

Compute the inverse retraction from `p` with respect to `q` on an [`AbstractPowerManifold`](@ref) `M`
using an [`AbstractInverseRetractionMethod`](@ref).
Then this method is performed elementwise, so the inverse
retraction method has to be one that is available on the base [`AbstractManifold`](@ref).
"""
inverse_retract(::AbstractPowerManifold, ::Any...)

function inverse_retract(
    M::AbstractPowerManifold,
    p,
    q,
    m::AbstractInverseRetractionMethod = default_inverse_retraction_method(M),
)
    X = allocate_result(M, inverse_retract, p, q)
    return inverse_retract!(M, X, p, q, m)
end

function inverse_retract!(
    M::AbstractPowerManifold,
    X,
    p,
    q,
    m::AbstractInverseRetractionMethod = default_inverse_retraction_method(M),
)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        inverse_retract!(
            M.manifold,
            _write(M, rep_size, X, i),
            _read(M, rep_size, p, i),
            _read(M, rep_size, q, i),
            m,
        )
    end
    return X
end
function inverse_retract!(
    M::PowerManifoldNestedReplacing,
    X,
    p,
    q,
    m::AbstractInverseRetractionMethod = default_inverse_retraction_method(M),
)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        X[i...] = inverse_retract(
            M.manifold,
            _read(M, rep_size, p, i),
            _read(M, rep_size, q, i),
            m,
        )
    end
    return X
end

@doc raw"""
    log(M::AbstractPowerManifold, p, q)

Compute the logarithmic map from `p` to `q` on the [`AbstractPowerManifold`](@ref) `M`,
which can be computed using the base manifolds logarithmic map elementwise.
"""
log(::AbstractPowerManifold, ::Any...)

function log!(M::AbstractPowerManifold, X, p, q)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        log!(
            M.manifold,
            _write(M, rep_size, X, i),
            _read(M, rep_size, p, i),
            _read(M, rep_size, q, i),
        )
    end
    return X
end
function log!(M::PowerManifoldNestedReplacing, X, p, q)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        X[i...] = log(M.manifold, _read(M, rep_size, p, i), _read(M, rep_size, q, i))
    end
    return X
end

@doc raw"""
    manifold_dimension(M::PowerManifold)

Returns the manifold-dimension of an [`PowerManifold`](@ref) `M`
$=\mathcal N = (\mathcal M)^{n_1,…,n_d}$, i.e. with $n=(n_1,…,n_d)$ the array
size of the power manifold and $d_{\mathcal M}$ the dimension of the base manifold
$\mathcal M$, the manifold is of dimension

````math
\dim(\mathcal N) = \dim(\mathcal M)\prod_{i=1}^d n_i = n_1n_2\cdot…\cdot n_d \dim(\mathcal M).
````
"""
function manifold_dimension(M::PowerManifold{𝔽,<:AbstractManifold,TSize}) where {𝔽,TSize}
    return manifold_dimension(M.manifold) * prod(size_to_tuple(TSize))
end

function mid_point!(M::AbstractPowerManifold, q, p1, p2)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        mid_point!(
            M.manifold,
            _write(M, rep_size, q, i),
            _read(M, rep_size, p1, i),
            _read(M, rep_size, p2, i),
        )
    end
    return q
end
function mid_point!(M::PowerManifoldNestedReplacing, q, p1, p2)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        q[i...] =
            mid_point(M.manifold, _read(M, rep_size, p1, i), _read(M, rep_size, p2, i))
    end
    return q
end

@doc raw"""
    norm(M::AbstractPowerManifold, p, X)

Compute the norm of `X` from the tangent space of `p` on an
[`AbstractPowerManifold`](@ref) `M`, i.e. from the element wise norms the
Frobenius norm is computed.
"""
function LinearAlgebra.norm(M::AbstractPowerManifold, p, X)
    sum_squares = zero(number_eltype(X))
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        sum_squares +=
            norm(M.manifold, _read(M, rep_size, p, i), _read(M, rep_size, X, i))^2
    end
    return sqrt(sum_squares)
end

function parallel_transport_direction!(M::AbstractPowerManifold, Y, p, X, d)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        parallel_transport_direction!(
            M.manifold,
            _write(M, rep_size, Y, i),
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            _read(M, rep_size, d, i),
        )
    end
    return Y
end
function parallel_transport_direction(M::AbstractPowerManifold, p, X, d)
    Y = allocate_result(M, vector_transport_direction, p, X, d)
    return parallel_transport_direction!(M, Y, p, X, d)
end

function parallel_transport_direction!(M::PowerManifoldNestedReplacing, Y, p, X, d)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        Y[i...] = parallel_transport_direction(
            M.manifold,
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            _read(M, rep_size, d, i),
        )
    end
    return Y
end
function parallel_transport_direction(M::PowerManifoldNestedReplacing, p, X, d)
    Y = allocate_result(M, parallel_transport_direction, p, X, d)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        Y[i...] = parallel_transport_direction(
            M.manifold,
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            _read(M, rep_size, d, i),
        )
    end
    return Y
end

function parallel_transport_to(M::AbstractPowerManifold, p, X, q)
    Y = allocate_result(M, vector_transport_to, p, X)
    return parallel_transport_to!(M, Y, p, X, q)
end
function parallel_transport_to!(M::AbstractPowerManifold, Y, p, X, q)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        vector_transport_to!(
            M.manifold,
            _write(M, rep_size, Y, i),
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            _read(M, rep_size, q, i),
        )
    end
    return Y
end
function parallel_transport_to!(M::PowerManifoldNestedReplacing, Y, p, X, q)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        Y[i...] = parallel_transport_to(
            M.manifold,
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            _read(M, rep_size, q, i),
        )
    end
    return Y
end

@doc raw"""
    power_dimensions(M::PowerManifold)

return the power of `M`,
"""
function power_dimensions(::PowerManifold{𝔽,<:AbstractManifold,TSize}) where {𝔽,TSize}
    return size_to_tuple(TSize)
end

@doc raw"""
    project(M::AbstractPowerManifold, p)

Project the point `p` from the embedding onto the [`AbstractPowerManifold`](@ref) `M`
by projecting all components.
"""
project(::AbstractPowerManifold, ::Any)

function project!(M::AbstractPowerManifold, q, p)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        project!(M.manifold, _write(M, rep_size, q, i), _read(M, rep_size, p, i))
    end
    return q
end
function project!(M::PowerManifoldNestedReplacing, q, p)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        q[i...] = project(M.manifold, _read(M, rep_size, p, i))
    end
    return q
end

@doc raw"""
    project(M::AbstractPowerManifold, p, X)

Project the point `X` onto the tangent space at `p` on the
[`AbstractPowerManifold`](@ref) `M` by projecting all components.
"""
project(::AbstractPowerManifold, ::Any, ::Any)

function project!(M::AbstractPowerManifold, Z, q, Y)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        project!(
            M.manifold,
            _write(M, rep_size, Z, i),
            _read(M, rep_size, q, i),
            _read(M, rep_size, Y, i),
        )
    end
    return Z
end
function project!(M::PowerManifoldNestedReplacing, Z, q, Y)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        q[i...] = project(M.manifold, _read(M, rep_size, q, i), _read(M, rep_size, Y, i))
    end
    return Z
end

Base.@propagate_inbounds @inline function _read(
    M::AbstractPowerManifold,
    rep_size::Tuple,
    x::AbstractArray,
    i::Int,
)
    return _read(M, rep_size, x, (i,))
end

Base.@propagate_inbounds @inline function _read(
    ::Union{PowerManifoldNested,PowerManifoldNestedReplacing},
    rep_size::Tuple,
    x::AbstractArray,
    i::Tuple,
)
    return x[i...]
end

@generated function rep_size_to_colons(rep_size::Tuple)
    N = length(rep_size.parameters)
    return ntuple(i -> Colon(), N)
end

@doc raw"""
    retract(M::AbstractPowerManifold, p, X, method::AbstractRetractionMethod)

Compute the retraction from `p` with tangent vector `X` on an [`AbstractPowerManifold`](@ref) `M`
using a [`AbstractRetractionMethod`](@ref).
Then this method is performed elementwise, so the retraction
method has to be one that is available on the base [`AbstractManifold`](@ref).
"""
retract(::AbstractPowerManifold, ::Any...)

function retract(
    M::AbstractPowerManifold,
    p,
    X,
    m::AbstractRetractionMethod = default_retraction_method(M),
)
    q = allocate_result(M, retract, p, X)
    return retract!(M, q, p, X, m)
end

function retract!(
    M::AbstractPowerManifold,
    q,
    p,
    X,
    m::AbstractRetractionMethod = default_retraction_method(M),
)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        retract!(
            M.manifold,
            _write(M, rep_size, q, i),
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            m,
        )
    end
    return q
end
function retract!(
    M::PowerManifoldNestedReplacing,
    q,
    p,
    X,
    m::AbstractRetractionMethod = default_retraction_method(M),
)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        q[i...] = retract(M.manifold, _read(M, rep_size, p, i), _read(M, rep_size, X, i), m)
    end
    return q
end

"""
    set_component!(M::AbstractPowerManifold, q, p, idx...)

Set the component of a point `q` on an [`AbstractPowerManifold`](@ref) `M` at index `idx`
to `p`, which itself is a point on the [`AbstractManifold`](@ref) the power manifold is build on.
"""
function set_component!(M::AbstractPowerManifold, q, p, idx...)
    rep_size = representation_size(M.manifold)
    return copyto!(_write(M, rep_size, q, idx), p)
end
function set_component!(::PowerManifoldNestedReplacing, q, p, idx...)
    return q[idx...] = p
end
"""
    setindex!(q, p, M::AbstractPowerManifold, i::Union{Integer,Colon,AbstractVector}...)
    q[M::AbstractPowerManifold, i...] = p

Set the element(s) at index `[i...]` of a point `q` on an [`AbstractPowerManifold`](@ref)
`M` by linear or multidimensional indexing to `q`.
See also [Array Indexing](https://docs.julialang.org/en/v1/manual/arrays/#man-array-indexing-1) in Julia.
"""
Base.@propagate_inbounds function Base.setindex!(
    q::AbstractArray,
    p,
    M::AbstractPowerManifold,
    I::Union{Integer,Colon,AbstractVector}...,
)
    return set_component!(M, q, p, I...)
end

function Base.show(io::IO, M::PowerManifold{𝔽,TM,TSize,TPR}) where {𝔽,TM,TSize,TPR}
    return print(
        io,
        "PowerManifold($(M.manifold), $(TPR()), $(join(TSize.parameters, ", ")))",
    )
end

function Base.show(
    io::IO,
    mime::MIME"text/plain",
    B::CachedBasis{𝔽,T,D},
) where {T<:AbstractBasis,D<:PowerBasisData,𝔽}
    println(io, "$(T()) for a power manifold")
    for i in Base.product(map(Base.OneTo, size(B.data.bases))...)
        println(io, "Basis for component $i:")
        show(io, mime, _access_nested(B.data.bases, i))
        println(io)
    end
    return nothing
end

function vector_transport_direction!(
    M::AbstractPowerManifold,
    Y,
    p,
    X,
    d,
    m::AbstractVectorTransportMethod = default_vector_transport_method(M),
)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        vector_transport_direction!(
            M.manifold,
            _write(M, rep_size, Y, i),
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            _read(M, rep_size, d, i),
            m,
        )
    end
    return Y
end
function vector_transport_direction(
    M::AbstractPowerManifold,
    p,
    X,
    d,
    m::AbstractVectorTransportMethod = default_vector_transport_method(M),
)
    Y = allocate_result(M, vector_transport_direction, p, X, d)
    return vector_transport_direction!(M, Y, p, X, d, m)
end

function vector_transport_direction!(
    M::PowerManifoldNestedReplacing,
    Y,
    p,
    X,
    d,
    m::AbstractVectorTransportMethod = default_vector_transport_method(M),
)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        Y[i...] = vector_transport_direction(
            M.manifold,
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            _read(M, rep_size, d, i),
            m,
        )
    end
    return Y
end
function vector_transport_direction(
    M::PowerManifoldNestedReplacing,
    p,
    X,
    d,
    m::AbstractVectorTransportMethod = default_vector_transport_method(M),
)
    Y = allocate_result(M, vector_transport_direction, p, X, d)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        Y[i...] = vector_transport_direction(
            M.manifold,
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            _read(M, rep_size, d, i),
            m,
        )
    end
    return Y
end

@doc raw"""
    vector_transport_to(M::AbstractPowerManifold, p, X, q, method::AbstractVectorTransportMethod)

Compute the vector transport the tangent vector `X`at `p` to `q` on the
[`PowerManifold`](@ref) `M` using an [`AbstractVectorTransportMethod`](@ref) `m`.
This method is performed elementwise, i.e. the method `m` has to be implemented on the
base manifold.
"""
vector_transport_to(
    ::AbstractPowerManifold,
    ::Any,
    ::Any,
    ::Any,
    ::AbstractVectorTransportMethod,
)
function vector_transport_to(
    M::AbstractPowerManifold,
    p,
    X,
    q,
    m::AbstractVectorTransportMethod = default_vector_transport_method(M),
)
    Y = allocate_result(M, vector_transport_to, p, X)
    return vector_transport_to!(M, Y, p, X, q, m)
end
function vector_transport_to!(
    M::AbstractPowerManifold,
    Y,
    p,
    X,
    q,
    m::AbstractVectorTransportMethod = default_vector_transport_method(M),
)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        vector_transport_to!(
            M.manifold,
            _write(M, rep_size, Y, i),
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            _read(M, rep_size, q, i),
            m,
        )
    end
    return Y
end
function vector_transport_to!(
    M::PowerManifoldNestedReplacing,
    Y,
    p,
    X,
    q,
    m::AbstractVectorTransportMethod = default_vector_transport_method(M),
)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        Y[i...] = vector_transport_to(
            M.manifold,
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            _read(M, rep_size, q, i),
            m,
        )
    end
    return Y
end

"""
    view(p, M::AbstractPowerManifold, i::Union{Integer,Colon,AbstractVector}...)

Get the view of the element(s) at index `[i...]` of a point `p` on an
[`AbstractPowerManifold`](@ref) `M` by linear or multidimensional indexing.
"""
function Base.view(
    p::AbstractArray,
    M::AbstractPowerManifold,
    I::Union{Integer,Colon,AbstractVector}...,
)
    rep_size = representation_size(M.manifold)
    return _write(M, rep_size, p, I...)
end

@inline function _write(M::AbstractPowerManifold, rep_size::Tuple, x::AbstractArray, i::Int)
    return _write(M, rep_size, x, (i,))
end

@inline function _write(
    ::AbstractPowerManifold,
    rep_size::Tuple,
    x::AbstractArray,
    i::Tuple,
)
    return view(x[i...], rep_size_to_colons(rep_size)...)
end
@inline function _write(::PowerManifoldNested, ::Tuple{}, x::AbstractArray, i::Tuple)
    return view(x, i...)
end

function zero_vector!(M::AbstractPowerManifold, X, p)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        zero_vector!(M.manifold, _write(M, rep_size, X, i), _read(M, rep_size, p, i))
    end
    return X
end
function zero_vector!(M::PowerManifoldNestedReplacing, X, p)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        X[i...] = zero_vector(M.manifold, _read(M, rep_size, p, i))
    end
    return X
end


text:
mutable struct Section <: Element
    contents::Vector{Element}
end

mutable struct Heading <: Element
    level::Integer
    keyword::Union{AbstractString, Nothing}
    priority::Union{AbstractString, Nothing}
    title::Vector{Object}
    tags::Vector{AbstractString}
    section::Union{Section, Nothing}
    planning::Union{Planning, Nothing}
    properties::Union{PropertyDrawer, Nothing}
end


text:
"""
Typical geodynamic simulations involve a large number of material parameters that have units that are often inconvenient to be directly used in numerical models
This package has two main features that help with this:
- Create a nondimensionalization object, which can be used to transfer dimensional to non-dimensional parameters (usually better for numerical solvers)
- Create an object in which you can specify material parameters employed in the geodynamic simulations

The material parameter object is designed to be extensible and can be passed on to the solvers, such that new creep laws or features can be readily added.
We also implement some typically used creep law parameters, together with tools to plot them versus and compare our results with those of published papers (to minimize mistakes).
"""
__precompile__()
module GeoParams

using Parameters        # helps setting default parameters in structures
using Unitful           # Units
using BibTeX            # references of creep laws
using Requires          # To only add plotting routines if Plots is loaded

export
        @u_str, uconvert, upreffered, unit, ustrip, NoUnits,  #  Units
        GeoUnit, GeoUnits, GEO_units, SI_units, NO_units, AbstractGeoUnit,
        nondimensionalize, dimensionalize,
        superscript, upreferred, GEO, SI, NONE, isDimensional, Value, NumValue, Unit, UnitValue, isdimensional,
        km, m, cm, mm, μm, Myrs, yr, s, MPa, Pa, kbar, Pas, K, C, g, kg, mol, J, kJ, Watt, μW, Quantity

export AbstractGeoUnit1, GeoUnit1

#
abstract type AbstractMaterialParam end                                    # structure that holds material parmeters (density, elasticity, viscosity)
abstract type AbstractMaterialParamsStruct end                             # will hold all info for a phase
abstract type AbstractPhaseDiagramsStruct <: AbstractMaterialParam end    # will hold all info for phase diagrams
function PerpleX_LaMEM_Diagram end                                         # necessary as we already use this function in Units, but only define it later in PhaseDiagrams
function param_info end
export AbstractMaterialParam, AbstractMaterialParamsStruct, AbstractPhaseDiagramsStruct

# note that this throws a "Method definition warning regarding superscript"; that is expected & safe
#  as we add a nicer way to create output of superscripts. I have been unable to get rid of this warning,
#  as I am indeed redefining a method originally defined in Unitful
include("Units.jl")
using .Units
export @unpack_units, @unpack_val
export compute_units

# Define Material Parameter structure
include("MaterialParameters.jl")
using .MaterialParameters
export MaterialParams, SetMaterialParams, No_MaterialParam, MaterialParamsInfo

# Phase Diagrams
using .MaterialParameters.PhaseDiagrams
export PhaseDiagram_LookupTable, PerpleX_LaMEM_Diagram

# Density
using .MaterialParameters.Density
export compute_density,                                # computational routines
        compute_density!,
        param_info,
        AbstractDensity,
        No_Density,
        ConstantDensity,
        PT_Density,
        Compressible_Density,
        PhaseDiagram_LookupTable,
        Read_LaMEM_Perple_X_Diagram

# Creep laws
using .MaterialParameters.CreepLaw
export computeCreepLaw_EpsII, computeCreepLaw_TauII, CreepLawVariables,
        LinearViscous, PowerlawViscous,
        DislocationCreep, SetDislocationCreep

# Gravitational Acceleration
using .MaterialParameters.GravitationalAcceleration
export compute_gravity,                                # computational routines
        ConstantGravity

# Energy parameters: Heat Capacity, Thermal conductivity, latent heat, radioactive heat
using .MaterialParameters.HeatCapacity
export compute_heatcapacity,
        compute_heatcapacity!,
        ConstantHeatCapacity,
        T_HeatCapacity_Whittacker

using .MaterialParameters.Conductivity
export compute_conductivity,
        compute_conductivity!,
        ConstantConductivity,
        T_Conductivity_Whittacker,
        TP_Conductivity,
        Set_TP_Conductivity

using .MaterialParameters.LatentHeat
export compute_latent_heat,
        ConstantLatentHeat

using .MaterialParameters.RadioactiveHeat
export compute_radioactive_heat,
        ConstantRadioactiveHeat

using .MaterialParameters.Shearheating
export compute_shearheating!, compute_shearheating,
        ConstantShearheating

# Seismic velocities
using .MaterialParameters.SeismicVelocity
export compute_pwave_velocity, compute_swave_velocity,
        compute_pwave_velocity!, compute_swave_velocity!,
        ConstantSeismicVelocity

# Add melting parameterizations
include("./MeltFraction/MeltingParameterization.jl")
using .MeltingParam
export compute_meltfraction, compute_meltfraction!,       # calculation routines
        MeltingParam_Caricchi

# Add plotting routines - only activated if the "Plots.jl" package is loaded
function __init__()
        @require Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" begin
                print("Adding plotting routines of GeoParam")
                @eval include("./Plotting.jl")
        end
end

#Set functions aliases using @use
include("aliases.jl")

end # module


text:
# The julia Coverage count all files in src by default. However, our
# CUDA components are only tested locally because the Travis-CI does
# not have GPU devices. This caused the coverage of Mocha to be very
# low. We add a simple hacking that allow us to ignore the "cuda"
# sub-directory to fix the coverage reporting.
using Coverage

function process_folder(folder="src", ignore=[])
  source_files=Any[]
  filelist = readdir(folder)
  for file in filelist
    if in(file, ignore)
      continue
    end

    fullfile = joinpath(folder,file)
    println(fullfile)
    if isfile(fullfile)
      try
        new_sf = Coveralls.process_file(fullfile)
        push!(source_files, new_sf)
      catch e
        println("Skipped $fullfile")
      end
    else isdir(fullfile)
      append!(source_files, process_folder(fullfile))
    end
  end
  return source_files
end


text:
struct NullInvariants end
Base.getindex(::NullInvariants, i...) = error("Invariants were indexed but the invariants are `nothing`. You likely forgot to pass invariants to the Equation!")
Base.iterate(::NullInvariants) = error("Invariants were indexed but the invariants are `nothing`. You likely forgot to pass invariants to the Equation!")

struct NullParameters end
Base.getindex(::NullParameters, i...) = error("Parameters were indexed but the parameters are `nothing`. You likely forgot to pass parameters to the Equation!")
Base.iterate(::NullParameters) = error("Parameters were indexed but the parameters are `nothing`. You likely forgot to pass parameters to the Equation!")

struct NullPeriodicity end
Base.getindex(::NullPeriodicity, i...) = error("Periodicity was indexed but the periodicity is `nothing`. You likely forgot to pass periodicity to the Equation!")
Base.iterate(::NullPeriodicity) = error("Periodicity was indexed but the periodicity is `nothing`. You likely forgot to pass periodicity to the Equation!")

const OptionalArray{arrayType} = Union{Nothing, arrayType} where {arrayType <: AbstractArray}

const OptionalAbstractArray = Union{Nothing, AbstractArray}
const OptionalCallable      = Union{Nothing, Callable}
const OptionalFunction      = Union{Nothing, Function}
const OptionalNamedTuple    = Union{Nothing, NamedTuple}
const OptionalTuple         = Union{Nothing, Tuple}

const OptionalInvariants    = Union{NamedTuple, NullInvariants}
const OptionalParameters    = Union{NamedTuple, NullParameters}
const OptionalPeriodicity   = Union{AbstractArray, NullPeriodicity}

const AbstractData = Union{Number, AbstractArray{<:Number}}

const State{DT <: Number} = AbstractArray{DT}
const StateVector{DT,VT} = VT where {DT, VT <: AbstractVector{<:State{DT}}}

Base.zero(X::ST) where {DT, VT, ST <: StateVector{DT,VT}} = VT[zero(x) for x in X]

SolutionVector{DT} = Union{Vector{DT}, Vector{TwicePrecision{DT}}}


text:
module ClimateMARGO

if VERSION < v"1.5"
    @error "ClimateMARGO requires Julia v1.5 or newer."
end

include("Models/Models.jl")
include("Utils/Utils.jl")
include("Diagnostics/Diagnostics.jl")
include("Optimization/Optimization.jl")
include("IO/IO.jl")
include("PolicyResponse/PolicyResponse.jl")
include("Plotting/Plotting.jl")

end


text:
import DataStructures
import AVLTrees
using BenchmarkTools

include("../src/soalight.jl")
include("../src/allocator.jl")
include("../src/tree_iter.jl")

function testbase1(n)
    tree = Set{Int}()
    for j in 1:n[]
        push!(tree, j)
        push!(tree, -j)
    end
    @assert !(0 in tree)
    for j in 1:n[]
        @assert j in tree
        @assert -j in tree
    end
    @assert !(n[] + 1 in tree)
    @assert !(-n[] - 1 in tree)
    for j in 1:n[]
        @assert j in tree
        @assert -j in tree
        pop!(tree, j)
        pop!(tree, -j)
        @assert !(j in tree)
        @assert !(-j in tree)
    end
end

function testbase2(n)
    tree = DataStructures.SortedSet{Int}()
    for j in 1:n[]
        push!(tree, j)
        push!(tree, -j)
    end
    @assert !haskey(tree, 0)
    for j in 1:n[]
        @assert haskey(tree, j)
        @assert haskey(tree, -j)
    end
    @assert !haskey(tree, n[] + 1)
    @assert !haskey(tree, -n[] - 1)
    for j in 1:n[]
        @assert haskey(tree, j)
        @assert haskey(tree, -j)
        pop!(tree, j)
        pop!(tree, -j)
        @assert !haskey(tree, j)
        @assert !haskey(tree, -j)
    end
end

function testbase3(n)
    tree = DataStructures.AVLTree{Int}()
    for j in 1:n[]
        insert!(tree, j)
        insert!(tree, -j)
    end
    @assert !haskey(tree, 0)
    for j in 1:n[]
        @assert haskey(tree, j)
        @assert haskey(tree, -j)
    end
    @assert !haskey(tree, n[] + 1)
    @assert !haskey(tree, -n[] - 1)
    for j in 1:n[]
        @assert haskey(tree, j)
        @assert haskey(tree, -j)
        delete!(tree, j)
        delete!(tree, -j)
        @assert !haskey(tree, j)
        @assert !haskey(tree, -j)
    end
end

function testbase4(n)
    tree = DataStructures.RBTree{Int}()
    for j in 1:n[]
        insert!(tree, j)
        insert!(tree, -j)
    end
    @assert !haskey(tree, 0)
    for j in 1:n[]
        @assert haskey(tree, j)
        @assert haskey(tree, -j)
    end
    @assert !haskey(tree, n[] + 1)
    @assert !haskey(tree, -n[] - 1)
    for j in 1:n[]
        @assert haskey(tree, j)
        @assert haskey(tree, -j)
        delete!(tree, j)
        delete!(tree, -j)
        @assert !haskey(tree, j)
        @assert !haskey(tree, -j)
    end
end

function testbase5(n)
    tree = DataStructures.SplayTree{Int}()
    for j in 1:n[]
        push!(tree, j)
        push!(tree, -j)
    end
    @assert !haskey(tree, 0)
    for j in 1:n[]
        @assert haskey(tree, j)
        @assert haskey(tree, -j)
    end
    @assert !haskey(tree, n[] + 1)
    @assert !haskey(tree, -n[] - 1)
    for j in 1:n[]
        @assert haskey(tree, j)
        @assert haskey(tree, -j)
        delete!(tree, j)
        delete!(tree, -j)
        @assert !haskey(tree, j)
        @assert !haskey(tree, -j)
    end
end

function testbase6(n)
    tree = AVLTrees.AVLSet{Int}()
    for j in 1:n[]
        push!(tree, j)
        push!(tree, -j)
    end
    @assert !(0 in tree)
    for j in 1:n[]
        @assert j in tree
        @assert -j in tree
    end
    @assert !(n[] + 1 in tree)
    @assert !(-n[] - 1 in tree)
    for j in 1:n[]
        @assert j in tree
        @assert -j in tree
        delete!(tree, j)
        delete!(tree, -j)
        @assert !(j in tree)
        @assert !(-j in tree)
    end
end

function testbase7(n)
    tree = nil()
    for j in 1:n[]
        tree = insert(tree, j, nothing)
        tree = insert(tree, -j, nothing)
    end
    @assert !haskey(tree, 0)
    for j in 1:n[]
        @assert haskey(tree, j)
        @assert haskey(tree, -j)
    end
    @assert !haskey(tree, n[] + 1)
    @assert !haskey(tree, -n[] - 1)
    for j in 1:n[]
        @assert haskey(tree, j)
        @assert haskey(tree, -j)
        tree = delete(tree, j)
        tree = delete(tree, -j)
        @assert !haskey(tree, j)
        @assert !haskey(tree, -j)
    end
end

function testalloc(n, alloc)
    tree = nil(alloc)
    for j in 1:n[]
        tree = insert(tree, j, nothing, alloc)
        tree = insert(tree, -j, nothing, alloc)
    end
    @assert !haskey(tree, 0, alloc)
    for j in 1:n[]
        @assert haskey(tree, j, alloc)
        @assert haskey(tree, -j, alloc)
    end
    @assert !haskey(tree, n[] + 1, alloc)
    @assert !haskey(tree, -n[] - 1, alloc)
    for j in 1:n[]
        @assert haskey(tree, j, alloc)
        @assert haskey(tree, -j, alloc)
        tree = delete(tree, j, alloc)
        tree = delete(tree, -j, alloc)
        @assert !haskey(tree, j, alloc)
        @assert !haskey(tree, -j, alloc)
    end
    emptyend!(alloc)
end

function testfree(n, alloc)
    tree = nil(alloc)
    for j in 1:n[]
        tree = insert(tree, j, nothing, alloc)
        tree = insert(tree, -j, nothing, alloc)
    end
    @assert !haskey(tree, 0, alloc)
    for j in 1:n[]
        @assert haskey(tree, j, alloc)
        @assert haskey(tree, -j, alloc)
    end
    @assert !haskey(tree, n[] + 1, alloc)
    @assert !haskey(tree, -n[] - 1, alloc)
    for j in 1:n[]
        @assert haskey(tree, j, alloc)
        @assert haskey(tree, -j, alloc)
        tree = delete(tree, j, alloc)
        tree = delete(tree, -j, alloc)
        @assert !haskey(tree, j, alloc)
        @assert !haskey(tree, -j, alloc)
    end
    @assert isempty(alloc)
end

ans = nothing
for n in [100000, 1000, 10]
    println(n, " elements")

    print("  Set                                    ")
    @btime ($ans = testbase1(Ref($n)))

    GC.gc()

    print("  DataStructures.SortedSet               ")
    @btime ($ans = testbase2(Ref($n)))

    GC.gc()

    print("  DataStructures.AVLTree                 ")
    @btime ($ans = testbase3(Ref($n)))

    GC.gc()

    print("  DataStructures.RBTree                  ")
    @btime ($ans = testbase4(Ref($n)))

    GC.gc()

    print("  DataStructures.SplayTree               ")
    @btime ($ans = testbase5(Ref($n)))

    GC.gc()

    print("  AVLTrees.AVLSet                        ")
    @btime ($ans = testbase6(Ref($n)))

    GC.gc()

    print("  Tree without allocator                 ")
    @btime ($ans = testbase7(Ref($n)))

    GC.gc()

    print("  with fixed allocator                   ")
    alloc = Allocator{TreeNode{Int, Int, Nothing}, Int}(2 * n)
    @btime ($ans = testalloc(Ref($n), $alloc))

    GC.gc()

    print("  with resizable allocator               ")
    alloc = Allocator{TreeNode{Int, Int, Nothing}, Int}(nothing)
    @btime ($ans = testalloc(Ref($n), $alloc))

    GC.gc()

    print("  with fixed free list allocator         ")
    alloc = FreeListAllocator{TreeNode{Int, Int, Nothing}, Int}(2 * n)
    @btime ($ans = testfree(Ref($n), $alloc))

    GC.gc()

    print("  with resizable free list allocator     ")
    alloc = FreeListAllocator{TreeNode{Int, Int, Nothing}, Int}(nothing)
    @btime ($ans = testfree(Ref($n), $alloc))

    GC.gc()

    print("  with fixed SOA allocator               ")
    store = TupleVector{TreeNode{Int, Int, Nothing}}(2 * n)
    alloc = SOAllocator{TreeNode{Int, Int, Nothing}, Int, typeof(store)}(store, 2 * n)
    @btime ($ans = testalloc(Ref($n), $alloc))

    GC.gc()

    print("  with resizable SOA allocator           ")
    store = TupleVector{TreeNode{Int, Int, Nothing}}(N)
    alloc = SOAllocator{TreeNode{Int, Int, Nothing}, Int, typeof(store)}(store, nothing)
    @btime ($ans = testalloc(Ref($n), $alloc))

    GC.gc()

    print("  with fixed free list SOA allocator     ")
    store = TupleVector{TreeNode{Int, Int, Nothing}}(2 * n)
    alloc = FreeListSOAllocator{TreeNode{Int, Int, Nothing}, Int, typeof(store)}(store, 2 * n)
    @btime ($ans = testfree(Ref($n), $alloc))

    GC.gc()

    print("  with resizable free list SOA allocator ")
    store = TupleVector{TreeNode{Int, Int, Nothing}}(N)
    alloc = FreeListSOAllocator{TreeNode{Int, Int, Nothing}, Int, typeof(store)}(store, nothing)
    @btime ($ans = testfree(Ref($n), $alloc))
end