# ----------------------------------------------------------------------
# Numenta Platform for Intelligent Computing (NuPIC)
# Copyright (C) 2014, Numenta, Inc.  Unless you have an agreement
# with Numenta, Inc., for a separate license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero Public License version 3 as
# published by the Free Software Foundation.
#
# 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 Affero Public License for more details.
#
# You should have received a copy of the GNU Affero Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------

import copy
import numpy as np
import os
import pandas
import shutil
import tempfile
import unittest

import nab.corpus
from nab.util import recur



class CorpusTest(unittest.TestCase):


  @classmethod
  def setUpClass(cls):
    depth = 3

    cls.root = recur(os.path.dirname, os.path.realpath(__file__), depth)
    cls.corpusSource = os.path.join(cls.root, "tests", "test_data")


  def setUp(self):
    self.corpus = nab.corpus.Corpus(self.corpusSource)


  def testGetDataFiles(self):
    """
    Test the getDataFiles() function, specifically check if corpus.dataFiles
    is a dictionary containing DataFile objects containing pandas.DataFrame
    objects to represent the underlying data.
    """
    for df in self.corpus.dataFiles.values():
      self.assertIsInstance(df, nab.corpus.DataFile)
      self.assertIsInstance(df.data, pandas.DataFrame)
      self.assertEqual(set(df.data.columns.values),
        set(["timestamp", "value"]))


  def testAddColumn(self):
    """
    Test the addColumn() function, specificially check if a new column named
    "test" is added.
    """
    columnData = {}
    for relativePath, df in self.corpus.dataFiles.iteritems():
      rows, _ = df.data.shape
      columnData[relativePath] = pandas.Series(np.zeros(rows))

    self.corpus.addColumn("test", columnData, write=False)

    for df in self.corpus.dataFiles.values():
      self.assertEqual(set(df.data.columns.values),
        set(["timestamp", "value", "test"]))


  def testRemoveColumn(self):
    """
    Test the removeColumn() function, specifically check if an added column
    named "test" is removed.
    """
    columnData = {}
    for relativePath, df in self.corpus.dataFiles.iteritems():
      rows, _ = df.data.shape
      columnData[relativePath] = pandas.Series(np.zeros(rows))

    self.corpus.addColumn("test", columnData, write=False)

    self.corpus.removeColumn("test", write=False)

    for df in self.corpus.dataFiles.values():
      self.assertEqual(set(df.data.columns.values),
        set(["timestamp", "value"]))


  def testCopy(self):
    """
    Test the copy() function, specifically check if it copies the whole corpus
    to another directory and that the copied corpus is the exact same as the
    original.
    """
    copyLocation = os.path.join(tempfile.mkdtemp(), "test")
    self.corpus.copy(copyLocation)

    copyCorpus = nab.corpus.Corpus(copyLocation)

    for relativePath in self.corpus.dataFiles.keys():
      self.assertIn(relativePath, copyCorpus.dataFiles.keys())

      self.assertTrue(
        all(self.corpus.dataFiles[relativePath].data == \
            copyCorpus.dataFiles[relativePath].data))

    shutil.rmtree(copyLocation)


  def testAddDataSet(self):
    """
    Test the addDataSet() function, specifically check if it adds a new
    data file in the correct location in directory and into the dataFiles
    attribute.
    """
    copyLocation = os.path.join(tempfile.mkdtemp(), "test")
    copyCorpus = self.corpus.copy(copyLocation)

    for relativePath, df in self.corpus.dataFiles.iteritems():
      newPath = relativePath + "_copy"
      copyCorpus.addDataSet(newPath, copy.deepcopy(df))

      self.assertTrue(all(copyCorpus.dataFiles[newPath].data == df.data))

    shutil.rmtree(copyLocation)


  def testGetDataSubset(self):
    """
    Test the getDataSubset() function, specifically check if it returns only
    dataFiles with relativePaths that contain the query given.
    """
    query1 = "realAWSCloudwatch"
    subset1 = self.corpus.getDataSubset(query1)

    self.assertEqual(len(subset1), 2)
    for relativePath in subset1.keys():
      self.assertIn(query1, relativePath)

    query2 = "artificialWithAnomaly"
    subset2 = self.corpus.getDataSubset(query2)

    self.assertEqual(len(subset2), 1)

    for relativePath in subset2.keys():
      self.assertIn(query2, relativePath)


if __name__ == '__main__':
  unittest.main()
