import random
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
from loguru import logger

from src.dataset.qvhighlights import QVHighlights
from src.dataset.utils import get_irrelevant_windows
from src.utils.rw_utils import load_jsonl

# Split from https://github.com/wjun0830/CGDETR/blob/main/data/youtube_uni/youtube_splits.py
YOUTUBE_SPLITS = {
    "dog": {
        "train": [
            "BsjTtq337mM",
            "eGCD1F74iy8",
            "x2Za-t1yHtI",
            "iyYiqa0QZXM",
            "azy9ijU6f9I",
            "NNtSZ6cPiwA",
            "U9CBalvFfbM",
            "AZDkqJaOgJU",
            "-olTgMPAyMI",
            "i35F1Ec3Ats",
            "6bS6-GVLBeM",
            "ZGszTEn28v8",
            "EEb8iSMqwj4",
            "p2hYGNkRMCw",
            "3kbptPDIz4U",
            "iLHRqR-M9HQ",
            "zyooMDuAgCA",
            "dOVsQ63N0gg",
            "7H_qqQvPUzY",
            "Z5BEFsaYIS4",
            "iWO6io44-Fs",
            "vVmGisWK0QI",
            "L10kN7Btk90",
            "2yql1mvWbDs",
            "Iu2nbtr_Uuk",
            "NSmOKAauZpM",
            "PAhQGoURAro",
            "uJ81Us4mBOc",
            "1krGVyfIaOw",
            "p9yW6FxsrJ4",
            "DLGRJfpGmCQ",
            "0XTXKe2TOAg",
            "qpc4OSqeV7I",
            "q_PJFuBOk7k",
            "0Uu53hCnKQ4",
            "-szRD9kyNug",
            "rUPxwWmJYpg",
            "hseONiKKx_8",
            "BLaQcOcDfjo",
            "nW5JulWYEc8",
            "rMvH1SMGwwI",
            "l6KlvTJkTgk",
            "O8j4U3NjNvs",
            "8AJTZeEeStk",
        ],
        "val": [
            "a2nj7XCo2Rk",
            "9rP5yF9EC3Y",
            "OxSsRZqPfyk",
            "bZzP2MieC1c",
            "PcvdX5OVgfQ",
            "p0oxRJD1GUk",
            "msjK8nHZHZ0",
            "hSRyclcZyGM",
            "dlH2K9N_jSM",
            "OCVXhRG2fEA",
            "MkBdHvXPocc",
            "yN7h90Y-04g",
            "PWqLJKZeBC8",
            "9D_Q8l_ruQk",
            "Mp8Pz86J660",
            "1gjntnYm8NA",
            "O3XxuutEvoo",
            "wf_qlAizlSM",
            "fXx44D1sqUw",
            "P0MnXh6bnKk",
            "sTd06idFa0E",
            "ppNjl3I3iJs",
            "Om5mczkpcVg",
            "xZIN_s-qhbU",
        ],
    },
    "gymnastics": {
        "train": [
            "Wfv90YJ2YtA",
            "MbD5OIR9yWc",
            "fZwCJWkC_Qw",
            "AyRI1CioQfY",
            "xV_5YCdVqSM",
            "19UO7T32DJI",
            "o2gAP2Clg_s",
            "ewyfAOrBzjQ",
            "CMTKpA683Ig",
            "aNjphhjTgqs",
            "dmJ0Nq4DF2w",
            "57IQ6EudvGU",
            "BAlUYtPUsVI",
            "_UU4XqYVDqE",
            "Kq4OhBiQk_E",
            "D6nyvx9kEac",
            "g-m4-zeCisU",
            "_45vTFtcduE",
            "9L-Pocc_u70",
            "0636XaURL-A",
            "GCabQyaHSMg",
            "vUi1Scb35fQ",
            "eK-Yuoou_1I",
            "kkS7TgNZwJI",
            "2EFkINKg3nA",
            "eKvALYDh7RU",
            "Hyp3Hpk6dyA",
            "9rpzf3sgQkw",
            "kHNAnpewyeo",
            "ydQij10qrZM",
            "41u2V_ZAKto",
            "6NSWsMKAgEU",
            "kUs_yUR-C2k",
            "bs3ZBcfhvKA",
        ],
        "val": [
            "2AuigNFEsTM",
            "rPsKpHKzUso",
            "tzq5cJQ9NQA",
            "DyZ0gZ5xmxI",
            "PEKRfJYYEgU",
            "affAIVH9uRA",
            "FT7yIi3-tG0",
            "T_zWyrVzyvw",
            "RoiLzMA_ilA",
            "nBZiGSccsTg",
            "z3cNtOMKK7A",
            "EwQ-aMK2sKg",
            "Rq0BpciuvBM",
            "s6LNwTThBgs",
            "-hE9v3izo4c",
            "KldEfRhv7H0",
            "eUyuw2J5FaE",
            "E0aRE1_ea8E",
            "BU7YlQAOBkM",
            "iDJM9j11U-c",
            "zr5LSPMBpiI",
            "NAfBa7lqg2Q",
            "eB4Toq9dUWs",
            "YPd7RDN5CkE",
            "86YLsw7efDM",
            "iQRMMFiYAUw",
            "lzEhLAPxZyQ",
            "PAjJbT1DRnY",
        ],
    },
    "parkour": {
        "train": [
            "qz1UnnxlWhI",
            "MzODICzycHs",
            "0swXWs9yWA4",
            "Nnv22OW_PaI",
            "LUhZJLY2uKc",
            "yZz8z1l3XJU",
            "3dvjtdMC2ls",
            "e27ppPer9XY",
            "HJNn2WlKFhM",
            "j4OxlxnapNI",
            "rhABvn7VjSQ",
            "3PCwXpwYqLs",
            "LECL1bIpi5w",
            "w0ouP79iZWc",
            "z6aKQPMJUC0",
            "kATlFTwxBVY",
            "3SM6a8eyuVA",
            "v-Sfc4COqRQ",
            "64eu8pwuIUE",
            "7WKm0XDk3og",
            "2F5Sc0Jgk4g",
        ],
        "val": [
            "TFdbCRkVeIA",
            "uGLs9atTvNc",
            "qlGPuopK3CI",
            "ucTkpjZO_o4",
            "4-4BgyGphLQ",
            "08k4ysX_XJE",
            "6sMNnWqa_as",
            "oT6g0I2Ok9o",
            "Be4IlnKeBOo",
            "yUjJq0kvxcw",
            "fLek7GRIxjE",
        ],
    },
    "skating": {
        "train": [
            "7owXLUkpoNY",
            "1OLM0_Jzt5M",
            "b1LXb0Sbiy0",
            "3fGux6-ttlA",
            "HQvRun80GyA",
            "a8M-5nTrll8",
            "bA3CxZllhsI",
            "AUAsfZtcB4E",
            "FG57uCJvQLw",
            "jXIuv5uFPTI",
            "eG-hdYLoS98",
            "2SdJBl251PU",
            "2PHJqqrGC80",
            "EtZkkFhniRw",
            "jUiwyguxzIw",
            "FL6mXlaF78Q",
            "BdemklZtYWI",
            "ATk_ncI1-BA",
            "4wiKDfq3X8U",
            "BN7GBjVlFTo",
            "JiMZvMkkbRo",
            "2DIXYkSnRf4",
            "dZ3i-HuhQXM",
            "7jZydh62m8M",
        ],
        "val": [
            "2oOe2_Ew6Ao",
            "DGcO0QgcXtw",
            "ixsKaNplm6o",
            "7TQbqKWjLcI",
            "CQZNrEstSag",
            "g1WbAIzkw80",
            "4cyx1VpDjc4",
            "BGZaaqFjoRY",
            "AJ98A2y1dVw",
            "1n7Afe5AZCM",
            "8x8ESK5MnR0",
        ],
    },
    "skiing": {
        "train": [
            "6Usy87KaF-A",
            "DtjKkp_4KDQ",
            "4Wt7TM2wDxI",
            "iKnzSGFwdbc",
            "nALCc6HPQNs",
            "WL4TA--CVcA",
            "dFrfsgW1M98",
            "x6qmrVojcYc",
            "pvcmQ9J_BYw",
            "S3VEYFAP_pk",
            "pU57a3jYMEk",
            "33TrLdo3ook",
            "xLhHU8uo2aY",
            "fAHBmka6Psc",
            "9HYzZk5kiJA",
            "T0gjqYbeU1g",
            "7o628W-bFy0",
            "YKDm_PCa-HM",
            "R3DV2zDnNqg",
            "NCe9YeXTvHo",
            "5tXxvscmZ-Y",
            "thNiPQLbi5w",
            "1TtJy8cSzqA",
            "zDRzOsmwa08",
            "gCI4gArPjNA",
            "uw0i26NHucs",
            "1giAsZC_ywQ",
            "OvgaPTfEnqo",
            "bFD_p5znoq4",
            "uKmqaAvjKgw",
            "5ivw_sdCTCU",
            "iwCSAYGwPq4",
            "HmmOPntPlRA",
            "FHCEyiM-NoY",
            "EUSFMmoE_jI",
            "igvSxtdsT8w",
            "zEgMYFiEaX4",
            "0K2FKccDp9A",
            "tdyz6h4ZtYs",
            "PO7GEbi2z3c",
            "mmiu7rRmSAU",
            "qL6Kic-CdTo",
            "0fNCsOY1WGk",
            "V3J26hr1ZSE",
            "GS-qBunN3B4",
            "ZLNvg8025Nw",
            "puAxGH6aWMY",
            "h-SlvHubhs8",
            "AdovZ4OAS8I",
            "UDvA1XMa1m4",
            "qdo3d7mR_9s",
            "qAinbyORWIw",
            "v1JpJueAElY",
            "TjH29fdjcqI",
            "f76B1uucoyo",
            "DNPPDcOd5eQ",
            "-GX95udKKm8",
            "YRO_RQ3aBgg",
            "1ptV2E7lm9U",
            "qa7dtf1Qcew",
            "_UJTkqYNrpA",
            "md14DNKq2_o",
            "tpewrb9dDyo",
            "yGoWYi_dHLY",
            "DZ3NRjDHwy8",
            "aMFcEuJUqpk",
            "6fT9KLuE7no",
            "lPdQMMAuOZo",
        ],
        "val": [
            "SSlv7qJK5zA",
            "_BYqZjuKpKA",
            "ZueaKXReGjU",
            "mGST8ZekCZc",
            "JJSu7Lh9rvs",
            "IyoD3G5igY0",
            "MXyv-Ut9HRg",
            "Z8X9WIojH1U",
            "vT33-8KUb2Q",
            "HW6_sPym938",
            "9wtXO2lF6hM",
            "mRdthCqe6Nk",
            "RGxiOb9hlS0",
            "ruySf5zL7Kw",
            "I7wFmP6P7p0",
            "0AHkDElk3ws",
            "zqXd4EgUFhE",
            "91lDbBHUx0w",
            "iaHbK6ogafc",
            "jRbst8kjWW8",
            "drHPy6wSZGs",
            "5VaY6LgIqDs",
            "bXq9rRSbI3c",
            "hjZLa2DTuqs",
            "Ka2qcp3jmWo",
            "ZnA4-ggkFu8",
            "iXdt4v42mbs",
            "8aWN-0NZErI",
            "09v0HNf81J0",
            "YJCR2q-WRhQ",
            "RjagI4pAUpw",
            "_10CbYdTG5M",
            "lhgmIgzBQxs",
            "2pstGBM4p0w",
            "b53-VPsWom4",
            "x-G4r153n6o",
            "qBbqK5qlVSM",
            "XamrS9XyHuQ",
            "u_n7jMS1vlw",
            "AO6p0jlOd6U",
            "m-W-lcTkBQ0",
            "bMuyPVIlXW8",
            "kAAvTAKkIy4",
            "U6vnbCurZQA",
            "dHE8q7sZ70U",
            "w7fzLVRPSUc",
            "FLYkD7zHuHQ",
            "nhOhI24P7dM",
            "n5q2KhfoiWw",
            "7Hcyse0h9HE",
            "6_BPy_VaPSY",
        ],
    },
    "surfing": {
        "train": [
            "Ai9FwQGn5ds",
            "hBl0Sm3_auw",
            "LMxMeg407Vg",
            "D3fk8doVui4",
            "Y9pxmLg6ti8",
            "p_JsivYdbgQ",
            "UokX-hcXQeo",
            "VYe5QfM5ecE",
            "I48VJ92ouTQ",
            "Tn-ebtUnq6E",
            "eWae-nWocPU",
            "-Yamat_0tbw",
            "c2Fy-rdXJy4",
            "xQ4NAp4vWbI",
            "g9kXCIjIjoE",
            "A96Jx6gv6_4",
            "e427qElqqN0",
            "tTcA5hiViPo",
            "wMdXzj_3aA0",
            "fqNzMz1n6uA",
            "jKVOA7RFCUo",
            "TJBJrk9iPPA",
            "_C8EjMxrS2s",
            "yj7abHfZTQQ",
            "NDcqgpsyWaU",
            "UJjwoivaGNo",
            "GZ_XS8EnnWo",
            "kJUBIcBjUZ0",
            "lWoLyR7lDAU",
            "FilbyF_PGjI",
            "fapRkcOe4vE",
            "t05r50PQqww",
            "QgStLppe610",
            "2TY8Q2WXUyk",
            "9y_ED3DyNhE",
            "CGwtinVGkVU",
            "nOuRhrAMaIw",
            "UN4TwjDajtQ",
            "-FHmVZWWgcE",
            "ksx0_BfpsLg",
            "agOBPDsQrTM",
            "XqggBwFOmFU",
            "orNzj1J8i-4",
            "6ZbTCHwt1gk",
            "0un3wh_pQAc",
            "4u6OURBLZDs",
            "us0agAKuvEM",
            "mVQYl7Q-TQs",
            "cB2SdlGHLMQ",
            "WK5t4To0zlA",
            "NNEuH_juUHI",
            "KTU7xfVOat0",
            "Y1nhbNaY1ZY",
            "YlXJnZe575s",
            "SH7Ns0ANzJU",
            "3TbZfeokCkE",
        ],
        "val": [
            "o0on6yIXJQE",
            "4RsZz_8d8Ro",
            "p8VUjcZyK70",
            "0P2PZXUa0Bg",
            "p2eU5z647Mw",
            "mSVxaAJcNJQ",
            "bcmXVyFbsRg",
            "Eiq8GHi4kEo",
            "H5FEdJYokO4",
            "Mkyp0z_Cgig",
            "NB5Ez5kJfMU",
            "Xa0y6b6Vm6U",
            "gVcCGUtpA90",
            "0-fstXuo_Pw",
            "-d72e4v9skA",
            "lbp6_wCXqvw",
            "9GpZHq1n8ps",
            "CefGXyYu_zU",
            "SI2JbS48Upg",
            "hdklRTNrq0I",
            "J-P-t6g19SM",
            "K0f_DpVOjfA",
            "lw_1fEY9QTo",
            "uUuYnKLETLw",
            "HwKv3Xc5MAE",
            "wvQ0h5Nwsxc",
            "l8ME6z_EWKE",
            "s9dTu2fcbNg",
            "GS09SevPYT4",
            "YbwdDCzVczU",
            "jaCOI_VwIjc",
            "3Y1Jp1_fFLQ",
            "82OzgxT2tH8",
            "IjQhHPlTfdE",
            "KzQcJrT91jU",
            "t05AD0c08zE",
            "rGxWxX6nYO4",
            "QGp0kRzKiAc",
            "pK9gDWoOyko",
            "Srjd4pe6vck",
            "twGcxuhCXoU",
            "AshLUHPEb8M",
            "8En3M5CUc2E",
            "8sTJfTUk1d0",
            "o-bubyWTw60",
            "NctbssxGCtU",
            "L09Qo1ql0nM",
        ],
    },
}


# pylint: disable=too-many-instance-attributes
class YouTube(QVHighlights):  # noqa: WPS230,WPS214
    """Dataset for YouTube data.

    YouTube dataset, see more information on https://github.com/aliensunmin/DomainSpecificHighlight

    Code from https://github.com/wjun0830/CGDETR/blob/main/cg_detr/start_end_dataset.py with some adaptations

    """
    # pylint: disable=too-many-arguments
    def __init__(
        self,
        domain: str,
        data_path: str,
        video_feat_dir: str,
        query_feat_dir: str,
        audio_feat_dir: Optional[str] = None,
        max_query_length: int = 32,
        max_video_length: int = 75,
        data_ratio: float = 1,
        normalize_video: bool = True,
        normalize_query: bool = True,
        use_tef: bool = True,
        clip_len: int = 2,
        max_windows: int = 10,
    ) -> None:
        """
        Initialize the dataset.

        Args:
            domain (str): Youtube Domain
            data_path (str): Path to the data.
            video_feat_dir (str): Path to the video features.
            query_feat_dir (str): Path to the query features.
            audio_feat_dir (Optional[str]): Path to the audio features.
            max_query_length (int): Max length of the query tensor. Defaults to 32.
            max_video_length (int): Max length of the video tensor. Defaults to 75.
            data_ratio (float): Portion of the data to use. Defaults to 1.0.
            normalize_video (bool): Whether to norm the video emb or not. Defaults to True.
            normalize_query (bool): Whether to norm the query emb or not. Defaults to True.
            use_tef (bool): Whether to use time positional features or not. Defaults to True.
            clip_len (int): Length of the clip in secs. Defaults to 2.
            max_windows (int): Maximum number of windows to use as labels. Defaults to 10.
        """
        self.domain = domain
        super().__init__(
            data_path,
            video_feat_dir,
            query_feat_dir,
            audio_feat_dir,
            max_query_length,
            max_video_length,
            data_ratio,
            normalize_video,
            normalize_query,
            use_tef,
            clip_len,
            max_windows,
        )

    def load_data(self) -> List[Dict[str, Any]]:
        """Load data from data_path, and filter by data_ratio.

        Returns:
            List[Dict[str, Any]]: List of data.
        """
        datalist = load_jsonl(self.data_path)
        domain_data = []
        for sample in datalist:
            # take data only from the desired domain
            if (  # noqa: WPS337
                sample["qid"] in YOUTUBE_SPLITS[self.domain]["train"]
                or sample["qid"] in YOUTUBE_SPLITS[self.domain]["val"]  # noqa: W503
            ):
                domain_data.append(sample)
        datalist = domain_data
        if self.data_ratio != 1:
            n_examples = int(len(datalist) * self.data_ratio)
            datalist = datalist[:n_examples]
            data_portion = self.data_ratio * 100
            logger.info(f"Using {data_portion}% of the data: {n_examples} examples")
        return datalist

    def get_saliency_labels_all_youtube(
        self,
        gt_scores: np.ndarray,
        ctx_l: int,
        max_n: int = 1,
        add_easy_negative: bool = False,
    ) -> Tuple[List[int], List[int], np.ndarray]:
        """
        Determine saliency labels for video clips by aggregating annotation scores and selecting pos and neg samples.

        This function sums the scores from three annotations for each clip.
        It then selects clips with the highest and lowest scores as positive and negative samples, respectively.
        Additionally, it can sample easy negatives from clips outside the set of relevant clips.

        Args:
            gt_scores (List[List[int]]): A list of lists containing scores from three annotations for each clip.
            ctx_l (int): The total number of clips.
            max_n (int): The number of clips to use as positive and negative. Defaults to 1.
            add_easy_negative (bool): If True, samples easy negatives from clips not in `relevant_clip_ids`.

        Returns:
            Tuple[List[int], List[int], np.ndarray]: A tuple containing three elements:
                - A list of positive clip indices.
                - A list of negative clip indices.
                - An ndarray of aggregated scores, where the score of each clip is at the index corresponding to its ID.

        Note:
            The function uses random sampling which can lead to different results on different executions.
        """
        # Youtube-hl only have binary score
        agg_scores = np.array(gt_scores)[:, 0]  # (L, 1) --> (L, )
        score_array = agg_scores * 1

        sort_indices = np.argsort(agg_scores)  # increasing

        hard_pos_clip_indices = [min(idx, ctx_l - 1) for idx in sort_indices[-max_n:]]
        hard_neg_clip_indices = [min(idx, ctx_l - 1) for idx in sort_indices[:max_n]]
        easy_pos_clip_indices = []  # type: ignore
        easy_neg_clip_indices = []
        if add_easy_negative:
            # conside that relevant clips are those where the average score is more than 1
            rel_clip_ids = np.where(score_array > 1)[0]
            easy_neg_pool = list(set(range(ctx_l)) - set(rel_clip_ids))
            if len(easy_neg_pool) >= max_n:
                easy_pos_clip_indices = random.sample(rel_clip_ids, k=max_n)  # type: ignore
                easy_neg_clip_indices = random.sample(easy_neg_pool, k=max_n)
            else:  # copy the hard ones
                easy_pos_clip_indices = hard_pos_clip_indices
                easy_neg_clip_indices = hard_neg_clip_indices

        pos_clip_indices = hard_pos_clip_indices + easy_pos_clip_indices
        neg_clip_indices = hard_neg_clip_indices + easy_neg_clip_indices

        return pos_clip_indices, neg_clip_indices, score_array

    def __getitem__(self, index: int) -> Dict[str, Any]:
        """Get item by index.

        Args:
            index (int): Index of the item.

        Returns:
            Dict[str, Any]: Dict of meta and model inputs.
        """
        meta = self.data[index]
        meta["irrelevant_windows"] = get_irrelevant_windows(meta["relevant_windows"], meta["duration"])
        model_inputs: Dict[str, Any] = {}
        model_inputs["query_feat"] = self.get_query_feat_by_qid(meta["qid"])
        video_emb = self.get_video_feat_by_vid(meta["vid"])
        audio_emb = self.get_audio_feat_by_vid(meta["vid"])
        video_emb, audio_emb = self.add_fpn_padding(video_emb, audio_emb)
        number_of_clips = len(video_emb)
        meta["duration"] = number_of_clips * self.clip_len
        model_inputs["video_feat"] = video_emb
        if audio_emb is not None:
            model_inputs["audio_feat"] = audio_emb

        if "relevant_windows" in meta:
            # moment retrieval
            model_inputs["span_labels"] = self.get_span_labels(
                meta.get("relevant_windows"),  # type: ignore
                meta.get("duration"),  # type: ignore
            )
            # highlight detection
            saliency_pos, saliency_neg, saliency_all = self.get_saliency_labels_all_youtube(
                meta["saliency_scores"],
                number_of_clips,
            )
            padding_size = video_emb.shape[0] - len(saliency_all)
            saliency_all = np.concatenate([saliency_all, np.zeros(padding_size)])

            model_inputs["saliency_pos_labels"] = saliency_pos
            model_inputs["saliency_neg_labels"] = saliency_neg
            model_inputs["saliency_all_labels"] = saliency_all
            model_inputs["relevant_clip_ids"] = meta["relevant_clip_ids"]
        model_inputs["vid"] = meta["vid"]
        model_inputs["qid"] = meta["qid"]
        return {"meta": meta, "model_inputs": model_inputs}
