{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyabc2.sources import load_example\n",
    "\n",
    "tune = load_example(\"For the Love of Music\")\n",
    "tune"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "## Trajectory\n",
    "\n",
    "Something simple we can do is plot the trajectory, as a sort of time series.\n",
    "Ignoring note duration, that looks like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "y = np.array([n.value for n in tune.iter_notes()])\n",
    "x = np.arange(len(y)) * 1/8\n",
    "\n",
    "plt.figure(figsize=(7, 3), layout=\"constrained\")\n",
    "plt.axis(\"off\")\n",
    "plt.title(tune.title)\n",
    "plt.plot(x, y);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "Or, considering duration:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(7, 3), layout=\"constrained\")\n",
    "\n",
    "y = np.array([n.value for n in tune.iter_notes()])\n",
    "x = np.arange(1, len(y) + 1) * 1/8  # shift for consistency\n",
    "\n",
    "plt.plot(x, y, label=\"ignored\")\n",
    "\n",
    "data = np.array(\n",
    "    [\n",
    "        [n.value, float(n.duration)]\n",
    "        for n in tune.iter_notes()\n",
    "    ]\n",
    ")\n",
    "x = data[:,1].cumsum()  # ends of notes\n",
    "y = data[:,0]\n",
    "assert len(x) == len(y)\n",
    "\n",
    "plt.plot(x, y, label=\"considered\")\n",
    "plt.axis(\"off\")\n",
    "plt.title(tune.title)\n",
    "plt.legend(title=\"duration\", loc=\"upper left\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "The divergence occurs due to the 16th notes in the B part."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "## Histogram\n",
    "\n",
    "We can make a histogram of the notes, again considering duration or not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = [\n",
    "    [n.value, float(n.duration), n.to_pitch().unicode()]\n",
    "    for n in tune.iter_notes()\n",
    "]\n",
    "\n",
    "df = pd.DataFrame(data, columns=[\"value\", \"duration\", \"pitch\"])\n",
    "\n",
    "_, ax = plt.subplots(figsize=(6, 3.5), layout=\"constrained\")\n",
    "\n",
    "count = (\n",
    "    df.groupby(\"value\")\n",
    "    .aggregate({\"pitch\": \"first\", \"duration\": \"size\"})\n",
    "    .rename(columns={\"duration\": \"unweighted\"})\n",
    "    .assign(\n",
    "        weighted=(\n",
    "            df.assign(w=df[\"duration\"] * 8)\n",
    "            .groupby(\"value\")[\"w\"].sum()\n",
    "        )\n",
    "    )\n",
    ")\n",
    "\n",
    "count.plot.bar(\n",
    "    x=\"pitch\",\n",
    "    rot=0,\n",
    "    xlabel=\"Pitch\",\n",
    "    ylabel=\"Count\",\n",
    "    title=tune.title,\n",
    "    ax=ax,\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "```{note}\n",
    "F♯₄ is skipped in this plot, which is a bit misleading.\n",
    "It would be better to have an empty place for it.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "(polar-tune-plot)=\n",
    "## Polar\n",
    "\n",
    "We can combine the two (trajectory and histogram) in a polar plot (featured in the readme)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {
    "mystnb": {
     "code_prompt_show": "Some tools"
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "def quadratic_bezier(p1, p2, p3, *, n=200):\n",
    "    \"\"\"Quadratic Bezier curve from start, middle, and end points.\"\"\"\n",
    "    # based on https://stackoverflow.com/a/61385858\n",
    "\n",
    "    (xa, ya), (xb, yb), (xc, yc) = p1, p2, p3\n",
    "\n",
    "    def rect(x1, y1, x2, y2):\n",
    "        a = (y1 - y2) / (x1 - x2)\n",
    "        b = y1 - a * x1\n",
    "        return (a, b)\n",
    "\n",
    "    x1, y1, x2, y2 = xa, ya, xb, yb\n",
    "    a1, b1 = rect(xa, ya, xb, yb)\n",
    "    a2, b2 = rect(xb, yb, xc, yc)\n",
    "\n",
    "    x = np.full((n,), np.nan)\n",
    "    y = x.copy()\n",
    "    for i in range(n):\n",
    "        a, b = rect(x1, y1, x2, y2)\n",
    "\n",
    "        x[i] = i*(x2 - x1)/n + x1\n",
    "        y[i] = a*x[i] + b\n",
    "\n",
    "        x1 += (xb - xa)/n\n",
    "        y1 = a1*x1 + b1\n",
    "        x2 += (xc - xb)/n\n",
    "        y2 = a2*x2 + b2\n",
    "\n",
    "    return x, y\n",
    "\n",
    "\n",
    "def get_polar_ax(key):\n",
    "    from pyabc2 import PitchClass\n",
    "\n",
    "    _, ax = plt.subplots(subplot_kw={\"projection\": \"polar\"})\n",
    "\n",
    "    chromatic_scale_degrees = np.arange(12)\n",
    "    note_labels = []\n",
    "    for csd in chromatic_scale_degrees:\n",
    "        pc = PitchClass(csd + key.tonic.value)\n",
    "        sd = pc.scale_degree_in(key, acc_fmt=\"unicode\")\n",
    "        if pc in key.scale:\n",
    "            s = f\"{sd} ({pc.name})\"\n",
    "        else:\n",
    "            s = sd\n",
    "        note_labels.append(s)\n",
    "\n",
    "    ax.set_theta_offset(np.pi/2)\n",
    "    ax.set_theta_direction(-1)\n",
    "    ax.set_xticks(chromatic_scale_degrees/12 * 2*np.pi)\n",
    "    ax.set_xticklabels(note_labels)\n",
    "    ax.set_rlabel_position(225)  # deg.\n",
    "    ax.xaxis.set_tick_params(pad=7)\n",
    "    ax.set_rlim(0, 3.35)  # TODO: configurable (and/or based on Tune)\n",
    "    ax.set_rticks([1, 2, 3])\n",
    "    ax.set_yticklabels([])\n",
    "\n",
    "    return ax\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "\n",
    "from pyabc2 import Pitch\n",
    "\n",
    "ref = Pitch.from_name(\"G4\")\n",
    "\n",
    "# Compute relative pitch values\n",
    "v = [n.value - ref.value for n in tune.iter_notes()]\n",
    "v = np.array(v)\n",
    "\n",
    "# Set up ax\n",
    "ax = get_polar_ax(tune.key)\n",
    "ax.set_title(tune.title)\n",
    "\n",
    "# Compute and plot trajectory\n",
    "r0, v0 = np.divmod(v, 12)\n",
    "# TODO: try spiral-y version with continuous r (v/12)\n",
    "r = r0 - r0.min() + 1\n",
    "t = v0 * 2 * np.pi / 12\n",
    "cmap = plt.get_cmap(\"plasma\")\n",
    "colors = [cmap(x) for x in np.linspace(0, 1, len(r) - 1)]\n",
    "traj_count = Counter()\n",
    "for c, t1, r1, t2, r2 in zip(colors[:], t[:-1], r[:-1], t[1:], r[1:]):\n",
    "    p1, p3 = (r1*np.cos(t1), r1*np.sin(t1)), (r2*np.cos(t2), r2*np.sin(t2))\n",
    "    if p1 == p3:\n",
    "        continue\n",
    "\n",
    "    traj = ((t1, r1), (t2, r2))\n",
    "    traj_count[traj] += 1\n",
    "\n",
    "    # Calculate isos triangle vertex point\n",
    "    # (to move the traj line out of the way of previous same motions)\n",
    "    mx = (p1[0] + p3[0]) / 2\n",
    "    my = (p1[1] + p3[1]) / 2\n",
    "    mt = np.arctan2(p3[1] - p1[1], p3[0] - p1[0])\n",
    "    h = 0.05 + 0.07 * (traj_count[traj] - 1)\n",
    "    # TODO: ^ would be nice to know total count beforehand, so to set a max for h\n",
    "    p2 = (mx + h*np.cos(mt - np.pi/2), my + h*np.sin(mt - np.pi/2))\n",
    "\n",
    "    # Plot Bezier curve\n",
    "    xb, yb = quadratic_bezier(p1, p2, p3)\n",
    "    rb = np.sqrt(xb**2 + yb**2)\n",
    "    tb = np.arctan2(yb, xb)\n",
    "    ax.plot(tb, rb, c=c, lw=2, alpha=0.4)\n",
    "\n",
    "# Compute and plot histogram data\n",
    "# TODO: optionally weight with duration\n",
    "vc = Counter(v)\n",
    "rc0, vc0 = np.divmod(list(vc), 12)\n",
    "rc = rc0 - rc0.min() + 1\n",
    "tc = vc0 * 2 * np.pi / 12\n",
    "s = np.array(list(vc.values())) * 50\n",
    "ax.scatter(\n",
    "    tc,\n",
    "    rc,\n",
    "    s=s,\n",
    "    marker=\"o\",\n",
    "    zorder=10,\n",
    "    alpha=0.4,\n",
    ");"
   ]
  }
 ],
 "metadata": {
  "execution": {
   "timeout": 60
  },
  "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}