import itertools
import pathlib
import argparse
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from PIL import Image
from tqdm.auto import tqdm

from IPython import embed

def extract_time_freq_range_from_filename(img_path):
    file_name_str, time_span_str, freq_span_str = str(img_path.with_suffix('').name).split('__')
    time_span_str = time_span_str.replace('s', '')
    freq_span_str = freq_span_str.replace('Hz', '')

    t0, t1 = np.array(time_span_str.split('-'), dtype=float)
    f0, f1 = np.array(freq_span_str.split('-'), dtype=float)
    return file_name_str, t0, t1, f0, f1

def bbox_to_data(img_path, t_min, t_max, f_min, f_max):
    label_path = img_path.parent.parent / 'labels' / img_path.with_suffix('.txt').name

    annotations = np.loadtxt(label_path, delimiter=' ')
    if len(annotations.shape) == 1:
        annotations = np.array([annotations])
    if annotations.shape[1] == 0:
        print('no rises detected in this window')
        return [], []


    boxes = np.array([[x[1] - x[3] / 2, 1 - (x[2] + x[4] / 2), x[1] + x[3] / 2, 1 - (x[2] - x[4] / 2)] for x in annotations]) # x0, y0, x1, y1
    boxes[:, 0] = boxes[:, 0] * (t_max - t_min) + t_min
    boxes[:, 2] = boxes[:, 2] * (t_max - t_min) + t_min

    boxes[:, 1] = boxes[:, 1] * (f_max - f_min) + f_min
    boxes[:, 3] = boxes[:, 3] * (f_max - f_min) + f_min

    scores = annotations[:, 5]

    return boxes, scores

def load_wavetracker_data(raw_path):
    fund_v = np.load(raw_path.parent / 'fund_v.npy')
    ident_v = np.load(raw_path.parent / 'ident_v.npy')
    idx_v = np.load(raw_path.parent / 'idx_v.npy')
    times = np.load(raw_path.parent / 'times.npy')

    return fund_v, ident_v, idx_v, times


def assign_rises_to_ids(raw_path, time_frequency_bboxes, bbox_groups):
    def identify_most_likely_rise_id(possible_ids, t0, t1, f0, f1, fund_v, ident_v, times, idx_v):

        mean_id_box_f_rel_to_bbox = []
        for id in possible_ids:
            id_box_f = fund_v[(ident_v == id) & (times[idx_v] >= t0) & (times[idx_v] <= t1)]
            id_box_f_rel_to_bbox = (id_box_f - f0) / (f1 - f0)
            mean_id_box_f_rel_to_bbox.append(np.mean(id_box_f_rel_to_bbox))
            # print(id, np.mean(id_box_f), f0, f1, np.mean(id_box_f_rel_to_bbox))

        most_likely_id = possible_ids[np.argsort(mean_id_box_f_rel_to_bbox)[0]]
        return most_likely_id

    fund_v, ident_v, idx_v, times = load_wavetracker_data(raw_path)

    fig, ax = plt.subplots()
    ax.plot(times[idx_v[~np.isnan(ident_v)]], fund_v[~np.isnan(ident_v)], '.')

    mask = time_frequency_bboxes['file_name'] == raw_path.parent.name
    for index, bbox in time_frequency_bboxes[mask].iterrows():
        name, t0, f0, t1, f1, score = (bbox[0], *bbox[1:-2].astype(float))
        if bbox_groups[index] == 0:
            color = 'tab:green'
        elif bbox_groups[index] > 0:
            color = 'tab:red'
        else:
            color = 'k'

        ax.add_patch(
            Rectangle((t0, f0),
                      (t1 - t0),
                      (f1 - f0),
                      fill=False, color=color, linestyle='--', linewidth=2, zorder=10)

        )
        ax.text(t1, f1, f'{score:.1%}', ha='right', va='bottom')

        possible_ids = np.unique(
            ident_v[~np.isnan(ident_v) &
                    (t0 <= times[idx_v]) &
                    (t1 >= times[idx_v]) &
                    (f0 <= fund_v) &
                    (f1 >= fund_v)]
        )
        if len(possible_ids) == 1:
            assigned_id = possible_ids[0]
            time_frequency_bboxes.at[index, 'id'] = assigned_id
        elif len(possible_ids) > 1:
            assigned_id = identify_most_likely_rise_id(possible_ids, t0, t1, f0, f1, fund_v, ident_v, times, idx_v)
            time_frequency_bboxes.at[index, 'id'] = assigned_id
            # rise_id[index] = identify_most_likely_rise_id(possible_ids, t0, t1, f0, f1, fund_v, ident_v, times, idx_v)
        else:
            continue

        rise_start_freq_th = f0 + (f1 - f0) * 0.37
        wavetracker_mask = np.arange(len(fund_v))[
            (times[idx_v] >= t0) &
            (times[idx_v] <= t1) &
            (ident_v == assigned_id)
        ]
        if np.sum(fund_v[wavetracker_mask] > rise_start_freq_th) > 0:
            # rise start time = moment where rise freq exceeds 37% of bbox freq range ...
            rise_start_idx = wavetracker_mask[fund_v[wavetracker_mask] > rise_start_freq_th][0]
            rise_time = times[idx_v[rise_start_idx]]
        else:
            ### if this is never the case use the largest slope
            rise_start_idx = wavetracker_mask[np.argmax(np.diff(fund_v[wavetracker_mask]))]
            rise_time = times[idx_v[rise_start_idx]]
        time_frequency_bboxes.at[index, 'event_time'] = rise_time

        ax.plot(rise_time, fund_v[rise_start_idx], 'ok')
        # embed()
        # quit()
    # time_frequency_bboxes['id'] = rise_id
    # embed()
    # plt.show()
    plt.close()
    return time_frequency_bboxes


def find_overlapping_bboxes(df_collect):
    file_names = np.array(df_collect)[:, 0]
    bboxes = np.array(df_collect)[:, 1:].astype(float)

    overlap_bbox_idxs = []

    for file_name in tqdm(np.unique(file_names)):
        file_bbox_idxs = np.arange(len(file_names))[file_names == file_name]
        for ind0, ind1 in itertools.combinations(file_bbox_idxs, r=2):
            bb0 = bboxes[ind0]
            bb1 = bboxes[ind1]
            t0_0, f0_0, t0_1, f0_1 = bb0[:-1]
            t1_0, f1_0, t1_1, f1_1 = bb1[:-1]

            bb_times = np.array([t0_0, t0_1, t1_0, t1_1])
            bb_time_associate = np.array([0, 0, 1, 1])
            time_helper = bb_time_associate[np.argsort(bb_times)]

            if time_helper[0] == time_helper[1]:
                # no temporal overlap
                continue

            # check freq overlap
            bb_freqs = np.array([f0_0, f0_1, f1_0, f1_1])
            bb_freq_associate = np.array([0, 0, 1, 1])

            freq_helper = bb_freq_associate[np.argsort(bb_freqs)]

            if freq_helper[0] == freq_helper[1]:
                continue

            overlap_bbox_idxs.append((ind0, ind1))

    return np.asarray(overlap_bbox_idxs)


def main(args):
    img_paths = sorted(list(pathlib.Path(args.annotations).absolute().rglob('*.png')))
    df_collect = []

    for img_path in img_paths:
        # convert to time_frequency
        file_name_str, t_min, t_max, f_min, f_max = extract_time_freq_range_from_filename(img_path)

        boxes, scores = bbox_to_data(img_path, t_min, t_max, f_min, f_max ) # t0, t1, f0, f1
        # store values in df
        if not len(boxes) == 0:
            for (t0, f0, t1, f1), s in zip(boxes, scores):
                df_collect.append([file_name_str, t0, f0, t1, f1, s])

    df_collect = np.array(df_collect)

    overlap_bbox_idxs = find_overlapping_bboxes(df_collect)

    bbox_groups = delete_double_boxes(overlap_bbox_idxs, df_collect)

    time_frequency_bboxes = pd.DataFrame(data= np.array(df_collect), columns=['file_name', 't0', 'f0', 't1', 'f1', 'score'])
    time_frequency_bboxes['id'] = np.full(len(time_frequency_bboxes), np.nan)
    time_frequency_bboxes['event_time'] = np.full(len(time_frequency_bboxes), np.nan)

    ###########################################
    # for file_name in time_frequency_bboxes['file_name'].unique():
    #     fig, ax = plt.subplots()
    #
    #     mask = time_frequency_bboxes['file_name'] == file_name
    #     for index, bbox in time_frequency_bboxes[mask].iterrows():
    #         name, t0, f0, t1, f1 = (bbox[0], *bbox[1:-1].astype(float))
    #         if bbox_groups[index] == 0:
    #             color = 'tab:green'
    #         elif bbox_groups[index] > 0:
    #             color = 'tab:red'
    #         else:
    #             color = 'k'
    #
    #         ax.add_patch(
    #             Rectangle((t0, f0),
    #                       (t1 - t0),
    #                       (f1 - f0),
    #                       fill=False, color=color, linestyle='--', linewidth=2, zorder=10)
    #         )
    #     # ax.set_xlim(float(time_frequency_bboxes[mask]['t0'].min()), float(time_frequency_bboxes[mask]['t1'].max()))
    #     ax.set_xlim(0, float(time_frequency_bboxes[mask]['t1'].max()))
    #     # ax.set_ylim(float(time_frequency_bboxes[mask]['f0'].min()), float(time_frequency_bboxes[mask]['f1'].max()))
    #     ax.set_ylim(400, 1200)
    #     plt.show()
    # exit()
    ###########################################

    if args.tracking_data_path:
        file_paths = sorted(list(pathlib.Path(args.tracking_data_path).absolute().rglob('*.raw')))
        for raw_path in file_paths:
            if not raw_path.parent.name in time_frequency_bboxes['file_name'].to_list():
                continue
            time_frequency_bboxes = assign_rises_to_ids(raw_path, time_frequency_bboxes, bbox_groups)
        for raw_path in file_paths:
            # mask = (time_frequency_bboxes['file_name'] == raw_path.parent.name)
            mask = ((time_frequency_bboxes['file_name'] == raw_path.parent.name) & (~np.isnan(time_frequency_bboxes['id'])))
            save_df = pd.DataFrame(time_frequency_bboxes[mask][['t0', 't1', 'f0', 'f1', 'score', 'id', 'event_time']].values, columns=['t0', 't1', 'f0', 'f1', 'score', 'id', 'event_time'])
            save_df['label'] = np.ones(len(save_df), dtype=int)
            save_df.to_csv(raw_path.parent / 'risedetector_bboxes.csv', sep = ',', index = False)
    quit()

def delete_double_boxes(overlap_bbox_idxs, df_collect, overlap_th = 0.2):
    def get_connected(non_regarded_bbox_idx, overlap_bbox_idxs):
        mask = np.array((np.array(overlap_bbox_idxs) == non_regarded_bbox_idx).sum(1), dtype=bool)
        affected_bbox_idxs = np.unique(overlap_bbox_idxs[mask])
        return affected_bbox_idxs

    handled_bbox_idxs = []
    bbox_groups = np.zeros(len(df_collect))
    # detele_bbox_idxs = []

    for Coverlapping_bbox_idx in tqdm(np.unique(overlap_bbox_idxs)):
        if Coverlapping_bbox_idx in handled_bbox_idxs:
            continue

        regarded_bbox_idxs = [Coverlapping_bbox_idx]
        mask = np.array((np.array(overlap_bbox_idxs) == Coverlapping_bbox_idx).sum(1), dtype=bool)
        affected_bbox_idxs = np.unique(overlap_bbox_idxs[mask])

        non_regarded_bbox_idxs = list(set(affected_bbox_idxs) - set(regarded_bbox_idxs))
        # non_regarded_bbox_idxs = list(set(non_regarded_bbox_idxs) - set(handled_bbox_idxs))
        while len(non_regarded_bbox_idxs) > 0:
            non_regarded_bbox_idxs_cp = np.copy(non_regarded_bbox_idxs)
            for non_regarded_bbox_idx in non_regarded_bbox_idxs_cp:
                Caffected_bbox_idxs = get_connected(non_regarded_bbox_idx, overlap_bbox_idxs)
                affected_bbox_idxs = np.unique(np.append(affected_bbox_idxs, Caffected_bbox_idxs))

                regarded_bbox_idxs.append(non_regarded_bbox_idx)
                non_regarded_bbox_idxs = list(set(affected_bbox_idxs) - set(regarded_bbox_idxs))

        bbox_idx_group = np.array(regarded_bbox_idxs)
        bbox_scores = df_collect[bbox_idx_group][:, -1].astype(float)

        bbox_groups[bbox_idx_group] = np.max(bbox_groups) + 1

        remove_idx_combinations = [()]
        remove_idx_combinations_scores = [0]
        for r in range(1, len(bbox_idx_group)):
            remove_idx_combinations.extend(list(itertools.combinations(bbox_idx_group, r=r)))
            remove_idx_combinations_scores.extend(list(itertools.combinations(bbox_scores, r=r)))
        for enu, combi_score in enumerate(remove_idx_combinations_scores):
            remove_idx_combinations_scores[enu] = np.sum(combi_score)
        if len(bbox_idx_group) > 1:
            remove_idx_combinations = [remove_idx_combinations[ind] for ind in np.argsort(remove_idx_combinations_scores)]
            remove_idx_combinations_scores = [remove_idx_combinations_scores[ind] for ind in np.argsort(remove_idx_combinations_scores)]


        for remove_idx in remove_idx_combinations:
            select_bbox_idx_group = list(set(bbox_idx_group) - set(remove_idx))
            time_overlap_pct, freq_overlap_pct = (
                compute_time_frequency_overlap_for_bbox_group(select_bbox_idx_group,df_collect))

            if np.all(np.min([time_overlap_pct, freq_overlap_pct], axis=0) < overlap_th):
                break

        if len(remove_idx) > 0:
            bbox_groups[np.array(remove_idx)] *= -1

        handled_bbox_idxs.extend(bbox_idx_group)

    return bbox_groups

def compute_time_frequency_overlap_for_bbox_group(bbox_idx_group, df_collect):
    time_overlap_pct = np.zeros((len(bbox_idx_group), len(bbox_idx_group)))
    freq_overlap_pct = np.zeros((len(bbox_idx_group), len(bbox_idx_group)))


    for i, j in itertools.product(range(len(bbox_idx_group)), repeat=2):
        if i == j:
            continue
        bb0_idx = bbox_idx_group[i]
        bb1_idx = bbox_idx_group[j]

        bb0_t0, bb0_t1 = df_collect[bb0_idx][1].astype(float), df_collect[bb0_idx][3].astype(float)
        bb1_t0, bb1_t1 = df_collect[bb1_idx][1].astype(float), df_collect[bb1_idx][3].astype(float)

        bb0_f0, bb0_f1 = df_collect[bb0_idx][2].astype(float), df_collect[bb0_idx][4].astype(float)
        bb1_f0, bb1_f1 = df_collect[bb1_idx][2].astype(float), df_collect[bb1_idx][4].astype(float)

        bb_times_idx = np.array([0, 0, 1, 1])
        bb_times = np.array([bb0_t0, bb0_t1, bb1_t0, bb1_t1])
        sorted_bb_times_idx = bb_times_idx[bb_times.argsort()]

        if sorted_bb_times_idx[0] == sorted_bb_times_idx[1]:
            time_overlap_pct[i, j] = 0
        elif sorted_bb_times_idx[1] == sorted_bb_times_idx[2] == 0:
            time_overlap_pct[i, j] = 1
        elif sorted_bb_times_idx[1] == sorted_bb_times_idx[2] == 1:
            time_overlap_pct[i, j] =  (bb1_t1 - bb1_t0) / (bb0_t1 - bb0_t0)
        else:
            time_overlap_pct[i, j] = np.diff(sorted(bb_times)[1:3])[0] / ((bb0_t1 - bb0_t0))

        bb_freqs_idx = np.array([0, 0, 1, 1])
        bb_freqs = np.array([bb0_f0, bb0_f1, bb1_f0, bb1_f1])
        sorted_bb_freqs_idx = bb_freqs_idx[bb_freqs.argsort()]

        if sorted_bb_freqs_idx[0] == sorted_bb_freqs_idx[1]:
            freq_overlap_pct[i, j] = 0
        elif sorted_bb_freqs_idx[1] == sorted_bb_freqs_idx[2] == 0:
            freq_overlap_pct[i, j] = 1
        elif sorted_bb_freqs_idx[1] == sorted_bb_freqs_idx[2] == 1:
            freq_overlap_pct[i, j] =  (bb1_f1 - bb1_f0) / (bb0_f1 - bb0_f0)
        else:
            freq_overlap_pct[i, j] = np.diff(sorted(bb_freqs)[1:3])[0] / ((bb0_f1 - bb0_f0))

    return time_overlap_pct, freq_overlap_pct


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Extract time, frequency and identity association of bboxes')
    parser.add_argument('annotations', nargs='?', type=str, help='path to annotations')
    parser.add_argument('-t', '--tracking_data_path', type=str, help='path to tracking dataa')
    args = parser.parse_args()
    main(args)