text:
using TensorOperations
using EllipsisNotation
using DataStructures
import Base: size, length, getindex

export MPSState, size, length, getindex

"""Data structure to represent an MPS state"""
struct MPSState{T, N} <: AbstractArray{T, N}
    # stores the data from each mps node
    data_tensors::Array{Array{T}, 2}
    # used to store the meta data for each mps node with index locations
    nodes::Array{Node, 1}
    # the order mapping of output qubits to classical measurement bits
    ordering::Array{Int, 1}
end

"""
    function MPSState(network::TensorNetworkCircuit, mps_nodes::Array{Symbol, 1})

Constructor for MPS state which takes an tensor network circuit and list of mps
nodes and populates fields which enables one ot efficiently retrieve amplitudes
via the array interface
"""
function MPSState(::Type{T}, network::TensorNetworkCircuit, mps_nodes::Array{Symbol, 1}) where {T <: Number}
    n = length(mps_nodes)
    # we create arrays to store node data structure which contains the meta data
    # and data arrays which we create for each mps tensor
    nodes = Array{Node, 1}(undef, n)
    data_tensors = Array{Array{T}, 2}(undef, (n, 2))
    output_positions = Array{Int64, 1}(undef, length(mps_nodes))
    for i = 1:length(mps_nodes)
        node_label = mps_nodes[i]
        # output_index = network.output_qubits[i]
        output_positions[i] = findfirst(x -> x in network.output_qubits,
                                    network.nodes[node_label].indices)
        other_positions = collect(
                    setdiff(
                        OrderedSet(1:length(network.nodes[node_label].indices)),
                        output_positions[i])
                    )
        tensor = permutedims(
                    load_tensor_data(network, node_label),
                                     (output_positions[i], other_positions...))
        nodes[i] = Node(network.nodes[node_label].indices[[other_positions...]],
                        network.nodes[node_label].dims[[other_positions...]],
                        Symbol("n_$i"))
        data_tensors[i, :] = [tensor[x,..] for x in 1:2]
    end
    ordering = Array{Int, 1}(undef, n)
    for (i, e) in zip(output_positions, network.qubit_ordering)
        ordering[e] = i
    end
    MPSState{T, n}(data_tensors, nodes, ordering)
end

function MPSState(network::TensorNetworkCircuit, mps_nodes::Array{Symbol, 1})
    MPSState(ComplexF32, network::TensorNetworkCircuit, mps_nodes::Array{Symbol, 1})
end

"""
    function size(a::MPSState{T, N}) where {T, N}

Return the dimensions in each dimension
TODO: generalise to N != 2
"""
function size(a::MPSState{T, N}) where {T, N}
    Tuple([2 for _ in 1:N])
end

"""
    function length(a::MPSState{T, N}) where {T, N}

Return the total size of the all amplitudes
TODO: generalise to N != 2
"""
function length(a::MPSState{T, N}) where {T, N}
    2^N
end

"""
    function getindex(a::MPSState{T, N}, i::Vararg{Int, N}) where {T, N}

Override the getindex function for abstract array which allows to use the
array notation [i_1, i_2,..., i_N] to access amplitudes
"""
function getindex(a::MPSState{T, N}, i::Vararg{Int, N}) where {T, N}
    n1 = a.nodes[1]
    conf_map = x -> (x % 2 == 1) ? 1 : 2
    n1_data = a.data_tensors[1, conf_map(i[a.ordering[1]])]
    for idx in 2:length(a.nodes)
        n2 = a.nodes[idx]
        common, remaining = sort_indices(n1, n2)
        indices = create_ncon_indices(n1, n2, common, remaining)
        n1 = Node(remaining, ones(Int64, length(remaining)), :n1) # TODO: correct dummy dims
        n2_data = a.data_tensors[idx, conf_map(i[a.ordering[idx]])]
        n1_data = ncon([n1_data, n2_data],  indices)
    end
    scalar(n1_data)
end

"""
    function getindex(a::MPSState, conf::String)

Override the getindex method to accept strings to enable retrieval of amplitudes
from strings
"""
function getindex(a::MPSState, conf::String)
    a[[x == '0' ? 1 : 2 for x in conf]...]
end


text:
# Load sample, Calculate and Plot STA

# load sample (stim and ρ)
include("sample.jl")

# Calculate STA
include("sta.jl")
num_timesteps = 150

sta_val = sta(stim, ρ, num_timesteps)

# Plot
using PyPlot

time = 1-length(sta_val):0

plot(time*2, sta_val, color="red", linewidth=1.0, linestyle="-")

xlabel("Time (ms)");
ylabel("Stimulus");
title("Spike Triggered Average")

ext: jl

max_stars_repo_name: voutasaurus/comp-neuro

max_issues_repo_name: voutasaurus/comp-neuro

max_forks_repo_name: voutasaurus/comp-neuro


text:
# VCF
# ===

include("record.jl")
include("metainfo.jl")
include("header.jl")
include("reader.jl")
include("writer.jl")


text:
Term(i...) = TermCombination(Set([Dict(j=>1 for j in i)]))

@test htesttag(x->x, [1,2]) == Input()
@test htesttag(x->x[1], [1,2]) == Term(1)

# Tuple / struct
@test htesttag(x->(x[1],x[2])[2], [1,2]) == Term(2)

# 1-arg linear
@test htesttag(x->deg2rad(x[1]), [1,2]) == Term(1)

# 1-arg nonlinear
@test htesttag(x->sin(x[1]), [1,2]) == (Term(1) * Term(1))

# 2-arg (true,true,true)
@test htesttag(x->x[1]+x[2], [1,2]) == Term(1)+Term(2)

# 2-arg (true,true, false)
@test htesttag(x->x[1]*x[2], [1,2]) == Term(1)*Term(2)

# 2-arg (true,false,false)
@test htesttag(x->x[1]/x[2], [1,2]) == Term(1)*Term(2)*Term(2)

# 2-arg (false,true,false)
@test htesttag(x->x[1]\x[2], [1,2]) == Term(1)*Term(1)*Term(2)

# 2-arg (false,false,false)
@test htesttag(x->hypot(x[1], x[2]), [1,2]) == (Term(1) + Term(2)) * (Term(1) + Term(2))

### Array operations

# copy
@test htesttag(x->copy(x)[1], [1,2]) == Term(1)
@test htesttag(x->x[:][1], [1,2]) == Term(1)
@test htesttag(x->x[1:1][1], [1,2]) == Term(1)

# tests `iterate`
function mysum(x)
    s = 0
    for a in x
        s += a
    end
    s
end
@test htesttag(mysum, [1,2]).terms == (Term(1) + Term(2)).terms
@test htesttag(mysum, [1,2.]).terms == (Term(1) + Term(2)).terms

using LinearAlgebra

# integer dot product falls back to generic
@test htesttag(x->dot(x,x), [1,2,3]) == sum(Term(i)*Term(i) for i=1:3)

# reroutes to generic implementation (blas.jl)
@test htesttag(x->dot(x,x), [1,2,3.]) == sum(Term(i)*Term(i) for i=1:3)
@test htesttag(x->x'x, [1,2,3.]) == sum(Term(i)*Term(i) for i=1:3)

# broadcast
@test htesttag(x->sum(x[1] .+ x[2]), [1,2,3.]) == Term(1) + Term(2)
@test htesttag(x->sum(x .+ x), [1,2,3.]) == sum(Term(i) for i=1:3)
@test htesttag(x->sum(x .* x), [1,2,3.]) == sum(Term(i)*Term(i) for i=1:3)


text:
"""
    MITgcm(p::String)

Run MITgcm; download it if needed.
"""
function MITgcm(p::String="./",c::Cmd=`./testreport -t hs94.cs-32x32x5`)
   d=pwd()
   cd(p)
   test=~isdir("MITgcm")
   test ? run(`git clone https://github.com/MITgcm/MITgcm`) : nothing
   cd("MITgcm/verification/")
   run(c)
   cd(d)
end

"""
    StartWorkers(nwrkrs::Int)

Start workers if needed.
"""
function StartWorkers(nwrkrs::Int)
   set_workers = nwrkrs
   nworkers() < set_workers ? addprocs(set_workers) : nothing
   nworkers()
end

"""
    TaskDriver(indx,fn)

Broacast / distribute task (fn; e.g. task1_loop) over indices (indx; e.g. file indices)

Examples:

```
using ClimateTasks, Distributed, SparseArrays
TaskDriver(1,task1_loop)

StartWorkers(4)
@everywhere using ClimateTasks, SparseArrays
TaskDriver(1:4,task1_loop)
```

Visualize results:

```
using FortranFiles, Plots
k=1
recl=720*360*4
fil="diags_interp/ETAN/ETAN.0000000732.data"
f =  FortranFile(fil,"r",access="direct",recl=recl,convert="big-endian")
tmp=read(f,rec=k,(Float32,(720,360))); close(f)
heatmap(tmp)
```
"""
function TaskDriver(indx::Union{UnitRange{Int},Array{Int,1},Int},fn::Function)
    i=collect(indx)
    length(i)>1 ? i=distribute(i) : nothing
    isa(i,DArray) ? println(i.indices) : nothing
    fn.(i)
end

"""
    task1_loop(indx::Int)

Interpolate all variables for one record
"""
function task1_loop(indx::Int)
    task=YAML.load(open("task.yml"))
    M=load(task["Specs"]["file"])
    MetaFile=task1_loop(indx,M)
end

"""
    task1_loop(indx,M)

Loop over a subset of model output files (`filList[indx]`), apply
`MatrixInterp` (`M`) as a postprocessing step, and write the result
to file (one subfolder for each variable)
"""
function task1_loop(indx,M)
   task=YAML.load(open("task.yml"))
   dirIn=task["InputDir"][1]
   filIn=task["InputFile"][1]
   dirOut=task["OutputDir"]
   !isdir(dirOut) ? mkdir(dirOut) : nothing
   MTRX=M["MTRX"]
   msk=M[task["Specs"]["mask"]]
   siz=Tuple(task["OutputSize"])

   tmp1=readdir(dirIn)
   tmp1=filter(x -> occursin(filIn,x),tmp1)
   filList=filter(x -> occursin(".data",x),tmp1)
   maximum(indx)>length(filList) ? error("missing files: "*filIn*"*") : nothing
   filList=filList[indx]

   !isa(filList,Array) ? filList=[filList] : nothing
   nf=length(filList)
   MetaFile=filList[1]
   MetaFile=dirIn*MetaFile[1:end-5]*".meta"
   MetaFile=MetaFileRead(MetaFile)

   nv=length(MetaFile["fldList"])
   nd=MetaFile["nDims"]
   dims=Int.(MetaFile["dimList"][:,1])
   prec=MetaFile["dataprec"]

   for ff=1:nf
      fil=dirIn*filList[ff]
      #println(fil)
      fid = open(fil)
      for vv=1:nv
         tmp=Array{Float32,2}(undef,(90,1170))
         nd==3 ? tmp=Array{Float32,3}(undef,(90,1170,50)) : nothing
         read!(fid,tmp)
         tmp = hton.(tmp)
         !ismissing(msk) ? tmp=tmp.*msk : nothing
         tmp=MatrixInterp(tmp,MTRX,siz)
         tmp=Float32.(tmp)
         #
         filOut=dirOut*strip(MetaFile["fldList"][vv])*"/"
         !isdir(filOut) ? mkdir(filOut) : nothing
         filOut=filOut*strip(MetaFile["fldList"][vv])
         filOut=filOut*fil[length(dirIn*filIn)+1:end]
         #println(filOut)
         #
         nd==3 ? recl=720*360*50*4 : recl=720*360*4
         f =  FortranFile(filOut,"w",access="direct",recl=recl,convert="big-endian")
         write(f,rec=1,tmp)
         close(f)
         #to re-read file:
         #f =  FortranFile(filOut,"r",access="direct",recl=recl,convert="big-endian");
         #tmp1=read(f,rec=1,(Float32,(720,360))); close(f);
      end
      close(fid)
   end

   return MetaFile
end


text:
header = html("
Open a file.
")
function file_buttons(http)
    cd(PUBLIC * "/texts")
    for dir in readdir()
        button = """$dir"""
        write(http, "   " * button)
    end
    cd("../../")
    ""
end
te = Page([header, fn(file_buttons)], "Toolips Text Editor")

function opener(http)
    args = getargs(http)
    write(http, "
$args
")
    text = read(PUBLIC * "/texts/" * args, String)
    textbox = TextArea("textle", maxlength = 1000, cols = 50, rows = 50,
    text = text)
    btton = Button("fname", label = "Save text", value = args)
    txtboxform = Form(textbox, btton, action = "/text-editor/save")
    txtboxform.f(http)
end

texteditsc = Page([fn(opener)])
function saver(http)
    args = getargs(http)
    bothvals = split(args, "&")
    nameval = split(bothvals[2], "=")
    write(http, "Saving file " * nameval[2])
    fname = nameval[2]
    nameval = split(bothvals[1], "=")
    out = replace(nameval[2], "%" => " ")
    replace(nameval[2], "+" => " ")
    open(PUBLIC * "/texts/" * fname, "w") do o
        write(o, out)
    end
    write(http, "File saved!")
    sleep(2)
    """<meta http-equiv="Refresh" content="0; url='/text-editor/open?$fname'" />"""
end

max_issues_repo_name: emmettgb/ToolipsApp

max_forks_repo_name: emmettgb/ToolipsApp


text:
"""
    GeometricConicForm{T, AT, VT, C} <: MOI.ModelLike

Represents an optimization model of the form:
```
sense ⟨c, x⟩ + c0
s.t.  b_i - A_i x ∈ C_i ∀ i
```
with each `C_i` a cone defined in MOI.
"""
mutable struct GeometricConicForm{T, AT, VB, VC, C} <: MOI.ModelLike
    num_rows::Vector{Int}
    dimension::Dict{Int, Int}
    sense::MOI.OptimizationSense
    objective_constant::T # The objective
    A::Union{Nothing, AT} # The constraints
    b::VB          # `b - Ax in cones`
    c::VC          # `sense c'x + objective_constant`
    cone_types::C
    cone_types_dict::Dict{DataType, Int}

    function GeometricConicForm{T, AT, VB, VC}(cone_types) where {T, AT, VB, VC}
        model = new{T, AT, VB, VC, typeof(cone_types)}()
        model.cone_types = cone_types
        model.cone_types_dict = Dict{DataType, Int}(
            s => i for (i, s) in enumerate(cone_types)
        )
        model.num_rows = zeros(Int, length(cone_types))
        model.dimension = Dict{Int, Int}()
        model.A = nothing
        return model
    end
end

function GeometricConicForm{T, AT, VT}(cone_types) where {T, AT, VT}
    return GeometricConicForm{T, AT, VT, VT}(cone_types)
end

_set_type(::MOI.ConstraintIndex{F,S}) where {F,S} = S

MOI.is_empty(model::GeometricConicForm) = model.A === nothing

function MOI.empty!(model::GeometricConicForm{T}) where {T}
    empty!(model.dimension)
    fill!(model.num_rows, 0)
    model.A = nothing
    model.sense = MOI.FEASIBILITY_SENSE
    model.objective_constant = zero(T)
end

function MOI.supports_constraint(
    model::GeometricConicForm{T},
    ::Type{MOI.VectorAffineFunction{T}},
    ::Type{S}) where {T, S <: MOI.AbstractVectorSet}
    return haskey(model.cone_types_dict, S)
end

function _allocate_variables(model::GeometricConicForm{T, AT, VT}, vis_src, idxmap) where {T, AT, VT}
    model.A = AT(length(vis_src))
    for (i, vi) in enumerate(vis_src)
        idxmap[vi] = MOI.VariableIndex(i)
    end
    return
end

function rows(model::GeometricConicForm{T}, ci::CI{MOI.VectorAffineFunction{T}}) where T
    return ci.value .+ (1:model.dimension[ci.value])
end

function MOI.set(::GeometricConicForm, ::MOI.VariablePrimalStart,
                 ::MOI.VariableIndex, ::Nothing)
end
function MOI.set(model::GeometricConicForm{T}, ::MOI.VariablePrimalStart,
                 vi::MOI.VariableIndex, value::T) where {T}
    model.primal[vi.value] = value
end
function MOI.set(::GeometricConicForm, ::MOI.ConstraintPrimalStart,
                 ::MOI.ConstraintIndex, ::Nothing)
end
function MOI.set(model::GeometricConicForm, ::MOI.ConstraintPrimalStart,
                 ci::MOI.ConstraintIndex, value)
    offset = constroffset(model, ci)
    model.slack[rows(model, ci)] .= value
end
function MOI.set(::GeometricConicForm, ::MOI.ConstraintDualStart,
                  ::MOI.ConstraintIndex, ::Nothing)
end
function MOI.set(model::GeometricConicForm, ::MOI.ConstraintDualStart,
                  ci::MOI.ConstraintIndex, value)
    offset = constroffset(model, ci)
    model.dual[rows(model, ci)] .= value
end
function MOI.set(model::GeometricConicForm, ::MOI.ObjectiveSense, sense::MOI.OptimizationSense)
    model.sense = sense
end
variable_index_value(t::MOI.ScalarAffineTerm) = t.variable.value
variable_index_value(t::MOI.VectorAffineTerm) = variable_index_value(t.scalar_term)
function MOI.set(model::GeometricConicForm{T}, ::MOI.ObjectiveFunction,
                 f::MOI.ScalarAffineFunction{T}) where {T}
    c = Vector(sparsevec(variable_index_value.(f.terms), MOI.coefficient.(f.terms),
                         model.A.n))
    model.objective_constant = f.constant
    model.c = c
    return nothing
end

function _allocate_constraint(model::GeometricConicForm, src, indexmap, cone_id, ci)
    # TODO use `CanonicalConstraintFunction`
    func = MOI.get(src, MOI.ConstraintFunction(), ci)
    func = MOIU.is_canonical(func) ? func : MOI.Utilities.canonical(func)
    allocate_terms(model.A, indexmap, func)
    offset = model.num_rows[cone_id]
    model.num_rows[cone_id] = offset + MOI.output_dimension(func)
    return ci, offset, func
end

function _allocate_constraints(model::GeometricConicForm{T}, src, indexmap, cone_id, ::Type{S}) where {T, S}
    cis = MOI.get(src, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{T}, S}())
    return map(cis) do ci
        _allocate_constraint(model, src, indexmap, cone_id, ci)
    end
end

function _load_variables(model::GeometricConicForm, nvars::Integer)
    m = sum(model.num_rows)
    model.A.m = m
    model.b = zeros(m)
    model.c = zeros(model.A.n)
    allocate_nonzeros(model.A)
end

function _load_constraints(model::GeometricConicForm, src, indexmap, cone_offset, i, cache)
    for (ci_src, offset_in_cone, func) in cache
        offset = cone_offset + offset_in_cone
        set = MOI.get(src, MOI.ConstraintSet(), ci_src)
        load_terms(model.A, indexmap, func, offset)
        copyto!(model.b, offset + 1, func.constants)
        model.dimension[offset] = MOI.output_dimension(func)
        indexmap[ci_src] = typeof(ci_src)(offset)
    end
end

function MOI.copy_to(dest::GeometricConicForm{T}, src::MOI.ModelLike) where T
    MOI.empty!(dest)

    vis_src = MOI.get(src, MOI.ListOfVariableIndices())
    idxmap = MOIU.IndexMap()

    has_constraints = BitSet()
    for (F, S) in MOI.get(src, MOI.ListOfConstraintTypesPresent())
        i = get(dest.cone_types_dict, S, nothing)
        if i === nothing || F != MOI.VectorAffineFunction{T}
            throw(MOI.UnsupportedConstraint{F, S}())
        end
        push!(has_constraints, i)
    end

    _allocate_variables(dest, vis_src, idxmap)

    # Allocate constraints
    caches = map(collect(has_constraints)) do i
        _allocate_constraints(dest, src, idxmap, i, dest.cone_types[i])
    end

    # Load variables
    _load_variables(dest, length(vis_src))

    # Set variable attributes
    MOIU.pass_attributes(dest, src, idxmap, vis_src)

    # Set model attributes
    MOIU.pass_attributes(dest, src, idxmap)

    # Load constraints
    offset = 0
    for (i, cache) in zip(has_constraints, caches)
        _load_constraints(dest, src, idxmap, offset, i, cache)
        offset += dest.num_rows[i]
    end

    final_touch(dest.A)

    return idxmap
end


text:
"""
    `Metropolis(β, J, L)` - Set up the Metropolis sampler

### Fields
* `β` - Inverse temperature
* `J` - Coupling constant>0
* `L` - Lattice size, L×L
"""
struct Metropolis{TF<:AbstractFloat, TI<:Integer} <:AbstractSampler
    β::TF
    J::TF
    L::TI
end
mutable struct MetropolisState{TI<:Integer, Tx<:AbstractArray{TI}} <:AbstractSamplerState
    x::Tx
    x_proposal::Tx
end

function InitState!(x, sampler::Metropolis)

    return MetropolisState(x, deepcopy(x));
end

function InitState(x₀, sampler::Metropolis)

    return MetropolisState(deepcopy(x₀), deepcopy(x₀));
end

function UpdateState!(state::MetropolisState, sampler::Metropolis)

    i = rand(1:sampler.L);
    j = rand(1:sampler.L);
    state.x_proposal[i,j] *= -1;

    e = 0.0;

    k,l = mod1(i+1,sampler.L), j;
    e += - 0.5 * sampler.J * state.x_proposal[i,j] * state.x_proposal[k,l];
    k,l = mod1(i-1,sampler.L), j;
    e += - 0.5 * sampler.J * state.x_proposal[i,j] * state.x_proposal[k,l];
    k,l = i, mod1(j+1,sampler.L);
    e += - 0.5 * sampler.J * state.x_proposal[i,j] * state.x_proposal[k,l];
    k,l = i, mod1(j-1,sampler.L);
    e += - 0.5 * sampler.J * state.x_proposal[i,j] * state.x_proposal[k,l];

    # double once for double counting of bonds and again for the sign change
    ΔE = 2 * 2 * e;

    a = min(1, exp(-sampler.β * ΔE))

    if rand() <a
        state.x[i,j] *= -1;
    else
        state.x_proposal[i,j] *= -1;
    end
    state
end


text:
"""
    Normal(μ,σ)

The *Normal distribution* with mean `μ` and standard deviation `σ≥0` has probability density function

```math
f(x; \\mu, \\sigma) = \\frac{1}{\\sqrt{2 \\pi \\sigma^2}}
\\exp \\left( - \\frac{(x - \\mu)^2}{2 \\sigma^2} \\right)
```

Note that if `σ == 0`, then the distribution is a point mass concentrated at `μ`.
Though not technically a continuous distribution, it is allowed so as to account for cases where `σ` may have underflowed,
and the functions are defined by taking the pointwise limit as ``σ → 0``.

```julia
Normal()          # standard Normal distribution with zero mean and unit variance
Normal(mu)        # Normal distribution with mean mu and unit variance
Normal(mu, sig)   # Normal distribution with mean mu and variance sig^2

params(d)         # Get the parameters, i.e. (mu, sig)
mean(d)           # Get the mean, i.e. mu
std(d)            # Get the standard deviation, i.e. sig
```

External links

* [Normal distribution on Wikipedia](http://en.wikipedia.org/wiki/Normal_distribution)

"""
struct Normal{T<:Real} <: ContinuousUnivariateDistribution
    μ::T
    σ::T
    Normal{T}(µ::T, σ::T) where {T<:Real} = new{T}(µ, σ)
end

function Normal(μ::T, σ::T; check_args=true) where {T <: Real}
    check_args && @check_args(Normal, σ >= zero(σ))
    return Normal{T}(μ, σ)
end

#### Outer constructors
Normal(μ::Real, σ::Real) = Normal(promote(μ, σ)...)
Normal(μ::Integer, σ::Integer) = Normal(float(μ), float(σ))
Normal(μ::T) where {T <: Real} = Normal(μ, one(T))
Normal() = Normal(0.0, 1.0, check_args=false)

const Gaussian = Normal

# #### Conversions
convert(::Type{Normal{T}}, μ::S, σ::S) where {T <: Real, S <: Real} = Normal(T(μ), T(σ))
convert(::Type{Normal{T}}, d::Normal{S}) where {T <: Real, S <: Real} = Normal(T(d.μ), T(d.σ), check_args=false)

@distr_support Normal -Inf Inf

#### Parameters

params(d::Normal) = (d.μ, d.σ)
@inline partype(d::Normal{T}) where {T<:Real} = T

location(d::Normal) = d.μ
scale(d::Normal) = d.σ

eltype(::Normal{T}) where {T} = T

#### Statistics

mean(d::Normal) = d.μ
median(d::Normal) = d.μ
mode(d::Normal) = d.μ

var(d::Normal) = abs2(d.σ)
std(d::Normal) = d.σ
skewness(d::Normal{T}) where {T<:Real} = zero(T)
kurtosis(d::Normal{T}) where {T<:Real} = zero(T)

entropy(d::Normal) = (log2π + 1)/2 + log(d.σ)

#### Evaluation

# Helpers
"""
    xval(d::Normal, z::Real)

Computes the x-value based on a Normal distribution and a z-value.
"""
function xval(d::Normal, z::Real)
    if isinf(z) && iszero(d.σ)
        d.μ + one(d.σ) * z
    else
        d.μ + d.σ * z
    end
end
"""
    zval(d::Normal, x::Real)

Computes the z-value based on a Normal distribution and a x-value.
"""
zval(d::Normal, x::Real) = (x - d.μ) / d.σ

gradlogpdf(d::Normal, x::Real) = -zval(d, x) / d.σ
# logpdf
function logpdf(d::Normal, x::Real)
    if iszero(d.σ)
        d.μ == x ? Inf : -Inf
    else
        -(zval(d, x)^2 + log2π) / 2 - log(d.σ)
    end
end
# pdf
function pdf(d::Normal, x::Real)
    if iszero(d.σ)
        d.μ == x ? Inf : 0.0
    else
        exp(-zval(d, x)^2 / 2) * invsqrt2π / d.σ
    end
end
# logcdf
function logcdf(d::Normal, x::Real)
    if iszero(d.σ)
        d.μ ≤ x ? 0.0 : -Inf
    else
        z = zval(d, x)
        if z < -1.0
            log(erfcx(-z * invsqrt2)/2) - abs2(z)/2
        else
            log1p(-erfc(z * invsqrt2)/2)
        end
    end
end
# logccdf
function logccdf(d::Normal, x::Real)
    if iszero(d.σ) && x == d.μ
        z = zval(Normal(zero(d.μ), d.σ), one(x))
    else
        z = zval(d, x)
    end
    z > 1.0 ?
    log(erfcx(z * invsqrt2) / 2) - z^2 / 2 :
    log1p(-erfc(-z * invsqrt2) / 2)
end
# cdf
function cdf(d::Normal{T}, x::S) where {T, S <: Real}
    if iszero(d.σ)
        convert(promote_type(T, S), d.μ ≤ x)
    else
        erfc(-zval(d, x) * invsqrt2) / 2
    end
end
# ccdf
function ccdf(d::Normal, x::Real)
    z = iszero(d.σ) && x == d.μ ?
        zval(Normal(zero(d.μ), d.σ), one(x)) :
        zval(d, x)
    erfc(z * invsqrt2) / 2
end
# invlogcdf
"""
    norminvlogcdf(lp::Real)

Helper function that calls `_norminvlogcdf_impl` used for `invlogccdf` with the Normal distributions.
"""
norminvlogcdf(lp::Real) = _norminvlogcdf_impl(convert(Float64, lp))
norminvlogcdf(lp::Union{Float16,Float32}) = convert(typeof(lp), _norminvlogcdf_impl(convert(Float64, lp)))
invlogcdf(d::Normal, lp::Real) = xval(d, norminvlogcdf(lp))
# invlogccdf
invlogccdf(d::Normal, lp::Real) = xval(d, -norminvlogcdf(lp))
# quantile
function quantile(d::Normal, p::Real)
    if iszero(d.σ)
        if iszero(p)
            -Inf
        elseif isone(p)
            Inf
        else
            0.5
        end
    end
    xval(d, -erfcinv(2p) * sqrt2)
end
# cquantile
function cquantile(d::Normal, q::Real)
    if iszero(d.σ)
        if iszero(q)
            Inf
        elseif isone(q)
            -Inf
        else
            0.5
        end
    end
    xval(d, erfcinv(2q) * sqrt2)
end

# norminvcdf & norminvlogcdf implementation
"""
    _norminvlogcdf_impl(lp::Float64)

Uses `_qnorm_ker1` and `_qnorm_ker2` to obtain the ational approximations for the inverse cdf and its logarithm, from:
Wichura, M.J. (1988) Algorithm AS 241: The Percentage Points of the Normal Distribution
  Journal of the Royal Statistical Society. Series C (Applied Statistics), Vol. 37, No. 3, pp. 477-484
"""
function _norminvlogcdf_impl(lp::Float64)
    if isfinite(lp) && lp < 0.0
        q = exp(lp) - 0.5
        # qnorm_kernel(lp, q, true)
        if abs(q) ≤ 0.425
            _qnorm_ker1(q)
        else
            r = sqrt(q < 0 ? -lp : -log1mexp(lp))
            return copysign(_qnorm_ker2(r), q)
        end
    elseif lp ≥ 0.0
        iszero(lp) ? Inf : NaN
    else # lp is -Inf or NaN
        lp
    end
end
"""
    _qnorm_ker1(r::Float64)

Along with `_qnorm_ker2` these helpers enable the rational approximations for the inverse cdf and its logarithm
"""
function _qnorm_ker1(q::Float64)
    # pre-condition: abs(q) ≤ 0.425
    r = 0.180625 - q^2
    return q * @horner(r,
                       3.38713_28727_96366_6080e0,
                       1.33141_66789_17843_7745e2,
                       1.97159_09503_06551_4427e3,
                       1.37316_93765_50946_1125e4,
                       4.59219_53931_54987_1457e4,
                       6.72657_70927_00870_0853e4,
                       3.34305_75583_58812_8105e4,
                       2.50908_09287_30122_6727e3) /
    @horner(r,
            1.0,
            4.23133_30701_60091_1252e1,
            6.87187_00749_20579_0830e2,
            5.39419_60214_24751_1077e3,
            2.12137_94301_58659_5867e4,
            3.93078_95800_09271_0610e4,
            2.87290_85735_72194_2674e4,
            5.22649_52788_52854_5610e3)
end
"""
    _qnorm_ker2(r::Float64)

Along with `_qnorm_ker1` these helpers enable the rational approximations for the inverse cdf and its logarithm
"""
function _qnorm_ker2(r::Float64)
    if r < 5.0
        r -= 1.6
        @horner(r,
                1.42343_71107_49683_57734e0,
                4.63033_78461_56545_29590e0,
                5.76949_72214_60691_40550e0,
                3.64784_83247_63204_60504e0,
                1.27045_82524_52368_38258e0,
                2.41780_72517_74506_11770e-1,
                2.27238_44989_26918_45833e-2,
                7.74545_01427_83414_07640e-4) /
        @horner(r,
                1.0,
                2.05319_16266_37758_82187e0,
                1.67638_48301_83803_84940e0,
                6.89767_33498_51000_04550e-1,
                1.48103_97642_74800_74590e-1,
                1.51986_66563_61645_71966e-2,
                5.47593_80849_95344_94600e-4,
                1.05075_00716_44416_84324e-9)
    else
        r -= 5.0
        @horner(r,
                6.65790_46435_01103_77720e0,
                5.46378_49111_64114_36990e0,
                1.78482_65399_17291_33580e0,
                2.96560_57182_85048_91230e-1,
                2.65321_89526_57612_30930e-2,
                1.24266_09473_88078_43860e-3,
                2.71155_55687_43487_57815e-5,
                2.01033_43992_92288_13265e-7) /
        @horner(r,
                1.0,
                5.99832_20655_58879_37690e-1,
                1.36929_88092_27358_05310e-1,
                1.48753_61290_85061_48525e-2,
                7.86869_13114_56132_59100e-4,
                1.84631_83175_10054_68180e-5,
                1.42151_17583_16445_88870e-7,
                2.04426_31033_89939_78564e-15)
    end
end

mgf(d::Normal, t::Real) = exp(t * d.μ + d.σ^2 / 2 * t^2)
cf(d::Normal, t::Real) = exp(im * t * d.μ - d.σ^2 / 2 * t^2)

#### Sampling

rand(rng::AbstractRNG, d::Normal{T}) where {T} = d.μ + d.σ * randn(rng, T)

#### Fitting

struct NormalStats <: SufficientStats
    s::Float64    # (weighted) sum of x
    m::Float64    # (weighted) mean of x
    s2::Float64   # (weighted) sum of (x - μ)^2
    tw::Float64    # total sample weight
end

function suffstats(::Type{<:Normal}, x::AbstractArray{T}) where T<:Real
    n = length(x)

    # compute s
    s = x[1]
    for i = 2:n
        @inbounds s += x[i]
    end
    m = s / n

    # compute s2
    s2 = abs2(x[1] - m)
    for i = 2:n
        @inbounds s2 += abs2(x[i] - m)
    end

    NormalStats(s, m, s2, n)
end

function suffstats(::Type{<:Normal}, x::AbstractArray{T}, w::AbstractArray{Float64}) where T<:Real
    n = length(x)

    # compute s
    tw = w[1]
    s = w[1] * x[1]
    for i = 2:n
        @inbounds wi = w[i]
        @inbounds s += wi * x[i]
        tw += wi
    end
    m = s / tw

    # compute s2
    s2 = w[1] * abs2(x[1] - m)
    for i = 2:n
        @inbounds s2 += w[i] * abs2(x[i] - m)
    end

    NormalStats(s, m, s2, tw)
end

# Cases where μ or σ is known

struct NormalKnownMu <: IncompleteDistribution
    μ::Float64
end

struct NormalKnownMuStats <: SufficientStats
    μ::Float64      # known mean
    s2::Float64     # (weighted) sum of (x - μ)^2
    tw::Float64     # total sample weight
end

function suffstats(g::NormalKnownMu, x::AbstractArray{T}) where T<:Real
    μ = g.μ
    s2 = abs2(x[1] - μ)
    for i = 2:length(x)
        @inbounds s2 += abs2(x[i] - μ)
    end
    NormalKnownMuStats(g.μ, s2, length(x))
end

function suffstats(g::NormalKnownMu, x::AbstractArray{T}, w::AbstractArray{Float64}) where T<:Real
    μ = g.μ
    s2 = abs2(x[1] - μ) * w[1]
    tw = w[1]
    for i = 2:length(x)
        @inbounds wi = w[i]
        @inbounds s2 += abs2(x[i] - μ) * wi
        tw += wi
    end
    NormalKnownMuStats(g.μ, s2, tw)
end

struct NormalKnownSigma <: IncompleteDistribution
    σ::Float64

    function NormalKnownSigma(σ::Float64)
        σ > 0 || throw(ArgumentError("σ must be a positive value."))
        new(σ)
    end
end

struct NormalKnownSigmaStats <: SufficientStats
    σ::Float64      # known std.dev
    sx::Float64      # (weighted) sum of x
    tw::Float64     # total sample weight
end

function suffstats(g::NormalKnownSigma, x::AbstractArray{T}) where T<:Real
    NormalKnownSigmaStats(g.σ, sum(x), Float64(length(x)))
end

function suffstats(g::NormalKnownSigma, x::AbstractArray{T}, w::AbstractArray{T}) where T<:Real
    NormalKnownSigmaStats(g.σ, dot(x, w), sum(w))
end

# fit_mle based on sufficient statistics

fit_mle(::Type{<:Normal}, ss::NormalStats) = Normal(ss.m, sqrt(ss.s2 / ss.tw))
fit_mle(g::NormalKnownMu, ss::NormalKnownMuStats) = Normal(g.μ, sqrt(ss.s2 / ss.tw))
fit_mle(g::NormalKnownSigma, ss::NormalKnownSigmaStats) = Normal(ss.sx / ss.tw, g.σ)

# generic fit_mle methods

function fit_mle(::Type{<:Normal}, x::AbstractArray{T}; mu::Float64=NaN, sigma::Float64=NaN) where T<:Real
    if isnan(mu)
        if isnan(sigma)
            fit_mle(Normal, suffstats(Normal, x))
        else
            g = NormalKnownSigma(sigma)
            fit_mle(g, suffstats(g, x))
        end
    else
        if isnan(sigma)
            g = NormalKnownMu(mu)
            fit_mle(g, suffstats(g, x))
        else
            Normal(mu, sigma)
        end
    end
end

function fit_mle(::Type{<:Normal}, x::AbstractArray{T}, w::AbstractArray{Float64}; mu::Float64=NaN, sigma::Float64=NaN) where T<:Real
    if isnan(mu)
        if isnan(sigma)
            fit_mle(Normal, suffstats(Normal, x, w))
        else
            g = NormalKnownSigma(sigma)
            fit_mle(g, suffstats(g, x, w))
        end
    else
        if isnan(sigma)
            g = NormalKnownMu(mu)
            fit_mle(g, suffstats(g, x, w))
        else
            Normal(mu, sigma)
        end
    end
end