module newton_cooling_hotJ
    
    !==========================================================================================
    ! Purpose 
    ! -------
    ! Compute a Newtonian cooling scheme for Hot Jupiters
    ! for scenario 1 of the MOCHA intercomparaison project
    ! (add citation to protocol paper here when it's live).
    ! Check this paper's equations (1) and (4):https://iopscience.iop.org/article/10.3847/PSJ/ac9dfe/pdf
    ! 
    ! We aim at having a generic code but you never know, it might need improving at some point.
    ! The current (at time of writing) newtrelax.F90 routine is hardcoded for telluric temperate planets and untested. 
    ! Thus, we don't use it and use this one instead.
    !
    ! Authors
    ! -------
    ! Lucas Teinturier (2024)
    !
    !==========================================================================================
    implicit none 

    ! Module variables
    real, allocatable, save :: T0(:)
    real, allocatable, save :: tau_relax(:)
    real, allocatable, save  :: delta_Teq(:)
    real, allocatable, save :: Trelax(:,:)
    character(100),save :: planetary_suffix
    !$OMP THREADPRIVATE(planetary_suffix, Trelax,tau_relax,T0,delta_Teq)
    
    contains

    subroutine newtcool_MOCHA(ngrid,nlayer,coslon,coslat,temp,pplay,firstcall,lastcall,dtrad)

        ! use callkeys_mod, only: planetary_suffix ! this is to know which profiles to load for the T0, the delta Teq and the tau_rad
        use mod_phys_lmdz_para, only : is_master, bcast ! for OpenMP stuff
        implicit none 

        ! Inputs
        integer, intent(in) :: ngrid,nlayer 
        logical, intent(in) :: firstcall ! is it the first call of physiq_mod ?
        logical, intent(in) :: lastcall !is it the last call of physiq_mod ?
        real, intent(in) :: coslon(ngrid) !cosine of the longitude
        real, intent(in) :: coslat(ngrid) ! cosine of the latitude
        real, intent(in) :: temp(ngrid,nlayer) ! Temperature at each layer (K)
        real, intent(in) :: pplay(ngrid,nlayer) ! Pressure mid-layers (Pa)

        ! Output 
        real, intent(out) :: dtrad(ngrid,nlayer) ! Tendency on temperature dT/dt (K/s)

        !! Internal variable 
        integer ig,l
        character(100) :: filename

        if (firstcall) then 
          ! Allocation of the dynamical arrays 
          allocate(T0(nlayer))
          allocate(tau_relax(nlayer))
          allocate(delta_Teq(nlayer))
          allocate(Trelax(ngrid,nlayer))

            if (is_master) then 
                print*,'-----------------------------------------------------'
                print*,'| ATTENTION: You are using a Newtonian cooling scheme'
                print*,'| for the radiative transfer. This means that ALL'
                print*,'| other physics subroutines must be switched off.'
                print*,'| Check that you have the required files in the '
                print*,'| simulation directory !'
                print*,'-----------------------------------------------------'
                print*,"the planetary suffix is ",planetary_suffix

                !! We load the data using the subroutine load_input

                ! Loading T0 
                filename = trim(planetary_suffix) // "T0.dat"
                ! print*,"filename = ",filename
                call read_input(nlayer,filename,T0)
                print*,"I successfully read",filename

                ! Loading tau_relax 
                filename = trim(planetary_suffix) // "tau_relax.dat"
                call read_input(nlayer,filename,tau_relax)
                print*,"I successfully read",filename

                ! Loading delta_Teq 
                filename = trim(planetary_suffix) // "delta_Teq.dat"
                call read_input(nlayer,filename,delta_Teq)
                print*,"I successfully read",filename

            endif ! of is_master

            ! Broadcast tau_relax and Trelax to everyone
            call bcast(tau_relax)
            call bcast(Trelax)
            call bcast(T0)
            call bcast(delta_Teq)

            ! now initialising Trelax depending on day or night side
            do l=1,nlayer
                do ig=1,ngrid
                    ! if we're on the day-side (the sub-stellar point is at lon =0, dayside is where the coslon >=0)
                    if (coslon(ig) .ge. 0) then 
                        Trelax(ig,l) = T0(l)+delta_Teq(l)*(ABS(coslon(ig)*coslat(ig))-0.5)
                    else !we're on the night-side 
                        Trelax(ig,l) = T0(l)-0.5*delta_Teq(l)
                    endif 
                enddo !ig=1,ngrid
            enddo ! l=1,nlayer

            ! deallocate T0 and delta_Teq, we don't need them anymore
            if (allocated(T0)) deallocate(T0)
            if (allocated(delta_Teq)) deallocate(delta_Teq)
        endif ! of firstcall

        ! call writediagfi(ngrid,"Trelax","Relaxation temperature ","K",3,Trelax)
        ! Calculation of the radiative forcing 
        do l=1,nlayer 
            do ig=1,ngrid 
                if (pplay(ig,l) .le. 1.0e6) then 
                    ! if pressure is lower than 10 bar
                    dtrad(ig,l) = (Trelax(ig,l)-temp(ig,l))/tau_relax(l)
                else 
                    ! Deeper than 10 bar, no relaxation, dtrad = 0
                    dtrad(ig,l) = 0.
                endif !(pplay(ig,l) .le. 1.e6) 
            enddo !ig =1,ngrid 
        enddo !l = 1,nlayer
        
        if (lastcall) then 
            deallocate(tau_relax)
            deallocate(Trelax)
        endif 

    end subroutine newtcool_MOCHA

    subroutine read_input(nlayer,filename, field)

        !========================================
        ! Purpose 
        ! -------
        ! Read the input file for this module 
        ! Each file starts with an integer that should
        ! be equal to nlayer (stops if that's not true)
        !
        ! Author 
        ! ------
        ! Lucas Teinturier(2024)
        !
        !========================================

        implicit none 

        ! Inputs
        integer,intent(in) :: nlayer
        character(100),intent(in) :: filename

        ! Output 
        real, intent(out) :: field(nlayer)

        !! Internal variables 
        integer ierr, nline, ii

        ! Opening the file 
        open(401,form='formatted',status='old',file=trim(filename) ,iostat=ierr)
            if (ierr /=0) then 
                print*,"Problem in newton_cooling_hotJ.F90"
                print*,"I have an issue opening file ",trim(filename)
                stop
            endif 
            ! Checking that we have the right number of atmospheric layers 
            read(401,*) nline 
            if (nline /= nlayer) then 
                print*,"Error, you're not using the right # of atmospheric layers in ",trim(filename)
                stop 
            endif 
            ! Now reading the content of the file
            do ii = 1,nline
                read(401,*) field(ii)
            enddo 
        close(401)
    end subroutine read_input

end module newton_cooling_hotJ