from fortuna.typing import Batch, Params, Mutable, CalibMutable, CalibParams, Outputs, Targets
from typing import List, Optional, Tuple, Union, Any, Callable
from jax._src.prng import PRNGKeyArray
import jax.numpy as jnp
from fortuna.utils.random import WithRNG
from fortuna.likelihood.base import Likelihood


class Loss(WithRNG):
    def __init__(self, likelihood: Likelihood, loss_fn: Callable[[Outputs, Targets], jnp.ndarray]):
        self.likelihood = likelihood
        self.loss_fn = loss_fn

    def __call__(
        self,
        params: Params,
        batch: Batch,
        mutable: Optional[Mutable] = None,
        calib_params: Optional[CalibParams] = None,
        calib_mutable: Optional[CalibMutable] = None,
        return_aux: Optional[List[str]] = None,
        train: bool = False,
        outputs: Optional[jnp.ndarray] = None,
        rng: Optional[PRNGKeyArray] = None,
        **kwargs
    ) -> Union[jnp.ndarray, Tuple[jnp.ndarray, Any]]:
        if return_aux is None:
            return_aux = []
        supported_aux = ["outputs", "mutable", "calib_mutable"]
        unsupported_aux = [s for s in return_aux if s not in supported_aux]
        if sum(unsupported_aux) > 0:
            raise AttributeError(
                """The auxiliary objects {} is unknown. Please make sure that all elements of `return_aux` 
                            belong to the following list: {}""".format(
                    unsupported_aux, supported_aux
                )
            )
        if train and outputs is not None:
            raise ValueError(
                """When `outputs` is available, `train` must be set to `False`."""
            )
        if "mutable" in return_aux and outputs is not None:
            raise ValueError(
                """When `outputs` is available, `return_aux` cannot contain 'mutable'`."""
            )
        if not train and "mutable" in return_aux:
            raise ValueError(
                "Returning an auxiliary mutable is supported only during training. Please either set `train` to "
                "`True`, or remove 'mutable' from `return_aux`."
            )
        if "mutable" in return_aux and mutable is None:
            raise ValueError(
                "In order to be able to return an auxiliary mutable, an initial mutable must be passed as `mutable`. "
                "Please either remove 'mutable' from `return_aux`, or pass an initial mutable as `mutable`."
            )
        if "mutable" not in return_aux and mutable is not None and train is True:
            raise ValueError(
                """You need to add `mutable` to `return_aux`. When you provide a (not null) `mutable` variable during 
                training, that variable will be updated during the forward pass."""
            )

        inputs, targets = batch
        if outputs is None:
            outs = self.likelihood.model_manager.apply(
                params,
                inputs,
                train=train,
                mutable=mutable,
                rng=rng,
            )
            if "mutable" in return_aux:
                outputs, aux = outs
                mutable = aux["mutable"]
            else:
                outputs = outs

        aux = dict()
        if self.likelihood.output_calib_manager is not None:
            outs = self.likelihood.output_calib_manager.apply(
                params=calib_params["output_calibrator"]
                if calib_params is not None
                else None,
                mutable=calib_mutable["output_calibrator"]
                if calib_mutable is not None
                else None,
                outputs=outputs,
                calib="calib_mutable" in return_aux,
            )
        if (
            calib_mutable is not None
            and calib_mutable["output_calibrator"] is not None
            and "calib_mutable" in return_aux
        ):
            outputs, aux["calib_mutable"] = outs
            aux["calib_mutable"] = dict(output_calibrator=aux["calib_mutable"])
        else:
            outputs = outs
            if "calib_mutable" in return_aux:
                aux["calib_mutable"] = dict(output_calibrator=None)

        if "outputs" in return_aux:
            aux["outputs"] = outputs
        if "mutable" in return_aux:
            aux["mutable"] = mutable

        loss = self.loss_fn(outputs, targets)
        if len(return_aux) > 0:
            return loss, aux
        return loss
