MODULE CODE6
 USE CODE5, ONLY : Sample_Gamma,Sample_Normal,Sample_Dirichlet,Sample_Multivariate_Normal, &
                                                Sample_Uniform
IMPLICIT NONE

TYPE Univariate_normal_mixture
INTEGER :: n              ! Observations
INTEGER :: k              ! Mixture components
REAL(8), POINTER :: mu(:)     ! Mixture means
REAL(8), POINTER :: tau(:)    ! Mixture precisions
REAL(8), POINTER :: l(:)      ! Mixture weights
REAL(8) :: prior_mu_mu,prior_mu_p,prior_tau_a,prior_tau_b
REAL(8), POINTER :: prior_l_a(:)
LOGICAL, POINTER :: linc(:)   
LOGICAL, POINTER :: lgroup(:,:)
INTEGER, POINTER :: ngroup(:)
LOGICAL :: zero_rest,scale_mix,mean_mix,norm_rest,first_scale
CHARACTER(len=25) :: title
CHARACTER(len=100) :: output_file
INTEGER :: output_file_number
REAL(8), POINTER :: par_mat(:,:)
END TYPE Univariate_normal_mixture

CONTAINS

SUBROUTINE init_Univariate_normal_mixture(this,n,k,mui,taui,li,title,cum)
IMPLICIT NONE
TYPE(Univariate_normal_mixture), INTENT(INOUT) :: this
INTEGER, INTENT(IN) :: n,k,cum
REAL(8), INTENT(IN), target :: mui(k),taui(k),li(k)
CHARACTER(*), INTENT(IN) :: title
this%n=n
this%k=k
ALLOCATE( this%lgroup(n,k),this%prior_l_a(k),this%ngroup(k),this%mu(k), &
                this%tau(k),this%l(k), this%linc(n) )
this%title = title
this%mu = mui
this%tau = taui
this%l = li
this%ngroup=0
this%ngroup(1)=n
this%output_file_number=cum
END SUBROUTINE init_Univariate_normal_mixture


subroutine update_Uni_mixture_parameters(this,y)
implicit none
type(Univariate_normal_mixture), intent(inout) :: this
real(8), intent(in) :: y(this%n)
integer ::j,k1,i
real(8) :: bv(this%k),av(this%k),pa,pb(this%k),p2bs,p2b(this%k), &
        ag,bg,al(this%k),mean_mu(this%k-1),cv_mu(this%k-1,this%k-1),ua(this%k),w
IF (this%norm_rest) THEN
        this%mu = 0.0d0
        ag = 0.5d0*DBLE(this%n) + this%prior_tau_a
        bg = 0.5d0*SUM(y*y) + this%prior_tau_b
        this%tau(1) = Sample_gamma(ag,bg)
ELSE
        IF (this%scale_mix) THEN
                this%mu=0.0d0
        ELSE
                DO j=1,this%k
                        bv(j) = 1.0d0/( this%ngroup(j)*this%tau(j) + this%prior_mu_p )
                        av(j) = bv(j)*( this%tau(j)*sum( y , this%lgroup(:,j) ) + &
                                this%prior_mu_p*this%prior_mu_mu )
                END DO
                IF (this%zero_rest) THEN
                        k1 = this%k-1
                        pa = sum( this%l*av )
                        pb = this%l*bv
                        p2b = (this%l**2)*bv
                        p2bs = sum( p2b )
                        mean_mu = av(1:k1) - pb(1:k1)*(pa/p2bs)
                        forall(i=1:k1) cv_mu(i,i) = bv(i)*(1.0d0 -  (p2b(i)/p2bs))
                        forall(i=1:k1,j=1:k1,i > j) cv_mu(i,j) = -pb(i)*pb(j)/p2bs
                        forall(i=1:k1,j=1:k1,i < j) cv_mu(i,j) = cv_mu(j,i)
                        this%mu(1:k1) = Sample_Multivariate_Normal(mean_mu,cv_mu)
                        this%mu(this%k) = - sum( this%mu(1:k1)*this%l(1:k1) )/this%l(this%k)
                ELSE
                        DO j = 1, this%k
                                this%mu(j) = Sample_Normal(av(j),bv(j))
                        END DO
                END IF
        END IF
        IF (this%mean_mix) THEN
                this%tau = 1.0d0  
        ELSE IF (this%first_scale) THEN
                this%tau(1) = 1.0d0  
                DO j=2,this%k
                        bg = 0.5d0*sum( (y - this%mu(j))**2, this%lgroup(:,j) ) + this%prior_tau_b
                        ag = 0.5d0*DBLE(this%ngroup(j)) + this%prior_tau_a
                        this%tau(j) = Sample_Gamma(ag,bg)
                END DO  
        ELSE
                DO j=1,this%k
                        bg = 0.5d0*sum( (y - this%mu(j))**2, this%lgroup(:,j) ) + this%prior_tau_b
                        ag = 0.5d0*DBLE(this%ngroup(j)) + this%prior_tau_a
                        this%tau(j) = Sample_Gamma(ag,bg)
                END DO
        END IF
        al = this%ngroup + this%prior_l_a
        this%l=Sample_Dirichlet(this%k,al)
END IF
END SUBROUTINE update_Uni_mixture_parameters

SUBROUTINE update_Uni_nor_mixture_groups(this,y)
IMPLICIT NONE
TYPE(Univariate_normal_mixture), intent(inout) :: this
real(8), intent(in) :: y(this%n)
real(8) :: lm_p(this%n,this%k), m_p(this%n,this%k),sm_p(this%n),cp(this%n,this%k+1),u
integer :: j,n,i
n = this%n
DO j=1,this%k
        lm_p(:,j) = dlog(this%l(j)) + 0.5d0*dlog(this%tau(j)) - 0.5d0*this%tau(j)*(y - this%mu(j))**2
END DO
WHERE (lm_p > -500.0d0)
        m_p = DEXP(lm_p)
ELSEWHERE
        m_p = 0.0d0
ENDWHERE
! normalize
sm_p = sum(m_p,dim=2)
DO j=1,this%k
        m_p(:,j) = m_p(:,j)/sm_p
END DO

this%lgroup = .false.
cp(:,1)=0.0d0
cp(:,2)=m_p(:,1)
DO j=2,this%k-1
        cp(:,j+1)=cp(:,j) + m_p(:,j)
END DO
cp(:,this%k+1)=1.0d0
DO i = 1, n
        u = Sample_Uniform(0.0d0,1.0d0)
        DO j = 1, this%k
                IF ((cp(i,j) < u) .and. (cp(i,j+1) > u) .and. (this%linc(i))) this%lgroup(i,j)=.true.   
        END DO
END DO
! update group counts
this%ngroup = count( this%lgroup , dim=1 )
END SUBROUTINE update_Uni_nor_mixture_groups


SUBROUTINE update_Uni_normal_mixture(this,y)
implicit none
type(Univariate_normal_mixture), intent(inout) :: this
real(8), intent(in) :: y(this%n)
IF (.NOT.this%norm_rest) call update_Uni_nor_mixture_groups(this,y)
call update_Uni_mixture_parameters(this,y)
END SUBROUTINE update_Uni_normal_mixture

SUBROUTINE write_output_Uni_nor_mixture(this)
implicit none
type(Univariate_normal_mixture), intent(in) :: this
integer k
integer nf
k = this%k
IF (this%norm_rest) THEN
        write(unit=this%output_file_number,fmt='(40f16.8)') this%mu,this%tau
ELSE
        write(unit=this%output_file_number,fmt='(40f16.8)') this%mu,this%tau,this%l
END IF
END SUBROUTINE write_output_Uni_nor_mixture

END MODULE CODE6
