"""
Functions to retrieve and interpolate atmosphere profiles.
"""
import numpy as np
from astropy.units import Quantity
from scipy.interpolate import interp1d

from ctapipe.utils import get_table_dataset

__all__ = ['get_atmosphere_profile_table', 'get_atmosphere_profile_functions']


def get_atmosphere_profile_table(atmosphere_name='paranal'):
    """
    Get an atmosphere profile table

    Parameters
    ----------
    atmosphere_name: str
        identifier of atmosphere profile

    Returns
    -------
    astropy.table.Table  containing atmosphere profile with at least columns
    'altitude' (m), and 'thickness' (g cm-2) as well as others.

    """
    table_name = '{}.atmprof'.format(atmosphere_name)
    table = get_table_dataset(table_name=table_name,
                              role='dl0.arr.svc.atmosphere')
    return table


def get_atmosphere_profile_functions(atmosphere_name="paranal",
                                     with_units=True):
    """
    Gives atmospheric profile as a continuous function thickness(
    altitude), and it's inverse altitude(thickness)  in m and g/cm^2
    respectively.

    Parameters
    ----------
    atmosphere_name: str
        identifier of atmosphere profile
    with_units: bool
       if true, return functions that accept and return unit quantities.
       Otherwise assume units are 'm' and 'g cm-2'

    Returns
    -------
    functions: thickness(alt), alt(thickness)
    """
    tab = get_atmosphere_profile_table(atmosphere_name)
    alt = tab['altitude'].to('m')
    thick = (tab['thickness']).to("g cm-2")

    alt_to_thickness = interp1d(x=np.array(alt), y=np.array(thick))
    thickness_to_alt = interp1d(x=np.array(thick), y=np.array(alt))

    if with_units:
        def thickness(a):
            return Quantity(alt_to_thickness(a.to('m')), 'g cm-2')

        def altitude(a):
            return Quantity(thickness_to_alt(a.to('g cm-2')), 'm')

        return thickness, altitude

    return alt_to_thickness, thickness_to_alt
