__author__ = 'DGriffith, ARamkilowan'

""" This module provides a variety of utilities related to radiometry, including, but not limited to
- Creating, writing and reading of MODTRAN-format .flt (filter) files commonly used to to represent
    spectral power/density functions such as transmittances, reflectance or spectral responsivity functions (SRF).
    The original purpose was to specify sensor spectral response functions (i.e. more than just a "filter").
-
"""

import numpy as np
import pandas as pd
import xarray as xr
import io
import matplotlib.pyplot as plt
import easygui  # for simple file/open dialogs and such
import re

import libraddask.rad.xd as xd

""" This module provides much of functionality related radiometry required by MORTICIA.
Included here is functionality for :
1) Interfacing to radiative transfer codes.
2) Dealing with spectral filtering and convolution
3) Creation, reading and writing of MODTRAN-style flt filter/SRF definitions

TO DO : Check that MODTRAN can read the .flt files created using the Flt class here
"""

#_micronsymbol = u"\u00B5"
# _micrometres = _micronsymbol + u"m"

# Micron symbol can be encoded in UTF-8 with
# _micronsymbol = u'\xB5'.encode('UTF-8')
_micronsymbol = '\x00B5'
_micrometres = _micronsymbol + "m"

def srfgen(center, fwhm, n=101, shape='gauss', yedge=0.001, wvmin=None, wvmax=None, centerflat=0.0,
           oob=np.nextafter(0.0, 1), peakval=1.0, units='nm'):
    """ Generate a spectral filter or spectral response function of various shapes

    :param center: center wavelength of the filter in nm
    :param fwhm: full width at half maximum of the filter  in nm
    :param n: number of spectral samples in the filter (minimum of 3, default 101), should be odd
    :param shape: filter shape, one of 'gauss', 'bartlett', 'welch', 'cosine', 'box', 'cos^2' (default 'gauss')
    :param yedge: minimum y-value at the limits of the filter (default 0.001). No filter values below this threshold
    :param wvmin: extend spectral definition by adding a single point at (wvmin, oob)
    :param wvmax: exptend spectral definition by adding a single point at (wvmax, oob)
    :param centerflat: opens a flat region in the centre of filter having a width of centerflat nm
    :param oob: out-of-band leakage, default 0.0, must be <= yedge
    :param peakval: simply scales the peak of the filter function to this value (default 1.0)
    :param units: spectral axis units, either 'nm', 'cm^-1' or 'um' (nanometers, wavenumber per cm or microns)
        will be returned
    :return: wvlnm, y, wvn, wvlum (wavelengths in nm, filter values, wavenumbers per cm and wavelengths in microns)
    """
    if n < 3:
        raise ValueError('Input parameter n to rad.srfgen must be greater than 3.')
    if oob > yedge:
        raise ValueError('Input parameter oob to rad.srfgen must be smaller than paramter yedge (default 0.001).')
    # force an odd number of points
    if not np.mod(n, 2.0):
        n += 1
    n2 = int((n-1)/2)  # mid-point
    x = np.arange(-n2, n2+1, 1.0)  # relative x-coordinates for filter
    shape = shape.lower()
    if shape == 'gauss' or shape == 'gaussian':
        sigma = np.sqrt(-n2**2.0 / 2.0 / np.log(yedge))  # scalar standard deviation
        nfwhm = 2.0 * np.sqrt(2.0 * np.log(2.0)) * sigma  # normalized fwhm (scalar)
        y = np.exp(-x**2.0 / 2.0 / sigma**2)
    elif shape == 'bartlett' or shape == 'triangle':
        fudge = n2 * (1+yedge+(yedge/100.0))
        nfwhm = n2 + (yedge * fudge)  # normalized fwhm
        y = 1.0 - (abs(x) / nfwhm)  # bartlett
        y[y<0] = np.nextafter(0.0, 1)  # in case of negative samples, make very small positive value
    elif shape == 'welch':
        fudge = n2 * (0.5+(yedge*0.378))
        alph = n2 + (yedge * fudge)
        nfwhm = np.sqrt(2.0) * alph  # normalized fwhm
        y = 1 - ((x ** 2.0) / (alph ** 2.0))  # welch
        y[y<0] = np.nextafter(0.0, 1)  # in case of negative samples, make very small positive value
    elif shape == 'cosine':
        fudge = n2 * (.6365+(yedge*.5))
        alph = n2 + (yedge * fudge)
        nfwhm = (4.0/3.0) * alph  # normalized fwhm
        y = np.cos((np.pi * x) / (2.0 * alph))
        y[y<0] = np.nextafter(0.0, 1)  # in case of negative samples, make very small positive value
    elif shape == 'cos^2' or shape == 'cosquared':
        fudge = n2 * (.6365+(yedge*.5))
        alph = n2 + (yedge * fudge)  # calculate alpha
        nfwhm = alph  # normalized fwhm
        y = np.cos((np.pi * x) / (2.0 * alph))**2
    elif shape == 'box' or shape == 'tophat':
        y = np.ones(x.shape)
        y[0] = np.nextafter(0.0, 1)
        y[-1] = np.nextafter(0.0, 1)
        nfwhm = np.trapz(y)
    else:
        raise ValueError('Unknown SRF/filter shape input to rad.srfgen')
    delta = fwhm / nfwhm  # Determine sample delta (scalar)
    y[y<yedge] = oob  # suppress values dropping below the edge threshold
    wvl = (np.ones(x.shape, dtype=np.float) * center) + (delta * x)
    if centerflat:  # Open a central flat region simply by shifting the wvl-coordinates
        # print(n2, type(n2))
        # print(centerflat, type(centerflat))
        # print(wvl[:n2+1]-centerflat/2, wvl[n2:]+centerflat/2)
        # print(type(wvl[:n2+1]-centerflat/2), type(wvl[n2:]+centerflat/2))

        wvl = np.hstack((wvl[:n2+1]-centerflat/2, wvl[n2:]+centerflat/2))
        y = np.hstack((y[:n2+1], y[n2:]))
    if wvmin and wvmin < wvl[0]:  # insert a point at lower wavelength extremity
        wvl = np.hstack((wvmin, wvl))
        y = np.hstack((oob, y))
    if wvmax and wvmax > wvl[-1]:  # insert a point at upper wavelength extremity
        wvl = np.hstack((wvl, wvmax))
        y = np.hstack((y, oob))
    if units == 'cm^-1':  # wavenumber per cm
        wvn = wvl
        wvlnm = 1.0e7 / wvn
        wvlum = wvlnm / 1000.0
    elif units == 'nm':  # wavelength in nm
        wvlnm = wvl
        wvn = 1.0e7 / wvlnm
        wvlum = wvlnm / 1000.0
    elif units == 'um' or units == _micrometres:  # wavelength in microns
        wvlum = wvl
        wvlnm = 1000.0 * wvlum
        wvn = 1e7 / wvlnm
    else:
        raise ValueError('Wavelength/wavenumber units for MODTRAN .flt definitions must be "cm^1", "nm" or "um".')
    y *= peakval  # finally, scale the peak value
    return wvlnm, y, wvn, wvlum

def tophat(center, fwhm, delta=0.0, wvmin=None, wvmax=None, oob=0.0, units='nm'):
    """ Return tophat/box filter defined by 6 points
    Can also specify out-of-band values and extreme limits, which adds upper and lower bound points

    :param center: center wavelength in nm
    :param fwhm: full width at half max in nm
    :param delta: smallest x-coordinate increment, default see np.nextafter
    :param wvmin: minimum wavelength to reach default None
    :param wvmax: maximum wavelength to reach default None
    :param oob: out-of-band leakage
    :param units: spectral axis units, either 'nm', 'cm^-1' or 'um' (nanometers, wavenumber per cm or microns)
        will be returned
    :return: wvlnm, y, wvn, wvlum (wavelengths in nm, filter values, wavenumbers per cm and wavelengths in microns)
    """
    edgelo = center - fwhm/2.0
    edgehi = center + fwhm/2.0
    if delta != 0.0:
        wvl = np.array([edgelo-delta, edgelo, edgelo+delta, edgehi-delta, edgehi, edgehi+delta])
    else:
        wvl = np.array([np.nextafter(edgelo, -1), edgelo, np.nextafter(edgelo, 1),
                        np.nextafter(edgehi, -1), edgehi, np.nextafter(edgehi, 1)])
    y = np.array([oob, 0.5, 1.0, 1.0, 0.5, oob])
    if wvmin:
        wvl = np.hstack((wvmin, wvl))
        y = np.hstack((oob, y))
    if wvmax:
        wvl = np.hstack((wvl, wvmax))
        y = np.hstack((y, oob))
    if units == 'cm^-1':  # wavenumber per cm
        wvn = wvl
        wvlnm = 1.0e7 / wvn
        wvlum = wvlnm / 1000.0
    elif units == 'nm':  # wavelength in nm
        wvlnm = wvl
        wvn = 1.0e7 / wvlnm
        wvlum = wvlnm / 1000.0
    elif units == 'um' or units == _micrometres:  # wavelength in microns
        wvlum = wvl
        wvlnm = 1000.0 * wvlum
        wvn = 1e7 / wvlnm
    else:
        raise ValueError('Wavelength/wavenumber units for MODTRAN .flt definitions must be "cm^1", "nm" or "um".')
    return wvl, y, wvn, wvlum

class Flt(object):
    """ Encapsulates a MODTRAN-style .flt spectral response function/filter function definition file.

    """

    def __init__(self, name, units='nm', filterheaders=[], filters=[], centers=[],
                 fwhms=[], shapes=['gauss'], yedges=[0.001], centerflats=[0.0], peakvals=[1.0], wvmins=[], wvmaxs=[],
                 oobs=[np.nextafter(0.0, 1)]):
        """ Create a filter definition object (MODTRAN flt style)
        Input name is mandatory. All other inputs are optional, but the filterheaders list must have the same number of
        string elements as the number of filters. Also, either the filters are given explicitly in the filters input, or
        a list of filter definitions are provided for use with rad.srfgen().
        If not empty, inputs centers through to oobs must be either scalar or have the same number of list elements as
        the filterheaders list. If scalar, the value will be replicated up to the number of filterheader values.

        :param name: Name of the set of filters. If the filterheaders input to this constructor function
            is empty, an attempt will be made to read the data from a file called name + '.flt'
        :param units: Spectral coordinate units for the filter, 'nm', 'um' or 'cm^-1'
        :param filterheaders: List of strings, one header for each filter/SRF in the set.
        :param filters: A list of numpy arrays. Each list element must comprise a 2-column numpy array with the
            spectral coordinate (wavelength in nm or micron or wavenumber per cm) in the first column and the filter
            magnitude in the second column.
        :param centres: rather than provide filters, the inputs to rad.srfgen can be provided, this is a list of center
            wavelengths in nm
        :param fwhms: list of full width at half maximum in nm
        :param shapes: list of strings providing the shapes of the filters (see rad.srfgen for alternatives)
        :param yedges: list of yedge values (see rad.srfgen)
        :param centerflats: list of centerflat values (see rad.srfgen). Note that giving a centerflat value adds this
            amount ot the fwhm of the filters (broadens the resulting width to centerflat + fwhm)
        :param peakvals: list of peak values of the filters
        :param wvmins: list of minimum wavelength limits in nm
        :param wvmaxs: list of maximum wavelength limits in nm
        :param oobs: list of out-of-band leakage values
        :return: MODTRAN-style filter/SRF object
        """
        self.name = name
        if name[-4:].lower() != '.flt':
            self.filename = name + '.flt'
        else:
            self.filename = name
        if not filterheaders:  # Attempt to read from a file
            self.read(self.filename)
            return
        # Check the units input and set up the file header
        if units == 'cm^-1':
            self.unitsheader = 'W' # wavenumber
        elif units == 'nm':
            self.unitsheader = 'N'
        elif units == 'um' or units == _micrometres:
            self.unitsheader = 'M'
        else:
            raise ValueError('Wavelength/wavenumber units for MODTRAN .flt definitions must be "cm^1", "nm" or "um".')
        self.units = units
        self.fileheader = self.unitsheader + ' ' + self.name
        #if len(filterheaders) != len(filters):
        #    raise ValueError('Number of filterheaders must equal number of filters when creating .flt objects.')
        self.filterheaders = filterheaders
        if filters and centers:
            raise ValueError('Do not give both explicit filter definitions and rad.srfgen definitions for rad.flt.')
        nfilters = len(filterheaders)
        self.nfilters = nfilters
        if filters:
            if len(filters) != nfilters:
                raise ValueError('Number of filters must equal number of filter headers in rad.flt instantiation.')
            self.filters = filters  # now fix the filters into the correct column order
            for ifilt in range(nfilters):
                if self.unitsheader == 'W':  # wavenumber
                    # column order is wvnm, y, wvn, wvum
                    self.filters[ifilt] = np.vstack((1.0e7/filters[ifilt][:,0], filters[ifilt][:,1],
                                                     filters[ifilt][:,0], 1.0e7/filters[ifilt][:,0] / 1000.0)).T
                elif self.unitsheader == 'N':
                    self.filters[ifilt] = np.vstack((filters[ifilt][:,0], filters[ifilt][:,1],
                                                     1.0e7 / filters[ifilt][:,0], filters[ifilt][:,0]/1000.0)).T
                elif self.unitsheader == 'M':
                    self.filters[ifilt] = np.vstack((1000.0*filters[ifilt][:,0], filters[ifilt][:,1],
                                                     1.0e10 /filters[ifilt][:,0], filters[ifilt][:,0])).T
            return  # filter definitions have been given explicitly
        # Check that parameters are scalar or multiply them up to size
        checklist = ['centers', 'fwhms', 'shapes', 'yedges', 'centerflats', 'peakvals', 'oobs']
        for checkitem in checklist:
            checkcall = 'use_' + checkitem + ' = Flt.checkparm("' + checkitem + '", ' + checkitem + ', nfilters)'
            exec(checkcall)
        if wvmins and len(wvmins) == 1:
            use_wvmins = wvmins * nfilters
        elif wvmins and len(wvmins) != nfilters:
            raise ValueError('Number of wvmins must equal number of filterheaders in rad.flt instantiation')
        else:
            use_wvmins = [[] for ie in range(nfilters)]
        # Check wvmaxs
        use_filters = []
        if wvmaxs and len(wvmaxs) == 1:
            use_wvmaxs = wvmaxs * nfilters
        elif wvmaxs and len(wvmaxs) != nfilters:
            raise ValueError('Number of wvmaxs must equal number of filterheaders in rad.flt instantiation')
        else:
            use_wvmaxs = [[] for ie in range(nfilters)]
        for ifilt in range(nfilters):
            # print(ifilt)
            # print(len(centers ))
            # print(len(fwhms ))
            # print(len(shapes ))
            # print(len(yedges ))
            # print(len(wvmins ))
            # print(len(use_wvmaxs ))
            # print(len(centerflats ))
            # print(len(oobs ))
            # print(len(peakvals ))
            wvl, y, wvn, wu = srfgen(center=centers[ifilt], fwhm=fwhms[ifilt], shape=shapes[ifilt],
                                 yedge=yedges[ifilt], wvmin=wvmins[ifilt], wvmax=use_wvmaxs[ifilt],
                                 centerflat=centerflats[ifilt], oob=oobs[ifilt], peakval=peakvals[ifilt])
            use_filters.append(np.vstack((wvl, y, wvn, wu)).T)
        self.filters = use_filters

    @staticmethod  # Some input parameter checking for Flt constructor
    def checkparm(parmname, parm, nfilters):
        """ Input parameter checking for Flt constructor

        :param parmname: Name of parameter fpr checking
        :param parm: Parameter value
        :param nfilters: Number of filters
        :return: Checked parameter
        """
        if len(parm) == 1:
            parm *= nfilters
        elif len(parm) != nfilters:
            raise ValueError('Number of ' + parmname + ' must equal number of filterheaders in rad.flt instantiation')
        return parm

    def read(self, filename, name='Unknown'):
        """ Read a .flt format spectral band filter definitions file (MODTRAN format)

        :param filename:
        :return: object of class Flt, if the file is a well-formatted MODTRAN-style .flt file
        """
        isfloatnum = '^[-+]?[0-9]*\.?[0-9]+([eE][-+]?[0-9]+)?$'  # regular expression to match a floating point number
        if filename == '.flt':
            filename = easygui.fileopenbox(msg='Please select a .flt file.', filetypes=["*.flt"])
        self.filename = filename
        self.name = name
        with open(filename, 'rt') as fltfil:
            fileheader = fltfil.readline()
            self.unitsheader = fileheader[0].upper()  # Must be W, N or M for wavenumber per cm, nm or microns
            if self.unitsheader == 'W':
                self.units = 'cm^-1'
            elif self.unitsheader == 'N':
                self.units = 'nm'
            elif self.unitsheader == 'M':
                self.units = _micrometres
            else:
                raise ValueError('File header for ' + filename + ' does not start with N, M or W as required for'
                                                                 ' .flt files.')
            self.name = fileheader[1:].strip()
            self.fileheader = fileheader.strip()
            self.filterheaders = []
            self.filters = []
            ifilter = 0  # count the filters
            nexlin = fltfil.readline()
            while nexlin:
                self.filterheaders.append(nexlin.strip())
                nexlin = fltfil.readline()
                if not nexlin:
                    break  # all done
                slin = nexlin.split()  # split line up into a list of tokens at whitespace
                thisfilterdata = np.array([])
                self.filters.append([])
                # if all tokens match
                while all([re.match(isfloatnum, tok) for tok in slin]):  # all tokens a are numbers-this is a line of data
                    # accumulate the data
                    thelinedata = np.array(slin).astype(np.float)
                    if len(thisfilterdata) != 0:
                        thisfilterdata = np.vstack((thisfilterdata, thelinedata))
                    else:
                        thisfilterdata = thelinedata
                    nexlin = fltfil.readline()
                    if not nexlin:
                        break  # all done
                    slin = nexlin.split()  # split line up into a list of tokens at whitespace
                # process the accumulated data depending on what the units are
                if self.unitsheader == 'W':
                    self.filters[ifilter] = np.vstack((1e7/thisfilterdata[:,1], thisfilterdata[:,1],
                                                       thisfilterdata[:,0], 1e9/thisfilterdata[:,1])).T
                elif self.unitsheader == 'N':
                    self.filters[ifilter] = np.vstack((thisfilterdata[:,0], thisfilterdata[:,1],
                                                       1e7/thisfilterdata[:,0], thisfilterdata[:,0]/1000.0)).T
                elif self.unitsheader == 'M':
                    self.filters[ifilter] = np.vstack((thisfilterdata[:,0]*1000.0, thisfilterdata[:,1],
                                                       1e9/thisfilterdata[:,0], thisfilterdata[:,0])).T
                ifilter += 1
            self.nfilters = ifilter

    def __repr__(self, format='  %f'):
        """ Return representation of a .flt MODTRAN-style spectral filter or SRF. This is the same as the format in
         the .flt file is written.

        :param format: format in which to present the numeric data, default is '  %f'
        :return: File representation of .flt object
        """
        selfrep = self.fileheader + '\n'
        for (ifilt, filter) in enumerate(self.filters):
            selfrep = selfrep + self.filterheaders[ifilt] + '\n'
            # Create a string buffer
            # strbuff = io.StringIO.StringIO()
            strbuff = io.StringIO()
            # column ordering depends on original unit specification
            if self.unitsheader == 'W':
                np.savetxt(strbuff, self.filters[ifilt][:,[2,1,0]], fmt=format)
            elif self.unitsheader == 'N':
                np.savetxt(strbuff, self.filters[ifilt][:,[0,1,2]], fmt=format)
            elif self.unitsheader == 'M':
                np.savetxt(strbuff, self.filters[ifilt][:,[3,1,2]], fmt=format)
            selfrep = selfrep + strbuff.getvalue()
            strbuff.close()  # delete the buffer
        return selfrep

    def write(self, filename=None, format='  %f'):
        """ Write a MODTRAN-style .flt file for this filter/SRF set

        :param filename: Optional filename without extension. If not given, the filter name will be used with
        :param format: Format specifier as for np.savetext for writing the data, default is '  %f'
            extension .flt
        :return: None
        """
        selfrep = self.__repr__(format=format)
        # Write the string to the file
        if filename:
            if filename[-4:].lower != '.flt':
                filename = filename + '.flt'
        else:
            filename = self.filename
        with open(filename, 'wt') as fltfile:
            fltfile.write(selfrep)

    def plot(self, filter_numbers = None):
        """ Plot a MODTRAN-style set of filter/SRF curves.

        :param filter_numbers: List of filter numbers to plot, defaults to all filters. Filter indices start at 0.

        :return: None
        """
        # plt.hold(True)
        if filter_numbers is None:
            filter_numbers = range(self.nfilters)
        for ifilt in filter_numbers:
            if self.unitsheader == 'W':
                plt.plot(self.filters[ifilt][:, 2], self.filters[ifilt][:,1])
            elif self.unitsheader == 'N':
                plt.plot(self.filters[ifilt][:, 0], self.filters[ifilt][:,1])
            elif self.unitsheader == 'M':
                plt.plot(self.filters[ifilt][:, 3], self.filters[ifilt][:,1])
        plt.title(self.name)
        plt.ylabel('Spectral Response/Transmission')
        if self.units == 'cm^-1':
            plt.xlabel('Wavenumber [' + self.units + ']')
        else:
            plt.xlabel('Wavelength [' + self.units + ']')
        if len(self.filterheaders) <= 12:
            plt.legend(self.filterheaders, loc='best')
        plt.grid()
        # plt.hold(False)

    def flt_as_xd(self):
        """ Convert an Flt class object to a list of xarray DataArray objects

        :return: The set of Flt filters as a list of xarray DataArray objects, with a wavelength coordinate
            axis ('wvl', long_name = 'Wavelength'
        """
        flt_list = []
        for ifilt in range(self.nfilters):
            wvl = xd.xd_identity(self.filters[ifilt][:,0], 'wvl', self.units)
            #print(wvl)
            flt_list.append(xr.DataArray(self.filters[ifilt][:,1], [wvl], name='srf',
                                                attrs={'long_name': long_name['srf'],
                                                       'labels': self.filterheaders[ifilt],
                                                       'units': default_units['srf'],
                                                       'title': self.name},
                                                ))
        return flt_list

    def flt_as_xd_harmonised(self, quantity_name='srf', chn_start_index=0):
        """ Convert the Flt class object into a single, wavelength-harmonised xr.DataArray.

        :param quantity_name: The name of the quantity as defined the long_names variable in moglo.py. Defaults to
            'srf', a Spectral Response Function, but could also be a transmission functions 'trn' or other spectral
            quantity known to moglo.py.
        :param chn_start_index: Use this parameter to select the starting channel number. Defaults to zero.
        :return: The set of Flt filters as a single, wavelength-harmonised xr.DataArray object. The filter
            headers are returned in an attribute called 'labels'. The fileheader of the Flt object is
            returned in an attribute called 'title' (netCDF recommendation)
        """
        flt_list = self.flt_as_xd()  # Create a list of xr.DataArray objects
        # Harmonise the wavelength axes
        flt_list_harmonised = xd.xd_harmonise_interp(flt_list)
        xd.xd_attrs_update(flt_list_harmonised)  # Update the attribute
        chn_indices = range(chn_start_index, chn_start_index + len(flt_list))
        # Compile the list into a single object
        flt_data = np.vstack([flt_list_harmonised[ifilt].data for ifilt in range(len(flt_list))])
        flt_harmonised = xr.DataArray(flt_data.T, [(flt_list_harmonised[0]['wvl']),
                                                      ('chn', chn_indices,
                                                         {'labels': self.filterheaders})],
                                           name=quantity_name,
                                           attrs={'long_name': long_name[quantity_name],
                                                  'units': default_units[quantity_name],
                                                  'title': self.name})
        return flt_harmonised

# Definitions for Kato spectral channels, [centre_wavelength, lower_wavelength, upper_wavelength]
# Kato wavelengths are in nm
kato_channels = np.array([
    [ 256.300,  240.1185,  272.4815],
    [ 277.948,  272.4815,  283.4140],
    [ 295.127,  283.4140,  306.8408],
    [ 317.306,  306.8408,  327.7722],
    [ 345.136,  327.7722,  362.5000],
    [ 385.000,  362.5000,  407.5000],
    [ 429.773,  407.5000,  452.0458],
    [ 484.863,  452.0458,  517.6806],
    [ 528.840,  517.6806,  540.0000],
    [ 544.800,  540.0000,  549.5000],
    [ 558.050,  549.5000,  566.6000],
    [ 585.800,  566.6000,  605.0000],
    [ 615.000,  605.0000,  625.0000],
    [ 645.850,  625.0000,  666.7000],
    [ 675.439,  666.7000,  684.1772],
    [ 694.313,  684.1772,  704.4486],
    [ 723.531,  704.4486,  742.6139],
    [ 767.046,  742.6139,  791.4788],
    [ 817.968,  791.4788,  844.4581],
    [ 866.714,  844.4581,  888.9693],
    [ 931.938,  888.9693,  974.9063],
    [1010.320,  974.9063, 1045.7440],
    [1119.970, 1045.7440, 1194.1880],
    [1355.060, 1194.1880, 1515.9400],
    [1564.700, 1515.9400, 1613.4510],
    [1789.120, 1613.4510, 1964.7980],
    [2059.130, 1964.7980, 2153.4640],
    [2214.330, 2153.4640, 2275.1900],
    [2638.540, 2275.1900, 3001.8930],
    [3318.660, 3001.8930, 3635.4170],
    [3813.210, 3635.4170, 3991.0030],
    [4298.320, 3991.0030, 4605.6540]])
kato_units = 'nm'

fu_channels = np.array([
    [ 0.55,   0.20000,      0.68966],
    [ 1.00,   0.68966,      1.29870],
    [ 1.60,   1.29870,      1.90476],
    [ 2.20,   1.90476,      2.50000],
    [ 3.00,   2.50000,      3.50877],
    [ 3.70,   3.50877,      4.00000],
    [ 4.90,   4.54545,      5.26316],
    [ 5.60,   5.26316,      5.88235],
    [ 6.50,   5.88235,      7.14286],
    [ 7.60,   7.14286,      8.00000],
    [ 8.50,   8.00000,      9.09091],
    [ 9.60,   9.09091,     10.20410],
    [11.30,  10.20410,     12.50000],
    [13.70,  12.50000,     14.92540],
    [16.60,  14.92540,     18.51850],
    [21.50,  18.51850,     25.00000],
    [30.00,  25.00000,     35.71430],
    [70.00,  35.71430,  10000.00000]])
fu_units = 'um'

avhrr_kratz_channels = np.array([
    # channel 1 (band 5 actually contributes to channel 1 and 2)
    [  0.569917,  0.561798,  0.578035],    # avhrr15.f
    [  0.590222,  0.578035,  0.602410],    # avhrr14.f
    [  0.613705,  0.602410,  0.625000],    # avhrr13.f
    [  0.657328,  0.625000,  0.689655],    # avhrr12.f
    [  0.720768,  0.689655,  0.751880],    # avhrr11.f
    # channel 2
    [  0.763537,  0.751880,  0.775194],    # avhrr24.f
    [  0.842143,  0.775194,  0.909091],    # avhrr23.f
    [  0.944741,  0.909091,  0.980392],    # avhrr22.f
    [  1.011030,  0.980392,  1.041667],    # avhrr21.f
    # channel 3
    [  3.55005,  3.496503,  3.603604],    # avhrr35.f
    [  3.65023,  3.603604,  3.696858],    # avhrr34.f
    [  3.74957,  3.696858,  3.802281],    # avhrr32.f
    [  3.85427,  3.802281,  3.906250],    # avhrr32.f
    [  3.96116,  3.906250,  4.016064],    # avhrr31.f
    # channel 4
    [ 10.8365,  10.309278, 11.363636],    # avhrr41.f
    # channel 5
    [ 11.9318,  11.363636, 12.500000]])   # avhrr51.f
avhrr_kratz_units = 'um'

class SpectralDistribution(object):
    """ The SpectralDistribution class defines any band-limited spectral distribution function. This could be
    the spectral response functions of a sensor, or the spectral transmittance of an optical filter, the spectral
    radiance, irradiance or any other band-limited spectral quantity.

    """
    # SpectralChannels is a list of channels that can be indexed in the usual way, by the global channel index
    # _channel_counter = 0  # This is a class global counter, incremented for each channel, so that every
    #                       # defined distribution function gets a unique number
    # _channel_list = []  # The global list of spectral channels indexed self.ichn
    # _channel_groups = []  # Global list of channel group names
    # _channel_group_dict = {}  # Dictionary of channels indexed by group name
    def __init__(self, extreme_limits, in_band_limits=None, in_band_threshold=0.01):
        """ Create a basic, spectral distribution function with certain spectral band limits

        By default, wavelenths are specified in nm for SpectralDistribution and SpectralSpace objects.

        :return: A SpectralDistribution object
        """
        pass

    # Use @classmethod, where the class of object instance is passed implicitly, instead of self
    # This syntax is used to create alternative constructors
    @classmethod
    def kato(cls, i_channel, resolution=0.001):
        """ Return a Kato correlated-k channel definition as a SpectralDistribution. Only a single channel can
        be represented. Use a SpectralSpace to get multiple Kato channels

        :param i_channel: the single kato channel to be obtained. Integer 1 to 32
        :param resolution: this is the spectral resolution in nm of the band edges. Default 0.001 nm
          which is typically adequate.
        :return: A SpectralDistribution object defining the requested Kato channel

        .. seealso: the libRadtran manual
        """

        obj = cls()
        return obj

    @classmethod
    def fu(cls, i_channel, resolution=0.001):
        """ Return a Fu correlated-k channel as a SpectralDistribution

        :param i_channel: the single Fu channel to be obtained. Integer 1 to 18
        :param resolution: this is the spectral resolution in nm of the band edges. Default 0.001 nm
        :return: A SpectralDistribution object defining the requested Fu channel

        .. seealso: the libRadtran manual
        """
        obj = cls()
        return obj

    @classmethod
    def avhrr_kratz(cls, i_channel, resolution=0.001):
        """

        :param i_channel: The single AVHRR Kratz channel to obtain. Integer 1 to 16
        :param resolution: this is the spectral resolution in nm of the band edges. Default 0.001 nm
        :return: A SpectralDistribution object defining the requested avhrr_kratz channel

        .. seealso: the libRadtran manual
        """
        obj = cls()
        return obj

    @classmethod
    def spectral_slice(cls, extreme_limits, in_band_limits, resolution=0.001, oob_leakage=0.0):
        """ Create a SpectralDistribution object as a simple spectral slice

        :param extreme_limits:
        :param in_band_limits:
        :param resolution:
        :param oob_leakage:
        :return:
        """

    @classmethod
    def sensor_channel(cls, platform_series, platform_name, sensor_name, channel_name):
        """ Load a spectral response file from the libRadtran-compatible library as a SpectralDistribution

        The available response function filter files can be viewed in the sub-directory rad/radata/filter.
        Note that only a single channel can be loaded using this method. To load multiple channels, use
        SpectralSpace.sensor_channels().

        :param platform_series: Name of the series of platforms on which the sensor is carried e.g. 'landsat'
        :param platform_name: Name of the specific satellite/platform on which the sensor is carried
            e.g. 'landsat7'
        :param sensor_name: Name of the specific sensor on the platform e.g. 'tm' or 'etm'
        :param channel_name: Name of the specific spectral channel on the given sensor e.g.
        :return:
        """

        obj = cls()
        return obj

    def decimate_resolution(self):
        """ Reduce the number of sample points representing the spectral distribution in an intelligent way.

        This procedure only works for positive spectral power distributions.

        :return:
        """

class SpectralSpace(object):
    """ A SpectralSpace is a set (represented as a list) of SpectralDistribution objects.

    For example, the full Kato correlated-k list of spectral slices constitutes a SpectralSpace.
    A sub-range of of Kato or other correlated-k channels (Fu or avhrr_kratz) also qualify.
    The spectral response functions of a sensor can also be represented using a SpectralSpace.

    An important use of SpectralSpaces is to compute "projections" which could also be thought of as
    "dot products". The projection of one SpectralSpace into another comprises multplying the
    distribution functions and integrating over wavelength to obtain a set of weights. The integral is
    typically also normalised to retain equivalent units.

    Here are some examples:
    A set of spectral end-member functions can be represented as a SpectralSpace. These end-members might
    be propagated to calculate an end-member response at a camera focal plane. The end-members are projected
    onto the spectral response functions of the sensor to get the sensor responses to each of the end-members.

    Notes: A SpectralSpace is not generally orthogonal unless specifically designed so by the user.
    """

    @classmethod
    def from_flt_file(cls, filename, re_select=None):
        """ Create a SpectralSpace list of SpectralDistribution objects by reading a MODTRAN-compatible
        "filter function" .flt file.

        :param filename: The MODTRAN-compatible .flt file from which to read the filter functions.
        :param re_select: Select a sub-set of the channels using a regular expression filter. Only the channel
            names/descriptions that match the regular expression will be included. The default is to read
            all channels in the file.
        :return:
        """

# Some simple OpenEXR utility functions for integration with Numpy
# For more complex OpenEXR handling, use a class
def readOpenEXR(filename):
    """ Simple read function for an OpenEXR file

    Use of this function requires that the OpenEXR package be installed.
    :param filename: The name of the OpenEXR file
    :return channel_names: List of image channel names found in the OpenEXR file
    :return im_dict: All image data in a dictionary keyed by channel names or channel groups. Channels are grouped
        into RGB triplets if the channel names have the form prefix.R, prefix.G and prefix.B.
        All image data is returned as numpy arrays.
    :return header: OpenEXR header as a dictionary.
    """
    import OpenEXR
    import Imath
    if not OpenEXR.isOpenExrFile(filename):
        raise IOError('OpenEXR file was not found, or file is not OpenEXR.')
    exr_file = OpenEXR.InputFile(filename)
    header = exr_file.header()  # Returns a dict
    channel_names = header['channels'].keys()
    data_window = header['dataWindow']
    size = (data_window.max.x - data_window.min.x + 1, data_window.max.y - data_window.min.y + 1)
    # Read all channels
    im_FLOAT = np.array([], dtype=np.float32)
    im_HALF = np.array([], dtype=np.float16)
    im_UINT = np.array([], dtype=np.uint32)
    im_dict = {}
    triplets = {}
    for i_chan, chan_name in enumerate(channel_names):
        pixel_type = str(header['channels'][chan_name]).split()[0]
        if pixel_type == 'FLOAT':
            imath_type = Imath.PixelType(Imath.PixelType.FLOAT)
        elif pixel_type == 'HALF':
            imath_type = Imath.PixelType(Imath.PixelType.HALF)
        elif pixel_type == 'UINT':
            imath_type = Imath.PixelType(Imath.PixelType.UINT)
        chan_data_str = exr_file.channel(chan_name, imath_type)
        if pixel_type == 'FLOAT':
            chan_data = np.fromstring(chan_data_str, dtype=np.float32).reshape((size[1], size[0]))
            if im_FLOAT.size == 0:
                im_FLOAT = chan_data
            else:
                im_FLOAT= np.dstack((im_FLOAT, chan_data))
        elif pixel_type == 'HALF':
            chan_data = np.fromstring(chan_data_str, dtype=np.float16).reshape((size[1], size[0]))
            if im_HALF.size == 0:
                im_HALF = chan_data
            else:
                im_HALF = np.dstack((im_HALF, chan_data))
        elif pixel_type == 'UINT':
            chan_data = np.fromstring(chan_data_str, dtype=np.uint32).reshape((size[1], size[0]))
            if im_UINT.size == 0:
                im_UINT = chan_data
            else:
                im_UINT = np.dstack((im_UINT, chan_data))
        else:
            raise ValueError('Unknown image pixel type in OpenEXR file.')
        im_dict[chan_name] = chan_data
        rgbya = chan_name[-1]  # red, green, blue, luminance, transparency
        prefix = chan_name.split('.')[0]
        if prefix == rgbya:
            prefix = 'rgb'
        if rgbya in 'RGB':
            if (prefix in triplets):
                triplets[prefix][rgbya] = chan_name
            else:
                triplets[prefix] = {rgbya: chan_name}
    for triplet in triplets:
        im_dict[triplet] = np.dstack((im_dict[triplets[triplet]['R']],
                                      im_dict[triplets[triplet]['G']],
                                      im_dict[triplets[triplet]['B']]))
    return channel_names, im_dict, header


import numpy as np
from astropy.coordinates import SkyCoord
from astropy import units as u
from pandas import read_csv
import h5py

import stan_utility

__all__ = ['Model', 'Direction', 'uv_to_coord', 'coord_to_uv']

Mpc_to_km = 3.086E19

class Model():
    """
    Simple wrapper for models defined in Stan.
    """
    def __init__(self,
                 model_filename=None,
                 sim_filename=None,
                 include_paths=None):
        """
        Simple wrapper for models defined in Stan.

        :param model_filename: location of the stan code for model
        :param sim_filename: locaiton of the stan code for simulation
        """

        self.model_filename = model_filename
        self.sim_filename = sim_filename
        self.include_paths = include_paths

        self.simulation = None

    def compile(self, reset=False):
        """
        Compile and cache the necessary Stan models if not already done.

        :param reset: Rerun the compilation
        """

        if self.model_filename:
            self.model = stan_utility.compile_model(
                filename=self.model_filename,
                model_name='model',
                include_paths=self.include_paths,
                reset=reset)

        if self.sim_filename:
            self.simulation = stan_utility.compile_model(
                filename=self.sim_filename,
                model_name='sim',
                include_paths=self.include_paths,
                reset=reset)

    def input(self,
              B=None,
              kappa=None,
              F_T=None,
              f=None,
              L=None,
              F0=None,
              alpha=None,
              Eth=None,
              ptype=None):
        """
        Get simulation inputs.

        :param F_T: total flux [# km-^2 yr^-1]
        :param f: associated fraction
        :param kappa: deflection parameter
        :param B: rms B field strength [nG]
        :param alpha: source spectral index
        :param Eth: threshold energy of study [EeV]
        :param ptype: element of composition
        """
        self.F_T = F_T
        self.f = f
        self.kappa = kappa
        self.B = B
        self.L = L
        self.F0 = F0
        self.alpha = alpha
        self.Eth = Eth
        self.Eth_sim = None  # To be set by Analysis
        self.ptype = ptype

    def _get_properties(self):
        """
        Convenience function to pack object into dict.
        """

        self.properties = {}
        self.properties['F_T'] = self.F_T
        self.properties['f'] = self.f
        self.properties['kappa'] = self.kappa
        self.properties['B'] = self.B
        self.properties['L'] = self.L
        self.properties['F0'] = self.F0
        self.properties['F0'] = self.F0
        self.properties['alpha'] = self.alpha
        self.properties['Eth'] = self.Eth
        self.properties['Eth_sim'] = self.Eth_sim

        self.properties['sim_filename'] = self.sim_filename
        self.properties['model_filename'] = self.model_filename
        self.properties['include_paths'] = self.include_paths

    def save(self, file_handle):
        """
        Save to the passed H5py file handle,
        i.e. something that cna be used with
        file_handle.create_dataset()

        :param file_handle: file handle
        """

        self._get_properties()

        for key, value in self.properties.items():
            try:
                file_handle.create_dataset(key, data=value)
            except:
                pass

class Direction():
    """
    Input the unit vector vMF samples and
    store x, y, and z and galactic coordinates
    of direction in Mpc.
    """
    def __init__(self, unit_vector_3d):
        """
        Input the unit vector samples and
        store x, y, and z and galactic coordinates
        of direction in Mpc.

        :param unit_vector_3d: a 3-dimensional unit vector.
        """

        self.unit_vector = unit_vector_3d
        transposed_uv = np.transpose(self.unit_vector)
        self.x = transposed_uv[0]
        self.y = transposed_uv[1]
        self.z = transposed_uv[2]
        self.d = SkyCoord(self.x,
                          self.y,
                          self.z,
                          unit='mpc',
                          representation_type='cartesian',
                          frame='icrs')
        self.d.representation_type = 'spherical'
        self.glons = self.d.galactic.l.wrap_at(360 * u.deg).deg
        self.glats = self.d.galactic.b.wrap_at(180 * u.deg).deg

        self.ras = self.d.ra.deg
        self.decs = self.d.dec.deg

def uv_to_coord(uv):
    """
    Convert unit vector array into SkyCoord object in the ICRS frame.

    :param uv: array of 3D unit vectors
    :return: astropy SkyCoord object
    """
    transposed_uv = np.transpose(uv)
    x = transposed_uv[0]
    y = transposed_uv[1]
    z = transposed_uv[2]

    c = SkyCoord(x,
                 y,
                 z,
                 unit='Mpc',
                 representation_type='cartesian',
                 frame='icrs')

    return c

def coord_to_uv(coord):
    """
    Convert SkyCoord object into array of unit vecotrs in the ICRS frame.
    Used for input into Stan programs.

    :param coord: astropy SkyCoord object
    :return: an array of 3D unit vectors
    """
    c = coord.icrs
    ds = [c.cartesian.x, c.cartesian.y, c.cartesian.z]
    uv = [d / np.linalg.norm(d) for d in np.transpose(ds)]

    return uv

def convert_scale(D, alpha_T, eps, F0=None, L=None, to_stan=True):
    """
    Convenience function to convert parameters
    to O(1) scale for sampling in Stan.
    D [Mpc] -> (D * 3.086) / 100
    alpha_T [km^2 yr] -> alpha_T / 1000
    eps [km^2 yr] -> eps / 1000
    F [# km^-2 yr^-1] -> F * 1000
    L [# yr^-1] -> L / 1e39

    Can also convert back by setting to_stan = False
    """

    # Convert from physical units to Stan units
    if to_stan:

        D = [(d * 3.086) / 100 for d in D]
        alpha_T = alpha_T / 1000.0
        eps = [e / 1000.0 for e in eps]

        if F0:
            F0 = F0 * 1000.0

        if isinstance(L, (list, np.ndarray)):
            L = L / 1.0e39

    # Convert from Stan units to physical units
    else:

        D = [(d / 3.086) * 100 for d in D]
        alpha_T = alpha_T * 1000.0
        eps = [e * 1000.0 for e in eps]

        if F0:
            F0 = F0 / 1000.0

        if isinstance(L, (list, np.ndarray)):
            L = L * 1.0e39

    if F0 and isinstance(L, (list, np.ndarray)):

        return D, alpha_T, eps, F0, L

    else:

        return D, alpha_T, eps

def get_simulation_input(Nsim, f, D, M, alpha_T):
    """
    For a given associated fraction and
    detector exposure, find the background flux and
    source luminosity as input to the simulation.

    :param Nsim: N simulated, ignoring exposure effects.
    :param f: Associated fraction.
    :param D: List of distances to sources [Mpc].
    :param M: Integral over the angular exposure [sr].
    :param alpha_T: Total exposure [km^2 sr yr].
    """

    FT = (Nsim * M) / alpha_T  # km^-2 yr^-1
    Fs = f * FT
    F0 = (1 - f) * FT

    # Assume equal luminosities
    L = (Fs / (sum([1 / (4 * np.pi * (d * Mpc_to_km)**2) for d in D])))
    L = np.tile(L, len(D))  # yr^-1

    return L, F0


"""
Spectral Synthesis Module of SME
"""
import logging
import uuid

import numpy as np
from scipy.constants import speed_of_light
from scipy.interpolate import interp1d
from scipy.ndimage.filters import convolve
from tqdm import tqdm

from . import broadening
from .atmosphere.interpolation import AtmosphereInterpolator
from .continuum_and_radial_velocity import (
    apply_radial_velocity_and_continuum,
    match_rv_continuum,
    null_result,
)
from .iliffe_vector import Iliffe_vector
from .large_file_storage import setup_lfs
from .sme_synth import SME_DLL

logger = logging.getLogger(__name__)

clight = speed_of_light * 1e-3  # km/s

__DLL_DICT__ = {}
__DLL_IDS__ = {}

class Synthesizer:
    def __init__(self, config=None, lfs_atmo=None, lfs_nlte=None, dll=None):
        self.config, self.lfs_atmo, self.lfs_nlte = setup_lfs(
            config, lfs_atmo, lfs_nlte
        )
        # dict: internal storage of the adaptive wavelength grid
        self.wint = {}
        # dll: the smelib object used for the radiative transfer calculation
        self.dll = dll if dll is not None else SME_DLL()
        self.dll = self.get_dll_id()
        self.atmosphere_interpolator = None
        # This stores a reference to the currently used sme structure, so we only log it once
        self.known_sme = None
        logger.info("Don't forget to cite your sources. Use sme.citation()")

    def get_atmosphere(self, sme):
        """
        Return an atmosphere based on specification in an SME structure

        sme.atmo.method defines mode of action:
            "grid"
                interpolate on atmosphere grid
            "embedded"
                No change
            "routine"
                calls sme.atmo.source(sme, atmo)

        Parameters
        ---------
            sme : SME_Struct
                sme structure with sme.atmo = atmosphere specification

        Returns
        -------
        sme : SME_Struct
            sme structure with updated sme.atmo
        """

        # Handle atmosphere grid or user routine.
        atmo = sme.atmo

        if atmo.method == "grid":
            if self.atmosphere_interpolator is None:
                self.atmosphere_interpolator = AtmosphereInterpolator(
                    depth=atmo.depth,
                    interp=atmo.interp,
                    geom=atmo.geom,
                    lfs_atmo=self.lfs_atmo,
                )
            else:
                self.atmosphere_interpolator.depth = atmo.depth
                self.atmosphere_interpolator.interp = atmo.interp
                self.atmosphere_interpolator.geom = atmo.geom

            atmo = self.atmosphere_interpolator.interp_atmo_grid(
                atmo.source, sme.teff, sme.logg, sme.monh
            )
        elif atmo.method == "routine":
            atmo = atmo.source(sme, atmo)
        elif atmo.method == "embedded":
            # atmo structure already extracted in sme_main
            pass
        else:
            raise AttributeError("Source must be 'grid', 'routine', or 'embedded'")

        sme.atmo = atmo
        return sme

    @staticmethod
    def get_wavelengthrange(wran, vrad, vsini):
        """
        Determine wavelengthrange that needs to be calculated
        to include all lines within velocity shift vrad + vsini
        """
        # 30 km/s == maximum barycentric velocity
        vrad_pad = 30.0 + 0.5 * np.clip(vsini, 0, None)  # km/s
        vbeg = vrad_pad + np.clip(vrad, 0, None)  # km/s
        vend = vrad_pad - np.clip(vrad, None, 0)  # km/s

        wbeg = wran[0] * (1 - vbeg / clight)
        wend = wran[1] * (1 + vend / clight)
        return wbeg, wend

    @staticmethod
    def new_wavelength_grid(wint):
        """Generate new wavelength grid within bounds of wint"""
        # Determine step size for a new model wavelength scale, which must be uniform
        # to facilitate convolution with broadening kernels. The uniform step size
        # is the larger of:
        #
        # [1] smallest wavelength step in WINT_SEG, which has variable step size
        # [2] 10% the mean dispersion of WINT_SEG
        # [3] 0.05 km/s, which is 1% the width of solar line profiles

        wbeg, wend = wint[0], wint[-1]
        wmid = 0.5 * (wend + wbeg)  # midpoint of segment
        wspan = wend - wbeg  # width of segment
        diff = wint[1:] - wint[:-1]
        jmin = np.argmin(diff)
        vstep1 = diff[jmin] / wint[jmin] * clight  # smallest step
        vstep2 = 0.1 * wspan / (len(wint) - 1) / wmid * clight  # 10% mean dispersion
        vstep3 = 0.05  # 0.05 km/s step
        vstep = max(vstep1, vstep2, vstep3)  # select the largest

        # Generate model wavelength scale X, with uniform wavelength step.
        nx = int(
            np.abs(np.log10(wend / wbeg)) / np.log10(1 + vstep / clight) + 1
        )  # number of wavelengths
        if nx % 2 == 0:
            nx += 1  # force nx to be odd

        # Resolution
        # IDL way
        # resol_out = 1 / ((wend / wbeg) ** (1 / (nx - 1)) - 1)
        # vstep = clight / resol_out
        # x_seg = wbeg * (1 + 1 / resol_out) ** np.arange(nx)

        # Python way (not identical, as IDL endpoint != wend)
        # difference approx 1e-7
        x_seg = np.geomspace(wbeg, wend, num=nx)
        resol_out = 1 / np.diff(np.log(x_seg[:2]))[0]
        vstep = clight / resol_out
        return x_seg, vstep

    @staticmethod
    def check_segments(sme, segments):
        if isinstance(segments, str) and segments == "all":
            segments = range(sme.nseg)
        else:
            segments = np.atleast_1d(segments)
            if np.any(segments < 0) or np.any(segments >= sme.nseg):
                raise IndexError("Segment(s) out of range")

        if sme.mask is not None:
            segments = [
                seg
                for seg in segments
                if not np.all(sme.mask[seg] == sme.mask_values["bad"])
            ]
        return segments

    @staticmethod
    def apply_radial_velocity_and_continuum(
        wave, spec, wmod, smod, cmod, vrad, cscale, cscale_type, segments
    ):
        smod = apply_radial_velocity_and_continuum(
            wave, wmod, smod, vrad, cscale, cscale_type, segments
        )
        cmod = apply_radial_velocity_and_continuum(
            wave, wmod, cmod, vrad, None, None, segments
        )
        return smod, cmod

    @staticmethod
    def integrate_flux(mu, inten, deltav, vsini, vrt, osamp=1):
        """
        Produces a flux profile by integrating intensity profiles (sampled
        at various mu angles) over the visible stellar surface.

        Intensity profiles are weighted by the fraction of the projected
        stellar surface they represent, apportioning the area between
        adjacent MU points equally. Additional weights (such as those
        used in a Gauss-Legendre quadrature) can not meaningfully be
        used in this scheme.  About twice as many points are required
        with this scheme to achieve the precision of Gauss-Legendre
        quadrature.
        DELTAV, VSINI, and VRT must all be in the same units (e.g. km/s).
        If specified, OSAMP should be a positive integer.

        Parameters
        ----------
        mu : array(float) of size (nmu,)
            cosine of the angle between the outward normal and
            the line of sight for each intensity spectrum in INTEN.
        inten : array(float) of size(nmu, npts)
            intensity spectra at specified values of MU.
        deltav : float
            velocity spacing between adjacent spectrum points
            in INTEN (same units as VSINI and VRT).
        vsini : float
            maximum radial velocity, due to solid-body rotation.
        vrt : float
            radial-tangential macroturbulence parameter, i.e.
            np.sqrt(2) times the standard deviation of a Gaussian distribution
            of turbulent velocities. The same distribution function describes
            the radial motions of one component and the tangential motions of
            a second component. Each component covers half the stellar surface.
            See 'The Observation and Analysis of Stellar Photospheres', Gray.
        osamp : int, optional
            internal oversampling factor for convolutions.
            By default convolutions are done using the input points (OSAMP=1),
            but when OSAMP is set to higher integer values, the input spectra
            are first oversampled by cubic spline interpolation.

        Returns
        -------
        value : array(float) of size (npts,)
            Disk integrated flux profile.

        Note
        ------------
            If you use this algorithm in work that you publish, please cite
            Valenti & Anderson 1996, PASP, currently in preparation.
        """
        """
        History
        -----------
        Feb-88  GM
            Created ANA version.
        13-Oct-92 JAV
            Adapted from G. Marcy's ANA routi!= of the same name.
        03-Nov-93 JAV
            Switched to annular convolution technique.
        12-Nov-93 JAV
            Fixed bug. Intensity compo!=nts not added when vsini=0.
        14-Jun-94 JAV
            Reformatted for "public" release. Heavily commented.
            Pass deltav instead of 2.998d5/deltav. Added osamp
            keyword. Added rebinning logic at end of routine.
            Changed default osamp from 3 to 1.
        20-Feb-95 JAV
            Added mu as an argument to handle arbitrary mu sampling
            and remove ambiguity in intensity profile ordering.
            Interpret VTURB as np.sqrt(2)*sigma instead of just sigma.
            Replaced call_external with call to spl_{init|interp}.
        03-Apr-95 JAV
            Multiply flux by pi to give observed flux.
        24-Oct-95 JAV
            Force "nmk" padding to be at least 3 pixels.
        18-Dec-95 JAV
            Renamed from dskint() to rtint(). No longer make local
            copy of intensities. Use radial-tangential instead
            of isotropic Gaussian macroturbulence.
        26-Jan-99 JAV
            For NMU=1 and VSINI=0, assume resolved solar surface#
            apply R-T macro, but supress vsini broadening.
        01-Apr-99 GMH
            Use annuli weights, rather than assuming ==ual area.
        07-Mar-12 JAV
            Force vsini and vmac to be scalars.
        """

        # Make local copies of various input variables, which will be altered below.
        # Force vsini and especially vmac to be scalars. Otherwise mu dependence fails.

        if np.size(vsini) > 1:
            vsini = vsini[0]
        if np.size(vrt) > 1:
            vrt = vrt[0]

        # Determine oversampling factor.
        os = round(np.clip(osamp, 1, None))  # force integral value > 1

        # Convert input MU to projected radii, R, of annuli for a star of unit radius
        #  (which is just sine, rather than cosine, of the angle between the outward
        #  normal and the line of sight).
        rmu = np.sqrt(1 - mu ** 2)  # use simple trig identity

        # Sort the projected radii and corresponding intensity spectra into ascending
        #  order (i.e. from disk center to the limb), which is equivalent to sorting
        #  MU in descending order.
        isort = np.argsort(rmu)
        rmu = rmu[isort]  # reorder projected radii
        nmu = np.size(mu)  # number of radii
        if nmu == 1:
            if vsini != 0:
                logger.warning(
                    "Vsini is non-zero, but only one projected radius (mu value) is set. No rotational broadening will be performed."
                )
                vsini = 0  # ignore vsini if only 1 mu

        # Calculate projected radii for boundaries of disk integration annuli.  The n+1
        # boundaries are selected such that r(i+1) exactly bisects the area between
        # rmu(i) and rmu(i+1). The in!=rmost boundary, r(0) is set to 0 (disk center)
        # and the outermost boundary, r(nmu) is set to 1 (limb).
        if nmu > 1 or vsini != 0:  # really want disk integration
            r = np.sqrt(
                0.5 * (rmu[:-1] ** 2 + rmu[1:] ** 2)
            )  # area midpoints between rmu
            r = np.concatenate(([0], r, [1]))

            # Calculate integration weights for each disk integration annulus.  The weight
            # is just given by the relative area of each annulus, normalized such that
            # the sum of all weights is unity.  Weights for limb darkening are included
            # explicitly in the intensity profiles, so they aren't needed here.
            wt = r[1:] ** 2 - r[:-1] ** 2  # weights = relative areas
        else:
            wt = np.array([1.0])  # single mu value, full weight

        # Generate index vectors for input and oversampled points. Note that the
        # oversampled indicies are carefully chosen such that every "os" finely
        # sampled points fit exactly into one input bin. This makes it simple to
        # "integrate" the finely sampled points at the end of the routine.
        npts = inten.shape[1]  # number of points
        xpix = np.arange(npts, dtype=float)  # point indices
        nfine = os * npts  # number of oversampled points
        xfine = (0.5 / os) * (
            2 * np.arange(nfine, dtype=float) - os + 1
        )  # oversampled points indices

        # Loop through annuli, constructing and convolving with rotation kernels.

        yfine = np.empty(nfine)  # init oversampled intensities
        flux = np.zeros(nfine)  # init flux vector
        for imu in range(nmu):  # loop thru integration annuli

            #  Use external cubic spline routine (adapted from Numerical Recipes) to make
            #  an oversampled version of the intensity profile for the current annulus.
            ypix = inten[isort[imu]]  # extract intensity profile
            if os == 1:
                # just copy (use) original profile
                yfine = ypix
            else:
                # spline onto fine wavelength scale
                yfine = interp1d(xpix, ypix, kind="cubic")(xfine)

            # Construct the convolution kernel which describes the distribution of
            # rotational velocities present in the current annulus. The distribution has
            # been derived analytically for annuli of arbitrary thickness in a rigidly
            # rotating star. The kernel is constructed in two pieces: o!= piece for
            # radial velocities less than the maximum velocity along the inner edge of
            # the annulus, and one piece for velocities greater than this limit.
            if vsini > 0:
                # nontrivial case
                r1 = r[imu]  # inner edge of annulus
                r2 = r[imu + 1]  # outer edge of annulus
                dv = deltav / os  # oversampled velocity spacing
                maxv = vsini * r2  # maximum velocity in annulus
                nrk = 2 * int(maxv / dv) + 3  ## oversampled kernel point
                # velocity scale for kernel
                v = dv * (np.arange(nrk, dtype=float) - ((nrk - 1) / 2))
                rkern = np.zeros(nrk)  # init rotational kernel
                j1 = np.abs(v) < vsini * r1  # low velocity points
                rkern[j1] = np.sqrt((vsini * r2) ** 2 - v[j1] ** 2) - np.sqrt(
                    (vsini * r1) ** 2 - v[j1] ** 2
                )  # generate distribution

                j2 = (np.abs(v) >= vsini * r1) & (np.abs(v) <= vsini * r2)
                rkern[j2] = np.sqrt(
                    (vsini * r2) ** 2 - v[j2] ** 2
                )  # generate distribution

                rkern = rkern / np.sum(rkern)  # normalize kernel

                # Convolve the intensity profile with the rotational velocity kernel for this
                # annulus. Pad each end of the profile with as many points as are in the
                # convolution kernel. This reduces Fourier ringing. The convolution may also
                # be do!= with a routi!= called "externally" from IDL, which efficiently
                # shifts and adds.
                if nrk > 3:
                    yfine = convolve(yfine, rkern, mode="nearest")

            # Calculate projected sigma for radial and tangential velocity distributions.
            muval = mu[isort[imu]]  # current value of mu
            sigma = os * vrt / np.sqrt(2) / deltav  # standard deviation in points
            sigr = sigma * muval  # reduce by current mu value
            sigt = sigma * np.sqrt(1.0 - muval ** 2)  # reduce by np.sqrt(1-mu**2)

            # Figure out how many points to use in macroturbulence kernel.
            nmk = int(10 * sigma)
            nmk = np.clip(nmk, 3, (nfine - 3) // 2)

            # Construct radial macroturbulence kernel with a sigma of mu*VRT/np.sqrt(2).
            if sigr > 0:
                xarg = np.linspace(-nmk, nmk, 2 * nmk + 1) / sigr
                xarg = np.clip(-0.5 * xarg ** 2, -20, None)
                mrkern = np.exp(xarg)  # compute the gaussian
                mrkern = mrkern / np.sum(mrkern)  # normalize the profile
            else:
                mrkern = np.zeros(2 * nmk + 1)  # init with 0d0
                mrkern[nmk] = 1.0  # delta function

            # Construct tangential kernel with a sigma of np.sqrt(1-mu**2)*VRT/np.sqrt(2).
            if sigt > 0:
                xarg = np.linspace(-nmk, nmk, 2 * nmk + 1) / sigt
                xarg = np.clip(-0.5 * xarg ** 2, -20, None)
                mtkern = np.exp(xarg)  # compute the gaussian
                mtkern = mtkern / np.sum(mtkern)  # normalize the profile
            else:
                mtkern = np.zeros(2 * nmk + 1)  # init with 0d0
                mtkern[nmk] = 1.0  # delta function

            # Sum the radial and tangential components, weighted by surface area.
            area_r = 0.5  # assume equal areas
            area_t = 0.5  # ar+at must equal 1
            mkern = area_r * mrkern + area_t * mtkern  # add both components

            # Convolve the total flux profiles, again padding the spectrum on both ends to
            # protect against Fourier ringing.
            yfine = convolve(
                yfine, mkern, mode="nearest"
            )  # add the padding and convolve

            # Add contribution from current annulus to the running total.
            flux = flux + wt[imu] * yfine  # add profile to running total

        flux = np.reshape(flux, (npts, os))  # convert to an array
        flux = np.pi * np.sum(flux, axis=1) / os  # sum, normalize
        return flux

    def sequential_synthesize_segments(
        self,
        sme,
        segments,
        wmod,
        smod,
        cmod,
        reuse_wavelength_grid,
        dll_id=None,
    ):
        for il in tqdm(segments, desc="Segments", leave=False):
            wmod[il], smod[il], cmod[il] = self.synthesize_segment(
                sme,
                il,
                reuse_wavelength_grid,
                il != segments[0],
                dll_id=dll_id,
            )
        return wmod, smod, cmod

    def get_dll_id(self, dll=None):
        if dll is None:
            dll = self.dll
        if dll in __DLL_IDS__:
            dll_id = __DLL_IDS__[dll]
        elif dll in __DLL_DICT__:
            dll_id = dll
        else:
            dll_id = uuid.uuid4()
            __DLL_DICT__[dll_id] = dll
            __DLL_IDS__[dll] = dll_id
        return dll_id

    def get_dll(self, dll_id=None):
        if dll_id is None:
            dll_id = self.dll
        if dll_id in __DLL_DICT__:
            return __DLL_DICT__[dll_id]
        else:
            return dll_id

    def parallel_synthesize_segments(
        self,
        sme,
        segments,
        wmod,
        smod,
        cmod,
        reuse_wavelength_grid,
        dll_id=None,
    ):
        # Make sure the dll is recorded in the global variables
        dll = self.get_dll(dll_id)
        dll_id = self.get_dll_id(dll)

        # We calculate the first segment sequentially
        with tqdm(desc="Segments", total=len(segments), leave=False) as progress:
            il = segments[0]
            wmod[il], smod[il], cmod[il] = self.synthesize_segment(
                sme, il, reuse_wavelength_grid, False, dll_id=dll_id
            )
            progress.update(1)
            # and then all others in parrallel
            # since we can keep the line opacities from the calculation of the first segment
            # TODO: do the line opacities also in parallel?

            # For multiple Processes we need to pickle all the components
            # BUT we can not pickle the smelib, since it has pointers (in the state)
            # Therefore we cheat by putting the library in a global variable
            # but only with a unqiue id, that should be unique to this library

            def parallel(il):
                return self.synthesize_segment(
                    sme,
                    il,
                    reuse_wavelength_grid,
                    True,
                    method="parallel",
                    dll_id=dll_id,
                )

            # Sequential version for debugging
            data = [None for _ in segments[1:]]
            for i, seg in enumerate(segments[1:]):
                data[i] = self.synthesize_segment(
                    sme,
                    seg,
                    reuse_wavelength_grid,
                    True,
                    method="sequential",
                    dll_id=dll_id,
                )
                progress.update(1)

        # data_seq = [None for _ in segments[1:]]
        # What is sticking around in the library that is not part of the state?
        # for seg in segments[1:]:
        #     i = seg-1
        #     data_seq[i] = self.synthesize_segment(
        #         sme,
        #         seg,
        #         reuse_wavelength_grid,
        #         True,
        #         method="sequential",
        #         dll_id=dll_id,
        #     )

        #     if not np.all(data[i][0] == data_seq[i][0]):
        #         print("What")
        #     if not np.all(data[i][1] == data_seq[i][1]):
        #         print("The")
        #     if not np.all(data[i][2] == data_seq[i][2]):
        #         print("Hell")

        # Pathos version crashes for some reason
        # with ThreadPool() as pool:
        #     data = pool.map(parallel, segments[1:])

        # Use "default" ThreadPool instead
        # data = [None for _ in segments[1:]]
        # with ThreadPoolExecutor() as executor:
        #     futures = {executor.submit(parallel, il): il for il in segments[1:]}
        #     for future in as_completed(futures):
        #         il = futures[future] - 1
        #         data[il] = future.result()

        for i, seg in enumerate(segments[1:]):
            wmod[seg] = data[i][0]
            smod[seg] = data[i][1]
            cmod[seg] = data[i][2]

        return wmod, smod, cmod

    def synthesize_spectrum(
        self,
        sme,
        segments="all",
        passLineList=True,
        passAtmosphere=True,
        passNLTE=True,
        updateStructure=True,
        updateLineList=False,
        reuse_wavelength_grid=False,
        radial_velocity_mode="robust",
        method="sequential",
        dll_id=None,
    ):
        """
        Calculate the synthetic spectrum based on the parameters passed in the SME structure
        The wavelength range of each segment is set in sme.wran
        The specific wavelength grid is given by sme.wave, or is generated on the fly if sme.wave is None

        Will try to fit radial velocity RV and continuum to observed spectrum, depending on vrad_flag and cscale_flag

        Other important fields:
        sme.iptype: instrument broadening type

        Parameters
        ----------
        sme : SME_Struct
            sme structure, with all necessary parameters for the calculation
        setLineList : bool, optional
            wether to pass the linelist to the c library (default: True)
        passAtmosphere : bool, optional
            wether to pass the atmosphere to the c library (default: True)
        passNLTE : bool, optional
            wether to pass NLTE departure coefficients to the c library (default: True)
        reuse_wavelength_grid : bool, optional
            wether to use sme.wint as the output grid of the function or create a new grid (default: False)

        Returns
        -------
        sme : SME_Struct
            same sme structure with synthetic spectrum in sme.smod
        """

        if sme is not self.known_sme:
            logger.debug("Synthesize spectrum")
            logger.debug("%s", sme)
            self.known_sme = sme

        # Define constants
        n_segments = sme.nseg
        cscale_degree = sme.cscale_degree

        # fix impossible input
        if "spec" not in sme:
            sme.vrad_flag = "none"
            sme.cscale_flag = "none"
        else:
            if "mask" not in sme:
                sme.mask = np.full(sme.spec.size, sme.mask_values["line"])
            for i in range(sme.nseg):
                mask = ~np.isfinite(sme.spec[i])
                mask |= sme.uncs[i] == 0
                sme.mask[i][mask] = sme.mask_values["bad"]

        if radial_velocity_mode != "robust" and (
            "cscale" not in sme or "vrad" not in sme
        ):
            radial_velocity_mode = "robust"

        segments = self.check_segments(sme, segments)

        # Prepare arrays
        vrad, _, cscale, _ = null_result(sme.nseg, sme.cscale_degree, sme.cscale_type)

        wave = [np.zeros(0) for _ in range(n_segments)]
        smod = [[] for _ in range(n_segments)]
        cmod = [[] for _ in range(n_segments)]
        wmod = [[] for _ in range(n_segments)]

        # If wavelengths are already defined use those as output
        if "wave" in sme:
            wave = [w for w in sme.wave]

        if method == "parallel" and not self.get_dll(dll_id).parallel:
            # display only once
            if (
                not hasattr(self, "_warning_parallel_mode")
                or not self._warning_parallel_mode
            ):
                self._warning_parallel_mode = True
                logger.warning(
                    "Parallel mode was requested, but the library in use is a sequential version. Running in sequential mode instead"
                )
            method = "sequential"

        if method == "parallel":
            dll = self.get_dll(dll_id).copy()
            dll_id = self.get_dll_id(dll)
        else:
            dll = self.get_dll(dll_id)

        # Input Model data to C library
        dll.SetLibraryPath()
        if passLineList:
            dll.InputLineList(sme.linelist)
        if updateLineList:
            # TODO Currently Updates the whole linelist, could be improved to only change affected lines
            dll.UpdateLineList(sme.atomic, sme.species, np.arange(len(sme.linelist)))
        if passAtmosphere:
            sme = self.get_atmosphere(sme)
            dll.InputModel(sme.teff, sme.logg, sme.vmic, sme.atmo)
            dll.InputAbund(sme.abund)
            dll.Ionization(0)
            dll.SetVWscale(sme.gam6)
            dll.SetH2broad(sme.h2broad)
        if passNLTE:
            sme.nlte.update_coefficients(sme, dll, self.lfs_nlte)

        # Loop over segments
        #   Input Wavelength range and Opacity
        #   Calculate spectral synthesis for each
        #   Interpolate onto geomspaced wavelength grid
        #   Apply instrumental and turbulence broadening

        # TODO Parallelization
        # This requires changes in the C code however, since SME uses global parameters
        # for the wavelength range (and opacities) which change within each segment
        if dll.parallel:
            self.parallel_synthesize_segments(
                sme,
                segments,
                wmod,
                smod,
                cmod,
                reuse_wavelength_grid,
                dll_id=dll,
            )
        else:
            self.sequential_synthesize_segments(
                sme,
                segments,
                wmod,
                smod,
                cmod,
                reuse_wavelength_grid,
                dll_id=dll,
            )

        for il in segments:
            if "wave" not in sme or len(sme.wave[il]) == 0:
                # trim padding
                wbeg, wend = sme.wran[il]
                itrim = (wmod[il] > wbeg) & (wmod[il] < wend)
                # Force endpoints == wavelength range
                wave[il] = np.concatenate(([wbeg], wmod[il][itrim], [wend]))

        if sme.specific_intensities_only:
            return wmod, smod, cmod

        # Fit continuum and radial velocity
        # And interpolate the flux onto the wavelength grid
        if radial_velocity_mode == "robust":
            cscale, cscale_unc, vrad, vrad_unc = match_rv_continuum(
                sme, segments, wmod, smod
            )
            logger.debug("Radial velocity: %s", str(vrad))
            logger.debug("Continuum coefficients: %s", str(cscale))
        elif radial_velocity_mode == "fast":
            cscale, vrad = sme.cscale, sme.vrad
        else:
            raise ValueError("Radial Velocity mode not understood")

        smod, cmod = self.apply_radial_velocity_and_continuum(
            wave,
            sme.spec,
            wmod,
            smod,
            cmod,
            vrad,
            cscale,
            sme.cscale_type,
            segments,
        )

        # Merge all segments
        # if sme already has a wavelength this should be the same
        if updateStructure:
            if "wave" not in sme:
                # TODO: what if not all segments are there?
                sme.wave = wave
            if "synth" not in sme:
                sme.synth = smod
            if "cont" not in sme:
                sme.cont = cmod

            for s in segments:
                sme.wave[s] = wave[s]
                sme.synth[s] = smod[s]
                sme.cont[s] = cmod[s]

            if sme.cscale_type in ["spline", "spline+mask"]:
                sme.cscale = cscale
                sme.cscale_unc = cscale_unc
            elif sme.cscale_flag not in ["fix", "none"]:
                for s in segments:
                    sme.cscale[s] = cscale[s]
                sme.cscale_unc = cscale_unc

            sme.vrad = np.asarray(vrad)
            sme.vrad_unc = np.asarray(vrad_unc)
            sme.nlte.flags = dll.GetNLTEflags()
            result = sme
        else:
            wave = Iliffe_vector(values=wave)
            smod = Iliffe_vector(values=smod)
            cmod = Iliffe_vector(values=cmod)
            result = wave, smod, cmod

        # Cleanup
        return result

    def synthesize_segment(
        self,
        sme,
        segment,
        reuse_wavelength_grid=False,
        keep_line_opacity=False,
        method="sequential",
        dll_id=None,
    ):
        """Create the synthetic spectrum of a single segment

        Parameters
        ----------
        sme : SME_Struct
            The SME strcuture containing all relevant parameters
        segment : int
            the segment to synthesize
        reuse_wavelength_grid : bool
            Whether to keep the current wavelength grid for the synthesis
            or create a new one, depending on the linelist. Default: False
        keep_line_opacity : bool
            Whether to reuse existing line opacities or not. This should be
            True iff the opacities have been calculated in another segment.

        Returns
        -------
        wgrid : array of shape (npoints,)
            Wavelength grid of the synthesized spectrum
        flux : array of shape (npoints,)
            The Flux of the synthesized spectrum
        cont_flux : array of shape (npoints,)
            The continuum Flux of the synthesized spectrum
        """
        logger.debug("Segment %i out of %i", segment, sme.nseg)
        if method == "parallel":
            dll = self.get_dll(dll_id).copy()
            dll_id = self.get_dll_id(dll)
        else:
            dll = self.get_dll(dll_id)

        # Input Wavelength range and Opacity
        vrad_seg = sme.vrad[segment] if sme.vrad[segment] is not None else 0
        wbeg, wend = self.get_wavelengthrange(sme.wran[segment], vrad_seg, sme.vsini)

        dll.InputWaveRange(wbeg, wend)
        dll.Opacity()

        # Reuse adaptive wavelength grid in the jacobians
        if reuse_wavelength_grid and segment in self.wint.keys():
            wint_seg = self.wint[segment]
        else:
            wint_seg = None

        # Only calculate line opacities in the first segment
        #   Calculate spectral synthesis for each
        _, wint, sint, cint = dll.Transf(
            sme.mu,
            sme.accrt,  # threshold line opacity / cont opacity
            sme.accwi,
            keep_lineop=keep_line_opacity,
            wave=wint_seg,
        )
        # Store the adaptive wavelength grid for the future
        # if it was newly created
        if wint_seg is None:
            self.wint[segment] = wint

        if not sme.specific_intensities_only:
            # Create new geomspaced wavelength grid, to be used for intermediary steps
            wgrid, vstep = self.new_wavelength_grid(wint)

            logger.debug("Integrate specific intensities")
            # Radiative Transfer Integration
            # Continuum
            cint = self.integrate_flux(sme.mu, cint, 1, 0, 0)
            cint = np.interp(wgrid, wint, cint)

            # Broaden Spectrum
            y_integrated = np.empty((sme.nmu, len(wgrid)))
            for imu in range(sme.nmu):
                y_integrated[imu] = np.interp(wgrid, wint, sint[imu])

            # Turbulence broadening
            # Apply macroturbulent and rotational broadening while integrating intensities
            # over the stellar disk to produce flux spectrum Y.
            sint = self.integrate_flux(sme.mu, y_integrated, vstep, sme.vsini, sme.vmac)
            wint = wgrid

            # instrument broadening
            if "iptype" in sme:
                logger.debug("Apply detector broadening")
                ipres = sme.ipres if np.size(sme.ipres) == 1 else sme.ipres[segment]
                sint = broadening.apply_broadening(
                    ipres, wint, sint, type=sme.iptype, sme=sme
                )

        # Divide calculated spectrum by continuum
        if sme.normalize_by_continuum:
            sint /= cint

        return wint, sint, cint

def synthesize_spectrum(sme, segments="all"):
    synthesizer = Synthesizer()
    return synthesizer.synthesize_spectrum(sme, segments)


# coding: utf-8
import logging

from amuse.units import nbody_system, units, constants

from amuse.couple.bridge import CalculateFieldForCodesUsingReinitialize

from amuse.community.hermite0.interface import Hermite
from amuse.community.ph4.interface import ph4
from amuse.community.rebound.interface import Rebound
from amuse.community.fi.interface import Fi
from amuse.community.gadget2.interface import Gadget2
from amuse.community.fastkick.interface import FastKick
from amuse.community.seba.interface import SeBa
from amuse.community.kepler_orbiters.interface import Kepler
from amuse.community.bonsai.interface import Bonsai

import numpy as np

logger = logging.getLogger(__name__)

def new_code_gravity(
        converter,
        epsilon,
        p,
        ):
    if p.codes_gravity == "ph4":
        gravity = new_code_gravity_ph4(
                converter,
                epsilon,
                p,
                )
    elif p.codes_gravity == "Hermite":
        gravity = new_code_gravity_hermite(
                converter,
                epsilon,
                p,
                )
    elif p.codes_gravity == "Rebound":
        gravity = new_code_gravity_rebound(
                converter,
                epsilon,
                p,
                )
    else:
        logger.error(
                "No gravity code %s known"%(
                    p.codes_gravity,
                    ),
                )
        exit()
    logger.info(
            "Gravity code %s added"%(
                p.codes_gravity,
                ),
            )
    return gravity

def new_code_gravity_rebound(
        converter,
        epsilon,
        p,
        ):
    result  = Rebound(
            converter,
            redirection     = "file",
            redirect_file   = (
                p.dir_codelogs + "/gravity_code.log"
                )
            )
    result.initialize_code()
    print(result.parameters)
    result.parameters.epsilon_squared   = p.particles_epsilon**2
    result.parameters.integrator    = p.codes_gravity_integrator
    result.parameters.timestep      = p.timestep_integrator
    result.parameters.solver        = p.codes_gravity_solver
    result.parameters.boundary      = p.codes_gravity_boundary
    result.parameters.boundary_size = p.codes_gravity_boundary_size
    result.parameters.exact_finish_time = 0
    return result

def new_code_gravity_bonsai(
        converter,
        epsilon,
        p,
        ):
    result  = Bonsai(
            converter,
            redirection     = "file",
            redirect_file   = (
                p.dir_codelogs + "/gravity_code.log"
                )
            )
    result.initialize_code()
    result.parameters.timestep          = p.timestep_integrator
    return result

def new_code_gravity_ph4(
        converter,
        epsilon,
        p,
        ):
    result  = ph4(
            converter,
            redirection     = "file",
            mode            = "gpu" if p.codes_gravity == "gpu" else "cpu",
            redirect_file   = (
                p.dir_codelogs + "/gravity_code.log"
                )
            )
    result.initialize_code()
    result.parameters.epsilon_squared   = epsilon**2
    return result

def new_code_gravity_hermite(
        converter,
        epsilon,
        p,
        ):
    result  = Hermite(
            converter,
            redirection     = "file",
            redirect_file   = (
                p.dir_codelogs + "/gravity_code.log"
                )
            )
    result.parameters.epsilon_squared           = epsilon**2
    result.parameters.end_time_accuracy_factor  = 0
    return result

class kepler_for_bridge(object):
    def __init__(self, converter):
        self.code = Kepler(converter)
        self.model_time = self.code.model_time
        self.central_particle = self.code.central_particle
        self.orbiters = self.code.orbiters

    def evolve_model(self, t_end):
        self.code.evolve_model(t_end)
        self.model_time = self.code.model_time

    def get_gravity_at_point(self, radius, x, y, z):
        mass = self.central_particle.mass
        xc, yc, zc = self.central_particle[0].position
        dr2 = ((x-xc)**2+(y-yc)**2+(z-zc)**2+radius**2)
        dr = dr2**0.5
        ax = -mass*(x-xc)/(dr2*dr)
        ay = -mass*(y-yc)/(dr2*dr)
        az = -mass*(z-zc)/(dr2*dr)

        for body in self.orbiters:
            mass = body.mass
            xc, yc, zc = body.position
            dr2 = ((x-xc)**2 + (y-yc)**2 + (z-zc)**2 + radius**2)
            dr = dr2**0.5
            ax -= mass*(x-xc)/(dr2*dr)
            ay -= mass*(y-yc)/(dr2*dr)
            az -= mass*(z-zc)/(dr2*dr)
        ax *= constants.G
        ay *= constants.G
        az *= constants.G
        return ax,ay,az

    def get_potential_at_point(self, radius, x, y, z):
        mass = self.central_particle.mass
        xc, yc, zc = self.central_particle.position
        dr2 = ((x-xc)**2+(y-yc)**2+(z-zc)**2+radius**2)
        dr = dr2**0.5
        phi = -mass/dr

        for body in self.orbiters:
            mass = body.mass
            xc, yc, zc = body.position
            dr2 = ((x-xc)**2+(y-yc)**2+(z-zc)**2+radius**2)
            dr = dr2**0.5
            phi -= mass/dr
        return phi

class non_central_kepler(object):
    #################
    #### Testing ####
    #################
    def __init__(self, converter):
        self.code = Kepler(converter)
        self.model_time = self.code.model_time
        self.central_particle = self.code.central_particle
        self.orbiters = self.code.orbiters
        self.particles = self.orbiters

    #def add_particles(particles):
    #    code.particles.add_particles(particles)
    #
    #def remove_particles():

    def evolve_model(self, t_end):
        self.pc = self.central_particle.position
        self.vc = self.central_particle.velocity

        self.central_particle.position -= self.pc
        self.orbiters.position -= self.pc
        self.central_particle.velocity -= self.vc
        self.orbiters.velocity -= self.vc

        self.code.evolve_model(t_end)
        self.model_time = self.code.model_time

        self.central_particle.position += self.pc
        self.orbiters.position += self.pc
        self.central_particle.velocity += self.vc
        self.orbiters.velocity += self.vc

    def stop(self):
        self.code.stop()

    def get_gravity_at_point(self, radius, x, y, z):
        self.code.get_gravity_at_point(radius, x, y, z)

    def get_potential_at_point(self, radius, x, y, z):
        self.code.get_potential_at_point(radius, x, y, z)

class advance_without_selfgravity(object):
    """
    Code by Lucie Jilkóva
    to advance particles
    """
    def __init__(self, particles, time= 0 |units.Myr):
        self.particles = particles
        self.model_time = time

    def evolve_model(self, t_end):
        dt = t_end - self.model_time
        self.particles.position += self.particles.velocity*dt
        self.model_time= t_end

    @property
    def potential_energy(self):
        return quantities.zero

    @property
    def kinetic_energy(self):
        return (0.5*self.particles.mass*self.particles.velocity.lengths()**2).sum()

    def get_gravity_at_point(self, radius, x, y, z):
        fr = (0|units.kg) * constants.G / (1 | units.m**2)
        ax = -0*fr*x/(1|units.m)
        ay = -0*fr*x/(1|units.m)
        az = -0*fr*x/(1|units.m)
        return ax,ay,az

class Star(object):
    def __init__(
            self,
            M = 1.0 | units.MSun,
            R = 1.0 | units.RSun,
            ):
        self.R = R
        self.M = M

    def get_gravity_at_point(self, eps, x, y, z):
        r2  = x**2+y**2+z**2
        r   = r2**0.5
        fr  = constants.G*self.M/r2
        ax  = -fr*x/r
        ay  = -fr*y/r
        az  = -fr*z/r
        return ax,ay,az

    def get_potential_at_point(self,eps,x,y,z):
        r2=x**2+y**2+z**2
        r=r2**0.5
        c=constant.G*self.M
        phi=c/(r*r2)
        return phi

    def vcirc(self,r):
        vc=(constants.G*self.M/r)**0.5
        return vc

def hill_radius(
        eccentricity,
        semimajor_axis,
        minor_mass,
        major_mass,
        ):

    return ((1 - eccentricity) * semimajor_axis * (minor_mass/(3*major_mass))**(1./3) )

def cylindrical_from_xyz(
        particles,
        ):
    particles.r       = particles.position.lengths()
    particles.theta   = np.arccos(particles.z/particles.r)
    particles.phi     = np.arctan2(
            particles.y.value_in(units.AU),
            particles.x.value_in(units.AU),
            )
    return

def xyz_from_cylindrical(
        orbiters,
        ):
    x = (orbiters.r * np.cos(orbiters.phi))
    y = (orbiters.r * np.sin(orbiters.phi))
    z = 0*x
    return [x,y,z]

def kepler_orbit(
        orbiters,
        centre,
        a_over_r = 1.,
        ):
    r = orbiters.position.lengths()
    a = a_over_r * r
    mu = constants.G * centre.mass
    vx +=  (np.sin(orbiters.phi)*(mu * (2./orbiters.r - 1./a))**0.5)
    vy += -(np.cos(orbiters.phi)*(mu * (2./orbiters.r - 1./a))**0.5)
    vz += 0|units.kms
    orbiters.initial_orbital_period = 2 * np.pi * (a**3 / mu)**0.5
    return r, [vx,vy,vz]

def eccentricity(
        orbiters,
        centre,
        ):
    length_unit = units.AU
    speed_unit = units.kms
    rmag = (orbiters.position - centre.position).lengths()
    r = (orbiters.position - centre.position)
    rs = (orbiters.position - centre.position).value_in(length_unit)
    v = (orbiters.velocity - centre.velocity)
    vs = (orbiters.velocity - centre.velocity).value_in(speed_unit)
    h = np.cross(rs,vs)
    mu = (constants.G * centre.mass).value_in(length_unit * speed_unit**2)
    tmp = np.cross(vs,h)/mu
    e = np.array([tmp[:,0]-r[:,0]/rmag,tmp[:,1]-r[:,1]/rmag,tmp[:,2]-r[:,2]/rmag])

    emag = (e[0,:]**2+e[1,:]**2+e[2,:]**2)**0.5

    return e, emag

def orbital_periods(
        allorbiters,
        planet,
        ):
    orbiters = allorbiters
    r   = orbiters.position - planet.position
    v   = orbiters.velocity - planet.velocity
    rabs = r.lengths()
    vabs = v.lengths()
    a   = (constants.G * planet[0].mass) * rabs / (
            (2 * constants.G * planet[0].mass) - rabs*vabs*vabs )
    orbital_periods = (2 * np.pi * (a**3 / (constants.G * planet[0].mass))).sqrt()

    return orbital_periods

def kepler_periastron_velocity(
        mass    = 1.0 | units.MSun,
        epsilon = 0.0,
        a       = 1.0 | units.AU,
        ):
    mu      = constants.G * mass
    v_per   = np.sqrt( ((1+epsilon)*mu)/((1-epsilon)*a))
    return v_per

def kepler_apastron_velocity(
        mass    = 1.0 | units.MSun,
        epsilon = 0.0,
        a       = 1.0 | units.AU,
        ):
    mu      = constants.G * mass
    v_apo   = np.sqrt( ((1-epsilon)*mu)/((1+epsilon)*a))
    return v_apo