001/*-
002 *******************************************************************************
003 * Copyright (c) 2011, 2016 Diamond Light Source Ltd.
004 * All rights reserved. This program and the accompanying materials
005 * are made available under the terms of the Eclipse Public License v1.0
006 * which accompanies this distribution, and is available at
007 * http://www.eclipse.org/legal/epl-v10.html
008 *
009 * Contributors:
010 *    Peter Chang - initial API and implementation and/or initial documentation
011 *******************************************************************************/
012
013package org.eclipse.january.dataset;
014
015import java.util.ArrayList;
016import java.util.Arrays;
017import java.util.List;
018
019/**
020 * Class to run over an array of integer datasets and return its items
021 */
022public class IntegersIterator extends IndexIterator {
023        final private int[] ishape; // shape of input
024        final private int irank; // rank of input shape
025        final private int[] oshape; // shape of output
026        final private int orank; // rank of output shape
027        private int offset; // offset of index subspace in new position 
028        private int srank; // rank of subspace
029
030        final private IndexIterator it;
031
032        /**
033         * position in input shape
034         */
035        final private int[] ipos;
036        /**
037         * position in output shape
038         */
039        final private int[] opos;
040        final private List<Object> indexes;
041
042        /**
043         * Constructor for an iterator over the items of an array of objects
044         * @param shape of entire data array
045         * @param index an array of integer dataset, boolean dataset, slices or null entries (same as full slices)
046         */
047        public IntegersIterator(final int[] shape, final Object... index) {
048                this(false, shape, index);
049        }
050
051        /**
052         * Constructor for an iterator over the items of an array of objects
053         * @param restrict1D if true, allow only one 1D integer datasets otherwise they must match shape
054         * @param shape of entire data array
055         * @param index an array of integer dataset, boolean dataset, slices or null entries (same as full slices)
056         */
057        public IntegersIterator(final boolean restrict1D, final int[] shape, final Object... index) {
058                if (shape == null) {
059                        ishape = null;
060                        irank = 0;
061                }  else {
062                        ishape = shape.clone();
063                        irank = shape.length;
064                }
065                if (irank < index.length) {
066                        throw new IllegalArgumentException("Number of index datasets is greater than rank of dataset");
067                }
068                indexes = new ArrayList<Object>();
069                for (Object i : index) {
070                        if (i instanceof BooleanDataset) { // turn boolean datasets into integer ones
071                                for (IntegerDataset id : Comparisons.nonZero((Dataset) i)) {
072                                        indexes.add(id);
073                                }
074                        } else if (i == null || i instanceof Slice) {
075                                indexes.add(i);
076                        } else if (i instanceof IntegerDataset) {
077                                Dataset id = (Dataset) i;
078                                int r = id.getRank();
079                                if (restrict1D && r > 1) {
080                                        throw new IllegalArgumentException("Integer datasets were restricted to zero or one dimensions");
081                                }
082                                if (r == 0) { // avoid zero-rank datasets
083                                        i = id.reshape(1);
084                                }
085                                indexes.add(i);
086                        } else {
087                                throw new IllegalArgumentException("Unsupported object for indexing");
088                        }
089                }
090                if (indexes.size() < irank) { // pad out index list
091                        for (int i = indexes.size(); i < irank; i++) {
092                                indexes.add(null);
093                        }
094                } else if (indexes.size() > irank) {
095                        throw new IllegalArgumentException("Too many indices (a boolean dataset may have too many dimensions)");
096                }
097
098                int ilength = -1;
099                int[] cshape = null;
100                int first = -1; // index of first null or slice after non-null index
101                boolean intact = true;
102                srank = 0;
103                for (int i = 0; i < irank; i++) { // see if shapes are consistent and subspace is intact
104                        Object obj = indexes.get(i);
105                        if (obj instanceof IntegerDataset && !restrict1D) {
106                                IntegerDataset ind = (IntegerDataset) obj;
107                                if (first > 0) {
108                                        intact = false;
109                                }
110
111                                int l = ind.size;
112                                if (ilength < l) {
113                                        ilength = l;
114                                        cshape = null;
115                                } else if (l != 1 && l != ilength) {
116                                        throw new IllegalArgumentException("Index datasets do not have same size");
117                                }
118                                if (cshape == null) {
119                                        cshape = ind.shape;
120                                        srank = cshape.length;
121                                        offset = i;
122                                } else if (l > 1 && !Arrays.equals(ind.shape, cshape)) { // broadcast
123                                        throw new IllegalArgumentException("Index datasets do not have same shape");
124                                }
125                        } else {
126                                if (cshape != null) {
127                                        if (first < 0)
128                                                first = i;
129                                }
130                        }
131                }
132
133                List<Integer> oShape = new ArrayList<Integer>(irank);
134
135                if (intact) { // get new output shape list
136                        boolean used = false;
137                        for (int i = 0; i < irank; i++) {
138                                Object obj = indexes.get(i);
139                                if (obj instanceof IntegerDataset) {
140                                        if (restrict1D || !used) {
141                                                used = true;
142                                                int[] lshape = restrict1D ? ((IntegerDataset) obj).shape : cshape;
143                                                for (int j : lshape) {
144                                                        oShape.add(j);
145                                                }
146                                        }
147                                } else if (obj instanceof Slice) {
148                                        Slice s = (Slice) obj;
149                                        int l = ishape[i];
150                                        s.setLength(l);
151                                        oShape.add(s.getNumSteps());
152                                } else {
153                                        oShape.add(ishape[i]);
154                                }
155                        }
156                } else {
157                        assert cshape != null;
158                        for (int j : cshape) {
159                                oShape.add(j);
160                        }
161                        for (int i = 0; i < irank; i++) {
162                                Object obj = indexes.get(i);
163                                if (obj == null) {
164                                        oShape.add(ishape[i]);
165                                } else if (obj instanceof Slice) {
166                                        Slice s = (Slice) obj;
167                                        int l = ishape[i];
168                                        s.setLength(l);
169                                        oShape.add(s.getNumSteps());
170                                }
171                        }
172                }
173                orank = oShape.size();
174                oshape = orank == 0 && ishape == null ? null : new int[orank];
175                for (int i = 0; i < orank; i++) {
176                        oshape[i] = oShape.get(i);
177                }
178
179                for (int i = 0; i < irank; i++) { // check input indexes for out of bounds
180                        Object obj = indexes.get(i);
181                        if (obj instanceof IntegerDataset) {
182                                IntegerDataset ind = (IntegerDataset) obj;
183                                if (ind.getSize() > 0) {
184                                        int l = ishape[i];
185                                        if (ind.min().intValue() < -l || ind.max().intValue() >= l) {
186                                                throw new IllegalArgumentException("A value in index datasets is outside permitted range");
187                                        }
188                                }
189                        }
190                }
191
192                ipos = new int[irank];
193                it = new PositionIterator(oshape);
194                opos = it.getPos();
195        }
196
197        @Override
198        public int[] getShape() {
199                return oshape;
200        }
201
202        @Override
203        public boolean hasNext() {
204                if (it.hasNext()) {
205                        int i = 0;
206                        for (; i < offset; i++) {
207                                Object obj = indexes.get(i);
208                                if (obj == null) {
209                                        ipos[i] = opos[i];
210                                } else if (obj instanceof Slice) {
211                                        Slice s = (Slice) obj;
212                                        ipos[i] = s.getPosition(opos[i]); // overwrite position
213                                } else {
214                                        throw new IllegalStateException("Bad state: index dataset before offset");
215                                }
216                        }
217                        int[] spos = srank > 0 ? Arrays.copyOfRange(opos, i, i+srank) : opos;
218                        if (spos == opos) {
219                                for (; i < irank; i++) {
220                                        Object obj = indexes.get(i);
221                                        if (obj == null) {
222                                                ipos[i] = opos[i];
223                                        } else if (obj instanceof Slice) {
224                                                Slice s = (Slice) obj;
225                                                ipos[i] = s.getPosition(opos[i]); // overwrite position
226                                        } else if (obj instanceof IntegerDataset) { // allowed when restricted to 1D
227                                                int p = ((Dataset) obj).getInt(opos[i]);
228                                                ipos[i] = p < 0 ? p + ishape[i] : p;
229                                        } else {
230                                                throw new IllegalStateException("Bad state: index dataset after subspace");
231                                        }
232                                }
233                        } else {
234                                for (int j = 0; j < irank; j++) {
235                                        Object obj = indexes.get(j);
236                                        if (obj instanceof IntegerDataset) {
237                                                IntegerDataset ind = (IntegerDataset) obj;
238                                                int p = ind.size > 1 ? ind.get(spos) : ind.getAbs(0); // broadcasting
239                                                ipos[i] = p < 0 ? p + ishape[i] : p;
240                                                i++;
241                                        }
242                                }
243                                int o = orank - irank;
244                                for (; i < irank; i++) {
245                                        Object obj = indexes.get(i);
246                                        if (obj == null) {
247                                                ipos[i] = opos[i+o];
248                                        } else if (obj instanceof Slice) {
249                                                Slice s = (Slice) obj;
250                                                ipos[i] = s.getPosition(opos[i+o]); // overwrite position
251                                        } else {
252                                                throw new IllegalStateException("Bad state: index dataset after subspace");
253                                        }
254                                }
255                        }
256//                      System.err.println(Arrays.toString(opos) + ", " + Arrays.toString(spos) + ", " + Arrays.toString(ipos));
257                        return true;
258                }
259                return false;
260        }
261
262        @Override
263        public int[] getPos() {
264                return ipos;
265        }
266
267        @Override
268        public void reset() {
269                it.reset();
270                Arrays.fill(ipos, 0);
271                index = 0;
272        }
273}