MODULE moduleProbe
  USE moduleSpecies

  TYPE:: probeDistFunc
    INTEGER:: id
    REAL(8), DIMENSION(1:3):: r
    CLASS(speciesGeneric), POINTER:: species
    REAL(8), ALLOCATABLE, DIMENSION(:):: vi, vj, vk
    REAL(8), ALLOCATABLE, DIMENSION(:,:,:):: f
    INTEGER, DIMENSION(1:3):: nv
    REAL(8), DIMENSION(1:3):: vrange
    REAL(8):: dvInv
    INTEGER:: every
    LOGICAL:: update
    INTEGER(KIND=OMP_LOCK_KIND):: lock
    REAL(8):: maxR
    CONTAINS
      PROCEDURE, PASS:: init
      PROCEDURE, PASS:: findLowerIndex
      PROCEDURE, PASS:: calculate
      PROCEDURE, PASS:: output
      
  END TYPE probeDistFunc

  INTEGER:: nProbes = 0
  TYPE(probeDistFunc), ALLOCATABLE:: probe(:)

  CONTAINS
    !Functions for probeDistFunc type
    SUBROUTINE init(self, id, speciesName, r, v1, v2, v3, points, everyTimeStep)
      USE moduleCaseParam
      USE moduleRefParam
      USE moduleSpecies
      USE moduleMesh
      USE moduleErrors
      IMPLICIT NONE

      CLASS(probeDistFunc), INTENT(out):: self
      INTEGER, INTENT(in):: id
      CHARACTER(:), ALLOCATABLE, INTENT(in):: speciesName
      REAL(8), INTENT(in):: r(1:3)
      REAL(8), INTENT(in):: v1(1:2), v2(1:2), v3(1:2)
      INTEGER, INTENT(in):: points(1:3)
      REAL(8), INTENT(in):: everyTimeStep
      INTEGER:: sp, i
      REAL(8):: dv(1:3)

      !Assign id
      self%id = id

      !Get species
      sp = speciesName2Index(speciesName)
      self%species => species(sp)%obj

      !Find cell
      self%r = r/L_ref

      !Allocates velocity grid
      self%nv = points
      ALLOCATE(self%vi(1:self%nv(1)))
      ALLOCATE(self%vj(1:self%nv(2)))
      ALLOCATE(self%vk(1:self%nv(3)))

      self%vrange(1) = (v1(2) - v1(1))/v_ref
      self%vrange(2) = (v2(2) - v2(1))/v_ref
      self%vrange(3) = (v3(2) - v3(1))/v_ref

      !Creates grid
      dv = self%vrange / REAL(self%nv - 1)

      self%dvInv = 1.D0/(dv(1)*dv(2)*dv(3))

      DO i = 1, self%nv(1)
        self%vi(i) = dv(1)*REAL(i - 1) + v1(1)/v_ref

      END DO

      DO i = 1, self%nv(2)
        self%vj(i) = dv(2)*REAL(i - 1) + v2(1)/v_ref

      END DO

      DO i = 1, self%nv(3)
        self%vk(i) = dv(3)*REAL(i - 1) + v3(1)/v_ref

      END DO

      !Allocates distribution function
      ALLOCATE(self%f(1:self%nv(1), &
                      1:self%nv(2), &
                      1:self%nv(3)))

      !Number of iterations between output
      IF (everyTimeStep == 0.D0) THEN
        self%every = 1

      ELSE
        self%every = NINT(everyTimeStep/ tauMin / ti_ref)

      END IF
      
      !Maximum radius
      !TODO: Make this an input parameter
      self%maxR = 1.D-2/L_ref

      !Init the probe lock
      CALL OMP_INIT_LOCK(self%lock)

    END SUBROUTINE init

    SUBROUTINE findLowerIndex(self, vp, i, j, k, inside)
      IMPLICIT NONE
      
      CLASS(probeDistFunc), INTENT(in):: self
      REAL(8), INTENT(in):: vp(1:3)
      INTEGER, INTENT(out):: i, j, k
      LOGICAL, INTENT(out):: inside

      inside = .TRUE.

      i = FLOOR((vp(1) - self%vi(1))/self%vrange(1)*(REAL(self%nv(1) - 1)) + 1.D0)
      IF (i >= self%nv(1) .OR. i <  1) inside = .FALSE.
      j = FLOOR((vp(2) - self%vj(1))/self%vrange(2)*(REAL(self%nv(2) - 1)) + 1.D0)
      IF (j >= self%nv(2) .OR. j <  1) inside = .FALSE.
      k = FLOOR((vp(3) - self%vk(1))/self%vrange(3)*(REAL(self%nv(3) - 1)) + 1.D0)
      IF (k >= self%nv(3) .OR. k <  1) inside = .FALSE.

    END SUBROUTINE findLowerIndex

    SUBROUTINE calculate(self, part)
      USE moduleSpecies
      USE moduleList
      IMPLICIT NONE

      CLASS(probeDistFunc), INTENT(inout):: self
      TYPE(particle), INTENT(in):: part
      REAL(8):: deltaR
      INTEGER:: i, j, k
      LOGICAL:: inside
      REAL(8):: weight
      REAL(8):: fi, fi1
      REAL(8):: fj, fj1
      REAL(8):: fk, fk1

      !If particle is of desired type, include in the distribution function
      IF (part%species%n == self%species%n) THEN
        !Calculate distance between particle and probe
        deltaR = NORM2(self%r - part%r)

        !Only include particle if it is inside the maximum radius
        ! IF (deltaR < self%maxR) THEN
          !find lower index for all dimensions
          CALL self%findLowerIndex(part%v, i, j, k, inside)

          !If particle is inside the velocity grid, add it to the distribution function
          IF (inside) THEN
            !Calculate weights
            fi  = self%vi(i+1) - part%v(1)
            fi1 = part%v(1)    - self%vi(i)
            fj  = self%vj(j+1) - part%v(2)
            fj1 = part%v(2)    - self%vj(j)
            fk  = self%vk(k+1) - part%v(3)
            fk1 = part%v(3)    - self%vk(k)

            weight = part%weight * DEXP(-deltaR/self%maxR)
            ! weight = part%weight

            !Lock the probe
            CALL OMP_SET_LOCK(self%lock)
            
            !Assign particle weight to distribution function
            self%f(i  , j  , k  ) = self%f(i  , j  , k  ) + fi  * fj  * fk  * weight
            self%f(i+1, j  , k  ) = self%f(i+1, j  , k  ) + fi1 * fj  * fk  * weight
            self%f(i  , j+1, k  ) = self%f(i  , j+1, k  ) + fi  * fj1 * fk  * weight
            self%f(i+1, j+1, k  ) = self%f(i+1, j+1, k  ) + fi1 * fj1 * fk  * weight
            self%f(i  , j  , k+1) = self%f(i  , j  , k+1) + fi  * fj  * fk1 * weight
            self%f(i+1, j  , k+1) = self%f(i+1, j  , k+1) + fi1 * fj  * fk1 * weight
            self%f(i  , j+1, k+1) = self%f(i  , j+1, k+1) + fi  * fj1 * fk1 * weight
            self%f(i+1, j+1, k+1) = self%f(i+1, j+1, k+1) + fi1 * fj1 * fk1 * weight

            !Unlock the probe
            CALL OMP_UNSET_LOCK(self%lock)

          END IF

        ! END IF

      END IF
          
    END SUBROUTINE calculate

    SUBROUTINE output(self)
      USE moduleOutput
      USE moduleRefParam
      USE moduleCaseParam, ONLY: timeStep
      IMPLICIT NONE

      CLASS(probeDistFunc), INTENT(inout):: self
      CHARACTER (LEN=iterationDigits):: tstring
      CHARACTER (LEN=3):: pstring
      CHARACTER(:), ALLOCATABLE:: filename
      INTEGER:: i, j, k

      !Divide by the velocity cube volume
      self%f = self%f * self%dvInv

      WRITE(tstring, iterationFormat) timeStep
      WRITE(pstring, "(I3.3)") self%id
      fileName='Probe_' // tstring// '_f_' // pstring // '.dat'
      WRITE(*, "(6X,A15,A)") "Creating file: ", fileName
      OPEN (10, file = path // folder // '/' //  fileName)
      WRITE(10, "(A1, 1X, A)")               "# ", self%species%name
      WRITE(10, "(A6, 3(ES15.6E3), A2)") "# r = ", self%r(:)*L_ref, " m"
      WRITE(10, "(A6, ES15.6E3, A2)")    "# t = ", REAL(timeStep)*tauMin*ti_ref, " s"
      WRITE(10, "(A1, A19, 3(A20))")     "#", "v1 (m s^-1)", "v2 (m s^-1)", "v3 (m s^-1)", "f"
      DO i = 1, self%nv(1)
        DO j = 1, self%nv(2)
          DO k = 1, self%nv(3)
            WRITE(10, "(4(ES20.6E3))") self%vi(i)*v_ref, &
                                       self%vj(j)*v_ref, &
                                       self%vk(k)*v_ref, &
                                       self%f(i, j, k)

          END DO
          WRITE(10, *)

        END DO

      END DO

      CLOSE(10)

      !Reset distribution function
      self%f = 0.D0

    END SUBROUTINE output

    SUBROUTINE doProbes(part)
      IMPLICIT NONE

      TYPE(particle), INTENT(in):: part
      INTEGER:: i

      !Do it so they are only calculated when output
      DO i = 1, nProbes
        IF (probe(i)%update) THEN
          CALL probe(i)%calculate(part)

        END IF

      END DO

    END SUBROUTINE doProbes

    SUBROUTINE outputProbes()
      IMPLICIT NONE

      INTEGER:: i

      DO i = 1, nProbes
        IF (probe(i)%update) THEN
          CALL probe(i)%output()

        END IF

      END DO

    END SUBROUTINE outputProbes

    SUBROUTINE resetProbes()
      USE moduleCaseParam, ONLY: timeStep
      IMPLICIT NONE

      INTEGER:: i

      DO i = 1, nProbes
        probe(i)%f = 0.D0
        probe(i)%update = timeStep == tFinal .OR. timeStep == tInitial .OR. MOD(timeStep, probe(i)%every) == 0

      END DO

    END SUBROUTINE resetProbes

END MODULE moduleProbe