#!/usr/bin/python

import sys
import os.path
import re
import time
import numpy as np
from scipy.stats import norm , gamma
from scipy import special
from scipy import signal
from scipy import ndimage
#from numba import njit, prange

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


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

# POPG molar ratio
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_n_CH  = 2
CONST_n_CH2 = 28
CONST_n_CH3 = 2

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

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

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

# Lipid-volume temperature-dependencies a0 + a1*T (nm^3)
CONST_a0_V_POPG = 1.17881068602663
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_p0_VW    = 0.0299218
CONST_p1_VW    = -2.25941e-06
CONST_p2_VW    = 2.5675e-07
CONST_p3_VW    = -1.69661e-09
CONST_p4_VW    = 6.52029e-12
# polynome coefficient for T-dependency of bulk-water molar concentration (Cw)
#Units in degree Celsius
CONST_p0_Cw    = 55.5052
CONST_p1_Cw    = 0.00131894
CONST_p2_Cw    = -0.000334396
CONST_p3_Cw    = 9.10861e-07

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

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

### Extra molecules

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


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

#########################################################
#@njit(parallel=True)
def PDF_normal(x, mu, sig) :
    return np.exp(-(x-mu)**2 / (2*sig**2) ) /( sig*np.sqrt(2*np.pi))

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

#########################################################
def 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)

#########################################################
def 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

#########################################################
#@njit(parallel=True)
def FTreal_erf(q, mu, d, sig) :
    """ FTreal_erf(q, mu, d, sig) """
    return np.where(q==0, 1, np.sin(q*d/2.)/(q*d/2.) * np.exp(-(q*sig)**2/2.) * np.cos(q*mu)  )

#########################################################
#@njit(parallel=True)
def FTreal_gauss(q, mu, sig) :
    """ FTreal_gauss(q, mu, sig) """
    return np.exp(-(q*sig)**2/2.) * np.cos(q*mu)

#########################################################
def Slab(x, mu, L, sig) :
    return 0.5 * ( special.erf( (x - (mu-L/2.))/(np.sqrt(2)*sig) ) - special.erf( (x - (mu+L/2))/(np.sqrt(2)*sig) ) )

#########################################################
def Gauss( x, V, mu, sig, A_L ) :
    return V * PDF_normal(x, mu, sig) / A_L

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

#@njit(parallel=True)
def mu4(q, Z, a) :
    return np.where(q==0, a**4*(Z+1)*(Z+2)*(Z+3)*(Z+3), a**2 * (Z+2)*(Z+1) * ( 1 - (1+4*q*q*a*a)**(-(Z+3)/2.) * np.cos((Z+3)*np.arctan(2*q*a)) ) / ( 2*q**2 )  )

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

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

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

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

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

class SDP_POPC_RecBuf:

##################
    def __init__(self, q, PAR) :
        self.q = q
        [self.Norm, self.nv,
        self.Rm, self.Z,
        self.n_TR, self.d_TR, self.s_TR,
        self.d_Chol, self.s_Chol, self.d_PCN, self.s_PCN, self.d_CG, self.s_CG,
        self.A_L, self.s_D_C,
        self.s_CH2, self.d_CH, self.s_CH, self.s_CH3,
        self.r_PCN, self.r_CG, self.r12, self.r32,
        self.T, self.V_BW,
        self.Con] = PAR

        # [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*self.T + CONST_p2_Cw*self.T**2 + CONST_p3_Cw*self.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
        self.V_L = lipid_volume(self.T)
        V_HW = water_volume(self.T)
        V_HC = self.V_L - ( (1-CONST_x_PG) * CONST_V_HL_PC + CONST_x_PG * CONST_V_HL_PG )

        # Calculation of mean D_C
        self.D_C    = V_HC / self.A_L

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

        self.V_CG      = CONST_V_HL_PC    *    self.r_CG                # Volume of CG group
        self.V_PCN     = CONST_V_HL_PC    *     self.r_PCN                 # Volume of PCN group
        self.V_Chol    = CONST_V_HL_PC    *   (1-self.r_PCN-self.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/self.V_Chol    + CONST_x_PG*CONST_b_PG2/CONST_V_PG2 )  - rho_sol
        drho_PCN       = ( (1-CONST_x_PG)*CONST_b_PCN/self.V_PCN    + CONST_x_PG*CONST_b_PG1/CONST_V_PG1 )    - rho_sol
        drho_CG        = CONST_b_CG  / self.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    / self.V_BW                                                      - rho_sol

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

        # Parameter dependent integration bounds
        # HC_min, HC_max = self.D_C-3*self.s_D_C, self.D_C+3*self.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 = np.linspace(HC_first, HC_last, N)
        Normal = PDF_normal(HC_array, self.D_C, self.s_D_C)

        ############### calculating scattering amplitude    -----------------------------------------------
        self.Am = np.zeros([N,self.q.shape[0]],dtype=float)
        c_CH = np.zeros(N,dtype=float)
        c_CH3 = np.zeros(N,dtype=float)

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

        for hc in range(N):
            c_CH[hc]    = V_CH    * CONST_n_CH    / (V_HC / HC_array[hc] )
            c_CH3[hc]   = V_CH3   * CONST_n_CH3   / (V_HC / HC_array[hc] )

        for hc in range(N):

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

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

        ############### Ensemble average
        # multiply each columns of Am by Normal and sum along the columns,
        # then multiply by integration step
        self.I = np.einsum('ij,i->j', self.Am**2, Normal) * integration_step

##################
    def intensity(self):
        alp = self.Rm/(self.Z+1)
        return ( self.Norm * self.nv*1e-6 ) * self.I * ( 16*np.pi**2*mu4(self.q,self.Z,alp) ) + self.Con*( 0.99*(1./(1+np.exp(-8*(self.q-1.)))) + 0.01 )

##################
    def negative_water(self):

        self.check = 0
        z_array = np.linspace(0.,4.,81)

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

        for i in(BW) :
            if i <-0.001 : self.check+= 1

        return self.check

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

class SDP_POPC_RecBuf_LogNormal:

##################
    def __init__(self, q, PAR) :
        self.q = q
        [self.Norm, self.nv,
        self.Rm, self.Z,
        self.n_TR, self.d_TR, self.s_TR,
        self.d_Chol, self.s_Chol, self.d_PCN, self.s_PCN, self.d_CG, self.s_CG,
        self.A_L, self.s_D_C,
        self.s_CH2, self.d_CH, self.s_CH, self.s_CH3,
        self.r_PCN, self.r_CG, self.r12, self.r32,
        self.T, self.V_BW,
        self.Con] = PAR

        # [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*self.T + CONST_p2_Cw*self.T**2 + CONST_p3_Cw*self.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
        self.V_L = lipid_volume(self.T)
        V_HW = water_volume(self.T)
        V_HC = self.V_L - ( (1-CONST_x_PG) * CONST_V_HL_PC + CONST_x_PG * CONST_V_HL_PG )

        # Calculation of mean D_C
        self.D_C    = V_HC / self.A_L

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

        self.V_CG      = CONST_V_HL_PC    *    self.r_CG                # Volume of CG group
        self.V_PCN     = CONST_V_HL_PC    *     self.r_PCN                 # Volume of PCN group
        self.V_Chol    = CONST_V_HL_PC    *   (1-self.r_PCN-self.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/self.V_Chol    + CONST_x_PG*CONST_b_PG2/CONST_V_PG2 )  - rho_sol
        drho_PCN       = ( (1-CONST_x_PG)*CONST_b_PCN/self.V_PCN    + CONST_x_PG*CONST_b_PG1/CONST_V_PG1 )    - rho_sol
        drho_CG        = CONST_b_CG  / self.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    / self.V_BW                                                      - rho_sol

        ############### D_C polydispersity
        N = 200
        HC_array = np.linspace(self.D_C-3*self.s_D_C, self.D_C+3*self.s_D_C, N)
        LogNormal = PDF_lognormal(HC_array, self.D_C, self.s_D_C)

        ############### calculating scattering amplitude    -----------------------------------------------
        self.Am = np.zeros([HC_array.shape[0],self.q.shape[0]],dtype=float)
        c_CH = np.zeros(HC_array.shape[0],dtype=float)
        c_CH3 = np.zeros(HC_array.shape[0],dtype=float)

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

        for hc in range(HC_array.shape[0]):
            c_CH[hc]    = V_CH    * CONST_n_CH    / (V_HC / HC_array[hc] )
            c_CH3[hc]   = V_CH3   * CONST_n_CH3   / (V_HC / HC_array[hc] )

        for hc in range(HC_array.shape[0]):

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

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

        ############### Ensemble average
        self.I = np.zeros(self.q.shape[0], dtype=float)

        for hc in range(HC_array.shape[0]):
            if hc==0 or hc==N-1 :    self.I+= self.Am[hc]**2 * LogNormal[hc] / 2
            else :                   self.I+= self.Am[hc]**2 * LogNormal[hc]

        self.I*= 6*self.s_D_C/(N-1)

##################
    def intensity(self):
        alp = self.Rm/(self.Z+1)
        return ( self.Norm * self.nv*1e-6 ) * self.I * ( 16*np.pi**2*mu4(self.q,self.Z,alp) ) + self.Con*( 0.99*(1./(1+np.exp(-8*(self.q-1.)))) + 0.01 )

##################
    def negative_water(self):

        self.check = 0
        z_array = np.linspace(0.,4.,81)

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

        for i in(BW) :
            if i <-0.001 : self.check+= 1

        return self.check

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

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