#!/usr/bin/env python3

"""
Script to compare the output of two PyBDSF runs. The generated FITS images,
FITS catalog, and text files will be compared.
"""

import argparse
import filecmp
import logging
import os
import sys
import textwrap
import warnings

import astropy.io.fits
import numpy as np
import difflib

import astropy.wcs
import re


# Suppress UserWarning messages (mostly generated by matplotlib) and
# FITSFixedWarning messages (generated by harmless WCS header issues)
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=astropy.wcs.FITSFixedWarning)

# Global Logger object, will be initialized in init_logger()
logger = None


def check_all_files_present(dcmp):
    """
    Checks that all files are present in both compared directories.

    :param dircmp dcmp: the result of `dircmp(left, right)`
    :return bool: True if all files are present, False if not
    """
    logger.debug("Comparing directory contents of '%s' and '%s'", dcmp.left, dcmp.right)
    agree = True
    if dcmp.left_only or dcmp.right_only:
        agree = False
        msg = "The following files are present in '%s' but not in '%s': %s"
        if dcmp.left_only:
            logger.error(msg, dcmp.left, dcmp.right, dcmp.left_only)
        if dcmp.right_only:
            logger.error(msg, dcmp.right, dcmp.left, dcmp.right_only)
    return agree


def text_files_are_similar(file1, file2, verbosity):
    """
    Compare text files ``file1`` and ``file2``, using `difflib`.
    :param str file1: First text file
    :param str file2: Second text file
    :param int verbosity: Verbosity level, higher is more verbose
    :return bool: True if text files are similar, else False
    """
    logger.debug("Comparing '%s' and '%s'", file1, file2)
    d = difflib.Differ()
    with open(file1) as f1, open(file2) as f2:
        diff = [
            line for line in d.compare(f1.readlines(), f2.readlines())
            if line.startswith(('-', '+')) and not line[2:].startswith('#')
        ]
    agree = False if diff else True
    if agree:
        logger.info("Text files '%s' and '%s' are similar", file1, file2)
    else:
        if verbosity > 0:
            logger.error(
                "Text files '%s' and '%s' differ:\n%s",
                file1,
                file2,
                diff
            )
        else:
            logger.error(
                "Text files '%s' and '%s' differ", file1, file2
            )
    return agree


def fits_files_are_similar(file1, file2, rtol, verbosity):
    """
    Compare FITS files ``file1`` and ``file2``, using `FITSDiff` in `astropy`.
    FITS files are considered similar if none of the values have a relative
    difference exceeding `rtol`.
    :param str file1: First FITS file
    :param str file2: Second FITS file
    :param float rtol: Relative tolerance threshold
    :param int verbosity: Verbosity level, higher is more verbose
    :return bool: True if FITS files are similar, else False
    """
    logger.debug("Comparing '%s' and '%s'", file1, file2)
    diff = astropy.io.fits.FITSDiff(file1, file2, rtol=rtol, ignore_keywords=["COMMENT"])
    agree = diff.identical
    if agree:
        logger.info("FITS files '%s' and '%s' are similar", file1, file2)
    else:
        if verbosity > 0:
            logger.error(
                "FITS files '%s' and '%s' differ (rtol=%s):\n%s",
                file1,
                file2,
                rtol,
                diff.report(),
            )
        else:
            logger.error(
                "FITS files '%s' and '%s' differ (rtol=%s)", file1, file2, rtol
            )
    return agree


def image_files_are_similar(file1, file2, rtol, verbosity):
    """
    Compare two FITS image files, by comparing the RMS noise values.
    Files are considered similar if the the noise ratio differs less than
    ``rtol`` from 1.
    :param str file1: Filename of first FITS image
    :param str file2: Filename of second FITS image
    :param float rtol: Relative tolerance threshold
    :param int verbosity: Verbosity level, higher is more verbose
    :return bool: True if images files are similar, else False
    """

    def image_noise(fname):
        """
        Calculate the RMS value of the noise in the image file ``fname``.
        :param str fname: Filename of the FITS image
        :return float: RMS noise in the image in Jy/b
        """
        image = FITSImage(fname)
        image.calc_noise()
        return image.noise

    logger.debug("Comparing noise in images '%s' and '%s'", file1, file2)
    noise1 = image_noise(file1)
    noise2 = image_noise(file2)
    agree = np.isclose(noise1, noise2, rtol=rtol)
    if agree:
        logger.info("Image files '%s' and '%s' are similar", file1, file2)
    if not agree:
        if verbosity > 0:
            logger.error(
                "Image files '%s' and '%s' differ (rtol=%s)\n"
                "Noise: %.3f mJy/b vs %.3f mJy/b",
                file1,
                file2,
                rtol,
                1000 * noise1,
                1000 * noise2,
            )
        else:
            logger.error(
                "Image files '%s' and '%s' differ (rtol=%s)", file1, file2, rtol
            )
    return agree


def compare_results(dcmp, rtol, verbosity, check_images=True, check_catalogs=False, check_other=False):
    """
    Compare the results between the files produced by each PyBDSF run.
    :param dircmp dcmp: the directory contents of the two workflows
    :param float rtol: relative tolerance threshold
    :param int verbosity: verbosity level, higher means more verbose
    """
    # Initialize logger and set log level depending on verbosity level
    init_logger(verbosity)

    logger.info("*** Comparing results in '%s' and '%s' ***", dcmp.left, dcmp.right)
    logger.info(f"Checking images: {check_images}")
    logger.info(f"Checking catalogs: {check_catalogs}")
    logger.info(f"Checking other: {check_other}")
    agree = True
    for dname, sub_dcmp in dcmp.subdirs.items():
        logger.debug(
            "dname: %s, sub_dcmp.left: %s, sub_dcmp.right: %s",
            dname,
            sub_dcmp.left,
            sub_dcmp.right,
        )
        check_all_files_present(sub_dcmp)
        for fname in sub_dcmp.common_files:
            left = os.path.join(sub_dcmp.left, fname)
            right = os.path.join(sub_dcmp.right, fname)
            root, ext = os.path.splitext(fname)
            logger.debug("fname: %s, root: %s, ext: %s", fname, root, ext)
            if check_images:
                if "fits" in ext.lower() and "catalogues" not in dname:
                    if not image_files_are_similar(
                        left, right, rtol=rtol, verbosity=verbosity
                    ):
                        agree = False
            if check_catalogs:
                if "fits" in ext.lower() and "catalogues" in dname:
                    if not fits_files_are_similar(
                        left, right, rtol=rtol, verbosity=verbosity
                    ):
                        agree = False
            if check_other:
                if "fits" not in ext.lower():
                    if not text_files_are_similar(
                        left, right, verbosity=verbosity
                    ):
                        agree = False
    return agree


class FITSImage(object):
    """
    The FITSImage class is used for processing/manipulation of FITS images

    This class is a simplified version of the FITSImage class located at:
        https://git.astron.nl/RD/rapthor/-/blob/master/rapthor/lib/fitsimage.py

    Parameters
    ----------
    imagefile : str
        Filename of the FITS file
    """
    def __init__(self, imagefile):
        self.imagefile = imagefile
        self.name = 'FITS_image'
        self.vertices_file = None
        self.scale = 1.
        self.shift = 0.
        self.noise = 1.

        self.header = astropy.io.fits.getheader(imagefile)
        self.find_beam()
        self.find_freq()
        self.flatten()
        self.ra = self.img_hdr['CRVAL1']
        self.dec = self.img_hdr['CRVAL2']

    def find_beam(self):
        """
        Find the primary beam headers following AIPS convenction
        """
        if ('BMAJ' in self.header) and ('BMIN' in self.header) and ('PA' in self.header):
            pass
        elif 'HISTORY' in self.header:
            for hist in self.header['HISTORY']:
                if 'AIPS   CLEAN BMAJ' in hist:
                    # remove every letter from the string
                    bmaj, bmin, pa = re.sub(' +', ' ', re.sub('[A-Z ]*=', '', hist)).strip().split(' ')
                    self.header['BMAJ'] = float(bmaj)
                    self.header['BMIN'] = float(bmin)
                    self.header['BPA'] = float(pa)
        self.beam = [float(self.header['BMAJ']), float(self.header['BMIN']), float(self.header['BPA'])]

    def find_freq(self):
        """
        Find frequency value in most common places of a fits header
        """
        self.freq = None
        if not self.header.get('RESTFREQ') is None and not self.header.get('RESTFREQ') == 0:
            self.freq = float(self.header.get('RESTFREQ'))
        elif not self.header.get('FREQ') is None and not self.header.get('FREQ') == 0:
            self.freq = float(self.header.get('FREQ'))
        else:
            for i in range(5):
                type_s = self.header.get('CTYPE%i' % i)
                if type_s is not None and type_s[0:4] == 'FREQ':
                    self.freq = float(self.header.get('CRVAL%i' % i))

    def flatten(self):
        """ Flatten a fits file so that it becomes a 2D image. Return new header and data """
        fits_data = astropy.io.fits.getdata(self.imagefile)

        naxis = self.header['NAXIS']
        if naxis < 2:
            raise RuntimeError('Can\'t make map from this')
        if naxis == 2:
            self.img_hdr = self.header
            self.img_data = fits_data
        else:
            w = astropy.wcs.WCS(self.header)
            wn = astropy.wcs.WCS(naxis=2)
            wn.wcs.crpix[0] = w.wcs.crpix[0]
            wn.wcs.crpix[1] = w.wcs.crpix[1]
            wn.wcs.cdelt = w.wcs.cdelt[0:2]
            wn.wcs.crval = w.wcs.crval[0:2]
            wn.wcs.ctype[0] = w.wcs.ctype[0]
            wn.wcs.ctype[1] = w.wcs.ctype[1]

            header = wn.to_header()
            header["NAXIS"] = 2
            header["NAXIS1"] = self.header['NAXIS1']
            header["NAXIS2"] = self.header['NAXIS2']
            header["FREQ"] = self.freq
            header['RESTFREQ'] = self.freq
            header['BMAJ'] = self.beam[0]
            header['BMIN'] = self.beam[1]
            header['BPA'] = self.beam[2]
            copy = ('EQUINOX', 'EPOCH')
            for k in copy:
                r = self.header.get(k)
                if r:
                    header[k] = r

            dataslice = []
            for i in range(naxis, 0, -1):
                if i <= 2:
                    dataslice.append(np.s_[:],)
                else:
                    dataslice.append(0)
            self.img_hdr = header
            self.img_data = fits_data[tuple(dataslice)]
        self.min_value = float(np.nanmin(self.img_data))
        self.max_value = float(np.nanmax(self.img_data))
        self.mean_value = float(np.nanmean(self.img_data))
        self.median_value = float(np.nanmedian(self.img_data))

    def calc_noise(self, sampling=4):
        """
        Return the rms of all the pixels in an image
        sampling : sampling interval to use to speed up the noise calculation (e.g.,
                   sampling = 4 means use every forth pixel)
        """
        sampling = int(sampling)
        if sampling < 1:
            sampling = 1
        data = self.img_data[::sampling]  # sample array
        data = data[np.isfinite(data)]
        rms = np.nanstd(data)
        self.noise = float(rms)


def parse_arguments():
    """
    Parse command-line arguments.
    :return argparse.Namespace: parsed command-line arguments
    """
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=textwrap.dedent(
            """
            Compare the results of two PyBDSF runs.
            """
        ),
    )
    parser.add_argument("first", help="Directory containing results of the first run")
    parser.add_argument("second", help="Directory containing results of the second run")
    parser.add_argument(
        "--rtol",
        type=float,
        default=1e-3,
        help="Relative tolerance threshold for comparing data (default=%(default)s)",
    )
    verb_group = parser.add_mutually_exclusive_group()
    verb_group.add_argument(
        "-q", "--quiet", action="store_true", help="Quiet mode; only report errors"
    )
    verb_group.add_argument(
        "-v",
        "--verbose",
        action="count",
        default=0,
        help="Verbose mode; multiple -v options increase verbosity",
    )

    args = parser.parse_args()

    return args


def init_logger(verbosity):
    """
    Initialize a global logger object.
    Set the loglevel depending on the `verbosity`.
    :param int verbosity: Verbosity level, higher means more verbose logging
    """
    global logger
    logger = logging.getLogger()

    ch = logging.StreamHandler()
    if verbosity == -1:
        ch.setLevel(logging.ERROR)
    elif verbosity == 0:
        ch.setLevel(logging.WARNING)
    elif verbosity == 1:
        ch.setLevel(logging.INFO)
    elif verbosity >= 2:
        ch.setLevel(logging.DEBUG)

    fmt = "%(levelname)s [%(name)s]: %(message)s"
    ch.setFormatter(logging.Formatter(fmt))
    logger.addHandler(ch)


def main(args):
    """
    Compare the images, text files and catalogs in the ``args.first`` directory
    to those in the ``args.second`` directory.
    :param argparse.Namespace args: parsed command-line
    :return bool: True if all files in ``args.first`` are similar to those in
                  ``args.second``, else False.
    """
    # Set level of verbosity to -1 if command-line option `-q` is given, else
    # to the number of times the command-line option `-v` is given.
    verbosity = -1 if args.quiet else args.verbose

    agree = True
    if not compare_results(
        filecmp.dircmp(args.first, args.second),
        rtol=args.rtol,
        verbosity=verbosity,
    ):
        agree = False

    return agree


if __name__ == "__main__":
    success = main(parse_arguments())
    sys.exit(0 if success else 1)
