import torch
from torch import nn
from torch.nn import functional as F
from torch.optim import Adam
from pytorch_lightning.core.lightning import LightningModule

class FrustaNetRegression(LightningModule):

    def __init__(self, n_features, n_estimators=1):
        super().__init__()

        """Every model starts with linear estimator.
        All subsquent models are added with non-linearities.

            Args:
                n_features: dimensions of input
                n_estimators: number of weak learners to fit
        """
        self.n_estimators = n_estimators

        self.linear = torch.nn.Linear(n_features, 1)

        if n_estimators > 1:
            self.estimators = nn.ModuleList()

            for n in range(n_estimators):
                estim_layers = nn.ModuleList()
                estim_layers.append(torch.nn.Linear(n_features, n_features*2))
                estim_layers.append(torch.nn.Sigmoid())
                estim_layers.append(torch.nn.Linear(n_features*2, 1))
                self.estimators.append(estim_layers)

    def forward(self, x):
        """
            Take in a value of x and return regression
            value.

            Args:
                x: model inputs
            Returns:
                out: list of model predictions
        """

        out_linear = self.linear(x)

        estim_out = [out_linear]
        
        if self.n_estimators > 1:
            for i, estim in enumerate(self.estimators):
                layer_in = x
                for layer in estim:
                    layer_in = layer(layer_in)
                estim_out.append(layer_in)

        return estim_out

    def training_step(self, batch, batch_idx):
        """
            The first estimator of this model is a linear
            model. Every subsequent model seeks to predict
            the residual of the estimator before it. The loss
            function is the sum over the individual squared
            residuals of all the estimators.

            Args: 
                batch
                batch_idx

            Returns:
                loss

        """
        x, y = batch
        estim_out = self(x)
        # Start with the linear residual
        # If the is only one estimator selected
        # The model collapses to a linear regression
        # With squared error loss
        resids = [y - estim_out[0]]
        if self.n_estimators > 1:
            for i, e_out in enumerate(estim_out):
                if i > 0:
                    resids.append(resids[i-1] - e_out)

        loss = 0
        for i, r in enumerate(resids):
            l = torch.mean(r ** 2)
            self.log('Loss Estimator '+str(i), l, on_step=True)
            loss += l

        self.log('Loss', loss, on_step=True)
        return loss

    def configure_optimizers(self):
        return Adam(self.parameters(), lr=0.01, weight_decay=0.001)

    def validation_step(self, batch, batch_idx):
        """
            On validation we want to compute the 
            mean squared error of the model.
            The prediction is generated by summing 
            up all the individual estimators predictions.
        """
        mseloss = nn.MSELoss()
        x, y = batch
        pred = self.predict(x)
        loss = mseloss(y, pred)
        self.log('Val Loss', loss)
        return loss

    def predict(self, x):
        estim_out = self(x)
        pred = 0
        for e in estim_out:
            pred += e
        return pred
