package com.bc.zarr;

import com.bc.zarr.chunk.ChunkReaderWriter;
import com.bc.zarr.chunk.Compressor;
import com.bc.zarr.ucar.NetCDF_Util;
import com.bc.zarr.ucar.Partial2dDataCopier;
import ucar.ma2.Array;
import ucar.ma2.InvalidRangeException;

import java.io.IOException;
import java.nio.file.Path;
import java.util.Arrays;
import java.util.Collections;
import java.util.Map;
import java.util.TreeMap;

public class ZarrReaderWriter implements ZarrWriter, ZarrReader {

    private final int[] _shape;
    private final int[] _chunks;
    private final Path _dataPath;
    private final ChunkReaderWriter _chunkReaderWriter;
    private final Map<Object, Object> _locks;
    private final ZarrDataType _dataType;

    public ZarrReaderWriter(Path dataPath, int[] shape, int[] chunkShape, ZarrDataType dataType, Number fillValue, Compressor compressor) {
        _dataPath = dataPath;
        _shape = shape;
        _chunks = chunkShape;
        _dataType = dataType;
        _chunkReaderWriter = ChunkReaderWriter.create(compressor, dataType, chunkShape, fillValue);
        _locks = Collections.synchronizedMap(new TreeMap<>());
    }

    @Override
    public ZarrDataType getDataType() {
        return _dataType;
    }

    @Override
    public int[] getShape() {
        return Arrays.copyOf(_shape, _shape.length);
    }

    @Override
    public int[] getChunks() {
        return Arrays.copyOf(_chunks, _chunks.length);
    }

    @Override
    public void write(Object dataBuffer, int[] bufferShape, int[] to) throws IOException, InvalidRangeException {
        final int[][] chunkIndices = ZarrUtils.computeChunkIndices(_shape, _chunks, bufferShape, to);
        final Array source = Array.factory(dataBuffer).reshapeNoCopy(bufferShape);

        for (int[] chunkIndex : chunkIndices) {
            final String chunkFilename = ZarrUtils.createChunkFilename(chunkIndex);
            synchronized (_locks) {
                if (!_locks.containsKey(chunkFilename)) {
                    _locks.put(chunkFilename, chunkFilename);
                }
            }
            final Path chunkFilePath = _dataPath.resolve(chunkFilename);
            final int[] fromBufferPos = computeFrom(chunkIndex, to, false);
            synchronized (_locks.get(chunkFilename)) {
                final Array targetChunk = _chunkReaderWriter.read(chunkFilePath);
                final Array read = Partial2dDataCopier.copy(fromBufferPos, source, targetChunk);
                _chunkReaderWriter.write(chunkFilePath, read);
            }
        }
    }

    @Override
    public void read(Object targetBuffer, int[] bufferShape, int[] from) throws IOException, InvalidRangeException {
        final int[][] chunkIndices = ZarrUtils.computeChunkIndices(_shape, _chunks, bufferShape, from);
        final Array target = NetCDF_Util.createArrayWithGivenStorage(targetBuffer, bufferShape);

        for (int[] chunkIndex : chunkIndices) {
            final String chunkFilename = ZarrUtils.createChunkFilename(chunkIndex);
            final Path chunkFilePath = _dataPath.resolve(chunkFilename);
            final int[] fromChunkPos = computeFrom(chunkIndex, from, true);
            final Array sourceChunk = _chunkReaderWriter.read(chunkFilePath);
            Partial2dDataCopier.copy(fromChunkPos, sourceChunk, target);
        }

    }

    private int[] computeFrom(int[] chunkIndex, int[] to, boolean read) {
        int[] from = {0, 0};
        for (int i = 0; i < chunkIndex.length; i++) {
            int index = chunkIndex[i];
            from[i] = index * _chunks[i];
            from[i] -= to[i];
        }
        if (read) {
            for (int i1 = 0; i1 < from.length; i1++) {
                from[i1] *= -1;
            }
        }
        return from;
    }
}







@@ -1895,240 +1895,387 @@
     
-synchronized (_locks) %7B%0A                if (!_locks.containsKey(chunkFilename)) %7B%0A                    _locks.put(chunkFilename, chunkFilename);%0A                %7D%0A            %7D%0A            final Path chunkFilePath = _dataPath.resolve
+final Path chunkFilePath = _dataPath.resolve(chunkFilename);%0A            final int%5B%5D fromBufferPos = computeFrom(chunkIndex, to, false);%0A            if (partialCopyingIsNotNeeded(bufferShape, fromBufferPos)) %7B%0A                _chunkReaderWriter.write(chunkFilePath, source);%0A            %7D else %7B%0A                synchronized (_locks) %7B%0A                    if (!_locks.containsKey
 (chu
@@ -2277,33 +2277,35 @@
 y(chunkFilename)
-;
+) %7B
 %0A            fin
@@ -2305,72 +2305,106 @@
     
-final int%5B%5D fromBufferPos = computeFrom(chunkIndex, to, false);%0A
+            _locks.put(chunkFilename, chunkFilename);%0A                    %7D%0A                %7D%0A    
     
@@ -2462,32 +2462,36 @@
                 
+    
 final Array targ
@@ -2532,32 +2532,36 @@
 chunkFilePath);%0A
+    
                 
@@ -2649,32 +2649,36 @@
                 
+    
 _chunkReaderWrit
@@ -2708,16 +2708,34 @@
  read);%0A
+                %7D%0A
         
@@ -2995,105 +2995,8 @@
 om);
-%0A        final Array target = NetCDF_Util.createArrayWithGivenStorage(targetBuffer, bufferShape);
 %0A%0A  
@@ -3368,79 +3368,827 @@
     
-Partial2dDataCopier.copy(fromChunkPos, sourceChunk, target);%0A        %7D%0A
+if (partialCopyingIsNotNeeded(bufferShape, fromChunkPos)) %7B%0A                System.arraycopy(sourceChunk.getStorage(), 0, targetBuffer, 0, (int) sourceChunk.getSize());%0A            %7D else %7B%0A                final Array target = NetCDF_Util.createArrayWithGivenStorage(targetBuffer, bufferShape);%0A                Partial2dDataCopier.copy(fromChunkPos, sourceChunk, target);%0A            %7D%0A        %7D%0A    %7D%0A%0A    public boolean partialCopyingIsNotNeeded(int%5B%5D bufferShape, int%5B%5D position) %7B%0A        return isZeroOffset(position) && isBufferShapeEqualChunkShape(bufferShape);%0A    %7D%0A%0A    public boolean isBufferShapeEqualChunkShape(int%5B%5D bufferShape) %7B%0A        return Arrays.equals(bufferShape, _chunks);%0A    %7D%0A%0A    public boolean isZeroOffset(int%5B%5D position) %7B%0A        return Arrays.equals(position, new int%5Bposition.length%5D);
 %0A   

improve chunk reading and writing by checking of cases, where partial copying of data is not needed