"""
Python-based grid handler, not to be confused with the SWIG-handler

Author: Matthew Turk <matthewturk@gmail.com>
Affiliation: KIPAC/SLAC/Stanford
Homepage: http://yt.enzotools.org/
License:
  Copyright (C) 2007-2008 Matthew Turk.  All Rights Reserved.

  This file is part of yt.

  yt is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 3 of the License, or
  (at your option) any later version.

  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 General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""

from yt.lagos import *
#import yt.enki, gc
from yt.funcs import *

class EnzoGridBase(EnzoData):
    _spatial = True
    _num_ghost_zones = 0
    """
    Class representing a single Enzo Grid instance.
    """
    _grids = None

    def __init__(self, id, filename=None, hierarchy = None):
        """
        Returns an instance of EnzoGrid with *id*, associated with
        *filename* and *hierarchy*.
        """
        #EnzoData's init function is a time-burglar.
        #All of the field parameters will be passed to us as needed.
        #EnzoData.__init__(self, None, [])
        self.data = {}
        self.field_parameters = {}
        self.fields = []
        self.start_index = None
        self.id = id
        if hierarchy: self.hierarchy = weakref.proxy(hierarchy)
        if filename: self.set_filename(filename)
        self.overlap_masks = [None, None, None]
        self._overlap_grids = [None, None, None]
        self._file_access_pooling = False

    def __len__(self):
        return na.prod(self.ActiveDimensions)

    def _generate_field(self, field):
        if fieldInfo.has_key(field):
            # First we check the validator
            try:
                fieldInfo[field].check_available(self)
            except NeedsGridType, ngt_exception:
                # This is only going to be raised if n_gz > 0
                n_gz = ngt_exception.ghost_zones
                f_gz = ngt_exception.fields
                gz_grid = self.retrieve_ghost_zones(n_gz, f_gz)
                temp_array = fieldInfo[field](gz_grid)
                sl = [slice(n_gz,-n_gz)] * 3
                self[field] = temp_array[sl]
            else:
                self[field] = fieldInfo[field](self)
        else: # Can't find the field, try as it might
            raise exceptions.KeyError, field

    def get_data(self, field):
        """
        Returns a field or set of fields for a key or set of keys
        """
        if not self.data.has_key(field):
            if field in self.hierarchy.field_list:
                conv_factor = 1.0
                if fieldInfo.has_key(field):
                    conv_factor = fieldInfo[field]._convert_function(self)
                try:
                    self[field] = self.readDataFast(field) * conv_factor
                except self._read_exception:
                    if field in fieldInfo:
                        if fieldInfo[field].particle_type:
                            self[field] = na.array([],dtype='int64')
                        if fieldInfo[field].not_in_all:
                            self[field] = na.zeros(self.ActiveDimensions, dtype='float64')
                    else: raise
            else:
                self._generate_field(field)
        return self.data[field]

    def clear_all_grid_references(self):
        """
        This clears out all references this grid has to any others, as
        well as the hierarchy.  It's like extra-cleaning after clear_data.
        """
        self.clear_all_derived_quantities()
        if hasattr(self, 'hierarchy'):
            del self.hierarchy
        if hasattr(self, 'Parent'):
            if self.Parent != None:
                self.Parent.clear_all_grid_references()
            del self.Parent
        if hasattr(self, 'Children'):
            for i in self.Children:
                if i != None:
                    del i
            del self.Children

    def _prepare_grid(self):
        """
        Copies all the appropriate attributes from the hierarchy
        """
        # This is definitely the slowest part of generating the hierarchy
        # Now we give it pointers to all of its attributes
        # Note that to keep in line with Enzo, we have broken PEP-8
        h = self.hierarchy # cache it
        self.Dimensions = h.gridDimensions[self.id-1]
        self.StartIndices = h.gridStartIndices[self.id-1]
        self.EndIndices = h.gridEndIndices[self.id-1]
        self.LeftEdge = h.gridLeftEdge[self.id-1]
        self.RightEdge = h.gridRightEdge[self.id-1]
        self.Level = h.gridLevels[self.id-1,0]
        self.Time = h.gridTimes[self.id-1,0]
        self.NumberOfParticles = h.gridNumberOfParticles[self.id-1,0]
        self.ActiveDimensions = (self.EndIndices - self.StartIndices + 1)
        self.Children = h.gridTree[self.id-1]
        pID = h.gridReverseTree[self.id-1]
        if pID != None and pID != -1:
            self.Parent = weakref.proxy(h.grids[pID - 1])
        else:
            self.Parent = None

    def _setup_dx(self):
        # So first we figure out what the index is.  We don't assume
        # that dx=dy=dz , at least here.  We probably do elsewhere.
        self.dx = self.hierarchy.gridDxs[self.id-1,0]
        self.dy = self.hierarchy.gridDys[self.id-1,0]
        self.dz = self.hierarchy.gridDzs[self.id-1,0]
        self.data['dx'] = self.dx
        self.data['dy'] = self.dy
        self.data['dz'] = self.dz
        self._corners = self.hierarchy.gridCorners[:,:,self.id-1]

    def _guess_properties_from_parent(self):
        """
        We know that our grid boundary occurs on the cell boundary of our
        parent.  This can be a very expensive process, but it is necessary
        in some hierarchys, where yt is unable to generate a completely
        space-filling tiling of grids, possibly due to the finite accuracy in a
        standard Enzo hierarchy file.
        """
        le = self.LeftEdge
        self.dx = self.Parent.dx/2.0
        self.dy = self.Parent.dy/2.0
        self.dz = self.Parent.dz/2.0
        ParentLeftIndex = na.rint((self.LeftEdge-self.Parent.LeftEdge)/self.Parent.dx)
        self.start_index = 2*(ParentLeftIndex + self.Parent.get_global_startindex()).astype('int64')
        self.LeftEdge = self.Parent.LeftEdge + self.Parent.dx * ParentLeftIndex
        self.RightEdge = self.LeftEdge + \
                         self.ActiveDimensions*na.array([self.dx,self.dy,self.dz])
        self.hierarchy.gridDxs[self.id-1,0] = self.dx
        self.hierarchy.gridDys[self.id-1,0] = self.dy
        self.hierarchy.gridDzs[self.id-1,0] = self.dz
        self.hierarchy.gridLeftEdge[self.id-1,:] = self.LeftEdge
        self.hierarchy.gridRightEdge[self.id-1,:] = self.RightEdge
        self.hierarchy.gridCorners[:,:,self.id-1] = na.array([ # Unroll!
            [self.LeftEdge[0], self.LeftEdge[1], self.LeftEdge[2]],
            [self.RightEdge[0], self.LeftEdge[1], self.LeftEdge[2]],
            [self.RightEdge[0], self.RightEdge[1], self.LeftEdge[2]],
            [self.RightEdge[0], self.RightEdge[1], self.RightEdge[2]],
            [self.LeftEdge[0], self.RightEdge[1], self.RightEdge[2]],
            [self.LeftEdge[0], self.LeftEdge[1], self.RightEdge[2]],
            [self.RightEdge[0], self.LeftEdge[1], self.RightEdge[2]],
            [self.LeftEdge[0], self.RightEdge[1], self.LeftEdge[2]],
            ], dtype='float64')
        self.__child_mask = None
        self.__child_indices = None
        self._setup_dx()

    def get_global_startindex(self):
        """
        Return the integer starting index for each dimension at the current
        level.
        """
        if self.start_index != None:
            return self.start_index
        if self.Parent == None:
            start_index = self.LeftEdge / na.array([self.dx, self.dy, self.dz])
            return na.rint(start_index).astype('int64').ravel()
        pdx = na.array([self.Parent.dx, self.Parent.dy, self.Parent.dz]).ravel()
        start_index = (self.Parent.get_global_startindex()) + \
                       na.rint((self.LeftEdge - self.Parent.LeftEdge)/pdx)
        self.start_index = (start_index*2).astype('int64').ravel()
        return self.start_index

    def _generate_overlap_masks(self, axis, LE, RE):
        """
        Generate a mask that shows which cells overlap with arbitrary arrays
        *LE* and *RE*) of edges, typically grids, along *axis*.
        Use algorithm described at http://www.gamedev.net/reference/articles/article735.asp
        """
        x = x_dict[axis]
        y = y_dict[axis]
        cond = self.RightEdge[x] >= LE[:,x]
        cond = na.logical_and(cond, self.LeftEdge[x] <= RE[:,x])
        cond = na.logical_and(cond, self.RightEdge[y] >= LE[:,y])
        cond = na.logical_and(cond, self.LeftEdge[y] <= RE[:,y])
        return cond
   
    def __repr__(self):
        return "Grid_%04i" % (self.id)

    def __int__(self):
        return self.id

    def clear_data(self):
        """
        Clear out the following things: child_mask, child_indices,
        all fields, all field parameters.
        """
        self._del_child_mask()
        self._del_child_indices()
        if hasattr(self, 'coarseData'):
            del self.coarseData
        if hasattr(self, 'retVal'):
            del self.retVal
        EnzoData.clear_data(self)
        self._setup_dx()

    def set_filename(self, filename):
        """
        Intelligently set the filename.
        """
        if self.hierarchy._strip_path:
            self.filename = os.path.join(self.hierarchy.directory,
                                         os.path.basename(filename))
        elif filename[0] == os.path.sep:
            self.filename = filename
        else:
            self.filename = os.path.join(self.hierarchy.directory, filename)
        return

    def find_max(self, field):
        """
        Returns value, index of maximum value of *field* in this gird
        """
        coord1d=(self[field]*self.child_mask).argmax()
        coord=na.unravel_index(coord1d, self[field].shape)
        val = self[field][coord]
        return val, coord

    def find_min(self, field):
        """
        Returns value, index of minimum value of *field* in this gird
        """
        coord1d=(self[field]*self.child_mask).argmin()
        coord=na.unravel_index(coord1d, self[field].shape)
        val = self[field][coord]
        return val, coord

    def get_position(self, index):
        """
        Returns center position of an *index*
        """
        pos = (index + 0.5) * self.dx + self.LeftEdge
        return pos

    def clear_all(self):
        """
        Clears all datafields from memory and calls
        :meth:`clear_derived_quantities`.
        """
        for key in self.keys():
            del self.data[key]
        del self.data
        if hasattr(self,"retVal"):
            del self.retVal
        self.data = {}
        self.clear_derived_quantities()

    def clear_derived_quantities(self):
        """
        Clears coordinates, child_indices, child_mask.
        """
        # Access the property raw-values here
        del self.child_mask
        del self.child_ind

    def __get_enzo_grid(self):
        """
        **DO NOT USE**

        This attempts to get an instance of this particular grid from the SWIG
        interface.  Note that it first checks to see if the ParameterFile has
        been instantiated.
        """
        if self.hierarchy.eiTopGrid == None:
            self.hierarchy.initializeEnzoInterface()
        p=re.compile("Grid = %s\n" % (self.id))
        h=open(self.hierarchyFilename,"r").read()
        m=re.search(p,h)
        h=open(self.hierarchyFilename,"r")
        retVal = yt.enki.EnzoInterface.fseek(h, long(m.end()), 0)
        self.eiGrid=yt.enki.EnzoInterface.grid()
        cwd = os.getcwd() # Hate doing this, need to for relative pathnames
        os.chdir(self.hierarchy.directory)
        self.eiGrid.ReadGrid(h, 1)
        os.chdir(cwd)
        mylog.debug("Grid read with SWIG")

    def _set_child_mask(self, newCM):
        if self.__child_mask != None:
            mylog.warning("Overriding child_mask attribute!  This is probably unwise!")
        self.__child_mask = newCM

    def _set_child_indices(self, newCI):
        if self.__child_indices != None:
            mylog.warning("Overriding child_indices attribute!  This is probably unwise!")
        self.__child_indices = newCI

    def _get_child_mask(self):
        if self.__child_mask == None:
            self.__generate_child_mask()
        return self.__child_mask

    def _get_child_indices(self):
        if self.__child_indices == None:
            self.__generate_child_mask()
        return self.__child_indices

    def _del_child_indices(self):
        try:
            del self.__child_indices
        except AttributeError:
            pass
        self.__child_indices = None

    def _del_child_mask(self):
        try:
            del self.__child_mask
        except AttributeError:
            pass
        self.__child_mask = None

    #@time_execution
    def __generate_child_mask(self):
        """
        Generates self.child_mask, which is zero where child grids exist (and
        thus, where higher resolution data is available.)
        """
        self.__child_mask = na.ones(self.ActiveDimensions, 'int32')
        for child in self.Children:
            # Now let's get our overlap
            startIndex = na.rint((child.LeftEdge - self.LeftEdge)/self.dx)
            endIndex = na.rint((child.RightEdge - self.LeftEdge)/self.dx)
            self.__child_mask[startIndex[0]:endIndex[0],
                              startIndex[1]:endIndex[1],
                              startIndex[2]:endIndex[2]] = 0
        self.__child_indices = (self.__child_mask==0) # bool, possibly redundant

    def _get_coords(self):
        if self.__coords == None: self._generate_coords()
        return self.__coords

    def _set_coords(self, newC):
        if self.__coords != None:
            mylog.warning("Overriding coords attribute!  This is probably unwise!")
        self.__coords = newC

    def _del_coords(self):
        del self.__coords
        self.__coords = None

    def _generate_coords(self):
        """
        Creates self.coords, which is of dimensions (3,ActiveDimensions)
        """
        #print "Generating coords"
        ind = na.indices(self.ActiveDimensions)
        LE = na.reshape(self.LeftEdge,(3,1,1,1))
        self['x'], self['y'], self['z'] = (ind+0.5)*self.dx+LE

    __child_mask = None
    __child_indices = None

    child_mask = property(fget=_get_child_mask, fdel=_del_child_mask)
    child_indices = property(fget=_get_child_indices, fdel = _del_child_indices)

    def retrieve_ghost_zones(self, n_zones, fields, all_levels=False,
                             smoothed=False):
        # We will attempt this by creating a datacube that is exactly bigger
        # than the grid by nZones*dx in each direction
        new_left_edge = self.LeftEdge - n_zones * self.dx
        new_right_edge = self.RightEdge + n_zones * self.dx
        # Something different needs to be done for the root grid, though
        level = self.Level
        if all_levels:
            level = self.hierarchy.max_level + 1
        args = (level, new_left_edge, new_right_edge)
        kwargs = {'dims': self.ActiveDimensions + 2*n_zones,
                  'num_ghost_zones':n_zones,
                  'use_pbar':False, 'fields':fields}
        if smoothed:
            cube = self.hierarchy.smoothed_covering_grid(*args, **kwargs)
        else:
            cube = self.hierarchy.covering_grid(*args, **kwargs)
        return cube

    def _save_data_state(self):
        self.__current_data_keys = self.data.keys()
        if self.__child_mask != None:
            self.__current_child_mask == True
        else:
            self.__current_child_mask = False

        if self.__child_indices != None:
            self.__current_child_indices == True
        else:
            self.__current_child_indices = False

    def _restore_data_state(self):
        if not self.__current_child_mask:
            self._del_child_mask()
        if not self.__current_child_indices:
            self._del_child_indices()
        for key in data.keys():
            if key not in self.__current_data_keys:
                del self.data[key]
