"""
AMReX/Boxlib data-file handling functions



"""

#-----------------------------------------------------------------------------
# Copyright (c) 2013, yt Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------

import os
import numpy as np
from collections import defaultdict

from yt.utilities.io_handler import \
    BaseIOHandler
from yt.funcs import mylog
from yt.frontends.chombo.io import parse_orion_sinks
from yt.geometry.selection_routines import GridSelector

def _remove_raw(all_fields, raw_fields):
    centered_fields = set(all_fields)
    for raw in raw_fields:
        centered_fields.discard(raw)
    return list(centered_fields)

class IOHandlerBoxlib(BaseIOHandler):

    _dataset_type = "boxlib_native"

    def __init__(self, ds, *args, **kwargs):
        super(IOHandlerBoxlib, self).__init__(ds)

    def _read_fluid_selection(self, chunks, selector, fields, size):
        chunks = list(chunks)
        if any(( not (ftype == "boxlib" or ftype == 'raw') for ftype, fname in fields)):
            raise NotImplementedError
        rv = {}
        raw_fields = []
        for field in fields:
            if field[0] == "raw":
                nodal_flag = self.ds.nodal_flags[field[1]]
                num_nodes = 2**sum(nodal_flag)
                rv[field] = np.empty((size, num_nodes), dtype="float64")
                raw_fields.append(field)
            else:
                rv[field] = np.empty(size, dtype="float64")
        centered_fields = _remove_raw(fields, raw_fields)
        ng = sum(len(c.objs) for c in chunks)
        mylog.debug("Reading %s cells of %s fields in %s grids",
                    size, [f2 for f1, f2 in fields], ng)
        ind = 0
        for chunk in chunks:
            data = self._read_chunk_data(chunk, centered_fields)
            for g in chunk.objs:
                for field in fields:
                    if field in centered_fields:
                        ds = data[g.id].pop(field)
                    else:
                        ds = self._read_raw_field(g, field)
                    nd = g.select(selector, ds, rv[field], ind)
                ind += nd
                data.pop(g.id)
        return rv

    def _read_raw_field(self, grid, field):
        field_name = field[1]
        base_dir = self.ds.index.raw_file

        nghost = self.ds.index.raw_field_nghost[field_name]
        box_list = self.ds.index.raw_field_map[field_name][0]
        fn_list = self.ds.index.raw_field_map[field_name][1]
        offset_list = self.ds.index.raw_field_map[field_name][2]

        lev = grid.Level        
        filename = base_dir + "Level_%d/" % lev + fn_list[grid.id]
        offset = offset_list[grid.id]
        box = box_list[grid.id]

        lo = box[0] - nghost
        hi = box[1] + nghost
        shape = hi - lo + 1
        with open(filename, "rb") as f:
            f.seek(offset)
            f.readline()  # always skip the first line
            arr = np.fromfile(f, 'float64', np.product(shape))
            arr = arr.reshape(shape, order='F')
        return arr[[slice(nghost[dim],-nghost[dim]) for dim in range(self.ds.dimensionality)]]

    def _read_chunk_data(self, chunk, fields):
        data = {}
        grids_by_file = defaultdict(list)
        if len(chunk.objs) == 0: return data
        for g in chunk.objs:
            if g.filename is None:
                continue
            grids_by_file[g.filename].append(g)
        dtype = self.ds.index._dtype
        bpr = dtype.itemsize
        for filename in grids_by_file:
            grids = grids_by_file[filename]
            grids.sort(key = lambda a: a._offset)
            f = open(filename, "rb")
            for grid in grids:
                data[grid.id] = {}
                local_offset = grid._get_offset(f) - f.tell()
                count = grid.ActiveDimensions.prod()
                size = count * bpr
                for field in self.ds.index.field_order:
                    if field in fields:
                        # We read it ...
                        f.seek(local_offset, os.SEEK_CUR)
                        v = np.fromfile(f, dtype=dtype, count=count)
                        v = v.reshape(grid.ActiveDimensions, order='F')
                        data[grid.id][field] = v
                        local_offset = 0
                    else:
                        local_offset += size
        return data

    def _read_particle_coords(self, chunks, ptf):
        for rv in self._read_particle_fields(chunks, ptf, None):
            yield rv

    def _read_particle_fields(self, chunks, ptf, selector):
        for chunk in chunks: # These should be organized by grid filename
            for g in chunk.objs:
                for ptype, field_list in sorted(ptf.items()):
                    npart = g._pdata[ptype]["NumberOfParticles"]
                    if npart == 0:
                        continue

                    fn = g._pdata[ptype]["particle_filename"]
                    offset = g._pdata[ptype]["offset"]
                    pheader = self.ds.index.particle_headers[ptype]

                    with open(fn, "rb") as f:
                        # read in the position fields for selection
                        f.seek(offset + 
                               pheader.particle_int_dtype.itemsize * npart)
                        rdata = np.fromfile(f, pheader.real_type, pheader.num_real * npart)
                        x = np.asarray(rdata[0::pheader.num_real], dtype=np.float64)
                        y = np.asarray(rdata[1::pheader.num_real], dtype=np.float64)
                        if (g.ds.dimensionality == 2):
                            z = np.ones_like(y)
                            z *= 0.5*(g.LeftEdge[2] + g.RightEdge[2])
                        else:
                            z = np.asarray(rdata[2::pheader.num_real], dtype=np.float64)

                        if selector is None:
                            # This only ever happens if the call is made from
                            # _read_particle_coords.
                            yield ptype, (x, y, z)
                            continue
                        mask = selector.select_points(x, y, z, 0.0)
                        if mask is None:
                            continue
                        for field in field_list:
                            # handle the case that this is an integer field
                            int_fnames = [fname for _, fname in pheader.known_int_fields]
                            if field in int_fnames:
                                ind = int_fnames.index(field)
                                f.seek(offset)
                                idata = np.fromfile(f, pheader.int_type,
                                                    pheader.num_int * npart)
                                data = np.asarray(idata[ind::pheader.num_int],
                                                  dtype=np.float64)
                                yield (ptype, field), data[mask].flatten()

                            # handle case that this is a real field
                            real_fnames = [fname for _, fname in pheader.known_real_fields]
                            if field in real_fnames:
                                ind = real_fnames.index(field)                            
                                data = np.asarray(rdata[ind::pheader.num_real],
                                                  dtype=np.float64)
                                yield (ptype, field), data[mask].flatten()

class IOHandlerOrion(IOHandlerBoxlib):
    _dataset_type = "orion_native"

    _particle_filename = None
    @property
    def particle_filename(self):
        fn = self.ds.output_dir + "/StarParticles"
        if not os.path.exists(fn):
            fn = self.ds.output_dir + "/SinkParticles"
        self._particle_filename = fn
        return self._particle_filename

    _particle_field_index = None
    @property
    def particle_field_index(self):

        index = parse_orion_sinks(self.particle_filename)

        self._particle_field_index = index
        return self._particle_field_index

    def _read_particle_selection(self, chunks, selector, fields):
        rv = {}
        chunks = list(chunks)

        if isinstance(selector, GridSelector):

            if not (len(chunks) == len(chunks[0].objs) == 1):
                raise RuntimeError

            grid = chunks[0].objs[0]

            for ftype, fname in fields:
                rv[ftype, fname] = self._read_particles(grid, fname)

            return rv

        rv = {f: np.array([]) for f in fields}
        for chunk in chunks:
            for grid in chunk.objs:
                for ftype, fname in fields:
                    data = self._read_particles(grid, fname)
                    rv[ftype, fname] = np.concatenate((data, rv[ftype, fname]))
        return rv

    def _read_particles(self, grid, field):
        """
        parses the Orion Star Particle text files

        """

        particles = []

        if grid.NumberOfParticles == 0:
            return np.array(particles)

        def read(line, field):
            entry = line.strip().split(' ')[self.particle_field_index[field]]
            return np.float(entry)

        try:
            lines = self._cached_lines
            for num in grid._particle_line_numbers:
                line = lines[num]
                particles.append(read(line, field))
            return np.array(particles)
        except AttributeError:
            fn = self.particle_filename
            with open(fn, 'r') as f:
                lines = f.readlines()
                self._cached_lines = lines
                for num in grid._particle_line_numbers:
                    line = lines[num]
                    particles.append(read(line, field))
            return np.array(particles)
