{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jmapel/miniconda3/envs/knoten/lib/python3.7/site-packages/ale/__init__.py:22: UserWarning: ALESPICEROOT environment variable not set, Spice Drivers will not function correctly\n",
      "  warnings.warn('ALESPICEROOT environment variable not set, Spice Drivers will not function correctly')\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from plio.io import io_controlnetwork\n",
    "from knoten.csm import create_csm\n",
    "from scipy import sparse\n",
    "import ale\n",
    "import csmapi\n",
    "import numpy as np\n",
    "from collections import OrderedDict\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from knoten.bundle import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cubes = 'data/cubes.lis'\n",
    "sensors = generate_sensors(cubes)\n",
    "\n",
    "network = 'data/hand_dense.net'\n",
    "cnet = io_controlnetwork.from_isis(network)\n",
    "sensors = {sn: sensors[sn] for sn in cnet[\"serialnumber\"].unique()}\n",
    "cnet = compute_apriori_ground_points(cnet, sensors) # autoseed did not generate ground points, calculate and repopulate the data frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "all_parameters = {sn: get_sensor_parameters(sensor) for sn, sensor in sensors.items()}\n",
    "for sn, parameters in all_parameters.items():\n",
    "    print(f\"Image: {sn}\")\n",
    "    for param in parameters:\n",
    "        print(f\"  {param.name} | {param.index} | {param.value}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve for angles and angular rates\n",
    "solve_parameters = {sn: params[6:12] for sn, params in all_parameters.items()}\n",
    "for sn, parameters in solve_parameters.items():\n",
    "    print(f\"Image: {sn}\")\n",
    "    for param in parameters:\n",
    "        print(f\"  {param.name} | {param.index} | {param.value}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "column_dict = compute_coefficient_columns(cnet, sensors, solve_parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_observations = 2 * len(cnet)\n",
    "W_observations = np.eye(num_observations) # this is a place holder until Jesse adds his calculations\n",
    "W_params = compute_parameter_weights(cnet, sensors, solve_parameters, column_dict)\n",
    "dof = W_observations.shape[0] - W_params.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_iterations = 10\n",
    "tol = 1e-10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_sensor_params = sum([len(parameters) for parameters in solve_parameters.values()])\n",
    "num_point_params = 3 * len(cnet[\"id\"].unique())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sensors = generate_sensors(cubes) # generate sensors\n",
    "cnet = io_controlnetwork.from_isis(network) # load in network\n",
    "cnet = compute_apriori_ground_points(cnet, sensors) # calculate ground points\n",
    "\n",
    "column_dict = compute_coefficient_columns(cnet, sensors, solve_parameters)\n",
    "num_parameters = max(col_range[1] for col_range in column_dict.values())\n",
    "num_observations = 2 * len(cnet)\n",
    "W_observations = np.eye(num_observations)\n",
    "W_params = compute_parameter_weights(cnet, sensors, solve_parameters, column_dict)\n",
    "\n",
    "iteration = 0\n",
    "V = compute_residuals(cnet, sensors)\n",
    "dX = np.zeros(W_params.shape[0]) #initialize for sigma calculatioN\n",
    "sigma0 = compute_sigma(V, dX, W_params, W_observations)\n",
    "print(f'iteration {iteration}: sigma0 = {sigma0}\\n')\n",
    "\n",
    "total_correction_dense = np.zeros(num_parameters)\n",
    "for i in range(max_iterations):   \n",
    "    iteration += 1\n",
    "    old_sigma0 = sigma0\n",
    "    \n",
    "    J = compute_jacobian(cnet, sensors, solve_parameters, column_dict)    \n",
    "    N = J.T.dot(W_observations).dot(J) + W_params # calculate the normal equation\n",
    "    C = J.T.dot(W_observations).dot(V) - W_params.dot(total_correction_dense)\n",
    "    dX = np.linalg.inv(N).dot(C) #calculate change in camera parameters and ground points\n",
    "    total_correction_dense += dX\n",
    "    print(f'corrections: mean = {dX.mean()} min = {dX.min()} max = {dX.max()}')\n",
    "    \n",
    "    update_parameters(sensors, solve_parameters, cnet, dX, column_dict)\n",
    "    \n",
    "    V = compute_residuals(cnet, sensors)\n",
    "    sigma0 = compute_sigma(V, dX, W_params, W_observations)\n",
    "    sigma0 = np.sqrt((V.dot(W_observations).dot(V) + dX.dot(W_params).dot(dX))/dof)\n",
    "    print(f'iteration {iteration}: sigma0 = {sigma0}\\n')\n",
    "    \n",
    "    if (abs(sigma0 - old_sigma0) < tol):\n",
    "        print(f'change in sigma0 of {abs(sigma0 - old_sigma0)} converged!')\n",
    "        break\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sensors = generate_sensors(cubes) # generate sensors\n",
    "cnet = io_controlnetwork.from_isis(network) # load in network\n",
    "cnet = compute_apriori_ground_points(cnet, sensors) # calculate ground points\n",
    "\n",
    "# This is setup once per bundle and then accumulates\n",
    "total_corrections = np.zeros(num_sensor_params + num_point_params)\n",
    "\n",
    "# These are computed once per bundle\n",
    "W_CC_sparse = sparse.csc_matrix((num_sensor_params, num_sensor_params))\n",
    "W_PP = {}\n",
    "\n",
    "# Compute image param weight matrices\n",
    "for sn, params in solve_parameters.items():\n",
    "    coeff_indices = column_dict[sn]\n",
    "    coeff_range = np.arange(coeff_indices[0], coeff_indices[1])\n",
    "    num_image_coeffs = coeff_indices[1] - coeff_indices[0]\n",
    "    W_CC_data = compute_image_weight(sensors[sn], params).ravel()\n",
    "    W_CC_row = np.repeat(coeff_range, num_image_coeffs)\n",
    "    W_CC_column = np.tile(coeff_range, num_image_coeffs)\n",
    "    W_CC_sparse += sparse.coo_matrix((W_CC_data, (W_CC_row, W_CC_column)), (num_sensor_params, num_sensor_params)).tocsc()\n",
    "    \n",
    "# Compute point param weight matrices\n",
    "for point_id in cnet['id'].unique():\n",
    "    W_PP[point_id] = compute_point_weight(cnet, point_id)\n",
    "    \n",
    "V = compute_residuals(cnet, sensors)\n",
    "sigma0 = compute_sigma_sparse(V, np.zeros(total_corrections.shape), W_CC_sparse, W_PP, W_observations, column_dict)\n",
    "    \n",
    "# Start iteration logic\n",
    "for i in range(max_iterations):\n",
    "    old_sigma0 = sigma0\n",
    "    \n",
    "    H_CC_sparse = sparse.csc_matrix((num_sensor_params, num_sensor_params))\n",
    "    H_CC_sparse += W_CC_sparse\n",
    "\n",
    "    g_C_sparse = np.zeros(num_sensor_params)\n",
    "    g_C_sparse -= W_CC_sparse.dot(total_corrections[:num_sensor_params])\n",
    "    # g_C_sparse += W_CC_sparse.dot(sensor_corrections)\n",
    "\n",
    "    # Q = H_PP^-1 * H_PC \n",
    "    Q_mats = {}\n",
    "    # NIC = H_PP^-1 * g_P\n",
    "    NIC_vecs = {}\n",
    "\n",
    "    updates = np.zeros(num_sensor_params + num_point_params)\n",
    "\n",
    "    for point_id, group in cnet.groupby('id'):\n",
    "        ground_pt = group.iloc[0][[\"adjustedX\", \"adjustedY\", \"adjustedZ\"]]\n",
    "        H_CP = sparse.csc_matrix((num_sensor_params, 3))\n",
    "        H_PP = np.zeros((3, 3))\n",
    "        g_P = np.zeros(3)\n",
    "        for measure_idx, row in group.iterrows():\n",
    "            serial = row[\"serialnumber\"]\n",
    "            sensor = sensors[serial]\n",
    "            point_partials = compute_ground_partials(sensor, ground_pt)\n",
    "            sensor_partials = compute_sensor_partials(sensor, solve_parameters[serial], ground_pt)\n",
    "            coeff_indices = column_dict[serial]\n",
    "            coeff_range = np.arange(coeff_indices[0], coeff_indices[1])\n",
    "            num_image_coeffs = coeff_indices[1] - coeff_indices[0]\n",
    "\n",
    "            H_CC_point_data = np.dot(sensor_partials.T, sensor_partials).ravel()\n",
    "            H_CC_point_row = np.repeat(coeff_range, num_image_coeffs)\n",
    "            H_CC_point_column = np.tile(coeff_range, num_image_coeffs)\n",
    "            H_CC_sparse += sparse.coo_matrix((H_CC_point_data, (H_CC_point_row, H_CC_point_column)), (num_sensor_params, num_sensor_params)).tocsc()\n",
    "\n",
    "            H_CP_point_data = np.dot(sensor_partials.T, point_partials).ravel()\n",
    "            H_CP_point_row = np.repeat(coeff_range, 3)\n",
    "            H_CP_point_column = np.tile(np.arange(0, 3), num_image_coeffs)\n",
    "            H_CP += sparse.coo_matrix((H_CP_point_data, (H_CP_point_row, H_CP_point_column)), (num_sensor_params, 3)).tocsc()\n",
    "\n",
    "            H_PP += np.dot(point_partials.T, point_partials)\n",
    "\n",
    "            g_C_sparse[coeff_indices[0]:coeff_indices[1]] += np.dot(sensor_partials.T, V[2*measure_idx:2*measure_idx+2])\n",
    "            g_P += np.dot(point_partials.T, V[2*measure_idx:2*measure_idx+2])\n",
    "        point_param_range = column_dict[point_id]\n",
    "        g_P -= W_PP[point_id].dot(total_corrections[point_param_range[0]:point_param_range[1]])\n",
    "        H_PP += W_PP[point_id]\n",
    "        H_PP_inv = sparse.csc_matrix(np.linalg.inv(H_PP))\n",
    "        Q_mats[point_id] = H_PP_inv.dot(H_CP.transpose())\n",
    "        NIC_vecs[point_id] = H_PP_inv.dot(g_P)\n",
    "        H_CC_sparse -= H_CP.dot(Q_mats[point_id])\n",
    "        g_C_sparse -= H_CP.dot(NIC_vecs[point_id])\n",
    "\n",
    "    updates[:num_sensor_params] = sparse.linalg.spsolve(H_CC_sparse, g_C_sparse)\n",
    "\n",
    "    for point_id in Q_mats:\n",
    "        point_param_indices = column_dict[point_id]\n",
    "        updates[point_param_indices[0]:point_param_indices[1]] = NIC_vecs[point_id] - Q_mats[point_id].dot(updates[:num_sensor_params])\n",
    "        \n",
    "    print(f'corrections: mean = {updates.mean()} min = {updates.min()} max = {updates.max()}')\n",
    "\n",
    "    total_corrections += updates\n",
    "\n",
    "    update_parameters(sensors, solve_parameters, cnet, updates, column_dict)\n",
    "    V = compute_residuals(cnet, sensors)\n",
    "    sigma0 = compute_sigma_sparse(V, updates, W_CC_sparse, W_PP, W_observations, column_dict)\n",
    "    print(f'iteration {i+1}: sigma0 = {sigma0}\\n')\n",
    "    \n",
    "    if (abs(sigma0 - old_sigma0) < tol):\n",
    "        print(f'change in sigma0 of {abs(sigma0 - old_sigma0)} converged!')\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Sensor diff\")\n",
    "print(np.min(total_correction_dense[:num_sensor_params].flatten() - total_corrections[:num_sensor_params]))\n",
    "print(np.max(total_correction_dense[:num_sensor_params].flatten() - total_corrections[:num_sensor_params]))\n",
    "print(\"Point diff\")\n",
    "print(np.min(total_correction_dense[num_sensor_params:].flatten() - total_corrections[num_sensor_params:]))\n",
    "print(np.max(total_correction_dense[num_sensor_params:].flatten() - total_corrections[num_sensor_params:]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
