import os
import numpy as np
import healpy as hp
from astropy.io import fits
from astropy import units as u
from astropy.coordinates import SkyCoord


class ExtinctionMW(object):
    def __init__(self):
        self.Rv = 3.1
        self.lamb = np.arange(2000, 11001+0.5, 0.5)

    @staticmethod
    def radec2pix(ra, dec, NSIDE=2048):
        return hp.pixelfunc.ang2pix(nside=NSIDE, theta=ra, phi=dec, lonlat=True)

    def init_ext_model(self, model_name="odonnell", Av=1.0, Rv=3.1, lamb=None):
        self.model_name = model_name
        self.Av = Av
        self.Rv = Rv
        if lamb is not None:
            self.lamb = lamb

        if self.model_name == "odonnell":
            alpha = np.zeros(self.lamb.shape)
            beta = np.zeros(self.lamb.shape)
            x = 1.e4 / self.lamb
            def f_alpha_1(x): return 0.574 * (x**1.61)
            def f_beta_1(x): return -0.527 * (x**1.61)

            def f_alpha_2(x):
                y = x - 1.82
                return 1 + 0.104*y - 0.609*(y**2) + 0.701*(y**3) + \
                    1.137*(y**4) - 1.718*(y**5) - 0.827 * \
                    (y**6) + 1.647*(y**7) - 0.505*(y**8)

            def f_beta_2(x):
                y = x - 1.82
                return 1.952*y + 2.908*(y**2) - 3.989*(y**3) - 7.985 * \
                    (y**4) + 11.102*(y**5) + 5.491 * \
                    (y**6) - 10.805*(y**7) + 3.347*(y**8)

            def f_alpha_3(x): return 1.752 - 0.316*x - \
                0.104 / ((x - 4.67)**2 + 0.341)

            def f_beta_3(x): return -3.090 + 1.825*x + \
                1.206 / ((x - 4.62)**2 + 0.262)

            def f_alpha_4(x): return f_alpha_3(x) + -0.04473 * \
                (x - 5.9)**2 - 0.009779 * (x - 5.9)**3

            def f_beta_4(x): return f_beta_3(x) + 0.2130 * \
                (x - 5.9)**2 + 0.1207 * (x - 5.9)**3

            def f_alpha_5(x): return -1.073 - 0.628*(x - 8) + \
                0.137*(x - 8)**2 - 0.070*(x - 8)**3

            def f_beta_5(x): return 13.670 + 4.257*(x - 8) - \
                0.420*(x - 8)**2 + 0.374 * (x - 8)**3

            alpha = np.piecewise(
                x, [x <= 1.1, (x > 1.1)*(x <= 3.3), (x > 3.3)*(x <= 5.9), (x > 5.9)*(x <= 8.), (x > 8.)], [f_alpha_1, f_alpha_2, f_alpha_3, f_alpha_4, f_alpha_5])
            beta = np.piecewise(
                x, [x <= 1.1, (x > 1.1)*(x <= 3.3), (x > 3.3)*(x <= 5.9), (x > 5.9)*(x <= 8.), (x > 8.)], [f_beta_1, f_beta_2, f_beta_3, f_beta_4, f_beta_5])

            self.ext = (alpha + beta / self.Rv) * self.Av

    def load_Planck_ext(self, file_path):
        hdu = fits.open(file_path)
        self.ebv_planck = hdu[1].data["EBV"]

    def Av_from_Planck(self, ra, dec):
        if not hasattr(self, 'ebv_planck'):
            raise ValueError(
                "Need to load planck dust map first")
        # Convert to galactic coordinates
        c = SkyCoord(ra=ra * u.degree, dec=dec *
                     u.degree, frame='icrs').galactic
        l, b = c.l.radian, c.b.radian
        NSIDE = hp.pixelfunc.get_nside(self.ebv_planck)
        pix = hp.pixelfunc.ang2pix(nside=NSIDE, theta=np.pi/2. - b, phi=l)
        return self.ebv_planck[pix] * self.Rv

    def apply_extinction(self, spec, Av=1.0):
        if len(spec) != len(self.lamb):
            raise ValueError(
                "Dimension of spec do not match that of ExtinctionMW.lamb: try initilize (init_ext_model) by the actual size of spec which you want to apply the extinction to")
        if not hasattr(self, 'ext'):
            raise ValueError(
                "Need to initialize the extinction model (init_ext_model) first")
        scale = 10**(-.4 * self.ext * Av)
        print("scale = ", scale)
        spec *= scale
        return spec