{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ctypes\n",
    "import datetime\n",
    "import glob\n",
    "import gzip\n",
    "import json\n",
    "import os\n",
    "import shutil\n",
    "import tarfile\n",
    "\n",
    "import csmapi\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from pylab import rcParams\n",
    "\n",
    "import ale\n",
    "import gdal\n",
    "import osr\n",
    "import pvl\n",
    "import spiceypy as spice\n",
    "from ale.drivers import JsonEncoder\n",
    "from ale.drivers.kaguya_drivers import KaguyaTcPds3NaifSpiceDriver\n",
    "\n",
    "import knoten\n",
    "from knoten import csm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Decompress .sl2 files \n",
    "files = glob.glob('*.sl2')\n",
    "for f in files:\n",
    "    tar = tarfile.open(f)\n",
    "    \n",
    "    # Extract the spc and the ctg files\n",
    "    for m in tar.getmembers():\n",
    "        tar.extract(m, \"images\")\n",
    "    imageid = os.path.splitext(os.path.basename(f))[0]\n",
    "    \n",
    "    # decompress .igz file containing image data\n",
    "    with gzip.open('images/{}.igz'.format(imageid), 'rb') as f_in:\n",
    "        with open('images/{}.img'.format(imageid), 'wb') as f_out:\n",
    "            shutil.copyfileobj(f_in, f_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "image = \"images/TC2W2B0_01_00366S490E1640.img\"\n",
    "gtifimage = os.path.splitext(image)[0]+\".tif\"\n",
    "alelabel = os.path.splitext(image)[0]+\".json\"\n",
    "vrtfile = os.path.splitext(image)[0]+\".vrt\"\n",
    "\n",
    "label = pvl.load(image)\n",
    "driver = KaguyaTcPds3NaifSpiceDriver(image)\n",
    "\n",
    "with driver as d:\n",
    "    aledict = d.to_dict()\n",
    "\n",
    "json.dump(aledict, open(alelabel, \"w\"), cls=JsonEncoder)\n",
    "\n",
    "nlines, nsamples = aledict[\"image_lines\"], aledict[\"image_samples\"]\n",
    "\n",
    "model=\"USGS_ASTRO_LINE_SCANNER_SENSOR_MODEL\"\n",
    "plugin = csmapi.Plugin.getList()[0]\n",
    "isd = csmapi.Isd(alelabel)\n",
    "camera = plugin.constructModelFromISD(isd, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 11588 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# a/b for the moon\n",
    "semi_major = aledict[\"radii\"][\"semimajor\"] * 1000 # scale to meters\n",
    "semi_minor = aledict[\"radii\"][\"semiminor\"] * 1000 # scale to meters\n",
    "\n",
    "imgbounds = csm.generate_boundary((nlines, nsamples))\n",
    "gcps = csm.generate_gcps(camera, imgbounds, semi_major, semi_minor)\n",
    "vrt = csm.generate_vrt((nsamples, nlines), gcps, gtifimage, proj='+proj=longlat +a={} +b={} +no_defs'.format(semi_major, semi_minor))\n",
    "with open(vrtfile, \"w\") as f:\n",
    "    f.write(vrt)\n",
    "    \n",
    "# plot footprint\n",
    "lats, lons, alts = csm.generate_latlon_boundary(camera, imgbounds, semi_major, semi_minor)\n",
    "plt.plot(lats, lons)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 83524 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# GDAL cannot currently read kaguya images, so we need to open them manually for visualization \n",
    "imagebytes = open(image, \"rb\").read()\n",
    "arr = np.frombuffer(imagebytes[aledict[\"label\"][\"^IMAGE\"].value-1:], dtype=\">i2\")\n",
    "\n",
    "scale = pvl.load(image)[\"IMAGE\"][\"SCALING_FACTOR\"]\n",
    "arr = arr*scale\n",
    "arr = arr.reshape(nlines, nsamples)\n",
    "\n",
    "# plot unprojected \n",
    "plt.imshow(arr)\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def array2raster(rasterfn, array, newRasterfn):\n",
    "    \"\"\"\n",
    "    Writes an array to a GeoDataset using another dataset as reference. Borrowed  \n",
    "    from: https://pcjericks.github.io/py-gdalogr-cookbook/raster_layers.html\n",
    "    \n",
    "    Parameters \n",
    "    ----------\n",
    "    rasterfn : str, GeoDataset\n",
    "               Dataset or path to the dataset to use as a reference. Geotransform \n",
    "               and spatial reference information is copied into the new image. \n",
    "               \n",
    "    array : np.array \n",
    "            Array to write \n",
    "            \n",
    "    newRasterfn : str \n",
    "                  Filename for new raster image \n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    : GeoDataset \n",
    "      File handle for the new raster file\n",
    "      \n",
    "    \"\"\"\n",
    "    naxis = len(array.shape)\n",
    "    assert naxis == 2 or naxis == 3      \n",
    "    \n",
    "    nbands = 1\n",
    "    \n",
    "    raster = gdal.Open(rasterfn, gdal.GA_ReadOnly)\n",
    "    geotransform = raster.GetGeoTransform()\n",
    "    originX = geotransform[0]\n",
    "    originY = geotransform[3]\n",
    "    pixelWidth = geotransform[1]\n",
    "    pixelHeight = geotransform[5]\n",
    "    cols = array.shape[1]\n",
    "    rows = array.shape[0]\n",
    "\n",
    "    driver = gdal.GetDriverByName('GTiff')\n",
    "    outRaster = driver.Create(newRasterfn, cols, rows, nbands, gdal.GDT_Byte)\n",
    "    outRaster.SetGeoTransform((originX, pixelWidth, 0, originY, 0, pixelHeight))\n",
    "    \n",
    "\n",
    "    outband = outRaster.GetRasterBand(1)\n",
    "    # Bands use indexing starting at 1\n",
    "    outband.WriteArray(array)\n",
    "    outband.FlushCache()\n",
    "\n",
    "    outRasterSRS = osr.SpatialReference()\n",
    "    outRasterSRS.ImportFromWkt(raster.GetProjectionRef())\n",
    "    outRaster.SetProjection(outRasterSRS.ExportToWkt())\n",
    "    del outRaster\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 85280 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# save geotiff\n",
    "array2raster(vrtfile, arr, gtifimage)\n",
    "fd = gdal.Warp(\"{}.tif\".format(image), vrtfile, srcNodata=0, dstNodata=-9999999)\n",
    "\n",
    "# draw projected image\n",
    "image_arr = fd.ReadAsArray()\n",
    "# set no data to np.nan so it doesn't affect the stretch\n",
    "image_arr[image_arr == -9999999] = np.nan\n",
    "plt.imshow(image_arr)\n",
    "plt.show()"
   ]
  }
 ],
 "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
