text:
"""Code to construct puncture initial data for single black hole."""
import sys
from numpy import zeros, size, sqrt, linspace
import scipy.linalg as la

class EllipticSolver:
    """Class Elliptic solves Poisson-type elliptic equations of the form:
          D^2 sol + fct sol = rhs
    where
        - D^2 is the flat Laplace operator
        - fct and rhs are user-supplied functions of the coordinates x, y, z,
        - and sol is the solution.

    To use this class:
        - initialize the class, providing Cartesian coordinates x, y, and z
        - call setup_matrix(fct) to set up the operator
        - call setup_rhs(rhs) to set up the right-hand side
        - then a call to solve() returns the solution sol
    """

    function __init__(self, x, y, z)
        """Constructor - provide Cartesian coordinates, all of length n_grid,
        as arguments.
        """

        print(" Setting up Poisson solver...")
        self.n_grid = size(x)
        self.Δ = x[2] - x[begin]

        # set up storage for matrix, solution, r.h.s.
        # Note: "sol" and "rhs" will store functions in 3d format, while
        # "sol_1d" and "rhs_1d" will store functions in 1d format using
        # super-index
        nnn = self.n_grid ^ 3
        self.rhs_1d = zeros(nnn)
        self.A = zeros(nnn, nnn)
        self.sol = zeros(self.n_grid, self.n_grid, self.n_grid)
        self.rad = zeros(self.n_grid, self.n_grid, self.n_grid)

        # compute radius
        for i in 1:self.n_grid
            for j in 1:self.n_grid
                for k in 1:self.n_grid
                    rad2 = x[i] ^ 2 + y[j] ^ 2 + z[k] ^ 2
                    self.rad[i,j,k] = sqrt(rad2)

    function setup_matrix(self, fct)
        """Set up matrix A."""

        n_grid = self.n_grid

        # Use Robin boundary conditions (B.30) to set up boundaries
        i = 0  # lower x-boundary
        for j in 1:n_grid
            for k in 1:n_grid
                index = self.super_index(i, j, k)
                self.A[index,index] = self.rad[i,j,k]
                self.A[index,index+1] = -self.rad[i+1,j,k]

        i = n_grid - 1  # upper x-boundary
        for j in 1:n_grid
            for k in 1:n_grid
                index = self.super_index(i, j, k)
                self.A[index,index] = self.rad[i,j,k]
                self.A[index,index-1] = -self.rad[i-1,j,k]

        j = 0  # lower y-boundary
        for i in 2:n_grid
            for k in 1:n_grid
                index = self.super_index(i, j, k)
                self.A[index,index] = self.rad[i,j,k]
                self.A[index,index+n_grid] = -self.rad[i,j+1,k]

        j = n_grid - 1  # upper y-boundary
        for i in 2:n_grid
            for k in 1:n_grid
                index = self.super_index(i, j, k)
                self.A[index,index] = self.rad[i,j,k]
                self.A[index,index-n_grid] = -self.rad[i,j-1,k]

        k = 0  # lower z-boundary
        for i in 2:n_grid
            for j in 2:n_grid
                index = self.super_index(i, j, k)
                self.A[index,index] = self.rad[i,j,k]
                self.A[index,index+n_grid*n_grid] = -self.rad[i,j,k+1]

        k = n_grid - 1  # upper z-boundary
        for i in 2:n_grid
            for j in 2:n_grid
                index = self.super_index(i, j, k)
                self.A[index,index] = self.rad[i,j,k]
                self.A[index,index-n_grid*n_grid] = -self.rad[i,j,k-1]

        # use (B.29) to fill matrix in interior
        for i in 2:n_grid
            for j in 2:n_grid
                for k in 2:n_grid
                    index = self.super_index(i, j, k)

                    # diagonal element
                    self.A[index,index] = -6. + self.Δ ^ 2 * fct[i,j,k]

                    # off-diagonal elements
                    self.A[index,index-1] = 1.0
                    self.A[index,index+1] = 1.0
                    self.A[index,index-n_grid] = 1.0
                    self.A[index,index+n_grid] = 1.0
                    self.A[index,index-n_grid*n_grid] = 1.0
                    self.A[index,index+n_grid*n_grid] = 1.0

    function setup_rhs(self, rhs)
        """Setup right-hand side of matrix equation"""

        n_grid = self.n_grid
        for i in 2:n_grid
            for j in 2:n_grid
                for k in 2:n_grid
                    index = self.super_index(i, j, k)
                    self.rhs_1d[index] = self.Δ ^ 2 * rhs[i,j,k]

    function solve(self)
        """Interface to scipy.linalg matrix solver,
        returns sol (in 3d format)."""

        # solve matrix using scipy.linalg interface...
        sol_1d = la.solve(self.A, self.rhs_1d)

        # ... then translate from superindex to 3d
        for i in 1:self.n_grid
            for j in 1:self.n_grid
                for k in 1:self.n_grid
                    index = self.super_index(i, j, k)
                    self.sol[i,j,k] = sol_1d[index]

        return self.sol

    function super_index(self, i, j, k)
        """Compute super index, see (B.28)."""
        return i + self.n_grid * (j + self.n_grid * k)

class Puncture:
    """Class that handles construction of puncture data.

    To use this class,
       - initialize class with physical parameters as arguments
       - then call construct_solution.
    """

    function __init__(self, bh_loc, lin_mom, n_grid, x_out)
        """Arguments to constructor specify physical parameters:
        - location of puncture (bh_loc)
        - linear momentum (lin_mom)
        - size of grid (n_grid)
        - outer boundary (x_out).
        """
        self.bh_loc = bh_loc
        self.lin_mom = lin_mom
        # echo out parameters
        print(" Constructing class Puncture for single black hole")
        print("    at bh_loc = (", bh_loc[begin], ",", bh_loc[2], ",",
                  bh_loc[3], ")")
        print("    with momentum p = (", lin_mom[begin], ",",
                  lin_mom[3], ",", lin_mom[3], ")")
        print(" Using", n_grid,"\b^3 gridpoints with outer boundary at", x_out)
        # set up grid
        self.n_grid = n_grid
        self.x_out = x_out
        self.Δ = 2.0 * x_out / n_grid

        # set up coordinates: use cell-centered grid covering (-x_out, x_out)
        # in each dimension; see (B.14)
        half_Δ = self.Δ / 2.0
        self.x = LinRange(half_Δ - x_out, x_out -
                          half_Δ, n_grid)
        self.y = LinRange(half_Δ - x_out, x_out -
                          half_Δ, n_grid)
        self.z = LinRange(half_Δ - x_out, x_out -
                          half_Δ, n_grid)

        # allocate elliptic solver
        self.solver = EllipticSolver(self.x, self.y, self.z)

        # allocate memory for functions u, α, β, and residual
        self.α = zeros(n_grid, n_grid, n_grid)
        self.β = zeros(n_grid, n_grid, n_grid)
        self.u = zeros(n_grid, n_grid, n_grid)
        self.res = zeros(n_grid, n_grid, n_grid)

    function construct_solution(self, tol, it_max)
        """Construct solution iteratively, provide tolerance and maximum
        number of iterations as arguments."""

        self.setup_α_β()
        residual_norm = self.residual()
        print(" Initial Residual = ", residual_norm)
        print(" Using up to", it_max, "iteration steps to reach tolerance of",
                  tol)

        # now iterate...
        it_step = 0
        while residual_norm > tol and it_step < it_max
            it_step += 1
            self.update_u()
            residual_norm = self.residual()
            print(" Residual after", it_step, "iterations :", residual_norm)
        if (residual_norm < tol)
            print(" Done!")
        else
            print(" Giving up...")

    function update_u(self)
        """Function that updates u using Poisson solver;
        takes one iteration step.
        """

        # set up linear term and right-hand side for SolvePoisson...
        n_grid = self.n_grid
        fct = zeros(n_grid, n_grid, n_grid)
        rhs = zeros(n_grid, n_grid, n_grid)

        for i in 2:n_grid
            for j in 2:n_grid
                for k in 2:n_grid
                    # compute h' from (B.39)
                    temp = self.α[i,j,k] * (1.0 + self.u[i,j,k]) + 1.0
                    fct[i,j,k] = (-7.0 * self.β[i,j,k] *
                                    self.α[i,j,k] / temp ^ 8)
                    rhs[i,j,k] = -self.res[i,j,k]

        # now update Poisson solver
        self.solver.setup_matrix(fct)

        # set up right-hand side
        self.solver.setup_rhs(rhs)

        # solve to find Δ_u, see (B.36)
        Δ_u = self.solver.solve()

        # update u
        self.u += Δ_u

    function residual(self)
        """Evaluate residual, see (B.35)."""

        residual_norm = 0.0
        for i in 2:self.n_grid
            for j in 2:self.n_grid
                for k in 2:self.n_grid

                    # compute left-hand side: Laplace operator
                    ddx = (self.u[i+1,j,k] - 2.0 * self.u[i,j,k] +
                           self.u[i-1,j,k])
                    ddy = (self.u[i,j+1,k] - 2.0 * self.u[i,j,k] +
                           self.u[i,j-1,k])
                    ddz = (self.u[i,j,k+1] - 2.0 * self.u[i,j,k] +
                           self.u[i,j,k-1])
                    lhs = (ddx + ddy + ddz) / self.Δ ^ 2

                    # compute right-hand side,
                    # recall h = - β/(α + α u + 1)^7
                    temp = self.α[i,j,k] * (1.0 + self.u[i,j,k]) + 1.0
                    rhs = -self.β[i,j,k] / temp ^ 7

                    # then compute difference to get residual, see (B.35)
                    self.res[i,j,k] = lhs - rhs
                    residual_norm += self.res[i,j,k] ^ 2

        residual_norm = sqrt(residual_norm) * self.Δ ^ 3
        return residual_norm

    function setup_α_β(self)
        """Set up functions α and β."""

        n_grid = self.n_grid
        p_x = self.lin_mom[begin]
        p_y = self.lin_mom[2]
        p_z = self.lin_mom[3]

        for i in 1:n_grid
            for j in 1:n_grid
                for k in 1:n_grid
                    s_x = self.x[i] - self.bh_loc[begin]
                    s_y = self.y[j] - self.bh_loc[2]
                    s_z = self.z[k] - self.bh_loc[3]
                    s2 = s_x ^ 2 + s_y ^ 2 + s_z ^ 2
                    s_bh = sqrt(s2)
                    l_x = s_x / s_bh
                    l_y = s_y / s_bh
                    l_z = s_z / s_bh
                    lP = l_x * p_x + l_y * p_y + l_z * p_z

                    # construct extrinsic curvature, see Eq. (3.43)
                    fac = 3.0 / (2.0 * s2)
                    A_xx = fac * (2.0 * p_x * l_x - (1.0 - l_x * l_x) * lP)
                    A_yy = fac * (2.0 * p_y * l_y - (1.0 - l_y * l_y) * lP)
                    A_zz = fac * (2.0 * p_z * l_z - (1.0 - l_z * l_z) * lP)
                    A_xy = fac * (p_x * l_y + p_y * l_x + l_x * l_y * lP)
                    A_xz = fac * (p_x * l_z + p_z * l_x + l_x * l_z * lP)
                    A_yz = fac * (p_y * l_z + p_z * l_y + l_y * l_z * lP)

                    # compute A_{ij} A^{ij}
                    A2 = (
                        A_xx ^ 2 + A_yy ^ 2 + A_zz ^ 2 +
                        2.0*(A_xy ^ 2 + A_xz ^ 2 + A_yz ^ 2)
                        )

                    # now compute α and β from (3.47) and (3.49)
                    self.α[i,j,k] = 2.0 * s_bh
                    self.β[i,j,k] = self.α[i,j,k] ^ 7 * A2 / 8.0

    function write_to_file(self)
        """Function that writes solution to file."""

        filename = "Puncture_" + str(self.n_grid) + "_" + str(self.x_out)
        filename = filename + ".data"
        out = open(filename, "w")
        if out
            k = self.n_grid // 2
            out.write(
                "# Data for black hole at x = (%f,%f,%f)\n"
                % (self.bh_loc[begin], self.bh_loc[3], self.bh_loc[3])
            )
            out.write("# with linear momentum P = (%f, %f, %f)\n" %
                          (self.lin_mom))
            out.write("# in plane for z = %e \n" % (self.z[k]))
            out.write("# x            y              u              \n")
            out.write("#============================================\n")
            for i in 1:self.n_grid
                for j in 1:self.n_grid
                    out.write("%e  %e  %e\n" % (self.x[i], self.y[j],
                                                self.u[i,j,k]))
                out.write("\n")
            out.close()
        else
            print(" Could not open file", filename,"in write_to_file()")
            print(" Check permissions?")
#
#=====================================================================
# Main routine: defines parameters, sets up puncture solver, and
# then finds solution
#=====================================================================
#
function main()
    """Main routine..."""
    print(" -------------------------------------------------------")
    print(" --- puncture.py --- use flag -h for list of options ---")
    print(" -------------------------------------------------------")
    #
    # set default values for variables
    #
    # location of black hole:
    loc_x = 0.0
    loc_y = 0.0
    loc_z = 0.0
    # momentum of black hole:
    p_x = 1.0
    p_y = 0.0
    p_z = 0.0
    # number of grid points
    n_grid = 16
    # location of outer boundary
    x_out = 4.0
    # tolerance and maximum number of iterations
    tol = 1.0e-12
    it_max = 50
    #
    # now look for flags to overwrite default values
    #
1:length(sys.argv)
        if sys.argv[i] == "-h"
            usage()
            return
        if sys.argv[i] == "-n_grid"
            n_grid = int(sys.argv[i+1])
        if sys.argv[i] == "-x_out"
            x_out = float(sys.argv[i+1])
        if sys.argv[i] == "-loc_x"
            loc_x = float(sys.argv[i+1])
        if sys.argv[i] == "-loc_y"
            loc_y = float(sys.argv[i+1])
        if sys.argv[i] == "-loc_z"
            loc_z = float(sys.argv[i+1])
        if sys.argv[i] == "-p_x"
            p_x = float(sys.argv[i+1])
        if sys.argv[i] == "-p_y"
            p_y = float(sys.argv[i+1])
        if sys.argv[i] == "-p_z"
            p_z = float(sys.argv[i+1])
        if sys.argv[i] == "-tol"
            tol = float(sys.argv[i+1])
        if sys.argv[i] == "-it_max"
            it_max = int(sys.argv[i+1])

    # location of puncture
    bh_loc = ( loc_x, loc_y, loc_z )
    # linear momentum
    lin_mom = ( p_x, p_y, p_z )
    #
    # set up Puncture solver
    black_hole = Puncture(bh_loc, lin_mom, n_grid, x_out)
    #
    # and construct solution
    black_hole.construct_solution(tol, it_max)
    #
    # and write results to file
    black_hole.write_to_file()

function usage()
    print("Constructs puncture initial data for single black hole.")
    print("")
    print("The following options can be used to over-write default parameters")
    print("\t-n_grid: number of grid points [default:16]")
    print("\t-x_out: location of outer boundary [4.0]")
    print("\t-loc_x, -loc_y, -loc_z: location of black hole [(0.0,0.0,0.0)]")
    print("\t-p_x, -p_y, -p_z: lin. momentum of black hole [(1.0,0.0,0.0)]")
    print("\t-tol: tolerance for elliptic solver [1.e-12]")
    print("\t-it_max: maximum number of iterations [51]")
    print("For example, to construct data with x_out = 6.0, call")
    print("\tpython3 puncture.py -x_out 6.0")

if __name__ == '__main__'
    main()


text:
module AccurateArithmetic

export add_acc, add_acc_hilo, add₊, add_hilo₊,
       sub_acc, sub_acc_hilo, sub_acc_lohi, sub₊, sub_hilo₊, sun_lohi₊,
       sqr_acc, mul_acc, sqr₊, mul₊,
       inv_acc, div_acc, inv₊, div₊,
       sqrt_acc, sqrt₊

import Base.IEEEFloat # Union{Float64, Float32, Float16}

end # module AccurateArithmetic

ext: jl


text:
using Base.Test
using GpSelection

using Distributions
using MatrixDepot
using Combinatorics

include("../src/utils.jl")
include("../src/sampledata.jl")
include("../src/varBayes.jl")

n = 100;  K = 5; num_sig = 2; df = 1.0; p_beta = 2
dat = generateData(K, p_beta, n, df, num_sig)
options = Options(100)
priors = PriorPars(5, 10, 10, fill(5, K), fill(10, K), 0.0, zeros(K, K))

@elapsed beta_E, sigma2_E, gamma_E, H_E = varBayes(dat, priors, options)


text:
# # BAT.jl plotting tutorial

using BAT
using Distributions
using IntervalSets

# ## Generate samples to be plotted
likelihood = params -> begin

    r1 = logpdf.(
    MixtureModel(Normal[
    Normal(-10.0, 1.2),
    Normal(0.0, 1.8),
    Normal(10.0, 2.5)], [0.1, 0.3, 0.6]), params.a)

    r2 = logpdf.(
    MixtureModel(Normal[
    Normal(-5.0, 2.2),
    Normal(5.0, 1.5)], [0.3, 0.7]), params.b[1])

    r3 = logpdf.(Normal(2.0, 1.5), params.b[2])

    return LogDVal(r1+r2+r3)
end

prior = BAT.NamedTupleDist(
    a = Normal(-3, 4.5),
    b = [-20.0..20.0, -10..10]
)

posterior = PosteriorDensity(likelihood, prior);

samples, chains = bat_sample(posterior, MCMCSampling(mcalg = MetropolisHastings(), nsteps = 10^5));

# ## Set up plotting
# Set up plotting using the [Plots.jl](https://github.com/JuliaPlots/Plots.jl) package:
using Plots

# The BAT.jl plotting recipes have been tested with `pyplot()`, `gr()` and `plotly()` backends.
# All available plotting options work correctly with `pyplot()`.
# Only colored 2D contour plots are currently not correctly supported by the `gr()` and `plotly()` backends.
# The backend can be chosen using:

#gr() # Julia's default
#pyplot()
#plotly()

# ## Examples for 1D plots
# BAT.jl includes different plotting styles (seriestypes) for 1D representations of samples and priors:
#   * `:smallest_intervals` (alias `:HDR`): highlighting the smallest intervals containing a certain probability (highest density regions) by color
#   * `:central_intervals`: highlighting the central intervals containing a certain probability by color
#   * `:histogram` (alias `:steppost`): filled histogram
#   * `:stephist`: step histogram
# Examples of (marginalized) 1D representations of samples and priors are shown below.

# ### Default 1D plot of samples:
# Samples can be plotted either by their index or by using the parameter name (as a symbol or expression) as specified in the prior:

plot(samples, :a) #default seriestype = :smallest_intervals (alias :HDR)
#or: plot(samples, 2)
# The default seriestype for plotting samples is `:smallest_intervals` (alias `:HDR`), highlighting the smallest intervals (the highest density region) containing 68.3, 95.5 and 99.7 perecent of the posterior probability. By default, the marginal mode(s) of the histogram is(are) indicated as dotted black line(s).

# ### Default 1D plot of prior:
# Priors can be plotted either by their index or by using the parameter name:

plot(prior, :a)
#or: plot(prior, 1)

# ## Knowledge update plot
# The knowledge update after performing the sampling can be visualized by plotting the prior and the samples of the psterior together in one plot using `plot!()`:
plot(samples, :(b[1]))
plot!(prior, :(b[1]))

# ## Further seriestypes for 1D plots (samples and prior):
# ### central intervals:
plot(samples, :a, seriestype = :central_intervals)

# ### filled histogram:
plot(samples, :a, seriestype = :histogram) # alias :hist

# ### step histogram:
plot(samples, :a, seriestype = :stephist)

# ## Customizing 1D plots:

# ### Keyword arguments for [attributes supported by *Plots.jl*](https://docs.juliaplots.org/latest/attributes/#attributes-1) can be passed:
plot(samples, :a, seriestype = :stephist, bins=50, linecolor = :red, linewidth = 5, linealpha=0.4, xlim=(-10,10))

# ### Customizing interval plots:
# For `:smallest_intervals` and `:central_intervals` plot, the probability enclosed in the intervals to be highlighted can be specified using the `intervals` keyword.
# The keyword `interval_labels` allows to specify the legend entries for the corresponding intervals (in same order).
# The interval colors need to be specified (in same order) using the `colors` keyword argument.
plot(samples, :a, seriestype=:smallest_intervals,
    intervals=[0.5, 0.1, 0.3, 0.99],
    interval_labels = ["HDR 50%", "HDR 10%", "HDR 30%", "HDR 99%"],
    colors=[:grey, :red, :blue, :orange]
)

# ### Plotting estimators in 1D sample plots:
# It is possible to indicate *mean*, *standard deviation*, *marginalmode*, and *globalmode* when plotting samples:
plot(samples, :(b[1]), globalmode=true, marginalmode=true, mean=true, std=true)

# #### The style of the estimators can be customized:
# By passing `true`, the point estimators are plotted using their default styles shown above.
# The styles can be modified by passing dictionaries specifying `linestyle`, `linecolor`, `linewidth` and `alpha` for *mean*, *globalmode* and *marginalmode*.
# The style of the standard deviation can be modified by specifying `fillcolor` and `fillalpha`.
plot(samples, :a, marginalmode=false,
    mean=Dict("linestyle" => :dot, "linecolor"=> :red, "linewidth"=>2, "alpha" => 0.7),
    std=Dict("fillcolor" => :red, "fillalpha"=>0.2))

# ## Examples for 2D plots of samples and priors
# Below, all available seriestypes and plotting features for 2D representations of samples and priors are shown
#   * `:smallest_intervals` (alias `:HDR`): histogram of the smallest intervals containing a certain probability (highest density regions)
#   * `:smallest_intervals_contour`: (colored) contours of the smallest intervals containing a certain probability (highest density regions)
#   * `:smallest__intervals_contourf`: (colored) filled contours of the smallest intervals containing a certain probability (highest density regions)
#   * `:histogram2d` (alias `:histogram`, `:hist`): 2D histogram

# ### Default 2D plot  of samples:
pyplot()
plot(samples, (:a, :(b[1])), mean=true, std=true) #default seriestype = :smallest_intervals (alias :HDR)
# The default seriestype for plotting samples is a 3-color heatmap showing the smallest intervals (highest density regions) containing 68.3%, 95.5% and 99.7% of the posterior probability. By default, the marginal mode
# of the histogram is indicated by a black square.

# ### Default 2D plot  of priors:
plot(prior, (:a, :(b[1]))) # default seriestype :smallest_intervals_contour

# ## 2D knowledge update plot
plot(samples, (:a, :(b[2])))
plot!(prior, (:a, :(b[2])))

# ## Further seriestypes for 2D plots (samples and prior):
# ### histogram:
plot(samples, (:a,:(b[2])), seriestype = :histogram)

# ### smallest intervals as colored contour lines:
# (currently only correctly supported with `pyplot()` backend)
plot(samples, (:a,:(b[2])), seriestype=:smallest_intervals_contour, bins=40)

# ### smallest intervals as filled contours with smoothing:
# (currently only correctly supported with `pyplot()` backend)
plot(samples, (:a,:(b[2])), seriestype=:smallest_intervals_contourf, bins=40, smoothing=1)

# ### Customizing smallest interval plots:
# The probability intervals to be highlighted can be specified using the `intervals` keyword.
# The keyword `interval_labels` allows to specify the legend entries for the corresponding intervals (in same order).
# The interval colors need to be specified (in same order) using the `colors` keyword argument.
plot(samples, (:a,:(b[2])), seriestype=:smallest_intervals, bins=200,
intervals=[0.7, 0.2],
interval_labels = ["HDR 70%", "HDR 20%"],
colors=[:blue, :red])

# ### marginal plot:
plot(samples, (:a, :(b[2])), seriestype = :marginal)

# ### scatter plot:
# (Only plotting the first 10.000 sample points, as for a large number of samples plot takes a lot of time and resources.)
gr()
plot(samples[1:10^4], (:a,:(b[2])), seriestype = :scatter)

# ## Customizing 2D plots:
# ### Plotting point estimators in 2D plots:
plot(samples, (:a,:(b[2])), seriestype=:smallest_intervals, nbins=200,
    mean=true, std=true, marginalmode=true, globalmode=true)

# #### It is possible to customize the style of the estimators:
# By passing `true`, the point estimators are plotted using their default styles shown above.
# The style of the point estimators *mean*, *globalmode* and *marginalmode* can be modified by passing a dictionary specifying `markershape`, `markercolor`, `markersize`, `markeralpha`, `markerstrokecolor`, `markerstrokestyle`, `markerstrokewidth` and `markerstrokealpha`.
# If `std==true`, the standard deviation of the mean value will be displayed as x- and y-errorbars.
plot(samples, (:a,:(b[2])), seriestype=:smallest_intervals, bins=(200, -4:0.2:8),
    marginalmode=Dict("markershape"=> :diamond, "markeralpha"=>1, "markercolor"=>:red, "markersize"=>5),
    mean = true, std=true
)

# ### Customizing marginal plots:
# Marginal plots can be modified by passing dictionaries to the keyword arguments `upper`, `right` and `diagonal`.
# The dictionaries for `upper` and `right` can contain the seriestypes and plot options for 1D distributions shown above.
# The dictionary for `diagonal` can use the seriestypes and plot options for 2D plots shown above.
plot(samples, (:a,:(b[2])), seriestype = :marginal,
    diagonal = Dict("seriestype"=>:histogram),
    upper=Dict("seriestype" => :smallest_intervals, "colors"=>[:blue, :grey, :orange]),
    right=Dict("seriestype" => :stephist)
)

# ## Overview plots
# Without specifying the parameters to be plotted, the 1D and 2D plots of (up to) the first 5 parameters are displayed as a grid.
plot(samples)

plot(prior)

plot(samples)
plot!(prior)

# The keyword argument `vsel` allows to specify which parameters to consider in the overview plot by passing the names or indeces:
plot(samples, vsel=(:a, :(b[2])))
# or: plot(samples, vsel=(1, 3))

# The binning of the overview plots can be specified individually for each parameter:
plot(samples, bins=200) # all paramaters have 200 bins (default)
plot(samples, bins=(10, -5:0.5:8, 500)) #  specify different bin numbers or bin edges for each parameter (only those specified with vsel)
plot(prior, bins=(a=50,)) # use 50 bins for parameter :a and the default value (200) for all others

# ### Customizing overview plots:
# The overview plots can be modified by passing dictionaries to the keyword arguments `upper`, `lower` and `diagonal`.
# The dictionaries for `upper` and `lower` can contain the 2D seriestypes and plot options shown above.
# The dictionary for `diagonal` can use the 1D seriestypes and plot options shown above.
# Nested dictonaries are possible (e.g. for modifying point estimators)
plot(samples, mean=true, globalmode=true, legend=true, bins = (a = -15:2:15, ),
    diagonal=Dict("seriestype"=>:stephist, "mean"=>Dict("linecolor" => :green, "linewidth" => 8)),
    lower = Dict("mean" => false, "colors"=>[:orange, :green, :grey]))

# ## Plots for MCMC diagnostics
# The diagnostic plots allow to plot the samples, a kernel density estimate, the trace and the autocorrelation function for each parameter and each chain:
# The parameters to be considered can be chosen with the `vsel` keyword.
diagnostics = BAT.MCMCDiagnostics(samples, chains)
pyplot()
plot(diagnostics, vsel=[1])

# ### Customizing diagnostics plots:
# #### available keyword arguments:
# * `vsel` - list of parameters to be plotted
# * `chains` - list of chains to be plotted
# * `diagnostics` - list of MCMC diagnostics to be plotted
#   * `:histogram` - 1D histograms of samples
#   * `:kde` - Kernel density estimate (using [*KernelDensity.jl*](https://github.com/JuliaStats/KernelDensity.jl))
#   * `:trace` - Trace plot
#   * `:acf` - Autocorrelation function (using [*StatsBase.autocor*](http://juliastats.github.io/StatsBase.jl/stable/signalcorr/#StatsBase.autocor))
# * `description::Bool = true` - show description (current chain, parameter, number of samples) as first column of plots
# * `histogram::Dict` - options for histogram plots (supports all arguments for 1D plots for samples)
# * `kde::Dict` - options for kde plots
# * `trace::Dict` - options for trace plots
# * `acf::Dict` - options for acf plots
# #### special options arguments for `:kde` (see [*KernelDensity.jl*](https://github.com/JuliaStats/KernelDensity.jl))
# * `npoints`: number of interpolation points to use (default: npoints = 2048)
# * `boundary`: lower and upper limits of the kde as a tuple
# * `kernel`: the distributional family from [*Distributions.jl*](https://github.com/JuliaStats/Distributions.jl) to use as the kernel (default = Distributions.Normal)
# * `bandwidth`: bandwidth of the kernel
# #### special keyword arguments for `:acf` (see [*StatsBase.autocor*](http://juliastats.github.io/StatsBase.jl/stable/signalcorr/#StatsBase.autocor))
# * `lags` - list of lags to be considered for ACF plots
# * `demean` - denotes whether the mean should be subtracted before computing the ACF
plot(diagnostics,
    vsel=[1, 2],
    chains=[1, 2],
    diagnostics = [:histogram, :kde, :trace, :acf],
    histogram = Dict("seriestype" => :smallest_intervals, "legend" => :false),
    trace = Dict("linecolor" => :red),
    acf = Dict("lags" => collect(1:20), "title"=> "Autocorrelation"),
    description = true)


text:
export Frequency, Time, Distance, Velocity
export toHz, tos, tom, tomps

# ----- Frequency and Time ----------------------
const Frequency = Union{Real, Unitful.Frequency}

toHz(f::Real) = Float64(f)
toHz(f::Unitful.Frequency) = ustrip(Float64, Hz, f)
toHz(fa::AbstractArray{<:Frequency}) = toHz.(fa)
toHz(fr::AbstractRange{<:Unitful.Frequency}) = toHz(minimum(fr)):toHz(step(fr)):toHz(maximum(fr))

const Time = Union{Real, Unitful.Time}

tos(t::Real) = Float64(t)
tos(t::Unitful.Time) = ustrip(Float64, s, t)
tos(ta::AbstractArray{<:Time}) = tos.(ta)
tos(tr::AbstractRange{<:Unitful.Time}) = tos(minimum(tr)):tos(step(tr)):tos(maximum(tr))

const Distance = Union{Real, Unitful.Length}

tom(s::Real) = Float64(s)
tom(s::Unitful.Length) = ustrip(Float64, u"m", s)
tom(sa::AbstractArray{<:Distance}) = tom.(sa)
tom(sr::AbstractRange{<:Unitful.Length}) = tom(minimum(sr)):tom(step(sr)):tom(maximum(sr))

const Velocity = Union{Real, Unitful.Velocity}

tomps(v::Velocity) = Float64(v)
tomps(v::Unitful.Velocity) = ustrip(Float64, u"m/s", v)
tomps(va::AbstractArray{<:Velocity}) = tomps.(va)
tomps(vr::AbstractRange{<:Unitful.Velocity}) = tomps(minimum(vr)):tomps(step(vr)):tomps(maximum(vr))


text:
include("./src/GBDTsfuzzy.jl")
include("./MyGrammar.jl")

using .GBDTsfuzzy
using .MyGrammar
using ArgParse
using MultivariateTimeSeries
using Printf
import JLD
using Random; Random.seed!(1)

println("Ready to accept")
while !eof(stdin)
    input = readline(stdin)
    splitted = split(input)
    data = String(splitted[1])
    model_path = String(splitted[2])
    println("Reading data")
    X, y = read_data_labeled(data)
    println("Finished reading")

    println("Reading model")

    m = JLD.load(model_path)

    println(m)
    global v = m["v"]
    model = m["model"]
    mode = m["mode"]

    println("Model loaded")

    println("Classify")
#    if mode == "normal"
#        pred = classify(model, X)
#    else
    pred = classifyfuzzy(model, X)
#    end

    println("Classification done")
    println(pred)

    T = Vector{Float64}(undef, length(X))

#    if mode == "normal"
#        for i in eachindex(T)
#            if pred[i] == model.void_label
#                T[i] = 1.0
#            else
#                T[i] = 0.0
#            end
#        end
#    else
    for i in eachindex(T)
        val = get(pred[i], model.void_label, 0.0)
        T[i] = val
    end
 #   end

    println("Wrongness:")
    println(T)
    println("Done")
end


text:
## PART 6

println("\nPart 6");
# Calculate the true covariance matrix for the data

println("\n\tCalculating the true covariance matrix, ĉ:");
D = v - ones(M) * transpose(v̄̂);
ĉ = transpose(D) * D ./ M;

println("\n\tCalculating the normalized covariance matrix, ρ̂:");
ρ̂ = ĉ ./ ( σ * transpose(σ) )
writedlm("p6-covariance.dat", ρ̂)


text:
# logarithmic barrier function

immutable LogBarrier{T <: Real} <: ProximableFunction
  a::T
  b::T
  mu::T
  function LogBarrier(a::T, b::T, mu::T)
    if mu <= 0
      error("parameter mu must be positive")
    else
      new(a, b, mu)
    end
  end
end

"""
  LogBarrier(a::Real=1.0, b::Real=0.0, mu::Real=1.0)

Returns the function `g(x) = -mu*sum(log(a*x_i+b), i=1,...,n)`.
"""

LogBarrier{T <: Real}(a::T=1.0, b::T=0.0, mu::T=1.0) = LogBarrier{T}(a, b, mu)

function (f::LogBarrier){T <: Real}(x::AbstractArray{T,1})
  sumf = 0.0
  v = 0.0
  for i in eachindex(x)
    v = f.a*x[i]+f.b
    if v <= 0.0
      return +Inf
    end
    sumf += log(v)
  end
  return -f.mu*sumf
end

function prox!{T <: Real}(f::LogBarrier, x::AbstractArray{T,1}, y::AbstractArray{T}, gamma::Real=1.0)
  par = 4*gamma*f.mu*f.a*f.a
  sumf = 0.0
  z = 0.0
  v = 0.0
  for i in eachindex(x)
    z = f.a*x[i] + f.b
    v = (z + sqrt(z*z + par))/2
    y[i] = (v - f.b)/f.a
    sumf += log(v)
  end
  return -f.mu*sumf
end

fun_name(f::LogBarrier) = "logarithmic barrier"
fun_dom(f::LogBarrier) = "AbstractArray{Real}"
fun_expr(f::LogBarrier) = "x ↦ -μ * sum( log(a*x_i+b), i=1,...,n )"
fun_params(f::LogBarrier) = "a = $(f.a), b = $(f.b), μ = $(f.mu)"

function prox_naive{T <: Real}(f::LogBarrier, x::AbstractArray{T,1}, gamma::Real=1.0)
  asqr = f.a*f.a
  z = f.a*x + f.b
  y = ((z + sqrt.(z.*z + 4*gamma*f.mu*asqr))/2 - f.b)/f.a
  fy = -f.mu * sum(log.(f.a*y+f.b))
  return y, fy
end


text:
export AbstractNonlinearDevice, χ3Device
export setup_χ!, solve
export IterativeMethod, IterativeMethodGN, IterativeMethodB

@enum IterativeMethod IterativeMethodGN=1 IterativeMethodB=2

abstract type AbstractNonlinearDevice{D} <: AbstractDevice{D} end

mutable struct χ3Device{D} <: AbstractNonlinearDevice{D}
    grid::Grid{D}
    ϵᵣ::Array{Complex}
    χ::Array{Float}
    src::Array{Complex}
    ω::AbstractVector{Float}
    modes::Array{Mode}
end

"    χ3Device(grid::Grid, ω::AbstractVector{Float})"
function χ3Device(grid::Grid, ω::AbstractVector{Float})
    ϵᵣ = ones(Complex, size(grid));
    χ = zeros(Float, size(grid));
    src = zeros(Complex, size(grid));
    return χ3Device(grid, ϵᵣ, χ, src, ω, Array{Mode}(0))
end

"    χ3Device(grid::Grid, ω::Float)"
χ3Device(grid::Grid, ω::Float) = χ3Device(grid::Grid, [ω])

"    setup_χ!(d::χ3Device, shapes::AbstractVector{<:Shape})"
setup_χ!(d::χ3Device, shapes::AbstractVector{<:Shape}) = _compose_shapes!(d.χ, d.grid, shapes)

"    setup_χ!(d::χ3Device, region, value)"
setup_χ!(d::χ3Device, region, value) = _mask_values!(d.χ, d.grid, region, value)

"    solve(d::χ3Device, which_method::IterativeMethod)"
function solve(d::χ3Device, which_method::IterativeMethod)
    (ϵ₀, μ₀, c₀) = normalize_parameters(d);

    Nω = length(d.ω);
    fields = Array{FieldTM}(Nω);

    for i in eachindex(d.ω)
        @info "Frequency: $i/$Nω"
        ω = d.ω[i];

        length(d.modes) > 0 && ( d.src=zeros(Complex, size(d.grid)) ) # Only reset if we are using modes
        for mode in d.modes
            # TODO: handle the polarization here
            setup_mode!(d, TM, ω, mode.neff, mode.pt, mode.dir, mode.width);
        end

        Tϵ = sparse(Diagonal((ϵ₀*d.ϵᵣ[:])));

        Hx = zeros(ComplexF64, size(d.grid));
        Hy = zeros(ComplexF64, size(d.grid));
        Ez = zeros(ComplexF64, size(d.grid));

        (Sxf, Sxb, Syf, Syb) = S_create(d.grid, ω);

        # Construct derivates
        δxb = Sxb*δ(x̂, Backward, d.grid);
        δxf = Sxf*δ(x̂, Forward,  d.grid);
        δyb = Syb*δ(ŷ, Backward, d.grid);
        δyf = Syf*δ(ŷ, Forward,  d.grid);

        A = δxf*μ₀^-1*δxb + δyf*μ₀^-1*δyb + ω^2*Tϵ;
        b = 1im*ω*d.src[:];
        @info "Solving linear system"
        ez = dolinearsolve(A, b, CNSym);

        coeff = ω^2*ϵ₀*3*d.χ[:]/d.grid.L₀;

        if which_method == IterativeMethodB
            @info "Starting nonlinear iteration using Born"
            (ez, err) = _doborn(ez, A, b, coeff);
        end
        if which_method == IterativeMethodGN
            @info "Starting nonlinear iteration using Gauss-Newton"
            (ez, err) = _donewton(ez, A, b, coeff);
        end

        hx = -1/1im/ω/μ₀*δyb*ez;
        hy = 1/1im/ω/μ₀*δxb*ez;

        fields[i] = FieldTM(d.grid, ω, ez, hx, hy)
    end

    Nω == 1 && return fields[1]
    return fields
end

function _doborn(ez, A, b, coeff; tol = 1e-12, maxiterations = 50)
    i = 1;
    err = [1.0];
    while err[end] > tol && i <= maxiterations
        @debug "iteration number: $i"
        ez_new = dolinearsolve(A + sparse(Diagonal(coeff.*ez.*conj.(ez)), b, CNSym));
        append!(err, norm(ez_new - ez)/norm(ez));

        ez = ez_new;
        i += 1;

        @debug "step error: $(err[end])"
    end
    return (ez, err)
end

function _donewton(ez, A, b, coeff; tol = 1e-12, maxiterations = 50)
    i = 1;
    err = [1.0];
    M = length(ez);

    ez = [ez; conj.(ez)];
    while err[end] > tol && i <= maxiterations
        @debug "iteration number: $i"
        F = (A + sparse(Diagonal(coeff.*ez[1:M].*conj.(ez[1:M]))))*ez[1:M] - b;
        J1  = A + sparse(Diagonal(2*coeff.*conj.(ez[1:M]).*ez[1:M]));
        J2 = sparse(Diagonal(coeff.*ez[1:M].*ez[1:M]));
        J = [J1 J2; conj.(J2) conj.(J1)];
        δez = dolinearsolve(J, [F; conj.(F)], CNSym);
        normez = norm(ez);
        ez = ez - δez;
        append!(err, norm(δez)/normez);
        i += 1;
        @debug "step error: $(err[end])"
    end
    ez = ez[1:M];
    return (ez, err)
end


text:
#DBG import Knet.Ops21: mmul
using Knet.KnetArrays: DevArray
using CUDA.CUBLAS: CUBLAS, cublasDgemm_v2, cublasSgemm_v2, cublasOperation_t
using AutoGrad: AutoGrad, @primitive1

function mmul(x1::R, x2::R; dims=1) where {T,R<:DevArray{T}}
    @assert ndims(x1) > dims "ndims(w) must be > dims"
    ntuple(i->size(x1,ndims(x1)-dims+i),dims) === ntuple(i->size(x2,i),dims) || throw(DimensionMismatch("mmul: w=$(size(x1)) x=$(size(x2)) dims=$dims"))
    t1,t2 = cublasOperation_t.((0,0))
    msize = (size(x1,i) for i in 1:ndims(x1)-dims)
    nsize = (size(x2,i) for i in 1:dims)
    ksize = (size(x2,i) for i in dims+1:ndims(x2))
    m,n,k = prod.((msize, nsize, ksize))
    y = similar(x1, (msize..., ksize...))
    if T<:Float64
        cublasDgemm_v2(CUBLAS.handle(), t1, t2, m, n, k, 1, x1, m, x2, n, 0, y, m)
    elseif T<:Float32
        cublasSgemm_v2(CUBLAS.handle(), t1, t2, m, n, k, 1, x1, m, x2, n, 0, y, m)
    else
        error("CUBLAS does not support $T")
    end
    return y
end

function _mmul1(x1::R, x2::R; dims=1) where {T,R<:DevArray{T}} # x1 transposed
    @assert ntuple(i->size(x1,i),dims) === ntuple(i->size(x2,i),dims)
    t1,t2 = cublasOperation_t.((1,0))
    msize = (size(x1,i) for i in dims+1:ndims(x1))
    nsize = (size(x1,i) for i in 1:dims)
    ksize = (size(x2,i) for i in dims+1:ndims(x2))
    m,n,k = prod.((msize, nsize, ksize))
    y = similar(x1, (msize..., ksize...))
    if T<:Float64
        cublasDgemm_v2(CUBLAS.handle(), t1, t2, m, n, k, 1, x1, m, x2, n, 0, y, m)
    elseif T<:Float32
        cublasSgemm_v2(CUBLAS.handle(), t1, t2, m, n, k, 1, x1, m, x2, n, 0, y, m)
    else
        error("CUBLAS does not support $T")
    end
    return y
end

function _mmul2(x1::R, x2::R; dims=1) where {T,R<:DevArray{T}} # x2 transposed
    @assert ndims(x1) > dims && ndims(x2) > dims
    @assert ntuple(i->size(x1,ndims(x1)-dims+i),dims) === ntuple(i->size(x2,ndims(x2)-dims+i),dims)
    t1,t2 = cublasOperation_t.((0,1))
    msize = (size(x1,i) for i in 1:ndims(x1)-dims)
    nsize = (size(x1,i) for i in 1+ndims(x1)-dims:ndims(x1))
    ksize = (size(x2,i) for i in 1:ndims(x2)-dims)
    m,n,k = prod.((msize, nsize, ksize))
    y = similar(x1, (msize..., ksize...))
    if T<:Float64
        cublasDgemm_v2(CUBLAS.handle(), t1, t2, m, n, k, 1, x1, m, x2, n, 0, y, m)
    elseif T<:Float32
        cublasSgemm_v2(CUBLAS.handle(), t1, t2, m, n, k, 1, x1, m, x2, n, 0, y, m)
    else
        error("CUBLAS does not support $T")
    end
    return y
end

function _mmul3(x1::R, x2::R; dims=1) where {T,R<:DevArray{T}} # x1,x2 transposed
    @assert ndims(x2) > dims
    @assert ntuple(i->size(x1,i),dims) === ntuple(i->size(x2,ndims(x2)-dims+i),dims)
    t1,t2 = cublasOperation_t.((1,1))
    msize = (size(x1,i) for i in dims+1:ndims(x1))
    nsize = (size(x1,i) for i in 1:dims)
    ksize = (size(x2,i) for i in 1:ndims(x2)-dims)
    m,n,k = prod.((msize, nsize, ksize))
    y = similar(x1, (msize..., ksize...))
    if T<:Float64
        cublasDgemm_v2(CUBLAS.handle(), t1, t2, m, n, k, 1, x1, m, x2, n, 0, y, m)
    elseif T<:Float32
        cublasSgemm_v2(CUBLAS.handle(), t1, t2, m, n, k, 1, x1, m, x2, n, 0, y, m)
    else
        error("CUBLAS does not support $T")
    end
    return y
end

@primitive1   mmul(x1,x2;dims=1),dy  _mmul2(dy,x2; dims=ndims(x2)-dims)  _mmul1(x1,dy; dims=ndims(x1)-dims)
@primitive1 _mmul1(x1,x2;dims=1),dy  _mmul2(x2,dy; dims=ndims(x2)-dims)    mmul(x1,dy; dims=ndims(x1)-dims)
@primitive1 _mmul2(x1,x2;dims=1),dy    mmul(dy,x2; dims=ndims(x2)-dims)  _mmul1(dy,x1; dims=ndims(x1)-dims)
@primitive1 _mmul3(x1,x2;dims=1),dy  _mmul3(x2,dy; dims=ndims(x2)-dims)  _mmul3(dy,x1; dims=ndims(x1)-dims)


text:
## Classical derivative-based, iterative, root-finding algorithms: Newton and Halley

## Newton
function newton_itr(f, fp, x0;
                    xtol=4*eps(), xtolrel=4*eps(), ftol=4*eps(),
                    maxsteps=100, maxfnevals=100)
    update =  (o) -> begin
        xn = o.xn[end]
        fxn = o.fxn[end]
        fpxn = o.fp(xn)
        incfn(o)

        xn1 = xn - fxn / fpxn
        fxn1 = o.f(xn1)
        incfn(o)

        push!(o.xn, xn1)
        push!(o.fxn, fxn1)

    end

    x, fx = promote(float(x0), f(float(x0)))
    out = ZeroType(f, fp, nothing, update, [x], [fx],
                   xtol, xtolrel, ftol,:not_converged,
                   0, maxsteps, 1, maxfnevals)

    out
end

function newton_method{T}(f, fp, x0::T;
                       xtol=4*eps(), xtolrel=4*eps(), ftol=4eps(),
                       maxsteps::Int=100, maxfnevals=100,
                       verbose::Bool=false)

    o = newton_itr(f, fp, x0;
                  xtol=xtol, xtolrel=xtolrel, ftol=ftol,
                  maxsteps=maxsteps, maxfnevals=maxfnevals)

    for x in o
        nothing
    end
    verbose && verbose_output(o)
    o.xn[end]
end
newton_method{T<:AbstractFloat}(f, x::T; kwargs...) = newton_method(f, D(f), x; kwargs...)

"""

Implementation of Newton's method: `x_n1 = x_n - f(x_n)/ f'(x_n)`

Arguments:

* `f::Function` -- function to find zero of

* `fp::Function=D(f)` -- derivative of `f`. Defaults to automatic derivative

* `x0::Real` -- initial guess

Keyword arguments:

* `ftol`. Stop iterating when |f(xn)| <= max(1, |xn|) * ftol.

* `xtol`. Stop iterating when |xn+1 - xn| <= xtol + max(1, |xn|) * xtolrel

* `xtolrel`. Stop iterating when |xn+1 - xn| <= xtol + max(1, |xn|) * xtolrel

* `maxeval`. Stop iterating if more than this many steps, throw error.

* `maxfneval`. Stop iterating if more than this many function calls, throw error.

* `verbose::Bool=false` Set to `true` to see trace.

"""
newton{T<:Number}(f, fp, x0::T; kwargs...) = newton_method(f, fp, float(x0); kwargs...)
newton{T<:Real}(f, x0::T; kwargs...) = newton(f, D(f), float(x0); kwargs...)

## Halley's method (cubic convergence)
function halley_itr(f, fp, fpp, x0::Real;
                   xtol=4*eps(), xtolrel=4*eps(), ftol=4*eps(),
                   maxsteps=100, maxfnevals=100)

    update = (o) -> begin
        xn = o.xn[end]
        fxn = o.fxn[end]
        fpxn = o.fp(xn)
        fppxn = o.fpp(xn)

        xn1 = xn - 2fxn*fpxn / (2*fpxn*fpxn - fxn * fppxn)
        fxn1 = o.f(xn1)

        incfn(o, 3)
        push!(o.xn, xn1)
        push!(o.fxn, fxn1)

    end

    x, fx = promote(float(x0), f(float(x0)))
    out = ZeroType(f, fp, fpp, update, [x], [fx],
                   xtol, xtolrel, ftol,:not_converged,
                   0, maxsteps, 1, maxfnevals)

    out
end

function halley_method{T<:AbstractFloat}(f, fp, fpp, x0::T;
                       xtol=4*eps(), xtolrel=4*eps(), ftol=4eps(),
                       maxsteps::Int=100, maxfnevals=100,
                       verbose::Bool=false)

    o = halley_itr(f, fp, fpp, float(x0);
                  xtol=xtol, xtolrel=xtolrel, ftol=ftol,
                  maxsteps=maxsteps, maxfnevals=maxfnevals)

    for x in o
        nothing
    end
    verbose && verbose_output(o)
    o.xn[end]
end

"""

Implementation of Halley's method. `xn1 = xn - 2f(xn)*f'(xn) / (2*f'(xn)^2 - f(xn) * f''(xn))`

Arguments:

* `f::Function` -- function to find zero of

* `fp::Function=D(f)` -- derivative of `f`. Defaults to automatic derivative

* `fpp:Function=D(f,2)` -- second derivative of `f`.

* `x0::Real` -- initial guess

Keyword arguments:

* `ftol`. Stop iterating when |f(xn)| <= max(1, |xn|) * ftol.

* `xtol`. Stop iterating when |xn+1 - xn| <= xtol + max(1, |xn|) * xtolrel

* `xtolrel`. Stop iterating when |xn+1 - xn| <= xtol + max(1, |xn|) * xtolrel

* `maxeval`. Stop iterating if more than this many steps, throw error.

* `verbose::Bool=false` Set to `true` to see trace.

"""
halley(f,fp, fpp, x::Number; kwargs...) = halley_method(f, fp, fpp, float(x); kwargs...)
halley(f, fp, x::Real; kwargs...) = halley(f, fp, D(fp), x; kwargs...)
halley(f, x::Real; kwargs...) = halley(f, D(f), D(f,2), x; kwargs...)


text:
# This file is a part of UpROOT.jl, licensed under the MIT License (MIT).

function pyjaggedarray2jl(x::PyObject)
    from = convert(Vector{Int}, x.starts)::Vector{Int}
    until = convert(Vector{Int}, x.stops)::Vector{Int}

    data = x.content
    T = eltype(data)

    @assert from[2:end] == until[1:end-1]
    n = size(from, 1)

    elem_ptr = Vector{Int}(undef, n + 1)
    elem_ptr[1:end-1] = from
    elem_ptr[end] = until[end]
    elem_ptr .+= 1

    kernel_size = Vector{Tuple{}}(undef, n)

    VectorOfArrays(data, elem_ptr, kernel_size, ArraysOfArrays.no_consistency_checks)
end

py2jl(x::Any) = x

function py2jl(x::PyObject)
    if pybuiltin(:isinstance)(x, awkward.JaggedArray)
        pyjaggedarray2jl(x)
    elseif pybuiltin(:isinstance)(x, awkward.array.base.AwkwardArray)
        Table(_dict2nt(x._contents))
    elseif pybuiltin(:isinstance)(x, uproot.rootio.ROOTDirectory)
        TDirectory(x)
    elseif pybuiltin(:isinstance)(x, uproot.tree.TTreeMethods)
        TTree(x)
    elseif pybuiltin(:isinstance)(x, uproot.tree.TBranchMethods)
        TBranch(x)
    else
        pytypename = pybuiltin(:type)(x).__name__
        @warn "Conversion of python type $pytypename to a Julia type not supported"
        x
    end
end

_dict2nt(d::Dict) = NamedTuple{Tuple(Symbol.(keys(d)))}(py2jl.(values(d)))


text:
# Functions for plotting
#
# improting Gadfly is terribly slow, so avoid this whenever it is
# possible; here do it in another file

using Gadfly
function plot_sing_vals(svs::Array{Array{Float64,1},1})
    num = length(svs)
    plts = Array{Any, 1}(num)
    for i in 1:length(svs)
        plts[i]=plot(y=svs[i], x=collect(1:length(svs[i])),
                     Geom.point, Scale.x_discrete, #Scale.y_continuous,
                     Scale.y_log10,
                     # Coord.cartesian(fixed=true),
                     # Geom.subplot_grid(free_y_axis=true),
                     Guide.title(string("Mode ", i)));
    end
    # set_default_plot_size(30cm, 6cm)
    hstack(plts...)
end


text:
###############################################################################
#
#  Direct product decomposition
#
###############################################################################

function direct_product_decomposition(G::GAP.GapObj, ab::Tuple{Int, Int})

  decompositions = Tuple{GAP.GapObj, GAP.GapObj}[]
  if GAP.Globals.IsAbelian(G)
    return ab, (1, 1), 1, 1
  end
  n = ab[1]
  subs = GAP.gap_to_julia(Vector{GAP.GapObj}, GAP.Globals.NormalSubgroups(G))
  sort!(subs, by = x -> GAP.Globals.Size(x))
  #First, I collect all the possible decompositions
  decompositions = Tuple{GAP.GapObj, GAP.GapObj}[]
  for i = 1:length(subs)
    g1 = subs[i]
    if isone(GAP.Globals.Size(g1))
      continue
    end
    o1 = GAP.Globals.Size(g1)
    for j = 1:length(subs)
      g2 = subs[j]
      o2 = GAP.Globals.Size(g2)
      if o2 < o1
        continue
      end
      if o1*o2 != n
        continue
      end
      if isone(GAP.Globals.Size(GAP.Globals.Intersection(g1, g2)))
        push!(decompositions, (g1, g2))
      end
    end
  end
  if isempty(decompositions)
    return ab, (1, 1) , 1, 1
  end

  #We pass to the list with the group ID
  grp_id_list = Vector{Tuple{Tuple{Int, Int}, Tuple{Int, Int}}}(undef, length(decompositions))
  for i = 1:length(grp_id_list)
    grp_id_list[i] = (GAP.gap_to_julia(Tuple{Int, Int}, GAP.Globals.IdGroup(decompositions[i][1])), GAP.gap_to_julia(Tuple{Int, Int}, GAP.Globals.IdGroup(decompositions[i][2])))
  end

  possible_decompositions = Set(grp_id_list)
  #First, I want to choose which decomposition to consider.
  #In general, a balanced one is the best, so that pair maximising the minimum of the two orders.
  #However, we also want to maximise the order of the abelian factor, if it exists.
  res = first(possible_decompositions)
  has_abelian_decomposition = 0
  if GAP.Globals.IsAbelian(GAP.Globals.SmallGroup(res[1]...))
    has_abelian_decomposition = 1
  elseif GAP.Globals.IsAbelian(GAP.Globals.SmallGroup(res[2]...))
    has_abelian_decomposition = 2
  end
  for x in possible_decompositions
    if x == res
      continue
    end
    if iszero(has_abelian_decomposition)
      if min(x[1][1], x[2][1]) > min(res[1][1], res[2][1])
        res = x
      end
    else
      has_abelian_decx = 0
      if GAP.Globals.IsAbelian(GAP.Globals.SmallGroup(x[1]...))
        has_abelian_decx = 1
      elseif GAP.Globals.IsAbelian(GAP.Globals.SmallGroup(x[2]...))
        has_abelian_decx = 2
      end
      if !iszero(has_abelian_decx)
        if x[has_abelian_decx][1] > res[has_abelian_decomposition][1]
          res = x
        end
      elseif min(x[1][1], x[2][1]) > min(res[1][1], res[2][1])
        res = x
      end
    end
  end
  res1 = res[1]
  res2 = res[2]
  #I count the redundancy, i.e. the number of possible decompositions of the same type.
  red = 0
  for i = 1:length(grp_id_list)
    l1 = grp_id_list[i][1]
    l2 = grp_id_list[i][2]
    if l1 == res1 && l2 == res2
      red += 1
    end
  end
  #Finally, I count the numer of times a single subgroup appears in the lists.
  redfirst = 1
  ind = 1
  for i = 1:length(decompositions)
    if grp_id_list[i] == (res1, res2)
      ind = i
      break
    end
  end
  for i = ind+1:length(decompositions)
    if decompositions[i][1] == decompositions[ind][1]
      redfirst += 1
    end
  end
  return res1, res2, red, redfirst
end

###############################################################################
#
#  Merging fields
#
###############################################################################

function _to_composite(x::FieldsTower, y::FieldsTower, abs_disc::fmpz)
  Kns, mx, my = NumberField(x.field, y.field, cached = false, check = false)
  OKns = maximal_order(Kns)
  if abs(discriminant(OKns)) > abs_disc
    return false, x
  end
  K, mK = simplified_simple_extension(Kns, cached = false)
  OK = maximal_order(K)

  Hecke._assure_has_inverse_data(mK)

  # Now, I have to translate the automorphisms.
  # Easy thing: first, I write the automorphisms of the non simple extension
  # Then translating them is straightforward.
  autK = Vector{NfToNfMor}(undef, length(x.generators_of_automorphisms)+ length(y.generators_of_automorphisms))
  el = image_primitive_element(mK)
  for i = 1:length(x.generators_of_automorphisms)
    ima = mx(image_primitive_element(x.generators_of_automorphisms[i]))
    autns = hom(Kns, Kns, NfAbsNSElem[ima, gens(Kns)[2]], check = false)
    ima = mK\(autns(el))
    autK[i] = hom(K, K, ima, check = false)
  end
  for j = 1:length(y.generators_of_automorphisms)
    ima = my(image_primitive_element(y.generators_of_automorphisms[j]))
    autns = hom(Kns, Kns, NfAbsNSElem[gens(Kns)[1], ima], check = false)
    ima = mK\(autns(el))
    autK[j+length(x.generators_of_automorphisms)] = hom(K, K, ima, check = false)
  end

  #Last thing: I have to add the maps of the subfields!
  #I want to merge the information for the last embedding.
  i = 1
  while codomain(x.subfields[i]) != x.field
    i += 1
  end
  emb_subx = x.subfields[i]
  i = 1
  while codomain(y.subfields[i]) != y.field
    i += 1
  end
  emb_suby = y.subfields[i]
  lsub, m1, m2 = number_field(domain(emb_subx), domain(emb_suby), cached = false, check = false)
  Seemb, mSeemb = simple_extension(lsub, check = false)
  ev = nf_elem[mK\(mx(image_primitive_element(emb_subx))), mK\(my(image_primitive_element(emb_suby)))]
  embs = NfToNfMor[hom(Seemb, K, evaluate(mSeemb(gen(Seemb)).data, ev))]
  for j = 1:length(x.subfields)
    if codomain(x.subfields[j]) != domain(emb_subx)
      push!(embs, x.subfields[j])
    else
      push!(embs, hom(domain(x.subfields[j]), Seemb, mSeemb\(m1(image_primitive_element(x.subfields[j])))))
    end
  end
  for j = 1:length(y.subfields)
    if codomain(y.subfields[j]) != domain(emb_suby)
      push!(embs, y.subfields[j])
    else
      push!(embs, hom(domain(y.subfields[j]), Seemb, mSeemb\(m2(image_primitive_element(y.subfields[j])))))
    end
  end
  return true, FieldsTower(K, autK, embs)
end

#merge function when all the fields are automatically linearly disjoint
function _easy_merge(list1, list2, absolute_bound::fmpz)

  res = FieldsTower[]
  @vprint :Fields 1 "Number of candidates = $(length(list1)*length(list2)) \n"
  ind = 0
  for x in list1
    ind += 1
    @vprint :Fields 1 "Doing field $(ind)/$(length(list1))"
    for y in list2
      check_bound_disc(x, y, absolute_bound) || continue
      fl, composite = _to_composite(x, y, absolute_bound)
      if fl
        push!(res, composite)
      end
    end
    @vprint :Fields 1 "$(Hecke.set_cursor_col())$(Hecke.clear_to_eol())"
  end
  return res

end

function _disjoint_ab_subs(list1::Vector{FieldsTower}, list2::Vector{FieldsTower})

  deg_ab_sub1 = 1
  deg_ab_sub2 = 1
  for i = 1:length(list1[1].subfields)
    if degree(domain(list1[1].subfields[i])) == 1
      deg_ab_sub1 *= degree(codomain(list1[1].subfields[i]))
    end
  end
  for i = 1:length(list2[1].subfields)
    if degree(domain(list2[1].subfields[i])) == 1
      deg_ab_sub2 *= degree(codomain(list2[1].subfields[i]))
    end
  end
  return deg_ab_sub1, deg_ab_sub2, gcd(deg_ab_sub1, deg_ab_sub2) == 1

end

function check_bound_disc(K::FieldsTower, L::FieldsTower, bound::fmpz)
  #First, I check the bound for the fields
  if !new_check_disc(K, L, bound)
    @hassert :Fields 1 !_to_composite(K, L, bound)[1]
    return false
  end
  #In the other cases, before returning true, I try to check if the
  #maximal abelian subextension is admissible
  Kab = maximal_abelian_subextension(K)
  Lab = maximal_abelian_subextension(L)
  disc_comp = abs(discriminant(maximal_order(Kab[1])))
  deg = degree(Kab[1])
  for i = 2:length(Kab)
    disc_new = abs(discriminant(maximal_order(Kab[i])))
    d = gcd(disc_new, disc_comp)
    if issquare(d)
      d = root(d, 2)
    end
    n = ppio(disc_comp, d)[2]^(degree(Kab[i]))*ppio(disc_new, d)[2]^deg
    ld = factor(d)
    for (q, v) in ld
      vnum = degree(Kab[i])*valuation(disc_comp, q) + deg*valuation(disc_new, q)
      vden = 2*(deg*degree(Kab[i])-1)*valuation(d, q)
      vfirst = vnum - vden
      v1 = degree(Kab[i])*valuation(disc_comp, q)
      v2 = deg*valuation(disc_new, q)
      vsecond = max(v1, v2)
      vd = max(vfirst, vsecond)
      n *= q^vd
    end
    disc_comp = n
    deg *= degree(Kab[i])
  end
  for i = 1:length(Lab)
    disc_new = abs(discriminant(maximal_order(Lab[i])))
    d = gcd(disc_new, disc_comp)
    if issquare(d)
      d = root(d, 2)
    end
    n = ppio(disc_comp, d)[2]^(degree(Lab[i]))*ppio(disc_new, d)[2]^deg
    ld = factor(d)
    for (q, v) in ld
      vnum = degree(Lab[i])*valuation(disc_comp, q) + deg*valuation(disc_new, q)
      vden = 2*(deg*degree(Lab[i])-1)*valuation(d, q)
      vfirst = vnum - vden
      v1 = degree(Lab[i])*valuation(disc_comp, q)
      v2 = deg*valuation(disc_new, q)
      vv = max(v1, v2)
      vd = max(vfirst, vv)
      n *= q^vd
    end
    disc_comp = n
    deg *= degree(Lab[i])
  end
  bound_max_ab_subext = root(bound, divexact(degree(K.field)*degree(L.field), deg))
  if disc_comp <= bound_max_ab_subext
    return true
  else
    @hassert :Fields 1 !_to_composite(K, L, bound)[1]
    return false
  end
end

function new_check_disc(K::FieldsTower, L::FieldsTower, absolute_bound::fmpz)
  Kf = K.field
  OKf = maximal_order(Kf)
  Lf = L.field
  OLf = maximal_order(Lf)
  d1 = discriminant(maximal_order(Kf))
  d2 = discriminant(maximal_order(Lf))
  g1 = gcd(d1, d2)
  wild, g = ppio(g1, fmpz(degree(Kf)*degree(Lf)))
  disc1 = abs(ppio(d1, g1)[2])^degree(Lf)
  disc2 = abs(ppio(d2, g1)[2])^degree(Kf)
  disc = disc1*disc2
  if disc > absolute_bound
    return false
  end
  lf = factor(g)
  ramification_indices = Vector{Tuple{fmpz, Int}}(undef, length(lf))
  ind = 1
  for (p, v) in lf
    pd1 = prime_decomposition_type(OKf, Int(p))
    pd2 = prime_decomposition_type(OLf, Int(p))
    ramification_indices[ind] = (p, lcm(pd1[1][2], pd2[1][2]))
    ind += 1
  end

  for i = 1:length(ramification_indices)
    expo = (ramification_indices[i][2] - 1)*divexact(degree(Kf)*degree(Lf), ramification_indices[i][2])
    disc *= ramification_indices[i][1]^expo
  end
  if disc > absolute_bound
    return false
  end
  lwild = factor(wild)
  for (q, v) in lwild
    v1 = valuation(d1, q)
    v2 = valuation(d2, q)
    valnum = degree(Lf)*v1 + degree(Kf)*v2
    valden = valuation(g1, q)
    if iseven(valden)
      valden = divexact(valden, 2)
    end
    valden *= (2*degree(Lf)*degree(Kf) -2)
    if valnum - valden > max(degree(Lf)*v1, degree(Kf)*v2)
      disc *= q^(valnum-valden)
    else
      disc *= q^max(degree(Lf)*v1, degree(Kf)*v2)
    end
  end
  return disc <= absolute_bound
end

function maximal_abelian_subextension(F::FieldsTower)
  fields = AnticNumberField[]
  for x in F.subfields
    if degree(domain(x)) == 1
      push!(fields, codomain(x))
    end
  end
  return fields
end

function check_norm_group_and_disc(lfieldsK::Vector{AnticNumberField}, lfieldsL::Vector{AnticNumberField}, bound::fmpz)

  target_deg = prod(degree(x) for x in lfieldsK) * prod(degree(x) for x in lfieldsL)
  discK = lcm([discriminant(maximal_order(x)) for x in lfieldsK])
  discL = lcm([discriminant(maximal_order(x)) for x in lfieldsL])
  n_quo1 = lcm(Int[degree(x) for x in lfieldsK])
  n_quo2 = lcm(Int[degree(x) for x in lfieldsL])
  exp_rcf = lcm(n_quo1, n_quo2)
  modulo = lcm(discK, discL)
  lf = factor(modulo)
  modulo_int = 1
  for (p, v) in lf
    if iscoprime(p, exp_rcf)
      modulo_int *= Int(p)
    else
      modulo_int *= Int(p)^v
    end
  end
  y = PolynomialRing(QQ, "y", cached = false)[2]
  K = NumberField(y-1, cached = false)[1]
  O = maximal_order(K)
  r, mr = Hecke.ray_class_groupQQ(O, modulo_int, true, exp_rcf)
  Kt = PolynomialRing(K, "t", cached = false)[1]
  h = change_base_ring(K, lfieldsK[1].pol, parent = Kt)
  S, mS = norm_group(h, mr, cached = false)
  for i = 2:length(lfieldsK)
    h = map_coefficients(K, lfieldsK[i].pol, parent = Kt)
    s, ms = norm_group(h, mr, cached = false)
    S, mS = intersect(ms, mS)
  end
  for i = 1:length(lfieldsL)
    h = map_coefficients(K, lfieldsL[i].pol, parent = Kt)
    s, ms = norm_group(h, mr, cached = false)
    S, mS = intersect(ms, mS)
  end
  Q, mQ = cokernel(mS, false)
  if order(Q) != target_deg
    return false
  else
    C = ray_class_field(mr, mQ)
    return Hecke.discriminant_conductorQQ(O, C, modulo_int, bound)
  end

end

function _first_sieve(list1::Vector{FieldsTower}, list2::Vector{FieldsTower}, absolute_bound::fmpz, redfirst::Int)
  ab1, ab2, fl = _disjoint_ab_subs(list1, list2)
  bound_max_ab_sub = root(absolute_bound, divexact(degree(list1[1])*degree(list2[1]), ab1*ab2))
  D = Dict{Tuple{Set{fmpz}, Set{fmpz}, Bool}, Vector{Tuple{Int, Int}}}() #The boolean true means real
  for i1 = 1:length(list1)
    @vprint :Fields 1 "$(Hecke.set_cursor_col())$(Hecke.clear_to_eol())Combinations with field $(i1)/$(length(list1)) of the first list"
    @vprint :FieldsNonFancy 1 "Fields $(i1)/$(length(list1))\n"
    K = list1[i1]
    DK = Dict{Tuple{Set{fmpz}, Set{fmpz}}, Vector{Int}}()
    rK = ramified_primes(K)
    lfieldsK = maximal_abelian_subextension(K)
    rK1 = Set(fmpz[])
    for x in lfieldsK
      rK1 = union(rK1, Hecke.ramified_primes(maximal_order(x)))
    end
    #First, I only check for disjointness
    for i2 = 1:length(list2)
      L = list2[i2]
      check_bound_disc(K, L, absolute_bound) || continue
      #First, I check if the discriminants are coprime
      rL = Hecke.ramified_primes(L)
      lfieldsL = maximal_abelian_subextension(L)
      rL1 = Set(fmpz[])
      for x in lfieldsL
        rL1 = union(rL1, Hecke.ramified_primes(maximal_order(x)))
      end
      fl = false
      if isempty(intersect(rL1, rK1))
        fl = true
      elseif length(lfieldsL) == 1 && length(lfieldsK) == 1
        if degree(lfieldsL[1]) == 2 && degree(lfieldsK[1]) == 2 && discriminant(maximal_order(lfieldsK[1])) != discriminant(maximal_order(lfieldsL[1]))
          fl = true
        elseif Hecke.islinearly_disjoint(lfieldsK[1], lfieldsL[1])
          fl = true
        end
      elseif check_norm_group_and_disc(lfieldsK, lfieldsL, bound_max_ab_sub)
        fl = true
      end
      if fl
        k =  Set(vcat(rK, rL))
        k1 = union(rK1, rL1)
        if haskey(DK, (k, k1))
          push!(DK[(k, k1)],  i2)
        else
          DK[(k, k1)] = Int[i2]
        end
      end
    end
    #Now, I sieve also by infinite places.
    for (k, v) in DK
      if length(v) < redfirst
        continue
      end
      if !istotally_real(K.field)
        ar = Vector{Tuple{Int, Int}}(undef, length(v))
        for j = 1:length(v)
          ar[j] = (i1, v[j])
        end
        if haskey(D, (k[1], k[2], false))
          append!(ar, D[(k[1], k[2], false)])
        end
        D[(k[1], k[2], false)] = ar
      else
        ar_real = Vector{Tuple{Int, Int}}()
        ar_not_real = Vector{Tuple{Int, Int}}()
        for j = 1:length(v)
          if istotally_real(list2[v[j]].field)
            push!(ar_real, (i1, v[j]))
          else
            push!(ar_not_real, (i1, v[j]))
          end
        end
        if !isempty(ar_real)
          if haskey(D, (k[1], k[2], true))
            append!(ar_real, D[(k[1], k[2], true)])
          end
          D[(k[1], k[2], true)] = ar_real
        end
        if !isempty(ar_not_real)
          if haskey(D, (k[1], k[2], false))
            append!(ar_not_real, D[(k[1], k[2], false)])
          end
          D[(k[1], k[2], false)] = ar_not_real
        end
      end
    end
  end
  return D
end

function _some_combinatorics(l::Vector{Tuple{Int, Int}}, red1, red2)
  #every entry must appear multiples of red1 or red2 times.
  while true
    entries1 = Dict{Int, Int}()
    entries2 = Dict{Int, Int}()
    for i = 1:length(l)
      if !haskey(entries1, l[i][1])
        entries1[l[i][1]] = 1
      else
        entries1[l[i][1]] += 1
      end
      if !haskey(entries2, l[i][2])
        entries2[l[i][2]] = 1
      else
        entries2[l[i][2]] += 1
      end
    end
    done = false
    for (k, v) in entries1
      if red1 <= v
        continue
      end
      done = true
      #remove all the tuples
      filter!(x -> x[1] != k, l)
    end
    for (k, v) in entries2
      if red1 > v
        done = true
        filter!(x -> x[2] != k, l)
      end
    end
    if !done
      break
    end
  end
  return l
end

function sieve_by_discriminant(list1, list2, v)

  d1 = degree(list1[1].field)
  d2 = degree(list2[1].field)
  D = Dict{fmpz, Vector{Int}}()

  for i = 1:length(v)
    candidate = abs(discriminant(maximal_order(list1[v[i][1]].field))^d2)*abs(discriminant(maximal_order(list2[v[i][2]].field))^d1)
    found = false
    for (d, l) in D
      if _differ_by_square(d, candidate)
        push!(l, i)
        found = true
        break
      end
    end
    if !found
      D[candidate] = Int[i]
    end
  end
  res = Vector{Vector{Tuple{Int, Int}}}()
  for val in values(D)
    to_push = Vector{Tuple{Int, Int}}(undef, length(val))
    for i = 1:length(val)
      to_push[i] = v[val[i]]
    end
    push!(res, to_push)
  end
  return res
end

function _differ_by_square(n::fmpz, m::fmpz)
  return issquare(divexact(fmpq(n), fmpq(m)))
end

function sieve_by_norm_group(list1::Vector{FieldsTower}, list2::Vector{FieldsTower}, v::Vector{Tuple{Int, Int}}, ramified_primes::Vector{Int})

  target_deg = prod(degree(x) for x in maximal_abelian_subextension(list1[1])) * prod(degree(x) for x in maximal_abelian_subextension(list2[1]))
  expo = lcm(vcat([degree(x) for x in maximal_abelian_subextension(list1[1])], [degree(x) for x in maximal_abelian_subextension(list2[1])]))
  modulo = 1
  for p in ramified_primes
    if !divisible(expo, p)
      modulo *= p
    else
      bound_disc = valuation_bound_discriminant(ppio(expo, p)[1], p)
      bound_exp = div(bound_disc, (p-1)*p^(valuation(expo, p)-1))
      modulo *= p^bound_exp
    end
  end
  K = rationals_as_number_field()[1]
  O = maximal_order(K)
  r, mr = Hecke.ray_class_groupQQ(O, modulo, true, expo)
  Kt = PolynomialRing(K, "t", cached = false)[1]
  norm_groups = Vector{GrpAbFinGenMap}(undef, length(v))
  for i = 1:length(v)
    lfieldsK = maximal_abelian_subextension(list1[v[i][1]])
    lfieldsL = maximal_abelian_subextension(list2[v[i][2]])
    h = change_base_ring(K, lfieldsK[1].pol, parent = Kt)
    S, mS = norm_group(h, mr, cached = false)
    for i = 2:length(lfieldsK)
      h = change_base_ring(K, lfieldsK[i].pol, parent = Kt)
      s, ms = norm_group(h, mr, cached = false)
      S, mS = intersect(ms, mS, false)
    end
    for i = 1:length(lfieldsL)
      h = change_base_ring(K, lfieldsL[i].pol, parent = Kt)
      s, ms = norm_group(h, mr, cached = false)
      S, mS = intersect(ms, mS, false)
    end
    norm_groups[i] = mS
  end
  #Now that I have the norm groups, I need to compare them.
  done = falses(length(v))
  res = Vector{Vector{Tuple{Int, Int}}}()
  for i = 1:length(v)
    if done[i]
      continue
    end
    done[i] = true
    mS = norm_groups[i]
    new_v = Vector{Tuple{Int, Int}}()
    push!(new_v, v[i])
    for j = i+1:length(v)
      if done[j]
        continue
      end
      if order(intersect(mS, norm_groups[j], false)[1]) == order(domain(mS))
        done[j] = true
        push!(new_v, v[j])
      end
    end
    push!(res, new_v)
  end
  return res
end

function refine_clusters(list1, list2, clusters, red, redfirst, redsecond)
  new_clusters = Vector{Vector{Tuple{Int, Int}}}()
  for i = 1:length(clusters)
    v1 = _some_combinatorics(clusters[i], redfirst, redsecond)
    if length(v1) < red
      continue
    end
    push!(new_clusters, v1)
  end
  return new_clusters
end

function _merge(list1::Vector{FieldsTower}, list2::Vector{FieldsTower}, absolute_bound::fmpz, red::Int, redsecond::Int, g1::Tuple{Int, Int}, g2::Tuple{Int, Int})

  G1 = GAP.Globals.SmallGroup(g1[1], g1[2])
  G2 = GAP.Globals.SmallGroup(g2[1], g2[2])
  mas1 = GAP.gap_to_julia(Vector{Int}, GAP.Globals.AbelianInvariants(G1))
  mas2 = GAP.gap_to_julia(Vector{Int}, GAP.Globals.AbelianInvariants(G2))
  if gcd(prod(mas1), prod(mas2)) == 1
    #All the fields are automatically linearly disjoint
    @vprint :Fields 1 "All the fields are linearly disjoint, easy case \n"
    @vprint :FieldsNonFancy 1 "All the fields are linearly disjoint, easy case \n"
    return _easy_merge(list1, list2, absolute_bound)
  end

  redfirst = divexact(red, redsecond)
  #Bad case: a mess.
  #I check that the maximal abelian subextensions are linearly disjoint.
  #Working with polynomials may be more expensive.
  #Since I am working with abelian extensions, I can compute the norm groups and see there if the norm groups are correct.
  @vprint :Fields 1 "Merging fields\n"
  @vprint :Fields 1 "Number of fields: $(length(list1)*length(list2))\n"
  @vprint :FieldsNonFancy 1 "Merging fields\n"
  @vprint :FieldsNonFancy 1 "Number of fields: $(length(list1)*length(list2))\n"
  @vprint :Fields 1 "Redundancy: $(red), $(redfirst), $(redsecond)\n"
  @vprint :FieldsNonFancy 1 "Redundancy: $(red), $(redfirst), $(div(red, redfirst))\n"

  res = FieldsTower[]
  D1 = _first_sieve(list1, list2, absolute_bound, redfirst)
  if isempty(D1)
    return res
  end
  clusters = Vector{Tuple{Int, Int}}[x for x in values(D1)]
  @vprint :Fields 1 "$(Hecke.set_cursor_col())$(Hecke.clear_to_eol())Candidates after first sieve: $(sum(length(x) for x in clusters))\n"

  @vprint :Fields 1 "Sieving by discriminant\n"
  #Now, I sieve by discriminant
  clusters1 = Vector{Vector{Tuple{Int, Int}}}()
  for v in clusters
    vnew =  sieve_by_discriminant(list1, list2, v)
    for x in vnew
      if !isempty(x)
        push!(clusters1, x)
      end
    end
  end
  @vprint :Fields 1 "Candidates: $(sum(length(x) for x in clusters1))\n"
  @vprint :Fields 1 "Sieving by prime_splitting\n"
  fields_to_be_computed = _sieve_by_prime_splitting(list1, list2, clusters1, red, redfirst, redsecond)

  @vprint :Fields 1 "Computing maximal order of $(length(fields_to_be_computed)) fields\n"
  for i = 1:length(fields_to_be_computed)
    @vprint :Fields 1 "Doing $(i) / $(length(fields_to_be_computed))"
    pair = fields_to_be_computed[i]
    fl, candidate = _to_composite(list1[pair[1]], list2[pair[2]], absolute_bound)
    if fl
      push!(res, candidate)
    end
    @vprint :Fields 1 "$(Hecke.set_cursor_col())$(Hecke.clear_to_eol())"
  end
  @vprint :Fields 1 "$(Hecke.set_cursor_col())$(Hecke.clear_to_eol())Fields found: $(length(res))\n"
  @vprint :FieldsNonFancy 1 "Fields found: $(length(res))\n"
  return res
end

function _sieve_by_prime_splitting(list1, list2, clusters, red, redfirst, redsecond)
  fields_to_be_computed = Vector{Tuple{Int, Int}}()
  d1 = degree(list1[1].field)
  d2 = degree(list2[1].field)
  total = length(clusters)
  @vprint :Fields 1 "\n Number of clusters: $(total)\n"
  nclu = 0
  for v in clusters
    nclu += 1
    #=
    if length(v) == red
      push!(fields_to_be_computed, v[1])
      continue
    end
    =#
    if length(v) < red
      continue
    end
    @vprint :Fields 1 "$(Hecke.set_cursor_col())$(Hecke.clear_to_eol()) Doing cluster $(nclu)/$(total) of length $(length(v))"
    p = next_prime(1000)
    iso_classes = Vector{Int}[Int[i for i = 1:length(v)]]
    cnt = 0
    while true
      splitting_types = Dict{Tuple{Int, Int}, Vector{Int}}()
      for i = 1:length(v)
        pd1 = prime_decomposition_type(maximal_order(list1[v[i][1]].field), p)
        if pd1[1][2] != 1
          break
        end
        f1 = pd1[1][1]
        pd2 = prime_decomposition_type(maximal_order(list2[v[i][2]].field), p)
        if pd2[1][2] != 1
          break
        end
        f2 = pd2[1][1]
        f = lcm(f1, f2)
        r = divexact(d1*d2, f)
        if haskey(splitting_types, (r, f))
          push!(splitting_types[(r, f)], i)
        else
          splitting_types[(r, f)] = Int[i]
        end
      end
      cnt += 1
      if all(x -> length(x) <= red, values(splitting_types)) && cnt >= 10
        #I intersect the data with the one that I already have
        for vals in values(splitting_types)
          to_be = intersect_infos(vals, iso_classes)
          for j in to_be
            if length(j) == red
              push!(fields_to_be_computed, v[j[1]])
            end
          end
        end
        break
      else
        #I save the information provided by the splitting
        new_iso_classes = Vector{Int}[]
        for vals in values(splitting_types)
          append!(new_iso_classes, intersect_infos(vals, iso_classes))
        end
        iso_classes = new_iso_classes
        if all(x -> length(x) <= red, iso_classes) && cnt >= 10
          for j in iso_classes
            if length(j) == red
              push!(fields_to_be_computed, v[j[1]])
            end
          end
          break
        else
          p = next_prime(p)
        end
      end
    end
  end
  @vprint :Fields 1 "$(Hecke.set_cursor_col())$(Hecke.clear_to_eol())"
  return fields_to_be_computed
end

function intersect_infos(v::Vector{Int}, iso_classes::Vector{Vector{Int}})
  res = Vector{Vector{Int}}()
  for w in iso_classes
    r = collect(intersect(v, w))
    if !isempty(r)
      push!(res, r)
    end
  end
  return res
end


text:
using Destruct
using Test
using Random
import Random.rand

rand(rng::AbstractRNG, T::Type{String}) = randstring(rng)
types = [Int32, Float64, Complex{Float64}, Bool]

@testset "NTuple" begin
for T=types, N=[1,2,3,4,5]
    sz = fill(5, N)
    a = rand(T, sz...); b = rand(T, sz...); z = collect(zip(a,b))
    @assert size(z) == size(a) # make sure collect preserved array shape
    @test (a,b) == destruct(z)
end end

@testset "Tuple" begin
for T1=types, T2=types, N=[1,2,3]
    sz = fill(3,N)
    a = rand(T1, sz...); b = rand(T2, sz...); z = collect(zip(a,b))
    @assert size(z) == size(a) # make sure collect preserved array shape
    @test (a,b) == destruct(z)
end end

@testset "Tuple, Non square" begin
for T1=types, T2=types, N=[2,3]
    sz = 2:N+1
    a = rand(T1, sz...); b = rand(T2, sz...); z = collect(zip(a,b))
    @assert size(z) == size(a) # make sure collect preserved array shape
    @test (a,b) == destruct(z)
end end


text:
using Test
import LessOLS: Sample, add_intercept, sample_factory#, Process, make_sample

@testset "Sample struct - constructor" begin
    x = [[1,2,3] [4,5,6]]
    y = [1,2,3]
    sam = Sample(x, y)
    @test sam.X == x
    @test sam.Y == y
    @test_throws DimensionMismatch Sample([1,2], [1,2,3])
    @test_throws DimensionMismatch Sample(y, x)
end

@testset "add_intercept() works on array and Sample" begin
    @test add_intercept([1,2,3]) == [[1,1,1] [1,2,3]]
    sam = Sample([1,2], [3,4])
    @test add_intercept(sam).X == [[1,1] [1,2]]
    @test add_intercept(sam).Y == sam.Y
end

@testset "data generating process make_sample()" begin
    fct = sample_factory(x_process = n->collect(1:n),
                         y_process = x->2*x,
                         error_process = x->[0.1 for _ in x]
                         )
    sam = fct(3)
    @test sam.X == [1, 2, 3]
    @test sam.Y == [2.1, 4.1, 6.1]
end


text:
using ConstrainedDynamics

# Parameters
ex = [1.;0.;0.]

h = 1.
r = .05
b1 = Cylinder(r, h, h, color = RGBA(1., 0., 0.))

vert11 = [0.;0.;h / 2]
vert12 = -vert11

# Initial orientation
phi = pi / 4
q1 = Quaternion(RotX(phi))

# Links
N = 6

origin = Origin{Float64}()
links = [Body(b1) for i = 1:N]

# Constraints
cf = 0.2
ineqcs = [InequalityConstraint(Friction(links[i], [0;0;1.0],cf)) for i = 2:N-1]

jointb1 = EqualityConstraint(OriginConnection(origin, links[1]))
constraints = [jointb1;[EqualityConstraint(Revolute(links[i - 1], links[i], vert12, vert11, ex)) for i = 2:N]]

shapes = [b1]

mech = Mechanism(origin, links, constraints, ineqcs;tend = 20.,Δt = 0.01, shapes = shapes)
setPosition!(mech,origin,links[1],p2 = vert11, Δx=[0;0;N*1.], Δq = q1)
previd = links[1].id
for body in Iterators.drop(mech.bodies, 1)
    global previd
    setPosition!(mech, ConstrainedDynamics.getbody(mech, previd), body, p1 = vert12, p2 = vert11)
    previd = body.id
end

simulate!(mech,save = true,debug=false)
visualize!(mech)


text:
using SatelliteToolbox

# If toolbox doesn't exist
# using Pkg
# Pkg.add("SatelliteToolbox")

# Read example TLE data
tle_strs = """
       SCD 1
       1 22490U 93009B   18359.76217587  .00000186  00000-0  84512-6 0  9998
       2 22490  24.9694 116.1709 0043211  90.3968  62.0083 14.44539396366163
       SCD 2
       1 25504U 98060A   18360.24309362  .00000218  00000-0  10518-4 0  9996
       2 25504  24.9967 319.8664 0017431 121.3810   9.7939 14.44057429 65541
       """;

tles = read_tle_from_string(tle_strs)


text:
# using Stars
using Statistics

@testset "st_Ops" begin
    r1 = GeoArray(rand(4, 4, 3))
    r2 = GeoArray(rand(4, 4, 3))

    @test (r1 + r2).A == r1.A + r2.A
    @test (r1 - r2).A == r1.A - r2.A
    @test (r1 * r2).A == r1.A .* r2.A
    @test (r1 / r2).A == r1.A ./ r2.A
end

@testset "st_Ops Statistics" begin
    r = GeoArray(rand(4, 4, 3))

    @test mean(r).A == mean(r.A, dims = 3)
    @test sum(r).A == sum(r.A, dims = 3)
    @test minimum(r).A == minimum(r.A, dims = 3)
    @test maximum(r).A == maximum(r.A, dims = 3)
end


text:
using ArnoldiMethod
using SparseArrays
using Base: OneTo
using LinearAlgebra
using Random
using BenchmarkTools
using Plots
using Arpack

function helloworld(n = 40, nev = 4)
    A = randn(n, n)
    @time S, hist = partialschur(A, nev=nev, which = LR(), tol = 1e-6)
    @time eigs(A, nev=nev, which = :LR, tol = 1e-6)
    @show norm(A * S.Q - S.Q * S.R)
    println(hist)
    θs = eigvals(S.R)
    @time λs = eigvals(A)

    scatter(real(λs), imag(λs), aspect_ratio = :equal)
    scatter!(real(θs), imag(θs), mark = :+)
end


text:
using LinearAlgebra
using StaticArrays
using Plots; pyplot()
using FFTW
using FastClosures
using NPZ
using LaTeXStrings

include("/Users/tfiguei/Work/Magnus/MagnusExpansion.jl")
include("AuxFunctions.jl")
include("HarmonicCorrectionSnap.jl")
include("snap.jl")

σ = pauli_matrices()
χ_GHz = (-8281.0, 0.0, 0.0)#(-8281.0, 48.8, 0.5)
χ = χ_GHz./abs(χ_GHz[1])
N = 10
levels = SA[0, 4]#collect(0:1:9)
α = (pi/2, pi/3)#pi*(rand(10) .- 0.5)
tf = 50.0
Ncorr = 10

system_params = Dict("chi"=>χ, "N"=>N, "levels"=>levels, "alpha"=>α,
					 "tf"=>tf, "Ncorr"=>Ncorr)

UI = U(system_params, σ; ideal_dynamics=false, frame="interaction", matrix=true)
Identity = Matrix{Complex{Float64}}(I, 2*N, 2*N)
println(1 - fidelity(UI, Identity, 2*N))

Uideal = U(system_params, σ; ideal_dynamics=true, matrix=true)
Ur = U(system_params, σ; ideal_dynamics=false, matrix=true)
println(1 - fidelity(Ur, Uideal, 2*N))

# "First" order correction
magnus_order = 1
max_order = 2
coeffs = get_coeffs(magnus_order, max_order, Ncorr, system_params)
println("---")
Ucorr = U(system_params, σ, coeffs, max_order; frame="qubit", matrix=true)
println(1 - fidelity(Ucorr, Uideal, 2*N))
# "Second" order correction
magnus_order = 2
coeffs2 = get_coeffs(coeffs, magnus_order, max_order, Ncorr, system_params)
println("---")
coeffs2 = coeffs + coeffs2
Ucorr = U(system_params, σ, coeffs2, max_order; frame="qubit", matrix=true)
println(1 - fidelity(Ucorr, Uideal, 2*N))

# In the lab frame
println("Results in the lab frame")
Uideal = U(system_params, σ; frame="lab", ideal_dynamics=true, matrix=true)
Unideal = U(system_params, σ; frame="lab", ideal_dynamics=false, matrix=true)
Ucorr = U(system_params, σ, coeffs2, max_order; frame="lab", matrix=true)
println(1 - fidelity(Unideal, Uideal, 2*N))
println(1 - fidelity(Ucorr, Uideal, 2*N))

#FFT results
Npoints = 1000
data = zeros(Npoints, 5)
for i = 1: Npoints
	t = (i-1)*tf/(Npoints-1)
	data[i,1] = t
	data[i, 2: 3] .= pulse_envelope(t, tf, χ, levels, α)
	data[i, 4: 5] .= data[i, 2: 3]
	data[i, 4: 5] .+= correction_pulse_envelope(t, tf, χ, coeffs2, max_order, Ncorr)
end

x = data[:, 1]
y = data[:, [2, 4]]
p1 = plot(x, y,
	 linewidth = 2,
	 labels = ["Initial\npulse" "Corrected\npulse"],
	 layout = (2, 1),
	 title = ["x quadrature" " "])
plot!(xlim=(0, tf),
	  tickfont = font(12, "Serif"),
	  xlabel = "time "*L"t/t_\mathrm{f}",
	  ylabel = "pulse "*L"\times |\chi|^{-1}",
      guidefont = font(14, "Serif"),
	  legendfont = font(12, "Serif"),
	  margin = 5Plots.mm)

x = data[:, 1]
y = data[:, [3, 5]]
p2 = plot(x, y,
	 linewidth = 2,
	 labels = ["Initial\npulse" "Corrected\npulse"],
	 layout = (2, 1),
	 title = ["y quadrature" " "])
plot!(xlim=(0, tf),
	  tickfont = font(12, "Serif"),
	  xlabel = "time "*L"t/t_\mathrm{f}",
	  ylabel = "pulse "*L"\times |\chi|^{-1}",
	  guidefont = font(14, "Serif"),
	  legendfont = font(12, "Serif"),
	  background_color_legend = false,
	  margin = 5Plots.mm)
plot(p1, p2, layout=(1,2), size = (800, 800))
savefig("pulse.pdf")
# Add padding to the pulse
data = vcat(zeros(Npoints, 5), data, zeros(Npoints, 5))

fft_data = zeros(size(data))
fft_data[:, 1] = fftfreq(size(data)[1], tf/(Npoints-1))
for i = 2: 5
	fft_pulse = fft(data[:, i])
	fft_data[:, i] = abs.(fft_pulse)
end

a = div(size(data)[1], 2)
x = fft_data[1: a, 1]
y = fft_data[1: a, [2, 4]]
p1 = plot(x, y,
		  linewidth = 2,
		  labels = ["Initial\npulse" "Corrected\npulse"],
		  linestyle = [:solid :dash])
plot!(xlim=(0, 10),
	  tickfont = font(12, "Serif"),
	  xlabel = "Frequency "*L"\omega/|\chi|",
	  ylabel = L"|S_x(x)|",
	  guidefont = font(14, "Serif"),
	  legendfont = font(12, "Serif"),
	  background_color_legend = false,
	  margin = 5Plots.mm)

x = fft_data[1: a, 1]
y = fft_data[1: a, [3, 5]]
p2 = plot(x, y,
	 linewidth = 2,
	 labels = ["Initial\npulse" "Corrected\npulse"],
	 linestyle = [:solid :dash])
plot!(xlim=(0, 10),
	  tickfont = font(12, "Serif"),
	  xlabel = "Frequency "*L"\omega/|\chi|",
	  ylabel = L"|S_y(x)|",
	  guidefont = font(14, "Serif"),
	  legendfont = font(12, "Serif"),
	  background_color_legend = false,
	  margin = 5Plots.mm)

plot(p1, p2, layout=(1,2), size = (800, 400))

npzwrite("fft.npy", fft_data[1: div(Npoints,2), :])
savefig("fft.pdf")

max_stars_repo_name: thalesfr/MagnusCorrection

max_issues_repo_name: thalesfr/MagnusCorrection

max_forks_repo_name: thalesfr/MagnusCorrection


text:
using ConstrainedDynamics
using LinearAlgebra

include("examples/chain_in_chain.jl")

ConstrainedDynamics.setentries!(mech)
A1, _, b1 = ConstrainedDynamics.densesystem(mech)
x1 = A1\b1

ConstrainedDynamics.factor!(mech.graph,mech.ldu)
ConstrainedDynamics.solve!(mech)

A2, x2, b2 = ConstrainedDynamics.densesystem(mech)

@test isapprox(norm(x1-x2), 0.0; atol = 2e-5)