text:
import ForwardDiff
import MacroTools
import LinearAlgebra
import Parameters: @with_kw

#######################
# trace transform DSL #
#######################

# See this math writeup for an understanding of how this code works:
# docs/tex/mcmc.pdf

"""
    TraceTransformDSLProgram

A program compiled from the [Trace Transform DSL](@ref).
"""
mutable struct TraceTransformDSLProgram
    fn!::Function
    inverse::Union{Nothing,TraceTransformDSLProgram}
end

"""
    pair_bijections!(f1::TraceTransformDSLProgram, f2::TraceTransformDSLProgram)

Assert that a pair of bijections contsructed using the [Trace Transform DSL](@ref) are inverses of one another.
"""
function pair_bijections!(f1::TraceTransformDSLProgram, f2::TraceTransformDSLProgram)
    f1.inverse = f2
    f2.inverse = f1
    return nothing
end

"""
    is_involution!(f::TraceTransformDSLProgram)

Assert that a bijection constructed with the [Trace Transform DSL](@ref) is its own inverse.
"""
function is_involution!(f::TraceTransformDSLProgram)
    f.inverse = f
    return nothing
end

"""
    b::TraceTransformDSLProgram = inverse(a::TraceTransformDSLProgram)

Obtain the inverse of a bijection that was constructed with the [Trace Transform DSL](@ref).

The inverse must have been associated with the bijection either via [`pair_bijections!`](@ref) or [`is_involution!`])(@ref).
"""
function inverse(bijection::TraceTransformDSLProgram)
    if isnothing(bijection.inverse)
        error("inverse bijection was not defined")
    end
    return bijection.inverse
end

struct ModelInputTraceToken{T}
    args::T
end

struct AuxInputTraceToken{T}
    args::T
end

struct ModelInputTraceRetValToken
end

struct AuxInputTraceRetValToken
end

struct ModelOutputTraceToken
end

struct AuxOutputTraceToken
end

struct ModelInputAddress{T}
    addr::T
end

struct AuxInputAddress{T}
    addr::T
end

struct ModelOutputAddress{T}
    addr::T
end

struct AuxOutputAddress{T}
    addr::T
end

Base.getindex(::ModelInputTraceToken, addr) = ModelInputAddress(addr) # model_in[addr]
Base.getindex(::ModelOutputTraceToken, addr) = ModelOutputAddress(addr) # model_out[addr]
Base.getindex(::AuxInputTraceToken, addr) = AuxInputAddress(addr) # aux_in[addr]
Base.getindex(::AuxOutputTraceToken, addr) = AuxOutputAddress(addr) # aux_out[addr]
Base.getindex(::ModelInputTraceToken) = ModelInputTraceRetvalToken() # model_in[]
Base.getindex(::AuxInputTraceToken) = AuxInputTraceRetValToken() # aux_in[]
get_args(token::ModelInputTraceToken) = token.args # get_args(model_in)
get_args(token::AuxInputTraceToken) = token.args # get_args(aux_in)

const bij_state = gensym("bij_state")

"""
    @transform f[(params...)] (in1 [,in2]) to (out1 [,out2])
        ..
    end

Write a program in the [Trace Transform DSL](@ref).
"""
macro transform(f_expr, src_expr, to_symbol::Symbol, dest_expr, body)
    syntax_err = "valid syntactic forms:\n@transform f (..) to (..) begin .. end\n@transform f(..) (..) to (..) begin .. end"
    err = false
    if MacroTools.@capture(f_expr, f_(args__))
    elseif MacroTools.@capture(f_expr, f_)
        args = []
    else
        err = true
    end
    err = err || (to_symbol != :to)
    if MacroTools.@capture(src_expr, (model_in_, aux_in_))
    elseif MacroTools.@capture(src_expr, (model_in_))
        aux_in = gensym("dummy_aux")
    else
        err = true
    end
    if MacroTools.@capture(dest_expr, (model_out_, aux_out_))
    elseif MacroTools.@capture(src_expr, (model_out_))
        aux_out = gensym("dummy_aux")
    else
        err = true
    end

    fn! = gensym("$(esc(f))_fn!")

    return quote

        # mutates the state
        function $fn!(
                $(esc(bij_state))::Union{FirstPassState,JacobianPassState},
                $(map(esc, args)...))
            model_args = get_args($(esc(bij_state)).model_trace)
            aux_args = get_args($(esc(bij_state)).aux_trace)
            $(esc(model_in)) = ModelInputTraceToken(model_args)
            $(esc(model_out)) = ModelOutputTraceToken()
            $(esc(aux_in)) = AuxInputTraceToken(aux_args)
            $(esc(aux_out)) = AuxOutputTraceToken()
            $(esc(body))
            return nothing
        end

        Core.@__doc__ $(esc(f)) = TraceTransformDSLProgram($fn!, nothing)

    end
end

macro tcall(ex)
    MacroTools.@capture(ex, f_(args__)) || error("expected syntax: f(..)")
    return quote $(esc(f)).fn!($(esc(bij_state)), $(map(esc, args)...)) end
end

# handlers

struct DiscreteAnn end
struct ContinuousAnn end

const DISCRETE = [:discrete, :disc]
const CONTINUOUS = [:continuous, :cont]

function typed(annotation::Symbol)
    if annotation in DISCRETE
        return DiscreteAnn()
    elseif annotation in CONTINUOUS
        return ContinuousAnn()
    else
        error("error")
    end
end

"""
    @read(<source>, <annotation>)

Macro for reading the value of a random choice from an input trace in the [Trace Transform DSL](@ref).

<source> is of the form <trace>[<addr>] where <trace> is an input trace, and <annotation> is either :discrete or :continuous.
"""
macro read(src, ann::QuoteNode)
    return quote read($(esc(bij_state)), $(esc(src)), $(esc(typed(ann.value)))) end
end

"""
    @write(<destination>, <value>, <annotation>)

Macro for writing the value of a random choice to an output trace in the [Trace Transform DSL](@ref).

<destination> is of the form <trace>[<addr>] where <trace> is an input trace, and <annotation> is either :discrete or :continuous.
"""
macro write(dest, val, ann::QuoteNode)
    return quote write($(esc(bij_state)), $(esc(dest)), $(esc(val)), $(esc(typed(ann.value)))) end
end

"""
    @copy(<source>, <destination>)

Macro for copying the value of a random choice (or a whole namespace of random choices) from an input trace to an output trace in the [Trace Transform DSL](@ref).

<destination> is of the form <trace>[<addr>] where <trace> is an input trace, and <annotation> is either :discrete or :continuous.
"""
macro copy(src, dest)
    return quote copy($(esc(bij_state)), $(esc(src)), $(esc(dest))) end
end

# TODO make more consistent by allowing us to read any hierarchical address,
# including return values of intermediate calls, not just the top-level call.

# TODO add haskey(model_in, addr), and haskey(aux_in, addr)

################################
# first pass through transform #
################################

struct FirstPassResults

    "subset of the output model choice map ``t'``"
    constraints::ChoiceMap

    "output proposal choice map ``u'``"
    u_back::ChoiceMap

    t_cont_reads::Dict
    u_cont_reads::Dict
    t_cont_writes::Dict
    u_cont_writes::Dict
    t_copy_reads::DynamicSelection
    u_copy_reads::DynamicSelection
end

function FirstPassResults()
    return FirstPassResults(
        choicemap(), choicemap(),
        Dict(), Dict(), Dict(), Dict(),
        DynamicSelection(), DynamicSelection())
end

struct FirstPassState

    "trace containing the input model choice map ``t``"
    model_trace::Trace
    model_choices::ChoiceMap

    "the input proposal choice map ``u``"
    aux_trace::Trace
    aux_choices::ChoiceMap

    results::FirstPassResults
end

function FirstPassState(model_trace, aux_trace)
    return FirstPassState(
        model_trace, get_choices(model_trace),
        aux_trace, get_choices(aux_trace), FirstPassResults())
end

function run_first_pass(transform::TraceTransformDSLProgram, model_trace, aux_trace)
    state = FirstPassState(model_trace, aux_trace)
    transform.fn!(state) # TODO allow for other args to top-level transform function
    return state.results
end

function read(state::FirstPassState, src::ModelInputTraceRetValToken, ::DiscreteAnn)
    return get_retval(state.model_trace)
end

function read(state::FirstPassState, src::AuxInputTraceRetValToken, ::DiscreteAnn)
    return get_retval(state.aux_trace)
end

function read(state::FirstPassState, src::ModelInputAddress, ::DiscreteAnn)
    addr = src.addr
    return state.model_trace[addr]
end

function read(state::FirstPassState, src::ModelInputAddress, ::ContinuousAnn)
    addr = src.addr
    state.results.t_cont_reads[addr] = state.model_choices[addr]
    return state.model_trace[addr]
end

function read(state::FirstPassState, src::AuxInputAddress, ::DiscreteAnn)
    addr = src.addr
    return state.aux_trace[addr]
end

function read(state::FirstPassState, src::AuxInputAddress, ::ContinuousAnn)
    addr = src.addr
    state.results.u_cont_reads[addr] = state.aux_choices[addr]
    return state.aux_trace[addr]
end

function write(state::FirstPassState, dest::ModelOutputAddress, value, ::DiscreteAnn)
    addr = dest.addr
    state.results.constraints[addr] = value
    return value
end

function write(state::FirstPassState, dest::ModelOutputAddress, value, ::ContinuousAnn)
    addr = dest.addr
    has_value(state.results.constraints, addr) && error("Model address $addr already written to")
    state.results.constraints[addr] = value
    state.results.t_cont_writes[addr] = value
    return value
end

function write(state::FirstPassState, dest::AuxOutputAddress, value, ::DiscreteAnn)
    addr = dest.addr
    state.results.u_back[addr] = value
    return value
end

function write(state::FirstPassState, dest::AuxOutputAddress, value, ::ContinuousAnn)
    addr = dest.addr
    has_value(state.results.u_back, addr) && error("Proposal address $addr already written to")
    state.results.u_back[addr] = value
    state.results.u_cont_writes[addr] = value
    return value
end

function copy(state::FirstPassState, src::ModelInputAddress, dest::ModelOutputAddress)
    from_addr, to_addr = src.addr, dest.addr
    model_choices = get_choices(state.model_trace)
    push!(state.results.t_copy_reads, from_addr)
    if has_value(model_choices, from_addr)
        state.results.constraints[to_addr] = model_choices[from_addr]
    else
        set_submap!(state.results.constraints, to_addr, get_submap(model_choices, from_addr))
    end
    return nothing
end

function copy(state::FirstPassState, src::ModelInputAddress, dest::AuxOutputAddress)
    from_addr, to_addr = src.addr, dest.addr
    model_choices = get_choices(state.model_trace)
    push!(state.results.t_copy_reads, from_addr)
    if has_value(model_choices, from_addr)
        state.results.u_back[to_addr] = model_choices[from_addr]
    else
        set_submap!(state.results.u_back, to_addr, get_submap(model_choices, from_addr))
    end
    return nothing
end

function copy(state::FirstPassState, src::AuxInputAddress, dest::AuxOutputAddress)
    from_addr, to_addr = src.addr, dest.addr
    push!(state.results.u_copy_reads, from_addr)
    aux_choices = get_choices(state.aux_trace)
    if has_value(aux_choices, from_addr)
        state.results.u_back[to_addr] = aux_choices[from_addr]
    else
        set_submap!(state.results.u_back, to_addr, get_submap(aux_choices, from_addr))
    end
    return nothing
end

function copy(state::FirstPassState, src::AuxInputAddress, dest::ModelOutputAddress)
    from_addr, to_addr = src.addr, dest.addr
    push!(state.results.u_copy_reads, from_addr)
    aux_choices = get_choices(state.aux_trace)
    if has_value(aux_choices, from_addr)
        state.results.constraints[to_addr] = aux_choices[from_addr]
    else
        set_submap!(state.results.constraints, to_addr, get_submap(aux_choices, from_addr))
    end
    return nothing
end

#####################################################################
# second pass through transform (gets automatically differentiated) #
#####################################################################

struct JacobianPassState{T<:Real}
    model_trace
    aux_trace
    input_arr::AbstractArray{T}
    output_arr::Array{T}
    t_key_to_index::Dict
    u_key_to_index::Dict
    cont_constraints_key_to_index::Dict
    cont_u_back_key_to_index::Dict
end

function read(state::JacobianPassState, src::ModelInputTraceRetValToken, ::DiscreteAnn)
    return get_retval(state.model_trace)
end

function read(state::JacobianPassState, src::AuxInputTraceRetValToken, ::DiscreteAnn)
    return get_retval(state.aux_trace)
end

function read(state::JacobianPassState, src::ModelInputAddress, ::DiscreteAnn)
    addr = src.addr
    return state.model_trace[addr]
end

function read(state::JacobianPassState, src::AuxInputAddress, ::DiscreteAnn)
    addr = src.addr
    return state.aux_trace[addr]
end

function _read_continuous(input_arr, addr_info::Int)
    return input_arr[addr_info]
end

function _read_continuous(input_arr, addr_info::Tuple{Int,Int})
    # TODO to handle things other than vectors, store shape in addr info and reshape?
    (start_idx, len) = addr_info
    return input_arr[start_idx:start_idx+len-1]
 end

function read(state::JacobianPassState, src::ModelInputAddress, ::ContinuousAnn)
    addr = src.addr
    if haskey(state.t_key_to_index, addr)
        return _read_continuous(state.input_arr, state.t_key_to_index[addr])
    else
        return state.model_trace[addr]
    end
end

function read(state::JacobianPassState, src::AuxInputAddress, ::ContinuousAnn)
    addr = src.addr
    if haskey(state.u_key_to_index, addr)
        return _read_continuous(state.input_arr, state.u_key_to_index[addr])
    else
        return state.aux_trace[addr]
    end
end

function write(state::JacobianPassState, dest::ModelOutputAddress, value, ::DiscreteAnn)
    return value
end

function write(state::JacobianPassState, dest::AuxOutputAddress, value, ::DiscreteAnn)
    return value
end

function _write_continuous(output_arr, addr_info::Int, value)
    return output_arr[addr_info] = value
end

function _write_continuous(output_arr, addr_info::Tuple{Int,Int}, value)
    (start_idx, len) = addr_info
    return output_arr[start_idx:start_idx+len-1] = value
end

function write(state::JacobianPassState, dest::AuxOutputAddress, value, ::ContinuousAnn)
    addr = dest.addr
    return _write_continuous(state.output_arr, state.cont_u_back_key_to_index[addr], value)
end

function write(state::JacobianPassState, dest::ModelOutputAddress, value, ::ContinuousAnn)
    addr = dest.addr
    return _write_continuous(state.output_arr, state.cont_constraints_key_to_index[addr], value)
end

function copy(state::JacobianPassState, src, dest)
    return nothing
end

#################################
# computing jacobian correction #
#################################

discard_skip_read_addr(addr, discard::ChoiceMap) = !has_value(discard, addr)
discard_skip_read_addr(addr, discard::Nothing) = false

function store_addr_info!(dict::Dict, addr, value::Real, next_index::Int)
    dict[addr] = next_index
    return 1 # number of elements of array
end

function store_addr_info!(dict::Dict, addr, value::AbstractArray{<:Real}, next_index::Int)
    len = length(value)
    dict[addr] = (next_index, len)
    return len # number of elements of array
end

function assemble_input_array_and_maps(
        t_cont_reads, t_copy_reads, u_cont_reads, u_copy_reads, discard::Union{ChoiceMap,Nothing})

    input_arr = Vector{Float64}()
    next_input_index = 1

    t_key_to_index = Dict()
    for (addr, v) in t_cont_reads
        if addr in t_copy_reads
            continue
        end
        if discard_skip_read_addr(addr, discard)
            # note: only happens when the model is unchanged
            continue
        end
        next_input_index += store_addr_info!(t_key_to_index, addr, v, next_input_index)
        append!(input_arr, v)
    end

    u_key_to_index = Dict()
    for (addr, v) in u_cont_reads
        if addr in u_copy_reads
            continue
        end
        next_input_index += store_addr_info!(u_key_to_index, addr, v, next_input_index)
        append!(input_arr, v)
    end

    return (t_key_to_index, u_key_to_index, input_arr)
end

function assemble_output_maps(t_cont_writes, u_cont_writes)
    next_output_index = 1

    cont_constraints_key_to_index = Dict()
    for (addr, v) in t_cont_writes
        next_output_index += store_addr_info!(cont_constraints_key_to_index, addr, v, next_output_index)
    end

    cont_u_back_key_to_index = Dict()
    for (addr, v) in u_cont_writes
        next_output_index += store_addr_info!(cont_u_back_key_to_index, addr, v, next_output_index)
    end

    return (cont_constraints_key_to_index, cont_u_back_key_to_index, next_output_index-1)
end

function jacobian_correction(transform::TraceTransformDSLProgram, prev_model_trace, proposal_trace, first_pass_results, discard)

    # create input array and mappings input addresses that are needed for Jacobian
    # exclude addresses that were copied explicitly to another address
    (t_key_to_index, u_key_to_index, input_arr) = assemble_input_array_and_maps(
        first_pass_results.t_cont_reads,
        first_pass_results.t_copy_reads,
        first_pass_results.u_cont_reads,
        first_pass_results.u_copy_reads, discard)

    # create mappings for output addresses that are needed for Jacobian
    (cont_constraints_key_to_index, cont_u_back_key_to_index, n_output) = assemble_output_maps(
        first_pass_results.t_cont_writes,
        first_pass_results.u_cont_writes)

    # this function is the partial application of the continuous part of the
    # transform, with inputs corresponding to a particular superset of the
    # columns of the reduced Jacobian matrix
    function f_array(input_arr::AbstractArray{T}) where {T <: Real}

        # closing over:
        # - trace, u
        # - u_key_to_index, t_key_to_index, cont_constraints_key_to_index, cont_u_back_key_to_index
        # - proposal_args, proposal_retval

        output_arr = Vector{T}(undef, n_output)

        jacobian_pass_state = JacobianPassState(
            prev_model_trace, proposal_trace, input_arr, output_arr,
            t_key_to_index, u_key_to_index,
            cont_constraints_key_to_index,
            cont_u_back_key_to_index)

        # mutates the state
        transform.fn!(jacobian_pass_state)

        # return the output array
        output_arr
    end

    # compute Jacobian matrix of f_array, where columns are inputs, rows are outputs
    J = ForwardDiff.jacobian(f_array, input_arr)
    @assert size(J)[2] == length(input_arr)
    num_outputs = size(J)[1]
    if size(J) != (num_outputs, num_outputs)
        error("Jacobian was not square (size was $(size(J)); the function may not be an bijection")
    end

    # log absolute value of Jacobian determinant
    correction = LinearAlgebra.logabsdet(J)[1]
    if isinf(correction)
        @error "Weight correction is infinite; the function may not be an bijection"
    end

    return correction
end

function check_round_trip(trace, trace_rt)
    choices = get_choices(trace)
    choices_rt = get_choices(trace_rt)
    if !isapprox(choices, choices_rt)
        @error("choices: $(sprint(show, "text/plain", choices))")
        @error("choices after round trip: $(sprint(show, "text/plain", choices_rt))")
        error("transform round trip check failed")
    end
    return nothing
end

function check_round_trip(
            prev_model_trace, prev_model_trace_rt,
            forward_proposal_trace, forward_proposal_trace_rt)
    check_round_trip(prev_model_trace, prev_model_trace_rt)
    check_round_trip(forward_proposal_trace, forward_proposal_trace_rt)
    return nothing
end

################################
# DeterministicTraceTranslator #
################################

"""
    translator = DeterministicTraceTranslator(;
        p_new::GenerativeFunction, p_args::Tuple=();
        new_observations::ChoiceMap=EmptyChoiceMap()
        f::TraceTransformDSLProgram)

Constructor for a deterministic trace translator.

Run the translator with:

    (output_trace, log_weight) = translator(input_trace)
"""
@with_kw struct DeterministicTraceTranslator
    p_new::GenerativeFunction
    p_args::Tuple = ()
    new_observations::ChoiceMap = EmptyChoiceMap()
    f::TraceTransformDSLProgram # a bijection
end

function deterministic_trace_translator_run_transform(
        f::TraceTransformDSLProgram, new_observations::ChoiceMap,
        prev_model_trace::Trace, p_new::GenerativeFunction, p_new_args::Tuple)
    first_pass_results = run_first_pass(f, prev_model_trace, nothing)
    log_abs_determinant = jacobian_correction(
        f, prev_model_trace, nothing, first_pass_results, nothing)
    constraints = merge(first_pass_results.constraints, new_observations)
    (new_model_trace, _) = generate(p_new, p_new_args, constraints)
    return (new_model_trace, log_abs_determinant)
end

"""
    (new_trace, log_weight) = (translator::DeterministicTraceTranslator)(trace)

Apply a trace translator.
"""
function (translator::DeterministicTraceTranslator)(
        prev_model_trace::Trace; check=false, prev_observations=EmptyChoiceMap())

    # apply trace transform
    (new_model_trace, log_abs_determinant) = deterministic_trace_translator_run_transform(
        translator.f, translator.new_observations, prev_model_trace, translator.p_new, translator.p_new_args)

    # compute log weight
    prev_model_score = get_score(prev_model_trace)
    new_model_score = get_score(new_model_trace)
    log_weight = new_model_score - prev_model_score + log_abs_determinant

    if check
        check_observations(get_choices(new_model_trace), observations)
        (prev_model_trace_rt, _) = deterministic_trace_translator_run_transform(
            inverse(translator.f), prev_observations, new_model_trace,
            get_gen_fn(prev_model_trace), get_args(prev_model_trace))
        check_round_trip(prev_model_trace, prev_model_trace_rt)
    end

    return (new_model_trace, log_weight)
end

##########################
# GeneralTraceTranslator #
##########################

"""
    translator = GeneralTraceTranslator(;
        p_new::GenerativeFunction,
        p_new_args::Tuple = (),
        new_observations::ChoiceMap = EmptyChoiceMap(),
        q_forward::GenerativeFunction,
        q_forward_args::Tuple  = (),
        q_backward::GenerativeFunction,
        q_backward_args::Tuple  = (),
        f::TraceTransformDSLProgram)

Constructor for a general trace translator.

Run the translator with:

    (output_trace, log_weight) = translator(input_trace; check=false, prev_observations=EmptyChoiceMap())

Use `check` to enable a bijection check (this requires that the transform `f` has been paired with its inverse using [`pair_bijections!](@ref) or [`is_involution`](@ref)).

If `check` is enabled, then `prev_observations` is a choice map containing the observed random choices in the previous trace.
"""
@with_kw struct GeneralTraceTranslator
    p_new::GenerativeFunction
    p_new_args::Tuple = ()
    new_observations::ChoiceMap = EmptyChoiceMap()
    q_forward::GenerativeFunction
    q_forward_args::Tuple  = ()
    q_backward::GenerativeFunction
    q_backward_args::Tuple  = ()
    f::TraceTransformDSLProgram # a bijection
end

function general_trace_translator_run_transform(
        f::TraceTransformDSLProgram, new_observations::ChoiceMap,
        prev_model_trace::Trace, forward_proposal_trace::Trace,
        p_new::GenerativeFunction, p_new_args::Tuple,
        q_backard::GenerativeFunction, q_backward_args::Tuple)
    first_pass_results = run_first_pass(f, prev_model_trace, forward_proposal_trace)
    log_abs_determinant = jacobian_correction(
        f, prev_model_trace, forward_proposal_trace, first_pass_results, nothing)
    constraints = merge(first_pass_results.constraints, new_observations)
    (new_model_trace, _) = generate(p_new, p_new_args, constraints)
    backward_proposal_trace, = generate(
        q_backward, (new_model_trace, q_backward_args...), first_pass_results.u_back)
    return (new_model_trace, backward_proposal_trace, log_abs_determinant)
end

function (translator::GeneralTraceTranslator)(
        prev_model_trace::Trace; check=false, prev_observations=EmptyChoiceMap())

    # sample auxiliary trace
    forward_proposal_trace = simulate(proposal, (prev_model_trace, translator.q_forward_args...,))

    # apply trace transform
    (new_model_trace, backward_proposal_trace, log_abs_determinant) = general_trace_translator_run_transform(
        translator.f, prev_model_trace, forward_proposal_trace, translator.new_observations,
        translator.p_new, translator.p_new_args, translator.q_backward, translator.q_backward_args)

    # compute log weight
    prev_model_score = get_score(prev_model_trace)
    new_model_score = get_score(new_model_trace)
    forward_proposal_score = get_score(forward_proposal_trace)
    backward_proposal_score = get_score(backward_proposal_trace)
    log_weight = new_model_score - prev_model_score + backward_proposal_score + forward_proposal_score + log_abs_determinant

    if check
        forward_proposal_choices = get_choices(forward_proposal_trace)
        f_inv = inverse(translator.f)
        (prev_model_trace_rt, forward_proposal_trace_rt, _) = general_trace_translator_run_transform(
            inverse(translator.f), prev_observations,
            new_model_trace, backward_proposal_trace,
            get_gen_fn(prev_model_trace), get_args(prev_model_trace),
            translator.q_forward, translator.q_forward_args)
        check_round_trip(
            prev_model_trace, prev_model_trace_rt,
            forward_proposal_trace, forward_proposal_trace_rt)
    end

    return (new_model_trace, log_weight)
end

##################################
# SimpleExtendingTraceTranslator #
##################################

"""
    translator = SimpleExtendingTraceTranslator(;
        p_new_args::Tuple = (),
        argdiffs::Tuple = (),
        new_obs::ChoiceMap = EmptyChoiceMap(),
        q_fwd::GenerativeFunction,
        q_fwd_args::Tuple  = ())

Constructor for a simple extending trace translator.

Run the translator with:

    (output_trace, log_weight) = translator(input_trace)
"""
@with_kw struct SimpleExtendingTraceTranslator
    p_new_args::Tuple = ()
    argdiffs::Tuple = ()
    new_obs::ChoiceMap = EmptyChoiceMap()
    q_fwd::GenerativeFunction
    q_fwd_args::Tuple  = ()
end

function (translator::SimpleExtendingTraceTranslator)(prev_model_trace::Trace)

    # simulate from auxiliary program
    forward_proposal_trace = simulate(translator.q_fwd, (prev_model_trace, translator.q_fwd_args...,))
    forward_proposal_score = get_score(forward_proposal_trace)

    # computing the new trace via update
    constraints = merge(get_choices(forward_proposal_trace), translator.new_obs)
    (new_model_trace, log_model_weight, _, discard) = update(
        prev_model_trace, translator.p_new_args,
        translator.argdiffs, constraints)

    if !isempty(discard)
        @error("can only extend the trace with random choices, cannot remove random choices")
        error("Invalid SimpleExtendingTraceTranslator")
    end

    log_weight = log_model_weight - forward_proposal_score
    return (new_model_trace, log_weight)
end

############################
# SymmetricTraceTranslator #
############################

"""
    translator = SymmetricTraceTranslator(;
        q::GenerativeFunction,
        q_args::Tuple = (),
        involution::Union{TraceTransformDSLProgram,Function})

Constructor for a symmetric trace translator.

The involution is either constructed via the [`@transform`](@ref) macro (recommended), or can be provided as a Julia function.

Run the translator with:

    (output_trace, log_weight) = translator(input_trace; check=false, observations=EmptyChoiceMap())

Use `check` to enable the involution check (this requires that the transform `f` has been marked with [`is_involution`](@ref)).

If `check` is enabled, then `observations` is a choice map containing the observed random choices, and the check will additionally ensure they are not mutated by the involution.
"""
@with_kw struct SymmetricTraceTranslator{T <: Union{TraceTransformDSLProgram,Function}}
    q::GenerativeFunction
    q_args::Tuple = ()
    involution::T # an involution
end

function symmetric_trace_translator_run_transform(
        involution::TraceTransformDSLProgram,
        prev_model_trace::Trace, forward_proposal_trace::Trace,
        q::GenerativeFunction, q_args::Tuple)
    first_pass_results = run_first_pass(involution, prev_model_trace, forward_proposal_trace)
    (new_model_trace, log_model_weight, _, discard) = update(
        prev_model_trace, get_args(prev_model_trace),
        map((_) -> NoChange(), get_args(prev_model_trace)),
        first_pass_results.constraints)
    log_abs_determinant = jacobian_correction(
        involution, prev_model_trace, forward_proposal_trace, first_pass_results, discard)
    backward_proposal_trace, = generate(
        q, (new_model_trace, q_args...), first_pass_results.u_back)
    return (new_model_trace, backward_proposal_trace, log_abs_determinant)
end

function (translator::SymmetricTraceTranslator{TraceTransformDSLProgram})(
        prev_model_trace::Trace; check=false, observations=EmptyChoiceMap())

    # simulate from auxiliary program
    forward_proposal_trace = simulate(translator.q, (prev_model_trace, translator.q_args...,))

    # apply trace transform
    (new_model_trace, backward_proposal_trace, log_abs_determinant) = symmetric_trace_translator_run_transform(
        translator.involution, prev_model_trace, forward_proposal_trace, translator.q, translator.q_args)

    # compute log weight
    prev_model_score = get_score(prev_model_trace)
    new_model_score = get_score(new_model_trace)
    forward_proposal_score = get_score(forward_proposal_trace)
    backward_proposal_score = get_score(backward_proposal_trace)
    log_weight = new_model_score - prev_model_score + backward_proposal_score - forward_proposal_score + log_abs_determinant

    if check
        check_observations(get_choices(new_model_trace), observations)
        forward_proposal_choices = get_choices(forward_proposal_trace)
        (prev_model_trace_rt, forward_proposal_trace_rt, _) = symmetric_trace_translator_run_transform(
            translator.involution, new_model_trace, backward_proposal_trace, translator.q, translator.q_args)
        check_round_trip(
            prev_model_trace, prev_model_trace_rt,
            forward_proposal_trace, forward_proposal_trace_rt)
    end

    return (new_model_trace, log_weight)
end

function (translator::SymmetricTraceTranslator{<:Function})(
        prev_model_trace::Trace; check=false, observations=EmptyChoiceMap())

    forward_trace = simulate(translator.q, (prev_model_trace, translator.q_args...,))
    forward_score = get_score(forward_trace)
    forward_choices = get_choices(forward_trace)
    forward_retval = get_retval(forward_trace)
    (new_model_trace, backward_choices, log_weight) = translator.involution(
        prev_model_trace, forward_choices, forward_retval, translator.q_args)
    (backward_score, backward_retval) = assess(translator.q, (new_model_trace, translator.q_args...), backward_choices)

    log_weight += (backward_score - forward_score)

    if check
        check_observations(get_choices(new_model_trace), observations)
        (prev_model_trace_rt, forward_choices_rt, _) = translator.involution(
            new_model_trace, backward_choices, backward_retval, translator.q_args)
        (forward_trace_rt, _) = generate(translator.q, (prev_model_trace, translator.q_args...), forward_choices_rt)
        check_round_trip(
            prev_model_trace, prev_model_trace_rt,
            forward_trace, forward_trace_rt)
    end

    return (new_model_trace, log_weight)
end

export @transform
export @read, @write, @copy, @tcall
export TraceTransformDSLProgram, pair_bijections!, is_involution!, inverse
export DeterministicTraceTranslator, SymmetricTraceTranslator, SimpleExtendingTraceTranslator, GeneralTraceTranslator

max_stars_repo_name: georgematheos/Gen

max_issues_repo_name: georgematheos/Gen

max_forks_repo_name: georgematheos/Gen


text:
using Pkg
Pkg.add("HTTP")
using Pkg
Pkg.add("Optim")
using Pkg
Pkg.add("GLM")
using Pkg
Pkg.add("LinearAlgebra")
using Pkg
Pkg.add("Random")
using Pkg
Pkg.add("Statistics")
using Pkg
Pkg.add("DataFrames")
using Pkg
Pkg.add("CSV")
using Pkg
Pkg.add("FreqTables")
using Optim, HTTP, GLM, LinearAlgebra, Random, Statistics, DataFrames, CSV, FreqTables
function ps3()
    #Question1
    url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2021/master/ProblemSets/PS3-gev/nlsw88w.csv"
df = CSV.read(HTTP.get(url).body, DataFrame)
X = [df.age df.white df.collgrad]
Z = hcat(df.elnwage1, df.elnwage2, df.elnwage3, df.elnwage4,
df.elnwage5, df.elnwage6, df.elnwage7, df.elnwage8)
y = df.occupation

function mlogit(Beta, X, Z, y)

    K = size(X,2)
    J = length(unique(y))
    N = length(y)
    bigY = zeros(N,J)
    for j=1:J
        bigY[:,j] = y.==j
    end
    bigBeta = [reshape(Beta,K,J-1) zeros(K)]

    bigZ = zeros(N,J)
    for j=1:J
        bigZ[:,j] = Z[:,j]-Z[:,J]
    end

    num = zeros(N,J)
    dem = zeros(N)
    for j=1:J
        XZ=cat(X,bigZ[:,j],dims=3)
        num[:,j] = exp.(XZ*bigBeta[:,j])
        dem .+= num[:,j]
    end

    P = num./repeat(dem,1,J)

    loglike = -sum( bigY.*log.(P) )

    return loglike
end

mlogit_hat_optim = optimize(b-> mlogit(b,X,Z, y), rand(7*(size(X,2))), LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
println(mlogit_hat_optim.minimizer)
#Question2
#Gamma represents the change in utility in percentage terms#
#Question3#
function nested_logit(alpha, X, Z, y, nesting_structure)
    beta = alpha[1:end-3]
    lambda = alpha[end-2:end-1]
    gamma = alpha[end]
    K = size(X,2)
    J = length(unique(y))
    N = length(y)
    bigY = zeros(N,J)
    for j=1:J
        bigY[:,j] = y.==j
    end
    bigBeta = [repeat(beta[1:K],1,length(nesting_structure[1])) repeat(beta[K+1:2K],1,length(nesting_structure[2])) zeros(K)]
end
ps3()

ext: jl


text:
function performance_func(prob, ctrl_amps, timeslices)
    λ = back_target_op(prob, ctrl_amps, timeslices)
    ρ = density_op(prob, ctrl_amps, timeslices)
    return expval(λ, ρ)
end

function C1_cost(init_unitary, tar_unitary, dim)
        return 1 - abs2(tr(tar_unitary'*init_unitary)/dim)
end

function C2_cost(init_state, final_state)
        return 1 - abs2(final_state'*init_state)
end

function density_op_fid(prob, ctrl_amps, init_density_op, target_density_op)
        ρt = density_op(prob, ctrl_amps, prob.n_steps)
        return expval(target_density_op, ρt)
end

function gate_infid(drift_Ham, ctrl_Hams, n_steps, transfer_time, ctrl_amps, tar_gate)
    KT = tar_gate
    KN = evolve(drift_Ham, ctrl_Hams, n_steps, transfer_time, ctrl_amps, n_steps)
    D = size(drift_Ham)[1]
    return 1 - abs2(tr(KT'*KN)/D)
end

function state_infid(drift_Ham, ctrl_Hams, n_steps, transfer_time, ctrl_amps, init_state, tar_state)
    ψT = tar_state
    ψN = evolve(drift_Ham, ctrl_Hams, n_steps, transfer_time, ctrl_amps, n_steps)*init_state
    return 1 - abs2(ψT'ψN)
end


text:
struct COSMOSAC10Param <: EoSParam
    Pnhb::SingleParam{Vector{Float64}}
    POH::SingleParam{Vector{Float64}}
    POT::SingleParam{Vector{Float64}}
    V::SingleParam{Float64}
    A::SingleParam{Float64}
end

abstract type COSMOSAC10Model <: COSMOSAC02Model end

struct COSMOSAC10{c<:EoSModel} <: COSMOSAC10Model
    components::Array{String,1}
    icomponents::UnitRange{Int}
    params::COSMOSAC10Param
    puremodel::Vector{c}
    absolutetolerance::Float64
    references::Array{String,1}
end

@registermodel COSMOSAC10
export COSMOSAC10

function COSMOSAC10(components; puremodel=PR,
    userlocations=String[],
     verbose=false)
    params = getparams(components, ["Activity/COSMOSAC/COSMOSAC10_like.csv"]; userlocations=userlocations, verbose=verbose)
    Pnhb  = COSMO_parse_Pi(params["Pnhb"])
    POH  = COSMO_parse_Pi(params["POH"])
    POT  = COSMO_parse_Pi(params["POT"])
    A  = params["A"]
    V  = params["V"]
    icomponents = 1:length(components)

    init_puremodel = [puremodel([components[i]]) for i in icomponents]
    packagedparams = COSMOSAC10Param(Pnhb,POH,POT,V,A)
    references = String[]
    model = COSMOSAC10(components,icomponents,packagedparams,init_puremodel,1e-12,references)
    return model
end

function lnγ_res(model::COSMOSAC10Model,V,T,z)
    x = z ./ sum(z)

    aeff = 7.5
    A = model.params.A.values
    n = A ./ aeff

    Pnhb  = model.params.Pnhb.values
    PSnhb = sum(x[i]*Pnhb[i] for i ∈ @comps) ./ sum(x[i]*A[i] for i ∈ @comps)

    POH = model.params.POH.values
    PSOH = sum(x[i]*POH[i] for i ∈ @comps) ./ sum(x[i]*A[i] for i ∈ @comps)

    POT = model.params.POT.values
    PSOT = sum(x[i]*POT[i] for i ∈ @comps) ./ sum(x[i]*A[i] for i ∈ @comps)

    (lnΓSnhb, lnΓSOH, lnΓSOT)= @f(lnΓ,PSnhb,PSOH,PSOT)
    lnΓi = [@f(lnΓ,Pnhb[i]./A[i],POH[i]./A[i],POT[i]./A[i]) for i ∈ @comps]

    lnγ_res_ =  [n[i]*(sum(Pnhb[i][v]/A[i]*(lnΓSnhb[v]-lnΓi[i][1][v]) for v ∈ 1:51)
                      +sum(POH[i][v]/A[i]*(lnΓSOH[v]-lnΓi[i][2][v]) for v ∈ 1:51)
                      +sum(POT[i][v]/A[i]*(lnΓSOT[v]-lnΓi[i][3][v]) for v ∈ 1:51)) for i ∈ @comps]
    return lnγ_res_
end

function lnΓ(model::COSMOSAC10Model,V,T,z,Pnhb,POH,POT)
    _TYPE = typeof(V+T+first(z))
    Γ0nhb = ones(_TYPE,length(Pnhb))
    Γ0OH = ones(_TYPE,length(POH))
    Γ0OT = ones(_TYPE,length(POT))

    σ   = -0.025:0.001:0.025
    Γnhb_old = exp.(-log.(sum(Pnhb[i]*Γ0nhb[i]*exp.(-ΔW.(σ,σ[i],1,1,T)./T) for i ∈ 1:51)
                         +sum(POH[i]*Γ0OH[i]*exp.(-ΔW.(σ,σ[i],2,1,T)./T) for i ∈ 1:51)
                         +sum(POT[i]*Γ0OT[i]*exp.(-ΔW.(σ,σ[i],3,1,T)./T) for i ∈ 1:51)))
    ΓOH_old = exp.(-log.(sum(Pnhb[i]*Γ0nhb[i]*exp.(-ΔW.(σ,σ[i],1,2,T)./T) for i ∈ 1:51)
                         +sum(POH[i]*Γ0OH[i]*exp.(-ΔW.(σ,σ[i],2,2,T)./T) for i ∈ 1:51)
                         +sum(POT[i]*Γ0OT[i]*exp.(-ΔW.(σ,σ[i],3,2,T)./T) for i ∈ 1:51)))
    ΓOT_old = exp.(-log.(sum(Pnhb[i]*Γ0nhb[i]*exp.(-ΔW.(σ,σ[i],1,3,T)./T) for i ∈ 1:51)
                         +sum(POH[i]*Γ0OH[i]*exp.(-ΔW.(σ,σ[i],2,3,T)./T) for i ∈ 1:51)
                         +sum(POT[i]*Γ0OT[i]*exp.(-ΔW.(σ,σ[i],3,3,T)./T) for i ∈ 1:51)))
    Γnhb_new = deepcopy(Γnhb_old)
    ΓOH_new = deepcopy(ΓOH_old)
    ΓOT_new = deepcopy(ΓOT_old)
    tol = 1
    i = 1

    while tol>sqrt(model.absolutetolerance)
        Γnhb_new .= exp.(-log.(sum(Pnhb[i]*Γnhb_old[i]*exp.(-ΔW.(σ,σ[i],1,1,T)./T) for i ∈ 1:51)
                         +sum(POH[i]*ΓOH_old[i]*exp.(-ΔW.(σ,σ[i],2,1,T)./T) for i ∈ 1:51)
                         +sum(POT[i]*ΓOT_old[i]*exp.(-ΔW.(σ,σ[i],3,1,T)./T) for i ∈ 1:51)))
        ΓOH_new .= exp.(-log.(sum(Pnhb[i]*Γnhb_old[i]*exp.(-ΔW.(σ,σ[i],1,2,T)./T) for i ∈ 1:51)
                         +sum(POH[i]*ΓOH_old[i]*exp.(-ΔW.(σ,σ[i],2,2,T)./T) for i ∈ 1:51)
                         +sum(POT[i]*ΓOT_old[i]*exp.(-ΔW.(σ,σ[i],3,2,T)./T) for i ∈ 1:51)))
        ΓOT_new .= exp.(-log.(sum(Pnhb[i]*Γnhb_old[i]*exp.(-ΔW.(σ,σ[i],1,3,T)./T) for i ∈ 1:51)
                         +sum(POH[i]*ΓOH_old[i]*exp.(-ΔW.(σ,σ[i],2,3,T)./T) for i ∈ 1:51)
                         +sum(POT[i]*ΓOT_old[i]*exp.(-ΔW.(σ,σ[i],3,3,T)./T) for i ∈ 1:51)))

        Γnhb_new .= (Γnhb_new .+ Γnhb_old) ./2
        ΓOH_new .= (ΓOH_new .+ ΓOH_old) ./2
        ΓOT_new .= (ΓOT_new .+ ΓOT_old) ./2

        tol = (cosmo_tol(Γnhb_new,Γnhb_old) + cosmo_tol(ΓOH_new,ΓOH_old) + cosmo_tol(ΓOT_new,ΓOT_old))/3
        Γnhb_old .= Γnhb_new
        ΓOH_old .= ΓOH_new
        ΓOT_old .= ΓOT_new
        i+=1
    end
    lnΓnhb = log.(Γnhb_old)
    lnΓOH = log.(ΓOH_old)
    lnΓOT = log.(ΓOT_old)

    return (lnΓnhb,lnΓOH,lnΓOT)
end

function ΔW(σm,σn,t,s,T)
    ces  = 6525.69+1.4859e8/T^2
    if t==2 && s==2 && σm*σn<0
        chb = 4013.78
    elseif t==3 && s==3 && σm*σn<0
        chb = 932.31
    elseif t==3 && s==2 && σm*σn<0
        chb = 3016.43
    else
        chb = 0.
    end
    R    = 0.001987
    return (ces*(σm+σn)^2-chb*(σm-σn)^2)/R
end


text:
"""
$(TYPEDEF)

A generic function was called with arguments that should be valid but for which
there is no existing implementation.
"""
struct NotImplementedError <: Exception
	f
	args::Tuple
    msg::Union{String, Nothing}
end

NotImplementedError(f, args::Tuple) = NotImplementedError(f, args, nothing)

function Base.showerror(io::IO, e::NotImplementedError)
    print(io, "NotImplementedError: ", e.f, e.args)
    if !isnothing(e.msg)
        print(io, " ", e.msg)
    end
end

is_notimplemented(::Exception) = false
is_notimplemented(::Exception, f) = false
is_notimplemented(::NotImplementedError) = true
is_notimplemented(e::NotImplementedError, f) = e.f == f

"""
$(SIGNATURES)

Define a function/method which raises [`NotImplementedError`](@ref) when called.
"""
macro notimplemented(expr, msg::Union{String, Nothing}=nothing)
    Meta.isexpr(expr, :call) || error("Expression must be function call")

    f = esc(expr.args[1])

    args = Symbol[]
    for a in expr.args[2:end]
        Meta.isexpr(a, :parameters) && error("Keyword arguments not supported")
        if a isa Symbol
            push!(args, a)
        elseif Meta.isexpr(a, :(::))
            push!(args, a.args[1])
        else
            error("Expected function parameter, got $a")
        end
    end

    return quote
        $(esc(expr)) = throw($NotImplementedError($f, ($(esc.(args)...),), $msg))
    end
end

"""
$(SIGNATURES)

Return result of evaluating `expr`, but catch [`NotImplementedError`](@ref) and return `default` instead.
"""
macro notimplemented_default(expr, default)
	return quote
		try
			$(esc(expr))
		catch e
			$is_notimplemented(e) ? $(esc(default)) : throw(e)
		end
	end
end


text:
#=##############################################################################
# DESCRIPTION
    Methods for generating a wind turbine geometry.
# AUTHORSHIP
  * Author    : Eduardo J Alvarez
  * Email     : Edo.AlvarezR@gmail.com
  * Created   : Jul 2018
  * License   : MIT License
=###############################################################################

################################################################################
# WIND TURBINE
################################################################################
"""
`generate_windturbine(Rtip::Float64, h::Float64, blade_name::String,
                              hub_name::String, tower_name::String;
                              nblades::Int64=3, data_path::String=def_data_path,
                              save_path=nothing, file_name="windturbine",
                              paraview=true)`

  Generates and returns a WindTurbine object scaled by `Rtip` with a tower of
height `h`, where the blade will be read from the `blade_name.jld` in
`data_path` containing the TriangularGrid geometry of the blade, samething for
`hub_name` and `tower_name`.

See the notebook under `docs/geometry_processing` for an example on how to
generate these files.
"""
function generate_windturbine(Rtip::Float64, h::Float64, blade_name::String,
                              hub_name::String, tower_name::String;
                              nblades::Int64=3, data_path::String=def_data_path,
                              save_path=nothing, file_name="windturbine",
                              paraview=true,
                              rot=nothing, random_rot::Bool=true, pitch=70)

  # Read gridded geometries
  blade_grid = JLD.load(joinpath(data_path, blade_name*".jld"), "blade_grid")
  hub_grid, Rhub, Thub = JLD.load(joinpath(data_path, hub_name*".jld"),
                                                    "hub_grid", "Rhub", "Thub")
  tower_grid = JLD.load(joinpath(data_path, tower_name*".jld"), "tower_grid")
  # Scales dimensions by Rtip
  blade_grid.orggrid.nodes .*= Rtip
  hub_grid.orggrid.nodes .*= Rtip
  for i in 1:size(tower_grid.orggrid.nodes, 2)
    # println([Rtip, h, Rtip])
    tower_grid.orggrid.nodes[:, i] .*= [Rtip, h, Rtip]
      # println(tower_grid.orggrid.nodes[:, i])
  end
  Rhub *= Rtip
  Thub *= Rtip

  # Rotor center
  C = [Thub*5/12, 0, h+Rhub/2]

  # Initiates rotor multigrid
  rotor = gt.MultiGrid(3)

  # Aligns and add hub to rotor
  gt.lintransform!(hub_grid, gt.rotation_matrix(0, 90, 0), zeros(3))
  gt.addgrid(rotor, "hub", hub_grid)

  # Pitches the blade
  gt.lintransform!(blade_grid, gt.rotation_matrix(0, -pitch, 0), zeros(3))

  # Rotates the rotor
  if rot != nothing
    gt.lintransform!(blade_grid, gt.rotation_matrix(0, 0, rot), zeros(3))
  elseif random_rot
    gt.lintransform!(blade_grid, gt.rotation_matrix(0, 0, rand()*360), zeros(3))
  end

  # Creates every blade and adds them to the rotor
  gt.lintransform!(blade_grid, gt.rotation_matrix(0, 0, -90), zeros(3))
  rotM = gt.rotation_matrix(0, 0, 360/nblades)
  for i in 1:nblades
    this_blade = deepcopy(blade_grid)
    gt.lintransform!(this_blade, eye(3), [-Thub*5/6, 0, 0])
    gt.addgrid(rotor, "blade$i", this_blade)

    gt.lintransform!(blade_grid, rotM, zeros(3))
  end

  # Starts multigrid of the wind turbine
  windturbine = gt.MultiGrid(3)

  # Aligns and adds tower
  gt.lintransform!(tower_grid, gt.rotation_matrix(0, 0, -90), zeros(3))
  gt.addgrid(windturbine, "tower", tower_grid)

  # Translates and adds rotor
  # gt.lintransform!(rotor, eye(3), C-*[Thub/2, 0, 0])
  gt.lintransform!(rotor, eye(3), C)
  gt.addgrid(windturbine, "rotor", rotor)

  if save_path!=nothing
    gt.save(windturbine, file_name; path=save_path)

    if paraview
      strn = ""
      strn *= file_name*"_tower.vtk;"
      strn *= file_name*"_rotor_hub.vtk;"
      for i in 1:nblades
        strn *= file_name*"_rotor_blade$i.vtk;"
      end

      run(`paraview --data=$save_path/$strn`)
    end

  end

  return windturbine::gt.MultiGrid
end

"""
`generate_loft(bscale::Real, b_low::Real, b_up::Real, b_NDIVS::Int64,
                        chords::Array{T,2}, twists::Array{T,2},
                        LE_x::Array{T,2}, LE_z::Array{T,2},
                        airfoils::Array{Tuple{T,Array{T,2}}, 1};
                        # MORE GEOMETRIC OPTIONS
                        tilt_z=nothing,
                        # SPLINE OPTIONS
                        spline_k::Int64=5, spline_bc::String="extrapolate",
                        spline_s::Real=0.001, verify_spline::Bool=true,
                        # OUTPUT OPTIONS
                        save_path=nothing, paraview::Bool=true,
                        file_name::String="myloft"
                       ) where{T<:Real}`

  Generates a lofted geometry.

  **Arguments**
  * `bscale::Float64`         : Semi-span scale.
  * `b_low::Float64`          : Scaled low bound of the span.
  * `b_up::Float64`           : Scaled upper bound of the span. To generate
                                a symmetric wing, give it b_low=-1, b_up=1; for
                                a semi-span, give it b_low=0, b_up=1. All
                                y/bscale value in the following arguments must
                                go from 0 to 1.
  * `b_NDIVS::Int64`          : Number of division along span.
  * `chords::Array{Float64,2}`: Chord distribution along the span in the form
                                [(y/bscale, c/bscale)].
  * `twists::Array{Float64,2}`: Twist distribution along the span in the form
                                [(y/bscale, deg)].
  * `LE_x::Array{Float64,2}`  : x-position (chordwise) of leading edge along the
                                span in the form [(y/bscale, x/bscale)].
  * `LE_z::Array{Float64,2}`  : z-position (dihedral-wise) of leading edge along
                                the span in the form [(y/bscale, z/bscale)].
  * `airfoils`                : Airfoil cross sections along the span in the
                                form [(y/bscale, airfoil_contour)], where
                                `airfoil_contour` is a matrix that contains
                                all points of the airfoil contour indexed by row.

  **Optional Arguments**
  * `tilt_z::{Array{Float64,2}}`            : Tilting about the z-axis of
                                              every span cross section in the
                                              form [(y/bscale, deg)].
  * `spline_k`, `spline_bc`, `spline_s`     : Spline parameters.

"""
function generate_loft(bscale::Real, b_low::Real, b_up::Real, b_NDIVS::Int64,
                        chords::Array{T,2}, twists::Array{T,2},
                        LE_x::Array{T,2}, LE_z::Array{T,2},
                        airfoils::Array{Tuple{T,Array{T,2}}, 1};
                        # MORE GEOMETRIC OPTIONS
                        tilt_z=nothing,
                        # SPLINE OPTIONS
                        spline_k::Int64=5, spline_bc::String="extrapolate",
                        spline_s::Real=0.001, verify_spline::Bool=true,
                        # OUTPUT OPTIONS
                        save_path=nothing, paraview::Bool=true,
                        file_name::String="myloft"
                       ) where{T<:Real}

  rfl_NDIVS = size(airfoils[1][2], 1)-1
  for (pos, rfl) in airfoils
    if rfl_NDIVS!=size(rfl,1)-1
      error("All airfoil sections must have the same number of points.")
    end
  end

  # ----------------- PARAMETRIC GRID ------------------------------------------
  P_min = [0, b_low, 0]            # Lower boundary arclength, span, dummy
  P_max = [1, b_up, 0 ]            # Upper boundary arclength, span, dummy
  loop_dim = 1                     # Loops the arclength dimension

  # Adds dummy division
  NDIVS = [rfl_NDIVS, b_NDIVS, 0]

  # Generates parametric grid
  grid = gt.Grid(P_min, P_max, NDIVS, loop_dim)

  # ----------------- GEOMETRY SPLINES -----------------------------------------
  # Splines all distributions for a smooth geometry
  _spl_chord = Dierckx.Spline1D(chords[:, 1], chords[:, 2];
                      k= size(chords,1)>=spline_k ? spline_k : size(chords,1)-1,
                                s=spline_s, bc=spline_bc)
  _spl_twist = Dierckx.Spline1D(twists[:, 1], twists[:, 2];
                      k= size(twists,1)>=spline_k ? spline_k : size(twists,1)-1,
                                s=spline_s, bc=spline_bc)
  _spl_LE_x = Dierckx.Spline1D(LE_x[:, 1], LE_x[:, 2];
                      k= size(LE_x,1)>=spline_k ? spline_k : size(LE_x,1)-1,
                                s=spline_s, bc=spline_bc)
  _spl_LE_z = Dierckx.Spline1D(LE_z[:, 1], LE_z[:, 2];
                      k= size(LE_z,1)>=spline_k ? spline_k : size(LE_z,1)-1,
                                s=spline_s, bc=spline_bc)
  if tilt_z!=nothing
    _spl_tlt_z = Dierckx.Spline1D(tilt_z[:, 1], tilt_z[:, 2];
                      k= size(tilt_z,1)>=spline_k ? spline_k : size(tilt_z,1)-1,
                                s=spline_s, bc=spline_bc)
  end

  # ----------------- SPLINE VERIFICATION --------------------------------------
  if verify_spline
    nnodesspan = gt.get_ndivsnodes(grid)[2]    # Number of nodes along span
    y_poss = [gt.get_node(grid, [1,i])[2] for i in 1:nnodesspan]  # Span positions

    fig = plt.figure("spl_verif", figsize=(7*2,5*1))

    plt.subplot(121)
    plt.plot(LE_x[:,1], LE_x[:,2], "og", label="Org LE x", alpha=0.5)
    plt.plot(LE_z[:,1], LE_z[:,2], "ob", label="Org LE z", alpha=0.5)
    plt.plot(y_poss, [_spl_LE_x(y) for y in y_poss], "--g", label="Spline LE x")
    plt.plot(y_poss, [_spl_LE_z(y) for y in y_poss], "--b", label="Spline LE z")
    plt.xlabel(plt.L"y/b_{scale}")
    plt.ylabel(plt.L"x/b_{scale}, z/b_{scale}")
    plt.grid(true, color="0.8", linestyle="--")
    plt.legend(loc="best")

    plt.subplot(122)
    p1 = plt.plot(twists[:,1], twists[:,2], "og", label="Org Twist", alpha=0.5)
    p2 = plt.plot(y_poss, [_spl_twist(y) for y in y_poss], "--g",
                                                          label="Spline twist")
    pextra = []
    if tilt_z!=nothing
      pextra1 = plt.plot(tilt_z[:,1], tilt_z[:,2], "or", label="Org tilt z",
                                                                      alpha=0.5)
      pextra2 = plt.plot(y_poss, [_spl_tlt_z(y) for y in y_poss], "--r",
                                                          label="Spline tilt z")
      pextra = vcat(pextra, [pextra1[1], pextra2[1]])
    end
    plt.ylabel("Twist (deg)")

    plt.grid(true, color="0.8", linestyle="--")
    plt.xlabel(plt.L"y/b_{scale}")

    plt.twinx()
    p3 = plt.plot(chords[:,1], chords[:,2], "ob", label="Org Chord", alpha=0.5)
    p4 = plt.plot(y_poss, [_spl_chord(y) for y in y_poss], "--b",
                                                          label="Spline chord")
    plt.ylabel(plt.L"c/b_{scale}")

    ps = vcat([p1[1], p2[1], p3[1], p4[1]], pextra)
    plt.legend(ps, [p[:get_label]() for p in ps], loc="best")
  end

  # ----------------- SURFACE GRID ---------------------------------------------
  # Auxiliary function for weighting values across span
  function calc_vals(span, array)

    # Finds bounding airfoil position
    val_in, val_out = nothing, array[1]
    for val in array[2:end]
        val_in = val_out
        val_out = val
        if val[1]>=abs(span); break; end
    end
    pos_in = val_in[1]
    val_in = val_in[2]
    pos_out = val_out[1]
    val_out = val_out[2]

    weight = (abs(span)-pos_in)/(pos_out-pos_in)

    return weight, val_in, val_out
  end

  # Space transformation function
  function my_space_transform(X, inds)
    span = X[2]                     # y/bscale span position
    chord = _spl_chord(abs(span))   # c/bscale chord length
    twist = _spl_twist(abs(span))   # twist
    le_x = _spl_LE_x(abs(span))     # x/bscale LE position
    le_z = _spl_LE_z(abs(span))     # z/bscale LE position

    # Merges airfoil contours at this span position
    weight, rfl_in, rfl_out = calc_vals(span, airfoils)

    # Point over airfoil contour
    point = weight*rfl_out[inds[1], :]+(1-weight)*rfl_in[inds[1], :]
    point = vcat(point, 0)

    # Scales the airfoil contour by the normalized chord length
    point = chord*point

    # Applies twist to the airfoil point
    tlt_z = tilt_z!=nothing ?  _spl_tlt_z(abs(span)) : 0.0
    point = gt.rotation_matrix(-twist, -tlt_z, 0)*point

    # Places the point relative to LE and scales by span scale
    point = [point[1]+le_x, span+point[3], point[2]+le_z]*bscale

    return point
  end

  # Transforms the quasi-two dimensional parametric grid into the wing surface
  gt.transform3!(grid, my_space_transform)

  # Splits the quadrialateral panels into triangles
  dimsplit = 2              # Dimension along which to split
  triang_grid = gt.GridTriangleSurface(grid, dimsplit)

  if save_path!=nothing
    # Outputs a vtk file
    gt.save(triang_grid, file_name; path=save_path)

    # Outputs a jld file
    JLD.save(joinpath(save_path, "$file_name.jld"), "triang_grid", triang_grid)

    if paraview
      # Calls paraview
      run(`paraview --data=$save_path$file_name.vtk`)
    end
  end

  return triang_grid
end

"""
  `generate_blade(Rtip, Rhub, r_NDIVS, blade_name; data_path=def_data_path,
    optargs...)`

  Reads the blade geometry in the data path under `blade_name`, and generates
a lofted geometry of the blade. Returns a TriangularGrid object of the geometry.

  **Arguments**
  * `Rtip::Float64`           : Blade radius.
  * `Rhub::Float64`           : Hub radius.
  * `r_NDIVS::Float64`        : Number of divisions along blade.
  * `blade_name::String`      : Blade geometry identifier.

"""
function generate_blade(Rtip::Real, Rhub::Real, r_NDIVS::Int64,
                        blade_name::String; data_path::String=def_data_path,
                        optargs...)

  # Reads all data
  data_airfoil = CSV.read(joinpath(data_path, blade_name*"_airfoilsections.csv"))
  data_chord = CSV.read(joinpath(data_path, blade_name*"_chord.csv"))
  data_twist = CSV.read(joinpath(data_path, blade_name*"_twist.csv"))
  data_lex = CSV.read(joinpath(data_path, blade_name*"_lex.csv"))
  data_lez = CSV.read(joinpath(data_path, blade_name*"_lez.csv"))

  airfoils = Tuple{Float64, Array{Float64, 2}}[]
  for i in 1:size(data_airfoil, 1)
    pos = data_airfoil[1][i]
    file_name = blade_name*"_"*data_airfoil[2][i]

    # rfl_contour = Array(CSV.read(joinpath(data_path, "airfoils/$file_name")))
    rfl_contour = CSV.read(joinpath(data_path, "airfoils/$file_name"))

    # Convert DataFrames to Arrays
    rfl_contour = hcat(rfl_contour[1], rfl_contour[2])
    rfl_contour = [val for val in rfl_contour]

    push!(airfoils, (pos, rfl_contour))
  end

  # Convert DataFrames to Arrays
  data_chord = hcat(data_chord[1], data_chord[2])
  data_twist = hcat(data_twist[1], data_twist[2])
  data_lex = hcat(data_lex[1], data_lex[2])
  data_lez = hcat(data_lez[1], data_lez[2])

  # Get rid of the Missing type
  data_chord = [val for val in data_chord]
  data_twist = [val for val in data_twist]
  data_lex = [val for val in data_lex]
  data_lez = [val for val in data_lez]

  return generate_loft(Rtip, Rhub/Rtip, 1.0, r_NDIVS, data_chord,
                  data_twist, data_lex, data_lez,
                  airfoils; file_name=blade_name, optargs...)
end
# ------------ END OF WIND TURBINE ---------------------------------------------

max_issues_repo_name: byuflowlab/wind-farm-utilities

max_forks_repo_name: byuflowlab/wind-farm-utilities


text:
struct cmps
    Q::TensorMap{ComplexSpace, 1, 1}
    R::TensorMap{ComplexSpace, 2, 1}
end

"""
    rrule(::typeof(cmps), Q::TensorMap{ComplexSpace, 1, 1}, R::TensorMap{ComplexSpace, 2, 1})

    rrule for the constructor of `cmps`.
"""
function rrule(::typeof(cmps), Q::TensorMap{ComplexSpace, 1, 1}, R::TensorMap{ComplexSpace, 2, 1})
    function cmps_pushback(f̄wd)
        return NoTangent(), f̄wd.Q, f̄wd.R
    end
    return cmps(Q, R), cmps_pushback
end

"""
    cmps(f, chi::Integer, d::Integer) -> cmps

    generate a cmps with vitual dim `chi` and physical dim `d` using function `f` to generate the parameters.
    Example: `cmps(rand, 2, 4)`
"""
function cmps(f, chi::Integer, d::Integer)
    Q = TensorMap(f, ComplexF64, ℂ^chi, ℂ^chi)
    R = TensorMap(f, ComplexF64, ℂ^chi*ℂ^d, ℂ^chi)
    return cmps(Q, R)
end

+(psi::cmps, phi::cmps) = cmps(psi.Q + phi.Q, psi.R + phi.R)
-(psi::cmps, phi::cmps) = cmps(psi.Q - phi.Q, psi.R - phi.R)
*(psi::cmps, x::Number) = cmps(psi.Q * x, psi.R * x)
*(x::Number, psi::cmps) = cmps(psi.Q * x, psi.R * x)
similar(psi::cmps) = cmps(similar(psi.Q), similar(psi.R))
length(iter::cmps) = length(iter.Q.data) + length(iter.R.data)
iterate(iter::cmps) = (iter.Q[1], 1)
function iterate(iter::cmps, state)
    next_state = state + 1
    len_Q = length(iter.Q.data)
    len_R = length(iter.R.data)
    if next_state <= len_Q
        next_elem = iter.Q[next_state]
    elseif next_state <= len_Q + len_R
        next_elem = iter.R[next_state - len_Q]
    else
        return nothing
    end
    return (next_elem, next_state)
end
function getindex(psi::cmps, ix::Integer)
    len_Q = length(psi.Q.data)
    len_R = length(psi.R.data)
    if ix <= len_Q && ix > 0
        return psi.Q[ix]
    elseif ix <= len_Q + len_R
        return psi.Q[ix]
    else
        throw(BoundsError())
    end
end

@inline get_chi(psi::cmps) = get_chi(psi.R)
@inline get_d(psi::cmps) = get_d(psi.R)

"""
    get_matrices(psi::cmps) -> tuple

    return the Q and R matrices in cMPS `psi`.
    Q and R will be in the form of a tuple of `TensorMap`s.
"""
get_matrices(psi::cmps) = (psi.Q, psi.R)

"""
    rrule(::typeof(get_matrices), psi::cmps)
"""
function rrule(::typeof(get_matrices), psi::cmps)
    fwd = (psi.Q, psi.R)
    function get_Q_pushback(f̄wd)
        (Q̄, R̄) = f̄wd
        (Q̄ isa ZeroTangent) && (Q̄ = TensorMap(zeros, ComplexF64, space(psi.Q)))
        (R̄ isa ZeroTangent) && (R̄ = TensorMap(zeros, ComplexF64, space(psi.R)))
        p̄si = cmps(Q̄, R̄)
        return NoTangent(), p̄si
    end
    return fwd, get_Q_pushback
end

"""
    convert_to_cmps(arr::Array{<:Complex, 3})

    Convert a chi*(d+1)*chi array to a cMPS.
    `chi` is the virtual bond dimension, `d+1` is the physical bond dimension.
"""
function convert_to_cmps(arr::Array{<:Complex, 3})
    Q = convert_to_tensormap(arr[:, 1, :], 1)
    R = convert_to_tensormap(arr[:, 2:end, :], 2)
    return cmps(Q, R)
end

"""
    convert_to_cmps(arr::Array{<:Complex, 1}, χ::Integer, d::Integer)

    Convert a vector to a left-canonical cMPS.
    ```
    Q = - R†R / 2 - i D
    ```
    where `D` is a real diagonal matrix.
"""
function convert_to_cmps(arr::Array{<:Complex, 1}, χ::Integer, d::Integer)
    D = diagm(real(arr[1:χ]))
    R = reshape(arr[χ+1:end], (χ*d, χ))
    Q = -0.5 * R' * R - 1im * D
    ψ_arr = reshape([Q; R], (χ, d+1, χ))
    return convert_to_cmps(ψ_arr)
end

"""
    convert_to_cmps(arr::TensorMap{ComplexSpace, 2, 1})

    Convert a ℂ^chi*ℂ^(d+1) ← ℂ^chi TensorMap to a cMPS.
"""
function convert_to_cmps(arr::TensorMap{ComplexSpace, 2, 1})
    return convert_to_cmps(convert_to_array(arr))
end

"""
    convert_to_array(psi::cmps)

    Convert a cmps to an array.
"""
function convert_to_array(psi::cmps)
    chi = get_chi(psi)
    Q_arr, R_arr = convert_to_array(psi.Q), convert_to_array(psi.R)
    Q_arr = reshape(Q_arr, (chi, 1, chi))
    arr = cat(Q_arr, R_arr, dims=2)
    return arr
end

"""
    convert_to_tensormap(psi::cmps)

    Convert a cmps to a tensormap.
"""
convert_to_tensormap(psi::cmps) = convert_to_tensormap(convert_to_array(psi), 2)

"""
    transf_mat(psi::cmps, phi::cmps) -> Function

    Obtain the transfer matrix of <psi|phi> as a linear operator.
"""
function transf_mat(psi::cmps, phi::cmps)
    function lop(v::TensorMap{ComplexSpace, 1, 1})
        @tensor Tv[-1; -2] := v[-1, 1] * psi.Q'[1, -2] +
                            phi.Q[-1, 1] * v[1, -2] +
                            phi.R[-1, 3, 1] * v[1, 2] * psi.R'[2, -2, 3]
        return Tv
    end
    return lop
end

"""
    transf_mat_T(psi::cmps, phi::cmps) -> Function

    Obtain the Hermtian conjugate of the transfer matrix of <psi|phi> as a linear operator.
"""
function transf_mat_T(psi::cmps, phi::cmps)
    function lop_T(v::TensorMap{ComplexSpace, 1, 1})
        @tensor Tv[-1; -2] := v[-1, 1] * psi.Q[1, -2] +
                            phi.Q'[-1, 1] * v[1, -2] +
                            phi.R'[-1, 1, 3] * psi.R[2, 3, -2] * v[1, 2]
        return Tv
    end
    return lop_T
end

"""
    leftorth(A::cmps) -> AL, Tϵ

    QR decomposition of cMPS local tensor.

    To avoid confusion (as we have used Q, R to denote the matrices inside the cMPS local tensor),
    we use notations A = AL T, where A is the input cMPS local tensor, AL is the left-orthogonalized cMPS local tensor (AL^{\\dagger} AL = I), and T = I + ϵ Tϵ is the upper triangular matrix.

    Returns AL and Tϵ.
"""
function leftorth(A::cmps)
    chi = get_chi(A)

    # calculate Tϵ
    Tϵ = A.Q + A.Q' + A.R' * A.R
    Tϵ_arr = triu(Tϵ.data) - 0.5 * Diagonal(Tϵ.data)
    Tϵ = TensorMap(Tϵ_arr, ℂ^chi, ℂ^chi)

    # calculate QL, and RL
    QL = A.Q - Tϵ
    RL = A.R

    return cmps(QL, RL), Tϵ
end

"""
    left_canonical(psi::cmps)

    Convert the input cmps into the left-canonical form.
    Return the gauge transformation matrix X and the left-canonicalized cmps.
"""
function left_canonical(psi::cmps)
    chi = get_chi(psi)

    # solve the fixed point equation
    vl = TensorMap(rand, ComplexF64, ℂ^chi, ℂ^chi)
    lop_T = transf_mat_T(psi, psi)
    w, vl = eigsolve(lop_T, vl, 1, :LR)
    w = w[1]
    vl = vl[1]

    # obtain gauge transformation matrix X
    _, s, u = tsvd(vl)
    X = sqrt(s) * u
    Xinv = u' * sqrt(inv(s))

    # update Q and R
    Q = X * psi.Q * Xinv - 0.5 * w * id(ℂ^chi) # normalized
    @tensor R[-1, -2; -3] := X[-1, 1] * psi.R[1, -2, 2] * Xinv[2, -3]

    return X, cmps(Q, R)
end

"""
    right_canonical(psi::cmps) -> Y::Matrix, psiR::cmps

    Convert the input cmps into the right-canonical form.
    Return the gauge transformation matrix Y and the right-canonicalized cmps.
"""
function right_canonical(psi::cmps)
    chi = get_chi(psi)

    # solve the fixed point equation
    vr = TensorMap(rand, ComplexF64, ℂ^chi, ℂ^chi)
    lop = transf_mat(psi, psi)
    w, vr = eigsolve(lop, vr, 1, :LR)
    w = w[1]
    vr = vr[1]

    # obtain gauge transformation matrix Yinv
    _, s, u = tsvd(vr)
    Y = sqrt(s) * u
    Yinv = u' * sqrt(inv(s))

    # update Q and R
    Q = Yinv' * psi.Q * Y' - 0.5 * w * id(ℂ^chi)
    @tensor R[-1, -2; -3] := Yinv'[-1, 1] * psi.R[1, -2, 2] * Y'[2, -3]

    return Y, cmps(Q, R)
end

"""
    entanglement_spectrum(psi::cmps) -> Vector

    return the entanglement spectrum of an infinite cmps.
"""
function entanglement_spectrum(psi::cmps)
    _, psi = left_canonical(psi)
    Y, _ = right_canonical(psi)
    _, s, _ = tsvd(Y')
    s = diag(s.data) .^ 2
    return s ./ sum(s)
end

"""
    entanglement_entropy(psi::cmps) -> Float64

    return the entanglement entropy.
"""
function entanglement_entropy(psi::cmps)
    s = entanglement_spectrum(psi)
    return -sum(s .* log.(s))
end

"""
    expand(psi::cmps, chi::Integer, perturb::Float64=1e-3) -> cmps

    expand the cmps `psi` to a target bond dimension `chi` by adding small numbers of size `perturb`.
"""
function expand(psi::cmps, chi::Integer, perturb::Float64=1e-3)
    chi0, d = get_chi(psi), get_d(psi)
    if chi <= chi0
        @warn "chi not larger than current bond D, not expanded "
        return psi
    end

    Q_arr = perturb*rand(ComplexF64, chi, chi)
    R_arr = perturb*rand(ComplexF64, chi, d, chi)
    Q_arr[1:chi0, 1:chi0] += toarray(psi.Q)
    R_arr[1:chi0, :, 1:chi0] += toarray(psi.R)
    Q = TensorMap(Q_arr, ℂ^chi, ℂ^chi)
    R = arr_to_TensorMap(R_arr)

    return cmps(Q, R)
end

"""
    act(op::cmpo, psi::cmps) -> cmps

    Act a cmpo to a cmps.
"""
function act(T::cmpo, psi::cmps)
    chi_cmpo, chi_psi = get_phy(T), get_chi(psi)
    chi_tot = chi_cmpo * chi_psi
    t_fuse = isomorphism(ℂ^chi_tot, ℂ^chi_cmpo*ℂ^chi_psi)

    @tensor Q[-1; -2] := t_fuse[-1, 1, 3] * T.Q[1, 2] * t_fuse'[2, 3, -2] +
                         t_fuse[-1, 3, 1] * psi.Q[1, 2] * t_fuse'[3, 2, -2] +
                         t_fuse[-1, 1, 2] * T.L'[1, 4, 3] * psi.R[2, 3, 5] * t_fuse'[4, 5, -2]
    @tensor R[-1, -2; -3] := t_fuse[-1, 1, 3] * T.R[1, -2, 2] * t_fuse'[2, 3, -3]

    return cmps(Q, R)
end

"""
    rrule(::typeof(act), T::cmpo, psi::cmps)

    Reverse diff rule for `act(op::cmpo, psi::cmps)`.
    We only calculate the adjoint of `psi`.
"""
function rrule(::typeof(act), T::cmpo, psi::cmps)
    chi_cmpo, chi_psi = get_phy(T), get_chi(psi)
    chi_tot = chi_cmpo * chi_psi
    t_fuse = isomorphism(ℂ^chi_tot, ℂ^chi_cmpo*ℂ^chi_psi)
    fwd = act(T, psi)

    function act_pushback(f̄wd)
        @tensor Q̄_psi[-1; -2] := f̄wd.Q[1, 2] * t_fuse'[3, -1, 1] * t_fuse[2, 3, -2]
        @tensor R̄_psi[-1, -2; -3] := f̄wd.Q[1, 2] * t_fuse'[3, -1, 1] * t_fuse[2, 4, -3] * T.L[4, -2, 3]
        p̄si = cmps(Q̄_psi, R̄_psi)
        return NoTangent(), NoTangent(), p̄si
    end
    return fwd, act_pushback
end

"""
    K_mat(phi::cmps, psi::cmps) -> Kmat::TensorMap{ComplexSpace, 2, 2}

    calculate the K_mat from two cmpses `phi` and `psi`. order of indices:

        -1 -->--  phi' -->-- -3
                   |
                 1 ^
                   |
        -2 --<--  psi  --<-- -4

"""
function K_mat(phi::cmps, psi::cmps)
    Id_phi = id(ℂ^get_chi(phi))
    Id_psi = id(ℂ^get_chi(psi))
    @tensor Kmat[-1, -2; -3, -4] := phi.Q'[-3, -1] * Id_psi[-2, -4] +
                                    Id_phi'[-3, -1] * psi.Q[-2, -4] +
                                    phi.R'[-3, -1, 1] * psi.R[-2, 1, -4]
    return Kmat
end

"""
    rrule(::typeof(K_mat), phi::cmps, psi::cmps) -> fwd::TensorMap{ComplexSpace, 2, 2}, K_mat_pushback::Function

    The reverse rule for function K_mat.
"""
function rrule(::typeof(K_mat), phi::cmps, psi::cmps)
    Id_phi = id(ℂ^get_chi(phi))
    Id_psi = id(ℂ^get_chi(psi))
    fwd = K_mat(phi, psi)

    function K_mat_pushback(f̄wd)
        f̄wd = permute(f̄wd, (3, 2), (1, 4))
        @tensor Q̄_phi[-1; -2] := conj(f̄wd[-2, 1, -1, 1])
        @tensor Q̄_psi[-1; -2] := f̄wd[1, -1, 1, -2]
        @tensor R̄_phi[-1, -2; -3] := conj(f̄wd[-3, 1, -1, 2]) * conj(psi.R'[2, 1, -2])
        @tensor R̄_psi[-1, -2; -3] := f̄wd[2, -1, 1, -3] * phi.R[1, -2, 2]
        p̄hi = cmps(Q̄_phi, R̄_phi)
        p̄si = cmps(Q̄_psi, R̄_psi)
        return NoTangent(), p̄hi, p̄si
    end
    return fwd, K_mat_pushback
end

"""
    log_ovlp(phi::cmps, psi::cmps, beta::Real) -> ComplexF64

    Caculate the log of overlap for two finite uniform cmps `phi` and `psi`.
    `beta` is the length of the uniform cmps.
"""
function log_ovlp(phi::cmps, psi::cmps, beta::Real; sym::Bool=false)
    t_trans = convert_to_array(K_mat(phi, psi))
    tot_dim = get_chi(phi) * get_chi(psi)
    t_trans = reshape(t_trans, (tot_dim, tot_dim))
    sym && (t_trans = Hermitian(0.5*(t_trans+t_trans')))
    w = eigvals(t_trans)
    return logsumexp(beta*w)
end

"""
    convergence_meaure(T::cmpo, ψ::cmps, beta::Real) -> Real
"""
function convergence_measure(T::cmpo, ψ::cmps, beta::Real)
    Tψ = act(T, ψ)
    return -2*log_ovlp(Tψ, ψ, beta; sym=true) + log_ovlp(Tψ, Tψ, beta; sym=true) + log_ovlp(ψ, ψ, beta; sym=true) |> real
end

"""
    finite_env(t::TensorMap{ComplexSpace}, L::Real)

    For a cMPS transfer matrix `t` (see `K_mat`), calculate the environment block for the finite length `L`.
    The input `t` matrix should look like
    ```
        -1 -->--  phi' -->-- -3
                   |
                   ^
                   |
        -2 --<--  psi  --<-- -4
    ```
    This function will calculate exp(t) / tr(exp(t)) by diagonalizing `t`.
    To make the indexing consistent with the arrows in the result, we can permute `res` as `permute(res, (2, 3), (4, 1))` .
"""
function finite_env(t::TensorMap{ComplexSpace}, L::Real)
    W, UR = eig(t)
    UL = inv(UR)

    Wvec = diag(W.data)
    Wvec = Wvec .- logsumexp(Wvec * L) / L #normalize
    expW = convert_to_tensormap(diagm(exp.(Wvec * L)), 1)

    return UR * expW * UL
end

"""
    rrule(::typeof(finite_env), t::TensorMap{ComplexSpace}, L::Real)

    Backward rule for `finite_env`.
    See https://math.stackexchange.com/a/3868894/488003 and https://doi.org/10.1006/aama.1995.1017 for the gradient of exp(t)
"""
function rrule(::typeof(finite_env), t::TensorMap{ComplexSpace}, L::Real)
    W, UR = eig(t)
    UL = inv(UR)

    Wvec = diag(W.data)
    Wvec = Wvec .- logsumexp(Wvec * L) / L #normalize
    expW = convert_to_tensormap(diagm(exp.(Wvec * L)), 1)
    fwd = UR * expW * UL

    function finite_env_pushback(f̄wd)
        function coeff(a::Number, b::Number)
            if a ≈ b
                return L*exp(a*L)
            else
                return (exp(a*L) - exp(b*L)) / (a - b)
            end
        end
        M = UR' * f̄wd * UL'
        M1 = similar(M)
        copyto!(M1.data, M.data .* coeff.(Wvec', conj.(Wvec)))
        t̄ = UL' * M1 * UR' - L * tr(f̄wd * fwd') * fwd'

        return NoTangent(), t̄, NoTangent()
    end
    return fwd, finite_env_pushback
end

#"""
#    observe(O::TensorMap{ComplexSpace, 2, 2}, env::TensorMap{ComplexSpace, 2, 2}) -> Number
#
#    Calculate the local observable `O` for a cMPS `ψ` with length `L`.
#    The indices for `O` and `env` should be ordered as:
#    ```
#        -4 -->--- --->-- -2
#                 |
#               O or env
#                 |
#        -1 --<--- ---<-- -3
#
#    ```
#"""
#function observe(O::TensorMap{ComplexSpace, 2, 2}, env::TensorMap{ComplexSpace, 2, 2})
#    #if L == Inf # thermodynamic limit
#    #    lop = transf_mat(ψ, ψ)
#    #    lopT = transf_mat_T(ψ, ψ)
#    #    _, vR = eigsolve(lop, TensorMap(rand, ComplexF64, ℂ^χ, ℂ^χ), 1, :LR)
#    #    vR = vR[1]
#    #    _, vL = eigsolve(lopT, TensorMap(rand, ComplexF64, ℂ^χ, ℂ^χ), 1, :LR)
#    #    vL = vL[1]
#
#    #    @tensor E = vL'[4, 1] * tensorE[1, 2, 3, 4] * vR[3, 2]
#    #    E /= tr(vL' * vR)
#    #end
#
#    # expectation value
#    @tensor E = env[4, 3, 1, 2] * O[1, 2, 4, 3]
#
#    return E
#end

"""
    kinetic(ψ::cmps) -> TensorMap{ComplexSpace, 2, 2}

    Construct the tensor for kinetic energy density `(dψ† / dx)(dψ / dx)` measurement.
"""
function kinetic(ψ::cmps)
    Q, R = get_matrices(ψ)
    @tensor O[-1, -2; -3, -4] := Q[-1, 1] * R[1, 3, -3] * Q'[2, -4] * R'[-2, 2, 3] +
                                       R[-1, 3, 1] * Q[1, -3] * R'[2, -4, 3] * Q'[-2, 2] -
                                       Q[-1, 1] * R[1, 3, -3] * R'[2, -4, 3] * Q'[-2, 2] -
                                       R[-1, 3, 1] * Q[1, -3] * Q'[2, -4] * R'[-2, 2, 3]
    return O
end

"""
    density(ψ::cmps) -> TensorMap{ComplexSpace, 2, 2}

    Construct the tensor for particle density `ψ†ψ` measurement.
"""
function particle_density(ψ::cmps)
    _, R = get_matrices(ψ)
    @tensor O[-1, -2; -3, -4] := R[-1, 1, -3] * R'[-2, -4, 1]
    return O
end

"""
    point_interaction(ψ::cmps) -> TensorMap{ComplexSpace, 2, 2}

    Construct the tensor for point interaction potential `ψ†ψ†ψψ` measurement.
"""
function point_interaction(ψ::cmps)
    _, R = get_matrices(ψ)
    @tensor O[-1, -2; -3, -4] := R[-1, 3, 1] * R[1, 4, -3] * R'[2, -4, 3] * R'[-2, 2, 4]
    return O
end

"""
    build_I_R(ψ::cmps) -> TensorMap{ComplexSpace, 2, 1}
"""
function build_I_R(ψ::cmps)
    χ, d = get_chi(ψ), get_d(ψ)
    I_R = zeros(ComplexF64, (χ, d, χ))
    for ix in 1:d
        I_R[:, ix, :] = Matrix{ComplexF64}(I, (χ, χ))
    end
    I_R = convert_to_tensormap(I_R, 2)
end

@non_differentiable build_I_R(ψ::cmps)

"""
    pairing(ψ::cmps) -> TensorMap{ComplexSpace, 2, 2}
"""
function pairing(ψ::cmps)
    _, R = get_matrices(ψ)
    I_R = build_I_R(ψ)
    @tensor O[-1, -2; -3, -4] := R[-1, 3, 1] * R[1, 4, -3] * I_R'[2, -4, 3] * I_R'[-2, 2, 4] +
                                I_R[-1, 3, 1] * I_R[1, 4, -3] * R'[2, -4, 3] * R'[-2, 2, 4]
    return O
end

"""
    optimize_conv_meas(T::cmpo, ψ::cmps, beta::Real)

    Optimize the convergence measure by gradient optimization.
"""
function optimize_conv_meas(T::cmpo, ψ::cmps, beta::Real, Niter::Integer)
    ψarr = convert_to_array(ψ)

    function _f(x::Array{ComplexF64, 3})
        ϕ = convert_to_cmps(x)
        return convergence_measure(T, ϕ, beta)
    end
    function _g!(gx::Array{ComplexF64, 3}, x::Array{ComplexF64, 3})
        gx .= gradient(_f, x)[1]
    end

    res = optimize(_f, _g!, ψarr, Optim.Options(show_trace=false, iterations=Niter))
    return convert_to_cmps(Optim.minimizer(res))
end

"""
    normalize(psi::cmps, beta::Real)

    Normalize the given cmps `psi`.
"""
function normalize(psi::cmps, beta::Real; sym=true)
    log_psi_norm = log_ovlp(psi, psi, beta; sym=sym)
    id_psi = id(ℂ^(get_chi(psi)))
    #@show log_psi_norm
    Q = psi.Q - log_psi_norm / beta / 2  * id_psi
    return cmps(Q, psi.R)
end

"""
    compress(psi::cmps, chi::Integer, beta::Real) -> Bool, Real, cmps

    Compress the cmps `psi` to the target bond dimension `chi`.
    Only implemented the case for symmetric cmps.
"""
function compress(ψ::cmps, chi::Integer, beta::Real; Niter::Integer=100, tol::Real=1e-12, init=nothing)
    chi_ψ = get_chi(ψ)
    if chi_ψ <= chi
        @warn "target bond dimension $(chi) larger than the current one $(chi_ψ)."
        return true, 0.0, ψ
    end

    K = K_mat(ψ, ψ)
    K = 0.5*(K + K')
    W, U = eigh(K)

    log_norm_ψ = logsumexp(diag(W.data)*beta) / beta

    if init === nothing
        Wmax = findmax(W.data)[1]
        expK = U * exp((W - Wmax * id(ℂ^chi_ψ^2)) * beta / 2) * U'
        expK = permute(expK, (2, 3), (1, 4))

        _, _, Ua = tsvd(expK, (4, 1, 2), (3,))
        #Ub, Sb, _ = tsvd(expK, (1,), (2, 3, 4))
        #Uc, Sc, _ = tsvd(expK, (2,), (3, 4, 1))
        #_, Sd, Ud = tsvd(expK, (1, 2, 3), (4,))
        Pa_arr = Ua.data[1:chi, :]
        Pa = TensorMap(Pa_arr, ℂ^chi, ℂ^chi_ψ)
        Q = Pa * ψ.Q * Pa'
        @tensor R[-1, -2; -3] := Pa[-1, 1] * ψ.R[1, -2, 2] * Pa'[2, -3]
        ϕ = cmps(Q, R)
    else
        ϕ = init
    end

    log_fidel = 2*log_ovlp(ϕ, ψ, beta) - log_ovlp(ϕ, ϕ, beta) - log_norm_ψ * beta |> real

    if log_fidel < -tol || init !== nothing
        function _f(x::Array{ComplexF64, 3})
            ϕx = convert_to_cmps(x)
            return - log_ovlp(ϕx, ψ, beta; sym=true) - log_ovlp(ψ, ϕx, beta; sym=true) + log_ovlp(ϕx, ϕx, beta; sym=true) |> real
        end
        function _g!(gx::Array{ComplexF64, 3}, x::Array{ComplexF64, 3})
            gx .= gradient(_f, x)[1]
        end

        res = optimize(_f, _g!, convert_to_array(ϕ), LBFGS(), Optim.Options(show_trace=false, iterations=Niter, f_abstol=tol))
        ϕ = convert_to_cmps(Optim.minimizer(res))
    end
    ϕ = normalize(ϕ, beta)

    log_fidel = 2*log_ovlp(ϕ, ψ, beta) - log_norm_ψ * beta |> real
    status = (log_fidel > -tol)
    return status, log_fidel, ϕ
end

"""
    truncation_measure(psi::cmps, psi1::cmps, beta::Real)

    Check the quality of periodic cMPS truncation from the following aspects:
    - fidelity
    - entanglement spectra
    - TODO. correlation function
"""
function truncation_check(psi::cmps, psi1::cmps, beta::Real)
    # normalize psi and psi_truncated
    chi, chi1 = get_chi(psi), get_chi(psi1)
    Q = psi.Q - log_ovlp(psi, psi, beta) / beta / 2 * id(ℂ^chi)
    psi = cmps(Q, psi.R)
    Q = psi1.Q - log_ovlp(psi1, psi1, beta) / beta / 2 * id(ℂ^chi1)
    psi1 = cmps(Q, psi1.R)

    # fidelity
    fidelity = 2*real(log_ovlp(psi, psi1, beta))

    # entanglement spectra
    K = K_mat(psi, psi)
    W, U = eigh(K)
    expK = U * exp(W*beta/2) * U'
    expK = permute(expK, (2, 3), (1, 4))
    _, SK, _ = tsvd(expK, (4, 1), (2, 3))
    entangle_spect = diag(SK.data) / SK.data[1]

    K1 = K_mat(psi, psi1)
    W1, U1 = eigh(K1)
    expK1 = U1 * exp(W1*beta/2) * U1'
    expK1 = permute(expK1, (2, 3), (1, 4))
    _, SK1, _ = tsvd(expK1, (4, 1), (2, 3))
    entangle_spect1 = diag(SK1.data) / SK1.data[1]

    return fidelity, entangle_spect[1:chi1^2], entangle_spect1

end

"""
    gauge_fixing_proj(ψ::cmps, L::Real; gauge::Symbol=:periodic) -> proj::TensorMap

    Fix the gauge for a tangent vector tensor. `gauge` can be chosen between `:periodic` and `:left` .
    This function generates a projector `proj` that can be used to
    - parametrize a fixed-gauge tangent vector
    - eleminate the gauge degrees of freedom in an existing tangent vector
"""
function gauge_fixing_proj(ψ::cmps, L::Real; gauge::Symbol=:periodic)
    # tensor `A` is connected to impurity tensor in the tangent vector
    χ = get_chi(ψ)
    A = cmps(id(ℂ^χ), copy(ψ.R))
    A = convert_to_tensormap(A)

    if gauge == :periodic
        # density matrix ρ
        K = K_mat(ψ, ψ)
        ρ = finite_env(K, L)
        ρ = permute(ρ, (2, 3), (4, 1))

        # construct the projector
        @tensor Aρ[-1, ; -2, -3] := ρ[-1, 1, -2, 2] * A'[2, 1, -3]
        proj = rightnull(Aρ)'

    elseif gauge == :left
        # left dominant eigenvector
        lopT = transf_mat_T(ψ, ψ)
        Vl = eigsolve(lopT, TensorMap(rand, ComplexF64, ℂ^χ, ℂ^χ), 1, :LR)[2][1]

        # projector
        @tensor Aρ[-1, ; -2, -3] := Vl'[1, -2] * A'[-1, 1, -3]
        proj = rightnull(Aρ)'
    end
    return proj
end

"""
    tangent_map(ψ::cmps, L::Real; gauge::Symbol=:periodic) -> Function

    tangent map. act on the parameter space obtained after gauge elimination.
"""
function tangent_map(ψ::cmps, L::Real, p::Real=0; gauge::Symbol=:periodic)

    χ, d = get_chi(ψ), get_d(ψ)

    # diagonalize K matrix, normalize W according to length
    K = K_mat(ψ, ψ)
    W, UR = eig(K)
    UL = inv(UR)

    # gauge fixing projector
    proj = gauge_fixing_proj(ψ, L; gauge=gauge)

    # calculate coefficient matrix
    Wvec = diag(W.data)
    normψ = logsumexp(L .* Wvec)
    Wvec .-= normψ / L
    @tullio coeffW[ix, iy] := theta2(L, Wvec[ix], Wvec[iy] - p*im)
    #coeffW = similar(W)
    #function coeff(a::Number, b::Number)
    #    # when using coeff.(avec, bvec'), avec is column vector, bvec is row vector
    #    # in this way the index ordering is consistent
    #    if a ≈ b'
    #        return L*exp(a*L)
    #    else
    #        return (exp(a*L) - exp(b'*L)) / (a - b')
    #    end
    #end
    #copyto!(coeffW.data, coeff.(Wvec, (Wvec .- p*im)'))

    # A tensor that is connected to G
    A = cmps(id(ℂ^get_chi(ψ)), copy(ψ.R))
    A = convert_to_tensormap(A)

    D = id(ℂ^(d+1))
    D.data[1] = 0

    idW = id(ℂ^(χ^2))

    # calculate the map
    function f(V::TensorMap{ComplexSpace, 1, 1})
        # crossing term
        @tensor M0[-1, -2; -3, -4] := V[1, -4] * proj[-2, 2, 1] * A'[-3, -1, 2]
        M = UL * M0 * UR
        #M = elem_mult(M, coeffW)
        #M = elem_mult_f1(M, (ix, iy) -> theta2(L, Wvec[ix], Wvec[iy] - p*im))
        @tullio M.data[ix, iy] = M.data[ix, iy] * coeffW[ix, iy]
        M = UR * M * UL
        M = permute(M, (2, 3), (4, 1))
        @tensor Vc[-1; -2] := M[1, 3, 2, -2] * A[2, 4, 1] * proj'[-1, 3, 4]
        # diagonal term
        ρ = UR * exp(W*L - idW*normψ) * UL
        ρ = permute(ρ, (2, 3), (4, 1))
        @tensor Vd[-1; -2] := ρ[3, 5, 2, -2] * proj[2, 4, 1] * V[1, 3] * D[6, 4] * proj'[-1, 5, 6]
        return L*(Vc + Vd)
    end

    return f
end

mutable struct preconditioner
    map::Function
    invmap::Function
    proj::AbstractTensorMap
end

"""
    preconditioner(ψ::cmps, L::Real)

    construct the preconditioner for cmps `ψ` of length `L`.
"""
function preconditioner(ψ::cmps, ψgrad::TensorMap{ComplexSpace, 2, 1}, L::Real; gauge::Symbol=:periodic)
    χ, d = get_chi(ψ), get_d(ψ)

    map = tangent_map(ψ, L)
    proj = gauge_fixing_proj(ψ, L; gauge=gauge)

    V = TensorMap(rand, ComplexF64, ℂ^(χ*d), ℂ^χ)
    δ = 100*norm(proj' * ψgrad)^2
    @show norm(ψgrad), sqrt(δ)/10, δ
    #δ = norm(map(V)) / norm(V) * tol

    # sometimes KrylovKit raises error " LoadError: operator does not appear to be hermitian" here
    # it seems that δ cannot be too small.
    try
        global Ws, Vs
        Ws, Vs = eigsolve(map, V, d*χ^2; tol=0.1*δ, krylovdim=d*χ^2, ishermitian=true)
    catch e
        @show e
        Ws, Vs = eigsolve(map, V, d*χ^2; tol=0.1*δ, krylovdim=d*χ^2)
    end

    function invmap(V::TensorMap{ComplexSpace, 1, 1})
        f = (Wx, Vx) -> Vx * tr(Vx' * V) * (1 / (Wx + δ) - 1 / δ)
        return sum(f.(Ws, Vs)) + V / δ
    end

    return preconditioner(map, invmap, proj)
end

"""
    precond_grad(P::preconditioner, ψgrad_arr::Array{ComplexF64, 3}) -> Array{ComplexF64, 3}

    Act the preconditioner on the gradient.
"""
function precond_grad(P::preconditioner, ψgrad_arr::Array{ComplexF64, 3})

    χ = size(ψgrad_arr)[1]
    d = size(ψgrad_arr)[2] - 1

    # modulate tangent map: P.map -> P.map + δ * id
    #V = TensorMap(rand, ComplexF64, ℂ^(χ*d), ℂ^χ)
    #δ = norm(P.map(V)) / norm(V) * 1e-4#; @show δ
    #f = (V -> P.map(V) + δ * V )

    # project out gauge freedoms
    G = convert_to_tensormap(ψgrad_arr, 2)
    VG = P.proj' * G

    # inverse of the metric
    #VG1 = linsolve(f, VG, VG)[1]
    #Wfs, Vfs = eigsolve(f, VG, d*χ^2; krylovdim=d*χ^2, ishermitian=true)
    #invf = (Wf, Vf) -> Vf * tr(Vf' * VG) / Wf
    G1 = P.proj * P.invmap(VG)#sum(invf.(Wfs, Vfs))
    return convert_to_array(G1)
end

"""
    dot(x::Array{ComplexF64, 3}, P::preconditioner, y::Array{ComplexF64, 3})

    util function for Optim.jl preconditioner.
    calculates the inner product <x, P*y>
"""
function dot(x::Array{ComplexF64, 3}, P::preconditioner, y::Array{ComplexF64, 3})
    x_tnmap = convert_to_tensormap(x, 2)
    y_tnmap = convert_to_tensormap(y, 2)

    Py = P.proj * P.map(P.proj' * y_tnmap)
    return tr(x_tnmap' * Py)
end

"""
    ldiv!(x::Array{ComplexF64, 3}, P::preconditioner, y::Array{ComplexF64, 3})

    util function for Optim.jl preconditioner.
    copy inv(P)*y to x.
"""
function ldiv!(x::Array{ComplexF64, 3}, P::preconditioner, y::Array{ComplexF64, 3})
    copyto!(x, precond_grad(P, y))
end

function elem_mult_f1(a::TensorMap{ComplexSpace, 1, 1}, f::Function)
    c = zero(a)
    @tullio c.data[ix, iy] = a.data[ix, iy] * f(ix, iy)
    return c
end

function elem_mult_f2(a::TensorMap{ComplexSpace, 1, 1}, b::TensorMap{ComplexSpace, 1, 1}, f::Function)
    c = zero(a)
    @tullio c.data[ix, iy] = a.data[ix, im] * b.data[im, iy] * f(ix, im, iy)
    return c
end

function theta2(L::Real, a::Number, b::Number)
    if a ≈ b
        return L*exp(a*L)
    else
        return (exp(a*L) - exp(b*L)) / (a - b)
    end
end
function theta3(L::Real, a::Number, b::Number, c::Number)
    if a ≈ b ≈ c
        return 0.5 * exp(L*b) * L^2
    elseif a ≈ b && !(b ≈ c)
        return -1 * (exp(L*b) - exp(L*c) - L*exp(L*b)*(b-c) ) / (b - c)^2
    elseif b ≈ c && !(c ≈ a)
        return -1 * (exp(L*c) - exp(L*a) - L*exp(L*c)*(c-a) ) / (c - a)^2
    elseif c ≈ a && !(a ≈ b)
        return -1 * (exp(L*a) - exp(L*b) - L*exp(L*a)*(a-b) ) / (a - b)^2
    else
        return (a * (exp(L*b) - exp(L*c)) + b * (exp(L*c) - exp(L*a)) + c * (exp(L*a) - exp(L*b))) / ((a-b)*(b-c)*(c-a))
    end
end

function lieb_liniger_h_tangent_map(ψ::cmps, p::Real, q::Real, c::Real, L::Real, μ::Real; k0::Real=1, gauge::Symbol=:periodic)

    χ, d = get_chi(ψ), get_d(ψ)

    # diagonalize K matrix, normalize W according to length
    K = K_mat(ψ, ψ)
    W, UR = eig(K)
    UL = inv(UR)
    Wvec = diag(W.data)
    normψ = logsumexp(L .* Wvec)
    Wvec .-= normψ / L
    copyto!(W.data, diagm(Wvec))

    # gauge fixing projector
    P = gauge_fixing_proj(ψ, L; gauge=gauge)

    A = cmps(id(ℂ^χ), copy(ψ.R))
    A = convert_to_tensormap(A)
    DR = id(ℂ^(d+1))
    DR.data[1] = 0
    DQ = zero(DR)
    DQ.data[1] = 1
    DQRRQ = zeros(d+1, d+1, d+1, d+1)
    DRQQR = zeros(d+1, d+1, d+1, d+1)
    for ix in 2:d+1
        DQRRQ[1, ix, ix, 1] = 1
        DRQQR[ix, 1, 1, ix] = 1
    end
    DQRRQ = TensorMap(DQRRQ, ℂ^(d+1)*ℂ^(d+1), ℂ^(d+1)*ℂ^(d+1))
    DRQQR = TensorMap(DRQQR, ℂ^(d+1)*ℂ^(d+1), ℂ^(d+1)*ℂ^(d+1))
    @tensor QR_comm[-1, -2; -3, -4] := DQ[-1, -3] * DR[-2, -4] + DR[-1, -3] * DQ[-2, -4] -
                                       DQRRQ[-1, -2, -3, -4] - DRQQR[-1, -2, -3, -4]
    @tensor QR_comm2[-1, -2; -3, -4] := DQ[-1, -3] * DR[-2, -4] - DRQQR[-1, -2, -3, -4]#DQR[-1, -3] * DRQ[-2, -4]
    @tensor QR_comm3[-1, -2; -3, -4] := DQ[-1, -3] * DR[-2, -4] - DQRRQ[-1, -2, -3, -4]#DQR[-1, -3] * DRQ[-2, -4]

    ψtn = convert_to_tensormap(ψ)

    # coefficient arrays constructed from Wvec.
    # for x != y != x', Mx * My * coeffs
    @tullio coeff_W3_xy[ix, iz, iy] := theta3(L, Wvec[ix], Wvec[iz] - im*p, Wvec[iy] - im*(p+q))
    # for x != y != x', My * Mx * coeffs
    @tullio coeff_W3_yx[ix, iz, iy] := theta3(L, Wvec[ix] + im*(p+q), Wvec[iz] + im*p, Wvec[iy])
    # for y != x == x'
    @tullio coeff_W2_q[ix, iy] := theta2(L, Wvec[ix], Wvec[iy] - im*q)
    # for x == y != x'
    @tullio coeff_W2_p1[ix, iy] := theta2(L, Wvec[ix] + im*(p+q), Wvec[iy])
    # for x != y == x'
    @tullio coeff_W2_p[ix, iy] := theta2(L, Wvec[ix] + im*p, Wvec[iy])

    function f(V::TensorMap{ComplexSpace, 1, 1})
        @tensor Mx[-1, -2; -3, -4] := P[-2, 2, 1] * V[1, -4] * A'[-3, -1, 2]
        Mx = UL * Mx * UR
        Mtmp = similar(Mx)

        ## y != x and y != x'
        @tensor My[-1, -2; -3, -4] := -μ * ψtn'[-3, -1, 1] * DR[1, 2] * ψtn[-2, 2, -4] +
                                      k0 * ψtn'[1, -1, 3] * ψtn'[-3, 1, 4] * QR_comm[3, 4, 5, 6] * ψtn[-2, 5, 2] * ψtn[2, 6, -4] +
                                      c * ψtn'[1, -1, 3] * ψtn'[-3, 1, 4] * DR[3, 5] * DR[4, 6] * ψtn[-2, 5, 2] * ψtn[2, 6, -4]
        My = UL * My * UR
        #M = elem_mult_f2(Mx, My, (ix, iz, iy) -> theta3(L, Wvec[ix], Wvec[iz] - im*p, Wvec[iy] - im*(p+q))) +
        #    elem_mult_f2(My, Mx, (ix, iz, iy) -> theta3(L, Wvec[ix] + im*(p+q), Wvec[iz] + im*p, Wvec[iy]))

        #### y != x == x'
        @tullio Mtmp.data[ix, iy] = My.data[ix, iy] * coeff_W2_q[ix, iy]
        M = permute(UR * Mtmp * UL, (2, 3), (4, 1))
        @tensor Vout[-1; -2] := M[2, 4, 3, -2] * P[3, 6, 1] * V[1, 2] * DR[5, 6] * P'[-1, 4, 5]

        #### y != x != x'
        @tullio Mtmp.data[ix, iy] = Mx.data[ix, iz] * My.data[iz, iy] * coeff_W3_xy[ix, iz, iy] +
                                    My.data[ix, iz] * Mx.data[iz, iy] * coeff_W3_yx[ix, iz, iy]

        ## x == y != x'
        @tensor My[-1, -2; -3, -4] := -μ * ψtn'[-3, -1, 2] * DR[2, 3] * P[-2, 3, 1] * V[1, -4] +
                                      k0 * ψtn'[3, -1, 4] * ψtn'[-3, 3, 5] * QR_comm[4, 5, 6, 7] * P[-2, 6, 1] * V[1, 2] * ψtn[2, 7, -4] +
                                      k0 * ψtn'[3, -1, 4] * ψtn'[-3, 3, 5] * QR_comm[4, 5, 6, 7] * ψtn[-2, 6, 2] * P[2, 7, 1] * V[1, -4] +
                                      k0 * ψtn'[3, -1, 4] * ψtn'[-3, 3, 5] * DQ[4, 6] * DR[5, 7] * im*p*A[-2, 6, 2] * P[2, 7, 1] * V[1, -4] -
                                      k0 * ψtn'[3, -1, 4] * ψtn'[-3, 3, 5] * DRQQR[4, 5, 6, 7] * im*p*A[-2, 6, 2] * P[2, 7, 1] * V[1, -4] +
                                      c * ψtn'[1, -1, 4] * ψtn'[-3, 1, 5] * DR[4, 6] * DR[5, 7] * ψtn[-2, 6, 3] * P[3, 7, 2] * V[2, -4] +
                                      c * ψtn'[1, -1, 4] * ψtn'[-3, 1, 5] * DR[4, 6] * DR[5, 7] * P[-2, 6, 2] * V[2, 3] * ψtn[3, 7, -4]
        My = UL * My * UR
        #M += elem_mult_f1(My, (ix, iy) -> theta2(L, Wvec[ix] + im*(p+q), Wvec[iy]))
        @tullio Mtmp.data[ix, iy] += My.data[ix, iy] * coeff_W2_p1[ix, iy]
        M = permute(UR * Mtmp * UL, (2, 3), (4, 1))
        @tensor Vout[-1; -2] = Vout[-1, -2] + M[1, 3, 2, -2] * A[2, 4, 1] * P'[-1, 3, 4]

        ## x != y == x'
        #M = elem_mult_f1(Mx, (ix, iy) -> theta2(L, Wvec[ix] + im*p, Wvec[iy]))
        @tullio Mtmp.data[ix, iy] = Mx.data[ix, iy] * coeff_W2_p[ix, iy]
        M = permute(UR * Mtmp * UL, (2, 3), (4, 1))
        @tensor Vout[-1; -2] = Vout[-1, -2] - μ * M[1, 3, 2, -2] * ψtn[2, 5, 1] * DR[4, 5] * P'[-1, 3, 4] +
                                              k0 * M[1, 4, 2, 5] * P'[-1, 4, 6] * ψtn'[5, -2, 7] * QR_comm[6, 7, 8, 9] * ψtn[2, 8, 3] * ψtn[3, 9, 1] +
                                              k0 * M[1, 4, 2, -2] * ψtn'[5, 4, 6] * P'[-1, 5, 7] * QR_comm[6, 7, 8, 9] * ψtn[2, 8, 3] * ψtn[3, 9, 1] +
                                              k0 * M[1, 4, 2, -2] * (-im*(p+q))*A'[5, 4, 6] * P'[-1, 5, 7] * DQ[6, 8] * DR[7, 9] * ψtn[2, 8, 3] * ψtn[3, 9, 1] -
                                              k0 * M[1, 4, 2, -2] * (-im*(p+q))*A'[5, 4, 6] * P'[-1, 5, 7] * DQRRQ[6, 7, 8, 9] * ψtn[2, 8, 3] * ψtn[3, 9, 1] +
                                              c * M[1, 4, 2, -2] * ψtn'[5, 4, 6] * P'[-1, 5, 7] * DR[6, 8] * DR[7, 9] * ψtn[2, 8, 3] * ψtn[3, 9, 1] +
                                              c * M[1, 5, 2, 4] * P'[-1, 5, 6] * ψtn'[4, -2, 7] * DR[6, 8] * DR[7, 9] * ψtn[2, 8, 3] * ψtn[3, 9, 1]

        ## x == y == x'
        M = permute(UR * exp(W*L) * UL, (2, 3), (4, 1))
        @tensor Vout[-1; -2] = Vout[-1, -2] - μ * M[2, 4, 3, -2] * P[3, 6, 1] * V[1, 2] * DR[5, 6] * P'[-1, 4, 5] +
                                            + k0 * M[1, 5, 2, 6] * P'[-1, 5, 7] * ψtn'[6, -2, 8] * QR_comm[7, 8, 9, 10] * P[2, 9, 3] * V[3, 4] * ψtn[4, 10, 1] +
                                            + k0 * M[1, 5, 2, 6] * P'[-1, 5, 7] * ψtn'[6, -2, 8] * QR_comm[7, 8, 9, 10] * ψtn[2, 9, 3] * P[3, 10, 4] * V[4, 1] +
                                            + k0 * M[1, 5, 2, -2] * ψtn'[6, 5, 7] * P'[-1, 6, 8] * QR_comm[7, 8, 9, 10] * ψtn[2, 9, 3] * P[3, 10, 4] * V[4, 1] +
                                            + k0 * M[1, 5, 2, -2] * ψtn'[6, 5, 7] * P'[-1, 6, 8] * QR_comm[7, 8, 9, 10] * P[2, 9, 3] * V[3, 4] * ψtn[4, 10, 1] +
                                            + k0 * M[1, 6, 2, 5] * P'[-1, 6, 7] * ψtn'[5, -2, 8] * QR_comm2[7, 8, 9, 10] * (im*p)*A[2, 9, 3] * P[3, 10, 4] * V[4, 1] +
                                            + k0 * M[1, 5, 2, -2] * ψtn'[6, 5, 7] * P'[-1, 6, 8] * QR_comm2[7, 8, 9, 10] * (im*p)*A[2, 9, 3] * P[3, 10, 4] * V[4, 1] +
                                            + k0 * M[1, 5, 2, -2] * (-im*(p+q))*A'[6, 5, 7] * P'[-1, 6, 8] * QR_comm3[7, 8, 9, 10] * ψtn[2, 9, 3] * P[3, 10, 4] * V[4, 1] +
                                            + k0 * M[1, 5, 2, -2] * (-im*(p+q))*A'[6, 5, 7] * P'[-1, 6, 8] * QR_comm3[7, 8, 9, 10] * P[2, 9, 3] * V[3, 4] * ψtn[4, 10, 1] +
                                            + k0 * M[1, 5, 2, -2] * (-im*(p+q))*A'[6, 5, 7] * P'[-1, 6, 8] * DQ[7, 9] * DR[8, 10] * (im*p)*A[2, 9, 3] * P[3, 10, 4] * V[4, 1] +
                                            + c * M[4, 5, 3, -2] * ψtn'[6, 5, 7] * P'[-1, 6, 8] * DR[7, 9] * DR[8, 10] * ψtn[3, 9, 2] * P[2, 10, 1] * V[1, 4] +
                                            + c * M[4, 6, 3, 5] * P'[-1, 6, 7] * ψtn'[5, -2, 8] * DR[7, 9] * DR[8, 10] * ψtn[3, 9, 2] * P[2, 10, 1] * V[1, 4] +
                                            + c * M[3, 5, 4, -2] * ψtn'[6, 5, 7] * P'[-1, 6, 8] * DR[7, 9] * DR[8, 10] * P[4, 9, 1] * V[1, 2] * ψtn[2, 10, 3] +
                                            + c * M[3, 6, 4, 5] * P'[-1, 6, 7] * ψtn'[5, -2, 8] * DR[7, 9] * DR[8, 10] * P[4, 9, 1] * V[1, 2] * ψtn[2, 10, 3]
        return L * Vout
    end
    return f
end