#!/usr/bin/env python
# This file contains code to chop up a CSV file containing 3D information
# (e.g., as outputted by a Keyence VHX microscope) according to the objects
# identified by segment.
#
# NOTE: Assumes the CSV file does not contain headings or indices
#
# Written by Allison Hsiang

import numpy as np
import pandas
import os
from pathlib import Path
from progress.bar import Bar
import sys


def process_csv(object_list, csv_path, settings, image_size, image_dir_path=None):
    '''
    Loops through segment-identified objects and cuts provided
    corresponding CSV file at bounding box of each object.
    '''
    print('INFO: Loading CSV file...', end=' ', flush=True)
    csv_matrix = pandas.read_csv(csv_path, header=None).to_numpy()
    print(u'\u2713')

    if settings['transpose_csv']:
        csv_matrix = csv_matrix.transpose()

    # Check to ensure CSV and image sizes match (within 5 pixel buffer)
    print('INFO: Checking that CSV and image dimensions match...', end=' ', flush=True)
    csv_size = csv_matrix.shape
    buffer = 5
    if not (csv_size[0] in range(image_size[0]-buffer, image_size[0]+buffer) and csv_size[1] in range(image_size[1]-buffer, image_size[1]+buffer)):
        sys.exit('Error: CSV and image dimensions do not match! Exiting...')
    print(u'\u2713')

    print('INFO: Cutting CSV...')
    bar = Bar(suffix = '%(percent).1f%% - %(eta)ds', max=len(object_list))
    for obj_num, box in enumerate(object_list):
        csv_object = cut(csv_matrix, box)

        if settings['single_plane']:
            csv_output_path = os.path.join(image_dir_path, 'csv_output')
        else:
            csv_output_path = os.path.join(settings['full_output'], 'csv_output')

        if settings['single_plane']:
            subpath = Path(image_dir_path).stem
        else:
            subpath = settings['unique_id']

        object_name = '%s_obj%05d' % (subpath, obj_num+1)
        save(csv_object, csv_output_path, object_name)

        bar.next()


def cut(matrix, indices):
    '''
    Subsets matrix by provided object indices.
    '''
    # [y1, x1, y2, x2]
    object = matrix[:,indices[1]:indices[3]][indices[0]:indices[2],:]
    return object


def save(object, output_path, object_name):
    '''
    Saves cropped CSV object to file.
    '''
    final_object_file = '%s.%s' % (object_name,'csv')
    final_path = os.path.join(output_path,final_object_file)
    pandas.DataFrame(object).to_csv(final_path, index=False, header=False)
