# Functions to read FITS array (includes image) and bintable files.
# All the functions are in this file in order of increasing dependency,
# so the "main," readFITS, is at the bottom.  Other functions in the file
# can be used for stand-alone header reading and parsing, image read, 
# and bintable read.
#
# Author: A. Harris, Univ. MD Astronomy
# Version 1.3  2008-07-19
# Comments and bug reports always welcome.
#
# Changes:
# 2008-07-19: Simplified cvec input for generating axis value vectors.  Included
#             changing readFITSim.r to return data frame with axis information,
#             simplified return values in list.  AH
# 
# Refs: http://fits.gsfc.nasa.gov/
#       Hanisch et al., Astr. Ap. 376, 359-380 (2001) 
#
# It is always a good idea to know what to expect in the files, perhaps
# by using fv or another fits-reader utility to read the headers or to 
# get an overview.  readFITS returns the parsed header for the image 
# or table, but not in a way that is optimized for humans to read.  See
# example below for extracting values for computations.  
#
######## Use ###############################################
#
# Return values and options:
#
# Call as
#  X <- readFITS(filename, numext=numext)
#    where numext is the number of the extension (e.g. numext=3 reads the 
#    third header and image pair in a file with mulitple images).  If 
#    no value is given for numext the default is the first data set. 
#    Modify the readFITS.r code below as needed for more complicated 
#    structures (e.g. primary header, data header, data, supplemental
#    header, supplemental header, data header, data, supplemental header, etc.
#
# Returns
#  X$imDat for image data with dimensions equal to data array.
#         Data are scaled by BSCALE and offset by BZERO or defaults
#         1, 0 as appropriate. 
# or  
#  X$col[[n]] to read dat from nth column in bintable.  If column
#         has depth then X$col[[n]] will be an array with number of
#         rows equal to data array rows and columns equal to "depth" 
#         of cells.
#         Data are scaled by TSCALEn and offset by TZEROn or defaults 
#         1, 0 as appropriate. 
#
# Summary(X) shows other variables the function returns.  Useful ones may
# be X$colNames, with names, and X$colUnits, with units, and X$axDat with
# axis information.  
# 
# Set phdu=0 in call if NAXIS!=0 and NAXIS1!=0 but header is in secondary 
# header unit:
#   readFITS(filname, phdu=0) 
#
# Function cvec.r makes axis vector from FITS file data
#
######## Notes ##############################################
#
# Functions below contain notes on known incompletes, mainly:
#  * Only support for arrays, images, and bintables (but other extensions
#    should be straightforward to code using array and binatble code as
#    a starting point).
#  * For bintables, bit, complex, and array descriptor data types
#    are not yet implemented.
# 
######## Examples ###########################################
#
# Example calls using data at http://fits.gsfc.nasa.gov/fits_samples.html:
# 
# A simple image
# source("readFITS.r")
# X <- readFITS("WFPC2ASSNu5780205bx.fits")
# summary(X)
# image(log(X$imDat+3), col=gray((0:128)/128), xaxt="n", yaxt="n")
# In the data the column data increments faster than row data.  This is
# good for presenting images.  To get back to [row, column] notation 
# for a 2D numerical data array, D <- t(D) (see r function 'perm' when
# the data arrays have more than 2 dimensions).
#
# A 4-deep data cube example 
# source("readFITS.r") 
# X <- readFITS("WFPC2u5780205r_c0fx.fits")
# image(X$imDat[,,2], col=gray((0:128)/128), xaxt="n", yaxt="n")
#
# The first and third images in a file with multiple images:
# X1 <- readFITS("NICMOSn4hk12010_mos.fits", numext=1)
# X3 <- readFITS("NICMOSn4hk12010_mos.fits", numext=3)
# 
# A spectrum from a bintable
# source("readFITS.r")
# X <- readFITS("IUElwp25637mxlo.fits")
# summary(X)
# X$colNames  # lists column names
# X$col[[1]]  # lists contents of the first data array column
# variable <- "NET"; yc <- which(X$colNames == variable)
# plot(X$col[[yc]], type="l", ylab=paste(X$colNames[yc], " [", X$colUnits[yc], "]", sep=""))
#
# Extract other keyword=value data from the header with e.g. 
#    tmp <- X$hdr[which(X$hdr=="BZERO")+1]
#      BZERO <- ifelse(length(tmp)!=1, 0, as.numeric(tmp))
# Here the 0 in the ifelse supplies the default parameter if the variable 
# is undefined.
#
# To make axes for an image:
# cv1 = cvec(1, X$axDat)
# cv2 = cvec(2, X$axDat)
# image(cv1, cv2/1.e9, X$imDat, col=heat.colors(128),
#      xlab=X$axDat$ctype[1], ylab=X$axDat$ctype[2])
#
#
#############  R functions #################################################

########################################
## Header parser

fitsHdrParse <- function(hdr_dat) {
# Function parses FITS header
#
# Refs: http://fits.gsfc.nasa.gov/
#       Hanisch et al., Astr. Ap. 376, 359-380 (2001) 
#
# Author: A. Harris, Univ. MD Astronomy, 3/20/08
#
# Requires:
#   called by readFITSheader   

  num <- 36  # number of card images
  z <- character(num)

# Break header data into card images
# Then separate comments and break at keyword = value
  for (i in 1:num) {
    start <- (i-1)*80+1
    end <- start+79
    z[i] <- strsplit(strsplit(substr(hdr_dat, start, end), "/")[[1]][1], "=")
  }

# Convert to simple vector
  z <- unlist(z)

# Strip spaces and single quotes
  for (i in 1:length(z)) {
    z[i] <- gsub(" ", "", z[i])
    z[i] <- gsub("'", "", z[i])
  }

# Check for end in header
  foundEnd <- as.logical(length(which(z=="END")))

  list(hdrInfo=z, foundEnd=foundEnd)
}

##################################################
## Read header, an integral number of header units

readFITSheader <- function(zz){
# Function gets FITS header
#
# Refs: http://fits.gsfc.nasa.gov/
#       Hanisch et al., Astr. Ap. 376, 359-380 (2001) 
#
# Author: A. Harris, Univ. MD Astronomy, 3/21/08
#
# Requires:
#   fitsHdrParse to parse header

  foundEnd <- FALSE
  hdr <- character()
  i <- 0
  maxunits <- 35*10  # maximum number of header units per header
  while(!foundEnd) {  
    # Each header is a set of 36 80-column card images
    tmp <- fitsHdrParse(readChar(zz, 2880)) 
    hdr <- c(hdr, tmp$hdrInfo)
    foundEnd <- tmp$foundEnd
    if (i > maxunits) stop("Haven't found END in header after ", 
        maxunits," header lines")
    i <- i+1
  }
  return(hdr)
}

##################################################
## Read FITS arrays

readFITSarray <- function(zz, hdr, verbose=FALSE) {
# Reader for FITS multidimentsional arrays, including 2D images
#
# Refs: http://fits.gsfc.nasa.gov/
#       Hanisch et al., Astr. Ap. 376, 359-380 (2001) 
#
# Author: A. Harris, Univ. MD Astronomy, 4/17/08
  # Changed to return axis information in data frame, deleted
  # return of individual vectors   AH 7/19/08
#

# Determine number of array dimensions
  naxis <- as.numeric(hdr[which(hdr=="NAXIS")+1]) 
# Find the right data type
  switch(hdr[which(hdr=="BITPIX")+1],
    "-64" = {bsize <- 8; btype <- numeric(); bsign <- FALSE},
    "-32" = {bsize <- 4; btype <- numeric(); bsign <- FALSE},
    "32" = {bsize <- 4; btype <- integer(); bsign <- TRUE},
    "16" = {bsize <- 2; btype <- integer(); bsign <- TRUE},
     "8" = {bsize <- 1; btype <- integer(); bsign <- FALSE}, 
           stop("Unknown BITPIX request in readFITSarray") )

# Put array information into vectors
  NAXISn <- integer(naxis)
  CRPIXn <- integer(naxis)
  CRVALn <- numeric(naxis)
  CDELTn <- numeric(naxis)
  CTYPEn <- character(naxis)
  CUNITn <- character(naxis)
  PTYPEn <- character(naxis)
  PSCALn <- numeric(naxis)
  PZEROn <- numeric(naxis)  
  numwords <- 1
  for(i in 1:naxis){
    tmp <- as.numeric(hdr[which(hdr==paste("NAXIS", i, sep=""))+1])
      numwords <- numwords*tmp
      NAXISn[i] <- tmp
    tmp <- hdr[which(hdr==paste("CRPIX", i, sep=""))+1]
      CRPIXn[i] <- ifelse(length(tmp)!=1, NA, as.numeric(tmp))
    tmp <- hdr[which(hdr==paste("CRVAL", i, sep=""))+1]
      CRVALn[i] <- ifelse(length(tmp)!=1, NA, as.numeric(tmp))
    tmp <- hdr[which(hdr==paste("CDELT", i, sep=""))+1]
      CDELTn[i] <- ifelse(length(tmp)!=1, NA, as.numeric(tmp))
    tmp <- hdr[which(hdr==paste("CTYPE", i, sep=""))+1]
      CTYPEn[i] <- ifelse(length(tmp)!=1, "", tmp)
    tmp <- hdr[which(hdr==paste("CUNIT", i, sep=""))+1]
      CUNITn[i] <- ifelse(length(tmp)!=1, "", tmp)
  }

  if(verbose){
    cat("naxis", naxis, "\n")
    cat(bsize, class(btype), "\n")
    cat(NAXISn, "\n", CRPIXn, "\n", CRVALn, "\n", 
        CDELTn, "\n", CTYPEn, "\n", CUNITn, "\n", 
        PTYPEn, "\n", PSCALn, "\n", PZEROn)
  }

# Read data into array.  Column data increments faster than row 
# data.  To get back to [row, column] notation for 2D array, D <- t(D).
  D <- array(readBin(zz, what=btype, n=numwords, size=bsize, 
            signed=bsign, endian="big"), dim=NAXISn)
# Finish reading block
  nbyte <- prod(NAXISn)*bsize
  nbyte <- ifelse(nbyte%%2880==0, 0, (1-(nbyte/2880)%%1)*2880)
  tmp <- readChar(zz, nbyte)

# Scale image to physical units if needed
  tmp <- hdr[which(hdr=="BUNIT")+1]
    BUNIT <- ifelse(length(tmp)!=1, "", tmp)
  tmp <- hdr[which(hdr=="BSCALE")+1]
    BSCALE <- ifelse(length(tmp)!=1, 1, as.numeric(tmp))
  tmp <- hdr[which(hdr=="BZERO")+1]
    BZERO <- ifelse(length(tmp)!=1, 0, as.numeric(tmp))
  if (BSCALE!=1 || BZERO!=0) D <- D*BSCALE + BZERO

# Make data frame with axis data
  axDat <- data.frame(CRPIXn, CRVALn, CDELTn, dim(D), CTYPEn, CUNITn)
  names(axDat) <- c('crpix', 'crval', 'cdelt', 'len', 'ctype', 'cunit')

# Return structure with data and image information
  list(imDat=D, axDat=axDat, hdr=hdr)
}

##################################################
## Read FITS BINTABLE

readFITSbintable <- function(zz, hdr, verbose=FALSE) {
# Reader for FITS BINTABLE arrays
#
# Refs: http://fits.gsfc.nasa.gov/
#       Hanisch et al., Astr. Ap. 376, 359-380 (2001) 
#
# Author: A. Harris, Univ. MD Astronomy, 4/22/08
#
# Changed stop to warning if pcount and gcount are incorrect 5/4/08 AH
#
# Known incomplete:
#  * Data types X, C, M, P unimplemented (bit, complex, double complex, 
#      array descriptor)


# Determine array dimensions
  naxis1 <- as.numeric(hdr[which(hdr=="NAXIS1")+1])     # bytes per row
  naxis2 <- as.numeric(hdr[which(hdr=="NAXIS2")+1])     # rows
  tfields <-  as.numeric(hdr[which(hdr=="TFIELDS")+1])  # columns

# Get additional memory allocation information
  tmp <- hdr[which(hdr=="PCOUNT")+1]
    pcount <- ifelse(length(tmp)!=1, 0, as.numeric(tmp))
    if (pcount != 0) warning("pcount must be 0 in a bintable")
  tmp <- hdr[which(hdr=="GCOUNT")+1]
    gcount <- ifelse(length(tmp)!=1, 1, as.numeric(tmp))
    if (gcount != 1) warning("gcount must be 1 in a bintable")

# Put array information into vectors
  TFORMn <- character(tfields)
  TTYPEn <- character(tfields)  # field names
  TUNITn <- character(tfields)
  TNULLn <- integer(tfields)
  TSCALn <- numeric(tfields)
  TZEROn <- numeric(tfields)
  TDISPn <- character(tfields)
  THEAPn <- integer(tfields)
  TDIMn <- character(tfields)
  for (i in 1:tfields) {
    tmp <- hdr[which(hdr==paste("TFORM", i, sep=""))+1]
      TFORMn[i] <- tmp
    tmp <- hdr[which(hdr==paste("TTYPE", i, sep=""))+1]
      TTYPEn[i] <- ifelse(length(tmp)!=1, "", tmp)
    tmp <- hdr[which(hdr==paste("TUNIT", i, sep=""))+1]
      TUNITn[i] <- ifelse(length(tmp)!=1, "", tmp)
    tmp <- hdr[which(hdr==paste("TNULL", i, sep=""))+1]
      TNULLn[i] <- ifelse(length(tmp)!=1, NA, as.numeric(tmp))
    tmp <- hdr[which(hdr==paste("TSCAL", i, sep=""))+1]
      TSCALn[i] <- ifelse(length(tmp)!=1, 1, as.numeric(tmp))
    tmp <- hdr[which(hdr==paste("TZERO", i, sep=""))+1]
      TZEROn[i] <- ifelse(length(tmp)!=1, 0, as.numeric(tmp))
    tmp <- hdr[which(hdr==paste("TDISP", i, sep=""))+1]
      TDISPn[i] <- ifelse(length(tmp)!=1, "", tmp)
    tmp <- hdr[which(hdr==paste("THEAP", i, sep=""))+1]
      THEAPn[i] <- ifelse(length(tmp)!=1, NA, as.numeric(tmp))
    tmp <- hdr[which(hdr==paste("TDIM", i, sep=""))+1]
      TDIMn[i] <- ifelse(length(tmp)!=1, "", tmp)
  }

# Work out formats and set up storage for each column.
  bsize <- integer(tfields)
  btype <- integer(tfields)
  bsign <- logical(tfields) 
  mult <- integer(tfields)
  col <- vector("list", tfields)
  for (i in 1:tfields) {
    # Separate multiplier and type.
    nc <- nchar(TFORMn[i])
    tmp <- substr(TFORMn[i], 1, nc-1)
      mult[i] <- ifelse(tmp=="", 1, as.numeric(tmp))
    # Format; btype: 1 character, 2 logical, 3 integer, 4 numeric, 5 complex
    form <- tolower(substr(TFORMn[i], nc, nc))
    switch (form, 
        l = {bsize[i] <- 1; btype[i] <- 2; bsign[i] <- FALSE},  # logical
        b = {bsize[i] <- 1; btype[i] <- 3; bsign[i] <- FALSE},  # 8-bit unsigned int
        i = {bsize[i] <- 2; btype[i] <- 3; bsign[i] <- TRUE},   # 16-bit signed int (single)
        j = {bsize[i] <- 4; btype[i] <- 3; bsign[i] <- TRUE},   # 32-bit signed int (double)
        a = {bsize[i] <- 1; btype[i] <- 1; bsign[i] <- FALSE},  # 8-bit character
        e = {bsize[i] <- 4; btype[i] <- 4; bsign[i] <- FALSE},  # 32-bit float (single)
        d = {bsize[i] <- 8; btype[i] <- 4; bsign[i] <- FALSE},  # 64-bit float (double)
        stop("X, C, M, P not yet implemented \n") 
    )
    # Set up storage arrays: rows = number of table rows, columns = 
    # multiplier (depth) of the cells in each column.  Characters are an
    # exception since they return strings. 
    if (btype[i] == 1) {  # character reads
      col[[i]] <- array("", dim=c(naxis2, 1))
    } else {
      col[[i]] <- array(switch(btype[i], "", FALSE, NA, NA, NA),
                     dim=c(naxis2, mult[i]))
    }
  }
 
# Read data, row by row
  for(i in 1:naxis2) {
    for (j in 1:tfields) {
      if (btype[j] <= 2) {  # character reads
        col[[j]][i,] <- readChar(zz, nchars=mult[j])
      } else {
        what <- switch(btype[j], character(), logical(), integer(), 
                    numeric(), complex())
        col[[j]][i,] <- readBin(zz, what=what, n=mult[j], 
                       size=bsize[j], signed=bsign[j], endian="big")
      }
    }
  }
# Finish reading block
  nbyte <- naxis1*naxis2
  nbyte <- ifelse(nbyte%%2880==0, 0, (1-(nbyte/2880)%%1)*2880)
  tmp <- readChar(zz, nbyte)

# Clean up before returning
  for (i in 1:tfields) {
  # Apply scaling and offset where appropriate
    if (btype[i]>=3 && (TSCALn[i]!=1 || TZEROn[i]!=0)) {
      col[[i]] <- (col[[i]])*TSCALn[i] + TZEROn[i]
    } 
  # Convert 1D arrays to vectors for easier plotting  
    if (nrow(col[[i]])==1 || ncol(col[[i]])==1 || btype[i]<=2) {
      col[[i]] <- as.vector(col[[i]])
    }
  # Terminate character strings that end in ascii 0
    if (btype[i]<=2) {
      lcol <- length(col[[i]])
      txttmp <- character(lcol)
      for (j in 1:lcol) txttmp[j] <- strsplit((col[[i]][j]), 0)
      col[[i]] <- unlist(txttmp) 
    }
  }

# Return data structure: column data plus ancillary data.  
  list(col=col, hdr=hdr, colNames=TTYPEn, colUnits=TUNITn,
      TNULLn=TNULLn, TSCALn=TSCALn, TZEROn=TZEROn, TDISPn=TDISPn) 
}

##################################################
## Main segment for reader function

readFITS <- function(filename="R.fits", numext=1, phdu=1) {
# Simple reader for FITS bintable, array, and image files
#
# Refs: http://fits.gsfc.nasa.gov/
#       Hanisch et al., Astr. Ap. 376, 359-380 (2001) 
#
# Author: A. Harris, Univ. MD Astronomy, 4/22/08
#
# Requires:
#   readFITSheader to get header data 
#   fitsHdrParse to parse header (called by readFITSheader)
# 
# Set phdu=0 in call if NAXIS!=0 but header is in secondary header unit 

# Open file
  zz <- file(filename, "rb")

# Read primary header unit
  hdr <- readFITSheader(zz)

# Determine number of array dimensions, select appropriate extension
  tmp <- hdr[which(hdr=="NAXIS")+1] 
    if (tmp == "") tmp <- '0'
  if (as.numeric(tmp) > 0) {   
    tmp <- as.numeric(hdr[which(hdr=="NAXIS1")+1]) 
      if (tmp == 0) phdu <- 0
  }
  if((as.numeric(tmp)*phdu) > 0) { # Data are an array off primary header
    D <- readFITSarray(zz, hdr) 
    close(zz)
    return (D)
  } else {                         # Data are off an extension header
    for (i in 1:numext) {          # Brute-force read to the numext-th header
      hdr <- readFITSheader(zz)
      switch(tolower(hdr[which(hdr=="XTENSION")+1]),
           bintable = {D <- readFITSbintable(zz, hdr)},
           image = {D <- readFITSarray(zz, hdr)},
           stop("Current version only supports bintable and image")
      )
    }
    close(zz)  
    return(D)
  } 
} 

########################################################################
# 
# Use is analogous to readFITS, but it applies only to fits binary
# table, and returns a data frame
#
# Written by Eric H. Neilsen, Jr., FNAL
#
# For example, to plot the positions of objects in an SDSS frame:
#
# > objs <- readFrameFromFITS("http://das.sdss.org/DR6/data/imaging/inchunk_best/stripe15_mu551256_1/3/tsObj-003836-3-41-0254.fit",1)
# > plot(objs$ra,objs$dec,pch=20)
#

nameFrameColumn <- function(nframe,datacol,colname) {
  if (length(dim(datacol)) < 1) {
    names(nframe) <- colname
  } else {
    names(nframe) <- sprintf("%s.%d",colname,c(1:length(datacol[1,])))
  } 
  nframe
}

readFrameFromFITS <- function(fileName,hdu) {
  fitsContent <- readFITS(fileName,hdu)
  dataFrame <- data.frame(fitsContent$col[[1]])
  newframe <- nameFrameColumn(dataFrame,fitsContent$col[[1]],fitsContent$colNames[1])
  for (i in 2:length(fitsContent$colNames)) {
    newframe <- data.frame(fitsContent$col[[i]])
    newframe <- nameFrameColumn(newframe,fitsContent$col[[i]],fitsContent$colNames[i])
    dataFrame <- c(dataFrame,newframe)
  }
  dataFrame
}

##################################################

cvec = function(nax=1, axDat) {

# Function makes axis vector from FITS file data
  ## nax is axis number
  ##  axDat is axis data from data frame produced by readFITSim
  ##  crpix, crval, cdelt, ctype, cunit have the usual FITS meanings
  ##
  ## AH 7/19/08

  cvec = seq(
    from = (1-axDat$crpix[nax])*axDat$cdelt[nax] + axDat$crval[nax],
    by = axDat$cdelt[nax],
    length = axDat$len[nax]
    )
  return(cvec)

}