"""
Test a bunch of functions that serve as an interface to standard stellar data
table
"""
import numpy as np
import logging

from astropy.io import fits
from astropy.table import Table
# from astropy.units.core import UnitConversionError
try:
    import exceptions
except ImportError:
    import builtins as exceptions # python 3 consistent

import sys
sys.path.insert(0, '..')
from chronostar.component import SphereComponent
from chronostar.synthdata import SynthData
from chronostar import tabletool
from chronostar import coordinate
from chronostar import transform

HIST_FILE_NAME = 'sample_data/small_historical_sample_table.fits'
CURR_FILE_NAME = 'sample_data/small_current_sample_table.fits'

# -----------------------------------------------
# -- To check correctness, have copied over    --
# -- previous implementation into this script  --
# -----------------------------------------------
# retired function, put here for comparison reasons
def transformAstrCovsToCartesian(astr_covs, astr_arr):
    """
    Converts a covariance matrix from astrometric coords to LSR XYZUVW
    Parameters
    ----------
    astr_covs: ([nstars, 6, 6] array)
        values in the diagaonal are the squared errors of
        (ra, dec, plx, pm_ra, pm_dec, rv), with the offdiagonals the product
        of the correlation (valued between -1 and 1) and the two
        intersecting coordinates.
    astr_arr: ([nstars, 6] array)
        the measured (mean) astrometric values
        (ra, dec, plx, pm_ra, pm-dec, rv)
    """
    nstars = astr_arr.shape[0]
    xyzuvw_covs = np.zeros((nstars, 6, 6))
    for ix in range(nstars):
        xyzuvw_covs[ix] = transform.transform_covmatrix(
            astr_covs[ix], coordinate.convert_astrometry2lsrxyzuvw, astr_arr[ix],
            dim=6
        )
    return xyzuvw_covs

# retired function, put here for comparison reasons
def convertAstrErrsToCovs(err_arr):
    """
    Converts astrometry errors for each star into covariance matrices

    Note that a negligible error is inserted for ra and dec

    Parameters
    ----------
    err_arr : ([nstars, 6] float array), astrometry errors with placeholder
        0's for ra and dec: (ra, dec, pi, pm_ra, pm_dec, rv)

    Returns
    -------
    astr_covs : ([nstars, 6, 6] float array), covariance matrix made by
        inserting the square of the errors into the diagonal
    """
    err_arr_cp = np.copy(err_arr)
    nstars = err_arr_cp.shape[0]

    err_arr_cp[:, :2] = 1e-1
    logging.info("Angular position error is: {}".format(err_arr_cp[0,0]))
    print("Angular position error is: {}".format(err_arr_cp[0,0]))
    print("Changed!")

    astr_covs = np.zeros((nstars, 6, 6))
    for ix in range(nstars):
        astr_covs[ix] = np.eye(6) * np.tile(err_arr_cp[ix], (6, 1))**2
    return astr_covs

# Retired function, only here for comparison reasons
def convertTableToArray(star_table):
    nstars = star_table['radeg'].shape[0]
    measured_vals = np.vstack((
        star_table['radeg'],
        star_table['dedeg'],
        star_table['plx'],
        star_table['pmra'],
        star_table['pmde'],
        star_table['rv'],
    )).T

    errors = np.vstack((
        np.zeros(nstars),
        np.zeros(nstars),
        star_table['e_plx'],
        star_table['e_pmra'],
        star_table['e_pmde'],
        star_table['e_rv'],
    )).T
    return measured_vals, errors


# Retired funciton, only here for comparison reasons
def buildMeanAndCovMatFromRow(row):
    """
    Build a covariance matrix from a row

    Paramters
    ---------
    row : astropy Table row
        Entries: {X, Y, Z, U, V, W, dX, dY, ..., cXY, cXZ, ...}

    Return
    ------
    cov_mat : [6,6] numpy array
        Diagonal elements are dX^2, dY^2, ...
        Off-diagonal elements are cXY*dX*dY, cXZ*dX*dZ, ...
    """
    dim = 6
    CART_COL_NAMES = ['X', 'Y', 'Z', 'U', 'V', 'W',
                      'dX', 'dY',   'dZ',   'dU',   'dV',   'dW',
                            'c_XY', 'c_XZ', 'c_XU', 'c_XV', 'c_XW',
                                    'c_YZ', 'c_YU', 'c_YV', 'c_YW',
                                            'c_ZU', 'c_ZV', 'c_ZW',
                                                    'c_UV', 'c_UW',
                                                            'c_VW']
    mean = np.zeros(dim)
    for i, col_name in enumerate(CART_COL_NAMES[:6]):
        mean[i] = row[col_name]

    std_vec = np.zeros(dim)
    for i, col_name in enumerate(CART_COL_NAMES[6:12]):
        std_vec[i] = row[col_name]
    corr_tri = np.zeros((dim,dim))
    # Insert upper triangle (top right) correlations
    for i, col_name in enumerate(CART_COL_NAMES[12:]):
        corr_tri[np.triu_indices(dim,1)[0][i],np.triu_indices(dim,1)[1][i]] \
            =row[col_name]
    # Build correlation matrix
    corr_mat = np.eye(6) + corr_tri + corr_tri.T
    # Multiply through by standard deviations
    cov_mat = corr_mat * std_vec * std_vec.reshape(6,1)
    return mean, cov_mat


# Retired funciton, put here for comparison reasons
def loadDictFromTable(table):
    """
    Takes the data in the table, builds dict with array of mean and cov matrices

    Paramters
    ---------
    table : Astropy.Table (or str)

    assoc_name : str
        One of the labels in Moving group column:
         118 Tau, 32 Orionis, AB Doradus, Carina, Carina-Near, Columba,
         Coma Ber, Corona Australis, Hyades, IC 2391, IC 2602,
         Lower Centaurus-Crux, Octans, Platais 8, Pleiades, TW Hya, Taurus,
         Tucana-Horologium, Upper Centaurus Lupus, Upper CrA, Upper Scorpius,
         Ursa Major, beta Pictoris, chi{ 1 For (Alessi 13), epsilon Cha,
         eta Cha, rho Ophiuci

    Returns
    -------
    dict :
        xyzuvw : [nstars,6] array of xyzuvw means
        xyzuvw_cov : [nstars,6,6] array of xyzuvw covariance matrices
        file_name : str
            if table loaded from file, the pathway is stored here
        indices : [nstars] int array
            the table row indices of data converted into arrays
        gaia_ids : [nstars] int array
            the gaia ids of stars successfully converted
    """
    star_pars = {}
    if type(table) is str:
        file_name = table
        star_pars['file_name'] = file_name
        table = Table.read(table)

    gaia_col_name = 'source_id'
    if gaia_col_name not in table.colnames:
        gaia_col_name = 'gaia_dr2'


    xyzuvw = []
    xyzuvw_cov = []
    indices = []
    gaia_ids = []
    nrows = len(table[gaia_col_name])
    for ix, row in enumerate(table):
        if nrows > 10000 and ix % 1000==0:
            print("Done {:7} of {} | {:.2}%".format(ix, nrows,
                                                    float(ix)/nrows*100))
        if np.isfinite(row['U']):
            mean, cov = buildMeanAndCovMatFromRow(row)
            xyzuvw.append(mean)
            xyzuvw_cov.append(cov)
            indices.append(ix)
            gaia_ids.append(row[gaia_col_name])

    star_pars['xyzuvw']     = np.array(xyzuvw).astype(np.float64)
    star_pars['xyzuvw_cov'] = np.array(xyzuvw_cov).astype(np.float64)
    star_pars['indices']    = np.array(indices)
    star_pars['gaia_ids']   = np.array(gaia_ids)
    star_pars['table']      = table
    return star_pars


# Retired funciton, put here for comparison reasons
def convertManyRecToArray(data):
    """
    Convert many Fits Records in astrometry into mean and covs (astro)