{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 124,
   "id": "3d049ef3-2630-4064-a2bd-ee28c7a89e9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "using LinearAlgebra, Zygote, Tullio"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfa5d466-5354-4ef4-a4ab-2bf2ed69626e",
   "metadata": {},
   "source": [
    "### notes\n",
    "\n",
    "Notice how this approach estimates policy and the value function in a single go.\n",
    "While you could eliminate the value function partials, it would be much more complex.\n",
    "Note that in RL, both must be iterated on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2bc6f9bd-8308-40a0-bee4-acbb1be4672d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Model dimensions\n",
    "N_constellations = 5\n",
    "N_debris = 1\n",
    "N_states = N_constellations + N_debris"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15670d3b-9651-4fc5-b98a-29ba60b4b38e",
   "metadata": {},
   "source": [
    "## Built Payoff functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "id": "1e99275a-9ab5-4c9e-921b-9587145b1fb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "stocks = rand(1:N_constellations,N_constellations);\n",
    "debris = rand(1:3);\n",
    "payoff = 3*I + ones(5,5) .+ [1,0,0,0,0]; #TODO: move this into a struct\n",
    "β=0.95\n",
    "launches = ones(N_constellations);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 157,
   "id": "a4e4bbaf-61de-4860-adc2-1e91f0626ead",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Define the market profit function\n",
    "F(stocks,debris,payoff,launches) = payoff*stocks  + 3.0*launches .+ (debris*-0.2)\n",
    "\n",
    "#create derivative functions\n",
    "∂f_∂stocks(st,debris,payoff,launches) = Zygote.jacobian(s -> F(s,debris,payoff,launches),st)[1];\n",
    "∂f_∂debris(stocks,d,payoff,launches) = Zygote.jacobian(s -> F(stocks,s,payoff,launches),d)[1];\n",
    "∂f_∂launches(stocks,debris,payoff,launches) = Zygote.jacobian(l -> F(stocks,debris,payoff,l),launches)[1];\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "id": "ad73e64a-2fa5-4bdf-a4e4-22357b53002c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5×5 Matrix{Float64}:\n",
       " 5.0  2.0  2.0  2.0  2.0\n",
       " 1.0  4.0  1.0  1.0  1.0\n",
       " 1.0  1.0  4.0  1.0  1.0\n",
       " 1.0  1.0  1.0  4.0  1.0\n",
       " 1.0  1.0  1.0  1.0  4.0"
      ]
     },
     "execution_count": 158,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "∂f_∂stocks(stocks,debris,payoff,launches) \n",
    "#Rows are constellations, columns are derivatives indexed by constellations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 159,
   "id": "a6159740-8ef2-40e0-b10f-c451cf78418d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5-element Vector{Float64}:\n",
       " -0.2\n",
       " -0.2\n",
       " -0.2\n",
       " -0.2\n",
       " -0.2"
      ]
     },
     "execution_count": 159,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "∂f_∂debris(stocks,debris,payoff,launches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 160,
   "id": "b46da3d9-1765-4427-b3de-bf0a9fdd6576",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5×5 Matrix{Float64}:\n",
       " 3.0  0.0  0.0  0.0  0.0\n",
       " 0.0  3.0  0.0  0.0  0.0\n",
       " 0.0  0.0  3.0  0.0  0.0\n",
       " 0.0  0.0  0.0  3.0  0.0\n",
       " 0.0  0.0  0.0  0.0  3.0"
      ]
     },
     "execution_count": 160,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "∂f_∂launches(stocks,debris,payoff,launches)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a3a09cc-e48d-4c27-9e4f-02e45a5fc5c1",
   "metadata": {},
   "source": [
    "## Building Physical Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "638ed4ac-36b0-4efb-bcb4-13d345470dd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "struct BasicModel\n",
    "    #rate at which debris hits satellites\n",
    "    debris_collision_rate\n",
    "    #rate at which satellites of different constellations collide\n",
    "    satellite_collision_rates\n",
    "    #rate at which debris exits orbits\n",
    "    decay_rate\n",
    "    #rate at which satellites\n",
    "    autocatalysis_rate\n",
    "    #ratio at which a collision between satellites produced debris\n",
    "    satellite_collision_debris_ratio\n",
    "    #Ratio at which launches produce debris\n",
    "    launch_debris_ratio\n",
    "end\n",
    "\n",
    "#Getting loss parameters together.\n",
    "loss_param = 2e-3;\n",
    "loss_weights = loss_param*(ones(N_constellations,N_constellations) - I);\n",
    "\n",
    "#orbital decay rate\n",
    "decay_param = 0.01;\n",
    "\n",
    "#debris generation parameters\n",
    "autocatalysis_param = 0.001;\n",
    "satellite_loss_debris_rate = 5.0;\n",
    "launch_debris_rate = 0.05;\n",
    "\n",
    "#Todo, wrap physical model as a struct with the parameters\n",
    "bm = BasicModel(\n",
    "    loss_param\n",
    "    ,loss_weights\n",
    "    ,decay_param\n",
    "    ,autocatalysis_param\n",
    "    ,satellite_loss_debris_rate\n",
    "    ,launch_debris_rate\n",
    ");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "id": "ae04a4f2-7401-450e-ba98-bfec375b7646",
   "metadata": {},
   "outputs": [],
   "source": [
    "#percentage survival function\n",
    "function survival(stocks,debris,physical_model) \n",
    "    exp.(-physical_model.satellite_collision_rates*stocks .- (physical_model.debris_collision_rate*debris));\n",
    "end\n",
    "\n",
    "#stock update rules\n",
    "function G(stocks,debris,launches, physical_model)\n",
    "    return diagm(survival(stocks,debris,physical_model) .- physical_model.decay_rate)*stocks + launches\n",
    "end;\n",
    "\n",
    "#stock evolution wrt various things\n",
    "∂G_∂launches(stocks,debris,launches,bm) = Zygote.jacobian(x -> G(stocks,debris,x,bm),launches)[1];\n",
    "∂G_∂debris(stocks,debris,launches,bm) = Zygote.jacobian(x -> G(stocks,x,launches,bm),debris)[1];\n",
    "∂G_∂stocks(stocks,debris,launches,bm) = Zygote.jacobian(x -> G(x,debris,launches,bm),stocks)[1];\n",
    "\n",
    "#debris evolution \n",
    "function H(stocks,debris,launches,physical_model)\n",
    "    #get changes in debris from natural dynamics\n",
    "    natural_debris_dynamics = (1-physical_model.decay_rate+physical_model.autocatalysis_rate) * debris \n",
    "    \n",
    "    #get changes in debris from satellite loss\n",
    "    satellite_loss_debris = physical_model.satellite_collision_debris_ratio * (1 .- survival(stocks,debris,physical_model))'*stocks \n",
    "    \n",
    "    #get changes in debris from launches\n",
    "    launch_debris = physical_model.launch_debris_ratio*sum(launches)\n",
    "    \n",
    "    #return total debris level\n",
    "    return natural_debris_dynamics + satellite_loss_debris + launch_debris\n",
    "end;\n",
    "\n",
    "#get jacobians of debris dynamics\n",
    "∂H_∂launches(stocks,debris,launches,physical_model) = Zygote.jacobian(x -> H(stocks,debris,x,physical_model),launches)[1];\n",
    "∂H_∂debris(stocks,debris,launches,physical_model) = Zygote.jacobian(x -> H(stocks,x,launches,physical_model),debris)[1];\n",
    "∂H_∂stocks(stocks,debris,launches,physical_model) = Zygote.jacobian(x -> H(x,debris,launches,physical_model),stocks)[1];"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 166,
   "id": "c6519b6e-f9ef-466e-a9ac-746b8f36b9e7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5×5 Matrix{Float64}:\n",
       " 1.0  0.0  0.0  0.0  0.0\n",
       " 0.0  1.0  0.0  0.0  0.0\n",
       " 0.0  0.0  1.0  0.0  0.0\n",
       " 0.0  0.0  0.0  1.0  0.0\n",
       " 0.0  0.0  0.0  0.0  1.0"
      ]
     },
     "execution_count": 166,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "∂G_∂launches(stocks,debris,launches,bm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 167,
   "id": "1d46c09d-2c91-4302-8e8d-1c0abeb010aa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5-element Vector{Float64}:\n",
       " -0.003843157756609293\n",
       " -0.009665715046375067\n",
       " -0.009665715046375067\n",
       " -0.009665715046375067\n",
       " -0.003843157756609293"
      ]
     },
     "execution_count": 167,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "∂G_∂debris(stocks,debris,launches,bm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 168,
   "id": "1f791c8a-5457-4ec2-bf0b-68eb6bdd578f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5×5 Matrix{Float64}:\n",
       "  0.950789    -0.00384316  -0.00384316  -0.00384316  -0.00384316\n",
       " -0.00966572   0.956572    -0.00966572  -0.00966572  -0.00966572\n",
       " -0.00966572  -0.00966572   0.956572    -0.00966572  -0.00966572\n",
       " -0.00966572  -0.00966572  -0.00966572   0.956572    -0.00966572\n",
       " -0.00384316  -0.00384316  -0.00384316  -0.00384316   0.950789"
      ]
     },
     "execution_count": 168,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "∂G_∂stocks(stocks,debris,launches,bm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "id": "d97f9af0-d3a9-4739-b569-fa100a1590f3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1×5 Matrix{Float64}:\n",
       " 0.05  0.05  0.05  0.05  0.05"
      ]
     },
     "execution_count": 169,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "∂H_∂launches(stocks,debris,launches,bm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "id": "1a7c857f-7eee-4273-b954-3a7b84a3b3e0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1-element Vector{Float64}:\n",
       " 1.174417303261719"
      ]
     },
     "execution_count": 170,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "∂H_∂debris(stocks,debris,launches,bm) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "id": "9ea44146-740c-47c0-9286-846585a4db39",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1×5 Matrix{Float64}:\n",
       " 0.360254  0.302231  0.302231  0.302231  0.360254"
      ]
     },
     "execution_count": 171,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "∂H_∂stocks(stocks,debris,launches,bm)\n",
    "#columns are derivatives"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17038bfe-c7c2-484b-9aaf-6b608bbbbfab",
   "metadata": {},
   "source": [
    "## Build optimality conditions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef3c7625-e5ee-4220-8bed-b289d0baea1e",
   "metadata": {},
   "source": [
    "## Build transition conditions\n",
    "\n",
    "I've built the transition conditions below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "id": "c3ef6bcc-37ca-4f6c-b876-8ea2a068cbe1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5×5 Matrix{Float64}:\n",
       " 2.0  1.0  1.0  1.0  1.0\n",
       " 1.0  2.0  1.0  1.0  1.0\n",
       " 1.0  1.0  2.0  1.0  1.0\n",
       " 1.0  1.0  1.0  2.0  1.0\n",
       " 1.0  1.0  1.0  1.0  2.0"
      ]
     },
     "execution_count": 172,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "W_∂stocks= ones(5,5) + I\n",
    "# columns are stocks\n",
    "# rows are derivatives"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "id": "5441641b-308b-4fec-a946-2568e79a1203",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5×1 Matrix{Float64}:\n",
       " 2.0\n",
       " 2.0\n",
       " 2.0\n",
       " 2.0\n",
       " 2.0"
      ]
     },
     "execution_count": 174,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    " #temporary value partials\n",
    "W_∂debris = 2*ones(5,1) #temporary value partials\n",
    "#columns are\n",
    "#rows are"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "id": "819df0f9-e994-4461-9329-4c4607de8779",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5×5 Matrix{Float64}:\n",
       "  0.966286    -0.00195257  -0.00195257  -0.00195257  -0.00195257\n",
       " -0.00785729   0.972161    -0.00785729  -0.00785729  -0.00785729\n",
       " -0.00588119  -0.00588119   0.970199    -0.00588119  -0.00588119\n",
       " -0.00195257  -0.00195257  -0.00195257   0.966286    -0.00195257\n",
       " -0.00391296  -0.00391296  -0.00391296  -0.00391296   0.96824"
      ]
     },
     "execution_count": 106,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "∂G_∂stocks(stocks,debris,launches,bm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "id": "5cc19637-9727-479e-aa60-a30c71e9e8b8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5×5 Matrix{Float64}:\n",
       " 1.91695  1.91695  1.91695  1.91695  1.91695\n",
       " 1.88146  1.88146  1.88146  1.88146  1.88146\n",
       " 1.89335  1.89335  1.89335  1.89335  1.89335\n",
       " 1.91695  1.91695  1.91695  1.91695  1.91695\n",
       " 1.90518  1.90518  1.90518  1.90518  1.90518"
      ]
     },
     "execution_count": 108,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = ∂f_∂stocks(stocks,debris,payoff,launches)\n",
    "b = ∂G_∂stocks(stocks,debris,launches,bm) * W_∂stocks #This last bit should eventually get replaced with a NN\n",
    "#Need to check dimensionality above. Not sure which direction is derivatives and which is functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "id": "65917e4c-bd25-4683-8553-17bc841e594a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5×1 Matrix{Float64}:\n",
       " -0.019525714195158184\n",
       " -0.07857288258866406\n",
       " -0.05881192039840531\n",
       " -0.019525714195158184\n",
       " -0.039129609402048404"
      ]
     },
     "execution_count": 119,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c = (∂G_∂debris(stocks,debris,launches,bm) .* ones(1,5)) * W_∂debris"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "dd5ca5b4-c910-449c-a322-b24b3ef4ecf2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "49.45445380487922"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss(stocks,debris,payoff,launches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "8ead6e86-3d03-48ff-bc3e-b3fa9808f981",
   "metadata": {},
   "outputs": [],
   "source": [
    "#TODO: create a launch model in flux and see if I can get it to do pullbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "658ba7f0-dfbc-42da-bb20-7f8bccbb257d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.6.2",
   "language": "julia",
   "name": "julia-1.6"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}