import unittest
import tsaugmentation as tsag
from gpforecaster.model.gpf import GPF
import shutil


class TestModel(unittest.TestCase):

    def setUp(self):
        self.data = tsag.preprocessing.PreprocessDatasets('police', top=10, freq='D').apply_preprocess()
        self.n = self.data['predict']['n']
        self.s = self.data['train']['s']
        self.gpf = GPF('police', self.data, log_dir="..")

    def test_predict_shape_xvalidation(self):
        model, like = self.gpf.train(epochs=11, n_splits=2)
        preds, preds_scaled = self.gpf.predict(model, like)
        self.assertTrue(preds[0].shape == (self.n, self.s))

    def test_predict_shape(self):
        model, like = self.gpf.train(epochs=11, n_splits=2, cross_validation=False)
        preds, preds_scaled = self.gpf.predict(model, like)
        self.assertTrue(preds[0].shape == (self.n, self.s))
