const ⊗ = kron

"""
Build Laplacian matrix using Kronecker products
"""
function get_Laplacian3d_kron( LF::LF3dGrid )
  IIx = speye(LF.Nx)
  IIy = speye(LF.Ny)
  IIz = speye(LF.Nz)
  return LF.LFx.D2jl ⊗ IIy ⊗ IIz +
         IIx ⊗ LF.LFy.D2jl ⊗ IIz +
         IIx ⊗ IIy ⊗ LF.LFz.D2jl
end

"""
A naive implementation to build Laplacian by manually, by searching for non-zero
indices and their values.
"""
function get_Laplacian3d_naive( LF::LF3dGrid )
  N = LF.Nx * LF.Ny * LF.Nz

  NNZ = calc_NNZ(LF)
  Tnz = [false, false, false]
  # Nabla2 matrix
  values = zeros(Float64,NNZ)
  column = zeros(Int64,NNZ)
  row    = zeros(Int64,NNZ)
  inz    = 0
  for ip2 = 1:N^3
    for ip1 = ip2:N^3
      Tnz[:] = false
      i1 = LF.lin2xyz[1,ip1]
      i2 = LF.lin2xyz[1,ip2]
      #
      j1 = LF.lin2xyz[2,ip1]
      j2 = LF.lin2xyz[2,ip2]
      #
      k1 = LF.lin2xyz[3,ip1]
      k2 = LF.lin2xyz[3,ip2]
      #
      nabla2 = 0.0
      #
      if j1 == j2 && k1 == k2
        nabla2 = nabla2 + LF.LFx.D2jl[i1,i2]
        Tnz[1] = true
      end
      #
      if i1 == i2 && k1 == k2
        nabla2 = nabla2 + LF.LFy.D2jl[j1,j2]
        Tnz[2] = true
      end
      #
      if i1 == i2 && j1 == j2
        nabla2 = nabla2 + LF.LFz.D2jl[k1,k2]
        Tnz[3] = true
      end
      #
      if any(Tnz) # update non-zero element
        inz = inz + 1
        # Notice that we do not exploit symmetry here
        values[inz] = nabla2
        column[inz] = ip1
        row[inz]    = ip2
        #@printf("%5d %5d: Non zero found, inz = %8d\n", ip1, ip2, )
        if ip1 != ip2 # not a diagonal element
          inz = inz + 1 # update NNZ again
          values[inz] = nabla2
          column[inz] = ip2
          row[inz]    = ip1
        end
      end # if Tnz
    end # for ip1
  end # for ip2
  return sparse(row, column, values)
end

function calc_NNZ( LF::LF3dGrid ; verbose=false )
  Npoints = LF.Nx * LF.Ny * LF.Nz
  #
  NNZ = 0
  Tnz = [false, false, false]
  # Nabla2 matrix
  for ip2 = 1:Npoints
    for ip1 = ip2:Npoints
      Tnz[:] = false
      i1 = LF.lin2xyz[1,ip1]
      i2 = LF.lin2xyz[1,ip2]
      #
      j1 = LF.lin2xyz[2,ip1]
      j2 = LF.lin2xyz[2,ip2]
      #
      k1 = LF.lin2xyz[3,ip1]
      k2 = LF.lin2xyz[3,ip2]
      #
      nabla2 = 0.0
      #
      if (j1 == j2 && k1 == k2) || (i1 == i2 && k1 == k2) || (i1 == i2 && j1 == j2)
        NNZ = NNZ + 1
        #@printf("%5d %5d: Non zero found, inz = %8d\n", ip1, ip2, NNZ)
        if ip1 != ip2 # not a diagonal element
          NNZ = NNZ + 1
          #@printf("%5d %5d: Non zero found, inz = %8d\n", ip1, ip2, NNZ)
        end
      end
    end # ip1
  end # ip2
  if verbose
    @printf("NNZ = %8d, nonzero percentage = %7.2f%%\n", NNZ, NNZ/(Npoints^2) * 100.0 )
  end
  return NNZ
end

function old_calc_NNZ( LF::LF3dGrid ; verbose=false )
  Npoints = LF.Nx * LF.Ny * LF.Nz
  #
  NNZ = 0
  Tnz = [false, false, false]
  # Nabla2 matrix
  for ip2 = 1:Npoints
    for ip1 = ip2:Npoints
      Tnz[:] = false
      i1 = LF.lin2xyz[1,ip1]
      i2 = LF.lin2xyz[1,ip2]
      #
      j1 = LF.lin2xyz[2,ip1]
      j2 = LF.lin2xyz[2,ip2]
      #
      k1 = LF.lin2xyz[3,ip1]
      k2 = LF.lin2xyz[3,ip2]
      #
      nabla2 = 0.0
      #
      if j1 == j2 && k1 == k2
        Tnz[1] = true
      end
      #
      if i1 == i2 && k1 == k2
        Tnz[2] = true
      end
      #
      if i1 == i2 && j1 == j2
        Tnz[3] = true
      end
      if any(Tnz)
        NNZ = NNZ + 1
        #@printf("%5d %5d: Non zero found, inz = %8d\n", ip1, ip2, NNZ)
        if ip1 != ip2 # not a diagonal element
          NNZ = NNZ + 1
          #@printf("%5d %5d: Non zero found, inz = %8d\n", ip1, ip2, NNZ)
        end
      end
    end # ip1
  end # ip2
  if verbose
    @printf("NNZ = %8d, nonzero percentage = %7.2f%%\n", NNZ, NNZ/(Npoints^2) * 100.0 )
  end
  return NNZ
end


"""
    create_surfaces_cylinder(cells, setup)

Create surface triangulation of [inner and] outer cylinders [and ECS].

The ground surface is triangulated first, before the walls are "extruded" and
the top surface is copied from the ground surface.
"""
function create_surfaces_cylinder(cells, setup::Setup)
    @unpack radii, centers = cells
    @unpack ncell, rmin, rmax, height, include_in, in_ratio, ecs_shape, ecs_ratio = setup

    # Choose approximate cylinder side length
    nside = 30
    nside_min = 12

    include_ecs = ecs_shape != "no_ecs"
    nboundary = (2 * include_in + 1 + include_ecs) * ncell + include_ecs
    rmean = (rmin + rmax) / 2

    # Create points on boundary circles
    create_nside(r) = max(nside_min, round(Int, nside * r / rmean))
    create_angles(n) = 2π * collect(0:n-1) / n
    create_circle(r, angles, center) = r * [cos.(angles'); sin.(angles')] .+ center
    create_circle_edges(n) = [(1:n-1)' n; (2:n)' 1]

    centers_columns = reshape([centers[:, i] for i = 1:ncell], (1, ncell))

    # Create points for IN compartments
    if include_in
        radii_in = in_ratio * radii
        nside_in = create_nside.(radii_in)
        angles_in = create_angles.(nside_in)
        circles_in = create_circle.(radii_in, angles_in, centers_columns)
        circle_edges_in = create_circle_edges.(nside_in)
        points_in = hcat(circles_in...)
        edges_in = hcat(circle_edges_in...)
        nprevious = 0
        for i = 1:ncell
            ncurrent = nprevious + nside_in[i]
            edges_in[:, nprevious+1:ncurrent] .+= nprevious
            nprevious = ncurrent
        end
    else
        # Create empty arrays
        nside_in = zeros(Int, 1, 0)
        points_in = zeros(2, 0)
        edges_in = zeros(Int, 2, 0)
    end

    # Create points of OUT compartments
    nside_out = create_nside.(radii)
    angles_out = create_angles.(nside_out)
    circles_out = create_circle.(radii, angles_out, centers_columns)
    circle_edges_out = create_circle_edges.(nside_out)
    points_out = hcat(circles_out...)
    edges_out = hcat(circle_edges_out...)
    nprevious = 0
    for i = 1:ncell
        ncurrent = nprevious + nside_out[i]
        edges_out[:, nprevious+1:ncurrent] .+= nprevious
        nprevious = ncurrent
    end

    if include_ecs
        radii_ecs = radii .+ ecs_ratio .* rmean
        nside_ecs = create_nside.(radii_ecs)
        angles_ecs = create_angles.(nside_ecs)
        circles_ecs = create_circle.(radii_ecs, angles_ecs, centers_columns)
        points_ecs = hcat(circles_ecs...)

        if ecs_shape == "box"
            # Determine bounds of domain
            pmin = minimum(points_ecs, dims = 2)
            pmax = maximum(points_ecs, dims = 2)

            # Define four corners of box ECS and their corresponding edges
            points_ecs = [
                pmin[1] pmax[1] pmax[1] pmin[1]
                pmin[2] pmin[2] pmax[2] pmax[2]
            ]
            edges_ecs = [
                1 2 3 4
                2 3 4 1
            ]
        elseif ecs_shape == "convex_hull"
            # Extract points defining convex hull of ECS
            edges_ecs, points_ecs = convexhull(points_ecs)
            npoint_ecs = size(points_ecs, 2)
        elseif ecs_shape == "tight_wrap"
            error("unimplemented")
            # ashape = alphashape(points_ecs)
            # edges_ecs = boundary(hull)
            # boundary_ecs = points_ecs[:, unique(edges_ecs)]
        end
        nedge_ecs = size(points_ecs, 2)
    else
        # Create empty arrays
        points_ecs = zeros(2, 0)
        edges_ecs = zeros(Int, 2, 0)
        nedge_ecs = zeros(Int, 1, 0)
    end

    # Create global collection of points and edges
    # Adjust numbering accordingly
    npoint_in = size(points_in, 2)
    npoint_out = size(points_out, 2)
    npoint_ecs = size(points_ecs, 2)
    points = [points_in points_out points_ecs]
    edges = [edges_in edges_out .+ npoint_in edges_ecs .+ npoint_in .+ npoint_out]
    npoint = size(points, 2)
    nedge = size(edges, 2)

    # Perform Delaynay triangulation of entire domain
    triin = TriangulateIO()
    triin.pointlist = points
    triin.segmentlist = edges
    triin.segmentmarkerlist = 1:nedge
    triout, = triangulate("pQ", triin)
    triangles = eachcol(triout.trianglelist)
    ntriangle = length(triangles)

    boundary_bounds = cumsum([0 nside_in nside_out nedge_ecs], dims = 2)

    ncell_in = ncell * include_in

    find_boundary(node) =
        findfirst(boundary_bounds[1:end-1] .+ 1 .≤ node .≤ boundary_bounds[2:end])

    # Create boundary numbers
    facetmarkers = zeros(Int, ntriangle)
    for (i, t) ∈ enumerate(triangles)
        b = find_boundary.(t)
        b_in = b .≤ ncell_in
        b_out = ncell_in .< b .≤ ncell_in + ncell
        if all(b_in)
            # Inside IN compartment
            facetmarkers[i] = ncell_in + include_ecs * ncell + b[1]
        elseif any(b_in)
            # Triangle touches IN compartment, but is not inside. It is thus the corresponding OUT compartment.
            facetmarkers[i] = ncell_in + include_ecs * ncell + b[.!b_in][1]
        elseif all(b_out) && b[1] == b[2] == b[3]
            # Triangle lies fully within the same OUT comparment
            facetmarkers[i] = ncell_in + include_ecs * ncell + b[1]
        else
            # Triangle lies in the ECS
            facetmarkers[i] = nboundary
        end
    end

    facets = hcat(triangles...)

    # Copy points, edges and markers two the two 3D planes top and bottom
    z = fill(height / 2, 1, npoint)
    points = [points points; -z z]

    function create_sidefacets(le, ri)
        [
            (le:ri-1)' ri (le+1:ri)'.+npoint le.+npoint
            (le:ri-1)'.+npoint ri.+npoint (le:ri-1)'.+npoint ri.+npoint
            (le+1:ri)' le (le+1:ri)' le
        ]
    end

    sidefacets =
        hcat(create_sidefacets.(boundary_bounds[1:end-1] .+ 1, boundary_bounds[2:end])...)

    sidefacetmarkers_in = 1:ncell_in
    sidefacetmarkers_in =
        vcat([repeat([m, m], nside_in[i]) for (i, m) ∈ enumerate(sidefacetmarkers_in)]...)

    sidefacetmarkers_out = (1+!include_ecs)ncell_in+1:(1+!include_ecs)ncell_in+ncell
    sidefacetmarkers_out =
        vcat([repeat([m, m], nside_out[i]) for (i, m) ∈ enumerate(sidefacetmarkers_out)]...)

    sidefacetmarkers_ecs = fill(nboundary, 2npoint_ecs)

    facets = [facets facets .+ npoint sidefacets]
    facetmarkers = [
        facetmarkers
        facetmarkers
        sidefacetmarkers_in
        sidefacetmarkers_out
        sidefacetmarkers_ecs
    ]

    if include_in
        regions_in = [centers; zeros(1, ncell)]
        regions_out = regions_in + [1; 0; 0] * (radii_in + radii) / 2
    else
        regions_in = zeros(3, 0)
        regions_out = [centers; zeros(1, ncell)]
    end
    if include_ecs
        xmin = argmin(points_ecs[1, :])
        region_ecs = [points_ecs[:, xmin]; 0] + [ecs_ratio * rmean / 2; 0; 0]
    else
        region_ecs = zeros(3, 0)
    end
    regions = [regions_in regions_out region_ecs]

    (; points, facets, facetmarkers, regions)
end


function synthesise_coupler(n::Int, _K, _BW_REL, POLYNOMIAL = :Butterworth)
    # POLYNOMIAL = :Butterworth
    # POLYNOMIAL = :Chebyshev

    # Parse inputs to desired precision (input string in ord)
    K = isa(_K, String) ? parse(BigFloat, _K) : convert(BigFloat, _K)
    BW_REL = isa(_BW_REL, String) ? parse(BigFloat, _BW_REL) : convert(BigFloat, _BW_REL)

    # Polynomial for pole extraction
    P_x̄ = -design_polynomial(POLYNOMIAL, n-1, BigFloat, :x̄)

    # Calculate related constants
    tc = tand(45 * (1 - BW_REL/2))      # Always a real number
	xc = real(x(im*tc))                 # Always a real number
    K_ = parse(BigFloat,"0.5") * K * P_x̄.coeffs[end];

	# Design function for calculations
	Γ_T(t) = parse(BigFloat,"0.5") * K * t * P_x̄(x(t)/xc)

    ## REFLECTION COEFFICIENT
    # Polynomial roots
    r̄ = roots(P_x̄)

    # Numerator zeros
    Z_n = begin
        # Preallocate zeros (for both functions and add )
        local Z_n = zeros(Complex{BigFloat}, 2n-1)
        for i in eachindex(r̄)
            Z_n[2i:(2i+1)] = [im, -im] .* sqrt((1-xc*r̄[i])/(1+xc*r̄[i]))
        end
        Z_n
    end

    # Denominator zeros
    LHS = Polynomial(xc .* BigFloat[1,0,-1], :t)^(n-1)
    RHS = begin
        local RHS_ = K_ * Polynomial(BigFloat[0,1], :t);
        for i in eachindex(r̄)
            RHS_ = RHS_ *
                (Polynomial(BigFloat[1,0,1], :t) - xc*r̄[i]*Polynomial(BigFloat[1,0,-1], :t))
        end
        RHS_
    end
    Z_d = begin
        Z_d_a = roots(LHS - RHS);
        Z_d_b = roots(LHS + RHS);
        # Join roots from both parentheses
        Z_d_ = vcat(Z_d_a, Z_d_b);
        # Select only left half-plane
        Z_d = Z_d_[sortperm(real.(Z_d_))[1:2n-1]]
    end

    # Coefficient reconstruction
    Γ_N = -1 * Polynomial(real.(fromroots(Complex{BigFloat}.(Z_n)).coeffs), :t)
    Γ_D = Polynomial(real.(fromroots(Complex{BigFloat}.(Z_d)).coeffs), :t)
    Γ(t) = (Γ_N(t)/Γ_D(t))

    ## TRANSMISSION COEFFICIENT
    _Kt = sqrt(abs(((LHS+RHS)*(LHS-RHS)).coeffs[end]))
	T(t) = -LHS(t)/(_Kt * Γ_D(t))

    ## DARLINGTON SYNTHESIS
    # Preallocate vectors
    a = Vector{BigFloat}(undef, n)
    b = Vector{BigFloat}(undef, n-1)

    # Initial admittance functions
    Y_N = Γ_D - Γ_N;    # Numerator of admittance
    Y_D = Γ_D + Γ_N;    # Denominator of admittance
    dY_in = polynomial_fraction_derivative(Y_N, Y_D) # Derivative of admittance

    # Extract first branch as an open-circuit stub
    a[1] = dY_in(1)
    a[end] = a[1]

    # Update input admittance and derivative
    Y_N = Y_N - a[1] * Polynomial(BigFloat[0,1], :t) * Y_D
    Y_D = Y_D
    dY_in = polynomial_fraction_derivative(Y_N, Y_D)

    # Iterate through all "transmission line followed by shunt open-circuit stub" sections
    for i in 2:Int(floor(n/2)+1)
        # Extract transmission line
        b[i-1] = Y_N(1)/Y_D(1)
        b[end-i+2] = b[i-1]

        # Get new admittance
        # Calculate numerator and denominator naïvely
        Y_N_ = Y_N * Polynomial(BigFloat[1,0,1], :t) * b[i-1] -
            2 * Y_D * Polynomial(BigFloat[0,1], :t) * (b[i-1]^2)
        Y_D_ = Y_D * Polynomial(BigFloat[1,0,1], :t) * b[i-1] -
            2 * Y_N * Polynomial(BigFloat[0,1], :t)
        # Remove common roots
        rn = roots(Y_N_) # Roots of naïve numerator
        rd = roots(Y_D_) # Roots of naïve denominator
        _goodi = fill(true,length(rn)); # Flags for non-common root
        for j in eachindex(rn)
            # Find first root in common within the range of precision
            #   of reconstructed function
            k = findfirst(isapprox.(rn[j], rd, atol = eps(Float64(abs(rn[j])))))
            # If there exists, pop from denominator and flag as common root for numerator
            isnothing(k) || (_goodi[j] = false; popat!(rd,k))
        end
        rn = rn[_goodi]; # Remove numerator common roots
        # Get coefficients with potentially spurious imaginary parts
        c_Y_N = (Y_N_.coeffs[end] * fromroots(Complex{BigFloat}.(rn); var = :t)).coeffs
        c_Y_D = (Y_D_.coeffs[end] * fromroots(Complex{BigFloat}.(rd); var = :t)).coeffs
        # Build admittance functions without spurious imaginary parts
        Y_N = Polynomial(real.(c_Y_N), :t)
        Y_D = Polynomial(real.(c_Y_D), :t)
        dY_in = polynomial_fraction_derivative(Y_N, Y_D)

        # Extract shunt line
        a[i] = dY_in(1)
        a[end-i+1] = a[i]

        # Prepare next admittance
        Y_N = Y_N - a[i] * Polynomial(BigFloat[0,1], :t) * Y_D
        Y_D = Y_D
        dY_in = polynomial_fraction_derivative(Y_N, Y_D)
    end

    return Coupler(a, b)
end

@inline x(t) = (1 + t^2) / (1 - t^2)

@inline t²(x) = (x-1)/(x+1)

@inline t(x) = sqrt(Complex(t²(x)))

function design_polynomial(family, degree, type_precision = Float64, polyvar = :x)
	p = zeros(type_precision, degree+1)
	p[end] = 1
    if family == :Butterworth
        return Polynomial(p, polyvar)
    elseif family == :Chebyshev
        return convert(Polynomial{type_precision}, ChebyshevT(p, polyvar))
    end
end

function polynomial_fraction_derivative(N, D)
	dN = derivative(N)
	dD = derivative(D)
	return (t) -> ((dN(t)*D(t) - dD(t)*N(t))/((D(t))^2))
end


include(joinpath(@__DIR__, "../../src/FLS.jl"))

"""
Cartesian Product
"""

pairschema = @acset LabeledGraph begin
    V = 2
    E = 2
    vlabel = [:s, :s2]
    elabel = [:p1, :p2]
    src = [2,2]
    tgt = [1,1]
end

paird= @acset LabeledGraph begin
    V = 2
    vlabel = [:s, :s]
end

pairs = FLS(:pairs, pairschema, [
    Cone(lconed, :s2, [1=>:p1,2=>:p2])], [])


@testset "time integration" begin
    D = [-1.0 2.0*π; -2.0*π -1.0]
    x0 = [1.0 0.0]
    test_data = (x0*exp(D))[:]
    e = ForwardEuler(1e-3)
    heuns = Heuns(1e-3)
    ssprk3 = StableRK3(1e-3)
    ssprk4 = StableRK4(1e-3)
    @test test_data≈integrate_time(x0,D,1.0,e) atol=e.step_size*10.0
    @test test_data≈integrate_time(x0,D,1.0,heuns) atol=heuns.step_size*10.0
    @test test_data≈integrate_time(x0,D,1.0,ssprk3) atol=ssprk3.step_size*10.0
    @test test_data≈integrate_time(x0,D,1.0,ssprk3) atol=ssprk4.step_size*10.0

    @testset "limiter -- MUSCL" begin
        # using DiscretisedFluidQueues
        one_D_T = [0.0][:,:]
        one_D_c = [1.0]
        one_D_model = BoundedFluidQueue(one_D_T,one_D_c,2.0)

        nodes = collect(0.0:2.0/50.0:2.0)

        mesh = DGMesh(nodes,7)#FRAPMesh(nodes,7)#

        dq = DiscretisedFluidQueue(one_D_model,mesh)
        B = build_full_generator(dq)
        # B[2,(end-n_bases_per_cell(mesh)+1):end] =
        #     B[n_bases_per_cell(mesh)+2,2:n_bases_per_cell(mesh)+1]
        # B[2,1] = 0.0
        t = 1.0
        f(x,i) = Float64(x<0.5)#sin(π*x)#
        d0 = SFMDistribution(f,dq)# interior_point_mass(-0.5,1,dq)#

        pdf0 = pdf(d0)
        pdf0_limit = pdf(limit(d0))
        x_vals = (-0.55:0.001:-0.45).+1
        @test any(pdf0.(x_vals,1).>1.01)
        @test !any(pdf0_limit.(x_vals,1).>1.01)
        @test any(pdf0.(x_vals,1).<-0.01)
        @test !any(pdf0_limit.(x_vals,1).<-0.01)

        dt = integrate_time(d0,B,t,StableRK4(0.01))
        dt_no_limit = integrate_time(d0,B,t,StableRK4(0.01); limiter=NoLimiter)

        pdft = pdf(dt)
        pdft_no_limit = pdf(limit(dt_no_limit))
        x_vals = (-0.05:0.001:0.55).+1
        @test !any(pdft.(x_vals,1).>1.01)
        @test any(pdft_no_limit.(x_vals,1).>1.01)
        @test !any(pdft.(x_vals,1).<-0.01)
        @test any(pdft_no_limit.(x_vals,1).<-0.01)
    end
end


# Optimization coordinates
function opt(known_pts, known_unf, to_unf, search, neigh, guess)
	idxs, dists = get_neighbors(known_pts, to_unf, search, neigh, true)
	out_coords  = zeros(Float64, size(to_unf))

	Threads.@threads for i in 1:length(idxs)
		locs = view(known_unf, :, idxs[i])
		initguess = guess[:,i]
        opt = optimize(x->mse_coords(x, locs, dists[i]), initguess)
        res = Optim.minimizer(opt)
        out_coords[:,i] .= res
	end
    out_coords
end

# Get surface normals
function getnormals(ref_surf::AbstractMatrix, search, neigh)
	idxs, dists = get_neighbors(ref_surf, search, neigh)
	n = length(idxs)

	ref_normals = zeros(Float64, size(ref_surf))
	ignore = Int[]

	for i in 1:n
		M  = fit(PCA, view(ref_surf, :, idxs[i]), maxoutdim=3, pratio=1)
		ev = projection(M)
		ndims = size(ev,2)

		if ndims <= 1
			push!(ignore,i)
		elseif ndims == 2
			ref_normals[:,i] .= cross(view(ev,:,1),view(ev,:,2))
		else
			ref_normals[:,i] .= view(ev,:,3)
		end
	end

	# get a sequential path to visit all nodes
	# might be possible to use unf indices to create a path in a simpler way in the future
	good = setdiff(1:n, ignore)
	ok = false
	n  = length(good)
	to_loop, visited, refx = Int64[], Int64[], Int64[]

	xi = good[floor(Int,n/2)]
	append!(visited,xi)
	append!(to_loop, idxs[xi])
	append!(refx, [xi for x in idxs[xi]])
	extra_loop = setdiff(to_loop,visited)
	while ok==false
		xj = extra_loop[1]
		append!(visited,xj)
		append!(to_loop, idxs[xj])
		append!(refx, [xj for x in idxs[xj]])
		extra_loop = setdiff(to_loop,visited)
		if length(extra_loop)==0
			ok=true
		end
	end

	ids_unique = unique(i -> to_loop[i], 1:length(to_loop))
	path = [p=>q for (p,q) in zip(refx[ids_unique],to_loop[ids_unique])]

	# Iterate through all normals and check for consistency
	prev = Int[]
	for x in 1:n
		i, j = path[x]

		prev = union(prev,i)
		chk = intersect(prev,idxs[j])
		nx = length(chk)

		d = [evaluate(CosineDist(), ref_normals[:,k], ref_normals[:,j]) for k in chk]
		sum(d .>= 1.0) >= ceil(nx/2) && (ref_normals[:,j] .*= -1)
		prev = union(prev,j)
	end

	ref_normals, good
end

# Error function
function mse_coords(x, locations, distances)
    mse = 0.0
    for k in 1:length(distances)
        loc, d = view(locations,:,k), distances[k]
        distance_calculated = euclidean([x[1], x[2], x[3]], loc)
        mse += (distance_calculated - d)^2
    end
    mse / length(distances)
end

# Guess initial XYZ
function firstguess(to_unf, ref_pts, unf_ref, ref_normals=nothing)
	# get the closest point as initial guess
	idxs, dists = get_neighbors(ref_pts, to_unf, :knn, 1)
	n   = length(idxs)
	ids = [idxs[x][1] for x in 1:n]
	guess = unf_ref[:,ids]

	# if reference surface normals are informed, define z as the signed distance
	if !isnothing(ref_normals)
		for (x, id) in enumerate(ids)
			refcoords = view(ref_pts, :, id)
			refnormal = view(ref_normals, :, id)
			ref_vect = view(to_unf, :, x) - refcoords
			dotn_val = dot(refnormal, ref_vect)
			guess[3,x] = dotn_val < 0 ? -1*dists[x][1] : dists[x][1]
		end
	end

	guess
end

"""
	errors(mode, coords, unf_coords; search=:knn, neigh=16, maxerr=5)

Unfolding distorts the original distances between neighbor points. This
function returns two types of results depending on `mode`.

If `mode = :ids`, this gives the IDs of the points above and below the `maxerr`
threshold as a tuple such that (ids_passed_the_tests, ids_not_passed_the_tests)

If `mode = :dists`, this function output the difference of the expected distance
for each pair analyzed in the neighborhood. Can be used as input in a boxplot to
verify distortions.

## Parameters:

* `mode`       - type of error, :ids or :dists.
* `coords`     - coordinate matrix of the points before unfolding.
* `unf_coords` - coordinate matrix of the points after unfolding.
* `search`     - type of neighborhood, :knn or :radius.
* `neigh`      - number of neighbors or radius used to make the validations.
* `maxerr`     - the maximum accepted absolute difference of the distances
  if mode = :ids.
"""
function errors(mode::Symbol, coords::AbstractMatrix, unf_coords::AbstractMatrix;
	            search=:knn, neigh=16, maxerr=5)
	idxs, dists = get_neighbors(coords, search, neigh, true)
	if mode == :ids
		error_ids(coords, unf_coords, search, neigh, maxerr)
	elseif mode == :dists
		error_dists(coords, unf_coords, search, neigh)
	else
		throw(ArgumentError("mode should be :ids or :dists"))
	end
end

function error_ids(coords, unf_coords, search=:knn, neigh=16, maxerr=5)
	idxs, dists = get_neighbors(coords, search, neigh, true)
	bad_ids = Int[]
	n = length(idxs)

	for x in 1:n
		p0 = view(unf_coords,:,x)
		tdists = [euclidean(p0, view(unf_coords,:,y)) for y in idxs[x]]
		tdists .-= dists[x]
		ids = findall(abs.(tdists) .> maxerr)

		if length(ids) > 0
			push!(bad_ids,x)
			append!(bad_ids,idxs[x][ids])
		end
	end

	bad  = unique(bad_ids)
	good = setdiff(collect(1:n), bad)
	good, bad
end

function error_dists(coords, unf_coords, search=:knn, neigh=16)
	!(coords[1] isa Float64) && (coords = Float64.(coords))
	idxs, dists = get_neighbors(coords, search, neigh, true)
	error = Float64[]

	for x in 1:length(idxs)
		p0 = view(unf_coords,:,x)
		tdists = [euclidean(p0, view(unf_coords,:,y)) for y in idxs[x]]
		tdists .-= dists[x]
		tdists .= abs.(tdists)
		mask = idxs[x] .!= x

		append!(error,tdists[mask])
	end

	error
end


#MDDatasets: Measure clock signals
#-------------------------------------------------------------------------------
#=NOTE:
 - Many functions make use of "cross" functions.
 - Naming assumes x-values are "time".
=#

#==
===============================================================================#

#measperiod: Measure period between successive zero-crossings:
#-------------------------------------------------------------------------------
function measperiod(d::DataF1; nmax::Integer=0, tstart::Real=-Inf,
	xing::CrossType=CrossType(), shiftx=true)

	dx = xcross(d, nmax=nmax, xstart=tstart, allow=xing)
	return delta(dx, shiftx=shiftx)
end

#measfreq: Measure 1/period between successive zero-crossings:
#-------------------------------------------------------------------------------
function measfreq(d::DataF1; nmax::Integer=0, tstart::Real=-Inf,
	xing::CrossType=CrossType(), shiftx=true)

	T = measperiod(d, nmax=nmax, tstart=tstart, xing=xing, shiftx=shiftx)
	return DataF1(T.x, 1 ./ T.y)
end

#measduty: Measure duty cycle of a periodic signal @ zero-crossings:
#TODO: Improve: Could break in the presence of noise.
#-------------------------------------------------------------------------------
function measduty(d::DataF1; nmax::Integer=0, tstart::Real=-Inf, shiftx=true)
	xrise1 = xcross1(d, xstart=tstart, allow=CrossType(:rise))
	xrise = xveccross(d, nmax, tstart, CrossType(:rise))
	xfall = xveccross(d, nmax, xrise1, CrossType(:fall))
	#Need 1 more rise to get period
	numpts = min(length(xrise)-1, length(xfall))
	x = shiftx ? meanadj(xrise[1:(numpts+1)]) : xrise[1:numpts]
	duty = similar(xrise, numpts)

	for i in 1:numpts
		T = xrise[i+1]-xrise[i]
		duty[i] = (xfall[i]-xrise[i])/T
	end

	return DataF1(x, duty)
end

#measckstats: Measure clock statistics @ zero-crossings:
#-------------------------------------------------------------------------------
#=NOTE
A bit awkward to implement because broadcast system does not work with this
return type (Dict).  So, to work with sweeps (DataHR/DataRS), we must only
use high-level (broadcastable) functions.  We cannot use functions that
return arbitrary vectors, or tuples, or ...

OR:
We could loop across all sweep dimensions manually....

TODO:
 - Figure out if stats are defined correctly.
 - Should there be a function for relative values vs referenced to an ideal clock?
 - Improve: Could break in the presence of noise.
=#
function measckstats(d::DataMD; tstart=-Inf, tck=nothing)
	if nothing==tck
		msg = "measckstats: Must provide value to tck."
		throw(ArgumentError(msg))
	end
	xrise1 = xcross1(d, xstart=tstart, allow=CrossType(:rise))
	xrise = xcross(Event, d, xstart=tstart, allow=CrossType(:rise))
	xfall = xcross(Event, d, xstart=xrise1, allow=CrossType(:fall))

	#Want to match up rise/fall events:
	maxevent = min(maximum(xval(xrise)-1), maximum(xval(xfall)))
	#TODO: cnovert to fixed-point to avoid roundoff errors?
	xrise = clip(xrise, xmax=maxevent+1)
	xfall = clip(xfall, xmax=maxevent)

	#Cycle-to-cycle jitter:
	period = delta(xrise, shiftx=false)
	jitc2c = std(period)

	#Match up # of events:
	xrise = clip(xrise, xmax=maxevent)

	tph = xfall - xrise
	tpl = period - tph

	duty = tph/period
	dcd = duty-0.5 #Is this off by a factor of 2??

	#Measure mean values, relative to absolute clock:
#-----------
	#Wrap crossings back around first point:
	_tshift = xval(xrise)*tck
	xrise = xrise-_tshift
	xfall = xfall-_tshift

	#Jitter relative to absolute reference:
	jit = std(xrise)

	#Mean rise/fall times:
	μxr = mean(xrise)
	μxf = mean(xfall)

	#Ensure rise happens first for all runs:
	risefirst = μxf>μxr #Probably not needed (see xrise1)
	xfall = xfall+(1-risefirst)*tck #High pulse
	μxf = mean(xfall)

	#Mean pulse high/low values:
	μtph = μxf-μxr
	μtpl = tck - μtph

	#Mean duty cycle values:
	μduty = μtph / tck
	μdcd = μduty-0.5 #Is this off by a factor of 2??

	return Dict{Symbol, Any}(
		:period => tck,
		:jitter => jit,
		:jitter_c2c => jitc2c,

		:std_tph => std(tph),
		:std_tpl => std(tpl),
		:std_duty => std(duty),
		:std_dcd => std(dcd),

		:min_tph => minimum(tph),
		:min_tpl => minimum(tpl),
		:min_duty => minimum(duty),
		:min_dcd => minimum(dcd),

		:max_tph => maximum(tph),
		:max_tpl => maximum(tpl),
		:max_duty => maximum(duty),
		:max_dcd => maximum(dcd),

		#Curretntly from absolute measurements only:
		:mean_tph => μtph,
		:mean_tpl => μtpl,
		:mean_duty => μduty,
		:mean_dcd => μdcd,
	)
end

#Last line


# class constructors
Frac(::Type{R}) where R<:Ring = Frac{R}
Frac(::Type{R}) where R<:Integer = QQ{R}

# construction
category_trait(Z::Type{<:Frac}) = category_trait_fraction(category_trait(basetype(Z)))
basetype(::Type{<:Frac{T}}) where T = T
depth(::Type{<:Frac{T}}) where T = depth(T) + 1
copy(a::Frac) = typeof(a)(a.num,a.den, NOCHECK)

numerator(a::FractionRing) = a.num
denominator(a::FractionRing) = a.den

issimpler(a::T, b::T) where T<:Frac = issimpler(a.num, b.num)

Frac{T}(a::Frac{T}) where T = a
Frac{T}(a::Frac{S}) where {T,S} = Frac{T}(T(a.num), T(a.den), NOCHECK)

Frac{T}(a::Integer) where T = Frac{T}(T(a), one(T), NOCHECK)
Frac{T}(a::Rational) where T = Frac{T}(T(a.num), T(a.den), NOCHECK)
Frac{T}(a::Ring) where T = Frac{T}(T(a), one(T), NOCHECK)
Frac(a::T) where T<:Ring  = Frac{T}(a)
Frac(a::T) where T<:Integer = Frac{ZZ{T}}(a)
Frac(a::Rational{T}) where T<:Integer = Frac{ZZ{T}}(a)
Frac{T}(a::Integer,b::Integer) where T = Frac(T(a), T(b))
function Frac(a::T, b::T) where T<:Polynomial
    cab = content(a) // content(b)
    a = primpart(a)
    b = primpart(b)
    g = pgcd(a, b)
    a /= g
    b /= g
    a *= numerator(cab)
    b *= denominator(cab)
    g = pgcd(a, b)
    a /= g
    b /= g
    s = lcunit(b)
    b /= s
    a /= s
    Frac{T}(a, b, NOCHECK)
end
function Frac(a::T, b::T) where T<:ZZ
    g = pgcd(a, b)
    a /= g
    b /= g
    s = lcunit(b)
    b /= s
    a /= s
    Frac{T}(a, b, NOCHECK)
end
//(a::T, b::T) where T<:Ring = Frac(a, b)
Frac{T}(a, b) where T = Frac(T(a), T(b))

_promote_rule(::Type{Frac{T}}, ::Type{Frac{S}}) where {S,T} = Frac{promote_type(S,T)}
_promote_rule(::Type{Frac{T}}, ::Type{S}) where {S<:Ring,T} = Frac{promote_type(S,T)}
promote_rule(::Type{Frac{T}}, ::Type{S}) where {S<:Integer,T} = Frac{promote_type(S,T)}
promote_rule(::Type{Frac{T}}, ::Type{Rational{S}}) where {S<:Integer,T} = Frac{promote_type(S,T)}

lcunit(a::Frac) = inv(lcunit(a.den))

# induced homomorphism
function (h::Hom{F,R,S})(a::Frac{<:R}) where {F,R,S}
    Frac(h.f(a.num), h.f(a.den))
end

Base.isless(p::T, q::T) where T<:Frac = isless(p.num * q.den, q.num * p.den)

# operations for Frac

function +(x::T, y::T) where T<:Frac
    a, b, c, d = x.num, x.den, y.num, y.den
    h = pgcd(b, d)
    b /= h
    d /= h
    n = a * d + b * c
    g = pgcd(n, h)
    T(n / g, h / g * b * d, NOCHECK)
end

function *(x::T, y::T) where T<:Frac
    a, b, c, d = x.num, x.den, y.num, y.den
    g = pgcd(a, d)
    a /= g
    d /= g
    g = pgcd(b, c)
    b /= g
    c /= g
    T(a * c, b * d, NOCHECK)
end
function inv(x::T) where T<:Frac
    T(x.den, x.num, NOCHECK)
end

==(a::T, b::T) where T<:Frac = iszero(a - b)
/(a::T, b::T) where T<:Frac = a * inv(b)
-(a::T, b::T) where T<:Frac = +(a, -b)
-(a::Frac{T}) where T = Frac{T}(-a.num, a.den, NOCHECK)
divrem(a::T, b::T) where T<:Frac = (a / b, zero(T))
div(a::T, b::T) where T<:Frac = a / b
rem(a::T, b::T) where T<:Frac = zero(T)

isunit(a::Frac) = !iszero(a.num)
isone(a::Frac) = a.num == a.den
iszero(a::Frac) = iszero(a.num)
zero(::Type{Frac{T}}) where T = Frac(zero(T), one(T))
one(::Type{Frac{T}}) where T = Frac(one(T), one(T))
hash(a::Frac, h::UInt) = hash(a.den, hash(a.num, h))

evaluate(p::Frac, a) = Frac(evaluate(p.num, a), evaluate(p.den, a))
(p::Frac)(a, b...) = evaluate(p, a, b...)

function show(io::IO, a::Frac)
    if isone(a.den)
        show(io, a.num)
    else
        print(io, '(', a.num, ")/(", a.den, ')')
    end
end


export MagPhase, zerophase
export SignalElement

import Base: Complex, promote_rule, abs, angle

struct MagPhase{T<:Real} <: Number
    mag::T
    phi::T
    function MagPhase{T}(mag::T, phi::T) where {T <: Real}
        mag < 0 && throw(ArgumentError("negative magnitude!"))
        new(mag, phi)
    end
end

const SignalElement{T} = Union{Complex{T}, MagPhase{T}} where T <: Real

MagPhase(x::T, y::T) where {T <: Real} = MagPhase{T}(x, y)
MagPhase(x::Real, y::Real) = MagPhase(promote(x,y)...)
MagPhase(x::Real) = MagPhase(x, zero(x))
MagPhase(z::Complex) = MagPhase(abs(z), angle(z))
MagPhase{T}(z::C) where {T, C <: Complex{<:T}} = MagPhase{T}(abs(z), angle(z))
MagPhase{T}(x) where T = MagPhase(convert(T, x))

Base.Complex{T}(z::MagPhase{T}) where T = Complex{T}(z.mag * cos(z.phi), z.mag * sin(z.phi))
Base.Complex(z::MagPhase{T}) where T = Complex{T}(z)

Base.abs(z::MagPhase) = z.mag
Base.angle(z::MagPhase) = z.phi

zerophase(z::MagPhase) = MagPhase(abs(z))
zerophase(z::Complex) = Complex(abs(z))

Base.promote_rule(::Type{Complex{T}}, ::Type{MagPhase{T}}) where T = Complex{T}

function Base.show(io::IO, z::MagPhase)
    print(io, "(")
    show(io, z.mag)
    print(io, ", ")
    show(io, z.phi)
    print(io, ")")
end


function moment0(spectrum::Vector{<:Real}, channel_width::Real)
    return sum(spectrum) * channel_width;
end

function moment0(cube::AbstractSpectralCube)
    m0 = mapslices(i->moment0(i, cube.z_resolution), cube.data; dims=[3])[:, :, 1]
    header = header2d_of(cube, cube.header["LINE"], "Moment0")
    return create_map(m0, header)
end


for N in [Float64, Rational{Int}, Float32]
    # random ball
    rand(Ball1)

    # 1D Ball1
    b = Ball1(N[0], N(1))
    # dimension
    @test dim(b) == 1
    # support vector
    d = N[1]
    @test σ(d, b) == N[1]
    d = N[-1]
    @test σ(d, b) == N[-1]

    # 2D Ball1
    b = Ball1(N[0, 0], N(1))
    # dimension
    @test dim(b) == 2
    # support vector
    d = N[1, 0]
    @test σ(d, b) == N[1, 0]
    d = N[-1, 0]
    @test σ(d, b) == N[-1, 0]
    d = N[0, 1]
    @test σ(d, b) == N[0, 1]
    d = N[0, -1]
    @test σ(d, b) == N[0, -1]
    d = N[0, 0]
    @test σ(d, b) == N[0, 0]

    # 2D Ball1 not 0-centered
    b = Ball1(N[1, 2], N(1))
    # dimension
    @test dim(b) == 2
    # support vector
    d = N[1, 0]
    @test σ(d, b) == N[2, 2]
    d = N[-1, 0]
    @test σ(d, b) == N[0, 2]
    d = N[0, 1]
    @test σ(d, b) == N[1, 3]
    d = N[0, -1]
    @test σ(d, b) == N[1, 1]

    # 2D Ball1 radius != 1
    b = Ball1(N[0, 0], N(2))
    # dimension
    @test dim(b) == 2
    # support vector
    d = N[1, 0]
    @test σ(d, b) == N[2, 0]
    d = N[-1, 0]
    @test σ(d, b) == N[-2, 0]
    d = N[0, 1]
    @test σ(d, b) == N[0, 2]
    d = N[0, -1]
    @test σ(d, b) == N[0, -2]

    # center
    @test center(b) == N[0, 0]

    # boundedness
    @test isbounded(b)

    # isempty
    @test !isempty(b)

    # isuniversal
    answer, w = isuniversal(b, true)
    @test !isuniversal(b) && !answer && w ∉ b

    # an_element & membership function
    e = an_element(b)
    answer = e ∈ b
    @test answer == ∈(e, b, true) == ∈(e, b, false) == true
    e += N[2 * b.radius, 1]
    answer = e ∈ b
    @test answer == ∈(e, b, true) == ∈(e, b, false) == false

    # translation
    @test translate(b, N[1, 2]) == Ball1(N[1, 2], N(2))

    # vertices_list
    vl = vertices_list(b)
    @test ispermutation(vl, [N[2, 0], N[0, 2], N[-2, 0], N[0, -2]])

    # check that vertices_list for zero radius doesn't repeat vertices
    b = Ball1(N[1, 2], N(0))
    vl = vertices_list(b)
    @test vl == [center(b)]

    # list of constraints
    b = Ball1(N[0, 0], N(1))
    clist = constraints_list(b)
    @test ispermutation(clist, [HalfSpace(N[1, 1], N(1)),  # x + y <= 1
                                HalfSpace(N[-1, 1], N(1)),  # -x + y <= 1
                                HalfSpace(N[1, -1], N(1)), # x - y <= 1
                                HalfSpace(N[-1, -1], N(1))])  # x + y >= -1
end


module SimplePolynomials
using Mods, Primes, Multisets

import Base: getindex, (==), show, zero, one, eltype, adjoint
import Base: numerator, denominator, big, inv

import Polynomials: degree, Polynomial, coeffs, roots, derivative

export SimplePolynomial, degree, coeffs, getx, Polynomial, roots
export derivative, integral, lead, make_function

# IntegerX is any sort of real or Gaussian integer
IntegerX = Union{S,Complex{S},AbstractMod} where {S<:Integer}

# RationalX is a Rational or Complex Rational based on integers
RationalX = Union{Rational{S},Complex{Rational{S}}} where {S<:Integer}

# CoefX is the type we allow for coefficients ... all exact
CoefX = Union{IntegerX,RationalX}

_enlarge(x::Integer) = big(x)
_enlarge(x::Rational) = big(x)
_enlarge(x::Complex) = big(x)
_enlarge(x::AbstractMod) = x

struct SimplePolynomial
    data::Vector
    func::Function
    function SimplePolynomial(list::Vector{T}) where {T<:CoefX}
        list = _enlarge.(_chomp(list))
        f = x -> @evalpoly(x, list...)
        new(_chomp(list), f)
    end
end

function SimplePolynomial(c...)
    n = length(c)
    T = typeof(sum(c))
    data = zeros(T, n)
    for j = 1:n
        data[j] = c[j]
    end
    SimplePolynomial(data)
end

"""
`make_function(p::SimplePolynomial)` or
`make_function(f::SimpleRationalFunction)`
returns a Julia `Function` that evaluates `p` or `f`.
Of course one can use `p(x)` or `f(x)`, but the function returned
can then be passed to (for example) `plot`.

Please note that `p(x)` [or `f(x)`] is already quite efficient.
"""
make_function(p::SimplePolynomial) = p.func

SimplePolynomial() = SimplePolynomial([0])
SimplePolynomial(p::SimplePolynomial) = SimplePolynomial(p.data)

# equality checking

(==)(p::SimplePolynomial, q::SimplePolynomial) = p.data == q.data
(==)(p::SimplePolynomial, a::T) where {T} = p == SimplePolynomial(a)
(==)(a::T, p::SimplePolynomial) where {T} = p == SimplePolynomial(a)

coeffs(p::SimplePolynomial) = copy(p.data)
eltype(p::SimplePolynomial) = eltype(p.data)

"""
`lead(p::SimplePolynomial)` returns the coefficient of the highest
power of `x` in the polynomial `p`. This is nonzero unless `p` is
the zero polynomial.
"""
lead(p::SimplePolynomial) = p.data[end]

# conversion to/from Polynomial type
SimplePolynomial(P::Polynomial) = SimplePolynomial(P.coeffs)
Polynomial(p::SimplePolynomial) = Polynomial(small.(p.data))

"""
`_chomp(list)` creates a new list by removing all trailing zeros
unless that would remove everything, in which case we leave a single zero.
"""
function _chomp(list::Vector{T})::Vector{T} where {T<:Number}
    if length(list) == 0 || all(list .== 0)
        return T[0]
    end
    k = findlast(list .!= 0)
    return list[1:k]
end

"""
`_chomp!(list)` modifies list by removing all trailing zeros, unless
the list is entirely zeros in which case we leave a single 0.
"""
function _chomp!(list::Vector{T}) where {T<:Number}
    while length(list) > 1
        n = length(list)
        if list[n] != 0
            return
        end
        deleteat!(list, n)
    end
end

"""
`degree(p::SimplePolynomial)` returns the degree of `p`; that is,
the highest exponent with a nonzero coefficient. If `p`
is the zero polynomial, then `-1` is returned.
"""
function degree(p::SimplePolynomial)
    n = length(p.data)
    if n > 1 || p.data[1] != 0
        return n - 1
    end
    return -1
end

"""
For a `SimplePolynomial`, `p`, use `p[k]` to read the coefficient
of `x^k`. In particular, `p[0]` is the constant term. If `k` is
larger than the degree of the polynomial, `0` is returned.
"""
function getindex(p::SimplePolynomial, k::Int)
    n = length(p.data)
    if k < 0
        error("index [$k] must be nonnegative")
    end

    if k >= n
        return 0
    end

    return p.data[k+1]
end

# This implements evaluation
function (p::SimplePolynomial)(x)
    try
        return p.func(x)
    catch
    end
    result = 0
    n = degree(p)
    for j = n:-1:0
        result *= x
        result += p[j]
    end
    return result
end

# special case when arg is a matrix
function (p::SimplePolynomial)(x::AbstractMatrix{T}) where {T}
    r, c = size(x)
    if r != c
        error("number of rows $r must equal the number of columns $c")
    end
    S = eltype(p)
    ST = promote_type(S, T)
    id = zeros(ST, r, r)
    for i = 1:r
        id[i, i] = 1
    end
    result = zeros(ST, r, r)
    n = degree(p)
    for j = n:-1:0
        result *= x
        result += p[j] * id
    end
    return result
end

"""
`getx()` is a convenient way to create `SimplePolynomial(0,1)`.
Typical use:
```
x = getx()
p = 2 + x - 3*x^2
```
Use `getx(T)` for coefficients of type `T`, e.g., `getx(Mod{13})`.
"""
function getx(T = Int)
    return SimplePolynomial(T(0), T(1))
end

zero(::Type{SimplePolynomial}) = SimplePolynomial(zero(Int))
one(::Type{SimplePolynomial}) = SimplePolynomial(one(Int))

roots(p::SimplePolynomial) = roots(Polynomial(p))

"""
`derivative(P::SimplePolynomial)` returns the derivative of `P`.
May also be found as `P'`.
"""
function derivative(P::SimplePolynomial)
    if degree(P) < 1
        return SimplePolynomial(0)
    end
    data = [k * P[k] for k = 1:degree(P)]
    return SimplePolynomial(data)
end
adjoint(P::SimplePolynomial) = derivative(P)

"""
`integral(P::SimplePolynomial)` returns the integral of `P`
with constant term equal to zero.
"""
function integral(P::SimplePolynomial)
    if P == 0
        return P
    end
    data = [P[k-1] // k for k = 1:degree(P)+1]
    prepend!(data, 0)
    return SimplePolynomial(data)
end

include("SimpleRationalFunctions.jl")
include("show.jl")
include("arithmetic.jl")
include("rat_arith.jl")
include("gcd.jl")
include("small.jl")
include("rational_roots.jl")
include("newton.jl")

end


n = 0

while n < 10
    global n = n + 1
    println(n)

    for i = 1:5
        println("oi")
    end
end

a = zeros(n,n)

for i = 1:n
  for j = 1:n
    if i==j

      a[i,j] = -4

      elseif i == (j-3) || i == (j+3)

        a[i,j] = 1

      elseif (mod(i,3) != 0 && j == (i+1))

        a[i,j] = 1

      elseif (mod(i,3) != 1 && j == (i-1))

        a[i,j] = 1

      end
  end
end

println("The A matrix is: ")

println(a)

for i = 1:n

  println(a[i,:])

end


#hideall
using Plots, Joe, Random
using Plots.PlotMeasures
Plots.reset_defaults()
gr()
default(
    titlefont  = font("JuliaMono", default(:titlefontsize),  ),
    guidefont  = font("JuliaMono",  default(:guidefontsize),  ),
    tickfont   = font("JuliaMono", default(:tickfontsize),   ),
    legendfont = font("JuliaMono",  default(:legendfontsize), ),
    left_margin = 30px,
    right_margin = 30px,
    bottom_margin = 30px
)
# ### 例64
# Bostonデータセットを決定木で回帰分析します。
# $\alpha$の値は0で、最低要素数だけを指定します。
using Joe
using ScikitLearn
@sk_import datasets: load_boston
df = load_boston()
X =df["data"];y=df["target"];
node = Joe.decisiontree(X,y;n_min=50);
# グラフ化はのちほど検討。
# ### 例65
# 10-fold CVで$\alpha \in [0,1.5]$の範囲での最適化を検討します。
α_seq = 0:0.1:1.5
n = 100
X65 = X[begin:n,:]; y65 = y[begin:n];
s = Int(n/10)
out = zeros(length(α_seq))
for (i,α) in enumerate(α_seq)
    global n, X65, y65, out # hide
    local node # hide
    SS = 0
    for h in 0:9
        test = h*s+1:(h+1)*s
        train = setdiff(1:n,test)
        node = Joe.decisiontree(X65[train,:],y65[train];α=α)
        for t in test
            SS += (y[t] - Joe.dtvalue(X65[t,:],node))^2
        end
    end
    out[i] = SS/n
end
p65_1 = plot(α_seq,out,xlabel="α",ylabel="二乗誤差",
            title="CVで最適なα(N=100)",label=false);
# 同様にn_minも$n_{min} \in \{ 1,2,\dots, 15\}$の範囲で最適化を検討します。

n_min_seq = 1:15
out2 = zeros(length(n_min_seq))
for (i,n_min) in enumerate(n_min_seq)
    global n, X65, y65, out2 # hide
    local node # hide
    SS = 0
    for h in 0:9
        test = h*s+1:(h+1)*s
        train = setdiff(1:n,test)
        node = Joe.decisiontree(X65[train,:],y65[train];n_min=n_min)
        for t in test
            SS += (y[t] - Joe.dtvalue(X65[t,:],node))^2
        end
    end
    out2[i] = SS/n
end
p65_2 = plot(n_min_seq,out2,xlabel="n_min",ylabel="二乗誤差",
            title="CVで最適なn_min(N=100)",label=false);

p65 = plot(p65_1,p65_2,layout=(1,2))
savefig(p65,joinpath(@OUTPUT,"fig7-5.svg")) # hide
# \fig{fig7-5}
# ## 分類の決定木
# ### 例66 Fisherのあやめ（何回目？）
iris = dataset("datasets","iris")
x = iris[!,1:4] |> Matrix
targets=unique(iris.Species)
# あやめの種類をIntに変換します。
label = Dict(target=>i for (i,target) in enumerate(targets))
y = [label[i] for i in iris.Species];


using Plots, SimpleDrawing

export life_run, life_movie

"""
`life_run(A)` runs the game of life starting with the matrix `A`. It is
visualized on the screen. It will run forever unless:
* there are no live cells remaining,
* the current state is the same as the previous state, or
* the current state is the same as the state two steps earlier.
If `life_run` halts, it returns the number of iterations.

Optional named arguments:
* `pause=0.0`: number of extra seconds between frames
* `wrap=false`: determine if the board wraps (is toroidal)
* `counter=false`: show the step number below the image
* `max_steps`: maximum number of steps
"""
function life_run(A::Matrix{Int};
        pause=0.0,
        wrap::Bool=false,
        max_steps::Int=typemax(Int),
        counter::Bool=false
    )
    @assert _life_check(A) "Matrix may only contain 0s and 1s"
    step = 0
    back_1 = 0*A    # previous A
    back_2 = 0*A    # previous previous A

    while true
        p = my_spy(A)
        if counter
            plot!(xlabel=string(step))
        end
        display(p)
        sleep(pause)

        back_2 = back_1
        back_1 = A
        A = life_step(A,wrap)

        if A==back_1 || A==back_2 || sum(A) == 0
            return step
        end
        if step >= max_steps
            return step
        end
        step += 1

    end
end

"""
`life_movie(A)` creates a GIF animation of the game of life.

Optional named arguments:
* `file_name="life.gif"`: name of file to hold the animation
* `rate=5`: frame rate (frames per second)
* `wrap=false`: determine if the board wraps (is toroidal)
* `max_steps`: maximum number of steps; if omitted, runs forever
"""
function life_movie(A::Matrix{Int};
        file_name="life.gif",
        wrap::Bool=false,
        max_steps::Int=typemax(Int),
        rate = 5
    )
    @assert _life_check(A) "Matrix may only contain 0s and 1s"
    step = 0
    back_1 = 0*A    # previous A
    back_2 = 0*A    # previous previous A

    movie = @animate for k=1:max_steps
        print(".")
        p = my_spy(A)
        back_2 = back_1
        back_1 = A
        A = life_step(A,wrap)

        if A==back_1 || A==back_2 || sum(A) == 0
            break
        end
        if step >= max_steps
            break
        end
        step += 1

    end
    println()
    G = gif(movie, file_name, fps = rate)
    return G
end


# unit under test
include("../../src/image/mask.jl")

# test data

array1 = ones(3, 3)
array2 = ones(5, 5)

mask1 = BitArray(
[0 0 0;
 0 1 0;
 0 0 0])

mask1Frisket = BitArray(
 [1 1 1;
  1 0 1;
  1 1 1])

mask2 = BitArray(
[0 0 0 0 0;
 0 0 0 0 0;
 0 0 1 0 0;
 0 0 0 0 0;
 0 0 0 0 0])

mask2Frisket = BitArray(
[0 0 0 0 0;
 0 1 1 1 0;
 0 1 0 1 0;
 0 1 1 1 0;
 0 0 0 0 0])

 mask2Frisket2 = BitArray(
 [1 1 1 1 1;
  1 1 1 1 1;
  1 1 0 1 1;
  1 1 1 1 1;
  1 1 1 1 1])

result = frisketMaskAroundMask(array1, mask1, 1)
@assert result == mask1Frisket

result = frisketMaskAroundMask(array2, mask2, 1)
@assert result == mask2Frisket

result = frisketMaskAroundMask(array2, mask2, 2)
@assert result == mask2Frisket2

# frisketDepth out of bounds of array
result = frisketMaskAroundMask(array2, mask2, 3)
@assert result == mask2Frisket2


using DelimitedFiles

function ReadData(fname::AbstractString)
    dataStr = readdlm(fname, String)

    data = zeros(Int, length(dataStr), length(dataStr[1]))
    for i = 1:length(dataStr)
        data[i,:] = map(x -> parse(Int, x), split(dataStr[i],""))
    end

    return data
end

function StepData!(data::Matrix)
    # need to do slicing to modify data in place
    data[:,:] += ones(Int, size(data)...)

    flashes = findall(x -> x == 10, data)
    while ! isempty(flashes)
        newFlashes = Array{CartesianIndex{2},1}(undef, 0)

        for flash in flashes
            x, y = flash[1], flash[2]
            xMin, xMax = x == 1 ? x : x-1, x == size(data)[1] ? x : x+1
            yMin, yMax = y == 1 ? y : y-1, y == size(data)[2] ? y : y+1
            data[xMin:xMax, yMin:yMax] += ones(Int, xMax-xMin+1, yMax-yMin+1)

            for i=xMin:xMax, j=yMin:yMax
                if data[i,j] == 10
                    push!(newFlashes, CartesianIndex(i,j))
                end
            end
        end

        flashes = newFlashes
    end

    map!(x -> x>=10 ? 0 : x, data, data)
end

function Part1(data::Matrix)
    nSteps = 100
    flashes = 0

    for i = 1:nSteps
        StepData!(data)
        flashes += count(iszero, data)
    end

    return flashes
end

function Part2(data::Matrix)
    steps = 0
    while ! all(iszero, data)
        StepData!(data)
        steps += 1
    end

    return steps
end

#datafile = "../data/test11.txt"
datafile = "../data/day11.txt"

data = ReadData(datafile)
println(Part1(data))

# need to re-read data since Part1 modifies it
data = ReadData(datafile)
println(Part2(data))


using Agents, Random
using GLMakie
using InteractiveDynamics

model, agent_step!, model_step! = Models.social_distancing(isolated = 0.8)

# Various possibilities for colors/sizes/markers
sir_colors(a) = a.status == :S ? "#2b2b33" : a.status == :I ? "#bf2642" : "#338c54"
# sir_sizes(a) = 10*randn()
# sir_sizes(a) = 5*(mod1(a.id, 3)+1)
sir_sizes = 10

# sir_shape(a) = rand(('🐑', '🐺', '🌳'))
# sir_shape(a) = rand(('😹', '🐺', '🌳'))
# sir_shape(a) = rand(('π', '😹', '⚃', '◑', '▼'))
# sir_shape(a) = rand((:diamond, :circle))
sir_shape(a) = a.status == :S ? :circle : a.status == :I ? :diamond : :rect

# function sir_shape(b)
#     φ = atan(b.vel[2], b.vel[1])
#     xs = [(i ∈ (0, 3) ? 2 : 1)*cos(i*2π/3 + φ) for i in 0:3]
#     ys = [(i ∈ (0, 3) ? 2 : 1)*sin(i*2π/3 + φ) for i in 0:3]
#     poly(xs, ys)
# end

fig, abmstepper = abm_plot(model;
    ac = sir_colors, as = sir_sizes, am = sir_shape,
)

display(fig)

# %% same, but interactive
abm_play(model, agent_step!, model_step!;
    ac = sir_colors, as = sir_sizes, am = sir_shape,
)

# %% same, but video
abm_video("socialdist.mp4", model, agent_step!, model_step!;
    ac = sir_colors, as = sir_sizes, am = sir_shape,
    spf = 2, frames = 100,
)

# %% Interactive data plot
model, agent_step!, model_step! = Models.social_distancing(isolated = 0.8)

infected(x) = count(i == :I for i in x)
recovered(x) = count(i == :R for i in x)
adata = [(:status, infected), (:status, recovered)]
alabels = ["I", "R"]
mdata = [nagents]
mlabels = ["N"]

params = Dict(
    :death_rate => 0.02:0.001:1.0,
    :reinfection_probability => 0:0.01:1.0,
    :infection_period => 24:24:2400,
)

when = (model, s) -> s % 50 == 0

fig, adf, mdf = abm_data_exploration(model, agent_step!, model_step!, params;
ac = sir_colors, as = sir_sizes, am = sir_shape,
when = when, mdata = mdata, adata=adata, alabels=alabels, mlabels=mlabels)


using RobustAndOptimalControl, ControlSystems

G = ssrand(1,1,3, proper=true)
Gd = add_low_frequency_disturbance(G)
@test Gd.nx == 4
@test rank(obsv(Gd)) == 4
@test rank(ctrb(Gd)) == 3
@test any(isapprox(0, atol=eps()), pole(Gd))

Gd = add_low_frequency_disturbance(G, measurement=true)
Gd.C[end] == 1
@test Gd.nx == 4
@test rank(obsv(Gd)) == 4
@test rank(ctrb(Gd)) == 3
@test any(isapprox(0, atol=eps()), pole(Gd))

G = ssrand(1,1,3, proper=true)
Gd = add_resonant_disturbance(G, 1, 0, 3)
@test Gd.nx == 5
@test rank(obsv(Gd)) == 5
@test rank(ctrb(Gd)) == 3
@test any(isapprox(1, atol=eps()), imag.(pole(Gd)))

Gd = add_resonant_disturbance(G, 1, 0, 1, measurement=true)
@test Gd.nx == 5
@test rank(obsv(Gd)) == 5
@test rank(ctrb(Gd)) == 3
@test any(isapprox(1, atol=eps()), imag.(pole(Gd)))

##

G = ssrand(1,1,3, proper=true, Ts=1)
GD = ssrand(1,1,3, proper=false, Ts=1)

## Diff
Gd = add_output_differentiator(G)
Gd2 = [tf(1,1); tf([1, -1], [1], 1)]*tf(G)
@test tf(Gd) ≈ Gd2
# @test hinfnorm(Gd-Gd2)[1] < 1e-10 hinfnorm not robust

## Int
Gd = add_output_integrator(G)
Gd2 = [tf(1,1); tf(1, [1, -1], 1)]*G
@test Gd ≈ Gd2
@test sminreal(Gd[1,1]) == G # Exact equivalence should hold here
@test Gd.nx == 4 # To guard agains changes in realization of tf as ss

Gd = add_input_integrator(G)
@test sminreal(Gd[1,1]) == G # Exact equivalence should hold here
@test Gd.nx == 4 # To guard agains changes in realization of tf as ss
@test tf(sminreal(Gd[2,1])) == tf(1, [1,-1], 1)

G = ssrand(2,1,3, proper=true, Ts=1)
Gd = add_input_integrator(G)
@test sminreal(Gd[1:2,1]) == G # Exact equivalence should hold here
@test Gd.nx == 4 # To guard agains changes in realization of tf as ss
@test tf(sminreal(Gd[3,1])) == tf(1, [1,-1], 1)

G = ssrand(1,2,3, proper=true, Ts=1)
Gd = add_input_integrator(G)
@test sminreal(Gd[1,1:2]) == G # Exact equivalence should hold here
@test Gd.nx == 4 # To guard agains changes in realization of tf as ss
@test tf(sminreal(Gd[2,1])) == tf(1, [1,-1], 1)
@test tf(sminreal(Gd[2,2])) == tf(0, 1)

Gd = add_input_integrator(G, 2) # other input
@test sminreal(Gd[1,1:2]) == G # Exact equivalence should hold here
@test Gd.nx == 4 # To guard agains changes in realization of tf as ss
@test tf(sminreal(Gd[2,2])) == tf(1, [1,-1], 1)
@test tf(sminreal(Gd[2,1])) == tf(0, 1)

# Both
Gd = add_output_differentiator(G)
Gd = add_output_integrator(Gd, 1)
Gd2 = [tf(1,1); tf([1, -1], [1], 1); tf(1, [1, -1], 1)]*tf(G)
w = exp10.(LinRange(-2, 2, 100))

# These are harder to compare, the freqresp method seems most reliable
f1 = freqresp(Gd, w)
f2 = freqresp(Gd2, w)
@test f1 ≈ f2
# bodeplot([Gd, Gd2])

# Input diff
nx,nu,ny = G.nx, G.nu, G.ny
Gd = add_input_differentiator(G, 1:G.nu)
@test Gd.A[nx+1:end, nx+1:end] == 0I
@test Gd.B[nx+1:end, :] == I
@test Gd.C[ny+1:end, nx+1:end] == -I
@test Gd.D[ny+1:end, :] == I


## test_circulant_embedding.jl : test circulant embedding for GRF generation

@testset "circulant embedding      " begin

    ## 1d Exponential ##
    cov = CovarianceFunction(1,Exponential(0.1))
    pts = 0:1//1001:1
    grf = GaussianRandomField(cov,CirculantEmbedding(),pts)
    @test isa(grf,GaussianRandomField)
    @test isa(grf.cov,CovarianceFunction)
    @test isa(grf.cov.cov,Exponential)
    @test ndims(grf.cov) == 1
    @test isa(grf,GaussianRandomField{CirculantEmbedding})
    @test length(grf.pts[1]) == length(pts)
    @test length(sample(grf)) == length(pts)

    ## 1d Matern with padding ##
    m = CovarianceFunction(1,Matern(1.0,1.0))
    grf = GaussianRandomField(m,CirculantEmbedding(),pts,minpadding=8000)
    @test isa(grf,GaussianRandomField)
    @test isa(grf.cov,CovarianceFunction)
    @test isa(grf.cov.cov,Matern)
    @test ndims(grf.cov) == 1
    @test isa(grf,GaussianRandomField{CirculantEmbedding})
    @test length(grf.pts[1]) == length(pts)
    @test length(sample(grf)) == length(pts)

    ## 2d Anisotropic ##
    A = [1000 0; 0 1000]
    a = AnisotropicExponential(A)
    c = CovarianceFunction(2,a)
    pts1 = range(0,stop = 1,length = 128)
    pts2 = range(0,stop = 1,length = 256)
    @suppress grf = GaussianRandomField(c,CirculantEmbedding(),pts1,pts2,minpadding=(128,256))
    @test isa(grf,GaussianRandomField)
    @test isa(grf.cov,CovarianceFunction)
    @test isa(grf.cov.cov,AnisotropicExponential)
    @test ndims(grf.cov) == 2
    @test isa(grf,GaussianRandomField{CirculantEmbedding})
    @test length(grf.pts) == 2
    @test length(grf.pts[1]) == length(pts1)
    @test length(grf.pts[2]) == length(pts2)
    @test size(sample(grf),1) == length(pts1)
    @test size(sample(grf),2) == length(pts2)

    ## 3d AnisotropicExponential ##
    A = [100 0 0; 0 100 0; 0 0 100]
    a = AnisotropicExponential(A)
    c = CovarianceFunction(3,a)
    pts = 0:0.025:1
    @suppress grf = GaussianRandomField(c,CirculantEmbedding(),pts,pts,pts,measure=false)
    @test isa(grf,GaussianRandomField)
    @test isa(grf.cov,CovarianceFunction)
    @test isa(grf.cov.cov,AnisotropicExponential)
    @test ndims(grf.cov) == 3
    @test isa(grf,GaussianRandomField{CirculantEmbedding})
    @test length(grf.pts) == 3
    @test length(grf.pts[1]) == length(pts)
    @test length(grf.pts[2]) == length(pts)
    @test length(grf.pts[3]) == length(pts)
    @test size(sample(grf),1) == length(pts)
    @test size(sample(grf),2) == length(pts)
    @test size(sample(grf),3) == length(pts)

    ## Larger domain ##
    A = [1 0.8; 0.8 1]
    m = AnisotropicExponential(A)
    c = CovarianceFunction(2,m)
    pts1 = range(-5,stop = 5,length = 128)
    pts2 = range(10,stop = 0,length = 128)
    @suppress grf = GaussianRandomField(c,CirculantEmbedding(),pts1,pts2,minpadding=128,primes=false)
    @test isa(grf,GaussianRandomField)
    @test isa(grf.cov,CovarianceFunction)
    @test isa(grf.cov.cov,AnisotropicExponential)
    @test ndims(grf.cov) == 2
    @test isa(grf,GaussianRandomField{CirculantEmbedding})
    @test length(grf.pts) == 2
    @test length(grf.pts[1]) == length(pts1)
    @test length(grf.pts[2]) == length(pts2)
    @test size(sample(grf),1) == length(pts1)
    @test size(sample(grf),2) == length(pts2)

    ## Dedicated plotting commands ##
    m = Matern(0.5,2)
    c = CovarianceFunction(3,m)
    pts1 = range(0,stop = 1,length = 16)
    pts2 = range(0,stop = 1,length = 8)
    pts3 = range(0,stop = 1,length = 4)
    #plot_covariance_matrix(c,pts1,pts2,pts3); close()
    m = Matern(0.1,1)
    c = CovarianceFunction(2,m)
    pts = range(0,stop = 1,length = 512)
    g = GaussianRandomField(c,CirculantEmbedding(),pts,pts,minpadding=512,measure=false)
    #contourf(g); close()
    #plot_eigenvalues(g); close()

end


export minFrobNormspg

function minFrobNormspg(L::Tx, b, p, eta, indi, Lprev, weight, iter, fac)where {ETx<:Number, Tx<:AbstractMatrix{ETx}}
    ###### add for weighted part
    if weight!=1.0
        L*= weight
        L += (1/weight- 1) * Lprev * ( Lprev'*L)
    end

    ###### Matrix Dimension Setup
    undefine_1,m = size(b);
    undefine_2,r = size(L);
    R1 = Array{Any}(undef,m,r);

    ###### Solve for each row of R independently
        R1 = @distributed vcat for k = 1:length(indi)

            minnorm(L,b,k,eta, r, iter);

        end

        R = conj(permutedims(reshape(R1,r,m),[2,1]));

    return R
end


"""
    ht_p2_p1_veclap

The manufactured-solution colliding flow example from Elman et al 2014. The
Hood-Taylor formulation with quadratic triangles for the velocity and continuous
pressure on linear triangles.

The formulation is the one derived in Elman, et al., Finite elements and fast
iterative solvers, p. 132. In other words, it is the vector Laplacian version.
"""
module ht_p2_p1_veclap

using LinearAlgebra
using StaticArrays
using MeshCore: nrelations, nentities, ir_identity, attribute, VecAttrib
using MeshSteward: T6block, T6toT3
using MeshSteward: Mesh, attach!, baseincrel, boundary
using MeshSteward: vselect, geometry, summary, transform
using MeshSteward: vtkwrite
using Elfel.RefShapes: manifdim, manifdimv
using Elfel.FElements: FEH1_T6, FEH1_T3, refshape, Jacobian
using Elfel.FESpaces: FESpace, ndofs, setebc!, nunknowns, doftype
using Elfel.FESpaces: numberfreedofs!, numberdatadofs!, numberdofs!
using Elfel.FESpaces: scattersysvec!, makeattribute, gathersysvec!, edofcompnt
using Elfel.FESpaces: highestfreedofnum, highestdatadofnum
using Elfel.FEIterators: FEIterator, ndofsperel, elnodes, eldofs, eldofvals
using Elfel.FEIterators: jacjac, location
using Elfel.QPIterators: QPIterator, bfun, bfungrad, weight
using Elfel.Assemblers: SysmatAssemblerSparse, start!, finish!, assemble!
using Elfel.Assemblers: SysvecAssembler
using Elfel.LocalAssemblers: LocalMatrixAssembler, LocalVectorAssembler, init!
using UnicodePlots

mu = 1.0 # dynamic viscosity
A = 1.0 # half of the length of the side of the square
trueux = (x, y) -> 20 * x * y ^ 3
trueuy = (x, y) -> 5 * x ^ 4 - 5 * y ^ 4
truep = (x, y) -> 60 * x ^ 2 * y - 20 * y ^ 3

function genmesh(N)
    # Hood-Taylor pair of meshes is needed
    # This mesh will be for the velocities
    vmesh = Mesh()
    attach!(vmesh, T6block(2 * A, 2 * A, N, N), "velocity")
    ir = baseincrel(vmesh)
    transform(ir, x -> x .- A)
    # This mesh will be used for the pressures. Notice that it needs to be
    # "compatible" with the velocity mesh in the sense that they need to share
    # the nodes at the corners of the triangles.
    pmesh = Mesh()
    attach!(pmesh, T6toT3(baseincrel(vmesh, "velocity")), "pressure")
    return vmesh, pmesh
end

function assembleK(Uxh, Uyh, Ph, tndof, mu)
    function integrate!(ass, elits, qpits, mu)
        uxnedof, uynedof, pnedof = ndofsperel.(elits)
        kuxux = LocalMatrixAssembler(uxnedof, uxnedof, 0.0)
        kuyuy = LocalMatrixAssembler(uynedof, uynedof, 0.0)
        kuxp = LocalMatrixAssembler(uxnedof, pnedof, 0.0)
        kuyp = LocalMatrixAssembler(uynedof, pnedof, 0.0)
        for el in zip(elits...)
            uxel, uyel, pel = el
            init!(kuxux, eldofs(uxel), eldofs(uxel))
            init!(kuyuy, eldofs(uyel), eldofs(uyel))
            init!(kuxp, eldofs(uxel), eldofs(pel))
            init!(kuyp, eldofs(uyel), eldofs(pel))
            for qp in zip(qpits...)
                uxqp, uyqp, pqp = qp
                Jac, J = jacjac(pel, pqp)
                JxW = J * weight(pqp)
                gradNp = bfungrad(pqp, Jac)
                gradNux = bfungrad(uxqp, Jac)
                gradNuy = bfungrad(uyqp, Jac)
                Np = bfun(pqp)
                for j in 1:uxnedof, i in 1:uxnedof
                    kuxux[i, j] += (mu * JxW) * dot(gradNux[i], gradNux[j])
                end
                for j in 1:uynedof, i in 1:uynedof
                    kuyuy[i, j] += (mu * JxW) * dot(gradNuy[i], gradNuy[j])
                end
                for j in 1:pnedof, i in 1:uxnedof
                    kuxp[i, j] += (-JxW) * (gradNux[i][1] * Np[j])
                end
                for j in 1:pnedof, i in 1:uynedof
                    kuyp[i, j] += (-JxW) * (gradNuy[i][2] * Np[j])
                end
            end
            assemble!(ass, kuxux)
            assemble!(ass, kuyuy)
            assemble!(ass, kuxp)
            assemble!(ass, transpose(kuxp))
            assemble!(ass, kuyp)
            assemble!(ass, transpose(kuyp))
        end
        return ass
    end

    elits = (FEIterator(Uxh), FEIterator(Uyh), FEIterator(Ph))
    qargs = (kind = :default, npts = 3,)
    qpits = (QPIterator(Uxh, qargs), QPIterator(Uyh, qargs), QPIterator(Ph, qargs))
    ass = SysmatAssemblerSparse(0.0)
    start!(ass, tndof, tndof)
    @time integrate!(ass, elits, qpits, mu)
    return finish!(ass)
end

function solve!(U, K, F, nu)
    KT = K * U
    U[1:nu] = K[1:nu, 1:nu] \ (F[1:nu] - KT[1:nu])
end

function evaluate_pressure_error(Ph)
    function integrate!(elit, qpit, truep)
        pnedof = ndofsperel(elit)
        E = 0.0
        for el in elit
            dofvals = eldofvals(el)
            for qp in qpit
                Jac, J = jacjac(el, qp)
                JxW = J * weight(qp)
                Np = bfun(qp)
                pt = truep(location(el, qp)...)
                pa = 0.0
                for j in 1:pnedof
                    pa += (dofvals[j] * Np[j])
                end
                E += (JxW) * (pa - pt)^2
            end
        end
        return sqrt(E)
    end

    elit = FEIterator(Ph)
    qargs = (kind = :default, npts = 3,)
    qpit = QPIterator(Ph, qargs)
    return integrate!(elit, qpit, truep)
end

function evaluate_velocity_error(Uxh, Uyh)
    function integrate!(elits, qpits, trueux, trueuy)
        uxnedof, uynedof = ndofsperel.(elits)
        E = 0.0
        for el in zip(elits...)
            uxel, uyel = el
            uxdofvals = eldofvals(uxel)
            uydofvals = eldofvals(uyel)
            for qp in zip(qpits...)
                uxqp, uyqp = qp
                Jac, J = jacjac(uxel, uxqp)
                JxW = J * weight(uxqp)
                Np = bfun(uxqp)
                uxt = trueux(location(uxel, uxqp)...)
                uyt = trueuy(location(uyel, uyqp)...)
                uxa = 0.0
                uya = 0.0
                for j in 1:uxnedof
                    uxa += (uxdofvals[j] * Np[j])
                    uya += (uydofvals[j] * Np[j])
                end
                E += (JxW) * ((uxa - uxt)^2 + (uya - uyt)^2)
            end
        end
        return sqrt(E)
    end

    elits = (FEIterator(Uxh), FEIterator(Uyh),)
    qargs = (kind = :default, npts = 3,)
    qpits = (QPIterator(Uxh, qargs), QPIterator(Uyh, qargs),)
    return integrate!(elits, qpits, trueux, trueuy)
end

function run(N)
    vmesh, pmesh = genmesh(N)
    # Velocity spaces
    Uxh = FESpace(Float64, vmesh, FEH1_T6(), 1)
    Uyh = FESpace(Float64, vmesh, FEH1_T6(), 1)
    locs = geometry(vmesh)
    inflate = A / N / 100
    # The entire boundary
    boxes = [[-A A -A -A], [-A -A -A A], [A A -A A], [-A A A A]]
    for box in boxes
        vl = vselect(locs; box = box, inflate = inflate)
        for i in vl
            setebc!(Uxh, 0, i, 1, trueux(locs[i]...))
            setebc!(Uyh, 0, i, 1, trueuy(locs[i]...))
        end
    end
    # Pressure space
    Ph = FESpace(Float64, pmesh, FEH1_T3(), 1)
    atcenter = vselect(geometry(pmesh); nearestto = [0.0, 0.0])
    setebc!(Ph, 0, atcenter[1], 1, 0.0)
    # Number the degrees of freedom
    numberdofs!([Uxh, Uyh, Ph])
    tndof = ndofs(Uxh) + ndofs(Uyh) + ndofs(Ph)
    tnunk = nunknowns(Uxh) + nunknowns(Uyh) + nunknowns(Ph)
    # Assemble the coefficient matrix
    K = assembleK(Uxh, Uyh, Ph, tndof, mu)
    # p = spy(K, canvas = DotCanvas)
    # display(p)
    # Solve the system
    U = fill(0.0, tndof)
    gathersysvec!(U, [Uxh, Uyh, Ph])
    F = fill(0.0, tndof)
    solve!(U, K, F, tnunk)
    scattersysvec!([Uxh, Uyh, Ph], U)
    # Postprocessing
    makeattribute(Ph, "p", 1)
    makeattribute(Uxh, "ux", 1)
    makeattribute(Uyh, "uy", 1)
    ep = evaluate_pressure_error(Ph)
    ev = evaluate_velocity_error(Uxh, Uyh)
    vtkwrite("ht_p2_p1_veclap-p", baseincrel(pmesh), [(name = "p",), ])
    vtkwrite("ht_p2_p1_veclap-v", baseincrel(vmesh), [(name = "ux",), (name = "uy",)])
    geom = geometry(Ph.mesh)
    ir = baseincrel(Ph.mesh)
    pt = VecAttrib([truep(geom[i]...) for i in 1:length(geom)])
    ir.right.attributes["pt"] = pt
    vtkwrite("ht_p2_p1_veclap-pt", baseincrel(pmesh), [(name = "pt",), ])
    return (ep, ev)
end

end

using .ht_p2_p1_veclap
let
    N = 2
    for loop in 1:5
        ep, ev = ht_p2_p1_veclap.run(N)
        @show N, ep, ev
        N = N * 2
    end
end


module Day6
export parse_input, part_1, part_2

function parse_input(input::String)
    lines = split(strip(input), "\n")
    coords = map(
        l -> map(x -> parse(Int, x), split(l, ", ")),
        lines
    )
    return hcat(coords...)'
end

function manhattan(x1, y1, x2, y2)
    return abs(x2 - x1) + abs(y2 - y1)
end

function part_1(coords)
    xmin, ymin = minimum(coords, dims=1)
    xmax, ymax = maximum(coords, dims=1)

    n_x = xmax - xmin
    n_y = ymax - ymin

    grid = zeros(Int, (n_x, n_y))

    for i in 1:n_x
        for j in 1:n_y
            x1 = i + xmin - 1
            y1 = j + ymin - 1

            min_distance = 2^31
            for k in 1:size(coords, 1)
                x2, y2 = coords[k, :]
                dist = manhattan(x1, y1, x2, y2)
                if dist < min_distance
                    grid[i, j] = k
                    min_distance = dist
                elseif dist == min_distance
                    grid[i, j] == -1
                end
            end
        end
    end

    areas = zeros(Int, size(coords, 1))
    for i in 1:n_x
        for j in 1:n_y

            if (grid[i, j] != -1) & (areas[grid[i, j]] != -1)
                if (i == 1) | (i == n_x) | (j == 1) | (j == n_y)
                    areas[grid[i, j]] = -1
                else
                    areas[grid[i, j]] += 1
                end
            end
        end
    end
    return maximum(areas)
end

function part_2(coords)
    max_total_distance = 10000
    xmin, ymin = minimum(coords, dims=1)
    xmax, ymax = maximum(coords, dims=1)

    n_x = xmax - xmin
    n_y = ymax - ymin

    grid = zeros(Int, (n_x, n_y))

    safe_size = 0

    for i in 1:n_x
        for j in 1:n_y
            x1 = i + xmin - 1
            y1 = j + ymin - 1

            total_distance = 0
            for k in 1:size(coords, 1)
                x2, y2 = coords[k, :]
                dist = manhattan(x1, y1, x2, y2)
                total_distance += dist
            end
            if total_distance < max_total_distance
                safe_size += 1
            end
        end
    end
    return safe_size
end

end


mutable struct RhoVsFixedthreshold <: AbstractModel
    # params::Vector{Float64}
    params::NamedTuple{(:threshold,:patch_aplha,:beta,:reset,:bias), NTuple{5,Float64}}
end

function init(::Type{WaltonForaging.RhoVsFixedthreshold})
    return RhoVsFixedthreshold((
        threshold = rand(0.0:0.0001:0.1), #=alpha_environment = [0,0.1] learning rate of environment=#
        patch_aplha = rand(0.0:0.0001:0.2), #=alpha_patch = [0, 0.2] learning rate of the patch=#
        beta = rand(2:0.0001:10), #=beta = [2, 10] noise parameter or inverse temperature=#
        reset = rand(0.0:0.0001:0.5), #=reset = [0, 0.5] patch env rew rate at the entrance of a trials, because it has decayed in the last trial=#
        bias = 0.1 #=bias = [2, 10] bias towards staying in the softmax=#
        ))
end

Distributions.params(m::RhoVsFixedthreshold) = collect(m.params)

function updateparams!(m::RhoVsFixedthreshold,v::Vector{Float64})
    (length(v) == 5) || error("Incorrect number of params for $(typeof(m))")
    threshold,patch_aplha,beta,reset,bias = tuple(v...)
    m.params = (;threshold,patch_aplha,beta,reset,bias)
end

function Pstay(m::RhoVsFixedthreshold, patch_rew_rate::Float64)
    threshold, patch_alpha, beta, reset_val, bias = tuple(params(m)...)
    return exp(bias + beta*patch_rew_rate)/(exp(bias + beta*patch_rew_rate) + exp(beta*threshold))
end

function Pleave(m::RhoVsFixedthreshold, patch_rew_rate::Float64)
    threshold, patch_alpha, beta, reset_val, bias = tuple(params(m)...)
    return exp(beta*threshold)/(exp(beta*threshold) + exp(bias + beta*patch_rew_rate))
end

function Poutcome(m::RhoVsFixedthreshold, threshold::Float64, patch_rew_rate::Float64, leave::Int64)
    if leave == 1
        Pleave(m, threshold, patch_rew_rate)
    elseif leave == 0
        Pstay(m, threshold, patch_rew_rate)
    else
        error("Wrong argument for Poutcome calculation: $leave")
    end
end

function Likelihood(m::RhoVsFixedthreshold,df)
    # if I start using subdf for patches conditions I will need:
    # an initial env reward rate before splitting per session (maybe a separate variable to feed to loglikelihood)
    # how do I minimize the same patch reset value for the whole session?

    threshold, patch_alpha, beta, reset_val, bias = tuple(params(m)...)
    final_env_rho = sum(df.REWARD/df[end,:TIME])
    transform!(df, [:REWARD,:BIN,:NEWTRIAL] => ((r,b,n) -> patch_get_rew_rate(r, b, reset_val, n, patch_alpha)) => :PatchRewRate)
    transform!(df, [:PatchRewRate,:LEAVE] => ByRow((p,l) -> Poutcome(m,p,l)) => :Probability)
    return filter(r -> r.SIDE == "travel", df).Probability
end


"""
"""
abstract type AbstractGF end

"""
"""
struct TimeOrderedGF <: AbstractGF
    L   # Lesser
    G   # Greater
    R   # Retarded
    ts  # time-grid

    TimeOrderedGF(L, G, ts) = new(L, G, LowerTriangular(G - L), ts)
end

Base.:*(a::Number, g::TimeOrderedGF) = TimeOrderedGF(a * lesser(g), a * greater(g), g.ts)
Base.:+(g1::TimeOrderedGF, g2::TimeOrderedGF) = TimeOrderedGF(lesser(g1) + lesser(g2), greater(g1) + greater(g2), g1.ts)

"""
"""
struct TimeOrderedConvolution <: AbstractGF
    A::AbstractGF # Left GF
    B::AbstractGF # Right GF
    dts::AbstractMatrix # Trapezoidal rule is achieved by a multiplication with this matrix
end

"""
"""
function conv(A::TimeOrderedGF, B::TimeOrderedGF;
    dts = reduce(hcat, ([calculate_weights(A.ts[1:i], ones(Int64, i-1)); zeros(length(A.ts)-i)] for i in eachindex(A.ts)))
    )
    c = TimeOrderedConvolution(A, B, dts)
    return TimeOrderedGF(lesser(c), greater(c), B.ts) # no need for lazy `TimeOrderedConvolution`
end
const ⋆ = conv # left-associative operator


"""
    solve(
        equation::AdvectionEquation{PeriodicIntervalDomain{T},F},
        u,
        tlist,
        M,
        N;
        method = "filterfirst",
        subspacedim = 10,
        solver = QNDF(),
        reltol = 1e-4,
        abstol = 1e-6,
        C̄ = nothing,
    ) where {T,F}

Solve `equation` from `tlist[1]` to `tlist[2]` with initial conditions `u` and a
discretization of `M` and `N` points. If `method` is `"filterfirst"`, the equation is filtered then
discretized. If `method` is `"discretizefirst"`, the equation is discretized then filtered.
The parameter `subspacedim` controls the accuracy of the linear exponential timestepping.
"""
function solve(
    equation::AdvectionEquation{PeriodicIntervalDomain{T},F},
    u,
    tlist,
    M,
    N;
    method = "filterfirst",
    subspacedim = 10,
    solver = QNDF(),
    reltol = 1e-4,
    abstol = 1e-6,
    degmax = 10,
    λ = 1e-6,
    λ_ridge = 1e-8,
    C̄ = nothing,
) where {T,F}
    @unpack domain, filter = equation

    x = discretize(domain, M)
    ξ = discretize(domain, N)

    ū = apply_filter(u, filter, domain)
    if filter isa TopHatFilter
        h = filter.width
        dh(x) = derivative(h, x)
        α(x) = 1 / 3 * dh(x) * h(x)
    end

    # Get matrices
    C = advection_matrix(domain, N)
    D = diffusion_matrix(domain, N)

    if method == "filterfirst"
        F === TopHatFilter ||
            error("Method \"filterfirst\" is only implemented for TopHatFilter")
        ūₕ = ū.(ξ)
        A = spdiagm(α.(ξ))
        J = DiffEqArrayOperator(-C + A * D)
        p = (; J)
        Mdu_taylor!(du, u, p, t) = mul!(du, p.J, u)
        odefunction =
            ODEFunction(Mdu_taylor!, jac = (J, u, p, t) -> (J .= p.J), jac_prototype = p.J)
        problem = ODEProblem(odefunction, ūₕ, tlist, p)
        solution = OrdinaryDiffEq.solve(problem, solver; reltol, abstol)
    elseif method == "filterfirst_exp"
        F === TopHatFilter ||
            error("Method \"filterfirst\" is only implemented for TopHatFilter")
        # M == N || @warn "Method \"filterfirst\" will only use M"
        A = spdiagm(α.(ξ))
        J = DiffEqArrayOperator(-C + A * D)
        ūₕ = ū.(ξ)
        problem = ODEProblem(J, ūₕ, tlist)
        solution = OrdinaryDiffEq.solve(
            problem,
            LinearExponential(krylov = :simple, m = subspacedim),
        )
    elseif method == "discretizefirst"
        # W = filter_matrix(filter, domain, M, N; degmax, λ)
        # R = reconstruction_matrix(filter, domain, M, N; degmax, λ)
        W, R = get_W_R(filter, domain, M, N; degmax, λ)
        ūₕ = ū.(x)
        # ūₕ = W * u.(ξ)
        p = (; J = -W * C * R)
        Mdu!(du, u, p, t) = mul!(du, p.J, u)
        odefunction = ODEFunction(
            Mdu!,
            jac = (J, u, p, t) -> (J .= p.J),
            jac_prototype = p.J,
            # mass_matrix = W * R,
        )
        problem = ODEProblem(odefunction, ūₕ, tlist, p)
        solution = OrdinaryDiffEq.solve(problem, solver; reltol, abstol)
    elseif method == "discretizefirst-without-R"
        W = filter_matrix(filter, domain, M, N; degmax, λ = λ_ridge)
        # println(W[M ÷ 2, :])
        # A = factorize(W'W + λ * I)
        A = lu(W'W + λ * I)
        J = -W * C
        ūₕ = ū.(x)
        # ūₕ = W * u.(ξ)
        p = (; u☆ = zeros(N), utmp = zeros(N), J, A)
        function Mdu!_without_R(dū, ū, p, t)
            @unpack u☆, utmp, J, A = p
            mul!(utmp, W', ū)
            ldiv!(u☆, A, utmp)
            mul!(dū, J, u☆)
        end
        odefunction = ODEFunction(
            Mdu!_without_R,
            # jac = (J, u, p, t) -> (J .= p.J),
            # jac_prototype = p.J,
            # mass_matrix = W * R,
        )
        problem = ODEProblem(odefunction, ūₕ, tlist, p)
        solution = OrdinaryDiffEq.solve(problem, solver; reltol, abstol)
    elseif method == "discretizefirst-fit-Cbar"
        ūₕ = ū.(x)
        p = (; J = -C̄)
        Mdu!_C̄(du, u, p, t) = mul!(du, p.J, u)
        odefunction = ODEFunction(
            Mdu!_C̄,
            jac = (J, u, p, t) -> (J .= p.J),
            jac_prototype = p.J,
        )
        problem = ODEProblem(odefunction, ūₕ, tlist, p)
        solution = OrdinaryDiffEq.solve(problem, solver; reltol, abstol)
    else
        error("Unknown method")
    end

    solution
end


module HH

import BP

#in-place builder of the q x q HH matrix in mom. sp.
momsphhmat!(A::Matrix{Complex{Float64}}, kx0::Float64,ky::Float64) =
   momsphhmat!(A, kx0,ky, 1)

function momsphhmat!(A::Matrix{Complex{Float64}},
                     kx0::Float64,ky::Float64,
                     p::Int)
    q::Int = size(A,1)
    A[1,q] = -exp(-im*q*ky)
    A[q,1] = -exp(im*q*ky)
    #upper subdiagonal
    for i in diagind(A,1)
        A[i] = -one(Complex{Float64})
    end
    #lower subdiagonal
    for i in diagind(A,-1)
        A[i] = -one(Complex{Float64})
    end
    #main diagonal
    α = p/q
    for (j,i) in enumerate(diagind(A))
        A[i] = -2*cos(kx0 + 2*π*α*j)
    end
    nothing
end

# Computes the q energy levels E(p_y)
hhladder(q::Int) = hhladder(1,q)
function hhladder(p::Int, q::Int)
    ky = linspace(-π, π, 100)
    kx₀= linspace(-π/q, π/q, 100)
    E = Array(Float64, (length(ky),length(kx₀),q))
    H = zeros(Complex{Float64}, (q,q))
    for col=1:length(kx₀), row=1:length(ky)
        momsphhmat!(H, kx₀[col], ky[row], p)
        E[row,col,:] = eigvals(H)
    end
    return E
end
hhladder!(E::Array{Float64,3}) = hhladder!(E, 1)
function hhladder!(E::Array{Float64, 3}, p::Int)
    ky = linspace(-π, π, size(E,1))
    q = size(E,3)
    kx₀= linspace(-π/q, π/q, size(E,2))
    H = zeros(Complex{Float64}, (q,q))
    for col=1:length(kx₀), row=1:length(ky)
        momsphhmat!(H, kx₀[col], ky[row], p)
        E[row,col,:] = eigvals(H)
    end
    nothing
end

# ground state energy of the HH hamiltonian
hhgrstate!(ve::Matrix{Float64}, q::Int) = hhgrstate!(ve, 1, q)
function hhgrstate!(ve::Matrix{Float64}, p::Int, q::Int)
    ky = linspace(-π, π, size(ve,1))
    kx₀= linspace(-π/q, π/q, size(ve,2))
    H = zeros(Complex{Float64}, (q,q))
    for col=1:length(kx₀), row=1:length(ky)
        momsphhmat!(H, kx₀[col], ky[row], p)
        ve[row,col] = eigmin(H)
    end
    nothing
end

# zero point energy error
ηzpe(q::Int, κ::Float64) = ηzpe(1, q, κ)
ηzpe(p::Int, q::Int, κ::Float64) = (N=15; A =
   spzeros(Complex{Float64}, N^2,N^2); ηzpe(A, p, q, κ))

function ηzpe(M::SparseMatrixCSC{Complex{Float64},Int}, p::Int,
   q::Int, κ::Float64)
    N::Int = sqrt(size(M,1))
    α::Float64 = p/q
    gs = Array(Float64, 25,25)
    hhgrstate!(gs, p, q)
    e1 = mean(gs)
    et = e1 + 1/2*κ/(2π*α)
    BP.buildham_exact!(M, N,α,κ)
    er = real(eigs(M, nev=1, which=:SR, ritzvec=false)[1][1])
    return 4π*α/κ * (er - et)
end

ηzpe(q::Int, κs::Vector{Float64}) = vec(ηzpe(collect(q), κs))
ηzpe(qs::UnitRange{Int}, κ::Float64) = ηzpe(collect(qs), κ)
ηzpe(qs::Vector{Int}, κ::Float64) = vec(ηzpe(qs, collect(κ)))
ηzpe(qs::UnitRange{Int}, κs::Vector{Float64}) = ηzpe(collect(qs), κs)

function ηzpe(qs::Vector{Int}, κs::Vector{Float64})
    N=15
    A = spzeros(Complex{Float64}, N^2,N^2)
    η = Array(Float64, length(qs), length(κs))
    for col=1:length(κs), row=1:length(qs)
        η[row,col] = ηzpe(A, 1, qs[row], κs[col])
    end
    return η
end

# level error
ηlev(q::Int, κ::Float64) = ηlev(1, q, κ)
ηlev(p::Int, q::Int, κ::Float64) = (N=15; A =
   spzeros(Complex{Float64}, N^2,N^2); ηlev(A, p, q, κ))

function ηlev(M::SparseMatrixCSC{Complex{Float64},Int}, p::Int,
   q::Int, κ::Float64)
    N::Int = sqrt(size(M,1))
    α::Float64 = p/q
    BP.buildham_exact!(M, N,α,κ)
    # get first 4 levels of spectrum of HH+trap
    bilevel = real(eigs(M, nev=4, which=:SR, ritzvec=false)[1])
    #calculate level spacing
    return 2π*α/κ*diff(bilevel[1:2])[1] - 1
end
function ηlev(qs::UnitRange{Int}, κ::Float64, lowβ::Int, highβ::Int)
    N=15 # system size is NxN
    # initialize (sparse) Hamiltonian matrix
    M = spzeros(Complex{Float64}, N^2,N^2)
    diffs = Array(Float64, length(qs))
    for (i,q) in enumerate(qs)
        # construct disipationless HH+trap Hamiltonian in Landau gauge
        # q dependent, need to loop over all q!
        BP.buildham_exact!(M, N,1/q,κ)
        # get first 5 levels of spectrum of M
        levels = real(eigs(M, nev=5, which=:SR, ritzvec=false)[1])
        #calculate level spacing
        diffs[i] = 2π/(q*κ) * (levels[highβ+1] - levels[lowβ+1]) - 1.0
    end
    return diffs
end

ηlev(q::Int, κs::Vector{Float64}) = vec(ηlev([q], κs))
ηlev(qs::UnitRange{Int}, κ::Float64) = ηlev([qs], κ)
ηlev(qs::Vector{Int}, κ::Float64) = vec(ηlev(qs, [κ]))
ηlev(qs::UnitRange{Int}, κs::Vector{Float64}) = ηlev([qs], κs)

function ηlev(qs::Vector{Int}, κs::Vector{Float64})
    N=15
    A = spzeros(Complex{Float64}, N^2,N^2)
    η = Array(Float64, length(qs), length(κs))
    for col=1:length(κs), row=1:length(qs)
        η[row,col] = ηlev(A, 1, qs[row], κs[col])
    end
    return η
end

bwidth(qs::UnitRange{Int}) = bwidth([qs])
bwidth(qs::Vector{Int}) = [bwidth(q) for q in qs]
bwidth(q::Int) = bwidth(1, q)

function bwidth(p::Int, q::Int)
    gstate = Array(Float64, 25, 25)
    hhgrstate!(gstate, p, q)
    a,b = extrema(gstate)
    return b-a
end

bgap(qs::UnitRange{Int}) = bgap([qs])
bgap(qs::Vector{Int}) = [bgap(q) for q in qs]
bgap(q::Int) = bgap(1, q)

function bgap(p::Int, q::Int)
    ladder = Array(Float64, (25,25,q))
    hhladder!(ladder, p)
    e1 = mean(ladder[:,:,1])
    e2 = mean(ladder[:,:,2])
    return e2-e1
end
end #module


struct RobustAdaptiveMetropolis{d, FT<:AbstractFloat, T <:
    Union{MVector{d,FT}, Vector{FT}}, SST<:StepSize,
    CT<:Cholesky{FT}} <: AdaptState
    L::CT
    α_opt::FT
    dx::T
    n_u::T
    step::SST
end

"""
    r = RobustAdaptiveMetropolis(x0, [α_opt, [step]])

Constructor for RobustAdaptiveMetropolis state.

# Arguments
- `x0`: The initial state vector
- `α_opt`: Desired mean accept rate; default `0.234`.
- `step`: Step size object; default `RAMStepSize(0.66,d)` where
  `d` is state dimension.

If `s` is `RWMState`, then proposal samples may be drawn calling
 `draw!(s, r)` and adaptation is performed with `adapt!(r, s, α)`.
"""

function RobustAdaptiveMetropolis(x0::T, α_opt::FT=FT(0.234),
  step::StepSize=RAMStepSize(FT(0.66),length(x0))) where {
      FT <: AbstractFloat, T <: AbstractVector{FT}}
    d = length(x0)
    L = Cholesky(eye(FT, d), :L, 0)
    dx = zeros(FT, d)
    n_u = zeros(FT, d)
    RobustAdaptiveMetropolis{d,FT,typeof(dx),typeof(step),typeof(L)}(L,
      α_opt, dx, n_u, step)
end

function RobustAdaptiveMetropolis(x0::T, α_opt::FT=FT(0.234),
  step::StepSize=RAMStepSize(FT(0.66),d)) where {d,
      FT <: AbstractFloat, T<:StaticArray{Tuple{d}, FT}}
    L = Cholesky(MMatrix{d,d}(eye(FT, d)), :L, 0)
    dx = zeros(MVector{d,FT})
    n_u = zeros(MVector{d,FT})
    RobustAdaptiveMetropolis{d,FT,typeof(dx),typeof(step),typeof(L)}(L,
      α_opt, dx, n_u, step)
end

@inline function adapt!(s::RobustAdaptiveMetropolis{d}, sr::RWMState{d},
      α::FT, k::Real) :: Nothing where {d, FT <: AbstractFloat}
    γ::FT = get(s.step, k)
    norm_u::FT = BLAS.nrm2(sr.u)
    norm_u = (norm_u == zero(FT)) ? one(FT) : norm_u

    s.n_u .= sr.u ./ norm_u
    #copy!(s.n_u, sr.u)
    #lmul!(one(FT)/norm_u, s.n_u)

    #mul!(s.dx, s.L.L, s.n_u)
    #lmul!(sqrt(γ*abs(α-s.α_opt)), s.dx)
    lowerTriInplaceMultiplyAdd!(s.dx, sqrt(γ*abs(α-s.α_opt)), s.L.factors,
              s.n_u, zero(FT), s.dx)
    if α >= s.α_opt
        lowrankupdate!(s.L, s.dx)
    else
        lowrankdowndate!(s.L, s.dx)
    end
    nothing
end

@inline function draw!(sr::RWMState, s::RobustAdaptiveMetropolis) :: Nothing
    draw!(sr, s.L)
    nothing
end


#################### G and K Distribution ####################

immutable GK <: ContinuousUnivariateDistribution
  A::Float64
  B::Float64
  g::Float64
  k::Float64
  c::Float64

  function GK(A::Real, B::Real, g::Real, k::Real, c::Real)
    ## check args
    0.0 <= c < 1.0 || throw(Argumenterror("c must be 0 <= c < 1"))

    ## distribution
    new(A, B, g, k, c)
  end
  GK(A::Real, B::Real, g::Real, k::Real) = GK(A, B, g, k, 0.8)
end

## Parameters

location(d::GK) = d.A
scale(d::GK) = d.B
skewness(d::GK) = d.g
kurtosis(d::GK) = d.k
asymmetry(d::GK) = d.c

params(d::GK) = (d.A, d.B, d.g, d.k, d.c)

minimum(d::GK) = -Inf
maximum(d::GK) = Inf

insupport{T <: Real}(d::GK, x::AbstractVector{T}) = true

function quantile(d::GK, p::Float64)
  z = quantile(Normal(), p)
  z2gk(d, z)
end

function Qinv(d::GK, x::Real)
  z = gk2z(d, x)
  cdf(Normal(), z)
end

function validate(d::GK)
  f = z -> -Fgkz(d, z)
  optres = optimize(f, 0, 20 / abs(skewness(d)))
  max(optres.f_minimum, 1.0) <= 1 / asymmetry(d)  &&
    scale(d) > 0 && kurtosis(d) > -0.5
end

function pdf(d::GK, x::Real)
  exp(logpdf(d, x))
end

function logpdf(d::GK, x::Real)
  if !validate(d)
    return -Inf
  end

  z = gk2z(d, x)

  term0 = exp(-skewness(d) * z)
  term1 = log(scale(d)) + log(asymmetry(d)) + 0.5 * log(2 * pi) +
    z^2 / 2 + kurtosis(d) * log(1 + z^2)
  term2 = (1 / asymmetry(d) + (1 - term0) / (1 + term0))
  term3 = (1 + (2 * kurtosis(d) + 1) * z^2) / (1 + z^2)
  term4 = 2 * skewness(d) * z * term0 / (1 + term0)^2

  logQp = term1 + log(term2 * term3 + term4)

  return -logQp
end

function logpdf{T <: Real}(d::GK, x::AbstractVector{T})
  if !validate(d)
    return -Inf
  end

  value = 0.0
  for i in 1:length(x)
    z = gk2z(d, x[i])

    term0 = exp(-skewness(d) * z)
    term1 = log(scale(d)) + log(asymmetry(d)) + 0.5 * log(2 * pi) +
      z^2 / 2 + kurtosis(d) * log(1 + z^2)
    term2 = (1 / asymmetry(d) + (1 - term0) / (1 + term0))
    term3 = (1 + (2 * kurtosis(d) + 1) * z^2) / (1 + z^2)
    term4 = 2 * skewness(d) * z * term0 / (1 + term0)^2

    value -= term1 + log(term2 * term3 + term4)
  end
  return value
end

function fit_mle{T <: Real}(d::GK, x::AbstractVector{T})
  loglikelihood = theta ->
    - logpdf(GK(theta[1], theta[2], theta[3], theta[4]), x)

  A0 = median(x)
  B0 = diff(quantile(x, [0.25,0.75]))[1]
  g0 = mean((x - A0) .^ 3) / (var(x) ^ (3/2))
  k0 = mean((x - A0) .^ 4) / (sumabs2(x - A0) / length(x))^2 - 3

  optres = optimize(loglikelihood, [A0, B0, g0, k0])

  return GK(optres.minimum...)
end

function gk2z(d::GK, x::Real; max_expands::Integer = 100, expand::Integer = 10)
  y = (x - location(d)) / scale(d)
  if y == 0
    return 0
  elseif isfinite(y)
    d0 = GK(0, 1, skewness(d), kurtosis(d), asymmetry(d))
    inter = [-abs(y) / (1 - asymmetry(d)), abs(y)]
    iter = 0
    a, b = z2gk(d0, inter[1]) - y, z2gk(d0, inter[2]) - y
    while a * b > 0
      idx = findmin(abs([a,b]))[2]
      inter[idx] -= sign(a)*expand
      iter += 1
      iter < max_expands || break
    end
    return fzero(z -> z2gk(d0, z) - y, inter)
  else
    return y
  end
end

function Fgkz(d::GK, z::Float64)
  term1 = exp(-skewness(d) * z)
  term2 = 2 * skewness(d) * z * term1
  term3 = (1 + (2 * kurtosis(d) + 1) * z ^ 2) / (1 + z ^ 2)

  (1 - term1) / (1 + term1) + term2 / ((1 + term1) * term3)
end

function z2gk(d::GK, z::Float64)
  term1 = exp(-skewness(d) * z)
  term2 = (1.0 + asymmetry(d) * (1.0 - term1) / (1.0 + term1))
  term3 = (1.0 + z ^ 2) ^ kurtosis(d)
  location(d) + scale(d) * z * term2 * term3
end


using SimplexSplitting

"""
Return a triangulation of the canonical simplex. Used for all the invariant
maps in this file.
"""
function canonical_simplex_triangulation(; dim::Int = 3, split_factor::Int = 3)
    # Define vertices of canonical simplex
    canonical_simplex_vertices = zeros(dim + 1, dim)
    canonical_simplex_vertices[2:(dim+1), :] = eye(dim)
    simplex_indices = zeros(Int, 1, dim + 1)
    simplex_indices[1, :] = collect(1:dim+1)

    refined = refine_triangulation(
        canonical_simplex_vertices,
        simplex_indices, [1],
        split_factor)

    points, simplex_inds = refined[1], refined[2]
    centroids, radii = SimplexSplitting.centroids_radii2(points, simplex_inds)
    orientations = SimplexSplitting.orientations(points, simplex_inds)
    volumes = abs.(orientations)

    Triangulation(
        points = points,
        simplex_inds = simplex_inds,
        centroids = centroids,
        radii = radii,
        orientations = orientations,
        volumes = volumes
    )
end

"""
    invariant_map_bull(t::Triangulation)

Invariant map given by

Ψ(x, y, z) = (x/2 + z/2 * (1-y), (y+x)/2, 1-z),

where x, y and z are columns of the t.points array. These points are guaranteed to lie in the unit cube, because each component <= 1.
"""
function invariant_map_bull(t::Triangulation)
    hcat(
        (t.points[:, 1]/2 + t.points[:, 3]/3) .* (1 - t.points[:, 2]),
        (t.points[:, 1] + t.points[:, 2])/2,
        1 - t.points[:, 3]
    )
end

function invariant_bullmap_on_cube(;dim::Int = 3, split_factor::Int = 4)

    if dim == 3
        t = canonical_simplex_triangulation(dim = dim, split_factor = split_factor)
        t.impoints = invariant_map_bull(t)
        t.centroids_im, t.radii_im = SimplexSplitting.centroids_radii2(t.impoints, t.simplex_inds)
        t.orientations_im = SimplexSplitting.orientations(t.impoints, t.simplex_inds)
    else
        error("map not defined for dim = ", dim)
    end

    return t
end


using Test
using LinearAlgebra

using CholeskyDecomposition

@testset "Cholesky decomposition" begin
    for i in 1:5
        n = rand(1:25)
        # Generates a random hermitian positive definite matrix
        A = begin
            L = rand(Complex{Float64}, n, n)

            # Makes it diagonally dominant
            (L + L') / sqrt(8.0) + I(n) * n
        end

        @assert ishermitian(A) && isposdef(A)

        @test begin
            cholesky_outer_product_decompose(A) ≈
            cholesky_banachiewicz_decompose(A) ≈
            cholesky_crout_decompose(A) ≈
            cholesky(A).L
        end
    end
end


#=
BFSLevelsMultiSource:
- Julia version: 1.6.2
- Author: samuel
- Date: 2021-09-11
=#
using SuiteSparseGraphBLAS
using SparseArrays
using LinearAlgebra

include("/Users/samuel/Documents/Jupyter/GraphBlasAlgorithmsFinal/src/BFSLevelsMultiSource.jl")
dim = 20000
sourcdim = 3
matrix =  sprand(Bool, dim, dim, 0.005)
matrix = Array(matrix)
matrix = GBMatrix(matrix)
sources =  sprand(Bool, sourcdim, dim, 0.15)
sources = Array(sources)
sources = GBMatrix(sources)
@time ms_bfsl(matrix, sources, dim, sourcdim)


function update_weights(data, e, same, other, weights, weights_mult, neigh_weights_same, neigh_weights_other, m, max_f_vals, min_f_vals)

	# Penalty term
	penalty = sum(neigh_weights_same.*(abs.(e .- same)./(max_f_vals .- min_f_vals .+ eps(Float64))), dims=1)

	# Reward term
	reward = sum(neigh_weights_other.*(weights_mult .* (abs.(e .- other)./(max_f_vals .- min_f_vals .+ eps(Float64)))), dims=1)

	# Weights update
	weights = weights .- penalty./(m*size(same, 1) + eps(Float64)) .+ reward./(m*size(other, 1) + eps(Float64))

	# Return updated weights.
	return vec(weights)
end


using HPAT
using MPI

@acc hpat function calcKmeans(k::Int64, file_name)
    points = DataSource(Matrix{Float64},HDF5,"/points",file_name)
    clusters = Kmeans(points, k,10)
    return clusters
end

K = 3

if (length(ARGS) > 0)
	K = parse(Int, ARGS[1])
end

c2 = calcKmeans(K, HPAT.getDefaultDataPath()*"/kmeans_small.hdf5")
rank = MPI.Comm_rank(MPI.COMM_WORLD)
pes = MPI.Comm_size(MPI.COMM_WORLD)

MPI.Barrier(MPI.COMM_WORLD)
t1 = time_ns()
c2 = calcKmeans(K, HPAT.getDefaultDataPath()*"/kmeans_large.hdf5")
t2 = time_ns()

if rank==0
	println("nodes: ", pes, "\ntime: ", (t2-t1)/1.0e9, "\nCentroids: ", c2)
end

MPI.Barrier(MPI.COMM_WORLD)


using LinearAlgebra
include("FEM1D.jl")
include("containers.jl")
include("Iterations.jl")

function ρ_error(dimension::Integer, k::Integer, iteration::Function, batch_size::Integer)
    x = randn((dimension, batch_size))
    y = x
    for i = 1:k
        y = iteration(y)
    end
    ρ = 0
    n_y, n_x = 0, 0
    for j = 1:batch_size
        ρ += (norm(y[:, j])/norm(x[:, j]))^(1/k)
    end
    ρ = ρ/batch_size
end

function ρ_F(dimension::Integer, k::Integer, iteration::Function, batch_size::Integer)
    x = rand((-1, 1), (dimension, batch_size))
    for l = 1:k
        x = iteration(x)
    end
    ρ = (norm(x)^2/batch_size)^(1/(2*k))
end

function prec_loss_function(parameters::AbstractArray{<:Real, 1}, prec::preconditioner, params::optimization_parameters)
    M = assemble_matrix(parameters, prec)
    iteration(x) = Richardson(x, prec.A, x -> M*x, parameters[end]; prec=prec.symmetry)
    dimension = size(prec.A)[1]
    if params.error == "Frobenius"
        loss = ρ_F(dimension, params.k_sweeps, iteration, params.k_batch)
    elseif params.error == "Spectral"
        loss = ρ_error(dimension, params.k_sweeps, iteration, params.k_batch)
    end
    return loss
end

function spectral_loss_function(parameters::AbstractArray{<:Real, 1}, prec::preconditioner, params::optimization_parameters)
    M = assemble_matrix(parameters, prec)
    iteration(x) = Richardson(x, prec.A, x -> M*x, 1/spectral_radius(parameters, prec, params); prec=prec.symmetry)
    dimension = size(prec.A)[1]
    if params.error == "Frobenius"
        loss = ρ_F(dimension, params.k_sweeps, iteration, params.k_batch)
    elseif params.error == "Spectral"
        loss = ρ_error(dimension, params.k_sweeps, iteration, params.k_batch)
    end
    return loss
end

function scalar_loss_function(parameters::AbstractArray{<:Real, 1}, M::SparseMatrixCSC{Float64,Int64}, prec::preconditioner, params::optimization_parameters)
    iteration(x) = Richardson(x, prec.A, x -> M*x, parameters[end]; prec=prec.symmetry)
    dimension = size(prec.A)[1]
    if params.error == "Frobenius"
        loss = ρ_F(dimension, params.k_sweeps, iteration, params.k_batch)
    elseif params.error == "Spectral"
        loss = ρ_error(dimension, params.k_sweeps, iteration, params.k_batch)
    end
    return loss
end

function ρ_error_T(dimension::Tuple{Integer, Integer}, k::Integer, iteration::Function, batch_size::Integer)
    n_y, n_x = 0, 0
    for j = 1:batch_size
        x = randn(dimension)
        n_x += norm(x)
        y = x
        for i = 1:k
            y = iteration(y)
        end
        n_y += norm(y)
    end
    ρ = (n_y/n_x)^(1/k)
end

function ρ_F_T(dimension::Tuple{Integer, Integer}, k::Integer, iteration::Function, batch_size::Integer)
    n_x = 0
    for j = 1:batch_size
        x = rand((-1, 1), dimension)
        for l = 1:k
            x = iteration(x)
        end
        n_x += norm(x)^2
    end
    ρ = (n_x/batch_size)^(1/(2*k))
end

function prec_loss_function(parameters::AbstractArray{<:Real, 1}, prec::T_preconditioner, params::optimization_parameters)
    N = 2^prec.L+1
    DOF_x, DOF_y = N - 2 + sum(prec.BCs[1:2] .== "N"), N - 2 + sum(prec.BCs[3:4] .== "N")
    dimension = (DOF_y, DOF_x)
    iteration(x) = Richardson_T(x, prec.C, x -> apply_T_preconditioner(x, parameters, prec), parameters[end]; prec=prec.symmetry)
    if params.error == "Frobenius"
        loss = ρ_F_T(dimension, params.k_sweeps, iteration, params.k_batch)
    elseif params.error == "Spectral"
        loss = ρ_error_T(dimension, params.k_sweeps, iteration, params.k_batch)
    end
    return loss
end

function spectral_loss_function(parameters::AbstractArray{<:Real, 1}, prec::T_preconditioner, params::optimization_parameters)
    N = 2^prec.L+1
    DOF_x, DOF_y = N - 2 + sum(prec.BCs[1:2] .== "N"), N - 2 + sum(prec.BCs[3:4] .== "N")
    dimension = (DOF_y, DOF_x)
    iteration(x) = Richardson_T(x, prec.C, x -> apply_T_preconditioner(x, parameters, prec), 1/spectral_radius(parameters, prec, params); prec=prec.symmetry)
    if params.error == "Frobenius"
        loss = ρ_F_T(dimension, params.k_sweeps, iteration, params.k_batch)
    elseif params.error == "Spectral"
        loss = ρ_error_T(dimension, params.k_sweeps, iteration, params.k_batch)
    end
    return loss
end

function Frobenius_loss_function(parameters::AbstractArray{<:Real, 1}, prec::T_preconditioner, params::optimization_parameters)
    N = 2^prec.L+1
    DOF_x, DOF_y = N - 2 + sum(prec.BCs[1:2] .== "N"), N - 2 + sum(prec.BCs[3:4] .== "N")
    dimension = (DOF_y, DOF_x)
    iteration(x) = x
    if prec.symmetry == "symmetric"
        iteration(x) = x - apply_T_preconditioner(reshape(prec.A*reshape(apply_T_preconditioner(x, parameters, prec), (DOF_y*DOF_x, )), dimension), parameters, prec)
    else
        iteration(x) = x - apply_T_preconditioner(reshape(prec.A*reshape(x, (DOF_y*DOF_x, )), dimension), parameters, prec)
    end
    loss = ρ_F_T(dimension, 1, iteration, params.k_batch)
end

function Frobenius_loss_function(parameters::AbstractArray{<:Real, 1}, prec::preconditioner, params::optimization_parameters)
    M = assemble_matrix(parameters, prec)
    iteration(x) = x
    if prec.symmetry == "left"
        iteration(x) = x - M*(prec.A*x)
    else
        iteration(x) = x - M*(prec.A*(M*x))
    end
    loss = ρ_F(size(prec.A)[1], 1, iteration, params.k_batch)
end

function scalar_loss_function(parameters::AbstractArray{<:Real, 1}, M::SparseMatrixCSC{Float64,Int64}, prec::T_preconditioner, params::optimization_parameters)
    iteration(x) = Richardson(x, prec.A, x -> M*x, parameters[end]; prec=prec.symmetry)
    dimension = size(prec.A)[1]
    if params.error == "Frobenius"
        loss = ρ_F(dimension, params.k_sweeps, iteration, params.k_batch)
    elseif params.error == "Spectral"
        loss = ρ_error(dimension, params.k_sweeps, iteration, params.k_batch)
    end
    return loss
end

function spectral_radius(parameters::AbstractArray{<:Real, 1}, prec::preconditioner, params::optimization_parameters)
    function iteration(x::AbstractArray{<:Real, 2}, M::SparseMatrixCSC{<:Real, <:Integer}, A::SparseMatrixCSC{<:Real, <:Integer}, symmetry::String)
        if symmetry == "symmetric"
            return M*A*M*x
        elseif symmetry == "left"
            return M*A*x
        end
    end
    M = assemble_matrix(parameters, prec)
    dimension = size(prec.A)[1]
    if params.error == "Frobenius"
        loss = ρ_F(dimension, params.k_sweeps, x -> iteration(x, M, prec.A, prec.symmetry), params.k_batch)
    elseif params.error == "Spectral"
        loss = ρ_error(dimension, params.k_sweeps, x -> iteration(x, M, prec.A, prec.symmetry), params.k_batch)
    end
    return loss
end

function spectral_radius(parameters::AbstractArray{<:Real, 1}, prec::T_preconditioner, params::optimization_parameters)
    function iteration(x::AbstractArray{<:Real, 2}, M::SparseMatrixCSC{<:Real, <:Integer}, A::SparseMatrixCSC{<:Real, <:Integer}, symmetry::String)
        if symmetry == "symmetric"
            return M*A*M*x
        elseif symmetry == "left"
            return M*A*x
        end
    end
    M = assemble_matrix(parameters, prec)
    dimension = size(prec.A)[1]
    if params.error == "Frobenius"
        loss = ρ_F(dimension, params.k_sweeps, x -> iteration(x, M, prec.A, prec.symmetry), params.k_batch)
    elseif params.error == "Spectral"
        loss = ρ_error(dimension, params.k_sweeps, x -> iteration(x, M, prec.A, prec.symmetry), params.k_batch)
    end
    return loss
end

function test_ρ_1(k::Integer, N::Integer, L::Integer)
    A, b, x = BVP_DD(L, x -> one(x), x -> zero(x), 1, 1, x -> x)
    ρ = maximum(abs.(eigvals(Array(A))))
    ρ_err = ρ_error(2^L-1, k, x -> A*x)
    ρ_Frob = ρ_F(2^L-1, k, N, x -> A*x)
    println("Spectral radius ", ρ)
    println("Estimation with error ", ρ_err)
    println("Estimation with Frobenius ", ρ_Frob)
end

function test_ρ_2(k::Integer, L::Integer)
    A, b, x = BVP_DD(L, x -> 1 + x^2, x -> exp(x), 1, 1, x -> x)
    ρ = maximum(abs.(eigvals(Array(A))))
    ρ_estimated = ρ_error(2^L-1, k, x -> A*x)
    println("Spectral radius ", ρ)
    println("Estimation ", ρ_estimated)
end