MODULE compute_geometry_ico_mod
   USE geometry
   USE dimensions

   USE domain_mod, ONLY: t_domain, &
                         domain, ndomain, assigned_domain, &
                         domain_glo, ndomain_glo, domloc_glo_ind, swap_needed
   USE omp_para, ONLY: is_omp_level_master, is_master
   USE transfert_mod, ONLY: req_i0, req_i1, t_message, transfert_request, transfert_message, init_message

   USE spherical_geom_mod, ONLY: xyz2lonlat, circumcenter, compute_centroid, &
                                 surf_triangle, dist_cart, div_arc_bis, schmidt_transform
   USE vector, ONLY: norm, cross_product2

   IMPLICIT NONE
   PRIVATE
   SAVE

   PUBLIC :: set_geometry, update_domain

CONTAINS

   SUBROUTINE update_circumcenters
      REAL(rstd) :: x1(3), x2(3)
      REAL(rstd) :: vect(3, 6)
      REAL(rstd) :: centr(3)
      INTEGER :: ind, i, j, n, k
      TYPE(t_message), SAVE :: message0, message1
      LOGICAL, SAVE :: first = .TRUE.
!$OMP THREADPRIVATE(first)

      IF (first) THEN
         CALL init_message(geom%xyz_i, req_i0, message0)
         CALL init_message(geom%xyz_i, req_i1, message1)
         first = .FALSE.
      END IF

      CALL transfert_message(geom%xyz_i, message0)
      CALL transfert_message(geom%xyz_i, message1)

      DO ind = 1, ndomain

         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         DO j = jj_begin, jj_end
            DO i = ii_begin, ii_end
               n = (j - 1)*iim + i
               DO k = 0, 5
                  x1(:) = xyz_i(n + t_pos(k + 1), :)
                  x2(:) = xyz_i(n + t_pos(MOD(k + 1, 6) + 1), :)
                  if (norm(x1 - x2) < 1e-16) x2(:) = xyz_i(n + t_pos(MOD(k + 2, 6) + 1), :)
                  CALL circumcenter(xyz_i(n, :), x1, x2, xyz_v(n + z_pos(k + 1), :))
               END DO
            END DO
         END DO
      END DO

   END SUBROUTINE update_circumcenters

   SUBROUTINE remap_schmidt_loc
      USE math_const, ONLY: pi
      USE getin_mod, ONLY: getin
      INTEGER :: ind, i, j, n
      REAL(rstd) :: schmidt_factor, schmidt_lon, schmidt_lat

      ! Schmidt transform parameters
      schmidt_factor = 1.
      CALL getin('schmidt_factor', schmidt_factor)
      schmidt_factor = schmidt_factor**2.

      schmidt_lon = 0.
      CALL getin('schmidt_lon', schmidt_lon)
      schmidt_lon = schmidt_lon*pi/180.

      schmidt_lat = 45.
      CALL getin('schmidt_lat', schmidt_lat)
      schmidt_lat = schmidt_lat*pi/180.

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         DO j = jj_begin, jj_end
            DO i = ii_begin, ii_end
               n = (j - 1)*iim + i
               CALL schmidt_transform(xyz_i(n, :), schmidt_factor, schmidt_lon, schmidt_lat)
            END DO
         END DO
      END DO
   END SUBROUTINE remap_schmidt_loc

   SUBROUTINE optimize_geometry
      USE getin_mod, ONLY: getin
      INTEGER :: nb_it = 0
      TYPE(t_domain), POINTER :: d
      INTEGER :: ind, it, i, j, n, k
      REAL(rstd) :: x1(3), x2(3)
      REAL(rstd) :: vect(3, 6)
      REAL(rstd) :: centr(3)
      REAL(rstd) :: sum
      LOGICAL    :: check

      CALL getin('optim_it', nb_it)

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
         d => domain(ind)
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         xyz_i(:, 1) = 0; xyz_i(:, 2) = 0; xyz_i(:, 3) = 1

         DO j = jj_begin, jj_end
            DO i = ii_begin, ii_end
               n = (j - 1)*iim + i
               xyz_i(n, :) = d%xyz(:, i, j)
            END DO
         END DO
      END DO

      CALL update_circumcenters

!ym you get discrepency if you try to read metric from file, vertex are not saved in file
!    DO ind=1,ndomain
!      IF (.NOT. assigned_domain(ind)  .OR. .NOT. is_omp_level_master ) CYCLE
!      d=>domain(ind)
!      CALL swap_dimensions(ind)
!      CALL swap_geometry(ind)
!      DO j=jj_begin,jj_end
!        DO i=ii_begin,ii_end
!          n=(j-1)*iim+i
!          DO k=0,5
!            x1(:) = xyz_v(n+z_pos(k+1),:)
!            x2(:) = d%vertex(:,k,i,j)
!            IF (norm(x1-x2)>1e-10) THEN
!              PRINT*,"vertex diff ",ind,i,j,k
!              PRINT*,x1
!              PRINT*,x2
!            ENDIF
!          ENDDO
!        ENDDO
!      ENDDO
!   ENDDO

      DO it = 1, nb_it
         IF (MOD(it, 100) == 0) THEN
            check = is_master
         ELSE
            check = .FALSE.
         END IF

         sum = 0
         DO ind = 1, ndomain
            IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
            CALL swap_dimensions(ind)
            CALL swap_geometry(ind)
            DO j = jj_begin, jj_end
               DO i = ii_begin, ii_end
                  n = (j - 1)*iim + i
                  vect(:, 1) = xyz_v(n + z_rup, :)
                  vect(:, 2) = xyz_v(n + z_up, :)
                  vect(:, 3) = xyz_v(n + z_lup, :)
                  vect(:, 4) = xyz_v(n + z_ldown, :)
                  vect(:, 5) = xyz_v(n + z_down, :)
                  vect(:, 6) = xyz_v(n + z_rdown, :)
                  CALL compute_centroid(vect, 6, centr)
                  IF (check) THEN
                     sum = MAX(sum, norm(xyz_i(n, :) - centr(:)))
                  END IF
                  xyz_i(n, :) = centr(:)
               END DO
            END DO

         END DO

         IF (check) PRINT *, "it = ", it, "  diff centroid circumcenter ", sum

         CALL update_circumcenters

      END DO

   END SUBROUTINE optimize_geometry

   SUBROUTINE update_domain
      ! copy position of generators and vertices back into domain(:)%xyz/vertex
      ! so that XIOS/create_header_gen gets the right values
      USE transfert_mod, ONLY: gather_field, bcast_field

      INTEGER :: ind, i, j, k, n
      TYPE(t_domain), POINTER :: d
      TYPE(t_field), POINTER, SAVE :: xyz_glo(:), xyz_loc(:), vertex_glo(:), vertex_loc(:)
      REAL(rstd), POINTER :: xyz(:, :), vertex(:, :)

      CALL allocate_field(xyz_loc, field_t, type_real, 3)
      CALL allocate_field(vertex_loc, field_z, type_real, 3)

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         xyz = xyz_loc(ind)
         xyz(:, :) = xyz_i(:, :)
         vertex = vertex_loc(ind)
         vertex(:, :) = xyz_v(:, :)
      END DO

!$OMP BARRIER
!$OMP MASTER
      CALL allocate_field_glo(xyz_glo, field_t, type_real, 3)
      CALL allocate_field_glo(vertex_glo, field_z, type_real, 3)

      CALL gather_field(xyz_loc, xyz_glo)
      CALL gather_field(vertex_loc, vertex_glo)

      CALL bcast_field(xyz_glo)
      CALL bcast_field(vertex_glo)

      DO ind = 1, ndomain_glo
         d => domain_glo(ind)
         xyz = xyz_glo(ind)
         vertex = vertex_glo(ind)
         DO j = d%jj_begin, d%jj_end
            DO i = d%ii_begin, d%ii_end
               n = (j - 1)*d%iim + i
               d%xyz(:, i, j) = xyz(n, :)
               DO k = 0, 5
                  d%vertex(:, k, i, j) = vertex(n + d%z_pos(k + 1), :)
               END DO
            END DO
         END DO
      END DO

      CALL deallocate_field_glo(vertex_glo)
      CALL deallocate_field_glo(xyz_glo)
!$OMP END MASTER
!$OMP BARRIER
      CALL deallocate_field(vertex_loc)
      CALL deallocate_field(xyz_loc)

   END SUBROUTINE update_domain

   SUBROUTINE set_geometry
      USE metric
      USE grid_param, ONLY: iim_glo, jjm_glo
      USE earth_const, ONLY: omega, radius
      REAL(rstd) :: surf(6)
      REAL(rstd) :: surf_v(6)
      REAL(rstd) :: vect(3, 6)
      REAL(rstd) :: centr(3)
      REAL(rstd) :: vet(3), vep(3), vertex(3)
      INTEGER :: ind, i, j, k, n
      TYPE(t_domain), POINTER :: d
      REAL(rstd) ::  S12
      REAL(rstd) :: w(6)
      REAL(rstd) :: lon, lat
      INTEGER :: ii_glo, jj_glo
      REAL(rstd) :: S

      CALL optimize_geometry
      CALL remap_schmidt_loc
      CALL update_circumcenters
      ! copy position of generators and vertices back into domain_glo(:)%xyz/vertex
      ! so that write_field gets the right values
      CALL update_domain

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
         d => domain(ind)
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         lon_i(:) = 0; lat_i(:) = 0
         lon_e(:) = 0; lat_e(:) = 0
         DO j = jj_begin - 1, jj_end + 1
            DO i = ii_begin - 1, ii_end + 1
               n = (j - 1)*iim + i

               DO k = 0, 5
                  ne(n, k + 1) = d%ne(k, i, j)
               END DO

               vect(:, 1) = xyz_v(n + z_rup, :)
               vect(:, 2) = xyz_v(n + z_up, :)
               vect(:, 3) = xyz_v(n + z_lup, :)
               vect(:, 4) = xyz_v(n + z_ldown, :)
               vect(:, 5) = xyz_v(n + z_down, :)
               vect(:, 6) = xyz_v(n + z_rdown, :)
               CALL compute_centroid(vect, 6, centr)
               centroid(n, :) = centr(:)

               CALL xyz2lonlat(xyz_v(n + z_up, :), lon, lat)
               fv(n + z_up) = 2*sin(lat)*omega
               CALL xyz2lonlat(xyz_v(n + z_down, :), lon, lat)
               fv(n + z_down) = 2*sin(lat)*omega

               CALL dist_cart(xyz_i(n, :), xyz_i(n + t_right, :), de(n + u_right))
               CALL dist_cart(xyz_i(n, :), xyz_i(n + t_lup, :), de(n + u_lup))
               CALL dist_cart(xyz_i(n, :), xyz_i(n + t_ldown, :), de(n + u_ldown))

               CALL div_arc_bis(xyz_i(n, :), xyz_i(n + t_right, :), 0.5, xyz_e(n + u_right, :))
               CALL div_arc_bis(xyz_i(n, :), xyz_i(n + t_lup, :), 0.5, xyz_e(n + u_lup, :))
               CALL div_arc_bis(xyz_i(n, :), xyz_i(n + t_ldown, :), 0.5, xyz_e(n + u_ldown, :))

               CALL dist_cart(xyz_v(n + z_rdown, :), xyz_v(n + z_rup, :), le(n + u_right))
               CALL dist_cart(xyz_v(n + z_up, :), xyz_v(n + z_lup, :), le(n + u_lup))
               CALL dist_cart(xyz_v(n + z_ldown, :), xyz_v(n + z_down, :), le(n + u_ldown))

               le_de(n + u_right) = le(n + u_right)/de(n + u_right) ! NaN possible but should be harmless
               le_de(n + u_lup) = le(n + u_lup)/de(n + u_lup)
               le_de(n + u_ldown) = le(n + u_ldown)/de(n + u_ldown)

               Ai(n) = 0
               DO k = 0, 5
                  CALL surf_triangle(xyz_i(n, :), xyz_i(n + t_pos(k + 1), :), xyz_i(n + t_pos(MOD((k + 1 + 6), 6) + 1), :), &
                                     surf_v(k + 1))
                  CALL surf_triangle(xyz_i(n, :), xyz_v(n + z_pos(MOD((k - 1 + 6), 6) + 1), :), xyz_v(n + z_pos(k + 1), :), &
                                     surf(k + 1))
                  Ai(n) = Ai(n) + surf(k + 1)
                  IF (i == ii_end .AND. j == jj_begin) THEN
                     IF (Ai(n) < 1e20) THEN
                     ELSE
                        PRINT *, "PB !!", Ai(n), k, surf(k + 1)
                        PRINT *, xyz_i(n, :), xyz_v(n + z_pos(MOD((k - 1 + 6), 6) + 1), :), xyz_v(n + z_pos(k + 1), :)
                     END IF
                  END IF
               END DO

               ! Sign convention : Ringler et al., JCP 2010, eq. 21 p. 3071
               ! Normal component is along outgoing normal vector if ne=1

               CALL cross_product2(xyz_v(n + z_rdown, :), xyz_v(n + z_rup, :), vep)
               IF (norm(vep) > 1e-30) THEN
                  vep(:) = vep(:)/norm(vep)                         ! Inward normal vector
                  CALL cross_product2(vep, xyz_e(n + u_right, :), vet) ! Counter-clockwise tangent vector
                  vet(:) = vet(:)/norm(vet)
                  ep_e(n + u_right, :) = -vep(:)*ne(n, right)
                  et_e(n + u_right, :) = vet(:)*ne(n, right)
               END IF

               CALL cross_product2(xyz_v(n + z_up, :), xyz_v(n + z_lup, :), vep)
               IF (norm(vep) > 1e-30) THEN
                  vep(:) = vep(:)/norm(vep)
                  CALL cross_product2(vep, xyz_e(n + u_lup, :), vet)
                  vet(:) = vet(:)/norm(vet)
                  ep_e(n + u_lup, :) = -vep(:)*ne(n, lup)
                  et_e(n + u_lup, :) = vet(:)*ne(n, lup)
               END IF

               CALL cross_product2(xyz_v(n + z_ldown, :), xyz_v(n + z_down, :), vep)
               IF (norm(vep) > 1e-30) THEN
                  vep(:) = vep(:)/norm(vep)
                  CALL cross_product2(vep, xyz_e(n + u_ldown, :), vet)
                  vet(:) = vet(:)/norm(vet)
                  ep_e(n + u_ldown, :) = -vep(:)*ne(n, ldown)
                  et_e(n + u_ldown, :) = vet(:)*ne(n, ldown)
               END IF

               CALL xyz2lonlat(xyz_i(n, :), lon, lat)
               lon_i(n) = lon
               lat_i(n) = lat
               elon_i(n, 1) = -sin(lon)
               elon_i(n, 2) = cos(lon)
               elon_i(n, 3) = 0
               elat_i(n, 1) = -cos(lon)*sin(lat)
               elat_i(n, 2) = -sin(lon)*sin(lat)
               elat_i(n, 3) = cos(lat)

               CALL xyz2lonlat(xyz_e(n + u_right, :), lon, lat)
               lon_e(n + u_right) = lon
               lat_e(n + u_right) = lat
               elon_e(n + u_right, 1) = -sin(lon)
               elon_e(n + u_right, 2) = cos(lon)
               elon_e(n + u_right, 3) = 0
               elat_e(n + u_right, 1) = -cos(lon)*sin(lat)
               elat_e(n + u_right, 2) = -sin(lon)*sin(lat)
               elat_e(n + u_right, 3) = cos(lat)

               CALL xyz2lonlat(xyz_e(n + u_lup, :), lon, lat)
               lon_e(n + u_lup) = lon
               lat_e(n + u_lup) = lat
               elon_e(n + u_lup, 1) = -sin(lon)
               elon_e(n + u_lup, 2) = cos(lon)
               elon_e(n + u_lup, 3) = 0
               elat_e(n + u_lup, 1) = -cos(lon)*sin(lat)
               elat_e(n + u_lup, 2) = -sin(lon)*sin(lat)
               elat_e(n + u_lup, 3) = cos(lat)

               CALL xyz2lonlat(xyz_e(n + u_ldown, :), lon, lat)
               lon_e(n + u_ldown) = lon
               lat_e(n + u_ldown) = lat
               elon_e(n + u_ldown, 1) = -sin(lon)
               elon_e(n + u_ldown, 2) = cos(lon)
               elon_e(n + u_ldown, 3) = 0
               elat_e(n + u_ldown, 1) = -cos(lon)*sin(lat)
               elat_e(n + u_ldown, 2) = -sin(lon)*sin(lat)
               elat_e(n + u_ldown, 3) = cos(lat)

               DO k = 0, 5
                  CALL surf_triangle(xyz_i(n, :), xyz_v(n + z_pos(k + 1), :), xyz_i(n + t_pos(k + 1), :), &
                                     S1(n, k + 1))
                  CALL surf_triangle(xyz_i(n, :), xyz_v(n + z_pos(k + 1), :), xyz_i(n + t_pos(MOD(k + 1 + 6, 6) + 1), :), &
                                     S2(n, k + 1))
                  S12 = .5*(S1(n, k + 1) + S2(n, k + 1))
                  Riv(n, k + 1) = S12/Ai(n)
                  Riv2(n, k + 1) = S12/surf_v(k + 1)
               END DO

               DO k = 1, 6
                  IF (ABS(surf_v(k)) < 1e-30) THEN
                     Riv(n, k) = 0.
                  END IF
               END DO

               Av(n + z_up) = surf_v(vup) + 1e-100
               Av(n + z_down) = surf_v(vdown) + 1e-100

            END DO
         END DO

         DO j = jj_begin, jj_end
            DO i = ii_begin, ii_end
               n = (j - 1)*iim + i

               CALL compute_wee(n, right, w)
               wee(n + u_right, 1:5) = w(1:5)
               CALL compute_wee(n + t_right, left, w)
               wee(n + u_right, 6:10) = w(1:5)

               CALL compute_wee(n, lup, w)
               wee(n + u_lup, 1:5) = w(1:5)
               CALL compute_wee(n + t_lup, rdown, w)
               wee(n + u_lup, 6:10) = w(1:5)

               CALL compute_wee(n, ldown, w)
               wee(n + u_ldown, 1:5) = w(1:5)
               CALL compute_wee(n + t_ldown, rup, w)
               wee(n + u_ldown, 6:10) = w(1:5)

            END DO
         END DO

         DO j = jj_begin, jj_end
            DO i = ii_begin, ii_end
               n = (j - 1)*iim + i
               ii_glo = d%ii_begin_glo - d%ii_begin + i
               jj_glo = d%jj_begin_glo - d%jj_begin + j

               IF (ii_glo == 1 .AND. jj_glo == 1) THEN
                  le(n + u_ldown) = 0
                  le_de(n + u_ldown) = 0
                  xyz_v(n + z_ldown, :) = xyz_v(n + z_down, :)

               END IF

               IF (ii_glo == iim_glo .AND. jj_glo == 1) THEN
                  le(n + u_right) = 0
                  le_de(n + u_right) = 0
                  xyz_v(n + z_rdown, :) = xyz_v(n + z_rup, :)
               END IF

               IF (ii_glo == iim_glo .AND. jj_glo == jjm_glo) THEN
                  le(n + u_rup) = 0
                  le_de(n + u_rup) = 0
                  xyz_v(n + z_rup, :) = xyz_v(n + z_up, :)
               END IF

               IF (ii_glo == 1 .AND. jj_glo == jjm_glo) THEN
                  le(n + u_lup) = 0
                  le_de(n + u_lup) = 0
                  xyz_v(n + z_up, :) = xyz_v(n + z_lup, :)
               END IF

            END DO
         END DO

         DO j = jj_begin - 1, jj_end + 1
            DO i = ii_begin - 1, ii_end + 1
               n = (j - 1)*iim + i
               xyz_i(n, :) = xyz_i(n, :)*radius
               xyz_v(n + z_up, :) = xyz_v(n + z_up, :)*radius
               xyz_v(n + z_down, :) = xyz_v(n + z_down, :)*radius
               de(n + u_right) = de(n + u_right)*radius
               de(n + u_lup) = de(n + u_lup)*radius
               de(n + u_ldown) = de(n + u_ldown)*radius
               xyz_e(n + u_right, :) = xyz_e(n + u_right, :)*radius
               xyz_e(n + u_lup, :) = xyz_e(n + u_lup, :)*radius
               xyz_e(n + u_ldown, :) = xyz_e(n + u_ldown, :)*radius
               le(n + u_right) = le(n + u_right)*radius
               le(n + u_lup) = le(n + u_lup)*radius
               le(n + u_ldown) = le(n + u_ldown)*radius
               Ai(n) = Ai(n)*radius**2
               Av(n + z_up) = Av(n + z_up)*radius**2
               Av(n + z_down) = Av(n + z_down)*radius**2
            END DO
         END DO

      END DO

      CALL transfert_request(geom%Ai, req_i1)
      CALL transfert_request(geom%centroid, req_i1)

!    CALL surf_triangle(d%xyz(:,ii_begin,jj_begin),d%xyz(:,ii_begin,jj_end),d%xyz(:,ii_end,jj_begin),S)

   END SUBROUTINE set_geometry

   SUBROUTINE compute_wee(n, pos, w)
      IMPLICIT NONE
      INTEGER, INTENT(IN) :: n
      INTEGER, INTENT(IN) :: pos
      REAL(rstd), INTENT(OUT) ::w(6)

      REAL(rstd) :: ne_(0:5)
      REAL(rstd) :: Riv_(6)
      INTEGER    :: k

      DO k = 0, 5
         ne_(k) = ne(n, MOD(pos - 1 + k + 6, 6) + 1)
         Riv_(k + 1) = Riv(n, MOD(pos - 1 + k + 6, 6) + 1)
      END DO

      w(1) = -ne_(0)*ne_(1)*(Riv_(1) - 0.5)
      w(2) = -ne_(2)*(ne_(0)*Riv_(2) - w(1)*ne_(1))
      w(3) = -ne_(3)*(ne_(0)*Riv_(3) - w(2)*ne_(2))
      w(4) = -ne_(4)*(ne_(0)*Riv_(4) - w(3)*ne_(3))
      w(5) = -ne_(5)*(ne_(0)*Riv_(5) - w(4)*ne_(4))
      w(6) = ne_(0)*ne_(5)*(Riv_(6) - 0.5)

!    IF ( ABS(w(5)-w(6))>1e-20) PRINT *, "pb pour wee : w(5)!=w(6)",sum(Riv_(:))

   END SUBROUTINE compute_wee

END MODULE compute_geometry_ico_mod
