#!/usr/bin/env python

'''
This is a script to run OCR on all the study folders in the input directory.
It replicates the directory structure in the output directory.
- The output directory structure also has info.csv in each study folder which lists the image type, and tag for the file
- The tags are read from a specified bounding box in the ultrasound image
- For cines, the tag is extracted from the middle frame of the video
- If the directory structure already exists, this won't overwrite the results
- Parallel processing is used whenever possible to speed up the process.

Author: Hina Shah
'''

from pathlib import Path
import os
import csv
from argparse import ArgumentParser
from pprint import pprint, pformat
import logging
import sys
import gc
import concurrent.futures
import time
import json
import numpy as np
import shutil

from FamliImgIO import dcmio
from FamliFileIO import taginfo
from FamliOCR import famlitesseract as tess
import utils

def getStudyOutputFolder(study, data_folder, out_parent_folder):
    subject_id = study.name 
    study_path = study.parent
    study_path_relative = study_path.relative_to(data_folder)
    out_dir = out_parent_folder  / study_path_relative / subject_id
    return out_dir

def extractImageArrayFromUS(file_path, out_dir=None):
    dcmobj = dcmio.DCMIO(file_path)
    # Create the output path
    out_path = None
    if out_dir is not None:
        out_path = out_dir / (file_path.stem + '.jpg')
    
    np_frame = dcmobj.get_repres_frame(out_path)
    us_type = dcmobj.get_type()
    us_model = dcmobj.get_model()
    n_rows  = dcmobj.get_numrows()
    n_cols = dcmobj.get_numcolumns()
    n_frames = dcmobj.get_numframes()
    del dcmobj
    
    return np_frame, us_type, us_model, n_rows, n_cols, n_frames

def extractTagForStudy(study, data_folder, out_images_dir, tag_list, non_tag_us, 
                        tag_bounding_box, server_path, greedy=False):
    
    logging.info("=========== PROCESSING SUBJECT: {} ===============".format(study.name))
    out_dir = getStudyOutputFolder(study, data_folder, out_images_dir)
    utils.checkDir(out_dir, delete = (not greedy))
    
    lock_file = out_dir / 'ocr_lock'
    if lock_file.exists():
        logging.info("SUBJECT is being processed by some other process, skipping")
        print("Study {} is being processed by some other process, skipping".format(study))
        return None
    else:
        lock_file.touch()

    # look for the info.csv file in the output folder. 
    # if exists, read the file names and the tags
     
    i=1
    file_paths = list( study.glob('**/*.dcm') )
    csvrows=[]
    unknown = 0
    tag_manager = taginfo.TagInfoFile(out_dir)
    tag_statistic = {}
    prev_tags = {}
    if greedy and tag_manager.exists():
        # If this is in greedy mode, try to get the names of the files
        # that have an undecided tag:
        tag_manager.read()
        # get a copy of the previously created tags for all files
        prev_tags = tag_manager.getDictWithNameKeys()
        # clear the tags stored in the file. (This does not delete the file, just empties the data structure)
        tag_manager.clear()
        tag_manager.deleteTagFile()
    elif tag_manager.exists():
        # File already exists, return.
        tag_manager.read()
        tag_statistic = tag_manager.tag_statistic
        log.info("Tag file {} already exists, skipping.".format(tag_manager.getFilePath()))
        print("Tag file {} already exists, skipping.".format(tag_manager.getFilePath()))
        lock_file.unlink()
        return tag_statistic
    for file_path in file_paths:
        logging.debug("FILE {}: {}".format(i, str(file_path)))

        # Make sure path exists
        if not file_path.exists():
            logging.info('File: {} does not exist, skipping'.format(file_path))
            continue

        # if in greedy mode, look at the tag, look at the tag and decide if we wamt to extract tag
        getTag = True
        tag = None
        us_type = None
        if greedy and len(prev_tags) > 0:
            name = file_path.name
            if name in prev_tags:
                tag = prev_tags[name][1]
                us_type = prev_tags[name][0]
                if tag in ['Unknown', 'Undecided', 'No tag']:
                    getTag = True
                else:
                     getTag = False

        if getTag:
            start = time.time()
            # Get the representative frame
            np_frame, us_type, capture_model, n_rows, n_cols, n_frames = extractImageArrayFromUS(file_path, out_dir=out_dir)
            end = time.time()
            logging.debug('Preprocessing took : {} seconds'.format(end-start))
            if len(capture_model)>0 and capture_model not in tag_bounding_box.keys():
                logging.info('US Model: {} not supported for file: {}'.format(capture_model, file_path))
                del np_frame
                continue
            # Extract the tag
            start = time.time()
            tag = 'Unknown'
            if np_frame is None and us_type not in non_tag_us:
                logging.warning('Failed to read file: {}, type: {}'.format(file_path, us_type))
                print('Failed to read file: {}, type: {}'.format(file_path, us_type))
            if np_frame is not None and \
                us_type not in non_tag_us and \
                capture_model in tag_bounding_box.keys():
                # Run tag extraction
                tag = tess.extractTagFromFrame(np_frame, tag_bounding_box[capture_model], tag_list, greedy=greedy)
            end = time.time()
            logging.debug('Tag extraction took : {} seconds'.format(end-start))
            del np_frame
            if tag in ['Unknown', 'Undecided', 'No tag']:
                unknown += 1
        else:
            logging.debug('Skipping the file: {}, tag: {}, type: {}, because it was known'.format(file_path, tag, us_type))

        tag_manager.addTag(file_path.parent, server_path, file_path.name, us_type, tag, write=True)
        i+=1
        gc.collect()

    tag_statistic = tag_manager.tag_statistic
    # If all unknown, delete the tag file. 
    if len(file_paths) != tag_manager.getNumFiles():
        print('WARNING: Study {} did not process ALL files. processed: {}, actual: {}'.format(study, tag_manager.getNumFiles(), len(file_paths)))
    if unknown == len(file_paths):
        try:
            with open(out_images_dir/'no_tags_found_studies.txt', "a+", newline='') as f:
                f.write(out_dir.name + os.linesep)
        except Exception as e:
            print('Error writing to no tags found directory\n' + e)
        tag_manager.deleteTagFile()
        tag_manager.clear()
    

    if len(file_paths) != tag_manager.getNumFiles():
        logging.warning('Number of files does not match rows in tag file for study: {}'.format(study))
     
    try:
        finished_study_file = out_images_dir/'finished_studies.txt'
        with open(finished_study_file, "a+", newline='') as f:
            f.write(out_dir.name + os.linesep)
    except Exception as e:
        print('Error writing to finished studies\n'+e)
    
    
    lock_file.unlink()
    return tag_statistic

def main(args):
    data_folder = Path(args.dir)
    out_images_dir = Path(args.out_dir)

    utils.checkDir(out_images_dir, False)    
    utils.setupLogFile(out_images_dir, args.debug)

    studies = []
    prev_list_file = out_images_dir/"original_us_dirs.txt"
    do_os_walk = True if not prev_list_file.exists() else False
    if not do_os_walk:
        print("Not searching recursively for directories, instead trying to read from a previously generated file")
        try:
            with open(prev_list_file, 'r') as f:
                studies = [ Path(s.strip().replace('\\','/')) for s in f.readlines() ]
        except Exception as e:
            print("Error reading the originally created list of files, wil resort to directory search")
            do_os_walk = True
        else:
            print("Found {} studies in the previously created list".format(len(studies)))
            do_os_walk = len(studies) == 0
            for s in studies:
                if not s.is_dir():
                    print("Did not find study: {}".format(s))
                    do_os_walk = True
                    break
    if do_os_walk:
        print("Will try to walk the U/S directory to get the list of all directories")
        for dirname, dirnames, __ in os.walk(str(data_folder)):
            if len(dirnames) == 0:
                studies.append(Path(dirname))

    # Reversing the list to process the newer ones first
    studies.reverse()

    logging.info('Found {} studies '.format(len(studies)))
    print('Found {} studies '.format(len(studies)))
    
    # read the list of acceptable tags in the ultrasound file
    tag_list = utils.getTagsList()
    tag_statistic = dict.fromkeys(tag_list, 0)
    tag_statistic['Unknown'] = 0
    tag_statistic['Undecided'] = 0
    tag_statistic['No tag'] = 0
    
    # Approximate bounding box of where the tag is written acoording to the 
    # us model
    tag_bounding_box = { 'V830':[[40,75], [265,190]],
                         'LOGIQe':  [[0,55], [210,160]],
                         'Voluson S': [[40,75], [265,190]],
                         'LOGIQeCine': [[0,0],[145,90]],
                         'Turbo': [[75,20],[240,80]],
                         'Voluson E8': [[40,75], [265,250]],
                         'iQ Probe': [[0,0], [300, 300]],
                         'C3HD': [[0,0], [384,384]]
                        }

    # list of ultrasound image types whose tags we do not care about right now.
    non_tag_us = ['Unknown', 'Secondary capture image report',
                    'Comprehensive SR', '3D Dicom Volume', 'ge kretz image']

    # Also read in study directories that might have been finished by a previous run - do not want to rerun them again
    # WARNING: finished_studies.txt should only contain names of the directories, not full paths
    finished_study_file = out_images_dir/'finished_studies.txt'
    finished_studies = None
    if finished_study_file.exists():
        try:
            with open(finished_study_file) as f:
                finished_studies = f.read().splitlines()
                finished_studies = [study for study in finished_studies if study.strip()]
            shutil.copyfile(str(finished_study_file), str(finished_study_file).replace('.txt', '_bkup.txt'))
        except Exception as e:
            logging.warning("Error reading the finished studies file\n" + e)
            print("Error: reading the finished studies file.")
            finished_studies = None
    if finished_studies is not None:
        logging.info('Found {} finished studies'.format(len(finished_studies)))
        cleaned_studies = [study for study in studies if study.name not in finished_studies]
        # Get statistics for the finished studies
        for study in finished_studies:
            logging.info('Will skip: {}'.format(study))
            try:
                study_dt = study.split('_')[1] # Assuming format StudyID_YYYYMMDD_HHMMSS
                infocsv_dir = out_images_dir/Path(study_dt[0:4]+'-'+study_dt[4:6])/study # Assuming a directory structure
                logging.info('Opening: {}'.format(infocsv_dir))
                tag_file_man =taginfo.TagInfoFile(infocsv_dir)
                tag_file_man.read()
                if tag_file_man.getNumFiles() > 0:
                    for tag in tag_file_man.tag_statistic:
                        if tag not in tag_statistic:
                            tag_statistic[tag] = 0
                        tag_statistic[tag] += tag_file_man.tag_statistic[tag]
            except Exception as e:
                logging.warning('Error reading previously created tags.csv for subject: {} \n {}'.format(study, e))
    else:
        cleaned_studies = studies
    del studies

    logging.info("Will process {} studies".format(len(cleaned_studies)))
    print("Will process {} studies".format(len(cleaned_studies)))

    if args.dir_list:
        try:
            logging.info("Getting a subset of directories from {}".format(args.dir_list))
            with open(args.dir_list, 'r') as f:
                dir_list_to_process = [r.strip() for r in f.readlines()]
        except Exception as e:
            logging.warning("Couldn't read the directory list. \n" + e )
            print("Couldn't read the directory list, Exiting.")
            return
     
        if len(dir_list_to_process) > 0:
            print("Found {} directoreies to process in the list {}".format(len(dir_list_to_process), args.dir_list))
        else:
            print("Did not find any directories to process. Exiting.")
            return

        list_is_studies = (len(dir_list_to_process[0].split('_')) == 1)
        study_ids = [study.name.split('_')[0] for study in cleaned_studies]

        if list_is_studies:
            print('Input list is a list of studies, getting directories based on the study names')
            to_process_studies =  [study for (study, iid) in zip(cleaned_studies, study_ids) if iid in dir_list_to_process]
        else:
            to_process_studies = [study for study in cleaned_studies if study.stem in dir_list_to_process]
        
        # to_process_studies = [ study for (study, iid) in zip(cleaned_studies, study_ids) if iid in dir_list_to_process ]
        cleaned_studies = to_process_studies
        print("Will run OCR on {} studies".format(len(cleaned_studies)))

    if args.use_threads:
        with concurrent.futures.ThreadPoolExecutor(max_workers=args.num_threads) as executor:
            # Start the load operations and mark each future with its URL
            {executor.submit(extractTagForStudy, study, 
                            data_folder, out_images_dir, tag_list,
                            non_tag_us, tag_bounding_box, 
                            Path(args.server_path), args.greedy ): study for study in cleaned_studies}
            # future_tags = {executor.submit(extractTagForStudy, study, 
            #                                 data_folder, out_images_dir, tag_list,
            #                                 non_tag_us, tag_bounding_box, 
            #                                 Path(args.server_path), args.greedy ): study for study in cleaned_studies}
            # for future in concurrent.futures.as_completed(future_tags, timeout=3600):
            #     d = future_tags[future] 
            #     logging.info('Finished processing: {}'.format(d))
            #     this_tag_statistic = future.result()
            #     #logging.info(future.result())
            #     try:
            #         if this_tag_statistic is not None:
            #             for key, value in this_tag_statistic.items():
            #                 tag_statistic[key] += value
            #             with open(finished_study_file, "a+") as f:
            #                 f.write(d.name+os.linesep)
            #     except Exception as e:
            #         logging.warning('Got an error while processing the result for {}'.format(d.name))

    else:
        i=1
        for study in cleaned_studies:
            this_tag_statistic = extractTagForStudy(study, data_folder, out_images_dir, 
                                                    tag_list, non_tag_us, tag_bounding_box, 
                                                    Path(args.server_path), args.greedy)
            logging.info('Finished processing: {}'.format(study))
            if this_tag_statistic is not None:
                for key, value in this_tag_statistic.items():
                    tag_statistic[key] += value
                endstr = "\n" if i%50 == 0 else "."
                print("f",end=endstr)
            i+=1
    
    pprint(tag_statistic)
    with open(out_images_dir/"NumberOfTags.json", "w") as outfile:
        json.dump(tag_statistic, outfile, indent=4) 
    logging.info(pformat(tag_statistic))
    logging.info('---- DONE ----')
    print('------DONE-----------')


if __name__=="__main__":
    
    parser = ArgumentParser()
    parser.add_argument('--dir', type=str, help='Directory with subject subfolders.'
                'Every lowest level subfolder will be considered as a study', required=True)
    parser.add_argument('--out_dir', type=str, help='Output directory location.'
                'The directory hierarchy will be copied to this directory', required=True)
    parser.add_argument('--server_path', type=str, help='Path to the server relative to which teh paths of files will be saved')
    parser.add_argument('--debug', action='store_true', help='Add debug info in log')
    parser.add_argument('--use_threads', action='store_true', help='Use threads to run the code')
    parser.add_argument('--greedy', action='store_true', help='If given, will retry to extract tags for files that were processed previously')
    parser.add_argument('--dir_list', type=str, help='List of directories to be processed')
    parser.add_argument('--num_threads', type=int, help='Number of worker threads', default=4)
    args = parser.parse_args()

    main(args)