/*
 * Copyright (c) 2014, Cloudera, Inc. All Rights Reserved.
 *
 * Cloudera, Inc. licenses this file to you under the Apache License,
 * Version 2.0 (the "License"). You may not use this file except in
 * compliance with the License. You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
 * CONDITIONS OF ANY KIND, either express or implied. See the License for
 * the specific language governing permissions and limitations under the
 * License.
 */

package com.cloudera.oryx.common.pmml;

import java.nio.file.Files;
import java.nio.file.Path;
import java.util.List;

import org.dmg.pmml.MiningFunction;
import org.dmg.pmml.Model;
import org.dmg.pmml.PMML;
import org.dmg.pmml.tree.CountingLeafNode;
import org.dmg.pmml.tree.Node;
import org.dmg.pmml.tree.TreeModel;
import org.junit.Test;

import com.cloudera.oryx.common.OryxTest;

public final class PMMLUtilsTest extends OryxTest {

  public static PMML buildDummyModel() {
    Node node = new CountingLeafNode().setRecordCount(123.0);
    TreeModel treeModel = new TreeModel(MiningFunction.CLASSIFICATION, null, node);
    PMML pmml = PMMLUtils.buildSkeletonPMML();
    pmml.addModels(treeModel);
    return pmml;
  }

  @Test
  public void testSkeleton() {
    PMML pmml = PMMLUtils.buildSkeletonPMML();
    assertEquals("Oryx", pmml.getHeader().getApplication().getName());
    assertNotNull(pmml.getHeader().getTimestamp());
  }

  @Test
  public void testReadWrite() throws Exception {
    Path tempModelFile = Files.createTempFile(getTempDir(), "model", ".pmml");
    PMML model = buildDummyModel();
    PMMLUtils.write(model, tempModelFile);
    assertTrue(Files.exists(tempModelFile));
    PMML model2 = PMMLUtils.read(tempModelFile);
    List<Model> models = model2.getModels();
    assertEquals(1, models.size());
    assertInstanceOf(models.get(0), TreeModel.class);
    TreeModel treeModel = (TreeModel) models.get(0);
    assertEquals(123.0, treeModel.getNode().getRecordCount().doubleValue());
    assertEquals(MiningFunction.CLASSIFICATION, treeModel.getMiningFunction());
  }

  @Test
  public void testToString() throws Exception {
    PMML model = buildDummyModel();
    model.getHeader().setTimestamp(null);
    assertEquals("<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"yes\"?>" +
                 "<PMML version=\"4.3\" xmlns=\"http://www.dmg.org/PMML-4_3\" " +
                 "xmlns:data=\"http://jpmml.org/jpmml-model/InlineTable\">" +
                 "<Header>" +
                 "<Application name=\"Oryx\"/>" +
                 "</Header>" +
                 "<TreeModel functionName=\"classification\">" +
                 "<Node recordCount=\"123.0\"/>" +
                 "</TreeModel>" +
                 "</PMML>",
                 PMMLUtils.toString(model));
  }

  @Test
  public void testFromString() throws Exception {
    PMML model = buildDummyModel();
    PMML model2 = PMMLUtils.fromString(PMMLUtils.toString(model));
    assertEquals(model.getHeader().getApplication().getName(),
                 model2.getHeader().getApplication().getName());
    assertEquals(model.getModels().get(0).getMiningFunction(),
                 model2.getModels().get(0).getMiningFunction());
  }

  @Test
  public void testPreviousPMMLVersion() throws Exception {
    String pmml42 =
        "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"yes\"?>\n" +
        "<PMML xmlns=\"http://www.dmg.org/PMML-4_2\" version=\"4.2.1\">\n" +
        "    <Header>\n" +
        "        <Application name=\"Oryx\"/>\n" +
        "    </Header>\n" +
        "    <TreeModel functionName=\"classification\">\n" +
        "        <Node recordCount=\"123.0\"/>\n" +
        "    </TreeModel>\n" +
        "</PMML>\n";
    PMML model = PMMLUtils.fromString(pmml42);
    // Actually transforms to latest version:
    assertEquals(PMMLUtils.VERSION, model.getVersion());
  }

}

@@ -903,16 +903,59 @@
 it.Test;
+%0Aimport static org.hamcrest.CoreMatchers.*;
 %0A%0Aimport
@@ -2348,23 +2348,75 @@
   assert
-Equals(
+That(PMMLUtils.toString(model), anyOf(is(%0A                 
 %22%3C?xml v
@@ -2920,17 +2920,22 @@
 %3C/PMML%3E%22
-,
+), is(
 %0A       
@@ -2948,32 +2948,545 @@
     
-PMMLUtils.toString(model
+%22%3C?xml version=%5C%221.0%5C%22 encoding=%5C%22UTF-8%5C%22 standalone=%5C%22yes%5C%22?%3E%22 +%0A                 %22%3CPMML version=%5C%224.3%5C%22 xmlns:data=%5C%22http://jpmml.org/jpmml-model/InlineTable%5C%22 %22 +%0A                 %22xmlns=%5C%22http://www.dmg.org/PMML-4_3%5C%22%3E%22 +%0A                 %22%3CHeader%3E%22 +%0A                 %22%3CApplication name=%5C%22Oryx%5C%22/%3E%22 +%0A                 %22%3C/Header%3E%22 +%0A                 %22%3CTreeModel functionName=%5C%22classification%5C%22%3E%22 +%0A                 %22%3CNode recordCount=%5C%22123.0%5C%22/%3E%22 +%0A                 %22%3C/TreeModel%3E%22 +%0A                 %22%3C/PMML%3E%22%0A                 )%0A    
 ));%0A

Fix flaky test PMMLUtilsTest.testToString (#368)