from itertools import accumulate, zip_longest
import os

import editdistance
import numpy as np

from fpua.data.general import aggregate_actions_and_lengths, extend_smallest_list, extract_last_action_and_full_length
from fpua.data.general import count_num_actions


def analyse_hierarchical_observations_and_predictions(predicted_fine_actions_per_step,
                                                      predicted_coarse_actions_per_step,
                                                      observed_fine_actions_per_frame,
                                                      observed_coarse_actions_per_frame,
                                                      unobserved_fine_actions_per_frame,
                                                      unobserved_coarse_actions_per_frame,
                                                      coarse_actions_per_frame_full,
                                                      save_path, save_file_name):
    """Analyse hierarchical model input, ground-truth output, and predicted output.

    Arg(s):
        predicted_fine_actions_per_step - The fine-level predictions generated by the hierarchical model. It is a list
            of tuples, each tuple being a pair containing the action and its predicted length.
        predicted_coarse_actions_per_step - Same as predicted_fine_actions_per_step but for the coarse level. Also,
            since the coarse level does not make a prediction at every time step, the non-predicted time steps contain
            None.
        observed_fine_actions_per_frame - A list containing the observed fine actions.
        observed_coarse_actions_per_frame - A list containing the observed coarse actions.
        unobserved_fine_actions_per_frame - A list containing the unobserved fine actions.
        unobserved_coarse_actions_per_frame - A list containing the unobserved coarse actions.
        coarse_actions_per_frame_full - Coarse actions per frame for the whole video.
        save_path - Directory to save the results of the analysis.
        save_file_name - Base name for the generated files to be saved.
    """
    obs_actions = list(zip(observed_coarse_actions_per_frame, observed_fine_actions_per_frame))
    obs_actions, obs_lengths = aggregate_actions_and_lengths(obs_actions)
    acc_obs_lengths = list(accumulate(obs_lengths))
    unobs_actions = list(zip(unobserved_coarse_actions_per_frame, unobserved_fine_actions_per_frame))
    unobs_actions, unobs_lengths = aggregate_actions_and_lengths(unobs_actions)
    acc_unobs_lengths = list(accumulate(unobs_lengths))
    num_frames_total = len(coarse_actions_per_frame_full)
    coarse_transition_length = extract_last_action_and_full_length(coarse_actions_per_frame_full,
                                                                   len(observed_coarse_actions_per_frame) - 1)[1]
    num_obs_coarse_actions = count_num_actions(observed_coarse_actions_per_frame)
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    with open(os.path.join(save_path, save_file_name), mode='w') as f:
        # OBSERVED
        f.write('Observed\n')
        last_coarse_action = None
        last_coarse_action_len = None
        running_num_obs_coarse_actions = 0
        for i, ((obs_coarse_action, obs_fine_action), obs_length) in enumerate(zip(obs_actions, obs_lengths)):
            if obs_coarse_action != last_coarse_action:
                last_coarse_action = obs_coarse_action
                global_cut = acc_obs_lengths[i] - 1
                _, obs_coarse_len = extract_last_action_and_full_length(observed_coarse_actions_per_frame,
                                                                        cut=global_cut)
                last_coarse_action_len = obs_coarse_len
                f.write(f'\t{obs_coarse_action:<24}\t{obs_coarse_len:>4}  ({obs_coarse_len / num_frames_total:.2f})\n')
                running_num_obs_coarse_actions += 1
            if running_num_obs_coarse_actions == num_obs_coarse_actions:
                f.write(f'\t\t{obs_fine_action:<24}\t{obs_length:>4d}  ({obs_length / coarse_transition_length:.2f})\n')
            else:
                f.write(f'\t\t{obs_fine_action:<24}\t{obs_length:>4d}  ({obs_length / last_coarse_action_len:.2f})\n')
        # UNOBSERVED
        f.write('\nUnobserved\n')
        for i, ((unobs_coarse_action, unobs_fine_action), unobs_length) in enumerate(zip(unobs_actions, unobs_lengths)):
            if unobs_coarse_action != last_coarse_action:
                last_coarse_action = unobs_coarse_action
                global_cut = acc_unobs_lengths[i] - 1
                _, unobs_coarse_len = extract_last_action_and_full_length(unobserved_coarse_actions_per_frame,
                                                                          cut=global_cut)
                f.write(f'\t{unobs_coarse_action:<24}\t{unobs_coarse_len:>4d}  '
                        f'({unobs_coarse_len / num_frames_total:.2f})\n')
            elif i == 0:
                global_cut = acc_unobs_lengths[i] - 1
                _, unobs_coarse_len = extract_last_action_and_full_length(unobserved_coarse_actions_per_frame,
                                                                          cut=global_cut)
                f.write(f'\t{unobs_coarse_action:<24}\t{unobs_coarse_len:>4d}  '
                        f'({unobs_coarse_len / num_frames_total:.2f})\n')
            global_cut = acc_obs_lengths[-1] + acc_unobs_lengths[i] - 1
            parent_len = extract_last_action_and_full_length(coarse_actions_per_frame_full, global_cut)[1]
            f.write(f'\t\t{unobs_fine_action:<24}\t{unobs_length:>4d}  ({unobs_length / parent_len:.2f})\n')
        # PREDICTIONS
        f.write('\nPredicted\n')
        if predicted_coarse_actions_per_step is None:
            predicted_coarse_actions_per_step = [(None, None) for _ in predicted_fine_actions_per_step]
        parent_len, is_first = None, True
        if predicted_coarse_actions_per_step[0][1] is not None:
            parent_len = extract_last_action_and_full_length(observed_coarse_actions_per_frame,
                                                             len(observed_coarse_actions_per_frame) - 1)[1]
            parent_len += predicted_coarse_actions_per_step[0][1]
        for predicted_coarse, predicted_fine in zip(predicted_coarse_actions_per_step, predicted_fine_actions_per_step):
            predicted_coarse_action, predicted_coarse_len = predicted_coarse
            if predicted_coarse_action is not None:
                if is_first:
                    is_first = False
                else:
                    parent_len = predicted_coarse_len
                f.write(f'\t{predicted_coarse_action:<24}\t{predicted_coarse_len:>4d}  '
                        f'({predicted_coarse_len / num_frames_total:.2f})\n')
            predicted_fine_action, predicted_fine_action_len = predicted_fine
            if predicted_fine_action is not None:
                if parent_len is not None:
                    f.write(f'\t\t{predicted_fine_action:<24}\t{predicted_fine_action_len:>4d}  '
                            f'({predicted_fine_action_len / parent_len:.2f})\n')
                else:
                    f.write(f'\t\t{predicted_fine_action:<24}\t{predicted_fine_action_len:>4d}\n')


def analyse_single_level_observations_and_predictions_per_step(predicted_actions_per_step, observed_actions_per_frame,
                                                               unobserved_actions_per_frame, num_frames,
                                                               save_path, save_file_name):
    """Analyse single level model input, ground-truth output, and predicted output.

    Arg(s):
        predicted_actions_per_step - The predictions generated by the model. It is a list
            of tuples, each tuple being a pair containing the action and its predicted length.
        observed_actions_per_frame - A list containing the observed  actions.
        unobserved_actions_per_frame - A list containing the unobserved  actions.
        num_frames - Number of frames in the video.
        save_path - Directory to save the results of the analysis.
        save_file_name - Base name for the generated files to be saved.
    """
    obs_actions, obs_lengths = aggregate_actions_and_lengths(observed_actions_per_frame)
    unobs_actions, unobs_lengths = aggregate_actions_and_lengths(unobserved_actions_per_frame)
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    with open(os.path.join(save_path, save_file_name), mode='w') as f:
        f.write('Observed\n')
        for i, (obs_action, obs_length) in enumerate(zip(obs_actions, obs_lengths)):
            f.write(f'\t{obs_action:<24}\t{obs_length:>4d}  ({obs_length / num_frames:.2f})\n')
        f.write('\nUnobserved\n')
        for i, (unobs_action, unobs_length) in enumerate(zip(unobs_actions, unobs_lengths)):
            f.write(f'\t{unobs_action:<24}\t{unobs_length:>4d}  ({unobs_length / num_frames:.2f})\n')
        f.write('\nPredicted\n')
        for predicted_action, predicted_length in predicted_actions_per_step:
            f.write(f'\t{predicted_action:<24}\t{predicted_length:>4d}  ({predicted_length / num_frames:.2f})\n')


def analyse_performance_per_future_action(predicted_actions_per_video, unobserved_actions_per_video,
                                          transition_action_per_video, save_path, mode='w', extra_str=''):
    performances = {}
    for predicted_actions, unobserved_actions, transition_action in zip(predicted_actions_per_video,
                                                                        unobserved_actions_per_video,
                                                                        transition_action_per_video):
        transition_action_is_not_finished = transition_action == unobserved_actions[0]
        _, unobserved_lengths = aggregate_actions_and_lengths(unobserved_actions.tolist())
        unobserved_actions_initial_frames = [0] + list(accumulate(unobserved_lengths))
        for action_idx, (initial_frame, final_frame) in enumerate(zip(unobserved_actions_initial_frames[:-1],
                                                                      unobserved_actions_initial_frames[1:])):
            unobserved = unobserved_actions[initial_frame:final_frame]
            predicted = predicted_actions[initial_frame:final_frame]
            accuracy = np.mean(unobserved == predicted).item()
            if transition_action_is_not_finished:
                performances.setdefault(action_idx, []).append(accuracy)
            else:
                performances.setdefault(action_idx + 1, []).append(accuracy)
    performances = {action_idx: (sum(perf_list) / len(perf_list), len(perf_list))
                    for action_idx, perf_list in performances.items()}
    suffixes = {1: 'st', 2: 'nd', 3: 'rd'}
    save_file_name = 'mfap_full_split.txt'
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    with open(os.path.join(save_path, save_file_name), mode=mode) as f:
        if extra_str:
            heading = 'Future ' + extra_str + ' Actions Average Accuracy:\n'
        else:
            heading = 'Future Actions Average Accuracy:\n'
        if mode == 'a':
            heading = '\n' + heading
        f.write(heading)
        for action_idx in range(200):
            perf = performances.get(action_idx)
            if perf is None:
                break
            fa_avg_acc, num_videos = perf
            fa_avg_acc = round(fa_avg_acc, 4)
            suffix = suffixes.get(action_idx)
            suffix = 'th' if suffix is None else suffix
            prefix = (str(action_idx) + suffix + ': ').rjust(6)
            f.write(prefix + str(fa_avg_acc).ljust(6, '0') + ' | ' + str(num_videos).rjust(3, ' ') + '\n')


def compute_mof(predicted_actions_per_frame, unobserved_actions_per_frame):
    paf, uaf = np.array(predicted_actions_per_frame), np.array(unobserved_actions_per_frame)
    accuracy = np.mean(paf == uaf).item()
    return accuracy


def compute_moc(predicted_actions, unobserved_actions, action_to_id):
    num_classes = len(action_to_id)
    correct_per_class = np.zeros(num_classes, dtype=np.float32)
    wrong_per_class = np.zeros(num_classes, dtype=np.float32)
    for predicted_action, unobserved_action in zip(predicted_actions, unobserved_actions):
        if predicted_action == unobserved_action:
            correct_per_class[action_to_id[unobserved_action]] += 1
        else:
            wrong_per_class[action_to_id[unobserved_action]] += 1
    moc = 0
    n = 0
    for cpc, wpc in zip(correct_per_class, wrong_per_class):
        if cpc + wpc > 0:
            moc += cpc / (cpc + wpc)
            n += 1
    return moc / n, correct_per_class, wrong_per_class


def compute_segmental_edit_score_single_video(unobserved_actions_per_frame, predicted_actions_per_frame):
    normalised_levenshtein_distance = editdistance.eval(unobserved_actions_per_frame, predicted_actions_per_frame)
    normalised_levenshtein_distance /= len(unobserved_actions_per_frame)  # predicted has the same length
    return 1 - normalised_levenshtein_distance


def compute_segmental_edit_score_multiple_videos(unobserved_actions_per_video, predicted_actions_per_video):
    seg_edit_scores = []
    for unobserved_actions_per_frame, predicted_actions_per_frame in zip(unobserved_actions_per_video,
                                                                         predicted_actions_per_video):
        seg_edit_score = compute_segmental_edit_score_single_video(unobserved_actions_per_frame,
                                                                   predicted_actions_per_frame)
        seg_edit_scores.append(seg_edit_score)
    return np.array(seg_edit_scores).mean()


def overlap_f1_single_video(unobserved_actions_per_frame, predicted_actions_per_frame, action_to_id, num_classes,
                            bg_class, overlap):
    true_intervals = np.array(segment_intervals(unobserved_actions_per_frame))
    true_labels = np.array(aggregate_actions_and_lengths(unobserved_actions_per_frame)[0])
    pred_intervals = np.array(segment_intervals(predicted_actions_per_frame))
    pred_labels = np.array(aggregate_actions_and_lengths(predicted_actions_per_frame)[0])

    # Remove background labels
    # if bg_class is not None:
    #     true_intervals = true_intervals[true_labels != bg_class]
    #     true_labels = true_labels[true_labels != bg_class]
    #     pred_intervals = pred_intervals[pred_labels != bg_class]
    #     pred_labels = pred_labels[pred_labels != bg_class]

    n_true = true_labels.shape[0]
    n_pred = pred_labels.shape[0]

    # We keep track of the per-class TPs, and FPs.
    # In the end we just sum over them though.
    TP = np.zeros(num_classes, dtype=np.float)
    FP = np.zeros(num_classes, dtype=np.float)
    true_used = np.zeros(n_true, dtype=np.float)

    for j in range(n_pred):
        # Compute IoU against all others
        intersection = np.minimum(pred_intervals[j, 1], true_intervals[:, 1]) - np.maximum(pred_intervals[j, 0],
                                                                                           true_intervals[:, 0])
        union = np.maximum(pred_intervals[j, 1], true_intervals[:, 1]) - np.minimum(pred_intervals[j, 0],
                                                                                    true_intervals[:, 0])
        IoU = (intersection / union) * (pred_labels[j] == true_labels)

        # Get the best scoring segment
        idx = IoU.argmax()

        # If the IoU is high enough and the true segment isn't already used
        # Then it is a true positive. Otherwise is it a false positive.
        action_id = action_to_id.get(pred_labels[j])
        if action_id is None:
            continue
        if IoU[idx] >= overlap and not true_used[idx]:
            TP[action_id] += 1
            true_used[idx] = 1
        else:
            FP[action_id] += 1

    TP = TP.sum()
    FP = FP.sum()
    # False negatives are any unused true segment (i.e. "miss")
    FN = n_true - true_used.sum()

    precision = TP / (TP + FP)
    recall = TP / (TP + FN)
    F1 = 2 * (precision * recall) / (precision + recall)

    # If the precision + recall == 0, it is a NaN. Set these to 0.
    F1 = np.nan_to_num(F1)

    return F1


def overlap_f1_multiple_videos(unobserved_actions_per_video, predicted_actions_per_video,
                               action_to_id, num_classes=0, bg_class=None, overlap=0.1):
    overlap_f1s = []
    for unobserved_actions_per_frame, predicted_actions_per_frame in zip(unobserved_actions_per_video,
                                                                         predicted_actions_per_video):
        overlap_f1 = overlap_f1_single_video(unobserved_actions_per_frame.tolist(),
                                             predicted_actions_per_frame.tolist(),
                                             action_to_id=action_to_id,
                                             num_classes=num_classes, bg_class=bg_class, overlap=overlap)
        overlap_f1s.append(overlap_f1)
    return np.array(overlap_f1s).mean()


def segment_intervals(actions_per_frame):
    actions, lengths = aggregate_actions_and_lengths(actions_per_frame)
    actions_initial_frames = [0] + list(accumulate(lengths))
    return list(zip(actions_initial_frames[:-1], actions_initial_frames[1:]))


def analyse_flushes_hierarchical(flushes_per_video, ground_truth_flushes_per_video,
                                 label_files, seq_len, save_path, encoder=True, extra_str=''):
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    if extra_str:
        save_file_name = 'flushes_' + extra_str + '_full_split.txt'
    else:
        save_file_name = 'flushes_full_split.txt'
    with open(os.path.join(save_path, save_file_name), mode='w') as f:
        for flushes, ground_truth_flushes, label_file in \
                zip(flushes_per_video, ground_truth_flushes_per_video, label_files):
            num_steps = len(flushes)
            flushes, ground_truth_flushes = extend_smallest_list(flushes, ground_truth_flushes, extension_val=0.0)
            flushes_ground_truth_positions = compute_flushes_positions(ground_truth_flushes)
            flushes_positions = compute_flushes_positions(flushes)
            f.write(label_file + '\n')
            if encoder:
                f.write('\tNumber of Input Steps: %d' % seq_len)
            else:
                f.write('\tNumber of Output Steps: %d' % seq_len)
                f.write('\n')
                f.write('\tNumber of Predicted Steps: %d' % num_steps)
            f.write('\n')
            f.write('\tGround-truth Flush Positions:')
            for flush_ground_truth_position in flushes_ground_truth_positions:
                f.write(' ' + str(flush_ground_truth_position).rjust(2))
            f.write('\n')
            f.write('\t   Predicted Flush Positions:')
            for flush_position in flushes_positions:
                f.write(' ' + str(flush_position).rjust(2))
            f.write('\n\n')


def compute_flushes_positions(flushes):
    positions = [t for t, flush in enumerate(flushes) if flush == 1.0]
    return positions


def write_results_per_video(f1_per_video, order_by, metric_name, save_path):
    if order_by is None:
        index, order_by = 1, 'level'
    else:
        index = 1 if order_by == 'coarse' else 2
    f1_per_video = sorted(f1_per_video, key=lambda x: x[index], reverse=True)
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    save_file_name = 'metric_' + metric_name + '_order_by_' + order_by + '.txt'
    with open(os.path.join(save_path, save_file_name), mode='w') as f:
        if order_by == 'level':
            f.write('File Name'.rjust(36) + 'Level'.rjust(8) + '\n')
            for label_file, metric in f1_per_video:
                f.write(label_file.rjust(36) + '{:8.4f}'.format(metric) + '\n')
        else:
            f.write('File Name'.rjust(36) + 'Coarse'.rjust(8) + 'Fine'.rjust(8) + '\n')
            for label_file, metric_coarse, metric_fine in f1_per_video:
                f.write(label_file.rjust(36) + '{:8.4f}'.format(metric_coarse) + '{:8.4f}'.format(metric_fine) + '\n')


def action_sequence_metrics(predicted_actions, target_actions):
    """Precision, Recall, and F1 regardless of action length and action class."""
    correct, attempt, total_target = 0, 0, len(target_actions)
    for predicted_action, target_action in zip_longest(predicted_actions, target_actions):
        if predicted_action is None:
            break
        if target_action is None:
            attempt += 1
            continue
        if predicted_action == target_action:
            correct += 1
        attempt += 1
    precision = correct / attempt
    recall = correct / total_target
    try:
        f1 = 2 * (precision * recall) / (precision + recall)
    except ZeroDivisionError:
        f1 = 0.0
    return precision, recall, f1


def write_sequence_results_per_video(metrics_per_video, save_path, order_by='file', level=None):
    order_index = {'file': 0, 'precision': 1, 'recall': 2, 'f1': 3}[order_by]
    metrics_per_video = sorted(metrics_per_video, key=lambda x: x[order_index], reverse=False)
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    if level is not None:
        save_file_name = f'metrics_action_sequence_{level}.txt'
    else:
        save_file_name = 'metrics_action_sequence.txt'
    with open(os.path.join(save_path, save_file_name), mode='w') as f:
        f.write('File Name'.rjust(36) + 'Prec.'.rjust(8) + 'Recall'.rjust(8) + 'F1'.rjust(8) + '\n')
        for label_file, precision, recall, f1 in metrics_per_video:
            f.write(f'{label_file:>36} {precision:>8.2f} {recall:>8.2f} {f1:>8.2f}\n')


def compute_metrics(predicted_actions_per_video_sub, unobserved_actions_per_video_sub, action_to_id):
    overlaps, overlap_f1_scores = [0.10, 0.25, 0.50], []
    for overlap in overlaps:
        overlap_f1_score = overlap_f1_multiple_videos(unobserved_actions_per_video_sub,
                                                      predicted_actions_per_video_sub,
                                                      action_to_id=action_to_id,
                                                      num_classes=len(action_to_id), overlap=overlap)
        overlap_f1_scores.append(overlap_f1_score)
    return overlaps, overlap_f1_scores
