""" Classes and methods for handling astronomical data """

import math

from astropy.io import fits
from astropy.table import Table
from astropy import wcs
from astropy.nddata import Cutout2D

from pathlib import Path
import numpy as np

import galprime

from tqdm import tqdm


class CutoutList:
    
    def __init__(self, cutouts=None, cutout_data=[], metadata={}):
        if cutouts is None:
            cutouts = []
        if metadata == None:
            metadata = {}
        
        self.cutouts = cutouts              # List to hold cutouts
        self.cutout_data = cutout_data      # This is the container for associated data to each cutout
        self.metadata = metadata            # Any extra needed metadata
    
        metadata["N_CUTOUTS"] = len(cutouts)
    
    def update_medatada(self):
        self.metadata["N_CUTOUTS"] = len(self.cutouts)
        
    def cutouts_to_file(self, filename="cutouts.fits", overwrite=False):
        
        if len(self.cutouts) == 0:
            return

        out_hdulist = fits.HDUList()
        head = fits.Header()
        
        # Add the metadata (if any) to the header
        for n in self.metadata:
            head[n] = self.metadata[n]

        # Now add the individual cutouts to the HDUList with associated data in the individual headers
        for i, cutout in enumerate(self.cutouts):
            this_head = head.copy()
            try:
                this_head.update(self.cutout_data[i])
            except ValueError:
                pass
            out_hdulist.append(fits.ImageHDU(data=cutout, header=this_head))
        
        out_hdulist.writeto(filename, overwrite=overwrite)



def get_image_filenames(images_directory):
    """
    Retrieves a list of all available filenames for a given directory, and a given band.

    :param images_directory: The top-level directory from which to get image filenames from.
    :type images_directory: str
    :param image_band: The image band, defaults to "i"
    :type image_band: str, optional
    :param check_band: Check to make sure the images are of a certain band, defaults to False.
        WARNING: Optimized for HSC filename format (ex: HSC-I_9813_4c3.fits).
    :type check_band: bool, optional
    """
    image_filenames = []
    images = Path(images_directory).rglob('*.fits')
    for image in images:
        image_filenames.append(str(image))
    return image_filenames


def get_closest_psf(psfs, ra, dec, ra_key="RA", dec_key="DEC"):
    ras = np.array([psfs.cutout_data[i]["RA"] for i in range(len(psfs.cutout_data))])
    decs = np.array([psfs.cutout_data[i]["DEC"] for i in range(len(psfs.cutout_data))])
    
    closest = np.argmin(((ras - ra) ** 2 + (decs - dec) ** 2))
        
    return psfs.cutouts[closest]


def get_wcs(fits_filename):
    """
        Finds and returns the WCS for an image. If Primary Header WCS no good, searches each index until a good one
        is found. If none found, raises a ValueError

        :param fits_filename: The filename for the FITS file.
        :type fits_filename: str
        :return: The WCS object.
        :raises: ValueError
    """
    # Try just opening the initial header
    wcs_init = wcs.WCS(fits_filename)
    ra, dec = wcs_init.axis_type_names
    if ra.upper() == "RA" and dec.upper() == "DEC":
        return wcs_init

    else:
        hdu_list = fits.open(fits_filename)
        for n in hdu_list:
            try:
                wcs_slice = wcs.WCS(n.header)
                ra, dec = wcs_slice.axis_type_names
                if ra.upper() == "RA" and dec.upper() == "DEC":
                    return wcs_slice
            except:
                continue
        hdu_list.close()

    raise ValueError


def gen_backgrounds(image_directory, config=None):
    n_cutouts = 10 if config is None else config["N_BACKGROUNDS"]
    
    image_files = get_image_filenames(image_directory)

    print(image_directory, image_files, n_cutouts)


def get_angular_size_dist(z, H0=71, WM=0.27):
    """
    Return the angular size distance in Megaparsecs.
    (Stripped down version of Cosmocalc by Ned Wright and Tom Aldcroft (aldcroft@head.cfa.harvard.edu))

    :param z: The redshift.
    :type z: float
    :param H0: The Hubble constant, defaults to 71 km/s/Mpc
    :type H0: float, optional
    :param WM: matter density parameter, defaults to 0.27
    :type WM: float, optional
    """
    try:
        c = 299792.458  # velocity of light in km/sec

        if z > 100:
            z /= 299792.458  # Values over 100 are in km/s

        WV = 1.0 - WM - 0.4165 / (H0 * H0)  # Omega(vacuum) or lambda
        age = 0.0  # age of Universe in units of 1/H0

        h = H0 / 100.
        WR = 4.165E-5 / (h * h)  # includes 3 massless neutrino species, T0 = 2.72528
        WK = 1 - WM - WR - WV
        az = 1.0 / (1 + 1.0 * z)
        n = 1000  # number of points in integrals
        for i in range(n):
            a = az * (i + 0.5) / n
            adot = math.sqrt(WK + (WM / a) + (WR / (a * a)) + (WV * a * a))
            age += 1. / adot

        DCMR = 0.0

        # do integral over a=1/(1+z) from az to 1 in n steps, midpoint rule
        for i in range(n):
            a = az + (1 - az) * (i + 0.5) / n
            adot = math.sqrt(WK + (WM / a) + (WR / (a * a)) + (WV * a * a))
            DCMR = DCMR + 1. / (a * adot)

        DCMR = (1. - az) * DCMR / n

        # tangential comoving distance
        ratio = 1.0
        x = math.sqrt(abs(WK)) * DCMR
        if x > 0.1:
            if WK > 0:
                ratio = 0.5 * (math.exp(x) - math.exp(-x)) / x
            else:
                ratio = math.math.sin(x) / x
        else:
            y = x * x
            if WK < 0:
                y = -y
            ratio = 1. + y / 6. + y * y / 120.
        DCMT = ratio * DCMR
        DA = az * DCMT
        Mpc = lambda x: c / H0 * x
        DA_Mpc = Mpc(DA)

        return DA_Mpc
    except:
        raise ValueError


def get_image_filenames(images_directory):
    """
    Retrieves a list of all available filenames for a given directory, and a given band.

    :param images_directory: The top-level directory from which to get image filenames from.
    :type images_directory: str
    :param image_band: The image band, defaults to "i"
    :type image_band: str, optional
    :param check_band: Check to make sure the images are of a certain band, defaults to False.
        WARNING: Optimized for HSC filename format (ex: HSC-I_9813_4c3.fits).
    :type check_band: bool, optional
    """
    image_filenames = []
    images = Path(images_directory).rglob('*.fits')
    for image in images:
        image_filenames.append(str(image))
    return image_filenames


def gen_cutouts(image_directory, config=None, buffer=10, progress_bar=False):
    n_cutouts = 10 if config is None else config["N_BACKGROUNDS"]
    cutout_size = 71 if config is None else config["SIZE"]
    if cutout_size % 2 == 0:
        cutout_size += 1 
    
    hw = cutout_size / 2
    
    image_files = get_image_filenames(image_directory)
    
    cutouts, cutout_data = [], []
    iterator = tqdm(range(n_cutouts), desc="Making Image Cutouts: ") if progress_bar else range(n_cutouts)
    for i in iterator:
        this_fn = np.random.choice(image_files)
        
        with fits.open(this_fn) as HDUList:
            for i in range(len(HDUList)):
                data = HDUList[i].data
                if type(data) == np.ndarray:
                    this_wcs = wcs.WCS(HDUList[i].header)
                    break
        x, y = np.random.randint(hw + buffer, data.shape[0] - hw - buffer), np.random.randint(hw + buffer, data.shape[1] - hw - buffer)
        ra, dec = this_wcs.wcs_pix2world(x,y,0)
        cutout = Cutout2D(data=data, position=[x,y], size=cutout_size, wcs=this_wcs)
        
        cutouts.append(cutout.data)
        
        cutout_data.append({"IMG":this_fn, "X":x, "Y":y, "RA":float(ra), "DEC":float(dec) })
    
    return CutoutList(cutouts=cutouts, cutout_data=cutout_data)


def cutouts_from_file(filename):
    cutouts, cutout_datasets = [], []
    with fits.open(filename) as HDUList:
        for hdu in HDUList:
            try:
                cutout, cutout_data = hdu.data, hdu.header
                cutouts.append(cutout)
                cutout_datasets.append(cutout_data)
            except:
                continue
        
    return CutoutList(cutouts=cutouts, cutout_data=cutout_datasets)


def ordered_dict(names, params):
    out_dict = {}
    for i, n in enumerate(names):
        out_dict[n] = params[i]
    return out_dict


def dict_extend(d, extension="_1"):
    """ Adds an extension onto the dict keys for a given input dictionary.
    Helpful when needing to add multiple dicts to a given metadata file

    :param d: An input dictionary
    :type d: dict
    :param extension: The extension to add to each key, defaults to "_1"
    :type extension: str, optional
    :return: Dictionary with added extensions onto each key
    :rtype: dict
    """
    d_new = {}
    for n in d:
        d_new[n + extension] = d[n]
    return d_new


def good_isolist_names():
     return ['sma', 'intens', 'intens_err', 'ellipticity', 'ellipticity_err', 
              'pa', 'pa_err', 'x0', 'x0_err', 'y0', 'y0_err']


def table_from_isolist(isolist, hdu = True, names=None):
    """ Generates an astropy table

    :param isolist: The photutils.isophote.isophotelist object 
    :type isolist: photutils.isophote.isophotelist
    :param names: Column names, defaults to galprime.good_isolist_names()
    :type names: list, optional
    :return: The table generated from the isolist
    :rtype: astropy.table.Table
    """

    if names is None:
        names = good_isolist_names()
    isolist_t = isolist.to_table()

    return Table([isolist_t[col] for col in names],  names=names)
