struct LinearStateSpaceModel{T}
    A::Matrix{T}
    B::Matrix{T}
    Q::Matrix{T}
end

function LinearStateSpaceModel(A::Matrix, Q::Matrix)
    LinearStateSpaceModel(A,zeros(size(A,1),1),Q)
end

struct Simulator <: Transducer
    lssm::LinearStateSpaceModel
end

function Transducers.start(rf::R_{Simulator}, result)
    private_state = (0, zeros(size(xform(rf).lssm.A,1)))
    return wrap(rf, private_state, start(inner(rf), result))
end

function step(l::LinearStateSpaceModel, x, input)
    l.A*x + l.B*input + Gaussian(l.Q)
end

function Transducers.next(rf::R_{Simulator}, result, input)
    wrapping(rf, result) do st, iresult
        (n, x) = st
        n += 1
        new_x = step(xform(rf).lssm,x,input)
        iresult = next(inner(rf), iresult, new_x)
        return (n, new_x), iresult
    end
end


@inline function Phi_s(Cell,s,z,Def,ϕ_tf,D,res0)
    """
    Solid Potential Transfer Function

    Phi_s(Cell,s,z,Def)

    """

 if Def == "Pos"
    Electrode = Cell.Pos #Electrode Length
 else
    Electrode = Cell.Neg #Electrode Length
 end

κ_eff = Cell.Const.κ*Electrode.ϵ_e^Electrode.κ_brug #Effective Electrolyte Conductivity
σ_eff = Electrode.σ*Electrode.ϵ_s^Electrode.σ_brug #Effective Electrode Conductivity
comb_cond_eff = κ_eff+σ_eff #combining into single variable

#Defining SOC
θ = Cell.Const.SOC * (Electrode.θ_100-Electrode.θ_0) + Electrode.θ_0

#Prepare for j0
cs0 = Electrode.cs_max * θ

#Current Flux Density
if Cell.Const.CellTyp == "Doyle_94"
   κ = Electrode.k_norm/Electrode.cs_max/Cell.Const.ce0^(1-Electrode.α)
   j0 = κ*(Cell.Const.ce0*(Electrode.cs_max-cs0))^(1-Electrode.α)*cs0^Electrode.α
else
   j0 = Electrode.k_norm*(Cell.Const.ce0*cs0*(Electrode.cs_max-cs0))^(1-Electrode.α)
end

#Resistance
Rtot = R*Cell.Const.T/(j0*F^2) + Electrode.RFilm
#Rtot = R*Cell.Const.T/(j0*Cell.Const.CC_A*F) + Electrode.RFilm

#∂Uocp_Def
∂Uocp_elc = Cell.Const.∂Uocp(Def,θ)/Electrode.cs_max

ν = @. Electrode.L*sqrt((Electrode.as/σ_eff+Electrode.as/κ_eff)/(Rtot+∂Uocp_elc*(Electrode.Rs/(F*Electrode.Ds))*(tanh(Electrode.β)/(tanh(Electrode.β)-Electrode.β)))) #Condensing Variable - eq. 4.13
ν_∞ = @. Electrode.L*sqrt((Electrode.as/κ_eff+Electrode.as/σ_eff)/(Rtot))

ϕ_tf .= @. (-Electrode.L*(κ_eff*(cosh(ν)-cosh((z-1)*ν)))-Electrode.L*(σ_eff*(1-cosh(z*ν)+z*ν*sinh(ν))))/(Cell.Const.CC_A*σ_eff*(comb_cond_eff)*ν*sinh(ν)) #Transfer Function - eq. 4.19
D .= @. (-Electrode.L*(κ_eff*(cosh(ν_∞)-cosh((z-1)*ν_∞)))-Electrode.L*(σ_eff*(1-cosh(z*ν_∞)+z*ν_∞*sinh(ν_∞))))/(Cell.Const.CC_A*σ_eff*(comb_cond_eff)*ν_∞*sinh(ν_∞)) # Contribution to D as G->∞
zero_tf = @. Electrode.L*(z-2)*z/(2*Cell.Const.CC_A*σ_eff)
ϕ_tf[:,findall(s.==0)] .= zero_tf[:,findall(s.==0)]
res0 .= zeros(length(z))

if Def == "Pos"
   ϕ_tf .= -ϕ_tf
   D .= -D
end

end


function calFlxCorrection!(
    ocn :: Ocean;
    τ   :: Float64 = 10 * 86400.0,
    cfgs...
)
    do_convadjust = cfgs[:do_convadjust]

    Δt = cfgs[:Δt]
    r = Δt / τ
    rr = r / (1.0 + r)

    calTsSsMixed!(ocn)

    # Euler backward method
    @. ΔT_mixed = rr * (ocn.in_fld.Tclim - ocn.Ts_mixed)
    @. ΔS_mixed = rr * (ocn.in_fld.Sclim - ocn.Ss_mixed)

    @loop_hor ocn i j let

        T_ML = ocn.T_ML[i, j]
        S_ML = ocn.S_ML[i, j]
        FLDO = ocn.FLDO[i, j]
        h_ML = ocn.h_ML[i, j]

        if FLDO == -1  # Whole ocean layer. In the future I want to restrict FLDO so that code can be cleaner.

            ΔT_ML = 0.0
            for k=1:ocn.Nz[i, j]
                ΔT_ML += ΔT_mixed[k, i, j] * ocn.hs[k, i, j]
            end
            ΔT_ML /= h_ML

            ΔS_ML = 0.0
            for k=1:ocn.Nz[i, j]
                ΔS_ML += ΔS_mixed[k, i, j] * ocn.hs[k, i, j]
            end
            ΔS_ML /= h_ML

            T_ML += ΔT_ML
            S_ML += ΔS_ML
            ocn.T_ML[i, j] = T_ML
            ocn.S_ML[i, j] = S_ML

            ocn.Ts[:, i, j] .= T_ML
            ocn.Ss[:, i, j] .= S_ML

        else
            # With in mixed-layer
            if FLDO == 1
                ocn.T_ML += ΔT_mixed[1, i, j]
                ocn.S_ML += ΔS_mixed[1, i, j]
            elseif FLDO > 1

                ΔT_ML = 0.0
                for k=1:FLDO-1
                    ΔT_ML += ΔT_mixed[k, i, j] * ocn.hs[k, i, j]
                end
                ΔT_ML += ΔT_mixed[FLDO, k, j] * (ocn.zs[FLDO, i, j] + h_ML)
                ΔT_ML /= h_ML

                ΔS_ML = 0.0
                for k=1:FLDO-1
                    ΔS_ML += ΔS_mixed[k, i, j] * ocn.hs[k, i, j]
                end
                ΔS_ML += ΔS_mixed[FLDO, k, j] * (ocn.zs[FLDO, i, j] + h_ML)
                ΔS_ML /= h_ML

                T_ML += ΔT_ML
                S_ML += ΔS_ML
                ocn.T_ML[i, j] = T_ML
                ocn.S_ML[i, j] = S_ML

                ocn.Ts[1:FLDO-1, i, j] .= T_ML
                ocn.Ss[1:FLDO-1, i, j] .= S_ML

            end

            # Deeper ocean
            if FLDO != -1
                @. ocn.Ts[FLDO:end, i, j] += ΔT_mixed[FLDO:end, i, j]
                @. ocn.Ss[FLDO:end, i, j] += ΔS_mixed[FLDO:end, i, j]
            end

        end

        OC_updateB!(ocn, i, j)

    end

    @. ocn.qflx_T_correction = ΔT_mixed * ocn.hs * ρc_sw   / Δt   # + => warming
    @. ocn.qflx_S_correction = ΔS_mixed * ocn.hs           / Δt   # + => saltier

    if do_convadjust
        @loop_hor ocn i j let
            OC_doConvectiveAdjustment!(ocn, i, j;)
        end
    end

end


using Multirate
using PyPlot

N𝜙              = 32
tapsPer𝜙        = 10
resampleRate    = 3.14
xƒ              = 0.1
Nx              = 25
x               = cos(2*pi*xƒ*(0:Nx-1))
# x             = zeros( Nx); x[1] = 1.0
cutoffFreq      = min( 0.45, resampleRate )
transitionWidth = 0.05
hLen            = tapsPer𝜙*N𝜙
h               = firdes( hLen, cutoffFreq, DSP.kaiser, samplerate = N𝜙 ) .* N𝜙
myfilter        = FIRFilter( h, resampleRate, N𝜙 )
δfilter         = (hLen-1)/(2*N𝜙)
y               = filt( myfilter, x )
tx              = [0:length(x)-1]
ty              = [0:length(y)-1]./resampleRate-δfilter

stem( tx, x, "r" )
hold(true)
stem( ty, y )
hold(false)


module SymbolicWedderburn

using LinearAlgebra
using SparseArrays
using Primes

using Cyclotomics
using GroupsCore
using PermutationGroups
using StarAlgebras

export symmetry_adapted_basis, WedderburnDecomposition
export basis,
    # degree, # too common name to export
    direct_summands,
    invariant_vectors,
    issimple,
    multiplicity

macro spawn_compat(expr)
    @static if VERSION < v"1.3.0"
        return :(@async $(esc(expr)))
    else
        return :(Threads.@spawn $(esc(expr)))
    end
end

include("Characters/Characters.jl")
using .Characters
import .Characters: row_echelon_form!
import .Characters.FiniteFields

include("util.jl")
include("actions.jl")
include("action_characters.jl")
include("matrix_projections.jl")
include("minimal_projections.jl")
include("direct_summands.jl")
include("sa_basis.jl")
include("wedderburn_decomposition.jl")

end # module


include("euler/euler.jl")
import StatsBase: countmap
using .BigIntegers: BigInteger
using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

function compute(n::Int, modulo::Int)::BigInteger
    function get_digits(k::Int, limit::Int)::Vector{Vector{Int}}
        result = Vector{Vector{Int}}()
        if k ≤ 0 return result end
        if k ≤ limit append!(result, [[k]]) end
        for i ∈ 1:limit for j ∈ get_digits(k - i, i) append!(result, [vcat(i, j)]) end end
        return result
    end
    factorials = Dict{Int, BigInteger}()
    all_digits = BigInteger(invmod(9, modulo)) * BigInteger((powermod(10, n, modulo) - 1)) % modulo
    total = BigInteger(0)
    for digit ∈ 1:9
        for other_digits ∈ get_digits(digit, digit)
            len_digits = length(other_digits)
            if n - len_digits - 1 < 0 continue end
            digits = countmap(vcat(digit, other_digits))
            for m ∈ keys(digits)
                partial_digits = copy(digits)
                partial_digits[m] -= 1
                if !haskey(factorials, len_digits) factorials[len_digits] = prod(BigInteger(n - i - 1) for i ∈ 0:len_digits - 1) end
                multinomial = factorials[len_digits]
                for p ∈ values(partial_digits) multinomial ÷= factorial(p) end
                total += all_digits * m * multinomial
            end
        end
        total %= modulo
    end
    return total
end

compute(3, 10_000_000_000_000_000)

compute(7, 10_000_000_000_000_000)

compute(2_020, 10_000_000_000_000_000)

@benchmark compute(2_020, 10_000_000_000_000_000)


using DiffEqFlux
using Plots

function autocatlysis!(du, u, (a, k₁, k₋₁), t)
    du .= k₁ * a * u - k₋₁ * u.^2
end

p = (10.0, 0.1, 0.01)

u0 = [0.1]
tspan = (0.0, 10.0)
a, k₁, k₋₁ = p
prob = ODEProblem(autocatlysis!,u0,tspan,p)
sol = solve(prob)
plot(sol)


# a simple LP example
#
#   maximize x + y
#
#   s.t. 50 x + 24 y <= 2400
#        30 x + 33 y <= 2100
#        x >= 45, y >= 5
#
#   solution: x = 45, y = 6.25, objv = 51.25

using Gurobi, Test

@testset "LP 01a" begin
    env = Gurobi.Env()
    setparam!(env, "OutputFlag", 0)

    method = getparam(env, "Method")

    model = Gurobi.Model(env, "lp_01", :maximize)

    # add variables
    add_cvar!(model, 1.0, 45., Inf)  # x
    add_cvar!(model, 1.0,  5., Inf)  # y
    update_model!(model)

    # add constraints
    add_constr!(model, [50., 24.], '<', 2400.)
    add_constr!(model, [30., 33.], '<', 2100.)
    update_model!(model)

    # perform optimization
    optimize(model)

    # show results
    info = get_optiminfo(model)

    @test get_solution(model) == [45, 6.25]
    @test get_objval(model) == 51.25

    GC.gc()  # test finalizers
end


module Scenes

export HitRecord, Sphere, Scene, TriangleMesh, ray_intersect, create_triangles, Portal
#export has_uvs, has_normals, get_vertex, get_uv, get_normal

using LinearAlgebra

#push!(LOAD_PATH, pwd())
using ..GfxBase
using ..OBJMeshes
using ..Materials

#####################################
###### Generic Scene Data Type ######
#####################################
struct Scene
    id::Int64
    background::RGB{Float32}
    objects::Array{Any,1}
    lights::Array{Any,1}
end
""" Structure to store data about an intersection
of a ray with an object (a "hit")."""
mutable struct HitRecord
    t::Float64
    intersection::Vec3
    normal::Vec3
    uv::Union{Vec2,Nothing}
    object
end

# Abstract ray-object intersection function:
# Needs to be implemented for each type of object to be rendered
""" Intersect a ray with an object.
Returns a HitRecord with info about the intersetion, or nothing if
the ray doesn't intersect with the object. """
function ray_intersect(ray::Ray, object) end

##################
##### Sphere #####
##################

# Data type:
struct Sphere
    center::Vec3
    radius::Float64
    material::Material
end

mutable struct Portal
    mesh::Array{Any, 1}
    center::Vec3
    buddy::Union{Portal, Nothing}
end

""" Ray-sphere intersection. """
function ray_intersect(ray::Ray, object::Sphere)
    # define constants to use later
    cent = object.center
    d = ray.direction
    p = ray.origin
    r = object.radius

    # find the value of t
    a = dot(d,d)
    b = dot(p - cent,d)
    c = dot(p-cent, p-cent)-r^2
    discrim = b*b-a*c
    dx = p[1] - cent[1]
    dy = p[2] - cent[2]
    dz = p[3] - cent[3]

    # not solns if discriminant is < 0
    if discrim < 0
        return nothing
    else
        num = -b-sqrt(discrim)
        if num > 0
            t = num/a
        else
            num = -b+sqrt(discrim)
            if num > 0
                t = num/a
            else
                return nothing
            end
        end
    end
    # intersection points
    x = p[1] + t*d[1]
    y = p[2] + t*d[2]
    z = p[3] + t*d[3]

    # to use for the uv coords, intersection - center
    dist = normalize(Vec3(x,y,z)-cent)

    normal = normalize((Vec3(x,y,z)- cent)/r)
    if isnothing(object.material.texture)
        return HitRecord(t,Vec3(x,y,z), normal, nothing, object)
    else
        theta = atan(dist[1],dist[3])
        phi = asin(-dist[2]) # negative as image was upside down
        # shift by .5 for middle of the texel
        u = .5 + (theta/ (2*π))
        v = .5 - (phi/π)
        uv = Vec2(u,v)
        return HitRecord(t,Vec3(x,y,z), normal, uv, object)
    end
end

###########################
###### Triangle Mesh ######
###########################

""" Data type: stores the OBJTriangle, a reference to its Mesh
object, and the material it should be rendered with. """
struct Triangle
    geometry::OBJTriangle
    mesh::OBJMesh
    material
end

""" Return an Array of all Triangles belonging to the given mesh, assigning
each one the given material. """
function create_triangles(mesh::OBJMesh, material)
    [Triangle(f, mesh, material) for f in mesh.triangles]
end

""" Some helper functions that make for easier access to triangle data: """
function get_vertex(tri::Triangle, i)
    tri.mesh.positions[tri.geometry.positions[i]]
end

function has_uvs(tri::Triangle)
    length(tri.geometry.uvs) == 3
end

function get_uv(tri::Triangle, i)
    tri.mesh.uvs[tri.geometry.uvs[i]]
end

function has_normals(tri::Triangle)
    length(tri.geometry.normals) == 3
end

function get_normal(tri::Triangle, i)
    tri.mesh.normals[tri.geometry.normals[i]]
end

function ray_intersect(ray::Ray, object::Triangle)

    # vertices
    x = get_vertex(object, 1)
    y = get_vertex(object, 2)
    z = get_vertex(object, 3)

    dir = ray.direction
    org = ray.origin

    # letters from the text
    a = x[1]-y[1]
    b = x[2]-y[2]
    c = x[3]-y[3]
    g = dir[1]
    h = dir[2]
    i = dir[3]
    j = x[1] - org[1]
    k = x[2] - org[2]
    l = x[3] - org[3]
    d = x[1]-z[1]
    e = x[2]-z[2]
    f = x[3]-z[3]

    # the main matrix
    m = (a * ((e * i) - (h * f))) + (b * ((g * f) - (d * i))) + (c * ((d * h) - (e * g)))
    if m == 0 # no solns
        return nothing
    end
    # cramers rule
    t = -(((f * (a * k - j * b)) + (e * (j * c - a * l)) + (d * (b * l - k * c))) / m)
    γ = ((i * (a * k - j * b)) + (h * (j * c - a * l)) + (g * (b * l - k * c))) / m
    if (γ > 1) || (γ < 0)
        return nothing
    end
    β = (j * (e * i - h * f) + k * (g * f - d * i) + l * (d * h - e * g)) / m
    if (β < 0) || (β > (1 -γ))
        return nothing
    end
    # intersection point
    dir *= t
    inter = org + dir
    if !has_normals(object)
        norm = normalize(cross(y-x, z-x))
    else
        # barycentric interpolation for normals
        n1 = get_normal(object, 1)
        n2 = get_normal(object, 2)
        n3 = get_normal(object, 3)
        x = n1[1]+(β*(n2[1]-n1[1]))+(γ*(n3[1]-n1[1]))
        y = n1[2]+(β*(n2[2]-n1[2]))+(γ*(n3[2]-n1[2]))
        z = n1[3]+(β*(n2[3]-n1[3]))+(γ*(n3[3]-n1[3]))
        norm = normalize(Vec3(x,y,z))
    end
    if isnothing(object.material.texture)
        return HitRecord(t, inter, norm, nothing, object)
    else
        if has_uvs(object)
            uv1 = get_uv(object, 1)
            uv2 = get_uv(object, 2)
            uv3 = get_uv(object, 3)

            # barycentric interpolation
            u = uv1[1]+(β*(uv2[1]-uv1[1]))+(γ*(uv3[1]-uv1[1]))
            v = uv1[2]+(β*(uv2[2]-uv1[2]))+(γ*(uv3[2]-uv1[2]))
            uvs = Vec2(u,v)
            return HitRecord(t, inter, norm, uvs, object)
        end
    end
end
end # module Scenes


# TODO this file should eventually go away and be folded into the
# infrastructure for all the other samplers.
#
# This samples the funnel distribution, then tests that the inferences are correct.
# Can trigger the warnings by modifying funnel_Lpstar; e.g. make the mean 0.8
# and it will detect the inferred mean is so far from 0 that there's a problem.
#
#Example output:
#julia> include("test_ss.jl")
#ls2 ess 2324.94, out of raw 99990
#ls2 mean 0.015439, std 2.994157, meanse 0.062097, z 0.248631
#ls2 var 8.964974, varse 0.262998, z -0.133179

include("../src/samplers/slice_sample.jl")

#log σ² ~ N(0, 3²)
#x_d ~ N(0, σ²) d=1...D

function funnel_Lpstar(xx)
    # First element is log(sigma^2)
    # Rest of the elements are x_d
    s2 = exp(xx[1])
    nx = length(xx) - 1
    -0.5*((xx[1]-0.0)^2/9.0 + dot(xx[2:end], xx[2:end])/s2 + nx*log(s2))
end

# log(s2)=10 is a poor initializer; and 100 takes a long time to get out
samples = slice_sample(funnel_Lpstar, [10, zeros(4)], 100000)
ls2 = samples[11:end,1]
#using Winston; plot(samples[:,1])

##ess = var(ls2) / MCMC.mcvar_imse(ls2) # This is an attempt to do the same thing as ess.jl.  However, it somewhat disagrees with coda effectiveSize() but I don't know if it's supposed to or I'm doing it wrong

# This is a hack, but MCMC.mcvar_imse is slow for big N
writedlm("bla",ls2)
ess = readall(`Rscript -e 'library(coda); cat(effectiveSize(scan("bla")))'`)
ess = parsefloat(ess)

# This part should instead have an MCMCChain object so it
# can just use ESS, quantile, etc. other calculations as implemented in
# stats/*.jl

@printf("ls2 ess %.2f, out of raw %d\n", ess, length(ls2))

zscore = mean(ls2) / (std(ls2)/sqrt(ess))
@printf("ls2 mean %f, std %f, meanse %f, z %f\n", mean(ls2), std(ls2), std(ls2)/sqrt(ess), zscore)
if abs(zscore) > 3
    warn("log(s2) posterior mean estimate differs from truth by more than 3 MC-stderr. Chance of this should be <0.3%.")
end

varse = var(ls2) * sqrt(2/(ess-1))
zscore = (var(ls2)-9) / varse
@printf("ls2 var %f, varse %f, z %f\n", var(ls2), varse, zscore)
if abs(zscore) > 3
    warn("log(s2) posterior variance estimate differs from truth by more than 3 MC-stderr. Chance of this should be <0.3%.")
end


const SUDOKU_BOX_SIZE = 3
const SUDOKU_GROUP_SIZE = SUDOKU_BOX_SIZE * SUDOKU_BOX_SIZE
const SUDOKU_CANDIDATE_SIZE = SUDOKU_GROUP_SIZE
const SUDOKU_CELL_SIZE = SUDOKU_GROUP_SIZE * SUDOKU_GROUP_SIZE
const SUDOKU_NO_CANDIDATES = 0
const SUDOKU_ALL_CANDIDATES = ((1 << SUDOKU_CANDIDATE_SIZE) - 1)
const SudokuCandidate=UInt8
const SudokuIndex=Int64
const SudokuMap=Array{SudokuCandidate, 2}

function overwrite_sudokumap(sudoku_map::SudokuMap, cell_index::SudokuIndex, candidate_index::SudokuIndex)
    sudoku_map[:, cell_index] .= zero(sudoku_map[1,1])
    sudoku_map[candidate_index, cell_index] = one(sudoku_map[1,1])
end

function string_to_sudokumap(line::String)
    sudoku_map = SudokuMap(ones(SudokuCandidate, SUDOKU_CANDIDATE_SIZE, SUDOKU_CELL_SIZE))
    for (index, c) in enumerate(line)
        if index > SUDOKU_CELL_SIZE
            break
        end

        num = Int(c) - Int('0')
        if num > 0 && num <= SUDOKU_CANDIDATE_SIZE
            overwrite_sudokumap(sudoku_map, index, num)
        end
    end
    sudoku_map
end

function to_string_sudokumap(sudoku_map::SudokuMap, one_line::Bool)
    result_str = ""

    for cell_index in 1:SUDOKU_CELL_SIZE
        for candidate in 1:SUDOKU_CANDIDATE_SIZE
            if sudoku_map[candidate, cell_index] != zero(sudoku_map[1,1])
                result_str *= string(candidate)
            end
        end

        if !one_line
            result_str *= ":"
        end

        if cell_index % SUDOKU_GROUP_SIZE == zero(sudoku_map[1,1]) && !one_line
            result_str *= "\n"
        end
    end

    result_str
end

function is_solved_group(sudoku_group)
    all(count -> (count == one(sudoku_group[1,1])), sum(sudoku_group, dims=(1))) && all(count -> (count == one(sudoku_group[1,1])), sum(sudoku_group, dims=(2)))
end

function is_solved_sudoku_rows(sudoku_map::SudokuMap)
    all(row ->
        (index = SUDOKU_GROUP_SIZE * (row - 1);
         is_solved_group(sudoku_map[:, index + 1: index + SUDOKU_GROUP_SIZE])),
        1:SUDOKU_GROUP_SIZE)
end

function is_solved_sudoku_columns(sudoku_map::SudokuMap)
    all(column ->
        (indexes = collect(column: SUDOKU_GROUP_SIZE: column + SUDOKU_GROUP_SIZE * (SUDOKU_GROUP_SIZE - 1));
         is_solved_group(view(sudoku_map, :, indexes))),
        1:SUDOKU_GROUP_SIZE)
end

function is_solved_sudoku_box_group(sudoku_map::SudokuMap, out_box_row::SudokuIndex, out_box_column::SudokuIndex)
    group = zeros(SudokuCandidate, SUDOKU_GROUP_SIZE, SUDOKU_CANDIDATE_SIZE)
    for in_box_index in 1:SUDOKU_BOX_SIZE
        base = (out_box_column - 1) * SUDOKU_BOX_SIZE
        base += (out_box_row - 1) * SUDOKU_GROUP_SIZE * SUDOKU_BOX_SIZE
        base += (in_box_index - 1) * SUDOKU_GROUP_SIZE

        sources = collect(1:SUDOKU_BOX_SIZE)
        sources .+= base
        dest = collect(1:SUDOKU_BOX_SIZE)
        dest .+= (in_box_index - 1) * SUDOKU_BOX_SIZE

        group[:, dest] = view(sudoku_map, :, sources)
    end

    is_solved_group(group)
end

function is_solved_sudoku_boxes(sudoku_map::SudokuMap)
    for out_box_row in 1:SUDOKU_BOX_SIZE
        for out_box_column in 1:SUDOKU_BOX_SIZE
            if !is_solved_sudoku_box_group(sudoku_map, out_box_row, out_box_column)
                return false
            end
        end
    end
    true
end

function is_solved_sudokumap(sudoku_map::SudokuMap)
    is_solved_sudoku_rows(sudoku_map) && is_solved_sudoku_columns(sudoku_map) && is_solved_sudoku_boxes(sudoku_map)
end

function is_consistent_sudokumap(sudoku_map::SudokuMap)
    all(count -> (count > zero(sudoku_map[1,1])), sum(sudoku_map, dims=(1)))
end

function has_unique_candidate_sudokumap(sudoku_map::SudokuMap, cell_index::SudokuIndex)
    sum(sudoku_map[:, cell_index]) == one(sudoku_map[1,1])
end

function set_unused_candidates_sudokumap(sudoku_map::SudokuMap, cell_index::SudokuIndex, other_index::SudokuIndex, candidates::Array{SudokuCandidate,1})
    if other_index != cell_index && has_unique_candidate_sudokumap(sudoku_map, other_index)
        candidates[:] .|= sudoku_map[:, other_index]
    end
end

function find_unique_candidates_sudokumap(sudoku_map::SudokuMap)
    for cell_index in 1:SUDOKU_CELL_SIZE
        if has_unique_candidate_sudokumap(sudoku_map, cell_index)
           continue
        end

        row = div(cell_index - 1, SUDOKU_GROUP_SIZE)
        column = (cell_index - 1) % SUDOKU_GROUP_SIZE
        box_base = div(row, SUDOKU_BOX_SIZE) * SUDOKU_BOX_SIZE * SUDOKU_GROUP_SIZE
        box_base += div(column, SUDOKU_BOX_SIZE) * SUDOKU_BOX_SIZE + 1

        candidates = zeros(SudokuCandidate, SUDOKU_CANDIDATE_SIZE)
        for index in 1:SUDOKU_GROUP_SIZE
            in_row_index = row * SUDOKU_GROUP_SIZE + index
            in_column_index = column + (index - 1) * SUDOKU_GROUP_SIZE + 1
            in_box_index = div(index - 1, SUDOKU_BOX_SIZE) * SUDOKU_GROUP_SIZE
            in_box_index += (index - 1) % SUDOKU_BOX_SIZE + box_base
            for other_index in [in_row_index, in_column_index, in_box_index]
                set_unused_candidates_sudokumap(sudoku_map, cell_index, other_index, candidates)
            end
        end

        for (index, candidate) in enumerate(candidates)
            if candidate != 0
                sudoku_map[index, cell_index] = zero(sudoku_map[1,1])
            end
        end

        if sum(sudoku_map[:, cell_index]) == 0
            return false
        end
    end

    true
end

function set_unused_candidates_sudokumap(sudoku_map::SudokuMap, cell_index::SudokuIndex, candidates::Array{SudokuCandidate,1})
    if sum(candidates) == SUDOKU_CANDIDATE_SIZE - 1
        index = findfirst(x -> x == 0, candidates)
        overwrite_sudokumap(sudoku_map, cell_index, index)
    end
end

function fill_unused_candidates_sudokumap(sudoku_map::SudokuMap)
    for cell_index in 1:SUDOKU_CELL_SIZE
        if has_unique_candidate_sudokumap(sudoku_map, cell_index)
           continue
        end

        cell_row = div((cell_index - 1), SUDOKU_GROUP_SIZE)
        candidates_row = zeros(SudokuCandidate, SUDOKU_CANDIDATE_SIZE)
        for column_index in 1:SUDOKU_GROUP_SIZE
            in_row_index = cell_row * SUDOKU_GROUP_SIZE + column_index
            if in_row_index != cell_index
                candidates_row[:] .|= sudoku_map[:, in_row_index]
            end
        end
        set_unused_candidates_sudokumap(sudoku_map, cell_index, candidates_row)

        cell_column = (cell_index - 1) % SUDOKU_GROUP_SIZE
        candidates_column = zeros(SudokuCandidate, SUDOKU_CANDIDATE_SIZE)
        for row_index in 1:SUDOKU_GROUP_SIZE
            in_column_index = cell_column + (row_index - 1) * SUDOKU_GROUP_SIZE + 1
            if in_column_index != cell_index
                candidates_column[:] .|= sudoku_map[:, in_column_index]
            end
        end
        set_unused_candidates_sudokumap(sudoku_map, cell_index, candidates_column)

        box_base = div(cell_row, SUDOKU_BOX_SIZE) * SUDOKU_BOX_SIZE * SUDOKU_GROUP_SIZE
        box_base += div(cell_column, SUDOKU_BOX_SIZE) * SUDOKU_BOX_SIZE + 1
        candidates_box = zeros(SudokuCandidate, SUDOKU_CANDIDATE_SIZE)
        for box_index in 1:SUDOKU_GROUP_SIZE
            in_box_index = div(box_index - 1, SUDOKU_BOX_SIZE) * SUDOKU_GROUP_SIZE
            in_box_index += (box_index - 1) % SUDOKU_BOX_SIZE + box_base
            if in_box_index != cell_index
                candidates_box[:] .|= sudoku_map[:, in_box_index]
            end
        end
        set_unused_candidates_sudokumap(sudoku_map, cell_index, candidates_box)

        if sum(sudoku_map[:, cell_index]) == zero(sudoku_map[1, 1])
            return false
        end
    end

    true
end

function find_backtracking_target(sudoku_map::SudokuMap)
    min_count = SUDOKU_CELL_SIZE
    found = false
    target_index = 0

    candidate_counts = sum(sudoku_map, dims=(1))
    for (index, count) in enumerate(candidate_counts)
        if count <= one(sudoku_map[1,1])
            continue
        end

        if min_count > count
            min_count = count
            found = true
        end
    end

    extract_row = (row -> view(candidate_counts,
                               ((row - 1) * SUDOKU_GROUP_SIZE + 1):(row * SUDOKU_GROUP_SIZE)))
    if found
        row_counts = map(row -> count(n -> (n == min_count), extract_row(row)), 1:SUDOKU_GROUP_SIZE)
        row = argmax(row_counts)
        target_index = findfirst(x -> (x == min_count), extract_row(row)) + (row - 1) * SUDOKU_GROUP_SIZE
    end

    target_index, found
end

function solve_sudokumap(sudoku_map::SudokuMap)
    n_candidates = sum(sudoku_map)
    while true
        if find_unique_candidates_sudokumap(sudoku_map) == false
            return (sudoku_map, false)
        end

        if fill_unused_candidates_sudokumap(sudoku_map) == false
            return (sudoku_map, false)
        end

        new_n_candidates = sum(sudoku_map)
        if n_candidates == new_n_candidates
            break
        end

        n_candidates = new_n_candidates
    end

    if is_solved_sudokumap(sudoku_map)
        return (sudoku_map, true)
    end

    if !is_consistent_sudokumap(sudoku_map)
        return (sudoku_map, false)
    end

    target_index, found = find_backtracking_target(sudoku_map)
    if !found
        return (sudoku_map, false)
    end

    for (index, candidate) in enumerate(sudoku_map[:, target_index])
        if candidate == zero(sudoku_map[1,1])
            continue
        end

        new_map = deepcopy(sudoku_map)
        overwrite_sudokumap(new_map, target_index, index)
        result_map, result = solve_sudokumap(new_map)
        if result
            return (result_map, result)
        end
    end

    return (sudoku_map, false)
end

function main()
    num = 0
    one_line = true
    total_result = true

    if one_line
        println("Solving in Julia")
    end

    for line in eachline(stdin)
        puzzle = string_to_sudokumap(line)
        solution, result = solve_sudokumap(puzzle)
        result = is_solved_sudokumap(solution)
        total_result &= result

        one_line_result = to_string_sudokumap(solution, one_line)
        print(one_line_result)
        if one_line
            println("")
        else
            if result
                println("Solved\n")
            else
                println("Not solved\n")
            end
        end
        num += 1
    end

    if total_result
        println("All ", num, " cases passed.")
    else
        println("Some of ", num, " cases are not solved correctly.")
    end
end

main()


## Defines functions to be used for active learning

# Fetch upper triangle of adj matrix and create parameter vector
# Ref: https://gist.github.com/tpapp/87e5675b87dbdd7a9a840e546eb20fae
function vec_triu_loop(adj_matrix::Array)
    m, n = size(adj_matrix)
    m == n || throw(error("not square"))
    len_theta = n*(n+1) ÷ 2

    # Array to hold values of adjacency matrix
    theta = ones(len_theta)

    # Create array of tuples to hold edges of corresponding values in theta
    edges_theta = Array{Tuple{Int, Int}}(undef, len_theta)

    k = 0
    @inbounds for i in 1:n
        for j in 1:i
            theta[k + j] = adj_matrix[j, i]
            edges_theta[k+j] = (j,i)
        end
        k += i
    end
    return len_theta, theta, edges_theta
end

function remove_node_from_edge(edge::Tuple, node::Integer)
    #=
    Functionality (1,2)\{2} = 1. Note that (1,1)\{1} = 1 as desired here.
    =#
    if edge[1] == node
        return edge[2]
    elseif edge[2] == node
        return edge[1]
    else
        return nothing
    end
end

# Function to calculate the fisher information matrix for a given query
function fisher_information_query(adj_matrix::Array,theta::Array,edges_theta::Array,sigma::Array,len_theta::Integer,num_spins::Integer)
    #=
    Inputs:
    adjacency matrix of the graphical model (estimated) -> corresponds to your theta
    sigma which is the particular query
    =#

    # Define Fisher Information matrix
    I_x = zeros((len_theta, len_theta))

    # Start filling up the upper triangle of I_x
    for a in 1:len_theta
        for b in a:len_theta
            # Get common node if exists
            set_common_nodes_temp = intersect(edges_theta[a],edges_theta[b])

            # if s is empty do nothing (i.e., I_x[a,b] = 0) else
            if !(isempty(set_common_nodes_temp))
                for s in set_common_nodes_temp
                    ind1 = remove_node_from_edge(edges_theta[a],s)
                    ind2 = remove_node_from_edge(edges_theta[b],s)

                    # set shared spin s to 1 as A_s = \sum_{u in ∂s} J_{su} σ_u + J_{ss}
                    sigma_temp = copy(sigma)
                    sigma_temp[s] = 1
                    A_s = sum(sigma_temp[i]*adj_matrix[s,i] for i=1:num_spins)

                    # We pick from sigma_temp instead of sigma for scenarios such as (1,1)∩(1,2)
                    I_x[a,b] += (1.0/num_spins)*sigma_temp[ind1]*sigma_temp[ind2]*(1.0 - tanh(A_s)^2)
                end
            end
        end
    end

    # Now copy over the upper triangle to lower triangle
    D = Diagonal(I_x)
    I_x_sut = I_x - D
    I_x = I_x_sut + transpose(I_x_sut) + D
    return I_x
end

# Function to calculate the fisher information matrix for a given query distribution
function fisher_information_query_distrn(adj_matrix::Array,theta::Array,edges_theta::Array,X_σ::Array,q_σ::Array,len_theta::Integer,num_spins::Integer)
    #=
    Inputs:
    X_σ is the set of queries we are considering
    q_σ is the query distribution over these queries
    =#

    # Make sure the given query distribution is a valid pdf
    #if sum(q_σ)!=1
    #    error("Given query distribution is not a valid pdf")
    #end

    I_q = zeros((len_theta, len_theta))
    for ind_σ=1:length(X_σ)
        sigma0 = 2*digits(X_σ[ind_σ], base=2, pad=num_spins).-1
        I_q += q_σ[ind_σ]*fisher_information_query(adj_matrix,theta,edges_theta,sigma0,len_theta,num_spins)
    end

    return I_q
end

# Function to calculate the fisher information matrix for a given query distribution where q is a Variable (Convex.jl)
function fisher_information_query_distrn_obj(adj_matrix::Array,theta::Array,edges_theta::Array,X_σ::Array,q_σ::Variable,len_theta::Integer,num_spins::Integer)
    #=
    Inputs:
    X_σ is the set of queries we are considering
    q_σ is the query distribution over these queries
    =#

    I_q = zeros((len_theta, len_theta))
    for ind_σ=1:length(X_σ)
        sigma0 = 2*digits(X_σ[ind_σ], base=2, pad=num_spins).-1
        I_q += q_σ[ind_σ]*fisher_information_query(adj_matrix,theta,edges_theta,sigma0,len_theta,num_spins)
    end

    return I_q
end

# Function to calculate the fisher information matrix for a given query distribution
function active_learning_M(adj_matrix::Array,X_σ::Array,num_spins::Integer)
    #=
    Inputs:
    adj_matrix is the current estimate of the parameters
    X_σ is the set of queries we are considering
    q_σ is the query distribution over these queries

    Ref: https://www.juliaopt.org/Convex.jl/dev/types/
    =#

    # Create the parameter vector: θ = [J_11, J_12, J_13, ..., J1n, ..., J_{i,i+1}, ..., J_nn]
    len_theta, theta, edges_theta = vec_triu_loop(adj_matrix)

    # Create the variables required for the SDP
    len_X = length(X_σ)
    q = Variable(len_X)  # query distribution
    t = Variable(len_theta)  # auxillary variables

    # Constraints for the query distribution
    constraints = [q>=0, q<=1, sum(q)==1]

    # Constraints involving the matrices
    U = 1.0*Matrix(I,len_theta,len_theta)
    for i=1:len_theta
        constraints += [fisher_information_query_distrn_obj(adj_matrix, theta, edges_theta, X_σ, q, len_theta, num_spins) U[i,:]; transpose(U[i,:]) t[i]] ⪰ 0
        #constraints += [t[i] transpose(U[i,:]); U[i,:] fisher_information_query_distrn_obj(adj_matrix, theta, edges_theta, X_σ, q, len_theta, num_spins)] ⪰ 0
    end

    # Solve the SDP
    problem = minimize(sum(t), constraints)
    #solve!(problem, SCS.Optimizer)
    solve!(problem, Mosek.Optimizer)

    return q.value
end

# Function to calculate the entropy of a given query
function entropy_query(adj_matrix::Array,sigma::Array,num_spins::Integer)
    #=
    Inputs:
    adjacency matrix of the graphical model (estimated) -> corresponds to your theta
    sigma which is the particular query
    =#

    # Define entropy
    entropy = 0.0

    # Start filling up the upper triangle of I_x
    for k in 1:num_spins
        sigma_temp = copy(sigma)
        sigma_temp[k] = 1
        A_k = sum(sigma_temp[i]*adj_matrix[k,i] for i=1:num_spins)

        entropy += log(2*cosh(A_k)) - A_k*tanh(A_k)
    end

    return entropy
end

# Function to calculate the entropy for a given set of queries/configurations
function entropy_configs(adj_matrix::Array,X_σ::Array,num_spins::Integer)
    #=
    Inputs:
    adjacency matrix of the graphical model (estimated) -> corresponds to your theta
    X_σ which is the set of queries
    =#

    # Define entropy
    entropy_configs = zeros(length(X_σ))

    for i in 1:length(X_σ)
        sigma_temp = 2*digits(X_σ[i], base=2, pad=num_spins) .- 1
        entropy_configs[i] = (1/length(X_σ))*entropy_query(adj_matrix, sigma_temp, num_spins)
    end

    return entropy_configs
end

# Unit test (by visualization) to check if sampling of M-regime with query was done correctly
function unit_test_sampling_M_regime_query(samples_U::Array, spin_number::Integer)
    config_number = 2^spin_number
    spin_configs = []
    for i = 0:(config_number - 1)
        spin_tmp = 2*digits(i, base=2, pad=spin_number).-1
        push!(spin_configs, string(spin_tmp))
    end
    dict_spin_configs = Dict(spin_configs[i+1] => i for i = 0:(config_number-1))

    bins_U_sigma0 = Array{Int}(undef, size(samples_U)[1],2)
    for i=1:size(samples_U)[1]
        sigma0_temp = samples_U[i,3:(n+2)]
        config0_temp = dict_spin_configs[string(sigma0_temp)]
        bins_U_sigma0[i,2] = config0_temp
    end

    bins_U_sigma0[:,1] = copy(samples_U[:,1])

    N_sigma0_U = bins_U_sigma0[:,1]
    config_sigma0_U = bins_U_sigma0[:,2]

    q_σ0 = zeros(config_number)
    for i=1:config_number
        idx_temp = findall(isequal(i-1),config_sigma0_U)
        q_σ0[i] = sum(N_sigma0_U[idx_temp])
    end
    q_σ0 = q_σ0/sum(q_σ0)

    return q_σ0
end


#Parameter file for instiating equation 31 in Kloden & Platen
include("problems.jl")

const avec = [0.8, 0.4, 0.9, 0.4, 0.2, 0.21, 0.6, 0.43, 0.21, -0.1]
const prob = Problem31(avec)

#Must be between -1 and 1!
const x0 = [-0.5, 0.5, 0.3, -0.2, 0.7, 0.85, 0.0, -0.3, -0.34, 0.62]
t = linspace(0., 1., 5001)


## Basis
"""
    Basis

A `basis` function is an element of a particular basis for a function space.
Every continuous function in the function space can be represented as a
linear combination of basis functions:
```math
f(x_1, \\cdots, x_n) = \\sum_i \\beta_i\\text{basis}_i(x_1, \\cdots, x_n)
```
So far the only implemented basis is [`Monomial`](@ref).
"""
abstract type Basis end

# evaluate : Matrix -> vector
function evaluate(X::Matrix{Float64}, basis::Basis)
    X_views, N, p = matrix2subviews(X)
    y = Vector{Float64}(N)
    evaluate!(y, X_views, basis, N, p)
    return y
end

# evaluate : Matrix -> Matrix
function evaluate{T<:Basis}(X::Matrix{Float64}, basis::Vector{T})
    X_views, N, p = matrix2subviews(X)
    q = length(basis)
    Y = Matrix{Float64}(N, q)
    Y_views, _, _ = matrix2subviews(Y)
    aux = Vector{Float64}(p)
    evaluate!(Y_views, X_views, basis, N, p)
    Y = subviews2matrix(Y_views, N, q)
    return Y
end

# evaluate! : Matrix -> Matrix (With Overide)
function evaluate!{T <: Basis}(
        Y::Vector{Vector{Float64}}, X::Vector{Vector{Float64}},
        basis::Vector{T}, N::Int, p::Int)
    for (i, y) in enumerate(Y)
        evaluate!(y, X, basis[i], N, p)
    end
end

# Print basis
function Base.show(io::IO, basis::Basis)
    names = [string("x_", i) for i=1:basis.n]
    str = function_display(basis, names)
    print(io, str)
end

# Separate basis terms containing certain terms
function detach_terms{T<:Basis}(basis::Vector{T},
                      list_of_inputs::BitVector)
    n_basis = length(basis)
    basis_containing_inputs = trues(n_basis)
    for (i, m) in enumerate(basis)
        if contains_inputs(basis[i], list_of_inputs)
            basis_containing_inputs[i] = false
        end
    end
    return (basis[basis_containing_inputs],
            basis[.~basis_containing_inputs])
end


using ADCME
using PyCall
using LinearAlgebra
using PyPlot
using Random
using DelimitedFiles
Random.seed!(233)

function mpi_tensor_transpose(row,col,ncol,val,n,rank,nt)
    require_mpi()
    mpi_tensor_transpose_ = load_op_and_grad("./build/libMPITensor.so","mpi_tensor_transpose", multiple=true)
    row,col,ncol,val,n,rank,nt = convert_to_tensor(Any[row,col,ncol,val,n,rank,nt], [Int32,Int32,Int32,Float64,Int64,Int64,Int64])
    indices, vals = mpi_tensor_transpose_(row,col,ncol,val,n,rank,nt)
end

mpi_init()

if mpi_rank()==0
    row = [0;1]
    col = [0;1;2;3;0;1;2;3]
    ncol = [4;4]
    val = Float64[1;2;3;4;5;6;7;8]
else
    row = [2,3]
    col = [0;1;2;3;0;1;2;3]
    ncol = [4;4]
    val = Float64[1;2;3;4;5;6;7;8] .+ 8.0
end

n = 2
rank = mpi_rank()
nt = 4

# TODO: specify your input parameters
u = mpi_tensor_transpose(row,col,ncol,val,n,rank,nt)
sess = Session(); init(sess)

indices, vals = run(sess, u)
sleep(mpi_rank())
println("===================================================")
writedlm(stdout, indices)
writedlm(stdout, vals)

mpi_finalize()


using jInv.Utils
using Test

# testing variousUtils

A = rand(0.:10., 10) + im*rand(0.:10., 10)

@test real(A) == complex2real(A)[1:2:end]
@test imag(A) == complex2real(A)[2:2:end]
@test A == real2complex(complex2real(A))


mutable struct AbstractTree
    node_counter::Int
    nodes::Vector{AbstractNode}
    processed::Vector{AbstractNode}
    best_bound::Real     # e.g., lower bound for minimization
    best_incumbent::Real # e.g., upper bound for minimization

    function AbstractTree(node_counter::Int = 0)
        return new(node_counter, [], [], -Inf, Inf)
    end
end

function run(tree::AbstractTree)
    while !termination(tree)
        node = next_node(tree)
        update_best_bound!(tree)

        bound!(node)
        heuristics!(node)
        processed!(tree, node)
        branch!(tree, node)
    end
end

# basic termination function
termination(tree::AbstractTree) = isempty(tree)
isempty(tree::AbstractTree) = Base.isempty(tree.nodes)

# add node to tree
function push!(tree::AbstractTree, node::AbstractNode)
    tree.node_counter += 1
    node.id = tree.node_counter
    Base.push!(tree.nodes, node)
    sort!(tree)
end

# add a set of nodes to tree
function push!(tree::AbstractTree, nodes::Vector{T}) where T<:AbstractNode
    for node in nodes
        push!(tree, node)
    end
end

# mark node as processed
processed!(tree::AbstractTree, node::AbstractNode) = Base.push!(tree.processed, node)

function update_best_bound!(tree::AbstractTree)
    if !isempty(tree)
        tree.best_bound = Base.minimum([node.bound for node in tree.nodes])
    end
end

function branch!(tree::AbstractTree, node::AbstractNode)
    children = branch(node)
    push!(tree, children)
end

# return the next search node
function next_node(tree::AbstractTree)
    sort!(tree)
    node = Base.pop!(tree.nodes)
    apply_changes!(node)
    return node
end

# best bound
sort!(tree::AbstractTree) = Base.sort!(tree.nodes, by=x->x.bound, rev=true)


"""
    truncated(d0::UnivariateDistribution; [lower::Real], [upper::Real])
    truncated(d0::UnivariateDistribution, lower::Real, upper::Real)

A _truncated distribution_ `d` of a distribution `d0` to the interval
``[l, u]=```[lower, upper]` has the probability density (mass) function:

```math
f(x; d_0, l, u) = \\frac{f_{d_0}(x)}{P_{Z \\sim d_0}(l \\le Z \\le u)}, \\quad x \\in [l, u],
```
where ``f_{d_0}(x)`` is the probability density (mass) function of ``d_0``.

The function throws an error if ``l > u``.

```julia
truncated(d0; lower=l)           # d0 left-truncated to the interval [l, Inf)
truncated(d0; upper=u)           # d0 right-truncated to the interval (-Inf, u]
truncated(d0; lower=l, upper=u)  # d0 truncated to the interval [l, u]
truncated(d0, l, u)              # d0 truncated to the interval [l, u]
```

The function falls back to constructing a [`Truncated`](@ref) wrapper.

# Implementation

To implement a specialized truncated form for distributions of type `D`, one or more of the
following methods should be implemented:
- `truncated(d0::D, l::T, u::T) where {T <: Real}`: interval-truncated
- `truncated(d0::D, ::Nothing, u::Real)`: right-truncated
- `truncated(d0::D, l::Real, u::Nothing)`: left-truncated
"""
function truncated end
function truncated(d::UnivariateDistribution, l::Real, u::Real)
    return truncated(d, promote(l, u)...)
end
function truncated(
    d::UnivariateDistribution;
    lower::Union{Real,Nothing}=nothing,
    upper::Union{Real,Nothing}=nothing,
)
    return truncated(d, lower, upper)
end
function truncated(d::UnivariateDistribution, ::Nothing, u::Real)
    # (log)ucdf = (log)tp = (log) P(X ≤ u) where X ~ d
    logucdf = logtp = logcdf(d, u)
    ucdf = tp = exp(logucdf)

    Truncated(d, promote(oftype(float(u), -Inf), u, zero(ucdf), ucdf, tp, logtp)...)
end
function truncated(d::UnivariateDistribution, l::Real, ::Nothing)
    # (log)lcdf = (log) P(X < l) where X ~ d
    loglcdf = if value_support(typeof(d)) === Discrete
        logsubexp(logcdf(d, l), logpdf(d, l))
    else
        logcdf(d, l)
    end
    lcdf = exp(loglcdf)

    # (log)tp = (log) P(l ≤ X) where X ∼ d
    logtp = log1mexp(loglcdf)
    tp = exp(logtp)

    Truncated(d, promote(l, oftype(float(l), Inf), lcdf, one(lcdf), tp, logtp)...)
end
truncated(d::UnivariateDistribution, ::Nothing, ::Nothing) = d
function truncated(d::UnivariateDistribution, l::T, u::T) where {T <: Real}
    l <= u || error("the lower bound must be less or equal than the upper bound")

    # (log)lcdf = (log) P(X < l) where X ~ d
    loglcdf = if value_support(typeof(d)) === Discrete
        logsubexp(logcdf(d, l), logpdf(d, l))
    else
        logcdf(d, l)
    end
    lcdf = exp(loglcdf)

    # (log)ucdf = (log) P(X ≤ u) where X ~ d
    logucdf = logcdf(d, u)
    ucdf = exp(logucdf)

    # (log)tp = (log) P(l ≤ X ≤ u) where X ∼ d
    logtp = logsubexp(loglcdf, logucdf)
    tp = exp(logtp)

    Truncated(d, promote(l, u, lcdf, ucdf, tp, logtp)...)
end

"""
    Truncated

Generic wrapper for a truncated distribution.
"""
struct Truncated{D<:UnivariateDistribution, S<:ValueSupport, T <: Real} <: UnivariateDistribution{S}
    untruncated::D      # the original distribution (untruncated)
    lower::T      # lower bound
    upper::T      # upper bound
    lcdf::T       # cdf of lower bound (exclusive): P(X < lower)
    ucdf::T       # cdf of upper bound (inclusive): P(X ≤ upper)

    tp::T         # the probability of the truncated part, i.e. ucdf - lcdf
    logtp::T      # log(tp), i.e. log(ucdf - lcdf)
    function Truncated(d::UnivariateDistribution, l::T, u::T, lcdf::T, ucdf::T, tp::T, logtp::T) where {T <: Real}
        new{typeof(d), value_support(typeof(d)), T}(d, l, u, lcdf, ucdf, tp, logtp)
    end
end

### Constructors of `Truncated` are deprecated - users should call `truncated`
@deprecate Truncated(d::UnivariateDistribution, l::Real, u::Real) truncated(d, l, u)

params(d::Truncated) = tuple(params(d.untruncated)..., d.lower, d.upper)
partype(d::Truncated) = partype(d.untruncated)
Base.eltype(::Type{Truncated{D, S, T} } ) where {D, S, T} = T

### range and support

islowerbounded(d::Truncated) = islowerbounded(d.untruncated) || isfinite(d.lower)
isupperbounded(d::Truncated) = isupperbounded(d.untruncated) || isfinite(d.upper)

minimum(d::Truncated) = max(minimum(d.untruncated), d.lower)
maximum(d::Truncated) = min(maximum(d.untruncated), d.upper)

function insupport(d::Truncated{D,<:Union{Discrete,Continuous}}, x::Real) where {D<:UnivariateDistribution}
    return d.lower <= x <= d.upper && insupport(d.untruncated, x)
end

### evaluation

quantile(d::Truncated, p::Real) = quantile(d.untruncated, d.lcdf + p * d.tp)

function pdf(d::Truncated, x::Real)
    result = pdf(d.untruncated, x) / d.tp
    return d.lower <= x <= d.upper ? result : zero(result)
end

function logpdf(d::Truncated, x::Real)
    result = logpdf(d.untruncated, x) - d.logtp
    return d.lower <= x <= d.upper ? result : oftype(result, -Inf)
end

function cdf(d::Truncated, x::Real)
    result = (cdf(d.untruncated, x) - d.lcdf) / d.tp
    return if x < d.lower
        zero(result)
    elseif x >= d.upper
        one(result)
    else
        result
    end
end

function logcdf(d::Truncated, x::Real)
    result = logsubexp(logcdf(d.untruncated, x), log(d.lcdf)) - d.logtp
    return if x < d.lower
        oftype(result, -Inf)
    elseif x >= d.upper
        zero(result)
    else
        result
    end
end

function ccdf(d::Truncated, x::Real)
    result = (d.ucdf - cdf(d.untruncated, x)) / d.tp
    return if x <= d.lower
        one(result)
    elseif x > d.upper
        zero(result)
    else
        result
    end
end

function logccdf(d::Truncated, x::Real)
    result = logsubexp(logccdf(d.untruncated, x), log1p(-d.ucdf)) - d.logtp
    return if x <= d.lower
        zero(result)
    elseif x > d.upper
        oftype(result, -Inf)
    else
        result
    end
end

## random number generation

function _rand!(rng::AbstractRNG, d::Truncated)
    d0 = d.untruncated
    if d.tp > 0.25
        while true
            r = _rand!(rng, d0)
            if d.lower <= r <= d.upper
                return r
            end
        end
    else
        return quantile(d0, d.lcdf + rand(rng) * d.tp)
    end
end

## show

function show(io::IO, d::Truncated)
    print(io, "Truncated(")
    d0 = d.untruncated
    uml, namevals = _use_multline_show(d0)
    uml ? show_multline(io, d0, namevals) :
          show_oneline(io, d0, namevals)
    if d.lower > -Inf
        if d.upper < Inf
            print(io, "; lower=$(d.lower), upper=$(d.upper))")
        else
            print(io, "; lower=$(d.lower))")
        end
    elseif d.upper < Inf
        print(io, "; upper=$(d.upper))")
    else
        print(io, ")")
    end
    uml && println(io)
end

_use_multline_show(d::Truncated) = _use_multline_show(d.untruncated)

### specialized truncated distributions

include(joinpath("truncated", "normal.jl"))
include(joinpath("truncated", "exponential.jl"))
include(joinpath("truncated", "uniform.jl"))
include(joinpath("truncated", "loguniform.jl"))