{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example Usage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example demonstrates an end-to-end use case of phylogenetic tracking to a population of custom `Organism` objects propagated through asexual reproduction.\n",
    "\n",
    "In this demo, organisms contain HSV-encoded color data and are selected for contrast with other population members.\n",
    "Organisms with the same hue are considered as members of the same taxon.\n",
    "The population is progressed through 1,000 rounds of tournament selection, then the systematics manager is serialized to create an alife-standard data file.\n",
    "After loading serialized data back into the notebook and converting to Biopython, we can visualize hue values over evolutionary history of the population."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Begin by importing necessary packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import colorsys\n",
    "from copy import copy\n",
    "from dataclasses import dataclass\n",
    "\n",
    "import alifedata_phyloinformatics_convert as apc\n",
    "from Bio import Phylo as BioPhylo\n",
    "from colormath import color_conversions, color_diff, color_objects\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "\n",
    "from phylotrackpy import systematics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reproducibility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(1)\n",
    "\n",
    "%load_ext watermark\n",
    "%watermark -iwbmuvg -iv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Patch *numpy* for compatibility with *colormath* package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://github.com/gtaylor/python-colormath/issues/104\n",
    "import numpy\n",
    "\n",
    "\n",
    "def patch_asscalar(a):\n",
    "    return a.item()\n",
    "\n",
    "\n",
    "setattr(numpy, \"asscalar\", patch_asscalar)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Write organism class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class Organism:\n",
    "    hue: float = 0.0\n",
    "    saturation: float = 0.0\n",
    "    value: float = 0.0\n",
    "\n",
    "    def mutate(self: \"Organism\") -> None:\n",
    "        \"\"\"Probabilistically tweak stored color information.\"\"\"\n",
    "        if np.random.rand() < 0.5:\n",
    "            self.hue = np.clip(self.hue + np.random.normal(0, 0.05), 0, 1)\n",
    "        if np.random.rand() < 0.5:\n",
    "            self.saturation = np.clip(\n",
    "                self.saturation + np.random.normal(0, 0.02), 0, 1\n",
    "            )\n",
    "        if np.random.rand() < 0.5:\n",
    "            self.value = np.clip(self.value + np.random.normal(0, 0.02), 0, 1)\n",
    "\n",
    "    def make_offspring(self: \"Organism\") -> \"Organism\":\n",
    "        \"\"\"Return copy of self with mutation applied.\"\"\"\n",
    "        offspring = copy(self)\n",
    "        offspring.mutate()\n",
    "        return offspring\n",
    "\n",
    "    def to_labcolor(self: \"Organism\") -> color_objects.LabColor:\n",
    "        \"\"\"Create colormath `LabColor` object representing stored color data.\"\"\"\n",
    "        as_hsv = color_objects.HSVColor(self.hue, self.saturation, self.value)\n",
    "        return color_conversions.convert_color(as_hsv, color_objects.LabColor)\n",
    "\n",
    "    def calc_distance(self: \"Organism\", other: \"Organism\") -> float:\n",
    "        \"\"\"Calculate color-theoretic distance between own color and other\n",
    "        `Organism`'s color.\"\"\"\n",
    "        return color_diff.delta_e_cie1976(\n",
    "            self.to_labcolor(), other.to_labcolor()\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate fitness values for population members, favoring `Organism`s that contrast other population members."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_fitnesses(organisms: list[Organism]) -> list[float]:\n",
    "    return [\n",
    "        np.max(\n",
    "            [\n",
    "                Organism.calc_distance(organism, other)\n",
    "                for other in np.random.choice(organisms, 10)\n",
    "            ],\n",
    "        )\n",
    "        for organism in organisms\n",
    "    ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up population tracking infrastructure.\n",
    "Use dummy founder to force common ancestry among all population members.\n",
    "\n",
    "Write `reproduce` and `remove` helpers to wrap systematics bookkeeping tasks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "population = [Organism() for _ in range(50)]\n",
    "\n",
    "syst = systematics.Systematics(lambda org: str(org.hue))\n",
    "founder_taxon = syst.add_org(Organism())\n",
    "taxa = {id(org): syst.add_org(org, founder_taxon) for org in population}\n",
    "syst.remove_org(founder_taxon)\n",
    "\n",
    "\n",
    "def reproduce(parent: Organism) -> Organism:\n",
    "    offspring = parent.make_offspring()\n",
    "    parent_taxon = taxa[id(parent)]\n",
    "    taxa[id(offspring)] = syst.add_org(offspring, parent_taxon)\n",
    "    return offspring\n",
    "\n",
    "\n",
    "def remove(org: Organism) -> None:\n",
    "    taxon = taxa[id(org)]\n",
    "    syst.remove_org(taxon)\n",
    "    del taxa[id(org)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run rolling evolutionary loop, one tournament at a time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TOURNAMENT_SIZE = 7\n",
    "NUM_UPDATES = 2000\n",
    "for update in tqdm(range(NUM_UPDATES)):\n",
    "    syst.set_update(update)  # track time in systematics manager\n",
    "\n",
    "    # do one tournament\n",
    "    fitnesses = calc_fitnesses(population)\n",
    "    target_idx = np.random.randint(len(population))\n",
    "    selection = max(\n",
    "        np.random.randint(len(population), size=TOURNAMENT_SIZE),\n",
    "        key=fitnesses.__getitem__,\n",
    "    )\n",
    "\n",
    "    # create offspring and replace target index\n",
    "    offspring = reproduce(population[selection])\n",
    "    remove(population[target_idx])\n",
    "    population[target_idx] = offspring"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export and Visualize"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Output phylogenetic history, including column storing taxon info (hue values)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "syst.add_snapshot_fun(systematics.Taxon.get_info, \"taxinfo\")\n",
    "syst.snapshot(\"/tmp/phylo.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load from file and convert to BioPython."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bp_tree = apc.alife_dataframe_to_biopython_tree(\n",
    "    pd.read_csv(\"/tmp/phylo.csv\"),\n",
    "    setattrs=[\"taxinfo\"],\n",
    "    setup_branch_lengths=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set branch colors and draw."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for node in bp_tree.find_elements():\n",
    "    if hasattr(node, \"taxinfo\"):\n",
    "        rgb_float = colorsys.hsv_to_rgb(node.taxinfo, 1.0, 0.5)\n",
    "        rgb_int = tuple(int(c * 255) for c in rgb_float)\n",
    "        node.color = rgb_int\n",
    "\n",
    "with plt.rc_context({\"lines.linewidth\": 5}):\n",
    "    BioPhylo.draw(bp_tree)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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": 2
}