{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test multi-target Mapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pixasonics.core import App, Mapper\n",
    "from pixasonics.features import MeanChannelValue\n",
    "from pixasonics.synths import Theremin, FilteredNoise, SimpleFM\n",
    "from IPython.display import display, Audio"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multi-target, single channel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "app0 = App()\n",
    "app0.load_image_file(\"images/test.jpg\")\n",
    "mean_red = MeanChannelValue(filter_channels=0, name=\"MeanRed\")\n",
    "app0.attach(mean_red)\n",
    "theremin0 = Theremin()\n",
    "app0.attach(theremin0)\n",
    "filtered_noise0 = FilteredNoise()\n",
    "app0.attach(filtered_noise0)\n",
    "red2freqs = Mapper(mean_red, [theremin0[\"frequency\"], filtered_noise0[\"cutoff\"]], name=\"Red2Freqs\")\n",
    "app0.attach(red2freqs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "red2freqs.in_low, red2freqs.in_high, red2freqs.out_low, red2freqs.out_high, red2freqs.exponent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "red2freqs.out_low = [60, 60]\n",
    "red2freqs.out_low"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "red2freqs.out_high = [4000, 4000]\n",
    "red2freqs.out_high"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "red2freqs.exponent = [2, 2]\n",
    "red2freqs.exponent = 2\n",
    "red2freqs.exponent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test NRT\n",
    "duration = 5\n",
    "my_timeline = [\n",
    "    (0, {\n",
    "        \"probe_width\": 50,\n",
    "        \"probe_height\": 50,\n",
    "        \"probe_x\": 0,\n",
    "        \"probe_y\": 0\n",
    "    }),\n",
    "    (duration, {\n",
    "        \"probe_x\": 499,\n",
    "        \"probe_y\": 499,\n",
    "    })\n",
    "]\n",
    "\n",
    "target_filename = \"multitarget_nrt_test.wav\"\n",
    "\n",
    "app0.render_timeline_to_file(my_timeline, target_filename)\n",
    "\n",
    "display(Audio(target_filename))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multi-target, multi-channel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "app1 = App()\n",
    "app1.load_image_file(\"images/test.jpg\")\n",
    "mean_pix = MeanChannelValue()\n",
    "app1.attach(mean_pix)\n",
    "num_channels = 3\n",
    "theremin1 = Theremin([440 for _ in range(num_channels)])\n",
    "app1.attach(theremin1)\n",
    "filtered_noise1 = FilteredNoise(cutoff=[440 for _ in range(num_channels)])\n",
    "app1.attach(filtered_noise1)\n",
    "red2freqs1 = Mapper(mean_pix, [theremin1[\"frequency\"], filtered_noise1[\"cutoff\"]], name=\"Red2Freqs1\")\n",
    "app1.attach(red2freqs1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test NRT\n",
    "duration = 5\n",
    "my_timeline = [\n",
    "    (0, {\n",
    "        \"probe_width\": 50,\n",
    "        \"probe_height\": 50,\n",
    "        \"probe_x\": 0,\n",
    "        \"probe_y\": 0\n",
    "    }),\n",
    "    (duration, {\n",
    "        \"probe_x\": 499,\n",
    "        \"probe_y\": 499,\n",
    "    })\n",
    "]\n",
    "\n",
    "target_filename = \"multitarget_nrt_test.wav\"\n",
    "\n",
    "app1.render_timeline_to_file(my_timeline, target_filename)\n",
    "\n",
    "display(Audio(target_filename))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test custom Feature + custom Mapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pixasonics.core import App, Mapper\n",
    "from pixasonics.features import MeanChannelValue, Feature\n",
    "from pixasonics.synths import Theremin, FilteredNoise, SimpleFM\n",
    "from IPython.display import display, Audio\n",
    "import numpy as np\n",
    "import os\n",
    "import json\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "from sklearn.neighbors import KDTree\n",
    "from sklearn.decomposition import IncrementalPCA\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "app = App()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "app.load_image_file(\"images/cellular_dataset/merged_8bit/Timepoint_001_220518-ST_C03_s1.jpg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# combine red and green channels and all layers\n",
    "img_folder = \"images/cellular_dataset/single_channel_16bit/\"\n",
    "img_files = os.listdir(img_folder)\n",
    "imgs_red = [f for f in img_files if f.endswith(\"w2.TIF\")] # only red channel images\n",
    "imgs_green = [f for f in img_files if f.endswith(\"w1.TIF\")] # only green channel images\n",
    "imgs = []\n",
    "for img_red, img_green in zip(imgs_red, imgs_green):\n",
    "    img_path_red = os.path.join(img_folder, img_red)\n",
    "    img_path_green = os.path.join(img_folder, img_green)\n",
    "    img_red = Image.open(img_path_red)\n",
    "    img_green = Image.open(img_path_green)\n",
    "    img_red = np.array(img_red)\n",
    "    img_green = np.array(img_green)\n",
    "    img = np.stack([img_red, img_green], axis=-1) # now the last dimension is the channel dimension\n",
    "    imgs.append(img)\n",
    "img = np.stack(imgs, axis=-1) # now the last dimension is the layer dimension\n",
    "print(img.shape)\n",
    "app.load_image_data(img) # load as numpy array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# synthmaps_pca_mel_json = \"/Volumes/T7RITMO/synthmaps_code/data/pca_mels_mean.json\"\n",
    "synthmaps_pca_mel_json = \"/Volumes/T7RITMO/synthmaps_code/data/pca_perceptual.json\"\n",
    "# synthmaps_pca_mel_json = \"/Volumes/T7RITMO/synthmaps_code/data/pca_encodec.json\"\n",
    "# synthmaps_pca_mel_json = \"/Volumes/T7RITMO/synthmaps_code/data/pca_clap.json\"\n",
    "with open(synthmaps_pca_mel_json, \"r\") as f:\n",
    "    pca_mel_data = json.load(f)\n",
    "print(pca_mel_data.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fluid_dataset2array(\n",
    "        dataset: dict,\n",
    ") -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Convert a json dataset to a numpy array.\n",
    "\n",
    "    Args:\n",
    "        dataset (dict): The json dataset to convert.\n",
    "\n",
    "    Returns:\n",
    "        np.ndarray: The numpy array.\n",
    "    \"\"\"\n",
    "    num_cols = dataset[\"cols\"]\n",
    "    num_rows = len(dataset[\"data\"])\n",
    "    out_array = np.zeros((num_rows, num_cols))\n",
    "    for i in range(num_rows):\n",
    "        out_array[i] = np.array(dataset[\"data\"][str(i)])\n",
    "    return out_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pca_mel_data_array = fluid_dataset2array(pca_mel_data)\n",
    "print(pca_mel_data_array.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "synthmaps_scaler = MinMaxScaler()\n",
    "pca_mel_data_scaled = synthmaps_scaler.fit_transform(pca_mel_data_array)\n",
    "print(pca_mel_data_scaled.shape)\n",
    "print(pca_mel_data_scaled.min(), pca_mel_data_scaled.max())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fm_params_json = \"/Volumes/T7RITMO/synthmaps_code/data/fm_params.json\"\n",
    "with open(fm_params_json, \"r\") as f:\n",
    "    fm_params_data = json.load(f)\n",
    "print(fm_params_data.keys())\n",
    "fm_params_data_array = fluid_dataset2array(fm_params_data)\n",
    "print(fm_params_data_array.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PCA2D(Feature):\n",
    "    def __init__(self, name=\"PCA2D\"):\n",
    "        super().__init__(name=name)\n",
    "        self.pca = None\n",
    "        self.pca_scaler = None\n",
    "        self._original_shape = None\n",
    "        self._transformed_points = None\n",
    "\n",
    "    def _reshape_for_pca(self, mat):\n",
    "        \"\"\"Helper to reshape 4D matrix (H, W, Ch, L) to 2D by concatenating the Channel and Layer dimensions\"\"\"\n",
    "        mat_reshaped = mat.reshape(mat.shape[0], mat.shape[1], -1)\n",
    "        return mat_reshaped.reshape(-1, mat_reshaped.shape[-1])\n",
    "\n",
    "    def process_image(self, mat):\n",
    "        self._original_shape = mat.shape\n",
    "        features = self._reshape_for_pca(mat)\n",
    "        self.pca = IncrementalPCA(n_components=2)\n",
    "        self.pca.fit(features)\n",
    "        self.pca_scaler = MinMaxScaler(feature_range=(0.1, 0.9))\n",
    "        self._transformed_points = self.pca.transform(features)\n",
    "        self.pca_scaler.fit(self._transformed_points)\n",
    "        projected_scaled = self.pca_scaler.transform(self._transformed_points)\n",
    "        return projected_scaled.T\n",
    "    \n",
    "    def compute(self, mat):\n",
    "        if self.pca is None:\n",
    "            raise ValueError(\"PCA model has not been fitted. Call process_image first.\")\n",
    "        features = self._reshape_for_pca(mat)\n",
    "        projected = self.pca.transform(features)\n",
    "        projected_scaled = self.pca_scaler.transform(projected)\n",
    "        projected_scaled_mean = projected_scaled.mean(axis=0, keepdims=False)\n",
    "        return projected_scaled_mean\n",
    "\n",
    "\n",
    "pca_2d = PCA2D()\n",
    "app.attach(pca_2d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot the pca space\n",
    "plt.figure(figsize=(10, 5))\n",
    "\n",
    "plt.subplot(121)\n",
    "\n",
    "points = pca_2d._transformed_points\n",
    "print(points.shape)\n",
    "plt.scatter(points[:, 0], points[:, 1], alpha=0.5)\n",
    "plt.title('PCA Space')\n",
    "plt.xlabel('PCA 1')\n",
    "plt.ylabel('PCA 2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fm = SimpleFM()\n",
    "app.attach(fm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PCA2FMParams(Mapper):\n",
    "    def __init__(self, source, target_fm, name=\"PCA2FMParams\"):\n",
    "        super().__init__(source, [target_fm[\"carrier_freq\"], target_fm[\"harm_ratio\"], target_fm[\"mod_index\"]], name=name)\n",
    "        self.kdtree = KDTree(pca_mel_data_scaled)\n",
    "\n",
    "    def map(self, in_data):\n",
    "        nearest_idx = self.kdtree.query(in_data.T, return_distance=False)[0][0]\n",
    "        fm_params = fm_params_data_array[nearest_idx]\n",
    "        return fm_params\n",
    "\n",
    "pca2fm_params = PCA2FMParams(pca_2d, fm)\n",
    "app.attach(pca2fm_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom synth example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pixasonics.core import App, Mapper\n",
    "from pixasonics.features import MeanChannelValue\n",
    "from pixasonics.synths import Synth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "app = App()\n",
    "app.load_image_file(\"images/test.jpg\")\n",
    "mean_pix = MeanChannelValue()\n",
    "app.attach(mean_pix)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Custom synth: AM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import signalflow as sf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleAMPatch(sf.Patch):\n",
    "    def __init__(self, carr_freq=440, mod_freq=1, mod_depth=0.25):\n",
    "        super().__init__()\n",
    "        carr_freq = self.add_input(\"carrier_freq\", carr_freq)\n",
    "        mod_freq = self.add_input(\"mod_freq\", mod_freq)\n",
    "        mod_depth = self.add_input(\"mod_depth\", mod_depth)\n",
    "        modulator = (sf.SineOscillator(mod_freq) * mod_depth) + (1 - mod_depth)\n",
    "        carrier = sf.SineOscillator(carr_freq)\n",
    "        out = carrier * modulator\n",
    "        self.set_output(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spec = SimpleAMPatch().to_spec()\n",
    "spec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### One-off version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a dict for params\n",
    "am_params = {\n",
    "    \"carrier_freq\": {\n",
    "        \"min\": 20,\n",
    "        \"max\": 8000,\n",
    "        \"unit\": \"Hz\",\n",
    "        \"scale\": \"log\"\n",
    "    },\n",
    "    \"mod_freq\": {\n",
    "        \"min\": 0.1,\n",
    "        \"max\": 100,\n",
    "        \"unit\": \"Hz\",\n",
    "        \"scale\": \"log\"\n",
    "    },\n",
    "    \"mod_depth\": {\n",
    "        \"min\": 0,\n",
    "        \"max\": 1,\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simple_am_synth = Synth(spec, am_params, name=\"SimpleAM\", add_amplitude=True, add_panning=True)\n",
    "app.attach(simple_am_synth)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Long-term version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleAM(Synth):\n",
    "    def __init__(\n",
    "            self,\n",
    "            carrier_frequency=440,\n",
    "            modulator_frequency=1,\n",
    "            modulator_depth=0.25,\n",
    "            name=\"SimpleAM\"\n",
    "    ):\n",
    "        _spec = SimpleAMPatch(\n",
    "                carr_freq=carrier_frequency,\n",
    "                mod_freq=modulator_frequency,\n",
    "                mod_depth=modulator_depth\n",
    "        ).to_spec()\n",
    "        _params = {\n",
    "            \"carrier_freq\": {\n",
    "                \"min\": 20,\n",
    "                \"max\": 8000,\n",
    "                \"unit\": \"Hz\",\n",
    "                \"scale\": \"log\"\n",
    "            },\n",
    "            \"mod_freq\": {\n",
    "                \"min\": 0.1,\n",
    "                \"max\": 100,\n",
    "                \"unit\": \"Hz\",\n",
    "                \"scale\": \"log\"\n",
    "            },\n",
    "            \"mod_depth\": {\n",
    "                \"min\": 0,\n",
    "                \"max\": 1,\n",
    "            }\n",
    "        }\n",
    "        # call the parent constructor\n",
    "        super().__init__(_spec, params_dict=_params, name=name, add_amplitude=True, add_panning=True)\n",
    "\n",
    "def __repr__(self):\n",
    "    return f\"SimpleAM {self.id}: {self.name}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simple_am = SimpleAM()\n",
    "app.attach(simple_am)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simple_am.ui"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom Mapper example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyRGB2AMMapper(Mapper):\n",
    "    def __init__(self, source, target_am, name=\"MyRGB2AMMapper\"):\n",
    "        super().__init__(source, [target_am[\"carrier_freq\"], target_am[\"mod_freq\"], target_am[\"mod_depth\"]], name=name)\n",
    "    \n",
    "    def map(self, in_data):\n",
    "        r, g, b = in_data\n",
    "        # map red to carrier frequency\n",
    "        carr_freq = np.interp(r, [0, 255], [20, 8000])\n",
    "        # map green to modulator frequency\n",
    "        mod_freq = np.interp(g, [0, 255], [0.1, 100])\n",
    "        # map blue to modulator depth\n",
    "        mod_depth = np.interp(b, [0, 2], [0, 1])\n",
    "        return [carr_freq, mod_freq, mod_depth]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_rgb_2_am = MyRGB2AMMapper(mean_pix, simple_am)\n",
    "app.attach(my_rgb_2_am)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test scaling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from pixasonics.core import App, Mapper\n",
    "from pixasonics.features import MeanChannelValue\n",
    "from pixasonics.synths import Theremin\n",
    "from pixasonics.utils import scale_array_exp\n",
    "from ipywidgets import interact\n",
    "import ipywidgets as widgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "app_scaletest = App()\n",
    "app_scaletest.load_image_file(\"images/test.jpg\")\n",
    "mean_red = MeanChannelValue(filter_channels=0, name=\"MeanRed\")\n",
    "app_scaletest.attach(mean_red)\n",
    "theremin = Theremin()\n",
    "app_scaletest.attach(theremin)\n",
    "red2freq = Mapper(mean_red, [theremin[\"frequency\"]], name=\"Red2Freqs\")\n",
    "app_scaletest.attach(red2freq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "red2freq.in_low, red2freq.in_high, red2freq.out_low, red2freq.out_high, red2freq.exponent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "red2freq.exponent = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_scale(in_val: float, exponent: float) -> float:\n",
    "    return scale_array_exp(\n",
    "        np.array([in_val], dtype=np.float64),\n",
    "        np.array([0], dtype=np.float64),\n",
    "        np.array([255], dtype=np.float64),\n",
    "        np.array([60], dtype=np.float64),\n",
    "        np.array([4000], dtype=np.float64),\n",
    "        exponent\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interact(test_scale, in_val=(0, 255, 0.01), exponent=(1, 2, 0.01))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pixasonics",
   "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
