module PLUV

import SpecialFunctions: erf

######################################################################################################
######################################################################################################
######################################################################################################
###############        POPC/POPG 95/5 mol/mol LUVs                                     ###############
###############        Reconstitution buffer: TRIS 20 mM, EDTA 2 mM                    ###############
###############        Buffer used for PLUV exp. in DESY in 2017:                      ###############
######################################################################################################
######################################################################################################
######################################################################################################

const parameter_names = [
  :norm, :nv, :rm, :z, :n_tr, :d_tr, :s_tr, :d_chol, :s_chol, :d_pcn, :s_pcn, :d_cg, :s_cg,
  :a_l, :s_d_c, :s_ch2, :d_ch, :s_ch, :s_ch3, :r_pcn, :r_cg, :r12, :r32, :t, :v_bw, :con
]

############### GLOBAL VARIABLES ###############

# POPG molar ratio
const CONST_x_PG = 0.05

#    ##############    POPC and POPG    ###############
#    1-palmitoyl-2-oleoyl-sn-glycero-3-phosphocholine
#    1-palmitoyl-2-oleoyl-sn-glycero-3-phosphoglycerol

# number of chain groups
const CONST_n_CH = 2
const CONST_n_CH2 = 28
const CONST_n_CH3 = 2

#X-ray scattering length chain groups (nm)
const CONST_b_CH = 1.97256E-05
const CONST_b_CH2 = 2.25435E-05
const CONST_b_CH3 = 2.53615E-05

### POPC
# Lipid-head volume
const CONST_V_HL_PC = 0.331 # 0.320
# X-ray scattering length of head groups (nm)
const CONST_b_PC = 2.73340E-04
const CONST_b_CG = 1.88802E-04
const CONST_b_PCN = 1.97256E-04
const CONST_b_Chol = 7.60844E-05
# Lipid-volume temperature-dependencies a0 + a1*T (nm^3)
const CONST_a0_V_POPC = 1.22810311835285
const CONST_a1_V_POPC = 0.000934915795086395

### POPG
# Lipid-head volume
const CONST_V_HL_PG = 0.289
# X-ray Scattering length of head groups (nm)
const CONST_b_PG = 2.47979E-04
const CONST_b_PG1 = 1.32443E-04
const CONST_b_PG2 = 1.15536E-04

# Lipid-volume temperature-dependencies a0 + a1*T (nm^3)
const CONST_a0_V_POPG = 1.17881068602663
const CONST_a1_V_POPG = 0.00108364914520327


###############    Other variables    ###############
### Water
# V_PW    = 24.5e-3 #(nm^3) Perkins 2001 (Hydration shell in proteins)
# polynome coefficient for T-dependency of bulk-water-molecule volume (V_HW)
# Units in degree Celsius
const CONST_p0_VW = 0.0299218
const CONST_p1_VW = -2.25941e-06
const CONST_p2_VW = 2.5675e-07
const CONST_p3_VW = -1.69661e-09
const CONST_p4_VW = 6.52029e-12
# polynome coefficient for T-dependency of bulk-water molar concentration (Cw)
#Units in degree Celsius
const CONST_p0_Cw = 55.5052
const CONST_p1_Cw = 0.00131894
const CONST_p2_Cw = -0.000334396
const CONST_p3_Cw = 9.10861e-07

const CONST_b_HW = 2.8179E-05
const CONST_d_shl = 0.31 # (nm) Perkins 2001 (Hydration shell in proteins)

# Composition of the Reconstitution buffer (M)
const CONST_ctris = 0.02
const CONST_cEDTA = 0.002

### Extra molecules

# TRIS buffer
const CONST_b_tris = 1.860E-04
const CONST_V_tris = 0.15147 # (nm^3)
# EDTA
const CONST_b_EDTA = 4.340E-04
const CONST_V_EDTA = 0.56430 # (nm^3)


##################################################################################################################
##################################################################################################################

#########################################################
#@njit(parallel=True)
function PDF_normal(x, mu, sig)
  return exp(-(x - mu)^2 / (2 * sig^2)) / (sig * sqrt(2 * pi))
end

#########################################################
#@njit(parallel=True)
function PDF_lognormal(x, mu_x, sig_x)
  # https://en.wikipedia.org/wiki/Log-normal_distribution
  mu = log(mu_x^2 / sqrt(mu_x^2 + sig_x^2))
  sig = sqrt(log(1 + sig_x^2 / mu_x^2))
  return exp(-(log(x) - mu)^2 / (2 * sig^2)) / (x * sig * sqrt(2 * pi))
end

#########################################################
function lipid_volume(T)
  return (1 - CONST_x_PG) * (CONST_a0_V_POPC + T * CONST_a1_V_POPC) + CONST_x_PG * (CONST_a0_V_POPG + T * CONST_a1_V_POPG)
end

#########################################################
function water_volume(T)
  return CONST_p0_VW + CONST_p1_VW * T + CONST_p2_VW * T^2 + CONST_p3_VW * T^3 + CONST_p4_VW * T^4
end

#########################################################
#@njit(parallel=True)
function FTreal_erf(q, mu, d, sig)
  """ FTreal_erf(q, mu, d, sig) """
  return sin(q * d / 2.0) / (q * d / 2.0) * exp(-(q * sig)^2 / 2.0) * cos(q * mu)
end

#########################################################
#@njit(parallel=True)
function FTreal_gauss(q, mu, sig)
  """ FTreal_gauss(q, mu, sig) """
  return exp(-(q * sig)^2 / 2.0) * cos(q * mu)
end

#########################################################
function Slab(x, mu, L, sig)
  return 0.5 * (erf((x - (mu - L / 2.0)) / (sqrt(2) * sig)) - erf((x - (mu + L / 2)) / (sqrt(2) * sig)))
end

#########################################################
function Gauss(x, V, mu, sig, A_L)
  return V * PDF_normal(x, mu, sig) / A_L
end

####################################        j0^2 MOMENTS (SCHULZ PDF)        #########################################
####################################        (checked!)                        #########################################

#@njit(parallel=True)
function mu4(q, Z, a)
  return a^2 * (Z + 2) * (Z + 1) * (1 - (1 + 4 * q^2 * a^2)^(-(Z + 3) / 2.0) * cos((Z + 3) * atan(2 * q * a))) / (2 * q^2)
end

##################################################################################################################
##################################################################################################################

################################     SYMMETRIC VESICLE FOR X-RAY SLDs    #########################################

########################################        SDP MODELLING            #########################################
####################################       SEPARATED FORM FACTOR         #########################################

##################################################################################################################
##################################################################################################################

#########################################################
######### Symmetric POPC bilayer ########################
######### liposomes and proteoliposomes #################
#########################################################

##################
function intensity(P::NamedTuple, q)
  return intensity(Float64[P...], q)
end

function intensity(P::AbstractArray{Float64}, q)

  Norm, nv, Rm, Z, n_TR, d_TR, s_TR, d_Chol, s_Chol, d_PCN, s_PCN, d_CG, s_CG,
  A_L, s_D_C, s_CH2, d_CH, s_CH, s_CH3, r_PCN, r_CG, r12, r32, T, V_BW, Con = P

  # [Example 1], fixed parameters
  # Norm  1e5    #    Normalization
  # n_TR  0.0    #    Tris fraction
  # d_TR  1.0    #    Tris width (nm)
  # s_TR  0.29   #    Tris position (nm)
  # d_CH  0.90   #    CH position (nm)
  # s_CH  0.305  #    CH width (nm)
  # r12   0.81   #    V_CH/V_CH2
  # T     37     #    Temperature (°C)

  # [example 1] fixed
  Cw = CONST_p0_Cw + CONST_p1_Cw * T + CONST_p2_Cw * T^2 + CONST_p3_Cw * T^3
  xtris = CONST_ctris / Cw # mole fraction of free TRIS in bulk
  xEDTA = CONST_cEDTA / Cw # mole fraction of free EDTA in bulk

  # Volumes
  # [example 1] fixed
  V_L = lipid_volume(T)
  V_HW = water_volume(T)
  V_HC = V_L - ((1 - CONST_x_PG) * CONST_V_HL_PC + CONST_x_PG * CONST_V_HL_PG)

  # Calculation of mean D_C
  D_C = V_HC / A_L

  # Quasi-molecular volumes
  # [example 1] r12 fixed
  V_CH2 = V_HC / (CONST_n_CH2 + CONST_n_CH * r12 + CONST_n_CH3 * r32)    # Volume of CH2 groups
  V_CH = V_CH2 * r12                             # Volume of CH groups
  V_CH3 = V_CH2 * r32                             # Volume of CH3 groups

  V_CG = CONST_V_HL_PC * r_CG                   # Volume of CG group
  V_PCN = CONST_V_HL_PC * r_PCN                  # Volume of PCN group
  V_Chol = CONST_V_HL_PC * (1 - r_PCN - r_CG)        # Volume of CholCH3 group

  # CONST
  CONST_V_PG1 = CONST_V_HL_PG * 0.16                 # Kucerka 2012
  CONST_V_PG2 = CONST_V_HL_PG * (1 - 0.51 - 0.16)   # Kucerka 2012

  ############### X-ray scattering lengths (nm)
  # [example 1] rho_sol fixed
  rho_sol = (CONST_b_HW + xtris * CONST_b_tris + xEDTA * CONST_b_EDTA) / V_HW
  drho_Chol = ((1 - CONST_x_PG) * CONST_b_Chol / V_Chol + CONST_x_PG * CONST_b_PG2 / CONST_V_PG2) - rho_sol
  drho_PCN = ((1 - CONST_x_PG) * CONST_b_PCN / V_PCN + CONST_x_PG * CONST_b_PG1 / CONST_V_PG1) - rho_sol
  drho_CG = CONST_b_CG / V_CG - rho_sol
  drho_TR = CONST_b_tris / CONST_V_tris - rho_sol
  drho_CH = CONST_b_CH / V_CH - rho_sol
  drho_CH2 = CONST_b_CH2 / V_CH2 - rho_sol
  drho_CH3 = CONST_b_CH3 / V_CH3 - rho_sol
  drho_HW = CONST_b_HW / V_BW - rho_sol

  ############### D_C polydispersity
  # Number of integration points
  N = 201

  # Parameter dependent integration bounds
  # HC_min, HC_max = D_C-3*s_D_C, D_C+3*s_D_C
  #
  # Hardcoded integration bounds for POPC, fixed T = 37, A_L in [0.598, 0.719]
  HC_min, HC_max = V_HC / 0.719 - 0.5, V_HC / 0.589 + 0.5
  # Minimum and maximum samples (offset by 1/2 of the integration step w.r.t integration bounds for the midpoint rule)
  integration_step = (HC_max - HC_min) / N
  HC_first, HC_last = HC_min + 0.5 * integration_step, HC_max - 0.5 * integration_step

  HC_array = range(HC_first, HC_last, length=N)
  Normal = PDF_normal.(HC_array, D_C, s_D_C)

  ############### c-prefactors
  c_Chol = ((1 - CONST_x_PG) * V_Chol + CONST_x_PG * CONST_V_PG2) / A_L
  c_PCN = ((1 - CONST_x_PG) * V_PCN + CONST_x_PG * CONST_V_PG1) / A_L
  c_CG = V_CG / A_L
  c_TR = CONST_V_tris * n_TR / A_L

  c_CH = V_CH * CONST_n_CH / V_HC * HC_array
  c_CH3 = V_CH3 * CONST_n_CH3 / V_HC * HC_array

  d_comp = d_CG + d_PCN + d_Chol + CONST_d_shl

  ############### calculating scattering amplitude    -----------------------------------------------
  Am = zeros(size(q, 1), N)

  for i in 1:N

    # Adding hydrocarbon-chain envelope
    Am[:, i] += 2 * drho_CH2 * HC_array[i] * FTreal_erf.(q, 0, 2 * HC_array[i], s_CH2)
    # Adding CH and CH3 groups
    Am[:, i] += 2 * (drho_CH - drho_CH2) * c_CH[i] * FTreal_gauss.(q, d_CH, s_CH)
    Am[:, i] += 2 * (drho_CH3 - drho_CH2) * c_CH3[i] * FTreal_gauss.(q, 0, s_CH3)

    # Adding hydration-water envelope
    Am[:, i] += 4 * drho_HW * d_comp * FTreal_erf.(q, (HC_array[i] + d_comp / 2.0), d_comp, s_CH2)
    # Adding CG, PCN and CholCH3 groups
    Am[:, i] += 2 * (drho_TR - drho_HW) * c_TR * FTreal_gauss.(q, (HC_array[i] + d_TR / 2.0), s_TR)
    Am[:, i] += 2 * (drho_CG - drho_HW) * c_CG * FTreal_gauss.(q, (HC_array[i] + d_CG / 2.0), s_CG)
    Am[:, i] += 2 * (drho_PCN - drho_HW) * c_PCN * FTreal_gauss.(q, (HC_array[i] + d_CG + d_PCN / 2.0), s_PCN)
    Am[:, i] += 2 * (drho_Chol - drho_HW) * c_Chol * FTreal_gauss.(q, (HC_array[i] + d_CG + d_PCN + d_Chol / 2.0), s_Chol)

    ############### Ensemble average
    # multiply each columns of Am by Normal and sum along the columns,
    # then multiply by integration step
    I = vec(sum(Am .^ 2 .* Normal'; dims=2)) * integration_step

    alp = Rm / (Z + 1)
    return @. (Norm * nv * 1e-6) * I * (16 * pi^2 * mu4(q, Z, alp)) + Con * (0.99 * (1.0 / (1 + exp(-8 * (q - 1.0)))) + 0.01)
  end
end

##################
function negative_water(P::NamedTuple)
  return negative_water([P...])
end

function negative_water(P::AbstractArray{Float64})

  _, _, _, _, n_TR, d_TR, s_TR, d_Chol, s_Chol, d_PCN, s_PCN, d_CG, s_CG,
  A_L, _, s_CH2, _, _, _, r_PCN, r_CG, _, _, T, _, _ = P

  # Volumes
  V_L = lipid_volume(T)
  V_HC = V_L - ((1 - CONST_x_PG) * CONST_V_HL_PC + CONST_x_PG * CONST_V_HL_PG)

  # Calculation of mean D_C
  D_C = V_HC / A_L

  # Quasi-molecular volumes
  V_CG = CONST_V_HL_PC * r_CG                   # Volume of CG group
  V_PCN = CONST_V_HL_PC * r_PCN                 # Volume of PCN group
  V_Chol = CONST_V_HL_PC * (1 - r_PCN - r_CG)   # Volume of CholCH3 group

  check = 0
  z_array = range(0.0, 4.0, length=81)

  CG = Gauss.(z_array, V_CG, D_C + d_CG / 2.0, s_CG, A_L)
  PCN = Gauss.(z_array, V_PCN, D_C + d_CG + d_PCN / 2.0, s_PCN, A_L)
  Chol = Gauss.(z_array, V_Chol, D_C + d_CG + d_PCN + d_Chol / 2.0, s_Chol, A_L)
  TRIS = Gauss.(z_array, n_TR * CONST_V_tris, D_C + d_TR / 2.0, s_TR, A_L)
  BW = Slab.(z_array, D_C + (d_CG + d_PCN + d_Chol + CONST_d_shl) / 2.0, d_CG + d_PCN + d_Chol + CONST_d_shl, s_CH2) - CG - PCN - Chol - TRIS

  for i in (BW)
    if i < -1e-3
      check += 1
    end
  end

  return check
end

function get_free_parameters_idx(metadata::Dict)
  m = metadata["parameters"]
  return filter(i -> !m[i]["fixed"], eachindex(m))
end

function reduce_to_free_parameters(metadata::Dict, params_init::NamedTuple, lb, ub, q)
  default_params = Float64[params_init...]
  free_idx = get_free_parameters_idx(metadata)
  P_reduced = view(default_params, free_idx)
  lb_reduced = lb[free_idx]
  ub_reduced = ub[free_idx]
  function reduced(P::AbstractArray{Float64})
    P_reduced .= P
    return intensity(default_params, q)
  end
  return reduced, P_reduced, lb_reduced, ub_reduced
end

function reduce_to_free_parameters(metadata::Dict, params::Vector{Float64})
  free_idx = get_free_parameters_idx(metadata)
  return params[free_idx]
end

end # module PLUV

# vim ts=4,sts=4,sw=4