{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "geographic-wilderness",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import combined as c\n",
    "import NeuralNetworkSpecifications as nns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "major-glucose",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[4., 9., 6.]],\n",
      "\n",
      "        [[0., 5., 4.]],\n",
      "\n",
      "        [[3., 1., 9.]],\n",
      "\n",
      "        [[6., 4., 8.]],\n",
      "\n",
      "        [[8., 7., 6.]]], grad_fn=<CatBackward>)\n"
     ]
    }
   ],
   "source": [
    "batch_size,states,choices = 5,3,1\n",
    "constellations = states -1 #determined by debris tracking\n",
    "max_start_state = 10\n",
    "\n",
    "stocks = torch.randint(max_start_state,(batch_size,1,constellations), dtype=torch.float32, requires_grad=True)\n",
    "debris = torch.randint(max_start_state,(batch_size,1,1), dtype=torch.float32, requires_grad=True)\n",
    "\n",
    "s = c.States(stocks, debris)\n",
    "\n",
    "print(s.values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "recognized-ability",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([5.6433e-07, 6.7631e-07], grad_fn=<MulBackward0>)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "constellation_collision_risk = 1e-6 * torch.rand(constellations, requires_grad=True)\n",
    "constellation_collision_risk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "elect-float",
   "metadata": {},
   "outputs": [],
   "source": [
    "debris_decay_rate = 0.1\n",
    "launch_debris = 0.05\n",
    "debris_autocatalysis_rate = 1.4\n",
    "\n",
    "benefit_weight0 = torch.tensor([1.0,-0.02], requires_grad=True)\n",
    "benefit_weight1 = torch.tensor([0.0,1.0], requires_grad=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "compressed-individual",
   "metadata": {},
   "outputs": [],
   "source": [
    "pm = c.PhysicalModel(10\n",
    "                , constellation_collision_risk #constellations_collision_risk #as tensor\n",
    "                , debris_decay_rate #debris_decay_rate\n",
    "                , launch_debris #launch_debris\n",
    "                , debris_autocatalysis_rate #debris_autocatalysis_rate\n",
    "               )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "tracked-bachelor",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LinearProfit():\n",
    "    \"\"\"\n",
    "    The simplest type of profit function available.\n",
    "    \"\"\"\n",
    "    def __init__(self, batch_size, constellation_number, discount_factor, benefit_weights, launch_cost, deorbit_cost=0, ):\n",
    "        self.batch_size = batch_size\n",
    "        \n",
    "        \n",
    "        #track which constellation this is.\n",
    "        self.constellation_number = constellation_number\n",
    "        \n",
    "        #get the number of constellations (pull from the benefit weight, in the dimension that counts across constellations)\n",
    "        self.number_of_constellations = benefit_weights.size()[0] -1\n",
    "\n",
    "        #parameters describing the agent's situation\n",
    "        self.discount_factor = discount_factor\n",
    "        self.benefit_weights = benefit_weights\n",
    "        self.launch_cost = launch_cost\n",
    "        self.deorbit_cost = deorbit_cost\n",
    "    \n",
    "    def _period_benefit(self,stocks,debris,launches):\n",
    "        # multiply benefits times stocks\n",
    "        # sum across constellations\n",
    "        # reshape to standard dimensions\n",
    "        # subtract launch costs. \n",
    "        profit = torch.tensordot(self.benefit_weights,stocks, [[0],[1]])[:,self.constellation_number] \\\n",
    "            - (self.launch_cost * launches)[:,self.constellation_number,0]\n",
    "        return profit.view(batch_size,1)\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "confidential-philippines",
   "metadata": {},
   "outputs": [],
   "source": [
    "launch_cost = 5\n",
    "ea0 = LinearProfit(\n",
    "                    batch_size\n",
    "                    ,0 #constellation index\n",
    "                    ,0.95 #discount\n",
    "                    ,benefit_weight0\n",
    "                    ,launch_cost #launch_cost\n",
    "                    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "religious-georgia",
   "metadata": {},
   "outputs": [],
   "source": [
    "enn = nns.EstimandNN(batch_size\n",
    "                 ,states\n",
    "                 ,choices\n",
    "                 ,constellations\n",
    "                 ,12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "painful-republican",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[0.],\n",
       "         [0.]],\n",
       "\n",
       "        [[0.],\n",
       "         [0.]],\n",
       "\n",
       "        [[0.],\n",
       "         [0.]],\n",
       "\n",
       "        [[0.],\n",
       "         [0.]],\n",
       "\n",
       "        [[0.],\n",
       "         [0.]]], grad_fn=<ReluBackward0>)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "launch_decisions = enn.forward(s.values).choices\n",
    "launch_decisions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "equal-raising",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(stocks,launches):\n",
    "    # multiply benefits times stocks\n",
    "    # sum across constellations\n",
    "    # reshape to standard dimensions\n",
    "    # subtract launch costs. \n",
    "    profit = torch.tensordot(benefit_weight0,stocks, [[0],[1]])[:,0] - (launch_cost * launch_decisions)[:,0,0]\n",
    "    return profit.view(batch_size,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "divine-editor",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = LinearProfit(batch_size #batch_size\n",
    "                    ,0 #constellation index\n",
    "                    ,0.95 #discount\n",
    "                    ,benefit_weight0\n",
    "                    ,launch_cost #launch_cost\n",
    "                    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "japanese-captain",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[3.9200],\n",
       "        [0.0000],\n",
       "        [2.9400],\n",
       "        [5.8800],\n",
       "        [7.8400]], grad_fn=<ViewBackward>)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test(stocks,launch_decisions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "determined-difference",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[3.9200],\n",
       "        [0.0000],\n",
       "        [2.9400],\n",
       "        [5.8800],\n",
       "        [7.8400]], grad_fn=<ViewBackward>)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t._period_benefit(s.stocks,s.debris,launch_decisions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "tribal-least",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[[[0.9800, 0.0000]],\n",
       " \n",
       "           [[0.0000, 0.0000]],\n",
       " \n",
       "           [[0.0000, 0.0000]],\n",
       " \n",
       "           [[0.0000, 0.0000]],\n",
       " \n",
       "           [[0.0000, 0.0000]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[0.0000, 0.0000]],\n",
       " \n",
       "           [[0.9800, 0.0000]],\n",
       " \n",
       "           [[0.0000, 0.0000]],\n",
       " \n",
       "           [[0.0000, 0.0000]],\n",
       " \n",
       "           [[0.0000, 0.0000]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[0.0000, 0.0000]],\n",
       " \n",
       "           [[0.0000, 0.0000]],\n",
       " \n",
       "           [[0.9800, 0.0000]],\n",
       " \n",
       "           [[0.0000, 0.0000]],\n",
       " \n",
       "           [[0.0000, 0.0000]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[0.0000, 0.0000]],\n",
       " \n",
       "           [[0.0000, 0.0000]],\n",
       " \n",
       "           [[0.0000, 0.0000]],\n",
       " \n",
       "           [[0.9800, 0.0000]],\n",
       " \n",
       "           [[0.0000, 0.0000]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[0.0000, 0.0000]],\n",
       " \n",
       "           [[0.0000, 0.0000]],\n",
       " \n",
       "           [[0.0000, 0.0000]],\n",
       " \n",
       "           [[0.0000, 0.0000]],\n",
       " \n",
       "           [[0.9800, 0.0000]]]]], grad_fn=<ViewBackward>),\n",
       " tensor([[[[[0.]],\n",
       " \n",
       "           [[0.]],\n",
       " \n",
       "           [[0.]],\n",
       " \n",
       "           [[0.]],\n",
       " \n",
       "           [[0.]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[0.]],\n",
       " \n",
       "           [[0.]],\n",
       " \n",
       "           [[0.]],\n",
       " \n",
       "           [[0.]],\n",
       " \n",
       "           [[0.]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[0.]],\n",
       " \n",
       "           [[0.]],\n",
       " \n",
       "           [[0.]],\n",
       " \n",
       "           [[0.]],\n",
       " \n",
       "           [[0.]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[0.]],\n",
       " \n",
       "           [[0.]],\n",
       " \n",
       "           [[0.]],\n",
       " \n",
       "           [[0.]],\n",
       " \n",
       "           [[0.]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[0.]],\n",
       " \n",
       "           [[0.]],\n",
       " \n",
       "           [[0.]],\n",
       " \n",
       "           [[0.]],\n",
       " \n",
       "           [[0.]]]]]),\n",
       " tensor([[[[[-5.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-0.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-0.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-0.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-0.],\n",
       "            [ 0.]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[-0.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-5.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-0.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-0.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-0.],\n",
       "            [ 0.]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[-0.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-0.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-5.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-0.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-0.],\n",
       "            [ 0.]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[-0.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-0.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-0.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-5.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-0.],\n",
       "            [ 0.]]]],\n",
       " \n",
       " \n",
       " \n",
       "         [[[[-0.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-0.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-0.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-0.],\n",
       "            [ 0.]],\n",
       " \n",
       "           [[-5.],\n",
       "            [ 0.]]]]]))"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#in this case, the debris isn't tracked because it isn't included, and launch_decisions has a similar issue.\n",
    "torch.autograd.functional.jacobian(t._period_benefit, (s.stocks,s.debris,launch_decisions), create_graph=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "statutory-lyric",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "naked-health",
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
