{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7742fa54",
   "metadata": {},
   "outputs": [],
   "source": [
    "# filter by duration, views and unique subtitle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4cdffaaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import json\n",
    "\n",
    "from suno_utils.harvest.utils import mp_scrape\n",
    "from suno_utils.harvest.youtube.collect import get_search_results, time_limit, parse_duration, TimeoutException\n",
    "from suno_utils.harvest.youtube.constants.base import SUPPORTED_LANGS\n",
    "from suno_utils.harvest.youtube.constants.harvest import HL_LANGS\n",
    "\n",
    "PROXY_URL = (\n",
    "     \"http://brd-customer-hl_98887cab-zone-us_proxy-route_err-block-country-us:\" +\n",
    "     \"3if5he8elfe7@zproxy.lum-superproxy.io:22225\"\n",
    ")\n",
    "\n",
    "DATA_DIR = \"/data2/suno/data/harvest/youtube_ml\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "10b28bbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import youtube_dl\n",
    "from contextlib import redirect_stderr, redirect_stdout\n",
    "import os\n",
    "\n",
    "DEFAULT_GET_SUBTITLES_TIMEOUT_S = 10.0\n",
    "YOUTUBE_BASE_URL = \"https://www.youtube.com/watch?v=\"\n",
    "\n",
    "def get_subtitles(youtube_id, timeout_s=DEFAULT_GET_SUBTITLES_TIMEOUT_S, proxy_url=None):\n",
    "    ydl_options = {\n",
    "        \"listallsubtitles\": True,\n",
    "        \"socket_timeout\": 5.0,\n",
    "    }\n",
    "    if proxy_url is not None:\n",
    "        ydl_options[\"proxy\"] = proxy_url\n",
    "    url = YOUTUBE_BASE_URL + youtube_id\n",
    "    with redirect_stderr(open(os.devnull, \"w\")):\n",
    "        with redirect_stdout(open(os.devnull, \"w\")):\n",
    "            with youtube_dl.YoutubeDL(ydl_options) as ydl:\n",
    "                if timeout_s is not None:\n",
    "                    with time_limit(timeout_s):\n",
    "                        info = ydl.extract_info(url, download=False)\n",
    "                else:\n",
    "                    info = ydl.extract_info(url, download=False)\n",
    "    return sorted(list(set(info[\"subtitles\"].keys())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b1641a2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get seed terms\n",
    "with open(os.path.join(DATA_DIR, \"query_terms\", \"unigrams.json\")) as f:\n",
    "    seed_terms = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2215cb5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# filter based on summary attributes\n",
    "DEFAULT_MIN_DURATION_S = 30\n",
    "DEFAULT_MAX_DURATION_S = 2 * 60 * 60\n",
    "DEFAULT_MIN_VIEWS = 50\n",
    "def filter_by_meta(\n",
    "    search_results, \n",
    "    min_duration_s=DEFAULT_MIN_DURATION_S, \n",
    "    max_duration_s=DEFAULT_MAX_DURATION_S, \n",
    "    min_views=DEFAULT_MIN_VIEWS,\n",
    "):\n",
    "    filtered_results = []\n",
    "    for m in search_results:\n",
    "        try:\n",
    "            duration_s = parse_duration(m[\"duration\"])\n",
    "            if duration_s < min_duration_s or duration_s > max_duration_s:\n",
    "                continue\n",
    "            views = int(re.sub(r\"[^0-9]\", \"\", m[\"views\"]))\n",
    "            if views < min_views:\n",
    "                continue\n",
    "            filtered_results.append(m)\n",
    "        except:\n",
    "            continue\n",
    "    return filtered_results\n",
    "\n",
    "# filter based on CC files\n",
    "DEFAULT_MAX_N_SUBTITLES = 5\n",
    "def filter_by_subtitles(search_results, lang_code, max_n_subtitles=DEFAULT_MAX_N_SUBTITLES):\n",
    "    filtered_results = []\n",
    "    for m in search_results:\n",
    "        try:\n",
    "            subtitle_langs = get_subtitles(m[\"id\"], proxy_url=PROXY_URL)\n",
    "        except:\n",
    "            # failed cause of timeout or age restiction etc\n",
    "            continue\n",
    "        if len(subtitle_langs) > max_n_subtitles:\n",
    "            continue\n",
    "        if not any([subtitle_lang.split(\"-\")[0] == lang_code for subtitle_lang in subtitle_langs]):\n",
    "            continue\n",
    "        filtered_results.append(m)\n",
    "    return filtered_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38470ff8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2 pages takes ~2h on single core with proxy\n",
    "import random\n",
    "import tqdm\n",
    "from collections import defaultdict, Counter\n",
    "\n",
    "\n",
    "MAX_IDS_PER_LANG = 10\n",
    "MAX_N_PAGES = 3\n",
    "\n",
    "random.seed(6006)\n",
    "data = defaultdict(list)\n",
    "counts_per_lang = Counter()\n",
    "# for n_iter in tqdm.tqdm(range(100_000)):\n",
    "for n_iter in tqdm.tqdm(range(10)):\n",
    "    for lang_code in SUPPORTED_LANGS:\n",
    "        if counts_per_lang[lang_code] >= MAX_IDS_PER_LANG:\n",
    "            continue\n",
    "        if len(seed_terms[lang_code]) <= n_iter:\n",
    "            continue\n",
    "        # TODO: do we randomize here or take most common term?\n",
    "        query_term = seed_terms[lang_code][n_iter]\n",
    "        # resolve into HL language\n",
    "        available_hl_langs = [k for k in HL_LANGS.keys() if k.split(\"-\")[0] == lang_code]\n",
    "        if len(available_hl_langs) == 0:\n",
    "            hl_lang = \"en\"\n",
    "        else:\n",
    "            hl_lang = random.choice(available_hl_langs)\n",
    "        try:\n",
    "            search_results = get_search_results(\n",
    "                query_term, hl=hl_lang, max_n_pages=MAX_N_PAGES, proxy_url=PROXY_URL\n",
    "            )\n",
    "        except TimeoutException:\n",
    "            continue\n",
    "        promising_results = filter_by_meta(search_results)\n",
    "        reliable_results = filter_by_subtitles(promising_results, lang_code)\n",
    "        # TODO: download vtt files to a directory?\n",
    "        reliable_ids = [m[\"id\"] for m in reliable_results]\n",
    "        counts_per_lang[lang_code] += len(reliable_ids)\n",
    "        data[lang_code].append((hl_lang, search_results, reliable_ids))\n",
    "    if all(counts_per_lang[lang_code] >= MAX_IDS_PER_LANG for lang_code in SUPPORTED_LANGS):\n",
    "        break\n",
    "    print(f\"STEP {n_iter}:\")\n",
    "    print(counts_per_lang)\n",
    "    print()\n",
    "    with open(\"tmp_data.json\", \"w\") as f:\n",
    "        json.dump(data, f)\n",
    "# TODO: add mp and retries/logging/etc\n",
    "# TODO: do we wanna add a global cache for which subtitles exist for each ID?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11d54927",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a3cd7b7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc3eec7b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab702a25",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1183096",
   "metadata": {},
   "outputs": [],
   "source": [
    "# YDL_OPTS = {\n",
    "# #     \"writesubtitles\": False,\n",
    "# #     \"allsubtitles\": True,\n",
    "# #     \"subtitlesformat\": \"best\",\n",
    "#     \"listallsubtitles\": True,\n",
    "# #     \"socket_timeout\": 5.0,\n",
    "#     \"socket_timeout\": 5.0,\n",
    "#     \"proxy\": US_PROXY,\n",
    "# }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b3d7fd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "YDL_OPTS = {\n",
    "    \"format\": \"worstaudio\",\n",
    "    \"outtmpl\": os.path.join(DATA_DIR, \"audio\", \"%(id)s.%(ext)s\"),\n",
    "    \"writesubtitles\": True,\n",
    "    \"allsubtitles\": True,\n",
    "    \"subtitlesformat\": \"best\",\n",
    "    \"socket_timeout\": 5.0,\n",
    "    \"proxy\": US_PROXY,\n",
    "}\n",
    "MAX_DOWNLOAD_TIME_S = 60 * 2\n",
    "\n",
    "YOUTUBE_BASE_URL = \"https://www.youtube.com/watch?v=\"\n",
    "\n",
    "def _resolve_youtube(youtube_id):\n",
    "    url = YOUTUBE_BASE_URL + youtube_id\n",
    "    out_data = {\n",
    "        \"id\": youtube_id,\n",
    "    }\n",
    "    t0 = time.time()\n",
    "    try:\n",
    "        with redirect_stderr(open(os.devnull, \"w\")):\n",
    "            with redirect_stdout(open(os.devnull, \"w\")):\n",
    "                with youtube_dl.YoutubeDL(YDL_OPTS) as ydl:\n",
    "                    # pre-download to check if we want it (duration, views etc)\n",
    "#                     info = ydl.extract_info(url, download=False)\n",
    "                    if MAX_DOWNLOAD_TIME_S is not None:\n",
    "                        with time_limit(MAX_DOWNLOAD_TIME_S):\n",
    "                            info = ydl.extract_info(url, download=True)\n",
    "                    else:\n",
    "                        info = ydl.extract_info(url, download=True)\n",
    "        _ = info.pop(\"formats\", None)\n",
    "        _ = info.pop(\"thumbnails\", None)\n",
    "        # TODO: figure out language info somehow\n",
    "        _ = info.pop(\"automatic_captions\", None)\n",
    "        out_data[\"success\"] = True\n",
    "        out_data[\"retry\"] = False\n",
    "        out_data[\"meta\"] = info\n",
    "        out_data[\"audio_filename\"] = f\"{info['id']}.{info['ext']}\"\n",
    "    except Exception as e:\n",
    "        out_data[\"success\"] = False\n",
    "        out_data[\"retry\"] = (\n",
    "            \"TimeoutException\" in str(type(e)) or\n",
    "            \"unable to download video data\" in str(e) or \n",
    "            \"No video formats found\" in str(e)\n",
    "        )\n",
    "        out_data[\"fail_type\"] = str(type(e))\n",
    "        out_data[\"fail_message\"] = str(e)\n",
    "    t1 = time.time()\n",
    "    out_data[\"runtime_s\"] = round(t1 - t0, 1)\n",
    "    return out_data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36542f0c",
   "metadata": {},
   "source": [
    "## Playground"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57daef7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TLDR: setting language hl helps but otherwise seed terms in language matter the most"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33ddc67d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: see what we need to do to get the highest fraction videos with ONLY german subtitles\n",
    "#  - proxy in country\n",
    "#  - specify location with google\n",
    "#  - sort by upload time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e5d308cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "import json\n",
    "import re\n",
    "import time\n",
    "import random\n",
    "import urllib.parse\n",
    "\n",
    "def _parse_youtube_search_results(html_str):\n",
    "    start_idx = html_str.index(\"ytInitialData\") + len(\"ytInitialData\") + 3\n",
    "    end_idx = html_str.index(\"};\", start_idx) + 1\n",
    "    json_str = html_str[start_idx:end_idx]\n",
    "    page_info = json.loads(json_str)\n",
    "    primary_content = page_info.get(\"contents\", {}).get(\"twoColumnSearchResultsRenderer\", {}).get(\"primaryContents\")\n",
    "    if primary_content is not None:\n",
    "        content_blocks = primary_content[\"sectionListRenderer\"][\"contents\"]\n",
    "    else:\n",
    "        continuation_content = page_info[\"onResponseReceivedCommands\"][0][\"appendContinuationItemsAction\"]\n",
    "        content_blocks = continuation_content[\"continuationItems\"]\n",
    "    results = []\n",
    "    for contents in content_blocks[:-1]:\n",
    "        for video in contents[\"itemSectionRenderer\"][\"contents\"]:\n",
    "            res = {}\n",
    "            if \"videoRenderer\" in video.keys():\n",
    "                video_data = video.get(\"videoRenderer\", {})\n",
    "                res[\"id\"] = video_data.get(\"videoId\", None)\n",
    "                res[\"title\"] = video_data.get(\"title\", {}).get(\"runs\", [[{}]])[0].get(\"text\", None)\n",
    "                res[\"long_desc\"] = video_data.get(\"descriptionSnippet\", {}).get(\"runs\", [{}])[0].get(\"text\", None)\n",
    "                res[\"channel\"] = video_data.get(\"longBylineText\", {}).get(\"runs\", [[{}]])[0].get(\"text\", None)\n",
    "                res[\"duration\"] = video_data.get(\"lengthText\", {}).get(\"simpleText\", None)\n",
    "                res[\"views\"] = video_data.get(\"viewCountText\", {}).get(\"simpleText\", None)\n",
    "                res[\"publish_time\"] = video_data.get(\"publishedTimeText\", {}).get(\"simpleText\", None)\n",
    "                res[\"url_suffix\"] = video_data.get(\n",
    "                    \"navigationEndpoint\", {}\n",
    "                ).get(\"commandMetadata\", {}).get(\"webCommandMetadata\", {}).get(\"url\", None)\n",
    "                results.append(res)\n",
    "    if \"continuationItemRenderer\" not in content_blocks[-1]:\n",
    "        continuation_token = None\n",
    "    else:\n",
    "        continuation_endpoint = content_blocks[-1][\"continuationItemRenderer\"][\"continuationEndpoint\"]\n",
    "        continuation_token = continuation_endpoint[\"continuationCommand\"][\"token\"]\n",
    "    return continuation_token, results\n",
    "\n",
    "def _get_youtube_search_results(\n",
    "    query_term, \n",
    "    max_n_pages=None,\n",
    "    sp=\"EgQQASgB\",  # videos, subtitles, by relevance\n",
    "    hl=\"en\", \n",
    "    gl=None,\n",
    "    proxy_url=None, \n",
    "):\n",
    "    # NOTE: this returns max ~500 results\n",
    "    query_term = urllib.parse.quote(query_term)\n",
    "    if proxy_url is not None:\n",
    "        proxies = {\n",
    "           \"http\": proxy_url,\n",
    "           \"https\": proxy_url,\n",
    "        }\n",
    "    else:\n",
    "        proxies = None\n",
    "    # prepare url\n",
    "    base_url = \"https://www.youtube.com/results?\"\n",
    "    urls_param_list = [f\"search_query={query_term}\"]\n",
    "    if sp is not None:\n",
    "        urls_param_list.append(f\"sp={sp}\")\n",
    "    if hl is not None:\n",
    "        urls_param_list.append(f\"hl={hl}\")\n",
    "    if gl is not None:\n",
    "        urls_param_list.append(f\"gl={gl}\")\n",
    "    urls_param_str = \"&\".join(urls_param_list)\n",
    "    search_url = base_url + urls_param_str\n",
    "    # grab result pages\n",
    "    all_results = []\n",
    "    seen_ids = set()\n",
    "    continuation_token = \"\"\n",
    "    n_pages = 1\n",
    "    while True:\n",
    "        if n_pages == 1:\n",
    "            harvest_url = search_url\n",
    "        else:\n",
    "            harvest_url = search_url + f\"&ctoken={continuation_token}&hl=de&sp=EgQQASgB\"\n",
    "        out = requests.get(harvest_url, proxies=proxies)\n",
    "        assert(out.ok)\n",
    "        continuation_token, results = _parse_youtube_search_results(out.text)\n",
    "        for m in results:\n",
    "            if m[\"id\"] in seen_ids:\n",
    "                continue\n",
    "            seen_ids.add(m[\"id\"])\n",
    "            all_results.append(m)\n",
    "        if continuation_token is None:\n",
    "            break\n",
    "        if max_n_pages is not None and n_pages >= max_n_pages:\n",
    "            break\n",
    "        n_pages += 1\n",
    "        time.sleep(0.3)\n",
    "    return all_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "572deaea",
   "metadata": {},
   "outputs": [],
   "source": [
    "import youtube_dl\n",
    "from contextlib import redirect_stderr, redirect_stdout\n",
    "import os\n",
    "\n",
    "DE_PROXY_IP = \"194.61.113.135\"\n",
    "PROXY_STR = (\n",
    "     f\"http://brd-customer-hl_98887cab-zone-data_center-route_err-block-country-de-ip-{DE_PROXY_IP}:\" +\n",
    "     \"ote0iqh3psfm@zproxy.lum-superproxy.io:22225\"\n",
    ")\n",
    "\n",
    "\n",
    "YDL_OPTS = {\n",
    "#     \"writesubtitles\": False,\n",
    "#     \"allsubtitles\": True,\n",
    "#     \"subtitlesformat\": \"best\",\n",
    "    \"listallsubtitles\": True,\n",
    "#     \"socket_timeout\": 5.0,\n",
    "#     \"proxy\": PROXY_STR,\n",
    "}\n",
    "\n",
    "YOUTUBE_BASE_URL = \"https://www.youtube.com/watch?v=\"\n",
    "\n",
    "def get_subtitles(youtube_id):\n",
    "    url = YOUTUBE_BASE_URL + youtube_id\n",
    "    with redirect_stderr(open(os.devnull, \"w\")):\n",
    "        with redirect_stdout(open(os.devnull, \"w\")):\n",
    "            with youtube_dl.YoutubeDL(YDL_OPTS) as ydl:\n",
    "                # pre-download to check if we want it (duration, views etc)\n",
    "                info = ydl.extract_info(url, download=False)\n",
    "    return sorted(list(set(info[\"subtitles\"].keys())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a4fec2b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm\n",
    "TEST_WORDS = [\n",
    "    \"banane\",\n",
    "    \"kürbis\",\n",
    "    \"schaukelpferd\",\n",
    "    \"christbaum\",\n",
    "    \"blind\",\n",
    "]\n",
    "cache_dict = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8669cbbb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████| 5/5 [06:40<00:00, 80.09s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "43 115 192 192\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "n_tot = 0\n",
    "n_has_any = 0\n",
    "n_has_de = 0\n",
    "n_only_de = 0\n",
    "for word in tqdm.tqdm(TEST_WORDS):\n",
    "    search_results = _get_youtube_search_results(\n",
    "        word, \n",
    "        max_n_pages=2, \n",
    "#         proxy_url=PROXY_STR,\n",
    "#         gl=\"DE\",\n",
    "        hl=\"de\",\n",
    "        sp=\"EgQQASgB\",\n",
    "#         sp=\"CAISAigB\",  # sort by upload_date\n",
    "    )\n",
    "    for m in search_results:\n",
    "        n_tot += 1\n",
    "        if m[\"id\"] in cache_dict:\n",
    "            subtitle_langs = cache_dict[m[\"id\"]]\n",
    "        else:\n",
    "            subtitle_langs = get_subtitles(m[\"id\"])\n",
    "            cache_dict[m[\"id\"]] = subtitle_langs\n",
    "        if len(subtitle_langs) >= 1:\n",
    "            n_has_any += 1\n",
    "        if \"de\" in subtitle_langs:\n",
    "            n_has_de += 1\n",
    "            if len(subtitle_langs) == 1:\n",
    "                n_only_de += 1\n",
    "print(n_only_de, n_has_de, n_has_any, n_tot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c65c111d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 43 115 192 192"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52ad98c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: test adding hl to page 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "250d2a8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# bare case\n",
    "# 24 95 192 192\n",
    "\n",
    "# hl=de\n",
    "# 41 116 191 191\n",
    "\n",
    "# hl=de_DE\n",
    "# 42 118 193 193\n",
    "\n",
    "# hl=de_DE, gl=DE\n",
    "# 42 123 191 191\n",
    "\n",
    "# hl=de_DE, proxy\n",
    "# 42 105 192 192\n",
    "\n",
    "# hl=de_DE, gl=DE, proxy\n",
    "# 44 113 197 197\n",
    "\n",
    "# hl=de_DE, gl=DE, proxy, upload_date\n",
    "# 44 89 159 159"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5f8f5a9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cce7110",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d7a3b0a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f24025a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}