{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Working directly with psydat files\n",
    "\n",
    "The recommended way to analyse your data in Python is to use the provided pandas DataFrame of data and statistics.\n",
    "\n",
    "However, if you need to you can access the raw data from which these DataFrames are constructed directly as shown in the examples below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from psychopy.misc import fromFile\n",
    "from shapely.geometry import LineString\n",
    "from shapely.geometry import Polygon\n",
    "from shapely.ops import polygonize\n",
    "from shapely.ops import unary_union"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import\n",
    "\n",
    "A psydat file can be imported using the psychopy `fromFile` function: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psydata = fromFile(\"example.psydat\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Contents\n",
    "\n",
    "This returns a Python object that contains all of the trial conditions and results as attributes.\n",
    "These attributes can be listed using the Python `vars` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for var in vars(psydata):\n",
    "    print(var, end=\" | \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conditions\n",
    "\n",
    "The trial conditions are in `trialList`, each element in this list is a dict of trial conditions that defines a trial:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psydata.trialList"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This can be more easily viewed if converted to a pandas DataFrame:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame(psydata.trialList)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `weight` of a trial is how many times it should be repeated.\n",
    "This information is also stored in the `trialWeights` list, so for example `trialList[i]` will be repeated `trialWeights[i]` times"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psydata.trialWeights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A block consists of doing a trial with each condition in `trialList` `weight` times, where `weight` can be a different number for each trial, so the total number of trials done in a block is then given by\n",
    "\n",
    "$$\n",
    "nTrials = \\sum_i^{nConditions} weight[i]\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This block is then repeared `nReps` times:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psydata.nReps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The total number of trials `nTotal` is then given by\n",
    "\n",
    "$$\n",
    "nTotal = nReps \\times \\left( \\sum_i^{nConditions} weight[i] \\right)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psydata.nTotal"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The condition used for a given trial `iTrial` and repetition number `iRep` is given by `sequenceIndices[iTrial][iRep]`, which gives the index of the conditions used in the `trialList` list:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psydata.sequenceIndices"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `method` specifies the order in which the trials were done:\n",
    "\n",
    "- sequential\n",
    "  - the same order as `trialList`\n",
    "- random\n",
    "  - order of trials shuffled within each block\n",
    "- fullRandom\n",
    "  - order of trials fully shuffled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psydata.method"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Results\n",
    "\n",
    "The results of the trials are in `data` which contains a dict of numpy arrays of recorded data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key, value in psydata.data.items():\n",
    "    print(key, value.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "colors = [\"blue\", \"green\", \"red\", \"cyan\", \"magenta\", \"yellow\", \"black\", \"orange\"]\n",
    "nTrials, nReps = psydata.sequenceIndices.shape\n",
    "\n",
    "\n",
    "def get_axs_row():\n",
    "    row = 0\n",
    "    for condition in range(len(psydata.trialList)):\n",
    "        if not psydata.trialList[condition][\"automove_cursor_to_center\"]:\n",
    "            row += 2 * psydata.trialWeights[condition]\n",
    "        else:\n",
    "            row += psydata.trialWeights[condition]\n",
    "    return row\n",
    "\n",
    "\n",
    "def get_fig_height():\n",
    "    add = 0\n",
    "    for condition in range(len(psydata.trialList)):\n",
    "        if not psydata.trialList[condition][\"automove_cursor_to_center\"]:\n",
    "            add += psydata.trialWeights[condition]\n",
    "    return 6 * (nTrials + add) * nReps\n",
    "\n",
    "\n",
    "fig, axs = plt.subplots(get_axs_row(), nReps, figsize=(6, get_fig_height()))\n",
    "axs = np.reshape(\n",
    "    axs, (get_axs_row(), nReps)\n",
    ")  # ensure axs is a 2d-array even if nTrials or nReps is 1\n",
    "for trial in range(nTrials):\n",
    "    for rep in range(nReps):\n",
    "        loc = (trial, rep)\n",
    "        condition = psydata.sequenceIndices[loc]\n",
    "        target_radius = psydata.trialList[condition][\"target_size\"]\n",
    "        central_target_radius = psydata.trialList[condition][\"central_target_size\"]\n",
    "        ax = axs[loc]\n",
    "        ax.set_title(f\"Trial {trial}, Rep {rep} [Condition {condition}]\")\n",
    "        for positions, target_pos, color in zip(\n",
    "            psydata.data[\"to_target_mouse_positions\"][loc],\n",
    "            psydata.data[\"target_pos\"][loc],\n",
    "            colors,\n",
    "        ):\n",
    "            ax.plot(positions[:, 0], positions[:, 1], color=color)\n",
    "            ax.add_patch(\n",
    "                plt.Circle(\n",
    "                    target_pos,\n",
    "                    target_radius,\n",
    "                    edgecolor=\"none\",\n",
    "                    facecolor=color,\n",
    "                    alpha=0.1,\n",
    "                )\n",
    "            )\n",
    "\n",
    "        # if condition \"automove_cursor_to_center\" is deselected, plot the line to center, fill the enclosed area and output the area\n",
    "        if not psydata.trialList[condition][\"automove_cursor_to_center\"]:\n",
    "            loc_area = (trial + nTrials, rep)\n",
    "            ax_area = axs[loc_area]\n",
    "            ax_area.set_xbound(-0.5, 0.5)\n",
    "            ax_area.set_ybound(-0.5, 0.5)\n",
    "            ax_area.set_title(\n",
    "                f\"area plot for Trial {trial}, Rep {rep} [Condition {condition}]\"\n",
    "            )\n",
    "            print(\"---------------------------------------------\")\n",
    "            print(f\"area of Trial {trial}, Rep {rep} [Condition {condition}]\")\n",
    "            for (\n",
    "                to_target_mouse_positions,\n",
    "                to_center_mouse_positions,\n",
    "                target_pos,\n",
    "                color,\n",
    "            ) in zip(\n",
    "                psydata.data[\"to_target_mouse_positions\"][loc],\n",
    "                psydata.data[\"to_center_mouse_positions\"][loc],\n",
    "                psydata.data[\"target_pos\"][loc],\n",
    "                colors,\n",
    "            ):\n",
    "                ax.plot(\n",
    "                    to_center_mouse_positions[:, 0],\n",
    "                    to_center_mouse_positions[:, 1],\n",
    "                    color=color,\n",
    "                )\n",
    "                coords = np.concatenate(\n",
    "                    (to_target_mouse_positions, to_center_mouse_positions)\n",
    "                )\n",
    "                polygon = Polygon(coords)\n",
    "                lr_coords = np.concatenate((coords[:], to_target_mouse_positions[0:1]))\n",
    "                lr = LineString(lr_coords)\n",
    "                validation = lr.is_valid\n",
    "                multi_LineString = unary_union(lr)\n",
    "                area = 0\n",
    "                for pg in polygonize(multi_LineString):\n",
    "                    area += pg.area\n",
    "                    ax_area.plot(*pg.exterior.xy, color=color)\n",
    "                    ax_area.fill(*pg.exterior.xy, facecolor=color)\n",
    "                ax_area.add_patch(\n",
    "                    plt.Circle(\n",
    "                        target_pos,\n",
    "                        target_radius,\n",
    "                        edgecolor=\"none\",\n",
    "                        facecolor=color,\n",
    "                        alpha=0.1,\n",
    "                    )\n",
    "                )\n",
    "                print(f\"{color}, area: {area}\")\n",
    "\n",
    "            ax.add_patch(\n",
    "                plt.Circle(\n",
    "                    [0, 0],\n",
    "                    central_target_radius,\n",
    "                    edgecolor=\"none\",\n",
    "                    facecolor=\"black\",\n",
    "                    alpha=0.1,\n",
    "                )\n",
    "            )\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each of these is a `nTrials` x `nReps` 2d array, where each element of this array contains the results from the corresponding trial for this variable (which might itself be a single value, e.g. `target_pos`, or an array of values, e.g. `timestamps`)\n",
    "\n",
    "Which set of conditions was used is given by the `sequenceIndices` entry in the same location\n",
    "\n",
    "### Plot of results for each trial\n",
    "\n",
    "For example, a scatter plot of the mouse positions for each trial, labelled by the condition, trial number and repetition number:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot of all trials combined for each condition\n",
    "\n",
    "Here we instead make one plot for each set of conditions in `trialList`, and super-impose all of the corresponding results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "colors = [\"blue\", \"green\", \"red\", \"cyan\", \"magenta\", \"yellow\", \"black\", \"orange\"]\n",
    "nConditions = len(psydata.trialList)\n",
    "nTrials, nReps = psydata.sequenceIndices.shape\n",
    "fig, axs = plt.subplots(nConditions, 1, figsize=(6, 6 * nConditions))\n",
    "axs = np.reshape(axs, (nConditions))  # ensure axs is a 1d-array\n",
    "for trial in range(nTrials):\n",
    "    for rep in range(nReps):\n",
    "        loc = (trial, rep)\n",
    "        condition = psydata.sequenceIndices[loc]\n",
    "        target_radius = psydata.trialList[condition][\"target_size\"]\n",
    "        ax = axs[condition]\n",
    "        ax.set_title(f\"Condition {condition}\")\n",
    "        for positions, target_pos, color in zip(\n",
    "            psydata.data[\"to_target_mouse_positions\"][loc],\n",
    "            psydata.data[\"target_pos\"][loc],\n",
    "            colors,\n",
    "        ):\n",
    "            ax.plot(positions[:, 0], positions[:, 1], color=color)\n",
    "            ax.add_patch(\n",
    "                plt.Circle(\n",
    "                    target_pos,\n",
    "                    target_radius,\n",
    "                    edgecolor=\"none\",\n",
    "                    facecolor=color,\n",
    "                    alpha=0.1,\n",
    "                )\n",
    "            )\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot of mouse movements vs time\n",
    "\n",
    "A plot of `x`, `y`, and `distance from target` versus time for a single move to/from a target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dist(xys, xy0):\n",
    "    return np.sqrt(np.mean(np.power(xys - xy0, 2), axis=1))\n",
    "\n",
    "\n",
    "fig, axs = plt.subplots(1, 2, figsize=(16, 8))\n",
    "\n",
    "trial = 0\n",
    "rep = 0\n",
    "i_target = 3\n",
    "loc = (trial, rep)\n",
    "condition = psydata.sequenceIndices[loc]\n",
    "target_radius = psydata.trialList[condition][\"target_size\"]\n",
    "central_target_radius = psydata.trialList[condition][\"central_target_size\"]\n",
    "for dest, ax in zip([\"target\", \"center\"], axs):\n",
    "    positions = psydata.data[f\"to_{dest}_mouse_positions\"][loc][i_target]\n",
    "    target = psydata.data[\"target_pos\"][loc][i_target] if dest == \"target\" else [0, 0]\n",
    "    times = psydata.data[f\"to_{dest}_timestamps\"][loc][i_target]\n",
    "    ax.set_title(f\"Mouse movements to {dest}\")\n",
    "    ax.plot(times, positions[:, 0], label=\"x\")\n",
    "    ax.plot(times, positions[:, 1], label=\"y\")\n",
    "    ax.plot(times, dist(positions, target), label=\"distance from target\")\n",
    "    ax.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}