using Test, OffsetArrays
using VanFoFy
#=
ω1 = complex(1.0)
ω3 = exp(1im)
wei = WeierstrassData(ω1=ω1, ω3=ω3)

@testset "Weierstrass construction" begin
    @test wei.g2 ≈ -1.0321084791761130361 - 2.2551981702100225569im atol=1e-13
    @test wei.g3 ≈ 13.8271498100001721809 + 1.9710124059856877210im atol=1e-13
    @test wei.η1 ≈ 0.91777507848457062778 + 0.02427726872609958786im atol=1e-13
end

@testset "Weierstrass computation" begin
    z = 0.3ω1 + 0.7ω3
    output = OffsetArray{ComplexF64}(undef, -2:10)
    weierstrass_normalized!(z=z, norm_factor=1.0, w=wei,
                        upper_term=10, output=output)

    ref_array = OffsetVector(
              [0.10922980957434628990 - 0.71842740909634417168im,
              -0.8598534582241565 + 0.7026363517078360im,
               -0.0639302522525613 - 1.2338575450881468im,
               0.315957874556992 + 1.608024389830146im,
               -1.432308324486529 + 0.345694829053258im,
               1.963873759456078 - 0.492648412322271im,
               -0.683501620549705 + 1.453552820019121im,
               -0.898424448464008 - 2.273035157032416im,
               1.667509374595122 + 1.332241142926466im,
               -2.633101865120501 + 0.633539568271906im,
               2.077985541274223 - 2.171576392262631im,
               0.162780226883512 + 3.098964550714613im,
               -2.479278381517072 - 2.856788007534111im],
               -2:10)
    for n in 1:10  ref_array[n] *= n+1  end

    @test output[-2 : 10] ≈ ref_array[-2 : 10] atol=1e-10
    #####
    r = 0.23
    for n in -1:10  ref_array[n] *= r^(n+2)  end
    weierstrass_normalized!(z=z, norm_factor=r, w=wei,
                        upper_term=10, output=output)
    @test output[-2 : 10] ≈ ref_array[-2 : 10] atol=1e-10
end
=#


function ajusteArea(V, areaSup)

    fmin = 90000
    fmax = 110000

    cond = true
    while cond
        factor = (fmin + fmax)/2
        V_ = factor .* copy(V)
        areaCalc = poly2D.polyArea(V_)
        if abs(areaCalc - areaSup) <= 1
            cond = false
            return factor
        elseif areaCalc <= areaSup
            fmin = factor
            factor = (factor + fmax)/2
        else
            fmax = factor
            factor = (factor + fmin)/2
        end
    end

    return factor
end


# Turing.jl version of model at https://github.com/stan-dev/blob/master/misc/cluster/lda/lda.stan

# data {
#   int<lower=2> K;               // num topics
#   int<lower=2> V;               // num words
#   int<lower=1> M;               // num docs
#   int<lower=1> N;               // total word instances
#   int<lower=1,upper=V> w[N];    // word n
#   int<lower=1,upper=M> doc[N];  // doc ID for word n
#   vector<lower=0>[K] alpha;     // topic prior
#   vector<lower=0>[V] beta;      // word prior
# }
# parameters {
#   simplex[K] theta[M];   // topic dist for doc m
#   simplex[V] phi[K];     // word dist for topic k
# }
# model {
#   for (m in 1:M)
#     theta[m] ~ dirichlet(alpha);  // prior
#   for (k in 1:K)
#     phi[k] ~ dirichlet(beta);     // prior
#   for (n in 1:N) {
#     real gamma[K];
#     for (k in 1:K)
#       gamma[k] <- log(theta[doc[n],k]) + log(phi[k,w[n]]);
#     increment_log_prob(log_sum_exp(gamma));  // likelihood
#   }
# }

# V = ldastandata[1]["V"]
# K = ldastandata[1]["K"]
# M = ldastandata[1]["M"]
# theta = Vector{Vector{Real}}(undef, M)
# for m = 1:M
#   theta[m] = Vector{Real}(undef, K)
# end
# phi = Vector{Vector{Real}}(undef, K)
# for k = 1:K
#   phi[k] = Vector{Real}(undef, V)
# end

# phi_dot_theta = [Vector{Real}(undef, V) for m=1:M]

@model ldamodel(K, V, M, N, w, doc, beta, alpha) = begin
  theta = TArray{Vector{Real}, 1}(M)
  @simd for m = 1:M
    @inbounds theta[m] ~ Dirichlet(alpha)
  end

  phi = TArray{Vector{Real}, 1}(K)
  @simd for k = 1:K
    @inbounds phi[k] ~ Dirichlet(beta)
  end

  # z = tzeros(Int, N)
  # for n = 1:N
  #   z[n] ~ Categorical(theta[doc[n]])
  # end

  # phi_dot_theta = [log.([dot(map(p -> p[v], phi), theta[m]) for v = 1:V]) for m = 1:M]

  # @simd for m=1:M
  #   @inbounds phi_dot_theta[m] = log.([dot(map(p -> p[v], phi), theta[m]) for v = 1:V])
  # end

  #for n = 1:N
  #  # phi_dot_theta = [dot(map(p -> p[v], phi), theta[doc[n]]) for v = 1:V]
  #  # w[n] ~ Categorical(phi_dot_theta)
  #  Turing.acclogp!(vi, phi_dot_theta[doc[n]][w[n]])
  #end

  # lp = mapreduce(n->phi_dot_theta[doc[n]][w[n]], +, 1:N)

  # lp = phi_dot_theta[doc[1]][w[1]]
  # @simd for n = 2:N
  #   @inbounds lp += phi_dot_theta[doc[n]][w[n]]
  # end

  phi_dot_theta = log.(hcat(phi...) * hcat(theta...))
  _lp += mapreduce(n->phi_dot_theta[w[n], doc[n]], +, 1:N)

  # Turing.acclogp!(vi, lp)
end

@model ldamodel_vec(K, V, M, N, w, doc, beta, alpha) = begin
  theta = Matrix{Real}(undef, K, M)
  theta ~ [Dirichlet(alpha)]

  phi = Matrix{Real}(undef, V, K)
  phi ~ [Dirichlet(beta)]

  phi_dot_theta = log.(phi * theta)
  #for n = 1:N
  #  Turing.acclogp!(vi, phi_dot_theta[w[n], doc[n]])
  #end
  _lp += mapreduce(n->phi_dot_theta[w[n], doc[n]], +, 1:N)

end


# NMF (Non-negative Matrix Factorization)
# Julia実装版

module NMF

export nmf_euc, nmf_kl, nmf_is

# EUC基準NMF処理（通常版）
# V     分解対象の行列
# k     因子分解数
# iter  イテレーション回数
function nmf_euc(V::AbstractMatrix,k::Int64=3;iter::Int64=100)
    # 変数チェック処理
    @assert countnz(V.< 0) == 0 "Input Matrix V is not non-negative"
    @assert k > 0 "Factor Num : k must be greater than 0"
    @assert iter > 0 "Iteration Num : iter must be greater than 0"

    # W,Hを指定の形で乱数作成
    W = rand(size(V)[1],k)
    H = rand(k,size(V)[2])

    # 反復乗法更新
    for i in 1:iter
        W = W .* (V * H') ./ (W * H * H')
        H = H .* (W' * V) ./ (W' * W * H)
    end
    return W,H
end

# EUC基準NMF処理（テンプレートあり版）
# V     分解対象の行列
# W     テンプレート辞書行列
# iter  イテレーション回数
function nmf_euc(V::AbstractMatrix,W::AbstractMatrix,;iter::Int64=100)
    #変数チェック処理
    @assert countnz(V.< 0) == 0 "Input Matrix V is not non-negative."
    @assert countnz(W.< 0) == 0 "Input Matrix W is not non-negative."
    @assert iter > 0 "Iteration Num : iter must be greater than 0."
    @assert size(V,1) == size(W,1) "Dimensions of V:$(size(V)) and W:$(size(W)) do not match. Align the first dimention."
    @assert size(W,2) > 0 "Factor Num : size(W,2) is greater than 0."

    # 指定の形でHを作成
    H = rand(size(W,2),size(V)[2])

    # H だけを更新
    for i in 1:iter
        H = H .* (W' * V) ./ (W' * W * H)
    end
    return H
end

# KL基準NMF処理（通常版）
# V     分解対象の行列
# k     因子分解数
# iter  イテレーション回数
function nmf_kl(V::AbstractMatrix,k::Int64=3;iter::Int64=100)
    #変数チェック処理
    @assert countnz(V.< 0) == 0 "Input Matrix V is not non-negative"
    @assert k > 0 "Factor Num : k must be greater than 0"
    @assert iter > 0 "Iteration Num : iter must be greater than 0"

    # W,Hを指定の形で乱数作成
    W = rand(size(V)[1],k)
    H = rand(k,size(V)[2])

    # 反復乗法更新
    for i in 1:iter
        W = W .* ((V ./ (W*H)) * H') ./ sum(H',1)
        H = H .* (W' * (V ./ (W*H))) ./ sum(W',2)
    end
    return W,H
end

# KL基準NMF処理（テンプレートあり版）
# V     分解対象の行列
# W     テンプレート辞書行列
# iter  イテレーション回数
function nmf_kl(V::AbstractMatrix,W::AbstractMatrix;iter::Int64=100)
    #変数チェック処理
    @assert countnz(V.< 0) == 0 "Input Matrix V is not non-negative."
    @assert countnz(W.< 0) == 0 "Input Matrix W is not non-negative."
    @assert iter > 0 "Iteration Num : iter must be greater than 0."
    @assert size(V,1) == size(W,1) "Dimensions of V:$(size(V)) and W:$(size(W)) do not match. Align the first dimention."
    @assert size(W,2) > 0 "Factor Num : size(W,2) is greater than 0."

    # 指定の形でHを作成
    H = rand(size(W,2),size(V)[2])

    # H だけを更新
    for i in 1:iter
        H = H .* (W' * (V ./ (W*H))) ./ sum(W',2)
    end
    return H
end

# IS基準NMF処理（通常版）
# V     分解対象の行列
# k     因子分解数
# iter  イテレーション回数
function nmf_is(V::AbstractMatrix,k::Int64=3;iter::Int64=100)
    #変数チェック処理
    @assert countnz(V.< 0) == 0 "Input Matrix V is not non-negative"
    @assert k > 0 "Factor Num : k must be greater than 0"
    @assert iter > 0 "Iteration Num : iter must be greater than 0"

    # W,Hを指定の形で乱数作成
    W = rand(size(V)[1],k)
    H = rand(k,size(V)[2])

    # 反復乗法更新
    for i in 1:iter
        W = W .* sqrt.(((V ./ (W*H))*(H./sum(W*H,1))') ./(sum((H./sum(W*H,1))',1)))
        H = H.*sqrt.(((W ./ sum(W*H,2))' * (V./(W*H)) ./ (sum((W ./ sum(W*H,2))',2))))
    end
    return W,H
end

# IS基準NMF処理（テンプレートあり版）
# V     分解対象の行列
# W     テンプレート辞書行列
# iter  イテレーション回数
function nmf_is(V::AbstractMatrix,W::AbstractMatrix;iter::Int64=100)
    #変数チェック処理
    @assert countnz(V.< 0) == 0 "Input Matrix V is not non-negative."
    @assert countnz(W.< 0) == 0 "Input Matrix W is not non-negative."
    @assert iter > 0 "Iteration Num : iter must be greater than 0."
    @assert size(V,1) == size(W,1) "Dimensions of V:$(size(V)) and W:$(size(W)) do not match. Align the first dimention."
    @assert size(W,2) > 0 "Factor Num : size(W,2) is greater than 0."

    # 指定の形でHを作成
    H = rand(size(W,2),size(V)[2])

    # H だけを更新
    for i in 1:iter
        H = H.*sqrt.(((W ./ sum(W*H,2))' * (V./(W*H)) ./ (sum((W ./ sum(W*H,2))',2))))
    end
    return H
end

end


using DelimitedFiles: readdlm

read_data(fn::AbstractString) = vec(readdlm(fn, skipstart=1))
read_cpts(fn::AbstractString) = vec([0; readdlm(fn, Int, skipstart=1)])

dir = "Test_Files"

@testset "Test Normal mean change..." begin
    data = read_data("$(dir)/Normal_change_in_mean_data")

    r_cpts = read_cpts("$(dir)/Normal_change_in_mean_cpts")
    cpts, cost = @PELT data Normal(:?, 1.0) 2*log(length(data))

    @test r_cpts == cpts
end

# Normal var change
@testset "Test Normal variance change..." begin
    data = read_data("$(dir)/Normal_change_in_var_data")

    r_cpts = read_cpts("$(dir)/Normal_change_in_var_cpts")
    cpts, cost = @PELT data Normal(1.0, :?) 2*log(length(data))

    @test r_cpts == cpts
end

# Normal meanvar change
@testset "Test Normal mean and variance chaange..." begin
    data = read_data("$(dir)/Normal_change_in_meanvar_data")

    r_cpts = read_cpts("$(dir)/Normal_change_in_meanvar_cpts")
    cpts, cost = @PELT data Normal(:?, :?) 4*log(length(data))

    @test r_cpts == cpts
end

# Poisson change
@testset "Test Poisson change..." begin
    data = read_data("$(dir)/Poisson_data")

    r_cpts = read_cpts("$(dir)/Poisson_cpts")
    cpts, cost = @PELT data Poisson(:?) 2*log(length(data))

    @test r_cpts == cpts
end

# Exponential change
@testset "Test Exponential  change..." begin
    data = read_data("$(dir)/Exponential_data")

    r_cpts = read_cpts("$(dir)/Exponential_change_in_mean_cpts")
    cpts, cost = @PELT data Exponential(:?) 2*log(length(data))

    @test r_cpts == cpts
end

# Gamma rate changes
@testset "Test Gamma rate change..." begin
    data = read_data("$(dir)/Gamma_data")

    r_cpts = read_cpts("$(dir)/Gamma_cpts")
    cpts, cost = @PELT data Gamma(1.0,:?) 2*log(length(data))

    @test r_cpts == cpts
end

@testset "Test Nonparametric change..." begin
    data = read_data("$(dir)/Nonparametric_data")

    r_cpts = read_cpts("$(dir)/Nonparametric_cpts")
    cpts, cost = @PELT data Nonparametric(10) 4*log(length(data))

    @test r_cpts == cpts
end


function constant_features_edge(s::Integer, d::Integer, mapf::MAPF)
    g = mapf.graph
    outdeg_s = outdegree(g, s)
    indeg_s = indegree(g, s)
    outdeg_d = outdegree(g, d)
    indeg_d = indegree(g, d)
    rev_edge_exists = has_edge(g, d, s)
    return Float64[outdeg_s, outdeg_d]
end

function solution_features_edge(s::Integer, d::Integer, solution::Solution, mapf::MAPF)
    g = mapf.graph
    paths_crossing_s = 0
    paths_crossing_d = 0
    paths_crossing_e = 0
    for path in solution
        for ((t1, v1), (t2, v2)) in zip(path[1:(end - 1)], path[2:end])
            paths_crossing_s += (v1 == s) + (v2 == s)
            paths_crossing_d += (v1 == d) + (v2 == d)
            paths_crossing_e += (v1 == s && v2 == d)
        end
    end
    return Float64[paths_crossing_s, paths_crossing_d]
end

function all_features_edge(s::Integer, d::Integer, solution::Solution, mapf::MAPF)
    return vcat(
        constant_features_edge(s, d, mapf), solution_features_edge(s, d, solution, mapf)
    )
end

function edges_embedding(mapf::MAPF)
    solution = independent_dijkstra(mapf)
    x = reduce(
        hcat,
        all_features_edge(src(ed), dst(ed), solution, mapf) for ed in edges(mapf.graph)
    )
    s = std(x; dims=2)
    s[isapprox.(s, 0.0)] .= 1  # columns with zero variance
    x = (x .- mean(x; dims=2)) ./ s
    return x
end


import Pkg
Pkg.activate(".")

using MCDTS
using Test
using Random
using DynamicalSystemsBase
using DelayEmbeddings
import MCDTS.L

# Check Lorenz System
Random.seed!(1234)
ds = Systems.lorenz()
data = trajectory(ds,200)
data = data[2001:2:3750,:]

w1 = DelayEmbeddings.estimate_delay(data[:,1],"mi_min")
w2 = DelayEmbeddings.estimate_delay(data[:,2],"mi_min")
w3 = DelayEmbeddings.estimate_delay(data[:,3],"mi_min")

w = maximum(hcat(w1,w2,w3))

@testset "MCDTS embedding tests" begin

    include("MCDTS_attractor_reconstruction_tests.jl")
    include("MCDTS_convergent_cross_mapping_tests.jl")
    include("MCDTS_nearest_neighbor_prediction_test.jl")

end


function test_fft_cosmx(g::OneDGrid)
    f1, f1h, f1hr, f1hr_mul, f1h_th, f1hr_th = create_testfuncs(g)
    isapprox(f1h, f1h_th, rtol=rtol_fft)
end

function test_rfft_cosmx(g::OneDGrid)
    f1, f1h, f1hr, f1hr_mul, f1h_th, f1hr_th = create_testfuncs(g)
    isapprox(f1hr, f1hr_th, rtol=rtol_fft)
end

function test_rfft_mul_cosmx(g::OneDGrid)
    f1, f1h, f1hr, f1hr_mul, f1h_th, f1hr_th = create_testfuncs(g)
    isapprox(f1hr_mul, f1hr_th, rtol=rtol_fft)
end

function test_fft_cosmxcosny(g::TwoDGrid)
    f1, f2, f1h, f2h, f1hr, f2hr, f1hr_mul, f2hr_mul, f1h_th, f1hr_th, f2h_th, f2hr_th = create_testfuncs(g)
    isapprox(f1h, f1h_th, rtol=rtol_fft)
end

function test_rfft_cosmxcosny(g::TwoDGrid)
    f1, f2, f1h, f2h, f1hr, f2hr, f1hr_mul, f2hr_mul, f1h_th, f1hr_th, f2h_th, f2hr_th = create_testfuncs(g)
    isapprox(f1hr, f1hr_th, rtol=rtol_fft)
end

function test_fft_sinmxny(g::TwoDGrid)
    f1, f2, f1h, f2h, f1hr, f2hr, f1hr_mulPk, f2hr_mul, f1h_th, f1hr_th, f2h_th, f2hr_th = create_testfuncs(g)
    isapprox(f2h, f2h_th, rtol=rtol_fft)
end

function test_rfft_sinmxny(g::TwoDGrid)
    f1, f2, f1h, f2h, f1hr, f2hr, f1hr_mul, f2hr_mul, f1h_th, f1hr_th, f2h_th, f2hr_th = create_testfuncs(g)
    isapprox(f2hr, f2hr_th, rtol=rtol_fft)
end

function test_rfft_mul_cosmxcosny(g::TwoDGrid)
    f1, f2, f1h, f2h, f1hr, f2hr, f1hr_mul, f2hr_mul, f1h_th, f1hr_th, f2h_th, f2hr_th = create_testfuncs(g)
    isapprox(f1hr_mul, f1hr_th, rtol=rtol_fft)
end

function test_rfft_mul_sinmxny(g::TwoDGrid)
    f1, f2, f1h, f2h, f1hr, f2hr, f1hr_mul, f2hr_mul, f1h_th, f1hr_th, f2h_th, f2hr_th = create_testfuncs(g)
    isapprox(f2hr_mul, f2hr_th, rtol=rtol_fft)
end


"""
    create_surfaces_sphere(cells, setup)

Create surface triangulation of [inner and] outer spheres [and ECS].
"""
function create_surfaces_sphere(cells, setup::Setup)

    # Extract parameters
    @unpack radii, centers = cells
    @unpack ncell, rmin, rmax, include_in, in_ratio, ecs_shape, ecs_ratio = setup

    include_ecs = ecs_shape != "no_ecs"

    rmean = (rmin + rmax) / 2

    create_npoint(radius) = round(Int, 200 * (radius / rmean)^2)
    if include_in
        npoint_in = create_npoint.(in_ratio .* radii)
        points_in = [
            centers[:, i] .+ in_ratio * radii[i] * create_fibonacci_sphere(npoint_in[i]) for i = 1:ncell
        ]
        facets_in = [convexhull(points_in[i])[1] for i = 1:ncell]
        regions_in = centers
        nfacet_in = size.(facets_in, 2)
    else
        npoint_in = zeros(Int, 0)
        nfacet_in = zeros(Int, 0)
        points_in = zeros(3, 0)
        facets_in = zeros(Int, 3, 0)
        regions_in = zeros(3, 0)
    end

    npoint_out = create_npoint.(radii)
    points_out =
        [centers[:, i] .+ radii[i] * create_fibonacci_sphere(npoint_out[i]) for i = 1:ncell]
    facets_out = [convexhull(points_out[i])[1] for i = 1:ncell]
    regions_out = @. centers + (1 + in_ratio) / 2 * radii * [1; 0; 0]
    nfacet_out = size.(facets_out, 2)

    if include_ecs
        if ecs_shape == "box"
            # Determine bounds of domain
            emin = minimum(hcat(points_out...), dims = 2)
            emax = maximum(hcat(points_out...), dims = 2)

            # Extend bounds by ECS gap
            @. emin = emin - ecs_ratio * rmean
            @. emax = emax + ecs_ratio * rmean

            # Define eight corners of box ECS and their corresponding edges
            points_ecs = [
                emin[1] emax[1] emax[1] emin[1] emin[1] emax[1] emax[1] emin[1]
                emin[2] emin[2] emax[2] emax[2] emin[2] emin[2] emax[2] emax[2]
                emin[3] emin[3] emin[3] emin[3] emax[3] emax[3] emax[3] emax[3]
            ]
            facets_ecs = [
                1 1 1 1 2 2 3 3 4 4 5 5
                2 3 2 6 3 7 4 8 1 5 6 7
                3 4 6 5 7 6 8 7 5 8 7 8
            ]
            regions_ecs = emin .+ ecs_ratio / 2 * rmean
            npoint_ecs = 8
            nfacet_ecs = 12
        elseif ecs_shape == "convex_hull"
            npoint_ecs = create_npoint.((1 + ecs_ratio) .* radii)
            points_ecs = [
                centers[:, i] .+
                (1 + ecs_ratio) * radii[i] * create_fibonacci_sphere(npoint_ecs[i]) for
                i = 1:ncell
            ]
            points_ecs = hcat(points_ecs...)

            facets_ecs, points_ecs = convexhull(points_ecs)
            npoint_ecs = size(points_ecs, 2)
            nfacet_ecs = size(facets_ecs, 2)
            regions_ecs = centers[:, 1] + (1 + ecs_ratio / 2) * radii[1] * [1; 0; 0]
        end
    else
        npoint_ecs = zeros(Int, 0)
        nfacet_ecs = zeros(Int, 0)
        points_ecs = zeros(3, 0)
        facets_ecs = zeros(Int, 3, 0)
        regions_ecs = zeros(3, 0)
    end

    npoint = cumsum(vcat(0, npoint_in..., npoint_out..., npoint_ecs))
    nfacet = cumsum(vcat(0, nfacet_in..., nfacet_out..., nfacet_ecs))
    points = hcat(points_in..., points_out..., points_ecs)
    facets = hcat(facets_in..., facets_out..., facets_ecs)
    facetmarkers = zeros(Int, size(facets, 2))
    for i = 1:length(npoint)-1
        inds = nfacet[i]+1:nfacet[i+1]
        facets[:, inds] .+= npoint[i]
        facetmarkers[inds] .= i
    end
    regions = [regions_in regions_out regions_ecs]

    (; points, facets, facetmarkers, regions)
end


# # Hydraulic functions

# This tutorial shows how to specify the hydraulic functions
# used in Richard's equation. In particular,
# we show how to choose the formalism for matric potential and hydraulic
# conductivity, and how to make the hydraulic conductivity account for
# the presence of ice as well as the temperature dependence of the
# viscosity of liquid water.

# ClimateMachine's Land model allows the user to pick between two hydraulics models,
# that of van Genuchten [vanGenuchten1980](@cite) or that of Brooks and Corey, see [BrooksCorey1964](@cite) or [Corey1977](@cite). The
# same model is consistently used for the matric potential
# and hydraulic conductivity.

# # Preliminary setup

# External modules
using Plots

# ClimateMachine modules
using ClimateMachine
using ClimateMachine.Land
using ClimateMachine.Land.SoilWaterParameterizations

const FT = Float32;
# # Specifying a  `hydraulics` model
# The van Genuchten model requires two free parameters, `α` and `n`.
# A third parameter, `m`, is computed from `n`. Of these, only `α` carries
# units, of inverse meters. The Brooks and Corey model also uses
# two free parameters, `ψ_b`, the magnitude of the matric potential at saturation,
#  and a constant `M`. `ψ_b` carries units of meters. The hydraulic conductivity
# requires an additional parameter, `Ksat` (m/s), which is the hydraulic conductivity
# in saturated soil. This parameter is the same between the two models for a given
# soil type, and is not stored in the `hydraulics` model, but rather as part of the
# [`SoilParamFunctions`](@ref ClimateMachine.Land.SoilParamFunctions).

# Below we show how to create two concrete examples of these hydraulics models,
# for sandy loam ([Bonan19a](@cite)). Note that the parameters chosen are a function of soil type.
vg_α = FT(7.5) # m^-1
vg_n = FT(1.89)
hydraulics = vanGenuchten{FT}(α = vg_α, n = vg_n)

ψ_sat = 0.09 # m
Mval = 0.228
hydraulics_bc = BrooksCorey{FT}(ψb = ψ_sat, m = Mval);
# # Matric Potential
# The matric potential `ψ` represents how much water clings to soil. Drier soil
# holds onto water more tightly, impeding the flow of water. As soil becomes
# wetter, the matric potential decreases in magnitude, enabling water flow.

# The van Genuchten expression for matric potential is
# ``
# ψ = -\frac{1}{α} S_l^{-1/(nm)}\times (1-S_l^{1/m})^{1/n},
# ``

# and the Brooks and Corey expression is
# ``
# ψ = -ψ_b S_l^{-M}.
# ``

# Here `S_l` is the effective saturation of liquid water, `θ_l/ν`, where `ν` is
# porosity of the soil. We neglect the residual pore space in the CliMA model.

# In the CliMA code, we use [multiple dispatch](https://en.wikipedia.org/wiki/Multiple_dispatch).
# With multiple dispatch, a function can have many
# ways of executing (called methods), depending on the *type* of the
# variables passed in. A simple example of multiple dispatch is the division operation.
# Integer division takes two numbers as input, and returns an integer - ignoring the decimal.
# Float division takes two numbers as input, and returns a floating point number, including the decimal.
# In Julia, we might write these as:

# ```julia
# function division(a::Int, b::Int)
#      return floor(Int, a/b)
# end
# ```
# ```julia
# function division(a::Float64, b::Float64)
#      return a/b
# end
# ```

# We can see that `division` is now a function with two methods.

# ```julia
# julia> division
# division (generic function with 2 methods)
# ```

# Now, using the same function signature, we can carry out integer
# division or floating point division, depending on the types of the
# arguments:

# ```julia
# julia> division(1,2)
# 0
#
# julia> division(1.0,2.0)
# 0.5
# ```

# Here is a more pertinent example:
# `hydraulics` is of type `vanGenuchten{Float32}` based on our choice of `FT`:

# ```julia
# julia> typeof(hydraulics)
# vanGenuchten{Float32}
# ```

# but meanwhile,

# ```julia
# julia> typeof(hydraulics_bc)
# BrooksCorey{Float32}
# ```

# The function `matric_potential` will execute different methods
# depending on if we pass a hydraulics model of type `vanGenuchten` or
# `BrooksCorey`. In both cases, it will return the correct value
# for `ψ`.

# Let's plot the matric potential as a function of the effective saturation `S_l = θ_l/ν`,
# which can range from zero to one.
S_l = FT.(0.01:0.01:0.99)
ψ = matric_potential.(Ref(hydraulics), S_l)
ψ_bc = matric_potential.(Ref(hydraulics_bc), S_l)
plot(
    S_l,
    log10.(-ψ),
    xlabel = "effective saturation",
    ylabel = "Log10(|ψ|)",
    label = "van Genuchten",
)
plot!(S_l, log10.(-ψ_bc), label = "Brooks and Corey")
savefig("bc_vg_matric_potential.png")
# ![](bc_vg_matric_potential.png)

# The steep slope in
# `ψ` near saturated and completely dry soil are part of the reason
# why Richard's equation is such a challenging numerical problem.

# # Hydraulic conductivity
# The hydraulic conductivity is a more complex function than the matric potential,
# as it depends on the temperature of the water, the volumetric ice fraction, and
# the volumetric liquid water fraction. It also depends on the `hydraulics` model
# chosen.

# We represent the hydraulic conductivity `K` as the product of four factors:
# `Ksat`, an `impedance_factor` (which accounts for the effect of ice on conductivity)
# a `viscosity_factor` (which accounts for the effect of temperature on the
# viscosity of liquid water, and how that in turn affects conductivity)
# and a `moisture_factor` (which accounts for the effect of liquid water, and is determined by the `hydraulics` model).

#  Let's start with ice and temperature independence, but moisture dependence.
# Below
# we choose additional parameters, consistent with the `hydraulics` parameters
# for sandy loam ([Bonan19a](@cite)).
ν = FT(0.41)
Ksat = FT(4.42 / (3600 * 100))
moisture_choice = MoistureDependent{FT}()
viscosity_choice = ConstantViscosity{FT}()
impedance_choice = NoImpedance{FT}();

# We are going to calculate `K = Ksat × viscosity_factor × impedance_factor × moisture_factor`.
# In the code, each of these factors is a function with multiple methods, except for `Ksat`.
#  Our function `hydraulic_conductivity` calls each of these functions in turn,
# and these functions use multiple dispatch to provide the correct value for `K`.

# Like we defined new type
# classes for `vanGenuchten{FT}` and `BrooksCorey{FT}`, we also created new type classes
# for the impedance choice, the viscosity choice, and the moisture choice.
#  For example, the function
# called `viscosity_factor`, when passed an object of type `ConstantViscosity{FT}`, executes a method that always
# returns 1. The same is true for the function `impedance_factor`, using the type
# `NoImpedance{FT}`, and for `moisture_factor`, using the type `MoistureIndependent{FT}`.

# In the case where the `moisture_factor = MoistureDependent{FT}()`, either the van
# Genuchten or Brooks and Corey expression is used based on the type of `hydraulics` model
# passed.

# The `moisture_factor` for the van Genuchten model is (denoting it as ``K_m``)

# ``
#  K_m = \sqrt{S_l}[1-(1-S_l^{1/m})^m]^2,
# ``

# for ``S_l < 1``,

# and for the Brooks and Corey model it is

# ``
# K_m = S_l^{2M+3},
# ``

# also for ``S_l<1``. When ``S_l\geq 1``, ``K_m = 1`` for each model.

# One side effect of this flexibility is that `hydraulic_conductivity`
# requires all the arguments it could possibly need passed to it, which is
# why here we must supply a value `T` and `θ_i`, even though they are not used
# in this particular example.
T = FT(0.0)
θ_i = FT(0.0)

K =
    Ksat .*
    hydraulic_conductivity.(
        Ref(impedance_choice),
        Ref(viscosity_choice),
        Ref(moisture_choice),
        Ref(hydraulics),
        Ref(θ_i),
        Ref(ν),
        Ref(T),
        S_l,
    );
# Let's also compute `K` when we include the effects of temperature
# and ice on the hydraulic conductivity.
# In the cases where a `TemperatureDependentViscosity{FT}` or
# `IceImpedance{FT}` type is passed, the correct factors are calculated,
# based on the temperature `T` and volumetric ice fraction `θ_i`.
# In these cases, the `viscosity_factor`, denoted here as ``K_v``,
# evaluates as:

# ``
# K_v = e^{\gamma(T-T_{ref})},
# ``

# where ``\gamma = 0.0264 \mbox{K}^{-1}`` is an empirical factor,
# and ``T_{ref} = 288`` K,
# and the `impedance_factor`, denoted ``K_i``, evaluates as:

# ``
# K_i = 10^{-\Omega f_i}
# ``

# where ``\Omega = 7`` is an empirical factor and
# ``f_i`` is the ratio of the volumetric
# ice fraction to total volumetric water fraction  ([Lundin1990](@cite)).

viscosity_choice_T = TemperatureDependentViscosity{FT}()
T = FT(300.0)
K_T =
    Ksat .*
    hydraulic_conductivity.(
        Ref(impedance_choice),
        Ref(viscosity_choice_T),
        Ref(moisture_choice),
        Ref(hydraulics),
        Ref(θ_i),
        Ref(ν),
        Ref(T),
        S_l,
    )
ice_impedance_I = IceImpedance{FT}()
θ_i = FT(0.1)
S_i = θ_i / ν;
# The total volumetric water fraction cannot
# exceed unity, so the effective liquid water saturation
# should have a max of 1-S_i.
S_l_accounting_for_ice = FT.(0.01:0.01:(0.99 - S_i))
K_i =
    Ksat .*
    hydraulic_conductivity.(
        Ref(ice_impedance_I),
        Ref(viscosity_choice),
        Ref(moisture_choice),
        Ref(hydraulics),
        Ref(θ_i),
        Ref(ν),
        Ref(T),
        S_l_accounting_for_ice,
    )
plot(
    S_l,
    log10.(K),
    xlabel = "total effective saturation, (θ_i+θ_l)/ν",
    ylabel = "Log10(K)",
    label = "Base case",
    legend = :bottomright,
)
plot!(S_l, log10.(K_T), label = "Temperature Dependent Viscosity; no ice")
plot!(
    S_l_accounting_for_ice .+ S_i,
    log10.(K_i),
    label = "Ice Impedance; S_i=0.24",
)
savefig("T_ice_K.png")
# ![](T_ice_K.png)

# If the user is not considering phase transitions
# and does not add in Freeze/Thaw source terms, the default is for zero
# ice in the model, for all time and space. In this case the ice impedance
# factor evaluates to 1 regardless of which type is passed.

# We can also look and see how the Brooks and Corey moisture factor differs from the
# van Genuchten moisture factor by changing the `hydraulics` model passed:
T = FT(0.0)
θ_i = FT(0.0)

K_bc =
    Ksat .*
    hydraulic_conductivity.(
        Ref(impedance_choice),
        Ref(viscosity_choice),
        Ref(moisture_choice),
        Ref(hydraulics_bc),
        Ref(θ_i),
        Ref(ν),
        Ref(T),
        S_l,
    )
plot(
    S_l,
    log10.(K),
    xlabel = "effective saturation",
    ylabel = "Log10(K)",
    label = "van Genuchten",
)
plot!(
    S_l,
    log10.(K_bc),
    xlabel = "effective saturation",
    ylabel = "Log10(K)",
    label = "Brooks and Corey",
)
savefig("bc_vg_k.png")
# ![](bc_vg_k.png)
# # Other features
# The user also has the choice of making the conductivity constant by choosing
# `MoistureIndependent{FT}()` along with `ConstantViscosity{FT}()` and
# `NoImpedance{FT}()` . This is useful for debugging!
no_moisture_dependence = MoistureIndependent{FT}()
K_constant =
    Ksat .*
    hydraulic_conductivity.(
        Ref(impedance_choice),
        Ref(viscosity_choice),
        Ref(no_moisture_dependence),
        Ref(hydraulics),
        Ref(θ_i),
        Ref(ν),
        Ref(T),
        S_l,
    );
# ```julia
# julia> unique(K_constant)
# 1-element Array{Float32,1}:
#  1.2277777f-5
# ```

# Note that choosing this option does not mean the matric potential
# is constant, as a hydraulics model is still required and employed.

# And, lastly, you might be wondering why we left `Ksat` out of the function
# for `hydraulic_conductivity`. It turns out it is also useful for debugging
# to be able to turn off the flow of water, by setting `Ksat = 0`.

# # References
# - [vanGenuchten1980](@cite)
# - [BrooksCorey1964](@cite)
# - [Corey1977](@cite)
# - [Lundin1990](@cite)
# - [Bonan19a](@cite)


module CategoricalAlgebra

include("ShapeDiagrams.jl")
include("Matrices.jl")
include("FinSets.jl")
include("FinRelations.jl")
include("Permutations.jl")
include("CSets.jl")
include("Graphs.jl")

end


"""
    AbstractPowerRepresentation

An abstract representation type of points and tangent vectors on a power manifold.
"""
abstract type AbstractPowerRepresentation end

"""
    ArrayPowerRepresentation

Representation of points and tangent vectors on a power manifold using multidimensional
arrays where first dimensions are equal to [`representation_size`](@ref) of the
wrapped manifold and the following ones are equal to the number of elements in each
direction.

[`Torus`](@ref) uses this representation.
"""
struct ArrayPowerRepresentation <: AbstractPowerRepresentation end

"""
    NestedPowerRepresentation

Representation of points and tangent vectors on a power manifold using arrays
of size equal to `TSize` of a [`PowerManifold`](@ref).
Each element of such array stores a single point or tangent vector.

[`GraphManifold`](@ref) uses this representation.
"""
struct NestedPowerRepresentation <: AbstractPowerRepresentation end

@doc raw"""
    AbstractPowerManifold{M,TPR} <: Manifold

An abstract [`Manifold`](@ref) to represent manifolds that are build as powers
of another [`Manifold`](@ref) `M` with representation type `TPR`, a subtype of
[`AbstractPowerRepresentation`](@ref).
"""
abstract type AbstractPowerManifold{M<:Manifold,TPR<:AbstractPowerRepresentation} <:
              Manifold end

@doc raw"""
    PowerManifold{TM<:Manifold, TSize<:Tuple, TPR<:AbstractPowerRepresentation} <: AbstractPowerManifold{TM}

The power manifold $\mathcal M^{n_1× n_2 × … × n_d}$ with power geometry
 `TSize` statically defines the number of elements along each axis.

For example, a manifold-valued time series would be represented by a power manifold with
$d$ equal to 1 and $n_1$ equal to the number of samples. A manifold-valued image
(for example in diffusion tensor imaging) would be represented by a two-axis power
manifold ($d=2$) with $n_1$ and $n_2$ equal to width and height of the image.

While the size of the manifold is static, points on the power manifold
would not be represented by statically-sized arrays. Operations on small
power manifolds might be faster if they are represented as [`ProductManifold`](@ref).

# Constructor

    PowerManifold(M, N_1, N_2, ..., N_n)
    PowerManifold(M, NestedPowerRepresentation(), N_1, N_2, ..., N_n)
    M^(N_1,N_2, ..., N_n)

Generate the power manifold $M^{N_1 × N_2 × … × N_n}$.
By default, the [`ArrayPowerRepresentation`](@ref) of points
and tangent vectors is used, although a different one, for example
[`NestedPowerRepresentation`](@ref), can be given as the second argument to the
constructor.
"""
struct PowerManifold{TM<:Manifold,TSize,TPR<:AbstractPowerRepresentation} <:
       AbstractPowerManifold{TM,TPR}
    manifold::TM
end

function PowerManifold(M::Manifold, size::Int...)
    return PowerManifold{typeof(M),Tuple{size...},ArrayPowerRepresentation}(M)
end
function PowerManifold(
    M::Manifold,
    ::TPR,
    size::Int...,
) where {TPR<:AbstractPowerRepresentation}
    PowerManifold{typeof(M),Tuple{size...},TPR}(M)
end

@doc raw"""
    PowerMetric <: Metric

Represent the [`Metric`](@ref) on an [`AbstractPowerManifold`](@ref), i.e. the inner
product on the tangent space is the sum of the inner product of each elements
tangent space of the power manifold.
"""
struct PowerMetric <: Metric end

"""
    PowerRetraction(retraction::AbstractRetractionMethod)

Power retraction based on `retraction`. Works on [`AbstractPowerManifold`](@ref)s.
"""
struct PowerRetraction{TR<:AbstractRetractionMethod} <: AbstractRetractionMethod
    retraction::TR
end

"""
    InversePowerRetraction(inverse_retractions::AbstractInverseRetractionMethod...)

Power inverse retraction of `inverse_retractions`. Works on [`AbstractPowerManifold`](@ref)s.
"""
struct InversePowerRetraction{TR<:AbstractInverseRetractionMethod} <:
       AbstractInverseRetractionMethod
    inverse_retraction::TR
end

"""
    PowerPointDistribution(M::AbstractPowerManifold, distribution)

Power distribution on manifold `M`, based on `distribution`.
"""
struct PowerPointDistribution{TM<:AbstractPowerManifold,TD<:MPointDistribution,TX} <:
       MPointDistribution{TM}
    manifold::TM
    distribution::TD
    point::TX
end

"""
    PowerFVectorDistribution([type::VectorBundleFibers], [x], distr)

Generates a random vector at a `point` from vector space (a fiber of a tangent
bundle) of type `type` using the power distribution of `distr`.

Vector space type and `point` can be automatically inferred from distribution `distr`.
"""
struct PowerFVectorDistribution{
    TSpace<:VectorBundleFibers{<:VectorSpaceType,<:AbstractPowerManifold},
    TD<:FVectorDistribution,
    TX,
} <: FVectorDistribution{TSpace,TX}
    type::TSpace
    point::TX
    distribution::TD
end

"""
    PowerBasisData{TB<:AbstractArray}

Data storage for an array of basis data.
"""
struct PowerBasisData{TB<:AbstractArray}
    bases::TB
end

const POWER_BASIS_LIST_CACHED = [
    CachedBasis{<:AbstractBasis{ℝ},<:PowerBasisData},
    CachedBasis{<:ManifoldsBase.AbstractOrthogonalBasis{ℝ},<:PowerBasisData},
    CachedBasis{<:ManifoldsBase.AbstractOrthonormalBasis{ℝ},<:PowerBasisData},
    CachedBasis{<:AbstractBasis{ℂ},<:PowerBasisData},
]

const PowerManifoldMultidimensional =
    AbstractPowerManifold{<:Manifold,ArrayPowerRepresentation} where {TSize}
const PowerManifoldNested =
    AbstractPowerManifold{<:Manifold,NestedPowerRepresentation} where {TSize}

_access_nested(x, i::Int) = x[i]
_access_nested(x, i::Tuple) = x[i...]

function allocate_result(M::PowerManifoldNested, f, x...)
    return [allocate_result(M.manifold, f, map(y -> _access_nested(y, i), x)...) for i in get_iterator(M)]
end
function allocate_result(M::PowerManifoldNested, f::typeof(flat), w::TFVector, x)
    alloc = [allocate(_access_nested(w.data, i)) for i in get_iterator(M)]
    return FVector(CotangentSpace, alloc)
end
function allocate_result(M::PowerManifoldNested, f::typeof(get_vector), p, X)
    return [allocate_result(M.manifold, f, _access_nested(p, i)) for i in get_iterator(M)]
end
function allocate_result(M::PowerManifoldNested, f::typeof(sharp), w::CoTFVector, x)
    alloc = [allocate(_access_nested(w.data, i)) for i in get_iterator(M)]
    return FVector(TangentSpace, alloc)
end
function allocate_result(M::PowerManifoldNested, f::typeof(get_coordinates), p, X, B)
    return invoke(
        allocate_result,
        Tuple{Manifold,typeof(get_coordinates),Any,Any,typeof(B)},
        M,
        f,
        p,
        X,
        B,
    )
end

function allocation_promotion_function(M::AbstractPowerManifold, f, args::Tuple)
    return allocation_promotion_function(M.manifold, f, args)
end

^(M::Manifold, n) = PowerManifold(M, n...)

"""
    check_manifold_point(M::AbstractProductManifold, p; kwargs...)

Check whether `p` is a valid point on an [`AbstractPowerManifold`](@ref) `M`,
i.e. each element of `p` has to be a valid point on the base manifold.

The tolerance for the last test can be set using the `kwargs...`.
"""
function check_manifold_point(M::AbstractPowerManifold, p; kwargs...)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        imp = check_manifold_point(M.manifold, _read(M, rep_size, p, i); kwargs...)
        imp === nothing || return imp
    end
    return nothing
end

"""
    check_tangent_vector(M::AbstractPowerManifold, p, X; check_base_point = true, kwargs... )

Check whether `X` is a tangent vector to `p` an the [`AbstractPowerManifold`](@ref)
`M`, i.e. atfer [`check_manifold_point`](@ref)`(M, p)`, and all projections to
base manifolds must be respective tangent vectors.
The optional parameter `check_base_point` indicates, whether to call [`check_manifold_point`](@ref)  for `p`.
The tolerance for the last test can be set using the `kwargs...`.
"""
function check_tangent_vector(
    M::AbstractPowerManifold,
    p,
    X;
    check_base_point = true,
    kwargs...,
)
    if check_base_point
        mpe = check_manifold_point(M, p)
        mpe === nothing || return mpe
    end
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        imp = check_tangent_vector(
            M.manifold,
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i);
            kwargs...,
        )
        imp === nothing || return imp
    end
    return nothing
end

function det_local_metric(
    M::MetricManifold{<:AbstractPowerManifold,PowerMetric},
    p::AbstractArray,
)
    result = one(number_eltype(p))
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        result *= det_local_metric(M.manifold, _read(M, rep_size, p, i))
    end
    return result
end

@doc raw"""
    distance(M::AbstractPowerManifold, p, q)

Compute the distance between `q` and `p` on an [`AbstractPowerManifold`](@ref),
i.e. from the element wise distances the Forbenius norm is computed.
"""
function distance(M::AbstractPowerManifold, p, q)
    sum_squares = zero(number_eltype(p))
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        sum_squares +=
            distance(M.manifold, _read(M, rep_size, p, i), _read(M, rep_size, q, i))^2
    end
    return sqrt(sum_squares)
end

@doc raw"""
    exp(M::AbstractPowerManifold, p, X)

Compute the exponential map from `p` in direction `X` on the [`AbstractPowerManifold`](@ref) `M`,
which can be computed using the base manifolds exponential map elementwise.
"""
exp(::AbstractPowerManifold, ::Any...)

function exp!(M::AbstractPowerManifold, q, p, X)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        exp!(
            M.manifold,
            _write(M, rep_size, q, i),
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
        )
    end
    return q
end

@doc raw"""
    flat(M::AbstractPowerManifold, p, X::FVector{TangentSpaceType})

use the musical isomorphism to transform the tangent vector `X` from the tangent space at
`p` on an [`AbstractPowerManifold`](@ref) `M` to a cotangent vector.
This can be done elementwise for each entry of `X` (and `p`).
"""
flat(::AbstractPowerManifold, ::Any...)

function flat!(M::AbstractPowerManifold, ξ::CoTFVector, p, X::TFVector)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        flat!(
            M.manifold,
            FVector(CotangentSpace, _write(M, rep_size, ξ.data, i)),
            _read(M, rep_size, p, i),
            FVector(TangentSpace, _read(M, rep_size, X.data, i)),
        )
    end
    return ξ
end

function get_basis(M::AbstractPowerManifold, p, B::AbstractBasis)
    rep_size = representation_size(M.manifold)
    vs = [get_basis(M.manifold, _read(M, rep_size, p, i), B) for i in get_iterator(M)]
    return CachedBasis(B, PowerBasisData(vs))
end
function get_basis(M::AbstractPowerManifold, p, B::DiagonalizingOrthonormalBasis)
    rep_size = representation_size(M.manifold)
    vs = [
        get_basis(
            M.manifold,
            _read(M, rep_size, p, i),
            DiagonalizingOrthonormalBasis(_read(M, rep_size, B.frame_direction, i)),
        ) for i in get_iterator(M)
    ]
    return CachedBasis(B, PowerBasisData(vs))
end
for BT in ManifoldsBase.DISAMBIGUATION_BASIS_TYPES
    if BT == DiagonalizingOrthonormalBasis
        continue
    end
    eval(quote
        @invoke_maker 3 AbstractBasis get_basis(M::AbstractPowerManifold, p, B::$BT)
    end)
end

function get_coordinates(M::AbstractPowerManifold, p, X, B::DefaultOrthonormalBasis)
    rep_size = representation_size(M.manifold)
    vs = [
        get_coordinates(
            M.manifold,
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            B,
        ) for i in get_iterator(M)
    ]
    return reduce(vcat, reshape(vs, length(vs)))
end
function get_coordinates(
    M::AbstractPowerManifold,
    p,
    X,
    B::CachedBasis{<:AbstractBasis,<:PowerBasisData,𝔽},
) where {𝔽}
    rep_size = representation_size(M.manifold)
    vs = [
        get_coordinates(
            M.manifold,
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            _access_nested(B.data.bases, i),
        ) for i in get_iterator(M)
    ]
    return reduce(vcat, reshape(vs, length(vs)))
end
for BT in POWER_BASIS_LIST_CACHED
    eval(quote
        @invoke_maker 4 CachedBasis{<:AbstractBasis,<:PowerBasisData} get_coordinates(
            M::AbstractPowerManifold,
            p,
            X,
            B::$BT,
        )
    end)
end

function get_coordinates!(M::AbstractPowerManifold, Y, p, X, B::DefaultOrthonormalBasis)
    rep_size = representation_size(M.manifold)
    dim = manifold_dimension(M.manifold)
    v_iter = 1
    for i in get_iterator(M)
        # TODO: this view is really suboptimal when `dim` can be statically determined
        get_coordinates!(
            M.manifold,
            view(Y, v_iter:v_iter+dim-1),
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            B,
        )
        v_iter += dim
    end
    return Y
end
function get_coordinates!(
    M::AbstractPowerManifold,
    Y,
    p,
    X,
    B::CachedBasis{<:AbstractBasis{ℝ},<:PowerBasisData,ℝ},
)
    TypeTuple = Tuple{
        AbstractPowerManifold,
        Any,
        Any,
        Any,
        CachedBasis{<:AbstractBasis,<:PowerBasisData,ℝ},
    }
    return invoke(get_coordinates!, TypeTuple, M, Y, p, X, B)
end
function get_coordinates!(
    M::AbstractPowerManifold,
    Y,
    p,
    X,
    B::CachedBasis{<:AbstractBasis,<:PowerBasisData,𝔽},
) where {𝔽}
    rep_size = representation_size(M.manifold)
    dim = manifold_dimension(M.manifold)
    v_iter = 1
    for i in get_iterator(M)
        get_coordinates!(
            M.manifold,
            view(Y, v_iter:v_iter+dim-1),
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            _access_nested(B.data.bases, i),
        )
        v_iter += dim
    end
    return Y
end

get_iterator(M::PowerManifold{<:Manifold,Tuple{N}}) where {N} = 1:N
@generated function get_iterator(M::PowerManifold{<:Manifold,SizeTuple}) where {SizeTuple}
    size_tuple = size_to_tuple(SizeTuple)
    return Base.product(map(Base.OneTo, size_tuple)...)
end

function get_vector!(
    M::PowerManifold,
    Y,
    p,
    X,
    B::CachedBasis{<:AbstractBasis,<:PowerBasisData},
)
    dim = manifold_dimension(M.manifold)
    rep_size = representation_size(M.manifold)
    v_iter = 1
    for i in get_iterator(M)
        get_vector!(
            M.manifold,
            _write(M, rep_size, Y, i),
            _read(M, rep_size, p, i),
            X[v_iter:v_iter+dim-1],
            _access_nested(B.data.bases, i),
        )
        v_iter += dim
    end
    return Y
end
function get_vector!(M::AbstractPowerManifold, Y, p, X, B::DefaultOrthonormalBasis)
    dim = manifold_dimension(M.manifold)
    rep_size = representation_size(M.manifold)
    v_iter = 1
    for i in get_iterator(M)
        get_vector!(
            M.manifold,
            _write(M, rep_size, Y, i),
            _read(M, rep_size, p, i),
            X[v_iter:v_iter+dim-1],
            B,
        )
        v_iter += dim
    end
    return Y
end

@doc raw"""
    injectivity_radius(M::AbstractPowerManifold[, p])

the injectivity radius on an [`AbstractPowerManifold`](@ref) is for the global case
equal to the one of its base manifold. For a given point `p` it's equal to the
minimum of all radii in the array entries.
"""
function injectivity_radius(M::AbstractPowerManifold, p)
    radius = 0.0
    initialized = false
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        cur_rad = injectivity_radius(M.manifold, _read(M, rep_size, p, i))
        if initialized
            radius = min(cur_rad, radius)
        else
            radius = cur_rad
            initialized = true
        end
    end
    return radius
end
injectivity_radius(M::AbstractPowerManifold) = injectivity_radius(M.manifold)
eval(quote
    @invoke_maker 1 Manifold injectivity_radius(
        M::AbstractPowerManifold,
        rm::AbstractRetractionMethod,
    )
end)
eval(quote
    @invoke_maker 1 Manifold injectivity_radius(
        M::AbstractPowerManifold,
        rm::ExponentialRetraction,
    )
end)

@doc raw"""
    inverse_retract(M::AbstractPowerManifold, p, q, m::InversePowerRetraction)

Compute the inverse retraction from `p` with respect to `q` on an [`AbstractPowerManifold`](@ref) `M`
using an [`InversePowerRetraction`](@ref), which by default encapsulates a inverse retraction
of the base manifold. Then this method is performed elementwise, so the encapsulated inverse
retraction method has to be one that is available on the base [`Manifold`](@ref).
"""
inverse_retract(::AbstractPowerManifold, ::Any...)

function inverse_retract!(M::AbstractPowerManifold, X, p, q, method::InversePowerRetraction)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        inverse_retract!(
            M.manifold,
            _write(M, rep_size, X, i),
            _read(M, rep_size, p, i),
            _read(M, rep_size, q, i),
            method.inverse_retraction,
        )
    end
    return X
end

@doc raw"""
    inner(M::AbstractPowerManifold, p, X, Y)

Compute the inner product of `X` and `Y` from the tangent space at `p` on an
[`AbstractPowerManifold`](@ref) `M`, i.e. for each arrays entry the tangent
vector entries from `X` and `Y` are in the tangent space of the corresponding
element from `p`.
The inner product is then the sum of the elementwise inner products.
"""
function inner(M::AbstractPowerManifold, p, X, Y)
    result = zero(number_eltype(X))
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        result += inner(
            M.manifold,
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            _read(M, rep_size, Y, i),
        )
    end
    return result
end

default_metric_dispatch(::AbstractPowerManifold, ::PowerMetric) = Val(true)

function isapprox(M::AbstractPowerManifold, p, q; kwargs...)
    result = true
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        result &= isapprox(
            M.manifold,
            _read(M, rep_size, p, i),
            _read(M, rep_size, q, i);
            kwargs...,
        )
    end
    return result
end
function isapprox(M::AbstractPowerManifold, p, X, Y; kwargs...)
    result = true
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        result &= isapprox(
            M.manifold,
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            _read(M, rep_size, Y, i);
            kwargs...,
        )
    end
    return result
end

@doc raw"""
    log(M::AbstractPowerManifold, p, q)

Compute the logarithmic map from `p` to `q` on the [`AbstractPowerManifold`](@ref) `M`,
which can be computed using the base manifolds logarithmic map elementwise.
"""
log(::AbstractPowerManifold, ::Any...)

function log!(M::AbstractPowerManifold, X, p, q)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        log!(
            M.manifold,
            _write(M, rep_size, X, i),
            _read(M, rep_size, p, i),
            _read(M, rep_size, q, i),
        )
    end
    return X
end

@doc raw"""
    manifold_dimension(M::PowerManifold)

Returns the manifold-dimension of an [`PowerManifold`](@ref) `M`
$=\mathcal N = (\mathcal M)^{n_1,…,n_d}$, i.e. with $n=(n_1,…,n_d)$ the array
size of the power manifold and $d_{\mathcal M}$ the dimension of the base manifold
$\mathcal M$, the manifold is of dimension

````math
\dim(\mathcal N) = \dim(\mathcal M)\prod_{i=1}^d n_i = n_1n_2\cdot…\cdot n_d \dim(\mathcal M).
````
"""
function manifold_dimension(M::PowerManifold{<:Manifold,TSize}) where {TSize}
    return manifold_dimension(M.manifold) * prod(size_to_tuple(TSize))
end

@doc raw"""
    norm(M::AbstractPowerManifold, p, X)

Compute the norm of `X` from the tangent space of `p` on an
[`AbstractPowerManifold`](@ref) `M`, i.e. from the element wise norms the
Frobenius norm is computed.
"""
function norm(M::AbstractPowerManifold, p, X)
    sum_squares = zero(number_eltype(X))
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        sum_squares +=
            norm(M.manifold, _read(M, rep_size, p, i), _read(M, rep_size, X, i))^2
    end
    return sqrt(sum_squares)
end

@doc raw"""
    power_dimensions(M::PowerManifold)

return the power of `M`,
"""
function power_dimensions(M::PowerManifold{<:Manifold,TSize}) where {TSize}
    return size_to_tuple(TSize)
end

function rand(rng::AbstractRNG, d::PowerFVectorDistribution)
    fv = zero_vector(d.type, d.point)
    _rand!(rng, d, fv)
    return fv
end
function rand(rng::AbstractRNG, d::PowerPointDistribution)
    x = allocate_result(d.manifold, rand, d.point)
    _rand!(rng, d, x)
    return x
end

function _rand!(rng::AbstractRNG, d::PowerFVectorDistribution, v::AbstractArray)
    PM = d.type.manifold
    rep_size = representation_size(PM.manifold)
    for i in get_iterator(d.type.manifold)
        copyto!(d.distribution.point, _read(PM, rep_size, d.point, i))
        _rand!(rng, d.distribution, _read(PM, rep_size, v, i))
    end
    return v
end
function _rand!(rng::AbstractRNG, d::PowerPointDistribution, x::AbstractArray)
    M = d.manifold
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        _rand!(rng, d.distribution, _write(M, rep_size, x, i))
    end
    return x
end

@inline function _read(M::AbstractPowerManifold, rep_size::Tuple, x::AbstractArray, i::Int)
    return _read(M, rep_size, x, (i,))
end
@inline function _read(
    ::PowerManifoldMultidimensional,
    rep_size::Tuple,
    x::AbstractArray,
    i::Tuple,
)
    return view(x, rep_size_to_colons(rep_size)..., i...)
end
@inline function _read(
    ::PowerManifoldMultidimensional,
    rep_size::Tuple,
    x::HybridArray,
    i::Tuple,
)
    return x[rep_size_to_colons(rep_size)..., i...]
end
@inline function _read(::PowerManifoldNested, rep_size::Tuple, x::AbstractArray, i::Tuple)
    return view(x[i...], rep_size_to_colons(rep_size)...)
end

@generated function rep_size_to_colons(rep_size::Tuple)
    N = length(rep_size.parameters)
    return ntuple(i -> Colon(), N)
end

function representation_size(M::PowerManifold{<:Manifold,TSize}) where {TSize}
    return (representation_size(M.manifold)..., size_to_tuple(TSize)...)
end

@doc raw"""
    retract(M::AbstractPowerManifold, p, X, method::PowerRetraction)

Compute the retraction from `p` with tangent vector `X` on an [`AbstractPowerManifold`](@ref) `M`
using a [`PowerRetraction`](@ref), which by default encapsulates a retraction of the
base manifold. Then this method is performed elementwise, so the encapsulated retraction
method has to be one that is available on the base [`Manifold`](@ref).
"""
retract(::AbstractPowerManifold, ::Any...)

function retract!(M::AbstractPowerManifold, q, p, X, method::PowerRetraction)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        retract!(
            M.manifold,
            _write(M, rep_size, q, i),
            _read(M, rep_size, p, i),
            _read(M, rep_size, X, i),
            method.retraction,
        )
    end
    return q
end

@doc raw"""
    sharp(M::AbstractPowerManifold, p, ξ::FVector{CotangentSpaceType})

Use the musical isomorphism to transform the cotangent vector `ξ` from the tangent space at
`p` on an [`AbstractPowerManifold`](@ref) `M` to a tangent vector.
This can be done elementwise for every entry of `ξ` (and `p`).
"""
sharp(::AbstractPowerManifold, ::Any...)

function sharp!(M::AbstractPowerManifold, X::TFVector, p, ξ::CoTFVector)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        sharp!(
            M.manifold,
            FVector(TangentSpace, _write(M, rep_size, X.data, i)),
            _read(M, rep_size, p, i),
            FVector(CotangentSpace, _read(M, rep_size, ξ.data, i)),
        )
    end
    return X
end

function show(
    io::IO,
    mime::MIME"text/plain",
    B::CachedBasis{T,D,𝔽},
) where {T<:AbstractBasis,D<:PowerBasisData,𝔽}
    println(io, "$(T()) for a power manifold")
    for i in Base.product(map(Base.OneTo, size(B.data.bases))...)
        println(io, "Basis for component $i:")
        show(io, mime, _access_nested(B.data.bases, i))
        println(io)
    end
end
function show(io::IO, M::PowerManifold{TM,TSize,ArrayPowerRepresentation}) where {TM,TSize}
    print(io, "PowerManifold($(M.manifold), $(join(TSize.parameters, ", ")))")
end
function show(io::IO, M::PowerManifold{TM,TSize,TPR}) where {TM,TSize,TPR}
    print(io, "PowerManifold($(M.manifold), $(TPR()), $(join(TSize.parameters, ", ")))")
end

support(tvd::PowerFVectorDistribution) = FVectorSupport(tvd.type, tvd.point)
support(d::PowerPointDistribution) = MPointSupport(d.manifold)

@inline function _write(M::AbstractPowerManifold, rep_size::Tuple, x::AbstractArray, i::Int)
    return _write(M, rep_size, x, (i,))
end
@inline function _write(
    M::PowerManifoldMultidimensional,
    rep_size::Tuple,
    x::AbstractArray,
    i::Tuple,
)
    return view(x, rep_size_to_colons(rep_size)..., i...)
end
@inline function _write(M::PowerManifoldNested, rep_size::Tuple, x::AbstractArray, i::Tuple)
    return view(x[i...], rep_size_to_colons(rep_size)...)
end


module Polyopt

export MomentProb, momentprob, momentprob_chordalembedding, solve_mosek

include("polynomial.jl")
include("cliques.jl")

immutable MomentProb{T<:Number}
    order :: Int
    basis :: Array{Poly{Int},1}
    obj   :: Array{T,1}
    mom   :: Array{Any,1}
    eq    :: Array{Any,1}
end

include("mosek.jl")
include("latex.jl")

# compute all monomials of degree 'deg' or less
function monomials_unsorted{T<:Number}(deg::Int, vars::Array{Poly{T},1})
    if length(vars) == 1 return [ vars[1]^i for i=0:deg ] end

    x, tail = vars[1], vars[2:end]
    monoms = monomials_unsorted(deg, tail)

    for i=1:deg
        append!(monoms, [ x^i * m for m=monomials_unsorted(deg-i, tail) ])
    end
    monoms
end

monomials{T<:Number}(deg::Int, vars::Array{Poly{T},1}) = sort!(monomials_unsorted(deg, vars))

function moment(order::Int, syms::Symbols)
    v = monomials(order, variables(syms))
    v*v'
end

function moment(order::Int, syms::Symbols, I::Array{Int})
    v = monomials(order, variables(syms)[I])
    v*v'
end

function moment{T<:Number}(order::Int, syms::Symbols, p::Poly{T})
    v = monomials(order - (p.deg+1) >> 1, variables(syms))
    p*v*v'
end

function moment{T<:Number}(order::Int, syms::Symbols, p::Poly{T}, I::Array{Int})
    v = monomials(order - (p.deg+1) >> 1, variables(syms)[I])
    p*v*v'
end

function indexmap{T<:Number}(v::Array{Poly{T},1})
    Dict(zip([vi.alpha for vi=v],[1:length(v)]))
end

linear_index(imap::Dict{Array{Int,2},Int}, p::Poly) = Int[ imap[p.alpha[i,:]] for i=1:p.m ]

function vectorize{T<:Number}(p::Poly{T}, l::Int, imap::Dict{Array{Int,2},Int})
    a = zeros(T, l)
    a[linear_index(imap,p) ] = p.c
    a
end

function vectorize{T<:Number}(M::Array{Poly{T},2}, l::Int, imap::Dict{Array{Int,2},Int})
    m, n = size(M)
    m == n || throw(DimensionMismatch("Matrix is not square: dimensions are $(size(M))"))
    subi = Int[]
    subj = Int[]
    val  = T[]
    for j=1:m
        v = M[j,j].c
        push!(subi, (j+(j-1)*m)*ones(Int, length(v))...)
        push!(subj, linear_index(imap,M[j,j])...)
        push!(val, v...)
        for i=j+1:m
            I = linear_index(imap,M[i,j])
            v = M[i,j].c
            push!(subi, (i+(j-1)*m)*ones(Int, length(v))...)
            push!(subj, I...)
            push!(val, v...)
        end
    end
    A = sparse(subi, subj, val, m*m, l)
end

function momentprob{S,T,U,V,W}(order::Int, obj::Poly{S}, pineq::Array{Poly{T},1}, peq::Array{Poly{U},1}, ineq::Array{Poly{V},1}, eq::Array{Poly{W},1})
    v = monomials(2*order, variables(obj.syms))
    l = length(v)
    imap = indexmap(v)

    obj.deg <= 2*order || error("obj has degree higher than 2*order")

    p = vectorize(obj, l, imap)
    mom = Array(Any, length(pineq)+length(ineq)+1)
    mom[1] = vectorize(moment(order, obj.syms), l, imap)
    for k=1:length(pineq)
       pineq[k].deg <= 2*order || error("pineq[$(k)] has degree higher than 2*order")
       mom[k+1] = vectorize(moment(order, pineq[k].syms, pineq[k]), l, imap)
    end

    for k=1:length(ineq)
        ineq[k].deg <= 2*order || error("ineq[$(k)] has degree higher than 2*order")
        mom[k+length(pineq)+1] = sparse(vectorize(ineq[k], l, imap)')
    end

    momeq = Array(Any, length(peq)+length(eq))
    for k=1:length(peq)
        peq[k].deg <= 2*order || error("peq[$(k)] has degree higher than 2*order")
        momeq[k] = vectorize(moment(order, peq[k].syms, peq[k]), l, imap)
    end

    for k=1:length(eq)
        eq[k].deg <= 2*order || error("eq[$(k)] has degree higher than 2*order")
        momeq[k+length(peq)] = sparse(vectorize(eq[k], l, imap)')
    end

    MomentProb(order, v,  p, mom, momeq)
end

momentprob{S}(order::Int, obj::Poly{S}) =
    momentprob(order, obj, Poly{Int}[], Poly{Int}[], Poly{Int}[], Poly{Int}[])

momentprob{S,T}(order::Int, obj::Poly{S}, pineq::Array{Poly{T},1}) =
    momentprob(order, obj, pineq, Poly{Int}[], Poly{Int}[], Poly{Int}[])

momentprob{S,T,U}(order::Int, obj::Poly{S}, pineq::Array{Poly{T},1}, peq::Array{Poly{U},1}) =
    momentprob(order, obj, pineq, peq, Poly{Int}[], Poly{Int}[])

function momentprob_chordal{S,T,U,V,W}(order::Int, cliques::Array{Array{Int,1},1}, obj::Poly{S},
                                       pineq::Array{Poly{T},1}, pineq_index::Array{Int,1},
                                       peq::Array{Poly{U},1}, peq_index::Array{Int,1},
                                       ineq::Array{Poly{V},1}, eq::Array{Poly{W},1})

    v = monomials(2*order, variables(obj.syms))
    l = length(v)
    imap = indexmap(v)

    obj.deg <= 2*order || error("obj has degree higher than 2*order")

    p = vectorize(obj, l, imap)
    mom = Array(Any, length(pineq) + length(cliques) + length(ineq))
    for k=1:length(cliques)
        mom[k] = vectorize(moment(order, obj.syms, cliques[k]), l, imap)
    end

    for k=1:length(pineq)
        pineq[k].deg <= 2*order || error("pineq[$(k)] has degree higher than 2*order")
        mom[length(cliques)+k] = vectorize(moment(order, pineq[k].syms, pineq[k], cliques[pineq_index[k]]), l, imap)
    end

    for k=1:length(ineq)
        ineq[k].deg <= 2*order || error("ineq[$(k)] has degree higher than 2*order")
        mom[length(cliques)+length(pineq)+k] = sparse(vectorize(ineq[k], l, imap)')
    end

    momeq = Array(Any, length(peq)+length(eq))
    for k=1:length(peq)
        peq[k].deg <= 2*order || error("peq[$(k)] has degree higher than 2*order")
        momeq[k] = vectorize(moment(order, peq[k].syms, peq[k], cliques[peq_index[k]]), l, imap)
    end

    for k=1:length(eq)
        eq[k].deg <= 2*order || error("eq[$(k)] has degree higher than 2*order")
        momeq[length(peq)+k] = sparse(vectorize(eq[k], l, imap)')
    end

    MomentProb(order, v,  p, mom, momeq)
end

momentprob_chordal{S,T}(order::Int, cliques::Array{Array{Int,1},1}, obj::Poly{S}, pineq::Array{Poly{T},1}, pineq_index::Array{Int,1}) =
   lasserre_chordal(order, cliques, obj, pineq, pineq_index, Poly{Int}[], Array{Int}[], Poly{Int}[], Poly{Int}[])

function correlative_sparsity{S,T}(obj::Poly{S}, p::Array{Poly{T},1})
    A = eye(Int,obj.n,obj.n)

    for k=1:obj.m
        I = find(obj.alpha[k,:])
        A[I,I] = 1
    end

    for j=1:length(p)
        I = find(sum(p[j].alpha,1))
        A[I,I] = 1
    end

    sparse(A)
end

function momentprob_chordalembedding{S,T,U,V,W}(order::Int, obj::Poly{S},
                                                pineq::Array{Poly{T},1}, peq::Array{Poly{U},1},
                                                ineq::Array{Poly{V},1}, eq::Array{Poly{W},1})

    C = correlative_sparsity(obj, [pineq, ineq, peq, eq])
    cliques = chordal_embedding(C)

    momentprob_chordal(order, cliques, obj,
                       pineq, [ clique_index(cliques, find(sum(p.alpha,1))) for p = pineq ],
                       peq,   [ clique_index(cliques, find(sum(p.alpha,1))) for p = peq   ],
                       ineq, eq)
end

momentprob_chordalembedding{S,T,U}(order::Int, obj::Poly{S}, pineq::Array{Poly{T},1}, peq::Array{Poly{U},1}) =
    momentprob_chordalembedding(order, obj, pineq, peq, Poly{Int}[], Poly{Int}[])

momentprob_chordalembedding{S,T}(order::Int, obj::Poly{S}, pineq::Array{Poly{T},1}) =
    momentprob_chordalembedding(order, obj, pineq, Poly{Int}[], Poly{Int}[], Poly{Int}[])

momentprob_chordalembedding{S}(order::Int, obj::Poly{S}) =
    momentprob_chordalembedding(order, obj, Poly{Int}[], Poly{Int}[], Poly{Int}[], Poly{Int}[])

end


mutable struct Model
    diffusion_coeff::Array{Float64,2}   # Diffusion coefficient: depends on the sine of latitude
    heat_capacity::Array{Float64,2}     # Heat capacity: depends on the geography (land, ocean, ice, etc.)
    albedo::Array{Float64,2}            # Albedo coefficient: depends on the geography (land, ocean, ice, etc.)
    solar_forcing::Array{Float64,3}     # Time-dependent incoming solar radiation: depends on the orbital parameters and the albedo
    radiative_cooling_co2::Float64      # Constant outgoing long-wave radiation: depends on the CO2 concentration
    radiative_cooling_feedback::Float64 # Outgoing long-wave radiation (feedback effects): models the water vapor cyces, lapse rate and cloud cover
end

function Model(mesh, num_steps_year; co2_concentration = 315.0) # co2_concentration in [ppm], default value from year 1950
    @unpack nx,ny = mesh

    # Constants
    radiative_cooling_feedback = 2.15   #[W/m^2/°C]: sensitivity of the seasonal cycle and annual change in the forcing agents

    radiative_cooling_co2 = calc_radiative_cooling_co2(co2_concentration)

    # Read parameters

    geography = read_geography(joinpath(@__DIR__, "..", "input", "The_World.dat"),nx,ny)
    albedo    = read_albedo(joinpath(@__DIR__, "..", "input", "albedo.dat"),nx,ny)
    diffusion_coeff = calc_diffusion_coefficients(geography,nx,ny)
    heat_capacity, tau_land, tau_snow, tau_sea_ice, tau_mixed_layer = calc_heat_capacity(geography,radiative_cooling_feedback) # TODO: Remove unused variables

    co_albedo = 1.0 .- albedo

    ecc = eccentricity(1950)
    ob = obliquity(1950)
    per = perihelion(1950)
    solar_forcing = calc_solar_forcing(co_albedo, ecc=ecc, ob=ob, per=per) #TODO: Add arguments [nx, ny, num_steps_year]

    return Model(diffusion_coeff, heat_capacity, albedo, solar_forcing, radiative_cooling_co2, radiative_cooling_feedback)
end

function set_co2_concentration!(model, co2_concentration)
    model.radiative_cooling_co2 = calc_radiative_cooling_co2(co2_concentration)
end

Base.size(model::Model) = size(model.heat_capacity)

function Base.show(io::IO, model::Model)
  nx, ny = size(model)
  print(io, "Model() with ", nx, "×", ny, " degrees of freedom")
end

"""
calc_radiative_cooling_co2()
Computes the CO2 parameter depending on the CO2 concentration

* Default CO2 concentration is 315 ppm (equivalent to year 1950)
"""
function calc_radiative_cooling_co2(co2_concentration=315.0)

  # Define base values for co2_concentration and radiative_cooling_co2
  co2_concentration_base = 315.0
  radiative_cooling_co2_base = 210.3

  # Doesn't change as long as CP2ppm doesn't change
  radiative_cooling_co2=radiative_cooling_co2_base-5.35*log(co2_concentration/co2_concentration_base)

  return radiative_cooling_co2
end

"""
_calc_insolation()
auxiliar function of calc_solar_forcing
"""
function _calc_insolation(dt, ob, ecc, per, nt, nlat, siny, cosy, tany, s0)

  eccfac = 1.0 - ecc^2
  rzero  = (2.0*pi)/eccfac^1.5

  #  Solve the orbital equation for lambda as a function of time with
  #  a fourth-order Runge-Kutta method
  lambda = zeros(Float64,nt+1)
  solar = zeros(Float64,nlat,nt)

  for n in 2:nt+1
    nu = lambda[n-1] - per
    t1 = dt*(rzero*(1.0 - ecc * cos(nu))^2)
    t2 = dt*(rzero*(1.0 - ecc * cos(nu+0.5*t1))^2)
    t3 = dt*(rzero*(1.0 - ecc * cos(nu+0.5*t2))^2)
    t4 = dt*(rzero*(1.0 - ecc * cos(nu + t3))^2)
    lambda[n] = lambda[n-1] + (t1 + 2.0*t2 + 2.0*t3 + t4)/6.0
  end

  #  Compute the average daily solar irradiance as a function of
  #  latitude and longitude (time)

  for n in 1:nt
    nu = lambda[n] - per
    rhofac = ((1.0- ecc*cos(nu))/eccfac)^2
    sindec = sin(ob)*sin(lambda[n])
    cosdec = sqrt(1.0-sindec^2)
    tandec = sindec/cosdec
    for j in 1:nlat
      z = -tany[j]*tandec
      if z >= 1.0    # polar latitudes when there is no sunrise (winter)
        solar[j,n] = 0.0
      else
        if z <= -1.0                # when there is no sunset (summer)
          solar[j,n] = rhofac * s0 * siny[j] * sindec
        else
          h_zero = acos(z)
          solar[j,n] = rhofac/pi*s0*(h_zero*siny[j]*sindec+cosy[j]*cosdec*sin(h_zero))
        end
      end
    end
  end

  return lambda, solar
end

"""
calc_solar_forcing()

* Default s0 is 1371.685 [W/m²] (Current solar constant)
* Default orbital parameters of correspond to year 1950 AD:
    ecc = 0.016740
    ob  = 0.409253
    per = 1.783037
"""
function calc_solar_forcing(co_albedo, yr=0; solar_cycle=false, s0=1371.685, orbital=false, ecc=0.016740, ob=0.409253, per=1.783037)
  # Calculate the sin, cos, and tan of the latitudes of Earth from the
  # colatitudes, calculate the insolation

  # TODO: Add as input arguments
  nlatitude   = 65
  nlongitude  = 128
  ntimesteps  = 48

  dy = pi/(nlatitude-1.0)
  dt = 1.0 / ntimesteps

  siny = zeros(Float64,nlatitude)
  cosy = zeros(Float64,nlatitude)
  tany = zeros(Float64,nlatitude)

  if yr == 0
    for j in 1:nlatitude
      lat = pi/2.0 - dy*(j-1)    # latitude in radians
      siny[j] = sin(lat)
      if j == 1
        cosy[j] = 0.0
        tany[j] = 1000.0
      elseif j == nlatitude
        cosy[j] = 0.0
        tany[j] = -1000.0
      else
        cosy[j] = cos(lat)
        tany[j] = tan(lat)
      end
    end
    lambda, solar = _calc_insolation(dt, ob, ecc, per, ntimesteps, nlatitude, siny, cosy, tany, s0)
  elseif yr > 0 && (solar_cycle || orbital)
    lambda, solar = _calc_insolation(dt, ob, ecc, per, ntimesteps, nlatitude, siny, cosy, tany, s0)
  end

  # TODO: Not needed??????
  for j in 1:nlatitude
    sum = 0.0
    for ts in 1:ntimesteps
      sum=sum+solar[j,ts]
    end
  end

  solar_forcing = zeros(Float64,nlongitude,nlatitude,ntimesteps)
  # calcualte the seasonal forcing
  for ts in 1:ntimesteps
    for j in 1:nlatitude
      for i in 1:nlongitude
        solar_forcing[i,j,ts] = solar[j,ts]*co_albedo[i,j]
      end
    end
  end
  return solar_forcing

end

function calc_heat_capacity(geography,radiative_cooling_feedback=2.15)

  # Depths
  depth_atmos = 5000.         # meters #TODO: not used???
  depth_mixed_layer = 70.0    # meters
  depth_soil = 2.0            # meters
  depth_seaice = 2.5          # meters
  depth_snow = 2.0            # meters
  layer_depth = 0.5           # kilometers

  # Physical properties of atmosphere
  rho_atmos = 1.293           # kg m^-3  dry air (STP)
  csp_atmos = 1005.0          # J kg^-1 K^-1 (STP)
  scale_height = 7.6          # kilometers

  # Physical properties of water
  rho_water = 1000.0          # kg m^-3
  csp_water = 4186.0          # J kg^-1 K^-1

  # Physical properties of soil
  rho_soil = 1100.0           # kg m^-3
  csp_soil = 850.0            # J kg^-1 K^-1

  # Physical properties of sea ice
  rho_sea_ice = 917.0         # kg m^-3
  csp_sea_ice = 2106.0        # J kg^-1 K^-1

  # Physical properties of snow covered surface
  rho_snow = 400.0            # kg m^-3
  csp_snow = 1900.0           # J kg^-1 K^-1

  # Other constants
  sec_per_yr = 3.15576e7      # seconds per year
  days_per_yr = 365.2422      # days per year

  # atmosphere with exponentially decaying density
  sum = 0.0
  for n in 1:10
    z = (0.25 + layer_depth*real(n-1))/scale_height
    sum = sum + exp(-z)
  end

  c_atmos  	= csp_atmos*layer_depth*1000.0*rho_atmos*sum/sec_per_yr
  c_soil   	= depth_soil*rho_soil*csp_soil/sec_per_yr
  c_seaice 	= depth_seaice*rho_sea_ice*csp_sea_ice/sec_per_yr
  c_snow   	= depth_snow * rho_snow * csp_snow/sec_per_yr
  c_mixed_layer = depth_mixed_layer*rho_water*csp_water/sec_per_yr

  # Calculate radiative relaxation times for columns
  tau_land = (c_soil + c_atmos)/radiative_cooling_feedback * days_per_yr
  tau_snow = (c_snow + c_atmos)/radiative_cooling_feedback * days_per_yr
  tau_sea_ice = (c_seaice + c_atmos)/radiative_cooling_feedback * days_per_yr
  tau_mixed_layer = (c_mixed_layer + c_atmos)/radiative_cooling_feedback

  # define heatcap
  heatcap = zeros(size(geography,1),size(geography,2))

  # Assign the correct value of the heat capacity of the columns
  for j in 1:size(geography,2)
    for i in 1:size(geography,1)
      geo  = geography[i,j]
      if geo == 1                            # land
        heatcap[i,j] = c_soil + c_atmos
      elseif geo == 2                        # perennial sea ice
        heatcap[i,j] = c_seaice + c_atmos
      elseif geo == 3                        # permanent snow cover
        heatcap[i,j] = c_snow + c_atmos
      elseif geo == 4                        # lakes, inland seas
        heatcap[i,j] = c_mixed_layer/3.0 + c_atmos
      elseif geo in (5, 6, 7, 8)
        # 5: Pacific ocean
        # 6: Atlantic ocean
        # 7: Indian ocean
        # 8: Mediterranean
        heatcap[i,j] = c_mixed_layer + c_atmos
      end
    end
  end

  return heatcap, tau_land, tau_snow, tau_sea_ice, tau_mixed_layer
end

"""
calc_diffusion_coefficients()
Calculate the diffusion coefficients at finest grid level.
"""
function calc_diffusion_coefficients(geography,nlongitude=128,nlatitude=65)

  coeff_eq     = 0.65 # coefficinet for diffusion at equator
  coeff_ocean  = 0.40 # coefficient for ocean diffusion
  coeff_land   = 0.65 # coefficinet for land diffusion
  coeff_landNP = 0.28 # coefficinet for land diffusion (north pole)
  coeff_landSP = 0.20 # coefficinet for land diffusion (south pole)

  diffusion = zeros(Float64,nlongitude,nlatitude)

  j_equator = div(nlatitude,2) + 1

  for j = 1:nlatitude
      theta = pi*real(j-1)/real(nlatitude-1)
      colat = sin(theta)^5

      for i = 1:nlongitude
          let geo = geography[i,j]
              if geo >= 5 && geo <= 7 # oceans
                  diffusion[i,j] = (coeff_eq-coeff_ocean)*colat + coeff_ocean
              else # land, sea ice, etc
                  if j <= j_equator # northern hemisphere
                      diffusion[i,j] = (coeff_land-coeff_landNP)*colat + coeff_landNP
                  else # southern hemisphere
                      diffusion[i,j] = (coeff_land-coeff_landSP)*colat + coeff_landSP
                  end
                end
          end
      end
  end

  return diffusion
end

# Calculate the area weighted mean of the diffusion at the mid-point between
# the pole and the first ring of grid points.

# If you need the mean diffusion at the poles just do the following:
#   mean_diffusion_north = sum(diffusion_north)
#   mean_diffusion_south = sum(diffusion_south)

## function calc_diffusion_coefficients_poles(diffusion,nlongitude=128,nlatitude=65)

##   # Fractional areas for the poles
##   angle  = pi/real(nlatitude-1)
##   area_1 = 0.5*(1.0 - cos(0.5*angle))
##   area_2 = sin(0.5*angle)*sin(angle)/float(nlongitude)

##   total_area = area_1 + area_2

##   diffusion_north = zeros(Float64,nlongitude)
##   diffusion_south = zeros(Float64,nlongitude)

##   for i = 1:nlongitude
##       diffusion_north[i] = (area_1*diffusion[1,        1] + area_2*diffusion[i,          2])/total_area
##       diffusion_south[i] = (area_1*diffusion[1,nlatitude] + area_2*diffusion[i,nlatitude-1])/total_area
##   end

##   return diffusion_north,diffusion_south
## end

function read_albedo(filepath="albedo.dat",nlongitude=128,nlatitude=65)
  result = zeros(Float64,nlongitude,nlatitude)
  open(filepath) do fh
      for lat = 1:nlatitude
          if eof(fh) break end
          result[:,lat] = parse.(Float64,split(strip(readline(fh) ),r"\s+"))
      end
  end
  return result
end

function read_geography(filepath="The_World.dat",nlongitude=128,nlatitude=65)
  result = zeros(Int8,nlongitude,nlatitude)
  open(filepath) do fh
      for lat = 1:nlatitude
          if eof(fh) break end
          result[:,lat] = parse.(Int8,split(strip(readline(fh) ),r""))
      end
  end
  return result
end

# Table 4 from Berger 1978, Long-Term Variations of daily Insolation and Quaternary Climatic Changes
# Higher accuracy values from Gary Russel, https://data.giss.nasa.gov/modelE/ar5plots/srorbpar.html
const eccentricity_parameters_amplitude = [0.01860798, 0.01627522, -0.01300660, 0.00988829,
                                           -0.00336700, 0.00333077, -0.00235400, 0.00140015,
                                           0.00100700, 0.00085700, 0.00064990, 0.00059900,
                                           0.00037800, -0.00033700, 0.00027600, 0.00018200,
                                           -0.00017400, -0.00012400, 0.00001250]
const eccentricity_parameters_mean_rate = [4.2072050, 7.3460910, 17.8572630, 17.2205460, 16.8467330,
                                           5.1990790, 18.2310760, 26.2167580, 6.3591690, 16.2100160,
                                           3.0651810, 16.5838290, 18.4939800, 6.1909530, 18.8677930,
                                           17.4255670, 6.1860010, 18.4174410, 0.6678630]
const eccentricity_parameters_phase = [28.620089, 193.788772, 308.307024, 320.199637, 279.376984,
                                       87.195000, 349.129677, 128.443387, 154.143880, 291.269597,
                                       114.860583, 332.092251, 296.414411, 145.769910, 337.237063,
                                       152.092288, 126.839891, 210.667199, 72.108838]

# Table 1 from Berger 1978, Long-Term Variations of daily Insolation and Quaternary Climatic Changes
# Higher accuracy values from Gary Russel, https://data.giss.nasa.gov/modelE/ar5plots/srorbpar.html
const obliquity_parameters_amplitude = [-2462.2214466, -857.3232075, -629.3231835, -414.2804924,
                                        -311.7632587, 308.9408604, -162.5533601, -116.1077911,
                                        101.1189923, -67.6856209, 24.9079067, 22.5811241,
                                        -21.1648355, -15.6549876, 15.3936813, 14.6660938,
                                        -11.7273029, 10.2742696, 6.4914588, 5.8539148, -5.4872205,
                                        -5.4290191, 5.160957, 5.0786314, -4.0735782, 3.7227167,
                                        3.3971932, -2.8347004, -2.6550721, -2.5717867, -2.4712188,
                                        2.462541, 2.2464112, -2.0755511, -1.9713669, -1.8813061,
                                        -1.8468785, 1.8186742, 1.7601888, -1.5428851, 1.4738838,
                                        -1.4593669, 1.4192259, -1.181898, 1.1756474, -1.1316126,
                                        1.0896928]
const obliquity_parameters_mean_rate = [31.609974, 32.620504, 24.172203, 31.983787, 44.828336,
                                        30.973257, 43.668246, 32.246691, 30.599444, 42.681324,
                                        43.836462, 47.439436, 63.219948, 64.230478, 1.01053,
                                        7.437771, 55.782177, 0.373813, 13.218362, 62.583231,
                                        63.593761, 76.43831, 45.815258, 8.448301, 56.792707,
                                        49.747842, 12.058272, 75.27822, 65.241008, 64.604291,
                                        1.647247, 7.811584, 12.207832, 63.856665, 56.15599,
                                        77.44884, 6.801054, 62.209418, 20.656133, 48.344406,
                                        55.14546, 69.000539, 11.07135, 74.291298, 11.047742,
                                        0.636717, 12.844549]
const obliquity_parameters_phase = [251.9025, 280.8325, 128.3057, 292.7252, 15.3747, 263.7951,
                                    308.4258, 240.0099, 222.9725, 268.7809, 316.7998, 319.6024,
                                    143.805, 172.7351, 28.93, 123.5968, 20.2082, 40.8226, 123.4722,
                                    155.6977, 184.6277, 267.2772, 55.0196, 152.5268, 49.1382,
                                    204.6609, 56.5233, 200.3284, 201.6651, 213.5577, 17.0374,
                                    164.4194, 94.5422, 131.9124, 61.0309, 296.2073, 135.4894,
                                    114.875, 247.0691, 256.6114, 32.1008, 143.6804, 16.8784,
                                    160.6835, 27.5932, 348.1074, 82.6496]

# Table 5 from Berger 1978, Long-Term Variations of daily Insolation and Quaternary Climatic Changes
# Higher accuracy values from Gary Russel, https://data.giss.nasa.gov/modelE/ar5plots/srorbpar.html
const perihelion_parameters_amplitude = [7391.022589, 2555.1526947, 2022.7629188, -1973.6517951,
                                         1240.2321818, 953.8679112, -931.7537108, 872.3795383,
                                         606.3544732, -496.0274038, 456.9608039, 346.946232,
                                         -305.8412902, 249.6173246, -199.10272, 191.0560889,
                                         -175.2936572, 165.9068833, 161.1285917, 139.7878093,
                                         -133.5228399, 117.0673811, 104.6907281, 95.3227476,
                                         86.7824524, 86.0857729, 70.5893698, -69.9719343,
                                         -62.5817473, 61.5450059, -57.9364011, 57.1899832,
                                         -57.0236109, -54.2119253, 53.2834147, 52.1223575,
                                         -49.0059908, -48.3118757, -45.4191685, -42.235792,
                                         -34.7971099, 34.4623613, -33.8356643, 33.6689362,
                                         -31.2521586, -30.8798701, 28.4640769, -27.1960802,
                                         27.0860736, -26.3437456, 24.725374, 24.6732126, 24.4272733,
                                         24.0127327, 21.7150294, -21.5375347, 18.1148363,
                                         -16.9603104, -16.1765215, 15.5567653, 15.4846529,
                                         15.2150632, 14.5047426, -14.3873316, 13.1351419,
                                         12.8776311, 11.9867234, 11.9385578, 11.7030822, 11.6018181,
                                         -11.2617293, -10.4664199, 10.433397, -10.2377466,
                                         10.1934446, -10.1280191, 10.0289441, -10.0034259]
const perihelion_parameters_mean_rate = [31.609974, 32.620504, 24.172203, 0.636717, 31.983787,
                                         3.138886, 30.973257, 44.828336, 0.991874, 0.373813,
                                         43.668246, 32.246691, 30.599444, 2.147012, 10.511172,
                                         42.681324, 13.650058, 0.986922, 9.874455, 13.013341,
                                         0.262904, 0.004952, 1.142024, 63.219948, 0.205021,
                                         2.151964, 64.230478, 43.836462, 47.439436, 1.384343,
                                         7.437771, 18.829299, 9.500642, 0.431696, 1.16009,
                                         55.782177, 12.639528, 1.155138, 0.168216, 1.647247,
                                         10.884985, 5.610937, 12.658184, 1.01053, 1.983748,
                                         14.023871, 0.560178, 1.273434, 12.021467, 62.583231,
                                         63.593761, 76.43831, 4.28091, 13.218362, 17.818769,
                                         8.359495, 56.792707, 8.448301, 1.978796, 8.863925,
                                         0.186365, 8.996212, 6.771027, 45.815258, 12.002811,
                                         75.27822, 65.241008, 18.870667, 22.009553, 64.604291,
                                         11.498094, 0.578834, 9.237738, 49.747842, 2.147012,
                                         1.196895, 2.133898, 0.173168]
const perihelion_parameters_phase = [251.9025, 280.8325, 128.3057, 348.1074, 292.7252, 165.1686,
                                     263.7951, 15.3747, 58.5749, 40.8226, 308.4258, 240.0099,
                                     222.9725, 106.5937, 114.5182, 268.7809, 279.6869, 39.6448,
                                     126.4108, 291.5795, 307.2848, 18.93, 273.7596, 143.805,
                                     191.8927, 125.5237, 172.7351, 316.7998, 319.6024, 69.7526,
                                     123.5968, 217.6432, 85.5882, 156.2147, 66.9489, 20.2082,
                                     250.7568, 48.0188, 8.3739, 17.0374, 155.3409, 94.1709, 221.112,
                                     28.93, 117.1498, 320.5095, 262.3602, 336.2148, 233.0046,
                                     155.6977, 184.6277, 267.2772, 78.9281, 123.4722, 188.7132,
                                     180.1364, 49.1382, 152.5268, 98.2198, 97.4808, 221.5376,
                                     168.2438, 161.1199, 55.0196, 262.6495, 200.3284, 201.6651,
                                     294.6547, 99.8233, 213.5577, 154.1631, 232.7153, 138.3034,
                                     204.6609, 106.5938, 250.4676, 332.3345, 27.3039]

# Eqn. 4 from Berger 1978, Long-Term Variations of daily Insolation and Quaternary Climatic Changes
function e_sincos_pi(year)
  M    = eccentricity_parameters_amplitude
  g    = eccentricity_parameters_mean_rate
  beta = eccentricity_parameters_phase

  arg = deg2rad.(g/3600) * (year-1950) .+ deg2rad.(beta)

  return sum(M .* sin.(arg)), sum(M .* cos.(arg))
end

eccentricity(year, e_sin_pi, e_cos_pi) = sqrt(e_sin_pi^2 + e_cos_pi^2)

eccentricity(year) = eccentricity(year, e_sincos_pi(year)...)

# Eqn. 1 from Berger 1978, Long-Term Variations of daily Insolation and Quaternary Climatic Changes
function obliquity(year)
  A     = obliquity_parameters_amplitude
  f     = obliquity_parameters_mean_rate
  delta = obliquity_parameters_phase

  epsilon_star = 23.320556

  return deg2rad(epsilon_star + sum(A/3600 .* cos.(deg2rad.(f/3600) * (year-1950) .+ deg2rad.(delta))))
end

# Eqn. 6 from Berger 1978, Long-Term Variations of daily Insolation and Quaternary Climatic Changes
function perihelion(year)
  F     = perihelion_parameters_amplitude
  f     = perihelion_parameters_mean_rate
  delta = perihelion_parameters_phase

  pi_ = atan(e_sincos_pi(year)...)
  psi_tilde = 50.439273
  zeta = 3.392506

  psi = psi_tilde/3600 * (year-1950) + zeta + sum(F/3600 .* sin.(deg2rad.(f/3600) * (year-1950) .+ deg2rad.(delta)))

  return (pi_ + deg2rad(psi)) % (2*pi)
end


include("../src/delta.jl")

function test_generators_coefficients()
   print("generators.coefficients...")

   delta12 = delta_k_qexp(12, 100)
   @test delta_qexp(100) - delta12 == 0

   delta16 = delta_k_qexp(16, 1)
   @test delta16 == 0

   delta16 = delta_k_qexp(16, 10)
   S = parent(delta16)
   q = gen(S)
   @test delta16 - (q+216*q^2-3348*q^3+13888*q^4+52110*q^5-723168*q^6+2822456*q^7-4078080*q^8-3139803*q^9+O(q^10)) == 0

   delta16 = delta_k_qexp(16, 100)
   @test coeff(delta16, 98) == 695238414190488

   delta18 = delta_k_qexp(18, 10)
   S = parent(delta18)
   q = gen(S)
   @test delta18 - (q - 528*q^2 - 4284*q^3 + 147712*q^4 - 1025850*q^5 + 2261952*q^6 + 3225992*q^7 - 8785920*q^8 - 110787507*q^9 + O(q^10)) == 0

   delta20 = delta_k_qexp(20, 10)
   S = parent(delta20)
   q = gen(S)
   @test delta20 - (q + 456*q^2 + 50652*q^3 - 316352*q^4 - 2377410*q^5 + 23097312*q^6 - 16917544*q^7 - 383331840*q^8 + 1403363637*q^9 + O(q^10)) == 0

   delta22 = delta_k_qexp(22, 10)
   S = parent(delta22)
   q = gen(S)
   @test delta22 - (q - 288*q^2 - 128844*q^3 - 2014208*q^4 + 21640950*q^5 + 37107072*q^6 - 768078808*q^7 + 1184071680*q^8 + 6140423133*q^9 + O(q^10)) == 0

   delta26 = delta_k_qexp(26, 10)
   S = parent(delta26)
   q = gen(S)
   @test delta26 - (q - 48*q^2 - 195804*q^3 - 33552128*q^4 - 741989850*q^5 + 9398592*q^6 + 39080597192*q^7 + 3221114880*q^8 - 808949403027*q^9 + O(q^10)) == 0

   println("PASS")
end

function test_generators()
   test_generators_coefficients()

   println("")
end


export gmm,
    shrink,
    estimate_component_prob,
    estimate_var_μ,
    rand_gmm

using Distributions
using StatsBase

"""
fit Gaussian mixture model with all means set to 0.

A SNP contributing h2 to heritability has μ ~ N(0, var = N*h2)
and z|μ ~ N(μ, 1.0)

Here we decompose distribution of z into null component z ~ N(0; 1)
and components z ~ N(0; var = 1 + N*h2/M) where h2 is total heritability

Due to population stratification etc. the null component might have
variance > 1. Then z needs to be rescaled such that variance of null component ≈ 1

dp0: Dirichlet pseudo counts for null component

"""
function gmm(x, π0, σ0; ϵ = 1e-3, maxiter = 1000, dp0 = 0, verbose = false, forcenull = false)

    n_data = length(x)
    n_cluster = length(π0)

    # Dirichlet prior
    dp = zeros(Float64, 1, n_cluster)
    dp[1] = dp0

    π = π0[:]'
    σ = σ0[:]'
    π_old = π0[:]'
    σ_old = σ0[:]'
    pc = zeros(Float64, n_data, n_cluster)

    llold = -Inf
    converged = false
    n_iter = 0

    for iter = 1:maxiter
        n_iter += 1
        # E-step
        # P(c_j|x_j) propto P(x_j|c_j) * P(c_j)
        for i = 1:n_data
            pc[i,:] = pdf.(Normal.(0., σ), x[i]) .* π
        end # i
        pc .= pc ./ sum(pc; dims = 2)

        # M-step
        π .= (sum(pc; dims = 1) .+ dp) ./ (n_data + dp0)
        σ  .= sqrt.(sum(x.^2 .* pc; dims = 1) ./ sum(pc; dims = 1))
        if forcenull
            σ[1] = 1.0
        end

        # loglikelihood
        # sum_i log( sum_c P(x_i|c)*P(c) )
        ll = 0.0
        for i = 1:n_data
            ll += log(sum(pdf.(Normal.(0., σ), x[i]) .* π))
        end

        verbose && println("Iteration $iter : LogLikelihood $ll")

        if (maximum(abs.((π_old .- π) ./ π_old)) < ϵ) && (maximum(abs.((σ_old .- σ) ./ σ_old)) < ϵ)
            converged = true
            break
        end

        llold = ll
        π_old .= π
        σ_old = σ

    end

    (π = π, σ = σ, converged = converged, ll = llold, iter = n_iter)
end

"""
estimate component probability for z

pc[i, k] : probability that z[i] belongs to component (π[k], σ[k])
"""
function estimate_component_prob(z, π, σ)
    @assert ndims(π) == ndims(σ) == 2 "π and σ need to be row vectors"
    @assert ndims(z) == 1 "z needs to be a column vector"
    @assert size(π) == size(σ) "π and σ need to have same dimensions"
    n_data = length(z)
    n_cluster = size(π, 2)
    pc = zeros(Float64, n_data, n_cluster)
    for i = 1:n_data
        pc[i,:] = pdf.(Normal.(0., σ), z[i]) .* π
    end
    pc ./ sum(pc; dims = 2)
end

function rand_gmm(π, σ)
    K = length(π)
    k = sample(1:K, Weights(π[:]))
    rand(Normal(0.0, σ[k]), 1)[1]
end

"""
compute shrinkage factor for z

i.e. estimate μ|z  where z|μ ~ N(μ; sd = 1)
and prior for μ given by  P(μ) = sum_k π[k] N(0; var = σ[k]^2 - 1)

σ2_μ for first component is set to 0
"""
function shrink(z, π, σ)

    pc = estimate_component_prob(z, π, σ)
    σ2_μ = max.(σ .^2 .- 1.0, 0.0)
    σ2_μ[1] = 0.0

    pc * (σ2_μ ./ (1.0 .+ σ2_μ))'
end

"""
estimate var_μ = N h2 / M for each SNP

N  number of samples
h2 total heritability
M  number of SNPs

π: component probability as obtained from gmm
σ: standard deviation of component as obtained from gmm

1/var_μ needs to be added to diagonal of LD matrix
"""
function estimate_var_μ(z, π, σ)

    pc = estimate_component_prob(z, π, σ)
    σ2_μ = max.(σ .^2 .- 1.0, 0.0)

    pc * σ2_μ'
end

function test_gmm()

    x = vcat(rand(Normal(0, 1.5), 100),
             rand(Normal(0,1.), 800),
             rand(Normal(0, 5.), 100))

    σ0 = [1., 2., 3.]
    π0 = [0.9, 0.05, 0.05]

    res = gmm(x, π0, σ0)
end


using Convolution2D
using Test

@testset "Convolution2D" begin
    X = [
        1  4  5  4  4
        1  3  4  5  3
        2  3  5  4  4
        4  4  2  1  1
        3  3  3  2  4
    ]
    kern = [0 1 0
		1 -4 1
		0 1 0]
    out = conv2D(X, kern)
    @test out[3, 3] == -7

    X =  [5  2  5  3
        1  2  2  4
        1  1  1  4
        3  2  4  3]
    @test_broken conv2D(X, kern)
end


module G0
function f(q::Vector{T}, q_dot::Vector{T}, xi::T) where T
    l1 = q[1]
    l2 = q[2]
    l3 = q[3]
    l1_dot = q_dot[1]
    l2_dot = q_dot[2]
    l3_dot = q_dot[3]

    -3.02912483102257e-6*l1.^5 - 1.78242802467006e-5*l1.^4.*l2 + 3.29870469288301e-5*l1.^4.*l3 + 9.37009405046133e-8*l1.^4 - 2.7515585012037e-5*l1.^3.*l2.^2 + 0.000126439665488924*l1.^3.*l2.*l3 + 1.41766156856612e-7*l1.^3.*l2 - 0.00012928389889518*l1.^3.*l3.^2 - 5.17114922403405e-7*l1.^3.*l3 + 3.82562374888515e-10*l1.^3 - 4.68312335895801e-7*l1.^2.*l2.^3 + 8.41756449126506e-5*l1.^2.*l2.^2.*l3 - 3.5556688146475e-7*l1.^2.*l2.^2 - 0.00027422615773196*l1.^2.*l2.*l3.^2 + 2.84603410436225e-7*l1.^2.*l2.*l3 + 8.24471484365873e-10*l1.^2.*l2 + 0.000220912928627533*l1.^2.*l3.^3 + 6.34804124641222e-7*l1.^2.*l3.^2 - 1.97385538081975e-9*l1.^2.*l3 + 4.72058579046329e-13*l1.^2 + 1.03644793616222e-5*l1.*l2.^4 - 4.03814334300281e-5*l1.*l2.^3.*l3 - 2.54281526051452e-8*l1.*l2.^3 - 2.4037236653465e-5*l1.*l2.^2.*l3.^2 + 7.88483312239122e-7*l1.*l2.^2.*l3 - 7.16189502159426e-10*l1.*l2.^2 + 0.00019939210057147*l1.*l2.*l3.^3 - 1.07291993224609e-6*l1.*l2.*l3.^2 - 2.21211033699355e-10*l1.*l2.*l3 + 1.17459514671635e-12*l1.*l2 - 0.00016055210411278*l1.*l3.^4 - 6.65739999913559e-8*l1.*l3.^3 + 2.0885116396423e-9*l1.*l3.^2 - 2.1200173246156e-12*l1.*l3 - 2.90488121299558e-5*l1 + 2.92654312504004e-6*l2.^5 - 2.49701624716648e-5*l2.^4.*l3 + 6.91116586537353e-8*l2.^4 + 7.00070469556942e-5*l2.^3.*l3.^2 - 2.50342189017214e-7*l2.^3.*l3 - 7.11777798043519e-11*l2.^3 - 6.1726059406604e-5*l2.^2.*l3.^3 - 2.02753578297076e-8*l2.^2.*l3.^2 + 9.30013738403587e-10*l2.^2.*l3 - 4.12072853349556e-13*l2.^2 - 1.9256742517163e-5*l2.*l3.^4 + 3.72132609316059e-7*l2.*l3.^3 - 8.17394230662383e-10*l2.*l3.^2 - 3.53791896171088e-13*l2.*l3 + 1.77973142556746e-5*l2 + 3.60656416727372e-5*l3.^5 - 7.63807770899228e-8*l3.^4 - 4.2573148239698e-10*l3.^3 + 1.23925501369454e-12*l3.^2 + 1.14158920149061e-5*l3 - 2.4957191258213e-8
end
end


using Revise
using AdFem
using SparseArrays
using LinearAlgebra
using PyPlot
αm = 0.0
αf = 0.0
β2 = 0.5*(1 - αm + αf)^2
γ = 0.5 - αm + αf

m = 40
n = 20
h = 0.01
NT = 100
Δt = 1/NT
bdedge = []
for j = 1:n
    push!(bdedge, [(j-1)*(m+1)+m+1 j*(m+1)+m+1])
end
bdedge = vcat(bdedge...)

bdnode = Int64[]
for j = 1:n+1
    push!(bdnode, (j-1)*(m+1)+1)
end

M = compute_fem_mass_matrix1(m, n, h)
S = spzeros((m+1)*(n+1), (m+1)*(n+1))
M = [M S;S M]
H = diagm(0=>[1,1,0.5])
K = 0.1
σY = 0.03
# σY = 1000.

state = zeros(2(m+1)*(n+1),NT+1)
velo = zeros(2(m+1)*(n+1),NT+1)
acce = zeros(2(m+1)*(n+1),NT+1)
stress = zeros(NT+1, 4*m*n, 3)
internal_variable = zeros(NT+1, 4*m*n)

t = 0.0
for i = 1:NT
    @info i
        ##### STEP 1: Computes the external force ####

    T = eval_f_on_boundary_edge((x,y)->0.02*sin(2π*i*Δt), bdedge, m, n, h)
    # T = eval_f_on_boundary_edge((x,y)->0.0, bdedge, m, n, h)
    T = [zeros(length(T)) -T]
    T = compute_fem_traction_term(T, bdedge, m, n, h)
    f1 = eval_f_on_gauss_pts((x,y)->0., m, n, h)
    f2 = eval_f_on_gauss_pts((x,y)->0., m, n, h)
    # f2 = eval_f_on_gauss_pts((x,y)->0.1, m, n, h)
    F = compute_fem_source_term(f1, f2, m, n, h)
    fext = F+T

        ##### STEP 2: Extract Variables ####

    u = state[:,i]
    ∂∂u = acce[:,i]
    ∂u = velo[:,i]

    ε0 = eval_strain_on_gauss_pts(u, m, n, h)
    σ0 = stress[i,:,:]
    α0 = internal_variable[i,:]

        ##### STEP 3: Newton Iteration ####

    global t += (1 - αf)*Δt
    ∂∂up = ∂∂u[:]
    iter = 0
    while true
        iter += 1
        # @info iter
        up = (1 - αf)*(u + Δt*∂u + 0.5 * Δt^2 * ((1 - β2)*∂∂u + β2*∂∂up)) + αf*u
        global fint, stiff, α, σ = compute_planestressplasticity_stress_and_stiffness_matrix(
            up, ε0, σ0, α0, K, σY, H, m, n, h
        )
        res = M * (∂∂up *(1 - αm) + αm*∂∂u)  + fint - fext
        A = M*(1 - αm) + (1 - αf) * 0.5 * β2 * Δt^2 * stiff
        A, _ = fem_impose_Dirichlet_boundary_condition(A, bdnode, m, n, h)
        res[[bdnode; bdnode .+ (m+1)*(n+1)]] .= 0.0

        Δ∂∂u = A\res
        ∂∂up -= Δ∂∂u
        err = norm(res)
        # @info err
        if err<1e-8
            break
        end
    end
    global t += αf*Δt

        ##### STEP 3: Update State Variables ####

    u += Δt * ∂u + Δt^2/2 * ((1 - β2) * ∂∂u + β2 * ∂∂up)
    ∂u += Δt * ((1 - γ) * ∂∂u + γ * ∂∂up)
    stress[i+1,:,:] = σ
    internal_variable[i+1,:] = α

    state[:,i+1] = u
    acce[:,i+1] = ∂∂up
    velo[:,i+1] = ∂u
end

x = []
y = []
for j= 1:n+1
    for i = 1:m+1
        push!(x, (i-1)*h)
        push!(y, (j-1)*h)
    end
end

for i = 1:5:NT+1
    close("all")
    scatter(x+state[1:(m+1)*(n+1), i], y+state[(m+1)*(n+1)+1:end, i])
    scatter(x[m+1]+state[m+1, i],
            y[m+1]+state[(m+1)*(n+1)+m+1, i], color="red")
    xlabel("x")
    ylabel("y")
    k = string(i)
    k = repeat("0", 3-length(k))*k
    title("t = $k")
    ylim(-0.05,0.25)
    xlim(-0.01, 0.45)
    gca().invert_yaxis()
    savefig("u$k.png")

    close("all");
    plot(1:i, -state[(m+1)*(n+1)+m+1, 1:i])
    xlim(0, NT+2)
    ylim(0, 0.03)
    scatter(i, -state[(m+1)*(n+1)+m+1, i], color="red")
    savefig("du$k.png")
end

run(`convert -delay 10 -loop 0 u*.png plasticity_u.gif`)
run(`convert -delay 10 -loop 0 du*.png plasticity_du.gif`)


module TestNamespace

using Omega
using Omega.Inference: swapsinglesite, normalkernel
using UnicodePlots
using Statistics: mean

function test_ssmc_1()
  μ = uniform(0.0, 1.0)
  x = normal(μ, 5.0)
  samples = rand(μ, x ==ₛ 7.0, 10000; alg= SSMH)
  println(histogram([samples...]))
  println(mean(samples))
end

test_ssmc_1()

function test_ssmc_2()
  x = logistic(1.0, 2.0, (2, 3, 4))
  y = sum(x)
  rand(x, y ==ₛ 5.0; alg = SSMH)
end

test_ssmc_2()

function test_sshm_drift()
  x = normal(0.0, 1.0)
  y = normal(1.0, 1.0)
  z = rand(x - y, x ==ₛ y, 10000; alg=SSMH, proposal = (rng, ω) -> swapsinglesite(rng, ω) do x
    normalkernel(rng, x, .1)
  end)
  z = convert(Array{Float64}, z)
  println(histogram([z[300:end]...]))
  println(mean(z[300:end]))
end

test_sshm_drift()

end