#!/usr/bin/env python

#######################################################
# Copyright (c) 2019, ArrayFire
# All rights reserved.
#
# This file is distributed under 3-clause BSD license.
# The complete license agreement can be obtained at:
# http://arrayfire.com/licenses/BSD-3-Clause
########################################################

import os
import sys
sys.path.insert(0, '../common')
from idxio import read_idx

import arrayfire as af
from arrayfire.algorithm import where
from arrayfire.array import Array
from arrayfire.data import constant, lookup, moddims
from arrayfire.random import randu


def classify(arr, k, expand_labels):
    ret_str = ''
    if expand_labels:
        vec = arr[:, k].as_type(af.Dtype.f32)
        h_vec = vec.to_list()
        data = []

        for i in range(vec.elements()):
            data.append((h_vec[i], i))

        data = sorted(data, key=lambda pair: pair[0], reverse=True)

        ret_str = str(data[0][1])

    else:
        ret_str = str(int(arr[k].as_type(af.Dtype.f32).scalar()))

    return ret_str


def setup_mnist(frac, expand_labels):
    root_path = os.path.dirname(os.path.abspath(__file__))
    file_path = root_path + '/../../assets/examples/data/mnist/'
    idims, idata = read_idx(file_path + 'images-subset')
    ldims, ldata = read_idx(file_path + 'labels-subset')

    idims.reverse()
    numdims = len(idims)
    images = af.Array(idata, tuple(idims))

    R = af.randu(10000, 1);
    cond = R < min(frac, 0.8)
    train_indices = af.where(cond)
    test_indices = af.where(~cond)

    train_images = af.lookup(images, train_indices, 2) / 255
    test_images = af.lookup(images, test_indices, 2) / 255

    num_classes = 10
    num_train = train_images.dims()[2]
    num_test = test_images.dims()[2]

    if expand_labels:
        train_labels = af.constant(0, num_classes, num_train)
        test_labels = af.constant(0, num_classes, num_test)

        h_train_idx = train_indices.to_list()
        h_test_idx = test_indices.to_list()

        for i in range(num_train):
            train_labels[ldata[h_train_idx[i]], i] = 1

        for i in range(num_test):
            test_labels[ldata[h_test_idx[i]], i] = 1

    else:
        labels = af.Array(ldata, tuple(ldims))
        train_labels = labels[train_indices]
        test_labels = labels[test_indices]

    return (num_classes,
            num_train,
            num_test,
            train_images,
            test_images,
            train_labels,
            test_labels)


def display_results(test_images, test_output, test_actual, num_display, expand_labels):
    for i in range(num_display):
        print('Predicted: ', classify(test_output, i, expand_labels))
        print('Actual: ', classify(test_actual, i, expand_labels))

        img = (test_images[:, :, i] > 0.1).as_type(af.Dtype.u8)
        img = af.moddims(img, img.elements()).to_list()
        for j in range(28):
            for k in range(28):
                print('\u2588' if img[j * 28 + k] > 0 else ' ', end='')
            print()
        input()