'''
Read module.
'''

import os
import math
import configparser

import h5py
import numpy as np

from species.core import box
from species.core import constants as con
from species.data import database
from species.read import read_filter


def get_planck(temperature,
               radius,
               distance,
               wavelength,
               specres):
    '''
    :param temperature: Temperature (K)
    :type temperature: float
    :param radius: Radius (Rjup).
    :type radius: float
    :param distance: Distance (pc).
    :type distance: float
    :param wavelength: Wavelength range (micron).
    :type wavelength: tuple(float, float)
    :param specres: Spectral resolution
    :type specres: float

    :return: Box with the Planck spectrum.
    :rtype: species.core.box.SpectrumBox
    '''

    wl_points = [wavelength[0]]

    while wl_points[-1] <= wavelength[1]:
        wl_points.append(wl_points[-1] + wl_points[-1]/specres)

    wl_points = np.asarray(wl_points[:-1]) # [micron]
    wl_points *= 1e-6 # [m]

    planck1 = 2.*con.PLANCK*con.LIGHT**2/wl_points**5
    planck2 = np.exp(con.PLANCK*con.LIGHT/(wl_points*con.BOLTZMANN*temperature)) - 1.

    flux = 4.*math.pi * (radius*con.R_JUP/(distance*con.PARSEC))**2 * planck1/planck2 # [W m-2 m-1]
    flux *= 1e-6 # [W m-2 micron-1]

    wl_points *= 1e6 # [micron]

    return box.create_box(boxtype='spectrum',
                          spectrum='planck',
                          wavelength=wl_points,
                          flux=flux,
                          name=None)


class ReadSpectrum:
    '''
    Text
    '''

    def __init__(self,
                 spectrum,
                 filter_name):
        '''
        :param spectrum: Spectral library.
        :type spectrum: str
        :param filter_name: Filter name. Full spectrum is read if set to None.
        :type filter_name: str

        :return: None
        '''

        self.spectrum = spectrum
        self.filter_name = filter_name

        if filter_name is None:
            self.wl_range = None

        else:
            transmission = read_filter.ReadFilter(filter_name)
            self.wl_range = transmission.wavelength_range()

        config_file = os.path.join(os.getcwd(), 'species_config.ini')

        config = configparser.ConfigParser()
        config.read_file(open(config_file))

        self.database = config['species']['database']

    def get_spectrum(self,
                     ignore_nan=True,
                     sptype=None):
        '''
        :param ignore_nan: Ignore wavelength points for which the flux is NaN.
        :type ignore_nan: bool
        :param sptype: Spectral types to select. All spectra are retrieved if set to None.
        :type sptype: str

        :return: Box with the spectra.
        :rtype: species.core.box.SpectrumBox
        '''

        h5_file = h5py.File(self.database, 'r')

        try:
            h5_file['spectra/'+self.spectrum]

        except KeyError:
            h5_file.close()
            species_db = database.Database()
            species_db.add_spectrum(self.spectrum)
            h5_file = h5py.File(self.database, 'r')

        list_wavelength = []
        list_flux = []
        list_name = []
        list_simbad = []
        list_sptype = []
        list_distance = []

        for item in h5_file['spectra/'+self.spectrum]:
            data = h5_file['spectra/'+self.spectrum+'/'+item]

            wavelength = data[0, :] # [micron]
            flux = data[1, :] # [W m-2 micron-1]

            if data.shape[0] == 3:
                error = data[2, :] # [W m-2 micron-1]

            if ignore_nan:
                indices = np.isnan(flux)
                indices = np.logical_not(indices)
                indices = np.where(indices)[0]

                wavelength = wavelength[indices]
                flux = flux[indices]

                if data.shape[0] == 3:
                    error = error[indices]

            if self.wl_range is None:
                wl_index = np.arange(0, len(wavelength), 1)

            else:
                wl_index = (flux > 0.) & (wavelength > self.wl_range[0]) & \
                           (wavelength < self.wl_range[1])

            count = np.count_nonzero(wl_index)

            if count > 0:
                index = np.where(wl_index)[0]

                if index[0] > 0:
                    wl_index[index[0] - 1] = True

                if index[-1] < len(wl_index)-1:
                    wl_index[index[-1] + 1] = True

                list_wavelength.append(wavelength[wl_index])
                list_flux.append(flux[wl_index])

                attrs = data.attrs
                if 'name' in attrs:
                    list_name.append(data.attrs['name'])

                if 'simbad' in attrs:
                    list_simbad.append(data.attrs['simbad'])

                if 'sptype' in attrs:
                    list_sptype.append(data.attrs['sptype'])

                if 'distance' in attrs:
                    list_distance.append(data.attrs['distance'])

        specbox = box.SpectrumBox()

        specbox.spectrum = self.spectrum
        specbox.wavelength = np.asarray(list_wavelength)
        specbox.flux = np.asarray(list_flux)

        if list_name:
            specbox.name = np.asarray(list_name)

        if list_simbad:
            specbox.simbad = np.asarray(list_simbad)

        if list_sptype:
            specbox.sptype = np.asarray(list_sptype)

        if list_distance:
            specbox.distance = np.asarray(list_distance)

        if sptype is not None:
            indices = np.where(np.chararray.startswith(specbox.sptype, sptype))[0]

            specbox.wavelength = specbox.wavelength[indices]
            specbox.flux = specbox.flux[indices]
            specbox.name = specbox.name[indices]
            specbox.simbad = specbox.simbad[indices]
            specbox.sptype = specbox.sptype[indices]
            specbox.distance = specbox.distance[indices]

        return specbox
