import torch
from torch.autograd.functional import jacobian
import itertools
import math


############### CONSTANTS ###################
#Parameters
BETA = 0.95

#Constants determining iterations etc.
NUMBER_CONSTELLATIONS = 5
NUMBER_DEBRIS_TRACKERS = 1
NUMBER_OF_CHOICE_VARIABLES = 1
NUMBER_OF_REQUIRED_ITERATED_CONDITIONS = (NUMBER_CONSTELLATIONS+NUMBER_DEBRIS_TRACKERS+(NUMBER_OF_CHOICE_VARIABLES*NUMBER_CONSTELLATIONS))
NUMBER_OF_REQUIRED_ITERATIONS = math.ceil(NUMBER_OF_REQUIRED_ITERATED_CONDITIONS/NUMBER_CONSTELLATIONS)

############# COMPOSITION FUNCTIONS ###################

def compose(f,g):
    """
    this function composes two functions f and g in
    the order f(g(*args))
    
    it returns a function.
    """
    return lambda *args: f(g(*args))

def recursively_compose_functions(fn,n):
    """
    This function takes a function fn and composes it with itself n times.
    
    Returns an array/list that contains each of the n+1 composition levels, ordered
    from fn() to fn^n()
    """
    #Set base conditions
    out_func = None
    out_func_list = []

    #build the compositions of functions
    for f in itertools.repeat(fn, n):
        #if first iteration
        if out_func == None:
            out_func = f
        else:
            out_func = compose(f,out_func)

        #append the ou
        out_func_list.append(out_func)
    
    return out_func_list

############### Physical Model ###################
class PhysicalModel():
    """
    This is designed as a default, base class for the physical model of satellite interactions.

    It captures the constants that characterize interactions, and provides a set of
    function to calculate changes to the physical environment.
    """
    def __init__(self
            ,collision_debris
            ,constellations_collision_risk
            ,debris_decay_rate
            ,launch_debris
            ,debris_autocatalysis_rate
            )
    self.collision_debris = collision_debris 
    self.constellations_collision_risk = constellations_collision_risk
    self.debris_decay_rate = debris_decay_rate
    self.launch_debris = launch_debris
    self.debris_autocatalysis_rate = debris_autocatalysis_rate


    def survival(self):
        #returns the survival rate (not destruction rate) for the given constellation.
        return 1- torch.exp(-self.constellations_collision_risk * self.stock - self.debris)

    def transition_debris(self, state, launch_decisions):
        """
        This function transitions debris levels based off of a state and launch decision.
        """

        new_debris = (1-self.debris_decay_rate + self.debris_autocatalysis_rate) * state.debris \ #debris decay and autocatalysis
                    + self.launch_debris*launch_decisions.sum() \ #debris from launches
                    + self.collision_debris * (1-self.survival()) @ state.stocks
        return new_debris

    def transition_stocks(self, state, launch_decisions):

        new_stock = self.survival() * state.stocks + launch_decisions

        return new_stock

    def transition(self, state, launch_decisions):
        """
        This function takes a state and launch decision, and updates the state according to the physical laws of motion.

        It returns a State object.
        """
        d = self.transition_debris(state, launch_decisions)
        s = self.transition_stocks(state, launch_decisions)

        return States(s,d)

class States():
    """
    This is supposed to capture the state variables of the model, to create a common interface 
    when passing between functions.
    """
    def __init__(self, stocks,debris):
        self.stocks = stocks
        self.debris = debris

################ NEURAL NETWORK TOOLS ###################
    
class EstimandInterface():
    """
    This defines a clean interface for working with the estimand (i.e. thing we are trying to estimate).
    In general, we are trying to estimate the choice variables and the partial derivatives of the value functions.
    This 

    This class wraps output for the neural network (or other estimand), allowing me to 
        - easily substitute various types of launch functions by having a common interface
            - this eases testing
        - check dimensionality etc without dealing with randomness
            - again, easing testing
        - reason more cleanly about the component pieces
            - easing programming
        - provide a clean interface to find constellation level launch decisions etc.

    It takes inputs of two general categories:
        - the choice function results
        - the partial derivatives of the value function
    """
    def __init__(self, partials, launches, deorbits=None):
        self.partials = partials
        self.launches = launches
        self.deorbits = deorbits
    
    def launch_single(constellation):
        #returns the launch decision for the constellation of interest
        
        filter_tensor = torch.zeros(NUMBER_CONSTELLATIONS)
        filter_tensor[constellation] = 1.0
        
        return self.launches @ filter_tensor
    
    def launch_vector(constellation):
        #returns the launch decision for the constellation of interest as a vector
        
        filter_tensor = torch.zeros(NUMBER_CONSTELLATIONS)
        filter_tensor[constellation] = 1.0
        
        return self.launches * filter_tensor
    
    def partial_vector(constellation):
        #returns the partials of the value function corresponding to the constellation of interest
        
        filter_tensor = torch.zeros(NUMBER_CONSTELLATIONS)
        filter_tensor[constellation] = 1.0
        
        return self.partials @ filter_tensor
    
    def partial_matrix(constellation):
        #returns the partials of the value function corresponding to 
        #the constellation of interest as a matrix
        
        filter_tensor = torch.zeros(NUMBER_CONSTELLATIONS)
        filter_tensor[constellation] = 1.0
        
        return self.partials * filter_tensor
    
    def __str__(self):
        #just a human readable descriptor
        return "Launch Decisions and Partial Derivativs of value function with\n\t states\n\t\t {}\n\tPartials\n\t\t{}".format(self.states,self.partials)


############## ECONOMIC MODEL ############

class EconomicModel():
    """
    This class describes the set of profit functions involved in the value function iteration.
    """
    def __init__(self,discount_factor, profit_objects):
        self.discount_factor = discount_factor
        self.profit_objects = profit_objects #A list of Profit objects

    def constellation_period_profits(self, state, estimand_interface, constellation):
        """
        This function calculates the current period profits for a single given constellation.
        """
        return self.profit_objects[constellation].profit(state,estimand_interface,constellation)

    def period_profits(self, state, estimand_interface):
        """
        This function calculates the current period profits for each constellation.
        """
        profits = []
        for i,profit_object in self.profit_objectives:
            constellation_profit = self.constellation_period_profits(state, estimand_interface, i)
            profits.append(constellation_profit)
        return profits

    @property
    def number_constellations(self):
        return len(profit_objects)
            

#Abstract class describing profit. Each subclass will connect a profit function "style" to a specific instance of parameters.
class ProfitFunctions(metaclass=abc.ABCMetaclass):
    @abc.abstractmethod
    def profit(self,state,estimand_interface,constellation_number):
        pass
    #TODO: Should I attach the jacobian here? It may simplify things.
    #def jacobian()

class LinearProfit(ProfitFunctions):
    """
    The simplest type of profit function available.
    """
    def __init__(self, benefit_weight, launch_cost, deorbit_cost=0):
        self.benefit_weights = benefit_weight
        self.launch_cost = launch_cost
        self.deorbit_cost = deorbit_cost

    def profit(self, state, estimand_interface, constellation_number):
        #Calculate the profit
        profits =  self.benefit_weights @ state.stock \
                    - self.launch_cost * estimand_interface.launches[constellation_number] #\ 
                    #- deorbit_cost @ estimand_interface.deorbits[constellation_number]
        return profits

#other profit functions to implement
# price competition (substitution)
# military (complementarity)

############### TRANSITION AND OPTIMALITY FUNCTIONS #################
#Rewrite these to use the abstractiosn present earlier

def single_transition(physical_model, economic_model, states, estimand_interface):
    """
    This function represents the inverted envelope conditions.
    It allows us to describe the derivatives of the value function evaluated at time $t+1$ in terms based in time period $t$.
    
    It takes a few derivatives (jacobians), and due to the nature of the return as tuples
    they must be concatenated.
    
    It returns the transitioned values
    """
    #TODO: rewrite using the current abstractions
    #possibly move jacobians of profit functions and physical models to those classes

    pass

def transition_wrapper(data_in): # Identify a way to eliminate this maybe?
    """
    This function wraps the single transition and handles updating states etc.
    
    """
    #TODO: rewrite using current abstractions
    pass

#Optimality math
def optimality(stocks
                ,debris
                ,profit_fn
                ,laws_motion_fn
                ,neural_net
               ):
    """
    This function takes in the 
     - stock levels
     - debris levels
     - profit function
     - laws of motion
     - results from the neural network
    and returns the parts used to make up the optimality conditions
    """
    pass