{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "similar-ebony",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import combined as c\n",
    "import NeuralNetworkSpecifications as nns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "spread-hygiene",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[6., 1., 0.]],\n",
      "\n",
      "        [[2., 0., 4.]],\n",
      "\n",
      "        [[7., 6., 9.]],\n",
      "\n",
      "        [[3., 6., 9.]],\n",
      "\n",
      "        [[9., 1., 2.]]])\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)\n",
    "debris = torch.randint(max_start_state,(batch_size,1,1),dtype=torch.float32)\n",
    "\n",
    "s = c.States(stocks, debris)\n",
    "\n",
    "print(s.values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "attended-making",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([6.3344e-07, 4.6190e-07])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "constellation_collision_risk = 1e-6 * torch.rand(constellations)\n",
    "constellation_collision_risk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "strategic-american",
   "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])\n",
    "benefit_weight1 = torch.tensor([0.0,1.0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "hired-consent",
   "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": "copyrighted-tackle",
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "__init__() missing 1 required positional argument: 'launch_cost'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-6-54bb8ddad0e2>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mlaunch_cost\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m5\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m ea0 = c.LinearProfit(\n\u001b[0m\u001b[1;32m      3\u001b[0m                     \u001b[0;36m0\u001b[0m \u001b[0;31m#constellation index\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m                     \u001b[0;34m,\u001b[0m\u001b[0;36m0.95\u001b[0m \u001b[0;31m#discount\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m                     \u001b[0;34m,\u001b[0m\u001b[0mbenefit_weight0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: __init__() missing 1 required positional argument: 'launch_cost'"
     ]
    }
   ],
   "source": [
    "launch_cost = 5\n",
    "ea0 = c.LinearProfit(\n",
    "                    0 #constellation index\n",
    "                    ,0.95 #discount\n",
    "                    ,benefit_weight0\n",
    "                    ,launch_cost #launch_cost\n",
    "                    )\n",
    "ea1 = c.LinearProfit(\n",
    "                    1 #constellation index\n",
    "                    ,0.95 #discount\n",
    "                    ,benefit_weight1\n",
    "                    ,launch_cost #launch_cost\n",
    "                    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "accepted-namibia",
   "metadata": {},
   "outputs": [],
   "source": [
    "enn = nns.EstimandNN(batch_size\n",
    "                 ,states\n",
    "                 ,choices\n",
    "                 ,constellations\n",
    "                 ,12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "revolutionary-eight",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[0.0000],\n",
       "         [0.0000]],\n",
       "\n",
       "        [[0.0000],\n",
       "         [0.0021]],\n",
       "\n",
       "        [[0.1109],\n",
       "         [0.0835]],\n",
       "\n",
       "        [[0.0884],\n",
       "         [0.1051]],\n",
       "\n",
       "        [[0.0000],\n",
       "         [0.0000]]], grad_fn=<ReluBackward0>)"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "launch_decisions = enn.forward(s.values).choices\n",
    "launch_decisions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "abroad-mobile",
   "metadata": {},
   "outputs": [],
   "source": [
    "w = torch.tensor([[1.0,0],[0,-0.2]])\n",
    "ww = torch.tensor([1.0, -0.2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "seasonal-companion",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 1, 2])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stocks.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "jewish-zoning",
   "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 __str__(self):\n",
    "        return \"LinearProfit\\n Benefit weights:\\t{}\\n launch cost:\\t{}\\n Deorbit cost:\\t{}\".format(self.benefit_weights, self.launch_cost, self.deorbit_cost)\n",
    "\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": 17,
   "id": "surgical-diversity",
   "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(ww,stocks, [[0],[1]])[:,0] - (launch_cost * launch_decisions)[:,0,0]\n",
    "    return profit.view(batch_size,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "western-sixth",
   "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": 19,
   "id": "conscious-debut",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<__main__.LinearProfit at 0x7f0664fad4c0>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "eight-cheat",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[5.8800],\n",
       "        [1.9600],\n",
       "        [6.3054],\n",
       "        [2.4978],\n",
       "        [8.8200]], grad_fn=<ViewBackward>)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t._period_benefit(s.stocks,s.debris,launch_decisions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "juvenile-barcelona",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def f(stocks, debris, launches):\n",
    "    return torch.autograd.functional.jacobian(t._period_benefit\n",
    "                                   ,(stocks,debris,launches)\n",
    "                                   ,create_graph=True\n",
    "                                  )\n",
    "def ff(stocks, debris, launches):\n",
    "        return torch.autograd.functional.jacobian(f\n",
    "                                   ,(stocks,debris,launches)\n",
    "                                   ,create_graph=True\n",
    "                                  )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "freelance-publicity",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 1, 5, 1, 2, 5, 1, 2])"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "vocational-operator",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[6.0000, 1.0000],\n",
       "         [2.0000, 0.0000],\n",
       "         [7.0000, 6.0000],\n",
       "         [3.0000, 6.0000],\n",
       "         [9.0000, 1.0000]],\n",
       "\n",
       "        [[4.8000, 0.8000],\n",
       "         [1.6000, 0.0000],\n",
       "         [5.6000, 4.8000],\n",
       "         [2.4000, 4.8000],\n",
       "         [7.2000, 0.8000]]])"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.tensordot(torch.tensor([[1.0,-0.2],[0,1]]),stocks, [[0],[1]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "nuclear-alberta",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 1, 2])"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stocks.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "immune-machinery",
   "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
}