### A Pluto.jl notebook ###
# v0.15.0

using Markdown
using InteractiveUtils

# ╔═╡ f4cc5c93-624b-45ca-a4da-fa299c691c8c
begin
	using Pkg

	Pkg.activate()
end

# ╔═╡ 23dd09a8-d5d0-11eb-2029-c3cd330fdfa2
begin
	using SatisfiabilityInterface

	using Symbolics
	using Symbolics: Sym, Term

	using PlutoUI

	using LightGraphs, GraphPlot, Colors
	using Compose
end

# ╔═╡ a70a8996-513c-4bbd-aaf8-255d349ee739
using SimpleSATSolver

# ╔═╡ 25d094ff-bfe9-48af-b6b9-f2203ad49e47
using Latexify

# ╔═╡ e2ece52b-efbd-41de-b98e-b9e4e52bf033
using Luxor

# ╔═╡ fbcd65de-d0a4-4914-84f7-5bbbfacd0aa9
emph(s) = HTML("$s")

# ╔═╡ 533cff90-d873-4c99-8e63-3f0a77811cf6
# TableOfContents()

# ╔═╡ 1a9aad3a-bc69-4a77-bf28-c268a0bc1737
set_default_graphic_size(6*Compose.cm, 6*Compose.cm)  # size for Compose figures

# ╔═╡ 22297817-e948-467c-9a4a-8a1cee20c1e9
html"<button onclick=present()>Present</button>"

# ╔═╡ ec63ded2-3d52-4256-bf8e-638be5133902
md"""
# Solving discrete problems using $(emph("Boolean satisfiability")) with Julia
"""

# ╔═╡ 5c1ac720-70b4-495a-a66c-4ed952fbd137
md"""
### David P. Sanders

#### Facultad de Ciencias $(HTML("
")) Universidad Nacional Autónoma de México
#### Department of Mathematics & Julia Lab, MIT
"""

# ╔═╡ f4e269bd-1e67-4bb3-9dc0-fa4c51eae012
md"""
## Outline
"""

# ╔═╡ 67713a33-da7d-495e-911b-051fc34bf638
md"""

- A $(emph("symbolic")) modelling interface for discrete problems
"""

# ╔═╡ 9635873c-7362-4b38-ad9e-1e1eacb49220
md"""
- How to $(emph("encode")) problems as Boolean satisfiability problems
"""

# ╔═╡ fc64bb2e-7c4b-4467-ad68-8d141ae67484
md"""
- How to $(emph("solve")) Boolean satisfiability problems
"""

# ╔═╡ f7598c9b-2673-4b8b-9b2f-b44d648c2b66
HTML("
")

# ╔═╡ f42290e3-27f8-4a74-99f0-ed823c5dd16f
md"""
- Related work:

  - `JuliaConstraints` organisation

  - `ConstraintProgrammingExtensions.jl` for JuMP / MOI

"""

# ╔═╡ 7a4f1c53-53cf-44be-9ac0-fdc960bdae51
md"""
## Motivation: Colouring $(emph("graphs"))
"""

# ╔═╡ a1723672-5ce5-4fe8-a808-2cc1da3c8c60
md"""
- A $(emph("graph")) has a set `V` of $(emph("vertices")):
"""

# ╔═╡ d0477e2c-9be5-46de-a85f-bbf94c3f3fe7
V = [1, 2, 3]

# ╔═╡ 13a4d93a-114a-4ccd-aaee-28c80fb0606f
md"""
- Joined by $(emph("edges")) in the set `E`:
"""

# ╔═╡ 90789020-b24c-41a8-847b-e6309babd31f
E = [(1, 2), (2, 3)]

# ╔═╡ f47c8993-1d2a-4baf-88c7-fa5e605e34b5
gplot(V, E)

# ╔═╡ e5ea5a80-af37-48ce-91ca-4bc2652c1673
md"""
Graphs rendered using `LightGraphs.jl` and `GraphPlot.jl`
"""

# ╔═╡ 0f305f88-26ca-4acb-9696-ca0bc31ab979
md"""
## Graph $(emph("colouring"))
"""

# ╔═╡ 6a1342de-507b-4008-bedb-5f387f17a850
md"""
- The $(emph("graph colouring")) problem is:
"""

# ╔═╡ e6bd1f4e-e23a-45a0-be82-2f47c1f665cf
md"""

> - Assign a colour to each vertex
"""

# ╔═╡ 729aae42-b4f2-493a-ade8-71a7b790561f
md"""
>   - Such that two vertices joined by an edge have *distinct* colours
"""

# ╔═╡ d61c615e-c751-490c-9889-506c8dee30aa
md"""
## $(emph("Modelling")) the colouring problem
"""

# ╔═╡ c4cc188f-1476-4d39-b4cb-3f5632d23c79
md"""
- We want a mathematical $(emph("model")) for the colouring problem
"""

# ╔═╡ fd63036a-c674-4f30-b8cc-21b67528dd42
md"""
- We use symbolic $(emph("variables")) from `Symbolics.jl`
"""

# ╔═╡ 12628428-1bb1-4f45-8909-88a1d71b8ee7
md"""
- We call $\quad c_i$ = colour of vertex $i$
"""

# ╔═╡ 57e9e379-6ef1-4cc8-ae7b-bd17b4e8e11b
n = length(V)   # number of vertices

# ╔═╡ 4c2c7c62-3b60-48cb-96db-4020885dee1c
@variables c[1:n];    # from Symbolics.jl

# ╔═╡ 8ac35f21-1ed6-4b2e-b019-4ddfd5830851
[c[i] for i in 1:n]

# ╔═╡ 170aaa75-7a8b-4886-90ad-21f417aac55b
md"""
## Domain constraints
"""

# ╔═╡ 759127e1-0703-4fa1-b61f-f4c26c126cc6
md"""
- We need a $(emph("domain")) for each variable
"""

# ╔═╡ b5b9a579-3c53-4724-8f3c-339254fd522d
md"""
- Let's try to colour the graph with 2 colours:
"""

# ╔═╡ c54d16eb-d06a-4dbc-a3e7-c3513da710ee
k = 2

# ╔═╡ e5efb464-042e-48fc-9224-c6505d8e04a3
colours = [:red, :green, :yellow, :blue, :black]

# ╔═╡ 79370a37-d35c-42e5-b6dd-abfd60711fa9
md"""
- We can write domains $(emph("symbolically")):
"""

# ╔═╡ fec874c5-f32d-4988-a6e0-f449d362fc9c
c[1] ∈ colours[1:k]

# ╔═╡ 182e2eac-cf13-4225-a003-cc2c82553ae5
domains = [c[i] ∈ colours[1:k] for i in 1:n]

# ╔═╡ 3fb4ca38-7f53-4251-aa5e-d183b03376af
md"""
## Edge constraints
"""

# ╔═╡ 7ade9ef2-416f-438c-b126-421db2379776
md"""
- If nodes $i$ and $j$ are joined by an edge then they must have distinct colours:
"""

# ╔═╡ d1070d1c-f74d-4bd3-b0ce-862d848fd565
restrictions = [c[i] ≠ c[j] for (i, j) in E]

# ╔═╡ 6e821823-1f30-4b2e-a3e3-4de2a550dc19
md"""
- Let's collect all the constraints together:
"""

# ╔═╡ de3fd8b8-37c2-424e-98be-45824154edd3
constraints = domains ∪ restrictions

# ╔═╡ 0181743f-1a2b-4d44-85b1-77bf08edbf19
md"""
## Constraint satisfaction problem
"""

# ╔═╡ 58ce7a7f-8797-41c1-af47-314157dcd873
md"""
- A $(emph("constraint satisfaction problem")) (CSP) is a collection of variables, domains and constraints
"""

# ╔═╡ 49977007-7b86-426c-bdfa-4fc218bf1e7e
md"""
- We make a `DiscreteCSP` object from `SatisfiabilityInterface.jl` using the constraints:
"""

# ╔═╡ 8d214801-c65d-4d71-bc0c-b05852343428
prob = DiscreteCSP(constraints)

# ╔═╡ c8acfa87-a671-437d-bb9b-b776d014c491
md"""
## Solving the constraint satisfaction problem
"""

# ╔═╡ c600bd16-b655-422d-bb28-7dc124e0ba5a
md"""
- We call `solve` to solve the constraint satisfaction problem:
"""

# ╔═╡ 53a0cc3a-aee5-479e-9a56-9d6563c64c8a
status, results = solve(prob);

# ╔═╡ 760fda23-0027-412d-a3c5-aab300631bc8
results

# ╔═╡ 1df9c28b-54ef-4a1a-adbe-2610861a8559
gplot(V, E, prob, results)

# ╔═╡ d8092ba3-57b9-46aa-ae32-9f66444edb5b
md"""
## Ring graph
"""

# ╔═╡ bb584a01-1051-4b4d-936e-a64e83e58f17
md"""
- Let's make a graph in the form of a ring:
"""

# ╔═╡ 9f6ae168-e89f-4e29-bbf1-057d0dc66c1a
md"""
- Let's try to colour it with 2 colours:
"""

# ╔═╡ b8fe62d5-70e3-4bc7-9d59-9c8ac19ae56f
md"""
	- We see that this is $(emph("unsatisfiable"))!
	"""

# ╔═╡ 4bf219bd-dc77-44e8-91af-9aa67bb26c24
md"""

	## Colouring a ring again

	- But of course it can be done with *three* colours:
	"""

# ╔═╡ 7b86ec36-6e8d-45a7-bbc7-81c8eb0037e0
md"""
## Colouring a Penrose tiling -- @Cormullion
"""

# ╔═╡ f41caa13-d354-4a42-9bc4-f7ba12f9fffe
PlutoUI.LocalResource("penrose.png")

# ╔═╡ 0110824f-c664-4312-81a4-2cdb967f926b
md"""
## How do we actually $(emph("solve")) these problems?
"""

# ╔═╡ c9db9c18-d1f6-47f4-8e8f-2f158a6eedde
md"""
- We **encode** the problem to a $(emph("boolean satisfiability")) problem (SAT)
"""

# ╔═╡ ef59a7e9-c952-4081-86d0-e5091cf950e5
md"""
- Simple encoding of a variable $x$ with domain 1:$n$ :

  - Use $n$ **Boolean** variables $x_1, \ldots, x_n$

  - Restrict such that **exactly one** of the $x_i$ is true
"""

# ╔═╡ f6769053-14b4-444b-8761-a217ebf3fbf2
encoded = encode(prob)

# ╔═╡ 6a95e734-7f1c-4364-a28c-f7ac7807336a
md"""
## SAT problem in CNF form
"""

# ╔═╡ 5f8ae350-aa2b-4331-9443-af183e355529
md"""
- CNF = $(emph("Conjunctive Normal Form"))

- A conjunction (AND) of clauses

- Each clause is an OR
"""

# ╔═╡ 7442450f-8919-4fe9-ba94-36a27135997d
encoded.p

# ╔═╡ d6d38c75-d38a-4009-acbd-9195370bcbcc
md"""
## SAT solvers
"""

# ╔═╡ a59ffbb7-56a4-4723-a40a-561ed4b995a5
md"""
- The Boolean satisfiability problem is $(emph("NP-complete"))
"""

# ╔═╡ 647b6d87-62ca-43fa-b772-afa8ab9555ac
md"""
- But it can be solved efficiently in many cases, by one of many $(emph("SAT solvers")) that are available
"""

# ╔═╡ e891d669-9431-4ee2-8eb0-d4921c1b2d38
md"""
- Usually written in C or C++
"""

# ╔═╡ f2dbb4d7-5c49-46e7-9535-dc10500e7e61
md"""
- `PicoSAT.jl` wraps the popular `PicoSAT` solver
"""

# ╔═╡ ea0028ed-df33-4f05-adb0-036bc07a8eb2
md"""
- `SatisfiabilityInterface.jl` currently calls out to `cryptominisat5` by default
"""

# ╔═╡ c4e2bb24-462c-450b-a145-c1fea2bd9ee8
md"""
## `SimpleSATSolver.jl`
"""

# ╔═╡ 17f423cf-a241-48f3-a718-51b17ab6a268
md"""
- A very simple SAT solver in $(emph("pure Julia"))
"""

# ╔═╡ bb498e44-559f-414e-ab11-e055037a4b87
md"""
- Implements the $(emph("DPLL algorithm")) (1962)
"""

# ╔═╡ f7061757-93e3-4405-80d3-92f854697e5b
md"""
- Meant to be pedagogical
"""

# ╔═╡ 958cea58-82b1-43b3-975d-aade709c718d
md"""
- But seems to be performant for small systems
"""

# ╔═╡ 03f70163-6712-4e92-9a47-ff246e227813
typeof(encoded.p)

# ╔═╡ ae00c44f-b04b-4233-95cf-9c9a8c6dd799
solve(encoded.p, SimpleSAT())

# ╔═╡ 1528a881-8501-4122-be89-bf6b75203a60
md"""
## Algorithm for solving boolean satisfiability
"""

# ╔═╡ 44638a74-7b54-4ef2-8a64-c1e8da6cd553
md"""
- Start assigning variables to `true` or `false`
"""

# ╔═╡ fa9d0583-522b-43ea-90b1-70e902b9d55c
md"""
- Check if any constraints are violated and $(emph("backtrack")) (choose the opposite value) if so
	"""

# ╔═╡ 3932867d-74c0-49cc-840e-2d1421a3867a
md"""
- Apply $(emph("unit propagation")):

   - If $(x \vee y)$ and $x$ is false then $y$ must be true

"""

# ╔═╡ ddb8ef7d-6e90-4699-a635-5351622f79d6
md"""
## Discrete numerical problems
"""

# ╔═╡ 812ba4be-e5b2-42f0-be4e-7ff643284662
md"""
- We can also solve discrete numerical problems, e.g.:
"""

# ╔═╡ d1f99f24-d9fb-409d-8ef6-316797a5f88a
vars = @variables x, y, z

# ╔═╡ 955768ef-c550-47ee-b56a-2dc4bc5baad4
numerical_constraints = [
    x ∈ 1:2
    y ∈ 2:5
    z ∈ 3:7
    z == x * y
	x ≠ 1
]

# ╔═╡ 8ed7146d-ca65-4676-b0c9-f97cf67aff32
md"""
## Solving numerical problems
"""

# ╔═╡ 1edf43b4-418d-4615-82a7-00ff7b1ac41c
prob6 = DiscreteCSP(numerical_constraints);

# ╔═╡ 0b0da897-61bb-40cf-a3f9-b30d95c47563
solve(prob6)   # single solution

# ╔═╡ f03160be-b968-4b19-a369-d48cb7dc010a
all_solutions(prob6)

# ╔═╡ 601d3d5d-e613-4c5a-a164-805ccee5f73d
md"""
## Conclusions
"""

# ╔═╡ 2193ce7e-dfde-4bed-a6fb-e4b4831f36fa
md"""
- Encoding to SAT is a $(emph("powerful paradigm")) for solving discrete problems
"""

# ╔═╡ f9580d63-0bdf-4f3b-9dfb-98838f3b666f
md"""
- Our work aims to make this accessible from Julia
"""

# ╔═╡ 6be5dbf5-bee0-4935-a5c7-e5daaad7360d
md"""

- Our packages:

   - `SatisfiabilityInterface.jl`

   - `SimpleSATSolver.jl`

"""

# ╔═╡ 420da272-eeb9-431c-8bcb-4b8666746655
md"""
# Appendix
"""

# ╔═╡ 97b7b89a-0919-4f81-812a-23d3823b38ba

# ╔═╡ e556a339-a596-49dd-83c3-e0c5c72c7d4d
function convert_to_lightgraphs(V, E)

	g = Graph()
	add_vertices!(g, length(V))

	for (i, j) in E
		add_edge!(g, i, j)
	end

	return g
end

# ╔═╡ b0d6ba28-d60d-4aee-a489-79af197a29de
function GraphPlot.gplot(V, E, c, results; nodelabel=V)
	RGB_colours = Dict(col => parse(RGBA, col) for col in colours)
	colour_indices = [results[c[i]] for i in 1:length(c)]

	nodefillc = [RGB_colours[c] for c in colour_indices]

	g = convert_to_lightgraphs(V, E)

	gplot(g,
		nodelabel=nodelabel,
		nodefillc=nodefillc
	)

end

# ╔═╡ 403216ab-a405-43c3-adf6-926be6b063bc
function GraphPlot.gplot(V, E, prob::DiscreteCSP, results; nodelabel=V)
	gplot(V, E, sort(collect(keys(prob.varmap))), results; nodelabel=nodelabel)
end

# ╔═╡ b48b43ba-115c-4b69-9e9c-481e25f12581
function GraphPlot.gplot(V, E; nodelabel=V)
	p = gplot(convert_to_lightgraphs(V, E), nodelabel=nodelabel)
end

# ╔═╡ 80e4d7f7-eb31-480a-beac-abd4bd8b2798

"k is the number of colours"
function graph_colouring_problem(V, E, k=3)

    # colours = [:red, :green, :yellow, :blue, :black][1:k]
    # c = [Num(Variable(:c, i)) for i in 1:length(V)]   # colour variables

    @variables c[1:length(V)]

    constraints =
    [
        [c[i] ∈ colours[1:k] for i in 1:length(V)]

        [c[i] ≠ c[j] for (i, j) in E]  #  different_colours(E, c)
    ]

    return DiscreteCSP(constraints)
end

# ╔═╡ c1ff01b8-916c-4333-bc34-d0c04debb362
### Fix Latexify output:

# ╔═╡ f796d811-dbaa-4f5f-83a5-5911ac3dae43
begin
	Latexify.latexify(x::Sym)  = Latexify.latexify(Num(x))
	Latexify.latexify(x::Term) = Latexify.latexify(Num(x))
end

# ╔═╡ a16ee047-921f-46bb-bfda-71be201740c3
function ring_graph(n=11)
    V = 1:n
    E = [(i, mod1((i+1), n)) for i in 1:n]

    return V, E
end

# ╔═╡ 10f71d92-0575-4940-9c57-27ea52b5fbb1
V2, E2 = ring_graph(3)

# ╔═╡ a79fb891-8ab0-4c19-b801-b7881e5525ca
gplot(V2, E2)

# ╔═╡ 5ee64af1-a53b-4c66-b4b9-0c11b05cae78
prob2 = graph_colouring_problem(V2, E2, 2);

# ╔═╡ a1e31033-998c-4bfc-9d23-8295d5da742d
status2, results2 = solve(prob2)

# ╔═╡ d9b22906-e72a-43dd-81f2-2241e5dca6ee
prob3 = graph_colouring_problem(V2, E2, 3);

# ╔═╡ 26ef3f82-be05-4545-bc1e-cefe96bcddde
status3, results3 = solve(prob3);

# ╔═╡ 9bdaa197-cf4b-45e2-b231-8ec571b3fc95
gplot(V2, E2, prob3, results3)

# ╔═╡ f801d658-0562-4239-9806-e5eccf239fff
## Penrose tiling -- code by Cormullion

# ╔═╡ cd29bbff-a06c-433f-882d-e18b4987114a
struct PenroseTriangle
    red::Bool # not red, more of a flag
    pointA::Point
    pointB::Point
    pointC::Point
end

# ╔═╡ 307e96c3-a270-4b12-a0ac-f865cb20d49a
function subdivide(triangles;
        type=:P3)
    result = []
    for triangle in triangles
        A, B, C = triangle.pointA, triangle.pointB, triangle.pointC
        if triangle.red == true
            if type == :P3
                #  P3 rhombus
                P = A + (B - A) / MathConstants.golden
                push!(result, PenroseTriangle(true, C, P, B))
                push!(result, PenroseTriangle(false, P, C, A))
            else # P2 half kite
                Q = A + (B - A) / MathConstants.golden
                R = B + (C - B) / MathConstants.golden
                push!(result, PenroseTriangle(false, R, Q, B))
                push!(result, PenroseTriangle(true,  Q, A, R))
                push!(result, PenroseTriangle(true,  C, A, R))
            end
        else
            if type == :P3
                # P3 rhombus
                Q = B + (A - B) / MathConstants.golden
                R = B + (C - B) / MathConstants.golden
                push!(result, PenroseTriangle(true, R, Q, A))
                push!(result, PenroseTriangle(false, R, A, C))
                push!(result, PenroseTriangle(false, Q, R, B))
            else # P2 kite/dart
                P = C + (A - C) / MathConstants.golden
                push!(result, PenroseTriangle(false, B, P, A))
                push!(result, PenroseTriangle(true,  P, C, B))
            end
        end
    end
    return result
end

# ╔═╡ 17ae28f4-ba09-42bb-83f0-51f771d0bbc5
function PenroseTiling(centerpos::Point, radius, depth=4;
        type=:P3)
    triangles = PenroseTriangle[]
    A = centerpos
    n = 10 # a circle of triangles 360/10 -> 36°
    for i in 1:n
        phi = (i - 1) * (2π / n)
        C = A + polar(radius, phi)
        phi = (i) * (2π / n)
        B = A + polar(radius, phi)
        if type == :P3
            if i % 2 == 1
                triangle = PenroseTriangle(true, A, C, B)
            else
                triangle = PenroseTriangle(true, A, B, C)
            end
        else # P2
            if i % 2 == 1
                triangle = PenroseTriangle(true, B, A, C)
            else
                triangle = PenroseTriangle(true, C, A, B)
            end
        end
        push!(triangles, triangle)
    end
    for i in 1:depth
        triangles = subdivide(triangles, type=type)
    end
    return triangles
end

# ╔═╡ 759d23f1-360c-4a36-8cf6-0d1f6c6133cd

function incidence_graph(triangles)
    n = length(triangles)
    g = SimpleGraph(n)
    for i in 1:n
        for j in i+1:n
            ti = triangles[i]
            tj = triangles[j]

            points_i = [ti.pointA, ti.pointB, ti.pointC]
            points_j = [tj.pointA, tj.pointB, tj.pointC]

            coincidental = 0
            tol = 1e-2

            for k in 1:3, l in 1:3
                if distance(points_i[k], points_j[l]) < tol
                    coincidental += 1
                end
            end

            if coincidental == 2  # share edge
                add_edge!(g, i, j)
            end
        end
    end

    return g
end

# ╔═╡ 4a0baabf-2d98-4d1a-9906-b05c0654ce86

# ╔═╡ e4281582-9748-4801-972f-2880d6031f1b
tiling = PenroseTiling(O, 1, 4)

# ╔═╡ f32d038d-c126-4fda-bd1c-f33dc039d042
g = incidence_graph(tiling)

# ╔═╡ 0bb71243-3bb9-4ffb-b062-fd8cffa43701
VV = vertices(g)

# ╔═╡ 6f90a7ba-e4cb-4cb1-a60c-f1918642f2b6
EE = [(e.dst, e.src) for e in edges(g)]

# ╔═╡ 0fcbd390-bc8f-4284-b7d3-61428ab6ec82
prob4 = graph_colouring_problem(VV, EE, 4)

# ╔═╡ cf89ab12-e70f-4533-b7e4-52bdc253c18a
@time status4, results4 = solve(prob4)  # 0.05 seconds

# ╔═╡ 38aa8bc1-5b9a-48ae-90e6-5826c86aaff5
unique(values(results4))

# ╔═╡ ff73ed7c-ef11-4876-a345-15f265e753a8

# ╔═╡ 77645c06-1dde-4a29-86bb-5d62f06ee3b0
set_default_graphic_size(20*Compose.cm, 20*Compose.cm)  # size for Compose figures

# ╔═╡ cc04b735-3bff-4f60-a392-2c3d9e6b0220
gplot(VV, EE; nodelabel=nothing)

# ╔═╡ 70405d58-e2e5-4d55-b145-ddaf5a3ed36f
gplot(VV, EE, prob4, results4; nodelabel=nothing)

# ╔═╡ cb819540-f084-45af-bec2-6afb535ab5f9
g.fadjlist[70]

# ╔═╡ bf2a29da-52ed-44f0-b572-38cbb53bd260
results4[c[73]]

# ╔═╡ 31b3bb83-286b-4937-9e86-bed812991013
V5, E5 = ring_graph(31)

# ╔═╡ 7fcbbd88-785f-4c65-96b2-d74239de6e56
prob5 = graph_colouring_problem(V5, E5, 3)

# ╔═╡ c90239b3-d451-4379-adf5-56262ab20b4e
status5, results5 = solve(prob5)

# ╔═╡ 0f1b0462-e0af-4d01-953b-231de11fdf79
gplot(V5, E5, prob5, results5)

# ╔═╡ 8872ee87-e03f-4445-b706-a0b15d86ba07
sort(collect(keys(prob4.varmap)))

# ╔═╡ fa0c842a-b252-4421-90c9-dc713bb1a040
isless(string(c[100]),  string(c[10]))

# ╔═╡ 42429a49-0c05-4b41-96ed-67f58bb762c8
c[100] < c[10]

# ╔═╡ 875f7bd7-a1bd-44f6-a8bb-9a6464cbcc0f
typeof(c[100])

# ╔═╡ eb29d8b9-37b8-4822-a4ec-5a98bd777892
sort(collect(keys(prob4.varmap)))

# ╔═╡ ece5d0ae-8edb-42a0-a467-9effd1253af8
Base.isless(x::Term, y::Term) = isless(x.args[2], y.args[2])

# ╔═╡ 31cdbb1c-044c-4bb0-9fed-ecbb1f8b126e
Base.isless(x::Num, y::Num) = isless(x.val, y.val)

# ╔═╡ 37cb6adb-c6af-47b6-b663-2e89dc11fc0a
isless(c[1], c[100])

# ╔═╡ 1706b53e-0bab-4206-ad88-570039430839
Dump(d.val)

# ╔═╡ Cell order:
# ╠═f4cc5c93-624b-45ca-a4da-fa299c691c8c
# ╠═23dd09a8-d5d0-11eb-2029-c3cd330fdfa2
# ╠═fbcd65de-d0a4-4914-84f7-5bbbfacd0aa9
# ╠═533cff90-d873-4c99-8e63-3f0a77811cf6
# ╠═1a9aad3a-bc69-4a77-bf28-c268a0bc1737
# ╟─22297817-e948-467c-9a4a-8a1cee20c1e9
# ╟─ec63ded2-3d52-4256-bf8e-638be5133902
# ╟─5c1ac720-70b4-495a-a66c-4ed952fbd137
# ╟─f4e269bd-1e67-4bb3-9dc0-fa4c51eae012
# ╟─67713a33-da7d-495e-911b-051fc34bf638
# ╟─9635873c-7362-4b38-ad9e-1e1eacb49220
# ╟─fc64bb2e-7c4b-4467-ad68-8d141ae67484
# ╟─f7598c9b-2673-4b8b-9b2f-b44d648c2b66
# ╟─f42290e3-27f8-4a74-99f0-ed823c5dd16f
# ╟─7a4f1c53-53cf-44be-9ac0-fdc960bdae51
# ╟─a1723672-5ce5-4fe8-a808-2cc1da3c8c60
# ╠═d0477e2c-9be5-46de-a85f-bbf94c3f3fe7
# ╟─13a4d93a-114a-4ccd-aaee-28c80fb0606f
# ╠═90789020-b24c-41a8-847b-e6309babd31f
# ╠═f47c8993-1d2a-4baf-88c7-fa5e605e34b5
# ╟─e5ea5a80-af37-48ce-91ca-4bc2652c1673
# ╟─0f305f88-26ca-4acb-9696-ca0bc31ab979
# ╟─6a1342de-507b-4008-bedb-5f387f17a850
# ╟─e6bd1f4e-e23a-45a0-be82-2f47c1f665cf
# ╟─729aae42-b4f2-493a-ade8-71a7b790561f
# ╟─d61c615e-c751-490c-9889-506c8dee30aa
# ╟─c4cc188f-1476-4d39-b4cb-3f5632d23c79
# ╟─fd63036a-c674-4f30-b8cc-21b67528dd42
# ╟─12628428-1bb1-4f45-8909-88a1d71b8ee7
# ╠═57e9e379-6ef1-4cc8-ae7b-bd17b4e8e11b
# ╠═4c2c7c62-3b60-48cb-96db-4020885dee1c
# ╠═8ac35f21-1ed6-4b2e-b019-4ddfd5830851
# ╟─170aaa75-7a8b-4886-90ad-21f417aac55b
# ╟─759127e1-0703-4fa1-b61f-f4c26c126cc6
# ╟─b5b9a579-3c53-4724-8f3c-339254fd522d
# ╠═c54d16eb-d06a-4dbc-a3e7-c3513da710ee
# ╠═e5efb464-042e-48fc-9224-c6505d8e04a3
# ╟─79370a37-d35c-42e5-b6dd-abfd60711fa9
# ╠═fec874c5-f32d-4988-a6e0-f449d362fc9c
# ╠═182e2eac-cf13-4225-a003-cc2c82553ae5
# ╟─3fb4ca38-7f53-4251-aa5e-d183b03376af
# ╟─7ade9ef2-416f-438c-b126-421db2379776
# ╠═d1070d1c-f74d-4bd3-b0ce-862d848fd565
# ╟─6e821823-1f30-4b2e-a3e3-4de2a550dc19
# ╠═de3fd8b8-37c2-424e-98be-45824154edd3
# ╟─0181743f-1a2b-4d44-85b1-77bf08edbf19
# ╟─58ce7a7f-8797-41c1-af47-314157dcd873
# ╟─49977007-7b86-426c-bdfa-4fc218bf1e7e
# ╠═8d214801-c65d-4d71-bc0c-b05852343428
# ╟─c8acfa87-a671-437d-bb9b-b776d014c491
# ╟─c600bd16-b655-422d-bb28-7dc124e0ba5a
# ╠═53a0cc3a-aee5-479e-9a56-9d6563c64c8a
# ╠═760fda23-0027-412d-a3c5-aab300631bc8
# ╠═1df9c28b-54ef-4a1a-adbe-2610861a8559
# ╟─d8092ba3-57b9-46aa-ae32-9f66444edb5b
# ╟─bb584a01-1051-4b4d-936e-a64e83e58f17
# ╠═10f71d92-0575-4940-9c57-27ea52b5fbb1
# ╠═a79fb891-8ab0-4c19-b801-b7881e5525ca
# ╟─9f6ae168-e89f-4e29-bbf1-057d0dc66c1a
# ╠═5ee64af1-a53b-4c66-b4b9-0c11b05cae78
# ╠═a1e31033-998c-4bfc-9d23-8295d5da742d
# ╟─b8fe62d5-70e3-4bc7-9d59-9c8ac19ae56f
# ╟─4bf219bd-dc77-44e8-91af-9aa67bb26c24
# ╠═d9b22906-e72a-43dd-81f2-2241e5dca6ee
# ╠═26ef3f82-be05-4545-bc1e-cefe96bcddde
# ╠═9bdaa197-cf4b-45e2-b231-8ec571b3fc95
# ╟─7b86ec36-6e8d-45a7-bbc7-81c8eb0037e0
# ╟─f41caa13-d354-4a42-9bc4-f7ba12f9fffe
# ╟─0110824f-c664-4312-81a4-2cdb967f926b
# ╟─c9db9c18-d1f6-47f4-8e8f-2f158a6eedde
# ╟─ef59a7e9-c952-4081-86d0-e5091cf950e5
# ╠═f6769053-14b4-444b-8761-a217ebf3fbf2
# ╟─6a95e734-7f1c-4364-a28c-f7ac7807336a
# ╟─5f8ae350-aa2b-4331-9443-af183e355529
# ╠═7442450f-8919-4fe9-ba94-36a27135997d
# ╟─d6d38c75-d38a-4009-acbd-9195370bcbcc
# ╟─a59ffbb7-56a4-4723-a40a-561ed4b995a5
# ╟─647b6d87-62ca-43fa-b772-afa8ab9555ac
# ╟─e891d669-9431-4ee2-8eb0-d4921c1b2d38
# ╟─f2dbb4d7-5c49-46e7-9535-dc10500e7e61
# ╟─ea0028ed-df33-4f05-adb0-036bc07a8eb2
# ╟─c4e2bb24-462c-450b-a145-c1fea2bd9ee8
# ╟─17f423cf-a241-48f3-a718-51b17ab6a268
# ╟─bb498e44-559f-414e-ab11-e055037a4b87
# ╟─f7061757-93e3-4405-80d3-92f854697e5b
# ╟─958cea58-82b1-43b3-975d-aade709c718d
# ╠═a70a8996-513c-4bbd-aaf8-255d349ee739
# ╠═03f70163-6712-4e92-9a47-ff246e227813
# ╠═ae00c44f-b04b-4233-95cf-9c9a8c6dd799
# ╟─1528a881-8501-4122-be89-bf6b75203a60
# ╟─44638a74-7b54-4ef2-8a64-c1e8da6cd553
# ╟─fa9d0583-522b-43ea-90b1-70e902b9d55c
# ╟─3932867d-74c0-49cc-840e-2d1421a3867a
# ╟─ddb8ef7d-6e90-4699-a635-5351622f79d6
# ╟─812ba4be-e5b2-42f0-be4e-7ff643284662
# ╠═d1f99f24-d9fb-409d-8ef6-316797a5f88a
# ╠═955768ef-c550-47ee-b56a-2dc4bc5baad4
# ╟─8ed7146d-ca65-4676-b0c9-f97cf67aff32
# ╠═1edf43b4-418d-4615-82a7-00ff7b1ac41c
# ╠═0b0da897-61bb-40cf-a3f9-b30d95c47563
# ╠═f03160be-b968-4b19-a369-d48cb7dc010a
# ╟─601d3d5d-e613-4c5a-a164-805ccee5f73d
# ╟─2193ce7e-dfde-4bed-a6fb-e4b4831f36fa
# ╟─f9580d63-0bdf-4f3b-9dfb-98838f3b666f
# ╟─6be5dbf5-bee0-4935-a5c7-e5daaad7360d
# ╟─420da272-eeb9-431c-8bcb-4b8666746655
# ╠═97b7b89a-0919-4f81-812a-23d3823b38ba
# ╠═e556a339-a596-49dd-83c3-e0c5c72c7d4d
# ╠═b0d6ba28-d60d-4aee-a489-79af197a29de
# ╠═403216ab-a405-43c3-adf6-926be6b063bc
# ╠═b48b43ba-115c-4b69-9e9c-481e25f12581
# ╠═80e4d7f7-eb31-480a-beac-abd4bd8b2798
# ╠═c1ff01b8-916c-4333-bc34-d0c04debb362
# ╠═25d094ff-bfe9-48af-b6b9-f2203ad49e47
# ╠═f796d811-dbaa-4f5f-83a5-5911ac3dae43
# ╠═a16ee047-921f-46bb-bfda-71be201740c3
# ╠═f801d658-0562-4239-9806-e5eccf239fff
# ╠═e2ece52b-efbd-41de-b98e-b9e4e52bf033
# ╠═cd29bbff-a06c-433f-882d-e18b4987114a
# ╠═17ae28f4-ba09-42bb-83f0-51f771d0bbc5
# ╠═307e96c3-a270-4b12-a0ac-f865cb20d49a
# ╠═759d23f1-360c-4a36-8cf6-0d1f6c6133cd
# ╠═4a0baabf-2d98-4d1a-9906-b05c0654ce86
# ╠═e4281582-9748-4801-972f-2880d6031f1b
# ╠═f32d038d-c126-4fda-bd1c-f33dc039d042
# ╠═0bb71243-3bb9-4ffb-b062-fd8cffa43701
# ╠═6f90a7ba-e4cb-4cb1-a60c-f1918642f2b6
# ╠═0fcbd390-bc8f-4284-b7d3-61428ab6ec82
# ╠═cf89ab12-e70f-4533-b7e4-52bdc253c18a
# ╠═38aa8bc1-5b9a-48ae-90e6-5826c86aaff5
# ╠═ff73ed7c-ef11-4876-a345-15f265e753a8
# ╠═77645c06-1dde-4a29-86bb-5d62f06ee3b0
# ╠═cc04b735-3bff-4f60-a392-2c3d9e6b0220
# ╠═70405d58-e2e5-4d55-b145-ddaf5a3ed36f
# ╠═cb819540-f084-45af-bec2-6afb535ab5f9
# ╠═bf2a29da-52ed-44f0-b572-38cbb53bd260
# ╠═31b3bb83-286b-4937-9e86-bed812991013
# ╠═7fcbbd88-785f-4c65-96b2-d74239de6e56
# ╠═c90239b3-d451-4379-adf5-56262ab20b4e
# ╠═0f1b0462-e0af-4d01-953b-231de11fdf79
# ╠═8872ee87-e03f-4445-b706-a0b15d86ba07
# ╠═fa0c842a-b252-4421-90c9-dc713bb1a040
# ╠═42429a49-0c05-4b41-96ed-67f58bb762c8
# ╠═875f7bd7-a1bd-44f6-a8bb-9a6464cbcc0f
# ╠═eb29d8b9-37b8-4822-a4ec-5a98bd777892
# ╠═ece5d0ae-8edb-42a0-a467-9effd1253af8
# ╠═31cdbb1c-044c-4bb0-9fed-ecbb1f8b126e
# ╠═37cb6adb-c6af-47b6-b663-2e89dc11fc0a
# ╠═1706b53e-0bab-4206-ad88-570039430839


###############################################################################
#
# # ListofStates tutorial
#
# We illustrate here the use of ListofStates in dealing with a warm start
# procedure.
#
# ListofStates can also prove the user history over the iteration process.
#
# We compare the resolution of a convex unconstrained problem with 3 variants:
# - a steepest descent method
# - an inverse-BFGS method
# - a mix with 5 steps of steepest descent and then switching to BFGS with
#the history (using the strength of the ListofStates).
#
###############################################################################
using Stopping, NLPModels, LinearAlgebra, Test, Printf

import Stopping.armijo
function armijo(xk, dk, fk, slope, f)
  t = 1.0
  fk_new = f(xk + dk)
  while f(xk + t * dk) > fk + 1.0e-4 * t * slope
    t /= 1.5
    fk_new = f(xk + t * dk)
  end
  return t, fk_new
end

#Newton's method for optimization:
function steepest_descent(stp :: NLPStopping)

  xk = stp.current_state.x
  fk, gk = objgrad(stp.pb, xk)

  OK = update_and_start!(stp, fx = fk, gx = gk)

  @printf "%2s %9s %7s %7s %7s\n" "k" "fk" "||∇f(x)||" "t" "λ"
  @printf "%2d %7.1e %7.1e\n" stp.meta.nb_of_stop fk norm(stp.current_state.current_score)
  while !OK
    dk = - gk
    slope = dot(dk, gk)
    t, fk = armijo(xk, dk, fk, slope, x->obj(stp.pb, x))
    xk += t * dk
    fk, gk = objgrad(stp.pb, xk)

    OK = update_and_stop!(stp, x = xk, fx = fk, gx = gk)

    @printf "%2d %9.2e %7.1e %7.1e %7.1e\n" stp.meta.nb_of_stop fk norm(stp.current_state.current_score) t slope
  end
  return stp
end

function bfgs_quasi_newton_armijo(stp :: NLPStopping; Hk = nothing)

  xk = stp.current_state.x
  fk, gk = objgrad(stp.pb, xk)
  gm = gk

  dk, t = similar(gk), 1.
  if isnothing(Hk)
    Hk = I #start from identity matrix
  end

  OK = update_and_start!(stp, fx = fk, gx = gk)

  @printf "%2s %7s %7s %7s %7s\n" "k" "fk" "||∇f(x)||" "t" "cos"
  @printf "%2d %7.1e %7.1e\n" stp.meta.nb_of_stop fk norm(stp.current_state.current_score)

  while !OK
    if stp.meta.nb_of_stop != 0
      sk = t * dk
      yk = gk - gm
      ρk = 1/dot(yk, sk)
      #we need yk'*sk > 0 for instance yk'*sk ≥ 1.0e-2 * sk' * Hk * sk
      Hk = ρk ≤ 0.0 ? Hk : (I - ρk * sk * yk') * Hk * (I - ρk * yk * sk') + ρk * sk * sk'
      if norm(sk) ≤ 1e-14
        break
      end
      #H2 = Hk + sk * sk' * (dot(sk,yk) + yk'*Hk*yk )*ρk^2 - ρk*(Hk * yk * sk' + sk * yk'*Hk)
    end
    dk = - Hk * gk
    slope = dot(dk, gk) # ≤ -1.0e-4 * norm(dk) * gnorm
    t, fk = armijo(xk, dk, fk, slope, x->obj(stp.pb, x))

    xk = xk + t * dk
    gm = copy(gk)
    gk = grad(stp.pb, xk)

    OK = update_and_stop!(stp, x = xk, fx = fk, gx = gk)
    @printf "%2d %7.1e %7.1e %7.1e %7.1e\n" stp.meta.nb_of_stop fk norm(stp.current_state.current_score) t slope
  end
  stp.stopping_user_struct = Dict( :Hk => Hk)
  return stp
end
using Test

############ PROBLEM TEST #############################################
fH(x) = (x[2]+x[1].^2-11).^2+(x[1]+x[2].^2-7).^2
nlp = ADNLPModel(fH, [10., 20.])
stp = NLPStopping(nlp, optimality_check = unconstrained_check,
                 atol = 1e-6, rtol = 0.0, max_iter = 100)

reinit!(stp, rstate = true, x = nlp.meta.x0)
steepest_descent(stp)

@test status(stp) == :Optimal
@test stp.listofstates == VoidListofStates()

@show elapsed_time(stp)
@show nlp.counters

reinit!(stp, rstate = true, x = nlp.meta.x0, rcounters = true)
bfgs_quasi_newton_armijo(stp)

@test status(stp) == :Optimal
@test stp.listofstates == VoidListofStates()

@show elapsed_time(stp)
@show nlp.counters

NLPModels.reset!(nlp)
stp_warm = NLPStopping(nlp, optimality_check = unconstrained_check,
                      atol = 1e-6, rtol = 0.0, max_iter = 5,
                      n_listofstates = 5) #shortcut for list = ListofStates(5, Val{NLPAtX{Float64,Array{Float64,1},Array{Float64,2}}}()))
steepest_descent(stp_warm)
@test status(stp_warm) == :IterationLimit
@test length(stp_warm.listofstates) == 5

Hwarm = I
for i=2:5
  sk = stp_warm.listofstates.list[i][1].x - stp_warm.listofstates.list[i-1][1].x
  yk = stp_warm.listofstates.list[i][1].gx - stp_warm.listofstates.list[i-1][1].gx
  ρk = 1/dot(yk, sk)
  if ρk > 0.0
    global Hwarm = (I - ρk * sk * yk') * Hwarm * (I - ρk * yk * sk') + ρk * sk * sk'
  end
end

reinit!(stp_warm)
stp_warm.meta.max_iter = 100
bfgs_quasi_newton_armijo(stp_warm, Hk = Hwarm)
status(stp_warm)

@show elapsed_time(stp_warm)
@show nlp.counters

#=
 k        fk ||∇f(x)||       t       λ
 0 1.7e+05 3.2e+04
 1  2.73e+04 8.6e+03 1.0e-03 -1.1e+09
 2  1.80e+03 1.1e+03 2.3e-03 -7.3e+07
 3  1.24e+03 7.9e+02 1.2e-02 -1.3e+06
 4  6.37e+01 2.4e+01 1.2e-02 -6.3e+05
 5  1.34e+01 5.8e+01 2.0e-01 -8.3e+02
 6  5.87e+00 2.5e+01 1.3e-01 -3.5e+03
 7  2.88e+00 2.4e+01 2.6e-02 -6.7e+02
 8  2.42e+00 1.8e+01 1.7e-02 -6.1e+02
 9  6.58e-01 1.2e+01 1.2e-02 -6.1e+02
10  1.64e-01 5.3e+00 1.2e-02 -1.7e+02
11  4.96e-02 3.2e+00 1.2e-02 -4.4e+01
12  1.44e-02 1.6e+00 1.2e-02 -1.3e+01
13  4.35e-03 9.2e-01 1.2e-02 -3.9e+00
14  1.29e-03 5.0e-01 1.2e-02 -1.2e+00
15  3.87e-04 2.7e-01 1.2e-02 -3.5e-01
16  1.15e-04 1.5e-01 1.2e-02 -1.0e-01
17  3.45e-05 8.2e-02 1.2e-02 -3.1e-02
18  1.03e-05 4.5e-02 1.2e-02 -9.2e-03
19  3.08e-06 2.4e-02 1.2e-02 -2.8e-03
20  9.21e-07 1.3e-02 1.2e-02 -8.2e-04
21  2.75e-07 7.3e-03 1.2e-02 -2.5e-04
22  8.23e-08 4.0e-03 1.2e-02 -7.4e-05
23  2.46e-08 2.2e-03 1.2e-02 -2.2e-05
24  7.35e-09 1.2e-03 1.2e-02 -6.6e-06
25  2.20e-09 6.5e-04 1.2e-02 -2.0e-06
26  6.57e-10 3.6e-04 1.2e-02 -5.9e-07
27  1.96e-10 1.9e-04 1.2e-02 -1.8e-07
28  5.87e-11 1.1e-04 1.2e-02 -5.3e-08
29  1.75e-11 5.8e-05 1.2e-02 -1.6e-08
30  5.24e-12 3.2e-05 1.2e-02 -4.7e-09
31  1.57e-12 1.7e-05 1.2e-02 -1.4e-09
32  4.68e-13 9.5e-06 1.2e-02 -4.2e-10
33  1.40e-13 5.2e-06 1.2e-02 -1.3e-10
34  4.18e-14 2.8e-06 1.2e-02 -3.7e-11
35  1.25e-14 1.6e-06 1.2e-02 -1.1e-11
36  3.74e-15 8.5e-07 1.2e-02 -3.3e-12
elapsed_time(stp) = 0.7508440017700195
nlp.counters =   Counters:
             obj: ████████████████████ 889               grad: █⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 37                cons: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0
            jcon: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0                jgrad: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0                  jac: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0
           jprod: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0               jtprod: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0                 hess: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0
           hprod: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0               jhprod: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0

 k      fk ||∇f(x)||       t     cos
 0 1.7e+05 3.2e+04
 1 2.7e+04 8.6e+03 1.0e-03 -1.1e+09
 2 1.8e+04 4.5e+03 1.2e-02 -1.8e+06
 3 2.5e+03 1.3e+03 1.0e+00 -7.1e+04
 4 1.2e+03 8.5e+02 1.0e+00 -1.7e+03
 5 3.2e+02 3.3e+02 1.0e+00 -1.4e+03
 6 9.8e+01 1.4e+02 1.0e+00 -3.2e+02
 7 2.7e+01 6.0e+01 1.0e+00 -1.1e+02
 8 6.4e+00 2.4e+01 1.0e+00 -3.0e+01
 9 9.9e-01 7.9e+00 1.0e+00 -8.2e+00
10 6.3e-02 1.9e+00 1.0e+00 -1.5e+00
11 8.7e-04 3.2e-01 1.0e+00 -1.1e-01
12 3.6e-05 7.9e-02 1.0e+00 -1.6e-03
13 1.4e-05 4.2e-02 1.0e+00 -2.9e-05
14 2.0e-07 3.4e-03 1.0e+00 -2.6e-05
15 4.1e-09 4.9e-04 1.0e+00 -3.6e-07
16 2.9e-12 2.5e-05 1.0e+00 -8.1e-09
17 2.5e-15 6.3e-07 1.0e+00 -5.6e-12
elapsed_time(stp) = 0.017869949340820312
nlp.counters =   Counters:
             obj: ████████████████████ 91                grad: ████⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 18                cons: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0
            jcon: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0                jgrad: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0                  jac: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0
           jprod: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0               jtprod: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0                 hess: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0
           hprod: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0               jhprod: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0

 k        fk ||∇f(x)||       t       λ
 0 1.7e+05 3.2e+04
 1  2.73e+04 8.6e+03 1.0e-03 -1.1e+09
 2  1.80e+03 1.1e+03 2.3e-03 -7.3e+07
 3  1.24e+03 7.9e+02 1.2e-02 -1.3e+06
 4  6.37e+01 2.4e+01 1.2e-02 -6.3e+05
 5  1.34e+01 5.8e+01 2.0e-01 -8.3e+02
 6  5.87e+00 2.5e+01 1.3e-01 -3.5e+03
 k      fk ||∇f(x)||       t     cos
 0 5.9e+00 2.5e+01
 1 3.8e+00 2.7e+01 1.7e-02 -1.1e+03
 2 2.8e+00 2.4e+01 4.4e-01 -1.1e+01
 3 1.4e+00 1.2e+01 3.0e-01 -3.0e+01
 4 1.1e-02 1.3e+00 1.0e+00 -2.5e+00
 5 9.0e-05 9.2e-02 1.0e+00 -2.5e-02
 6 7.9e-08 3.9e-03 1.0e+00 -1.8e-04
 7 7.7e-10 3.8e-04 1.0e+00 -1.4e-07
 8 1.3e-19 4.2e-09 1.0e+00 -1.5e-09
elapsed_time(stp_warm) = 0.01520395278930664
nlp.counters =   Counters:
             obj: ████████████████████ 192               grad: ██⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 16                cons: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0
            jcon: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0                jgrad: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0                  jac: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0
           jprod: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0               jtprod: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0                 hess: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0
           hprod: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0               jhprod: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0

  Counters:
             obj: ████████████████████ 192               grad: ██⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 16                cons: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0
            jcon: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0                jgrad: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0                  jac: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0
           jprod: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0               jtprod: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0                 hess: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0
           hprod: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0               jhprod: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0

=#


###########################################################
# This is my module that contains some wrapper functions.
# Author: Tokuma Suzuki
###########################################################

using LinearAlgebra, SparseArrays

function speye(size, dtype=Float64)
    """
    Return sparse identity matrix.

    This function needs LinearAlgebra and SparseArrays.
    Arguments:
    size: Int64
    dtype: keywords, Float64 or Int64
    """

    return sparse(Matrix{dtype}(I, size, size))
end

function linspace(start,stop,size=50)
    """

    Return linealy spaced 1d array

    Arguments:
    start: Real
    end: Real
    size: Int
    """

    return collect(LinRange(start,stop,size))
end


using LombScargle, Distances

export IrregularLombScargle

"""
    IrregularLombScargle(t; tol = 1, n_total = 100000, n_acc = 50000, q = 1)

[`IrregularLombScargle`](@ref) surrogates for unevenly sampled time series
with supporting time steps `t`, generated using the simulated annealing algorithm
described in [^SchreiberSchmitz1999].

[`IrregularLombScargle`](@ref) surrogates (given enough iterations and a low enough
tolerance) preserve the periodogram and the amplitude
distribution of the original signal. For time series with equidistant time steps,
surrogates generated by this method result in surrogates similar to those produced
by the [`IAAFT`](@ref) method.

This algorithm starts with a random permutation of the original data. Then it iteratively
approaches the power spectrum of the original data by swapping two randomly selected values
in the surrogate data if the Minkowski distance of order `q` between the power spectrum of
the surrogate data and the original data is less than before. The iteration procedure ends
when the relative deviation between the periodograms is less than `tol` or when `n_total`
number of tries or `n_acc` number of actual swaps is reached.

[^SchmitzSchreiber1999]: A.Schmitz T.Schreiber (1999). "Testing for nonlinearity in unevenly sampled time series" [Phys. Rev E](https://journaIrregularLombScargle.aps.org/pre/pdf/10.1103/PhysRevE.59.4044)
"""
struct IrregularLombScargle{T<:AbstractVector, S<:Real} <: Surrogate
    t::T
    tol::S
    n_total::Int
    n_acc::Int
    q::Int
end
IrregularLombScargle(t; tol = 1.0, n_total = 100000, n_acc = 50000, q = 1) = IrregularLombScargle(t, tol, n_total, n_acc, q)

function surrogenerator(x, method::IrregularLombScargle, rng = Random.default_rng())
    # Plan Lombscargle periodogram. Use default flags.
    lsplan = LombScargle.plan(method.t, x, fit_mean = false)

    # Compute initial periodogram.
    x_ls = lombscargle(lsplan)

    # We have to copy the power vector here, because we are reusing `lsplan` later on
    xpower = copy(x_ls.power)

    # Use Minkowski distance of order q
    dist = Distances.Minkowski(method.q)

    # When re-computing the Lomb-Scargle periodogram, we will use the
    # `_periodogram!` method, which re-uses the lsplan with a shuffled
    # time vector. This is the same as shuffling the signal, so the
    # surrogate starts out as a shuffled version of `t`.
    s = shuffle(rng, method.t)

    # Initialize a candidate surrogate.
    candidate = zero(s)

    init = (
        lsplan = lsplan,
        xpower = xpower,
        n = length(x),
        dist = dist,
        candidate = candidate,
    )

    return SurrogateGenerator(method, x, s, init, rng)
end

function (sg::SurrogateGenerator{<:IrregularLombScargle})()
    lsplan, xpower, n, dist, candidate = sg.init
    t, tol = sg.method.t, sg.method.tol
    x, s, rng = sg.x, sg.s, sg.rng

    # Power spectrum for the randomly shuffled signal.
    spower = LombScargle._periodogram!(lsplan.P, s, lsplan)

    # Compare power spectra for original (`xpower`) and randomly shuffled signal (`spower`).
    lossold = Distances.evaluate(dist, xpower, spower)

    range = 1:n
    i = j = 0
    while i < sg.method.n_total && j < sg.method.n_acc
        if mod(i, 2000) == 0
            @info "iterations: $i, swaps: $j, loss: $lossold"
        end

        # Initially, the new surrogate is identical to the existing surrogate.
        copy!(candidate, s)

        # Swap two random points and re-compute power spectrum for the candidate.
        swap_elements!(candidate, s, rng, range)

        # This is the bottleneck.
        spower = LombScargle._periodogram!(lsplan.P, candidate, lsplan)

        # If spectra are more similar after the swap, accept the new
        # surrogate. Otherwise, do a new swap.
        lossnew = evaluate(dist, xpower, spower)

        if lossnew < lossold
            lossnew <= tol && break
            s = copy(candidate)
            lossold = lossnew
            j += 1
        end
        i += 1
    end
    @info "Terminated simulated annealing process after $i iterations and $j swaps. Loss: $lossold"

    # Use the permutation of the time vector to permute the signal vector
    # This gives us the inverse permutation from t to perm
    perm = sortperm(sortperm(s))

    # Check if this worked as expected
    @assert t[perm] == s

    # Re-scale back to original time series values.
    return x[perm]
end

function swap_elements!(candidate, s, rng, range)
    k = sample(rng, range)
    l = sample(rng, range)
    candidate[k] = s[l]
    candidate[l] = s[k]
end


"""
    lautat(surf, curfield, nsteps=500, dtstar=0.015, startflag=0,
writeflag=0, writeInterval=1000., delvort=delNone(), maxwrite=100,
nround=6)

Simulates potential flow for an airfoil undergoing unsteady motion
    using Large-Angle Unsteady Thin-Airfoil Theory.

    Inputs required:
    - `TwoDSurf` : based on airfoil geometry and kinematics
    - `TwoDFlowField` : initial flowfield

    `dtstar` should be in accordance with the kinematics; see `?find_tstep`.

    `nsteps` should be calculated according to total run time required and
    dtstar.

    `startflag=1` can be used to continue a simulation using the
    *resultsSummary* file in current directory.

    `writeflag=1` along with a suitable writeInterval is used to write flow
    data at intervals. These can be used to create plots and movies of the
    flowfield.

    `delvort` is used to control the vortex count. The simulation cost
    increases as O(n``^2``) with number of vortices in the flow field and so
    long run times can get very computationally expensive. Look at
    `?delVortDef` for algorthims to control vortex count.

    `maxwrite` is the maximum number of timestamps/directories that may be
    written.

    `nround` is the number of digits to which the output timestamps are
    rounded off to.

    The Large-Angle Unsteady Thin-Airfoil theory is described in: Ramesh,
    K. et al., "An unsteady airfoil theory applied to pitching motions
    validated against experiment and computation", Theor. Comput. Fluid
    Dyn. (2013) 27: 843. [Weblink](https://doi.org/10.1007/s00162-012-0292-8)

"""
function lautat(surf::TwoDSurf, curfield::TwoDFlowField, nsteps::Int64 = 500, dtstar::Float64 = 0.015, startflag = 0, writeflag = 0, writeInterval = 1000., delvort = delNone(); maxwrite = 50, nround=6, wakerollup=1)

    # If a restart directory is provided, read in the simulation data
    if startflag == 0
        mat = zeros(0, 8)
        t = 0.
    elseif startflag == 1
        dirvec = readdir()
        dirresults = map(x->(v = tryparse(Float64,x); typeof(v) == Nothing ? 0.0 : v),dirvec)
        latestTime = maximum(dirresults)
        mat = DelimitedFiles.readdlm("resultsSummary",header = true)
        t = mat[end,1]
    else
        throw("invalid start flag, should be 0 or 1")
    end
    mat = mat'

    dt = dtstar*surf.c/surf.uref

    # if writeflag is on, determine the timesteps to write at
    if writeflag == 1
        writeArray = Int64[]
        tTot = nsteps*dt
        for i = 1:maxwrite
            tcur = writeInterval*real(i)
            if t > tTot
                break
            else
                push!(writeArray, Int(round(tcur/dt)))
            end
        end
    end

    #Intialise flowfield
    for istep = 1:nsteps
        #Udpate current time
        t = t + dt

        #Update kinematic parameters
        update_kinem(surf, t)

        #Update bound vortex positions
        update_boundpos(surf, dt)

        #Add a TEV with dummy strength
        place_tev(surf,curfield,dt)

        #Solve for TEV strength to satisfy Kelvin condition
        kelv = KelvinCondition(surf,curfield)
        soln = nlsolve(not_in_place(kelv), [-0.01])
        curfield.tev[length(curfield.tev)].s = soln.zero[1]

        #Update adot
        update_a2a3adot(surf,dt)

        #Set previous values of aterm to be used for derivatives in next time step
        surf.a0prev[1] = surf.a0[1]
        for ia = 1:3
            surf.aprev[ia] = surf.aterm[ia]
        end

        #Update rest of Fourier terms
        update_a2toan(surf)

        #Calculate bound vortex strengths
        update_bv(surf)

        # Delete or merge vortices if required
        controlVortCount(delvort, surf.bnd_x[Int(round(surf.ndiv/2))], surf.bnd_z[Int(round(surf.ndiv/2))], curfield)

        #Wake rollup if flag on
        if wakerollup == 1
            wakeroll(surf, curfield, dt)
        end

        # Calculate force and moment coefficients
        cl, cd, cm = calc_forces(surf, [curfield.u[1], curfield.w[1]])

        # write flow details if required
        if writeflag == 1
            if istep in writeArray
                dirname = "$(round(t, sigdigits=nround))"
                writeStamp(dirname, t, surf, curfield)
            end
        end

        # for writing in resultsSummary
        mat = hcat(mat,[t, surf.kinem.alpha, surf.kinem.h, surf.kinem.u, surf.a0[1], cl, cd, cm])

    end

    mat = mat'

    f = open("resultsSummary", "w")
    println(f,"#time \t","alpha(deg) \t","h/c \t", "u/uref \t", "A0 \t", "Cl \t", "Cd \t", "Cm ")
    DelimitedFiles.writedlm(f, mat)
    close(f)

    mat, surf, curfield

end

function ldvm(surf::TwoDSurf, curfield::TwoDFlowField, nsteps::Int64 = 500, dtstar::Float64 = 0.015, startflag = 0, writeflag = 0, writeInterval = 1000., delvort = delNone(); maxwrite = 50, nround=6)

    # If a restart directory is provided, read in the simulation data
    if startflag == 0
        mat = zeros(0, 8)
        t = 0.
    elseif startflag == 1
        dirvec = readdir()
        dirresults = map(x->(v = tryparse(Float64,x); typeof(v) == Nothing ? 0.0 : v),dirvec)
        latestTime = maximum(dirresults)
        mat = DelimitedFiles.readdlm("resultsSummary",header = true)
        t = mat[end,1]
    else
        throw("invalid start flag, should be 0 or 1")
    end
    mat = mat'

    dt = dtstar*surf.c/surf.uref

    # if writeflag is on, determine the timesteps to write at
    if writeflag == 1
        writeArray = Int64[]
        tTot = nsteps*dt
        for i = 1:maxwrite
            tcur = writeInterval*real(i)
            if t > tTot
                break
            else
                push!(writeArray, Int(round(tcur/dt)))
            end
        end
    end

    # time stepping
    for istep = 1:nsteps
        #Udpate current time
        t = t + dt

        #Update external flowfield
        update_externalvel(curfield, t)

        #Update kinematic parameters
        update_kinem(surf, t)

        #Update bound vortex positions
        update_boundpos(surf, dt)

        #Add a TEV with dummy strength
        place_tev(surf,curfield,dt)

        #Update incduced velocities on airfoil
        update_indbound(surf, curfield)

        #Solve for TEV strength to satisfy Kelvin condition
        kelv = KelvinCondition(surf,curfield)
        soln = nlsolve(not_in_place(kelv), [-0.01])
        curfield.tev[length(curfield.tev)].s = soln.zero[1]

        #Update adot
        update_a2a3adot(surf,dt)

        lesp = surf.a0[1]

        #Check for LESP condition
        if abs(lesp)>surf.lespcrit[1]
            #2D iteration if LESP_crit is exceeded

            #Add a LEV with dummy strength
            place_lev(surf,curfield,dt)

            #Solve for TEV and LEV strengths to satisfy Kelvin condition and Kutta condition at leading edge
            kelvkutta = KelvinKutta(surf,curfield)
            soln = nlsolve(not_in_place(kelvkutta), [-0.01; 0.01])
            (curfield.tev[length(curfield.tev)].s, curfield.lev[length(curfield.lev)].s) = soln.zero[1], soln.zero[2]

            #set flag for levshedding=on
            surf.levflag[1] = 1
        else
            surf.levflag[1] = 0
        end

        #Update rest of Fourier terms
        update_a2toan(surf)

        #Set previous values of aterm to be used for derivatives in next time step
        surf.a0prev[1] = surf.a0[1]
        for ia = 1:3
            surf.aprev[ia] = surf.aterm[ia]
        end

        #Calculate bound vortex strengths
        update_bv(surf)

        # Delete or merge vortices if required
        controlVortCount(delvort, surf.bnd_x[Int(round(surf.ndiv/2))], surf.bnd_z[Int(round(surf.ndiv/2))], curfield)

        # free wake rollup
        wakeroll(surf, curfield, dt)

        # Calculate force and moment coefficients
        cl, cd, cm, cn = calc_forces(surf, [curfield.u[1], curfield.w[1]])

        # write flow details if required
        if writeflag == 1
            if istep in writeArray
                dirname = "$(round(t, sigdigits=nround))"
                writeStamp(dirname, t, surf, curfield)
            end
        end

        # for writing in resultsSummary
        mat = hcat(mat,[t, surf.kinem.alpha, surf.kinem.h, surf.kinem.u, surf.a0[1], cl, cd, cm])

    end

    mat = mat'

    f = open("resultsSummary", "w")
    println(f,"#time \t","alpha(deg) \t","h/c \t", "u/uref \t", "A0 \t", "Cl \t", "Cd \t", "Cm ")
    DelimitedFiles.writedlm(f, mat)
    close(f)

    mat, surf, curfield

end

function ldvm(surf::Vector{TwoDSurf}, curfield::TwoDFlowField, nsteps::Int64 = 500, dtstar::Float64 = 0.015, startflag = 0, writeflag = 0, writeInterval = 1000., delvort = delNone(); maxwrite = 50, nround=6)

    nsurf = length(surf)

    # If a restart directory is provided, read in the simulation data
    if startflag == 0
        mat = zeros(0, 7*nsurf+1)
        t = 0.
    elseif startflag == 1
        dirvec = readdir()
        dirresults = map(x->(v = tryparse(Float64,x); typeof(v) == Nothing ? 0.0 : v),dirvec)
        latestTime = maximum(dirresults)
        mat = DelimitedFiles.readdlm("resultsSummary",header = true)
        t = mat[end,1]
    else
        throw("invalid start flag, should be 0 or 1")
    end
    mat = mat'

    dt = dtstar*minimum(map(q->q.c/q.uref, surf))

    # if writeflag is on, determine the timesteps to write at
    if writeflag == 1
        writeArray = Int64[]
        tTot = nsteps*dt
        for i = 1:maxwrite
            tcur = writeInterval*real(i)
            if t > tTot
                break
            else
                push!(writeArray, Int(round(tcur/dt)))
            end
        end
    end

    lesp = zeros(nsurf)
    cl = zeros(nsurf)
    cd = zeros(nsurf)
    cm = zeros(nsurf)

    # time stepping
    for istep = 1:nsteps
        #Udpate current time
        t = t + dt

        #Update external flowfield
        update_externalvel(curfield, t)

        for is = 1:nsurf
            #Update kinematic parameters
            update_kinem(surf[is], t)

            #Update bound vortex positions
            update_boundpos(surf[is], dt)
        end

        #Add TEVs with dummy strength
        place_tev(surf,curfield,dt)

        for is = 1:nsurf
            #Update induced velocities on airfoil
            update_indbound(surf[is], curfield)
            for js = 1:nsurf
                if is != js
                    #add influence of other surfaces to downwash
                    add_indbound_b(surf[is], surf[js])
                end
            end
        end

        #Solve for TEV strength to satisfy Kelvin condition
        kelv = KelvinConditionMult(surf,curfield)
        soln = nlsolve(not_in_place(kelv), ones(nsurf)*-0.01)

        for i = 1:nsurf
            curfield.tev[end-nsurf+i].s = soln.zero[i]

            #Update adot
            update_a2a3adot(surf[i],dt)

            lesp[i] = surf[i].a0[1]
        end

        shed_ind = Int[]

        #Check for LESP condition
        for i = 1:nsurf
            if (abs(lesp[i])>surf[i].lespcrit[1])

                push!(shed_ind, i)
            end
        end
        nshed = length(shed_ind)

        if nshed > 0
            #2D iteration if LESP_crit is exceeded
            #Add LEVs with dummy strength
            place_lev(surf,curfield,dt,shed_ind)

            #Solve for TEV and LEV strengths to satisfy Kelvin condition and Kutta condition at leading edge
            kelvkutta = KelvinKuttaMult(surf,curfield,shed_ind)
            soln = nlsolve(not_in_place(kelvkutta), -0.01*ones(nshed+nsurf))

            for i = 1:nsurf
                if i in shed_ind
                    if shed_ind != [1] && length(shed_ind)==1 #incase LEV only formed on 2nd surface
                    surf[i].levflag[1] = 1
                    i=i-1
                    curfield.tev[end-nsurf+i].s = soln.zero[i]
                    curfield.lev[end-nsurf+shed_ind[i]].s = soln.zero[nsurf+i]
                else
                    curfield.tev[end-nsurf+i].s = soln.zero[i]
                    curfield.lev[end-nsurf+shed_ind[i]].s = soln.zero[nsurf+i]
                    surf[i].levflag[1] = 1
                else
                    curfield.tev[end-nsurf+i].s = soln.zero[i]
                    surf[i].levflag[1] = 0
                end
            end
        else
            for i = 1:nsurf
                surf[i].levflag[1] = 0
            end
        end

        for i = 1:nsurf
            #Set previous values of aterm to be used for derivatives in next time step
            surf[i].a0prev[1] = surf[i].a0[1]
            for ia = 1:3
                surf[i].aprev[ia] = surf[i].aterm[ia]
            end
        end

        # Delete or merge vortices if required
        mean_bndx = sum(map(q->q.bnd_x[Int(round(q.ndiv/2))], surf))/nsurf
        mean_bndz = sum(map(q->q.bnd_z[Int(round(q.ndiv/2))], surf))/nsurf
        controlVortCount(delvort, mean_bndx, mean_bndz, curfield)

        # free wake rollup
        wakeroll(surf, curfield, dt)

        # Calculate force and moment coefficients
        for i = 1:nsurf
            cl[i], cd[i], cm[i] = calc_forces(surf[i], [curfield.u[1], curfield.w[1]])
        end

        # write flow details if required
        if writeflag == 1
            if istep in writeArray
                dirname = "$(round(t, sigdigits=nround))"
                writeStamp(dirname, t, surf, curfield)
            end
        end

        # for writing in resultsSummary
        matvect = [t;]
        for is = 1:nsurf
            matvect = [matvect;[surf[is].kinem.alpha, surf[is].kinem.h,
                                surf[is].kinem.u, surf[is].a0[1], cl[is], cd[is], cm[is]]]
        end
        mat = hcat(mat,matvect)

    end

mat = mat'

f = open("resultsSummary", "w")
println(f,"#time \t","alpha(deg) \t","h/c \t", "u/uref \t", "A0 \t", "Cl \t", "Cd \t", "Cm ")
DelimitedFiles.writedlm(f, mat)
close(f)

mat, surf, curfield

end

function ldvmLin(surf::TwoDSurf, curfield::TwoDFlowField, nsteps::Int64 = 500, dtstar::Float64 = 0.015, startflag = 0, writeflag = 0, writeInterval = 1000., delvort = delNone(); maxwrite = 50, nround=6)

    # If a restart directory is provided, read in the simulation data
    if startflag == 0
        mat = zeros(0, 8)
        t = 0.
    elseif startflag == 1
        dirvec = readdir()
        dirresults = map(x->(v = tryparse(Float64,x); typeof(v) == Nothing ? 0.0 : v),dirvec)
        latestTime = maximum(dirresults)
        mat = DelimitedFiles.readdlm("resultsSummary",header=true)
        t = mat[end,1]
    else
        throw("invalid start flag, should be 0 or 1")
    end
    mat = mat'

    dt = dtstar*surf.c/surf.uref

    # if writeflag is on, determine the timesteps to write at
    if writeflag == 1
        writeArray = Int64[]
        tTot = nsteps*dt
        for i = 1:maxwrite
            tcur = writeInterval*real(i)
            if t > tTot
                break
            else
                push!(writeArray, Int(round(tcur/dt)))
            end
        end
    end

    vcore = 0.02*surf.c

    T1 = zeros(surf.ndiv)
    T2 = zeros(surf.ndiv)
    T3 = zeros(surf.ndiv)

    # time stepping
    for istep = 1:nsteps
        #Udpate current time
        t = t + dt

        #Update external flowfield
        update_externalvel(curfield, t)

        #Update kinematic parameters
        update_kinem(surf, t)

        #Update bound vortex positions
        update_boundpos(surf, dt)

        #Update induced velocities on airfoil
        update_indbound(surf, curfield)

        #Calculate downwash
        update_downwash(surf, [curfield.u[1],curfield.w[1]])

        #The integrals I1 and J1 are based on this downwash which doesnt
        #include vortives shed at current step (equal to term T1)
        #I1 has units if circulation, J1 is dimensionless
        T1[:] = surf.downwash[:]
        I1 = surf.c*simpleTrapz(T1.*(cos.(surf.theta) .- 1. ), surf.theta)
        J1 = -simpleTrapz(T1,surf.theta)/(surf.uref*pi)

        # T2 depends on recenetly shed TEV
        ntev = length(curfield.tev)

        if ntev == 0
            xloc_tev = surf.bnd_x[surf.ndiv] + 0.5*surf.kinem.u*dt
            zloc_tev = surf.bnd_z[surf.ndiv]
        else
            xloc_tev = surf.bnd_x[surf.ndiv]+(1. /3.)*(curfield.tev[ntev].x - surf.bnd_x[surf.ndiv])
            zloc_tev = surf.bnd_z[surf.ndiv]+(1. /3.)*(curfield.tev[ntev].z - surf.bnd_z[surf.ndiv])
        end

        for ib = 1:surf.ndiv
            xdist = surf.bnd_x[ib] - xloc_tev
            zdist = surf.bnd_z[ib] - zloc_tev
            distsq = xdist*xdist + zdist*zdist
            T2[ib] = (surf.cam_slope[ib]*zdist + xdist)/(2*pi*sqrt(distsq^2 + vcore^4))
        end

        #sig_prev = sum(map(q->q.s, curfield.tev)) + sum(map(q->q.s, curfield.lev))
        sig_prev = -surf.uref*surf.c*pi*(surf.a0prev[1] + surf.aprev[1]/2. )

        I2 = simpleTrapz(T2.*(cos.(surf.theta) .- 1. ), surf.theta)
        J2 = -simpleTrapz(T2, surf.theta)/(pi*surf.uref)

        tevstr = -(I1 + sig_prev)/(1 + I2)

        #Calc first 3 fourier coefficients and derivatives
        surf.a0[1] = J1 + J2*tevstr
        for ia = 1:surf.naterm
            surf.aterm[ia] = 2. *(simpleTrapz(T1.*cos.(ia*surf.theta), surf.theta) + tevstr*simpleTrapz(T2.*cos.(ia*surf.theta), surf.theta))/(pi*surf.uref)
        end

        #Calculate adot
        surf.a0dot[1] = (surf.a0[1] - surf.a0prev[1])/dt
        for ia = 1:surf.naterm
            surf.adot[ia] = (surf.aterm[ia]-surf.aprev[ia])/dt
        end

        #Check if LEV shedding is true
        if abs(surf.a0[1]) > surf.lespcrit[1]
            if (surf.a0[1] >= 0.)
                lesp_cond = surf.lespcrit[1]
            else
                lesp_cond = -surf.lespcrit[1]
            end

            # T3 depends on recenetly shed LEV
            nlev = length(curfield.lev)
            if surf.levflag[1] == 0
                le_vel_x = surf.kinem.u - surf.kinem.alphadot*sin(surf.kinem.alpha)*surf.pvt*surf.c + surf.uind[1]
                le_vel_z = -surf.kinem.alphadot*cos(surf.kinem.alpha)*surf.pvt*surf.c- surf.kinem.hdot + surf.wind[1]
                xloc_lev = surf.bnd_x[1] + 0.5*le_vel_x*dt
                zloc_lev = surf.bnd_z[1] + 0.5*le_vel_z*dt
            else
                xloc_lev = surf.bnd_x[1] + (1. /3.)*(curfield.lev[nlev].x - surf.bnd_x[1])
                zloc_lev = surf.bnd_z[1]+(1. /3.)*(curfield.lev[nlev].z - surf.bnd_z[1])
            end

            for ib = 1:surf.ndiv
                xdist = surf.bnd_x[ib] - xloc_lev
                zdist = surf.bnd_z[ib] - zloc_lev
                distsq = xdist*xdist + zdist*zdist
                T3[ib] = (surf.cam_slope[ib]*zdist + xdist)/(2*pi*sqrt(distsq^2 + vcore^4))
            end
            I3 = simpleTrapz(T3.*(cos.(surf.theta) .- 1. ), surf.theta)
            J3 = -simpleTrapz(T3, surf.theta)/(pi*surf.uref)

            det = J3*(I2 + 1. ) - J2*(I3 + 1. )

            tevstr = (-J3*(I1 + sig_prev) + (I3 + 1)*(J1 - lesp_cond))/det
            levstr = (J2*(I1 + sig_prev) - (I2 + 1)*(J1 - lesp_cond))/det

            #Recalculate required fourier terms
            surf.a0[1] = J1 + J2*tevstr + J3*levstr
            for ia = 1:3
                surf.aterm[ia] = 2. *(simpleTrapz(T1.*cos.(ia*surf.theta), surf.theta) +
                                      tevstr*simpleTrapz(T2.*cos.(ia*surf.theta), surf.theta) +
                                      levstr*simpleTrapz(T3.*cos.(ia*surf.theta), surf.theta))/(pi*surf.uref)
            end

            push!(curfield.tev, TwoDVort(xloc_tev, zloc_tev, tevstr, vcore, 0., 0.))
            push!(curfield.lev, TwoDVort(xloc_lev, zloc_lev, levstr, vcore, 0., 0.))

            for ia = 4:surf.naterm
                surf.aterm[ia] = 2. *(simpleTrapz(T1.*cos.(ia*surf.theta), surf.theta) +
                                      tevstr*simpleTrapz(T2.*cos.(ia*surf.theta), surf.theta) +
                                      levstr*simpleTrapz(T3.*cos.(ia*surf.theta), surf.theta))/(pi*surf.uref)
            end

            surf.levflag[1] = 1
        else
            push!(curfield.tev, TwoDVort(xloc_tev, zloc_tev, tevstr, vcore, 0., 0.))

            for ia = 4:surf.naterm
                surf.aterm[ia] = 2. *(simpleTrapz(T1.*cos.(ia*surf.theta), surf.theta) +
                                      tevstr*simpleTrapz(T2.*cos.(ia*surf.theta), surf.theta))/(pi*surf.uref)
            end

            surf.levflag[1] = 0
        end

#Set previous values of aterm to be used for derivatives in next time step
surf.a0prev[1] = surf.a0[1]
for ia = 1:3
    surf.aprev[ia] = surf.aterm[ia]
end

#Calculate bound vortex strengths
update_bv(surf)

# Delete or merge vortices if required
controlVortCount(delvort, surf.bnd_x[Int(round(surf.ndiv/2))], surf.bnd_z[Int(round(surf.ndiv/2))], curfield)

# free wake rollup
wakeroll(surf, curfield, dt)

# Calculate force and moment coefficients
cl, cd, cm = calc_forces(surf, [curfield.u[1], curfield.w[1]])

# write flow details if required
if writeflag == 1
    if istep in writeArray
        dirname = "$(round(t,sigdigits=nround))"
        writeStamp(dirname, t, surf, curfield)
    end
end

# for writing in resultsSummary
mat = hcat(mat,[t, surf.kinem.alpha, surf.kinem.h, surf.kinem.u, surf.a0[1], cl, cd, cm])

end

mat = mat'

f = open("resultsSummary", "w")
println(f,"#time \t","alpha(deg) \t","h/c \t", "u/uref \t", "A0 \t", "Cl \t", "Cd \t", "Cm ")
DelimitedFiles.writedlm(f, mat)
close(f)

mat, surf, curfield

end

function ldvm2DOF(surf::TwoDSurf, curfield::TwoDFlowField, strpar::TwoDOFPar, kinem::KinemPar2DOF, nsteps::Int64 = 500, dtstar::Float64 = 0.015, startflag = 0, writeflag = 0, writeInterval = 1000., delvort = delNone(); maxwrite = 50, nround=6)

    # If a restart directory is provided, read in the simulation data
    if startflag == 0
        mat = zeros(0, 8)
        t = 0.
    elseif startflag == 1
        dirvec = readdir()
        dirresults = map(x->(v = tryparse(Float64,x); typeof(v) == Nothing ? 0.0 : v),dirvec)
        latestTime = maximum(dirresults)
        mat = DelimitedFiles.readdlm("resultsSummary",header=true)
        t = mat[end,1]
    else
        throw("invalid start flag, should be 0 or 1")
    end
    mat = mat'

    dt = dtstar*surf.c/surf.uref

    # if writeflag is on, determine the timesteps to write at
    if writeflag == 1
        writeArray = Int64[]
        tTot = nsteps*dt
        for i = 1:maxwrite
            tcur = writeInterval*real(i)
            if t > tTot
                break
            else
                push!(writeArray, Int(round(tcur/dt)))
            end
        end
    end

    cl = 0.
    cm = 0.

    #Intialise flowfield

    for istep = 1:nsteps
        #Udpate current time
        t = t + dt

        #Update external flowfield
        update_externalvel(curfield, t)

        #Update kinematic parameters (based on 2DOF response)
        update_kinem2DOF(surf, strpar, kinem, dt, cl, cm)

        #Update bound vortex positions
        update_boundpos(surf, dt)

        #Add a TEV with dummy strength
        place_tev(surf,curfield,dt)

        #Update incduced velocities on airfoil
        update_indbound(surf, curfield)

        #Solve for TEV strength to satisfy Kelvin condition
        kelv = KelvinCondition(surf,curfield)
        soln = nlsolve(not_in_place(kelv), [-0.01])
        curfield.tev[length(curfield.tev)].s = soln.zero[1]

        #Update adot
        #update_a2a3adot(surf,dt)
        update_atermdot(surf, dt)

        lesp = surf.a0[1]

        #Check for LESP condition
        if (abs(lesp)>surf.lespcrit[1])
            #2D iteration if LESP_crit is exceeded
            #Remove the previous tev

            #Add a LEV with dummy strength
            place_lev(surf,curfield,dt)

            #Solve for TEV and LEV strengths to satisfy Kelvin condition and Kutta condition at leading edge
            kelvkutta = KelvinKutta(surf,curfield)
            soln = nlsolve(not_in_place(kelvkutta), [-0.01; 0.01])
            (curfield.tev[length(curfield.tev)].s, curfield.lev[length(curfield.lev)].s) = soln.zero[1], soln.zero[2]

            #set flag for levshedding=on
            surf.levflag[1] = 1
        else
            surf.levflag[1] = 0
        end

        #Update rest of Fourier terms
        update_a2toan(surf)

        #Set previous values of aterm to be used for derivatives in next time step
        surf.a0prev[1] = surf.a0[1]
        for ia = 1:surf.naterm
            surf.aprev[ia] = surf.aterm[ia]
        end

        #Calculate bound vortex strengths
        update_bv(surf)

        # Delete or merge vortices if required
        controlVortCount(delvort, surf.bnd_x[Int(round(surf.ndiv/2))], surf.bnd_z[Int(round(surf.ndiv/2))], curfield)

        # free wake rollup
        wakeroll(surf, curfield, dt)

        # Calculate force and moment coefficients
        cl, cd, cm = calc_forces(surf, [curfield.u[1], curfield.w[1]])

        # write flow details if required
        if writeflag == 1
            if istep in writeArray
                dirname = "$(round(t,sigdigits=nround))"
                writeStamp(dirname, t, surf, curfield)
            end
        end

        # for writing in resultsSummary
        mat = hcat(mat,[t, surf.kinem.alpha, surf.kinem.h, surf.kinem.u, surf.a0[1], cl, cd, cm])

    end

    mat = mat'

    f = open("resultsSummary", "w")
    println(f,"#time \t","alpha(deg) \t","h/c \t", "u/uref \t", "A0 \t", "Cl \t", "Cd \t", "Cm ")
    DelimitedFiles.writedlm(f, mat)
    close(f)

    mat, surf, curfield
end


# %% chaoticity map
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot, Random

using OrdinaryDiffEq
diffeq = (alg = Vern9(), abstol = 1e-9, reltol = 1e-9)

fig = figure(figsize = (figx/2, figy))
ax = gca()
hh = Systems.henonheiles()
ics = Systems.henonheiles_ics(0.13, 15)
cmap = matplotlib.cm.get_cmap("viridis")

scres = nothing
for ic in ics
    psos = poincaresos(hh, (1, 0.0), 2000; u0 = ic)
    λ = lyapunov(hh, 10000; u0 = ic, diffeq...)
    λmax = 0.06
    v = clamp(λ/λmax, 0, 1)
    global scres = ax.scatter(psos[:, 2], psos[:, 4];
        color = cmap(v), s = 5, edgecolors = "0.5", linewidths = 0.1,
        vmin = 0, vmax = 1,
    )
end

cb = colorbar(scres)
cb.set_ticks([0, 1])
cb.set_ticklabels(["0", "0.06"])
cb.set_label("chaoticity (\$\\lambda_1\$)"; labelpad=-50)
ax.set_xticks(-0.5:0.4:0.8)
ax.set_xlabel("\$y\$", labelpad = -20)
ax.set_yticks(-0.5:0.3:0.5)
ax.set_ylabel("\$p_y\$", labelpad = -10)
fig.tight_layout(; pad = 0.25)
fig.subplots_adjust(bottom = 0.16, top = 0.95, left = 0.17, wspace = 0.01)
wsave(plotsdir("3", "chaoticity"), fig)


cd(@__DIR__); using Pkg; Pkg.activate(".");

using Revise, StableDQMC, GenericSVD, JacobiSVD, LinearAlgebra, Plots, Statistics, LaTeXStrings, DataStructures
using BenchmarkTools, StatsPlots, DataFrames, Infiltrator

pyplot();
# PyPlot.plt.style.use("default")
PyPlot.plt.style.use("publication_tex")
default(
    framestyle = :grid,
    grid = true,
    guidefont=font(15),
#     framestyle = :box,
#     grid = false,
#     guidefont=font(18),
    legend = true,
    size = (1.6*360, 360),
    xtickfont=font(13),
    ytickfont=font(13),
    legendfont=font(12),
    linewidth = 1.5
);

## Model

# Hubbard
L = 2
nsites = L^2
n = nsites * 2 # spin
t = 1
U = 6
μ = 0
Δτ = 0.1

# Kinetic part (PBC)
T = diagm(-1 => fill(-t, n-1), 1 => fill(-t, n-1), 0 => fill(μ, n))
T[1,end] = -t; T[end,1] = -t;

# on-site interaction
λ = acosh(exp(0.5*U*Δτ))
eV = Diagonal(vcat(fill(exp(λ), nsites), fill(exp(-λ), nsites)))

B = exp(-Δτ .* T) * eV

println()
@info "Hubbard model: U = $U, L = $L"
@info "Condition number of B: $(cond(B))"

## Bchain product

Bbig = BigFloat.(B)
N = 500

setprecision(100) do
    global svs_genericsvd = calc_Bchain_svd(Bbig, N; svdalg = genericsvd)[2];
end
svs_qr_udt = calc_Bchain_qr(B, N)[2];
svs_gesvd  = calc_Bchain_svd(B, N; svdalg = gesvd)[2];
svs_gesdd  = calc_Bchain_svd(B, N; svdalg = gesdd)[2];
svs_gesvj  = calc_Bchain_svd(B, N; svdalg = gesvj)[2];

betas = range(1, N * 0.1, length=N)
plot(betas, svs_genericsvd, color="#CA0020", labels = reshape(["exact", "", "","", "","", "", ""], (1,8)))
plot!(betas, svs_qr_udt, color="#E66101", labels = reshape(["QR", "", "","", "","", "", ""], (1,8)))
plot!(betas, svs_gesvj, color="#FDB863", labels = reshape(["SVD (Jacobi)", "", "","", "","", "", ""], (1,8)))
plot!(betas, svs_gesvd, color="#008837", labels = reshape(["SVD", "", "","", "","", "", ""], (1,8)))
plot!(betas, svs_gesdd, color="#5E3C99", labels = reshape(["SVD (D\\&C)", "", "","", "","", "", ""], (1,8)))

ylabel!("log singular values")
xlabel!(L"inverse temperature $\beta$")
# ylims!(-160, 100)
xlims!(0., 50.)
# savefig("../paper/figures/decomp_comparison_simple.pdf")

## Equal-time Green's function

function accuracy_greens(B; svdinversion = inv_one_plus)
    Bbig = BigFloat.(B)

    Δs_qr    = SortedDict{Int, Float64}()
    Δs_gesvd = SortedDict{Int, Float64}()
    Δs_gesdd = SortedDict{Int, Float64}()
    Δs_gesvj = SortedDict{Int, Float64}()

    for beta in range(5, 40, step=5)
        N = Int(beta * 10)

        local G_genericsvd
        setprecision(1000) do
            F_genericsvd = calc_Bchain_svd(Bbig, N; svdalg = genericsvd)[1];
            G_genericsvd = inv(I + Matrix(F_genericsvd))
        end

        # Bchain
        F_qr_udt = calc_Bchain_qr(B, N)[1];
        F_gesvd  = calc_Bchain_svd(B, N; svdalg = gesvd)[1];
        F_gesdd  = calc_Bchain_svd(B, N; svdalg = gesdd)[1];
        F_gesvj  = calc_Bchain_svd(B, N; svdalg = StableDQMC.gesvj)[1];

        # inversion
        G_qr_udt = inv_one_plus(F_qr_udt)
        G_gesvd  = svdinversion(F_gesvd; svdalg = gesvd!)
        G_gesdd  = svdinversion(F_gesdd; svdalg = gesdd!)
        G_gesvj  = svdinversion(F_gesvj; svdalg = StableDQMC.gesvj!)

        Δs_qr[beta]    = maximum(abs.(G_qr_udt - G_genericsvd))
        Δs_gesvd[beta] = maximum(abs.(G_gesvd - G_genericsvd))
        Δs_gesdd[beta] = maximum(abs.(G_gesdd - G_genericsvd))
        Δs_gesvj[beta] = maximum(abs.(G_gesvj - G_genericsvd))
    end

    return Δs_qr, Δs_gesvd, Δs_gesdd, Δs_gesvj
end

# Δs_qr, Δs_gesvd, Δs_gesdd, Δs_gesvj = Juno.@enter accuracy_greens(B);
Δs_qr, Δs_gesvd, Δs_gesdd, Δs_gesvj = accuracy_greens(B);

##

betas = collect(keys(Δs_qr))
p = plot(betas, log.(10, values(Δs_qr)), color="#E66101", marker=true, label="QR")
plot!(betas, log.(10, values(Δs_gesvd)), color="#008837", marker=true, label="SVD")
plot!(betas, log.(10, values(Δs_gesdd)), color="#5E3C99", marker=true, label="SVD (D\\&C)")
plot!(betas, log.(10, values(Δs_gesvj)), color="#FDB863", marker=true, label="SVD (Jacobi)")

xlabel!(L"inverse temperature $\beta$")
ylabel!(L"\log(\textrm{max}(\textrm{abs}(G - G_{\textrm{exact}})))")
# ylims!(0., 1e-3)
savefig("../paper/figures/accuracy_greens_svd_regularinv_U_$(U)_L_$(L)_t_$(t).pdf")

## Time-displaced Green's function

# G(slice, 1) = [B(slice, 1)^-1 + B(beta, slice)]^-1
# G(slice, 1) = [B(slice, 1)^-1 + B(beta, slice)]^-1
function accuracy_tdgf(B; loh = false)
    Bbig = BigFloat.(B)

#     betas = range(5, 40, step=5)
    beta = 40
    M = Int(beta * 10)
    slices = 1:10:M

    df = DataFrame(slice=Int[], qr=Float64[], gesdd=Float64[], gesvd=Float64[], gesvj=Float64[])

    inv_sum_method = loh ? inv_sum_loh : inv_sum

    for s in slices
        local TDGF_genericsvd
        setprecision(1000) do
            TDGF_genericsvd = calc_tdgf_svd(Bbig, s, M-s, inv_sum_method = inv_sum_method, svdalg_chain = genericsvd, svdalg_inv = genericsvd!)
        end

        TDGF_qr = calc_tdgf_qr(B, s, M-s, inv_sum_method = inv_sum_method)
        TDGF_gesdd = calc_tdgf_svd(B, s, M-s, inv_sum_method = inv_sum_method, svdalg_chain = gesdd, svdalg_inv = gesdd!)
        TDGF_gesvd = calc_tdgf_svd(B, s, M-s, inv_sum_method = inv_sum_method, svdalg_chain = gesvd, svdalg_inv = gesvd!)
        TDGF_gesvj = calc_tdgf_svd(B, s, M-s, inv_sum_method = inv_sum_method, svdalg_chain = StableDQMC.gesvj, svdalg_inv = StableDQMC.gesvj!)

        Δ_qr    = maximum(abs, TDGF_qr - TDGF_genericsvd)
        Δ_gesdd = maximum(abs, TDGF_gesdd - TDGF_genericsvd)
        Δ_gesvd = maximum(abs, TDGF_gesvd - TDGF_genericsvd)
        Δ_gesvj = maximum(abs, TDGF_gesvj - TDGF_genericsvd)

        push!(df, [s, Δ_qr, Δ_gesdd, Δ_gesvd, Δ_gesvj])
    end

    return df
end

df_tdgf = accuracy_tdgf(B; loh = true)

##

@df df_tdgf plot(:slice .* 0.1, (x -> log.(10, x)).([:qr, :gesvd, :gesdd, :gesvj]),
            marker=true,
            color=permutedims(["#E66101", "#008837", "#5E3C99", "#FDB863"]), # #E66101 for QR
            label=permutedims(["QR", "SVD", "SVD (D\\&C)", "SVD (Jacobi)"]),
#             xlims=(0,10)
           )
xlabel!(L"imaginary time $\tau$")
ylabel!(L"\log(\textrm{max}(\textrm{abs}(G(\tau, 0) - G_{\textrm{exact}}(\tau, 0))))")
savefig("../paper/figures/accuracy_tdgf_loh_U_$(U)_L_$(L)_t_$(t).pdf")


# @inline function if_else(condition::Bool, @nospecialize(trueval), @nospecialize(falseval))
#            return ifelse(condition, trueval, falseval)
# end

# ModelingToolkit.@register if_else(x, trueval, falseval)
# # This is needed to bring if_else into MTK's scope.
# @eval ModelingToolkit using ..Main: if_else

# nasty(i) = if_else(i > 1, 1, 0)
# @variables t
# nasty(t)

function MassSpringOscillator(input)

  @parameters t
  @parameters k,c,m
  D = Differential(t)
  @variables pos(t) vel(t)

  eqs = [D(pos) ~ vel,
        D(vel) ~ (-1/m)*(c*vel + k*pos - input(t))
  ]

  ps = [k,c,m] .=>  [2.,1.,4.]
  ics = [pos, vel] .=> [1.,0.]
  od = ODESystem(eqs, t, first.(ics),first.(ps))
  tspan = (0.,100.)
  # prob = ODEProblem(od, ics, tspan, ps)
  return od, ics, tspan, ps
end

# od,ic,ps = create(t->0)
# tspan = (0.,10.)
# prob = ODEProblem(od,ic,tspan,ps)


# ---
# title: 970. Powerful Integers
# id: problem970
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Hash Table, Math
# link: <https://leetcode.com/problems/powerful-integers/description/>
# hidden: true
# ---
#
# Given two positive integers `x` and `y`, an integer is _powerful_  if it is
# equal to `x^i + y^j` for some integers `i >= 0` and `j >= 0`.
#
# Return a list of all _powerful_ integers that have value less than or equal to
# `bound`.
#
# You may return the answer in any order.  In your answer, each value should
# occur at most once.
#
#
#
# **Example 1:**
#
#
#
#     Input: x = 2, y = 3, bound = 10
#     Output: [2,3,4,5,7,9,10]
#     Explanation:
#     2 = 2^0 + 3^0
#     3 = 2^1 + 3^0
#     4 = 2^0 + 3^1
#     5 = 2^1 + 3^1
#     7 = 2^2 + 3^1
#     9 = 2^3 + 3^0
#     10 = 2^0 + 3^2
#
#
# **Example 2:**
#
#
#
#     Input: x = 3, y = 5, bound = 15
#     Output: [2,4,6,8,10,14]
#
#
#
#
# **Note:**
#
#   * `1 <= x <= 100`
#   * `1 <= y <= 100`
#   * `0 <= bound <= 10^6`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


## --- test ArrayStats.jl

    # Sorting, counting, matching
    A = rand(1:100.,100); B = sort(A)
    @test A[1:count_unique!(A)] == unique(B)
    @test findclosest(3.3:5.3,1:10) == 3:5
    @test findclosestbelow(3.3:5.3,1:10) == 3:5
    @test findclosestabove(3.3:5.3,1:10) == 4:6
    @test findmatches(40:60,1:100) == 40:60
    @test findnth(fill(true,50), 25) == 25

    # Interpolation
    @test cntr(0:2:100) == 1:2:99

    # Integration
    @test trapezoidalquadrature(1:10,fill(1,10)) == 9
    @test trapz(collect(1:10.),ones(10)) == 9
    @test midpointquadrature(1:10,ones(10)) == 10

    # Distributions
    A = draw_from_distribution(ones(100), 10000)::AbstractArray
    @test length(A) == 10000
    @test isapprox(nanmean(A), 0.5, atol=0.08)
    @test isapprox(nanstd(A), sqrt(1/12), atol=0.08)

    # Strings
    @test contains("JuliaLang is pretty cool!", "Julia")
    @test !contains("JuliaLang is pretty cool!", "julia")

    @test containsi("JuliaLang is pretty cool!", "Julia")
    @test containsi("JuliaLang is pretty cool!", "julia")
    @test !containsi("JuliaLang is pretty cool!", "tomatoes")

## ---


using TestParticle, Meshes, OrdinaryDiffEq, StaticArrays, Random
using Test

"Initial state perturbation for EnsembleProblem."
function prob_func(prob, i, repeat)
   remake(prob, u0=rand(MersenneTwister(i))*prob.u0)
end

@testset "TestParticle.jl" begin
   @testset "numerical field" begin
      x = range(-10, 10, length=15)
      y = range(-10, 10, length=20)
      z = range(-10, 10, length=25)
      B = fill(0.0, 3, length(x), length(y), length(z)) # [T]
      E = fill(0.0, 3, length(x), length(y), length(z)) # [V/m]

      B[3,:,:,:] .= 10e-9
      E[3,:,:,:] .= 5e-10

      Δx = x[2] - x[1]
      Δy = y[2] - y[1]
      Δz = z[2] - z[1]

      mesh = CartesianGrid((length(x)-1, length(y)-1, length(z)-1),
         (x[1], y[1], z[1]),
         (Δx, Δy, Δz))

      x0 = [0.0, 0.0, 0.0] # initial position, [m]
      u0 = [1.0, 0.0, 0.0] # initial velocity, [m/s]
      stateinit = [x0..., u0...]

      param = prepare(mesh, E, B)
      tspan = (0.0, 1.0)
      trace! = trace_numeric!

      prob = ODEProblem(trace!, stateinit, tspan, param)

      sol = solve(prob, Tsit5(); save_idxs=[1,2,3])

      x = getindex.(sol.u, 1)
      y = getindex.(sol.u, 2)
      z = getindex.(sol.u, 3)

      @test length(x) == 8 && x[end] ≈ 0.8540967226885379

      trajectories = 10
      prob = ODEProblem(trace!, stateinit, tspan, param)
      ensemble_prob = EnsembleProblem(prob, prob_func=prob_func)
      sol = solve(ensemble_prob, Tsit5(), EnsembleThreads();
         trajectories=trajectories, save_idxs=[1,2,3])

      x = getindex.(sol.u[10].u, 1)
      y = getindex.(sol.u[10].u, 2)
      z = getindex.(sol.u[10].u, 3)

      @test x[7] ≈ 0.09615629718624641 rtol=1e-6

      stateinit = SA[x0..., u0...]

      param = prepare(mesh, E, B)
      tspan = (0.0, 1.0)
      trace = trace_numeric

      prob = ODEProblem(trace, stateinit, tspan, param)

      sol = solve(prob, Tsit5(); save_idxs=[1,2,3])

      x = getindex.(sol.u, 1)
      y = getindex.(sol.u, 2)
      z = getindex.(sol.u, 3)

      @test length(x) == 8 && x[end] ≈ 0.8540967226885379
   end

   @testset "analytical field" begin
      # in-place version
      Ek = 5e7 # [eV]

      m = TestParticle.mᵢ
      q = TestParticle.qᵢ
      c = TestParticle.c
      Rₑ = TestParticle.Rₑ

      # initial velocity, [m/s]
      v₀ = TestParticle.sph2cart(c*sqrt(1-1/(1+Ek*q/(m*c^2))^2), 0.0, π/4)
      # initial position, [m]
      r₀ = TestParticle.sph2cart(2.5*Rₑ, 0.0, π/2)
      stateinit = [r₀..., v₀...]

      param = prepare(TestParticle.getE_dipole, TestParticle.getB_dipole)
      tspan = (0.0, 1.0)

      trace! = trace_analytic!
      prob = ODEProblem(trace!, stateinit, tspan, param)

      sol = solve(prob, Tsit5(); save_idxs=[1,2,3])

      x = getindex.(sol.u, 1)

      @test x[300] ≈ 1.2592311352654776e7 rtol=1e-6

      # static array version (results not identical with above: maybe some bugs?)
      stateinit = SA[r₀..., v₀...]

      trace = trace_analytic
      prob = ODEProblem(trace, stateinit, tspan, param)

      sol = solve(prob, Tsit5(); save_idxs=[1,2,3])

      x = getindex.(sol.u, 1)

      @test x[306] ≈ 1.2588844644203672e7 rtol=1e-6

   end
end


# Reexport SinkingParticles as they are useful outside too
@reexport module SinkingParticles

using Unitful
using LinearAlgebra, SparseArrays
using ..GridTools

"""
    buildPFD(w, DIV, Iabove)

Builds the particle flux divergence operator `PFD`
for a particle sinking speed w.

Schematic of a grid cell:
```
     top  ┌─────────────────────────────┐   ┬
          │      ↓ w            ↓ Φ     │   │
          │ (sinking speed)    (flux)   │   │
          │                             │   │
          │            x                │   dz
          │      (particle conc.)       │   │
          │                             │   │
          │                             │   │
  bottom  └─────────────────────────────┘   ┴
```

- dz is the height of grid cell [m]
- w is the particle sinking speed at the top of the grid cell [m s⁻¹]
- Φ is the flux at the top of the grid cell [mmol m⁻² s⁻¹]
- x is the particle concentration of the cell [mmol m⁻³]

The PFD operator is defined by

    PFD * x ≃ dΦ/dz,

i.e., so that applied to x it approximates dΦ/dz.
It is calculated as the matrix product of the DIV and FLUX operators:

    PFD = DIV * FLUX.

DIV, the divergence operator (d/dz only here), is defined by

    DIV * Φ = 1/dz * (Φ - Φ[below]) ≃ dΦ/dz.

And FLUX, the sinking particle flux operator, is defined by

    FLUX * x = w * x[above] ≃ Φ.

# Example

```julia-repl
julia> PFD = buildPFD(w)
```
"""
buildPFD(w, DIV, Iabove) = DIV * buildFLUX(w, Iabove)

function buildPFD(grid, wet3D; sinking_speed=1.0)
    w = sinking_speed
    iwet = findall(vec(wet3D))
    DIV = buildDIV(wet3D, iwet, grid)
    Iabove = buildIabove(wet3D, iwet)
    return DIV * buildFLUX(w, Iabove)
end

"""
    buildDIV(wet3D, iwet, grid)

Build the DIV operator such that

    DIV * Φ = 1/dz * (Φ - Φ[below]) ≃ dΦ/dz.
"""
function buildDIV(wet3D, iwet, grid)
    Ibelow = buildIbelow(wet3D, iwet)
    dz = ustrip.(grid.δz_3D[iwet])
    return sparse(Diagonal(dz.^(-1))) * (Ibelow - I) # divergence with positive downwards
end

"""
    buildFLUX(w, Iabove)

Build the `FLUX` operator for a particle sinking speed ``w(r)``

Note `w` is the sinking speed at the top of each grid cell.
(`w` can also be a constant (Float64 type).)

The `FLUX` operator is defined by

    FLUX * x = w * x(above) ≃ Φ.
"""
buildFLUX(w::Vector, Iabove) = sparse(Diagonal(w)) * Iabove
buildFLUX(w::Number, Iabove) = w * Iabove

export buildDIV, buildPFD, buildFLUX

end


# --------10--------20--------30--------40--------50--------60--------70--------80--------90--------100-------110-------120--------130
## envelope
function get_env(t::TF, pulse::Pulse{TL,TF})::TF where {TL,TF}
    @unpack tL, tR, T, env_type, ncyc = pulse
    if t <= tL || tR <= t
        return 0
    end
    if env_type == :sin2
        return sin((π * t) / (ncyc * T))^2
    elseif env_type == :sin4
        return sin((π * t) / (ncyc * T))^4
    else
        env_type == :flat || env_type == :delta
        return 1
    end
end
function get_env(t::TVF, pulse::Pulse)::TVF where {TVF <: AbstractVector{<:AbstractFloat}}
    env = zeros(eltype(t), size(t))
    for it in eachindex(t)
        env[it] = get_env(t[it], pulse)
    end
    return env
end
## electric field
function getEtAt(t::TF, pulse::Pulse{TL,TF})::Tuple{TF,TF} where {TL,TF}
    return (getEt(t, pulse), getAt(t, pulse))
end
function getEt(t::TF, pulse::Pulse{Et,TF})::TF where {TF}
    @unpack E, ω, ϕ = pulse
    return E * get_env(t, pulse) * sin(ω * t + ϕ)
end
function getEt(t::TF, pulse::Pulse{At,TF})::TF where {TF}
    @unpack tL, tR, T = pulse
    if t <= tL || tR <= t
        return zero(TF)
    end
    dt = T * 1e-5
    Et =
        (1 / 60dt) * (getAt(t + 3dt, pulse) - getAt(t - 3dt, pulse)) -
        (3 / 20dt) * (getAt(t + 2dt, pulse) - getAt(t - 2dt, pulse)) +
        (3 / 4dt) * (getAt(t + dt, pulse) - getAt(t - dt, pulse))
    return -Et
end
function getEt(t::TVF, pulse::Pulse)::TVF where {TVF <: AbstractVector{<:AbstractFloat}}
    Et = zeros(eltype(t), length(t))
    for it in eachindex(t)
        Et[it] = getEt(t[it], pulse)
    end
    return Et
end
## vector potential
function getAt(t::TF, pulse::Pulse{Et,TF})::TF where {TF}
    @unpack tL, tR, T = pulse
    if t <= tL || tR <= t
        return 0
    end
    nt = 1
    tprev = pulse.tprev[1]
    At = pulse.Atprev[1]
    if tprev > t
        tprev = 0
        At = 0
    end
    while (t - tprev) / nt > T / 100000
        nt *= 10
    end
    dt = (t - tprev) / nt
    for it = 1:nt
        t0 = tprev + (it - 1) * dt
        t1 = t0 + dt / 2
        t2 = t0 + dt
        At += (-dt / 6) * (getEt(t0, pulse) + 4getEt(t1, pulse) + getEt(t2, pulse))
    end
    pulse.tprev[1] = t
    pulse.Atprev[1] = At
    return At
end
function getAt(t::TF, pulse::Pulse{At,TF})::TF where {TF}
    @unpack tL, tR, T, A, ω, ϕ = pulse
    if t <= tL || tR <= t
        return 0
    end
    return A * get_env(t, pulse) * sin(ω * t + ϕ)
end
function getAt(t::TVF, pulse::Pulse)::TVF where {TVF <: AbstractVector{<:AbstractFloat}}
    At = zeros(eltype(t), size(t))
    for it = eachindex(t)
        At[it] = getAt(t[it], pulse)
    end
    return At
end


function build_gram_matrix(q::Vector{MOI.VariableIndex},
                           monos::AbstractVector{<:MP.AbstractMonomial})
    return build_gram_matrix([MOI.SingleVariable(vi) for vi in q], monos)
end
function build_gram_matrix(q::Vector,
                           monos::AbstractVector{<:MP.AbstractMonomial})
    return GramMatrix(MultivariateMoments.SymMatrix(q, length(monos)),
                         monos)
end

function add_gram_matrix(model::MOI.ModelLike, matrix_cone_type::Type,
                         monos::AbstractVector{<:MP.AbstractMonomial})
    Q, cQ = MOI.add_constrained_variables(model, matrix_cone(matrix_cone_type, length(monos)))
    q = build_gram_matrix(Q, monos)
    return q, Q, cQ
end

function build_moment_matrix(q::Vector,
                             monos::AbstractVector{<:MP.AbstractMonomial})
    return MomentMatrix(MultivariateMoments.SymMatrix(q, length(monos)),
                        monos)
end

struct SOSPolynomialSet{DT <: AbstractSemialgebraicSet,
                        MT <: MP.AbstractMonomial,
                        MVT <: AbstractVector{MT},
                        CT <: Certificate.AbstractCertificate} <: MOI.AbstractVectorSet
    domain::DT
    monomials::MVT
    certificate::CT
end
MOI.dimension(set::SOSPolynomialSet) = length(set.monomials)


#=
Simple example of use of montecarlo.jl
From the system prompt, execute
mpirun -np X julia pi.jl
setting X to the number of ranks you'd
like to use, subject to  X-1 being an even divisor
of 1e6, e.g., set X=5.
=#
using Pkg
Pkg.activate(".")
using MPI, Statistics, LinearAlgebra
include("montecarlo.jl")
# this is the function that gets Monte Carlo'd
function pi_wrapper()
    pihat = [4.0*float((norm(rand(2,1)) .< 1.))]
end
# this function reports intermediate results during MC runs
function pi_monitor(sofar, results)
    # examine results every 2.5*10^5 draws
    if mod(sofar,2.5e5)==0.
        m = mean(results[1:sofar,:],dims=1)
        println("reps so far: ", sofar)
        println("pihat: ", m)
        println()
        #if sofar == size(results,1)
        #    writedlm("mcresults.out", results)
        #end
    end
end
# do the monte carlo: 10^6 reps of single draws
function main()
    if !MPI.Initialized()
        MPI.Init()
    end
    comm = MPI.COMM_WORLD
    reps = Int(1e7)  # desired number of MC reps
    nreturns = 1
    pooled = Int(1e5)
    montecarlo(pi_wrapper, pi_monitor, comm, reps, nreturns, pooled)
end
main()


using Oceananigans
using Oceananigans.Advection: VelocityStencil
using Oceananigans.Models.ShallowWaterModels: VectorInvariantFormulation, ConservativeFormulation, shallow_water_velocities

Lx, Ly, Lz = 2π, 20, 10
Nx, Ny = 128, 128

advection(formulation::ConservativeFormulation) = WENO5()
advection(formulation::VectorInvariantFormulation) = WENO5(vector_invariant=VelocityStencil())

function run_shallow_water_regression(arch, formulation; regenerate_data = false)
    grid = RectilinearGrid(arch, size = (Nx, Ny),
                        x = (0, Lx), y = (-Ly/2, Ly/2),
                        topology = (Periodic, Bounded, Flat))

    gravitational_acceleration = 1
    coriolis = FPlane(f=1)

    model = ShallowWaterModel(; grid, coriolis, gravitational_acceleration, formulation,
                            timestepper = :RungeKutta3,
                            momentum_advection = advection(formulation))

    U = 1 # Maximum jet velocity
    f = coriolis.f
    g = gravitational_acceleration
    Δη = f * U / g  # Maximum free-surface deformation as dictated by geostrophy

    h̄(x, y, z) = Lz - Δη * tanh(y)
    ū(x, y, z) = U * sech(y)^2

    small_amplitude = 1e-4

    uⁱ(x, y, z) = ū(x, y, z) + small_amplitude * exp(-y^2) * cos(8π * x / Lx)
    uhⁱ(x, y, z) = uⁱ(x, y, z) * h̄(x, y, z)

    if formulation isa VectorInvariantFormulation
        set!(model, u = uⁱ, h = h̄)
    else
        set!(model, uh = uhⁱ, h = h̄)
    end

    stop_iteration = 20
    Δt = 1e-2

    simulation = Simulation(model; stop_iteration, Δt)

    ## Build velocities
    u, v = shallow_water_velocities(model)
    h = model.solution.h

    output_filename = "shallow_water_bickley_jet_regression_$(typeof(formulation)).jld2"

    if regenerate_data
        @warn "Generating new data for the ShallowWater Bickley Jet regression test."

        directory =  joinpath(dirname(@__FILE__), "data")
        outputs   = (; u, v, h)
        simulation.output_writers[:fields] = JLD2OutputWriter(model, outputs,
                                                              dir = directory,
                                                              schedule = IterationInterval(stop_iteration),
                                                              filename = output_filename,
                                                              overwrite_existing = true)
    end

    run!(simulation)

    u, v = shallow_water_velocities(model)
    h = model.solution.h

    # Test results
    test_fields = (
        u = Array(interior(u)),
        v = Array(interior(v)),
        h = Array(interior(h))
    )

    if !regenerate_data
        datadep_path = "regression_test_data/" * output_filename
        regression_data_path = @datadep_str datadep_path
        file = jldopen(regression_data_path)

        truth_fields = (
            u = file["timeseries/u/$stop_iteration"][:, :, 1],
            v = file["timeseries/v/$stop_iteration"][:, :, 1],
            h = file["timeseries/h/$stop_iteration"][:, :, 1],
        )

        close(file)

        summarize_regression_test(test_fields, truth_fields)

        @test all(test_fields.u .≈ truth_fields.u)
        @test all(test_fields.v .≈ truth_fields.v)
        @test all(test_fields.h .≈ truth_fields.h)
    end
end


function perform_linprog(A,b,c,maxit=-1,tol=1e-10)
    """
        LINPROG  This code uses the revised simplex method to solve the linear
               programming problem: Minimize the cost c'x subject to
               equations Ax=b and nonnegativity x >=0:

                        f =  Min { c'x; Ax = b, x >= 0 },

               You must define an m by n matrix A , a column vector b with
               m components, and a column vector c with n components.

               The output vector x gives the minimum cost, which is the output f.

                        [x,f,itn] = linprog(A,b,c,k,maxit,tol)

               At most "maxit" iterations (default 10*length(b)) are applied
               and the actual number of iterations is returned in "itn".

               If the optimal solution is unbounded or the constraints are
               inconsistent then a diagnostic is displayed.
               Bland's rule is used to resolve degeneracies.  In exact
               arithmetic cycling is not possible.  But in real life!!!
               If x has more than 20 components it is returned in the
               sparse format. Note that if A has many zeros it is worth
               passing it to linprog in sparse format.

               Although written for teaching purposes this routine has
               successfully solved some problems with size(A) = [50,100000]!

               Please report any difficulties to: idc@math.canterbury.ac.nz

               New version                  (c) I.D.Coope, 1988, 1993
    """

    c = vec(c)
    b = vec(b)
    m,n = size(A)
    if maxit == -1
        maxit = 10*m
    end

    it=0

    D= sign(sign(b)+.5)
    D = diagm(D)                       # initial (inverse) basis matrix
    A = [A D]                      # incorporate slack/artificial variables
    B = collect(n+1:n+m)                     # initial basis
    N = collect(1:n)                    # non-basis

    phase = 1
    xb = abs(b)
    s =[zeros(n); ones(m)]    # supercost

    while phase < 3
        df = -1
        t = Inf
        yb= D'*s[B]        # multipliers for Ax=b
        while (it < maxit)
            if isempty(N)
                break
            end                              # no freedom for minimization
            r = s[N] - A[:,N]'*yb          # reduced costs
            rmin,q = findmin(r)                                    # determine new basic variable
            if rmin >= -tol*(norm(s[N],Inf)+1)
                break # optimal!
            end
            it = it+1
            if df >= 0                         # apply Bland's rule to avoid cycling
                J = find(r.<0)
                Nq = minimum(N[J])
                q = find(N.==Nq)[1]
            end
            d = vec(D*A[:,N[q]])
            I = find(d .> tol)
            if isempty(I)
                print("Solution is unbounded")
                it = -it
                break
            end
            xbd=xb[I]./d[I]
            r,p = findmin(xbd)
            p = I[p]
            if df >= 0                         # apply Bland's rule to avoid cycling
                J = find(xbd .== r)
                Bp = minimum(B[I[J]])
                p = find(B .== Bp)[1]
            end
            xb= xb - r*d

            xb[p] = r                            # update x
            df=r*rmin;
            # change in f
            v = vec(D[p,:])./d[p] # row vector
            yb= yb + v.*(s[N[q]] - dot(d',s[B]))
            d[p] = d[p] - 1
            D = D - d*v'                 # update inverse basis matrix
            t = B[p]
            B[p] = N[q]

            if t > n
                deleteat!(N,q)
            else
                N[q] = t
            end

        end
        xb=xb+D*(b-A[:,B]*xb);  # iterative refinement
        I = find(xb .< 0)                  # must be due to rounding error
        if length(I) > 0
            xb[I]=xb[I]-xb[I]                      # so correct
        end

        if phase == 2 || it < 0
            break                                # B, xb,n,m,res=A(:,B)*xb-b
        end
        if sum(xb.*s[B]) > tol
            it=-it
            print("No feasible solution")
            break
        end
        phase=phase+1                            # re-initialise for Phase 2
        s=1e6*norm(c, Inf)*s
        s[1:n]=c
    end
    x = zeros(2*n)
    x[B]=xb
    x=x[1:n]
    f=c'*x
    if it >= maxit
        print("Too many iterations")
        it=-it
    end
    return x
end


data=readdlm("MIST_iso_1470351180.iso")

using CGS
rho_sun = MSUN/(4pi/3.*RSUN^3)
# Convert to density:
mass = vec(data[:,4])
density = mass./10.^(3.*data[:,12])*rho_sun
teff = vec(10.^data[:,11])

using PyPlot

fig,axes = subplots(2,1)
ax = axes[1,1]
ax[:plot](density,teff,"r-",alpha=0.6,label="t = 1.5 Gyr, [Fe/H]=0.0, MIST",linewidth=3)

rhoA=0.001167
srhoA=0.000052
tA = 4750.
stA = 250.
ax[:plot](rhoA+srhoA*[-1,1],tA*[1,1],"g-",linewidth=3,alpha=0.6)
ax[:plot](rhoA*[1,1],tA+stA*[-1,1],"g-",linewidth=3,alpha=0.6,label="Star A")

rhoB=0.0824
srhoB=0.0034
tB = 6500.
stB = 500.
ax[:semilogx](rhoB+srhoB*[-1,1],tB*[1,1],"b-",linewidth=3,alpha=0.6)
ax[:semilogx](rhoB*[1,1],tB+stB*[-1,1],"b-",linewidth=3,alpha=0.6,label="Star B")
ax[:axis]([1e-4,1e0,4e3,8e3])
#ax[:set_xlabel]("Density [g/cc]")
ax[:set_ylabel]("Teff [K]")
ax[:legend](loc = "upper left")

ax = axes[2,1]

ax[:semilogx](density,mass,"r-",alpha=0.6,label="t = 1.5 Gyr, [Fe/H]=0.0, MIST",linewidth=3)
mA = 1.9
smA = 0.27
ax[:semilogx](rhoA+srhoA*[-1,1],mA*[1,1],"g-",linewidth=3,alpha=0.6)
ax[:semilogx](rhoA*[1,1],mA+smA*[-1,1],"g-",linewidth=3,alpha=0.6,label="Star A")
ax[:semilogx](rhoB+srhoB*[-1,1,1,-1,-1],[0,0,10,10,0],"b-",linewidth=3,alpha=0.6,label="Star B")
ax[:axis]([1e-4,1e0,1.7,2.0])
ax[:set_xlabel]("Density [g/cc]")
ax[:set_ylabel]("Mass [Msun]")

PyPlot.savefig("isochrone.png",bbox_inches="tight")