from .atmosphere import *


class HotJupiter(DiscreteAtmosphere):
    def read_transmission(self, TiO=False, **kwargs):
        """
        Read a Hot Jupiter's transmission spectrum in from a file
        and define the required attributes:

            ._wavelengths (with units of wavelength)
            ._tau_zenith_reference (unitless)
            .H (with units of length)
            .radius (with units of length)
        """

        # load a transmission spectrum from Fortney et al. (2010)
        if TiO:
            choice = "lambda_1500K_g10_wTiOVO.dat"
        else:
            choice = "lambda_1500K_g10_noTiOVO.dat"
        filename = os.path.join(data_directory, f"fortney/{choice}")
        d = ascii.read(filename, comment="#")

        # define the wavelength grid
        self._wavelength = d["wavelength"].data * 1e3 * u.nm

        # store the effective transit radius
        radiusinkm = d["radiusinkm"].data * u.km
        self._transit_radius = radiusinkm.to("Rjup")

        # pick an (arbitrary-ish) reference radius
        reference_radius = np.min(radiusinkm)

        # define geometry of the atmosphere (how spherical?)
        mu = 2.3
        self.T = 1500 * u.K
        g = 10 * u.m / u.s ** 2
        self.H = (con.k_B * self.T / mu / g / con.m_p).to("km")
        self.radius = reference_radius

        # calculate the optical depth at zenith
        z_over_H = ((radiusinkm - reference_radius) / self.H).decompose()
        tau_slant = np.exp(z_over_H)
        self._tau_zenith_reference = (tau_slant / self.fortney_factor()).decompose()
