/* -*- Mode: C; tab-width: 2; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
/* vim:set sw=2 sts=2 et cin: */
/*
 * This file is part of the MUSE Instrument Pipeline
 * Copyright (C) 2008-2014 European Southern Observatory
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 */

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <string.h>

#include "muse_cplwrappers.h"
#include "muse_data_format_z.h"
#include "muse_raman.h"
#include "muse_sky.h"
#include "muse_utils.h"

/* Vibrational temperature for OH transitions */
#define MUSE_SKY_T_VIBR 200.0

/** @addtogroup muse_skysub */
/**@{*/

/*----------------------------------------------------------------------------*/
/**
   @private
   @brief Load the (optional) table with the OH transitions.
   @param aFile   file name
   @return @ref muse_sky_lines_oh_transitions_def "OH transition table"
 */
/*----------------------------------------------------------------------------*/
static cpl_table *
muse_sky_ohtransitions_load(const char *aFile)
{
  if (!cpl_fits_find_extension(aFile, "OH_TRANSITIONS")) {
    return NULL;
  }
  return muse_cpltable_load(aFile, "OH_TRANSITIONS",
                            muse_sky_lines_oh_transitions_def);
}

/*----------------------------------------------------------------------------*/
/**
   @private
   @brief  Recalculate the group indices
   @param  aLines   @ref muse_sky_lines_lines_def "Lines table".
   @retval CPL_ERROR_NONE Success
   @cpl_ensure_code{aLines != NULL, CPL_ERROR_NULL_INPUT}

   After calling this function, the group indices start with zero and go up
   continiously.
 */
/*----------------------------------------------------------------------------*/
static cpl_error_code
muse_sky_lines_reindex_groups(cpl_table *aLines)
{
  cpl_ensure_code(aLines, CPL_ERROR_NULL_INPUT);
  cpl_size n_rows = cpl_table_get_nrow(aLines);
  if (n_rows > 0) {
    int i;
    int i_max = cpl_table_get_column_max(aLines, "group");
    int new_ids[i_max + 1];
    for (i = 0; i <= i_max; i++) {
      new_ids[i] = -1;
    }
    i_max = 0;
    cpl_size i_row;
    for (i_row = 0; i_row < n_rows; i_row++) {
      int old_id = cpl_table_get_int(aLines, "group", i_row, NULL);
      if (new_ids[old_id] < 0) {
        new_ids[old_id] = i_max++;
      }
      cpl_table_set_int(aLines, "group", i_row, new_ids[old_id]);
    }
  }
  return CPL_ERROR_NONE;
}

/*----------------------------------------------------------------------------*/
/**
   @brief  Limit the lines in the table to a wavelength range.
   @param aLines     sky lines table
   @param aLow       lower limit, in Angstrom
   @param aHigh      higher limit, in Angstrom
   @retval CPL_ERROR_NONE Success
   @cpl_ensure_code{aSkyData != NULL, CPL_ERROR_NULL_INPUT}

   This function is used to improve the efficiency of the fits by removing
   lines outside of the data range.
 */
/*----------------------------------------------------------------------------*/
cpl_error_code
muse_sky_lines_set_range(cpl_table *aLines, double aLow, double aHigh)
{
  cpl_ensure_code(aLines, CPL_ERROR_NULL_INPUT);
#pragma omp critical(cpl_table_select)
  cpl_table_unselect_all(aLines);
  cpl_table_or_selected_double(aLines, "lambda", CPL_LESS_THAN, aLow);
  cpl_table_or_selected_double(aLines, "lambda", CPL_GREATER_THAN, aHigh);
  cpl_table_erase_selected(aLines);
  muse_sky_lines_reindex_groups(aLines);

  return CPL_ERROR_NONE;
}

/*----------------------------------------------------------------------------*/
/**
  @brief  Save sky lines table to file.
  @param  aProcessing   the processing structure
  @param  aLines        the sky lines table
  @param  aHeader       the FITS header to use for the primary HDU
  @return CPL_ERROR_NONE on success, another CPL error code on failure

  The table extension is marked with and EXTNAME of "LINES".

  @error{return CPL_ERROR_NULL_INPUT, one of the arguments is NULL}
  @error{return CPL_ERROR_ILLEGAL_INPUT, the output frame could not be created}
 */
/*----------------------------------------------------------------------------*/
cpl_error_code
muse_sky_lines_save(muse_processing *aProcessing, const cpl_table *aLines,
                    cpl_propertylist *aHeader)
{
  cpl_ensure_code(aProcessing && aLines && aHeader, CPL_ERROR_NULL_INPUT);
  cpl_frame *frame = muse_processing_new_frame(aProcessing, -1, aHeader,
                                               MUSE_TAG_SKY_LINES,
                                               CPL_FRAME_TYPE_TABLE);
  cpl_ensure_code(frame, CPL_ERROR_ILLEGAL_INPUT);
  const char *filename = cpl_frame_get_filename(frame);
  cpl_error_code rc = cpl_propertylist_save(aHeader, filename, CPL_IO_CREATE);
  rc = muse_cpltable_append_file(aLines, filename,
                                "LINES", muse_sky_lines_lines_def);
  if (rc == CPL_ERROR_NONE) {
#pragma omp critical(muse_processing_output_frames)
    cpl_frameset_insert(aProcessing->outframes, frame);
  } else {
    cpl_frame_delete(frame);
  }
  return rc;
}

/*----------------------------------------------------------------------------*/
/**
   @brief  Load the sky data files.
   @param  aProcessing   the processing structure
   @return Table with sky lines
   @cpl_ensure{aFile != NULL, CPL_ERROR_NULL_INPUT, NULL}

   The SKY_LINES file can contain up to two extensions containing the table of
   sky line groups ("LINES")  and the OH transition table ("OH_TRANSITIONS").
   The latter is usually missiny, but it it is present it is used to supplement
   the table with theoretically computed OH transition strengths.
 */
/*----------------------------------------------------------------------------*/
cpl_table *
muse_sky_lines_load(muse_processing *aProcessing)
{
  cpl_ensure(aProcessing, CPL_ERROR_NULL_INPUT, NULL);
  cpl_frameset *frames = muse_frameset_find(aProcessing->inframes,
                                            MUSE_TAG_SKY_LINES, 0, CPL_FALSE);
  cpl_errorstate es = cpl_errorstate_get();
  cpl_frame *frame = cpl_frameset_get_position(frames, 0);
  if (!frame) {
    cpl_frameset_delete(frames);
    cpl_errorstate_set(es); /* swallow the illegal input state */
    cpl_msg_warning(__func__, "No sky lines found in input frameset!");
    return NULL;
  }
  const char *fn = cpl_frame_get_filename(frame);
  cpl_table *oh_transitions = muse_sky_ohtransitions_load(fn);
  cpl_table *lines = muse_cpltable_load(fn, "LINES", muse_sky_lines_lines_def);
  if (!lines && !oh_transitions) {
    cpl_msg_warning(__func__, "Could not load sky lines from \"%s\"", fn);
    cpl_frameset_delete(frames);
    return NULL;
  }

  /* Output message, no matter if unit conversions happen           *
   * (so that the following messages can be related to a file name) */
  cpl_msg_info(__func__, "Loaded sky lines from \"%s\"", fn);
  /* Convert tables with unscaled flux into ones with scaled flux. */
  if (lines) {
    cpl_msg_indent_more();
    const char *unit = cpl_table_get_column_unit(lines, "flux");
    if (!unit) {
      cpl_msg_warning(__func__, "No flux unit given!");
    }
    if (unit && strcmp(unit, "erg/(s cm^2 arcsec^2)") == 0) {
      cpl_msg_info(__func__, "Scaling flux by 1e20.");
      cpl_table_multiply_scalar(lines, "flux", 1e20);
      cpl_table_set_column_unit(lines, "flux",
                                "10**(-20)*erg/(s cm^2 arcsec^2)");
    }
    if (unit && strcmp(unit, "10**(-20)*erg/(s cm^2 arcsec^2)") != 0) {
      cpl_msg_warning(__func__, "Unsupported flux unit \"%s\".",
                      cpl_table_get_column_unit(lines, "flux"));
    }
    cpl_msg_indent_less();
  } /* if lines */

  muse_processing_append_used(aProcessing, frame, CPL_FRAME_GROUP_CALIB, 1);
  cpl_frameset_delete(frames);

  /* If we got the basic line list but also the OH   *
   * transitions, compute and add them to the table. */
  cpl_table *sky_lines = muse_sky_lines_create(lines, oh_transitions,
                                               MUSE_SKY_T_VIBR);
  cpl_table_delete(oh_transitions);
  cpl_table_delete(lines);

  return sky_lines;
} /* muse_sky_lines_load() */

/*----------------------------------------------------------------------------*/
/**
   @brief Apply the line strengths to the lines.
   @param aLines      @ref muse_sky_lines_lines_def "List of emission lines".
   @param aStrength   Array of strengths, sorted by line group.
   @retval CPL_ERROR_NONE Everything went OK
   @cpl_ensure_code{aLines != NULL, CPL_ERROR_NULL_INPUT}
   @cpl_ensure_code{aStrength != NULL, CPL_ERROR_NULL_INPUT}
   @cpl_ensure_code{aLines['group'] != NULL, CPL_ERROR_ILLEGAL_INPUT}
   @cpl_ensure_code{aLines['flux'] != NULL, CPL_ERROR_ILLEGAL_INPUT}

   The line strengths are applied in-place, changing the emission lines list.
 */
/*----------------------------------------------------------------------------*/
cpl_error_code 
muse_sky_lines_apply_strength(cpl_table *aLines, const cpl_array *aStrength)
{
  cpl_ensure_code(aLines != NULL, CPL_ERROR_NULL_INPUT);
  cpl_ensure_code(aStrength != NULL, CPL_ERROR_NULL_INPUT);

  int *group = cpl_table_get_data_int(aLines, "group");
  cpl_ensure_code(group != NULL, CPL_ERROR_ILLEGAL_INPUT);
  double *flux = cpl_table_get_data_double(aLines, "flux");
  cpl_ensure_code(flux != NULL, CPL_ERROR_ILLEGAL_INPUT);
  cpl_size nRows = cpl_table_get_nrow(aLines);
  cpl_size i;
  for (i = 0; i < nRows; i++, group++, flux++) {
    *flux *= cpl_array_get(aStrength, *group, NULL);
  }
  return CPL_ERROR_NONE;
}

/*----------------------------------------------------------------------------*/
/**
   @brief Remove all lines below a certain flux limit.
   @param aLines     @ref muse_sky_lines_lines_def "List of emission lines".
   @param aMinflux   Minimal flux, in 10^-20 erg/(s cm^2 arcsec^2)
   @retval CPL_ERROR_NONE Everything went OK
   @cpl_ensure_code{aLines != NULL, CPL_ERROR_NULL_INPUT}

   The line strengths are applied in-place, changing the emission lines list.
 */
/*----------------------------------------------------------------------------*/
cpl_error_code
muse_sky_lines_cut(cpl_table *aLines, double aMinflux)
{
  cpl_ensure_code(aLines != NULL, CPL_ERROR_NULL_INPUT);
  cpl_table_select_all(aLines);
  cpl_table_and_selected_double(aLines, "flux", CPL_LESS_THAN, aMinflux);
  cpl_table_erase_selected(aLines);
  return CPL_ERROR_NONE;
}

/*----------------------------------------------------------------------------*/
/**
   @private
   @brief Get a list of lines from an OH transition table.
   @param aTransitions   @ref muse_sky_lines_oh_transitions_def "OH transition table"
   @param aTemperature   Transitional temperature
   @return @ref muse_sky_lines_lines_def "List of OH lines".
   @error{set CPL_ERROR_UNSUPPORTED_MODE\, return NULL,
          difference of upper transition levels "v_u" in aTransitions is >= 100}

   If the input table is <tt>NULL</tt>, then the function returns an
   empty table (default).

   @pseudocode
group transitions by upper vibrational level
for each group:
  Q = sum((2 * J_u + 1) * exp(-E_u / (k_B * t_vibr)))
  for each transition within the group:
    I = 2 * (2 * J_u + 1) * A * exp(-E_u / (k_B * t_vibr)) / Q@endpseudocode
 */
/*----------------------------------------------------------------------------*/
static cpl_table *
muse_sky_lines_create_oh(const cpl_table *aTransitions, double aTemperature)
{
  if (!aTransitions || (aTransitions && cpl_table_get_nrow(aTransitions) == 0)) {
    return muse_cpltable_new(muse_sky_lines_lines_def, 0);
  }

  const int max_level = cpl_table_get_column_max(aTransitions, "v_u");
  const int min_level = cpl_table_get_column_min(aTransitions, "v_u");
  cpl_ensure(max_level-min_level < 100, CPL_ERROR_UNSUPPORTED_MODE, NULL);
  const int n_transitions = cpl_table_get_nrow(aTransitions);

  cpl_table *lines = muse_cpltable_new(muse_sky_lines_lines_def, n_transitions);

  cpl_table_copy_data_string(lines, "name",
                             cpl_table_get_data_string_const(aTransitions,
                                                             "name"));
  cpl_table_copy_data_double(lines, "lambda",
                             cpl_table_get_data_double_const(aTransitions,
                                                             "lambda"));
  cpl_table_copy_data_int(lines, "group",
                          cpl_table_get_data_int_const(aTransitions, "v_u"));
  cpl_table_subtract_scalar(lines, "group", min_level);

  cpl_table_fill_column_window_double(lines, "flux", 0, n_transitions, 0.0);

  const double k_B = 1.3806504e-23; /* Boltzmann constant [eV/K] */
  const double hc = 1.98623e-8;   /* h * c [Angstrom*erg] */
  cpl_table_copy_data_double(lines, "flux",
                             cpl_table_get_data_double_const(aTransitions, "E_u"));
  cpl_table_divide_scalar(lines, "flux", -k_B * aTemperature);
  cpl_table_exponential_column(lines, "flux", CPL_MATH_E);

  cpl_table_duplicate_column(lines, "J_u", aTransitions, "J_u");
  cpl_table_multiply_scalar(lines, "J_u", 2.0);
  cpl_table_add_scalar(lines, "J_u", 1.0);
  cpl_table_multiply_columns(lines, "flux", "J_u");
  cpl_table_erase_column(lines, "J_u");

  cpl_table_fill_column_window_int(lines, "dq", 0, n_transitions, 0);

  double *flux = cpl_table_get_data_double(lines, "flux");
  const int *group = cpl_table_get_data_int_const(lines, "group");

  double q[max_level - min_level + 1];
  memset(q, 0, (max_level - min_level + 1)*sizeof(double));
  int i_transition;
  for (i_transition = 0; i_transition < n_transitions; i_transition++,
         group++, flux++) {
    q[*group] += *flux;
  }

  cpl_table_duplicate_column(lines, "A", aTransitions, "A");
  cpl_table_multiply_columns(lines, "flux", "A");
  cpl_table_erase_column(lines, "A");
  cpl_table_multiply_scalar(lines, "flux", 2.0 * hc);
  cpl_table_divide_columns(lines, "flux", "lambda");

  flux = cpl_table_get_data_double(lines, "flux");
  group = cpl_table_get_data_int_const(lines, "group");

  for (i_transition = 0; i_transition < n_transitions; i_transition++,
         flux++, group++) {
    *flux /= q[*group];
  }

  cpl_table_multiply_scalar(lines, "flux", 1e20); /* use 10^-20 scaling */
  /* ad-hoc correction to get flux in the correct range */
  cpl_table_divide_scalar(lines, "flux", 50);

  muse_sky_lines_reindex_groups(lines);
  return lines;
} /* muse_sky_lines_create_oh() */

/*----------------------------------------------------------------------------*/
/**
   @brief Create the emission lines from the OH transitions and other lines.
   @param aLines            non-OH lines table
   @param aOh_transitions   OH transition table
   @param aT_vibr           temperature for vibrational transitions
   @return the sky emission line table

   If the input parameters are <tt>NULL</tt>, then the function returns an
   unmodified table (default).

   @algorithm Create emission lines from OH transitions, append the additional
              emission lines, apply line strengths by group
 */
/*----------------------------------------------------------------------------*/
cpl_table *
muse_sky_lines_create(const cpl_table *aLines,
                      const cpl_table *aOh_transitions, double aT_vibr)
{
  int group_start = (aLines && cpl_table_get_nrow(aLines) > 0)
                  ? cpl_table_get_column_max(aLines, "group") + 1 : 0;
  cpl_table *res = muse_sky_lines_create_oh(aOh_transitions, aT_vibr);
  cpl_errorstate prestate = cpl_errorstate_get();
  if (aLines) {
    cpl_table_add_scalar(res, "group", group_start);
    cpl_table_insert(res, aLines, 0);
  }
  if (!cpl_errorstate_is_equal(prestate)) {
    cpl_msg_error(__func__, "while cpl_table_insert(): %s, %s",
                  cpl_table_get_column_unit(res, "flux"), 
                  cpl_table_get_column_unit(aLines, "flux"));
    cpl_errorstate_dump(prestate, CPL_FALSE, NULL);
    cpl_errorstate_set(prestate);
  }
  return res;
} /* muse_sky_lines_create() */

/**@}*/
