from ..imports import *
from ..units import determine_quantity
from ..plottingtools import setup_axes_with_rainbow
from .sunset import Sunset


class Atmosphere:
    def __repr__(self):
        """
        How should this object appear as a string?

        Returns
        -------
        s : str
            A simple string representation.
        """
        scale = (self.H / self.radius).decompose()
        # H/R={scale},
        return f"{self.__class__.__name__}Atmosphere ({self.zenith_angle} from zenith, {self.altitude} scale heights from reference)"

    # FIXME (maybe both spectrum and atmosphere should inherit from the same thing?)
    def wavelength(self, wavelength=None):
        """
        A wrapper to ensure at least some grid of wavelengths
        gets defined. A default grid will be assumed, unless
        any wavelength array at all is passed.
        """

        # make sure at least some wavelengths are defined
        if wavelength is None:
            wavelength = self.default_wavelengths
        return wavelength.to("nm")

    def plot(
        self,
        ax=None,
        wavelength=None,
        rainbow=True,
        color=None,
        style="dark_background",
        figsize=(5.0, 2.5),
        **kwargs,
    ):
        """
        A quick tool to plot the transmission through the atmosphere.

        Parameters
        ----------
        ax : matplotlib.axes._subplots.AxesSubplot
            Specify an axes object into which
            this plot should be drawn. This allows
            overplotting multiple spectra, for example
            with a structure like
                ```
                ax = Thermal(teff=5800*u.K, radius=1*u.Rsun).plot()
                ax = Sun.plot(ax)
                ```
            To overplot on current axes use `ax=plt.gca()`.

        wavelength : astropy.units.quantity.Quantity
            A grid of wavelengths on which the spectrum should
            be plotted. If None, the function defaults to
            covering visible wavelengths at 1nm resolution.

        rainbow : bool
            Should we add an extra rainbow above the plot,
            to indicate how wavelengths match to visible light?

        color : str
            The color for drawing the spectrum.
            'auto' represents the actual visible color.

        """

        # set up to use a dark background for the plot; make sure units match
        with plt.style.context(style), quantity_support():

            # setup the basic axes
            ax = setup_axes_with_rainbow(ax=ax, rainbow=rainbow, figsize=figsize)

            # make sure at least some wavelengths are defined
            w = self.wavelength(wavelength)

            # pull out the spectrum
            t = self.transmission(w)

            plt.plot(w, t * 100, color=color, label=self, **kwargs)

            # add the axis labels
            wunit = w.unit.to_string("latex_inline")
            plt.xlabel(f"Wavelength ({wunit})")
            plt.ylabel(f"Transmission (%)")

            plt.xlim(np.min(w).value, np.max(w).value)

        return ax

    def set_elevation_angle(self, elevation=90 * u.deg):
        """
        Set the elevation angle above the horizon.

        Parameters
        ----------
        elevation : astropy.units.quantity.Quantity
            The angle above the horizon along which
            the transmission of the atmosphere should
            be calculated.
        """
        self.set_zenith_angle(90 * u.deg - elevation)

    def set_zenith_angle(self, zenith_angle=0 * u.deg):
        """
        Set the angle from zenith.

        Parameters
        ----------
        zenith_angle : astropy.units.quantity.Quantity
            The angle away from zenith along which
            the transmission of the atmosphere should
            be calculated.
        """

        # set the zenith angle and airmass
        self.zenith_angle = np.minimum(zenith_angle, 90 * u.deg)
        self.airmass = 1 / np.cos(self.zenith_angle)

    def set_altitude(self, altitude=0):
        """
        Set how many atmospheric scale heights we are
        above or below the reference radius of this
        atmosphere.

        Parameters
        ----------
        altitude : float
            The altitude at which we should be floating
            in the atmosphere, in units of scale heights.
            For example, +1.0 is one scale height above
            the reference radius and therefore a more
            transparent atmosphere, and -1.0 is one scale
            height deeper than the reference radius and
            therefore a more opaque atmosphere.
            Negative altitudes are probably a weird concept
            if the reference radius is the surface of a
            rocky planet, but they're quite reasonable
            on gas-dominated planets.
        """
        self.altitude = altitude
        self.tau_zenith = self._tau_zenith_reference * np.exp(-altitude)

        # guess what is scattering vs. extinction
        self.guess_scattering()

    def transmit(self, spectrum):
        """
        Calculate the spectrum of light that results from
        transmitting a known light source through this
        atmosphere.

        Parameters
        ----------
        spectrum : Spectrum
            Any light source that you want to transmit.
        """

        # create a composite object, linking the source and the atmosphere
        return Sunset(source=spectrum, atmosphere=self)


class DiscreteAtmosphere(Atmosphere):
    """
    The DiscreteAtmosphere represents an atmosphere
    whose transmission as a function of wavelength
    can be represented as a grid of zenithward optical
    depths. Most pre-calculated models will fall
    into this category.
    """

    def __init__(self, zenith_angle=0.0 * u.deg, altitude=0.0, **kwargs):
        """
        In inherited classes, this initialization relies on the
        existence of a method called ".read_transmission"
        that will create the attributes:
            ._wavelengths (with units of wavelength)
            ._tau_zenith_reference (unitless)
            .H (with units of length)
            .radius (with units of length)
        """

        # read the transmission spectrum data
        self.read_transmission(**kwargs)
        self.default_wavelengths = self._wavelength

        # set the zenith angle (or fall back to the current setting)
        self.set_zenith_angle(zenith_angle)

        # set the altitude at which we're hovering
        self.set_altitude(altitude)

    def guess_scattering(self, visualize=False):
        """
        Make a guess at the Rayleigh scattering component of
        the atmospheric extinction, for use in estimating
        the sky color.

        This should be called once per every new altitude.
        """

        # figure out the appropriate normalization
        self._rayleigh_normalization = np.min(
            self.tau_zenith * self.default_wavelengths ** 4
        )

        w = self.default_wavelengths
        self.tau_zenith_scatter = self._rayleigh_normalization / w ** 4
        self.tau_zenith_absorb = self.tau_zenith - self.tau_zenith_scatter

        if visualize:

            plt.plot(w, self.tau_zenith, zorder=100, label="total")
            plt.plot(
                w,
                self.tau_zenith_scatter,
                alpha=0.5,
                label="scattering",
            )
            plt.plot(
                w,
                self.tau_zenith_absorb,
                alpha=0.5,
                label="absorbtion",
            )
            plt.legend()
            plt.xscale("log")
            plt.yscale("log")
            plt.xlabel("Wavelength ({})".format(w.unit.to_string("latex_inline")))
            plt.ylabel(r"$\tau_{zenith}$")

    """
    def tau_zenith_scatter(self):
        '''
        Estimate the optical depth to *absorption*.
        (at the default wavelengths of the grid)
        '''
        w = self.default_wavelengths
        return self._rayleigh_normalization/w**4

    def tau_zenith_absorb(self):
        '''
        Estimate the optical depth to *absorption*.
        (at the default wavelengths of the grid)
        '''
        # calculate the extinction from absorption
        return self.tau_zenith - self.tau_zenith_scatter()
    """

    def fortney_factor(self):
        """
        The ratio of optical depth between a horizontal slant path
        through an atmosphere to a vertical path, accounting
        for the sphericity of the planet. This is a simple approach
        to connect exoplanet transmission spectra to vertical
        optical depths through an atmosphere (see Fortney 2005).

        Returns
        -------
        fortney_factor : float
            The ratio of slant to vertical optical depths.
        """

        # calculate the fortney factor = ratio of slant/vertical optical depths
        return np.sqrt(2 * np.pi * self.radius / self.H).decompose()

    def transit_radius(self, wavelength=None):
        """
        Calculate the effective transit radius of the planet's atmosphere,
        as seen from infinitely far away.

        Parameters
        ----------
        wavelength : astropy.units.quantity.Quantity
            The wavelengths on which we want the transit radius.

        Returns
        -------
        transit_radius : astropy.units.quantity.Quantity
            The projected radius of the planet, with units.
        """
        # make sure at least some grid of wavelengths is defined
        w = self.wavelength(wavelength)

        # figure out the slant optical depth at the reference radius
        effective_airmass = self.fortney_factor()
        tau = self._tau_zenith_reference * effective_airmass

        # bin this spectrum to the particular wavelength grid
        # KLUDGE -- this should really be in transmission space!
        neww, tau_slant = bintogrid(
            self._wavelength.to("nm").value,
            tau,
            newx=w.to("nm").value,
            drop_nans=False,
        )

        # make sure the wavelengths match up
        assert np.all(neww == w.to("nm").value)

        # convert back to tau
        z_over_H = np.log(tau_slant)

        return self.radius + self.H * z_over_H

    def transmission(self, wavelength=None, zenith_angle=None, altitude=None):
        """
        Calculate the transmission through the atmosphere.

        Parameters
        ----------
        wavelength : astropy.units.quantity.Quantity
            The wavelengths on which we want the transmission.

        zenith_angle : astropy.units.quantity.Quantity
            The angle from zenith along which transmission
            should be calculated.

        altitude : float
            The altitude at which we should be floating
            in the atmosphere, in units of scale heights.

        Returns
        -------
        transmission : numpy.ndarray
            The fractional transmission through the atmosphere.
        """

        # update the zenith angle, if necessary
        if zenith_angle is not None:
            self.set_zenith_angle(zenith_angle)

        # update the altitude, if necessary
        if altitude is not None:
            self.set_altitude(altitude)

        # make sure at least some grid of wavelengths is defined
        w = self.wavelength(wavelength)

        # figure out the transmission at this altitude
        # FIXME -- this is a major kludge! do the integral!
        effective_airmass = np.minimum(self.fortney_factor() / 2.0, self.airmass)
        tau = self.tau_zenith * effective_airmass

        # bin this spectrum to the particular wavelength grid
        # FIXME: binning choice gets real scary with transmission
        neww, newt = bintogrid(
            self._wavelength.to("nm").value,
            np.exp(-tau),
            newx=w.to("nm").value,
            drop_nans=False,
        )

        # make sure the wavelengths match up
        assert np.all(neww == w.to("nm").value)

        # return the binned transmission
        return newt
