from amuse.test import amusetest
from io import StringIO
import textwrap
import os
import numpy


from amuse import io
from amuse.io import vtk
from amuse.units import units
from amuse.units import generic_unit_system
from amuse.units import generic_unit_system
from amuse import datamodel

class VtkStructuredGridTests(amusetest.TestCase):
    
    def test1(self):
        grid = datamodel. new_regular_grid([2,3,4], [1,1,1] | generic_unit_system.length)
        grid.rho = grid.x * (0.1 | generic_unit_system.mass / generic_unit_system.length ** 4)
        data_file = StringIO()
        instance = vtk.VtkStructuredGrid("test.vts", data_file, grid)
        instance.store()
        
        contents = data_file.getvalue()
        self.assertTrue(contents.find('WholeExtent="0 2 0 3 0 4"')> 0)
        self.assertTrue(contents.find('<DataArray type="Float64" NumberOfComponents="3">')> 0)
        self.assertTrue(contents.find('<DataArray type="Float64" NumberOfComponents="1" Name="rho">')> 0)
    

    def test2(self):
        grid = datamodel. new_regular_grid([4,5,6], [1,1,1] | generic_unit_system.length)
        grid.mass = generic_unit_system.density(numpy.random.rand(4,5,6))
        data_file = StringIO()
        instance = vtk.VtkStructuredGrid("test.vts", data_file, grid)
        instance.store()
        
        contents = data_file.getvalue()
        self.assertTrue(contents.find('<Piece Extent="0 4 0 5 0 6">')> 0)
        self.assertTrue(contents.find('<CellData>')> 0)
        self.assertTrue(contents.find('<PointData>')> 0)
        self.assertTrue(contents.find('<DataArray type="Float64" NumberOfComponents="3">')> 0)
        self.assertTrue(contents.find('<DataArray type="Float64" NumberOfComponents="1" Name="mass">')> 0)
    
class VtkUnstructuredGridTests(amusetest.TestCase):
    
    def test1(self):
        grid = datamodel. new_regular_grid([2,3,4], [1,1,1] | generic_unit_system.length)
        grid.rho = grid.x * (0.1 | generic_unit_system.mass / generic_unit_system.length ** 4)
        data_file = StringIO()
        instance = vtk.VtkUnstructuredGrid("test.vtu", data_file, grid)
        instance.store()
        contents = data_file.getvalue()
        
        self.assertTrue(contents.find('<Piece NumberOfPoints="60" NumberOfCells="24"')> 0)
        self.assertTrue(contents.find('DataArray type="Int32" NumberOfComponents="1" Name="connectivity">')> 0)
        self.assertTrue(contents.find('DataArray type="Int32" NumberOfComponents="1" Name="offsets">')> 0)
        self.assertTrue(contents.find('11 11 11')> 0)
        self.assertTrue(contents.find('8 16 24')> 0)
        self.assertTrue(contents.find('DataArray type="Float64" NumberOfComponents="3" Name="points">')> 0)
        
    
    def test2(self):
        grid = datamodel. new_regular_grid([2,2,2], [1,1,1] | generic_unit_system.length)
        grid.rho = grid.x * (0.1 | generic_unit_system.mass / generic_unit_system.length ** 4)
        data_file = StringIO()
        instance = vtk.VtkUnstructuredGrid("test.vtu", data_file, grid)
        instance.store()
        contents = data_file.getvalue()
        
        self.assertTrue(contents.find('<Piece NumberOfPoints="27" NumberOfCells="8"')> 0)
        self.assertTrue(contents.find('DataArray type="Int32" NumberOfComponents="1" Name="connectivity">')> 0)
        self.assertTrue(contents.find('DataArray type="Int32" NumberOfComponents="1" Name="offsets">')> 0)
        self.assertTrue(contents.find('11 11 11')> 0)
        self.assertTrue(contents.find('56 64')> 0)
        self.assertTrue(contents.find('DataArray type="Float64" NumberOfComponents="3" Name="points">')> 0)
        self.assertTrue(contents.find('<DataArray type="Float64" NumberOfComponents="1" Name="rho">')>0)
        
    
    def test3(self):
        grid1 = datamodel. new_regular_grid([2,2,2], [1,1,1] | generic_unit_system.length)
        grid1.rho = grid1.x * (0.1 | generic_unit_system.mass / generic_unit_system.length ** 4)
        grid2 = datamodel. new_regular_grid([2,2,2], [1,1,1] | generic_unit_system.length)
        grid2.position += [1.0,0.0,0.0] | generic_unit_system.length
        grid2.rho = grid2.x * (0.1 | generic_unit_system.mass / generic_unit_system.length ** 4)
        
        data_file = StringIO()
        instance = vtk.VtkUnstructuredGrid("test.vtu", data_file, [grid1, grid2])
        instance.store()
        contents = data_file.getvalue()
        index_of_piece_1 = contents.find('<Piece NumberOfPoints="27" NumberOfCells="8"')
        index_of_piece_2 = contents.find('<Piece NumberOfPoints="27" NumberOfCells="8"', index_of_piece_1 + 1)
        self.assertTrue(index_of_piece_1 > 0)
        self.assertTrue(index_of_piece_2 > 0)
        self.assertTrue(contents.find('DataArray type="Int32" NumberOfComponents="1" Name="connectivity">')> 0)
        self.assertTrue(contents.find('DataArray type="Int32" NumberOfComponents="1" Name="offsets">')> 0)
        self.assertTrue(contents.find('11 11 11')> 0)
        self.assertTrue(contents.find('56 64')> 0)
        self.assertTrue(contents.find('DataArray type="Float64" NumberOfComponents="3" Name="points">')> 0)
        self.assertTrue(contents.find('<DataArray type="Float64" NumberOfComponents="1" Name="rho">')>0)
        
