'''
A module containing some functions I found useful.
A few of which are used in the main fitting code
'''
from __future__ import division
import os
import sys

import numpy as np
import astropy.units as u
import matplotlib.pyplot as plt
import montage_wrapper as montage
from astropy.io import fits
from astropy.wcs import WCS
from astropy.table import Table, Column
from astropy.coordinates import SkyCoord


def adxy(header, ra, dec):
    '''
    Converts RA and Dec to x and y pixel positions
    '''
    ra = np.array(ra)
    dec = np.array(dec)
    x, y = WCS(header).all_world2pix(ra, dec, 0)
    return x, y


def angsep(ra0, dec0, ra, dec, frame='icrs'):
    '''
    Gives angular separation between the input position and the position after
    fitting in arcsec. Note: input should be in degrees
    '''
    ra0 = np.array(ra0)
    dec0 = np.array(dec0)
    ra = np.array(ra)
    dec = np.array(dec)
    c1 = SkyCoord(ra0 * u.deg, dec0 * u.deg, frame=frame)
    c2 = SkyCoord(ra * u.deg, dec * u.deg, frame=frame)
    return c1.separation(c2).value * 3600.


def data2table(arrays=[], names=[], units=[], filename='d2table.fits'):
    '''
    Given a list of arrays, names and units, constructs table and writes
    to file. Units (optional) should be given as astropy.units.
    If extenstion of filename is .dat (.fits) it will be saved as an
    ascii (FITS) table.
    Returns the table that it saved to file
    '''
    t = Table()
    if len(units) > 0:
        for info in zip(names, arrays, units):
            t.add_column(
                Column(name=info[0], data=info[1], unit=info[2]))
    else:
        for nm, arr in zip(names, arrays):
            t.add_column(Column(name=nm, data=arr))
    if os.path.splitext(filename)[1] == '.dat':
        t.write(filename, format='ipac')
    elif os.path.splitext(filename)[1] == '.fits':
        t.write(filename)
    else:
        print 'Unacceptable filename given. Table not being saved.'
    return t