{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "pleasant-equation",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.autograd.functional import jacobian\n",
    "import itertools\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "moved-christian",
   "metadata": {},
   "outputs": [],
   "source": [
    "import combined as c"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "pressed-slope",
   "metadata": {},
   "source": [
    "So, this contains a bunch of initial tests of my abstractions. I eventually need to change these to assertions and package them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "capable-equality",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Instantiate some objects\n",
    "pm = c.PhysicalModel(1.0,1e-6,0.01,2.0,1e-8)\n",
    "s = c.States(torch.tensor([1.0,2,3]), torch.tensor([0.0]))\n",
    "lp = c.LinearProfit(0,0.95,torch.tensor([1.0,0,0]), 5)\n",
    "est_int = c.EstimandInterface(torch.tensor([[1.0,2,3,2]\n",
    "                                           ,[4,5,6,2]\n",
    "                                           ,[7,8,9,2]\n",
    "                                           ,[1,3,5,7]]\n",
    "                                          ),torch.ones(3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "written-experience",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "stocks\ttensor([1., 2., 3.]) \n",
      "debris\t tensor([0.])\n",
      "3\n",
      "1\n"
     ]
    }
   ],
   "source": [
    "#test State object \n",
    "print(s)\n",
    "print(s.number_constellations)\n",
    "print(s.number_debris_trackers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "twelve-arthur",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Launch Decisions and Partial Derivativs of value function with\n",
      "\tlaunches\n",
      "\t\t tensor([1., 1., 1.])\n",
      "\tPartials\n",
      "\t\ttensor([[1., 2., 3., 2.],\n",
      "        [4., 5., 6., 2.],\n",
      "        [7., 8., 9., 2.],\n",
      "        [1., 3., 5., 7.]])\n",
      "tensor([1., 1., 1.]) tensor([[1., 2., 3., 2.],\n",
      "        [4., 5., 6., 2.],\n",
      "        [7., 8., 9., 2.],\n",
      "        [1., 3., 5., 7.]])\n",
      "tensor(1.)\n",
      "tensor([0., 1., 0.])\n",
      "tensor([2., 5., 8., 3.])\n",
      "tensor([[0., 2., 0., 0.],\n",
      "        [0., 5., 0., 0.],\n",
      "        [0., 8., 0., 0.],\n",
      "        [0., 3., 0., 0.]])\n"
     ]
    }
   ],
   "source": [
    "#Test estimand interface\n",
    "print(est_int)\n",
    "print(est_int.launches,est_int.partials)\n",
    "\n",
    "print(est_int.launch_single(1))\n",
    "print(est_int.launch_vector(1))\n",
    "print(est_int.partial_vector(1)) \n",
    "print(est_int.partial_matrix(1)) #TODO: double check orientation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "impressive-tribe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "1.0\n",
      "1e-06\n",
      "0.01\n",
      "2.0\n",
      "1e-08\n",
      "tensor([1.0133e-06, 2.0266e-06, 2.9802e-06])\n",
      "tensor([1., 2., 3.]) tensor([0.])\n",
      "tensor([1.0000, 1.0000, 1.0000]) tensor([12.0000])\n"
     ]
    }
   ],
   "source": [
    "#Test physical model methods\n",
    "print(pm)\n",
    "print(pm.survival(s))\n",
    "s2 = pm.transition(s,est_int)\n",
    "print(s.stocks,s.debris)\n",
    "print(s2.stocks,s2.debris)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "stretch-reward",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LinearProfit\n",
      " Benefit weights:\ttensor([1., 0., 0.])\n",
      " launch cost:\t5\n",
      " Deorbit cost:\t0\n",
      "tensor(-4.)\n"
     ]
    }
   ],
   "source": [
    "#test linear profit object\n",
    "print(lp)\n",
    "print(lp.period_benefit(s,est_int))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "advance-folder",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1., 0., 0., 0.])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lp._period_benefit_jacobian_wrt_states( s.stocks, s.debris, est_int.launches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "posted-subscriber",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-5.,  0.,  0.])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lp._period_benefit_jacobian_wrt_launches( s.stocks, s.debris, est_int.launches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "divine-agenda",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1., 0., 0., 0.])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lp.period_benefit_jacobian_wrt_states( s, est_int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "surgical-direction",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 2., 3., 2.],\n",
       "        [4., 5., 6., 2.],\n",
       "        [7., 8., 9., 2.],\n",
       "        [1., 3., 5., 7.]])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "est_int.partials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "mounted-roots",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([-7.5676e+05,  2.8893e+05,  4.6783e+05,  1.5236e+00]),\n",
       " <combined.States at 0x7f8c3c9c54f0>)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c.single_transition(pm,lp,s,est_int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "pediatric-iceland",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-718920.5625,  274490.1562,  444444.6250])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "minimand, iterated_partials, iterated_state = c.optimality(pm,lp,s,est_int)\n",
    "minimand"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "isolated-cleveland",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-2285563., -2285557., -2285557.])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "est_int2 = c.EstimandInterface(iterated_partials,torch.ones(3))\n",
    "\n",
    "minimand2, iterated_partials2, iterated_state2 = c.optimality(pm,lp,iterated_state,est_int2)\n",
    "minimand2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "relevant-romance",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-2405858.5000, -2405852.5000, -2405852.5000])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "est_int3 = c.EstimandInterface(iterated_partials2,torch.ones(3))\n",
    "\n",
    "minimand3, iterated_partials3, iterated_state3 = c.optimality(pm,lp,iterated_state2,est_int3)\n",
    "minimand3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "israeli-oracle",
   "metadata": {},
   "source": [
    "So, this succesfylly let me link the two. I'm going to move to another notebook, clean up, and start integrating the system"
   ]
  }
 ],
 "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}