
import numpy
import pytz
import datetime
import dateutil.parser
import logging
import re

import astropy.io.fits as pyfits

from tkp.accessors.dataaccessor import DataAccessor
from tkp.utility.coordinates import WCS

logger = logging.getLogger(__name__)

class FitsImage(DataAccessor):
    """
    Use PyFITS to pull image data out of a FITS file.

    Provide standard attributes, as per :class:`DataAccessor`. In addition, we
    provide a ``telescope`` attribute if the FITS file has a ``TELESCOP``
    header.
    """
    def __init__(self, url, plane=None, beam=None, hdu_index=0):
        super(FitsImage, self).__init__()
        self.url = url
        self.header = self._get_header(hdu_index)
        self.wcs = self.parse_coordinates()
        self.data = self.read_data(hdu_index, plane)
        self.taustart_ts, self.tau_time = self.parse_times()
        self.freq_eff, self.freq_bw = self.parse_frequency()
        self.pixelsize = self.parse_pixelsize()
        if beam:
            (bmaj, bmin, bpa) = beam
        else:
            (bmaj, bmin, bpa) = self.parse_beam()
        self.beam = self.degrees2pixels(
                bmaj, bmin, bpa, self.pixelsize[0], self.pixelsize[1]
            )
        self.centre_ra, self.centre_decl = self.calculate_phase_centre()

        # Bonus attribute
        if 'TELESCOP' in self.header:
            self.telescope = self.header['TELESCOP']

    def _get_header(self, hdu_index):
        with pyfits.open(self.url) as hdulist:
            hdu = hdulist[hdu_index]
        return hdu.header.copy()

    def read_data(self, hdu_index, plane):
        """
        Read and store data from our FITS file.

        NOTE: PyFITS reads the data into an array indexed as [y][x]. We
        take the transpose to make this more intuitively reasonable and
        consistent with (eg) ds9 display of the FitsFile. Transpose back
        before viewing the array with RO.DS9, saving to a FITS file,
        etc.
        """
        with pyfits.open(self.url) as hdulist:
            hdu = hdulist[hdu_index]
            data = numpy.float64(hdu.data.squeeze())
        if plane is not None and len(data.shape) > 2:
            data = data[plane].squeeze()
        n_dim = len(data.shape)
        if n_dim != 2:
            logger.warn("Loaded datacube with %s dimensions, assuming Stokes I and taking plane 0" % n_dim)
            data = data[0, :, :]
        data = data.transpose()
        return data


    def parse_coordinates(self):
        """Returns a WCS object"""
        header = self.header
        wcs = WCS()
        try:
            wcs.crval = header['crval1'], header['crval2']
            wcs.crpix = header['crpix1'] - 1, header['crpix2'] - 1
            wcs.cdelt = header['cdelt1'], header['cdelt2']
        except KeyError:
            msg = "Coordinate system not specified in FITS"
            logger.error(msg)
            raise TypeError(msg)
        try: