text:
# Plot the loss function on a random 1D line
# Authors: Ali Siahkoohi, alisk@gatech.edu
# Date: December 2020

using DrWatson
@quickactivate :ExtendedConv

using Flux
using PyPlot
using Seaborn
using Printf
using LinearAlgebra
using Optim
using JLD2
using Random
using ProgressMeter

matplotlib.use("Agg")
set_style("whitegrid")
rc("font", family="serif", size=10)
font_prop = matplotlib.font_manager.FontProperties(
    family="serif",
    style="normal",
    size=10
)
sfmt=matplotlib.ticker.ScalarFormatter(useMathText=true)
sfmt.set_powerlimits((0, 0))

sim_name = "sim_name"
fname = "batchsize=1_f_tol=1e-6_iterations=100_k=3_lambda=0.04_n_grid=21_nlayers=4_nx=16_ny=16_seed1=21_seed2=12_z_std=0.01_#2.jld2"

file_path = datadir(sim_name, fname)
file = jldopen(file_path, "r")

n_grid = file["n_grid"]
lambda = file["lambda"]
nx = file["nx"]
ny = file["ny"]
k = file["k"]
nlayers = file["nlayers"]
z_std = file["z_std"]
f_tol = file["f_tol"]
iterations = file["iterations"]
batchsize = file["batchsize"]
w_true = file["w"]
z = file["z"]
y = file["y"]
v1 = file["v1"]
v2 = file["v2"]
loss = file["loss"]
loss_ext = file["loss_ext"]
αs = file["αs"]
βs = file["βs"]
seed1 = file["seed1"]
seed2 = file["seed2"]

nc = 1

G = Chain([Conv((k, k), nc => nc, sigmoid, pad=(1, 1), stride=(1, 1))  for i=1:nlayers]...)
Flux.loadparams!(G, deepcopy(w_true))

w_true_ext = deepcopy(cat([w_true[i] for i=1:2:2*nlayers]..., dims=3)[:, :, :, 1])

Random.seed!(19)
G = Chain([Conv((k, k), nc => nc, sigmoid, pad=(1, 1), stride=(1, 1))  for i=1:nlayers]...)

α = 2f0
β = -2f0
w0 = deepcopy(w_true) + α*v1 + β*v2
Flux.loadparams!(G, deepcopy(w0))
w0_ext = cat([w0[i] for i=1:2:2*nlayers]..., dims=3)[:, :, :, 1]

C = extended_conv_op(nx, ny; nc=nc, batchsize=1, k=k, nlayers=nlayers, w=w0_ext)
w_ex_0 = cat(
    [randn(Float32, size(C[i].mask)).*C[i].mask for i=1:length(C)]...,
    dims=3
)

max_itr = 200
adam_lr = 0.05

G_ext(x) = reshape(C(reshape(x, :)), nx, ny, nc, 1)[end:-1:1, end:-1:1]
perm_y(y, idx) = reshape(y[end:-1:1, end:-1:1, :, idx], :)

opt_var_proj = Optim.Options(
    iterations=iterations,
    f_tol=1f-6,
    show_trace=false,
    store_trace=false,
    show_every=1
)

var_proj!(func_val, dw, w0_ext) = var_proj_op(
    func_val, dw, w0_ext, reshape(z, :), perm_y(y, 1), C, w_ex_0,
    lambda, opt_var_proj
)

ws_ext_lbfgs = []
fvals_ext_lbfgs = []
dummy_f = 0f0
dummy_w = zeros(Float32, size(w0_ext))
cbb = tr -> begin
            push!(ws_ext_lbfgs, tr[end].metadata["x"])
            fval = var_proj!(dummy_w, dummy_w,  tr[end].metadata["x"])
            push!(fvals_ext_lbfgs, deepcopy(fval))
            false
        end

opt_global = Optim.Options(
    iterations=max_itr,
    show_trace=true,
    store_trace=true,
    show_every=1,
    extended_trace=true,
    callback=cbb
)

opt_lbfgs = optimize(Optim.only_fg!(var_proj!), w0_ext, LBFGS(m=100), opt_global)

w0 = deepcopy(w_true) + α*v1 + β*v2
w0_ext = cat([w0[i] for i=1:2:2*nlayers]..., dims=3)[:, :, :, 1]
Flux.loadparams!(G, deepcopy(w0))

lr = adam_lr
opt = ADAM(lr)
ws_ext = []
var_proj!(func_val, dw, w0_ext) = var_proj_op(
    func_val, dw, w0_ext, reshape(z, :), perm_y(y, 1), C, w_ex_0,
    lambda, opt_var_proj
)

fvals_ext = []
func_val = 0f0
dw = zeros(Float32, size(w0_ext))
for j =1:max_itr

    loss_val = var_proj!(func_val, dw, w0_ext)
    push!(ws_ext, deepcopy(w0_ext))
    println(" Iteration: ", j, " | Objective = ", loss_val, "\r")

    push!(fvals_ext, loss_val)
    Flux.Optimise.update!(opt, w0_ext, dw)
end

w0 = deepcopy(w_true) + α*v1 + β*v2
w0_ext = cat([w0[i] for i=1:2:2*nlayers]..., dims=3)[:, :, :, 1]
Flux.loadparams!(G, deepcopy(w0))

lr = adam_lr
opt = ADAM(lr)
ws = []
fvals = []
for i in 1:max_itr

    loss_val, back = Flux.pullback(Flux.params(G)) do
        5f-1*norm(G(z) - y)^2f0
    end
    grads = back(1.0f0)
    println(" Iteration: ", i, " | Objective = ", loss_val, "\r")

    what = deepcopy(Flux.params(G))
    what_ = cat([what[i] for i=1:2:2*nlayers]..., dims=3)[:, :, :, 1]
    push!(ws, what_)

    push!(fvals, loss_val)
    for p in Flux.params(G)
        if length(size(p)) == 4
            Flux.Optimise.update!(opt, p, grads[p])
        end
    end

end

G = Chain([Conv((k, k), nc => nc, sigmoid, pad=(1, 1), stride=(1, 1))  for i=1:nlayers]...)
w0 = deepcopy(w_true) + α*v1 + β*v2
Flux.loadparams!(G, deepcopy(w0))
w0_ext = cat([w0[i] for i=1:2:2*nlayers]..., dims=3)[:, :, :, 1]

loss_flux!(func_val, dw, w0_ext) = loss_flux(func_val, dw, w0_ext, z, y, G)

ws_lbfgs = []
fvals_lbfgs = []
dummy_f = 0f0
dummy_w = zeros(Float32, size(w0_ext))
cbb = tr -> begin
            push!(ws_lbfgs, tr[end].metadata["x"])
            fval = loss_flux!(dummy_w, dummy_w,  tr[end].metadata["x"])
            push!(fvals_lbfgs, deepcopy(fval))
            false
        end

opt_global = Optim.Options(
    iterations=max_itr,
    show_trace=true,
    store_trace=true,
    show_every=1,
    extended_trace=true,
    callback=cbb
)

opt_lbfgs = optimize(Optim.only_fg!(loss_flux!), w0_ext, LBFGS(m=100), opt_global)

proj_ws_ext_lbfgs = sol_traj(v1, v2, ws_ext_lbfgs, w_true_ext)
proj_ws_ext = sol_traj(v1, v2, ws_ext, w_true_ext)
proj_ws = sol_traj(v1, v2, ws, w_true_ext)
proj_ws_lbfgs = sol_traj(v1, v2, ws_lbfgs, w_true_ext)

save_dict = @strdict loss loss_ext lambda n_grid f_tol iterations batchsize z_std nlayers nx ny k v1 v2 seed1 seed2 w_true z y αs βs max_itr adam_lr α β ws_ext_lbfgs ws_ext ws ws_lbfgs fvals_ext fvals_ext_lbfgs fvals fvals_lbfgs
@tagsave(
    datadir(sim_name, savename(save_dict, "jld2"; digits=12)),
    save_dict;
    safe=true
)

loss = (loss .- minimum(loss))./maximum(loss)
loss_ext = (loss_ext .- minimum(loss_ext))./maximum(loss_ext)

save_path = plotsdir(sim_name, savename(save_dict; digits=12))

fig, ax = plt.subplots(1, 1, dpi=300, figsize=(4, 4))
implt=ax.contourf(αs, βs, loss_ext, 25, cmap=Seaborn.ColorMap("mako"), linewidths=1.5,
    vmin=-1f-1, vmax=1f0)
ax.quiver(
    proj_ws_ext_lbfgs[1][1:end-1], proj_ws_ext_lbfgs[2][1:end-1],
    proj_ws_ext_lbfgs[1][2:end]-proj_ws_ext_lbfgs[1][1:end-1],
    proj_ws_ext_lbfgs[2][2:end]-proj_ws_ext_lbfgs[2][1:end-1],
    scale_units="xy", angles="xy", scale=1,
    color="y"
)
ax.set_title("Extended loss landscape")
fig.colorbar(implt, ax=ax, fraction=0.057, pad=0.01, format=sfmt)
ax.scatter(
    βs[argmin(loss_ext)[2]], αs[argmin(loss_ext)[1]], s=30.0,
    color="#b80000", marker="*"
)
ax.set_ylabel(L"$\beta$")
ax.set_xlabel(L"$\alpha$")
safesave(joinpath(save_path, "loss_ext_and_opt_traj_lbfgs.png"), fig)
close(fig)

fig, ax = plt.subplots(1, 1, dpi=300, figsize=(4, 4))
implt=ax.contourf(αs, βs, loss_ext, 25, cmap=Seaborn.ColorMap("mako"), linewidths=1.5,
    vmin=-1f-1, vmax=1f0)
ax.quiver(
    proj_ws_ext[1][1:end-1], proj_ws_ext[2][1:end-1],
    proj_ws_ext[1][2:end]-proj_ws_ext[1][1:end-1],
    proj_ws_ext[2][2:end]-proj_ws_ext[2][1:end-1],
    scale_units="xy", angles="xy", scale=1,
    color="y"
)
ax.set_title("Extended loss landscape")
fig.colorbar(implt, ax=ax, fraction=0.057, pad=0.01, format=sfmt)
ax.scatter(
    βs[argmin(loss_ext)[2]], αs[argmin(loss_ext)[1]], s=30.0,
    color="#b80000", marker="*"
)
ax.set_ylabel(L"$\beta$")
ax.set_xlabel(L"$\alpha$")
safesave(joinpath(save_path, "loss_ext_and_opt_traj.png"), fig)
close(fig)

fig, ax = plt.subplots(1, 1, dpi=300, figsize=(4, 4))
implt=ax.contourf(αs, βs, loss, 25, cmap=Seaborn.ColorMap("mako"), linewidths=1.5,
    vmin=-1f-1, vmax=1f0)
ax.quiver(
    proj_ws[1][1:end-1], proj_ws[2][1:end-1],
    proj_ws[1][2:end]-proj_ws[1][1:end-1],
    proj_ws[2][2:end]-proj_ws[2][1:end-1],
    scale_units="xy", angles="xy", scale=1,
    color="y"
)
ax.set_title("Loss landscape")
fig.colorbar(implt, ax=ax, fraction=0.057, pad=0.01, format=sfmt)
ax.scatter(
    βs[argmin(loss)[2]], αs[argmin(loss)[1]], s=30.0,
    color="#b80000", marker="*"
)
ax.set_ylabel(L"$\beta$")
ax.set_xlabel(L"$\alpha$")
safesave(joinpath(save_path, "loss_and_opt_traj.png"), fig)
close(fig)

fig, ax = plt.subplots(1, 1, dpi=300, figsize=(4, 4))
implt=ax.contourf(αs, βs, loss, 25, cmap=Seaborn.ColorMap("mako"), linewidths=1.5,
    vmin=-1f-1, vmax=1f0)
ax.quiver(
    proj_ws_lbfgs[1][1:end-1], proj_ws_lbfgs[2][1:end-1],
    proj_ws_lbfgs[1][2:end]-proj_ws_lbfgs[1][1:end-1],
    proj_ws_lbfgs[2][2:end]-proj_ws_lbfgs[2][1:end-1],
    scale_units="xy", angles="xy", scale=1,
    color="y"
)
ax.set_title("Loss landscape")
fig.colorbar(implt, ax=ax, fraction=0.057, pad=0.01, format=sfmt)
ax.scatter(
    βs[argmin(loss)[2]], αs[argmin(loss)[1]], s=30.0,
    color="#b80000", marker="*"
)
ax.set_ylabel(L"$\beta$")
ax.set_xlabel(L"$\alpha$")
safesave(joinpath(save_path, "loss_and_opt_traj_lbfgs.png"), fig)
close(fig)

fvals_ext_norm = fvals_ext./fvals_ext[1]
fvals_ext_norm_lbfgs = fvals_ext_lbfgs./fvals_ext_lbfgs[1]
fvals_norm = fvals./fvals[1]
fvals_norm_lbfgs = fvals_lbfgs./fvals_lbfgs[1]

fig, ax = plt.subplots(1, 1, figsize=(4, 7.5))
ax.semilogy(fvals_norm, label="conventional (Adam)")
ax.semilogy(fvals_norm_lbfgs, label="conventional (L-BFGS)")
ax.semilogy(fvals_ext_norm, label="extended (Adam)")
ax.semilogy(range(0, length(ws_ext_lbfgs)-1, length=length(fvals_ext_norm_lbfgs)), fvals_ext_norm_lbfgs, label="extended (L-BFGS)")
ax.legend(loc="upper right")
ax.set_xlabel("Iterations")
ax.set_xlim([0, max_itr])
ax.set_ylabel("Normalized loss")
ax.set_title("Extended vs conventional convergence")
safesave(joinpath(save_path, "loss-comparison.png"), fig)
close(fig)

true_dist_ext = [norm(ws_ext[i] - w_true_ext) for i =1:length(ws_ext)]./norm(w_true_ext)
true_dist = [norm(ws[i] - w_true_ext) for i =1:length(ws)]./norm(w_true_ext)

true_dist_ext_lbfgs = [norm(ws_ext_lbfgs[i] - w_true_ext) for i =1:length(ws_ext_lbfgs)]
true_dist_ext_lbfgs = true_dist_ext_lbfgs./norm(w_true_ext)

true_dist_lbfgs = [norm(ws_lbfgs[i] - w_true_ext) for i =1:length(ws_lbfgs)]
true_dist_lbfgs = true_dist_lbfgs./norm(w_true_ext)

fig, ax = plt.subplots(1, 1, figsize=(4, 7.5))
ax.plot(true_dist, label="conventional (Adam)")
ax.plot(true_dist_lbfgs, label="conventional (L-BFGS)")
ax.plot(true_dist_ext, label="extended (Adam)")
ax.plot(true_dist_ext_lbfgs, label="extended (L-BFGS)")
ax.legend(loc="lower right")
ax.set_xlabel("Iterations")
ax.set_ylabel("Relative weight recovery error")
ax.set_xlim([0, max_itr])
ax.set_title("error w.r.t. ground-truth weight")
safesave(joinpath(save_path, "true-error-comparison.png"), fig)
close(fig)


text:
using BasisMatrices

# construct 2d basis
basis = Basis(LinParams(15, -2, 2), LinParams(10, -1, 3))
S, (x, y) = nodes(basis)

# define function to approximate
func(x, y) = exp(-2*x) * sin(y)
f = func.(S[:, 1], S[:, 2])

### getting coefficient vector

## tensor form
bmt = BasisMatrix(basis, Tensor(), [x, y])

# bmt.vals[1] is all basis funcs for first dim, evaluated at nodes for first dim
# similar for 2
@assert bmt.vals[1] == evalbase(basis.params[1], x, 0)
@assert bmt.vals[2] == evalbase(basis.params[2], y, 0)

# to get coefficients we need to do kron(bmt.vals[2], bmt.vals[1]) \ f
c1 = kron(bmt.vals[2], bmt.vals[1]) \ f

# calling funfitxy with the BasisMatrix instance will use optimized version of
# that operation that never constructs full kronecker product
c2 = funfitxy(basis, bmt, f)[1]
@assert maximum(abs, c1 - c2) < 1e-14

## Direct form
bmd = BasisMatrix(basis, Direct(), S)

# bmd.vals[1] will be have size (length(x)*length(y), length(x)) and
# bmd.vals[2] will be have size (length(x)*length(y), length(y))
# Each column represents the evaluation at one basis function.

# The reason there are length(x)*length(y) rows in both entries of bmd.vals is
# that an expansion down the columns (i.e. at the evaluation points) has
# already been done. The ordering of the rows follows column major conventions,
# which means that the first dimension varies fastest. To see this more clearly
# look at the form of the matrix `S`.

# Now to get coefficients we will need to compute row_kron(bm.vals[2], bm.vals[2]) \ f
c3 = row_kron(bmd.vals[2], bmd.vals[1]) \ f

@assert c3 ≈ c2

# calling funfitxy with the BasisMatrix instance will use optimized version of
# that operation that never constructs row wise kronecker product
c4 = funfitxy(basis, bmd, f)[1]

@assert c4 ≈ c2

## Expanded form
bme = BasisMatrix(basis, Expanded(), S)

# to get coefficients we just need to do bme.vals[1] \ f
c5 = bme.vals[1] \ f

@assert c5 ≈ c2

# the reason this worked is that Expanded form computes the fully expanded
# version of the basis matrix. This means that the expansion down columns and
# across rows has already been done.

# the funfitxy method for the Expanded BasisMatrix does exactly this operation
c6 = funfitxy(basis, bme, f)[1]

@assert c6 ≈ c2

#=

Compare/contrast formats:

- Tensor is the most efficient in terms of memory, but you are limited to
evaluation at a full cartesian product of the individual vectors you pass in
- Direct is the middle in terms of efficiency, but doesn't suffer from the
flexibility issue we saw with Expanded: you can evaluate at arbitrary points in
the 2d space simply by passing in the points as rows to the matrix S. However,
if we do want the cartesian product of 1d vectors, we are better off using
Tensor because the Direct form will repeat calculations for the expanded grid
points (to see it, look at how the second column of S repeats the first
gridpoint of y a total of length(x) times. When you build the Direct form, we
don't assume any structure about the rows of S, so we will repeat the
computations needed to evaluate the BasisMatrices along the second dimension
length(x) times. The same holds for the first dimension (each evaluation is
repeated length(y) times), but it is easier to explain using the second.)
- Expanded is the least efficient in terms of memory and CPU time, but is the
easiest to reason about when thinking about how to construct a coefficient
vector. It also has the benefit of specifying arbitrary evaluation points as
rows of the matrix argument.

=#

### Conversion

# you can only convert "outword" in terms of expanding matrices. This means
# you can convert from Tensor to either Direct or Expanded, or you can convert
# direct to Expanded. No other combinations are possible.

## Tensor -> Direct
# conversion from Tensor -> Direct amounts to expanding the columns of the
# tensor form such that elements are repeated in column major order (i.e.
# the first column varies fastest)
Φd1 = (
    repeat(bmt.vals[1], length(y)),
    hcat([repeat(bmt.vals[2][:, i], inner=length(x)) for i in 1:length(y)]...)
    )

@assert bmd.vals[1] ≈ Φd1[1]
@assert bmd.vals[2] ≈ Φd1[2]

## Tensor -> expanded
# we've already seen this one as just kron(bmt.vals[2], bmt.vals[1]) (think
# about how we computed the c1 and c5)
Φe1 = kron(bmt.vals[2], bmt.vals[1])
@assert Φe1 ≈ bme.vals[1]

## Direct -> Expanded
# We've also seen that this is row_kron(bmd.vals[2], bmd.vals[1]) (check c3 and
# c5)
Φe2 = row_kron(bmd.vals[2], bmd.vals[1])
@assert Φe2 ≈ bme.vals[1]

# The julia methods `convert(Format, bm, [order])` do these operations for you,
# but return a fully formed BasisMatrix object instead of just a plain Julia
# AbstractMatrix subtype
Φd2 = convert(Direct, bmt)
@assert bmd.vals[1] ≈ Φd2.vals[1]
@assert bmd.vals[2] ≈ Φd2.vals[2]

Φe3 = convert(Expanded, bmt)
@assert Φe3.vals[1] ≈ bme.vals[1]

Φe4 = convert(Expanded, bmd)
@assert Φe4.vals[1] ≈ bme.vals[1]

### Evaluation off nodes

# the last stop on our tour is how we might consider evaluating the
# interpolated function at points different from the interpolation nodes.
# As an example, suppose we want to evaluate at all the original y points, but
# new x points. We will want the full combination of all these points
x2 = collect(range(-1.5; stop=1.5, length=40))

# for evaluation of our approximation, let's evaluate func at these points
f2 = vec(func.(x2, y'))

## Using Tensor form
# Let's try using tensor form first. Remember that bmt.vals[2] already contains
# the evaluation of all  the basis functions in the 2nd dimension at all the
# points in y. This means that we only need to evaluate a new basis matrix for
# the x2 vector
Φ1_x2 = evalbase(basis.params[1], x2, 0)

# then to evaluate our interpolated function we do...
ft1 = kron(bmt.vals[2], Φ1_x2)*c1

# the approximation actually doesn't fit func extremely well, so we will have a
# relatively loose sense of success in our interpolation
@assert maximum(abs, ft1 - f2) < 1.0

# if we didn't want to form the full `kron` we could have construct a
# BasisMatrix instance by hand (NOTE: to do this we had to do a relatively ugly
# thing and allocate vals first, then populate it and we had to pass type
# params to the BasisMatrix constructor. We can get around this, we just need
# to be more clever in how we accept arguments to BasisMatrix).
bmt2_vals = Array{typeof(Φ1_x2)}(undef,1, 2)
bmt2_vals[1] = Φ1_x2; bmt2_vals[2] = bmt.vals[2]
bmt2 = BasisMatrix{Tensor,typeof(Φ1_x2)}([0 0], bmt2_vals)
ft2 = funeval(c1, bmt2)

@assert ft1 ≈ ft2

## Using Direct form
# to use direct form, we first need to construct the expanded grid on x2 and y
S2 = BasisMatrices.gridmake(x2, y)
bmd2 = BasisMatrix(basis, Direct(), S2)

# now we can evaluate using row_kron(bmd2.vals[2], bmd2.vals[1]) * c
fd1 = row_kron(bmd2.vals[2], bmd2.vals[1]) * c1
@assert fd1 ≈ ft1

# we could also use the funeval method which would help us avoid building that
# full row_kron matrix
fd2 = funeval(c1, bmd2)
@assert fd2 ≈ ft1

## Using Expanded form
# we can also use S2 to construct the Expanded basis matrix
bme2 = BasisMatrix(basis, Expanded(), S2)

# evaluation now is juse bme2.vals[2] * c
fe1 = bme2.vals[1]*c1

@assert fe1 ≈ ft1

# Here the funeval version is identical to the operation from above
fe2 = funeval(c1, bme2)
@assert fe2 ≈ ft1

ext: jl


text:
"""
calcentropy \\

Calculates the requested entropies using the x and y variables given. \\
Arguments:  x = random variable, usually input \\
            y = random variable, usually output/classification \\
            xdims = [number of variables in x, the maximum number of possible values across the dimensions in x] \\
            ydims = [number of variables in y, the maximum number of possible values across the dimensions in y] \\
            calc = The entropies to calculate ex: ["HX", "HXY"] \\
            method = sampling bias correction method ex: "pt" \\
            sampling = method for sampling ex: "naive" \\
Returns:    dictionary, entropies listed in calc => the entropy values \\
"""

function calcentropy(x::Array{Int}, y::Array{Int}, xdims::Array{Int, 1}, ydims::Array{Int, 1}, calc::Array{String, 1}, method::String, sampling::String; base::Int=2)
  xn, xmax = xdims
  yn, ymax = ydims
  x = atleast2d(x)
  y = atleast2d(y)
  checkinputs(x, y, xmax, ymax, xn, yn)
  probs = requireprobs(x, xn, xmax, y, yn, ymax, calc, sampling)
  return requireents(probs, calc, base=base)
end

function calcentropy(x::Array{Int}, xdims::Array{Int, 1}, calc::Array{String, 1}, method::String, sampling::String; base::Int=2)
  xn, xmax = xdims
  x = atleast2d(x)
  checkinputs(x, xmax, xn)
  x_dist_dim = xmax ^ xn
  probs = Dict{String, Array{Float64}}("PX" => prob(decimalise(x, xn, xmax), x_dist_dim, sampling))
  return requireents(probs, calc, base=base)
end


text:
export chain_complex, isexact, homology, free_resolution

######################################################################
#
# Lift of homomorphisms
#
######################################################################
#=
  G
  | phi
  V
  F <- H
    psi
 and Im(phi) subset Im(psi), then G -> H can be constructed
=#

@doc Markdown.doc"""
    lift(phi::Map, psi::Map) -> Map
Given $\phi: G\to F$ and $\psi:H \to F$ s.th. $\Im(\phi) \subseteq \Im(\psi)$
return the map $G\to H$ to make the diagram commute.
"""
function lift(phi::Map, psi::Map)
  x = [haspreimage(psi, image(phi, g)) for g = gens(domain(phi))]
  @assert all(t -> t[1], x)
  return hom(domain(phi), domain(psi), [t[2] for t = x])
end

@doc Markdown.doc"""
    zero_map(G::GrpAbFinGen) -> Map
Create the map $G \to \{0\}$.
"""
function zero_map(G::GrpAbFinGen)
  Z = abelian_group([1])
  set_name!(Z, "Zero")
  return hom(G, Z, [Z[0] for i=1:ngens(G)])
end

######################################################################
#
# complex/ free resolution
#
######################################################################
function iszero(h::T) where {T <: Map{<:GrpAbFinGen, <:GrpAbFinGen}}
  return all(x -> iszero(h(x)), gens(domain(h)))
end

mutable struct ChainComplex{T}
  @declare_other
  maps::Array{<:Map, 1}
  direction::Symbol
  exact::Array{Bool, 1}
  function ChainComplex(A::S; check::Bool = true, direction:: Symbol = :left) where {S <:Array{<:Map{<:T, <:T}, 1}} where {T}
    if check
      @assert all(i-> iszero(A[i]*A[i+1]), 1:length(A)-1)
    end
    r = new{T}()
    r.maps = A
    r.direction = direction
    return r
  end
  function ChainComplex(X::Type, A::S; check::Bool = true, direction:: Symbol = :left) where {S <:Array{<:Map, 1}}
    if check
      @assert all(i-> iszero(A[i]*A[i+1]), 1:length(A)-1)
    end
    r = new{X}()
    r.maps = A
    r.direction = direction
    return r
  end

end

length(C::ChainComplex) = length(C.maps)
Base.map(C::ChainComplex, i::Int) = C.maps[i]
obj(C::ChainComplex, i::Int) = (i==0 ? domain(C.maps[1]) : codomain(C.maps[i]))

obj_type(C::ChainComplex{T}) where {T} = T
map_type(C::ChainComplex) = eltype(C.maps)

function show(io::IO, C::ChainComplex)
  @show_name(io, C)
  @show_special(io, C)

  Cn = get_special(C, :name)
  if Cn === nothing
    Cn = "C"
  end
  name_mod = String[]
  name_map = String[]
  mis_map = Tuple{Int, <:Map}[]
  mis_mod = Tuple{Int, <:Any}[]

  if C.direction == :left
    rng = 0:length(C)
    arr = ("--", "-->")
    dir = 1
  else
    rng = length(C)+1:-1:1
    arr = ("<--", "--")
    dir = 0
  end

  for i=1:length(C)
    phi = map(C, i)
    if get_special(phi, :name) !== nothing
      push!(name_map, get_special(phi, :name))
    else
      push!(name_map, "")
      push!(mis_map, (i, phi))
    end
  end
  for i=0:length(C)
    M = obj(C, i)
    if get_special(M, :name) !== nothing
      push!(name_mod, get_special(M, :name))
    else
      push!(name_mod, "$(Cn)_$i")
      push!(mis_mod, (i, M))
    end
  end

  io = IOContext(io, :compact => true)
  for i=rng
    if i == first(rng)
      print(io, name_mod[i+dir])
      continue
    end
    if name_map[i] != ""
      print(io, " ", arr[1], " ", name_map[i], " ", arr[2], " ", name_mod[i+dir])
    else
      print(io, " ", arr[1], arr[2], " ", name_mod[i+dir])
    end
  end
  if length(mis_mod) > 0 # || length(mis_map) > 0
    print(io, "\nwhere:\n")
    for (i, M) = mis_mod
      print(io, "\t$(Cn)_$i = ", M, "\n")
    end
#    for (i, phi) = mis_map
#      print(io, "\tphi_$i = ", phi, "\n")
#    end
  end
end

@doc Markdown.doc"""
    chain_complex(A::Map{GrpAbFinGen, GrpAbFinGen, <:Any, <:Any}...) -> ChainComplex{GrpAbFinGen}
Given maps $A_i$ s.th. $\Im(A_i) \subseteq \Kern(A_{i+1})$, this creates
the chain complex.
"""
function chain_complex(A::Map{GrpAbFinGen, GrpAbFinGen, <:Any, <:Any}...)
  return ChainComplex(collect(A))
end

function chain_complex(A::Array{<:Map{GrpAbFinGen, GrpAbFinGen, <:Any, <:Any}, 1})
  return ChainComplex(A)
end

Base.lastindex(C::ChainComplex) = length(C)
getindex(C::ChainComplex{T}, u::UnitRange) where {T} = ChainComplex(T, C.maps[u], check = false)

@doc Markdown.doc"""
    isexact(C::ChainComplex) -> Bool
Tests is the complex $A_i: G_i \to G_{i+1}$
is exact, ie. if $\Im(A_i) = \Kern(A_{i+1})$.
"""
function isexact(C::ChainComplex)
  return all(i->iseq(image(C.maps[i])[1], kernel(C.maps[i+1])[1]), 1:length(C)-1)
end

@doc Markdown.doc"""
    free_resolution(G::GrpAbFinGen) -> ChainComplex{GrpAbFinGen}
A free resultion for $G$, ie. a chain complex terminating in
$G \to \{0\}$ that is exact.
"""
function free_resolution(G::GrpAbFinGen)
  A = free_abelian_group(ngens(G))
  R = rels(G)
  B = free_abelian_group(nrows(R))
  h_A_G = hom(A, G, gens(G))
  h_B_A = hom(B, A, [GrpAbFinGenElem(A, R[i, :]) for i=1:ngens(B)])
  Z = abelian_group(Int[1])
  set_name!(Z, "Zero")
  return chain_complex(hom(Z, B, [B[0]]), h_B_A, h_A_G, hom(G, Z, [Z[0] for i = 1:ngens(G)]))
end

mutable struct ChainComplexMap{T} <: Map{ChainComplex{T}, ChainComplex{T}, HeckeMap, ChainComplexMap}
  header::MapHeader{ChainComplex{T}, ChainComplex{T}}
  maps::Array{<:Map{<:T, <:T}, 1}
  function ChainComplexMap(C::ChainComplex{T}, D::ChainComplex{T}, A::S; check::Bool = !true) where {S <: Array{<:Map{<:T, <:T}, 1}} where {T}
    r = new{T}()
    r.header = MapHeader(C, D)
    r.maps = A
    return r
  end
end

@doc Markdown.doc"""
    hom(C::ChainComplex{T}, D::ChainComplex{T}, phi::Map{<:T, <:T}) where {T} -> ChainComplexMap
Given chain complexes $C_i: G_i \to G_{i+1}$ and $D_i: H_i \to H_{i+1}$
as well as a map $\phi = \phi_n: G_n \to H_n$, lift $\phi$ to
the entire complex: $\phi_i: G_i \to H_i$ s.th. all squares commute.
"""
function hom(C::ChainComplex{T}, D::ChainComplex{T}, phi::Map{<:T, <:T}) where {T}
  @assert length(C) == length(D)
  @assert domain(C.maps[end]) == domain(phi)
  @assert domain(D.maps[end]) == codomain(phi)

  h = [phi]
  for i=length(C)-1:-1:1
    push!(h, lift(C.maps[i]*h[end], D.maps[i]))
  end
  return ChainComplexMap(C, D, reverse(h))
end

@doc Markdown.doc"""
    hom(C::ChainComplex{T}, G::T) -> ChainComplex{T}
Given a complex $A_i: G_i \to G_{i+1}$ and a module $G$,
compute the derived complex $\hom(G_i, G)$.
"""
function hom(C::ChainComplex{T}, G::T) where {T}
  A = map_type(C)[]
  H = [hom(domain(C.maps[1]), G)]
  H = vcat(H, [hom(codomain(f), G) for f = C.maps])

  R = map_type(C)[]
  for i=1:length(C)
    A = H[i+1][1] # hom(C_i+1, G)
    B = H[i][1]   # hom(C_i  , G)
    #need map from A -> B
    #   C.maps[i] : E -> D
    D = codomain(C.maps[i])
    E = domain(C.maps[i])
    #  H[2][i+1]: A -> Hom(D, G)
    #  H[2][i]  : B -> hom(E, G)
    g = elem_type(obj_type(C))[]
    for h = gens(A)
      phi = H[i+1][2](h) # D -> G
      psi = C.maps[i] * phi
      push!(g, preimage(H[i][2], psi))
    end
    push!(R, hom(A, B, g))
  end
  return ChainComplex(reverse(R))
end

@doc Markdown.doc"""
    hom(C::ChainComplex{T}, G::T) -> ChainComplex{T}
Given a complex $A_i: G_i \to G_{i+1}$ and a module $G$,
compute the derived complex $\hom(G, G_i)$.
"""
function hom(G::T, C::ChainComplex{T}) where {T}

  A = map_type(C)[]
  H = [hom(G, domain(C.maps[1]))]
  H = vcat(H, [hom(G, codomain(f)) for f = C.maps])

  R = map_type(C)[]
  for i=1:length(C)
    A = H[i+1][1] # hom(G, C_i+1)
    B = H[i][1]   # hom(G, C_i)
    #need map from A -> B
    #   C.maps[i] : E -> D
    D = codomain(C.maps[i])
    E = domain(C.maps[i])
    #  H[2][i+1]: A -> Hom(G, D)
    #  H[2][i]  : B -> hom(G, E)
    g = elem_type(obj_type(C))[]
    for h = gens(B)
      phi = H[i][2](h) # G -> E
      psi = phi * C.maps[i]
      push!(g, preimage(H[i+1][2], psi))
    end
    push!(R, hom(B, A, g))
  end
  return ChainComplex(R)
end

@doc Markdown.doc"""
    homology(C::ChainComplex{GrpAbFinGen}) -> Array{GrpAbFinGen, 1}
Given a complex $A_i: G_i \to G_{i+1}$,
compute the homology, ie. the modules $H_i = \Kern A_{i+1}/\Im A_i$
"""
function homology(C::ChainComplex)
  H = obj_type(C)[]
  for i=1:length(C)-1
    push!(H, quo(kernel(C.maps[i+1])[1], image(C.maps[i])[1])[1])
  end
  return H
end

function snake_lemma(C::ChainComplex{T}, D::ChainComplex{T}, A::Array{<:Map{T, T}, 1}) where {T}
  @assert length(C) == length(D) == 3
  @assert length(A) == 3
  @assert domain(A[1]) == obj(C,0) && codomain(A[1]) == obj(D, 1)
  @assert domain(A[2]) == obj(C,1) && codomain(A[2]) == obj(D, 2)
  @assert domain(A[3]) == obj(C,2) && codomain(A[3]) == obj(D, 3)

  ka, mka = kernel(A[1])
  kb, mkb = kernel(A[2])
  kc, mkc = kernel(A[3])
  ca, mca = cokernel(A[1])
  cb, mcb = cokernel(A[2])
  cc, mcc = cokernel(A[3])

  MT = map_type(C)
  res = MT(C)[]
  push!(res, MT(mka * map(C, 1) * inv(mkb)))
  push!(res, MT(mkb * map(C, 2) * inv(mkc)))
  #now the snake
  push!(res, MT(mkc * inv(map(C, 2)) * A[2] * inv(map(D, 2)) * mca))
  #and the boring rest
  push!(res, MT(inv(mca) * map(D, 2) * mcb))
  push!(res, MT(inv(mcb) * map(D, 3) * mcc))
  return chain_complex(res...)
end

@doc Markdown.doc"""
    tensor_product(C::ChainComplex{T}, G::T) -> ChainComplex{T}
Given a complex $A_i: G_i \to G_{i+1}$ and a module $G$,
compute the derived complex $G_i \otimes G$.
"""
function tensor_product(C::ChainComplex{T}, G::T) where {T}
  A = map_type(C)[]
  H = [tensor_product(domain(C.maps[1]), G, task  = :none)]
  H = vcat(H, [tensor_product(codomain(f), G, task = :none) for f = C.maps])

  R = map_type(C)[]
  I = identity_map(G)
  for i = 1:length(C)
    push!(R, hom(H[i], H[i+1], [C.maps[i], I]))
  end
  return chain_complex(R)
end


text:
# these functions return vectors of states, actions and observations, ordered according to state_index, action_index, etc.

"""
    ordered_actions(mdp)

Return an `AbstractVector` of actions ordered according to `action_index(mdp, a)`.

`ordered_actions(mdp)` will always return an `AbstractVector{A}` `v` containing all of the actions in `actions(mdp)` in the order such that `action_index(mdp, v[i]) == i`. You may wish to override this for your problem for efficiency.
"""
ordered_actions(mdp::Union{MDP,POMDP}) = ordered_vector(action_type(typeof(mdp)), a->action_index(mdp,a), iterator(actions(mdp)), n_actions(mdp), "action")

"""
    ordered_states(mdp)

Return an `AbstractVector` of states ordered according to `state_index(mdp, a)`.

`ordered_states(mdp)` will always return a `AbstractVector{A}` `v` containing all of the states in `states(mdp)` in the order such that `state_index(mdp, v[i]) == i`. You may wish to override this for your problem for efficiency.
"""
ordered_states(mdp::Union{MDP,POMDP}) = ordered_vector(state_type(typeof(mdp)), s->state_index(mdp,s), iterator(states(mdp)), n_states(mdp), "state")

"""
    ordered_observations(pomdp)

Return an `AbstractVector` of observations ordered according to `obs_index(pomdp, a)`.

`ordered_observations(mdp)` will always return a `AbstractVector{A}` `v` containing all of the observations in `observations(pomdp)` in the order such that `obs_index(pomdp, v[i]) == i`. You may wish to override this for your problem for efficiency.
"""
ordered_observations(pomdp::POMDP) = ordered_vector(obs_type(typeof(pomdp)), o->obs_index(pomdp,o), iterator(observations(pomdp)), n_observations(pomdp), "observation")

function ordered_vector(T::Type, index::Function, iterator, len, singular, plural=singular*"s")
    a = Array{T}(len)
    gotten = falses(len)
    for x in iterator
        id = index(x)
        if id > len || id < 1
            error("""
                  $(singular)_index(...) returned an index that was out of bounds for $singular $x.

                  index was $id.

                  n_$plural(...) was $len.
                  """)
        end
        a[id] = x
        gotten[id] = true
    end
    if !all(gotten)
        missing = find(.!gotten)
        warn("""
             Problem creating an ordered vector of $plural in ordered_$plural(...). There is likely a mistake in $(singular)_index(...) or n_$plural(...).

             n_$plural(...) was $len.

             $plural corresponding to the following indices were missing from iterator($plural(...)): $missing
             """)
    end
    return a
end

@POMDP_require ordered_actions(mdp::Union{MDP,POMDP}) begin
    P = typeof(mdp)
    @req action_index(::P, ::action_type(P))
    @req n_actions(::P)
    @req actions(::P)
    as = actions(mdp)
    @req iterator(::typeof(as))
end

@POMDP_require ordered_states(mdp::Union{MDP,POMDP}) begin
    P = typeof(mdp)
    @req state_index(::P, ::state_type(P))
    @req n_states(::P)
    @req states(::P)
    as = states(mdp)
    @req iterator(::typeof(as))
end

@POMDP_require ordered_observations(mdp::Union{MDP,POMDP}) begin
    P = typeof(mdp)
    @req obs_index(::P, ::obs_type(P))
    @req n_observations(::P)
    @req observations(::P)
    as = observations(mdp)
    @req iterator(::typeof(as))
end