from __future__ import print_function

import openpathsampling as paths
from openpathsampling.netcdfplus import (
    NetCDFPlus, 
    ObjectStore, 
    StorableObject,
    NamedObjectStore,
    UniqueNamedObjectStore,
    DictStore,
    ImmutableDictStore,
    VariableStore,
    StorableNamedObject
)
import numpy as np

# force numpy print options for test comparison
np.set_printoptions(precision=6, formatter={'float_kind': lambda x: "{:.6f}".format(x)})

class Node(StorableObject):
    def __init__(self, value):
        super(Node, self).__init__()
        self.value = value
        
    def __repr__(self):
        return 'Node(%s)' % self.value

class NamedNode(StorableNamedObject):
    def __init__(self, value):
        super(NamedNode, self).__init__()
        self.value = value
        
    def __repr__(self):
        return 'Node(%s)' % self.value

"""
Explanation: test netcdf+
This is a more extensive integration test, if all the features of netcdf+ work as expected.
End of explanation
"""


st = NetCDFPlus('test_netcdfplus.nc', mode='w')

"""
Explanation: Open new storage
try to create a new storage
End of explanation
"""


class NodeIntStore(VariableStore):
    def __init__(self):
        super(NodeIntStore, self).__init__(Node, ['value'])
    
    def initialize(self):
        super(VariableStore, self).initialize()

        # Add here the stores to be imported
        self.create_variable('value', 'int')    

st.create_store('nodesnamed', NamedObjectStore(NamedNode))
st.create_store('nodesunique', UniqueNamedObjectStore(NamedNode))
st.create_store('dict', DictStore())
st.create_store('dictimmutable', ImmutableDictStore())
st.create_store('varstore', NodeIntStore())

"""
Explanation: Create some stores
End of explanation
"""


st.create_store('nodes', ObjectStore(Node))

# NBVAL_IGNORE_OUTPUT
# lazy because output checking fails in Py3k tests -- why is that?
print(st.find_store(Node))

"""
Explanation: And the default store. The last store for a particular object is used as the default if no specific store is specified.
End of explanation
"""


st.finalize_stores()

v = st.variables['nodes_uuid']

v.chunking()

print(st.find_store(Node))

st.nodes.save(Node(10));

st.close()

"""
Explanation: Initialize the store
End of explanation
"""


st = NetCDFPlus('test_netcdfplus.nc', mode='a')

"""
Explanation: Reopen empty storage
End of explanation
"""


for store in st.stores:
    store.set_caching(10)

"""
Explanation: set caching of the new stores
End of explanation
"""


assert('nodes' in st.objects)

assert('stores' in st.objects)

assert(len(st.nodes) == 1)

assert(len(st.stores) == 7)

for store in st.stores:
    print('{:40} {:30}'.format(str(store), str(store.cache)))

"""
Explanation: Check if the stores were correctly loaded
End of explanation
"""


# NBVAL_IGNORE_OUTPUT
print(sorted(st.get_var_types()))

"""
Explanation: Create variables types
Get a list of all possible variable types
End of explanation
"""


st.create_dimension('pair', 2)

for var_type in st.get_var_types():
    st.create_variable(var_type, var_type, dimensions=('pair', 'pair', 'pair'))

st.update_delegates()

# NBVAL_IGNORE_OUTPUT
for var_name, var in sorted(st.variables.items()):
    print(var_name, var.dimensions)

for var in sorted(st.vars):
    print(var)

"""
Explanation: Make a dimension on length 2 to simplify dimension nameing.
Now we construct for each type a corresponding variable of dimensions 2x2x2.
End of explanation
"""


st.vars['bool'][:] = True

print(st.vars['bool'][:])

"""
Explanation: Bool
End of explanation
"""


st.vars['float'][1,1] = 1.0

print(st.vars['float'][:])

"""
Explanation: Float
End of explanation
"""


st.vars['index'][0,1,0] = 10
st.vars['index'][0,1,1] = -1
st.vars['index'][0,0] = None

print(st.vars['index'][0,1])
print(st.vars['index'][0,0])

"""
Explanation: Index
Index is special in the sense that it supports only integers that are non-negative. Negative values will be interpreted as None
End of explanation
"""


st.vars['int'][0,1,0] = 10
st.vars['int'][0,1,1] = -1

print(st.vars['int'][:])

"""
Explanation: Int
End of explanation
"""


st.vars['json'][0,1,1] = {'Hallo': 2, 'Test': 3}

# NBVAL_IGNORE_OUTPUT
print(st.vars['json'][0,1,1])

st.vars['json'][0,1,0] = Node(10)

# NBVAL_IGNORE_OUTPUT
print(st.variables['json'][0,1,:])

"""
Explanation: JSON
The variable type JSON encode the given object as a JSON string in the shortest possible way. This includes using referenes to storable objects.
End of explanation
"""


nn = Node(10)
st.vars['jsonobj'][1,0,0] = nn

print(st.variables['jsonobj'][1,0,0])

st.vars['jsonobj'][1,0,0]

"""
Explanation: All object types registered as being Storable by subclassing from openpathsampling.base.StorableObject.
JSONObj
A JSON serializable object. This can be normal very simple python objects, plus numpy arrays, and objects that implement to_dict and from_dict. This is almost the same as JSON except if the object to be serialized is a storable object itself, it will not be referenced but the object itself will be turned into a JSON representation.
End of explanation
"""


st.vars['numpy.float32'][:] = np.ones((2,2,2)) * 3.0
st.vars['numpy.float32'][0] = np.ones((2,2)) * 7.0

print(st.vars['numpy.float32'][:])

"""
Explanation: Numpy
End of explanation
"""


st.vars['obj.nodes'][0,0,0] = Node(1)
st.vars['obj.nodes'][0,1,0] = Node('Second')
st.vars['obj.nodes'][0,0,1] = Node('Third')

# st.vars['obj.nodes'][1] = Node(20)

# NBVAL_IGNORE_OUTPUT
print(st.variables['obj.nodes'][:])
print(st.variables['nodes_json'][:])

print(st.vars['obj.nodes'][0,0,0])
print(type(st.vars['obj.nodes'][0,0,0]))

"""
Explanation: Obj
You can store objects of a type which you have previously added. For loading you need to make sure that the class (and the store if set manually) is present when you load from the store.
End of explanation
"""


st.vars['lazyobj.nodes'][0,0,0] = Node('First')

"""
Explanation: lazy
Lazy loading will reconstruct an object using proxies. These proxies behave almost like the loaded object, but will delay loading of the object until it is accessed. Saving for lazy objects is the same as for regular objects. Only loading return a proxy object.
End of explanation
"""


# NBVAL_IGNORE_OUTPUT
proxy = st.vars['lazyobj.nodes'][0,0,0]
print('Type:   ', type(proxy))
print('Class:  ', proxy.__class__)
print('Content:', proxy.__subject__.__dict__)
print('Access: ', proxy.value)

"""
Explanation: The type of the returned object is LoaderProxy while the class is the actual class is the baseclass loaded by the store to not trigger loading when the __class__ attribute is accessed. The actual object can be accessed by __subject__ and doing so will trigger loading the object. All regular attributes will be delegated to __subject__.attribute and also trigger loading.
End of explanation
"""


print(st.nodes[:])

obj = Node('BlaBla')
st.nodes.save(obj);

"""
Explanation: Load/Save objects
Note that there are now 6 Node objects.
End of explanation
"""


print(len(st.nodes))
obj = Node('BlaBlaBla')
st.save(obj)
print(len(st.nodes))

"""
Explanation: Saving without specifying should use store nodes which was defined last.
End of explanation
"""


print(st.idx(obj))

"""
Explanation: Get the index of the obj in the storage
End of explanation
"""


print(st.nodes.variables['json'][st.idx(obj)])

"""
Explanation: And test the different ways to access the contained json
1. direct json using variables in the store
End of explanation
"""


print(st.variables['nodes_json'][st.idx(obj)])

"""
Explanation: 2. direct json using variables in the full storage
End of explanation
"""


print(st.nodes.vars['json'][st.idx(obj)])
print(st.nodes.vars['json'][st.idx(obj)] is obj)

"""
Explanation: 3. indirect json and reconstruct using vars in the store
End of explanation
"""


print(st.nodes[st.idx(obj)])
print(st.nodes[st.idx(obj)] is obj)

"""
Explanation: 4. using the store accessor  __getitem__ in the store
End of explanation
"""


n = NamedNode(3)

"""
Explanation: One importance difference is that a store like nodes has a cache (which we set to 10 before). Using vars will not use a store and hence create a new object!
ObjectStores
ObjectStores are resposible to save and load objects. There are now 6 types available.
ObjectStore
The basic store which we have used before
NamedObjectStore
Supports to give objects names
End of explanation
"""


print(n.name)

"""
Explanation: NamedObjects have a .name property, which has a default.
End of explanation
"""


n.name = 'OneNode'
print(n.name)
n.name = 'MyNode'
print(n.name)

"""
Explanation: and can be set.
End of explanation
"""


st.nodesnamed.save(n);

try:
    n.name = 'NewName'
except ValueError as e:
    print('# We had an exception')
    print(e)
else:
    raise RuntimeWarning('This should have produced an error')

"""
Explanation: Once the object is saved, the name cannot be changed anymore.
End of explanation
"""


n2 = NamedNode(9)
n2.name = 'MyNode'
st.nodesnamed.save(n2);

"""
Explanation: usually names are not unique (see next store). So we can have more than one object with the same name.
End of explanation
"""


# NBVAL_IGNORE_OUTPUT
print(st.nodesnamed.name_idx)

"""
Explanation: See the list of named objects
End of explanation
"""


st.nodesunique.save(n);

"""
Explanation: UniqueNamedObjectStore
The forces names to be unique
End of explanation
"""


try:
    st.nodesunique.save(n2)
except RuntimeWarning as e:
    print('# We had an exception')
    print(e)
else:
    raise RuntimeWarning('This should have produced an error')


"""
Explanation: Note here that an object can be store more than once in a storage, but only if more than one store supports the file type.
End of explanation
"""


# NBVAL_IGNORE_OUTPUT
print(st.nodesunique.name_idx)

"""
Explanation: As said before this can only happen if you have more than one store for the same object type.
End of explanation
"""


n3 = NamedNode(10)
n4 = NamedNode(12)
st.nodesunique.save(n3);
st.nodesunique.save(n4);

n5 = NamedNode(1)
n5.name = 'MyNode'

try:
    st.nodesunique.save(n5)
except RuntimeWarning as e:
    print('# We had an exception')
    print(e)
else:
    raise RuntimeWarning('This should have produced an error')

"""
Explanation: some more tests. First saving onnamed objects. This is okay. Only given names should be unique.
End of explanation
"""


st.nodesunique.save(n5, 'NextNode');

n6 = NamedNode(1)
n6.name = 'SecondNode'

try:
    st.nodesunique.save(n6, 'MyNode')
except RuntimeWarning as e:
    print('# We had an exception')
    print(e)
else:
    raise RuntimeWarning('This should have produced an error')

"""
Explanation: This works since it does a rename before saving.
End of explanation
"""


print(dict(st.dict))
print(st.dict.name_idx)

n1 = NamedNode(1)
n2 = NamedNode(2)
n3 = NamedNode(3)
st.dict['Number1'] = n1

for key in sorted(st.dict):
    obj = st.dict[key]
    idxs = sorted(st.dict.name_idx[key])
    print(key, ':', str(obj), idxs)

st.dict['Number2'] = n2

for key in sorted(st.dict):
    obj = st.dict[key]
    idxs = sorted(st.dict.name_idx[key])
    print(key, ':', str(obj), idxs)

st.dict['Number1'] = n3

for key in sorted(st.dict):
    obj = st.dict[key]
    idxs = sorted(st.dict.name_idx[key])
    print(key, ':', str(obj), idxs)

print(st.dict['Number1'])

print(st.dict.find('Number1'))

# NBVAL_IGNORE_OUTPUT
print('[', ', '.join(st.dict.variables['json'][:]), ']')

for key in sorted(st.dict):
    obj = st.dict[key]
    idxs = sorted(st.dict.name_idx[key])
    print(key, ':', str(obj), idxs)

"""
Explanation: DictStore
A dictstore works a like a dictionary on disk. The content is returned using dict()
End of explanation
"""


try:
    st.dictimmutable['Number1'] = n1
    st.dictimmutable['Number1'] = n2
except RuntimeWarning as e:
    print('# We had an exception')
    print(e)
else:
    raise RuntimeWarning('This should have produced an error')

"""
Explanation: ImmutableDictStore
This adds the check that already used names cannot be used again
End of explanation
"""


a = Node(30)
st.varstore.save(a);

"""
Explanation: VariableStore
Store a node with an int as we defined for our VariableStore
End of explanation
"""


st.varstore.clear_cache()

"""
Explanation: clear the cache
End of explanation
"""


assert(st.varstore[0].value == 30)

"""
Explanation: And try loading
End of explanation
"""


try:
    a = Node('test')
    print(st.varstore.save(a))
except ValueError as e:
    print('# We had an exception')
    print(e)
else:
    raise RuntimeWarning('This should have produced an error')

"""
Explanation: Try storing non int() parseable value
End of explanation
"""


st_uuid = NetCDFPlus('test_netcdfplus_uuid.nc', mode='w')
st_uuid.create_store('nodes', ObjectStore(Node))
st_uuid.finalize_stores()
st_uuid.save(st.nodes[0])
st_uuid.close()

st.close()

st_fb = NetCDFPlus('test_netcdfplus_fb.nc', mode='w', fallback=NetCDFPlus('test_netcdfplus_uuid.nc'))
st_fb.create_store('nodes', ObjectStore(Node))
st_fb.finalize_stores()

st_fb.exclude_from_fallback

assert(st_fb.fallback.nodes[0] in st_fb.fallback)

assert(st_fb.fallback.nodes[0] in st_fb)

assert(st.nodes[0] in st_fb)

assert(st.nodes[0] in st_fb.fallback)

"""
Explanation: Test fallback
End of explanation
"""


# NBVAL_IGNORE_OUTPUT
print(hex(st_fb.nodes.save(st_fb.fallback.nodes[0])))

assert(len(st_fb.nodes) == 0)

assert(st_fb.fallback.nodes[0] in st_fb)
assert(st_fb.fallback.nodes[0] in st_fb.fallback)
assert(st.nodes[0] in st_fb)
assert(st.nodes[0] in st_fb.fallback)

st_fb.fallback.close()
st_fb.close()

"""
Explanation: Try saving object in fallback
End of explanation
"""


