import numpy as np
from astropy import units as u
from scipy.interpolate import InterpolatedUnivariateSpline
from functools import partial

spline = partial(InterpolatedUnivariateSpline, k=1)

# Mapping from telescope name to site's extinction curve
# We have no curve for CPO, so use the CTIO one
telescope_sites = {'Gemini-North': 'MKO',
                   'Gemini-South': 'CTIO'}

# Extinction curves for supported sites, as a function of wavelength in nm
# Units are mag/airmass
extinction_curves = {
    # From Buton et al. (2013, A&A 549, A8), 310nm point from Gemini website
    'MKO': spline(np.arange(310, 1001, 10),
                  (1.37, 0.856, 0.588, 0.514, 0.448, 0.400, 0.359, 0.323,
                   0.292, 0.265, 0.241, 0.220, 0.202, 0.185, 0.171, 0.159,
                   0.147, 0.139, 0.130, 0.125, 0.119, 0.114, 0.113, 0.109,
                   0.106, 0.107, 0.108, 0.103, 0.098, 0.098, 0.092, 0.084,
                   0.078, 0.070, 0.065, 0.060, 0.056, 0.052, 0.048, 0.044,
                   0.042, 0.039, 0.037, 0.035, 0.033, 0.032, 0.030, 0.029,
                   0.028, 0.027, 0.026, 0.025, 0.024, 0.023, 0.023, 0.022,
                   0.021, 0.021, 0.020, 0.019, 0.019, 0.018, 0.018, 0.017,
                   0.017, 0.016, 0.016, 0.015, 0.015, 0.014)),
    # From Stritzinger et al. (2005, PASP 117, 810). We note the discussion
    # in Patat et al. (2011, A&A 527, A91) regarding the evolution of the
    # extinction above the Atacama desert and elect to use this curve instead
    # of the IRAF ctioextinct.dat
    'CTIO': spline((305.000, 308.465, 311.931, 315.396, 318.861, 322.327,
                    325.792, 329.257, 332.723, 336.188, 339.654, 343.119,
                    346.584, 350.050, 353.515, 356.980, 360.446, 363.911,
                    367.376, 370.842, 374.307, 377.772, 381.238, 384.703,
                    388.169, 391.634, 395.099, 398.565, 402.030, 405.495,
                    408.961, 412.426, 415.891, 419.357, 422.822, 426.287,
                    429.753, 433.218, 436.683, 440.149, 443.614, 447.079,
                    450.545, 454.010, 457.476, 460.941, 464.406, 467.872,
                    471.337, 474.802, 478.268, 481.733, 485.198, 488.664,
                    492.129, 495.594, 499.060, 502.525, 505.991, 509.456,
                    512.921, 516.387, 519.852, 523.317, 526.783, 530.248,
                    533.713, 537.179, 540.644, 544.109, 547.575, 551.040,
                    554.505, 557.971, 561.436, 564.902, 568.367, 571.832,
                    575.298, 578.763, 582.228, 585.694, 589.159, 592.624,
                    596.090, 599.555, 603.020, 606.486, 609.951, 613.417,
                    616.882, 620.347, 623.813, 627.278, 630.743, 634.209,
                    637.674, 641.139, 644.605, 648.070, 648.285, 653.538,
                    658.791, 664.044, 669.296, 674.549, 679.802, 685.055,
                    690.307, 695.560, 700.813, 706.065, 711.318, 716.571,
                    721.824, 727.076, 732.329, 737.582, 742.835, 748.087,
                    753.340, 758.593, 763.845, 769.098, 774.351, 779.604,
                    784.856, 790.109, 795.362, 800.615, 805.867, 811.120,
                    816.373, 821.625, 826.878, 832.131, 837.384, 842.636,
                    847.889, 853.142, 858.395, 863.647, 868.900, 874.153,
                    879.405, 884.658, 889.911, 895.164, 900.416, 905.669,
                    910.922, 916.175, 921.427, 926.680, 931.933, 937.185,
                    942.438, 947.691, 952.944, 958.196, 963.449, 968.702,
                    973.955, 979.207, 984.460, 989.713, 994.965, 1000.22,
                    1005.47, 1010.72, 1015.98, 1021.23, 1026.48, 1031.73,
                    1036.99, 1042.24, 1047.49, 1052.75, 1058.00, 1063.25,
                    1068.50, 1073.76, 1079.01, 1084.26, 1089.51, 1094.77,
                    1100.02),
                   (1.395, 1.283, 1.181, 1.088, 1.004, 0.929, 0.861, 0.801,
                    0.748, 0.700, 0.659, 0.623, 0.591, 0.564, 0.540, 0.520,
                    0.502, 0.487, 0.473, 0.460, 0.448, 0.436, 0.425, 0.414,
                    0.402, 0.391, 0.381, 0.370, 0.360, 0.349, 0.339, 0.330,
                    0.321, 0.313, 0.304, 0.296, 0.289, 0.281, 0.274, 0.267,
                    0.260, 0.254, 0.247, 0.241, 0.236, 0.230, 0.225, 0.220,
                    0.215, 0.210, 0.206, 0.202, 0.198, 0.194, 0.190, 0.187,
                    0.184, 0.181, 0.178, 0.176, 0.173, 0.171, 0.169, 0.167,
                    0.166, 0.164, 0.163, 0.162, 0.160, 0.159, 0.158, 0.158,
                    0.157, 0.156, 0.155, 0.155, 0.154, 0.153, 0.153, 0.152,
                    0.151, 0.151, 0.150, 0.149, 0.149, 0.148, 0.147, 0.146,
                    0.144, 0.143, 0.142, 0.140, 0.138, 0.136, 0.134, 0.132,
                    0.129, 0.126, 0.123, 0.120, 0.120, 0.115, 0.111, 0.107,
                    0.103, 0.099, 0.096, 0.092, 0.088, 0.085, 0.082, 0.078,
                    0.075, 0.072, 0.069, 0.066, 0.064, 0.061, 0.058, 0.056,
                    0.053, 0.051, 0.049, 0.047, 0.045, 0.043, 0.041, 0.039,
                    0.037, 0.035, 0.034, 0.032, 0.030, 0.029, 0.028, 0.026,
                    0.025, 0.024, 0.023, 0.022, 0.020, 0.019, 0.019, 0.018,
                    0.017, 0.016, 0.015, 0.015, 0.014, 0.013, 0.013, 0.012,
                    0.011, 0.011, 0.011, 0.010, 0.010, 0.009, 0.009, 0.009,
                    0.008, 0.008, 0.008, 0.007, 0.007, 0.007, 0.007, 0.007,
                    0.006, 0.006, 0.006, 0.006, 0.006, 0.006, 0.005, 0.005,
                    0.005, 0.005, 0.005, 0.005, 0.004, 0.004, 0.004, 0.004,
                    0.003, 0.003, 0.003)),
}


def extinction(wave, site=None, telescope=None):
    """
    This function returns the extinction (in mag/airmass) at the specified
    input wavelengths. Wavelengths should be Quantity objects but, if they
    are scalars, units of nm are assumed.

    Parameters
    ----------
    wave: float/array/Quantity
        wavelength(s) at which to derive extinction (if no units, the
        wavelength(s) are assumed to be in nm)
    site: str/None
        name of site (key for extinction curve)
    telescope: str/None
        name of telescope (maps to site)

    Returns
    -------
    arrray: extinction in magnitudes/airmass at requested wavelength(s)
    """
    if telescope in telescope_sites:
        site = telescope_sites[telescope]
    elif site not in extinction_curves:
        raise KeyError(f"Site {site} not recongized")
    try:
        wave_in_nm = wave.to(u.nm)
    except AttributeError:
        # Assume it's nm already
        wave_in_nm = wave
    return extinction_curves[site](wave_in_nm)