import tensorflow as tf
import tensorflow_probability as tfp
tfm = tfp.mcmc
from ..networks import BaseFullyConnectedNet, Discriminator, BayesianFullyConnectedNet
import numpy as np
from bayesgm.datasets import Gaussian_sampler
from bayesgm.utils.data_io import save_data
import dateutil.tz
import datetime
import os
from tqdm import tqdm

class CausalBGM(object):
    """Causal Bayesian Generative Model (CausalBGM) for causal inference.

    CausalBGM learns a latent-variable generative model for causal inference
    with treatment :math:`X`, outcome :math:`Y`, and high-dimensional covariates
    :math:`V`.  The latent variable :math:`Z` is partitioned into
    :math:`(Z_0, Z_1, Z_2, Z_3)` to disentangle confounding, outcome-specific,
    treatment-specific, and residual variation.

    Parameters
    ----------
    params : dict
        Configuration dictionary.  Required keys:

        - ``'v_dim'`` (int): Dimension of covariates :math:`V`.
        - ``'z_dims'`` (list[int]): Dimensions ``[z0, z1, z2, z3]`` of the
          four latent sub-vectors.
        - ``'binary_treatment'`` (bool): ``True`` for binary treatment,
          ``False`` for continuous.
        - ``'dataset'`` (str): Dataset name (used for checkpoint paths).
        - ``'output_dir'`` (str): Root directory for outputs.

        Optional keys (with defaults):

        - ``'use_bnn'`` (bool): Whether to use Bayesian neural networks. Default ``True``.
        - ``'g_units'`` (list[int]): Hidden-layer sizes for the generator network. Default ``[64, 64, 64, 64, 64]``.
        - ``'e_units'`` (list[int]): Hidden-layer sizes for the encoder network. Default ``[64, 64, 64, 64, 64]``.
        - ``'f_units'`` (list[int]): Hidden-layer sizes for the outcome network. Default ``[64, 32, 8]``.
        - ``'h_units'`` (list[int]): Hidden-layer sizes for the treatment network. Default ``[64, 32, 8]``.
        - ``'dz_units'`` (list[int]): Hidden-layer sizes for the latent discriminator. Default ``[64, 32, 8]``.
        - ``'lr'`` (float): Learning rate for EGM pre-training. Default ``0.0002``.
        - ``'lr_theta'`` (float): Learning rate for network parameters. Default ``0.0001``.
        - ``'lr_z'`` (float): Learning rate for latent-variable updates. Default ``0.0001``.
        - ``'g_d_freq'`` (int): Discriminator-to-generator update ratio. Default ``5``.
        - ``'save_model'`` (bool): Whether to save model checkpoints. Default ``False``.
        - ``'save_res'`` (bool): Whether to save results. Default ``True``.
        - ``'kl_weight'`` (float): KL-divergence weight when ``use_bnn`` is True. Default ``0.0001``.
    timestamp : str or None, optional
        Timestamp string for the run.  If ``None``, the current local time
        is used.
    random_seed : int or None, optional
        If provided, sets the global random seed for reproducibility.
    """

    def __init__(self, params, timestamp=None, random_seed=None):
        super(CausalBGM, self).__init__()
        self.params = params
        self.timestamp = timestamp
        if random_seed is not None:
            tf.keras.utils.set_random_seed(random_seed)
            os.environ['TF_DETERMINISTIC_OPS'] = '1'
            tf.config.experimental.enable_op_determinism()
        if self.params['use_bnn']:
            self.g_net = BayesianFullyConnectedNet(input_dim=sum(params['z_dims']),output_dim = params['v_dim']+1, 
                                           model_name='g_net', nb_units=params['g_units'])
            self.e_net = BayesianFullyConnectedNet(input_dim=params['v_dim'],output_dim = sum(params['z_dims']), 
                                            model_name='e_net', nb_units=params['e_units'])
            self.f_net = BayesianFullyConnectedNet(input_dim=params['z_dims'][0]+params['z_dims'][1]+1,
                                           output_dim = 2, model_name='f_net', nb_units=params['f_units'])
            self.h_net = BayesianFullyConnectedNet(input_dim=params['z_dims'][0]+params['z_dims'][2],
                                           output_dim = 2, model_name='h_net', nb_units=params['h_units'])
        else:
            self.g_net = BaseFullyConnectedNet(input_dim=sum(params['z_dims']),output_dim = params['v_dim']+1, 
                                           model_name='g_net', nb_units=params['g_units'])
            self.e_net = BaseFullyConnectedNet(input_dim=params['v_dim'],output_dim = sum(params['z_dims']), 
                                            model_name='e_net', nb_units=params['e_units'])
            self.f_net = BaseFullyConnectedNet(input_dim=params['z_dims'][0]+params['z_dims'][1]+1,
                                           output_dim = 2, model_name='f_net', nb_units=params['f_units'])
            self.h_net = BaseFullyConnectedNet(input_dim=params['z_dims'][0]+params['z_dims'][2],
                                           output_dim = 2, model_name='h_net', nb_units=params['h_units'])

        self.dz_net = Discriminator(input_dim=sum(params['z_dims']),model_name='dz_net',
                                        nb_units=params['dz_units'])

        self.g_pre_optimizer = tf.keras.optimizers.Adam(params['lr'], beta_1=0.9, beta_2=0.99)
        self.d_pre_optimizer = tf.keras.optimizers.Adam(params['lr'], beta_1=0.9, beta_2=0.99)
        self.z_sampler = Gaussian_sampler(mean=np.zeros(sum(params['z_dims'])), sd=1.0)

        self.g_optimizer = tf.keras.optimizers.Adam(params['lr_theta'], beta_1=0.9, beta_2=0.99)
        self.f_optimizer = tf.keras.optimizers.Adam(params['lr_theta'], beta_1=0.9, beta_2=0.99)
        self.h_optimizer = tf.keras.optimizers.Adam(params['lr_theta'], beta_1=0.9, beta_2=0.99)
        self.posterior_optimizer = tf.keras.optimizers.Adam(params['lr_z'], beta_1=0.9, beta_2=0.99)
        
        self.initialize_nets()
        if self.timestamp is None:
            now = datetime.datetime.now(dateutil.tz.tzlocal())
            self.timestamp = now.strftime('%Y%m%d_%H%M%S')
        
        self.checkpoint_path = "{}/checkpoints/{}/{}".format(
            params['output_dir'], params['dataset'], self.timestamp)

        if self.params['save_model'] and not os.path.exists(self.checkpoint_path):
            os.makedirs(self.checkpoint_path)

        self.save_dir = "{}/results/{}/{}".format(
            params['output_dir'], params['dataset'], self.timestamp)

        if self.params['save_res'] and not os.path.exists(self.save_dir):
            os.makedirs(self.save_dir)   

        self.ckpt = tf.train.Checkpoint(g_net = self.g_net,
                                    e_net = self.e_net,
                                    f_net = self.f_net,
                                    h_net = self.h_net,
                                    dz_net = self.dz_net,
                                    g_pre_optimizer = self.g_pre_optimizer,
                                    d_pre_optimizer = self.d_pre_optimizer,
                                    g_optimizer = self.g_optimizer,
                                    f_optimizer = self.f_optimizer,
                                    h_optimizer = self.h_optimizer,
                                    posterior_optimizer = self.posterior_optimizer)
        
        self.ckpt_manager = tf.train.CheckpointManager(self.ckpt, self.checkpoint_path, max_to_keep=5)                 

        if self.ckpt_manager.latest_checkpoint:
            self.ckpt.restore(self.ckpt_manager.latest_checkpoint)
            print ('Latest checkpoint restored!!') 

    def get_config(self):
        """Return the configuration of the CausalBGM model.

        Returns
        -------
        dict
            A dictionary with key ``"params"`` containing the full
            configuration dictionary passed at construction time.
        """

        return {
                "params": self.params,
        }

    def initialize_nets(self, print_summary = False):
        """Initialize all the networks in CausalBGM."""

        self.g_net(np.zeros((1, sum(self.params['z_dims']))))
        self.f_net(np.zeros((1, self.params['z_dims'][0]+self.params['z_dims'][1]+1)))
        self.h_net(np.zeros((1, self.params['z_dims'][0]+self.params['z_dims'][2])))
        if print_summary:
            print(self.g_net.summary())
            print(self.f_net.summary())    
            print(self.h_net.summary()) 

    def _sample_flat_weights_from_bnn(self, bayes_net):
        """Sample one concrete weight vector from a DenseFlipout network posterior."""
        flattened = []
        for layer in bayes_net.all_layers:
            kernel = layer.kernel_posterior.sample()
            bias = layer.bias_posterior.sample() if layer.bias_posterior is not None else tf.zeros([layer.units], dtype=kernel.dtype)
            flattened.append(tf.reshape(kernel, [-1]))
            flattened.append(tf.reshape(bias, [-1]))
        return tf.concat(flattened, axis=0)

    def sample_theta_weights(self, method='sample'):
        """Draw one theta sample (or posterior mean) for g/h/f networks."""
        if not self.params['use_bnn']:
            return None, None, None
        if method not in ('sample', 'mean'):
            raise ValueError("theta method must be 'sample' or 'mean'.")
        if method == 'sample':
            g_weights = self._sample_flat_weights_from_bnn(self.g_net)
            h_weights = self._sample_flat_weights_from_bnn(self.h_net)
            f_weights = self._sample_flat_weights_from_bnn(self.f_net)
        else:
            def _flat_mean_weights(net):
                chunks = []
                for layer in net.all_layers:
                    chunks.append(tf.reshape(layer.kernel_posterior.mean(), [-1]))
                    bias = layer.bias_posterior.mean() if layer.bias_posterior is not None else tf.zeros([layer.units], dtype=tf.float32)
                    chunks.append(tf.reshape(bias, [-1]))
                return tf.concat(chunks, axis=0)
            g_weights = _flat_mean_weights(self.g_net)
            h_weights = _flat_mean_weights(self.h_net)
            f_weights = _flat_mean_weights(self.f_net)
        return tf.cast(g_weights, tf.float32), tf.cast(h_weights, tf.float32), tf.cast(f_weights, tf.float32)

    def _bnn_forward_with_flat_weights(self, bayes_net, inputs, flattened_weights):
        """Deterministic forward pass of BayesianFullyConnectedNet with frozen weights."""
        # Keep deterministic frozen-theta forward consistent with
        # BayesianFullyConnectedNet.call() after input BN removal.
        x = inputs
        start = 0
        for layer in bayes_net.all_layers[:-1]:
            k_shape = layer.kernel_posterior.mean().shape
            b_shape = layer.bias_posterior.mean().shape
            k_size = int(np.prod(k_shape))
            b_size = int(np.prod(b_shape))
            kernel = tf.reshape(flattened_weights[start:start + k_size], k_shape)
            start += k_size
            bias = tf.reshape(flattened_weights[start:start + b_size], b_shape)
            start += b_size
            x = tf.matmul(x, kernel) + bias
            x = tf.keras.layers.LeakyReLU(alpha=0.2)(x)
        layer = bayes_net.all_layers[-1]
        k_shape = layer.kernel_posterior.mean().shape
        b_shape = layer.bias_posterior.mean().shape
        k_size = int(np.prod(k_shape))
        b_size = int(np.prod(b_shape))
        kernel = tf.reshape(flattened_weights[start:start + k_size], k_shape)
        start += k_size
        bias = tf.reshape(flattened_weights[start:start + b_size], b_shape)
        return tf.matmul(x, kernel) + bias

    # Update generative model for covariates V
    @tf.function
    def update_g_net(self, data_z, data_v, eps=1e-6):
        with tf.GradientTape() as gen_tape:
            g_net_output = self.g_net(data_z)
            mu_v = g_net_output[:,:self.params['v_dim']]
            if 'sigma_v' in self.params:
                sigma_square_v = self.params['sigma_v']**2
            else:
                sigma_square_v = tf.nn.softplus(g_net_output[:,-1]) + eps
            #loss = -log(p(x|z))
            loss_mse = tf.reduce_mean((data_v - mu_v)**2)
            loss_v = tf.reduce_sum((data_v - mu_v)**2, axis=1)/(2*sigma_square_v) + \
                    self.params['v_dim'] * tf.math.log(sigma_square_v)/2
            loss_v = tf.reduce_mean(loss_v)
            
            if self.params['use_bnn']:
                loss_kl = sum(self.g_net.losses)
                loss_v += loss_kl * self.params['kl_weight']

        # Calculate the gradients for generators and discriminators
        g_gradients = gen_tape.gradient(loss_v, self.g_net.trainable_variables)
        
        # Apply the gradients to the optimizer
        self.g_optimizer.apply_gradients(zip(g_gradients, self.g_net.trainable_variables))
        return loss_v, loss_mse
    
    # Update generative model for treatment X
    @tf.function
    def update_h_net(self, data_z, data_x, eps=1e-6):
        with tf.GradientTape() as gen_tape:
            data_z0 = data_z[:,:self.params['z_dims'][0]]
            data_z2 = data_z[:,sum(self.params['z_dims'][:2]):sum(self.params['z_dims'][:3])]
            h_net_output = self.h_net(tf.concat([data_z0, data_z2], axis=-1))
            mu_x = h_net_output[:,:1]
            if self.params['binary_treatment']:
                loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=data_x, 
                                                       logits=mu_x))
                loss_x =  loss
            else:
                if 'sigma_x' in self.params:
                    sigma_square_x = self.params['sigma_x']**2
                else:
                    sigma_square_x = tf.nn.softplus(h_net_output[:,-1]) + eps
                #loss = -log(p(x|z))
                loss = tf.reduce_mean((data_x - mu_x)**2)
                loss_x = tf.reduce_sum((data_x - mu_x)**2, axis=1)/(2*sigma_square_x) + \
                        tf.math.log(sigma_square_x)/2
                loss_x = tf.reduce_mean(loss_x)

            if self.params['use_bnn']:
                loss_kl = sum(self.h_net.losses)
                loss_x += loss_kl * self.params['kl_weight']
                
        # Calculate the gradients for generators and discriminators
        h_gradients = gen_tape.gradient(loss_x, self.h_net.trainable_variables)
        
        # Apply the gradients to the optimizer
        self.h_optimizer.apply_gradients(zip(h_gradients, self.h_net.trainable_variables))
        return loss_x, loss
    
    # Update generative model for outcome Y
    @tf.function
    def update_f_net(self, data_z, data_x, data_y, eps=1e-6):
        with tf.GradientTape() as gen_tape:
            data_z0 = data_z[:,:self.params['z_dims'][0]]
            data_z1 = data_z[:,self.params['z_dims'][0]:sum(self.params['z_dims'][:2])]
            f_net_output = self.f_net(tf.concat([data_z0, data_z1, data_x], axis=-1))
            mu_y = f_net_output[:,:1]
            if 'sigma_y' in self.params:
                sigma_square_y = self.params['sigma_y']**2
            else:
                sigma_square_y = tf.nn.softplus(f_net_output[:,-1]) + eps
            #loss = -log(p(y|z,x))
            loss_mse = tf.reduce_mean((data_y - mu_y)**2)
            loss_y = tf.reduce_sum((data_y - mu_y)**2, axis=1)/(2*sigma_square_y) + \
                    tf.math.log(sigma_square_y)/2
            loss_y = tf.reduce_mean(loss_y)
            
            if self.params['use_bnn']:
                loss_kl = sum(self.f_net.losses)
                loss_y += loss_kl * self.params['kl_weight']

        # Calculate the gradients for generators and discriminators
        f_gradients = gen_tape.gradient(loss_y, self.f_net.trainable_variables)
        
        # Apply the gradients to the optimizer
        self.f_optimizer.apply_gradients(zip(f_gradients, self.f_net.trainable_variables))
        return loss_y, loss_mse
    
    # Update posterior of latent variables Z
    @tf.function
    def update_latent_variable_sgd(self, data_x, data_y, data_v, data_z, eps=1e-6):
        with tf.GradientTape() as tape:
            
            data_z0 = data_z[:,:self.params['z_dims'][0]]
            data_z1 = data_z[:,self.params['z_dims'][0]:sum(self.params['z_dims'][:2])]
            data_z2 = data_z[:,sum(self.params['z_dims'][:2]):sum(self.params['z_dims'][:3])]
            
            # logp(v|z) for covariate model
            mu_v = self.g_net(data_z)[:,:self.params['v_dim']]
            if 'sigma_v' in self.params:
                sigma_square_v = self.params['sigma_v']**2
            else:
                sigma_square_v = tf.nn.softplus(self.g_net(data_z)[:,-1]) + eps
                
            loss_pv_z = tf.reduce_sum((data_v - mu_v)**2, axis=1)/(2*sigma_square_v) + \
                    self.params['v_dim'] * tf.math.log(sigma_square_v)/2
            loss_pv_z = tf.reduce_mean(loss_pv_z)
            
            # log(x|z) for treatment model
            mu_x = self.h_net(tf.concat([data_z0, data_z2], axis=-1))[:,:1]
            if 'sigma_x' in self.params:
                sigma_square_x = self.params['sigma_x']**2
            else:
                sigma_square_x = tf.nn.softplus(self.h_net(tf.concat([data_z0, data_z2], axis=-1))[:,-1]) + eps

            if self.params['binary_treatment']:
                loss_px_z = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=data_x, 
                                                       logits=mu_x))
            else:
                loss_px_z = tf.reduce_sum((data_x - mu_x)**2, axis=1)/(2*sigma_square_x) + \
                        tf.math.log(sigma_square_x)/2
                loss_px_z = tf.reduce_mean(loss_px_z)
                
            # log(y|z,x) for outcome model
            mu_y = self.f_net(tf.concat([data_z0, data_z1, data_x], axis=-1))[:,:1]
            if 'sigma_y' in self.params:
                sigma_square_y = self.params['sigma_y']**2
            else:
                sigma_square_y = tf.nn.softplus(self.f_net(tf.concat([data_z0, data_z1, data_x], axis=-1))[:,-1]) + eps

            loss_py_zx = tf.reduce_sum((data_y - mu_y)**2, axis=1)/(2*sigma_square_y) + \
                    tf.math.log(sigma_square_y)/2
            loss_py_zx = tf.reduce_mean(loss_py_zx)

            loss_prior_z =  tf.reduce_sum(data_z**2, axis=1)/2
            loss_prior_z = tf.reduce_mean(loss_prior_z)

            loss_postrior_z = loss_pv_z + loss_px_z + loss_py_zx + loss_prior_z
            #loss_postrior_z = loss_postrior_z/self.params['v_dim']

        # Calculate the gradients
        posterior_gradients = tape.gradient(loss_postrior_z, [data_z])
        # Apply the gradients to the optimizer
        self.posterior_optimizer.apply_gradients(zip(posterior_gradients, [data_z]))
        return loss_postrior_z
    
#################################### EGM initialization ###########################################
    @tf.function
    def train_disc_step(self, data_z, data_v):
        epsilon_z = tf.random.uniform([],minval=0., maxval=1.)
        with tf.GradientTape(persistent=True) as disc_tape:
            with tf.GradientTape() as gp_tape:
                data_z_ = self.e_net(data_v)
                data_z_hat = data_z*epsilon_z + data_z_*(1-epsilon_z)
                data_dz_hat = self.dz_net(data_z_hat)

            data_dz_ = self.dz_net(data_z_)
            data_dz = self.dz_net(data_z)
            dz_loss = -tf.reduce_mean(data_dz) + tf.reduce_mean(data_dz_)

            # Calculate gradient penalty 
            grad_z = gp_tape.gradient(data_dz_hat, data_z_hat)
            grad_norm_z = tf.sqrt(tf.reduce_sum(tf.square(grad_z), axis=1))
            gpz_loss = tf.reduce_mean(tf.square(grad_norm_z - 1.0))
            
            d_loss = dz_loss + 10 * gpz_loss

        # Calculate the gradients for generators and discriminators
        d_gradients = disc_tape.gradient(d_loss, self.dz_net.trainable_variables)
        
        # Apply the gradients to the optimizer
        self.d_pre_optimizer.apply_gradients(zip(d_gradients, self.dz_net.trainable_variables))
        return dz_loss, d_loss
    
    @tf.function
    def train_gen_step(self, data_z, data_v, data_x, data_y):
        with tf.GradientTape(persistent=True) as gen_tape:
            sigma_square_loss = 0
            data_v_ = self.g_net(data_z)[:,:self.params['v_dim']]
            sigma_square_loss += tf.reduce_mean(tf.square(self.g_net(data_z)[:,-1]))
            data_z_ = self.e_net(data_v)
            
            data_z0 = data_z_[:,:self.params['z_dims'][0]]
            data_z1 = data_z_[:,self.params['z_dims'][0]:sum(self.params['z_dims'][:2])]
            data_z2 = data_z_[:,sum(self.params['z_dims'][:2]):sum(self.params['z_dims'][:3])]

            data_z__= self.e_net(data_v_)
            data_v__ = self.g_net(data_z_)[:,:self.params['v_dim']]
            
            data_dz_ = self.dz_net(data_z_)
            
            l2_loss_v = tf.reduce_mean((data_v - data_v__)**2)
            l2_loss_z = tf.reduce_mean((data_z - data_z__)**2)
            
            e_loss_adv = -tf.reduce_mean(data_dz_)

            data_y_ = self.f_net(tf.concat([data_z0, data_z1, data_x], axis=-1))[:,:1]
            sigma_square_loss += tf.reduce_mean(
                tf.square(self.f_net(tf.concat([data_z0, data_z1, data_x], axis=-1))[:,-1]))
            data_x_ = self.h_net(tf.concat([data_z0, data_z2], axis=-1))[:,:1]
            sigma_square_loss += tf.reduce_mean(
                tf.square(self.h_net(tf.concat([data_z0, data_z2], axis=-1))[:,-1]))

            if self.params['binary_treatment']:
                l2_loss_x = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=data_x, 
                                                       logits=data_x_))
            else:
                l2_loss_x = tf.reduce_mean((data_x_ - data_x)**2)
            l2_loss_y = tf.reduce_mean((data_y_ - data_y)**2)
            g_e_loss = e_loss_adv+(l2_loss_v + self.params['use_z_rec']*l2_loss_z) \
                        + (l2_loss_x+l2_loss_y) + 0.001 * sigma_square_loss

        # Calculate the gradients for generators and discriminators
        g_e_gradients = gen_tape.gradient(g_e_loss, self.g_net.trainable_variables+self.e_net.trainable_variables+\
                                        self.f_net.trainable_variables+self.h_net.trainable_variables)
        
        # Apply the gradients to the optimizer
        self.g_pre_optimizer.apply_gradients(zip(g_e_gradients, self.g_net.trainable_variables+self.e_net.trainable_variables+\
                                            self.f_net.trainable_variables+self.h_net.trainable_variables))
        return e_loss_adv, l2_loss_v, l2_loss_z, l2_loss_x, l2_loss_y, g_e_loss
    

    def egm_init(self, data, egm_n_iter=10000, batch_size=32, egm_batches_per_eval=500, verbose=1):
        """Run Encoding Generative Modeling (EGM) initialization.

        Initializes network parameters by jointly fitting the generator,
        encoder, and discriminator before the main iterative optimization.
        This is optional and can be skipped by setting `use_egm_init` to False in :meth:`fit`.

        Parameters
        ----------
        data : tuple of np.ndarray
            A triplet ``(data_x, data_y, data_v)``.
        egm_n_iter : int, default=10000
            Number of EGM mini-batch iterations.
        batch_size : int, default=32
            Mini-batch size.
        egm_batches_per_eval : int, default=500
            Evaluate and log every this many iterations.
        verbose : int, default=1
            Verbosity level (0 = silent).
        """
        data_x, data_y, data_v = data
        
        print('EGM Initialization Starts ...')
        for batch_iter in range(egm_n_iter+1):
            # Update model parameters of Discriminator
            for _ in range(self.params['g_d_freq']):
                batch_idx = np.random.choice(len(data_x), batch_size, replace=False)
                batch_z = self.z_sampler.get_batch(batch_size)
                batch_v = data_v[batch_idx,:]
                dz_loss, d_loss = self.train_disc_step(batch_z, batch_v)

            # Update model parameters of G, H, F with SGD
            batch_z = self.z_sampler.get_batch(batch_size)
            batch_idx = np.random.choice(len(data_x), batch_size, replace=False)
            batch_x = data_x[batch_idx,:]
            batch_y = data_y[batch_idx,:]
            batch_v = data_v[batch_idx,:]
            e_loss_adv, l2_loss_v, l2_loss_z, l2_loss_x, l2_loss_y, g_e_loss = self.train_gen_step(batch_z, batch_v, batch_x, batch_y)
            if batch_iter % egm_batches_per_eval == 0:
                
                loss_contents = (
                    'EGM Initialization Iter [%d] : e_loss_adv [%.4f], l2_loss_v [%.4f], l2_loss_z [%.4f], '
                    'l2_loss_x [%.4f], l2_loss_y [%.4f], g_e_loss [%.4f], dz_loss [%.4f], d_loss [%.4f]'
                    % (batch_iter, e_loss_adv, l2_loss_v, l2_loss_z, l2_loss_x, l2_loss_y, g_e_loss, dz_loss, d_loss)
                )
                if verbose:
                    print(loss_contents)
                if self.params['save_res']:
                    causal_pre, mse_x, mse_y, mse_v, data_x_pred, data_y_pred, data_v_pred = self.evaluate(data = data)
                    causal_pre = causal_pre.numpy()
                    save_data('{}/causal_pre_egm_init_iter-{}.txt'.format(self.save_dir, batch_iter), causal_pre)
        print('EGM Initialization Ends.')
#################################### EGM initialization #############################################

    def fit(self, data,
            batch_size=32, epochs=100, epochs_per_eval=5, startoff=0,
            use_egm_init=True, egm_n_iter=30000, egm_batches_per_eval=500, save_format='txt', verbose=1):
        """Train the CausalBGM model on observed data.

        The training procedure consists of two phases:

        1. **EGM initialization** (optional) — warm-start by jointly training encoder and
           generator with adversarial losses to obtain a good starting point
           for the latent variables and model parameters. This phase is optional and can be skipped by setting `use_egm_init` to False.
        2. **Iterative optimization** — alternates between updating the
           generator network parameters :math:`\\theta` and the per-sample
           latent variables :math:`Z` via SGD.
           
        Parameters
        ----------
        data : tuple of np.ndarray
            A triplet ``(data_x, data_y, data_v)`` where

            - ``data_x`` has shape ``(n, 1)`` — treatment variable.
            - ``data_y`` has shape ``(n, 1)`` — outcome variable.
            - ``data_v`` has shape ``(n, v_dim)`` — covariates.
        batch_size : int, default=32
            Mini-batch size.
        epochs : int, default=100
            Number of training epochs for the iterative phase.
        epochs_per_eval : int, default=5
            Evaluate and (optionally) save every this many epochs.
        startoff : int, default=0
            Only start tracking the best model after this many epochs.
        use_egm_init : bool, default=True
            Whether to run EGM initialization before iterative training.
        egm_n_iter : int, default=30000
            Number of EGM initialization iterations.
        egm_batches_per_eval : int, default=500
            Evaluate EGM every this many iterations.
        save_format : str, default='txt'
            File format for saving causal estimates (``'txt'`` or ``'npy'``).
        verbose : int, default=1
            Verbosity level.
        """
        
        data_x, data_y, data_v = data
        
        if self.params['save_res']:
            f_params = open('{}/params.txt'.format(self.save_dir),'w')
            f_params.write(str(self.params))
            f_params.close()
        
        if use_egm_init:
            self.egm_init(data, egm_n_iter=egm_n_iter, egm_batches_per_eval=egm_batches_per_eval, batch_size=batch_size, verbose=verbose)
            print('Initialize latent variables Z with e(V)...')
            data_z_init = self.e_net(data_v)
        else:
            print('Random initialization of latent variables Z...')
            data_z_init = np.random.normal(0, 1, size = (len(data_x), sum(self.params['z_dims']))).astype('float32')

        self.data_z = tf.Variable(data_z_init, name="Latent Variable",trainable=True)
        
        best_loss = np.inf
        print('Iterative Updating Starts ...')
        for epoch in range(epochs+1):
            sample_idx = np.random.choice(len(data_x), len(data_x), replace=False)
            
            # Create a progress bar for batches
            with tqdm(total=len(data_x) // batch_size, desc=f"Epoch {epoch}/{epochs}", unit="batch") as batch_bar:
                for i in range(0,len(data_x) - batch_size + 1,batch_size): ## Skip the incomplete last batch
                    batch_idx = sample_idx[i:i+batch_size]
                    # Update model parameters of G, H, F with SGD
                    batch_z = tf.Variable(tf.gather(self.data_z, batch_idx, axis = 0), name='batch_z', trainable=True)
                    batch_x = data_x[batch_idx,:]
                    batch_y = data_y[batch_idx,:]
                    batch_v = data_v[batch_idx,:]
                    loss_v, loss_mse_v = self.update_g_net(batch_z, batch_v)
                    loss_x, loss_mse_x = self.update_h_net(batch_z, batch_x)
                    loss_y, loss_mse_y = self.update_f_net(batch_z, batch_x, batch_y)

                    # Update Z by maximizing a posterior or posterior mean
                    loss_postrior_z = self.update_latent_variable_sgd(batch_x, batch_y, batch_v, batch_z)

                    # Update data_z with updated batch_z
                    self.data_z.scatter_nd_update(
                        indices=tf.expand_dims(batch_idx, axis=1),
                        updates=batch_z                             
                    )
                    
                    # Update the progress bar with the current loss information
                    loss_contents = (
                        'loss_px_z: [%.4f], loss_mse_x: [%.4f], loss_py_z: [%.4f], '
                        'loss_mse_y: [%.4f], loss_pv_z: [%.4f], loss_mse_v: [%.4f], loss_postrior_z: [%.4f]'
                        % (loss_x, loss_mse_x, loss_y, loss_mse_y, loss_v, loss_mse_v, loss_postrior_z)
                    )
                    batch_bar.set_postfix_str(loss_contents)
                    batch_bar.update(1)
            
            # Evaluate the full training data and print metrics for the epoch
            if epoch % epochs_per_eval == 0:
                causal_pre, mse_x, mse_y, mse_v, data_x_pred, data_y_pred, data_v_pred = self.evaluate(data = data, data_z = self.data_z)
                causal_pre = causal_pre.numpy()
                
                if verbose:
                    print('Epoch [%d/%d]: MSE_x: %.4f, MSE_y: %.4f, MSE_v: %.4f\n' % (epoch, epochs, mse_x, mse_y, mse_v))

                if epoch >= startoff and mse_y < best_loss:
                    best_loss = mse_y
                    self.best_causal_pre = causal_pre
                    self.best_epoch = epoch
                    if self.params['save_model']:
                        ckpt_save_path = self.ckpt_manager.save(epoch)
                        print('Saving checkpoint for epoch {} at {}'.format(epoch, ckpt_save_path))
                if self.params['save_res']:
                    save_data('{}/causal_pre_at_{}.{}'.format(self.save_dir, epoch, save_format), causal_pre)

    @tf.function
    def evaluate(self, data, data_z=None, nb_intervals=200):
        """Evaluate the model on observed data and compute causal estimates.

        For **binary treatment**, the Individual Treatment Effect (ITE) is
        returned.  For **continuous treatment**, the Average Dose-Response
        Function (ADRF) is computed over a grid of treatment values.

        Parameters
        ----------
        data : tuple of np.ndarray
            A triplet ``(data_x, data_y, data_v)``.
        data_z : tf.Tensor or None, optional
            Latent variables.  If ``None``, inferred via the encoder.
        nb_intervals : int, default=200
            Number of grid points for the dose-response curve
            (continuous treatment only).

        Returns
        -------
        causal_estimate : tf.Tensor
            ITE with shape ``(n, 1)`` (binary) or ADRF with shape
            ``(nb_intervals,)`` (continuous).
        mse_x : tf.Tensor
            MSE for the treatment model.
        mse_y : tf.Tensor
            MSE for the outcome model.
        mse_v : tf.Tensor
            MSE for the covariate model.
        data_x_pred : tf.Tensor
            Predicted treatment values.
        data_y_pred : tf.Tensor
            Predicted outcome values.
        data_v_pred : tf.Tensor
            Predicted covariates.
        """
        data_x, data_y, data_v = data
        if data_z is None:
            data_z = self.e_net(data_v)
        data_z0 = data_z[:,:self.params['z_dims'][0]]
        data_z1 = data_z[:,self.params['z_dims'][0]:sum(self.params['z_dims'][:2])]
        data_z2 = data_z[:,sum(self.params['z_dims'][:2]):sum(self.params['z_dims'][:3])]
        data_v_pred = self.g_net(data_z)[:,:self.params['v_dim']]
        data_y_pred = self.f_net(tf.concat([data_z0, data_z1, data_x], axis=-1))[:,:1]
        data_x_pred = self.h_net(tf.concat([data_z0, data_z2], axis=-1))[:,:1]
        if self.params['binary_treatment']:
            data_x_pred = tf.sigmoid(data_x_pred)
        mse_v = tf.reduce_mean((data_v-data_v_pred)**2)
        mse_x = tf.reduce_mean((data_x-data_x_pred)**2)
        mse_y = tf.reduce_mean((data_y-data_y_pred)**2)
        
        if self.params['binary_treatment']:
            # Individual treatment effect (ITE) && average treatment effect (ATE)
            y_pred_pos = self.f_net(tf.concat([data_z0, data_z1, np.ones((len(data_x),1))], axis=-1))[:,:1]
            y_pred_neg = self.f_net(tf.concat([data_z0, data_z1, np.zeros((len(data_x),1))], axis=-1))[:,:1]
            ite_pre = y_pred_pos-y_pred_neg
            return ite_pre, mse_x, mse_y, mse_v, data_x_pred, data_y_pred, data_v_pred
        else:
            # Average dose response function (ADRF)
            # Compute 5% and 95% quantiles from data_x
            x_min = tfp.stats.percentile(data_x, 5.0)
            x_max = tfp.stats.percentile(data_x, 95.0)
            x_values = tf.linspace(x_min, x_max, nb_intervals)
            
            def compute_dose_response(x):
                data_x_tile = tf.fill([tf.shape(data_x)[0], 1], x)
                data_x_tile = tf.cast(data_x_tile, tf.float32)
                y_pred = self.f_net(tf.concat([data_z0, data_z1, data_x_tile], axis=-1))[:, :1]
                return tf.reduce_mean(y_pred)
        
            dose_response = tf.map_fn(compute_dose_response, x_values, fn_output_signature=tf.float32)
            
            return dose_response, mse_x, mse_y, mse_v, data_x_pred, data_y_pred, data_v_pred

    # Predict with MCMC sampling
    def predict(self, data, alpha=0.01, n_mcmc=3000, burn_in=5000, x_values=None, q_sd=1.0, sample_y=True, bs=5000,
                n_theta=1, theta_method='sample', num_leapfrog_steps=10):
        """
        Evaluate the model on the test data and provide both point estimates and posterior intervals for causal effects.
        - For binary treatment, the Individual Treatment Effect (ITE) is estimated.
        - For continuous treatment, the Average Dose Response Function (ADRF) is estimated.

        Parameters
        ----------
        data : list
            Input data containing [data_x, data_y, data_v].
        alpha : float, default=0.01
            Significance level for the posterior interval.
        n_mcmc : int, default=3000
            Number of posterior MCMC samples to draw.
        burn_in : int, default=5000
            Number of burn-in iterations for latent HMC sampling.
        x_values : list of floats or np.ndarray
            Treatment values for dose-response prediction.
            Required for continuous treatment.
        q_sd : float, default=1.0
            HMC step size for latent posterior sampling.
        sample_y : bool, default=True
            Whether to sample from the outcome variance model.
        bs : int, default=5000
            Batch size over test samples during posterior prediction.
        n_theta : int, default=1
            Number of outer theta samples. Each theta sample is paired with an
            inner latent MCMC run and all draws are pooled for intervals.
        theta_method : {'sample', 'mean'}, default='sample'
            Strategy to obtain theta from the BNN posterior.
        num_leapfrog_steps : int, default=10
            Number of HMC leapfrog steps for latent sampling.

        Returns
        -------
        effect : np.ndarray
            Binary treatment: point estimate of ITE with shape ``(n,)``.
            Continuous treatment: ADRF estimate with shape ``(len(x_values),)``.
        pos_int : np.ndarray
            Posterior interval with shape ``(n, 2)`` for ITE, or
            ``(len(x_values), 2)`` for ADRF.
        """
        assert 0 < alpha < 1, "The significance level 'alpha' must be greater than 0 and less than 1."

        if not self.params['binary_treatment']:
            if x_values is None:
                raise ValueError("For continous treatment, 'x_values' must not be None. Provide a list or a single treatment value.")

        if x_values is not None:
            if np.isscalar(x_values):
                x_values = np.array([x_values], dtype=float) 
            else:
                x_values = np.array(x_values, dtype=float)
            x_values_tensor = tf.convert_to_tensor(x_values, dtype=tf.float32)
        else:
            x_values_tensor = None

        data_x = tf.convert_to_tensor(data[0], dtype=tf.float32)
        data_y = tf.convert_to_tensor(data[1], dtype=tf.float32)
        data_v = tf.convert_to_tensor(data[2], dtype=tf.float32)
        n_test = int(data_x.shape[0])
        if bs is None or bs <= 0:
            bs = n_test
        bs = int(min(bs, n_test))

        n_outer = n_theta if self.params['use_bnn'] else 1
        if self.params['use_bnn']:
            theta_weights_list = [self.sample_theta_weights(method=theta_method) for _ in range(n_outer)]
        else:
            theta_weights_list = [(None, None, None)] * n_outer

        def _sample_batch_causal_effects(batch_data, reduce_subject_mean=False):
            """Draw pooled causal-effect samples for one test batch."""
            batch_theta_effects = []
            for g_weights, h_weights, f_weights in theta_weights_list:
                data_posterior_z, acceptance_rate = self.sample_latent_posterior_mcmc(
                    batch_data,
                    burn_in=burn_in,
                    n_keep=n_mcmc,
                    q_sd=q_sd,
                    num_leapfrog_steps=num_leapfrog_steps,
                    g_weights=g_weights,
                    h_weights=h_weights,
                    f_weights=f_weights
                )
                print(f"HMC Acceptance Rate: {float(acceptance_rate.numpy()):.4f}")
                effect = self.infer_from_latent_posterior(
                    data_posterior_z,
                    x_values=None if self.params['binary_treatment'] else x_values_tensor,
                    sample_y=sample_y,
                    f_weights=f_weights
                ).numpy()
                if reduce_subject_mean and not self.params['binary_treatment']:
                    effect = np.mean(effect, axis=-1)
                batch_theta_effects.append(effect)

            # Binary: (n_mcmc, batch_n) per theta
            # Continuous (full): (len(x_values), n_mcmc, batch_n) per theta
            # Continuous (reduced): (len(x_values), n_mcmc) per theta
            concat_axis = 0 if self.params['binary_treatment'] else 1
            return np.concatenate(batch_theta_effects, axis=concat_axis)

        print('MCMC Latent Variable Sampling ...')
        if self.params['binary_treatment']:
            ite = np.zeros((n_test,), dtype=np.float32)
            posterior_interval_upper = np.zeros((n_test,), dtype=np.float32)
            posterior_interval_lower = np.zeros((n_test,), dtype=np.float32)

            for start in range(0, n_test, bs):
                end = min(start + bs, n_test)
                batch_data = (data_x[start:end], data_y[start:end], data_v[start:end])
                batch_effects = _sample_batch_causal_effects(batch_data)

                ite[start:end] = np.mean(batch_effects, axis=0)
                posterior_interval_upper[start:end] = np.quantile(batch_effects, 1 - alpha / 2, axis=0)
                posterior_interval_lower[start:end] = np.quantile(batch_effects, alpha / 2, axis=0)

            ITE = ite
            pos_int = np.stack([posterior_interval_lower, posterior_interval_upper], axis=1)
            return ITE, pos_int
        else:
            # keep only (len(x_values), n_outer * n_mcmc) in memory.
            n_draws = n_outer * n_mcmc
            adrf_draw_sums = np.zeros((len(x_values), n_draws), dtype=np.float32)
            n_seen = 0

            for start in range(0, n_test, bs):
                end = min(start + bs, n_test)
                batch_data = (data_x[start:end], data_y[start:end], data_v[start:end])
                batch_n = end - start
                batch_effects_mean = _sample_batch_causal_effects(
                    batch_data,
                    reduce_subject_mean=True
                )
                adrf_draw_sums += batch_effects_mean * batch_n
                n_seen += batch_n

            causal_effects = adrf_draw_sums / float(n_seen)
            ADRF = np.mean(causal_effects, axis=1)
            posterior_interval_upper = np.quantile(causal_effects, 1-alpha/2, axis=1)
            posterior_interval_lower = np.quantile(causal_effects, alpha/2, axis=1)
            pos_int = np.stack([posterior_interval_lower, posterior_interval_upper], axis=1)
            return ADRF, pos_int

    @tf.function
    def infer_from_latent_posterior(self, data_posterior_z, x_values=None, sample_y=True, f_weights=None, eps=1e-6):
        """Infer causal estimate on the test data and give estimation interval and posterior latent variables. ITE is estimated for binary treatment and ADRF is estimated for continous treatment.
        data_posterior_z: (np.ndarray): Posterior latent variables with shape (n_mcmc, n, p), where p is the dimension of Z.
        x_values: (list of floats or np.ndarray): Number of intervals for the dose response function.
        sample_y: (bool): consider the variance function in outcome generative model.
        return (np.ndarray): 
            ITE with shape (n_mcmc, n) containing all the MCMC samples.
            ADRF with shape (len(x_values), n_mcmc, n) containing all the MCMC samples for each treatment value.
        """

        # Extract the components of Z for X,Y
        data_z0 = data_posterior_z[:,:,:self.params['z_dims'][0]]
        data_z1 = data_posterior_z[:,:,self.params['z_dims'][0]:sum(self.params['z_dims'][:2])]
        data_z2 = data_posterior_z[:,:,sum(self.params['z_dims'][:2]):sum(self.params['z_dims'][:3])]

        if self.params['binary_treatment']:
            
            # Extract mean and sigma^2 of positive samples both with shape (n_keep, n_test)
            y_out_pos_all = tf.map_fn(
                lambda z: (
                    self.f_net(tf.concat([z[:, :self.params['z_dims'][0]],
                                          z[:, self.params['z_dims'][0]:sum(self.params['z_dims'][:2])],
                                          tf.ones([tf.shape(z)[0], 1])], axis=-1))
                    if f_weights is None else
                    self._bnn_forward_with_flat_weights(
                        self.f_net,
                        tf.concat([z[:, :self.params['z_dims'][0]],
                                   z[:, self.params['z_dims'][0]:sum(self.params['z_dims'][:2])],
                                   tf.ones([tf.shape(z)[0], 1])], axis=-1),
                        f_weights
                    )
                ),
                data_posterior_z,
                fn_output_signature=tf.float32
            )
            mu_y_pos_all = y_out_pos_all[:,:,0]
            if 'sigma_y' in self.params:
                sigma_square_y_pos = self.params['sigma_y']**2
            else:
                sigma_square_y_pos = tf.nn.softplus(y_out_pos_all[:,:,1]) + eps
                
            if sample_y:
                y_pred_pos_all = tf.random.normal(
                    shape=tf.shape(mu_y_pos_all), mean=mu_y_pos_all, stddev=tf.sqrt(sigma_square_y_pos)
                )
            else:
                y_pred_pos_all = mu_y_pos_all
            
            # Extract mean and sigma^2 of negative samples both with shape (n_keep, n_test)
            y_out_neg_all = tf.map_fn(
                lambda z: (
                    self.f_net(tf.concat([z[:, :self.params['z_dims'][0]],
                                          z[:, self.params['z_dims'][0]:sum(self.params['z_dims'][:2])],
                                          tf.zeros([tf.shape(z)[0], 1])], axis=-1))
                    if f_weights is None else
                    self._bnn_forward_with_flat_weights(
                        self.f_net,
                        tf.concat([z[:, :self.params['z_dims'][0]],
                                   z[:, self.params['z_dims'][0]:sum(self.params['z_dims'][:2])],
                                   tf.zeros([tf.shape(z)[0], 1])], axis=-1),
                        f_weights
                    )
                ),
                data_posterior_z,
                fn_output_signature=tf.float32
            )
            mu_y_neg_all = y_out_neg_all[:,:,0]
            if 'sigma_y' in self.params:
                sigma_square_y_neg = self.params['sigma_y']**2
            else:
                sigma_square_y_neg = tf.nn.softplus(y_out_neg_all[:,:,1]) + eps
                
            if sample_y:
                y_pred_neg_all = tf.random.normal(
                    shape=tf.shape(mu_y_neg_all), mean=mu_y_neg_all, stddev=tf.sqrt(sigma_square_y_neg)
                )
            else:
                y_pred_neg_all = mu_y_neg_all
                
            ite_pred_all = y_pred_pos_all-y_pred_neg_all
            
            return ite_pred_all
        else:
            
            def compute_dose_response(x):
                data_x = tf.fill([tf.shape(data_posterior_z)[1], 1], x)
                data_x = tf.cast(data_x, tf.float32)
                y_out_all = tf.map_fn(
                    lambda z: (
                        self.f_net(tf.concat([z[:, :self.params['z_dims'][0]],
                                              z[:, self.params['z_dims'][0]:sum(self.params['z_dims'][:2])],
                                              data_x], axis=-1))
                        if f_weights is None else
                        self._bnn_forward_with_flat_weights(
                            self.f_net,
                            tf.concat([z[:, :self.params['z_dims'][0]],
                                       z[:, self.params['z_dims'][0]:sum(self.params['z_dims'][:2])],
                                       data_x], axis=-1),
                            f_weights
                        )
                    ),
                    data_posterior_z,
                    fn_output_signature=tf.float32
                )
                mu_y_all = y_out_all[:,:,0]
                if 'sigma_y' in self.params:
                    sigma_square_y = self.params['sigma_y']**2
                else:
                    sigma_square_y = tf.nn.softplus(y_out_all[:,:,1]) + eps
                
                if sample_y:
                    y_pred_all = tf.random.normal(
                        shape=tf.shape(mu_y_all), mean=mu_y_all, stddev=tf.sqrt(sigma_square_y)
                    )
                else:
                    y_pred_all = mu_y_all
                    
                return y_pred_all
            
            dose_response = tf.map_fn(compute_dose_response, x_values, fn_output_signature=tf.float32)
            
            return dose_response

    @tf.function
    def get_log_posterior(self, data_x, data_y, data_v, data_z, g_weights=None, h_weights=None, f_weights=None, eps=1e-6):
        """
        Calculate log posterior.
        data_x: (np.ndarray): Input data with shape (n, 1), where p is the dimension of X.
        data_y: (np.ndarray): Input data with shape (n, 1), where q is the dimension of Y.
        data_v: (np.ndarray): Input data with shape (n, p), where r is the dimension of V.
        data_z: (np.ndarray): Input data with shape (n, q), where q is the dimension of Z.
        return (np.ndarray): Log posterior with shape (n, ).
        """
        data_z0 = data_z[:,:self.params['z_dims'][0]]
        data_z1 = data_z[:,self.params['z_dims'][0]:sum(self.params['z_dims'][:2])]
        data_z2 = data_z[:,sum(self.params['z_dims'][:2]):sum(self.params['z_dims'][:3])]

        if g_weights is None:
            g_out = self.g_net(data_z)
        else:
            g_out = self._bnn_forward_with_flat_weights(self.g_net, data_z, g_weights)
        mu_v = g_out[:,:self.params['v_dim']]
        if 'sigma_v' in self.params:
            sigma_square_v = self.params['sigma_v']**2
        else:
            sigma_square_v = tf.nn.softplus(g_out[:,-1]) + eps

        h_in = tf.concat([data_z0, data_z2], axis=-1)
        if h_weights is None:
            h_out = self.h_net(h_in)
        else:
            h_out = self._bnn_forward_with_flat_weights(self.h_net, h_in, h_weights)
        mu_x = h_out[:,:1]
        if 'sigma_x' in self.params:
            sigma_square_x = self.params['sigma_x']**2
        else:
            sigma_square_x = tf.nn.softplus(h_out[:,-1]) + eps

        f_in = tf.concat([data_z0, data_z1, data_x], axis=-1)
        if f_weights is None:
            f_out = self.f_net(f_in)
        else:
            f_out = self._bnn_forward_with_flat_weights(self.f_net, f_in, f_weights)
        mu_y = f_out[:,:1]
        if 'sigma_y' in self.params:
            sigma_square_y = self.params['sigma_y']**2
        else:
            sigma_square_y = tf.nn.softplus(f_out[:,-1]) + eps

        loss_pv_z = tf.reduce_sum((data_v - mu_v)**2, axis=1)/(2*sigma_square_v) + \
                self.params['v_dim'] * tf.math.log(sigma_square_v)/2
        
        if self.params['binary_treatment']:
            loss_px_z = tf.squeeze(tf.nn.sigmoid_cross_entropy_with_logits(labels=data_x,logits=mu_x))
        else:
            loss_px_z = tf.reduce_sum((data_x - mu_x)**2, axis=1)/(2*sigma_square_x) + \
                    tf.math.log(sigma_square_x)/2

        loss_py_zx = tf.reduce_sum((data_y - mu_y)**2, axis=1)/(2*sigma_square_y) + \
                tf.math.log(sigma_square_y)/2

        loss_prior_z =  tf.reduce_sum(data_z**2, axis=1)/2

        loss_postrior_z = loss_pv_z + loss_px_z + loss_py_zx + loss_prior_z

        log_posterior = -loss_postrior_z
        return log_posterior

    @tf.function(reduce_retracing=True)
    def sample_latent_posterior_mcmc(self, data, initial_q_sd = 1.0, q_sd = None, burn_in = 5000, n_keep = 3000,
                                    num_leapfrog_steps=10, seed=None, g_weights=None, h_weights=None, f_weights=None):
        """
        Samples from P(Z|X,Y,V) using TFP Hamiltonian Monte Carlo (HMC).

        The method name is kept for backward compatibility with existing calls.

        Args:
            data (tuple): Tuple containing data_x, data_y, data_v.
            q_sd (float or None): HMC step size. If None or <= 0, falls back to
                `initial_q_sd`.
            initial_q_sd (float): Default HMC step size when q_sd is not set.
            burn_in (int): Number of burn-in iterations.
            n_keep (int): Number of samples retained after burn-in.
            num_leapfrog_steps (int): Number of leapfrog steps in HMC.
            seed (int or None): Random seed for sampler initialization.
            g_weights, h_weights, f_weights : tf.Tensor or None
                Optional flattened deterministic weights for conditional
                posterior sampling p(z|x,y,v,theta).

        Returns:
            tuple[tf.Tensor, tf.Tensor]:
                - Posterior samples with shape (n_keep, n, q), where q is
                  the dimension of Z.
                - Scalar acceptance rate averaged across sampling steps.
        """
        data_x, data_y, data_v = data
        data_x = tf.convert_to_tensor(data_x, dtype=tf.float32)
        data_y = tf.convert_to_tensor(data_y, dtype=tf.float32)
        data_v = tf.convert_to_tensor(data_v, dtype=tf.float32)
        burn_in = tf.convert_to_tensor(int(burn_in), dtype=tf.int32)
        n_keep = tf.convert_to_tensor(int(n_keep), dtype=tf.int32)
        num_leapfrog_steps = tf.convert_to_tensor(int(num_leapfrog_steps), dtype=tf.int32)

        base_step_size = tf.cast(initial_q_sd, tf.float32)
        proposed_step_size = base_step_size if q_sd is None else tf.cast(q_sd, tf.float32)
        step_size = tf.where(proposed_step_size > 0.0, proposed_step_size, base_step_size)

        if g_weights is not None:
            g_weights = tf.cast(g_weights, tf.float32)
        if h_weights is not None:
            h_weights = tf.cast(h_weights, tf.float32)
        if f_weights is not None:
            f_weights = tf.cast(f_weights, tf.float32)

        # Initialize chains from encoder output instead of random Gaussian noise.
        initial_state = tf.cast(self.e_net(data_v), tf.float32)
        if seed is not None:
            tf.random.set_seed(seed)

        def target_log_prob_fn(z):
            return self.get_log_posterior(data_x, data_y, data_v, z, g_weights, h_weights, f_weights)

        hmc_kernel = tfm.HamiltonianMonteCarlo(
            target_log_prob_fn=target_log_prob_fn,
            step_size=step_size,
            num_leapfrog_steps=num_leapfrog_steps
        )

        adaptive_kernel = tfm.SimpleStepSizeAdaptation(
            inner_kernel=hmc_kernel,
            num_adaptation_steps=tf.cast(tf.cast(burn_in, tf.float32) * 0.8, tf.int32),
            target_accept_prob=0.75
        )

        samples, acceptance_trace = tfm.sample_chain(
            num_results=n_keep,
            num_burnin_steps=burn_in,
            current_state=initial_state,
            kernel=adaptive_kernel,
            trace_fn=lambda _, pkr: tf.reduce_mean(
                tf.cast(pkr.inner_results.is_accepted, tf.float32)
            )
        )

        acceptance_rate = tf.reduce_mean(acceptance_trace)
        return samples, acceptance_rate
