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

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 iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el)
        el
    end
end

# ╔═╡ 45468d3a-3456-4e99-aec8-b3c41b20a426
let
	import Pkg
	Pkg.activate(".")
	Pkg.instantiate()
end

# ╔═╡ 5ad9f0fb-3688-4b15-94c1-a18e5f41eeed
begin
	using
		Random,
		Statistics,
		Distributions,
		DataFrames,
		Wavelets, WaveletsExt,
		MLJ,
		Gadfly,
		Plots,
		PlutoUI
end

# ╔═╡ aa7ba1e0-ff23-4406-96a1-4406b4e94399
begin
	include("./src/data.jl")
	include("./src/plots.jl")
	import .LDBDatasets: get_dataset, textures_label
	import .LDBPlots: plot_emap, plot_emaps
end

# ╔═╡ 45f88030-a821-11eb-0c6d-f5c7c82b7115
md"# Local Discriminant Basis
**Authors**: Shozen Dan, Zeng Fung Liew, Naoki Saito"

# ╔═╡ f3a892ef-0c5a-4d4a-86ab-036d913d9d97
md"## I. Introduction"

# ╔═╡ c195f5d9-2538-4278-9d27-c14446e7cb65
md"**Local Discriminant Basis (LDB)** is a wavelet based feature extraction method concieved by [Naoki Saito](https://www.math.ucdavis.edu/~saito). In 1992, [Victor Wickerhauser](https://www.math.wustl.edu/~victor/) had generalized the wavelet best basis algorithm such that it worked not only for a single signal, but for a collection of signals that share the same features. This so called Joint Best Basis (JBB) can be viewed as a time-frequency localized version of the Principle Component Analysis (PCA) or the Karhunen-Loève Basis (KLB).\
\
While JBB is good for signals belonging to the same class, it does not work for signal classifications in general. LDB sets out to solve this issue by replacing the original minimum entropy cost function used in the JBB with discriminant measures such as Kullback-Leibler divergence. In a nutshell, the LDB takes to following steps to extract discriminant features from a set of signals.

1. Decomposes individual signals into time-frequency dictionaries (wavelet packet decomposition).
2. Creates a time-frequency energy distribution for each class by accumulating these dictionaries.
3. A complete orthonormal basis called LDB, which encodes the time-frequency localized similarity and differences between signal classes, is computed using discriminant measures such as KL-divergence.
4. The coordinates that house the most distinguishing (the most divergent) feature of each signal class is identified.
5. The features at these coordinates are extracted from the time-frequency dictionaries of each individual signal and fed to classification algorithms such as Decision Trees or Linear Discriminant Analysis.\

Some resourses on LDB including the original paper can be found in the following links:
* [Local Discriminant Basis[1]](https://www.math.ucdavis.edu/~saito/contrib/#ldb)
* [On Local Orthonormal Bases for Classification and Regression[2]](http://math.ucdavis.edu/~saito/publications/saito_icassp95.pdf)
* [Local discriminant bases and their applications[3]](http://math.ucdavis.edu/~saito/publications/saito_ldb_jmiv.pdf)

The application capabilities of LDB has been studied in multiple studies, with a large number belonging to the bio-medical and neuroscience fields. Results have shown that LDB can successfully classify tumors (e.g. breast and brain), fungal infections, and brain waves using signals obtained from medical sensing techniques such as MRI and EEG[4][5][6].

**References**

[1]N. Saito and R. R. Coifman, “On local orthonormal bases for classification and regression,” in 1995 International Conference on Acoustics, Speech, and Signal Processing, May 1995, vol. 3, pp. 1529–1532 vol.3. doi: 10.1109/ICASSP.1995.479852.

[2]N. Saito, “Naoki’s Selected Scientific Contributions,” Naoki’s Selected Scientific Contributions. https://www.math.ucdavis.edu/~saito/contrib/#ldb

[3]N. Saito and R. R. Coifman, “Local discriminant bases and their applications,” J Math Imaging Vis, vol. 5, no. 4, pp. 337–358, Dec. 1995, doi: 10.1007/BF01250288.

[4]D. Li, W. Pedrycz, and N. J. Pizzi, “Fuzzy wavelet packet based feature extraction method and its application to biomedical signal classification,” IEEE Transactions on Biomedical Engineering, vol. 52, no. 6, pp. 1132–1139, Jun. 2005, doi: 10.1109/TBME.2005.848377.

[5]N. F. Ince, S. Arica, and A. Tewfik, “Classification of single trial motor imagery EEG recordings with subject adapted non-dyadic arbitrary time–frequency tilings,” J. Neural Eng., vol. 3, no. 3, pp. 235–244, Sep. 2006, doi: 10.1088/1741-2560/3/3/006.

[6]S. K. Davis, B. D. Van Veen, S. C. Hagness, and F. Kelcz, “Breast Tumor Characterization Based on Ultrawideband Microwave Backscatter,” IEEE Transactions on Biomedical Engineering, vol. 55, no. 1, pp. 237–246, Jan. 2008, doi: 10.1109/TBME.2007.900564.
"

# ╔═╡ a751cd87-80c5-48b1-b798-f1aecebc08a1
md"## II. Tutorial"

# ╔═╡ 4e5fe030-8a87-4f4a-88a4-b7b824157880
md"**Auto run:** Check the box to start the tutorial

$(@bind autorun CheckBox())"

# ╔═╡ b8077eb3-ec64-4a84-9dcc-3aafce015597
md"We begin by obtaining some signals to classify.

For 1D signals, a wrapper function for `WaveletsExt.generateclassdata()` is written for easier signal generations. Here, we can generate 2 different sets of signals (Cylinder-Bell-Funnel, Triangular waveform), each consisting of 3 classes of signals.

For 2D signals, one set of data is obtained from the [SIPI Image Database - Textures](https://sipi.usc.edu/database/database.php?volume=textures). Another set of data used in this experiment is the MNIST dataset, obtained from the MLDatasets.jl package."

# ╔═╡ e8182e69-8776-4ab5-a38e-bf2175ceb0ea
md"**Select** the dataset"

# ╔═╡ 910d24a0-8c00-42c5-8e11-13da2557a09d
@bind sigtype Radio(["Cylinder-Bell-Funnel","Triangular", "Textures", "MNIST"], default = "Cylinder-Bell-Funnel")

# ╔═╡ a0ae476e-7939-4bfe-83e4-9666d0ed366e
md"""
**Classes of textures for texture dataset**

*Ignore the message below if the selected dataset is not Textures. Otherwise, hold `Ctrl` + `Click` on the texture files to generate samples from. Note that each file selected corresponds to a unique class.*

$(sigtype != "Textures" || @bind textclass MultiSelect((sort ∘ collect ∘ keys)(LDBDatasets.textures_label),
	default = ["Brodatz - Grass (D9) - 512",
			   "Brodatz - Beach sand (D29) - 512",
			   "Brodatz - Plastic bubbles (D112) - 512"]))
"""

# ╔═╡ 705723ac-b0e0-4205-b3aa-8b516f9233d4
st = Dict(["Triangular" => :tri, "Cylinder-Bell-Funnel" => :cbf,
		   "Textures" => :textures, "MNIST" => :mnist]);

# ╔═╡ dc92cbff-ccf1-4355-8d60-0e2f39dac6b0
begin
	if sigtype ∈ ["Triangular", "Cylinder-Bell-Funnel"]
		args = [5, 0]
	elseif sigtype == "Textures"
		args = [textclass, 1, 0]
	else
		args = [1, 0]
	end
	(X₀, y₀), _ = get_dataset(st[sigtype], args...)
	Y₀ = coerce(y₀, Multiclass); # For compatibility with MLJ.jl
end;

# ╔═╡ 39f64d00-350d-43a6-bf57-06600aac2365
begin
	gr(size=(700,400))
	if sigtype ∈ ["Triangular", "Cylinder-Bell-Funnel"]
		p1 = wiggle(X₀[:,1:5], sc=0.5)
		Plots.plot!(xlab = "Class 1")
		p2 = wiggle(X₀[:,6:10], sc=0.5)
		Plots.plot!(xlab = "Class 2")
		p3 = wiggle(X₀[:,11:15], sc=0.5)
		Plots.plot!(xlab = "Class 3")
		Plots.plot(p1, p2, p3, layout = (3,1))
	else
		plts = []
		for (x, y) in zip(eachslice(X₀, dims=3), y₀)
			plt = heatmap(x', color = :greys, title = y, legend = false, yflip = true,
						  titlefontsize = 7, tickfontsize=5)
			push!(plts, plt)
		end
		Plots.plot(plts...)
	end
end

# ╔═╡ 3c077b0c-ad81-43bf-af45-32d7f48185c7
md"**Local Discriminant Basis (LDB)** is a feature extractor meaning that given a set of signals and their labels
$\mathbf{X} = (\mathbf{x_1},\mathbf{x_2},\mathbf{x_3},...), \mathbf{y} = (y_1,y_2,y_3,...)$ where $\mathbf{X} \subset \mathbb{R}^n$ and $y \in \{1,2,...,K\}$, LDB will find a map $f: \mathbf{X} \to \mathbf{Z} \subset \mathbb{R}^m$ such that $m \leq n$.
"

# ╔═╡ 25720fa1-ad95-451d-8143-ba34b6e0551b
md"In the context of LDB, the map $f$ can be expressed as

$f = \Theta_m \circ\Psi^T$

where, $\Psi$ is an $n$-dimensional orthogonal matrix selected from a **library of orthonormal bases**, and $\Theta_m$ is a selection rule that selects the most important $m$ features from $n$ features."

# ╔═╡ e0c92531-1e17-4757-885d-408d62829766
md"After the features have been extracted, we can use a conventional classifier such as Linear Discriminant Anlaysis (LDA), Classfication Trees (CART), or Multiclass Logistic Regression to map $\mathbf{Z} \to \mathbf{y}$."

# ╔═╡ dc5ba00d-1a5b-4233-96a6-73981882345a
md"### Library of Orthonormal Bases"

# ╔═╡ 01a254f4-9575-4ab2-af6a-27ad5ef8efde
md"The LDB algorithm begins by decomposing each signal into a dictionary of orthonormal bases. Each dictionary is a binary tree where each node is a subspace of $\mathbf{w}_{0,0} = \mathbb{R}^n$. Examples of dictionaries include wavelet packet bases, local trigonometric bases, and local Fourier bases. In the figure below, we illustrate the wavelet packet bases where $G$, $H$, and $f$ represent the scaling (lowpass) analysis filters, wavelet (highpass) analysis filters, and the frequency axis respectively. By decomposing a signal as such, information about the signal at a specific time and frequency can be encoded within each node, which can then be used to distinguish different classes of signals."

# ╔═╡ 9f7c2639-c455-425a-a2ab-0deac638b47f
img_url = "https://raw.githubusercontent.com/ShozenD/LDBExperiments/main/images/wpt-figure.png";

# ╔═╡ b2db449c-0fe5-482a-9e85-9062a218df03
md"""$(Resource(img_url, :width => 500))"""

# ╔═╡ 964f8fcd-0516-4c6f-a02a-6db5dd497520
em_img_url = "https://raw.githubusercontent.com/ShozenD/LDBExperiments/main/images/normalized-energy-map.png";

# ╔═╡ 1c4b794f-2b17-429d-809a-2f69f0a82e41
md"### Energy Map"

# ╔═╡ a3b05137-0499-45c1-bbce-79784dbf59dc
md"**Normalized energies**"

# ╔═╡ 55528ce0-3374-4d13-bb6f-61df9f854a39
md"
$V_i^{(y)} \triangleq
\frac{
	E\left[Z_i^2|Y=y\right]
}{
	\sum_{i=1}^n E\left[Z_i^2|Y=y\right]
}
\to
\hat{V}_i^{(y)} =

\frac{
	\sum_{k=1}^{N_y} \left\| \boldsymbol{w}_i \cdot \boldsymbol{x}_k^{(y)}\right\|^2
}{
	\sum_k^{N_y} \left\| \boldsymbol{x}_k^{(y)} \right\|^2
}$"

# ╔═╡ 8d2707f8-a982-4c83-b14a-1a2deb6743b4
md"""$(Resource(em_img_url))"""

# ╔═╡ fb91da71-303f-4c43-be7b-e39df1429355
md"**Probability density**

Another way to estimate $E\left[Z_i^2|Y=y\right]$ is to use kernel density estimators. The LDB algorithm in `WaveletsExt.jl` uses a method called Average Shifted Histograms(ASH)."

# ╔═╡ af1b235e-6fff-478f-a5c1-38fbc6c39b8f
md"
$q_i^{(y)}(z) \triangleq

\int_{\boldsymbol{w}_i \cdot \boldsymbol{x}=\boldsymbol{z}}
	p(\boldsymbol{x}|y)d\boldsymbol{x} \to \hat{q}_i^{(y)}(z)$"

# ╔═╡ d27634a5-e703-4fa2-bc1a-d7297b2388a3
begin
	pem_img_url = "https://raw.githubusercontent.com/ShozenD/LDBExperiments/main/images/probability-density-tensor.png";

	md"""$(Resource(pem_img_url))"""
end

# ╔═╡ e0a525ed-35f0-48cc-8403-ddfe03871074
md"**Select** the type of energy map to use"

# ╔═╡ 9eee7238-6e9c-4837-a30b-ebd09abdcca6
@bind e_measure Radio(
	[
		"Time Frequency",
		"Probability Density",
	],
	default = "Time Frequency"
)

# ╔═╡ fd63c142-ae62-40a2-b34f-986c803ddb72
em = Dict(["Time Frequency"=>TimeFrequency(), "Probability Density"=>ProbabilityDensity()]);

# ╔═╡ 4f8a7bb5-db64-4f82-8544-c961cca068db
md"### Discriminant Measure"

# ╔═╡ a0525192-f1d9-4173-960e-ea3c009e067b
md"For each cell within the energy map, we can quatify the difference between each pair of classes by choosing an appropriate chosen discriminant measure. `WaveletsExt.jl` implements four different types of discriminant measures: asymmetic relative entropy, symmetric relative entropy, squared euclidean distance, and Hellinger distance. The definition for each type of discriminant measure is shown below."

# ╔═╡ e77667ca-9bb8-4f30-b5ba-ff107eb9a568
md"Asymmetric Relative Entropy (Kullback-Leibler divergence):

$D_{KL}(p,q) = \int_{-\infty}^{\infty}p(x)\log_2\frac{p(x)}{q(x)}dx$"

# ╔═╡ 885cc8dd-dc5a-4d28-be72-2e26613ec252
md"Symmetric Relative Entropy

$$D_S(p,q) = D_{KL}(p,q) + D_{KL}(q,p)$$
"

# ╔═╡ ed92e98f-e823-45a6-a401-342f584c333e
md"Squared Euclidean Distance

$D_P(p,q) = \int_{-\infty}^{\infty}\left(p(x) - q(x)\right)^2dx$"

# ╔═╡ 0b12ee12-9229-486f-aa65-1da5c53955cc
md"Hellinger distance

$D_H(p,q)=\int_{-\infty}^{\infty}\left(\sqrt{p(x)} - \sqrt{q(x)}\right)^2dx$"

# ╔═╡ 05a9e8db-fce0-4d12-b67b-c0089621ae07
md"**Select** the type of discriminant measure to use"

# ╔═╡ 28604f68-a957-4a3c-92f5-13a0ff4ba158
@bind d_measure Radio(
	[
		"Asymmetric Relative Entropy",
		"Symmetric Relative Entropy",
		"Squared Euclidean Distance",
		"Hellinger Distance"
	],
	default = "Asymmetric Relative Entropy"
)

# ╔═╡ b27a4714-cbda-417e-85e1-26d7d98780ee
dm = Dict([
		"Asymmetric Relative Entropy" => AsymmetricRelativeEntropy(),
		"Squared Euclidean Distance" => LpDistance(),
		"Symmetric Relative Entropy" => SymmetricRelativeEntropy(),
		"Hellinger Distance" => HellingerDistance()
	]);

# ╔═╡ 7dc1ba75-fa54-4d59-9b54-93e01da7211e
dmm_img_url = "https://raw.githubusercontent.com/ShozenD/LDBExperiments/main/images/discriminant-measure-map.png";

# ╔═╡ 86df0e34-8f77-4d6e-8f40-6f4d8e706c15
md"By adding the discriminant measures between each pair of classes together, we can obtain a matrix where the most discriminant cells (i.e. features) have the largest values and the least discriminant cells have the lowest values."

# ╔═╡ ad3fe2dc-8003-451c-bf83-c3c7f24e7f0b
md"""$(Resource(dmm_img_url))"""

# ╔═╡ f0c4b67e-208e-41a4-9510-c47e04a65e20
prn_img_url = "https://raw.githubusercontent.com/ShozenD/LDBExperiments/main/images/pruning.png";

# ╔═╡ 9c4cf3a1-cd6a-4a42-acce-ceaca6c66df2
md"We can use this discriminant measure matrix to prune our binary tree. The pruning algorithm will start from the bottom of the tree and will eliminate the children nodes if the sum of their discriminant measure is smaller or equal to the discriminant measure of their parent node. This pruning process can not only identify features that are most disciminant, but also reduce the computational resources required for to perform LDB."

# ╔═╡ 299c91a6-c315-4299-97fc-57a6b6f419b5
md"""$(Resource(prn_img_url))"""

# ╔═╡ df7c5ef9-73ff-44b7-aacb-d5fa132d7c2b
md"After the pruning, we can sort the coefficients by their discriminant power and extract the top n features (wavelet coefficient) to perform dimensionality reduction. We can pass the selected features to ML classification algorithms such as logistic regression, support vector machine, random forests, etc."

# ╔═╡ f1c6268b-a6d5-445a-9e52-748898ec08da
md"**Select** the number of features to extract"

# ╔═╡ 9e523918-daaf-4c17-851a-7fac12b04cd3
begin
	sz = ndims(X₀) == 2 ? size(X₀,1) : size(X₀,1) * size(X₀,2)
	@bind dim Slider(1:sz, default=10, show_value=true)
end

# ╔═╡ 22a8dfdb-c046-4614-8e2b-aab22d12b205
md"To run the LDB algorithm using `WaveletsExt.jl`, we begin by calling `LocalDiscriminantBasis` constructor to specify the parameters. We need to pass 4 arguments to the constructor:

1. The type of wavelet filter to use
2. The type of discriminant measure to use
3. The type of energy map to use
4. The number of features to extract (can be changed at a later stage)"

# ╔═╡ 09bc8c83-2f25-44a0-81ab-9f0a5724673f
md"After the model parameters are specified, the LDB model is trained using the `fit!` method. This follows the syntax in `MLJ.jl`."

# ╔═╡ 2c9b5aef-ba63-44b6-8ef9-ca31cc682fad
if autorun
	ldb₀ = LocalDiscriminantBasis(wt = wavelet(WT.coif6),
								  dm = dm[d_measure],
								  en = em[e_measure],
								  n_features = dim)
	WaveletsExt.fit!(ldb₀, X₀, y₀)
end

# ╔═╡ ae624404-0770-41e9-962b-139006356ea6
md"`WaveletsExt.jl` provides interfaces to understand the results of LDB. For example we can visualize the energy/probability density maps for each class by accessing the `Γ` attribute in the fitted LDB object. The plots below shows the energy/probability density maps for each class, where darker colors indicate that the sub-band has a higher discriminant power.

*Note: When probability density maps are selected, a slider will appear below the energy maps below. Use the slider to select the position of the probability density function and compare the energy maps between classes.*"

# ╔═╡ f1727ea3-7353-4c6a-bc13-51622c615780
if autorun && ldb₀.en == ProbabilityDensity()
	pdf_len = st[sigtype] ∈ [:tri, :cbf] ? size(ldb₀.Γ,3) : size(ldb₀.Γ,4)
	@bind i Slider(1:pdf_len)
end

# ╔═╡ f7669f3f-9e34-4dc1-b3d4-7eda7fe6e479
begin
	if autorun
		if ldb₀.en == TimeFrequency()
			st[sigtype]∈[:tri,:cbf] ? gr(size=(700,600)) : gr(size=(600,1800))
			plot_emaps(ldb₀.Γ, ldb₀.tree)
		elseif ldb₀.en == ProbabilityDensity() && st[sigtype] ∈ [:tri, :cbf]
			gr(size=(700,600))
			plot_emaps(ldb₀.Γ[:,:,i,:], ldb₀.tree)
		elseif ldb₀.en == ProbabilityDensity() && st[sigtype] ∈ [:textures, :mnist]
			gr(size=(600,1800))
			plot_emaps(ldb₀.Γ[:,:,:,i,:], ldb₀.tree)
		end
	end
end

# ╔═╡ 82ffc65d-54ea-42ae-a7ef-dbe6216b0d1e
md"We can also visualize the discriminant measure map and the selected nodes/sub-bands using the `discriminant_measure` function and the `tree` attributes respectively. The first plot in the figure below displays the heatmap for the discriminant measure map and the second plot displays the selected nodes from the binary tree."

# ╔═╡ 121fd299-5e82-4159-8472-5d385c736c18
begin
	if autorun
		st[sigtype]≠:textures ? gr(size=(700,400)) : gr(size=(700,1800))
		dmap = plot_emap(discriminant_measure(ldb₀.Γ, dm[d_measure]), ldb₀.tree)
		tmap = st[sigtype]∈[:tri,:cbf] ? plot_tfbdry(ldb₀.tree) : plot_tfbdry2(ldb₀.tree)
		Plots.plot(dmap, tmap, layout=(2,1))
	end
end

# ╔═╡ 96a49a0c-4646-43f9-98c2-09ac484f6df8
md"## III. Signal Classification Experiment"

# ╔═╡ 406e7ffe-fa01-4622-ae09-aca5473bfe6c
md"In this section, we will evaluate the classification capabilities of LDB via a simulation experiment. We will generate 3 classes of signals, use LDB to extract the most discriminant features, and classify the signals using various ML methods."

# ╔═╡ a5126ad3-19b1-4b4e-b96f-ef6b5220854b
md"**Select** dataset"

# ╔═╡ f9a60263-7ebd-4df8-b33f-5f0e85719186
@bind sigtype2 Radio(["Cylinder-Bell-Funnel","Triangular", "Textures", "MNIST"], default = "Cylinder-Bell-Funnel")

# ╔═╡ 6fef2648-058a-4136-8108-38c1624a19ca
begin
	if st[sigtype2] == :tri
		@bind dim2 Slider(1:32, default=10, show_value=true)
	elseif st[sigtype2] == :cbf
		@bind dim2 Slider(1:128, default=10, show_value=true)
	elseif st[sigtype2] == :textures
		@bind dim2 Slider(1:128^2, default=10, show_value=true)
	elseif st[sigtype2] == :mnist
		@bind dim2 Slider(1:32^2, default=10, show_value=true)
	end
end

# ╔═╡ 8f8bb837-ed86-4a75-8254-913530ed8bc5
md"""
**Classes of textures for texture dataset**

*Ignore the message below if the selected dataset is not Textures. Otherwise, hold `Ctrl` + `Click` on the texture files to generate samples from. Note that each file selected corresponds to a unique class.*

$(sigtype2 != "Textures" || @bind textclass2 MultiSelect((sort ∘ collect ∘ keys)(LDBDatasets.textures_label),
	default = ["Brodatz - Grass (D9) - 512",
			   "Brodatz - Beach sand (D29) - 512",
			   "Brodatz - Plastic bubbles (D112) - 512"]))
"""

# ╔═╡ 2999257a-03bf-4313-8dd6-d2da0ed2ef9c
md"""**Select** the type of wavelet to use

$(@bind wavelet_type Select(
	["WT.haar",
	"WT.db1", "WT.db2", "WT.db3", "WT.db4", "WT.db5",
	"WT.db6", "WT.db7", "WT.db8", "WT.db9", "WT.db10",
	"WT.coif2", "WT.coif4", "WT.coif6", "WT.coif8",
	"WT.sym4", "WT.sym5", "WT.sym6", "WT.sym7", "WT.sym8", "WT.sym9", "WT.sym10",
	"WT.batt2", "WT.batt4", "WT.batt6"],
	default = "WT.coif6"
))"""

# ╔═╡ 65d45fbd-09bf-49e9-b027-e43751ce070f
md"**Select** the type of energy map"

# ╔═╡ bf8314d6-eb38-4594-afb0-eed5f3151389
@bind e_measure2 Radio(
	[
		"Time Frequency",
		"Probability Density",
	],
	default = "Time Frequency"
)

# ╔═╡ dde5cc7d-1840-49a9-bcd0-bf3ed6e66007
md"**Select** the type of discriminant measure"

# ╔═╡ 1b71478b-a386-416b-97dc-a2e5da1ce071
@bind d_measure2 Radio(
	[
		"Asymmetric Relative Entropy",
		"Symmetric Relative Entropy",
		"Lp Entropy",
		"Hellinger Distance"
	],
	default = "Asymmetric Relative Entropy"
)

# ╔═╡ 2bf768ee-7bc7-4039-85d2-dbdf9ed1f75a
md"**Select** the numbers of features to extract"

# ╔═╡ 60cffbcf-d539-4038-9a12-c40fa41d6880
md"**Auto Run**: Check the box after you are satisfied with the experiment parameters or when you want to re-run the experiment (uncheck and check again)

$(@bind autorun2 CheckBox())"

# ╔═╡ 7a9548a4-c528-41af-bba7-a99b0c91247b
begin
	if autorun2
		dfm = DataFrame(name=String[], model=[], predtype=Symbol[])
		machines = Dict() # Models
		X = (train=Dict(), test=Dict())
		y = (train=Dict(), test=Dict())
		df = DataFrame(Type = String[], TrainAcc = Float64[], TestAcc = Float64[])
	end
end;

# ╔═╡ 4774bfcf-9e50-428c-b51f-76a887031862
begin
	if autorun2
		if sigtype2 ∈ ["Triangular", "Cylinder-Bell-Funnel", "MNIST"]
			args2 = [33, 333]
		else
			args2 = [textclass2, 33, 333]
		end
		(X_train, y_train), (X_test, y_test) = get_dataset(st[sigtype2], args2...)

		nclass_train, nclass_test = (length∘unique)(y_train), (length∘unique)(y_test)
		X.train["STD"]= DataFrame(reshape(X_train,:,33*nclass_train)', :auto)
		X.test["STD"] = DataFrame(reshape(X_test,:,333*nclass_test)', :auto)
		y.train["STD"], y.test["STD"] = coerce(y_train, Multiclass), coerce(y_test, Multiclass)
	end
end;

# ╔═╡ 2d398c73-37bc-44d4-8559-e220de94624d
md"Next, we will load some machine learning models from `MLJ.jl`. We will include two very basic decision tree models(with and without pruning), Linear Discriminant classifier (LDA), Multinomial classifier with L1 regularization (i.e., LASSO), and finally a Random Forest classifier."

# ╔═╡ 7a7cae84-3272-4303-80fa-d56a8615b9ff
begin
	Tree = @load DecisionTreeClassifier pkg=DecisionTree
	LDA = @load LDA pkg=MultivariateStats
	MCLogit = @load MultinomialClassifier pkg=MLJLinearModels
	RForest = @load RandomForestClassifier pkg=DecisionTree
	SVC = @load SVC pkg=LIBSVM
end;

# ╔═╡ 54fdabe6-85ff-4928-ac1c-1555d89ce456
md"Intialize the ML models"

# ╔═╡ 9ddb4726-5a78-4adf-a3eb-0796636467c1
begin
	function add_model!(df::DataFrame,
						name::String,
						model,
						predtype=info(model).prediction_type)

		M = Dict(:name=>name, :model=>model, :predtype=>predtype)
		push!(df, M)
	end

	if autorun2
		add_model!(dfm, "FCT", Tree())
		add_model!(dfm, "PCT", Tree(post_prune=true, merge_purity_threshold=0.8))
		add_model!(dfm, "LDA", LDA())
		add_model!(dfm, "MCLogit", MCLogit(penalty=:l1, lambda=0.01))
		add_model!(dfm, "Rforest", RForest())
		add_model!(dfm, "SVC", SVC())
	end
end

# ╔═╡ 51d745c9-8c1a-41ef-8ee6-c5e9c35d39fe
md"### 1. Training models using the original signal"

# ╔═╡ b0e3e833-47d6-493e-bb51-940267e6f85d
md"To evaluate the LDB algorithm's performance we first train the models using the original signals as input (i.e., Euclidean coordinates). To evaluate the training loss, we will perform a 5 fold cross validation."

# ╔═╡ fded58ea-e7b1-4be1-b952-b7aa1358d5dd
function evaluate_model!(dfm::DataFrame,
						 df::DataFrame,
					     model::String,
						 dat::String,
						 X, y)
	name = model * "-" * dat

	# Training error
	train, test = partition(eachindex(X.train[dat][:,1]), 0.7, shuffle=true)
	M = first(dfm[dfm.name.==model,:model])
	predtype = first(dfm[dfm.name.==model,:predtype])

	mach₀ = machine(M,
					X.train[dat][train,:],
					y.train["STD"][train])
	MLJ.fit!(mach₀)
	if predtype == :deterministic
		ŷ₀ = predict(mach₀, X.train[dat][test,:])
	else
		ŷ₀ = predict_mode(mach₀, X.train[dat][test,:])
	end

	trainacc = Accuracy()(ŷ₀, y.train["STD"][test])

	mach = machine(M, X.train[dat], y.train["STD"])
	MLJ.fit!(mach)

	# Test error
	if predtype == :deterministic
		ŷ = predict(mach, X.test[dat])
	else
		ŷ = predict_mode(mach, X.test[dat])
	end

	testacc = Accuracy()(ŷ, y.test["STD"])

	push!(df, Dict(:Type=>name, :TrainAcc=>trainacc, :TestAcc=>testacc))
end

# ╔═╡ 19e7d3f3-970d-4d05-9664-8fe23009fb71
begin
	if autorun2
		for name in dfm.name
			evaluate_model!(dfm, df, name, "STD", X, y)
		end
	end
end

# ╔═╡ 4ffbeab9-67c5-46a0-8e09-449d91dfa34c
if autorun2
	df
end

# ╔═╡ 97516db4-c019-49a7-b826-64294fd14220
md"### Using LDB-k"

# ╔═╡ c47d64f3-12fc-4628-9162-21980066bd00
md"Next, we significantly reduce the dimensionality of the models by only using the top k most discriminant features obtained from LDB."

# ╔═╡ a828877d-1f49-4b76-b397-869bb11e40c5
begin
	if autorun2
		ldbk = "LDB-" * string(dim2)

		ldb = LocalDiscriminantBasis(wt = wavelet(eval(Meta.parse(wavelet_type))),
									 dm = dm[d_measure],
									 en = em[e_measure],
									 n_features = dim2)
		WaveletsExt.fit!(ldb, X_train, y_train)
		X.train[ldbk] = DataFrame(WaveletsExt.transform(ldb, X_train)', :auto);
		X.test[ldbk] = DataFrame(WaveletsExt.transform(ldb, X_test)', :auto);
	end
end;

# ╔═╡ 34ff82ef-e7a7-4df2-ab71-3280a5ef34fe
begin
	if autorun2
		for name in dfm.name
			evaluate_model!(dfm, df, name, ldbk, X, y)
		end
	end
end

# ╔═╡ 407cce96-73cb-4baf-90f9-b46d5d617018
if autorun2
	df
end

# ╔═╡ 7dd079af-0445-436c-9bd3-9550cfaa9b19
md"### 3. Using all LDB features"

# ╔═╡ b31e54a1-f1b7-44c4-b2bc-99123933c289
md"Finally, we use all the LDB features to train our models. Note that we do not include the LDA model because theoretically it is the same with using the euclidean coordinates."

# ╔═╡ a08dd1ea-a403-41b4-915c-56fde82222e7
if autorun2
	ldb₁ = LocalDiscriminantBasis(
		wt = wavelet(eval(Meta.parse(wavelet_type))),
		dm = dm[d_measure2],
		en = em[e_measure2]
	)
end

# ╔═╡ 3a374ac2-e225-4637-9dbd-6644cb80b629
begin
	if autorun2
		WaveletsExt.fit!(ldb₁, X_train, y_train)
		X.train["LDB"] = DataFrame(WaveletsExt.transform(ldb₁, X_train)', :auto)
		X.test["LDB"] = DataFrame(WaveletsExt.transform(ldb₁, X_test)', :auto)
		for machine in dfm.name
			evaluate_model!(dfm, df, machine, "LDB", X, y)
		end
	end
end

# ╔═╡ 3f3d4bcf-3f2b-4140-ba52-2246c5140303
if autorun2
	sort!(df, :TestAcc, rev=true)
end

# ╔═╡ 9b292713-1580-48ae-b9cc-05dca097a673
md"## Results"

# ╔═╡ c1b823e9-4a80-4cca-9527-5b0f2933933d
begin
	if autorun2
		set_default_plot_size(18cm, 16cm)
		Gadfly.plot(
			sort(df, :TestAcc),
			layer(
				x=:TestAcc, y=:Type,
				Geom.point,
				color=[colorant"#de425b"],
				size = [1.5mm],
				alpha = [0.7]
			),
			layer(
				x=:TrainAcc, y=:Type,
				Geom.point,
				color=[colorant"#488f31"],
				size = [1.5mm],
				alpha = [0.7]
			),
			layer(
				xintercept = [1/3],
				Geom.vline(color = "gray", size = [1mm])
			),
			Guide.title("Model Performance"),
			Guide.xlabel("Accuracy"),
			Guide.manual_color_key(
				"",
				["Test Acc","Train Acc", "Baseline"],
				[colorant"#de425b", colorant"#488f31", "gray"],
				size = [1.5mm],
				shape = [Gadfly.Shape.circle, Gadfly.Shape.circle, Gadfly.Shape.square]
			)
		)
	end
end

# ╔═╡ e58b5f6a-c4da-4125-be5a-7f9125d3f326

# ╔═╡ Cell order:
# ╟─45f88030-a821-11eb-0c6d-f5c7c82b7115
# ╠═45468d3a-3456-4e99-aec8-b3c41b20a426
# ╠═5ad9f0fb-3688-4b15-94c1-a18e5f41eeed
# ╠═aa7ba1e0-ff23-4406-96a1-4406b4e94399
# ╟─f3a892ef-0c5a-4d4a-86ab-036d913d9d97
# ╟─c195f5d9-2538-4278-9d27-c14446e7cb65
# ╟─a751cd87-80c5-48b1-b798-f1aecebc08a1
# ╟─4e5fe030-8a87-4f4a-88a4-b7b824157880
# ╟─b8077eb3-ec64-4a84-9dcc-3aafce015597
# ╟─e8182e69-8776-4ab5-a38e-bf2175ceb0ea
# ╟─910d24a0-8c00-42c5-8e11-13da2557a09d
# ╟─a0ae476e-7939-4bfe-83e4-9666d0ed366e
# ╟─705723ac-b0e0-4205-b3aa-8b516f9233d4
# ╟─dc92cbff-ccf1-4355-8d60-0e2f39dac6b0
# ╟─39f64d00-350d-43a6-bf57-06600aac2365
# ╟─3c077b0c-ad81-43bf-af45-32d7f48185c7
# ╟─25720fa1-ad95-451d-8143-ba34b6e0551b
# ╟─e0c92531-1e17-4757-885d-408d62829766
# ╟─dc5ba00d-1a5b-4233-96a6-73981882345a
# ╟─01a254f4-9575-4ab2-af6a-27ad5ef8efde
# ╟─9f7c2639-c455-425a-a2ab-0deac638b47f
# ╟─b2db449c-0fe5-482a-9e85-9062a218df03
# ╟─964f8fcd-0516-4c6f-a02a-6db5dd497520
# ╟─1c4b794f-2b17-429d-809a-2f69f0a82e41
# ╟─a3b05137-0499-45c1-bbce-79784dbf59dc
# ╟─55528ce0-3374-4d13-bb6f-61df9f854a39
# ╟─8d2707f8-a982-4c83-b14a-1a2deb6743b4
# ╟─fb91da71-303f-4c43-be7b-e39df1429355
# ╟─af1b235e-6fff-478f-a5c1-38fbc6c39b8f
# ╟─d27634a5-e703-4fa2-bc1a-d7297b2388a3
# ╟─e0a525ed-35f0-48cc-8403-ddfe03871074
# ╟─9eee7238-6e9c-4837-a30b-ebd09abdcca6
# ╟─fd63c142-ae62-40a2-b34f-986c803ddb72
# ╟─4f8a7bb5-db64-4f82-8544-c961cca068db
# ╟─a0525192-f1d9-4173-960e-ea3c009e067b
# ╟─e77667ca-9bb8-4f30-b5ba-ff107eb9a568
# ╟─885cc8dd-dc5a-4d28-be72-2e26613ec252
# ╟─ed92e98f-e823-45a6-a401-342f584c333e
# ╟─0b12ee12-9229-486f-aa65-1da5c53955cc
# ╟─05a9e8db-fce0-4d12-b67b-c0089621ae07
# ╟─28604f68-a957-4a3c-92f5-13a0ff4ba158
# ╟─b27a4714-cbda-417e-85e1-26d7d98780ee
# ╟─7dc1ba75-fa54-4d59-9b54-93e01da7211e
# ╟─86df0e34-8f77-4d6e-8f40-6f4d8e706c15
# ╟─ad3fe2dc-8003-451c-bf83-c3c7f24e7f0b
# ╟─f0c4b67e-208e-41a4-9510-c47e04a65e20
# ╟─9c4cf3a1-cd6a-4a42-acce-ceaca6c66df2
# ╟─299c91a6-c315-4299-97fc-57a6b6f419b5
# ╟─df7c5ef9-73ff-44b7-aacb-d5fa132d7c2b
# ╟─f1c6268b-a6d5-445a-9e52-748898ec08da
# ╟─9e523918-daaf-4c17-851a-7fac12b04cd3
# ╟─22a8dfdb-c046-4614-8e2b-aab22d12b205
# ╟─09bc8c83-2f25-44a0-81ab-9f0a5724673f
# ╠═2c9b5aef-ba63-44b6-8ef9-ca31cc682fad
# ╟─ae624404-0770-41e9-962b-139006356ea6
# ╟─f7669f3f-9e34-4dc1-b3d4-7eda7fe6e479
# ╟─f1727ea3-7353-4c6a-bc13-51622c615780
# ╟─82ffc65d-54ea-42ae-a7ef-dbe6216b0d1e
# ╟─121fd299-5e82-4159-8472-5d385c736c18
# ╟─96a49a0c-4646-43f9-98c2-09ac484f6df8
# ╟─406e7ffe-fa01-4622-ae09-aca5473bfe6c
# ╟─a5126ad3-19b1-4b4e-b96f-ef6b5220854b
# ╟─f9a60263-7ebd-4df8-b33f-5f0e85719186
# ╟─6fef2648-058a-4136-8108-38c1624a19ca
# ╟─8f8bb837-ed86-4a75-8254-913530ed8bc5
# ╟─2999257a-03bf-4313-8dd6-d2da0ed2ef9c
# ╟─65d45fbd-09bf-49e9-b027-e43751ce070f
# ╟─bf8314d6-eb38-4594-afb0-eed5f3151389
# ╟─dde5cc7d-1840-49a9-bcd0-bf3ed6e66007
# ╟─1b71478b-a386-416b-97dc-a2e5da1ce071
# ╟─2bf768ee-7bc7-4039-85d2-dbdf9ed1f75a
# ╟─60cffbcf-d539-4038-9a12-c40fa41d6880
# ╠═7a9548a4-c528-41af-bba7-a99b0c91247b
# ╠═4774bfcf-9e50-428c-b51f-76a887031862
# ╟─2d398c73-37bc-44d4-8559-e220de94624d
# ╠═7a7cae84-3272-4303-80fa-d56a8615b9ff
# ╟─54fdabe6-85ff-4928-ac1c-1555d89ce456
# ╠═9ddb4726-5a78-4adf-a3eb-0796636467c1
# ╟─51d745c9-8c1a-41ef-8ee6-c5e9c35d39fe
# ╟─b0e3e833-47d6-493e-bb51-940267e6f85d
# ╠═fded58ea-e7b1-4be1-b952-b7aa1358d5dd
# ╠═19e7d3f3-970d-4d05-9664-8fe23009fb71
# ╟─4ffbeab9-67c5-46a0-8e09-449d91dfa34c
# ╟─97516db4-c019-49a7-b826-64294fd14220
# ╟─c47d64f3-12fc-4628-9162-21980066bd00
# ╠═a828877d-1f49-4b76-b397-869bb11e40c5
# ╠═34ff82ef-e7a7-4df2-ab71-3280a5ef34fe
# ╠═407cce96-73cb-4baf-90f9-b46d5d617018
# ╟─7dd079af-0445-436c-9bd3-9550cfaa9b19
# ╟─b31e54a1-f1b7-44c4-b2bc-99123933c289
# ╠═a08dd1ea-a403-41b4-915c-56fde82222e7
# ╠═3a374ac2-e225-4637-9dbd-6644cb80b629
# ╠═3f3d4bcf-3f2b-4140-ba52-2246c5140303
# ╟─9b292713-1580-48ae-b9cc-05dca097a673
# ╠═c1b823e9-4a80-4cca-9527-5b0f2933933d
# ╠═e58b5f6a-c4da-4125-be5a-7f9125d3f326


text:
### Abstract parameter IOStreams

abstract ParameterIOStream{S<:ValueSupport} <: VariableIOStream

codegen(f::Symbol, iostream::ParameterIOStream, fnames::Vector{Symbol}) = codegen(Val{f}, iostream, fnames)

Base.open{S<:AbstractString}(iostream::ParameterIOStream, names::Vector{S}, mode::AbstractString="w") =
  iostream.open(iostream, names, mode)

close(iostream::ParameterIOStream) = iostream.close(iostream)

Base.mark(iostream::ParameterIOStream) = iostream.mark(iostream)

Base.reset(iostream::ParameterIOStream) = iostream.reset(iostream)

Base.flush(iostream::ParameterIOStream) = iostream.flush(iostream)

Base.write{S<:ValueSupport, F<:VariateForm}(iostream::ParameterIOStream{S}, state::ParameterState{S, F}) =
  iostream.write(iostream, state)

Base.writemime(io::IO, ::MIME"text/plain", iostream::ParameterIOStream) = show(io, iostream)

ext: jl


text:
@rule typeof(dot)(:out, Marginalisation) (m_in1::NormalDistributionsFamily, m_in2::PointMass, meta::AbstractCorrection) = begin
    return @call_rule typeof(dot)(:out, Marginalisation) (m_in1 = m_in2, m_in2 = m_in1, meta = meta)
end

@rule typeof(dot)(:out, Marginalisation) (m_in1::PointMass, m_in2::NormalDistributionsFamily, meta::AbstractCorrection) = begin
    A = mean(m_in1)
    in2_mean, in2_cov = mean_cov(m_in2)
    return NormalMeanVariance(dot(A, in2_mean), A' * in2_cov * A)
end


text:
export MachODynamicLinks, MachODynamicLink, MachORPath

struct MachODynamicLinks{H <: MachOHandle, C <: MachOLoadCmdRef{H,MachOLoadDylibCmd{H}}} <: DynamicLinks{H}
    handle::H
    cmds::Vector{C}
end

function MachODynamicLinks(lcs::MachOLoadCmds)
    ld_cmds = findall(lcs, [MachOLoadDylibCmd])
    return MachODynamicLinks(handle(lcs), ld_cmds)
end
MachODynamicLinks(oh::MachOHandle) = MachODynamicLinks(MachOLoadCmds(oh))

DynamicLinks(oh::MachOHandle) = MachODynamicLinks(oh)
DynamicLinks(lcs::MachOLoadCmds) = MachODynamicLinks(lcs)

lastindex(dls::MachODynamicLinks) = lastindex(dls.cmds)
function getindex(dls::MachODynamicLinks, idx)
    return MachODynamicLink(dls, dls.cmds[idx])
end
handle(dls::MachODynamicLinks) = dls.handle

"""
    MachODynamicLink

MachO type representing a dynamic link between a MachO file and one of its
dynamic dependencies.  Although Mach-O encodes more than just the path of the
dependency, in order to get at it you will need to dig into the LoadCmd that
describes it.
"""
struct MachODynamicLink{H <: MachOHandle} <: DynamicLink{H}
    dls::MachODynamicLinks
    cmd_ref::MachOLoadCmdRef{H,MachOLoadDylibCmd{H}}
end

DynamicLinks(dl::MachODynamicLink) = dl.dls
handle(dl::MachODynamicLink) = handle(DynamicLinks(dl))
path(dl::MachODynamicLink) = dylib_name(deref(dl.cmd_ref))

struct MachORPath{H <: MachOHandle} <: RPath{H}
    handle::H
    cmds::Vector
end

function MachORPath(lcs::MachOLoadCmds)
    cmds = findall(lcs, [MachORPathCmd])
    return MachORPath(handle(lcs), cmds)
end

RPath(oh::MachOHandle) = MachORPath(oh)
MachORPath(oh::MachOHandle) = MachORPath(MachOLoadCmds(oh))
handle(mrp::MachORPath) = mrp.handle

function rpaths(mrp::MachORPath)
    return [rpath(c) for c in mrp.cmds]
end


text:
using BinaryBuilder

cuda_version = v"10.2.89"   # NOTE: could be less specific

sources_linux_x64 = [
    ArchiveSource("https://developer.nvidia.com/compute/machine-learning/cudnn/secure/8.0.4/10.2_20200923/cudnn-10.2-linux-x64-v8.0.4.30.tgz",
                  "c12c69eb16698eacac40aa46b9ce399d4cd86efb6ff0c105142f8a28fcfb980e")
]
sources_linux_ppc64le = [
    ArchiveSource("https://developer.nvidia.com/compute/machine-learning/cudnn/secure/8.0.4/10.2_20200923/cudnn-10.2-linux-ppc64le-v8.0.4.30.tgz",
                  "32a5b92f9e1ef2be90e10f220c4ab144ca59d215eb6a386e93597f447aa6507e")
]
sources_windows = [
    ArchiveSource("https://developer.nvidia.com/compute/machine-learning/cudnn/secure/8.0.4/10.2_20200923/cudnn-10.2-windows10-x64-v8.0.4.30.zip",
                  "45dc48f6c19d62e4e4f56da216b72cf69454db374ffb6589aca0cf7c309762fa")
]

include("../common.jl")


text:
using Documenter, SumOfSquares

makedocs(
    format = :html,
    sitename = "SumOfSquares",
    pages = [
        "Index" => "index.md",
        "Sum-of-Squares Programming" => "sumofsquares.md",
        "Variables" => "variables.md",
        "Constraints" => "constraints.md",
    ],
    # The following ensures that we only include the docstrings from
    # this module for functions define in Base that we overwrite.
    modules = [SumOfSquares]
)

deploydocs(
    repo   = "github.com/JuliaOpt/SumOfSquares.jl.git",
    target = "build",
    osname = "linux",
    julia  = "1.0",
    deps   = nothing,
    make   = nothing
)


text:
"""
A node type that applies the swing equation and voltage dynmaics
to the frequency/angle dynamics.
In the following, we followed the implementation of the Extended Model/ third
order model including voltage dynamics according to K. Schmietendorf et. al.
"Self-Organized Synchronization and Voltage Stability in Networks of Synchronous
Machines", 2013.
The third order model is a special case of the fourth order model with e_d = 0
and X_d = X_d_dash.

In addition to the third order model the Type IEEEDC1A exciter was implemented.

Additionally to u, it has the internal dynamic variable ω representing the
frequencyof the rotator relative to the grid frequency Ω, i.e. the real
frequency ωr of the rotator is given as ωr=Ω+ω.

Keyword Arguments

Third Order Model
- H: inertia constant (given in [s]
- P: active (real) power output, also called the mechanical torque applied to the shaft, given in [pu]
- D: damping coefficient, (given in [s], see P. Sauer, eq. (5.156) where the damping torque is equal Dω)
- Ω: rated frequency in [1/s] of the power grid, often 2π⋅50Hz
- T_d_dash: time constant of d-axis, given in [s], see P. Sauer, chapter 3.7, p. 54 for a general explanation on time constants
- X_d_dash: transient reactance of d-axis, given in [pu]
- X_q_dash: transient reactance of q-axis, given in [pu]
- X_d: reactance of d-, given in [pu]
- X_d: reactance of q-axis, given in [pu]

IEEE DC1A Exciter Model
- K_e: Exciter constant related to self-excited field [pu]
- K_f: Excitation control system stabilizer gains [pu]
- K_a: Voltage Regulator gain [pu]
- U:
- U_ref: Reference value of the stator terminal voltage [pu]
- U_ref2: Reference value of the stator terminal voltage [pu]
- U_rmax: Voltage regulator maximum output [pu]
- U_rmin: Voltage regulator minimum output [pu]
- T_a: Time constant of the voltage regulator [s]
- T_f: Excitation control system stabilizer time constant [s]
- T_e: Exciter time constant, integration rate associated with exciter control [s]

"""

using PowerDynamics
using NetworkDynamics
import PowerDynamics: dimension, symbolsof, construct_vertex
import PowerDynamics: AbstractNode
import Base: @__doc__
include("ExciterSaturtionEq.jl")

@DynamicNode ThirdOrderEqExciterIEEED1A(H, P, D, Ω, T_d_dash, X_d_dash, X_q_dash, X_d, K_e, K_f, K_a, U, U_ref, U_ref2, U_rmax, U_rmin, T_a, T_f, T_e, S_E_max, S_E_tq, V_R_max) begin
    @assert H > 0 "inertia (H) should be >0"
    @assert D >= 0 "damping (D) should be >=0"
    @assert T_d_dash > 0 "time constant of d-axis (T_d_dash) should be >0"
    @assert X_d_dash >= 0 "transient reactance of d-axis (X_d_dash) should be >=0"
    @assert X_d >= 0 "reactance of d-axis (X_d_dash) should be >=0"

    Ω_H = Ω / (2 * H)
end [[θ,dθ],[ω, dω], [E_f, dE_f], [U_f, dU_f], [U_r, dU_r]] begin
    i_c = 1im * i * exp(-1im * θ)
    e_c = 1im * u * exp(-1im * θ)
    p = real(u * conj(i))
    e_d = 0           # simpilfication from K.Schmietendorf paper
    e_q = imag(e_c)
    i_d = real(i_c)
    i_q = imag(i_c)

    # Exciter Model, IEEEDC1A
	E_f_d_max, A_x, B_x = ExciterSaturationEq(K_e, S_E_max, S_E_tq, V_R_max) # Saturation modeling
	U_x = A_x * exp(B_x * E_f)

	dU_r = 1 / T_a * (K_a * (U_ref - U + U_ref2 - U_f) - U_r)
	dU_f = 1 / T_f * (K_f / T_e * (U_r - U_x - K_e * E_f) - U_f)

    if U_r > U_rmax
		println("Ur > Urmax")
		U_r2 = U_rmax

	elseif U_r < U_rmin
		println("Ur < Urmax")
		U_r2 = U_rmin
	else
		U_r2 = U_r
	end

	dE_f = 1 / T_e * ( U_r2 - U_x - K_e * E_f)

    # Third Order Model
	dθ = ω
    dω = (P - D * ω - p - (X_q_dash - X_d_dash) * i_d * i_q) * Ω_H
    de_q = (1 / T_d_dash) * (E_f - e_q + i_d * (X_d - X_d_dash))
    # -> u = e_q * exp(1im * θ)
    du = de_q * exp(1im * θ) + u * 1im * ω

end

export ThirdOrderEqExciterIEEEDoneA

max_stars_repo_name: luap-pik/PowerDynIC

max_issues_repo_name: luap-pik/PowerDynIC

max_forks_repo_name: luap-pik/PowerDynIC


text:
#=
Data structures and methods to efficiently represent and work with realizations of HAL basis functions
=#

import Base: keys, getindex, setindex!, BitMatrix, *, intersect
import StatsBase: sample
import Combinatorics: powerset
using InvertedIndices

BasisIndex = Set{Tuple{Int, Float64, Symbol}}
global const INTERCEPT = BasisIndex([(0,NaN,:right)])

struct BasisVector
    length::Int
    nzind::Vector{UInt32}
end

function *(a::Float64, x::BasisVector)::Vector{Float64}
    result = zeros(Float64, x.length)
    result[x.nzind] .= a
    return result
end
*(x::BasisVector, a::Float64) = a*x

function *(v::Vector{Float64}, x::BasisVector)::Float64
    if length(x.nzind) == 0
        return 0.0
    end
    return sum(v[i] for i in x.nzind)
end
*(x::BasisVector, v::Vector{Float64}) = v*x

hash(x::BasisVector) = hash(x.nzind)

struct FeatureVector
    n::Int # number of observations
    sorted::Vector{Float64} # S
    sorted_to_raw::Vector{UInt32} # an index I that makes S[i] = X[I[i]]
    raw_to_sorted::Vector{UInt32} # an index I: S[I[i]] = X[i]
end

function FeatureVector(X::AbstractVector)::FeatureVector
    n = length(X)
    sorted_to_raw = sortperm(X, rev=true)
    raw_to_sorted = sortperm(sorted_to_raw)
    sorted = X[sorted_to_raw]
    return FeatureVector(n, sorted, sorted_to_raw, raw_to_sorted)
end

function findfirst_gte_sorted(S::Vector{Float64}, x::Float64)
    # S should be sorted in reverse
    for (i,xi) in enumerate(S)
        if xi < x return i-1 end
    end
    return length(S)
end

function findfirst_lte_sorted(S::Vector{Float64}, x::Float64)
    for (i,xi) in Iterators.reverse(enumerate(S))
        if xi > x return i+1 end
    end
    return 1
end

function one_way(X::FeatureVector, x, d)
    if d == :right
        final = findfirst_gte_sorted(X.sorted, x)
        if final == 0 return Vector{UInt32}() end
        return X.sorted_to_raw[1:final]
    else
        start = findfirst_lte_sorted(X.sorted, x)
        if start == X.n+1 return Vector{UInt32}() end
        return X.sorted_to_raw[start:end]
    end
end

function build_basis(X::Vector{FeatureVector}, idx::BasisIndex)
    return BasisVector(
        X[1].n,
        intersect([
            j==0 ? (1:X[1].n) : one_way(X[j], x, d)
            for (j,x,d) in idx
        ]...)
    )
end

function sortperm_subset(X::FeatureVector, I)
    # given an index I of length b relative to the the unsorted vector X, return the
    # subset permutation index Ĩ for X so that X[I[Ĩ]] is sorted,
    # i.e. S[sort(raw_to_sorted[I])] = X[I[Ĩ]]
    return sortperm(X.raw_to_sorted[I])
end

changepoints(x::Vector{Float64}) = [i for i in 1:(length(x)-1) if x[i] ≠ x[i+1]]

function sse_indicator_ols(Y::Vector{Float64}, Y2::Vector{Float64}=Y.^2)
    # [SSE(Y, Ŷₖ) for k∈1:n] where Ŷₖ is the OLS prediction of Y regressed on [1,1,... 1ₖ, 0... 0]
    Y_sums = cumsum(Y)
    Y_sq_sums = cumsum(Y2)
    nnz = 1:length(Y)

    sses = (Y_sq_sums - (Y_sums.^2) ./ nnz) + (Y_sq_sums[end] .- Y_sq_sums)
    return sses, Y_sums
    # return ( (Y_sq_sums - (Y_sums.^2) ./ nnz) + (Y_sq_sums[end] .- Y_sq_sums) ) ./ abs.(Y_sums)
    # return -abs.(Y_sums)
end

function best_split_sorted(
    I, # index within unsorted: X[I] is sorted
    X::Vector{Float64},
    Y::Vector{Float64},
    Y2::Vector{Float64} = Y.^2,
    λ = -Inf,
    off_limits = Set{Float64}(),
)::Tuple{Any, Vector{Float64}, Vector{Float64}, Float64, Float64}
    sses, ρs = sse_indicator_ols(Y, Y2)
    βnz = (abs.(ρs) .> λ)
    sse_order = sortperm(sses)
    changes = changepoints(X)
    for ĩ in (i for i in sse_order if ((i ∈ changes) & βnz[i]))
        x = X[ĩ]
        if (x ∉ off_limits)
            return I[1:ĩ], Y[1:ĩ], Y2[1:ĩ], x, sses[ĩ]
        end
    end
    return I, Y, Y2, NaN, Inf
end

function interact_basis(
    I,
    X::Vector{FeatureVector},
    Y::Vector{Float64},
    Y2::Vector{Float64} = Y.^2;
    features = 1:length(X),
    λ = -Inf,
    off_limits::Set{Tuple{Int, Float64, Symbol}} = Set{Tuple{Int, Float64, Symbol}}(),
    right_only = true,
)
    results = []
    for j in features
        Ĩ = sortperm_subset(X[j], I)
        I_sorted = I[Ĩ]
        vector_args = (
            I_sorted,
            X[j].sorted[X[j].raw_to_sorted[I_sorted]],
            Y[Ĩ],
            Y2[Ĩ],
        )
        right_split = best_split_sorted(
            vector_args..., λ,
            Set(x for (k,x,d) in off_limits if (k==j)&(d==:right)),
        )
        if right_only
            push!(results, [right_split..., :right])
        else
            left_split = best_split_sorted(
                [reverse(v) for v in vector_args]..., λ,
                Set(x for (k,x,d) in off_limits if (k==j)&(d==:left)),
            )
            if right_split[end] ≤ left_split[end] # these are the sses
                push!(results, [right_split..., :right])
            else
                push!(results, [left_split..., :left])
            end
        end
    end
    Is, Ys, Y2s, xs, sses, ds = zip(results...)
    sse, j = findmin(sses)
    return Is[j], Ys[j], Y2s[j], (features[j], xs[j], ds[j]), sses[j]
end

filter_and_pare(Ω, ω) = Set(delete!(A,ω) for A in Ω if ω in A) # ω ∈ A ∈ Ω

function basis_search(
    X::Vector{FeatureVector},
    Y::Vector{Float64},
    λ;
    subsample_n::Int=length(Y),
    feat_n::Int=Int(ceil(sqrt(length(X)))),
    off_limits = Set(INTERCEPT), # iterable of [BasisIndex]es
)::Tuple{BasisIndex, BasisVector}
    #=
    Greedily attempt to find the basis vector that will be most useful to add to the lasso by interacting
    one-way basis vectors, one at a time.

    This "top-down" approach makes sense heuristically because we expect that low-order interactions are
    what's important for most real-world data-generating processes. There is also the fact that realizations
    of higher-order sections must have greater and greater sparsity since these are products of h∈{0,1}ⁿ.
    =#
    n = length(Y)
    I = sample(1:n, subsample_n, replace=false)
    Y = Y[I]
    Y2 = Y.^2

    current_sse = Inf
    basis_index = BasisIndex()
    orig_off_limits = deepcopy(off_limits)

    while true
        features = sample(1:length(X), feat_n, replace=false)
        off_limits_coords = Set(first(b) for b in off_limits if length(b)==1)

        Ĩ, Y, Y2, coord, sse = interact_basis(
            I, X, Y, Y2,
            features = features,
            λ = λ*(subsample_n/n),
            off_limits = BasisIndex(off_limits_coords ∪ basis_index),
            right_only = false,
        )
        if (current_sse ≤ sse) | (current_sse ≈ sse)
            if length(basis_index) == 0
                basis_index = INTERCEPT
            end
            if subsample_n == n
                return basis_index, BasisVector(n, I)
            else
                return basis_index, build_basis(X, basis_index)
            end
        else
            current_sse, I = sse, Ĩ
            push!(basis_index, coord)
            off_limits = filter_and_pare(off_limits, coord)
        end
    end
end

function basis_search_random(X::Vector{FeatureVector})::Tuple{BasisIndex, BasisVector}
    features = []
    n, p = X[1].n, length(X)
    while length(features) == 0
        features = (1:p)[sample([false,true], p)]
    end
    basis_index = Set((j, X[j].sorted[rand(1:n)], rand([:right,:left])) for j in features)
    return basis_index, build_basis(X, basis_index)
end

struct Bases
    dict::Dict{BasisIndex, BasisVector}
    set::Set{BasisVector}
    sum::Dict{BasisIndex, Int}
end

function Bases(
    X::Vector{FeatureVector};
    indices = Set{BasisIndex}([INTERCEPT])
)::Bases
    bases = Bases(
        Dict{BasisIndex, BasisVector}(),
        Set{BasisVector}(),
        Dict{BasisIndex, Int}(),
    )
    for idx in indices
        add_basis!(bases, idx, build_basis(X, idx))
    end
    return bases
end

keys(bases::Bases) = keys(bases.dict)
getindex(bases::Bases, b::BasisIndex)::BasisVector = bases.dict[b]

function *(X::Bases, β::Dict{BasisIndex, Float64})::Vector{Float64}
    total = zeros(Float64, X.sum[INTERCEPT])
    for (b, βi) in β
        # total[X[b].nzind] .+= βi
        for i in X[b].nzind
            total[i] += βi
        end
    end
    return total
end

function add_basis!(
    bases::Bases,
    index::BasisIndex,
    basis::BasisVector,
)::Bases
    s = basis.length
    if s ≠ 0 # don't add a null basis
        push!(bases.set, basis)
        bases.dict[index] = basis
        bases.sum[index] = s
    end
    return bases
end

function delete_basis!(
    bases::Bases,
    index::BasisIndex
)::Bases
    if index in keys(bases.dict)
        delete!(bases.set, bases.dict[index])
        delete!(bases.dict, index)
        delete!(bases.sum, index)
    end
    return bases
end

function filter_bases!(bases::Bases, indices)
    indices = push!(Set(indices), INTERCEPT)
    for b in setdiff(Set(keys(bases)), Set(indices))
        delete_basis!(bases, b)
    end
    return bases
end


text:
export maximal_integral_lattice, is_maximal_integral

################################################################################
#
#  String I/O
#
################################################################################

function Base.show(io::IO, L::HermLat)
  println(io, "Hermitian lattice of rank $(rank(L)) and degree $(degree(L))")
  println(io, "over")
  print(IOContext(io, :compact => true), base_ring(L))
end

################################################################################
#
#  Construction
#
################################################################################

function lattice(V::HermSpace, B::PMat ; check::Bool = true)
  E = base_ring(V)
  if check
    @req rank(matrix(B)) == min(nrows(B), ncols(B)) "B must be of full rank"
  end
  @req nf(base_ring(B)) == E "Incompatible arguments: B must be defined over E"
  @req ncols(B) == dim(V) "Incompatible arguments: the number of columns of B must be equal to the dimension of V"
  L = HermLat{typeof(E), typeof(base_field(E)), typeof(gram_matrix(V)), typeof(B), morphism_type(typeof(E))}()
  L.pmat = B
  L.base_algebra = E
  L.involution = involution(V)
  L.space = V
  return L
end

@doc Markdown.doc"""
    hermitian_lattice(E::NumField, B::PMat ; gram = nothing,
				             check::Bool = true) -> HermLat

Given a pseudo-matrix `B` with entries in a number field `E` of degree 2,
return the hermitian lattice spanned by the pseudo-matrix `B` inside the hermitian
space over `E` with Gram matrix `gram`.

If `gram` is not supplied, the Gram matrix of the ambient space will be the
identity matrix over `E` of size the number of columns of `B`.

By default, `B` is checked to be of full rank. This test can be disabled by setting
`check` to false.
"""
function hermitian_lattice(E::NumField, B::PMat ; gram = nothing, check::Bool = true)
  @req nf(base_ring(B)) == E "Incompatible arguments: B must be defined over E"
  @req degree(E) == 2 "E must be a quadratic extension"
  if gram === nothing
    V = hermitian_space(E, ncols(B))
  else
    @assert gram isa MatElem
    @req is_square(gram) "gram must be a square matrix"
    @req ncols(B) == nrows(gram) "Incompatible arguments: the number of columns of B must correspond to the size of gram"
    gram = map_entries(E, gram)
    V = hermitian_space(E, gram)
  end
  return lattice(V, B, check = check)
end

@doc Markdown.doc"""
    hermitian_lattice(E::NumField, basis::MatElem ; gram = nothing,
				                    check::Bool = true) -> HermLat

Given a matrix `basis` and a number field `E` of degree 2, return the hermitian lattice
spanned by the rows of `basis` inside the hermitian space over `E` with Gram matrix `gram`.

If `gram` is not supplied, the Gram matrix of the ambient space will be the identity
matrix over `E` of size the number of columns of `basis`.

By default, `basis` is checked to be of full rank. This test can be disabled by setting
`check` to false.
"""
hermitian_lattice(E::NumField, basis::MatElem ; gram = nothing, check::Bool = true) = hermitian_lattice(E, pseudo_matrix(basis), gram = gram, check = check)

@doc Markdown.doc"""
    hermitian_lattice(E::NumField, gens::Vector ; gram = nothing) -> HermLat

Given a list of vectors `gens` and a number field `E` of degree 2, return the hermitian
lattice spanned by the elements of `gens` inside the hermitian space over `E` with
Gram matrix `gram`.

If `gram` is not supplied, the Gram matrix of the ambient space will be the identity
matrix over `E` of size the length of the elements of `gens`.

If `gens` is empty, `gram` must be supplied and the function returns the zero lattice
in the hermitan space over `E` with Gram matrix `gram`.
"""
function hermitian_lattice(E::NumField, gens::Vector ; gram = nothing)
  if length(gens) == 0
    @assert gram !== nothing
    pm = pseudo_matrix(matrix(E, 0, nrows(gram), []))
    L = hermitian_lattice(E, pm, gram = gram, check = false)
    return L
  end
  @assert length(gens[1]) > 0
  @req all(v -> length(v) == length(gens[1]), gens) "All vectors in gens must be of the same length"

  if gram === nothing
    V = hermitian_space(E, length(gens[1]))
  else
    @assert gram isa MatElem
    @req is_square(gram) "gram must be a square matrix"
    @req length(gens[1]) == nrows(gram) "Incompatible arguments: the length of the elements of gens must correspond to the size of gram"
    gram = map_entries(E, gram)
    V = hermitian_space(E, gram)
  end
  return lattice(V, gens)
end

@doc Markdown.doc"""
    hermitian_lattice(E::NumField ; gram::MatElem) -> HermLat

Given a matrix `gram` and a number field `E` of degree 2, return the free hermitian
lattice inside the hermitian space over `E` with Gram matrix `gram`.
"""
function hermitian_lattice(E::NumField ; gram::MatElem)
  @req is_square(gram) "gram must be a square matrix"
  gram = map_entries(E, gram)
  B = pseudo_matrix(identity_matrix(E, ncols(gram)))
  return hermitian_lattice(E, B, gram = gram, check = false)
end

################################################################################
#
#  Absolute pseudo matrix
#
################################################################################

# Given a lattice for the quadratic extension E/K, return the pseudo matrix
# over Eabs, where Eabs/Q is an absolute field isomorphic to E
function absolute_pseudo_matrix(E::HermLat{S, T, U, V, W}) where {S, T, U, V, W}
  c = get_attribute(E, :absolute_pseudo_matrix)
  if c === nothing
    f = absolute_simple_field(ambient_space(E))[2]
    pm = _translate_pseudo_hnf(pseudo_matrix(E), pseudo_inv(f))
    set_attribute!(E, :absolute_pseudo_matrix => pm)
    return pm::PMat{elem_type(T), fractional_ideal_type(order_type(T))}
  else
    return c::PMat{elem_type(T), fractional_ideal_type(order_type(T))}
  end
end

################################################################################
#
#  Rational span
#
################################################################################

function rational_span(L::HermLat)
  if isdefined(L, :rational_span)
    return L.rational_span
  else
    G = gram_matrix_of_rational_span(L)
    V = hermitian_space(base_field(L), G)
    L.rational_span = V
    return V
  end
end

################################################################################
#
#  Involution
#
################################################################################

involution(L::HermLat) = L.involution

################################################################################
#
#  Hasse invariant
#
################################################################################

function hasse_invariant(L::HermLat, p)
  error("The lattice must be quadratic")
end

################################################################################
#
#  Witt invariant
#
################################################################################

function witt_invariant(L::HermLat, p)
  error("The lattice must be quadratic")
end

################################################################################
#
#  New lattice in same ambient space
#
################################################################################

# Used internally to return a lattice in the same ambient space of L with
# parameter m (that could be a basis matrix, a pseudo matrix, a list of gens
function lattice_in_same_ambient_space(L::HermLat, m::T) where T
  return lattice(ambient_space(L), m)
end

################################################################################
#
#  Norm
#
################################################################################

# TODO: Be careful with the +, ideal_trace gives fmpq and then this must be gcd
function norm(L::HermLat)
  if isdefined(L, :norm)
    return L.norm::fractional_ideal_type(base_ring(base_ring(L)))
  end
  G = gram_matrix_of_rational_span(L)
  v = involution(L)
  K = base_ring(G)
  R = base_ring(L)
  C = coefficient_ideals(L)
  to_sum = sum(G[i, i] * C[i] * v(C[i]) for i in 1:length(C))
  to_sum = to_sum + R * reduce(+, [tr(C[i] * G[i, j] * v(C[j])) for j in 1:length(C) for i in 1:(j-1)])
  n = minimum(numerator(to_sum))//denominator(to_sum)
  L.norm = n
  return n
end

################################################################################
#
#  Scale
#
################################################################################

function scale(L::HermLat)
  if isdefined(L, :scale)
    return L.scale::fractional_ideal_type(base_ring(L))
  end
  G = gram_matrix_of_rational_span(L)
  C = coefficient_ideals(L)
  to_sum = [ G[i, j] * C[i] * involution(L)(C[j]) for j in 1:length(C) for i in 1:j ]
  d = length(to_sum)
  for i in 1:d
    push!(to_sum, involution(L)(to_sum[i]))
  end
  s = sum(to_sum)
  L.scale = s
  return s
end

################################################################################
#
#  Rescale
#
################################################################################

function rescale(L::HermLat, a)
  if isone(a)
    return L
  end
  K = fixed_field(L)
  b = base_field(L)(K(a))
  gramamb = gram_matrix(ambient_space(L))
  return hermitian_lattice(base_field(L), pseudo_matrix(L),
                           gram = b * gramamb)
end

################################################################################
#
#  Bad primes
#
################################################################################

@doc Markdown.doc"""
    bad_primes(L::HermLat; discriminant = false) -> Vector{NfOrdIdl}

Given a hermitian lattice `L` over $E/K$, return the prime ideals of $\mathcal O_K$
dividing the scale or the volume of `L`. If `discriminant == true`, also the prime
ideals dividing the discriminant of $\mathcal O_E$ are returned.
"""
function bad_primes(L::HermLat; discriminant::Bool = false)
  s = scale(L)
  f = factor(norm(scale(L)))
  ff = factor(norm(volume(L)))
  for (p, e) in ff
    f[p] = 0
  end
  if discriminant
    for (p, ) in factor(Hecke.discriminant(base_ring(L)))
      f[p] = 0
    end
  end
  return collect(keys(f))
end

################################################################################
#
#  Dual lattice
#
################################################################################

function dual(L::HermLat)
  G = gram_matrix_of_rational_span(L)
  @req rank(G) == nrows(G) "Lattice must be non-degenerate"
  B = matrix(pseudo_matrix(L))
  C = coefficient_ideals(L)
  Gi = inv(G)
  new_bmat = Gi * B
  new_coeff =  eltype(C)[inv(involution(L)(c)) for c in C]
  pm = pseudo_matrix(new_bmat, new_coeff)
  return lattice(ambient_space(L), pm)
end

################################################################################
#
#  Jordan decomposition
#
################################################################################

function jordan_decomposition(L::HermLat, p)
  R = base_ring(L)
  E = nf(R)
  aut = involution(L)
  even = is_dyadic(p)

  S = local_basis_matrix(L, p)

  D = prime_decomposition(R, p)
  split = length(D) == 2
  ram = D[1][2] == 2
  n = rank(L)

  P = D[1][1]

  if split
    # I need a p-uniformizer
    pi = elem_in_nf(p_uniformizer(P))
    @assert valuation(pi, D[2][1]) == 0
  elseif ram
    pi = elem_in_nf(uniformizer(P))
  else
    pi = base_field(L)(elem_in_nf(uniformizer(p)))
    su = even ? _special_unit(P, p) : one(nf(order(P)))
  end

  oldval = inf
  blocks = Int[]
  exponents = Int[]

  F = gram_matrix(ambient_space(L))

  k = 1
  while k <= n
    G = S * F * transpose(_map(S, aut))
    X = Union{Int, PosInf}[ iszero(G[i, i]) ? inf : valuation(G[i, i], P) for i in k:n]
    m, ii = findmin(X)
    ii = ii + (k - 1)
    pair = (ii, ii)
    for i in k:n
      for j in k:n
        tmp = iszero(G[i, j]) ? inf : valuation(G[i, j], P)
        if tmp < m
          m = tmp
          pair = (i, j)
        end
      end
    end
    if m != oldval
      push!(blocks, k)
      oldval = m
      push!(exponents, m)
    end
    i, j = pair[1], pair[2]
    if (i != j) && !(ram && (even || isodd(m)))
      a = G[i, j]
      if split
        lambda = valuation(pi * a, P) == m ? pi : aut(pi)
      elseif ram
        @assert iseven(m)
        lambda = E(norm(pi)^(div(m ,2)))//a
      else
        lambda = pi^m//a * su
      end
      for l in 1:ncols(S)
        S[i, l] = S[i, l] + aut(lambda) * S[j, l]
      end
      G = S * F * transpose(_map(S, aut))
      @assert valuation(G[i, i], P) == m
      j = i
    end

    if i != j
      @assert i < j
      swap_rows!(S, i, k)
      swap_rows!(S, j, k + 1)
      SF = S * F
      X1 = SF * transpose(_map(view(S, k:k, 1:ncols(S)), aut))
      X2 = SF * transpose(_map(view(S, (k + 1):(k + 1), 1:ncols(S)), aut))
      for l in k+2:n
        den = norm(X2[k, 1]) - X1[k, 1] * X2[k + 1, 1]
        t1 = (X2[l, 1] * X1[k + 1, 1] - X1[l, 1] * X2[k + 1, 1])//den
        t2 = (X1[l, 1] * X2[k, 1] - X2[l, 1] * X1[k, 1])//den

        for o in 1:ncols(S)
          S[l, o] = S[l, o] - (t1 * S[k, o] + t2 * S[k + 1, o])
        end
      end
      k = k + 2
    else
      swap_rows!(S, i, k)
      X1 = S * F * transpose(_map(view(S, k:k, 1:ncols(S)), aut))
      for l in (k + 1):n
        for o in 1:ncols(S)
          S[l, o] = S[l, o] - X1[l, 1]//X1[k, 1] * S[k, o]
        end
      end
      k = k + 1
    end
  end

  if !ram
    G = S * F * transpose(_map(S, aut))
    @assert is_diagonal(G)
  end

  push!(blocks, n + 1)

  matrices = typeof(F)[ sub(S, blocks[i]:(blocks[i+1] - 1), 1:ncols(S)) for i in 1:(length(blocks) - 1)]
  return matrices, typeof(F)[ m * F * transpose(_map(m, aut)) for m in matrices], exponents
end

################################################################################
#
#  Local isometry
#
################################################################################

# base case for order(p) == base_ring(base_ring(L1))
function is_locally_isometric(L1::HermLat, L2::HermLat, p)
  # Test first rational equivalence
  return genus(L1, p) == genus(L2, p)
end

################################################################################
#
#  Maximal integral lattices
#
################################################################################

# Checks whether L is p-maximal integral. If not, a minimal integral
# over-lattice at p is returned
function _ismaximal_integral(L::HermLat, p)
  R = base_ring(L)
  E = nf(R)
  D = prime_decomposition(R, p)
  Q = D[1][1]
  e = valuation(discriminant(R), p)
  if e == 0
    s = one(E)
  else
    s = elem_in_nf(p_uniformizer(D[1][1]))^e
  end
  @assert valuation(s, D[1][1]) == valuation(discriminant(R), p)

  absolute_map = absolute_simple_field(ambient_space(L))[2]

  M = local_basis_matrix(L, p, type = :submodule)
  G = gram_matrix(ambient_space(L), M)
  F, h = ResidueField(R, D[1][1])
  hext = extend(h, E)
  sGmodp = map_entries(hext, s * G)
  Vnullity, V = kernel(sGmodp, side = :left)
  if Vnullity == 0
    return true, zero_matrix(E, 0, 0)
  end

  hprim = u -> elem_in_nf((h\u))

  @vprint :GenRep 1 """Enumerating projective points over $F
                       of length $(nrows(V))\n"""

  for x in enumerate_lines(F, nrows(V))
    v = map_entries(hprim, matrix(F, 1, nrows(V), x) * V)::dense_matrix_type(E)
    res = v * M
    resvec = elem_type(E)[res[1, i] for i in 1:ncols(res)]
    t = inner_product(ambient_space(L), resvec, resvec)
    valv = iszero(t) ? inf : valuation(t, Q)
    if valv >= 2
      # I don't want to compute the generators
      genL = generators(L)
      X = Vector{Union{Int, PosInf}}(undef, length(genL))
      for i in 1:length(genL)
        g = genL[i]
        ip = inner_product(ambient_space(L), resvec, g)
        if iszero(ip)
          X[i] = inf
        else
          X[i] = valuation(ip, Q)
        end
      end
      @assert minimum(X) >= 1 - e
      return false, v * M
    end
  end
  return true, zero_matrix(E, 0, 0)
end

# Check if L is maximal integral at p. If not, return either:
# - a minimal integral overlattice at p (minimal = true)
# - a maximal integral overlattice at p (minimal = false).
function _maximal_integral_lattice(L::HermLat, p, minimal = true)
  R = base_ring(L)
  # already maximal?
  if valuation(norm(volume(L)), p) == 0 && !is_ramified(R, p)
    return true, L
  end

  absolute_map = absolute_simple_field(ambient_space(L))[2]

  B, G, S = jordan_decomposition(L, p)
  D = prime_decomposition(R, p)
  P = D[1][1]
  is_max = true
  lS = length(S)

  invP = inv(P)

  if length(D) == 2 # split
    @assert S[end] != 0
    if minimal
      max = 1
      M = pseudo_matrix(B[lS][1, :], fractional_ideal_type(R)[invP])
    else
      max = S[end]
      coeff_ideals = fractional_ideal_type(R)[]
      _matrix = zero_matrix(nf(R), 0, ncols(B[1]))
      for i in 1:length(B)
        if S[i] == 0
          continue
        end
        _matrix = vcat(_matrix, B[i])
        for k in 1:nrows(B[i])
          push!(coeff_ideals, invP^(S[i]))
        end
      end
      M = pseudo_matrix(_matrix, coeff_ideals)
    end
    _new_pmat = _sum_modules_with_map(pseudo_matrix(L), M, absolute_map)
    LLL = invP^(max) * pseudo_matrix(L)
    _new_pmat = _intersect_modules_with_map(_new_pmat, LLL, absolute_map)
    return false, lattice(ambient_space(L), _new_pmat)
  elseif D[1][2] == 1 # The inert case
    if S[end] >= 2
      if minimal
        max = 1
        M = pseudo_matrix(B[lS][1, :], [invP^(div(S[end], 2))])
      else
        max = S[end]
        coeff_ideals = fractional_ideal_type(R)[]
        _matrix = zero_matrix(nf(R), 0, ncols(B[1]))
        for i in 1:length(B)
          if !(S[i] >= 2)
            continue
          end
          _matrix = vcat(_matrix, B[i])
          for k in 1:nrows(B[i])
            push!(coeff_ideals, invP^(div(S[i], 2)))
          end
        end
        M = pseudo_matrix(_matrix, coeff_ideals)
      end
      _new_pmat = _sum_modules_with_map(pseudo_matrix(L), M, absolute_map)
      LLL = invP^(max) * pseudo_matrix(L)
      _new_pmat = _intersect_modules_with_map(_new_pmat, LLL, absolute_map)
      L = lattice(ambient_space(L), _new_pmat)
      if minimal
        return false, L
      end
      B, G, S = jordan_decomposition(L, p)
      is_max = false
    end
    # new we look for zeros of ax^2 + by^2
    kk, h = ResidueField(R, P)
    while sum(Int[S[i] * nrows(B[i]) for i in 1:length(B)]) > 1
      k = 0
      for i in 1:(length(S) + 1)
        if S[i] == 1
          k = i
          break
        end
      end
      @assert nrows(B[k]) >= 2
      r = h\rand(kk)
      # The following might throw ...
      while valuation(G[k][1, 1] + G[k][2, 2] * elem_in_nf(norm(r)), P) < 2
        r = h\rand(kk)
      end
      M = pseudo_matrix(B[k][1, :] + elem_in_nf(r) * B[k][2, :], [invP])
      _new_pmat = _sum_modules_with_map(pseudo_matrix(L), M, absolute_map)
      LLL = invP * pseudo_matrix(L)
      _new_pmat = _intersect_modules_with_map(_new_pmat, LLL, absolute_map)
      L = lattice(ambient_space(L), _new_pmat)
      if minimal
        return false, L
      end
      is_max = false
      B, G, S = jordan_decomposition(L, p)
      @assert S[1] >= 0
    end
    return is_max, L
  else # ramified case
    if S[end] >= 2
      if minimal
        max = 1
        M = pseudo_matrix(B[lS][1, :], [invP^(div(S[end], 2))])
      else
        max = S[end]
        coeff_ideals = fractional_ideal_type(R)[]
        _matrix = zero_matrix(nf(R), 0, ncols(B[1]))
        for i in 1:length(B)
          if !(S[i] >= 2)
            continue
          end
          _matrix = vcat(_matrix, B[i])
          for k in 1:nrows(B[i])
            push!(coeff_ideals, invP^(div(S[i], 2)))
          end
        end
        M = pseudo_matrix(_matrix, coeff_ideals)
      end
      _new_pmat = _sum_modules_with_map(pseudo_matrix(L), M, absolute_map)
      LLL = invP^(max) * pseudo_matrix(L)
      _new_pmat = _intersect_modules_with_map(_new_pmat, LLL, absolute_map)
      L = lattice(ambient_space(L), _new_pmat)
      if minimal
        return false, L
      end
      B, G, S = jordan_decomposition(L, p)
    end
    v = valuation(volume(L), P)
    ok, x = _ismaximal_integral(L, p)
    while !ok
      LL = L
      LLL = pseudo_matrix(x, fractional_ideal_type(R)[invP])
      LpLLL = _sum_modules_with_map(pseudo_matrix(L), LLL, absolute_map)
      L = lattice(ambient_space(L), LpLLL)
      v = v - 2
      @assert v == valuation(volume(L), P)
      @assert valuation(norm(L), p) >= 0
      if minimal
        return false, L
      end
      is_max = false
      ok, x = _ismaximal_integral(L, p)
    end
    @assert iseven(v)
    v = div(v, 2)
    m = rank(L)
    e = valuation(discriminant(R), p)
    if isodd(m)
      valmax = - div(m - 1, 2) * e
    else
      valmax = -div(m, 2) * e
      disc = discriminant(ambient_space(L))
      if !is_local_norm(nf(R), disc, p)
        valmax += 1
      end
    end
    @assert v == valmax
    return is_max, L
  end
end

function is_maximal_integral(L::HermLat, p)
  valuation(norm(L), p) < 0 && return false, L
  return _maximal_integral_lattice(L, p, true)
end

function is_maximal_integral(L::HermLat)
  !is_integral(norm(L)) && throw(error("The lattice is not integral"))
  S = base_ring(L)
  f = factor(discriminant(S))
  ff = factor(norm(volume(L)))
  for (p, e) in ff
    f[p] = 0
  end
  bad = collect(keys(f))
  for p in bad
    ok, LL = _maximal_integral_lattice(L, p, true)
    if !ok
      return false, LL
    end
  end
  return true, L
end

function is_maximal(L::HermLat, p)
  #iszero(L) && error("The lattice must be non-zero")
  v = valuation(norm(L), p)
  x = elem_in_nf(p_uniformizer(p))^(-v)
  b, LL = is_maximal_integral(rescale(L, x), p)
  if b
    return b, L
  else
    return false, lattice(ambient_space(L), pseudo_matrix(LL))
  end
end

function maximal_integral_lattice(L::HermLat)
  !is_integral(norm(L)) && throw(error("The lattice is not integral"))
  S = base_ring(L)
  f = factor(discriminant(S))
  ff = factor(norm(volume(L)))
  for (p, e) in ff
    f[p] = 0
  end
  bad = collect(keys(f))
  for p in bad
    _, L = _maximal_integral_lattice(L, p, false)
  end
  return L
end

function maximal_integral_lattice(L::HermLat, p)
  valuation(norm(L), p) < 0 && throw(error("Lattice is not locally integral"))
  _, L = _maximal_integral_lattice(L, p, false)
  return L
end

function maximal_integral_lattice(V::HermSpace)
  L = lattice(V, identity_matrix(base_ring(V), rank(V)))
  fac = collect(factor(scale(L)))
  S = base_ring(L)
  s = one(nf(S)) * S
  while length(fac) > 0
    nrm = norm(fac[1][1])
    i = findfirst(i -> norm(fac[i][1]) == nrm, 2:length(fac))
    if i !== nothing
      i = i + 1 # findfirst gives the index and not the element
      @assert fac[i][2] == fac[1][2]
      s = s * fractional_ideal(S, fac[1][1])^fac[1][2]
      deleteat!(fac, i)
    else
      s = s * inv(fac[1][1])^(div(fac[1][2], 2))
    end
    deleteat!(fac, 1)
  end
  if !isone(s)
    L = s * L
  end
  return maximal_integral_lattice(L)
end


text:
using Statistics
using NCDatasets

using ArgParse
using JSON
using Formatting

include("LinearRegression.jl")
include("CESMReader.jl")

using .CESMReader

correlation = (x1, x2) -> x1' * x2 / (sum(x1.^2)*sum(x2.^2)).^0.5

function parse_commandline()

    s = ArgParseSettings()
    @add_arg_table s begin

        "--data-file-prefix"
            help = "Data filename prefix including folder and path until the timestamp. File extension `nc` is assumed."
            arg_type = String
            required = true

        "--data-file-timestamp-form"
            help = "Data filename timestamp form. Either `YEAR` or `YEAR_MONTH`."
            arg_type = String
            required = true

        "--domain-file"
            help = "Domain file."
            arg_type = String
            required = true

        "--output-file"
            help = "Output atm temperature file."
            arg_type = String
            required = true

        "--beg-year"
            help = "Year of begin."
            arg_type = Int64
            required = true

        "--end-year"
            help = "Year of end."
            arg_type = Int64
            required = true

    end

    return parse_args(ARGS, s)
end

parsed = parse_commandline()
print(json(parsed, 4))

output_file = parsed["output-file"]

let

    if parsed["data-file-timestamp-form"] == "YEAR"
        filename_format = format("{:s}{{:04d}}.nc", joinpath(parsed["data-file-prefix"]))
        form = :YEAR
    elseif parsed["data-file-timestamp-form"] == "YEAR_MONTH"
        filename_format = format("{:s}{{:04d}}-{{:02d}}.nc", joinpath(parsed["data-file-prefix"]))
        form = :YEAR_MONTH
    end

    fh = FileHandler(filename_format=filename_format, form=form)

    global TREFHT = getData(fh, "TREFHT", (parsed["beg-year"], parsed["end-year"]), (:, :))
    global (Nx, Ny, Nt) = size(TREFHT)
    global nyears = Int64(Nt / 12)

    if mod(Nt, 12) != 0
        ErrorException("Time record is not multiple of 12") |> throw
    end

end

Dataset(parsed["domain-file"], "r") do ds
    global mask = replace(ds["mask"], missing=>NaN)
    global lat  = replace(ds["yc"], missing=>NaN)

    global lnd_mask = mask .== 1
    global ocn_mask = mask .!= 1

    global w_GLB = cos.(lat * π/180.0)
    global w_LND = w_GLB .* lnd_mask
    global w_OCN = w_GLB .* ocn_mask

    global sum_w_GLB = sum(w_GLB)
    global sum_w_LND = sum(w_LND)
    global sum_w_OCN = sum(w_OCN)
end

TREFHTA    = zeros(Float64, Nx, Ny, Nt)
TREFHTAVAR = zeros(Float64, Nx, Ny, 12)
TREFHTMM   = zeros(Float64, Nx, Ny, 12)
x = collect(Float64, 1:Nt)
for i=1:Nx, j=1:Ny

    d = view(TREFHT, i, j, :)

    TREFHTMM[i, j, :] = mean( reshape(d, 12, :), dims=2 )[:, 1]
    TREFHTA[i, j, :]  = d - repeat( TREFHTMM[i, j, :], outer=nyears)

    for m = 1:12
        d_yy = view(TREFHTA, i, j, m:12:(m+nyears*12-1))
        TREFHTAVAR[i, j, m] = std(d_yy)
    end

end

TREFHT_GLB = zeros(Float64, Nt)
TREFHT_LND = zeros(Float64, Nt)
TREFHT_OCN = zeros(Float64, Nt)
for t=1:Nt

    T = view(TREFHT, :, :, t)

    TREFHT_GLB[t] = sum(w_GLB .*  T) / sum_w_GLB
    TREFHT_LND[t] = sum(w_LND .*  T) / sum_w_LND
    TREFHT_OCN[t] = sum(w_OCN .*  T) / sum_w_OCN

end

# Meridional distribution

TREFHT_ZM     = zeros(Float64, Ny)
TREFHT_ZVAR   = zeros(Float64, Ny)

TREFHT_ZMM    = zeros(Float64, Ny, 12)
TREFHT_ZMVAR  = zeros(Float64, Ny, 12)

for j=1:Ny

    d = mean(view(TREFHT, :, j, :), dims=1)[1, :]

    TREFHT_ZM[j]   = mean(d)
    TREFHT_ZVAR[j] = std(d)

    d_wrap = reshape(d, 12, :)
    TREFHT_ZMM[j, :]   = mean( d_wrap, dims=2 )[:, 1]
    TREFHT_ZMVAR[j, :] = std(  d_wrap, dims=2 )[:, 1]

end

Dataset(output_file, "c") do ds

    defDim(ds, "months", 12)
    defDim(ds, "time", Inf)
    defDim(ds, "Nx", Nx)
    defDim(ds, "Ny", Ny)

    ds.attrib["TREFHT_GLB_mean"] = Float64(mean(TREFHT_GLB))
    ds.attrib["TREFHT_OCN_mean"] = Float64(mean(TREFHT_OCN))
    ds.attrib["TREFHT_LND_mean"] = Float64(mean(TREFHT_LND))

    for (varname, vardata, vardim, attrib) in [
        ("TREFHTA",    TREFHTA,    ("Nx", "Ny", "time"), Dict()),
        ("TREFHTMM",   TREFHTMM,   ("Nx", "Ny", "months"), Dict()),
        ("TREFHTAVAR", TREFHTAVAR, ("Nx", "Ny", "months"), Dict()),
        ("TREFHT_GLB", TREFHT_GLB,   ("time",), Dict()),
        ("TREFHT_LND", TREFHT_LND,   ("time",), Dict()),
        ("TREFHT_OCN", TREFHT_OCN,   ("time",), Dict()),
        ("TREFHT_ZM",    TREFHT_ZM,    ("Ny",), Dict()),
        ("TREFHT_ZVAR",  TREFHT_ZVAR,  ("Ny",), Dict()),
        ("TREFHT_ZMM",   TREFHT_ZMM,   ("Ny", "months"), Dict()),
        ("TREFHT_ZMVAR", TREFHT_ZMVAR, ("Ny", "months"), Dict()),
    ]

        println("Doing var: ", varname)

        var = defVar(ds, varname, Float64, vardim)
        var.attrib["_FillValue"] = 1e20

        var = ds[varname]

        for (k, v) in attrib
            var.attrib[k] = v
        end

        rng = []
        for i in 1:length(vardim)-1
            push!(rng, Colon())
        end
        push!(rng, 1:size(vardata)[end])
        var[rng...] = vardata

    end

end

max_stars_repo_name: meteorologytoday/CESM-diagnostic

max_issues_repo_name: meteorologytoday/CESM-diagnostic

max_forks_repo_name: meteorologytoday/CESM-diagnostic


text:
# Generate Gauss quadrature rules on cubes (0,0,...) - (1,1,...) by doing an outer product
# over all dimensions

import Base.Cartesian: @nloops, @nref, @ntuple, @nexprs

for dim in (1,2,3)
    @eval begin
        function (::Type{QuadratureRule{HyperCube{$dim}}})(quad_type::GaussLegendre, order::Int)
            p,w = FastGaussQuadrature.gausslegendre(order)
            weights = Vector{Float64}(undef, order^($dim))
            points = Vector{Vec{$dim,Float64}}(undef, order^($dim))
            count = 1
            @nloops $dim i j->(1:order) begin
                t = @ntuple $dim q-> 0.5*(p[$(Symbol("i"*"_q"))] .+ 1.0)
                points[count] = Vec{$dim,Float64}(t)
                weight = 1.0
                @nexprs $dim j->(weight *= w[i_{j}])
                weights[count] = weight
                count += 1
            end
            return QuadratureRule{HyperCube{$dim},$dim,Float64}(weights./sum(weights), points)
        end
    end
end

max_stars_repo_name: Paulms/jFEMToools

max_issues_repo_name: Paulms/jFEMToools

max_forks_repo_name: Paulms/jFEMToools


text:
let
    method = ME_MAX_POL
    niders = length(FLX_IDERS)

    stds = map(EXPS) do exp
        DAT[method, :eperr, :flx, FLX_IDERS, exp] ./ abs.(DAT[method, :Hd, :flx, "GLC", exp])
    end

    # sGLC
    ids = [:D, :sAC, :sGLC, :sSUCC, :sFORM,  :DCW, :uAC, :uGLC, :uSUCC, :uFORM]
    id_valss = Dict(id => [Hd.val(id, exp) for exp in EXPS] for id in ids)
    push!(ids, :cgD_X)
    id_valss[:cgD_X] = [-Hd.ciD_X(:GLC, exp) for exp in EXPS]

    for id in ids
        p = scatter(;xlabel = string(id), ylabel = "std")
        id_vals = id_valss[id]
        sidx = sortperm(id_vals)
        for idx in sidx
            x = id_vals[idx]
            ys = stds[idx]
            scatter!(p, fill(x, niders), ys; label = "", color = :black, m = 8)
            scatter!(p, [x], [mean(ys)]; label = "", color = :red, m = 12)
        end
        mysavefig(p, "av_var_vs_"; id)
    end
end


text:
for N in [Float64, Rational{Int}, Float32]
    # π/2 trigonometric rotation
    b = BallInf(N[1, 2], N(1))
    M = N[0 -1 ; 1 0]
    # Test Construction
    lm1 = LinearMap(M, b)
    @test lm1.M == M
    @test lm1.X == b
    # Test Dimension
    @test dim(lm1) == 2
    # Test Support Vector
    d = N[1, 1]
    @test σ(d, lm1) == N[-1, 2]
    d = N[-1, 1]
    @test σ(d, lm1) == N[-3, 2]
    d = N[-1, -1]
    @test σ(d, lm1) == N[-3, 0]
    d = N[1, -1]
    @test σ(d, lm1) == N[-1, 0]

    # 2D -> 1D Projection
    b = BallInf(N[1, 2], N(1))
    M = N[1 0]
    lm = M*b
    # Test Dimension
    @test dim(lm) == 1
    # Test Support Vector
    d = N[1]
    @test σ(d, lm) == N[2]
    d = N[-1]
    @test σ(d, lm) == N[0]

    # scalar multiplication
    b = BallInf(N[0, 0], N(1))
    lm = N(2) * b
    @test lm == b * N(2)
    # repeated scalar multiplication
    lm2 = N(2) * lm
    # Test Dimension
    @test dim(lm) == 2
    # Test Support Vector
    d = N[1, 1]
    @test σ(d, lm) == N[2, 2]
    d = N[-1, 1]
    @test σ(d, lm) == N[-2, 2]
    d = N[-1, -1]
    @test σ(d, lm) == N[-2, -2]
    d = N[1, -1]
    @test σ(d, lm) == N[2, -2]

    # Nested construction
    lm1_copy = LinearMap(Matrix{N}(I, 2, 2), lm1)
    @test lm1_copy.M == lm1.M
    @test lm1_copy.X == lm1.X

    # boundedness
    @test isbounded(ones(N, 2, 2) * Singleton(N[1, 2]))  # bounded set
    @test isbounded(zeros(N, 2, 2) * HalfSpace(N[1, 1], N(1)))  # zero map
    @test !isbounded(N[2 3; 1 2] * HalfSpace(N[1, 1], N(1)))  # invertible matrix
    @test !isbounded(N[2 3; 0 0] * HalfSpace(N[1, 1], N(1)))  # singular matrix (expensive check)

    # isempty
    @test !isempty(lm)

    # an_element function
    lm = N(2) * BallInf(N[0, 0], N(1))
    an_element(lm)
    @test an_element(lm) ∈ lm

    # check linear map between vector and set
    X = BallInf(N[1], N(1.10445))
    a = N[-1, 2]
    @test a * X isa LinearMap{N,BallInf{N,Array{N,1}},N,Array{N,2}}

    # absorbing elements
    for neutral in [ZeroSet{N}(2), EmptySet{N}(2)]
        @test N[0 -1 ; 1 0] * neutral == neutral
    end

    # vertices_list
    b = BallInf(N[0, 0], N(1))
    M = N[1 2; 3 4]
    vlist = vertices_list(LinearMap(M, b))
    @test ispermutation(vlist, [N[3, 7], N[1, 1], N[-1, -1], N[-3, -7]])
    M = zeros(N, 2, 2)
    vlist = vertices_list(LinearMap(M, b))
    @test vlist == [zeros(N, 2)]

    # test that redundant vertices are removed by default (see #1355)
    X = Interval(N(0), N(1)) × ZeroSet{N}(1)
    M = N[0 1; 0 2]
    @test vertices_list(M * X) == [N[0, 0]]

    # concrete linear map of a LinearMap
    b = BallInf(N[0, 0], N(1))
    M = N[2 3; 1 2]
    L = LinearMap(M, b)
    V = linear_map(M, LinearMap(M, b))

    # concretize
    @test concretize(L) == linear_map(M, b)
end

# tests that only work with Float64
for N in [Float64]
    b = BallInf(N[0, 0], N(1))

    if test_suite_polyhedra
        # concrete intersection with lazy linear map
        M = N[2 3; 1 2]
        L = M * b
        Lb = intersection(L, b)
        @test M * an_element(b) ∈ Lb

        # constraints_list
        b = BallInf(N[0, 0], N(1))
        M = N[2 3; 1 2]  # invertible
        lm1 = LinearMap(M, b)
        clist = constraints_list(lm1)
        p1 = HPolygon(clist)
        M = N[2 3; 0 0]  # not invertible
        lm2 = LinearMap(M, b)
        clist = constraints_list(lm2)
        p2 = HPolygon(clist)
        for d in BoxDiagDirections{N}(2)
            @test ρ(d, lm1) ≈ ρ(d, p1)
            @test ρ(d, lm2) ≈ ρ(d, p2)
        end
    end

    # concrete linear map of a LinearMap and membership
    b = BallInf(N[0, 0], N(1))
    M = N[2 3; 1 2]
    L = LinearMap(M, b)
    V = linear_map(M, LinearMap(M, b))
    @test M * M * N[1, 1] ∈ V
end


text:
using Expect
using Base.Test
using Compat: readline

# like @test, but stops execution of the testset
macro require(expr)
    quote
        value = $(esc(expr))
        if value
            res = Base.Test.Pass(:test, nothing, nothing, nothing)
            Base.Test.record(Base.Test.get_testset(), res)
        else
            throw(Main.Base.AssertionError(string($(Expr(:inert, expr)))))
        end
    end
end


text:
using Test

include("../src/MiniMax.jl")
using .MiniMax: play_bot_turn!, State, MARK_P1, MARK_P2, MARK_TIE, MARK_EMPTY

@testset "first move x to 2,2" begin
    # algorithm should update the initially empty board to
    # _ _ _
    # _ x _
    # _ _ _
    empty_board = fill(Int8(0), 3, 3)
    state::State = play_bot_turn!(empty_board)

    @test state.player == MARK_P1

    empty_board[5] = MARK_P1

    @test state.last_entry == Int8(5)
    @test all(state.board .== empty_board)
end

@testset "make move x to 3,2 to win" begin
    # initial game board
    # o x o
    # _ x _
    # x _ o
    game_board = reshape(
        [
            MARK_P2,
            MARK_EMPTY,
            MARK_P1,
            MARK_P1,
            MARK_P1,
            MARK_EMPTY,
            MARK_P2,
            MARK_EMPTY,
            MARK_P2,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[6] = MARK_P1

    @test state.last_entry == Int8(6)
    @test all(state.board .== game_board)
    @test state.winner == MARK_P1
end

@testset "make move x to 3,2 to win, second" begin
    # initial game board
    # x x o
    # _ x _
    # o _ o
    game_board = reshape(
        [
            MARK_P1,
            MARK_EMPTY,
            MARK_P2,
            MARK_P1,
            MARK_P1,
            MARK_EMPTY,
            MARK_P2,
            MARK_EMPTY,
            MARK_P2,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[6] = MARK_P1

    @test state.last_entry == Int8(6)
    @test all(state.board .== game_board)
    @test state.winner == MARK_P1
end

@testset "make move x to 3,2 to prevent win" begin
    # initial game board
    # x o o
    # o o x
    # x _ _
    game_board = reshape(
        [
            MARK_P1,
            MARK_P2,
            MARK_P1,
            MARK_P2,
            MARK_P2,
            MARK_EMPTY,
            MARK_P2,
            MARK_P1,
            MARK_EMPTY,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[6] = MARK_P1

    @test state.last_entry == Int8(6)
    @test all(state.board .== game_board)
    @test state.winner == MARK_EMPTY
end

@testset "make move x to 3,2 to prevent win, second" begin
    # initial game board
    # x o _
    # _ o x
    # _ _ _
    game_board = reshape(
        [
            MARK_P1,
            MARK_EMPTY,
            MARK_EMPTY,
            MARK_P2,
            MARK_P2,
            MARK_EMPTY,
            MARK_EMPTY,
            MARK_P1,
            MARK_EMPTY,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[6] = MARK_P1

    @test state.last_entry == Int8(6)
    @test all(state.board .== game_board)
    @test state.winner == MARK_EMPTY
end

@testset "make move x to 1,1 to win" begin
    # initial game board
    # _ o x
    # _ x o
    # o _ x
    game_board = reshape(
        [
            MARK_EMPTY,
            MARK_EMPTY,
            MARK_P2,
            MARK_P2,
            MARK_P1,
            MARK_EMPTY,
            MARK_P1,
            MARK_P2,
            MARK_P1,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[1] = MARK_P1

    @test state.last_entry == Int8(1)
    @test all(state.board .== game_board)
    @test state.winner == MARK_P1
end

@testset "make move x to 1,1 to win, second" begin
    # initial game board
    # _ o _
    # _ x _
    # o _ x
    game_board = reshape(
        [
            MARK_EMPTY,
            MARK_EMPTY,
            MARK_P2,
            MARK_P2,
            MARK_P1,
            MARK_EMPTY,
            MARK_EMPTY,
            MARK_EMPTY,
            MARK_P1,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[1] = MARK_P1

    @test state.last_entry == Int8(1)
    @test all(state.board .== game_board)
    @test state.winner == MARK_P1
end

@testset "make move x to 1,1 to prevent win" begin
    # initial game board
    # _ o o
    # o x x
    # _ x _
    game_board = reshape(
        [
            MARK_EMPTY,
            MARK_P2,
            MARK_EMPTY,
            MARK_P2,
            MARK_P1,
            MARK_P1,
            MARK_P2,
            MARK_P1,
            MARK_EMPTY,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[1] = MARK_P1

    @test state.last_entry == Int8(1)
    @test all(state.board .== game_board)
    @test state.winner == MARK_EMPTY
end

@testset "make move x to 1,1 to prevent win, second" begin
    # initial game board
    # _ o o
    # o x x
    # o x _
    game_board = reshape(
        [
            MARK_EMPTY,
            MARK_P2,
            MARK_P2,
            MARK_P2,
            MARK_P1,
            MARK_P1,
            MARK_P2,
            MARK_P1,
            MARK_EMPTY,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[1] = MARK_P1

    @test state.last_entry == Int8(1)
    @test all(state.board .== game_board)
    @test state.winner == MARK_EMPTY
end

@testset "make move x to 2,3 to end the game to draw" begin
    # initial game board
    # x o o
    # o x _
    # x x o
    game_board = reshape(
        [
            MARK_P1,
            MARK_P2,
            MARK_P1,
            MARK_P2,
            MARK_P1,
            MARK_P1,
            MARK_P2,
            MARK_EMPTY,
            MARK_P2,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[8] = MARK_P1

    @test state.last_entry == Int8(8)
    @test all(state.board .== game_board)
    @test state.winner == MARK_TIE
end

@testset "make move x to 2,2 to end the game to draw" begin
    # initial game board
    # x o o
    # o _ x
    # x x o
    game_board = reshape(
        [
            MARK_P1,
            MARK_P2,
            MARK_P1,
            MARK_P2,
            MARK_EMPTY,
            MARK_P1,
            MARK_P2,
            MARK_P1,
            MARK_P2,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[5] = MARK_P1

    @test state.last_entry == Int8(5)
    @test all(state.board .== game_board)
    @test state.winner == MARK_TIE
end

@testset "make move x to 3,3 to win" begin
    # initial game board
    # x o x
    # o x _
    # o _ _
    game_board = reshape(
        [
            MARK_P1,
            MARK_P2,
            MARK_P2,
            MARK_P2,
            MARK_P1,
            MARK_EMPTY,
            MARK_P1,
            MARK_EMPTY,
            MARK_EMPTY,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[9] = MARK_P1

    @test state.last_entry == Int8(9)
    @test all(state.board .== game_board)
    @test state.winner == MARK_P1
end

@testset "make move x to 1,3 to win" begin
    # initial game board
    # o _ _
    # _ o x
    # o _ x
    game_board = reshape(
        [
            MARK_P2,
            MARK_EMPTY,
            MARK_P2,
            MARK_EMPTY,
            MARK_P2,
            MARK_EMPTY,
            MARK_EMPTY,
            MARK_P1,
            MARK_P1,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[7] = MARK_P1

    @test state.last_entry == Int8(7)
    @test all(state.board .== game_board)
    @test state.winner == MARK_P1
end

@testset "make move x to 1,3 to prevent win" begin
    # initial game board
    # o _ _
    # x o o
    # o x x
    game_board = reshape(
        [
            MARK_P2,
            MARK_P1,
            MARK_P2,
            MARK_EMPTY,
            MARK_P2,
            MARK_P1,
            MARK_EMPTY,
            MARK_P2,
            MARK_P1,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[7] = MARK_P1

    @test state.last_entry == Int8(7)
    @test all(state.board .== game_board)
    @test state.winner == MARK_EMPTY
end

@testset "find out that the opponent has won by diagonal line" begin
    # initial game board
    # o x _
    # x o x
    # o x o
    game_board = reshape(
        [
            MARK_P2,
            MARK_P1,
            MARK_P2,
            MARK_P1,
            MARK_P2,
            MARK_P1,
            MARK_EMPTY,
            MARK_P1,
            MARK_P2,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    @test state.last_entry == MARK_EMPTY
    @test all(state.board .== game_board)
    @test state.winner == MARK_P2
end

@testset "find out that the opponent has won by antidiagonal line" begin
    # initial game board
    # x x o
    # x o x
    # o _ o
    game_board = reshape(
        [
            MARK_P1,
            MARK_P1,
            MARK_P2,
            MARK_P1,
            MARK_P2,
            MARK_EMPTY,
            MARK_P2,
            MARK_P1,
            MARK_P2,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    @test state.last_entry == MARK_EMPTY
    @test all(state.board .== game_board)
    @test state.winner == MARK_P2
end

@testset "find out that the game has ended in draw" begin
    # initial game board
    # x o x
    # x o x
    # o x o
    game_board = reshape(
        [MARK_P1, MARK_P1, MARK_P2, MARK_P2, MARK_P2, MARK_P1, MARK_P1, MARK_P1, MARK_P2],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    @test state.last_entry == MARK_EMPTY
    @test all(state.board .== game_board)
    @test state.winner == MARK_TIE
end

@testset "make move x to 2,2 to win" begin
    # initial game board
    # _ x _
    # o _ _
    # _ x o
    game_board = reshape(
        [
            MARK_EMPTY,
            MARK_P2,
            MARK_EMPTY,
            MARK_P1,
            MARK_EMPTY,
            MARK_P1,
            MARK_EMPTY,
            MARK_EMPTY,
            MARK_P2,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[5] = MARK_P1

    @test state.last_entry == Int8(5)
    @test all(state.board .== game_board)
    @test state.winner == MARK_P1
end

@testset "make move x to 2,1 to win" begin
    # initial game board
    # o o _
    # _ x x
    # _ _ o
    game_board = reshape(
        [
            MARK_P2,
            MARK_EMPTY,
            MARK_EMPTY,
            MARK_P2,
            MARK_P1,
            MARK_EMPTY,
            MARK_EMPTY,
            MARK_P1,
            MARK_P2,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[2] = MARK_P1

    @test state.last_entry == Int8(2)
    @test all(state.board .== game_board)
    @test state.winner == MARK_P1
end

@testset "make move x to 2,1 to win, second" begin
    # initial game board
    # o _ _
    # _ x x
    # o _ o
    game_board = reshape(
        [
            MARK_P2,
            MARK_EMPTY,
            MARK_P2,
            MARK_EMPTY,
            MARK_P1,
            MARK_EMPTY,
            MARK_EMPTY,
            MARK_P1,
            MARK_P2,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[2] = MARK_P1

    @test state.last_entry == Int8(2)
    @test all(state.board .== game_board)
    @test state.winner == MARK_P1
end

@testset "make move x to 1,2 to win" begin
    # initial game board
    # _ _ o
    # o x x
    # o x _
    game_board = reshape(
        [
            MARK_EMPTY,
            MARK_P2,
            MARK_P2,
            MARK_EMPTY,
            MARK_P1,
            MARK_P1,
            MARK_P2,
            MARK_P1,
            MARK_EMPTY,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[4] = MARK_P1

    @test state.last_entry == Int8(4)
    @test all(state.board .== game_board)
    @test state.winner == MARK_P1
end

@testset "make move x to 1,2 to win, second" begin
    # initial game board
    # _ _ o
    # o x _
    # x x o
    game_board = reshape(
        [
            MARK_EMPTY,
            MARK_P2,
            MARK_P2,
            MARK_EMPTY,
            MARK_P1,
            MARK_P1,
            MARK_P2,
            MARK_EMPTY,
            MARK_P2,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[4] = MARK_P1

    @test state.last_entry == Int8(4)
    @test all(state.board .== game_board)
    @test state.winner == MARK_P1
end

@testset "make move x to 2,3 to prevent win" begin
    # initial game board
    # x o o
    # o x _
    # x _ o
    game_board = reshape(
        [
            MARK_P1,
            MARK_P2,
            MARK_P1,
            MARK_P2,
            MARK_P1,
            MARK_EMPTY,
            MARK_P2,
            MARK_EMPTY,
            MARK_P2,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[8] = MARK_P1

    @test state.last_entry == Int8(8)
    @test all(state.board .== game_board)
    @test state.winner == MARK_EMPTY
end

@testset "make move x to 2,3 to prevent win, second" begin
    # initial game board
    # x o o
    # o x _
    # _ x o
    game_board = reshape(
        [
            MARK_P1,
            MARK_P2,
            MARK_EMPTY,
            MARK_P2,
            MARK_P1,
            MARK_P1,
            MARK_P2,
            MARK_EMPTY,
            MARK_P2,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[8] = MARK_P1

    @test state.last_entry == Int8(8)
    @test all(state.board .== game_board)
    @test state.winner == MARK_EMPTY
end

@testset "make move x to 2,3 to prevent win, third" begin
    # initial game board
    # x o o
    # _ x _
    # _ x o
    game_board = reshape(
        [
            MARK_P1,
            MARK_EMPTY,
            MARK_EMPTY,
            MARK_P2,
            MARK_P1,
            MARK_P1,
            MARK_P2,
            MARK_EMPTY,
            MARK_P2,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[8] = MARK_P1

    @test state.last_entry == Int8(8)
    @test all(state.board .== game_board)
    @test state.winner == MARK_EMPTY
end

@testset "make move x to 2,3 to prevent win, fourth" begin
    # initial game board
    # _ _ o
    # _ x _
    # x _ o
    game_board = reshape(
        [
            MARK_EMPTY,
            MARK_EMPTY,
            MARK_P1,
            MARK_EMPTY,
            MARK_P1,
            MARK_EMPTY,
            MARK_P2,
            MARK_EMPTY,
            MARK_P2,
        ],
        (3, 3),
    )
    state::State = play_bot_turn!(game_board)

    game_board[8] = MARK_P1

    @test state.last_entry == Int8(8)
    @test all(state.board .== game_board)
    @test state.winner == MARK_EMPTY
end

@testset "play full game between bots, should end in draw" begin
    game_board = fill(Int8(0), 3, 3)
    state::State = play_bot_turn!(game_board)

    for j = 2:9
        game_board = copy(state.board)
        game_board .*= Int8(-1)
        state = play_bot_turn!(game_board)
    end

    @test state.winner == MARK_TIE
    @test sum(==(MARK_EMPTY), state.board) == 0
end

max_stars_repo_name: elmomoilanen/Tic-tac-toe

max_issues_repo_name: elmomoilanen/Tic-tac-toe

max_forks_repo_name: elmomoilanen/Tic-tac-toe