import sys
import numpy as np

# the following line is not required if BatchFlow is installed as a python package.
sys.path.append("../..")
from batchflow import Dataset, DatasetIndex, Batch

# number of items in the dataset
NUM_ITEMS = 10
# number of items in a batch when iterating
BATCH_SIZE = 3

"""
Explanation: How to generate batches from a dataset and work with batch components
End of explanation
"""


dataset = Dataset(index=NUM_ITEMS, batch_class=Batch)

"""
Explanation: Create a dataset
A dataset is defined by an index (a sequence of item ids) and a batch class (see the documentation for details).
In the simplest case an index is a natural sequence 0, 1, 2, 3, ...
So all you need to define the index is just a number of items in the dataset.
End of explanation
"""


for i, batch in enumerate(dataset.gen_batch(BATCH_SIZE, n_epochs=1)):
    print("batch", i, " contains items", batch.indices)

for i, batch in enumerate(dataset.gen_batch(BATCH_SIZE, n_iters=5)):
    print("batch", i, " contains items", batch.indices)

"""
Explanation: The dataset index
See the documentation for more info about how to create an index which fits your needs.
Here are the most frequent use cases:
client_index = DatasetIndex(my_client_ids)

images_index = FilesIndex(path="/path/to/images/*.jpg", no_ext=True)

Iterate with gen_batch(...)
gen_batch is a python generator.
End of explanation
"""


for i, batch in enumerate(dataset.gen_batch(BATCH_SIZE, n_epochs=1, drop_last=True)):
    print("batch", i, " contains items", batch.indices)

"""
Explanation: drop_last=True skips the last batch if it contains fewer than BATCH_SIZE items
End of explanation
"""


for i, batch in enumerate(dataset.gen_batch(BATCH_SIZE, n_iters=4, drop_last=True, shuffle=True)):
    print("batch", i, " contains items", batch.indices)

"""
Explanation: shuffle permutes items across batches
End of explanation
"""


for i, batch in enumerate(dataset.gen_batch(BATCH_SIZE, n_epochs=1, drop_last=True, shuffle=123)):
    print("batch", i, " contains items", batch.indices)

"""
Explanation: Run the cell above multiple times to see how batches change.
Shuffle can be bool, int (seed number) or a RandomState object
End of explanation
"""


for i in range(NUM_ITEMS * 3):
    try:
        batch = dataset.next_batch(BATCH_SIZE, shuffle=True, n_epochs=2, drop_last=True)
        print("batch", i + 1, "contains items", batch.indices)
    except StopIteration:
        print("got StopIteration")
        break

"""
Explanation: Run the cell above multiple times to see that batches stay the same across runs.
Iterate with next_batch(...)
While gen_batch is a generator, next_batch is an ordinary method.
Most of the time you will use gen_batch, but for a deeper control over training and a more sophisticated finetuning next_batch might be more convenient.
If too many iterations are made, StopIteration will be raised.
Check that there are NUM_ITEMS * 3 iterations (i.e. 3 epochs) in loop, but n_epochs=2 is specified inside next_batch() call.
End of explanation
"""


dataset.reset('iter')

"""
Explanation: And finally with shuffle=True, n_epochs=None and a variable batch size
Do not forget to reset iterator to start next_batch'ing from scratch
End of explanation
"""


for i in range(int(NUM_ITEMS * 1.3)):
    batch = dataset.next_batch(BATCH_SIZE + (-1)**i * i % 3, shuffle=True, n_epochs=None, drop_last=True)
    print("batch", i + 1, "contains items", batch.indices)

"""
Explanation: n_epochs=None allows for infinite iterations.
End of explanation
"""


data = np.arange(NUM_ITEMS).reshape(-1, 1) * 100 + np.arange(3).reshape(1, -1)
data

"""
Explanation: To get a deeper understanding of drop_last read very important notes in the API.
Working with data
For illustrative purposes let's create a small array which will serve as a raw data source.
End of explanation
"""


for batch in dataset.gen_batch(BATCH_SIZE, n_epochs=1):
    batch = batch.load(src=data)
    print("batch contains items with indices", batch.indices)
    print('and batch data is')
    print(batch.data)
    print()

"""
Explanation: Load data into a batch
After loading data is available as batch.data
End of explanation
"""


for batch in dataset.gen_batch(BATCH_SIZE, n_epochs=1):
    batch = batch.load(src=data)
    print("batch contains")
    for item in batch:
        print(item)
    print()

"""
Explanation: You can easily iterate over batch items too
End of explanation
"""


class MyBatch(Batch):
    components = 'features', 'labels'

"""
Explanation: Data components
Not infrequently, the batch stores a more complex data structures, e.g. features and labels or images, masks, bounding boxes and labels. To work with these you might employ data components. Just define a property as follows:
End of explanation
"""


features_array = np.arange(NUM_ITEMS).reshape(-1, 1) * 100 + np.arange(3).reshape(1, -1)
labels_array = np.random.choice(10, size=NUM_ITEMS)
data = features_array, labels_array

"""
Explanation: Let's generate some random data:
End of explanation
"""


dataset = Dataset(index=NUM_ITEMS, batch_class=MyBatch, preloaded=data)

"""
Explanation: Now create a dataset (preloaded handles data loading from data stored in memory)
End of explanation
"""


for i, batch in enumerate(dataset.gen_batch(BATCH_SIZE, n_epochs=1)):
    print("batch", i, " contains items", batch.indices)
    print("and batch data consists of features:")
    print(batch.features)
    print("and labels:", batch.labels)
    print()

"""
Explanation: Since components are defined, you can address them as batch and even item attributes (they are created and loaded automatically).
End of explanation
"""


for i, batch in enumerate(dataset.gen_batch(BATCH_SIZE, n_epochs=1)):
    print("Batch", i)
    for item in batch:
        print("item features:", item.features, "    item label:", item.labels)

    print()
    print("You can change batch data, even scalars.")
    for item in batch:
        item.features = item.features + 1000
        item.labels = item.labels + 100
    print("New batch features:\n", batch.features)
    print("and labels:", batch.labels)
    print()

"""
Explanation: You can iterate over batch items and change them on the fly
End of explanation
"""


dataset.split(0.8)

"""
Explanation: Splitting a dataset
For machine learning tasks you might need to split a dataset into train, test and validation parts.
End of explanation
"""


len(dataset.train), len(dataset.test)

dataset.split([.6, .2, .2])

len(dataset.train), len(dataset.test), len(dataset.validation)

"""
Explanation: Now the dataset is split into train / test in 80/20 ratio.
End of explanation
"""


dataset.split(0.7, shuffle=True)

dataset.train.indices, dataset.test.indices

"""
Explanation: Dataset may be shuffled before splitting.
End of explanation
"""


