{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "weird-cloud",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.autograd.functional import jacobian\n",
    "import itertools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ready-attention",
   "metadata": {},
   "source": [
    "# Setup Functions\n",
    "## General CompositionFunctions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "convertible-thinking",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Set up functions to compose functions \n",
    "# These functions will \n",
    "#   - compose two functions together\n",
    "#   - compose a function to itself n times.\n",
    "\n",
    "def compose(f,g):\n",
    "    \"\"\"\n",
    "    this function composes two functions f and g in\n",
    "    the order f(g(*args))\n",
    "    \n",
    "    it returns a function\n",
    "    \"\"\"\n",
    "    return lambda *args: f(g(*args))\n",
    "\n",
    "def compose_recursive_functions(fn,n):\n",
    "    \"\"\"\n",
    "    This function takes a function fn and composes it with itself n times\n",
    "    \n",
    "    Returns an array/list that contains each of the n composition levels, ordered\n",
    "    from fn() to fn^n()\n",
    "    \"\"\"\n",
    "\n",
    "    \n",
    "    #Set base conditions\n",
    "    out_func = None\n",
    "    out_func_list =[]\n",
    "\n",
    "    #build the compositions of functions\n",
    "    for f in itertools.repeat(fn, n):\n",
    "        #if first iteration\n",
    "        if out_func == None:\n",
    "            out_func = f\n",
    "        else:\n",
    "            out_func = compose(f,out_func)\n",
    "\n",
    "        #append the ou\n",
    "        out_func_list.append(out_func)\n",
    "    \n",
    "    return out_func_list"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cloudy-creek",
   "metadata": {},
   "source": [
    "# functions related to transitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "musical-virtue",
   "metadata": {},
   "outputs": [],
   "source": [
    "def single_transition(item_to_iterate, laws_motion_fn, profit_fn, stocks, debris, launch_fn):\n",
    "    \"\"\"\n",
    "    This function represents the inverted envelope conditions.\n",
    "    It allows us to describe the derivatives of the value function evaluated at time $t+1$ in terms based in time period $t$.\n",
    "    \n",
    "    It takes a few derivatives (jacobians), and due to the nature of the return as tuples\n",
    "    they must be concatenated.\n",
    "    \n",
    "    It returns the transitioned values\n",
    "    \"\"\"\n",
    "    \n",
    "    #Find the matrix to invert. \n",
    "    #it consists of the derivative of the laws of motion with respect to stocks and debris\n",
    "    \n",
    "    #Get the jacobian\n",
    "    a = jacobian(laws_motion_fn, (stocks,debris, launch_fn(stocks,debris)))\n",
    "    \n",
    "    #Reassemble the Jacobian nested tuples into the appropriate tensor\n",
    "    A = BETA * torch.cat((torch.cat((a[0][0],a[0][1]),dim=1),torch.cat((a[1][0],a[1][1]),dim=1)), dim=0)\n",
    "\n",
    "    #TODO: figure out some diagnostics for this section\n",
    "    #Possibilities include:\n",
    "    # - Det(A) ~= 0\n",
    "    # - EigVal(A) ~= 0\n",
    "    # - A.inverse() with a try catch system to record types of returns\n",
    "    #Alternatively, \n",
    "    #if abs(a.det())\n",
    "    \n",
    "    #Calculate the item to transition\n",
    "    f_jacobians = jacobian(profit_fn,(stocks, debris, launch_fn(stocks,debris)))\n",
    "\n",
    "    #issue with shape here: my launch function is for all launches, not just a single launch.\n",
    "    f_theta = torch.cat([f_jacobians[0][0], f_jacobians[1][0]],axis=0) \n",
    "\n",
    "    T = item_to_iterate - f_theta\n",
    "\n",
    "    #Includes rearranging the jacobian of profit.\n",
    "\n",
    "    #Return the transitioned values\n",
    "    return  ( A.inverse()/BETA ) @  T\n",
    "\n",
    "# This function wraps the single transition and handles updating dates etc.\n",
    "def transition_wrapper(data_in):\n",
    "    \"\"\"\n",
    "    \"\"\"\n",
    "    #unpack states and functions\n",
    "    stocks, debris,profit_fn, laws_motion_fn, item_to_transition,launch_fn = data_in\n",
    "    \n",
    "    #Calculate new states\n",
    "    new_stocks, new_debris = laws_motion_fn(stocks,debris, launch_fn(stocks,debris))\n",
    "    \n",
    "    #WARNING: RECURSION: You may break your head...\n",
    "    #This gets the transition of the value function derivatives over time.\n",
    "    transitioned = single_transition(\n",
    "                    item_to_transition,  #item to iterate, i.e. the derivatives of the value function\n",
    "                    #functions\n",
    "                    laws_motion_fn, \n",
    "                    profit_fn, \n",
    "                    stocks, debris, #states\n",
    "                    launch_fn #launch function\n",
    "                    )\n",
    "    \n",
    "    #collects the data back together for return, including the updated state variables\n",
    "    data_out = new_stocks, new_debris, profit_fn, laws_motion_fn, transitioned, launch_fn\n",
    "    \n",
    "    return data_out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "quick-destruction",
   "metadata": {},
   "source": [
    "## Setup functions related to the problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "strategic-classics",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "    \n",
    "\n",
    "### functions\n",
    "\n",
    "def survival(stock, debris):\n",
    "    \"\"\"\n",
    "    This is a basic, deterministic survival function. \n",
    "    It is based on the CDF of an exponential distribution.\n",
    "    \"\"\"\n",
    "    #SURVIVAL FUNCTION BASED ON AN EXPONENTIAL DISTRIBUTION\n",
    "    return 1 - torch.exp(-SCALING * stock - debris)\n",
    "\n",
    "def test_launch(stock, debris):\n",
    "    \"\"\"\n",
    "    Temporary launch function \n",
    "    \"\"\"\n",
    "    return torch.ones(5, requires_grad=True),\n",
    "\n",
    "def laws_of_motion(stock, debris, launch):\n",
    "    \"\"\"\n",
    "    This function updates state variables (stock and debris), according \n",
    "    to the laws of motion.\n",
    "    \n",
    "    It returns the state variables as \n",
    "    \"\"\"\n",
    "\n",
    "    s = survival(stock,debris)\n",
    "    #Notes: Survival is a global function.\n",
    "    \n",
    "    new_stock = stock*s + launch\n",
    "    \n",
    "    \n",
    "    #TODO: Currently Ignoring autocatalysis\n",
    "    new_debris = (1-DELTA)*debris + LAUNCH_DEBRIS_RATE * launch.sum() + COLLISION_DEBRIS_RATE*(1-s) @ stock\n",
    "    \n",
    "    return (new_stock, new_debris)\n",
    "\n",
    "#This is not a good specification of the profit function, but it will work for now.\n",
    "def profit(stock, debris, launches):\n",
    "    return UTIL_WEIGHTS @ stock - LAUNCH_COST*launches\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "hairy-finding",
   "metadata": {},
   "source": [
    "# Actual calculations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "renewable-handy",
   "metadata": {},
   "outputs": [],
   "source": [
    "#number of states\n",
    "N = 5\n",
    "\n",
    "#set states\n",
    "stocks = torch.ones(N)\n",
    "#Last one is different if there are more than 1 (just for variety, no theoretical reason)\n",
    "if N > 1:\n",
    "    stocks[-1] = 0.5\n",
    "#now add the tracking requirement in place\n",
    "stocks.requires_grad_()\n",
    "\n",
    "#Setup Debris\n",
    "debris = torch.tensor([2.2],requires_grad=True)\n",
    "\n",
    "#CHANGE LATER: Launch is currently a value, should be a function (i.e. neural network)\n",
    "launches = test_launch\n",
    "\n",
    "#Starting point\n",
    "# Stocks, debris, profit fn, laws of motion, item to transition, Launch function\n",
    "base_data = (stocks,debris, profit, laws_of_motion, torch.ones(N+1, requires_grad=True),launches)\n",
    "\n",
    "#Parameters\n",
    "SCALING = torch.ones(5)\n",
    "DELTA = 0.9\n",
    "LAUNCH_DEBRIS_RATE = 0.005\n",
    "LAUNCH_COST = 1.0\n",
    "COLLISION_DEBRIS_RATE = 0.0007\n",
    "UTIL_WEIGHTS = torch.tensor([1,-0.2,0,0,0])\n",
    "BETA = 0.95"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "engaged-acceptance",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(tensor([1.9592, 1.9592, 1.9592, 1.9592, 1.4664], grad_fn=<AddBackward0>), tensor([0.2451], grad_fn=<AddBackward0>), <function profit at 0x7fa6269e4550>, <function laws_of_motion at 0x7fa6269e44c0>, tensor([-0.4523,  0.8774,  0.6558,  0.6558,  0.7608, 11.0954],\n",
      "       grad_fn=<MvBackward>), <function test_launch at 0x7fa6269e4430>) \n",
      "\n",
      "\n",
      "\n",
      "(tensor([2.7431, 2.7431, 2.7431, 2.7431, 2.2016], grad_fn=<AddBackward0>), tensor([0.0503], grad_fn=<AddBackward0>), <function profit at 0x7fa6269e4550>, <function laws_of_motion at 0x7fa6269e44c0>, tensor([-25.6591, -23.1244, -23.5468, -23.5468, -29.4675, 123.8316],\n",
      "       grad_fn=<MvBackward>), <function test_launch at 0x7fa6269e4430>) \n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#Get the values from 5 transitions\n",
    "wt1 = compose_recursive_functions(transition_wrapper,2)\n",
    "for f in wt1:\n",
    "    result = f(base_data)\n",
    "    print(result, \"\\n\"*3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ready-quick",
   "metadata": {},
   "source": [
    "Also, maybe I can create a `Model` class that upon construction will capture the necesary constants, functions, etc.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "intensive-forum",
   "metadata": {},
   "source": [
    "# Optimatility conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "prospective-mambo",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Optimality math\n",
    "def optimality_parts(stocks,debris,profit_fn,laws_motion_fn,launch_fn):\n",
    "    #Derivative of the value function with respect to choice functions\n",
    "    #this returns derivatives with respect to every launch, so I've removed that\n",
    "    fx = jacobian(profit_fn, (stocks,debris,launch_fn(stocks,debris)))[-1][:,0]\n",
    "    \n",
    "    \n",
    "    #The following returns a tuple of tuples of tensors.\n",
    "    #the first tuple contains jacobians related to laws of motion for stocks\n",
    "    #the second tuple contains jacobians related to laws of motion for debris.\n",
    "    #we need the derivatives related to both\n",
    "    b = jacobian(laws_of_motion,(stocks,debris,launches(stocks,debris)))\n",
    "    B = torch.cat((b[0][2],b[1][2].T),axis=1)\n",
    "\n",
    "\n",
    "    return fx, B\n",
    "\n",
    "#Optimality condition\n",
    "def optimality(stocks,debris,profit_fn,laws_motion_fn,launch_fn, iterated_item):\n",
    "    #calculate portions of the optimality values\n",
    "    fx, B = optimality_parts(stocks,debris,profit_fn,laws_motion_fn,launch_fn)\n",
    "\n",
    "    return fx + BETA *B @ iterated_item"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "distinguished-consultancy",
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp_result = torch.tensor([1.0,2.0,3,4,5,6])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "stylish-plaintiff",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(49.4968)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sum(optimality(stocks,debris,profit,laws_of_motion,launches,tmp_result)**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "accredited-wales",
   "metadata": {},
   "source": [
    "## Now to set up the recursive set of optimatliy conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "steady-chart",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([-1.3770,  0.8862,  0.6757,  0.6757,  0.7754], grad_fn=<AddBackward0>)\n",
      "tensor(4.1957, grad_fn=<AddBackward0>)\n",
      "\n",
      "tensor([-24.7879, -21.3799, -21.7813, -21.7813, -27.4059],\n",
      "       grad_fn=<AddBackward0>)\n",
      "tensor(2771.4756, grad_fn=<AddBackward0>)\n",
      "\n",
      "tensor([-218.6896, -214.1374, -214.7294, -214.7294, -290.5158],\n",
      "       grad_fn=<AddBackward0>)\n",
      "tensor(270296.8438, grad_fn=<AddBackward0>)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "base_data = (stocks,debris, profit, laws_of_motion, torch.ones(6, requires_grad=True),launches)\n",
    "for f in compose_recursive_functions(transition_wrapper,3):\n",
    "    result = f(base_data)\n",
    "    \n",
    "    #unpack results\n",
    "    new_stocks, new_debris, profit_fn, laws_motion_fn, transitioned, launch_fn = result\n",
    "    \n",
    "    optimal = optimality(new_stocks,new_debris,profit_fn,laws_motion_fn,launch_fn,transitioned)\n",
    "    print(optimal)\n",
    "    print(sum(optimal**2))\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "distributed-mailing",
   "metadata": {},
   "outputs": [],
   "source": [
    "def optimality_wrapper(stocks,debris, launches):\n",
    "    return optimality(stocks,debris, profit, laws_of_motion, launches, torch.ones(6, requires_grad=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "considerable-settle",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.0452,  0.9548,  0.9548,  0.9548,  0.9548], grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "optimality_wrapper(stocks,debris,launches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "second-aquarium",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(1.)"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "innocent-visibility",
   "metadata": {},
   "source": [
    "Notes so far\n",
    " - This takes $\\frac{\\partial W_t}{\\partial{\\theta_t}$ as given. I need to reframe this to clean that out. \n",
    "     - a failed method was just to calculate the partials of W for t = t. That doesn't work.\n",
    "         - Becasue the B value is not invertible, you can't just solve for the partials.\n",
    "         - note the approach below.\n",
    "     - Use part of the launch NN to approximate the gradient conditions.\n",
    "         - Complicates things slightly as it would require adding more outputs to the NN\n",
    "         - Would allow for PDE solution for value function.\n",
    "             - This might be useful for free-entry conditions.\n",
    "             - Not very helpful otherwise.\n",
    "             - I don't think it adds much to the final analysis.\n",
    "     \n",
    "             \n",
    " - There is an issue in how to handle the multiple value functions/loss functions\n",
    "     - Current status\n",
    "         - The launch function (soon to be NN) returns launch choices for each constellation.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "operating-carpet",
   "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
}