import numpy as np
import pytest

import astropy.units as u
from astropy.modeling.fitting import LevMarLSQFitter

from dust_extinction.shapes import P92


def get_axav_cor_vals():

    # Milky Way observed extinction as tabulated by Pei (1992)
    # fmt: off
    MW_x = [0.21, 0.29, 0.45, 0.61, 0.80, 1.11, 1.43, 1.82,
            2.27, 2.50, 2.91, 3.65, 4.00, 4.17, 4.35, 4.57, 4.76,
            5.00, 5.26, 5.56, 5.88, 6.25, 6.71, 7.18, 7.60,
            8.00, 8.50, 9.00, 9.50, 10.00]
    MW_x = np.array(MW_x)
    MW_exvebv = [-3.02, -2.91, -2.76, -2.58, -2.23, -1.60, -0.78, 0.00,
                 1.00, 1.30, 1.80, 3.10, 4.19, 4.90, 5.77, 6.57, 6.23,
                 5.52, 4.90, 4.65, 4.60, 4.73, 4.99, 5.36, 5.91,
                 6.55, 7.45, 8.45, 9.80, 11.30]
    # fmt: on

    MW_exvebv = np.array(MW_exvebv)
    Rv = 3.08
    MW_axav = MW_exvebv / Rv + 1.0

    # add units
    x = MW_x / u.micron

    # correct values
    cor_vals = MW_axav

    return (x, cor_vals)


def test_extinction_P92_values():
    # get the correct values
    x, cor_vals = get_axav_cor_vals()

    # initialize extinction model
    tmodel = P92()

    # test
    np.testing.assert_allclose(tmodel(x), cor_vals, rtol=0.25, atol=0.01)


x_vals, axav_vals = get_axav_cor_vals()
test_vals = zip(x_vals, axav_vals)


@pytest.mark.parametrize("xtest_vals", test_vals)
def test_extinction_P92_single_values(xtest_vals):
    x, cor_val = xtest_vals

    # initialize extinction model
    tmodel = P92()

    # test
    np.testing.assert_allclose(tmodel(x), cor_val, rtol=0.25, atol=0.01)


def test_P92_fitting():

    # get an observed extinction curve to fit
    x_quant, y = get_axav_cor_vals()

    x = x_quant.value

    # change from defaults to make the best fit harder to find
    p92_init = P92()

    fit = LevMarLSQFitter()
    # accuracy set to avoid warning the fit may have failed
    p92_fit = fit(p92_init, x, y, acc=1e-2)

    fit_vals = p92_fit._parameters

    # fmt: off
    good_vals = [220.43887482191883, 0.04772654363744676, 89.13955305851606, 2.0,
                 17.53404746789474, 0.072502973523927, 1.6026274827983342,
                 5.454549320255999, 0.05379254892041606, 0.2187391000979195,
                 -1.9501411144824525, 2.0, 0.0, 13.0, 504.5387984216088, 2.0,
                 0.0, 15.0, -5192.423852001501, 2.0, 0.0, 30.0,
                 -10264.367792916015, 2.0]
    # fmt: on

    np.testing.assert_allclose(good_vals, fit_vals)
