import torch
import combined as c

"""
This module holds the neural networks I am going to use to estimate 
the functions of interest.
"""

class LaunchFnEstimand(torch.nn.Module):
    """
    This is used to estimate the launch function
    """
    def __init__(self, state_tensor_size,layers_size,number_constellations):
        super().__init__()
        self.number_constellations = number_constellations
        self.layers_size = layers_size
        self.state_tensor_size = state_tensor_size
        
        
        #Layers
        self.linear_1 = torch.nn.Linear(in_features=state_tensor_size, out_features=layers_size)
        self.relu = torch.nn.ReLU()
        self.linear_3 = torch.nn.Linear(in_features=layers_size, out_features=layers_size)
        self.linear_5 = torch.nn.Linear(in_features=layers_size, out_features=number_constellations)

        
    def forward(self, input_values):
        
        intermediate_values = self.relu(input_values) #states should be positive anyway.
        intermediate_values = self.linear_1(intermediate_values)
        intermediate_values = self.linear_3(intermediate_values)
        intermediate_values = self.linear_5(intermediate_values)
        intermediate_values = self.relu(intermediate_values) #launches are always positive
        
        return intermediate_values
    
class PartialDerivativesEstimand(torch.nn.Module):
    """
    This is used to estimate the partial derivatives of the value functions
    """
    def __init__(self,batch_size, number_constellations, number_states, layer_size=12):
        super().__init__()
        self.batch_size = batch_size
        self.number_constellations = number_constellations
        self.number_states = number_states
        self.layer_size = layer_size
        
        
        #preprocess (single linear layer in case there is anything that needs to happen to all states)
        self.preprocess = torch.nn.Sequential(
            torch.nn.ReLU() #cleanup as states must be positive
            ,torch.nn.Linear(in_features = self.number_states, out_features=self.number_states)
        )
        #upscale to get the basic dimensionality correct. From (batch,State) to (batch, constellation, state). Includes a reshape
        self.upsample = lambda x: torch.nn.Upsample(scale_factor=self.number_constellations)(x).view(self.batch_size
                                                                                            ,self.number_constellations
                                                                                            ,self.number_states)
        
        #sequential steps
        self.sequential = torch.nn.Sequential(
            torch.nn.Linear(in_features=number_states, out_features=layer_size)
            #who knows if a convolution might help here.
            ,torch.nn.Linear(in_features=layer_size, out_features=layer_size)
            ,torch.nn.Linear(in_features=layer_size, out_features=layer_size)
        )

        #reduce the feature axis to match expected results
        self.feature_reduction = torch.nn.Linear(in_features=layer_size, out_features=number_states)
        
    def forward(self, input_values):
        #Note that the input values are just going to be the state variables
        #TODO:check that input values match the prepared dimension?
        
        #preprocess
        intermediate = self.preprocess(input_values)
        
        #upscale the input values
        intermediate = self.upsample(intermediate)
        
        #intermediate processing
        intermediate = self.sequential(intermediate)
        
        #reduce feature axis to match the expected number of partials
        intermediate = self.feature_reduction(intermediate)
        
        
        return intermediate
    
class EstimandNN(torch.nn.Module):
    """
    This neural network takes the current states as input values and returns both
    the partial derivatives of the value function and the launch function.
    """
    def __init__(self, state_tensor_size,layers_size,number_constellations):
        super().__init__()
        
        #So, this next section constructs different layers within the NN
        #sinlge linear section
        pass
        #TODO:verify these are correct
        self.partials_estimator = PartialDerivativesEstimand(state_tensor_size,layers_size,number_constellations) #TODO
        self.launch_estimator = LaunchFnEstimand(state_tensor_size,layers_size,number_constellations)
        
    def forward(self, input_values):
        pass
        partials = self.partials_estimator(input_values)
        launch = self.launch_estimator(input_values)
        
        return c.EstimandInterface(partials,launch)