module tmsv

using Nemo
using LinearAlgebra

using ..tools


"""
    TMSV

Parameters for a Two-Mode Squeezed Vacuum (TMSV) entanglement source.

The TMSV is the simplest Gaussian entangled state. Both modes are produced by a
single parametric interaction, and the entanglement is characterized by the mean photon number μ.

# Fields
- `mean_photon::Real`         : Mean photon number per mode (default `1e-2`)
- `detection_efficiency::Real`: Detector efficiency, ∈ [0, 1] (default `1.0`)
"""
Base.@kwdef mutable struct TMSV
    mean_photon::Real = 1e-2
    detection_efficiency::Real = 1.0
end

# Global canonical position and momentum variables
const mds = 2 # Number of modes for our system

_qai = ["qa$i" for i in 1:mds]
_pai = ["pa$i" for i in 1:mds]
_qbi = ["qb$i" for i in 1:mds]
_pbi = ["pb$i" for i in 1:mds]
all_qps = hcat(_qai, _pai, _qbi, _pbi)
CC = ComplexField()
i = onei(CC) # Imaginary unit in CC ring
R, generators = polynomial_ring(CC, all_qps)
(qai, pai, qbi, pbi) = (generators[:,i] for i in 1:4)

# Define the alpha and beta vectors
α = (qai + i .* pai) / sqrt(2)
β = (qbi - i .* pbi) / sqrt(2)

"""
    covariance_matrix(μ::Real)

Construct the covariance matrix for a TMSV state.

# Parameters
- μ : The mean photon number of the TMSV state

# Returns
The covariance matrix for the TMSV state, in the qpqp ordering
"""
covariance_matrix(μ::Real) = [
    0.5 + μ        0               sqrt(μ*(μ+1))  0;
    0              0.5 + μ         0              -sqrt(μ*(μ+1));
    sqrt(μ*(μ+1))  0               0.5 + μ        0;
    0              -sqrt(μ*(μ+1))  0              0.5 + μ;
]
covariance_matrix(tmsv::TMSV) = covariance_matrix(tmsv.mean_photon)

"""
    loss_matrix_pgen(ηᵈ::Real)

Construct the loss contribution to the A-matrix for TMSV probability-of-success calculations.

# Parameters
- ηᵈ: Detection efficiency, ∈ [0, 1]

# Returns
8×8 `ComplexF64` loss matrix for use in `A = k_function_matrix(cov) + loss_matrix_pgen(ηᵈ)`.
"""
function loss_matrix_pgen(ηᵈ::Real)
    G = zeros(ComplexF64, 8, 8)

    for i in 1:2
        G[i,     i+2*mds] = ηᵈ - 1
        G[i,     i+3*mds] = -im*(ηᵈ - 1)
        G[i+mds, i+2*mds] = im*(ηᵈ - 1)
        G[i+mds, i+3*mds] = ηᵈ - 1
    end

    return (G + transpose(G) + I) / 2
end
loss_matrix_pgen(tmsv::TMSV) = loss_matrix_pgen(tmsv.detection_efficiency)

"""
    moment_vector(n::Int)

Construct the symbolic moment polynomial for order `n` of the TMSV coincidence measurement.

Returns the Nemo polynomial `(α₁α₂)ⁿ/n! · (β₁β₂)ⁿ/n!` in the global phase-space variables,
representing the n-photon coincidence moment. Evaluated at `n=1` for `probability_success`.

# Parameters
- n: Photon number order

# Returns
Nemo multivariate polynomial over `ComplexField`.
"""
function moment_vector(n::Int)
    (α[1]*α[2])^n / factorial(n) * (β[1]*β[2])^n / factorial(n)
end

"""
    probability_success(μ::Real, ηᵈ::Real)

Calculate the probability of photon-photon state generation with the given parameters.

# Parameters
- μ : The mean photon number of the TMSV state
- ηᵈ : Detection efficiency

# Returns
Probability of successful photon-photon state generation
"""
function probability_success(μ::Real, ηᵈ::Real)
    # Compute covariance matrix and reorder qpqp → qqpp
    cov = reorder(covariance_matrix(μ))

    A = k_function_matrix(cov) + loss_matrix_pgen(ηᵈ)
    Ainv = inv(A)
    Γ = cov + (1/2)*I
    detΓ = det(Γ)

    N1 = ηᵈ^2
    D1 = sqrt(det(A))
    D2 = detΓ^(1/4)
    D3 = conj(detΓ)^(1/4)
    Coef = N1/(D1*D2*D3)

    C = moment_vector(1)

    return real(Coef * W(C, Ainv))
end
probability_success(tmsv::TMSV) = probability_success(tmsv.mean_photon, tmsv.detection_efficiency)

end # module
