/* PSPP - a program for statistical analysis.
   Copyright (C) 2017, 2019 Free Software Foundation, Inc.

   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 3 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, see <http://www.gnu.org/licenses/>. */

#include <config.h>

#include "matrix-reader.h"

#include <stdbool.h>
#include <math.h>

#include "data/casegrouper.h"
#include "data/casereader.h"
#include "data/casewriter.h"
#include "data/data-out.h"
#include "data/dataset.h"
#include "data/dictionary.h"
#include "data/format.h"
#include "data/variable.h"
#include "language/command.h"
#include "language/lexer/lexer.h"
#include "libpspp/i18n.h"
#include "libpspp/message.h"
#include "libpspp/str.h"
#include "output/pivot-table.h"

#include "gettext.h"
#define _(msgid) gettext (msgid)
#define N_(msgid) msgid

struct lexer;

/*
This module interprets a "data matrix", typically generated by the command
MATRIX DATA.  The dictionary of such a matrix takes the form:

 s_0, s_1, ... s_m, ROWTYPE_, VARNAME_, v_0, v_1, .... v_n

where s_0, s_1 ... s_m are the variables defining the splits, and
v_0, v_1 ... v_n are the continuous variables.

m >= 0; n >= 0

The ROWTYPE_ variable is of type A8.
The VARNAME_ variable is a string type whose width is not predetermined.
The variables s_x are of type F4.0 (although this reader accepts any type),
and v_x are of any numeric type.

The values of the ROWTYPE_ variable are in the set {MEAN, STDDEV, N, CORR, COV}
and determine the purpose of that case.
The values of the VARNAME_ variable must correspond to the names of the varibles
in {v_0, v_1 ... v_n} and indicate the rows of the correlation or covariance
matrices.



A typical example is as follows:

s_0 ROWTYPE_   VARNAME_   v_0         v_1         v_2

0   MEAN                5.0000       4.0000       3.0000
0   STDDEV              1.0000       2.0000       3.0000
0   N                   9.0000       9.0000       9.0000
0   CORR       V1       1.0000        .6000        .7000
0   CORR       V2        .6000       1.0000        .8000
0   CORR       V3        .7000        .8000       1.0000
1   MEAN                9.0000       8.0000       7.0000
1   STDDEV              5.0000       6.0000       7.0000
1   N                   9.0000       9.0000       9.0000
1   CORR       V1       1.0000        .4000        .3000
1   CORR       V2        .4000       1.0000        .2000
1   CORR       V3        .3000        .2000       1.0000

*/

void
matrix_material_uninit (struct matrix_material *mm)
{
  gsl_matrix_free (mm->corr);
  gsl_matrix_free (mm->cov);
  gsl_matrix_free (mm->n);
  gsl_matrix_free (mm->mean_matrix);
  gsl_matrix_free (mm->var_matrix);
}

static const struct variable *
find_matrix_string_var (const struct dictionary *dict, const char *name)
{
  const struct variable *var = dict_lookup_var (dict, name);
  if (var == NULL)
    {
      msg (SE, _("Matrix dataset lacks a variable called %s."), name);
      return NULL;
    }
  if (!var_is_alpha (var))
    {
      msg (SE, _("Matrix dataset variable %s should be of string type."), name);
      return NULL;
    }
  return var;
}

struct matrix_reader *
matrix_reader_create (const struct dictionary *dict,
                      struct casereader *in_reader)
{
  const struct variable *varname = find_matrix_string_var (dict, "VARNAME_");
  const struct variable *rowtype = find_matrix_string_var (dict, "ROWTYPE_");
  if (!varname || !rowtype)
    return NULL;

  size_t varname_idx = var_get_dict_index (varname);
  size_t rowtype_idx = var_get_dict_index (rowtype);
  if (varname_idx < rowtype_idx)
    {
      msg (SE, _("Variable %s must precede %s in matrix file dictionary."),
           "ROWTYPE_", "VARNAME_");
      return NULL;
    }

  for (size_t i = 0; i < dict_get_n_vars (dict); i++)
    {
      const struct variable *v = dict_get_var (dict, i);
      if (!var_is_numeric (v) && v != rowtype && v != varname)
        {
          msg (SE, _("Matrix dataset variable %s should be numeric."),
               var_get_name (v));
          return NULL;
        }
    }

  size_t n_vars;
  const struct variable **vars;
  dict_get_vars (dict, &vars, &n_vars, DC_SCRATCH);

  /* Different kinds of variables. */
  size_t first_svar = 0;
  size_t n_svars = rowtype_idx;
  size_t first_fvar = rowtype_idx + 1;
  size_t n_fvars = varname_idx - rowtype_idx - 1;
  size_t first_cvar = varname_idx + 1;
  size_t n_cvars = n_vars - varname_idx - 1;
  if (!n_cvars)
    {
      msg (SE, _("Matrix dataset does not have any continuous variables."));
      free (vars);
      return NULL;
    }

  struct matrix_reader *mr = xmalloc (sizeof *mr);
  *mr = (struct matrix_reader) {
    .dict = dict,
    .grouper = casegrouper_create_vars (in_reader, &vars[first_svar], n_svars),
    .svars = xmemdup (vars + first_svar, n_svars * sizeof *mr->svars),
    .n_svars = n_svars,
    .rowtype = rowtype,
    .fvars = xmemdup (vars + first_fvar, n_fvars * sizeof *mr->fvars),
    .n_fvars = n_fvars,
    .varname = varname,
    .cvars = xmemdup (vars + first_cvar, n_cvars * sizeof *mr->cvars),
    .n_cvars = n_cvars,
  };
  free (vars);

  return mr;
}

bool
matrix_reader_destroy (struct matrix_reader *mr)
{
  if (mr == NULL)
    return false;
  bool ret = casegrouper_destroy (mr->grouper);
  free (mr->svars);
  free (mr->cvars);
  free (mr->fvars);
  free (mr);
  return ret;
}


/*
   Allocates MATRIX if necessary,
   and populates row MROW, from the data in C corresponding to
   variables in VARS. N_VARS is the length of VARS.
*/
static void
matrix_fill_row (gsl_matrix **matrix,
      const struct ccase *c, int mrow,
      const struct variable **vars, size_t n_vars)
{
  int col;
  if (*matrix == NULL)
    {
      *matrix = gsl_matrix_alloc (n_vars, n_vars);
      gsl_matrix_set_all (*matrix, SYSMIS);
    }

  for (col = 0; col < n_vars; ++col)
    {
      const struct variable *cv = vars [col];
      double x = case_num (c, cv);
      assert (col  < (*matrix)->size2);
      assert (mrow < (*matrix)->size1);
      gsl_matrix_set (*matrix, mrow, col, x);
    }
}

static int
find_varname (const struct variable **vars, int n_vars,
              const char *varname)
{
  for (int i = 0; i < n_vars; i++)
    if (!strcasecmp (var_get_name (vars[i]), varname))
      return i;
  return -1;
}

struct substring
matrix_reader_get_string (const struct ccase *c, const struct variable *var)
{
  struct substring s = case_ss (c, var);
  ss_rtrim (&s, ss_cstr (CC_SPACES));
  return s;
}

void
matrix_reader_set_string (struct ccase *c, const struct variable *var,
                          struct substring src)
{
  struct substring dst = case_ss (c, var);
  for (size_t i = 0; i < dst.length; i++)
    dst.string[i] = i < src.length ? src.string[i] : ' ';
}

bool
matrix_reader_next (struct matrix_material *mm, struct matrix_reader *mr,
                    struct casereader **groupp)
{
  struct casereader *group;
  if (!casegrouper_get_next_group (mr->grouper, &group))
    {
      *mm = (struct matrix_material) MATRIX_MATERIAL_INIT;
      if (groupp)
        *groupp = NULL;
      return false;
    }

  if (groupp)
    *groupp = casereader_clone (group);

  const struct variable **vars = mr->cvars;
  size_t n_vars = mr->n_cvars;

  *mm = (struct matrix_material) { .n = NULL };

  struct matrix
    {
      const char *name;
      gsl_matrix **m;
      size_t good_rows;
      size_t bad_rows;
    };
  struct matrix matrices[] = {
    { .name = "CORR", .m = &mm->corr },
    { .name = "COV", .m = &mm->cov },
  };
  enum { N_MATRICES = 2 };

  struct ccase *c;
  for (; (c = casereader_read (group)); case_unref (c))
    {
      struct substring rowtype = matrix_reader_get_string (c, mr->rowtype);

      gsl_matrix **v
        = (ss_equals_case (rowtype, ss_cstr ("N")) ? &mm->n
           : ss_equals_case (rowtype, ss_cstr ("MEAN")) ? &mm->mean_matrix
           : ss_equals_case (rowtype, ss_cstr ("STDDEV")) ? &mm->var_matrix
           : NULL);
      if (v)
        {
          if (!*v)
            *v = gsl_matrix_calloc (n_vars, n_vars);

          for (int x = 0; x < n_vars; ++x)
            {
              double n = case_num (c, vars[x]);
              if (v == &mm->var_matrix)
                n *= n;
              for (int y = 0; y < n_vars; ++y)
                gsl_matrix_set (*v, y, x, n);
            }
          continue;
        }

      struct matrix *m = NULL;
      for (size_t i = 0; i < N_MATRICES; i++)
        if (ss_equals_case (rowtype, ss_cstr (matrices[i].name)))
          {
            m = &matrices[i];
            break;
          }
      if (m)
        {
          struct substring varname_raw = case_ss (c, mr->varname);
          struct substring varname = ss_cstr (
            recode_string (UTF8, dict_get_encoding (mr->dict),
                           varname_raw.string, varname_raw.length));
          ss_rtrim (&varname, ss_cstr (CC_SPACES));
          varname.string[varname.length] = '\0';

          int y = find_varname (vars, n_vars, varname.string);
          if (y >= 0)
            {
              m->good_rows++;
              matrix_fill_row (m->m, c, y, vars, n_vars);
            }
          else
            m->bad_rows++;
          ss_dealloc (&varname);
        }
    }
  casereader_destroy (group);

  for (size_t i = 0; i < N_MATRICES; i++)
    if (matrices[i].good_rows && matrices[i].good_rows != n_vars)
      msg (SW, _("%s matrix has %zu columns but %zu rows named variables "
                 "to be analyzed (and %zu rows named unknown variables)."),
           matrices[i].name, n_vars, matrices[i].good_rows,
           matrices[i].bad_rows);

  return true;
}

int
cmd_debug_matrix_read (struct lexer *lexer, struct dataset *ds)
{
  if (lex_match_id (lexer, "NODATA"))
    {
      struct casereader *cr = casewriter_make_reader (
        mem_writer_create (dict_get_proto (dataset_dict (ds))));
      struct matrix_reader *mr = matrix_reader_create (dataset_dict (ds), cr);
      if (!mr)
        {
          casereader_destroy (cr);
          return CMD_FAILURE;
        }
      matrix_reader_destroy (mr);
      return CMD_SUCCESS;
    }

  struct matrix_reader *mr = matrix_reader_create (dataset_dict (ds),
                                                   proc_open (ds));
  if (!mr)
    return CMD_FAILURE;

  struct pivot_table *pt = pivot_table_create ("Debug Matrix Reader");

  enum mm_stat
    {
      MM_CORR,
      MM_COV,
      MM_N,
      MM_MEAN,
      MM_STDDEV,
    };
  const char *mm_stat_names[] = {
    [MM_CORR] = "Correlation",
    [MM_COV] = "Covariance",
    [MM_N] = "N",
    [MM_MEAN] = "Mean",
    [MM_STDDEV] = "Standard Deviation",
  };
  enum { N_STATS = sizeof mm_stat_names / sizeof *mm_stat_names };
  for (size_t i = 0; i < 2; i++)
    {
      struct pivot_dimension *d = pivot_dimension_create (
        pt,
        i ? PIVOT_AXIS_COLUMN : PIVOT_AXIS_ROW,
        i ? "Column" : "Row");
      if (!i)
        pivot_category_create_leaf_rc (d->root, pivot_value_new_text ("Value"),
                                       PIVOT_RC_CORRELATION);
      for (size_t j = 0; j < mr->n_cvars; j++)
        pivot_category_create_leaf_rc (
          d->root, pivot_value_new_variable (mr->cvars[j]),
          PIVOT_RC_CORRELATION);
    }

  struct pivot_dimension *stat = pivot_dimension_create (pt, PIVOT_AXIS_ROW,
                                                         "Statistic");
  for (size_t i = 0; i < N_STATS; i++)
    pivot_category_create_leaf (stat->root,
                                pivot_value_new_text (mm_stat_names[i]));

  struct pivot_dimension *split = pivot_dimension_create (
    pt, PIVOT_AXIS_ROW, "Split");

  int split_num = 0;

  struct matrix_material mm = MATRIX_MATERIAL_INIT;
  while (matrix_reader_next (&mm, mr, NULL))
    {
      pivot_category_create_leaf (split->root,
                                  pivot_value_new_integer (split_num + 1));

      const gsl_matrix *m[N_STATS] = {
        [MM_CORR] = mm.corr,
        [MM_COV] = mm.cov,
        [MM_N] = mm.n,
        [MM_MEAN] = mm.mean_matrix,
        [MM_STDDEV] = mm.var_matrix,
      };

      for (size_t i = 0; i < N_STATS; i++)
        if (m[i])
          {
            if (i == MM_COV || i == MM_CORR)
              {
                for (size_t y = 0; y < mr->n_cvars; y++)
                  for (size_t x = 0; x < mr->n_cvars; x++)
                    pivot_table_put4 (
                      pt, y + 1, x, i, split_num,
                      pivot_value_new_number (gsl_matrix_get (m[i], y, x)));
              }
            else
              for (size_t x = 0; x < mr->n_cvars; x++)
                {
                  double n = gsl_matrix_get (m[i], 0, x);
                  if (i == MM_STDDEV)
                    n = sqrt (n);
                  pivot_table_put4 (pt, 0, x, i, split_num,
                                    pivot_value_new_number (n));
                }
          }

      split_num++;
      matrix_material_uninit (&mm);
    }
  pivot_table_submit (pt);

  proc_commit (ds);

  matrix_reader_destroy (mr);
  return CMD_SUCCESS;
}
