import re
import os
import logging
import pandas as pd
import numpy as np
# multiprocess (pathos fork of multiprocessing) uses dill rather than pickle.
# This allows objects with loggers as attributes to be processed in parallel.
import multiprocess as mp
import astropy.wcs as wcs
from astropy import units as u
from astropy.coordinates import SkyCoord
import execnet
import tqdm
from astropy.table import Table
from sklearn.metrics import auc


class Str(str):
    """
    Subclass that extends the format string for strings to begin
    with an optional 'u' or 'l' to shift to upper or lower case.
    """

    def __format__(self, fmt):
        """

        Parameters
        ----------
            fmt : 

        Returns
        ----------

        """
        s = str(self)
        if fmt[0] == 'u':
            s = self.upper()
            fmt = fmt[1:]
        elif fmt[0] == 'l':
            s = self.lower()
            fmt = fmt[1:]
        return s.__format__(fmt)

    def __getitem__(self, item):
        """

        Parameters
        ----------
            item : 

        Returns
        ---------- The item as the form of another Str object.

        """
        return Str(str(self).__getitem__(item))


def flat(l):
    """

    Parameters
    ----------
    l : list of list

    Returns
    -------
    Flattened list

    """
    return [item for sublist in l for item in sublist]


def table_to_records(table):
    """
    Converts an Astropy Table to records.

    Parameters
    ----------
        table : 

    Returns
    ----------

    """
    bar = table_to_dict(table)
    df = pd.DataFrame(bar, columns=bar.keys())
    return df.to_records()


def table_to_dict(table):
    """Convert Astropy Table to Python dict.
    Numpy arrays are converted to lists. This Can work with multi-dimensional
    array columns, by representing them as list of list.
    e.g. This is useful in the following situation.
        foo = Table.read('foo.fits')
        foo.to_pandas() <- This will not work if columns are multi-dimensional.
        The alternative is:
        foo = Table.read('foo.fits')
        bar = table_to_dict(foo)
        df = pd.DataFrame(bar, columns=bar.keys()) <- The desired result.
    """
    total_data = {}
    multi_cols = []
    for i, _ in enumerate(table.columns):
        # This looks unusual, but it is the only way to iterate over columns.
        col = table.columns[i]
        data = table[col.name].tolist()
        total_data[col.name] = data
        if len(col.shape) == 2:
            multi_cols.append(col.name)
    return total_data, multi_cols


def fits_pandas(dir,
                name,
                drop_duplicates=False,
                duplicate_id=None,
                multi_cols=[],
                multi_startswith=None,
                missing=[-99999, -1.0]):
    """
    Astropy.Table is not able to convert straight to DataFrame due to
    multivalued elements of certain columns so must be converted to a
    dictionary first then to a DataFrame. Specifying the columns
    explicitly with pd.DataFrame() conserves the order of the columns
    rather than using pd.DataFrame.from_dict() where it is not possible.
    Parameters
    ----------
    dir : str
        Directory where fits file is located.
    name : str
        Name of fits file including extension.
    drop_duplicates : bool, optional
        True if you want to drop duplicate rows from the resulting DataFrame.
    duplicate_id : str, optional
        Column name on which to match (requires drop_duplicates=True).
    Returns
    -------
    df : pd.DataFrame
        DataFrame loaded from fits.
    """
    cat = os.path.join(dir, name)
    dict_table, multi_cols = table_to_dict(Table.read(cat))
    df = pd.DataFrame(dict_table, columns=dict_table.keys())

    if drop_duplicates is True:
        df = df.drop_duplicates(
            subset=duplicate_id, keep='first', inplace=False)

    # For some reason fits files cause bytes in str columns
    # 'O' is for objects, bytes are stored as 'objects' rather than bytes.

    # DataFrame columns whose elements contain lists are also dtype
    # 'object' so to avoid converting them to NaN they must be dropped prior
    #  to converting 'object' dtype columns to strings. Such columns to
    # avoid usually start with ALL_*.

    # Convert byte columns to strings. Avoid multi_cols
    if multi_cols == [] and multi_startswith is not None:
        multi_cols = df.columns.str.startswith(multi_startswith)
        str_df = df.loc[:, ~multi_cols].select_dtypes(['O'])

    elif isinstance(multi_cols, list):
        str_df = df.drop(multi_cols, axis=1).select_dtypes(['O'])

    else:
        str_df = df.select_dtypes(['O'])

    # Strip whitespace from eitherside of strings for each string column
    str_df = str_df.stack().str.decode('utf-8').unstack()
    for col in str_df:
        df[col] = str_df[col].str.strip()

    # Remove missing data
    all_df = df.loc[:, multi_cols]
    for col in all_df:
        df[col] = all_df[col].apply(
            lambda y: [x for x in y if x not in missing])

    return df


def regex_match(string, start, end):
    """Regex match to return the pattern between two substrings in a string.

    Example:
        If string contains the substring: 'uncomp_rings.v3.skycell.2517.054.stk'
        and if start = "skycell." and ".stk" then this function  will extract
        and return the skycell identifier "2517.054".

    Parameters
    ----------
        string (str): The string to search.
        start (str): Substring defining the first enclosing character to search
        from.
        end (str): Substring defining the final enclosing character to search
        up to.

    Returns
    ----------
        pattern (str): The string between the start and end substrings in the
        original string.
        None: If no pattern is found in the string.
    """
    try:
        pattern = re.search('%s(.*)%s' % (start, end), string).group(1)
        if pattern != []:
            return pattern
    except AttributeError:
        logging.debug("pattern not found in the string {}".format(string))
        return None


def filter_conversion(names, directory):

    areas = []
    for n in names[:2]:
        photo_filter = np.loadtxt(os.path.join(directory, n))
        areas.append(auc(photo_filter[0], photo_filter[1]))

    return -2.51 * np.log(areas[0] / areas[1])


def scale_image_data(img_header, img_data):
    """ Scales data as one would expect if "do_not_scale_image_data = False" in
    astropy.io.fits. The default behaviour in astropy.io.fits is to
    scale_image_data data automatically. This function provides a way to
    carry out scaling explicitly.

    Sometimes an image is scaled, i.e. the data stored in the file is not the
    image’s physical (true) values, but linearly transformed according to the
    equation:

    physical value = BSCALE * (storage value) + BZERO

    BSCALE and BZERO are stored as keywords of the same names in the header of
    the same HDU. The most common use of scaled image is to store unsigned
    16-bit integer data because FITS standard does not allow it.

    Parameters
    ----------
    img_header :
        Header from the Header Data Unit of a fits file image.
    img_data :  array-like
        Data array from the Header Data Unit of a fits file image.

    Returns
    ----------
    scaled_data: array-like
        The image data after carrying out the scaling transformation.

    """
    bzero, bscale = img_header['BZERO'], img_header['BSCALE']
    scaled_data = bscale * (img_data) + bzero
    return scaled_data


def assert_arraylike(x):
    # Ensure that x is array-like
    if isinstance(x, (int, float)):
        x = [x]

    # Ensure that x is a numpy array as np.where is used in this function.
    if isinstance(x, list):
        x = np.array(x)

    if isinstance(x, pd.Series):
        x = x.values

    assert isinstance(x, np.ndarray)

    return x


def call_python(module, path, func="main", chdir=None, args=[]):
    """

    Parameters
    ----------
        chdir : 
        module : 
        func : 
        args : 
        path : 

    Returns
    ----------

    """
    # Don't write bytecode prevents .pyc files from being created/kept.
    spec = "popen//python={0}//dont_write_bytecode".format(path)
    if chdir is not None:
        spec = spec.join("//chdir={0}".format(chdir))
    gw = execnet.makegateway(spec)
    # This option is extremley important. The default configuration of string
    # coercion (2str to 3str, 3str to 2unicode) is inconvient.
    gw.reconfigure(py3str_as_py2str=True)
    channel = gw.remote_exec("""
        from {0} import {1} as the_function
        channel.send(the_function(*channel.receive()))
    """.format(module, func))
    channel.send(args)
    return channel.receive()


def parallelize(func, inputs, n_proc, sync=True, async=False, loglevel=None):
    """Helper function to parrallelise any mappable function either
    synchously or asynchrously using python's multiprocessing module.

    Parameters
    ----------
        func :  Any mappable function.
        inputs :  The arguments to multiproces.
        n_proc :  The maximum number of python processes (engines) that
        can be running at any one time.
        sync (bool): Set to true if synchonous processing is required.
        async (bool): Set to true if asynchonous processing is required.

    Returns
    ----------
        None

    """
    if loglevel is not None:
        logger = mp.log_to_stderr()
        logger.setLevel(loglevel)

    # TODO: could improve by only having one sync/ a sync arg
    with mp.Pool(processes=n_proc) as p:
        # Run multiprocessing with progress bar.
        # TODO: other bands can have (quicker) async map as order no longer
        # matters.
        r = list(tqdm.tqdm(p.imap(func, inputs), total=len(inputs)))
        # inform the processor that no new tasks will be added the pool
        p.close()
        # wait for map to be completed before proceeding
        p.join()

    # TODO: clean this up.  # if sync is True:  #     p.map(func, inputs)  #
    #  inform the processor that no new tasks will be added the pool  #  #
    # pool.close()  #     # wait for map to be completed before proceeding  #
    #      pool.join()  # elif async is True:  #     pool.map_async(func,
    #  inputs)  # , callback=on_result)  #     # inform the processor that no
    #  new tasks will be added the pool  #     pool.close()  #     # wait for
    #  map to be completed before proceeding  #     pool.join()


def file_size(path, unit='B'):
    if os.path.isfile(path):
        size = os.path.getsize(path)
    else:
        size = 0
    n = {'B': 0, 'KB': 1, 'MB': 2}
    return size / pow(1024, n[unit])


def dist_great_circle(ra1, dec1, ra2=None, dec2=None):
    """ Computes rigorous great circle arc distances an returns the value in
    arcsecs.

    This function is based on the GCIRC IDL function.
    https://idlastro.gsfc.nasa.gov/ftp/pro/astro/gcirc.pro.

    """
    # Ensure that when only one catalog (set of ra and dec lists) is given
    # that the catalog is matched with itself.
    if ra2 is None and dec2 is None:
        ra2, dec2 = ra1, dec1
    sky1 = SkyCoord(ra=ra1.deg, dec=dec1.deg, unit='deg')
    sky2 = SkyCoord(ra=ra2.deg, dec=dec2.deg, unit='deg')
    return sky1.separation(sky2).to(u.arcsec)


def dist_circle(nx, ny, x_cen, y_cen):
    """

    Create a 2D numpy array where each value is its distance to a given
    center. "

    Parameters
    ----------
        nx : 
        ny : 
        x_cen : 
        y_cen : 

    Returns
    ----------
        mask:

   """
    import numpy as np
    mask = np.zeros([nx, ny])
    x = np.linspace(0, nx - 1, nx)
    y = np.linspace(0, ny - 1, ny)
    x, y = np.meshgrid(x, y)
    x, y = x.transpose(), y.transpose()
    x_cen, y_cen = y_cen, x_cen  # transpose centers well
    mask = np.sqrt((x - x_cen)**2 + (y - y_cen)**2)
    return mask


def wcs_pixel_info(ra, dec):
    """
    Construct a World Coordinate System (WCS) object and return relevant pixel
    information for creating a fits image from the (ra,dec) positions of
    sources.

    For more information see the following links.

    Build a WCS structure programmatically
    http://docs.astropy.org/en/v0.2/wcs/index.html#building-a-wcs
    -structure-programmatically

    Writing out images with WCS information
    http://docs.hyperion-rt.org/en/stable/tutorials/tutorial_images.html

    Parameters
    ----------
        ra : 
        dec : 

    Returns
    ----------
        w: The World Coordinate System (WCS) object.
        idx_sources: A tuple of pixel position values (sx, sy) converted from
        (ra, dec) of all the SExtracted sources within the
        boundaries of the dimensions of the frame (dim1, dim2).
        shape: The shape (a tuple of dimensions: (dim1, dim2)) that represents
        the shape of the axes for the frame.

    """
    # Get rough ima get dimensions
    # TODO: change to get rid of dependancy of dist_great_Circle
    dim_ra = dist_great_circle(
        np.nanmin(ra), np.median(dec), np.nanmax(ra), np.median(dec))
    dim_dec = dist_great_circle(
        np.median(ra), np.nanmin(dec), np.median(ra), np.nanmax(dec))

    # Number of pixels.
    n_ra_pix = dim_ra.arcsec / 2
    n_dec_pix = dim_dec.arcsec / 2

    # Create a new WCS object. The number of axes must be set.
    w = wcs.WCS(naxis=2)
    # Use the center of the image as projection center.
    w.wcs.crpix = [n_ra_pix / 2, n_dec_pix / 2]
    # Set the pixel scale (in deg/pix).
    scale = 2 * u.arcsec.to(u.deg)
    w.wcs.cdelt = np.array([-scale, scale])
    # Set the coordinates of the image center.
    w.wcs.crval = [np.median(ra.deg), np.median(dec.deg)]
    # Set the coordinate system
    w.wcs.ctype = ["RA---TAN", "DEC--TAN"]

    # Convert the ra and dec np arrays of to arrays of the corresponding
    # pixel values in the newly defined wcs.

    # Note the third argument, set to 1, which indicates whether the pixel
    # coordinates should be treated as starting from (1, 1) (as FITS files
    # do) or from (0, 0).
    sx, sy = w.wcs_world2pix(ra, dec, 1)

    # Make sure there is a whole number of pixels in each pixel array.
    if np.nanmin(sx) < 0.5:
        w.wcs.crpix[0] = w.wcs.crpix[0] + 1 - np.nanmin(sx)
    if np.nanmin(sy) < 0.5:
        w.wcs.crpix[1] = w.wcs.crpix[1] + 1 - np.nanmin(sy)

    # dim1 and dim2 represent the dimensions of the axes for the frame.
    # sx, sy represent the pixel values of all the positions (converted from
    # ra, dec) of the SExtracted sources within the boundaries of dim1 and dim2.

    # Rerun world to pix conversion with new whole number values.
    # Note the third argument, set to 1, which indicates whether the pixel
    # coordinates should be treated as starting from (1, 1) (as FITS files
    # do) or from (0, 0).
    sx, sy = w.wcs_world2pix(ra, dec, 1)
    # Pixel arrays now contain whole numbers but need to be ints.
    sx, sy = np.rint(sx).astype(int), np.rint(sy).astype(int)
    idx_sources = sx, sy
    # Set dimensions of new data from max values in the pixel arrays.
    dim1, dim2 = 1 + np.amax(sx), 1 + np.amax(sy)
    shape = (dim1, dim2)
    return w, idx_sources, shape


def coord_meshgrid(data, wcs, skycoord=True):
    """
    Provides an RA and DEC for every single pixel of a 2D array given a world
    coordinate system. Primarily designed with the SkyCoord class in mind.


    Parameters
    ----------
    data - This should
    data
    wcs
    skycoord

    Returns
    -------

    """
    # based on https://github.com/astropy/astropy/issues/1587
    h, w = data.shape
    x, y, = np.arange(h), np.arange(w)
    x_mesh, y_mesh = np.meshgrid(x, y)
    ra, dec = wcs.wcs_pix2world(x_mesh, y_mesh, 0)
    if skycoord is True:
        return SkyCoord(ra, dec, unit='deg')
    else:
        return np.dstack((ra, dec))


def create_circular_mask(shape, center=None, radius=None):
    """
    based on https://stackoverflow.com/questions/44865023/circular-masking
    -an-image-in-python-using-numpy-arrays

    Parameters
    ----------
    shape
    center
    radius

    Returns
    -------

    """
    h, w = shape
    if center is None:
        # use the middle of the image
        center = [int(w / 2), int(h / 2)]
    if radius is None:
        # use the smallest distance between the center and image walls
        radius = min(center[0], center[1], w - center[0], h - center[1])

    y, x = np.ogrid[:h, :w]
    dist_from_center = np.sqrt((x - center[0])**2 + (y - center[1])**2)

    mask = dist_from_center <= radius
    return mask


def circle_skyregions_mask(regions, wcs, shape):
    """

    Parameters
    ----------
    regions - array-like, list of CircularSkyRegions
    regions
    wcs
    shape

    Returns
    -------

    """
    final_mask = np.zeros(shape=shape, dtype=bool)
    for sky_region in regions:
        pixel_region = sky_region.to_pixel(wcs=wcs)
        radius = pixel_region.radius
        center = pixel_region.center.x, pixel_region.center.y
        mask = create_circular_mask(shape=shape, center=center, radius=radius)
        final_mask = final_mask | mask
    return final_mask


def region_list_to_compound(region_list):

    for i, r in enumerate(region_list):
        if i == 0:
            compound_region = r
        else:
            compound_region = compound_region | r

        return compound_region
