"""
Automated Pipeline Functions

@ Kyra Swanson 2023
"""
import os
import yaml
import pandas as pd
from pathlib import Path

from animl import (classification, detection, export, file_management,
                   video_processing, split)
from animl.utils import visualization


def from_paths(image_dir: str,
               detector_file: str,
               classifier_file: str,
               class_label: str = "class",
               sort: bool = False,
               visualize: bool = False,
               sequence: bool = False,
               detect_only: bool = False) -> pd.DataFrame:
    """
    This function is the main method to invoke all the sub functions
    to create a working directory for the image directory.

    Args:
        image_dir (str): directory path containing the images or videos.
        detector_file (str): file path of the MegaDetector model.
        classifier_file (str): file path of the classifier model.
        classlist_file (list): list of classes or species for classification.
        class_label: column in the class list that contains the label wanted
        batch_size (int): batch size for inference
        sort (bool): toggle option to create symlinks
        visualize (bool): if True, run visualization
        sequence (bool): if True, run sequence_classification

    Returns:
        pandas.DataFrame: Concatenated dataframe of animal and empty detections
    """
    # Create a working directory, build the file manifest from img_dir
    print("Searching directory...")
    working_dir = file_management.WorkingDirectory(image_dir)
    files = file_management.build_file_manifest(image_dir,
                                                out_file=working_dir.filemanifest,
                                                exif=True)
    # files["station"] = files["filepath"].apply(lambda x: x.split(os.sep)[-2])
    print(f"Found {len(files)} files.")

    # Obtain frames from videos
    all_frames = video_processing.extract_frames(files, frames=5, out_file=working_dir.imageframes)

    # Run detector
    print("Running images and video frames through detector...")
    if (file_management.check_file(working_dir.detections, output_type="Detections")):
        detections = file_management.load_data(working_dir.detections)
    else:
        detector = detection.load_detector(detector_file)
        md_results = detection.detect(detector,
                                      all_frames,
                                      resize_height=detection.MEGADETECTORv5_SIZE,
                                      resize_width=detection.MEGADETECTORv5_SIZE,
                                      checkpoint_path=working_dir.mdraw,
                                      checkpoint_frequency=1000)
        # Convert MD JSON to pandas dataframe, merge with manifest
        print("Converting MD JSON to dataframe and merging with manifest...")
        detections = detection.parse_detections(md_results, manifest=all_frames, out_file=working_dir.detections)

    if detect_only:
        print("Detection only flag set, skipping classification.")
        return detections

    # Extract animal detections from the rest
    animals = split.get_animals(detections)
    empty = split.get_empty(detections)

    # Use the classifier model to predict the species of animal detections
    print("Predicting species of animal detections...")
    classifier, class_list = classification.load_classifier(classifier_file)
    predictions_raw = classification.classify(classifier,
                                              animals,
                                              resize_height=classification.SDZWA_CLASSIFIER_SIZE,
                                              resize_width=classification.SDZWA_CLASSIFIER_SIZE,
                                              out_file=working_dir.predictions)
    if sequence:
        print("Classifying sequences...")
        manifest = classification.sequence_classification(animals, empty, predictions_raw,
                                                          class_list[class_label],
                                                          station_col='station',
                                                          empty_class="",
                                                          sort_columns=["station", "datetime", "frame"],
                                                          maxdiff=60)
    else:
        print("Classifying individual frames...")
        manifest = classification.single_classification(animals, empty, 
                                                        predictions_raw, 
                                                        class_list[class_label],
                                                        best=True)

    if sort:
        print("Sorting...")
        working_dir.activate_linkdir()
        manifest = export.export_folders(manifest, working_dir.linkdir)

    # Plot boxes
    if visualize:
        working_dir.activate_visdir()
        visualization.plot_all_bounding_boxes(manifest, working_dir.visdir, file_col='filepath', label_col='prediction')

    file_management.save_data(manifest, working_dir.results)
    print("Final Results in " + str(working_dir.results))

    return manifest


def from_config(config: str):
    """
    This function is the main method to invoke all the sub functions
    to create a working directory for the image directory.

    Args:
        config (str): path containing config file for inference

    Returns:
        pandas.DataFrame: Concatenated dataframe of animal and empty detections
    """
    print(f'Using config "{config}"')
    cfg = yaml.safe_load(open(config, 'r'))

    # get image dir and cuda defaults
    image_dir = cfg['image_dir']
    device = cfg.get('device', 'cpu')

    print("Searching directory...")
    # Create a working directory, default to image_dir
    working_dir = file_management.WorkingDirectory(cfg.get('working_dir', image_dir))
    files = file_management.build_file_manifest(image_dir,
                                                out_file=working_dir.filemanifest,
                                                exif=cfg.get('exif', True))
    print(f"Found {len(files)} files.")

    # Station Col
    station_dir = cfg.get('station_dir', None)
    if station_dir:
        files["station"] = files["filepath"].apply(lambda x: x.split(os.sep)[station_dir])

    # split out videos
    all_frames = video_processing.extract_frames(files, frames=5, out_file=working_dir.imageframes)

    print("Running images and video frames through detector...")
    if (file_management.check_file(working_dir.detections, output_type="Detections")):
        detections = file_management.load_data(working_dir.detections)
    else:
        detector = detection.load_detector(cfg['detector_file'])
        md_results = detection.detect(detector,
                                      all_frames,
                                      resize_height=detection.MEGADETECTORv5_SIZE,
                                      resize_width=detection.MEGADETECTORv5_SIZE,
                                      letterbox=cfg.get('letterbox', True),
                                      file_col=cfg.get('file_col_detection', 'filepath'),
                                      checkpoint_path=working_dir.mdraw,
                                      checkpoint_frequency=1000)
        # Convert MD JSON to pandas dataframe, merge with manifest
        print("Converting MD JSON to dataframe and merging with manifest...")
        detections = detection.parse_detections(md_results, manifest=all_frames, out_file=working_dir.detections)

    # Extract animal detections from the rest
    animals = split.get_animals(detections)
    empty = split.get_empty(detections)

    # Use the classifier model to predict the species of animal detections
    print("Predicting species...")
    # Load classifier
    classifier, class_list = classification.load_classifier(cfg['classifier_file'], cfg.get('class_list', None))

    predictions_raw = classification.classify(classifier,
                                              animals,
                                              resize_height=cfg.get('classifier_resize_height', classification.SDZWA_CLASSIFIER_SIZE),
                                              resize_width=cfg.get('classifier_resize_width', classification.SDZWA_CLASSIFIER_SIZE),
                                              file_col=cfg.get('file_col_classification', 'filepath'),
                                              out_file=working_dir.predictions)

    # Convert predictions to labels
    if station_dir:
        manifest = classification.sequence_classification(animals, empty, predictions_raw,
                                                          class_list[cfg.get('class_label_col', 'class')],
                                                          station_col='station',
                                                          empty_class=cfg['empty_class'],
                                                          sort_columns=["station", "datetime", "frame"],
                                                          file_col=cfg.get('file_col_classification', 'frame'),
                                                          maxdiff=60)
    else:
        manifest = classification.single_classification(animals, empty, 
                                                        predictions_raw, 
                                                        class_list[cfg.get('class_label_col', 'class')],
                                                        file_col=cfg.get('file_col_classification', 'filepath'),
                                                        best=cfg.get('best_only', True))

    if cfg.get('sort', True):
        print("Sorting...")
        working_dir.activate_linkdir()
        manifest = export.export_folders(manifest, working_dir.linkdir, copy=cfg.get('copy', False))

    # Plot boxes
    if cfg.get('visualize', False):
        working_dir.activate_visdir()
        visualization.plot_all_bounding_boxes(manifest, working_dir.visdir, file_col='filepath', label_col='prediction')

    file_management.save_data(manifest, working_dir.results)
    print("Final Results in " + str(working_dir.results))

    return manifest
