#
#  Copyright (C) 2000-2008  greg Landrum and Rational Discovery LLC
#    All Rights Reserved
#
""" classes to be used to help work with data sets

"""
from __future__ import print_function

import copy
import math

import numpy

from rdkit.six import integer_types


numericTypes = integer_types + (float, )


class MLDataSet(object):
  """ A data set for holding general data (floats, ints, and strings)

   **Note**
     this is intended to be a read-only data structure
     (i.e. after calling the constructor you cannot touch it)
  """

  def __init__(self, data, nVars=None, nPts=None, nPossibleVals=None, qBounds=None, varNames=None,
               ptNames=None, nResults=1):
    """ Constructor

      **Arguments**

        - data: a list of lists containing the data. The data are copied, so don't worry
              about us overwriting them.

        - nVars: the number of variables

        - nPts: the number of points

        - nPossibleVals: an list containing the number of possible values
                       for each variable (should contain 0 when not relevant)
                       This is _nVars_ long

        - qBounds: a list of lists containing quantization bounds for variables
                 which are to be quantized (note, this class does not quantize
                 the variables itself, it merely stores quantization bounds.
                 an empty sublist indicates no quantization for a given variable
                 This is _nVars_ long

        - varNames: a list of the names of the variables.
                 This is _nVars_ long

        - ptNames: the names (labels) of the individual data points
           This is _nPts_ long

        - nResults: the number of results columns in the data lists.  This is usually
                    1, but can be higher.
    """
    self.data = [x[:] for x in data]
    self.nResults = nResults
    if nVars is None:
      nVars = len(self.data[0]) - self.nResults
    self.nVars = nVars
    if nPts is None:
      nPts = len(data)
    self.nPts = nPts
    if qBounds is None:
      qBounds = [[]] * len(self.data[0])
    self.qBounds = qBounds
    if nPossibleVals is None:
      nPossibleVals = self._CalcNPossible(self.data)
    self.nPossibleVals = nPossibleVals
    if varNames is None:
      varNames = [''] * self.nVars
    self.varNames = varNames
    if ptNames is None:
      ptNames = [''] * self.nPts
    self.ptNames = ptNames

  def _CalcNPossible(self, data):
    """calculates the number of possible values of each variable (where possible)

      **Arguments**

         -data: a list of examples to be used

      **Returns**

         a list of nPossible values for each variable

    """
    nVars = self.GetNVars() + self.nResults
    nPossible = [-1] * nVars
    cols = list(range(nVars))
    for i, bounds in enumerate(self.qBounds):
      if len(bounds) > 0:
        nPossible[i] = len(bounds)
        cols.remove(i)

    for i, pt in enumerate(self.data):
      for col in cols[:]:
        d = pt[col]
        if type(d) in numericTypes:
          if math.floor(d) == d:
            nPossible[col] = max(math.floor(d), nPossible[col])
          else:
            nPossible[col] = -1
            cols.remove(col)
        else:
          nPossible[col] = -1
          cols.remove(col)
    return [int(x) + 1 for x in nPossible]

  def GetNResults(self):
    return self.nResults

  def GetNVars(self):
    return self.nVars

  def GetNPts(self):
    return self.nPts

  def GetNPossibleVals(self):
    return self.nPossibleVals

  def GetQuantBounds(self):
    return self.qBounds

  def __getitem__(self, idx):
    res = [self.ptNames[idx]] + self.data[idx][:]
    return res

  def __setitem__(self, idx, val):
    if len(val) != self.GetNVars() + self.GetNResults() + 1:
      raise ValueError('bad value in assignment')
    self.ptNames[idx] = val[0]
    self.data[idx] = val[1:]
    return val

  def GetNamedData(self):
    """ returns a list of named examples

     **Note**

       a named example is the result of prepending the example
        name to the data list

    """
    res = [None] * self.nPts
    for i in range(self.nPts):
      res[i] = [self.ptNames[i]] + self.data[i][:]
    return res

  def GetAllData(self):
    """ returns a *copy* of the data

    """
    return copy.deepcopy(self.data)

  def GetInputData(self):
    """ returns the input data

     **Note**

       _inputData_ means the examples without their result fields
        (the last _NResults_ entries)

    """
    v = self.GetNResults()
    return [x[:-v] for x in self.data]

  def GetResults(self):
    """ Returns the result fields from each example

    """
    if self.GetNResults() > 1:
      v = self.GetNResults()
      res = [x[-v:] for x in self.data]
    else:
      res = [x[-1] for x in self.data]
    return res

  def GetVarNames(self):
    return self.varNames

  def GetPtNames(self):
    return self.ptNames

  def AddPoint(self, pt):
    self.data.append(pt[1:])
    self.ptNames.append(pt[0])
    self.nPts += 1

  def AddPoints(self, pts, names):
    if len(pts) != len(names):
      raise ValueError("input length mismatch")
    self.data += pts
    self.ptNames += names
    self.nPts = len(self.data)


class MLQuantDataSet(MLDataSet):
  """ a data set for holding quantized data


    **Note**

      this is intended to be a read-only data structure
      (i.e. after calling the constructor you cannot touch it)

    **Big differences to MLDataSet**

      1) data are stored in a numpy array since they are homogenous

      2) results are assumed to be quantized (i.e. no qBounds entry is required)

  """

  def _CalcNPossible(self, data):
    """calculates the number of possible values of each variable

      **Arguments**

         -data: a list of examples to be used

      **Returns**

         a list of nPossible values for each variable

    """
    return [max(x) + 1 for x in numpy.transpose(data)]

  def GetNamedData(self):
    """ returns a list of named examples

     **Note**

       a named example is the result of prepending the example
        name to the data list

    """
    res = [None] * self.nPts
    for i in range(self.nPts):
      res[i] = [self.ptNames[i]] + self.data[i].tolist()
    return res

  def GetAllData(self):
    """ returns a *copy* of the data

    """
    return self.data.tolist()

  def GetInputData(self):
    """ returns the input data

     **Note**

       _inputData_ means the examples without their result fields
        (the last _NResults_ entries)

    """
    return (self.data[:, :-self.nResults]).tolist()

  def GetResults(self):
    """ Returns the result fields from each example

    """
    if self.GetNResults() > 1:
      v = self.GetNResults()
      res = [x[-v:] for x in self.data]
    else:
      res = [x[-1] for x in self.data]
    return res

  def __init__(self, data, nVars=None, nPts=None, nPossibleVals=None, qBounds=None, varNames=None,
               ptNames=None, nResults=1):
    """ Constructor

      **Arguments**

        - data: a list of lists containing the data. The data are copied, so don't worry
              about us overwriting them.

        - nVars: the number of variables

        - nPts: the number of points

        - nPossibleVals: an list containing the number of possible values
                       for each variable (should contain 0 when not relevant)
                       This is _nVars_ long

        - qBounds: a list of lists containing quantization bounds for variables
                 which are to be quantized (note, this class does not quantize
                 the variables itself, it merely stores quantization bounds.
                 an empty sublist indicates no quantization for a given variable
                 This is _nVars_ long

        - varNames: a list of the names of the variables.
                 This is _nVars_ long

        - ptNames: the names (labels) of the individual data points
           This is _nPts_ long

        - nResults: the number of results columns in the data lists.  This is usually
                    1, but can be higher.
    """
    self.data = numpy.array(data)
    self.nResults = nResults
    if nVars is None:
      nVars = len(data[0]) - self.nResults
    self.nVars = nVars
    if nPts is None:
      nPts = len(data)
    self.nPts = nPts
    if qBounds is None:
      qBounds = [[]] * self.nVars
    self.qBounds = qBounds
    if nPossibleVals is None:
      nPossibleVals = self._CalcNPossible(data)
    self.nPossibleVals = nPossibleVals
    if varNames is None:
      varNames = [''] * self.nVars
    self.varNames = varNames
    if ptNames is None:
      ptNames = [''] * self.nPts
    self.ptNames = ptNames


if __name__ == '__main__':
  from . import DataUtils
  examples = [[0, 0, 0, 0, 0], [0, 0, 0, 1, 0], [1, 0, 0, 0, 1], [2, 1, 0, 0, 1], [2, 2, 1, 0, 1]]
  varNames = ['foo1', 'foo2', 'foo3', 'foo4', 'res']
  ptNames = ['p1', 'p2', 'p3', 'p4', 'p5']
  dataset = MLQuantDataSet(examples, varNames=varNames, ptNames=ptNames)
  DataUtils.WritePickledData('test_data/test.qdat.pkl', dataset)
  print('nVars:', dataset.GetNVars())
  print('nPts:', dataset.GetNPts())
  print('nPoss:', dataset.GetNPossibleVals())
  print('qBounds:', dataset.GetQuantBounds())
  print('data:', dataset.GetAllData())
  print('Input data:', dataset.GetInputData())
  print('results:', dataset.GetResults())

  print('nameddata:', dataset.GetNamedData())

  examples = [
    ['foo', 1, 1.0, 1, 1.1],
    ['foo', 2, 1.0, 1, 2.1],
    ['foo', 3, 1.2, 1.1, 3.1],
    ['foo', 4, 1.0, 1, 4.1],
    ['foo', 5, 1.1, 1, 5.1],
  ]
  qBounds = [[], [], [], [], [2, 4]]
  varNames = ['foo1', 'foo2', 'foo3', 'foo4', 'res']
  ptNames = ['p1', 'p2', 'p3', 'p4', 'p5']
  dataset = MLDataSet(examples, qBounds=qBounds)
  DataUtils.WritePickledData('test_data/test.dat.pkl', dataset)
  print('nVars:', dataset.GetNVars())
  print('nPts:', dataset.GetNPts())
  print('nPoss:', dataset.GetNPossibleVals())
  print('qBounds:', dataset.GetQuantBounds())
  print('data:', dataset.GetAllData())
  print('Input data:', dataset.GetInputData())
  print('results:', dataset.GetResults())

  print('nameddata:', dataset.GetNamedData())
