!===================================!
! Magnetic field evolution routines !
!===================================!
! SUBROUTINE evolve_fff() !
!===================================!
MODULE fff_evolve
  IMPLICIT NONE

  CONTAINS

  !--------------------------------------------------------------!
  ! SUBROUTINE evolve_fff !
  !--------------------------------------------------------------!
  ! PURPOSE: !
  ! Evolves the magnetc field for 1 time step !
  ! INPUT: !
  ! none !
  ! OUTPUT: !
  ! none; indirect (bx, by, bz,local_l,conv_flag) !
  ! local_l = final value of l_value !
  ! conv_flag = set to TRUE if converged !
  ! REQUIRES: !
  ! nx,ny,nz,bx,by,bz,qsphere,imax,dt,abs_frac_diff: from vars !
  ! nx, nz, ny = the grid sizes !
  ! bx, by, bz = the magnetic field !
  ! qsphere = set to TRUE for spherical coordinates !
  ! imax = the max number of iterations !
  ! dt = step size in "Time" t !
  ! abs_frac_diff = the fractional change in the l_value !
  ! needed for convergence (the change in L !
  ! needs to be smaller than this) !
  ! xder_coeff,yder_coeff,zder_coeff: from vars !
  ! xint_coeff,yint_coeff,zint_coeff: from vars !
  ! HISTORY: !
  ! 22-may-2006: jmm, jimm@ssl.berkeley.edu !
  ! 13-feb-2007: jmm, added weight function, wf !
  ! Feb-2008: djb, modularized !
  !--------------------------------------------------------------!
  SUBROUTINE evolve_fff
    USE params, ONLY: i4, r8
    USE vars
    USE io
    USE fff_derivs
    USE fff_operators
    IMPLICIT NONE

    ! Local variable declarations
    INTEGER(KIND=i4) :: nxm1, nym1, nzm1, i, flhist, fdebug
    REAL(KIND=r8) :: l_value, delta_l, dtmin
    REAL(KIND=r8), DIMENSION(:, :, :), ALLOCATABLE :: jx, jy, jz
    REAL(KIND=r8), DIMENSION(:, :, :), ALLOCATABLE :: omega_x, omega_y
    REAL(KIND=r8), DIMENSION(:, :, :), ALLOCATABLE :: omega_z, div_b
    REAL(KIND=r8), DIMENSION(:, :, :), ALLOCATABLE :: omega2, b2
    REAL(KIND=r8), DIMENSION(:, :, :), ALLOCATABLE :: f_x, f_y, f_z

    ! Set some parameters and store often used values
    nxm1 = nx - 1
    nym1 = ny - 1
    nzm1 = nz - 1
    i = 0
    conv_flag = .FALSE.
    l_value = 1.0e20_r8 ! High number to start...
    delta_l = 1.0e20_r8
    dtmin = 1.0e-10_r8
    ! Calculate and store grid-based coefficients for derivative and
    ! integral routines
    ALLOCATE(xder_coeff(nx,3), yder_coeff(ny,3), zder_coeff(nz,3))
    ALLOCATE(xint_coeff(nx), yint_coeff(ny), zint_coeff(nz))
    CALL init_derivs
    CALL init_obj_funct

    ALLOCATE(b2(nx,ny,nz))
    b2 = bx*bx + by*by + bz*bz
    ALLOCATE(jx(nx,ny,nz), jy(nx,ny,nz), jz(nx,ny,nz))
    CALL curl_3d(bx, by, bz, jx, jy, jz)
    ALLOCATE(div_b(nx,ny,nz))
    CALL div_3d(bx, by, bz, div_b)

! CALL openfile(fdebug,'fff_deriv_coeffs.dat',IFORM='formatted', &
! ISTATUS='replace')
! WRITE(UNIT=fdebug,FMT=*) xder_coeff
! WRITE(UNIT=fdebug,FMT=*) yder_coeff
! WRITE(UNIT=fdebug,FMT=*) zder_coeff
! CALL closefile(fdebug)

!use the div_b variable for testing purposes
! div_b = deriv_dx(bx)
! CALL openfile(fdebug,'field_test_dbxdx.dat',IFORM='formatted', &
! ISTATUS='replace')
! WRITE(UNIT=fdebug,FMT=*) nx, ny, nz
! WRITE(UNIT=fdebug,FMT=20) div_b
! CALL closefile(fdebug)
! div_b = deriv_dy(by)
! CALL openfile(fdebug,'field_test_dbydy.dat',IFORM='formatted', &
! ISTATUS='replace')
! WRITE(UNIT=fdebug,FMT=*) nx, ny, nz
! WRITE(UNIT=fdebug,FMT=20) div_b
! CALL closefile(fdebug)
! div_b = deriv_dz(bz)
! CALL openfile(fdebug,'field_test_dbzdz.dat',IFORM='formatted', &
! ISTATUS='replace')
! WRITE(UNIT=fdebug,FMT=*) nx, ny, nz
! WRITE(UNIT=fdebug,FMT=20) div_b
! CALL closefile(fdebug)
! CALL openfile(fdebug,'field_test_b.dat',IFORM='formatted', &
! ISTATUS='replace')
! WRITE(UNIT=fdebug,FMT=*) nx, ny, nz
! WRITE(UNIT=fdebug,FMT=20) bx
! WRITE(UNIT=fdebug,FMT=20) by
! WRITE(UNIT=fdebug,FMT=20) bz
! WRITE(UNIT=fdebug,FMT=20) x
! WRITE(UNIT=fdebug,FMT=20) y
! WRITE(UNIT=fdebug,FMT=20) z
!20 FORMAT(5F16.7)
! STOP


    ALLOCATE(omega_x(nx,ny,nz), omega_y(nx,ny,nz), omega_z(nx,ny,nz))
    omega_x = jy*bz - jz*by
    omega_y = jz*bx - jx*bz
    omega_z = jx*by - jy*bx
    WHERE(b2 /= 0.0_r8)
      omega_x = (omega_x - div_b*bx)/b2
      omega_y = (omega_y - div_b*by)/b2
      omega_z = (omega_z - div_b*bz)/b2
    ELSEWHERE
      omega_x = 0.0_r8
      omega_y = 0.0_r8
      omega_z = 0.0_r8
    END WHERE
    ALLOCATE(omega2(nx,ny,nz))
    ! Scale omega_{x,y,z} and omega2 by weighting function
    omega2 = wf*(omega_x*omega_x + omega_y*omega_y + omega_z*omega_z)
    omega_x = wf*omega_x
    omega_y = wf*omega_y
    omega_z = wf*omega_z
    l_value = obj_funct_fff(b2, omega2)

    ALLOCATE(f_x(nx,ny,nz), f_y(nx,ny,nz), f_z(nx,ny,nz))
    CALL openfile(flhist,'lhist.dat',IFORM='formatted',ISTATUS='replace')
    DO
      IF((i >= imax) .OR. (conv_flag .EQV. .TRUE.) .OR. (dt <= dtmin)) EXIT
      i = i + 1
! IF(MOD(i,10).eq.1) THEN
          WRITE(UNIT=flhist,FMT=*) i, dt, l_value
          WRITE(UNIT=6,FMT=*) i, dt, l_value, delta_l
! END IF
      ! now get the force,
      ! maybe we can save memory, subtract omega cross j
      f_x = -(omega_y*jz - omega_z*jy)
      f_y = -(omega_z*jx - omega_x*jz)
      f_z = -(omega_x*jy - omega_y*jx)
      ! Now you can use b2 and jx,jy,jz for temp storage
      ! subtract grad omega dot B,
      ! Add omega times div_b !Add omega2 times B
      b2 = omega_x*bx + omega_y*by + omega_z*bz
      CALL grad_3d(b2, jx, jy, jz)
      f_x = f_x - jx + omega_x*div_b + omega2*bx
      f_y = f_y - jy + omega_y*div_b + omega2*by
      f_z = f_z - jz + omega_z*div_b + omega2*bz
! CALL openfile(fdebug,'field_test_boo.dat',IFORM='formatted', &
! ISTATUS='replace')
! WRITE(UNIT=fdebug,FMT=*) nx, ny, nz
! WRITE(UNIT=fdebug,FMT=20) jx
! WRITE(UNIT=fdebug,FMT=20) jy
! WRITE(UNIT=fdebug,FMT=20) jz
! CALL closefile(fdebug)
! STOP
! 20 FORMAT(5F16.7)
      ! Add curl of omega cross b, b2 will be the x component,
      ! div_b will be the y component, omega2 the z component,
      ! jx,jy,jz are the output
      b2 = omega_y*bz - omega_z*by
      div_b = omega_z*bx - omega_x*bz
      omega2 = omega_x*by - omega_y*bx
      CALL curl_3d(b2, div_b, omega2, jx, jy, jz)
      f_x = f_x + jx
      f_y = f_y + jy
      f_z = f_z + jz
      ! new B
      bx(2:nxm1,2:nym1,2:nzm1) = bx(2:nxm1,2:nym1,2:nzm1) &
                                 + dt*f_x(2:nxm1,2:nym1,2:nzm1)
      by(2:nxm1,2:nym1,2:nzm1) = by(2:nxm1,2:nym1,2:nzm1) &
                                 + dt*f_y(2:nxm1,2:nym1,2:nzm1)
      bz(2:nxm1,2:nym1,2:nzm1) = bz(2:nxm1,2:nym1,2:nzm1) &
                                 + dt*f_z(2:nxm1,2:nym1,2:nzm1)
      ! now new values for j, omega
      b2 = bx*bx + by*by + bz*bz
      CALL curl_3d(bx, by, bz, jx, jy, jz)
      CALL div_3d(bx, by, bz, div_b)
      omega_x = jy*bz - jz*by
      omega_y = jz*bx - jx*bz
      omega_z = jx*by - jy*bx
      WHERE(b2 /= 0.0_r8)
        omega_x = (omega_x - div_b*bx)/b2
        omega_y = (omega_y - div_b*by)/b2
        omega_z = (omega_z - div_b*bz)/b2
      ELSEWHERE
        omega_x = 0.0_r8
        omega_y = 0.0_r8
        omega_z = 0.0_r8
      END WHERE
      ! Scale omega_{x,y,z} and omega2 by weighting function
      omega2 = wf*(omega_x*omega_x + omega_y*omega_y + omega_z*omega_z)
      omega_x = wf*omega_x
      omega_y = wf*omega_y
      omega_z = wf*omega_z
      local_l = obj_funct_fff(b2, omega2)
      delta_l = (l_value - local_l)/l_value
      IF(delta_l > 0.0_r8) THEN ! Success
        dt = 1.01_r8*dt
        ! Convergence?
        l_value = local_l
        IF(ABS(delta_l) < abs_frac_diff) THEN
          conv_flag = .TRUE.
          WRITE(UNIT=6,FMT=*) 'conv_flag = ', conv_flag
        END IF
      ELSE ! Reset B, j, and omega
        dt = 0.5_r8*dt ! Reset dt
        ! Reset b
        bx(2:nxm1,2:nym1,2:nzm1) = bx(2:nxm1,2:nym1,2:nzm1) &
                                   - dt*f_x(2:nxm1,2:nym1,2:nzm1)
        by(2:nxm1,2:nym1,2:nzm1) = by(2:nxm1,2:nym1,2:nzm1) &
                                   - dt*f_y(2:nxm1,2:nym1,2:nzm1)
        bz(2:nxm1,2:nym1,2:nzm1) = bz(2:nxm1,2:nym1,2:nzm1) &
                                   - dt*f_x(2:nxm1,2:nym1,2:nzm1)
        ! Now new values for j, omega, b2, div_b, omega2
        b2 = bx*bx + by*by + bz*bz
        CALL curl_3d(bx, by, bz, jx, jy, jz)
        CALL div_3d(bx, by, bz, div_b)
        omega_x = jy*bz - jz*by
        omega_y = jz*bx - jx*bz
        omega_z = jx*by - jy*bx
        WHERE(b2 /= 0.0_r8)
          omega_x = (omega_x - div_b*bx)/b2
          omega_y = (omega_y - div_b*by)/b2
          omega_z = (omega_z - div_b*bz)/b2
        ELSEWHERE
          omega_x = 0.0_r8
          omega_y = 0.0_r8
          omega_z = 0.0_r8
        END WHERE
        ! Scale omega_{x,y,z} and omega2 by weighting function
        omega2 = wf*(omega_x*omega_x + omega_y*omega_y + omega_z*omega_z)
        omega_x = wf*omega_x
        omega_y = wf*omega_y
        omega_z = wf*omega_z
      END IF
    END DO
    WRITE(UNIT=flhist,FMT=*) i, dt, l_value
    WRITE(UNIT=6,FMT=*) i, dt, l_value, delta_l
    CALL closefile(flhist)
    IF(conv_flag .EQV. .TRUE.) THEN
       WRITE(UNIT=6,FMT=*) 'Successful convergence'
    ELSE
       WRITE(UNIT=6,FMT=*) 'Not converged'
    END IF
    DEALLOCATE(xint_coeff, yint_coeff, zint_coeff)
    DEALLOCATE(xder_coeff, yder_coeff, zder_coeff)
    DEALLOCATE(jx, jy, jz)
    DEALLOCATE(omega_x, omega_y, omega_z)
    DEALLOCATE(div_b, omega2, b2)
    DEALLOCATE(f_x, f_y, f_z)
  END SUBROUTINE evolve_fff

END MODULE fff_evolve
