!    Copyright (C) 2007 Imperial College London and others.
!    
!    Please see the AUTHORS file in the main source directory for a full list
!    of copyright holders.
!
!    Prof. C Pain
!    Applied Modelling and Computation Group
!    Department of Earth Science and Engineering
!    Imperial College London
!
!    amcgsoftware@imperial.ac.uk
!    
!    This library is free software; you can redistribute it and/or
!    modify it under the terms of the GNU Lesser General Public
!    License as published by the Free Software Foundation,
!    version 2.1 of the License.
!
!    This library 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
!    Lesser General Public License for more details.
!
!    You should have received a copy of the GNU Lesser General Public
!    License along with this library; if not, write to the Free Software
!    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307
!    USA

#include "fdebug.h"
module initialise_fields_module
use spud
use global_parameters, only: OPTION_PATH_LEN, PYTHON_FUNC_LEN, is_active_process
use futils
use parallel_tools
use fields
use coordinates
use tictoc
use vtk_cache_module
use climatology
use fluxes
use nemo_states_module
use physics_from_options
use load_netcdf_module

implicit none

  interface initialise_field_over_regions

    module procedure initialise_scalar_field_over_regions, initialise_vector_field_over_regions, &
          initialise_tensor_field_over_regions

  end interface

  interface initialise_field

    module procedure initialise_scalar_field, initialise_vector_field, initialise_tensor_field

  end interface

  interface apply_region_ids

    module procedure apply_region_ids_scalar, apply_region_ids_vector, apply_region_ids_tensor

  end interface
    
  private
  public initialise_field, initialise_field_over_regions, apply_region_ids
  
contains

  recursive subroutine initialise_scalar_field(field, path, position, time, phase_path)
    !!< Initialises field with values prescribed in option_path
    !!< This is used for initial conditions, prescribed fields and
    !!< setting a boundary condition surface_field, a.o.
    type(scalar_field), intent(inout) :: field
    character(len=*), intent(in) :: path
    type(vector_field), intent(in):: position
    !! if present use this time level, instead of that in the options tree
    real, optional, intent(in):: time
    character(len=*), intent(in), optional :: phase_path

    type(scalar_field), pointer:: read_field
    type(vector_field), pointer:: vtk_position
    type(vector_field) :: field_position
    real :: const
    character(len=OPTION_PATH_LEN) :: format, field_name, filename
    character(len=PYTHON_FUNC_LEN) :: func
    real :: current_time
    real :: gravity_magnitude
    
    real value
    integer nid

    ! Find out whether initial condition is constant or generated by a 
    ! python function (or comes from something else).    
    if(have_option(trim(path)//"/constant")) then
       call get_option(trim(path)//"/constant", const)
       ! Set field to constant
       call set(field, const)
    else if(have_option(trim(path)//"/python")) then
       call get_option(trim(path)//"/python", func)
       ! Get current time
       if (present(time)) then
         current_time=time
       else
         call get_option("/timestepping/current_time", current_time)
       end if
       ! Set initial condition from python function
       call set_from_python_function(field, trim(func), position, current_time)
    else if(have_option(trim(path)//"/generic_function")) then
       FLExit("Generic functions are obsolete. Please use a Python function.")
    else if(have_option(trim(path)//"/internally_calculated")) then
       continue
    else if(have_option(trim(path)//"/free_surface")) then
       call initialise_field(field, trim(path)//"/free_surface", position, time=time, phase_path=trim(phase_path))
       ! Scale the entered freesurface height by the magnitude of gravity to give the pressure initial condition
       ewrite(3, *) "Free-surface to pressure minmax pre-conversion ", minval(field), maxval(field)
       call convert_free_surface_to_pressure(field, phase_path)
       ewrite(3, *) "Free-surface to pressure minmax post-conversion ", minval(field), maxval(field)

    else if(have_option(trim(path) // "/from_file")) then
       
       if(is_active_process) then
       
         call tic(TICTOC_ID_IO_READ)
         
         call get_option(trim(path) // "/from_file/format/name", format)
         call get_option(trim(path) // "/from_file/file_name", filename)
         if(isparallel()) then
           filename = parallel_filename(trim_file_extension(filename), ".vtu")
         end if
         ewrite(2, *) "Initialising field " // trim(field%name) // " from file " // trim(filename)
         
         select case (format)
         case ("vtu")
            call get_option(trim(path) // "/from_file/format::vtu/field_name", field_name, default = field%name)
            read_field => vtk_cache_read_scalar_field(filename, field_name)            
            
            if (.not. field%mesh%periodic) then
              ! check that the vtk mesh is the same as the derived mesh from state
              ! by comparing its coordinate fields (currently not done for periodic meshes)
              
              vtk_position => vtk_cache_read_positions_field(filename)
              ! get the right field to compare it with
              if (field%mesh==position%mesh .or. element_degree(field,1)==0) then
                ! if the vtk mesh is the same mesh as the position%mesh, use position
                ! (for P0 the vtk_position is that of the linear vertex mesh, so we also can directly use it)
                field_position=position
                call incref(field_position)
              else
                ! otherwise the position mapped onto field%mesh
                field_position = get_remapped_coordinates(position, field%mesh)
              end if
              if (.not. mesh_compatible(vtk_position, field_position)) then
                ! repeat, this to make sure it ends up in any error log
                ewrite(-1,*) "Initialising field " // trim(field%name) // " from file " // trim(filename)
                ! give useful error message (make sure to blame the user)
                ewrite(-1,*) "Error: The mesh in the vtu file is not the same as mesh " // &
                   &trim(field%mesh%name) // &
                   &" specified under " // trim(field%mesh%option_path)
                ! treat this as a user error!
                call print_mesh_incompatibility(-1, vtk_position, field_position)
                FLExit("Mesh from file and in state are not compatible")
              end if
              call deallocate(field_position)
            end if

            call set(field, read_field)
            
            
         case ("climatology (Boyer2005)")
           
            do nid=1, node_count(position)
               if(field%name=="Temperature") then
                  call climatology_GetSurfaceValue("temperature", node_val(position, nid), value)
               else if(field%name=="Salinity") then
                  call climatology_GetSurfaceValue("salinity", node_val(position, nid), value)
               else
                  FLExit("No climatology data available for field: "//field%name)
               end if
               call set(field, nid, value)
            end do
            
         case default
            
            ewrite(-1,*) 'Format: ', trim(format)
            FLExit("Unknown file format")
            
         end select

         call toc(TICTOC_ID_IO_READ)

       end if

    else if(have_option(trim(path)//"/NEMO_data")) then

         call insert_nemo_scalar_field(field)

    else if(have_option(trim(path)//"/from_netcdf")) then

         ! This will set your initial field from data contained in a
         ! netcdf file. The file should be in netcdf format and in
         ! x,y,z format. This option will currently only work for
         ! free surface height (but in principal can easily be extended
         ! to a range of fields).
         call set_scalar_field_from_netcdf(field,path,position)

    else if (have_option(trim(path)//"/no_initial_condition")) then
      continue
    else
       ! This really shouldn't happen (with a valid flml)
       ewrite(-1,*) "Incorrect initial or boundary condition for field ", trim(field%name)
       ewrite(-1,*) "specified in the options at ", trim(path)
       FLAbort("Incorrect initial or boundary condition for field")
    end if

  end subroutine initialise_scalar_field

  subroutine initialise_vector_field(field, path, position, time, phase_path)
    !!< Initialises field with values prescribed in option_path
    !!< This is used for initial conditions, prescribed fields and
    !!< setting a boundary condition surface_field, a.o.
    type(vector_field), intent(inout) :: field
    character(len=*), intent(in) :: path
    type(vector_field), intent(in) :: position
    !! if present use this time level, instead of that in the options tree
    real, optional, intent(in):: time
    character(len=*), intent(in), optional :: phase_path

    type(vector_field), pointer:: read_field, vtk_position
    type(vector_field) :: field_position
    real, dimension(1:field%dim) :: const
    character(len=OPTION_PATH_LEN) :: format, field_name, filename, varname1, varname2
    character(len=PYTHON_FUNC_LEN) :: func
    real :: current_time
    integer :: i
    real :: longitude, latitude, scalars(2), x, y
    logical :: spherical_earth

    ! Find out whether initial condition is constant or generated by a
    ! a python function
    
    if(have_option(trim(path)//"/constant")) then
      call get_option(trim(path)//"/constant", const)
      call set(field, const)
    else if (have_option(trim(path)//"/python")) then
       call get_option(trim(path)//"/python", func)
       ! Get current time
       if (present(time)) then
         current_time=time
       else
         call get_option("/timestepping/current_time", current_time)
       end if
       ! Set initial condition from python function
       call set_from_python_function(field, trim(func), position, current_time)
    else if (have_option(trim(path)//"/generic_function")) then
       call get_option(trim(path)//"/generic_function", func)
       ! Python rules
       FLExit("Generic functions are obsolete. Try Python instead.")
    else if(have_option(trim(path)//"/internally_calculated")) then
       continue
    else if(have_option(trim(path) // "/from_file")) then
      
       if(is_active_process) then
      
         call tic(TICTOC_ID_IO_READ)
         
         call get_option(trim(path) // "/from_file/format/name", format)
         call get_option(trim(path) // "/from_file/file_name", filename)
         if(isparallel()) then
           filename = parallel_filename(trim_file_extension(filename), ".vtu")
         end if
         ewrite(2, *) "Initialising field " // trim(field%name) // " from file " // trim(filename)
         
         select case (format)
         case ("vtu")
            call get_option(trim(path) // "/from_file/format::vtu/field_name", field_name, default = field%name)            
            read_field => vtk_cache_read_vector_field(filename, field_name)        

            if (.not. field%mesh%periodic) then
              ! check that the vtk mesh is the same as the derived mesh from state
              ! by comparing its coordinate fields (currently not done for periodic meshes)
              
              ! get the right field to compare it with
              vtk_position => vtk_cache_read_positions_field(filename)
              if (field%mesh==position%mesh .or. element_degree(field,1)==0) then
                ! if the vtk mesh is the same mesh as the position%mesh, use position
                ! (for P0 the vtk_position is that of the linear vertex mesh, so we also can directly use it)
                field_position=position
                call incref(field_position)
              else
                ! otherwise the position mapped onto field%mesh
                field_position = get_remapped_coordinates(position, field%mesh)
              end if
              
              if (.not. mesh_compatible(vtk_position, field_position)) then
                ! repeat, this to make sure it ends up in any error log
                ewrite(-1,*) "Initialising field " // trim(field%name) // " from file " // trim(filename)
                ! give useful error message (make sure to blame the user)
                ewrite(-1,*) "Error: The mesh in the vtu file is not the same as mesh " // &
                   &trim(field%mesh%name) // &
                   &" specified under " // trim(field%mesh%option_path)
                ! treat this as a user error!
                call print_mesh_incompatibility(-1, vtk_position, field_position)
                FLExit("Mesh from file and in state are not compatible")
              end if
              
              call deallocate(field_position)
            end if
            
            call set(field, read_field)
            
         case default
            
            ewrite(-1,*) 'Format: ', trim(format)
            FLExit("Unknown file format")
            
         end select

         call toc(TICTOC_ID_IO_READ)

       end if

    else if(have_option(trim(path)//"/NEMO_data")) then

         call insert_nemo_vector_field(field)
       
    else if(have_option(trim(path) // "/from_netcdf")) then
             
       ! Get file name
       call get_option(trim(path)//"/from_netcdf/file_name", filename)
       call fluxes_registerdatafile(trim(filename))

       ! Get current time
       call get_option("/timestepping/current_time", current_time)
       call fluxes_settimeseconds(current_time)

       spherical_earth = have_option("/geometry/spherical_earth")

       call get_option(trim(path)//"/from_netcdf/east_west", varname1)
       call fluxes_addfieldofinterest(trim(varname1))

       call get_option(trim(path)//"/from_netcdf/north_south", varname2)
       call fluxes_addfieldofinterest(trim(varname2))
       
       do i=1, node_count(position)
          if(spherical_earth) then
             call LongitudeLatitude(node_val(position, i), &
                  longitude, latitude)
             
             call fluxes_getscalars(longitude, latitude, scalars)
          else
             x=node_val(position, 1, i)
             y=node_val(position, 2, i)
             call fluxes_getscalars(x, y, scalars)
          end if
       end do
    elseif(have_option(trim(path)//"/balanced")) then
       !This option exists in shallow_water schema
       !This will get reset in shallow_water after populate_state
       const = 0.0
       call set(field,const)
    else if (have_option(trim(path)//"/no_initial_condition")) then
      continue
    else
       ! This really shouldn't happen
       ewrite(-1,*) "Incorrect initial or boundary condition for field ", trim(field%name)
       FLExit("Incorrect initial or boundary condition for field")
    end if
 
  end subroutine initialise_vector_field

  subroutine initialise_tensor_field(field, path, position, time, phase_path)
    !!< Initialises field with values prescribed in option_path
    !!< This is used for initial conditions, prescribed fields and
    !!< setting a boundary condition surface_field, a.o.

    type(tensor_field), intent(inout) :: field
    character(len=*), intent(in) :: path
    type(vector_field), intent(in) :: position
    !! if present use this time level, instead of that in the options tree
    real, optional, intent(in):: time
    character(len=*), intent(in), optional :: phase_path
    
    integer :: i
    logical :: is_isotropic, is_diagonal, is_symmetric
    ! name of python function
    character(len=OPTION_PATH_LEN) :: func
    ! time and position needed for python function...
    real :: current_time    
    ! to read in constant initial value
    real:: const
    real, dimension(1:field%dim(1)) :: const_vec
    real, dimension(1:field%dim(1),1:field%dim(2)) :: const_array
    character(len=OPTION_PATH_LEN):: tpath
    ! Temporary field for calculating isotropic tensor fields from python...
    type(scalar_field) :: sfield
    type(vector_field) :: vfield
    
    ! Find out whether tensor is isotropic or symmetric or not
    is_isotropic=have_option(trim(path)//"/isotropic")
    is_diagonal=have_option(trim(path)//"/diagonal")
    is_symmetric=have_option(trim(path)//"/anisotropic_symmetric")

    if(is_isotropic) then
       
       ! Zero tensor
       call zero(field)
       
       tpath=trim(path)//"/isotropic"
       
       ! Find out if tensor diagonal components are constant or whether 
       ! the tensor is generated by a python function.
       
       if(have_option(trim(tpath)//"/constant")) then
          ! Allocate constant array
          const_array=0.0
          ! Put constant on diagonal of tensor
          call get_option(trim(tpath)//"/constant", const)
          do i=1, field%dim(1)
            const_array(i,i)=const
          end do
          call set(field, const_array)
       else if(have_option(trim(tpath)//"/python")) then
          call get_option(trim(tpath)//"/python", func)
          ! Get current time
          if (present(time)) then
            current_time=time
          else
            call get_option("/timestepping/current_time", current_time)
          end if
          call allocate(sfield, field%mesh, name="TemporaryIsotropic")
          call set_from_python_function(sfield, trim(func), position,&
                  & current_time) 
          ! Put isotropic value on diagonal of tensor
          do i=1, field%dim(1)
             call set(field, i, i, sfield)
          end do
          call deallocate(sfield)
       else if (have_option(trim(tpath)//"/generic_function")) then
          FLExit("Generic functions are obsolete. Use a Python function.")
       else if(have_option(trim(path)//"/internally_calculated")) then
          continue
       else
          FLExit("Incorrect initial condition for field")
       end if
       
    else if(is_diagonal) then
    
       ! Zero tensor
       call zero(field)
       
       tpath=trim(path)//"/diagonal"
       
       ! Find out if tensor diagonal components are constant or whether 
       ! the tensor is generated by a python function.
       
       if(have_option(trim(tpath)//"/constant")) then
          ! Allocate constant array
          const_array=0.0
          ! Put constant on diagonal of tensor
          call get_option(trim(tpath)//"/constant", const_vec)
          do i=1, minval(field%dim)
            const_array(i,i)=const_vec(i)
          end do
          call set(field, const_array)
       else if(have_option(trim(tpath)//"/python")) then
          call get_option(trim(tpath)//"/python", func)
          ! Get current time
          if (present(time)) then
            current_time=time
          else
            call get_option("/timestepping/current_time", current_time)
          end if
          call allocate(vfield, minval(field%dim), field%mesh, name="TemporaryDiagonal")
          call set_from_python_function(vfield, trim(func), position,&
                  & current_time) 
          ! Put values on diagonal of tensor
          call set(field, vfield)
          
          call deallocate(vfield)
       else if (have_option(trim(tpath)//"/generic_function")) then
          FLExit("Generic functions are obsolete. Use a Python function.")
       else
          FLExit("Incorrect initial condition for field")
       end if

    else

       ! Set path
       if(is_symmetric) then
          tpath=trim(path)//"/anisotropic_symmetric"
       else
          tpath=trim(path)//"/anisotropic_asymmetric"
       end if


       if(have_option(trim(tpath)//"/constant")) then
          call get_option(trim(tpath)//"/constant", const_array)
          call set(field, const_array)
       else if (have_option(trim(tpath)//"/python")) then
          call get_option(trim(tpath)//"/python", func)
          ! Get current time
          if (present(time)) then
            current_time=time
          else
            call get_option("/timestepping/current_time", current_time)
          end if
          call set_from_python_function(field, trim(func), position,&
                  & current_time) 
       else if (have_option(trim(tpath)//"/generic_function")) then
          FLExit("Generic functions are obsolete. Use a Python function.")
       else
          FLExit("Incorrect initial condition for field")
       end if

    end if

  end subroutine initialise_tensor_field

  subroutine initialise_scalar_field_over_regions(field, path, position, time, phase_path)
    !!< Wrapper to initialise_scalar_field for prescribed and prognostic
    !!< fields in case mesh regions are being used
    type(scalar_field), intent(inout) :: field
    !! path should point to either /"path to field"/prescribed/value 
    !! or /"path to field"/prognostic/initial_condition
    character(len=*), intent(in) :: path
    type(vector_field), intent(in):: position
    real, intent(in), optional :: time
    character(len=*), intent(in), optional :: phase_path

    type(scalar_field) :: tempfield
    integer :: value, nvalues

    call allocate(tempfield, field%mesh, field%name, field_type=field%field_type)

    nvalues = option_count(trim(path))

    do value = 0, nvalues-1
       call zero(tempfield)
       call initialise_field(tempfield, &
            trim(path)//'['//int2str(value)//']', &
            position, time=time, phase_path=trim(phase_path))
       call apply_region_ids(field, tempfield, &
            trim(path)//'['//int2str(value)//']')
    end do

    call deallocate(tempfield)

  end subroutine initialise_scalar_field_over_regions

  subroutine initialise_vector_field_over_regions(field, path, position, time, phase_path)
    !!< Wrapper to initialise_field for prescribed fields in case
    !!< mesh regions are being used
    type(vector_field), intent(inout) :: field
    !! path should point to either /"path to field"/prescribed/value 
    !! or /"path to field"/prognostic/initial_condition
    character(len=*), intent(in) :: path
    type(vector_field), intent(in):: position
    real, intent(in), optional :: time
    character(len=*), intent(in), optional :: phase_path

    type(vector_field) :: tempfield
    integer :: value, nvalues

    call allocate(tempfield, field%dim, field%mesh, field%name, &
      field_type=field%field_type)

    nvalues = option_count(trim(path))

    do value = 0, nvalues-1
       call zero(tempfield)
       call initialise_field(tempfield, &
            trim(path)//'['//int2str(value)//']', &
            position, time=time, phase_path=trim(phase_path))
       call apply_region_ids(field, tempfield, &
            trim(path)//'['//int2str(value)//']')
    end do

    call deallocate(tempfield)

  end subroutine initialise_vector_field_over_regions

  subroutine initialise_tensor_field_over_regions(field, path, position, time, phase_path)
    !!< Wrapper to initialise_field for prescribed fields in case
    !!< mesh regions are being used
    type(tensor_field), intent(inout) :: field
    !! path should point to either /"path to field"/prescribed/value 
    !! or /"path to field"/prognostic/initial_condition
    character(len=*), intent(in) :: path
    type(vector_field), intent(in):: position
    real, intent(in), optional :: time
    character(len=*), intent(in), optional :: phase_path
    
    type(tensor_field) :: tempfield
    integer :: value, nvalues

    call allocate(tempfield, field%mesh, field%name, field_type=field%field_type)

    nvalues = option_count(trim(path))

    do value = 0, nvalues-1
       call zero(tempfield)
       call initialise_field(tempfield, &
            trim(path)//'['//int2str(value)//']', &
            position, time=time, phase_path=trim(phase_path))
       call apply_region_ids(field, tempfield, &
            trim(path)//'['//int2str(value)//']')
    end do

    call deallocate(tempfield)

  end subroutine initialise_tensor_field_over_regions

  subroutine apply_region_ids_scalar(field, tempfield, path)

    type(scalar_field), intent(inout) :: field
    type(scalar_field), intent(in) :: tempfield
    character(len=*), intent(in) :: path

    integer :: stat, ele
    integer, dimension(2) :: shape_option
    integer, dimension(:), allocatable :: region_ids

    stat = 1
    if(have_option(trim(path)//"/region_ids")) then
       shape_option=option_shape(trim(path)//"/region_ids")
       allocate(region_ids(1:shape_option(1)))
       call get_option(trim(path)//"/region_ids", region_ids, stat)
    end if

    ! Set field
    if((associated(field%mesh%region_ids)).and.(stat==0)) then
       do ele = 1, element_count(field%mesh)
          if(any(region_ids==field%mesh%region_ids(ele))) then
             call set(field, ele_nodes(field, ele), ele_val(tempfield, ele))
          end if
       end do
    else
       call set(field, tempfield)
    end if

  end subroutine apply_region_ids_scalar

  subroutine apply_region_ids_vector(field, tempfield, path)

    type(vector_field), intent(inout) :: field
    type(vector_field), intent(in) :: tempfield
    character(len=*), intent(in) :: path

    integer :: stat, ele, dim
    integer, dimension(2) :: shape_option
    integer, dimension(:), allocatable :: region_ids

    stat = 1
    if(have_option(trim(path)//"/region_ids")) then
       shape_option=option_shape(trim(path)//"/region_ids")
       allocate(region_ids(1:shape_option(1)))
       call get_option(trim(path)//"/region_ids", region_ids, stat)
    end if

    ! Get dimension
    call get_option("/geometry/dimension", dim)

    ! Set field
    if((associated(field%mesh%region_ids)).and.(stat==0)) then
       do ele = 1, element_count(field%mesh)
          if(any(region_ids==field%mesh%region_ids(ele))) then
             call set(field, ele_nodes(field, ele), ele_val(tempfield, ele))
          end if
       end do
    else
       call set(field, tempfield)
    end if

  end subroutine apply_region_ids_vector

  subroutine apply_region_ids_tensor(field, tempfield, path)

    type(tensor_field), intent(inout) :: field
    type(tensor_field), intent(in) :: tempfield
    character(len=*), intent(in) :: path

    integer :: stat, ele
    integer, dimension(2) :: shape_option
    integer, dimension(:), allocatable :: region_ids

    stat = 1
    if(have_option(trim(path)//"/region_ids")) then
       shape_option=option_shape(trim(path)//"/region_ids")
       allocate(region_ids(1:shape_option(1)))
       call get_option(trim(path)//"/region_ids", region_ids, stat)
    end if

    ! Set field
    if((associated(field%mesh%region_ids)).and.(stat==0)) then
       do ele = 1, element_count(field%mesh)
          if(any(region_ids==field%mesh%region_ids(ele))) then
             call set(field, ele_nodes(field, ele), ele_val(tempfield, ele))
          end if
       end do
    else
       call set(field, tempfield)
    end if

  end subroutine apply_region_ids_tensor
    
end module initialise_fields_module
