{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49a6d466",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import dask.array as da\n",
    "from polsarpro.util import boxcar, C3_to_T3, S_to_C3, S_to_T3, C4_to_T4, T3_to_C3\n",
    "from polsarpro.io import open_netcdf_beam\n",
    "from polsarpro.classification import wishart_h_a_alpha\n",
    "# from polsarpro.classification import _reconstruct_matrix_from_ds\n",
    "from pathlib import Path\n",
    "\n",
    "# optional import for progress bar\n",
    "from dask.diagnostics import ProgressBar\n",
    "\n",
    "\n",
    "# change to your data paths\n",
    "# original dataset\n",
    "input_alos_data = Path(\"/data/psp/test_files/SAN_FRANCISCO_ALOS1_slc.nc\")\n",
    "\n",
    "# S2 inputs generated with C-PSP\n",
    "input_test_dir = Path(\"/data/psp/SAN_FRANCISCO_ALOS1/\")\n",
    "\n",
    "# # output files from C\n",
    "# output_test_dir = Path(\"/data/psp/res/vanzyl_c/\")\n",
    "# if not os.path.isdir(output_test_dir):\n",
    "#     os.mkdir(output_test_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c1d8707",
   "metadata": {},
   "outputs": [],
   "source": [
    "# uncomment to test on S matrix made with SNAP\n",
    "S = open_netcdf_beam(input_alos_data)\n",
    "label = wishart_h_a_alpha(S)\n",
    "# C3 = boxcar(S_to_C3(S), 5, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff9da699",
   "metadata": {},
   "outputs": [],
   "source": [
    "with ProgressBar():\n",
    "    label = label.compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "955d437f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(label.label, interpolation=\"none\")\n",
    "plt.colorbar(fraction=0.046, pad=0.04)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5a1499e",
   "metadata": {},
   "outputs": [],
   "source": [
    "shp = C3.m11.shape\n",
    "# to be replaced by h_a_alpha initial classes\n",
    "class_map = da.random.randint(1, 9, size=shp)#.chunks(C3.chunks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "332e81c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _trace_product_3_wish(M_inv, cov_ds):\n",
    "\n",
    "    m11 = cov_ds.m11.data\n",
    "    m12 = cov_ds.m12.data\n",
    "    m22 = cov_ds.m22.data\n",
    "    m13 = cov_ds.m13.data\n",
    "    m23 = cov_ds.m23.data\n",
    "    m33 = cov_ds.m33.data\n",
    "    \n",
    "    trace = (M_inv[0, 0] * m11 +\n",
    "             M_inv[0, 1] * da.conj(m12) +\n",
    "             M_inv[0, 2] * da.conj(m13) +\n",
    "             M_inv[1, 0] * m12 +\n",
    "             M_inv[1, 1] * m22 +\n",
    "             M_inv[1, 2] * da.conj(m23) +\n",
    "             M_inv[2, 0] * m13 +\n",
    "             M_inv[2, 1] * m23 +\n",
    "             M_inv[2, 2] * m33)\n",
    "    \n",
    "    return trace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fbe8ca5",
   "metadata": {},
   "outputs": [],
   "source": [
    "nclass = 8\n",
    "centers = []\n",
    "inv_centers = []\n",
    "dist_min = da.full_like(C3.m11, fill_value=da.inf)\n",
    "label_min = class_map.copy()\n",
    "for i in range(1, nclass + 1):\n",
    "    # class center\n",
    "    mask = class_map == i\n",
    "    npts = mask.sum()\n",
    "    # center = xr.where(class_map == i, C3.chunk(dict(x=1024,y=1024)), 0).sum() / npts\n",
    "    center = xr.where(class_map == i, C3, 0).sum() / npts\n",
    "    # TODO: hardcode chunk shape (C3 T3 -> 3x3, C4 T4 -> 4x4)\n",
    "    M = _reconstruct_matrix_from_ds(center).rechunk(\"auto\") + 1e-12*da.eye(3)\n",
    "\n",
    "    # compute distance components\n",
    "    meta = (np.array([], dtype=\"complex64\").reshape((0, 0)),)\n",
    "    M_inv = da.apply_gufunc(np.linalg.inv, \"(i,j)->(i,j)\", M, meta=meta)\n",
    "    M_det = da.apply_gufunc(np.linalg.det, \"(i,j)->()\", M, meta=meta)\n",
    "\n",
    "    # update classification \n",
    "    dist = da.log(M_det.real) + _trace_product_3_wish(M_inv, C3)\n",
    "    is_smaller = dist < dist_min\n",
    "    label_min = da.where(is_smaller, i, label_min)\n",
    "    dist_min = da.where(is_smaller, dist, dist_min)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddfd89cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "with ProgressBar():\n",
    "    # center = center.compute()\n",
    "    # centers = da.compute(centers)\n",
    "    label_map = label_min.compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb35c141",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(label_map[::8], interpolation=\"none\")\n",
    "plt.colorbar(fraction=0.046, pad=0.04)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d67dfa2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "psp",
   "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.12.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
