!
! astrofits - write WCS to FITS image & read star table
!
!
! Copyright © 2011-3, 2015-7, 2020 F.Hroch (hroch@physics.muni.cz)
!
! This file is part of Munipack.
!
! Munipack is free software: you can redistribute it and/or modify
! it under the terms of the GNU General Public License as published by
! the Free Software Foundation, either version 3 of the License, or
! (at your option) any later version.
!
! Munipack 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 General Public License for more details.
!
! You should have received a copy of the GNU General Public License
! along with Munipack.  If not, see <http://www.gnu.org/licenses/>.
!

module astrofits

  use titsio
  use astrotrafo
  use trajd
  use iso_fortran_env

  implicit none

  integer, parameter, private :: dbl = selected_real_kind(15)

contains

  subroutine readtable(filename, label_alpha, label_delta, label_pmalpha, &
       label_pmdelta, label_mag, key_epoch, alpha, delta, pmalpha, pmdelta, &
       refflux, rframe, catid, epoch,status)

    character(len=*), intent(in) :: filename, label_alpha, label_delta, &
         label_pmalpha, label_pmdelta, label_mag, key_epoch
    real(dbl), allocatable, dimension(:), intent(out) :: alpha, delta, &
         pmalpha, pmdelta, refflux
    character(len=*), intent(out) :: rframe, catid
    real(dbl), intent(out) :: epoch
    integer, intent(out) :: status

    real(dbl), parameter :: nullcoo = 1e33, clim = 0.1*nullcoo
    real(dbl), parameter :: nullmag = 99.999
    real(dbl), parameter :: mag25 = 25
    integer, parameter :: ncols = 5
    integer, dimension(ncols) :: col, statuses
    real(dbl), allocatable, dimension(:) :: ra,dec,pmra,pmdec,mag
    integer :: nrows,i,l,nr,n,frow,srows,stat
    logical :: anyf
    character(len=FLEN_VALUE), dimension(ncols) :: label,colname
    character(len=80) :: errmsg
    type(fitsfiles) :: fits

    label(1) = label_alpha
    label(2) = label_delta
    label(3) = label_pmalpha
    label(4) = label_pmdelta
    label(5) = label_mag

    status = 0

    ! open, and move, to a table extension
    call fits_open_table(fits,filename,FITS_READONLY,status)
    if( status /= 0 ) then
       write(error_unit,*) &
            'Error: failed to read a table in the file `',trim(filename),"'."
       return
    end if

    call fits_get_num_rows(fits,nrows,status)
    if( status /= 0 ) goto 666
    if( .not. (nrows > 0) )then
       write(error_unit,*) 'Error: encountered empty table in the file `', &
            trim(filename),"'."
       goto 666
    end if

    ! define reference frame and identification of catalogue
    rframe = ''
    catid = ''
    call fits_read_key(fits,'EXTNAME',catid,status)
    call fits_read_key(fits,key_epoch,epoch,status)
    if( status == FITS_KEYWORD_NOT_FOUND ) then
       status = 0
       epoch = yearjd(2000.0_dbl)   ! 2451545.0
    else
       epoch = yearjd(epoch)
    end if

    ! find columns by labels
    statuses = 0
    do i = 1, ncols
       colname(i) = ''
       if( label(i) /= '' ) then
          call fits_get_colname(fits,.false.,label(i),colname(i),col(i),statuses(i))
       else
          statuses(i) = -1
       end if
    end do
    if( any(statuses(1:2) /= 0) ) then
       write(error_unit,*) &
            'Error: Right Ascension and Declination columns are mandatory.'
       status = maxval(statuses(1:2))
       goto 666
    end if

    allocate(ra(nrows),dec(nrows),pmra(nrows),pmdec(nrows),mag(nrows),&
         stat=stat,errmsg=errmsg)
    if( stat /= 0 ) then
       write(error_unit,'(2a)') 'Error:', trim(errmsg)
       goto 666
    end if

    pmra = 0.0_dbl
    pmdec = 0.0_dbl
    mag = mag25

    call fits_get_rowsize(fits,srows,status)

    do i = 1, nrows, srows
       frow = i
       l = min(i+srows,nrows)

       ! coordinates
       call fits_read_col(fits,col(1),frow,nullcoo,ra(i:l),anyf,status)
       call fits_read_col(fits,col(2),frow,nullcoo,dec(i:l),anyf,status)

       ! proper motion
       if( all(statuses(3:4) == 0) ) then
          call fits_read_col(fits,col(3),frow,nullcoo,pmra(i:l),anyf,status)
          call fits_read_col(fits,col(4),frow,nullcoo,pmdec(i:l),anyf,status)
       end if

       ! magnitudes
       if( statuses(5) == 0 ) then
          call fits_read_col(fits,col(5),frow,nullmag,mag(i:l),anyf,status)
       end if

       if( status /= 0 ) goto 666
    end do

    pmra = pmra / 3.6e6_dbl  ! supposing mas/year
    pmdec = pmdec / 3.6e6_dbl


    nr = min(count(abs(ra) < clim), count(abs(dec) < clim), count(mag < 99))
    allocate(alpha(nr),delta(nr),pmalpha(nr),pmdelta(nr),refflux(nr), &
         stat=stat,errmsg=errmsg)
    if( stat /= 0 ) then
       write(error_unit,'(2a)') 'Error:', trim(errmsg)
       goto 666
    end if

    n = 0
    do i = 1, nrows
       if( abs(ra(i)) < clim .and. abs(dec(i)) < clim .and. mag(i) < 99 ) then
          n = n + 1
          alpha(n) = ra(i)
          delta(n) = dec(i)
          pmalpha(n) = pmra(i)
          pmdelta(n) = pmdec(i)
          refflux(n) = 10.0**((mag25 - mag(i))/2.5)
       end if
    end do
    deallocate(ra,dec,pmra,pmdec,mag)

    call fits_close_file(fits,status)
    return

666 continue

    if( allocated(ra) ) deallocate(ra,dec,pmra,pmdec,mag)
    call fits_close_file(fits,status)
    call fits_report_error(error_unit,status)

  end subroutine readtable


  subroutine readfile(filename, key_dateobs, width, height, xcoo, ycoo, flux, &
       crpix, jd, crval, sc, pa, refl, status)

    use phio

    character(len=*), intent(in) :: filename, key_dateobs
    real(dbl), intent(out) :: width, height
    real(dbl), allocatable, dimension(:), intent(out) :: xcoo, ycoo, flux
    real(dbl), dimension(:), intent(out) :: crpix, crval
    real(dbl), intent(out) :: jd, sc, pa, refl
    integer,intent(out) :: status

    integer :: naxis, bitpix, nrows, n, ns, i, l, srows, frow
    integer, dimension(2) :: naxes
    logical :: anyf
    integer, parameter :: extver = 0
    real(dbl), allocatable, dimension(:) :: x,y,cts
    real(dbl), parameter :: nullval = 0.0_dbl
    character(len=FLEN_CARD) :: colname1, colname2, dateobs, aplabel
    integer :: col1, col2
    type(AstroTrafoProj) :: t
    type(fitsfiles) :: fits

    status = 0
    call fits_open_image(fits,filename,FITS_READONLY,status)
    if( status /= 0 ) then
       write(error_unit,*) &
            'Error: failed to read an image in the file `',trim(filename),"'."
       return
    end if

    call fits_get_img_param(fits,bitpix,naxis,naxes,status)
    if( status /= 0 .or. naxis /= 2 ) then
       if( naxis /= 2 ) then
            write(error_unit,*) trim(filename), &
                 ": Sorry, just two dimensional images are implemented yet."
            status = -1
         end if
       goto 666
    end if
    width = naxes(1)
    height = naxes(2)
    crpix = naxes/2.0_dbl

    call fits_read_key(fits,key_dateobs,dateobs,status)
    if( status == 0 ) then
       jd = fits_jd(dateobs,status)
    else if( status == FITS_KEYWORD_NOT_FOUND ) then
       status = 0
       jd = 2451544.5_dbl
       write(error_unit,*) "Warning: the keyword `",trim(key_dateobs), &
            "' not found in `",trim(filename),"'. Set to J2000.0."
    end if
    if( status /= 0 ) goto 666

    ! get calibration
    call wcsget(fits,t,status)
    if( status == 0 ) then
       sc = t%scale
       pa = t%rot
       refl = t%refl
       crval = [ t%acen, t%dcen ]
       crpix = [ t%xcen, t%ycen ]
    else
       status = 0
       crval = -999.999
    end if

    ! FIND extension
    call fits_movnam_hdu(fits,FITS_BINARY_TBL,FINDEXTNAME,extver,status)
    if( status == FITS_BAD_HDU_NUM ) then
       write(error_unit,*) trim(FINDEXTNAME)// &
            ': extension not found. Coordinates of objects are not available.'
       goto 666
    end if
    call fits_get_num_rows(fits,nrows,status)
    if( status /= 0 ) goto 666

    allocate(x(nrows),y(nrows))

    ! the column labels are specific for MUNIPACK extension only
    colname1 = ''
    colname2 = ''
    call fits_get_colname(fits,.false.,FITS_COL_X,colname1,col1,status)
    call fits_get_colname(fits,.false.,FITS_COL_Y,colname2,col2,status)

    call fits_get_rowsize(fits,srows,status)
    do i = 1, nrows, srows
       frow = i
       l = min(i+srows,nrows)

       call fits_read_col(fits,col1,frow,nullval,x(i:l),anyf,status)
       call fits_read_col(fits,col2,frow,nullval,y(i:l),anyf,status)

       if( status /= 0 ) goto 666
    end do

    ! photometry extension
    call fits_movnam_hdu(fits,FITS_BINARY_TBL,APEREXTNAME,extver,status)
    if( status == FITS_BAD_HDU_NUM ) then
       status = 0
       ! aperture photometry is not available, no problem,
       ! we are continuing still and all ...?
       ns = size(x)
       allocate(xcoo(ns),ycoo(ns),flux(ns))
       xcoo = x
       ycoo = y
       flux = 1.0
    else if( status == 0 ) then
       ! aperture photometry is available, we can use photometry informations
       ! to improve reliability
       aplabel = trim(FITS_COL_APCOUNT)//"1"
       call fits_get_colname(fits,.false.,aplabel,colname1,col1,status)
       allocate(cts(nrows))
       frow = 1
       call fits_read_col(fits,col1,frow,nullval,cts,anyf,status)

       ! we're selecting stars with valid photometry only
       ns = count(cts > 0.0)
       allocate(xcoo(ns),ycoo(ns),flux(ns))
       ns = 0
       do n = 1,nrows
          if( cts(n) > 0.0 .and. ns < size(xcoo) ) then
             ns = ns + 1
             xcoo(ns) = x(n)
             ycoo(ns) = y(n)
             flux(ns) = cts(n)
          end if
       end do
       deallocate(cts)

    end if
    if( status /= 0 ) goto 666

    call fits_close_file(fits,status)

    deallocate(x,y)
    return

666 continue

    if( allocated(x) ) deallocate(x,y)
    if( allocated(cts) ) deallocate(cts)

    call fits_close_file(fits,status)
    call fits_report_error(error_unit,status)

  end subroutine readfile


  subroutine reffile(filename,key_dateobs,alpha,delta,rflux,jd,status)

    character(len=*), intent(in) :: filename, key_dateobs
    real(dbl), allocatable, dimension(:), intent(out) :: alpha, delta, rflux
    real(dbl), intent(out) :: jd
    integer, intent(out) :: status

    real(dbl) :: w, h, sc, rot, refl
    real(dbl), allocatable, dimension(:) :: x,y
    real(dbl), dimension(2) :: crpix,crval
    type(AstroTrafoProj) :: t


    call readfile(filename,key_dateobs,w,h,x,y,rflux,crpix,jd,crval,&
         sc,rot,refl,status)

    if( status == 0 .and. crval(1) > -666.0 .and. allocated(x) ) then

       allocate(alpha(size(x)),delta(size(x)))

       call trafo_init(t,"GNOMONIC",crval(1),crval(2),crpix(1),crpix(2),&
            scale=sc,rot=rot,refl=refl)
       call invtrafo(t,x,y,alpha,delta)

       deallocate(x,y)
    end if

  end subroutine reffile



!-----------------------------------------------------------------------------

  subroutine wcsupdate(fits,type,crval,sc,pa,dpa,fcrpix,crpixels,refl,rms,&
       com,status)

    use astrotrafo

    integer, parameter :: ndim = 2
    real(dbl), parameter :: rad = 57.295779513082322865_dbl

    type(fitsfiles) :: fits
    character(len=*), intent(in) :: type
    real(dbl), dimension(:), intent(in) :: crpixels,crval
    real(dbl), intent(in) :: sc,pa,dpa,rms,refl
    logical, intent(in) :: fcrpix
    character(len=*), dimension(:), intent(in) :: com
    integer, intent(in out) :: status

    integer :: naxis, bitpix
    integer, dimension(ndim) :: naxes
    real(dbl), dimension(ndim) :: crpix
    character(len=6), dimension(2), parameter :: keys = [ 'CTYPE1', 'CTYPE2' ]
    real(dbl) :: a,b,da,db,s,err
    real(dbl), dimension(2,2) :: mat,rmat
    type(AstroTrafoProj) :: tra
    character(len=5) :: un
    integer :: i

    if( type == " " ) then
       s = 1.0_dbl
       un = '[pix]'
    else
       s = -1.0_dbl
       un = '[deg]'
    end if

    a = cos(pa/rad)
    da = (dpa/rad)*sin(pa/rad)
    b = sin(pa/rad)
    db = (dpa/rad)*cos(pa/rad)
    err = rms/sqrt(2.0)

    ! clear possible previous calibration
    call wcsremove(fits,status)

    if( fcrpix ) then
       crpix = crpixels
    else
       call fits_get_img_param(fits,bitpix,naxis,naxes,status)
       if( status /= 0 ) return
       crpix = naxes/2.0_dbl
    end if

    call fits_update_key(fits,FITS_KEY_CREATOR,FITS_VALUE_CREATOR,FITS_COM_CREATOR,status)

    if( type == "GNOMONIC" )then
       call fits_update_key(fits,'CTYPE1', &
            'RA---TAN','the coordinate type for the first axis',status)
       call fits_update_key(fits,'CTYPE2', &
            'DEC--TAN','the coordinate type for the second axis',status)
    else if( type == " " ) then
       call fits_write_errmark
       do i = 1,size(keys)
          call fits_delete_key(fits,keys(i),status)
          if( status == FITS_KEYWORD_NOT_FOUND ) status = 0
       end do
       call fits_clear_errmark
    end if

    call fits_update_key(fits,'CRVAL1',crval(1),16, &
         un//' first axis value at reference pixel',status)
    call fits_update_key(fits,'CRVAL2',crval(2),16, &
         un//' second axis value at reference pixel',status)
    call fits_update_key(fits,'CRDER1',err,1, &
         un//' random error in first axis',status)
    call fits_update_key(fits,'CRDER2',err,1, &
         un//' random error in second axis',status)

    if( type == " " ) then
       call fits_update_key(fits,'CUNIT1','pix','units of first axis',status)
       call fits_update_key(fits,'CUNIT2','pix','units of second axis',status)
    else
       call fits_update_key(fits,'CUNIT1','deg','units of first axis',status)
       call fits_update_key(fits,'CUNIT2','deg','units of second axis',status)
    end if

    call trafo_init(tra,scale=sc,rot=pa,refl=refl)

    ! flip x-axis due to conversion rectangular -> spherical coordinates
    rmat(1,:) = [ -1.0_dbl,0.0_dbl ]
    rmat(2,:) = [ 0.0_dbl,1.0_dbl ]
    mat = matmul(rmat,tra%mat)

    call fits_update_key(fits,'CRPIX1',crpix(1),16, &
         'x-coordinate of reference pixel',status)
    call fits_update_key(fits,'CRPIX2',crpix(2),16, &
         'y-coordinate of reference pixel',status)
    call fits_update_key(fits,'CD1_1',mat(1,1),16, &
         'partial of first axis coordinate w.r.t. x',status)
    call fits_update_key(fits,'CD1_2',mat(1,2),16, &
         'partial of first axis coordinate w.r.t. y',status)
    call fits_update_key(fits,'CD2_1',mat(2,1),16, &
         'partial of second axis coordinate w.r.t. x',status)
    call fits_update_key(fits,'CD2_2',mat(2,2),16, &
         'partial of second axis coordinate w.r.t. y',status)

    call fits_write_comment(fits,BEGIN_ASTROMETRY,status)
    do i = 1, size(com)
       call fits_write_comment(fits,com(i),status)
    end do

    ! no fitting
    if( size(com) < 1 ) &
         call fits_write_comment(fits,"Astrometry calibration provided by user.",status)

    call fits_write_comment(fits,MUNIPACK_VERSION,status)
    call fits_write_comment(fits,END_ASTROMETRY,status)

  end subroutine wcsupdate


!-----------------------------------------------------------------------------


  subroutine wcsremove(fits,status)

    type(fitsfiles) :: fits
    integer, intent(in out) :: status
    integer :: i, nbegin, nend
    character(len=FLEN_CARD) :: card
    character(len=6), dimension(14) :: keys = [ &
         'CTYPE1', 'CTYPE2','CRVAL1','CRVAL2','CRDER1','CRDER2', &
         'CRPIX1','CRPIX2','CUNIT1','CUNIT2','CD1_1 ','CD1_2 ', &
         'CD2_1 ','CD2_2 ' ]

    call fits_write_errmark
    do i = 1, size(keys)
       call fits_delete_key(fits,keys(i),status)
       if( status == FITS_KEYWORD_NOT_FOUND ) status = 0
    end do
    call fits_clear_errmark

    i = 0
    nbegin = 0
    nend = 0
    card = ''
    do
       i = i + 1
       call fits_read_record(fits,i,card,status)
       if( status /= 0) exit
       if( card(9:) == BEGIN_ASTROMETRY ) nbegin = i
       if( card(9:) == END_ASTROMETRY ) nend = i
    end do
    if( status == 203 ) status = 0

    if( nbegin > 0 ) then
       do i = nbegin, nend
          ! looks strangle, but it's correct --
          ! records in header are moved up imediately following delete
          call fits_delete_record(fits,nbegin,status)
       end do
    end if

  end subroutine wcsremove

end module astrofits
