{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import pandas as pd\n",
    "import dotenv\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "dotenv.load_dotenv(\".env\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataset import load_hoho_dataset\n",
    "from process_sample import process_sample, read_colmap_rec\n",
    "from example_solutions_copy import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = load_hoho_dataset(testing=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hoho2025.metric_helper import hss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = []\n",
    "for i, sample in enumerate(dataset[\"train\"]):\n",
    "    gt = {\"wf_vertices\": sample[\"wf_vertices\"], \"wf_edges\": sample[\"wf_edges\"]}\n",
    "    \n",
    "    res = process_sample(sample, handle_error=False)\n",
    "    sample_good = convert_entry_to_human_readable(sample)\n",
    "    results.append(res)\n",
    "    print(\"LOSS\", hss(res[\"wf_vertices\"], res[\"wf_edges\"], gt[\"wf_vertices\"], gt[\"wf_edges\"]))\n",
    "    if i + 1 == 1:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = []\n",
    "for i, sample in enumerate(dataset[\"train\"]):\n",
    "    gt = {\"wf_vertices\": sample[\"wf_vertices\"], \"wf_edges\": sample[\"wf_edges\"]}\n",
    "    \n",
    "    res = process_sample(sample, handle_error=False)\n",
    "    sample_good = convert_entry_to_human_readable(sample)\n",
    "    results.append(res)\n",
    "    print(\"LOSS\", hss(res[\"wf_vertices\"], res[\"wf_edges\"], gt[\"wf_vertices\"], gt[\"wf_edges\"]))\n",
    "    if i + 1 == 1:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "good_entry = sample_good\n",
    "rec = sample_good['colmap_binary']\n",
    "for i, (gest, depth, K, R, t, img_id, ade_seg) in enumerate(\n",
    "    zip(\n",
    "        good_entry[\"gestalt\"],\n",
    "        good_entry[\"depth\"],\n",
    "        good_entry[\"K\"],\n",
    "        good_entry[\"R\"],\n",
    "        good_entry[\"t\"],\n",
    "        good_entry[\"image_ids\"],\n",
    "        good_entry[\"ade\"],  # Added ade20k segmentation\n",
    "    )\n",
    "):\n",
    "    if i == 1:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(contour[:, 0, 0], contour[:, 0, 1], c='r')\n",
    "plt.scatter(contour[:, 0, 0], contour[:, 0, 1], c='g')\n",
    "plt.show()\n",
    "angles = []\n",
    "for i in range(len(contour)):\n",
    "    vprev, vcur, vnext = contour[i-1, 0], contour[i, 0], contour[(i+1) % len(contour), 0]\n",
    "    vecprev, vecnext = vcur - vprev, vnext - vcur\n",
    "    vecprev = vecprev / np.linalg.norm(vecprev)\n",
    "    vecnext = vecnext / np.linalg.norm(vecnext)\n",
    "\n",
    "    # curvature angle\n",
    "    angle = np.degrees(np.arctan2(np.cross(vecprev, vecnext), np.dot(vecprev, vecnext)))\n",
    "    angles.append(angle)\n",
    "angles = np.array(angles)\n",
    "plt.plot(angles)\n",
    "plt.show()\n",
    "\n",
    "conv = 10\n",
    "angles_convolved = np.array([np.sum(angles[i: i+conv]) for i in range(len(angles))])\n",
    "plt.plot(angles_convolved)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "metadata": {},
   "outputs": [],
   "source": [
    "angle_len = cv2.arcLength(contour, True) / 50\n",
    "interesting_points = []\n",
    "\n",
    "for i in range(len(angles)):\n",
    "    j = i + 1\n",
    "    while True:\n",
    "        cur_len = cv2.arcLength(slice_arr(contour, i, j), False)\n",
    "        if cur_len > angle_len:\n",
    "            break\n",
    "        j += 1\n",
    "    # i:j is smaller than angle_len\n",
    "    turns = np.cumsum(slice_arr(angles, i, j))\n",
    "    if np.abs(turns).max() > 70:\n",
    "        interesting_points.append(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_turn_angles(contour):\n",
    "    angles = []\n",
    "    vcur = contour[:, 0] # (N, 2)\n",
    "    vprev = np.concatenate([vcur[-1, None], vcur[:-1]]) # (N, 2)\n",
    "    vnext = np.concatenate([vcur[1:], vcur[0, None]]) # (N, 2)\n",
    "\n",
    "    vecprev, vecnext = vcur - vprev, vnext - vcur\n",
    "    vecprev = vecprev / np.linalg.norm(vecprev, axis=1, keepdims=True)\n",
    "    vecnext = vecnext / np.linalg.norm(vecnext, axis=1, keepdims=True)\n",
    "\n",
    "    def dot(a, b):\n",
    "        return (a * b).sum(axis=-1)\n",
    "    angles = np.degrees(np.arctan2(np.cross(vecprev, vecnext), dot(vecprev, vecnext)))\n",
    "    return angles\n",
    "def slice_arr(arr, i, j):\n",
    "    if i <= j:\n",
    "        if j <= len(arr):\n",
    "            return arr[i:j]\n",
    "        else:\n",
    "            return np.concatenate([arr[i:], arr[:j-len(arr)]])\n",
    "    else:\n",
    "        return np.concatenate([arr[i:], arr[:j]])\n",
    "\n",
    "def group_segments(segments):\n",
    "    segments = sorted(segments, key=lambda x: x[0])\n",
    "    grouped = []\n",
    "    for i in range(len(segments)):\n",
    "        if i == 0:\n",
    "            grouped.append(segments[i])\n",
    "        else:\n",
    "            if segments[i][0] <= grouped[-1][1]:\n",
    "                grouped[-1] = (grouped[-1][0], max(grouped[-1][1], segments[i][1]))\n",
    "            else:\n",
    "                grouped.append(segments[i])\n",
    "    return grouped\n",
    "\n",
    "def get_contour_interesting_points_indices(contour):\n",
    "    angles = get_turn_angles(contour)\n",
    "    angle_len = cv2.arcLength(contour, True) / 20\n",
    "\n",
    "    interesting_segments = []\n",
    "    interesting_points = []\n",
    "    for i in range(len(angles)):\n",
    "        j = i + 1\n",
    "        while True:\n",
    "            cur_len = cv2.arcLength(slice_arr(contour, i, j), False)\n",
    "            if cur_len > angle_len:\n",
    "                break\n",
    "            j += 1\n",
    "        # i:j is smaller than angle_len\n",
    "        turns = np.cumsum(slice_arr(angles, i, j))\n",
    "        k = 2\n",
    "        if len(turns) > k and np.abs(turns[k:]).max() > 70:\n",
    "            matching_i = np.where(np.abs(turns[k:]) > 70)[0][0] + k + i\n",
    "            interesting_segments.append((i, int(matching_i)))\n",
    "            interesting_points.append(i)\n",
    "    \n",
    "    grouped_segments = group_segments(interesting_segments)\n",
    "    return [((i + j) // 2) % len(contour) for i, j in grouped_segments]\n",
    "    # return interesting_points\n",
    "def get_contour_interesting_wireframe(contour):\n",
    "    indices = get_contour_interesting_points_indices(contour)\n",
    "    connections = []\n",
    "    for i in range(len(indices)):\n",
    "        i1, i2 = indices[i], indices[(i+1) % len(indices)]\n",
    "        segment_len = np.linalg.norm(contour[i1, 0] - contour[i2, 0])\n",
    "        points_side1 = slice_arr(contour[:, 0], i1, i2)\n",
    "        points_side2 = slice_arr(contour[:, 0], i2, i1)\n",
    "        points_side1_distances = np.array([point_to_segment_dist(p, contour[i1, 0], contour[i2, 0]) for p in points_side1])\n",
    "        points_side2_distances = np.array([point_to_segment_dist(p, contour[i2, 0], contour[i1, 0]) for p in points_side2])\n",
    "        dist_side_1 = points_side1_distances.max()\n",
    "        dist_side_2 = points_side2_distances.max()\n",
    "        factor = 0.1\n",
    "        if dist_side_1 <= segment_len * factor or dist_side_2 <= segment_len * factor:\n",
    "            connections.append((i, (i + 1) % len(indices)))\n",
    "    return contour[indices, 0], np.array(connections)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "angles = get_turn_angles(contour)\n",
    "interesting_points, interesting_connections = get_contour_interesting_wireframe(contour)\n",
    "ip_draw = np.concatenate([interesting_points, [interesting_points[0]]])\n",
    "plt.axis('equal')\n",
    "plt.plot(contour[:, 0, 0], contour[:, 0, 1], c='r')\n",
    "for connection in interesting_connections:\n",
    "    plt.plot(ip_draw[connection, 0], ip_draw[connection, 1], c='g')\n",
    "for i in range(len(interesting_points)):\n",
    "    plt.text(interesting_points[i][0], interesting_points[i][1], str(i), c='b')\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.asarray(gest).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_contours = []\n",
    "keys_segments = [\"eave\", \"ridge\", \"rake\", \"valley\"]\n",
    "all_mask = combine_segs(keys_segments, np.asarray(gest))\n",
    "for key in keys_segments:\n",
    "    mask = combine_segs([key], np.asarray(gest))\n",
    "    contours, _ = cv2.findContours(\n",
    "        mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_KCOS\n",
    "    )\n",
    "    all_contours.extend(contours)\n",
    "# contours = contours[::-1]\n",
    "plt.imshow(all_mask)\n",
    "for contour in all_contours[3:4]:\n",
    "    area = cv2.contourArea(contour, oriented=True)\n",
    "    if area < 0:\n",
    "        contour = contour[::-1]\n",
    "\n",
    "    print()\n",
    "    # plt.plot(contour[:, 0, 0], contour[:, 0, 1], c='r')\n",
    "    # plt.scatter(contour[:, 0, 0], contour[:, 0, 1], c='g', s=3)\n",
    "    interesting_points, interesting_connections = get_contour_interesting_wireframe(contour)\n",
    "    if len(interesting_points) > 0 and len(interesting_connections) > 0:\n",
    "        for connection in interesting_connections:\n",
    "            plt.plot(interesting_points[connection, 0], interesting_points[connection, 1], c='b')\n",
    "    else:\n",
    "        continue\n",
    "    # plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = []\n",
    "\n",
    "for i, sample in enumerate(dataset[\"train\"]):\n",
    "    gt = {\"wf_vertices\": sample[\"wf_vertices\"], \"wf_edges\": sample[\"wf_edges\"]}\n",
    "    \n",
    "    res = process_sample(sample, handle_error=False)\n",
    "    sample = convert_entry_to_human_readable(sample)\n",
    "    results.append(res)\n",
    "    print(\"LOSS\", hss(res[\"wf_vertices\"], res[\"wf_edges\"], gt[\"wf_vertices\"], gt[\"wf_edges\"]))\n",
    "    if i + 1 == 1:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = []\n",
    "for i, sample in enumerate(dataset[\"train\"]):\n",
    "    gt = {\"wf_vertices\": sample[\"wf_vertices\"], \"wf_edges\": sample[\"wf_edges\"]}\n",
    "    \n",
    "    res = process_sample(sample, handle_error=False)\n",
    "    sample = convert_entry_to_human_readable(sample)\n",
    "    results.append(res)\n",
    "    print(\"LOSS\", hss(res[\"wf_vertices\"], res[\"wf_edges\"], gt[\"wf_vertices\"], gt[\"wf_edges\"]))\n",
    "    if i + 1 == 1:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hoho2025.viz3d import *\n",
    "def read_colmap_rec(colmap_data):\n",
    "    import pycolmap\n",
    "    import tempfile,zipfile\n",
    "    import io\n",
    "    with tempfile.TemporaryDirectory() as tmpdir:\n",
    "        with zipfile.ZipFile(io.BytesIO(colmap_data), \"r\") as zf:\n",
    "            zf.extractall(tmpdir)  # unpacks cameras.txt, images.txt, etc. to tmpdir\n",
    "        # Now parse with pycolmap\n",
    "        rec = pycolmap.Reconstruction(tmpdir)\n",
    "        return rec\n",
    "        \n",
    "fig3d = init_figure()\n",
    "plot_reconstruction(fig3d, read_colmap_rec(sample['colmap_binary']))\n",
    "plot_wireframe(fig3d, sample['wf_vertices'], sample['wf_edges'], None, color='red')\n",
    "plot_wireframe(fig3d, res['wf_vertices'], res['wf_edges'], None, color='blue')\n",
    "plot_bpo_cameras_from_entry(fig3d, sample)\n",
    "fig3d"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.10.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}