! ***********************************************************************
!
!   Copyright (C) 2010  Bill Paxton
!
!   this file is part of mesa.
!
!   mesa is free software; you can redistribute it and/or modify
!   it under the terms of the gnu general library public license as published
!   by the free software foundation; either version 2 of the license, or
!   (at your option) any later version.
!
!   mesa is distributed in the hope that it will be useful,
!   but without any warranty; without even the implied warranty of
!   merchantability or fitness for a particular purpose.  see the
!   gnu library general public license for more details.
!
!   you should have received a copy of the gnu library general public license
!   along with this software; if not, write to the free software
!   foundation, inc., 59 temple place, suite 330, boston, ma 02111-1307 usa
!
! ***********************************************************************

      module run_star_extras

      use star_lib
      use star_def
      use const_def
      use crlibm_lib

      ! Needed for May's routines
      use chem_def

      implicit none

      character(len=300):: save_prefix
      character(len=300):: save_point
      ! definition of custom-defined parameters that select the
      ! Xc values at which MESA profiles will be saved
      real(dp) :: Xc_save, Xc_save_step
      logical :: need_to_save

      ! these routines are called by the standard run_star check_model
      contains


      subroutine extras_controls(id, ierr)
        integer, intent(in) :: id
        integer, intent(out) :: ierr
        type (star_info), pointer :: s

        integer::narg,cptArg
        character(len=300)::name

        ierr = 0
        call star_ptr(id, s, ierr)
        if (ierr /= 0) return

        narg=command_argument_count()
        write(*,*)'Number of CL-Args: ',narg

        !         read(*,*)test1,test2
        ! this is the place to set any procedure pointers you want to change
        ! e.g., other_wind, other_mixing, other_energy  (see star_data.inc)

        ! Uncomment these lines if you wish to use the functions in this file,
        ! otherwise we use a null_ version which does nothing.
        s% extras_startup => extras_startup
        s% extras_check_model => extras_check_model
        s% extras_finish_step => extras_finish_step
        s% extras_after_evolve => extras_after_evolve
        s% how_many_extra_history_columns => how_many_extra_history_columns
        s% data_for_extra_history_columns => data_for_extra_history_columns
        s% how_many_extra_profile_columns => how_many_extra_profile_columns
        s% data_for_extra_profile_columns => data_for_extra_profile_columns

        ! Once you have set the function pointers you want,
        ! then uncomment this (or set it in your star_job inlist)
        ! to disable the printed warning message,
        s% job% warn_run_star_extras =.false.

        !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
        ! Parameters defined by Jordan Van Beeck, to select 
        ! Xc values at which profiles will be saved
        !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

        ! parameters that determine when profiles are saved throughout
        ! main-sequence evolution
        Xc_save = 0.675 ! start saving models when this Xc value is reached
        Xc_save_step = 0.335 ! interval in Xc for next save point
        ! (i.e. save profiles at Xc = [0.675, 0.340, 0.005])

        if(narg>0)then
            !loop across options
            do cptArg=1,narg
                call get_command_argument(cptArg,name)
                    select case(adjustl(name))
                        !First known args
                        case("--save_prefix")
                            call get_command_argument(cptArg+1,save_prefix)
                            write(*,*)'Save prefix: ',adjustl(save_prefix)
                        case("--save_point")
                            call get_command_argument(cptArg+1,save_point)
                            write(*,*)'Save point: ',adjustl(save_point)
                    end select
            end do

        endif


      end subroutine extras_controls

!       include 'standard_run_star_extras.inc'
      ! None of the following functions are called unless you set their
      ! function point in extras_control.

      integer function extras_startup(id, restart, ierr)
         integer, intent(in) :: id
         logical, intent(in) :: restart
         integer, intent(out) :: ierr
         type (star_info), pointer :: s

         ierr = 0
         call star_ptr(id, s, ierr)
         if (ierr /= 0) return
         extras_startup = 0
         if (.not. restart) then
            write(*,*)'STARTING INITIAL RUN (NO RESTART)'

         else ! it is a restart
            call unpack_extra_info(s)
         end if
      end function extras_startup

      ! returns either keep_going, retry, backup, or terminate.
      integer function extras_check_model(id, id_extra)
         integer, intent(in) :: id, id_extra
         integer :: ierr
         logical :: retry_needed
         type (star_info), pointer :: s
         ierr = 0
         call star_ptr(id, s, ierr)
         if (ierr /= 0) return
         extras_check_model = keep_going

         ! if you want to check multiple conditions, it can be useful
         ! to set a different termination code depending on which
         ! condition was triggered.  MESA provides 9 customizeable
         ! termination codes, named t_xtra1 .. t_xtra9.  You can
         ! customize the messages that will be printed upon exit by
         ! setting the corresponding termination_code_str value.
         ! termination_code_str(t_xtra1) = 'my termination condition'

         retry_needed = .false.
         ! check if MESA needs to save
         call save_at_Xc(id,Xc_save,need_to_save,ierr)

         ! retry if desired precision (0.0005) is not reached when saving a profile, or if target Xc was missed
         if ((need_to_save .and. abs(s% center_h1 - Xc_save)> 5d-4) .or. (s%center_h1 < Xc_save .and. .not. need_to_save)) then
            s%dt = 0.5*s%dt
            extras_check_model = retry
            retry_needed = .true.
         endif

         ! Adjust Xc_to_save if model is going to be saved
         if (need_to_save .and. .not. retry_needed) Xc_save = Xc_save-Xc_save_step

         ! by default, indicate where (in the code) MESA terminated
         if (extras_check_model == terminate) s% termination_code = t_extras_check_model
      end function extras_check_model

   ! %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
      ! EXTRA SUBROUTINES, FROM CODE WRITTEN BY MATHIAS MICHIELSEN 
      ! WITH ADAPTATIONS BY JORDAN VAN BEECK

      subroutine save_at_Xc(id, Xc_save,make_save,ierr)
         integer, intent(in) :: id
         real(dp), intent(in) :: Xc_save
         logical, intent(out) :: make_save
         integer, intent(out) :: ierr

         type(star_info), pointer :: s
         real(dp) :: Xc_hi, Xc_lo, Xc
         real(dp) :: dXc

         call star_ptr(id,s,ierr)
         if (ierr /= 0) then
            write(*,*) 'Error: runstarex_numerics: save_at_Xc: star_ptr failed'
            return
         endif

         if (Xc_save .le. 0) then
            make_save = .false.
            return
         endif

         ! Enforce finer Xc-step near the save values
         Xc = s%center_h1 
         if (ABS(Xc-Xc_save)> 0.04) then
            dXc = 0.005
         else
            if (ABS(Xc-Xc_save) > 0.01) then
               dXc = 0.0005
            else
               dXc = 0.00002
            endif
         endif

         ! set bounds of the Xc-grid used for saving a profile
         Xc_hi = Xc_save+0.08
         Xc_lo = Xc_save-0.002
         call set_grid_dXc(id, Xc_hi, Xc_lo, dXc, ierr)

         ! save a profile if the difference in Xc is smaller than 5d-5
         if (ABS(Xc-Xc_save) < 0.00005) then
            make_save = .true.
            return
         endif
         make_save = .false.
         return

      end subroutine save_at_Xc

      subroutine set_grid_dXc(id,Xc_hi,Xc_lo,dXc,ierr)
         integer, intent(in) :: id
         real(dp), optional, intent(in) :: Xc_hi, Xc_lo
         real(dp), intent(in) :: dXc
         integer, intent(out) :: ierr

         type(star_info), pointer :: s
         real(dp) :: Xc, Xc_next, min_Xc_chb, max_Xc_chb, delta_lg_XH_cntr_limit, lowest_delta_lg_XH_cntr_limit
         logical :: is_chb

         call star_ptr(id,s,ierr)
         if (ierr /= 0) then
            write(*,*) 'Error: runstarex numerics: set_grid_dXc: star_ptr failed'
            return
         endif

         call crlibm_init()
         Xc = s%center_h1

         ! resort back to controls.defaults values and return back
         if (present(Xc_hi) .and. present(Xc_lo)) then
            if (Xc .gt. Xc_hi .or. Xc .lt. Xc_lo) then
               s% dX_nuc_drop_min_X_limit = 1d-4 ! controls.defaults
               s% delta_lg_XH_cntr_min = -6 ! controls.defaults
               s% delta_lg_XH_cntr_max = -3 ! controls.defaults
               s% delta_lg_XH_cntr_limit = 0.05 ! controls.defaults
               return
            endif
         endif

         ! define boundaries of the main sequence in terms of Xc
         max_Xc_chb = (s% job% initial_h1 + s% job% initial_h2)
         min_Xc_chb = 1d-12
         ! check if the model is still on the main sequence / burning core hydrogen
         is_chb = ((Xc .ge. min_Xc_chb) .and. (Xc <= max_Xc_chb))
         ! modify the delta_lg_XH_cntr_limit only during the main sequence
         if (.not. is_chb) return
         ! check if the next step in Xc ends the main sequence
         Xc_next = Xc-dXc
         if (Xc_next .lt. min_Xc_chb) return

         ! change dXc when Xc > min_Xc_chb
         s% dX_nuc_drop_min_X_limit = min_Xc_chb 
         s% delta_lg_XH_cntr_min = log10_cr(min_Xc_chb)
         s% delta_lg_XH_cntr_max = 0d0

         ! parameter used to define time step change treshold
         lowest_delta_lg_XH_cntr_limit = 1d-5
         ! define time step change treshold
         delta_lg_XH_cntr_limit = log10_cr(Xc) - log10_cr(Xc_next)
         delta_lg_XH_cntr_limit = max(delta_lg_XH_cntr_limit,lowest_delta_lg_XH_cntr_limit)
         ! If max delta is greater than this, reduce the next timestep 
         ! by delta_lg_XH_cntr_limit/max_delta.
         s% delta_lg_XH_cntr_limit = delta_lg_XH_cntr_limit

         if (.false.) write(*,'(3(e18.6))') Xc, Xc_next, delta_lg_XH_cntr_limit
      end subroutine set_grid_dXc


      subroutine save_a_profile(id, id_extra, prefix, ierr)
         integer, intent(in) :: id, id_extra
         character(len=*), intent(in) :: prefix
         integer, intent(out) :: ierr
         character(len=256) :: prof_name_Xc, full_path

         type(star_info), pointer :: s

         ! extra subroutine to trigger MESA to save a profile
         call star_ptr(id, s, ierr)
         if (ierr /= 0) then
            write(*,*) 'Error: runstarex_numerics: save_a_profile: star_ptr failed'
            return
         endif
         ! define a parameter used to change the name of a MESA profile
         write(prof_name_Xc, '(a2, f6.4)') 'Xc', s% center_h1
         ! change the name of the saved profile
         full_path = trim(s% log_directory) // '/' // trim(prefix) // '_' // trim(prof_name_Xc) // '.profile'
         ! save the profile
         call star_write_profile_info(id, full_path, id_extra, ierr)
         if (ierr /= 0) then
            write(*,*) 'Error: runstarex_numerics: save_a_profile: star_Write_profile_info failed'
            return
         endif

      end subroutine save_a_profile

      subroutine save_gyre(id, prefix, ierr)
         integer, intent(in) :: id
         character(len=*), intent(in) :: prefix
         integer, intent(out) :: ierr
         character(len=256) :: gyre_name_Xc, full_path
         type(star_info), pointer :: s

         ! extra subroutine to trigger MESA to save GYRE-readable output
         call get_star_ptr(id, s, ierr)
         if (ierr /= 0) then
            write(*,*) 'Error: runstarex_numerics: save_gyre: get_star_ptr failed'
            return
         endif
         ! define a parameter used to change the name of the GYRE-readable output
         write(gyre_name_Xc, '(a2, f6.4)') 'Xc', s% center_h1 !4 numbers after decimal point (6 characters, including decimal point)
         ! change the name of the saved GYRE-readable output file
         full_path = trim(s% log_directory) // '/' // trim(prefix) // '_' // trim(gyre_name_Xc) // '.data.GYRE'
         ! save the GYRE-readable output
         call star_export_pulse_data(id, s% pulse_data_format, full_path, &
         s% add_center_point_to_pulse_data, s% keep_surface_point_for_pulse_data, & 
         s% add_atmosphere_to_pulse_data, ierr)

         if (ierr /= 0) then
            write(*,*) 'Error: runstarex_numerics: save_gyre: star_write_gyre failed'
            ierr = 0
            return
         endif

      end subroutine save_gyre

   ! %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


      integer function how_many_extra_history_columns(id, id_extra)
         integer, intent(in) :: id, id_extra
         integer :: ierr
         type (star_info), pointer :: s
         ierr = 0
         call star_ptr(id, s, ierr)
         if (ierr /= 0) return
         how_many_extra_history_columns = 2
      end function how_many_extra_history_columns



      subroutine data_for_extra_history_columns(id, id_extra, n, names, vals, ierr)
         integer, intent(in) :: id, id_extra, n
         integer :: nn
         character (len=maxlen_history_column_name) :: names(n)
         real(dp) :: vals(n)
         real(dp) :: average_m, average_r

         integer, intent(out) :: ierr
         type (star_info), pointer :: s
         integer :: k, l
         double precision :: frctn
         double precision, allocatable :: diff(:)

         ierr = 0
         call star_ptr(id, s, ierr)
         if (ierr /= 0) return

         ! -------------------------------------------------------------------------------
         ! Routine developed by Jordan Van Beeck to estimate the Roche break-up velocity.
         ! -------------------------------------------------------------------------------

         !! This routine computes the approximation of the break-up velocity (in rad/s) 
         !! following the Roche lobe criterion used in Van Beeck et al. (2020):
         ! Omega_crit = sqrt(8 * G * Mstar / (27 * Rstar^3))

         ! save the computed break-up velocity (in rad/s)
         names(1) = "Omega_crit"
         vals(1) = sqrt(8.d0 * standard_cgrav * msol * s% star_mass / (27.d0 * (rsol * s% photosphere_r)**3.d0))

         ! ----------------------------------------------------------------------------------------------
         ! Routine developed by Timothy Van Reeth and modified by Cole Johnston to compute the core mass.
         ! ----------------------------------------------------------------------------------------------
         names(2) = "mass_conv_core_custom"
         vals(2) = 0.d0

         allocate(diff(s%nz-1))

         do k=1,s%nz-1
             diff(k) = ( ( s% gradr(k) - s% grada(k) ) * ( s% gradr(k+1)- s% grada(k+1) ) )
         end do

         if ( s% q(2) > s% q(1) ) then
             do k=1,s%nz-1
                 if (diff(k) < 0) then
                     if ( (s% gradr(k+1) - s% grada(k+1)) < (s% gradr(k)- s% grada(k))  ) then
                         frctn = (0.d0-(s% gradr(k)-s% grada(k)))/((s% gradr(k+1)-s% grada(k+1))-(s% gradr(k)-s% grada(k)))
                         vals(2) = s% initial_mass * ( s% q(k+1)*frctn+ s% q(k)*(1.d0-frctn) )
                         exit
                     end if
                 end if
             end do

         else
             do l=1,s %nz-1
                 k = s% nz - l
                 if ( diff(k) < 0 ) then
                     if ( (s% gradr(k+1) - s% grada(k+1)) > (s% gradr(k)- s% grada(k)) ) then
                         frctn = (0.d0-(s% gradr(k)-s% grada(k)))/((s% gradr(k+1)-s% grada(k+1))-(s% gradr(k)-s% grada(k)))
                         vals(2) = s% initial_mass * ( s% q(k+1)*frctn+ s% q(k)*(1.d0-frctn) )
                         exit
                     end if
                 end if
             end do
         end if
         ! deallocate the variable
         deallocate(diff)

         ierr = 0

         !note: do NOT add the extras names to history_columns.list
         ! the history_columns.list is only for the built-in log column options.
         ! it must not include the new column names you are adding here.

      end subroutine data_for_extra_history_columns

      integer function how_many_extra_profile_columns(id, id_extra)
         use star_def, only: star_info
         integer, intent(in) :: id, id_extra
         integer :: ierr
         type (star_info), pointer :: s
         ierr = 0
         call star_ptr(id, s, ierr)
         if (ierr /= 0) return
         how_many_extra_profile_columns = 0
      end function how_many_extra_profile_columns


      subroutine data_for_extra_profile_columns(id, id_extra, n, nz, names, vals, ierr)
         use star_def, only: star_info, maxlen_profile_column_name
         use const_def, only: dp
         integer, intent(in) :: id, id_extra, n, nz
         character (len=maxlen_profile_column_name) :: names(n)
         real(dp) :: vals(nz,n)
         integer, intent(out) :: ierr
         type (star_info), pointer :: s
         integer :: k
         real(dp) :: rhs, lhs
         ierr = 0
         call star_ptr(id, s, ierr)
         if (ierr /= 0) return

      end subroutine data_for_extra_profile_columns


      integer function extras_finish_step(id, id_extra)
          integer, intent(in) :: id, id_extra
          integer :: ierr
          real(16), save :: tams_age  ! DMB included 'save' such that tams_age becomes a global variable for next iteration
          type (star_info), pointer :: s
          logical, save :: tams_age_saved = .false.  ! DMB included 'save' such that tams_age becomes a global varibale for next iteration
          character(len=256) :: prof_prefix
          ierr = 0

          call star_ptr(id, s, ierr)
          if (ierr /= 0) return
          extras_finish_step = keep_going
          call store_extra_info(s)

          ! Save profiles at certain core hydrogen mass fractions (adapted from code by Mathias Michielsen)
          ! using specific functions defined earlier in this run_star_extras.
          if (need_to_save) then
            prof_prefix = trim(s% profile_data_prefix)
            print*, 'MESA save_a_profile', s% model_number
            call save_a_profile(id, id_extra, prof_prefix, ierr)
            if (ierr /= 0) then
               write(*,*) 'Error: run_star_extras: extras_finish_step: save_a_profile failed'
               return
            endif

            if (s% write_pulse_data_with_profile) then
               print*, 'MESA save_gyre', s% model_number
               call save_gyre(id,prof_prefix,ierr)
               if (ierr /= 0) then
                  write(*,*) 'Error: run_star_extras: extras_finish_step: save_gyre failed'
                  return
               endif
            endif
          endif

          !!! for MS evolution
          if (s% center_h1 > 0.005) then

            extras_finish_step = keep_going

          !!! for post-MS evolution
          else
            s% max_years_for_timestep = 10**(7 - s% star_mass)  !  1.0d5   ! DMB changed from 1d4
            s% force_timestep_min_years = 1.0d3
            if ( .not. tams_age_saved ) then
              tams_age = s% star_age
              tams_age_saved = .true.
            else
              if (s% Teff > 6500) then   !(s% star_age .le. 1.03d0*tams_age) then
                extras_finish_step = keep_going
              else
                write(*,*)'TERMINATING: post_MS is colder than 6500'  !star_age => 3% beyond tams_age'
                extras_finish_step = terminate
              endif
            endif
          endif

           ! see extras_check_model for information about custom termination codes
           ! by default, indicate where (in the code) MESA terminated
           if (extras_finish_step == terminate) s% termination_code = t_extras_finish_step
        end function extras_finish_step

      subroutine extras_after_evolve(id, id_extra, ierr)
         integer, intent(in) :: id, id_extra
         integer, intent(out) :: ierr
         type (star_info), pointer :: s
         ierr = 0
         call star_ptr(id, s, ierr)
         if (ierr /= 0) return
      end subroutine extras_after_evolve

      ! routines for saving and restoring extra data so you can do restarts

         ! put these defs at the top and delete from the following routines
         !integer, parameter :: extra_info_alloc = 1
         !integer, parameter :: extra_info_get = 2
         !integer, parameter :: extra_info_put = 3


      subroutine alloc_extra_info(s)
         integer, parameter :: extra_info_alloc = 1
         type (star_info), pointer :: s
         call move_extra_info(s,extra_info_alloc)
      end subroutine alloc_extra_info


      subroutine unpack_extra_info(s)
         integer, parameter :: extra_info_get = 2
         type (star_info), pointer :: s
         call move_extra_info(s,extra_info_get)
      end subroutine unpack_extra_info


      subroutine store_extra_info(s)
         integer, parameter :: extra_info_put = 3
         type (star_info), pointer :: s
         call move_extra_info(s,extra_info_put)
      end subroutine store_extra_info


      subroutine move_extra_info(s,op)
         integer, parameter :: extra_info_alloc = 1
         integer, parameter :: extra_info_get = 2
         integer, parameter :: extra_info_put = 3
         type (star_info), pointer :: s
         integer, intent(in) :: op

         integer :: i, j, num_ints, num_dbls, ierr

         i = 0
         ! call move_int or move_flg
         num_ints = i

         i = 0
         ! call move_dbl

         num_dbls = i

         if (op /= extra_info_alloc) return
         if (num_ints == 0 .and. num_dbls == 0) return

         ierr = 0
         call star_alloc_extras(s% id, num_ints, num_dbls, ierr)
         if (ierr /= 0) then
            write(*,*) 'failed in star_alloc_extras'
            write(*,*) 'alloc_extras num_ints', num_ints
            write(*,*) 'alloc_extras num_dbls', num_dbls
            stop 1
         end if

         contains

         subroutine move_dbl(dbl)
            real(dp) :: dbl
            i = i+1
            select case (op)
            case (extra_info_get)
               dbl = s% extra_work(i)
            case (extra_info_put)
               s% extra_work(i) = dbl
            end select
         end subroutine move_dbl

         subroutine move_int(int)
            integer :: int
            i = i+1
            select case (op)
            case (extra_info_get)
               int = s% extra_iwork(i)
            case (extra_info_put)
               s% extra_iwork(i) = int
            end select
         end subroutine move_int

         subroutine move_flg(flg)
            logical :: flg
            i = i+1
            select case (op)
            case (extra_info_get)
               flg = (s% extra_iwork(i) /= 0)
            case (extra_info_put)
               if (flg) then
                  s% extra_iwork(i) = 1
               else
                  s% extra_iwork(i) = 0
               end if
            end select
         end subroutine move_flg

      end subroutine move_extra_info

      end module run_star_extras
