
"""
Truncated multivariate normal distribution without oracle access (ie. unknown truncation set)
"""

from re import I
import torch as ch
from torch import Tensor
from torch.distributions.multivariate_normal import MultivariateNormal
import cox 

from .censored_multivariate_normal import CensoredMultivariateNormalModel
from ..oracle import UnknownGaussian
from .distributions import distributions
from ..trainer import Trainer
from ..grad import TruncatedMultivariateNormalNLL
from ..utils.datasets import TruncatedNormalDataset, make_train_and_val_distr
from ..utils.helpers import Parameters, PSDError
from ..utils.defaults import check_and_fill_args, TRAINER_DEFAULTS, DELPHI_DEFAULTS, TRUNC_MULTI_NORM_DEFAULTS


class TruncatedMultivariateNormal(distributions):
    """
    Truncated multivariate normal distribution class.
    """
    def __init__(self,
            args: Parameters, 
            store: cox.store.Store=None):
        super(TruncatedMultivariateNormal).__init__()
        # instance variables 
        assert isinstance(args, Parameters), "args is type {}. expecting type delphi.utils.helper.Parameters.".format(type(args))
        assert store is None or isinstance(store, cox.store.Store), "store is type: {}. expecting cox.store.Store.".format(type(store))
        self.store = store 
        self.truncated = None
        # algorithm hyperparameters
        TRUNC_MULTI_NORM_DEFAULTS.update(DELPHI_DEFAULTS)
        TRUNC_MULTI_NORM_DEFAULTS.update(TRAINER_DEFAULTS)
        self.args = check_and_fill_args(args, TRUNC_MULTI_NORM_DEFAULTS)

    def fit(self, S: Tensor):
        
        assert isinstance(S, Tensor), "S is type: {}. expected type torch.Tensor.".format(type(S))
        assert S.size(0) > S.size(1), "input expected to be num samples by dimensions, current input is size {}.".format(S.size()) 
        while True:
            try:
                self.train_loader_, self.val_loader_ = make_train_and_val_distr(self.args, S, TruncatedNormalDataset)
                self.truncated = TruncatedMultivariateNormalModel(self.args, self.train_loader_.dataset)
                self.trainer = Trainer(self.truncated, self.args, store=self.store) 
        
                # run PGD for parameter estimation 
                self.trainer.train_model((self.train_loader_, self.val_loader_))
                return self
            except PSDError as psd:
                print(psd.message) 
                continue
            except Exception as e: 
                    raise e

    @property 
    def loc_(self): 
        """
        Returns the mean of the normal disribution.
        """
        return self.truncated.model.loc.clone()
  
    @property 
    def covariance_matrix_(self): 
        """
        Returns the standard deviation for the normal distribution.
        """
        return self.truncated.model.covariance_matrix.clone()


class TruncatedMultivariateNormalModel(CensoredMultivariateNormalModel):
    '''
    Model for truncated normal distributions to be passed into trainer.
    '''
    def __init__(self, args, train_ds):
        '''
        Args: 
            args (delphi.utils.Parameters) : parameter object holding hyperparameters
        '''
        super().__init__(args, train_ds)
        # initialiaze pseudo oracle for gaussians with unknown truncation 
        self.emp_loc, self.emp_covariance_matrix = None, None
        # initialize empirical estimates
        self.calc_emp_model()
        self.args.__setattr__('phi', UnknownGaussian(self.emp_loc, self.emp_covariance_matrix, self.train_ds.S, self.args.d))

        # exponent class
        self.exp_h = Exp_h(self.emp_loc, self.emp_covariance_matrix)

    def __call__(self, batch):
        '''
        Training step for defined model.
        Args: 
            batch (Iterable) : iterable of inputs
        '''
        loss = TruncatedMultivariateNormalNLL.apply(self.model.loc, self.model.covariance_matrix, *batch, self.args.phi, self.exp_h)
        return loss, None, None

    def calc_emp_model(self): 
        # initialize projection set
        self.emp_covariance_matrix = self.train_ds.covariance_matrix
        self.emp_loc = self.train_ds.loc
        self.model = MultivariateNormal(self.emp_loc, self.emp_covariance_matrix)

    def post_training_hook(self):
        # reparamterize distribution
        self.model.covariance_matrix.requires_grad, self.model.loc.requires_grad = False, False
        self.model.covariance_matrix.data = self.model.covariance_matrix.inverse()
        self.model.loc.data = (self.model.loc[None,...]  @ self.model.covariance_matrix).flatten()
        # set estimated distribution in membership oracle
        self.args.phi.dist = self.model


# HELPER FUNCTIONS
class Exp_h:
    def __init__(self, emp_loc, emp_cov):
        self.emp_loc = emp_loc
        self.emp_cov = emp_cov
        self.pi_const = (self.emp_loc.size(0) / 2.0) * ch.log(2.0 * Tensor([ch.acos(ch.zeros(1)).item() * 2]).unsqueeze(0))

    def __call__(self, u, B, x):
        """
        returns: evaluates exponential function
        """
        cov_term = ch.bmm(x.unsqueeze(1)@B, x.unsqueeze(2)).squeeze(1) / 2.0
        # trace_term = ch.trace((B - self.emp_cov) * (self.emp_cov + self.emp_loc[...,None]@self.emp_loc[None,...])).unsqueeze(0) / 2.0
        trace_term = ch.trace((B - ch.eye(u.size(0))) * (self.emp_cov + self.emp_loc[...,None]@self.emp_loc[None,...])).unsqueeze(0) / 2.0
        loc_term = (x - self.emp_loc)@u.unsqueeze(1)
        return ch.exp((cov_term - trace_term - loc_term + self.pi_const).double())
