import pytest

import copy
import gzip
import json
import numpy as np
import os
import shutil
import sys

from cloudvolume import exceptions
from cloudvolume.exceptions import AlignmentError, ReadOnlyException
from cloudvolume import CloudVolume, chunks
from cloudvolume.lib import Bbox, Vec, yellow, mkdir
import cloudvolume.sharedmemory as shm
from layer_harness import (
  TEST_NUMBER,  
  delete_layer, create_layer
)
from cloudvolume.datasource.precomputed.common import cdn_cache_control
from cloudvolume.datasource.precomputed.image.tx import generate_chunks

@pytest.fixture
def shard_vol():
  test_dir = os.path.dirname(os.path.abspath(__file__))
  test_dir = os.path.join(test_dir, "test_cv_sharded")
  return CloudVolume("file://" + test_dir)

@pytest.fixture
def shard_vol_data_cpso():
  import compresso
  test_dir = os.path.dirname(os.path.abspath(__file__))
  test_dir = os.path.join(test_dir, "test_cv_sharded")
  with open(os.path.join(test_dir, "connectomics_mip_1.npy.cpso"), "rb") as f:
    data = f.read()
  return compresso.CompressoArray(data)

def test_from_numpy():
  arr = np.random.randint(0, high=256, size=(128,128, 128))
  arr = np.asarray(arr, dtype=np.uint8)
  vol = CloudVolume.from_numpy(arr, max_mip=1)
  arr2 = vol[:,:,:]
  np.alltrue(arr == arr2)
  
  arr = np.random.randn(128,128, 128, 3)
  arr = np.asarray(arr, dtype=np.float32)
  vol = CloudVolume.from_numpy(arr, max_mip=1)
  arr2 = vol[:,:,:]
  np.alltrue(arr == arr2)
  shutil.rmtree('/tmp/image')

def test_cloud_access():
  CloudVolume('gs://seunglab-test/test_v0/image')
  CloudVolume('s3://seunglab-test/test_dataset/image')

def test_fill_missing():
  info = CloudVolume.create_new_info(
    num_channels=1, # Increase this number when we add more tests for RGB
    layer_type='image', 
    data_type='uint8', 
    encoding='raw',
    resolution=[ 1,1,1 ], 
    voxel_offset=[0,0,0], 
    volume_size=[128,128,64],
    mesh='mesh', 
    chunk_size=[ 64,64,64 ],
  )

  vol = CloudVolume('file:///tmp/cloudvolume/empty_volume', mip=0, info=info)
  vol.commit_info()

  vol.cache.flush()

  vol = CloudVolume('file:///tmp/cloudvolume/empty_volume', mip=0, fill_missing=True)
  assert np.count_nonzero(vol[:]) == 0

  vol = CloudVolume('file:///tmp/cloudvolume/empty_volume', mip=0, fill_missing=True, cache=True)
  assert np.count_nonzero(vol[:]) == 0
  assert np.count_nonzero(vol[:]) == 0

  vol.cache.flush()
  delete_layer('/tmp/cloudvolume/empty_volume')

def test_background_color():
  info = CloudVolume.create_new_info(
    num_channels=1, 
    layer_type='image', 
    data_type='uint8', 
    encoding='raw',
    resolution=[ 1,1,1 ], 
    voxel_offset=[0,0,0], 
    volume_size=[128,128,1],
    mesh='mesh', 
    chunk_size=[ 64,64,1 ],
  )

  vol = CloudVolume('file:///tmp/cloudvolume/empty_volume', mip=0, info=info)
  vol.commit_info()

  vol.cache.flush()

  vol = CloudVolume('file:///tmp/cloudvolume/empty_volume', 
                    mip=0, 
                    background_color=1, 
                    fill_missing=True)
  assert np.count_nonzero(vol[:] - 1) == 0

  vol = CloudVolume('file:///tmp/cloudvolume/empty_volume', 
                    mip=0, 
                    background_color=1, 
                    fill_missing=True,
                    bounded=False)
  assert np.count_nonzero(vol[0:129,0:129,0:1]-1) == 0

  vol = CloudVolume('file:///tmp/cloudvolume/empty_volume', 
                    mip=0, 
                    background_color=1, 
                    fill_missing=True,
                    bounded=False,
                    parallel=2)
  assert np.count_nonzero(vol[0:129,0:129,0:1]-1) == 0
  vol.cache.flush()
  delete_layer('/tmp/cloudvolume/empty_volume')

def test_has_data():
  delete_layer()
  cv, data = create_layer(size=(50,50,50,1), offset=(0,0,0))
  cv.add_scale((2,2,1))

  assert cv.image.has_data(0) == True
  assert cv.image.has_data(1) == False
  assert cv.image.has_data([1,1,1]) == True
  assert cv.image.has_data([2,2,1]) == False

  try:
    cv.image.has_data(2)
    assert False
  except exceptions.ScaleUnavailableError:
    pass

def image_equal(arr1, arr2, encoding):
  assert arr1.shape == arr2.shape
  if encoding == 'raw':
    return np.all(arr1 == arr2)

  # jpeg
  u1 = arr1.flatten().mean()
  u2 = arr2.flatten().mean()
  return np.abs(u1 - u2) < 1

@pytest.mark.parametrize('green', (True, False))
@pytest.mark.parametrize('encoding', ('raw', 'jpeg'))
@pytest.mark.parametrize('lru_bytes', (0,1e6,10e6))
def test_aligned_read(green, encoding, lru_bytes):
  delete_layer()
  cv, data = create_layer(size=(50,50,50,1), offset=(0,0,0), encoding=encoding)
  cv.green_threads = green
  cv.image.lru.resize(lru_bytes)
  # the last dimension is the number of channels
  assert cv[0:50,0:50,0:50].shape == (50,50,50,1)
  assert image_equal(cv[0:50,0:50,0:50], data, encoding)
  
  delete_layer()
  cv, data = create_layer(size=(128,64,64,1), offset=(0,0,0), encoding=encoding)
  cv.green_threads = green
  # the last dimension is the number of channels
  assert cv[0:64,0:64,0:64].shape == (64,64,64,1) 

  assert image_equal(cv[0:64,0:64,0:64], data[:64,:64,:64,:], encoding)

  delete_layer()
  cv, data = create_layer(size=(128,64,64,1), offset=(10,20,0), encoding=encoding)
  cv.green_threads = green
  cv.image.lru.resize(lru_bytes)
  cutout = cv[10:74,20:84,0:64]
  # the last dimension is the number of channels
  assert cutout.shape == (64,64,64,1) 
  assert image_equal(cutout, data[:64,:64,:64,:], encoding)
  # get the second chunk
  cutout2 = cv[74:138,20:84,0:64]
  assert cutout2.shape == (64,64,64,1) 
  assert image_equal(cutout2, data[64:128,:64,:64,:], encoding)

  assert cv[25, 25, 25].shape == (1,1,1,1)

@pytest.mark.parametrize('green', (True, False))
@pytest.mark.parametrize('encoding', ('raw', 'compresso', 'compressed_segmentation'))
@pytest.mark.parametrize('lru_bytes', (0,1e6,10e6))
def test_point_reads_sharded(shard_vol, shard_vol_data_cpso, green, encoding, lru_bytes):
  cv = shard_vol
  data = shard_vol_data_cpso
  cv.green_threads = green
  cv.image.lru.resize(lru_bytes)

  N = 25

  x = np.random.randint(cv.bounds.minpt[0]+1, cv.bounds.maxpt[0]-1, size=(200,))
  y = np.random.randint(cv.bounds.minpt[1]+1, cv.bounds.maxpt[1]-1, size=(200,))
  z = np.random.randint(cv.bounds.minpt[2]+1, cv.bounds.maxpt[2]-1, size=(200,))
  pts = np.stack([x,y,z]).T

  for pt in pts:
    assert cv[tuple(pt)] == data[tuple(pt)]

@pytest.mark.parametrize('green', (True, False))
@pytest.mark.parametrize('encoding', ('raw', 'compresso', 'compressed_segmentation'))
@pytest.mark.parametrize('lru_bytes', (0,1e6,10e6))
def test_point_reads(green, encoding, lru_bytes):
  delete_layer()
  cv, data = create_layer(
    size=(256,256,100,1), offset=(0,0,0), 
    encoding=encoding, layer_type="segmentation"
  )
  cv.green_threads = green
  cv.image.lru.resize(lru_bytes)  

  N = 200

  x = np.random.randint(cv.bounds.minpt[0]+1, cv.bounds.maxpt[0]-1, size=(200,))
  y = np.random.randint(cv.bounds.minpt[1]+1, cv.bounds.maxpt[1]-1, size=(200,))
  z = np.random.randint(cv.bounds.minpt[2]+1, cv.bounds.maxpt[2]-1, size=(200,))
  pts = np.stack([x,y,z]).T

  for pt in pts:
    assert cv[tuple(pt)] == data[tuple(pt)]

  cv[:] = 62

  for pt in pts:
    assert cv[tuple(pt)] == 62

  cv.config.parallel = 2
  cv[:] = 17  

  for pt in pts:
    assert cv[tuple(pt)] == 17

def test_save_images():
  delete_layer()
  cv, data = create_layer(size=(50,50,50,1), offset=(0,0,0))  

  img = cv[:]
  directory = img.save_images()

  for z, fname in enumerate(sorted(os.listdir(directory))):
    assert fname == str(z).zfill(3) + '.png'

  shutil.rmtree(directory)

def test_bbox_read():
  delete_layer()
  cv, data = create_layer(size=(50,50,50,1), offset=(0,0,0))

  x = Bbox( (0,1,2), (48,49,50) )
  # the last dimension is the number of channels
  assert cv[x].shape == (48,48,48,1)
  assert np.all(cv[x] == data[0:48, 1:49, 2:50])  

def test_number_type_read():
  delete_layer()
  cv, data = create_layer(size=(50,50,50,1), offset=(0,0,0))

  for datatype in (