# modsource.py -- Model the emission of a source at the LTE
# 
# This file is part of Weeds.

from numpy import *
import sys
import pyclass
from linedb.consts import *
import linedb
from linedb.db import blankPartfunc
from sicparse import OptionParser
import copy

class component:
    pass

def readmdl(filename):
    """
    Read the source model

    Arguments
    filename -- name of the source model (".mdl")
    
    """

    components = []
    linenumber = 0
    foreground_warning_flag = False

    try:
        f = open(filename)
    except:
        raise Exception("Can't open %s" % filename)

    # Set a parser for the optional fields
    parser = OptionParser()
    parser.add_option("-c", "--cache", action="store_false", dest="online",
                      default=True) # unused
    parser.add_option("-a", "--absorption", action="store_true", dest="absorption",
                      default=False)
    parser.add_option("-p", "--partfunc", dest = "partfunc", nargs = 1,
                      type = "float", default = None)

    try:
        for line in f.readlines():
            linenumber = linenumber + 1
            if len(line.strip()) == 0:
                continue
            if line[0] in ["#", "!"]:
                continue
            c = component()

            # Get the species name
            if line[0] == '"':
                c.species = line.split('"')[1]
                line = line.split('"', 3)[2]
            else:
                c.species = line.split(None, 1)[0]
                line = line.split(None, 1)[1]

            # Get other mandatory fields
            field = line.split(None)
            c.Ntot = float(field[0])
            c.Tex = float(field[1])
            c.theta = float(field[2])
            c.v_off = float(field[3])
            c.delta_v = float(field[4])
 
            # Get optional fields
            if len(field) > 5:
                (opts, args) = parser.parse_args(field[5:])
            else:
                (opts, args) = parser.parse_args([]) # to set default values
            if len(args) == 1:
                c.origin = args[0]
            else:
                c.origin = 'All'
            if len(args) > 1:
                raise ValueError
            if not(opts.online):
                pyclass.message(pyclass.seve.w, "MODSOURCE", "/CACHE option is obsolete (ignored).")
            c.absorption = opts.absorption
            if c.absorption and not(foreground_warning_flag):
                pyclass.message(pyclass.seve.i, "MODSOURCE", "Foreground layers are assumed to have a filling factor of 1."
                                " The source sizes for these layers are ignored.")
                foreground_warning_flag = True
            c.partfunc = opts.partfunc # log10 of the partition function
            c.keep_opacity = False

            components.append(c)

    except Exception:
        raise Exception("Incorrect input on line %i of %s" % (linenumber, filename))

    return components

def J(T, freq):
    """
    Returns the radiation temperature

    This function returns the radiation temperature corresponding to
    the given kinetic temperature at the given frequency. See Eq. 1.28
    in the Tools of Radio Astronony by Rohlfs & Wilson.

    Arguments:
    T    -- kinetic temperature, in K
    freq -- frequency, in MHz

    """

    if T == 0.:
        J = zeros(len(freq))
    else:
        J = planck_constant * freq * 1e6 / boltzmann_constant \
            / (exp((planck_constant * freq * 1e6) \
                   / (boltzmann_constant * T)) - 1)

    return J

def modsource(components, antsize, tcont, removecont, verbose):
    """
    Model the emission of a given source at the ETL

    """

    tcmb = 2.73

    # Import Gildas variables if not yet done

    if (not pyclass.gotgdict()):
        pyclass.get(verbose=False)

    if (not pyclass.ingdict('rx')):
        pyclass.message(pyclass.seve.e, "MODSOURCE", "No R spectrum in memory")
        pyclass.sicerror()
        return

    frequency = pyclass.gdict.frequency
    reference = pyclass.gdict.reference
    channels = pyclass.gdict.channels
    freq_step = pyclass.gdict.freq_step
    velocity = pyclass.gdict.velocity

    units = pyclass.gdict.set.las.unit   # Upper and lower axis units
    if units[0] != 'F'.encode():
        pyclass.message(pyclass.seve.e, "MODSOURCE", "Lower x-axis unit must be FREQUENCY")
        pyclass.sicerror()
        return
    (fmin,fmax) = pyclass.px_minmax("F")

    # Save RY and keep memory of the frequency axis of the observation

    pyclass.comm("memorize obs")
    frequency_obs = frequency.copy()
    reference_obs = reference.copy()
    channels_obs = channels.copy()
    freq_step_obs = freq_step.copy() # in the observatory frame
    freq_step_obs /= (1. + pyclass.gdict.r.head.spe.doppler) # in the rest frame

    # Determine the frequency sampling for the model. We make sure
    # that the sampling is ten times smaller than the smallest
    # linewidth, and smaller (or equal) than the resolution of the
    # observed spectra (depending on which one is the smallest).

    delta_v = []
    for c in components:
        delta_v.append(c.delta_v)
    min_delta_v = min(abs(array(delta_v)))
    min_delta_f = min_delta_v * 1e3 / speed_of_light * fmin # MHz
    if (min_delta_f / 10.) < abs(freq_step):
        model_freq_step = min_delta_f / 10.
    else:
        model_freq_step = abs(freq_step)

    # Define the rest frequency and velocity axis for the model. The
    # model has the same range than the current plot.

    freq = arange(fmin, fmax, model_freq_step)  # Numpy array
    vfunc = vectorize(lambda f: pyclass.rx_val(f,'F','V'))  # Numpy-Vectorized rx_val
    velo = vfunc(freq)  # Numpy array
    reference_channel = (frequency - freq[0]) / (freq[1] - freq[0]) + 1

    # Compute the telescope HPBW (which depends on the frequency) from
    # the antenna size. This simply assumes that the telescope is
    # diffraction limited, which should be fine for our purpose.

    theta_tel = 1.22 * speed_of_light / (freq * 1e6) / antsize \
            * 180 / pi * 3600. # arcsecs

    # Compute the antenna temperature the entire frequency range.
    # See the WEEDS documentation manual for the formula used.

    if size(tcont) == 1:  # numpy.size(), works also on standard Python scalars
      tcont = ones(len(freq)) * tcont
    elif size(tcont) == size(pyclass.gdict.ry):
      # User provides a continuum array (same size as RY): resample to our
      # internal X axis ('freq' array)
      # NB: this modifies the current spectrum intensities and its axis
      pyclass.gdict.ry = tcont
      pyclass.comm("resample %i %.16g %.16g %.16g frequency" % (len(freq),reference_channel,frequency,model_freq_step))
      tcont = pyclass.gdict.ry.copy()
    else:
      pyclass.message(pyclass.seve.e, "MODSOURCE", "/TCONT argument must be scalar or same size as RY (got %i instead of %i)" % (size(tcont),size(pyclass.gdict.ry)))
      pyclass.sicerror()
      return
    tb_grand_tot = tcont.copy() # All components, species and lines
    keep_opacity_flag = False
    main_group = True

    for c in components:

        # Get the frequencies, Einstein coefficients, statistical
        # weights and upper level energies for the lines in the
        # frequency range.

        lines = []
        try:
            lines = linedb.select("MODSOURCE", fmin, fmax, c.species, c.origin, 'All', -1, -1)
        except Exception as error:
            pyclass.message(pyclass.seve.e, "MODSOURCE", "search: {0}".format(error))
            pyclass.sicerror()
            return

        if len(lines) == 0:
            pyclass.message(pyclass.seve.w, "MODSOURCE", "No %s lines found in the frequency range" % c.species)
            continue
        else:
            pyclass.message(pyclass.seve.i, "MODSOURCE", "%i %s lines found in the frequency range" % (len(lines), c.species))

        # Get the partition function

        # If the database is in the voparis format, we must pass the
        # url of the partition function. Otherwise, we simply pass the
        # species name.

        # Compute the partition_function for each origin 

        # FixMe (SM): I don't understand this part of the code. It
        # looks like the partition function is kept in a dictionary,
        # perhaps to avoid fetching it multiple times.  But what is
        # the l.prev key? And why getting the partition function for
        # each line? All lines have the same partition function, since
        # they belong to the same species... Unless the database
        # contains lines of the same species from different origins?

        QOri = {}
        for l in lines:
            if not l.prev:
                l.prev = l.dbsource
            if not((l.species, l.origin, l.prev)) in QOri:
                if c.partfunc:
                    Q = 10**c.partfunc
                else:
                    try:
                        # If the origin is not set in the modsource file, the line
                        # can come from any of the currently opened databases... but
                        # which one?
                        modsourcedb = None
                        for db in linedb.dbin.databases:
                            if l.dbsource == db.name:
                                modsourcedb = db
                                break
                        if modsourcedb is None:
                            raise Exception("Database %s is not opened" % (l.origin))
                        # We do not want .cat files as input for modsource because JP wants to enforce
                        # .db as the format to work with, and .cat only as a possible exchange format.
                        if isinstance(modsourcedb, linedb.jplmod.Catfile):
                            pyclass.message(pyclass.seve.e, "MODSOURCE", "modsource not allowed on .cat files")
                            raise Exception
                        #
                        if modsourcedb.online and modsourcedb.protocol == "slap":
                            partfunc_url = lines[0].partition_function_link
                            if partfunc_url == blankPartfunc:
                                raise Exception("Partition function not found for {} in {}".format(l.species, l.dbsource))
                            Q = modsourcedb.partition_function(c.Tex, l.origin, l.prev, url = partfunc_url)
                        else:
                            Q = modsourcedb.partition_function(c.Tex, l.origin, l.prev, species = c.species)
                    except Exception as ex:
                        pyclass.message(pyclass.seve.e, "MODSOURCE", "partfunc: {0}".format(ex))
                        pyclass.sicerror()
                        return
                QOri[(l.species, l.origin, l.prev)] = Q
                if verbose:
                    if not(c.partfunc):
                        pyclass.message(pyclass.seve.i, "MODSOURCE", "log10 of the partition function at %.1f K from %s is %.4f" % 
                                        (c.Tex, l.dbsource, log10(Q)))
                    else:
                        pyclass.message(pyclass.seve.i, "MODSOURCE", "log10 of the partition function at %.1f K from mdlfile is %.4f" % 
                                        (c.Tex, log10(Q)))

        # Compute the total opacity for that species

        tau_tot = zeros(len(freq))
        for l in lines:

            # Line profile function
            freq_off = -c.v_off * 1e3 / speed_of_light * l.frequency # MHz
            sigma = l.frequency / (speed_of_light * sqrt(8 * log(2))) \
                * c.delta_v * 1e3 * 1e6 # Hz
            phi = 1 / (sigma * sqrt(2 * pi)) * exp (-((freq - l.frequency - freq_off) \
                                                          * 1e6)**2 / (2 * sigma**2))

            # Line opacity
            tau = speed_of_light**2 / (8 * pi * (freq * 1e6)**2) * l.einstein_coefficient \
                * c.Ntot * 1e4 * l.upper_level.statistical_weight \
                * exp(-l.upper_level.energy / c.Tex) \
                / QOri[(l.species, l.origin, l.prev)] * (exp(planck_constant * l.frequency * 1e6 \
                                                                 / (c.Tex * boltzmann_constant))-1) * phi

            # Opacity at line center
            l.tau0 = max(tau)
                
            tau_tot = tau_tot + tau

        if c.keep_opacity:
            tau_kept = tau_tot
            keep_opacity_flag = True

        # Compute the antenna temperature for that species

        if main_group and c.absorption:
            main_group = False # first foreground component
        if main_group:
            eta_source = c.theta**2 / (theta_tel**2 + c.theta**2)
            tb_grand_tot +=  eta_source * (J(c.Tex, freq) - J(tcmb, freq) - tcont) * (1 - exp(-tau_tot))
        else:
            eta_source = 1. # foreground layers are assumed to have a filling factor of 1
            tb_grand_tot = tb_grand_tot * exp(-tau_tot) + eta_source * (J(c.Tex, freq) - J(tcmb, freq)) * (1 - exp(-tau_tot))

        if verbose:
            linedb.llist(lines,extra=[["tau0","Tau","%8.2e"]])

    # Remove the continuum if requested
    if removecont:
        tb_grand_tot -= tcont

    # Make a CLASS spectrum with the modeled radiation temperature

    pyclass.comm("define double tb_model[%i]" % len(tb_grand_tot))
    pyclass.comm("define double velo_model[%i]" % len(velo))
    pyclass.gdict.tb_model = tb_grand_tot
    pyclass.gdict.velo_model = velo
    pyclass.comm("model tb_model velo_model /reg %.16g %.16g %.16g" %
                             (reference_channel, velocity, frequency_obs))
    pyclass.comm("resample %i %.16g %.16g %.16g frequency" %
                             (channels_obs, reference_obs, frequency_obs, freq_step_obs))
    pyclass.comm("memorize tb_model")
    pyclass.comm("retrieve obs")
    pyclass.comm("delete /var tb_model")
    
    # Make a CLASS spectrum with the modeled opacity, if requested

    if keep_opacity_flag:
        pyclass.comm("define double tau_model[%i]" % len(tb_grand_tot))
        pyclass.gdict.tau_model = tau_kept
        pyclass.comm("model tau_model velo_model /reg %.16g %.16g %.16g" %
                     (reference_channel, velocity, frequency_obs))
        pyclass.comm("resample %i %.16g %.16g %.16g frequency" %
                     (channels_obs, reference_obs, frequency_obs, freq_step_obs))
        pyclass.comm("memorize tau_model")
        pyclass.comm("retrieve obs")
        pyclass.comm("delete /var tau_model")
    pyclass.comm("delete /var velo_model")

    pyclass.message(pyclass.seve.i, "MODSOURCE", "Model has been stored in memory")

def main():
    """
    Main program for MODSOURCE

    """

    # Get arguments and options from the command line

    parser = OptionParser()
    parser.add_option("-b", "--background", dest = "background", nargs = 1,
                                      type = "float", default = None) # deprecated
    parser.add_option("-t", "--tcont", dest = "tcont", nargs = 1, type =
                      "float", default = 0.)
    parser.add_option("-v", "--verbose", action="store_true", dest="verbose",
                                      default=False)
    parser.add_option("-o", "--opacity", dest="keep_opacity", nargs = 1,
                      type = "int")
    parser.add_option("-k", "--keepcont", action="store_false", dest="remove_cont",
                    default=True)

    try:
        (options, args) = parser.parse_args()
    except:
        pyclass.message(pyclass.seve.e, "MODSOURCE", "Invalid option")
        pyclass.sicerror()
        return
    if options.background != None:
        pyclass.message(pyclass.seve.e, "MODSOURCE", "/BACKGROUND option is deprecated, please use /TCONT instead.")
        pyclass.sicerror()
        return        

    if len(args) == 2:
        mdlfile = pyclass.parsefile(args[0])
        antsize = args[1]
    else:
        pyclass.message(pyclass.seve.e, "MODSOURCE", "Incorrect number of arguments")
        pyclass.sicerror()
        return
    try:
        antsize = float(antsize)
    except:
        pyclass.message(pyclass.seve.e, "MODSOURCE", "Incorrect value for ANTSIZE")
        pyclass.sicerror()
        return

    # Read the source model

    try:
        components = readmdl(mdlfile)
    except Exception as ex:
        pyclass.message(pyclass.seve.e, "MODSOURCE readmdl", "{0}".format(ex))
        pyclass.sicerror()
        return

    # Order components by type (first emission then absorption)

    components_sorted = []
    for c in components:
        if not(c.absorption):
            components_sorted.append(c)
    for c in components:
        if c.absorption:
            components_sorted.append(c)
    components = components_sorted

    if options.keep_opacity != None:
        if options.keep_opacity >= 1 and options.keep_opacity <= len(components):
            components[options.keep_opacity - 1].keep_opacity = True
        else:
            pyclass.message(pyclass.seve.e, "MODSOURCE", "Invalid component number for /OPACITY")
            pyclass.sicerror()
            return

    # Compute the model
    
    modsource(components, antsize, options.tcont, options.remove_cont, options.verbose)

if __name__ == "__main__":
    main()
