{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "expired-austria",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.autograd.functional import jacobian\n",
    "import itertools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "damaged-accountability",
   "metadata": {},
   "source": [
    "# Setup Functions\n",
    "## General CompositionFunctions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "modular-memorabilia",
   "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": "direct-picture",
   "metadata": {},
   "source": [
    "# functions related to transitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "bridal-ordinary",
   "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": "miniature-karaoke",
   "metadata": {},
   "source": [
    "## Setup functions related to the problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "bright-minimum",
   "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"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "conservative-ukraine",
   "metadata": {},
   "source": [
    "# Actual calculations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "initial-mathematics",
   "metadata": {},
   "outputs": [],
   "source": [
    "#set states\n",
    "stocks = torch.ones(5)\n",
    "#Last one is different\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, \n",
    "base_data = (stocks,debris, profit, laws_of_motion, torch.ones(6, 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": 75,
   "id": "nuclear-definition",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([-0.4523,  0.8774,  0.6558,  0.6558,  0.7608, 11.0954],\n",
      "       grad_fn=<MvBackward>) \n",
      "\n",
      "\n",
      "\n",
      "tensor([-25.6591, -23.1244, -23.5468, -23.5468, -29.4675, 123.8316],\n",
      "       grad_fn=<MvBackward>) \n",
      "\n",
      "\n",
      "\n",
      "tensor([-236.0462, -232.3070, -232.9302, -232.9302, -312.7054, 1379.8452],\n",
      "       grad_fn=<MvBackward>) \n",
      "\n",
      "\n",
      "\n",
      "tensor([-1638.0408, -1632.9258, -1633.7783, -1633.7783, -2380.2834, 15341.8955],\n",
      "       grad_fn=<MvBackward>) \n",
      "\n",
      "\n",
      "\n",
      "tensor([ -9855.1680,  -9848.4297,  -9849.5527,  -9849.5527, -15547.5273,\n",
      "        170308.7812], grad_fn=<MvBackward>) \n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#Get the values from 5 transitions\n",
    "for f in compose_recursive_functions(transition_wrapper,5):\n",
    "    result = f(base_data)\n",
    "    print(result[4], \"\\n\"*3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "casual-annex",
   "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": "stopped-socket",
   "metadata": {},
   "source": [
    "# Optimatility conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 193,
   "id": "excessive-script",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Optimality condition\n",
    "def optimality(stocks,debris,profit_fn,laws_motion_fn,launch_fn, iterated_item):\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 + BETA * B @ iterated_item"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 195,
   "id": "unlikely-coverage",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(49.4968)"
      ]
     },
     "execution_count": 195,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sum(optimality(stocks,debris,profit,laws_of_motion,launches,tmp_result)**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "endless-occupation",
   "metadata": {},
   "source": [
    "## Now to set up the recursive set of optimatliy conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 179,
   "id": "valuable-bleeding",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.0300, 2.0300, 3.0300, 4.0300, 5.0300])"
      ]
     },
     "execution_count": 179,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "subjective-chassis",
   "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
}