INCLUDE "ChemPlugin.f90"
#ifdef GWB_WINDWOS
INCLUDE "mpi.f90"
#endif

PROGRAM ReactiveTransport
    Use ChemPluginModule
    USE OMP_LIB
    IMPLICIT NONE
    INCLUDE "mpif.h"
    TYPE(ChemPlugin) :: cp_inlet
    TYPE(ChemPlugin),Dimension(:),Allocatable :: cp
    TYPE(CpiLink) :: link0
    INTEGER :: file_unit=20, ios=0, error, nx,&
            i, scope = 0, tid, l
    REAL(8) :: time_end, delta_years, next_output, &
               length, deltax, deltay, deltaz, &
               veloc_in, velocity, flow, porosity, &
               diffcoef, dispersivity, dispcoef, trans,&
               tcond, ttrans, deltat, deltat_world
    CHARACTER(LEN=500) :: input_file, line, cmd, veloc_str
    INTEGER id, ierr, world_size, world_rank, chunk_size, argc
    
    argc = COMMAND_ARGUMENT_COUNT()
    
    CALL MPI_Init( ierr )
    CALL MPI_Comm_size(MPI_COMM_WORLD, world_size, ierr)
    CALL MPI_Comm_rank(MPI_COMM_WORLD, world_rank, ierr)

    IF (world_rank == 0) THEN
       WRITE(*,'("")')
       WRITE(*,'("Model reactive transport in one dimension using MPI and OPENMP")')
       WRITE(*,'("")')
    END IF

    nx = 4000 ! number of instances along x
    length = 100 ! m
    deltax = length / nx ! m
    deltay = 1.0  ! m
    deltaz = 1.0  ! m
    time_end = 10.0  ! years
    delta_years = time_end / 5 ! years
    next_output = 0
   
    ! Set up our omp step/schedule size so that
    ! actual work is split evenly among the omp threads.
    chunk_size = nx / world_size / omp_get_max_threads()

    ! Create the inlet and interior instances
    cp_inlet = ChemPlugin()
    ALLOCATE(cp(nx))
    !$OMP PARALLEL DO SCHEDULE(static,chunk_size)
    Do i = 1, nx
       cp(i) = ChemPlugin()
    END DO
    !$OMP END PARALLEL DO
    
    ! Assign instances to ranks
    error = cp_inlet%MpiAssign(0)
    !$OMP PARALLEL DO SCHEDULE(static,chunk_size)
    Do i = 1, nx
       error = cp(i)%MpiAssign((i-1) * world_size / nx)
    END DO
    !$OMP END PARALLEL DO

    error = cp_inlet%Console("stdout")
    error = cp(1)%Console("stdout")
    error = cp(1)%Config("pluses = banner")

    WRITE(cmd, '("volume =", F6.2, " m3; ", "time end =", F5.1, " years;")') deltax * deltay * deltaz, time_end

    !$OMP PARALLEL DO REDUCTION(+: error) SCHEDULE(static,chunk_size)
    Do i = 1, nx
       error = cp(i)%Config(cmd)
    END DO
    !$OMP END PARALLEL DO

    !Open input file
    IF ( argc < 1 ) THEN
       WRITE(*, '("You must specify an input file as the first program argument")')
       CALL exit_client(-1)
    END IF
    CALL GET_COMMAND_ARGUMENT( 1, input_file )
    OPEN(unit=file_unit, file=TRIM(input_file), action="read")
    Do WHILE(ios == 0) 
       READ(file_unit,'(A)',IOSTAT=ios) line
       IF (line == 'go') THEN
           EXIT
       ELSE IF (line == 'scope inlet') THEN
           scope = 1
       ELSE IF (line == 'scope initial') THEN
           scope = 2
       ELSE IF (scope == 1) THEN
           error = cp_inlet%Config(line)
       ELSE IF (scope == 2) THEN
           !$OMP PARALLEL DO REDUCTION(+: error) SCHEDULE(static,chunk_size)
           DO i = 1, nx
              error = cp(i)%Config(line)
           END DO
           !$OMP END PARALLEL DO
       END IF
    END DO
    CLOSE(file_unit) 
    ! Initialize the inlet and interior instances; write out initial conditions
    IF (cp_inlet%Initialize() /= 0) THEN
        WRITE(*,*) "Inlet failed to initialize"
        CALL exit_client(-1)
    END IF
    error = cp_inlet%PrintOutput("RTM.txt", "Inlet fluid")
    error = 0

    !$OMP PARALLEL DO REDUCTION(+: error) SCHEDULE(static,chunk_size)
    DO i = 1, nx
       IF (cp(i)%Initialize() /= 0) THEN
          error = error + 1
          WRITE(*, '("Instance", I2, " failed to initalize")') i
       END IF
    END DO
    !$OMP END PARALLEL DO
    If (error /= 0) THEN
        CALL exit_client(-1)
    END IF

    call write_results(cp, nx, delta_years, next_output)
    
    ! set velocity; calculate flow rate and transmissivities
    IF ( argc < 2 ) THEN
       WRITE(*, '("You must specify fluid velocity in m/yr as the second program argument.")')
       CALL exit_client(-1)
    END IF
    CALL GET_COMMAND_ARGUMENT( 2, veloc_str )
    READ(veloc_str,*) veloc_in
    velocity = veloc_in / 31557600 ! m/s

    porosity = cp(1)%MpiReport1("porosity")
    flow = deltay * deltaz * porosity * velocity ! m3/s

    diffcoef = 1d-10 ! m2/s
    dispersivity = 1.0 ! m
    dispcoef = velocity * dispersivity + diffcoef ! m2/s
    trans = deltay * deltaz * porosity * dispcoef / deltax  ! m3/s
    
    tcond = 2.0 ! W/m/K
    ttrans = deltay * deltaz * tcond / deltax

    ! Link the instances
    link0 = cp(1)%Link(cp_inlet)
    error = link0%Transmissivity(trans, "m3/s")
    error = link0%HeatTrans(ttrans, "W/K")
    error = link0%FlowRate(flow, "m3/s")
    DO i = 2, nx
       link0 = cp(i)%Link(cp(i-1))
       error = link0%Transmissivity(trans, "m3/s")
       error = link0%HeatTrans(ttrans, "W/K")
       error = link0%FlowRate(flow, "m3/s")
    END DO
    link0 = cp(nx)%Link()
    error = link0%FlowRate(-flow, "m3/s")

    DO WHILE(.TRUE.)
       error = 0
       deltat = 1e10

       !$OMP PARALLEL DO REDUCTION(min:deltat) SCHEDULE(static,chunk_size)
       DO i = 1, nx
          deltat = min(deltat, cp(i)%ReportTimeStep())
       END DO
       !$OMP END PARALLEL DO
       
       CALL MPI_Allreduce( deltat, deltat_world, 1, MPI_DOUBLE, MPI_MIN, MPI_COMM_WORLD, ierr)

       !$OMP PARALLEL DO REDUCTION(+:error) SCHEDULE(static,chunk_size)
       DO i = 1, nx
          IF (cp(i)%AdvanceTimeStep(deltat_world) /= 0) THEN
             error = error + 1
          END IF
       END DO  
       !$OMP END PARALLEL DO
       IF (error /= 0) THEN
          CALL exit_client(0)
       END IF

       DO i = 1, nx
          DO l = 0, cp(i)%nLinks() - 1
             IF (MpiUpdateLink(cp(i)%Link(l)) /= 0) THEN
                error = error + 1
             END IF
          END DO
       END DO
       IF (error /= 0) THEN
          CALL exit_client(-1)
       END IF
       
       !$OMP PARALLEL DO REDUCTION(+:error) SCHEDULE(static,chunk_size)
       DO i = 1, nx
          IF (cp(i)%AdvanceTransport() /= 0) THEN
              error = error + 1
          END IF
       END DO
       !$OMP END PARALLEL DO
       IF (error /= 0) THEN
          CALL exit_client(-1)
       END IF
      
       !$OMP PARALLEL DO REDUCTION(+:error) SCHEDULE(static,chunk_size)
       DO i = 1, nx
          IF (cp(i)%AdvanceHeatTransport() /= 0) THEN
              error = error + 1
          END IF
       END DO
       !$OMP END PARALLEL DO
       IF (error /= 0) THEN
          CALL exit_client(-1)
       END IF

       !$OMP PARALLEL DO REDUCTION(+:error) SCHEDULE(dynamic)
       DO i = 1, nx
          IF (cp(i)%AdvanceChemical() /= 0) THEN
              error = error + 1
          END IF
       END DO
       !$OMP END PARALLEL DO
       IF (error /= 0) THEN
          CALL exit_client(-1)
       END IF
       
       call write_results(cp, nx, delta_years, next_output)
    END DO
     
     
      
CONTAINS 
        SUBROUTINE write_results(cp, nx, gap, next_output) 
            IMPLICIT NONE 
            TYPE(ChemPlugin),dimension(:), INTENT(in) :: cp
            INTEGER, INTENT(in) :: nx
            REAL(8), INTENT(in) :: gap
            REAL(8), INTENT(inout) :: next_output
            CHARACTER(LEN=12) :: label
            INTEGER :: i , error, last_rank, ierr, step
            REAL(8) :: now

            now = cp(1)%MpiReport1("Time", "years")
            
            IF ((next_output - now) <= gap/1d4) THEN
                last_rank = cp(1)%MpiRank()
                step = nx / 100
                Do i = 1, nx, +step
                   IF ( cp(i)%MpiRank() /= last_rank ) THEN
                      last_rank = cp(i)%MpiRank()
                      CALL MPI_Barrier(MPI_COMM_WORLD,ierr)
                   END IF
                   WRITE(label, '("Instance ",I3)') i-1
                   error = cp(i)%PrintOutput("RTM.txt", label)
                END DO
                next_output = next_output + gap
            END IF
         
        END SUBROUTINE write_results

        SUBROUTINE exit_client(exit_status)
           IMPLICIT NONE
           INTEGER, INTENT(in) :: exit_status
           IF (exit_status /= 0) THEN
              CALL MPI_Abort(MPI_COMM_WORLD, ierr)
              STOP -1
           ELSE
              CALL MPI_Finalize(ierr)
              STOP
           END IF
        END SUBROUTINE

END PROGRAM ReactiveTransport