import cv2
import numpy as np
import os
import scipy.io as sio


def write_files(mat_input_file, output_dir):
    if not os.path.exists(mat_input_file):
        raise Exception('Path does not exist: ' + mat_input_file)

    if not os.path.exists(output_dir):
        os.mkdir(output_dir)

    data = sio.loadmat(mat_input_file)
    counts = np.empty((data['X'].shape[3],)).astype(np.uint8)

    for i in range(data['X'].shape[3]):
        truth = data['y'][i, 0]

        if truth == 10:
            truth = 0

        counts[i] = truth
        img = data['X'][..., i]
        cv2.imwrite(os.path.join(output_dir, str(i + 1) + '.png'), img)

    np.savetxt(os.path.join(output_dir, 'labels.csv'), counts, delimiter=",", fmt='%d')