{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimization of spectral model\n",
    "\n",
    "## Prelude"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "\n",
    "import numpy as np\n",
    "import spectra\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.ndimage\n",
    "import struct\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "from ipywidgets import interact\n",
    "from scipy.ndimage.filters import gaussian_filter\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "from scipy.interpolate import RegularGridInterpolator\n",
    "import os\n",
    "\n",
    "def plot_spectrum(coeffs):\n",
    "    \"\"\" Convenience routine for plotting fitted spectra \"\"\"\n",
    "    coeffs = np.asarray(coeffs).squeeze()\n",
    "    wavelengths = np.linspace(360, 830, 1000)\n",
    "    plt.plot(wavelengths, spectra.eval_spec(coeffs, wavelengths))\n",
    "    plt.suptitle(\"RGB value: \" + str(spectra.eval_rgb(coeffs)))\n",
    "\n",
    "# xy coordinates of the gamut\n",
    "def rgb2xy(value):\n",
    "    value = np.asarray(value)\n",
    "    xyz = (spectra.rgb_to_xyz() @ value[..., None])[..., 0]\n",
    "    return xyz[..., 0:2] / np.sum(xyz, axis=-1)[..., None]\n",
    "\n",
    "# Divide by largest component\n",
    "def normalize(x):\n",
    "    return (x / np.max(x, axis=-1)[..., np.newaxis]).squeeze()\n",
    "\n",
    "def pcolormesh(x, y, z, vmax):\n",
    "    if x.shape == z.shape:\n",
    "        pcol=plt.pcolormesh(x, y, z, vmin=0, vmax=vmax, linewidth=0,rasterized=True)\n",
    "    else:\n",
    "        c = np.array(z)[:-1, :-1, :]\n",
    "        c = c.reshape(c.shape[0] * c.shape[1], 3)\n",
    "        c = np.insert(c, 3, 1.0, axis=1)\n",
    "        plt.pcolormesh(x, y, z[..., 0], vmin=0, vmax=vmax, color=c, edgecolors='none', shading='gouraud')\n",
    "\n",
    "def minmax(x):\n",
    "    return np.min(x), np.max(x)\n",
    "\n",
    "def add_colorbar(im, ax, pad=0.01):\n",
    "    cax = make_axes_locatable(ax).append_axes('right', size='5%', pad=pad)        \n",
    "    return plt.colorbar(im, cax=cax)\n",
    "    \n",
    "def write_file(fname, scale, coeffs):\n",
    "    print('Writing \"%s\" ..' % fname)\n",
    "    with open(fname, \"wb\") as f:\n",
    "        f.write('SPEC'.encode('utf8'))\n",
    "        f.write(struct.pack('i', coeffs.shape[1]))\n",
    "        np.float32(scale).tofile(f)\n",
    "        np.float32(coeffs).tofile(f)\n",
    "\n",
    "def read_file(fname):\n",
    "    # Usage: res, coeffs, target = read_file(prefix + '/' + prefix + \".coeff\")\n",
    "    print('Loading \"%s\" ..' % fname)\n",
    "    with open(fname, \"rb\") as f:\n",
    "        assert f.read(4).decode('utf8') == 'SPEC'\n",
    "        res = struct.unpack('i', f.read(4))[0]\n",
    "        scale = np.fromfile(f, count=res, dtype=np.float32)\n",
    "        coeffs = np.fromfile(f, count=3*res*res*res*3, dtype=np.float32)\n",
    "        coeffs = coeffs.reshape(3, res, res, res, 3)\n",
    "        scale_bcast = scale[:, None, None, None]\n",
    "        t = np.linspace(0, 1, res)\n",
    "        x, y = np.meshgrid(t, t)\n",
    "        target = np.stack([\n",
    "            np.dstack((np.full(x.shape, 1.0), x, y))[np.newaxis, ...] * scale_bcast,\n",
    "            np.dstack((y, np.full(x.shape, 1.0), x))[np.newaxis, ...] * scale_bcast,\n",
    "            np.dstack((x, y, np.full(x.shape, 1.0)))[np.newaxis, ...] * scale_bcast\n",
    "        ])\n",
    "        return res, coeffs, target\n",
    "\n",
    "ndim = 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Select target gamut"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Target resolution\n",
    "res = 64\n",
    "\n",
    "# Let's fit sRGB spectra first\n",
    "# gamut = spectra.Gamut.SRGB\n",
    "# gamut = spectra.Gamut.ProPhotoRGB\n",
    "# gamut = spectra.Gamut.ACES2065_1\n",
    "gamut = spectra.Gamut.REC2020\n",
    "\n",
    "spectra.init_tables(gamut)\n",
    "prefix = str(gamut).split('.')[-1].lower()\n",
    "try:\n",
    "    os.mkdir(prefix)\n",
    "except:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Let's check if the solver works using a single RGB value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.zeros((1,1,ndim)) # Starting guess\n",
    "spectra.optimize(\n",
    "    [[[0.79264853,  0.4 , 0.63703843]]], x\n",
    ")\n",
    "plt.figure(figsize=(5, 3))\n",
    "plot_spectrum(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Compute cube of target RGB values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def smoothstep(x):\n",
    "    return x * x * (3 - 2 * x)\n",
    "\n",
    "t = np.linspace(0, 1, res)\n",
    "x, y = np.meshgrid(t, t)\n",
    "\n",
    "scale = np.linspace(0, 1, res)\n",
    "scale = smoothstep(smoothstep(scale))\n",
    "    \n",
    "scale_bcast = scale[:, None, None, None]\n",
    "\n",
    "target = np.stack([\n",
    "    np.dstack((np.full(x.shape, 1.0), x, y))[np.newaxis, ...] * scale_bcast,\n",
    "    np.dstack((y, np.full(x.shape, 1.0), x))[np.newaxis, ...] * scale_bcast,\n",
    "    np.dstack((x, y, np.full(x.shape, 1.0)))[np.newaxis, ...] * scale_bcast\n",
    "])\n",
    "\n",
    "fig, axs = plt.subplots(1, 3, figsize=(8, 3))\n",
    "for i in range(3):\n",
    "    axs[i].imshow(target[i, -1, ...])\n",
    "    axs[i].set_xticks([])\n",
    "    axs[i].set_yticks([])\n",
    "plt.tight_layout()\n",
    "fig.savefig(prefix + '/gamut-orig.pdf', bbox_inches='tight')\n",
    "plt.figure()\n",
    "plt.plot(scale)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Perform the actual optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "coeffs = np.zeros((3, res, res, res, ndim))\n",
    "error = np.zeros((3, res, res, res))\n",
    "start = res // 5\n",
    "\n",
    "with tqdm(total=3*res) as progress:\n",
    "    # Don't start fitting from black, but from some reasonably bright colors\n",
    "    for i in range(start, res):\n",
    "        coeffs[:, i, ...] = coeffs[:, i-1, ...]\n",
    "        for j in range(3):\n",
    "            error[j, i, ...] = spectra.optimize(target[j, i, ...], coeffs[j, i, ...],\n",
    "                                                cie_lab=True)\n",
    "            progress.update(1)\n",
    "\n",
    "    # Now, work backwards towards black\n",
    "    for i in range(start - 1, -1, -1):\n",
    "        coeffs[:, i, ...] = coeffs[:, start, ...]\n",
    "        for j in range(3):\n",
    "            error[j, i, ...] = spectra.optimize(target[j, i, ...], coeffs[j, i, ...],\n",
    "                                                cie_lab=True,\n",
    "                                                similarity_constraint=.01)\n",
    "            progress.update(1)\n",
    "            \n",
    "# Rescale polynomial from 0..1 to 360..830nm range\n",
    "A, B, C = coeffs[..., 0], coeffs[..., 1], coeffs[..., 2]\n",
    "c0, c1 = 360.0, 1.0 / (830.0 - 360.0)\n",
    "A2 = A*(c1**2)\n",
    "B2 = B*c1 - 2*A*c0*(c1**2)\n",
    "C2 = C - B*c0*c1 + A*((c0*c1)**2)\n",
    "coeffs_scaled = np.stack([A2, B2, C2], -1)\n",
    "\n",
    "write_file(prefix + '/' + prefix + \".coeff\", scale, coeffs_scaled)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Interactive plots to view the computed coefficients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def coeff_plot(i, save_pdf=False, filename='coeffs.pdf'):\n",
    "    fig, axs = plt.subplots(3, 4, figsize=(9, 5))\n",
    "    plots = []\n",
    "    for j in range(3):\n",
    "        plots_row = []\n",
    "        for k in range(3):\n",
    "            im = axs[j][k].imshow(coeffs[j, i, ..., k])\n",
    "            axs[j][k].set_xticklabels([])\n",
    "            axs[j][k].set_yticklabels([])\n",
    "            cbar = add_colorbar(im, axs[j][k], pad=.05)\n",
    "            plots_row.append((im, cbar))\n",
    "        im = axs[j][3].imshow(error[j, i, ...])\n",
    "        axs[j][3].set_xticklabels([])\n",
    "        axs[j][3].set_yticklabels([])\n",
    "        cbar = add_colorbar(im, axs[j][3], pad=.05)\n",
    "        plots_row.append((im, cbar))\n",
    "        plots.append(plots_row)\n",
    "    if save_pdf:\n",
    "        fig.savefig(filename, bbox_inches='tight')\n",
    "    return plots\n",
    "\n",
    "plots = coeff_plot(10)\n",
    "\n",
    "@interact(i = (0, res-1, 1))\n",
    "def update_coeff_plot(i):\n",
    "    for j in range(3):\n",
    "        for k in range(3):\n",
    "            im, cbar = plots[j][k]\n",
    "            im.set_data(coeffs[j, i, ..., k])\n",
    "            cbar.set_clim(*minmax(coeffs[j, i, ..., k]))\n",
    "            cbar.draw_all()\n",
    "        im, cbar = plots[j][3]\n",
    "        im.set_data(error[j, i, ...])\n",
    "        cbar.set_clim(*minmax(error[j, i, ...]))\n",
    "        cbar.draw_all()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Coefficient fetch routine (takes a linear RGB value as input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = np.linspace(0, 1, res)\n",
    "interp = [ [ RegularGridInterpolator((scale, t, t), coeffs[j, :, :, :, i]) for i in range(3) ] for j in range(3)]\n",
    "\n",
    "def fetch(c):\n",
    "    c = np.asarray(c)\n",
    "    tbl = interp[c.argmax()]\n",
    "    cn = c / (c.max() + 1e-10)\n",
    "    cn = np.array([np.maximum(scale[0], c.max()), cn[(c.argmax()+2) % 3], cn[(c.argmax()+1) % 3]])\n",
    "    cn = np.array([tbl[i](cn) for i in range(3)])\n",
    "    return cn.squeeze()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Interactive plots to view the computed spectra"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(5,3))\n",
    "wl = np.linspace(360, 830, 100)\n",
    "plot = ax.plot(wl, np.zeros(wl.shape))[0]\n",
    "ax.set_ylim(0,1.2)\n",
    "\n",
    "@interact(r=(0,1,0.001), g=(0,1,0.001), b=(0,1,0.001), scale=(0, 1, 0.001))\n",
    "def test_spectrum(r=0.27, g=0.61, b=1, scale=1):\n",
    "    rgb = np.array([r, g, b]) * scale\n",
    "    c = fetch(rgb)\n",
    "    eval_rgb = spectra.eval_rgb(c)\n",
    "    eval_spec = spectra.eval_spec(c, wl)\n",
    "    \n",
    "    plot.set_ydata(eval_spec)\n",
    "    plt.suptitle(\"RGB value: %.3f, %.3f, %.3f\" % tuple(eval_rgb))\n",
    "    fig.canvas.draw();"
   ]
  }
 ],
 "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",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
