'''
Tools for Saving, Loading, and Using Species Classifiers

@ Kyra Swanson 2023
'''
import json
from typing import Optional, Union
import pandas as pd
import numpy as np
from pathlib import Path
from time import time
from tqdm import tqdm
import onnxruntime

from animl import generator, file_management
from animl.utils.general import get_onnx_device, softmax


SDZWA_CLASSIFIER_SIZE = 299


def load_classifier(model_path: str,
                    classes: Union[int, str, Path, pd.DataFrame] = None):
    '''
    Creates a model instance and loads the latest model state weights.

    Args:
        model_path (str): file or directory path to model weights
        classes (int | str | Path | pd.DataFrame): number of classes, path to associated class list, or pd.DataFrame of class list
        device (str): specify to run on cpu or gpu
        architecture (str): expected model architecture

    Returns:
        model: model object of given architecture with loaded weights
        class_list: list of class names
    '''
    class_list = None
    model_path = Path(model_path)
    # check to make sure GPU is available if chosen
    providers = get_onnx_device()

    print(f'Loading model at {model_path}')
    model = onnxruntime.InferenceSession(model_path, providers=providers)
    model.framework = "onnx"

    # get number of classes
    if isinstance(classes, str) or isinstance(classes, Path):
        class_list = load_class_list(classes)
    elif isinstance(classes, pd.DataFrame):
        class_list = classes
    else:
        class_list = None
    # try to load class dict from metadata
    props = model.get_modelmeta().custom_metadata_map
    if "class_dict" in props:
        print("Loaded class_dict from ONNX metadata.")
        class_dict = json.loads(props["class_dict"])
        class_list = [class_dict[str(i)] for i in range(len(class_dict))]
        class_list = pd.DataFrame({'class': class_list})

        # no need to return epoch
    return model, class_list


def load_class_list(classlist_file):
    """
    Return classlist file as pd.DataFrame.

    Args:
        classlist_file (str): file path to class list

    Returns:
        pd.DataFrame of class list
    """
    if not Path(classlist_file).is_file():
        raise FileNotFoundError(f"Class list file not found at {classlist_file}")
    return pd.read_csv(classlist_file)


def classify(model,
             detections,
             resize_width: int = 480,
             resize_height: int = 480,
             file_col: str = 'filepath',
             crop: bool = True,
             normalize: bool = True,
             out_file: Optional[str] = None):
    """
    TODO: align with R version
    Predict species using classifier model.

    Args:
        model: preloaded classifier model
        detections (mult): dataframe of (animal) detections, list of filepaths or filepath str
        resize_width (int): image width input size
        resize_height (int): image height input size
        file_col (str): column name containing file paths
        crop (bool): use bbox to crop images before feeding into model
        normalize (bool): normalize the tensor before inference
        batch_size (int): data generator batch size
        num_workers (int): number of cores
        device (str): specify to run model on cpu or gpu, default to cpu
        out_file (str): path to save prediction results to

    Returns:
        detections (pd.DataFrame): MD detections with classifier prediction and confidence
    """
    if file_management.check_file(out_file, output_type="Classification results"):
        return file_management.load_data(out_file).to_numpy()

    # initialize lists
    raw_output = []

    # Manifest
    if isinstance(detections, pd.DataFrame):
        if file_col not in detections.columns:
            raise ValueError(f"file_col {file_col} not found in manifest columns")
        # no frame column, assume all images and set to 0
        if 'frame' not in detections.columns:
            print("Warning: 'frame' column not found in manifest columns. Defaulting to 0 assuming images.")
            detections['frame'] = 0

        dataset = generator.manifest_dataloader(detections, file_col=file_col, crop=crop,
                                                resize_width=resize_width, resize_height=resize_height,
                                                normalize=normalize)
    # Single File
    elif isinstance(detections, str):
        detections = pd.DataFrame({file_col: detections, 'frame': 0}, index=[0])
        dataset = generator.manifest_dataloader(detections, file_col=file_col, crop=False,
                                                resize_width=resize_width, resize_height=resize_height,
                                                normalize=normalize)
    # List of Files
    elif isinstance(detections, list):
        detections = pd.DataFrame({file_col: detections, 'frame': 0}, index=range(len(detections)))
        dataset = generator.manifest_dataloader(detections, file_col=file_col, crop=False,
                                                resize_width=resize_width, resize_height=resize_height,
                                                normalize=normalize)
    else:
        raise AssertionError("Input must be a data frame of crops, single file path or vector of file paths.")

    # Predict
    start_time = time()
    for _, batch in tqdm(enumerate(dataset), total=len(detections)):
        image = batch[0]
        output = model.run(None, {model.get_inputs()[0].name: image})[0]
        raw_output.extend(softmax(output))

    raw_output = np.vstack(raw_output)

    if out_file:
        file_management.save_data(pd.DataFrame(raw_output), out_file)

    print(f"\nFinished classification. Total images processed: {len(raw_output)} at {round(len(raw_output)/(time() - start_time), 1)} img/s.")

    return raw_output


def single_classification(animals: pd.DataFrame,
                          empty: Optional[pd.DataFrame],
                          predictions_raw: np.array,
                          class_list: list,
                          best: bool = False):
    """
    Get maximum likelihood prediction from softmaxed logits.

    Args:
        animals (pd.DataFrame): animal detections from manifest
        empty (Optional) (pd.DataFrame): empty detections from manifest
        predictions_raw (np.array): softmaxed logits from classify()
        class_list (list): class list associated with model
        best (bool): whether to return one prediction per file

    Returns:
        animals dataframe with "prediction" label an "confidence" columns
    """
    class_list = pd.Series(class_list)

    # convert None to empty dataframe fo concat
    if empty is None:
        empty = pd.DataFrame()

    if not animals.empty:
        animals = animals.reset_index(drop=True)
        animals["prediction"] = class_list.values[np.argmax(predictions_raw, axis=1)]
        animals["confidence"] = animals["conf"].mul(np.max(predictions_raw, axis=1))

    manifest = pd.concat([animals if not animals.empty else None, empty if not empty.empty else None]).reset_index(drop=True)

    # remove empties from videos
    files = manifest.groupby('filepath')
    for f, file in files:
        if file['extension'].iloc[0] in file_management.VIDEO_EXTENSIONS:
            predictions = file['prediction'].unique()
            if 'empty' in predictions and len(predictions) > 1:
                real_prediction = predictions[predictions != 'empty'][0]
                manifest.loc[manifest['filepath'] == f, 'prediction'] = real_prediction

    # best guess
    if best:
        # take most confident guess    
        manifest = manifest.sort_values("confidence", ascending=False)
        manifest = manifest.drop_duplicates(subset="filepath", keep="first")
    
    return manifest.reset_index(drop=True)


def sequence_classification(animals: pd.DataFrame,
                            empty: Optional[pd.DataFrame],
                            predictions_raw: np.array,
                            class_list: pd.DataFrame,
                            station_col: str,
                            empty_class: str = "",
                            sort_columns: list[str] = None,
                            file_col: str = "filepath",
                            maxdiff: int = 60):
    """
    Applies class labels to images based on sequential information.

    This function applies image classifications at a sequence level by leveraging information from
    multiple images. A sequence is defined as all images at the same camera and station where the
    time between consecutive images is <=maxdiff. This can improve classification accuracy, but
    assumes that only one species is present in each sequence. If you regularly expect multiple
    species to occur in an image or sequence don't use this function.

    Args:
        animals (pd.DataFrame): Sub-selection of all images that contain animals
        empty (Optional) (pd.DataFrame): Sub-selection of all images that do not contain animals
        predictions_raw (Numpy Array of Numpy Arrays): Logits of all entries in "animals"
        class_list (pd.DataFrame): class list associated with classifier model
        station_col (str): The name of the station column
        empty_class (str) (Optional): the name of class_list 'empty' label
        sort_columns (List of Strings): Defines sorting order for the DataFrame
        file_col (str): The name of the filepath column
        maxdiff (int): Maximum time difference in seconds between any two images in a sequence

    Returns:
        final_df (pd.DataFrame): Sequence classified data from both animals and empty

    Raises:
        Exception: If 'sort_columns' is not a list or is empty
        Exception: If 'station_col' is not a string or is empty
        Exception: If 'empty' is defined and is not a pandas DataFrame
        Exception: If maxdiff is defined and is not a positive number
    """
    if not isinstance(station_col, str) or station_col == '':
        raise Exception("'station_col' must be a non-empty string")

    # Sanity check to verify that empty is a Pandas DataFrame, if defined
    if empty is not None and not isinstance(animals, pd.DataFrame):
        raise Exception("'empty' must be a DataFrame")

    # Sanity check to verify that maxdiff is a positive number
    if not isinstance(maxdiff, (int, float)) or maxdiff < 0:
        raise Exception("'maxdiff' must be a number >= 0")

    if not {file_col}.issubset(animals.columns):
        raise ValueError(f"DataFrame must contain '{file_col}' column.")

    if not {"datetime"}.issubset(animals.columns):
        raise ValueError("DataFrame must contain 'datetime' column.")

    if "conf" not in animals.columns:
        animals["conf"] = 1

    if empty_class > "":
        empty_col = class_list[class_list == empty_class].index[0]
    else:
        empty_col = None

    if empty is not None or not empty.empty:
        empty["ID"] = range(0, empty.shape[0])
        predempty = empty.pivot(index="ID", columns="prediction", values="confidence")
        # Replace NaN with 0
        predempty = predempty.fillna(0)
        predempty = pd.concat([pd.DataFrame(np.zeros((empty.shape[0], len(class_list)))), predempty], axis=1)

        if empty_class > "":
            if 'empty' in predempty.columns:
                predempty[empty_col] = predempty["empty"]
                predempty = predempty.drop("empty", axis=1)
            class_list = pd.concat([class_list,
                                    pd.Series([x for x in empty["prediction"].unique() if x != "empty"])], ignore_index=True)

        else:
            class_list = pd.concat([class_list, pd.Series(empty["prediction"].unique())], ignore_index=True)
            empty_col = predempty.columns.get_loc("empty")

        # placeholders
        animals["prediction"] = list(class_list[np.argmax(predictions_raw, axis=1)])
        animals["confidence"] = animals["conf"].mul(np.max(predictions_raw, axis=1))

        empty["conf"] = 1
        animals_merged = pd.concat([animals, empty.iloc[:, :-1]]).reset_index(drop=True)  # dont add ID column
        predictions = np.hstack((predictions_raw,
                                 np.zeros((predictions_raw.shape[0], len(predempty.columns) - predictions_raw.shape[1]))))
        # concat
        predictions = np.vstack((predictions, np.array(predempty)))

    if sort_columns is None:
        sort_columns = [station_col, "datetime"]

    animals_merged['datetime'] = pd.to_datetime(animals_merged['datetime'], format="%Y-%m-%d %H:%M:%S")

    sort = animals_merged.sort_values(by=sort_columns).index
    animals_sort = animals_merged.loc[sort].reset_index(drop=True)
    predsort = predictions[sort]

    conf_placeholder = np.zeros(len(animals_sort))
    predict_placeholder = np.empty(len(animals_sort), dtype='U30')
    sequence_placeholder = np.zeros(len(animals_sort))

    i = 0
    s = 0
    while i < len(animals_sort):
        rows = [i]
        last_index = i+1

        while (last_index < len(animals_sort) and not pd.isna(animals_sort.loc[i, "datetime"]) and
               not pd.isna(animals_sort.loc[last_index, "datetime"]) and
               animals_sort.loc[last_index, station_col] == animals_sort.loc[i, station_col] and
               (animals_sort.loc[last_index, "datetime"] - animals_sort.loc[i, "datetime"]).total_seconds() <= maxdiff):
            rows.append(last_index)
            last_index += 1

        rows = np.array(rows)

        # multiple detections in sequence
        if len(rows) > 1:
            predclass = np.argmax(predsort[rows], axis=1)

            # no empties
            if empty_col is None or empty_col not in predclass:
                predsort_confidence = predsort[rows] * np.reshape(animals_sort.loc[rows, 'conf'].values, (-1, 1))
                predbest = np.mean(predsort_confidence, axis=0)
                conf_placeholder[rows] = np.max(predsort_confidence[:, np.argmax(predbest)])
                predict_placeholder[rows] = class_list[np.argmax(predbest)]
                sequence_placeholder[rows] = int(s)

            else:
                mask = pd.DataFrame((predclass == empty_col))
                filtered_animals = animals_sort.loc[rows, file_col].reset_index(drop=True)
                sum_values = mask.groupby(filtered_animals).sum()
                count_values = mask.groupby(filtered_animals).count()

                sel_all_empty = count_values == sum_values
                sel_mixed = np.where(filtered_animals.isin(sel_all_empty[~sel_all_empty[0]].index))[0]
                sel_no_empties = np.where(filtered_animals.isin(sel_all_empty[~sel_all_empty[0]].index) & (predclass != empty_col))[0]

                if len(sel_mixed) > 0 and len(sel_no_empties) > 0:
                    predsort_confidence = predsort[rows[sel_no_empties]] * np.reshape(animals_sort.loc[rows[sel_no_empties], 'conf'].values, (-1, 1))
                    predbest = np.mean(predsort_confidence, axis=0)
                    conf_placeholder[rows[sel_mixed]] = np.max(predsort_confidence[:, np.argmax(predbest)])
                    predict_placeholder[rows[sel_mixed]] = class_list[np.argmax(predbest)]
                    sequence_placeholder[rows[sel_mixed]] = int(s)
                for file in sel_all_empty[sel_all_empty[0]].index:
                    empty_row = np.where(animals_sort[file_col] == file)
                    predsort_confidence = predsort[empty_row] * np.reshape(animals_sort.loc[empty_row, 'conf'].values, (-1, 1))
                    predbest = np.mean(predsort_confidence, axis=0)
                    conf_placeholder[empty_row] = np.max(predsort_confidence[:, np.argmax(predbest)])
                    sequence_placeholder[empty_row] = int(s)
                    predict_placeholder[empty_row] = class_list[np.argmax(predbest)]
        # single row in sequence
        else:
            predbest = predsort[rows]
            conf_placeholder[rows] = np.max(predbest * animals_sort.loc[rows, 'conf'].values)
            predict_placeholder[rows] = class_list[np.argmax(predbest)]
            sequence_placeholder[rows] = int(s)

        i = last_index
        s += 1

    animals_sort['confidence'] = conf_placeholder
    animals_sort['prediction'] = predict_placeholder
    animals_sort['sequence'] = sequence_placeholder

    return animals_sort
