text:
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: jl:light,ipynb
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.5.1
#   kernelspec:
#     display_name: Julia 1.4.2
#     language: julia
#     name: julia-1.4
# ---

# +
"""
Tutorial: A reference implementation of orbital optimized second-order Moller-Plesset perturbation theory.
"""

__authors__   = ["D. Menendez", "Boyi Zhang"]
__credits__   = ["D. Menendez", "Boyi Zhang", "Justin M. Turney"]

__copyright_amp__ = "(c) 2014-2020, The Psi4Julia Developers"
__license__   = "BSD-3-Clause"
__date__      = "2020-08-03"
# -

# # Orbital-Optimized Second-Order Moller Plesset Perturbation Theory (OMP2)

# In this tutorial, we will implement the orbital-optimized second-order Moller-Plesset method in the spin orbital notation. The groundwork for working in the spin orbital notation has been laid out in "Introduction to the Spin Orbital Formulation of Post-HF methods" [tutorial](../08_CEPA0_and_CCD/8a_Intro_to_spin_orbital_postHF.ipynb). It is highly recommended to work through that introduction before starting this tutorial.

# ## I. Theoretical Overview

# ### The general orbital optimization procedure
#
# In orbital optimization methods, the energy is minimized with respect to(w.r.t) an orbital rotation parameter $\textbf{X}$ and can be expanded to second-order as:
#
# \begin{equation}
# E(\textbf{X}) = E(\textbf{X}) + \textbf{X}^\dagger \textbf{w} + \frac{1}{2}\textbf{X}^\dagger\textbf{A}\textbf{X}
# \end{equation}
#
# Here, $\textbf{w}$ is the orbital gradient (derivative of E w.r.t. $\textbf{X}^\dagger$ evaluated at zero and $\textbf{A}$ is the orbital Hessian matrix (second derivative of E w.r.t. $\textbf{X}^\dagger\textbf{X}$ evaluated at zero).
#
# It can be shown that $\textbf{X} = -\textbf{A}^{-1}\textbf{w}$, which gives us the equation used in the Newton-Raphson step of the orbital optimization.
#
# We define the unitary rotation matrix to be $\textbf{U} = exp(\textbf{X}-\textbf{X}^\dagger)$ and use this to rotate the orbitals (using the cofficient matrix).
#
# We then transform the 1 and 2-electron integrals using the new cofficient matrix and evaluate the energy.
#
# This process is repeated until the energy convergence satisfies a specified convergence parameter.
#
# A detailed algorithm for OMP2 is provided in the implementation section.
#

# ### A note on the MP2 amplitude equation
#
# The MP2 amplitude equation can be explicitly written as
#
# \begin{equation}
#  t_{ab}^{ij} = (\mathcal{E}_{ab}^{ij})^{-1} \left(
#      \bar{g}_{ab}^{ij} + P_{(a/b)}f'{}_{a}^{c} t_{cb}^{ij} -
#      P^{(i/j)}f'{}_k^it_{ab}^{kj} \right)
# \end{equation}
#
# where f' is the off-digonal Fock matrix.
#
# Indices p, q, r... are used to indicate arbitrary orbitals, indices a, b, c... are used to indicate virtual orbitals, and indices i, j, k... are used to indicate occupied orbitals.
#
# In conventional MP2, the use canonical orbitals result in a diagonal Fock matrix and the last two terms of the t amplitude equation goes to zero. In OMP2, however, the orbitals are no longer canonical due to orbital rotation, and we have to include these terms in the equation.
#

# ## II. Implementation

# As with previous tutorials, let's begin by importing Psi4, NumPy, TensorOperations, and LinearAlgebra and setting memory and output file options.
# Note that we will also be importing SciPy, which is another library that builds on NumPy and has additional capabilities that we will use.

# +
# ==> Import Psi4, NumPy, & TensorOperations <==
using PyCall: pyimport
psi4 = pyimport("psi4")
np   = pyimport("numpy") # used only to cast to Psi4 arrays
using TensorOperations: @tensor
import LinearAlgebra
eye(n) = LinearAlgebra.I(n)
using Formatting: printfmt

# ==> Set Basic Psi4 Options <==

# Memory specifications
psi4.set_memory(Int(2e9))
numpy_memory = 2

# Output options
psi4.core.set_output_file("output.dat", false)
# -

# We now define the molecule and set Psi4 options:

# +
mol = psi4.geometry("""
0 1
O
H 1 1.1
H 1 1.1 2 104
symmetry c1
""")

psi4.set_options(Dict("basis"         => "6-31g",
                      "scf_type"      => "pk",
                      "reference"     => "rhf",
                      "mp2_type"      => "conv",
                      "e_convergence" => 1e-8,
                      "d_convergence" => 1e-8))
# -

# We use Psi4 to compute the RHF energy and wavefunction and store them in variables `scf_e` and `scf_wfn`. We also check the memory requirements for computation:

# +
# Get the SCF wavefunction & energies
scf_e, scf_wfn = psi4.energy("scf", return_wfn=true)

# ==> Nuclear Repulsion Energy <==
E_nuc = mol.nuclear_repulsion_energy()

# Check memory requirements
nmo = scf_wfn.nmo()
I_size = nmo^4 * 8e-9
println("\nSize of the ERI tensor will be $I_size GB.")
memory_footprint = I_size * 1.5
if I_size > numpy_memory
    psi4.core.clean()
    throw(OutOfMemoryError("Estimated memory utilization ($memory_footprint GB) exceeds " *
                           "allotted memory limit of $numpy_memory GB."))
end
# -

# We need to set the maximum number of iterations for the OMP2 code as well as the energy convergence criteria:

# ==> Set default program options <==
# Maximum OMP2 iterations
MAXITER = 40
# Energy convergence criterion
E_conv = 1.0e-8

# We first obtain orbital information from our wavefunction. We also create an instance of MintsHelper to help build our molecular integrals:

# +
# Create instance of MintsHelper class
mints = psi4.core.MintsHelper(scf_wfn.basisset())

# Get basis and orbital information
nbf = mints.nbf()          # Number of basis functions
nalpha = scf_wfn.nalpha()  # Number of alpha electrons
nbeta = scf_wfn.nbeta()    # Number of beta electrons
nocc = nalpha + nbeta      # Total number of electrons
nso = 2nbf                 # Total number of spin orbitals
nvirt = nso - nocc         # Number of virtual orbitals
# -

# We now build our 2-electron integral, a 4D tensor, in the spin orbital formulation. We also convert it into physicist's notation and antisymmetrize for easier manipulation of the tensor later on.

# +
"""
Spin blocks 2-electron integrals
Using np.kron, we project I and I tranpose into the space of the 2x2 ide
The result is our 2-electron integral tensor in spin orbital notation
"""
function spin_block_tei(I)
    identity = eye(2)
    I = np.kron(identity, I)
    np.kron(identity, permutedims(I, reverse(1:4)))
end

I = np.asarray(mints.ao_eri())
I_spinblock = spin_block_tei(I)

# Convert chemist's notation to physicist's notation, and antisymmetrize
# (pq|rs) ---> ⟨pr|qs⟩
# ⟨pr||qs⟩ = ⟨pr|qs⟩ - ⟨pr|sq⟩
gao = permutedims(I_spinblock, (1, 3, 2, 4)) - permutedims(I_spinblock, (1, 3, 4, 2));
# -

# We get the core Hamiltonian from the reference wavefunction and build it in the spin orbital formulation. The Julia function `kron` is used to project the core Hamiltonian into the space of a 2x2 identity matrix. Note that `np.kron` was used for spin-blocking the 2-electron integral. In the current case, `kron` is only called once because the core Hamltonian is a 2D matrix.

# +
# ==> core Hamiltoniam <==

h = np.asarray(scf_wfn.H())

# Using np.kron, we project h into the space of the 2x2 identity
# The result is the core Hamiltonian in the spin orbital formulation
hao = kron(eye(2), h);
# -

# We get the orbital energies from alpha and beta electrons and append them together. We spin-block the coefficients obtained from the reference wavefunction and convert them into Julia arrays. There is a set corresponding to coefficients from alpha electrons and a set of coefficients from beta electrons. We then sort them according to the order of the orbital energies using `sortperm()`:

# +
# Get orbital energies, cast into Julia array, and extend eigenvalues
eps_a = np.asarray(scf_wfn.epsilon_a())
eps_b = np.asarray(scf_wfn.epsilon_b())
eps = vcat(eps_a, eps_b)

# Get coefficients, block, and sort
Ca = np.asarray(scf_wfn.Ca())
Cb = np.asarray(scf_wfn.Cb())
C = [Ca zero(Ca); zero(Cb) Cb]; # direct sum

# Sort the columns of C according to the order of orbital energies
C = C[:, sortperm(eps)];
# -

# We now define two functions that will transform the core Hamiltonian and the 2-electron integral from the AO basis into the MO basis using the coefficients:
#
# \begin{align}
# h_p^q &= \sum_{\mu \nu} C_{\mu p}^* h_{\mu \nu} C_{\nu q} \\
# \bar{g}_{pq}^{rs} &= \sum_{\mu \nu \rho \sigma}
#                 C_{\mu p}^* C_{\nu q}^* \langle \mu \nu || \rho \sigma \rangle C_{\rho r}C_{\sigma s}
# \end{align}
#
# Note that we transform the core Hamiltonian twice because it has two dimensions. We use these functions to transform the `hao` and `gao` previously defined:

# +
# ==> AO to MO transformation functions <==

"""
Transform hao, which is the core Hamiltonian in the spin orbital basis,
into the MO basis using MO coefficients
"""
function ao_to_mo(hao, C)

    @tensor begin
       hmo[P,Q] := hao[p,Q] * C[p,P]
       hmo[p,Q] := hmo[p,q] * C[q,Q]
    end
end

"""
Transform gao, which is the spin-blocked 4d array of physicist's notation,
antisymmetric two-electron integrals, into the MO basis using MO coefficients
"""
function ao_to_mo_tei(gao, C)

   @tensor begin
      gmo[P,Q,R,S] := gao[p,Q,R,S] * C[p,P]
      gmo[p,Q,R,S] := gmo[p,q,R,S] * C[q,Q]
      gmo[p,q,R,S] := gmo[p,q,r,S] * C[r,R]
      gmo[p,q,r,S] := gmo[p,q,r,s] * C[s,S]
   end
end

# Transform gao and hao into MO basis
hmo = ao_to_mo(hao, C)
gmo = ao_to_mo_tei(gao, C);
# -

# Here we define slices corresponding to the number and position of occupied and virtual indices. We will use these later in the code to access occupied and virtual blocks of relevant arrays. For example, to get $\bar{g}_{ab}^{ij}$, we call:
# ~~~julia
# gmo[v, v, o, o]
# ~~~

# Make slices
x = [CartesianIndex()]
o = [p ≤ nocc for p in 1:nso]
v = [p > nocc for p in 1:nso];

# **OMP2 iteration algorithm:**
#
# 1. Build the fock matrix
#
#    \begin{equation}
#    f_p^q = h_p^q +\bar{g}_{pi}^{qi}
#    \end{equation}
#
# 2. Build the off-diagonal Fock matrix and the orbital energies, where off-diagonal Fock matrix(`fprime`) is just the Fock matrix with its diagonal elements set to zero, and the orbital energies (`eps`) are just the diagonal elements of the Fock matrix
#
#     \begin{equation}
#     \epsilon_p = f_p^p
#     \end{equation}
#
#     \begin{equation}
#     f'{}_p^q =(1 - \delta_p^q)f_p^q
#     \end{equation}
#
# 3. Update the amplitudes (`t_amp`)
#
#     \begin{equation}
#      t_{ab}^{ij} = (\mathcal{E}_{ab}^{ij})^{-1} \left(
#      \bar{g}_{ab}^{ij} + P_{(a/b)}f'{}_{a}^{c} t_{cb}^{ij} -
#      P^{(i/j)}f'{}_k^it_{ab}^{kj} \right)
#     \end{equation}
#
#    Here, P is a permutation operator that permutes the indices indicated. For example, $P_{(a/b)}$ would give all    possible permutations of a and b. Thus,
#
#     \begin{equation}
#     P_{(a/b)}f'{}_{a}^{c} t_{cb}^{ij} = f'{}_{a}^{c} t_{cb}^{ij} - f'{}_{b}^{c} t_{ca}^{ij}
#     \end{equation}
#
#    where the minus sign arises as a result of antisymmetric properties due to the interchange of the two indices
#    The amplitudes terms in the code are assigned as `t1`, `t2`, and `t3`, respectively.
#
#    To take in account the permutation terms, we evaluate the term and then transpose the relevant indices.
#    For example, for the second term in the amplitude equation we first evaluate it as it:
#    ~~~julia
#    @tensor t2[a,b,i,j] := (fprime[v, v])[a,c] * t_amp[c,b,i,j]
#    ~~~
#    Then, to account for the permutation, we transpose the two dimensions corresponding to the permuted indices. Since    a and b are in the first two dimensions of `t2`, we switch 0 and 1:
#    ~~~julia
#    t2 = t2 - permutedims(t2, (2, 1, 3, 4))
#    ~~~
# 4. Build the one-particle density matrix (`opdm`)
#
#     \begin{equation}
#     \gamma_q^p = \tilde{\gamma}_q^p + \mathring{\gamma}_q^p
#     \end{equation}
#
#    The one-particle density matrix(opdm) is a sum of the reference opdm ($\mathring{\gamma}_q^p$) and a correlation opdm ($\tilde{\gamma}_q^p$).
#
#     $\mathring{\gamma}_q^p$ is assigned as the variable `odm_ref` and defined as:
#      \begin{align}
#      & \, \delta^i_j \, \text{for $p=i$, $q=j$}, \\
#      & 0 \,  \text{otherwise}
#     \end{align}
#
#     The virtual block of $\tilde{\gamma}_q^p$ (assigned as `odm_corr`) is defined as:
#     \begin{equation}
#     \tilde{\gamma}_b^a  = \frac{1}{2} t_{ij}^{ab*}t_{bc}^{ij}
#     \end{equation}
#
#      The occupied block of $\tilde{\gamma}_q^p$ is defined as:
#     \begin{equation}
#     \tilde{\gamma}_j^i  = -\frac{1}{2} t_{jk}^{ab*}t_{ab}^{ik}
#     \end{equation}
#
#     As seen before, we used our defined slices to pick out these specific blocks:
#     ~~~julia
#     @tensor (opdm_corr[v, v])[b,a] :=  0.5(permutedims(t_amp, reverse(1:4)))[i,j,a,c] * t_amp[b,c,i,j]
#     @tensor (opdm_corr[o, o])[j,i] := -0.5(permutedims(t_amp, reverse(1:4)))[j,k,a,b] * t_amp[a,b,i,k]
#     ~~~
#
# 5. Build the two-particle density matrix (`tpdm`)
#
#     \begin{equation}
#     \Gamma_{rs}^{pq} = \tilde{\Gamma}_{rs}^{pq} + P_{(r/s)}^{(p/q)}\tilde{\gamma}_r^p\mathring{\gamma}_s^q
#     +P_{(r/s)}\mathring{\gamma}_r^p\mathring{\gamma}_s^q
#     \end{equation}
#
#      where as before, P is the permutation operator
#
#  $\tilde{\Gamma}_{rs}^{pq}$ (`tdm_corr`) can be separated into two components:
#
#  \begin{align}
#  \tilde{\Gamma}_{ij}^{ab} = & t_{ij}^{ab*}\\
#  \tilde{\Gamma}_{ab}^{ij} = & t_{ab}^{ij}
#  \end{align}
#
# 6. Compute the Newton-Raphson step
#
#    First, form a generalized-Fock matrix using the one and two particle density matrices. This will be used to form the MO gradient matrix needed for the rotation matrix:
#
#    \begin{equation}
#    (\textbf{F})_p^q \equiv h_p^r \gamma_r^q + \frac{1}{2} \bar{g}_{pr}^{st}\Gamma_{st}^{qr}
#    \end{equation}
#
#    We have seen in the theoretical overview that the X matrix while paramtetrizes the orbital rotations can be expressed in terms of the orbital gradient matrix and orbital Hessian matrix. It can be shown that the individual elements of X can be computed by:
#
#    \begin{equation}
#     x_a^i = \frac{(\textbf{F} - \textbf{F}^\dagger)_a^i}{\epsilon_i - \epsilon_a}
#     \end{equation}
#
#      Here we only consider rotations between the occupied and virtual orbitals, since rotations within each block are redudant since energy is invariant to rotations within those spaces.
#
#      Rather than computing individual elements we can compute the whole virtual-occupied block:
#
#     \begin{equation}
#     \textbf{X}_v^o = (\textbf{F} - \textbf{F}^\dagger)_v^o (\mathcal{E}_v^o)^{-1}
#     \end{equation}
#     Translating this to code, this becomes:
#     ~~~julia
#     X[v, o] = ((F - F')[v, o]) ./ (-eps[v, x] .+ eps[x, o])
#     ~~~
# 7. We can now build the Newton-Raphson orbital rotation matrix from $\textbf{X}$:
#
#     \begin{equation}
#     \textbf{U} = exp(\textbf{X} - \textbf{X}^\dagger)
#     \end{equation}
#
# 8. Use the rotation matrix to rotate the MO coefficients
#    \begin{equation}
#    \textbf{C} \leftarrow \textbf{CU}
#    \end{equation}
#
# 9. Transform the 1-electron (`hmo`) and 2-electron (`gmo`) integrals to the MO basis using the new coefficient matrix. We can use our previously defined transformation functions for this step.
#
#     \begin{align}
#     h_p^q &= \sum_{\mu \nu} C_{\mu p}^* h_{\mu \nu} C_{\nu q} \\
#     \bar{g}_{pq}^{rs} &= \sum_{\mu \nu \rho \sigma}
#     C_{\mu p}^* C_{\nu q}^* \langle \mu \nu || \rho \sigma \rangle C_{\rho r}C_{\sigma s}
#     \end{align}
# 10. Evaluate the energy (`E_OMP2`)
#     \begin{equation}
#     E = h_p^q \gamma_q^p + \frac{1}{4} \bar{g}_{pq}^{rs}\Gamma_{rs}^{pq}
#     \end{equation}
#
# 11. If the energy is converged according to the convergence criterion defined above, quit. Otherwise, loop over the algorithm again.

# Before beginning the iterations, we initialize OMP2 energy and the t amplitudes $t_{ab}^{ij}$ (`t_amp`) to be zero. We also initialize the correlation and reference one-particle density matrix and the correlation two-particle density matrix. Finally we intialize `X`, which is the parameter used to optimize our orbitals in the Newton-Raphson step.
#

# +
# Intialize t amplitude and energy
t_amp = zeros(nvirt, nvirt, nocc, nocc)
E_OMP2_old = 0.0

# Initialize the correlation one particle density matrix
opdm_corr = zeros(nso, nso)

# Build the reference one particle density matrix
opdm_ref = zeros(nso, nso)
opdm_ref[o, o] = eye(nocc)

# Initialize two particle density matrix
tpdm_corr = zeros(nso, nso, nso, nso)

# Initialize the rotation matrix parameter
E_OMP2 = let hmo=hmo, gmo=gmo, tpdm_corr=tpdm_corr, opdm_corr=opdm_corr, opdm_ref=opdm_ref,
   E_OMP2_old = E_OMP2_old, t_amp=t_amp, C=C

   E_OMP2 = 0.0
   X = zeros(nso, nso)

   for iteration in 1:MAXITER

       # Build the Fock matrix
       @tensor f[p,q] := hmo[p,q] + (gmo[:, o, :, o])[p,i,q,i]

       # Build off-diagonal Fock Matrix and orbital energies
       fprime = copy(f)
       fprime[LinearAlgebra.diagind(fprime)] .= 0
       eps = LinearAlgebra.diag(f)

       # Update t amplitudes
       t1 = @view gmo[v, v, o, o]
       @tensor t2[a,b,i,j] := (fprime[v, v])[a,c] * t_amp[c,b,i,j]
       @tensor t3[a,b,i,j] := (fprime[o, o])[k,i] * t_amp[a,b,k,j]
       t_amp = t1 .+ t2 .- permutedims(t2, (2, 1, 3, 4)) .-
               t3 .+ permutedims(t3, (1, 2, 4, 3))

       # Divide by a 4D tensor of orbital energies
       @. t_amp /= (- eps[v, x, x, x] - eps[x, v, x, x] +
                      eps[x, x, o, x] + eps[x, x, x, o])

       # Build one particle density matrix
       @tensor (opdm_corr[v, v])[b,a] :=  0.5(permutedims(t_amp, reverse(1:4)))[i,j,a,c] * t_amp[b,c,i,j]
       @tensor (opdm_corr[o, o])[j,i] := -0.5(permutedims(t_amp, reverse(1:4)))[j,k,a,b] * t_amp[a,b,i,k]
       opdm = opdm_corr + opdm_ref

       # Build two particle density matrix
       tpdm_corr[v, v, o, o] = t_amp
       tpdm_corr[o, o, v, v] = permutedims(t_amp, reverse(1:4))
       @tensor tpdm2[r,s,p,q] := opdm_corr[r,p] * opdm_ref[s,q]
       @tensor tpdm3[r,s,p,q] := opdm_ref[r,p] * opdm_ref[s,q]
       tpdm = tpdm_corr +
           tpdm2 - permutedims(tpdm2, (2, 1, 3, 4)) -
           permutedims(tpdm2, (1, 2, 4, 3)) + permutedims(tpdm2, (2, 1, 4, 3)) +
           tpdm3 - permutedims(tpdm3, (2, 1, 3, 4))

       # Newton-Raphson step
       @tensor F[p,q] := hmo[p,r] * opdm[r,q] + 0.5gmo[p,r,s,t] * tpdm[s,t,q,r]
       X[v, o] = ((F - F')[v, o]) ./ (- eps[v, x] .+ eps[x, o])

       # Build Newton-Raphson orbital rotation matrix
       U = exp(X - X')

       # Rotate spin-orbital coefficients
       C = C * U

       # Transform one and two electron integrals using new C
       hmo = ao_to_mo(hao, C)
       gmo = ao_to_mo_tei(gao, C)

       # Compute the energy
       E_OMP2 = E_nuc + @tensor scalar(hmo[p,q] * opdm[q,p]) +
                  1/4 * @tensor scalar(gmo[p,q,r,s] * tpdm[r,s,p,q])
       printfmt("OMP2 iteration: {1:3d} Energy: {2:15.8f} dE: {3:2.5e}\n", iteration, E_OMP2, E_OMP2-E_OMP2_old)

       abs(E_OMP2-E_OMP2_old) < E_conv && break

       # Updating values
       E_OMP2_old = E_OMP2
   end
   E_OMP2
end
# -

# We compare the final energy with Psi4's OMP2 energy:

psi4.compare_values(psi4.energy("omp2"), E_OMP2, 6, "OMP2 Energy")

# ## References
#
# 1. Background paper:
#     >"Quadratically convergent algorithm for orbital optimization in the orbital-optimized
# coupled-cluster doubles method and in orbital-optimized second-order Møller-Plesset
# perturbation theory"[[Bozkaya:2011:135](http://aip.scitation.org/doi/10.1063/1.3631129)] U. Bozkaya, J. M. Turney, Y. Yamaguchi, H. F. Schaefer III, and C. D. Sherrill, *J. Chem. Phys.* **135**, 104103 (2011).
#
# 2. Useful notes on orbital rotation:
# 	> A. V. Copan, "Orbital Relaxation" accessed with https://github.com/CCQC/chem-8950/tree/master/2017/.
#
# 3. Algorithms from:
# 	> A. V. Copan, "OMP2" accessed with https://github.com/CCQC/chem-8950/tree/master/2017/programming.


text:
function get_fields
end

function get_records
end

function build_result
end

# interface for `Dict`s
get_fields(df :: Dict{Symbol, Vector{T} where T}) = collect(keys(df))
get_records(df :: Dict{Symbol, Vector{T} where T}) = zip((df[key] for key in keys(df))...)
function build_result(::Type{Dict{Symbol, Vector{T} where T}}, fields, typs, source :: Base.Generator)
    res = Tuple(typ[] for typ in typs)
    for each in source
        push!.(res, each)
    end
    Dict(zip(fields, res))
end

ext: jl


text:
module Retry

__precompile__(true)

export @repeat, @protected, efield, ecode

efield(x, f, default=nothing) = f in fieldnames(x) ? getfield(x, f) : default
ecode(x) = efield(x, :code)

include("repeat_try.jl")
include("protected_try.jl")

end # module


text:
"""
Wouldn't it be nice if we had a selection of integrators, like Thomas and I
discussed all those months ago!
"""

"""
Compute the piecewise constant propagator for a single time slice of duration
    dt. The function makes certain assumptions about the input arguments as described
    below.

    In order to make the multiplication easy the Hctrl array should be of the form:
    Hctrl = [sx, sy]
    and the control array should be an array of the corresponding length.

    The return type of this function will be the same type as the static Hamiltonian provided.
    We also assume that internally that everything maintains the same dimensions
"""
function pw_integrate(H0::T, Hctrl, control, dt)::T where T
    H::T = H0
    for i = 1:length(control)
        H += (Hctrl[i] .* control[i])
    end
    exp((-im*dt) .* H)
end

"""
This simply computes the full time evolution and returns the final propagator which
can be useful for GRAPE
"""
function pw_full_evolution(H0::T, Hctrl, control, dt, init::T)::T where T
    for i = 1:size(control)[2]
        init = pw_integrate(H0, Hctrl, control[:, i], dt) * init
    end
    init
end

# we should write this to use concrete_solve so it'll work with Zygote in future
function ode_integrate()

end

# include an example of a master equation solver
function me_integrate()

end


text:
# output = wx
immutable Linear{T,P<:Params} <: Learnable
    nin::Int
    nout::Int
    input::SumNode{T,1}
    output::OutputNode{T,1}
    params::P
end

function Linear{T}(::Type{T}, nin::Int, nout::Int,
                    θ::AbstractVector = zeros(T, nout*nin),
                    ∇::AbstractVector = zeros(T, nout*nin))
    input = InputNode(T, nin)
    output = OutputNode(T, nout)
    params = Params(θ, ∇, ((nout,nin),))
    w = params.views[1]
    initialize_weights!(w)
    Linear(nin, nout, input, output, params)
end
Linear(nin::Int, nout::Int, args...) = Linear(Float64, nin, nout, args...)

function Base.show(io::IO, t::Linear)
    print(io, "Linear{$(t.nin)-->$(t.nout)}")
end

function transform!(t::Linear)
    w = t.params.views[1]
    x = t.input.val
    y = t.output.val
    for o=1:t.nout
        y[o] = sum(w[o,i] * x[i] for i=1:t.nin)
    end
    y
end

function grad!{T}(t::Linear{T})
    w = t.params.views[1]
    ∇w = t.params.∇_views[1]
    x = t.input.val
    ∇x = t.input.∇
    ∇y = t.output.∇

    # ∇x, ∇w
    for i=1:t.nin
        ∇xᵢ = zero(T)
        for o=1:t.nout
            ∇xᵢ += w[o,i] * ∇y[o]
            ∇w[o,i] = x[i] * ∇y[o]
        end
        ∇x[i] = ∇xᵢ
    end
    return
end


text:
"""
    getPhenotype(tree::CGPGenotype, gpExperimentInfo::CGPInfo)

Obtains the phenotype of a given CGP genotype as an evaluable expresion (See [`Expr`](@ref)).

# Arguments
- `genotype::CGPGenotype`: genotype of the individual from which the phenotype is
    wanted.

# Returns
The phenotype of the individual as `Expr`.

See also: [`evaluate`](@ref)
"""
function getPhenotype(genotype::CGPGenotype)
    phenotype = Array{String}(undef, 0)
    representation = getRepresentation(genotype)
    visitedFromRoot = Array{Integer}(undef, 0, 2)

    nodeIndex = 1
    parent = 0

    while nodeIndex != 0

        if nodeIndex > size(visitedFromRoot)[1]
            visitedFromRoot = vcat(visitedFromRoot, zeros(Integer, 1, 2))
            visitedFromRoot[end, 2] = parent

            if typeof(representation[nodeIndex]) <: TerminalNode

                if typeof(representation[nodeIndex]) <: VariableNode
                    push!(phenotype, getVariableName(representation[nodeIndex]) * ",")

                elseif typeof(representation[nodeIndex]) <: ConstantNode
                    push!(phenotype, string(eval(representation[nodeIndex])) * ",")

                elseif typeof(representation[nodeIndex]) <: NoArgsFunctionNode
                    push!(phenotype, getNoArgsFunctionName(representation[nodeIndex]) * "(),")
                end

            else typeof(representation[nodeIndex]) <: FunctionNode
                push!(phenotype, getFunctionName(representation[nodeIndex]) * "(")
            end
        end

        if visitedFromRoot[nodeIndex, 1] < getArity(representation[nodeIndex])
            parent = nodeIndex
            nodeIndex = size(visitedFromRoot)[1] + 1
        else
            if typeof(representation[nodeIndex]) <: FunctionNode
                push!(phenotype, "),")
            end
            nodeIndex = visitedFromRoot[nodeIndex, 2]
            if nodeIndex != 0
                visitedFromRoot[nodeIndex, 1] += 1
            end
        end
    end

    phenotype[end] = phenotype[end][1:end-1]

    if length(phenotype) == 1
        phenotype = Meta.parse(phenotype[1])
        return quote $phenotype end
    else
        return Meta.parse(join(phenotype))
    end
end # function

max_stars_repo_name: SergioGmezM/GeneticJulia

max_issues_repo_name: SergioGmezM/GeneticJulia

max_forks_repo_name: SergioGmezM/GeneticJulia


text:
# ---
# title: 513. Find Bottom Left Tree Value
# id: problem513
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Tree, Depth-first Search, Breadth-first Search
# link: <https://leetcode.com/problems/find-bottom-left-tree-value/description/>
# hidden: true
# ---
#
# Given a binary tree, find the leftmost value in the last row of the tree.
#
# **Example 1:**
#
#
#
#     Input:
#
#         2
#        / \
#       1   3
#
#     Output:
#     1
#
#
# **Example 2:**
#
#
#
#     Input:
#
#             1
#            / \
#           2   3
#          /   / \
#         4   5   6
#            /
#           7
#
#     Output:
#     7
#
#
# **Note:** You may assume the tree (i.e., the given root node) is not **NULL**.
#
#
## @lc code=start
using LeetCode

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


text:
# ---
# title: 1218. Longest Arithmetic Subsequence of Given Difference
# id: problem1218
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Math, Dynamic Programming
# link: <https://leetcode.com/problems/longest-arithmetic-subsequence-of-given-difference/description/>
# hidden: true
# ---
#
# Given an integer array `arr` and an integer `difference`, return the length of
# the longest subsequence in `arr` which is an arithmetic sequence such that the
# difference between adjacent elements in the subsequence equals `difference`.
#
#
#
# **Example 1:**
#
#
#
#     Input: arr = [1,2,3,4], difference = 1
#     Output: 4
#     Explanation: The longest arithmetic subsequence is [1,2,3,4].
#
# **Example 2:**
#
#
#
#     Input: arr = [1,3,5,7], difference = 1
#     Output: 1
#     Explanation: The longest arithmetic subsequence is any single element.
#
#
# **Example 3:**
#
#
#
#     Input: arr = [1,5,7,8,5,3,4,2,1], difference = -2
#     Output: 4
#     Explanation: The longest arithmetic subsequence is [7,5,3,1].
#
#
#
#
# **Constraints:**
#
#   * `1 <= arr.length <= 10^5`
#   * `-10^4 <= arr[i], difference <= 10^4`
#
#
## @lc code=start
using LeetCode

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


text:
mutable struct DAQTask
    "Number of frames read"
    nread::Int
    "Is the device reading frames?"
    isreading::Bool
    "Stop data acquisition?"
    stop::Bool
    "Are we using threads?"
    thrd::Bool
    "Initial time, end time (ns) and number of frames"
    timing::NTuple{3, UInt64}
    "`Task` object executing the data acquisition"
    task::Task
    """
    `DAQTasq`

    Creates a structure that handles asynchronous data acquisition.

    This structure stores the number of samples read and the general state of data
    acquisition. It also provides timing measurements so that sampling frequency
    can estimated.

    """
    DAQTask() = new(0,false,false,false, (UInt64(0),UInt64(0),UInt64(0)), Task(()->0))
end

"""
`isreading(tsk)`

Is the daq device currently acquiring data?

See [`samplesread`](@ref) to see the number of samples already read.
"""
isreading(task::DAQTask) = task.isreading

samplesread(task::DAQTask) = task.nread

"""
`cleartask!(task)`

Clear the buffer. Set the DAQTask parameters to null state.
"""
function cleartask!(task::DAQTask)
    isreading(task) && error("Can not clear a task while it is reading!")

    task.isreading = false # Let's make sure...
    task.stop = false
    task.thrd = false
    task.nread = 0
    return
end

setdaqthread!(task::DAQTask, thrdstatus=false) = task.thrd=thrdstatus
daqthread(task::DAQTask) = task.thrd

setdaqtask!(task::DAQTask, jtsk::Task) = task.task = jtsk
daqtask(task::DAQTask) = task.task

stoptask(task::DAQTask) = task.stop
stoptask!(task::DAQTask, s=true) = task.stop = s

"""
    `samplingfreq(task)`

    Returns the measured sampling frequency achieved during data acquisition

"""
samplingfreq(task::DAQTask) = task.timing[3] / (1e-9 * (task.timing[2] - task.timing[1]))

"""
`settiming!(task, t1, t2, n)`

Updates timing information on current data acquisition

 * `task` The `DAQTask` object
 * `t1` Initial time of data acquisition
 * `t2` Last time of data acquisition
 * `n` Number of samples read between `t1` and `t2`

"""
settiming!(task, t1, t2, n) = task.timing = (UInt64(t1), UInt64(t2), UInt64(n))


text:
@doc raw"""
WIP
"""
function retrieve_surface(algorithm::Union{DeterministicCustom,Deterministic}, starting_shape::AbstractSyntheticShape, img::AbstractArray, integration_factor::Real = 10, smoothness::Real = 20; noise_val::Real = 0, radius=50)
    # Generate initial gradient and dipatch on desired function.
    p,q = synthetic_gradient(starting_shape, radius = radius, noise = noise_val)
    p,q = retrieve_surface(algorithm, img, p, q, integration_factor, smoothness)
    return p,q
end

function retrieve_surface(algorithm::DeterministicCustom, img::AbstractArray, pin::AbstractArray, qin::AbstractArray, integration_factor::Real = 10, smoothness::Real = 20)
    # Setup initial variables
    E = Array{Float64}(img)
    Eₘ = maximum(E)
    p = copy(pin)
    q = copy(qin)
    N, M = size(E)
    λᵢ = integration_factor
    λₛ = smoothness
    R = zeros(Float64,axes(E))
    Z = zeros(axes(E))
    Δϵ = zeros(Float64,(size(E)...,2))
    val = sqrt(2 * length(E))*2 + 10
    prevd = 0.0001

    # Loop untill norm of Δϵ is less then desired value
    while val > sqrt(2 * length(E))

        # Calculate the energy
        ϵ = 0.0
        for i = 1:(M-1)
            for j = 1:(N-1)
                ϵ += (Eₘ/sqrt(1+p[j,i]^2+q[j,i]^2)-E[j,i])^2 + λᵢ*((p[j,i+1]-p[j,i])-(q[j+1,i]-q[j,i]))^2 + λₛ*((p[j,i+1]-p[j,i])^2 + (p[j+1,i]-p[j,i])^2 + (q[j,i+1]-q[j,i])^2 + (q[j+1,i]-q[j,i])^2)
                if i>1 && j>1
                    Δϵ[j,i,1] = -2 * Eₘ * p[j,i] * (Eₘ/sqrt(1+p[j,i]^2+q[j,i]^2)-E[j,i]) * (1/(1+p[j,i]^2+q[j,i]^2)^(3/2)) - 2*λᵢ*((p[j+1,i]-p[j,i])-(q[j,i+1]-q[j,i])) - 2*λₛ*((p[j,i+1]-p[j,i]) + (p[j+1,i]-p[j,i])) + 2*λᵢ*((p[j,i]-p[j-1,i])-(q[j-1,i+1]-q[j-1,i])) + 2*λₛ*((p[j,i]-p[j,i-1]) + (p[j,i]-p[j-1,i]))
                    Δϵ[j,i,2] = -2 * Eₘ * q[j,i] * (Eₘ/sqrt(1+p[j,i]^2+q[j,i]^2)-E[j,i]) * (1/(1+p[j,i]^2+q[j,i]^2)^(3/2)) + 2*λᵢ*((p[j+1,i]-p[j,i])-(q[j,i+1]-q[j,i])) - 2*λₛ*((q[j,i+1]-q[j,i]) + (q[j+1,i]-q[j,i])) - 2*λᵢ*((p[j+1,i-1]-p[j,i-1])-(q[j,i]-q[j,i-1])) + 2*λₛ*((q[j,i]-q[j,i-1]) + (q[j,i]-q[j-1,i]))
                elseif i>1
                    Δϵ[j,i,1] = -2 * Eₘ * p[j,i] * (Eₘ/sqrt(1+p[j,i]^2+q[j,i]^2)-E[j,i]) * (1/(1+p[j,i]^2+q[j,i]^2)^(3/2)) - 2*λᵢ*((p[j+1,i]-p[j,i])-(q[j,i+1]-q[j,i])) - 2*λₛ*((p[j,i+1]-p[j,i]) + (p[j+1,i]-p[j,i])) + 2*λₛ*(p[j,i]-p[j,i-1])
                    Δϵ[j,i,2] = -2 * Eₘ * q[j,i] * (Eₘ/sqrt(1+p[j,i]^2+q[j,i]^2)-E[j,i]) * (1/(1+p[j,i]^2+q[j,i]^2)^(3/2)) + 2*λᵢ*((p[j+1,i]-p[j,i])-(q[j,i+1]-q[j,i])) - 2*λₛ*((q[j,i+1]-q[j,i]) + (q[j+1,i]-q[j,i])) - 2*λᵢ*((p[j+1,i-1]-p[j,i-1])-(q[j,i]-q[j,i-1])) + 2*λₛ*(q[j,i]-q[j,i-1])
                elseif j>1
                    Δϵ[j,i,1] = -2 * Eₘ * p[j,i] * (Eₘ/sqrt(1+p[j,i]^2+q[j,i]^2)-E[j,i]) * (1/(1+p[j,i]^2+q[j,i]^2)^(3/2)) - 2*λᵢ*((p[j+1,i]-p[j,i])-(q[j,i+1]-q[j,i])) - 2*λₛ*((p[j,i+1]-p[j,i]) + (p[j+1,i]-p[j,i])) + 2*λᵢ*((p[j,i]-p[j-1,i])-(q[j-1,i+1]-q[j-1,i])) +2*λₛ*(p[j,i]-p[j-1,i])
                    Δϵ[j,i,2] = -2 * Eₘ * q[j,i] * (Eₘ/sqrt(1+p[j,i]^2+q[j,i]^2)-E[j,i]) * (1/(1+p[j,i]^2+q[j,i]^2)^(3/2)) + 2*λᵢ*((p[j+1,i]-p[j,i])-(q[j,i+1]-q[j,i])) - 2*λₛ*((q[j,i+1]-q[j,i]) + (q[j+1,i]-q[j,i])) + 2*λᵢ*(q[j,i]-q[j-1,i])
                else
                    Δϵ[j,i,1] = -2 * Eₘ * p[j,i] * (Eₘ/sqrt(1+p[j,i]^2+q[j,i]^2)-E[j,i]) * (1/(1+p[j,i]^2+q[j,i]^2)^(3/2)) - 2*λᵢ*((p[j+1,i]-p[j,i])-(q[j,i+1]-q[j,i])) - 2*λₛ*((p[j,i+1]-p[j,i]) + (p[j+1,i]-p[j,i]))
                    Δϵ[j,i,2] = -2 * Eₘ * q[j,i] * (Eₘ/sqrt(1+p[j,i]^2+q[j,i]^2)-E[j,i]) * (1/(1+p[j,i]^2+q[j,i]^2)^(3/2)) + 2*λᵢ*((p[j+1,i]-p[j,i])-(q[j,i+1]-q[j,i])) - 2*λₛ*((q[j,i+1]-q[j,i]) + (q[j+1,i]-q[j,i]))
                end
            end
        end

        # Update norm value
        val = norm(Δϵ)

        # Line search to find minimizor of ϵ along the direction of gradient
        dmin = 1
        ϵmin = typemax(Float64)
        for d = prevd/10:prevd/10:prevd*10
            # Setup new gradient field
            ϵtemp = 0.0
            ptemp = copy(p)
            qtemp = copy(q)
            for i = 1:(M-1)
                for j = 1:(N-1)
                    ptemp[j,i] = p[j,i] - d*Δϵ[j,i,1]
                    qtemp[j,i] = q[j,i] - d*Δϵ[j,i,2]
                end
            end

            # Calculate new energy
            for i = 1:(M-1)
                for j = 1:(N-1)
                    ϵtemp += (Eₘ/sqrt(1+(ptemp[j,i])^2+qtemp[j,i]^2)-E[j,i])^2 + λᵢ*((ptemp[j,i+1]-ptemp[j,i])-(qtemp[j+1,i]-qtemp[j,i]))^2 + λₛ*((ptemp[j,i+1]-ptemp[j,i])^2 + (ptemp[j+1,i]-ptemp[j,i])^2 + (qtemp[j,i+1]-qtemp[j,i])^2 + (qtemp[j+1,i]-qtemp[j,i])^2)
                end
            end

            # Check if new energy is less then the minimum found
            if ϵtemp < ϵmin
                ϵmin = ϵtemp
                dmin = d
            end
        end

        # Set gradients to minimize energy
        for i = 1:(M-1)
            for j = 1:(N-1)
                p[j,i] = p[j,i] - dmin*Δϵ[j,i,1]
                q[j,i] = q[j,i] - dmin*Δϵ[j,i,2]
            end
        end
        prevd = dmin
        @show maximum(Δϵ)
        @show norm(Δϵ), ϵ, dmin
    end
    return p,q
end

# Energy function
function f(E, pq, λᵢ, λₛ)
    Eₘ = maximum(E)
    val = 0.0
    N, M = size(E)
    for i = 1:(M-1)
        for j = 1:(N-1)
            val += (Eₘ/sqrt(1+pq[j,i,1]^2+pq[j,i,2]^2)-E[j,i])^2 + λᵢ*((pq[j,i+1,1]-pq[j,i,1])-(pq[j+1,i,2]-pq[j,i,1]))^2 + λₛ*((pq[j,i+1,1]-pq[j,i,1])^2 + (pq[j+1,i,1]-pq[j,i,1])^2 + (pq[j,i+1,2]-pq[j,i,2])^2 + (pq[j+1,i,2]-pq[j,i,2])^2)
        end
    end
    return val
end

function retrieve_surface(algorithm::Deterministic, img::AbstractArray, pin::AbstractArray, qin::AbstractArray, integration_factor::Real = 10, smoothness::Real = 20)
    # Setup initial variables
    E = Float64.(img)
    x₀ = zeros(Float64,first(size(img)),last(size(img)),2)
    x₀[:,:,1] .= pin
    x₀[:,:,2] .= qin

    # Calcualte derivatives
    od = OnceDifferentiable(pq -> f(E,pq,integration_factor,smoothness), x₀; autodiff = :forward);
    a = zeros(size(x₀))

    # Run gradient descent
    res = optimize(od, x₀, GradientDescent(), Optim.Options(show_trace = true, g_tol = sqrt(2*length(E))))

    # Get minimised gradients
    out = Optim.minimizer(res)
    p = out[:,:,1]
    q = out[:,:,2]
    return p,q
end


text:
function horner(x, pu...)
    t = gensym(:t)
    ex = p[end]
    for i ∈ length(p)-1:-1:1
        ex = :(SIMDPirates.vmuladd($t, $ex, $(p[i])))
    end
    Expr(:block, :($t = $x), ex)
end

spfunc(f) = Expr(:(.), :SIMDPirates, QuoteNode(f))
function ret_pirate(ex::Expr, x::Symbol, i::Int)
    f = get(VECTOR_SYMBOLS, x, x)
    f === x || (ex.args[i] = spfunc(f))
    nothing
end
ret_pirate!(::Expr, ::Any, ::Int) = nothing
function ret_pirate!(ex::Expr, x::Expr, i::Int)
    length(x.args) == 0 && return
    f = first(x.args)
    x.args[1] = if f == :(Base.FastMath.add_fast)
        spfunc(:vadd)
    elseif f == :(Base.FastMath.sub_fast)
        spfunc(:vsub)
    elseif f == :(Base.FastMath.mul_fast)
        spfunc(:vmul)
    elseif f == :(Base.FastMath.div_fast)
        spfunc(:vfdiv)
    elseif f == :(Base.FastMath.sqrt)
        spfunc(:vsqrt)
    elseif f == :(Base.Math.muladd)
        spfunc(:vmuladd)
    else
        f
    end
    _pirate!(expr, mod)
    nothing
end

function _pirate!(expr, mod)
    for (i,ex) ∈ enumerate(expr.args)
        ret_pirate!(expr, ex, i)
    end
    esc(expr)
end

macro pirate(ex) _pirate!(contract_pass!(Expr(:(.), mod, QuoteNode(:SIMDPirates))), Symbol(__module__)) end

# struct vBitArray
#     ptr::Ptr{UInt64}
# end
# @inline VectorizationBase.vectorizable(b::BitArray) = vBitArray(pointer(b.chunks))
# @inline function Base.:+(i, b::vBitArray)
#     vBitArray((i >> 3) + b.ptr)
# end
# @inline function Base.:+(b::vBitArray, i)
#     vBitArray((i >> 3) + b.ptr)
# end
# @inline function Base.:-(b::vBitArray, i)
#     vBitArray(b.ptr - (i >> 3))
# end
# """
# The method on BitArrays is a hack, ignoring the type passed to Vec, to better support the LoopVectorization.jl implementation.
# """
# @generated function vload(::Type{Vec{N,T}}, b::vBitArray, i::Integer) where {N,T}
#     N < 8 && throw("Bit array vectors with $N < 8 not yet supported.")
#     utype = mask_type(N)
#     quote
#         $(Expr(:meta, :inline))
#         unsafe_load(Base.unsafe_convert(Ptr{$utype}, b.ptr + (i >> 3)))
#     end
# end
# @generated function vload(::Type{Vec{N,T}}, b::vBitArray) where {N,T}
#     N < 8 && throw("Bit array vectors with $N < 8 not yet supported.")
#     utype = mask_type(N)
#     quote
#         $(Expr(:meta, :inline))
#         unsafe_load(Base.unsafe_convert(Ptr{$utype}, b.ptr))
#     end
# end
# @generated function vload(::Type{SVec{N,T}}, b::vBitArray) where {N,T}
#     N < 8 && throw("Bit array vectors with $N < 8 not yet supported.")
#     utype = mask_type(N)
#     quote
#         $(Expr(:meta, :inline))
#         unsafe_load(Base.unsafe_convert(Ptr{$utype}, b.ptr))
#     end
# end
# @generated function vload(::Type{SVec{N,T}}, b::vBitArray, i) where {N,T}
#     N < 8 && throw("Bit array vectors with $N < 8 not yet supported.")
#     utype = mask_type(N)
#     quote
#         $(Expr(:meta, :inline))
#         unsafe_load(Base.unsafe_convert(Ptr{$utype}, b.ptr + (i >> 3)))
#     end
# end

# """
# Masks on vectorizable bit arrays are currently ignored.
# """
# @generated function vload(::Type{Vec{N,T}}, b::vBitArray, i::Integer, mask::Union{<:Unsigned,Vec{N,Bool}}) where {N,T}
#     N < 8 && throw("Bit array vectors with $N < 8 not yet supported.")
#     utype = mask_type(N)
#     quote
#         $(Expr(:meta, :inline))
#         unsafe_load(Base.unsafe_convert(Ptr{$utype}, b.ptr + (i >> 3)))
#     end
# end
# @generated function vload(::Type{Vec{N,T}}, b::vBitArray, mask::Union{<:Unsigned,Vec{N,Bool}}) where {N,T}
#     N < 8 && throw("Bit array vectors with $N < 8 not yet supported.")
#     utype = mask_type(N)
#     quote
#         $(Expr(:meta, :inline))
#         unsafe_load(Base.unsafe_convert(Ptr{$utype}, b.ptr))
#     end
# end
# @generated function vload(::Type{SVec{N,T}}, b::vBitArray, mask::Union{<:Unsigned,Vec{N,Bool}}) where {N,T}
#     N < 8 && throw("Bit array vectors with $N < 8 not yet supported.")
#     utype = mask_type(N)
#     quote
#         $(Expr(:meta, :inline))
#         unsafe_load(Base.unsafe_convert(Ptr{$utype}, b.ptr))
#     end
# end
# @generated function vload(::Type{SVec{N,T}}, b::vBitArray, i, mask::Union{<:Unsigned,Vec{N,Bool}}) where {N,T}
#     N < 8 && throw("Bit array vectors with $N < 8 not yet supported.")
#     utype = mask_type(N)
#     quote
#         $(Expr(:meta, :inline))
#         unsafe_load(Base.unsafe_convert(Ptr{$utype}, b.ptr + (i >> 3)))
#     end
# end


text:
export principal_generator, kernel_group

################################################################################
#
#  Map Types
#
################################################################################

mutable struct MapRayClassGroupAlg{S, T} <: Map{S, T, HeckeMap, MapRayClassGroupAlg}
  header::MapHeader{S, T}
  modulus#::AlgAssAbsOrdIdl{...}
  groups_in_number_fields::Vector{Tuple{S, MapRayClassGrp}}
  into_product_of_groups::GrpAbFinGenMap # the isomorphism between the domain and the product of the groups in groups_in_number_fields

  function MapRayClassGroupAlg{S, T}() where {S, T}
    return new{S, T}()
  end
end

function modulus(f::MapRayClassGroupAlg{S, T}) where {S, T}
  return f.modulus::elem_type(base_ring_type(T))
end

base_ring_type(::Type{FacElemMon{S}}) where {S} = S

mutable struct MapPicardGrp{S, T} <: Map{S, T, HeckeMap, MapPicardGrp}
  header::MapHeader{S, T}

  # Only used for picard groups of orders in number fields
  OO_mod_F_mod_O_mod_F::GrpAbFinGenToAbsOrdQuoRingMultMap

  # For the refined discrete logarithm: (only used for picard groups of orders in algebras)
  right_transform::fmpz_mat
  betas # Vector of factorized algebra elements
  gammas # the same type as betas
  ray_class_group_map::MapRayClassGroupAlg{S, FacElemMon{T}}

  function MapPicardGrp{S, T}() where {S, T}
    return new{S, T}()
  end
end

function show(io::IO, mP::MapPicardGrp)
  @show_name(io, mP)
  println(io, "Picard Group map of ")
  show(IOContext(io, :compact => true), codomain(mP))
end

################################################################################
#
#  Picard group
#
################################################################################

@doc Markdown.doc"""
    picard_group(O::AlgAssAbsOrd, prepare_ref_disc_log::Bool = false)
      -> GrpAbFinGen, MapPicardGroup

Given an order $O$ in a commutative algebra over $\mathbb Q$, this function
returns the picard group of $O$.
If `prepare_ref_disc_log` is `true`, then (possibly expensive) preparations
for the computation of refined discrete logarithms in non maximal orders are done.
"""
function picard_group(O::AlgAssAbsOrd, prepare_ref_disc_log::Bool = false)
  @assert iscommutative(O)
  if !prepare_ref_disc_log && isdefined(O, :picard_group)
    mp = O.picard_group::MapPicardGrp{GrpAbFinGen, parent_type(ideal_type(O))}
    return domain(mp), mp
  end

  if ismaximal(O)
    return _picard_group_maximal(O)
  end

  if prepare_ref_disc_log && isdefined(O, :picard_group)
    mP = O.picard_group::MapPicardGrp{GrpAbFinGen, parent_type(ideal_type(O))}
    if isdefined(mP, :betas) && isdefined(mP, :gammas) && isdefined(mP, :right_transform)
      return domain(mP), mP
    end
  end
  return _picard_group_non_maximal(O, prepare_ref_disc_log)
end

function _picard_group_maximal(O::AlgAssAbsOrd)
  A = algebra(O)
  fields_and_maps = as_number_fields(A)
  class_groups = Tuple{GrpAbFinGen, MapClassGrp}[ class_group(field) for (field, map) in fields_and_maps ]
  P = class_groups[1][1]
  for i = 2:length(class_groups)
    P = direct_product(P, class_groups[i][1]; task = :none)::GrpAbFinGen
  end
  S, StoP = snf(P)

  local disc_exp
  let A = A, StoP = StoP, fields_and_maps = fields_and_maps, class_groups = class_groups, O = O
    function disc_exp(x::GrpAbFinGenElem)
      p = StoP(x)
      basis_of_ideal = Vector{elem_type(A)}()
      offset = 1
      for i = 1:length(fields_and_maps)
        K, AtoK = fields_and_maps[i]
        C, CtoIdl = class_groups[i]
        c = C(sub(p.coeff, 1:1, offset:(offset + ngens(C) - 1)))
        I = CtoIdl(c)
        for b in basis(I)
          push!(basis_of_ideal, AtoK\K(b))
        end
        offset += ngens(C)
      end
      I = ideal_from_lattice_gens(A, O, basis_of_ideal)
      return I
    end
  end

  local disc_log
  let fields_and_maps = fields_and_maps, class_groups = class_groups, StoP = StoP, P = P
    function disc_log(x::AlgAssAbsOrdIdl)
      ideals = Vector{NfOrdIdl}()
      for i = 1:length(fields_and_maps)
        push!(ideals, _as_ideal_of_number_field(x, fields_and_maps[i][2]))
      end

      p = zero_matrix(FlintZZ, 1, 0)
      for i = 1:length(ideals)
        C, CtoIdl = class_groups[i]
        c = CtoIdl\ideals[i]
        p = hcat(p, c.coeff)
      end
      return StoP\P(p)
    end
  end

  Idl = IdealSet(O)
  StoIdl = MapPicardGrp{GrpAbFinGen, typeof(Idl)}()
  StoIdl.header = MapHeader(S, Idl, disc_exp, disc_log)
  O.picard_group = StoIdl
  return S, StoIdl
end

function _trivial_picard(O::AlgAssAbsOrd, R::GrpAbFinGen, mR)
  A = algebra(O)
  Idl = IdealSet(O)
  local disc_exp_triv
  let O = O
    function disc_exp_triv(x::GrpAbFinGenElem)
      return ideal(O, one(O))
    end
  end

  local disc_log_triv
  let R = R
    function disc_log_triv(x::AlgAssAbsOrdIdl)
      return R()
    end
  end

  Idl = IdealSet(O)
  fac_elem_mon_A = FacElemMon(A)
  RtoIdl = MapPicardGrp{GrpAbFinGen, typeof(Idl)}()
  RtoIdl.header = MapHeader(R, Idl, disc_exp_triv, disc_log_triv)
  RtoIdl.right_transform = zero_matrix(FlintZZ, 0, 0)
  RtoIdl.betas = Vector{elem_type(fac_elem_mon_A)}()
  RtoIdl.gammas = Vector{elem_type(fac_elem_mon_A)}()
  RtoIdl.ray_class_group_map = mR
  O.picard_group = RtoIdl
  return R, RtoIdl

end

# See Bley, Endres "Picard Groups and Refined Discrete Logarithms".
function _picard_group_non_maximal(O::AlgAssAbsOrd, prepare_ref_disc_log::Bool = false)
  A = algebra(O)
  OO = maximal_order(A)

  F = conductor(O, OO, :left)
  FOO = extend(F, OO)

  # We want to use the exact sequence
  # (O/F)^\times \to C_FOO(OO) \to Pic(O) \to 0.
  # where C_FOO(OO) is the ray class group of OO modulo FOO.

  # Firstly, we compute the groups.
  R, mR = ray_class_group(FOO)
  @assert issnf(R)

  Idl = IdealSet(O)
  fac_elem_mon_A = FacElemMon(A)
  if ngens(R) == 0
    return _trivial_picard(O, R, mR)
  end

  Q, OtoQ = quo(O, F)
  G, GtoQ = multiplicative_group(Q)

  if !prepare_ref_disc_log
    # If we don't need to compute refined discrete logarithms, we compute the
    # picard group as the quotient of groups.

    # Compute the image of the map from G to R, i. e. the kernel of the map from
    # R to Pic(O).
    GinR = Vector{GrpAbFinGenElem}()
    for i = 1:ngens(G)
      g = OO(OtoQ\(GtoQ(G[i])))
      r = mR\(ideal(OO, g))
      push!(GinR, r)
    end

    # Compute the Picard group
    P, RtoP = quo(R, GinR)
    S, StoP = snf(P)

    StoR = compose(StoP, pseudo_inv(RtoP))

    gens_snf = Vector{ideal_type(O)}(undef, ngens(S))
    for i = 1:ngens(S)
      r = RtoP\StoP(S[i])
      # Contraction does not commute with multiplication, so we have to evaluate.
      # TODO: Can we make this work?
      gens_snf[i] = contract(evaluate(mR(r)), O)
    end
  else
    # This is now Algorithm 3.3 in the Bley-Endres-paper.

    # Let [r_1,...,r_n] in R. Then mR([r_1,...,r_n]) is in general not equal to
    # prod_i mR(R[i])^r_i but just something equivalent. In the following, we
    # need this latter "discrete exponentiation". Therefore we collect the
    # generators here and don't use mR later.
    fac_elem_mon_idl_OO = FacElemMon(IdealSet(OO))
    gens_of_R = Vector{elem_type(fac_elem_mon_idl_OO)}()
    for i = 1:ngens(R)
      push!(gens_of_R, mR(R[i]))
    end

    # Compute the relations of the generators of G in R
    gens_of_G = Vector{elem_type(fac_elem_mon_A)}()
    D = zero_matrix(FlintZZ, ngens(R) + ngens(G), ngens(R))
    for i = 1:ngens(R)
      D[i, i] = R.snf[i]
    end
    for i = (ngens(R) + 1):(ngens(R) + ngens(G))
      g = OO(OtoQ\(GtoQ(G[i - ngens(R)])))
      gOO = ideal(OO, g)
      a, r = disc_log_generalized_ray_class_grp(gOO, mR)

      # Now we have gOO = a*mR(r), but we need gOO = c*gOO2
      gOO2 = fac_elem_mon_idl_OO()
      for j = 1:ngens(R)
        gOO2 *= gens_of_R[j]^r[j]
      end
      b, s = disc_log_generalized_ray_class_grp(gOO2, mR)
      @assert r == s
      c = a*inv(b)

      for j = 1:ngens(R)
        D[i, j] = r[j]
      end
      push!(gens_of_G, elem_in_algebra(g, copy = false)*inv(c))
    end

    # D is the relation matrix of the picard group
    H, W = hnf_with_transform(D)
    H = sub(H, 1:ngens(R), 1:ngens(R))

    S_rels, U, V = snf_with_transform(H)
    Vi = inv(V)

    # In principle S_rels is the SNF of GrpAbFinGen(H), but we can as well use R
    # here, since the transformation for the HNF is from the left.
    S, StoR = _reduce_snf(R, S_rels, V, Vi)

    # We need to change Vi, since we want integral generators.
    Q = similar(V) # Not exactly the same Q as in the Bley-Endres-paper
    C = deepcopy(Vi) # Will be Vi + Q*R.rels, we can't change Vi in place, since it is also part of StoR
    for i = 1:nrows(C)
      for j = 1:ncols(C)
        if C[i, j] >= 0
          continue
        end

        Q[i, j] = div(-C[i, j], R.snf[j]) + 1
        C[i, j] = addeq!(C[i, j], R.snf[j]*Q[i, j])
      end
    end

    # Compute the generators
    generators_in_OO = Vector{elem_type(fac_elem_mon_idl_OO)}()
    gens_snf = Vector{ideal_type(O)}()
    for i = 1:nrows(C)
      I = fac_elem_mon_idl_OO()
      for j = 1:ncols(C)
        I *= gens_of_R[j]^C[i, j]
      end
      push!(generators_in_OO, I)

      if !isone(S_rels[i, i])
        # Contraction does not commute with multiplication, so we have to evaluate.
        push!(gens_snf, contract(evaluate(I), O))
      end
    end
    @assert length(gens_snf) == ngens(S)

    # Compute the additional data needed for the refined discrete logarithm.
    princ_gens_ray = Vector{elem_type(fac_elem_mon_A)}()
    for i = 1:ngens(R)
      b, a = has_principal_generator_1_mod_m(gens_of_R[i]^R.snf[i], FOO)
      @assert b
      push!(princ_gens_ray, a)
    end

    gammas = Vector{elem_type(fac_elem_mon_A)}()
    for i = 1:ngens(R)
      g = fac_elem_mon_A()
      for j = 1:ngens(R)
        g *= princ_gens_ray[j]^Q[i, j]
      end
      push!(gammas, g)
    end

    # In principle this should be
    # (U | 0)       (* | *)
    # (-----) * W = (-----)
    # (0 | 0)       (0 | 0).
    # We don't do the multiplications with zero.
    UW = U*sub(W, 1:ngens(R), 1:ncols(W))

    betas = Vector{elem_type(fac_elem_mon_A)}()
    for i = 1:nrows(UW)
      t = fac_elem_mon_A()
      for j = 1:ngens(R)
        t *= princ_gens_ray[j]^UW[i, j]
      end
      for j = 1:ngens(G)
        t *= gens_of_G[j]^UW[i, j + ngens(R)]
      end
      if i <= length(gammas)
        t *= gammas[i]^S_rels[i, i]
      end
      push!(betas, t)
    end
  end

  local disc_exp
  let Idl = Idl, gens_snf = gens_snf
    function disc_exp(x::GrpAbFinGenElem)
      y = one(Idl)
      for i = 1:length(x.coeff)
        y *= gens_snf[i]^x[i]
      end
      return y
    end
  end

  local disc_log
  let F = F, OO = OO, mR = mR, StoR = StoR
    function disc_log(x::AlgAssAbsOrdIdl)
      if !isinvertible(x)[1]
        error("Ideal is not invertible")
      end
      if !isone(x + F)
        x, _ = _coprime_integral_ideal_class(x, F)
      end

      xOO = extend(x, OO)
      r = mR\xOO
      return StoR\r
    end
  end

  StoIdl = MapPicardGrp{GrpAbFinGen, typeof(Idl)}()
  StoIdl.header = MapHeader(S, Idl, disc_exp, disc_log)
  StoIdl.ray_class_group_map = mR

  if prepare_ref_disc_log
    StoIdl.right_transform = V
    StoIdl.betas = betas
    StoIdl.gammas = gammas
  end

  O.picard_group = StoIdl

  return S, StoIdl
end

# Algorithm 3.5 in Bley, Endres "Picard groups and refined discrete logarithms"
function refined_disc_log_picard_group(a::AlgAssAbsOrdIdl, mP::MapPicardGrp)
  if !isdefined(mP, :betas) || !isdefined(mP, :gammas) || !isdefined(mP, :right_transform)
    # I don't want to recompute the group here, because the user gave me the map.
    error("Data for the refined discrete logarithm is not available.")
  end

  if !isinvertible(a)[1]
    error("Ideal is not invertible")
  end

  O = order(a)
  A = algebra(O)
  OO = maximal_order(A)
  F = conductor(O, OO, :left)
  FOO = F*OO
  aOO = a*OO

  t = one(A)
  if !isone(a + F)
    # After this modification of a we do not need steps 1 to 4 of the algorithm
    a, t = _coprime_integral_ideal_class(a, F)
    aOO = a*OO
  end

  # Compute the refined disc log of aOO in the ray class group
  mR = mP.ray_class_group_map
  R = domain(mR)

  k1, r = disc_log_generalized_ray_class_grp(aOO, mR)
  aOO2 = FacElemMon(IdealSet(OO))()
  for i = 1:ngens(R)
    # We need this kind of discrete exponentiation
    aOO2 *= mR(R[i])^r[i]
  end

  k2, s = disc_log_generalized_ray_class_grp(aOO2, mR)
  @assert r == s
  k = k1*inv(k2)

  # Map r into the picard group
  rV = r.coeff*mP.right_transform

  P = domain(mP)
  rr = Vector{fmpz}(undef, ngens(R))
  num1 = ngens(R) - ngens(P) # the number of ones in the SNF of P
  for i = 1:num1
    rr[i] = rV[1, i]
  end

  p = Vector{fmpz}(undef, ngens(P))
  for i = 1:ngens(P)
    inum1 = i + num1
    q, r = divrem(rV[inum1], P.snf[i])

    # We need positive remainder
    if r < 0
      r += P.snf[i]
      q -= 1
    end
    rr[inum1] = q
    p[i] = r
  end

  beta = one(A)
  gamma = one(A)
  for i = 1:ngens(R)
    beta *= mP.betas[i]^rr[i]
    gamma *= mP.gammas[i]^rV[i]
  end

  return k*beta*inv(t*gamma), P(p)
end

@doc Markdown.doc"""
    principal_generator(a::AlgAssAbsOrdIdl) -> AlgAssAbsOrdElem

Given a principal ideal $a$ in an order $O$ in a commutative algebra over
$\mathbb Q$, this function returns a principal generator of $a$.
"""
function principal_generator(a::AlgAssAbsOrdIdl)
  a, g = isprincipal(a)
  if !a
    error("Ideal is not principal")
  end
  return g
end

function principal_generator_fac_elem(a::AlgAssAbsOrdIdl)
  @assert ismaximal(order(a)) "Not implemented"
  a, g = isprincipal_maximal_fac_elem(a)
  if !a
    error("Ideal is not principal")
  end
  return g
end

function isprincipal(a::AlgAssAbsOrdIdl)
  if ismaximal(order(a))
    return isprincipal_maximal(a)
  end
  return isprincipal_non_maximal(a)
end

function isprincipal_fac_elem(a::AlgAssAbsOrdIdl)
  @assert ismaximal(order(a)) "Not implemented"
  return isprincipal_maximal_fac_elem(a)
end

function isprincipal_maximal(a::AlgAssAbsOrdIdl)
  b, x = isprincipal_maximal_fac_elem(a)
  return b, order(a)(evaluate(x))
end

function isprincipal_maximal_fac_elem(a::AlgAssAbsOrdIdl)
  O = order(a)
  A = algebra(O)
  fields_and_maps = as_number_fields(A)

  idems = [ AtoK\one(K) for (K, AtoK) in fields_and_maps ]
  sum_idems = sum(idems)
  minus_idems = [ -one(A)*idem for idem in idems ]
  bases = Vector{elem_type(A)}()
  exps = Vector{fmpz}()
  for i = 1:length(fields_and_maps)
    K, AtoK = fields_and_maps[i]
    C, mC = class_group(K) # should be cached
    Hecke._assure_princ_gen(mC)
    ai = _as_ideal_of_number_field(a, AtoK)
    c, g = isprincipal_fac_elem(ai)
    if !c
      return false, FacElemMon(A)()
    end
    for (b, e) in g
      t = AtoK\b
      t = add!(t, t, sum_idems)
      t = add!(t, t, minus_idems[i])
      push!(bases, t)
      push!(exps, e)
    end
  end
  return true, FacElem(bases, exps)
end

# for isprincipal_non_maximal see NfOrd/PicardGroup.jl

################################################################################
#
#  Ray Class Group
#
################################################################################

# inf_pcl[i] may contain places for the field A.maps_to_numberfields[i][1]
function ray_class_group(m::AlgAssAbsOrdIdl, inf_plc::Vector{Vector{InfPlc}} = Vector{Vector{InfPlc}}())
  O = order(m)
  A = algebra(O)
  fields_and_maps = as_number_fields(A)

  # Compute the groups in the number fields
  groups = Vector{Tuple{GrpAbFinGen, MapRayClassGrp}}()
  for i = 1:length(fields_and_maps)
    mi = _as_ideal_of_number_field(m, fields_and_maps[i][2])
    if length(inf_plc) != 0
      r, mr = ray_class_group(mi, inf_plc[i])
    else
      r, mr = ray_class_group(mi)
    end
    push!(groups, _make_disc_exp_deterministic(mr))
    #push!(groups, ray_class_group(mi))
  end

  C = groups[1][1]
  for i = 2:length(groups)
    C = direct_product(C, groups[i][1], task = :none)::GrpAbFinGen
  end
  S, StoC = snf(C)

  one_ideals = _lift_one_ideals(O)
  fac_elem_mon = FacElemMon(parent(m))

  local disc_exp
  let StoC = StoC, groups = groups, O = O, one_ideals = one_ideals
    function disc_exp(x::GrpAbFinGenElem)
      z = StoC(x).coeff
      y = fac_elem_mon()
      j = 1
      for k = 1:length(groups)
        G, GtoIdl = groups[k]
        zz = sub(z, 1:1, j:(j + ngens(G) - 1))
        Ifac = GtoIdl(G(zz))
        y *= _as_ideal_of_algebra(Ifac, k, O, one_ideals)
        j += ngens(G)
      end
      return y
    end
  end

  local disc_log
  let fields_and_maps = fields_and_maps, C = C, StoC = StoC
    function disc_log(x::AlgAssAbsOrdIdl)
      ideals = Vector{NfOrdIdl}()
      for i = 1:length(fields_and_maps)
        push!(ideals, _as_ideal_of_number_field(x, fields_and_maps[i][2]))
      end

      c = zero_matrix(FlintZZ, 1, 0)
      for i = 1:length(ideals)
        G, GtoIdl = groups[i]
        g = GtoIdl\ideals[i]
        c = hcat(c, g.coeff)
      end
      return StoC\C(c)
    end

    function disc_log(x::FacElem)
      ideals = Vector{FacElem{NfOrdIdl, NfOrdIdlSet}}()
      for i = 1:length(fields_and_maps)
        base = Vector{NfOrdIdl}()
        exp = Vector{fmpz}()
        for (I, e) in x
          push!(base, _as_ideal_of_number_field(I, fields_and_maps[i][2]))
          push!(exp, e)
        end
        push!(ideals, FacElem(base, exp))
      end

      c = zero_matrix(FlintZZ, 1, 0)
      for i = 1:length(ideals)
        G, GtoIdl = groups[i]
        g = GtoIdl\ideals[i]
        c = hcat(c, g.coeff)
      end
      return StoC\C(c)
    end
  end

  StoIdl = MapRayClassGroupAlg{typeof(S), typeof(fac_elem_mon)}()
  StoIdl.header = MapHeader(S, fac_elem_mon, disc_exp, disc_log)
  StoIdl.modulus = m
  StoIdl.groups_in_number_fields = groups
  StoIdl.into_product_of_groups = StoC
  return S, StoIdl
end

# A simple way to make the discrete exponentiation in the ray class group (for
# number fields) deterministic.
function _make_disc_exp_deterministic(mR::MapRayClassGrp)
  R = domain(mR)
  S, StoR = snf(R)

  fac_elem_mon = codomain(mR)
  generators = Vector{elem_type(fac_elem_mon)}()
  for i = 1:ngens(S)
    r = StoR(S[i])
    push!(generators, mR(r))
  end

  local disc_exp
  let fac_elem_mon = fac_elem_mon, generators = generators
    function disc_exp(x::GrpAbFinGenElem)
      y = fac_elem_mon()
      for i = 1:ngens(parent(x))
        y *= generators[i]^x[i]
      end
      return y
    end
  end

  local disc_log
  let StoR = StoR, mR = mR
    function disc_log(x::FacElem)
      return StoR\(mR\x)
    end

    function disc_log(x::NfAbsOrdIdl)
      return StoR\(mR\x)
    end
  end

  mRR = MapRayClassGrp()
  mRR.header = MapHeader(S, fac_elem_mon, disc_exp, disc_log)
  for x in fieldnames(typeof(mR))
    if x != :header && isdefined(mR, x)
      setfield!(mRR, x, getfield(mR, x))
    end
  end
  return S, mRR
end

function has_principal_generator_1_mod_m(I::Union{ AlgAssAbsOrdIdl, FacElem{ <: AlgAssAbsOrdIdl, <: AlgAssAbsOrdIdlSet} }, m::AlgAssAbsOrdIdl)
  O = order(m)
  A = algebra(O)
  fields_and_maps = as_number_fields(A)

  idems = elem_type(A)[ AtoK\one(K) for (K, AtoK) in fields_and_maps ]
  sum_idems = sum(idems)
  minus_idems = elem_type(A)[ -one(A)*idem for idem in idems ]
  bases = Vector{elem_type(A)}()
  exps = Vector{fmpz}()
  for i = 1:length(fields_and_maps)
    K, AtoK = fields_and_maps[i]
    mi = _as_ideal_of_number_field(m, AtoK)
    Ii = _as_ideal_of_number_field(I, AtoK)
    a, g = has_principal_generator_1_mod_m(Ii, mi)
    if !a
      return false, FacElem(base, exps)
    end
    for (b, e) in g
      t = AtoK\b
      t = add!(t, t, sum_idems)
      t = add!(t, t, minus_idems[i])
      push!(bases, t)
      push!(exps, e)
    end
  end
  return true, FacElem(bases, exps)
end

function disc_log_generalized_ray_class_grp(I::FacElem{S, T}, mR::MapRayClassGroupAlg) where { S <: AlgAssAbsOrdIdl, T <: AlgAssAbsOrdIdlSet }
  m = modulus(mR)
  O = order(m)
  A = algebra(O)

  if isempty(I)
    return FacElemMon(A)(), domain(mR)()
  end

  fields_and_maps = as_number_fields(A)

  groups = mR.groups_in_number_fields

  ideals = Vector{FacElem{NfOrdIdl, NfOrdIdlSet}}()

  for i = 1:length(fields_and_maps)
    push!(ideals, _as_ideal_of_number_field(I, fields_and_maps[i][2]))
  end

  idems = elem_type(A)[ AtoK\one(K) for (K, AtoK) in fields_and_maps ]
  sum_idems = sum(idems)
  minus_idems = elem_type(A)[ -one(A)*idem for idem in idems ]
  bases = Vector{elem_type(A)}()
  exps = Vector{fmpz}()
  p = zero_matrix(FlintZZ, 1, 0)
  ideal_gens = Vector{elem_type(O)}()
  for i = 1:length(ideals)
    K, AtoK = fields_and_maps[i]
    d, lI = disc_log_generalized_ray_class_grp(ideals[i], groups[i][2])
    p = hcat(p, lI.coeff)
    for (b, e) in d
      t = AtoK\b
      t = add!(t, t, sum_idems)
      t = add!(t, t, minus_idems[i])
      push!(bases, t)
      push!(exps, e)
    end
  end
  RtoP = mR.into_product_of_groups
  P = codomain(RtoP)
  r = RtoP\P(p)

  return FacElem(bases, exps), r
end

function disc_log_generalized_ray_class_grp(I::S, mR::MapRayClassGroupAlg) where {S <: AlgAssAbsOrdIdl}
  m = modulus(mR)
  O = order(m)
  A = algebra(O)

  fields_and_maps = as_number_fields(A)

  groups = mR.groups_in_number_fields

  ideals = Vector{NfOrdIdl}()

  for i = 1:length(fields_and_maps)
    push!(ideals, _as_ideal_of_number_field(I, fields_and_maps[i][2]))
  end

  idems = elem_type(A)[ AtoK\one(K) for (K, AtoK) in fields_and_maps ]
  sum_idems = sum(idems)
  minus_idems = elem_type(A)[ -one(A)*idem for idem in idems ]
  bases = Vector{elem_type(A)}()
  exps = Vector{fmpz}()
  p = zero_matrix(FlintZZ, 1, 0)
  ideal_gens = Vector{elem_type(O)}()
  for i = 1:length(ideals)
    K, AtoK = fields_and_maps[i]
    d, lI = disc_log_generalized_ray_class_grp(ideals[i], groups[i][2])
    p = hcat(p, lI.coeff)
    for (b, e) in d
      t = AtoK\b
      t = add!(t, t, sum_idems)
      t = add!(t, t, minus_idems[i])
      push!(bases, t)
      push!(exps, e)
    end
  end
  RtoP = mR.into_product_of_groups
  P = codomain(RtoP)
  r = RtoP\P(p)

  return FacElem(bases, exps), r
end

################################################################################
#
#  Kernel group
#
################################################################################

@doc Markdown.doc"""
    kernel_group(O::AlgAssAbsOrd) -> GrpAbFinGen, MapPicardGroup

Given an order $O$ in a commutative algebra over $\mathbb Q$, this function
returns the group $D$ in the exact sequence $0 \to D \to Pic(O) \to Pic(O')$
where $O'$ is a maximal order containing $O$.
"""
function kernel_group(O::AlgAssAbsOrd)
  A = algebra(O)
  OO = maximal_order(A)

  # We use the short exact sequence
  # 0 \to D(O) \to Pic(O) \to Pic(OO) \to 0,
  # where D(O) is the kernel group of O.

  P, mP = picard_group(O)
  C, mC = picard_group(OO)

  Idl = IdealSet(O)
  if C == P
    D = GrpAbFinGen(fmpz[])
    local disc_exp_triv
    let O = O
      function disc_exp_triv(x::GrpAbFinGenElem)
        return ideal(O, one(O))
      end
    end

    local disc_log_triv
    let D = D
      function disc_log_triv(x::AlgAssAbsOrdIdl)
        return D()
      end
    end

    DtoIdl = MapPicardGrp{typeof(D), typeof(Idl)}()
    DtoIdl.header = MapHeader(D, Idl, disc_exp_triv, disc_log_triv)
    return D, DtoIdl
  end

  # Build the morphism from Pic(O) to Pic(OO)
  A = Vector{GrpAbFinGenElem}()
  B = Vector{GrpAbFinGenElem}()
  for i = 1:ngens(P)
    push!(A, P[i])
    p = mP(P[i])
    pOO = extend(p, OO)
    c = mC\pOO
    push!(B, c)
  end

  PtoC = hom(A, B)

  # The kernel group is the kernel of this morphism
  D, DtoP = kernel(PtoC)
  S, StoD = snf(D)

  gens_snf = Vector{elem_type(Idl)}()
  for i = 1:ngens(S)
    p = DtoP(StoD(S[i]))
    push!(gens_snf, mP(p))
  end

  local disc_exp
  let Idl = Idl, gens_snf = gens_snf
    function disc_exp(x::GrpAbFinGenElem)
      y = one(Idl)
      for i = 1:length(x.coeff)
        y *= gens_snf[i]^x.coeff[1, i]
      end
      return y
    end
  end

  local disc_log
  let mP = mP, DtoP = DtoP, StoD = StoD
    function disc_log(x::AlgAssAbsOrdIdl)
      p = mP\x
      b, g = haspreimage(DtoP, p)
      @assert b "Ideal not an element of the kernel group"
      return StoD\g
    end
  end

  StoIdl = MapPicardGrp{GrpAbFinGen, typeof(Idl)}()
  StoIdl.header = MapHeader(S, Idl, disc_exp, disc_log)
  return S, StoIdl
end

################################################################################
#
#  Coprime representatives
#
################################################################################

# Mostly taken from NfOrd/LinearAlgebra.jl
function _coprime_integral_ideal_class(a::AlgAssAbsOrdIdl, b::AlgAssAbsOrdIdl)
  O = order(b)
  @assert isone(denominator(b, O))
  d = denominator(a, O)
  a = d*a
  a_inv = inv(a)
  c = ideal(O, one(O))
  x = algebra(O)()
  check = true
  while check
    x = rand(a_inv, 100)
    c = x*a
    @assert denominator(c, O) == 1
    c.order = O
    if iszero(det(basis_matrix(c)))
      continue
    end
    isone(c + b) ? (check = false) : (check = true)
  end
  return c, x*d
end

################################################################################
#
#  Unit groups
#
################################################################################

function unit_group(O::AlgAssAbsOrd)
  @assert iscommutative(O)
  if isdefined(O, :unit_group)
    return domain(O.unit_group), O.unit_group
  end

  if ismaximal(O)
    U, mU = _unit_group_maximal(O)
  else
    U, mU = _unit_group_non_maximal(O)
  end

  O.unit_group = mU
  return U, mU
end

function _unit_group_maximal(O::AlgAssAbsOrd)
  A = algebra(O)
  fields_and_maps = as_number_fields(A)
  unit_groups = Tuple{GrpAbFinGen, MapUnitGrp{NfOrd}}[ unit_group(maximal_order(field)) for (field, map) in fields_and_maps ]
  G = unit_groups[1][1]
  for i = 2:length(unit_groups)
    G = direct_product(G, unit_groups[i][1], task = :none)::GrpAbFinGen
  end
  S, StoG = snf(G)

  local disc_exp
  let StoG = StoG, unit_groups = unit_groups, fields_and_maps = fields_and_maps
    function disc_exp(x::GrpAbFinGenElem)
      g = StoG(x)
      v = zero(O)
      offset = 1
      for i = 1:length(fields_and_maps)
        K, AtoK = fields_and_maps[i]
        U, mU = unit_groups[i]
        u = U(sub(g.coeff, 1:1, offset:(offset + ngens(U) - 1)))
        v += O(AtoK\elem_in_nf(mU(u), copy = false))
        offset += ngens(U)
      end
      return v
    end
  end

  local disc_log
  let fields_and_maps = fields_and_maps, unit_groups = unit_groups, StoG = StoG, G = G
    function disc_log(x::AlgAssAbsOrdElem)
      g = zero_matrix(FlintZZ, 1, 0)
      for i = 1:length(fields_and_maps)
        K, AtoK = fields_and_maps[i]
        U, mU = unit_groups[i]
        OK = codomain(mU)
        y = OK(AtoK(elem_in_algebra(x, copy = false)))
        @assert isunit(y)
        u = mU\y
        g = hcat(g, u.coeff)
      end
      return StoG\G(g)
    end
  end

  StoO = MapUnitGrp{typeof(O)}()
  StoO.header = MapHeader(S, O, disc_exp, disc_log)
  return S, StoO
end

# Returns the group (O_A/F)^\times/(O/F)^\times and a map from this group to
# O_A/F where F is the conductor of O in the maximal order O_A.
function OO_mod_F_mod_O_mod_F(O::AlgAssAbsOrd)
  OA = maximal_order(algebra(O))
  F = conductor(O, OA, :left)
  FOA = F*OA
  Q1, toQ1 = quo(OA, FOA)
  H1, H1toQ1 = unit_group(Q1)
  Q2, toQ2 = quo(O, F)
  H2, H2toQ2 = unit_group(Q2)
  H2inH1, _ = sub(H1, [ H1toQ1\(toQ1(OA(elem_in_algebra(toQ2\(H2toQ2(H2[i])), copy = false)))) for i = 1:ngens(H2) ])
  H, toH = quo(H1, H2inH1)
  S, StoH = snf(H)
  local _disc_log
  let toH = toH, StoH = StoH, Q1 = Q1
    function _disc_log(x)
      @assert parent(x) === Q1
      s = StoH\(toH(H1toQ1\x))
      return s.coeff
    end
  end

  StoQ1 = GrpAbFinGenToAbsOrdQuoRingMultMap(S, Q1, [ H1toQ1(toH\(StoH(S[i]))) for i = 1:ngens(S) ], _disc_log)
  return S, StoQ1, toQ1
end
# for _unit_group_non_maximal see NfOrd/PicardGroup.jl


text:
# created: 20210519
# 2D linear rate model with balanced amplification effect.

push!(LOAD_PATH, abspath(@__DIR__,".."))

using LinearAlgebra,Statistics,StatsBase
using Plots,NamedColors ; theme(:dark)
using SparseArrays
using SpikingRNNs; const global S = SpikingRNNs

## Parameters
ne = 1
ni = 1

τe = 1.
τi = 1.

α = 1.

# Define neuron types and population states

neuron_e = S.NTReLU(τe,α)
neuron_i = S.NTReLU(τi,α)
pse  = S.PSRate(neuron_e,ne)
psi  = S.PSRate(neuron_i,ni)

# Define the connections
function onesparsemat(w::Real)
  mat = Matrix{Float64}(undef,1,1)
  mat[1,1] = w
  return sparse(mat)
end

(w_ee,w_ie,w_ei,w_ii) = let w = 20. , k = 1.1
  onesparsemat.((w,w,-k*w,-k*w))
end

conn_ee = S.BaseConnection(w_ee)
conn_ei = S.BaseConnection(w_ei)
conn_ie = S.BaseConnection(w_ie)
conn_ii = S.BaseConnection(w_ii)

# inputs
in_state_e = S.PSSimpleInput(S.InputSimpleOffset(0.0))
in_state_i = S.PSSimpleInput(S.InputSimpleOffset(0.0))

# populations are population states, plus all incoming connections
# plus presynaptic population states

pop_e = S.Population(pse,(conn_ee,conn_ei,S.FakeConnection()),
  (pse,psi,in_state_e))
pop_i = S.Population(psi,(conn_ie,conn_ii,S.FakeConnection()),
  (pse,psi,in_state_i))

dt = 1E-4
T = 10.0
times = 0:dt:T
ntimes = length(times)
mynetwork = S.RecurrentNetwork(dt,(pop_e,pop_i))

e_out = Vector{Float64}(undef,ntimes)
i_out = Vector{Float64}(undef,ntimes)

##

# initial conditions
pse.state_now[1] = S.ioinv(10.0,pse)
psi.state_now[1] = S.ioinv(5.0,pse)

for (k,t) in enumerate(times)
  e_out[k] = S.iofunction(pse.state_now[1],pse)
  i_out[k] = S.iofunction(psi.state_now[1],psi)
  # rate model with constant input  does not really depend on absolute time (first argument)
  S.dynamics_step!(0.0,mynetwork)
end

plot(times,[e_out i_out];linewidth=4,leg=:topright,
  color=[colorant"Teal" colorant"Salmon"],label=["exc" "inh"])


text:
using .Networks

using DataStructures
using Dates
using JuMP
using Printf
using Parameters

@with_kw mutable struct EODAssessmentModel <: AbstractModelContainer
    model::Model = Model()

    # Uncertainty model
    #bus,ts
    uncertain_load::SortedDict{Tuple{String, DateTime},VariableRef} =
        SortedDict{Tuple{String, DateTime},VariableRef}();
    #limitableUnit,ts
    uncertain_prod::SortedDict{Tuple{String, DateTime},VariableRef} =
        SortedDict{Tuple{String, DateTime},VariableRef}();

    # Market model
    #unit, ts
    p_injected::SortedDict{Tuple{String, DateTime},VariableRef} =
        SortedDict{Tuple{String, DateTime},VariableRef}();
    b_in::SortedDict{Tuple{String, DateTime},VariableRef} =
        SortedDict{Tuple{String, DateTime},VariableRef}();
    b_marg::SortedDict{Tuple{String, DateTime},VariableRef} =
        SortedDict{Tuple{String, DateTime},VariableRef}();
    b_out::SortedDict{Tuple{String, DateTime},VariableRef} =
        SortedDict{Tuple{String, DateTime},VariableRef}();

    #ts :
    #cut prod
    p_cut_prod::SortedDict{DateTime,VariableRef} =
        SortedDict{DateTime,VariableRef}();
    #LoL
    p_loss_of_load::SortedDict{DateTime,VariableRef} =
        SortedDict{DateTime,VariableRef}();

    #objective
    loss_of_load_obj =  AffExpr(0.)
    prod_obj =  AffExpr(0.)
    cut_prod_obj =  AffExpr(0.)
    full_obj = AffExpr(0.)
end

function add_load_uncertainties_vars!(model_container::EODAssessmentModel,
                                    network, TS, assessment_uncertainties)
    for bus in Networks.get_buses(network)
        bus_id = Networks.get_id(bus)
        lb_l = get_assessment_uncertainties_lb(assessment_uncertainties, bus_id)
        ub_l = get_assessment_uncertainties_ub(assessment_uncertainties, bus_id)
        for ts in TS
            name =  @sprintf("uncertain_load[%s,%s]", bus_id, ts);
            model_container.uncertain_load[bus_id,ts] = @variable(model_container.model, base_name = name,
                                                                lower_bound = lb_l, upper_bound = ub_l)
        end
    end

    return model_container.uncertain_load
end
function add_prod_uncertainties_vars!(model_container::EODAssessmentModel,
                                    network, TS, assessment_uncertainties, tso_actions)
    for limitable_gen in Networks.get_generators_of_type(network, Networks.LIMITABLE)
        gen_id = Networks.get_id(limitable_gen)
        lb_l = get_assessment_uncertainties_lb(assessment_uncertainties, gen_id)
        for ts in TS
            # this allows to set p_inj[gen_id] == uncertain_prod without needing to express min(uncertain, limit) in the model
            ub_l = min(get_assessment_uncertainties_ub(assessment_uncertainties, gen_id),
                    get_limitation(tso_actions, gen_id, ts),
                    Networks.get_p_max(limitable_gen))
            name =  @sprintf("uncertain_prod[%s,%s]", gen_id, ts);
            model_container.uncertain_prod[gen_id, ts] = @variable(model_container.model, base_name=name,
                                                                lower_bound = lb_l, upper_bound = ub_l)
        end
    end

    return model_container.uncertain_prod
end
function add_uncertainties_vars!(model_container::EODAssessmentModel,
                                network, TS, assessment_uncertainties, tso_actions)
    add_load_uncertainties_vars!(model_container,
                                network, TS, assessment_uncertainties)
    add_prod_uncertainties_vars!(model_container,
                                network, TS, assessment_uncertainties, tso_actions)
    return model_container
end

function add_limitable_prod_vars!(model_container::EODAssessmentModel, network, TS)
    for limitable_gen in Networks.get_generators_of_type(network, Networks.LIMITABLE)
        gen_id = Networks.get_id(limitable_gen)
        for ts in TS
            name =  @sprintf("p_injected[%s,%s]", gen_id, ts)
            model_container.p_injected[gen_id,ts] = @variable(model_container.model, base_name=name)
        end
    end

    return model_container.p_injected
end

function add_use_limitables_constraints!(model_container::EODAssessmentModel, network, TS, tso_actions)
    for limitable_gen in Networks.get_generators_of_type(network, Networks.LIMITABLE)
        gen_id = Networks.get_id(limitable_gen)
        for ts in TS
            limit_l = get_limitation(tso_actions, gen_id, ts)
            @assert(upper_bound(model_container.uncertain_prod[gen_id, ts]) <= limit_l)
            @constraint(model_container.model,
                        model_container.p_injected[gen_id,ts] == model_container.uncertain_prod[gen_id, ts])
        end
    end
    return model_container
end

function add_pilotable_prod_vars!(model_container_p::EODAssessmentModel, network, TS, tso_actions)
    p_injected_l = model_container_p.p_injected
    b_in_l = model_container_p.b_in
    b_marg_l = model_container_p.b_marg
    b_out_l = model_container_p.b_out

    for pilotable_gen in Networks.get_generators_of_type(network, Networks.PILOTABLE)
        gen_id = Networks.get_id(pilotable_gen)
        for ts in TS
            pmin_l, pmax_l = safeget_imposition(tso_actions, gen_id, ts)
            name_l =  @sprintf("b_in[%s,%s]", gen_id, ts);
            b_in_l[gen_id,ts] = @variable(model_container_p.model, base_name = name_l, binary = true)
            name_l =  @sprintf("b_marginal[%s,%s]", gen_id, ts);
            b_marg_l[gen_id,ts] = @variable(model_container_p.model, base_name = name_l, binary = true)
            name_l =  @sprintf("b_out[%s,%s]", gen_id, ts);
            b_out_l[gen_id,ts] = @variable(model_container_p.model, base_name = name_l, binary = true)
            name_l =  @sprintf("prod[%s,%s]", gen_id, ts);
            p_injected_l[gen_id,ts] = @variable(model_container_p.model, base_name = name_l,
                                                lower_bound=0., upper_bound = pmax_l);

            @constraint(model_container_p.model,
                        p_injected_l[gen_id,ts] >= ( pmax_l * b_in_l[gen_id,ts] +  pmin_l * b_marg_l[gen_id,ts] ));
            @constraint(model_container_p.model,
                        p_injected_l[gen_id,ts] <= (1 - b_out_l[gen_id,ts]) * pmax_l);
            @constraint(model_container_p.model,
                        b_out_l[gen_id,ts] + b_marg_l[gen_id,ts] + b_in_l[gen_id,ts] == 1.);
        end
    end

    return model_container_p
end

function add_market_vars!(model_container::EODAssessmentModel, network, TS, assessment_uncertainties, tso_actions)
    add_limitable_prod_vars!(model_container, network, TS)
    add_pilotable_prod_vars!(model_container, network, TS, tso_actions)
end

function add_cut_prod_vars!(model_container::EODAssessmentModel, network, TS)
    for ts in TS
        name_l =  @sprintf("cut_prod[%s]", ts);
        model_container.p_cut_prod[ts] = @variable(model_container.model,
                                                    base_name = name_l, lower_bound=0.)

        prod_expr = sum(model_container.p_injected[gen_id,ts]
                        for gen_id in Networks.get_id.(Networks.get_generators(network)))
        @constraint(model_container.model, model_container.p_cut_prod[ts] <= prod_expr)
    end
    return model_container
end

function add_loss_of_load_vars!(model_container::EODAssessmentModel, network, TS)
    for ts in TS
        name_l =  @sprintf("loss_of_load[%s]", ts);
        model_container.p_loss_of_load[ts] = @variable(model_container.model,
                                                    base_name = name_l, lower_bound=0.)

        conso_expr = sum(model_container.uncertain_load[bus_id,ts]
                        for bus_id in Networks.get_id.(Networks.get_buses(network)))
        @constraint(model_container.model, model_container.p_loss_of_load[ts] <= conso_expr)
    end
    return model_container
end

function add_cheapest_prod_constraints!(model_container::EODAssessmentModel, network, TS)
    pilotables = Networks.get_generators_of_type(network, Networks.PILOTABLE)
    for (index, gen_1) in enumerate(pilotables)
        cost_1 = Networks.get_prop_cost(gen_1)
        gen_id_1 = Networks.get_id(gen_1)
        for gen_2 in pilotables[index:end]
            cost_2 = Networks.get_prop_cost(gen_2)
            if cost_1 < cost_2
                gen_id_2 = Networks.get_id(gen_2)
                for ts in TS
                    @constraint(model_container.model, model_container.b_in[gen_id_1,ts] >= model_container.b_in[gen_id_2,ts] );
                end
            end
        end
    end

    return model_container
end

function add_eod_constraint!(model_container_p::EODAssessmentModel, network, TS)

    for ts in TS

        supply_l = AffExpr(0.)
        for gen in Networks.get_generators(network)
            gen_id = Networks.get_id(gen)
            supply_l += model_container_p.p_injected[gen_id, ts];
        end
        supply_l -= model_container_p.p_cut_prod[ts];

        demand_l = AffExpr(0.)
        for bus in Networks.get_buses(network)
            bus_id = Networks.get_id(bus)
            demand_l += model_container_p.uncertain_load[bus_id, ts];
        end
        demand_l -= model_container_p.p_loss_of_load[ts];

        @constraint(model_container_p.model, supply_l == demand_l);
    end

    return model_container_p
end

function create_objective!(model_container::EODAssessmentModel, network, TS, loss_of_load_coeff, inj_prod_coeff, cut_prod_coeff)
    @assert(loss_of_load_coeff >= 0.)
    @assert(inj_prod_coeff >= 0.)
    @assert(cut_prod_coeff >= 0.)

    model_container.loss_of_load_obj = AffExpr(0.)
    for ts in TS
        add_to_expression!(model_container.loss_of_load_obj, loss_of_load_coeff * model_container.p_loss_of_load[ts])
    end

    model_container.cut_prod_obj = AffExpr(0.)
    for ts in TS
        add_to_expression!(model_container.cut_prod_obj, -cut_prod_coeff * model_container.p_cut_prod[ts])
    end

    model_container.prod_obj = AffExpr(0.)
    #exclude limitables to allow setting their uncertainties at a low level
    # for gen in Networks.get_generators_of_type(network, Networks.PILOTABLE)
    #     gen_id = Networks.get_id(gen)
    #     for ts in TS
    #         add_to_expression!(model_container.prod_obj, inj_prod_coeff * model_container.p_injected[gen_id, ts])
    #     end
    # end

    model_container.full_obj = model_container.loss_of_load_obj + model_container.cut_prod_obj + model_container.prod_obj
    return model_container.full_obj
end

#Only cut conso when no possible extra production
function add_loss_of_load_constraint!(model_container::EODAssessmentModel, network, TS, assessment_uncertainties)
    big_m = 0.
    for bus_id in Networks.get_id.(Networks.get_buses(network))
        big_m += get_assessment_uncertainties_ub(assessment_uncertainties, bus_id)
    end

    for pilotable_gen_id in Networks.get_id.(Networks.get_generators_of_type(network, Networks.PILOTABLE))
        for ts in TS
            b_in_var = model_container.b_in[pilotable_gen_id, ts]
            @constraint(model_container.model, model_container.p_loss_of_load[ts] <= big_m*b_in_var)
        end
    end
end

function add_uncertainties_constraints!(model_container_l::EODAssessmentModel, network, TS, assessment_uncertainties, configs)
    #FIXME constraints on uncertainties go here !
    #e.g | u[bus1] - u[bus2] | < 5
    #e.g ( 1 - u[bus1] / ub[limitable_1] ) <  ( 1 - u[limitable_1] / ub[limitable_1] )
end

function formulate_eod_assessment(network, TS, assessment_uncertainties, tso_actions, configs)
    model_container_l = EODAssessmentModel()

    add_uncertainties_vars!(model_container_l, network, TS, assessment_uncertainties, tso_actions)
    add_market_vars!(model_container_l, network, TS, assessment_uncertainties, tso_actions)
    add_cut_prod_vars!(model_container_l, network, TS)
    add_loss_of_load_vars!(model_container_l, network, TS)

    add_use_limitables_constraints!(model_container_l, network, TS, tso_actions)
    add_cheapest_prod_constraints!(model_container_l, network, TS)
    add_eod_constraint!(model_container_l, network, TS)
    add_loss_of_load_constraint!(model_container_l, network, TS, assessment_uncertainties)
    add_uncertainties_constraints!(model_container_l, network, TS, assessment_uncertainties, configs)

    create_objective!(model_container_l, network, TS,
                    configs.loss_of_load_coeff, configs.cut_prod_coeff, configs.inj_prod_coeff)
    @objective(get_model(model_container_l), Max, model_container_l.full_obj)

    return model_container_l
end

function is_validated(model_container::EODAssessmentModel)
    return value(model_container.loss_of_load_obj) < 1e-09
end

function has_positive_slack(model_container::EODAssessmentModel)::Bool
    return false # ( has_positive_value(model_container.lower.lol_model.p_loss_of_load)
           #  || has_positive_value(model_container.lower.lol_model.p_cut_prod) )
end

function get_assessment_uncertainties_lb(assessment_uncertainties, bus_or_limitable::String)
    return assessment_uncertainties[bus_or_limitable][1]
end
function get_assessment_uncertainties_ub(assessment_uncertainties, bus_or_limitable::String)
    return assessment_uncertainties[bus_or_limitable][2]
end

max_stars_repo_name: klorel/pscopf

max_issues_repo_name: klorel/pscopf

max_forks_repo_name: klorel/pscopf


text:
try
    using Coverage
catch error
    if isa(error, LoadError)
        using Pkg
        Pkg.add("Coverage")
        using Coverage
    else
        rethrow()
    end
end

# Process '*.cov' files
coverage = process_folder() # defaults to src/; alternatively, supply the folder name as argument
coverage = append!(coverage, process_folder("test"))
# Process '*.info' files
coverage = merge_coverage_counts(
    coverage,
    filter!(
        let prefixes = (joinpath(pwd(), "src", ""), joinpath(pwd(), "test", ""))
            c -> any(p -> startswith(c.filename, p), prefixes)
        end,
        LCOV.readfolder("test"),
    ),
)
# Get total coverage for all Julia files
covered_lines, total_lines = get_summary(coverage)
percent = round(Int, 100 * covered_lines / total_lines)
println(
    "Line coverage: $(percent)% ($(covered_lines) covered out of $(total_lines) total lines)",
)


text:
module CommonMetrics

include("common_metrics.jl")

end


text:
using LinearAlgebra
using Polynomials
using SemiLagrangian: Lagrange

struct Pol2{T}
    tab::Array{T,2}
end
function (fct::Pol2{T})(x, y) where {T}
    res = 0
    pc_x = [T(x)^(i - 1) for i = 1:size(fct.tab, 1)]
    pc_y = [T(y)^(i - 1) for i = 1:size(fct.tab, 2)]
    for i = 1:size(fct.tab, 1), j = 1:size(fct.tab, 2)
        res += fct.tab[i, j] * pc_x[i] * pc_y[j]
    end
    return res
end

function test_base_lagrange(order)
    lag = Lagrange(order, Rational{BigInt})
    tab = rationalize.(BigInt, rand(order+1), tol = 1 / 1000000)
    fct = Polynomial(tab)
    dec = div(order, 2)
    for i = 1:3
        value = rationalize.(BigInt, rand(), tol = 1 / 1000000)
        res = 0 // 1
        for j = 0:order
            res += lag.tabfct[j+1](value) * fct(j - dec)
        end
        @test res == fct(value)
    end
end
function test_base_lagrange2d(order)
    lag = Lagrange(order, Rational{BigInt})
    tab = rationalize.(BigInt, rand(order+1, order+1), tol = 1 / 1000000)
    fct = Pol2(tab)
    dec = div(order, 2)
    for i = 1:1
        val_x = rationalize.(BigInt, rand(), tol = 1 / 1000000)
        val_y = rationalize.(BigInt, rand(), tol = 1 / 1000000)
        res = 0 // 1
        for j = 0:order, k = 0:order
            res += lag.tabfct[j+1](val_x) * lag.tabfct[k+1](val_y) * fct(j - dec, k - dec)
        end
        resf = fct(val_x, val_y)
        #        println("order = $order res-fct=$(res-resf)")
        @test res == resf
    end
end
@time @testset "test base interpolation" begin
    for ord = 3:27
        test_base_lagrange(ord)
    end
end
@time @testset "test base interpolation2d" begin
    for ord = 3:21
        test_base_lagrange2d(ord)
    end
end


text:
module TypesBasis

export Basis,BasisFunctions
export Gaussian,GaussOrbitals

using TypesParticles

abstract type Basis end

"""
Datatype for storing an array of basis functions.
""" struct BasisFunctions <: Basis
    nbasisfunc::Int
    basisfunc::Array{Basis,1}
    function BasisFunctions(nbasisfunc,basisfunc)
        if nbasisfunc != length(basisfunc)
            error("Length of basis function array does not equal the number of identified basis functions.")
        else
            new(nbasisfunc,basisfunc)
        end
    end
end #BasisFunction

"""
Contains basis set info for Gaussian Functions

""" struct Gaussian <: Basis
       info::String
       nprims::Int
       coefs::Array{Float64,1}
       alphas::Array{Float64,1}
       function Gaussian(btype,nprims,coefs,alphas)
           if nprims != length(alphas)
               error("Number of alphas not equal to number of primitives.")
           elseif length(coefs) != length(alphas)
               error("Number of contractions not equal to number of alphas.")
           end
           new(btype,nprims,coefs,alphas)
       end
end #Gaussian

"""
Data type for Gaussian based orbitals centered on  ax, ay, az. No contractions or
polarization functiosn just simple function form:

\$\\phi(x) = A e^{-\\alpha\\left(x-x_o\\right)^{2}}\$

""" struct GaussOrbitals
     ax::Float64
     ay::Float64
     az::Float64
     coefs::Array{Float64,1}
     alphas::Array{Float64,1}
     norms::Array{Float64,1}
     function GaussOrbitals(ax,ay,az,basis)
          norms = (2.00e0.*basis.alphas./pi).^(3/4);
          new(ax,ay,az,basis.coefs,basis.alphas,norms);
      end
     function GaussOrbitals(atom::Atom,basis)
          norms = (2.00e0.* basis.alphas ./pi).^(3/4);
          new(atom.x,atom.y,atom.z,basis.coefs,basis.alphas,norms);
      end
end #GausSOrbitals

end #module


text:
function main(r::Robot, cell_size::Int)
    path = go_to_west_south_corner_and_return_path!(r)
    x=0; y=0
    horisontalDirection = East

    while !(isborder(r, North) && isborder(r, East))
        marker_special!(r, x, y, cell_size)
        if move_up_condition(r)
            move!(r, North)
            y += 1
            marker_special!(r, x, y, cell_size)
            horisontalDirection = inverse_side(horisontalDirection)
        end

        move!(r,horisontalDirection)
        (horisontalDirection == East) ? x += 1 : x -= 1
    end

    marker_special!(r, x, y, cell_size)

    go_to_west_south_corner_and_return_path!(r)
    go_by_path!(r, path)
end

function marker_special!(r, x, y, cell_size)
    if (mod(x, 2 * cell_size)) < cell_size && (mod(y, 2 * cell_size)) < cell_size ||
        (mod(x + cell_size, 2 * cell_size)) < cell_size && (mod(y, 2 * cell_size)) >= cell_size
        putmarker!(r)
    end
end

function move_up_condition(r)
    return isborder(r, East) || isborder(r, West) && !(isborder(r, South) && isborder(r, West))
end

max_stars_repo_name: Savko-fokin/mirea-progs

max_issues_repo_name: Savko-fokin/mirea-progs

max_forks_repo_name: Savko-fokin/mirea-progs


text:
import Base: length

struct Link{T<:Integer, U<:Integer}
    i::T
    j::T
    l::U
end

length(link::Link) = link.l
upn(link::Link) = link.i
dwn(link::Link) = link.j

@enum NodeType source sink merge diverge
struct Node{T<:Integer}
    id::T
    type::NodeType
end

struct Network{T<:Integer}
    # representations
    ## forward star
    fadjlist::Vector{Vector{T}}
    fpoints::Vector{T}

    ## backward star
    badjlist::Vector{Vector{T}}
    bpoints::Vector{T}

    # data
    ## links
    links::Vector{Link}

    ## nodes
    nodes::Vector{Union{Missing,Node}}
end

numnodes(net::Network) = length(net.nodes)
numlinks(net::Network) = length(net.links)

sources(net::Network) = filter(i -> length(net.badjlist[i]) == 0, 1:numnodes(net))
sinks(net::Network) = filter(i -> length(net.fadjlist[i]) == 0, 1:numnodes(net))
merges(net::Network) = filter(i -> length(net.badjlist[i]) > 1, 1:numnodes(net))
diverges(net::Network) = filter(i -> length(net.fadjlist[i]) > 1, 1:numnodes(net))

haslink(net::Network{T}, i::U, j::U) where {T<:Integer, U<:Integer} = (j in net.fadjlist[i])

idx(net::Network{T}, i::U, j::U) where {T<:Integer, U<:Integer} = haslink(net, i, j) ? net.fpoints[i]-1+searchsortedfirst(net.fadjlist[i],j) : -1
idx(net::Network{T}, l::Link) where {T<:Integer} = idx(net, l.i, l.j)
outneighbors(net::Network, i) = net.fadjlist[i]
inneighbors(net::Network, j) = net.badjlist[j]

outlinkids(net::Network, i) = collect(net.fpoints[i]:(net.fpoints[i+1]-1))
inlinkids(net::Network, j) = idx.(Ref(net), net.badjlist[j], Ref(j))

outlinks(net::Network, i) = net.links[outlinkids(net, i)]
inlinks(net::Network, j) = net.links[inlinkids(net, j)]

link(net::Network, i) = net.links[i]

function Network(linkdata::DataFrame)
    nnodes = max(maximum(linkdata.i), maximum(linkdata.j))
    net = Network(nnodes)

    for row in eachrow(linkdata)
        add_link!(net, Link(row.i, row.j, row.l))
    end

    for i in 1:nnodes
        if length(net.fpoints[i]) == 0
            net.nodes[i] = Node(i, sink)
        elseif length(net.bpoints[i]) == 0
            net.nodes[i] = Node(i, source)
        elseif length(net.fpoints[i]) > 1
            net.nodes[i] = Node(i, diverge)
        else
            net.nodes[i] = Node(i, merge)
        end
    end

    net
end

function Network(nnodes::T) where {T<:Integer}
    fadjlist = [Vector{T}() for _ in one(T):nnodes]
    fpoints = ones(T, nnodes+1)

    badjlist = [Vector{T}() for _ in one(T):nnodes]
    bpoints = ones(T, nnodes+1)

    links = Vector{Link}()

    nodes = fill(missing, nnodes)

    Network{T}(fadjlist, fpoints, badjlist, bpoints, links, nodes)
end

function add_link!(net, l::Link)
    i,j = l.i, l.j

    @inbounds list = net.badjlist[j]
    index = searchsortedfirst(list, i)
    insert!(list, index, i)
    net.bpoints[(j+1):end] .+= 1

    @inbounds list = net.fadjlist[i]
    index = searchsortedfirst(list, j)
    insert!(list, index, j)
    net.fpoints[(i+1):end] .+= 1

    # add to list
    insert!(net.links, net.fpoints[i]-1+index, l)
end

@with_kw struct LinkState
    trac::AbstractFloat # tracker - floating point value
    #frac::AbstractFloat # fraction of the flow already moved from the cohort of 'trac'
    sval::AbstractFloat # sending value
    rval::AbstractFloat # receiving value
end

tracker(ls::LinkState) = ls.trac
#fracmoved(ls::LinkState) = ls.frac
svalue(ls::LinkState) = ls.sval
rvalue(ls::LinkState) = ls.rval