module mod_harm2grid
    use mod_parameters
    use mod_pp_variables
    use mod_read_f
    implicit none
    
contains

    pure integer function jmm(j,m)
        integer, intent(in) :: j, m
        jmm = j*(j+1) + m + 1
    end function jmm

    pure integer function jml_v(j, m, l) 
        implicit none
        integer, intent(in) :: j, m, l
        jml_v = 3 * (j*(j+1)/2 + m) + l - j
    end function jml_v

    pure integer function jm(j, m)
        implicit none
        integer, intent(in) :: j, m
        jm = j*(j+1)/2 + m + 1
    end function jm

    integer function jm_f(j,m)
        implicit none
        integer :: j, m
        jm_f = j*(j+1)/2 + m + 1
    end function jm_f

    subroutine generate_theta(n)
        implicit none
        integer :: ith, n
        real(8) :: ht

        theta(1) = 0.0d0
        theta(n) = Pi
        ht = (theta(n)-theta(1)) / dble(n-1)
        do ith = 2, n-1
            theta(ith) = theta(ith-1) + ht
        end do
    end subroutine generate_theta


    subroutine generate_phi(n)
        implicit none
        integer :: iph, n
        real(8) :: hp

        phi(1) = 0.0d0
        phi(n) = 2.0d0 * Pi
        hp = (phi(n)-phi(1)) / dble(n-1)
        do iph = 2, n-1
            phi(iph) = phi(iph-1) + hp
        end do
    end subroutine generate_phi

    subroutine scalar2grid(lmax, sph_coeffs, grid_coeffs)
        implicit none
        integer                                   :: lmax
        complex(8), dimension(:,:)  , intent(in)  :: sph_coeffs
        real(8)   , dimension(:,:,:), intent(out) :: grid_coeffs
        integer                                   :: j, m, ith, iph, ir
        complex(8)                                :: YY
        real(8)                                   :: c

        grid_coeffs = zero
        
        do j = 1, lmax
            do m = 0, j
                c = 2.0d0 - dble(dirac(m,0))
                do ith = 1, size(theta)
                    do iph = 1, size(phi)
                        YY = yjm(j, m, theta(ith), phi(iph))
                        do ir = 1, size(sph_coeffs,1)
                            grid_coeffs(iph, ith, ir) = grid_coeffs(iph, ith, ir) + real(sph_coeffs(ir, jm(j,m)) * YY) * c
                        end do
                    end do
                end do
            end do
        end do
        print*, 's2g', maxval(abs(grid_coeffs))
    end subroutine scalar2grid


    subroutine m2_vector(lmax, vin, vout)
        implicit none
        integer                   , intent(in)  :: lmax
        complex(8), dimension(:,:), intent(in)  :: vin
        complex(8), dimension(:,:), intent(out) :: vout
        integer :: j, m

        vout = dzero
        do j = 1, lmax
            do m = 0, min(2,j), 2
                vout(:, jml_v(j,m,j-1)) = vin(:, jml2(j,m,j-1))
                vout(:, jml_v(j,m,j  )) = vin(:, jml2(j,m,j  ))
                vout(:, jml_v(j,m,j+1)) = vin(:, jml2(j,m,j+1))
            end do
        end do
        
    end subroutine m2_vector


    subroutine m2_scalar(lmax, vin, vout)
        implicit none
        integer                   , intent(in)  :: lmax
        complex(8), dimension(:,:), intent(in)  :: vin
        complex(8), dimension(:,:), intent(out) :: vout
        integer :: j, m

        vout = dzero
        do j = 0, lmax
            do m = 0, min(2,j), 2
                vout(:, jm(j,m)) = vin(:, jm2(j,m))
            end do
        end do
    end subroutine m2_scalar


    subroutine vector2vxyz(lmax, v, vx, vy, vz)
        implicit none
        complex(8), dimension(:), intent(in)  :: v
        complex(8), dimension(:), intent(out) :: vx, vy, vz
        real(8)  :: mo, pcc, cg, cx
        complex(8) :: vv, cy
        integer :: j2, m2, l2, ii, lmax
        complex(8),dimension(1:jmm(lmax+1, lmax+1)) :: x, y, z
        
        intent(in) :: lmax
        x  = dcmplx(0.0d0, 0.0d0)
        y  = dcmplx(0.0d0, 0.0d0)
        z  = dcmplx(0.0d0, 0.0d0)
        

        !===============================================================
        ! x - component
        !===============================================================
        cx  = -dsqrt(2.0d0)/2.0d0
        do j2 = 0, lmax
        do m2 = -j2, j2
        do l2 = abs(j2-1), j2+1
            mo  = (-1.0d0)**(j2 + l2)
            pcc  = -dsqrt(dble(2*j2 + 1)/dble(2*(2*l2+1)))

            vv  = v(jml_v(j2,abs(m2),l2))
            if(m2.lt.0) then
                vv = (-1)**(j2 + m2 + l2 + 1) * dconjg(vv)
            end if
            if(abs(m2+1).le.l2) then
                ii  = jmm(l2,m2+1)
                call cleb_1(j2,m2,1,1,l2,m2+1,cg)
                x(ii) = x(ii) + mo * pcc * cg * vv
            end if
            if(abs(m2-1).le.l2) then
                ii  = jmm(l2,m2-1)
                call cleb_1(j2,m2,1,-1,l2,m2-1,cg)
                x(ii) = x(ii) - mo * pcc * cg * vv
            end if
        end do
        end do
        end do

        do j2 = 0, lmax + 1
        do m2  = 0, j2
            vx(jm(j2, m2)) = x(jmm(j2, m2))
        end do
        end do
        !===============================================================
        ! y - component
        !===============================================================
        cy  = Im * dsqrt(2.0d0)/2.0d0
        do j2 = 0, lmax
        do m2 = -j2, j2
        do l2 = abs(j2-1), j2+1

            mo  = (-1.0d0)**(j2 + l2)
            pcc  = dsqrt(dble(2*j2 + 1)/dble(2*(2*l2+1)))

            vv  = v(jml_v(j2,abs(m2),l2))
            if(m2.lt.0) then
                vv = ((-1.0d0)**(j2 + m2 + l2 + 1) )* dconjg(vv)
            end if
            if(abs(m2+1).le.l2) then
                ii  = jmm(l2,m2+1)
                call cleb_1(j2,m2,1,1,l2,m2+1,cg)
                y(ii) = y(ii) + im * mo * pcc * cg * vv
            end if
            if(abs(m2-1).le.l2) then
                ii  = jmm(l2,m2-1)
                call cleb_1(j2,m2,1,-1,l2,m2-1,cg)
                y(ii) = y(ii) + im * mo * pcc * cg * vv
            end if
        end do
        end do
        end do

        do j2 = 0, lmax + 1
        do m2  = 0, j2
            vy(jm(j2, m2)) = y(jmm(j2, m2))
        end do
        end do
        !===============================================================
        ! z - component
        !===============================================================
        
        do j2 = 0, lmax!+1
        do m2 = -j2, j2
        do l2 = abs(j2-1), j2+1
            mo  = (-1)**(j2 + l2)
            pcc  = dsqrt(dble(2*j2 + 1)/dble(2*l2+1))

            vv  = v(jml_v(j2,abs(m2),l2))
            if(m2.lt.0) then
                vv = ((-1.0d0)**(j2 + m2 + l2 + 1)) * dconjg(vv)
            end if
            if(abs(m2).le.l2) then
                ii  = jmm(l2,m2)
                call cleb_1(j2,m2,1,0,l2,m2,cg)
                z(ii) = z(ii) +  mo * pcc * cg * vv
            end if
        end do
        end do
        end do

        do j2 = 0, lmax + 1
        do m2  = 0, j2
            vz(jm(j2, m2)) = z(jmm(j2, m2))
        end do
        end do
    end subroutine vector2vxyz



    complex(8) function yjm(j,m,th,ph)
        implicit none
        integer :: j,m
        real(8) :: th, ph

        
        if(m.eq.0) then
            yjm = normalized_legendre_pol_x(j,m,cos(th))*dcmplx(1.0d0,0.0d0)
        else
            yjm = normalized_legendre_pol_x(j,m,cos(th))*exp(im*m*ph)
        end if
    end function yjm

    subroutine yjm_array(lmax, th, ph, yjm_a)
        implicit none
        integer :: lmax
        real(8) :: th, ph
        complex(8) :: yjm_a(:)
        integer :: j, m

        do j = 0, lmax
            yjm_a(jm_f(j,0)) = dcmplx(normalized_legendre_pol_x(j,0,cos(th)),0.0d0)
            do m = 1, j
                yjm_a(jm_f(j,m)) = normalized_legendre_pol_x(j,m,cos(th))*exp(im*m*ph)
            end do
        end do
    end subroutine yjm_array

    real(8) function normalized_legendre_pol_x(l,m,x)
        implicit none
        integer :: l,m
        real(8) :: x
        real(8) :: xx, pmm, f, pmmp1, pll, c1, c2, c3, c4
        integer :: i, ll

        if((m.lt.0).or.(m.gt.l).or.(abs(x).gt.1.d0)) stop 'bad arguments in plgndr'
        pmm=1.0d0/dsqrt(4.0d0*pi) !P(0,0)
        if(m.gt.0)then
            xx = dsqrt(1.0d0-x**2)
            do i = 1, m
                f  = dsqrt(dble(2*i+1)/dble(2*i)) !(dble(2*m+1)/dble(2*m-1))*
                pmm = -pmm*f*xx
            end do
        end if
        if(l.eq.m) then
            normalized_legendre_pol_x = pmm
        else
            pmmp1=x*dsqrt(2.0d0*dble(m)+3.0d0)*pmm
            if(l.eq.(m+1)) then
                normalized_legendre_pol_x = pmmp1
            else
                do ll = m+2, l
                    c1 = dble((2*ll+1)*(2*ll-1))
                    c2 = dble((ll-m)*(ll+m))
                    c3 = dble((2*ll+1)*(ll-m-1)*(ll+m-1))
                    c4 = dble((2*ll-3)*(ll-m)*(ll+m))
                    pll = x*dsqrt(c1/c2)*pmmp1 - dsqrt(c3/c4)*pmm
                    pmm  =pmmp1
                    pmmp1=pll
                end do
                normalized_legendre_pol_x = pll
            end if
        end if
        normalized_legendre_pol_x = normalized_legendre_pol_x * (-1.0d0)**(m)
    end function normalized_legendre_pol_x


    subroutine ocean_velocity_xyz(itime)
        implicit none
        integer :: itime, lmax, nn, ir, lmax2
        real(8) :: time
        complex(8), dimension(:,:)  , allocatable :: vx, vy, vz
        real(8)   , dimension(:,:,:), allocatable :: vxg, vyg, vzg


        call read_velocity(itime, lmax, time)
        nn = layer(numoflayers-1)%n+1
        print*, 'nn:', nn
        lmax2 = lmax+1
        if(allocated(velocity_ocean_m2)) deallocate(velocity_ocean_m2)
        if(allocated(vx)) deallocate(vx, vy, vz)
        if(allocated(vxg)) deallocate(vxg, vyg, vzg)
        allocate(vx(nn+1, jml_v(lmax2, lmax2, lmax2+1)), vy(nn+1, jml_v(lmax2, lmax2, lmax2+1)), vz(nn+1, jml_v(lmax2, lmax2, lmax2+1)))
        allocate(velocity_ocean_m2(nn, jml_v(lmax, lmax, lmax+1)))
        call m2_vector(lmax, layer(numoflayers-1)%vel, velocity_ocean_m2)
        print*, 'vector2vxyz...'
        do ir = 1, nn+1
            call vector2vxyz(lmax, velocity_ocean_m2(ir,:), vx(ir,:), vy(ir,:), vz(ir,:))
        end do
        allocate(vxg(size(phi), size(theta), nn+1), vyg(size(phi), size(theta), nn+1), vzg(size(phi), size(theta), nn+1))
        print*, 'scalar2grid...'
        call scalar2grid(lmax2, vx, vxg)
        call scalar2grid(lmax2, vy, vyg)
        call scalar2grid(lmax2, vz, vzg)
        print*, 'writing vtk...'
        ! call write_vector_vtk(time, theta, phi, vxg, vyg, vzg)
        call write_r_theta(time, theta, phi, vxg, vyg, vzg)
    end subroutine ocean_velocity_xyz

    subroutine cleb_1(j1,m1,j2,m2,j,m,cg)
        !
        !      j m
        !     C
        !      j1 m1 1 m2
        !........................................................................
        implicit none
        integer        :: j1, m1, j2, m2, j, m
        real(8) :: cg, c, g, zn

        cg=0.0d0
        if(j2.ne.1) return
        if(abs(j1-j).gt.1.or.(j1+j).eq.0) return
        if(abs(m2).gt.1.or.abs(m1).gt.j1) return

        if(m.ne.(m1+m2)) return
        c=dble(j)
        g=dble(m)
        zn=1.0d0
        
        if(m2.lt.0)then
            if((j1-j).lt.0)then
                cg=(c-g-1.0d0)*(c-g)/(c+c-1.0d0)/(c+c)
            elseif((j1-j).eq.0)then
                cg=(c+g+1.0d0)*(c-g)/(c+1.0d0)/(c+c)
            else
                cg=(c+g+2.0d0)*(c+g+1.0d0)/(c+c+2.0d0)/(c+c+3.0d0)
            end if
        end if

        if(m2.eq.0)then
            if((j1-j).lt.0)then
                cg=(c+g)*(c-g)/(c+c-1.0d0)/c
            elseif((j1-j).eq.0)then
                cg=g/dsqrt(c*(c+1.0d0)); return
            else
                cg=(c+g+1.0d0)*(c-g+1.0d0)/(c+1.0d0)/(c+c+3.0d0)
                zn=-1.0d0
            end if
        end if

        if(m2.gt.0)then
            if((j1-j).lt.0)then
                cg=(c+g-1.0d0)*(c+g)/(c+c-1.0d0)/(c+c)
            elseif((j1-j).eq.0)then
                cg=(c+g)*(c-g+1.0d0)/(c+1.0d0)/(c+c)
                zn=-1.0d0
            else
                cg=(c-g+1.0d0)*(c-g+2.0d0)/(c+c+2.0d0)/(c+c+3.0d0)
            end if
        end if

        cg=zn*dsqrt(cg)
    end subroutine cleb_1


    subroutine write_vector_vtk(time, th, ph, vx, vy, vz)
        implicit none
        real(8) :: time
        integer :: id, ir, ith, iph, nr, nph
        character(300) :: filename
        character(80)  :: line
        real(8)        :: th(:), ph(:), x, y, z, er(3), eth(3), eph(3), r
        real(8)        :: vx(:,:,:), vy(:,:,:), vz(:,:,:)
        real(8), allocatable :: vr(:,:,:), vt(:,:,:), vp(:,:,:)
        real(8), allocatable :: vrxyz(:,:,:,:), vtxyz(:,:,:,:), vpxyz(:,:,:,:)
        id = 10
        nr = size(layer(numoflayers-1)%r_l)
        nph = size(ph)

        print*, size(vx,1), size(vx,2), size(vx,3)
        print*, size(th), size(ph)
        allocate(vr(size(vx,1), size(vx,2), size(vx,3)))
        allocate(vt(size(vx,1), size(vx,2), size(vx,3)))
        allocate(vp(size(vx,1), size(vx,2), size(vx,3)))
        allocate(vrxyz(3, size(vx,1), size(vx,2), size(vx,3)))
        allocate(vtxyz(3, size(vx,1), size(vx,2), size(vx,3)))
        allocate(vpxyz(3, size(vx,1), size(vx,2), size(vx,3)))

        write(filename, '(A,I0,A)') trim(dir_output)//'/on_grid/velocity/velocity_', nint(time), '.vtk'
        open(id, file=filename,form='unformatted',access='stream',status='replace')
        write(id) '# vtk DataFile Version 3.0'//new_line(' ')
        write(id) 'comment (max. 1023 chars)'//new_line(' ')
        write(id) 'BINARY'//new_line(' ')
        write(id) 'DATASET STRUCTURED_GRID'//new_line(' ')
    
        write(line, '(A,3(1X,I0))') 'DIMENSIONS', nph, size(th), nr
        write(id) trim(line)//new_line(' ')
        write(line, '(A,I0,A)') 'POINTS ', (nph)*(size(th))*(nr), ' float'
        write(id) trim(line)//new_line(' ')
        print*, 'min r..', minval(layer(numoflayers-1)%r_l), maxval(layer(numoflayers-1)%r_l)
        
        do ir = 1, nr
            r = layer(numoflayers-1)%r_l(ir)
            do ith = 1, size(th)
                do iph = 1, nph
                    er  = [sin(th(ith)) * cos(ph(iph)), sin(th(ith)) * sin(ph(iph)),  cos(th(ith))]
                    eth = [cos(th(ith)) * cos(ph(iph)), cos(th(ith)) * sin(ph(iph)), -sin(th(ith))]
                    eph = [        - sin(ph(iph)),           cos(ph(iph)),    0.0d0]

                    vr(iph, ith, ir) =  vx(iph, ith, ir) * sin(th(ith))*cos(ph(iph)) + & !vx(e_x.e_r)
                                        vy(iph, ith, ir) * sin(th(ith))*sin(ph(iph)) + & !vy(e_y.e_r)
                                        vz(iph, ith, ir) * cos(th(ith))             !vz(e_z.e_r)
                        
                    vt(iph, ith, ir) =  vx(iph, ith, ir) * cos(th(ith))*cos(ph(iph)) + & !vx(e_x.e_th)
                                        vy(iph, ith, ir) * cos(th(ith))*sin(ph(iph)) - & !vy(e_y.e_th)
                                        vz(iph, ith, ir) * sin(th(ith))             !vz(e_z.e_th)

                    vp(iph, ith, ir) = -vx(iph, ith, ir) * sin(ph(iph)) + & !vx(e_x.e_ph)
                                        vy(iph, ith, ir) * cos(ph(iph))     !vy(e_y.e_ph)

                    x = (r) * sin(th(ith)) * cos(ph(iph))
                    y = (r) * sin(th(ith)) * sin(ph(iph))
                    z = (r) * cos(th(ith))
                    write(id) swapEndian(x),swapEndian(y) ,swapEndian(z)
                    vrxyz(:,iph, ith, ir) = vr(iph, ith, ir)  * er  !vrxyz  = (vr)er
                    vtxyz(:,iph, ith, ir) = vt(iph, ith, ir) * eth !vthxyz = (vth)eth
                    vpxyz(:,iph, ith, ir) = vp(iph, ith, ir) * eph !vphxyz = (vph)eph
                end do
            end do
        end do
        write (line,'(a,i0)')   'POINT_DATA ',  (nph)*(size(th))*(nr)
        write (id) trim(line)//new_line(' ')
        write(line,'(a)')      'VECTORS v float' ! arbitrary name of vector field
        write (id) trim(line)//new_line(' ')
        write (id)  swapEndian(vrxyz + vtxyz + vpxyz)
        write(line,'(a)')      'VECTORS vr float' ! arbitrary name of vector field
        write (id) trim(line)//new_line(' ')
        write (id)  swapEndian(vrxyz)
        write(line,'(a)')      'VECTORS vt float' ! arbitrary name of vector field
        write (id) trim(line)//new_line(' ')
        write (id)  swapEndian(vtxyz)
        ! write(line,'(a)')      'VECTORS vp float' ! arbitrary name of vector field
        ! write (id) trim(line)//new_line(' ')
        ! write (id)  swapEndian(vpxyz)
        close(id)

        contains
    
        ELEMENTAL FUNCTION swapEndian(xx) RESULT (result)
        REAL(8),INTENT(IN) :: xx
        CHARACTER(4) :: result,zz
        CHARACTER(1),DIMENSION(4) :: za
        EQUIVALENCE (zz,za)
        zz=transfer(real(xx,4),zz)
        za(1:4)=za(4:1:-1)
        result=zz
        END FUNCTION
    end subroutine write_vector_vtk

    subroutine write_r_theta(time, th, ph, vx, vy, vz)
        implicit none
        integer :: ir, ith, nr, nth, iph
        real(8) :: r, th_val, x, y, time, th(:), ph(:)
        real(8), dimension(:,:,:), allocatable :: vr, vt, vp
        real(8), dimension(:,:,:)              :: vx, vy, vz
        character(300)                         :: filename
        nth = size(theta)
        nr  = size(layer(numoflayers-1)%r_l)


        allocate(vr(size(vx,1), size(vx,2), size(vx,3)))
        allocate(vt(size(vx,1), size(vx,2), size(vx,3)))
        allocate(vp(size(vx,1), size(vx,2), size(vx,3)))
        iph = 1
        
        write(filename, '(A,I0,A)') trim(dir_output)//'/on_grid/velocity/velocity_', nint(time), '.dat'
        open(12, file=filename)
        do ir = 1, nr
            r = layer(numoflayers-1)%r_l(ir)
            do ith = 1, nth/2
                th_val = th(ith)
                vr(iph, ith, ir) =  vx(iph, ith, ir) * sin(th(ith))*cos(ph(iph)) + & !vx(e_x.e_r)
                                    vy(iph, ith, ir) * sin(th(ith))*sin(ph(iph)) + & !vy(e_y.e_r)
                                    vz(iph, ith, ir) * cos(th(ith))             !vz(e_z.e_r)
                        
                vt(iph, ith, ir) =  vx(iph, ith, ir) * cos(th(ith))*cos(ph(iph)) + & !vx(e_x.e_th)
                                    vy(iph, ith, ir) * cos(th(ith))*sin(ph(iph)) - & !vy(e_y.e_th)
                                    vz(iph, ith, ir) * sin(th(ith))             !vz(e_z.e_th)

                vp(iph, ith, ir) = -vx(iph, ith, ir) * sin(ph(iph)) + & !vx(e_x.e_ph)
                                    vy(iph, ith, ir) * cos(ph(iph))     !vy(e_y.e_ph)
                x = r*cos(th_val)
                y = r*sin(th_val)
                write(12,'(4(F20.8))') x, y, vr(iph, ith, ir), vt(iph, ith, ir)
            end do
        end do
        close(10)
    end subroutine write_r_theta

end module mod_harm2grid