package org.simantics.pythonlink.test;

import static org.junit.Assert.*;

import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.simantics.pythonlink.NDArray;
import org.simantics.pythonlink.Python;
import org.simantics.pythonlink.PythonContext;

public class TestPythonlink {
    PythonContext python;

    @Before
    public void setUp() throws Exception {
        python = Python.openPythonContext();
    }

    @After
    public void tearDown() throws Exception {
        python.close();
    }

    @Test
    public void testInteger() {
        int input = 4;
        python.setPythonIntegerVariable( "foo", input );
        python.executePythonStatement( "bar = foo // 2" );
        int output = python.getPythonIntegerVariable( "bar" );
        assertEquals( input / 2, output, 0.0 );
    }

    @Test
    public void testIntegerArray() {
        int[] input = new int[] { 1, 2 };
        python.setPythonIntegerArrayVariable( "foo", input );
        python.executePythonStatement( "bar = foo\nbar.append(3)" );
        int[] output = python.getPythonIntegerArrayVariable( "bar" );
        assertArrayEquals( new int[] { 1, 2, 3 }, output );
    }

    @Test
    public void testDouble() {
        double input = 2.5;
        python.setPythonDoubleVariable( "foo", input );
        python.executePythonStatement( "bar = foo / 2" );
        double output = python.getPythonDoubleVariable( "bar" );
        assertEquals( input / 2, output, 0.0 );
    }

    @Test
    public void testDoubleArray() {
        double[] input = new double[] { 1.0, 2.0 };
        python.setPythonDoubleArrayVariable( "foo", input );
        python.executePythonStatement( "bar = foo\nbar.append(3.0)" );
        double[] output = python.getPythonDoubleArrayVariable( "bar" );
        assertArrayEquals( new double[] { 1.0, 2.0, 3.0 }, output, 0.0 );
    }
    
    @Test
    public void testString() {
        String input = "Foo";
        python.setPythonStringVariable( "foo", input );
        python.executePythonStatement( "bar = foo + 'bar!'" );
        String output = python.getPythonStringVariable( "bar" );
        assertEquals( "Foobar!", output );
    }    
    
    @Test
    public void testStringArray() {
        String[] input = new String[] { "Foo" };
        python.setPythonStringArrayVariable( "foo", input );
        python.executePythonStatement( "bar = foo + ['bar!']" );
        String[] output = python.getPythonStringArrayVariable( "bar" );
        assertArrayEquals( new String[] { "Foo", "bar!" }, output );
    }
    
    @Test
    public void testNumpyArrays() {
        NDArray data = new NDArray(new int[] { 2, 3 }, new double[] { 1, 2, 3,  4, 5, 6 });
        python.executePythonStatement( "import numpy" );
        python.setPythonNDArrayVariable( "foo", data );
        python.executePythonStatement( "bar = foo.cumsum(axis = 0)" );
        NDArray result = python.getPythonNDArrayVariable( "bar" );
        assertEquals(new NDArray(new int[] { 2, 3 }, new double[] { 1, 2, 3,  5, 7, 9 }), result);
        assertEquals(5, result.getValue( 1, 0 ), 0.0);
        assertEquals(3, result.getValue( 0, 2 ), 0.0);
        assertEquals(9, result.getValue( 1, 2 ), 0.0);
    }
    
    @Test
    public void testNumpyArrays3() {
        NDArray data = new NDArray(new int[] { 3, 2, 2 }, new double[] { 1, 2,  3, 4,   5, 6,  7, 8,   9, 10,  11, 12 });
        python.executePythonStatement( "import numpy" );
        python.setPythonNDArrayVariable( "foo", data );
        python.executePythonStatement( "bar = numpy.ndarray(shape = (3,2,2), order='f')" );
        python.executePythonStatement( "bar.flat = foo.cumsum(axis = 1).flat" );
        NDArray result = python.getPythonNDArrayVariable( "bar" );
        assertEquals(new NDArray(new int[] { 3, 2, 2 }, new double[] { 1, 2,  4, 6,   5, 6,  12, 14,   9, 10,  20, 22 }), result);
        assertEquals(10, result.getValue( 2, 0, 1 ), 0.0);
        assertEquals(12, result.getValue( 1, 1, 0 ), 0.0);
    }
}