text: from __future__ import annotations
  
  import math
  from abc import ABC, abstractmethod
  from typing import Any, Callable, Dict, Optional, Tuple, Type, Union
  import torch
  from botorch.acquisition.acquisition import AcquisitionFunction
  from botorch.acquisition.analytic import AnalyticAcquisitionFunction
  from botorch.acquisition.objective import IdentityMCObjective, MCAcquisitionObjective, ScalarizedObjective
  from botorch.acquisition.utils import prune_inferior_points
  from botorch.exceptions.errors import UnsupportedError
  from botorch.models.model import Model
  from botorch.sampling.samplers import MCSampler, SobolQMCNormalSampler
  from botorch.utils.transforms import (
      concatenate_pending_points,
      match_batch_shape,
      t_batch_mode_transform,
  )
  from torch import Tensor
  from botorch.acquisition.monte_carlo import MCAcquisitionFunction
  from torch.distributions import Normal
  from scipy.stats import norm
  
  class ExpectedQuantileImprovement(AnalyticAcquisitionFunction):
  
  
      def __init__(
          self,
          model: Model,
          train: Union[float, Tensor],
          objective: Optional[ScalarizedObjective] = None,
          maximize: bool = True,
          noise_level:float = 0.01,
          beta:float = 0.5,
          budget:float = 1
      ) -> None:
  
          super().__init__(model=model, objective=objective)
          self.maximize = maximize
          if not torch.is_tensor(train):
              train = torch.tensor(train)
          self.register_buffer("train_x", train)
          # self.noise_level = torch.tensor([noise_level])
          self.beta = torch.tensor([norm.ppf(beta)])
          self.budget = budget
  
      @t_batch_mode_transform(expected_q=1, assert_output_shape=False)
      def forward(self, X: Tensor) -> Tensor:
          # calculate qmin
          self.train_x = self.train_x.to(X)
          self.beta = self.beta.to(X)
          beta = self.beta
          posterior = self._get_posterior(X=self.train_x)
          mean = posterior.mean
          sigma = posterior.variance.clamp_min(1e-9).sqrt()
          
          qmin = mean + sigma * beta
          qmin = torch.max(qmin).detach()
          # mq, sq and q(x)
          posterior = self._get_posterior(X=X)
          mean = posterior.mean
          view_shape = mean.shape[:-2] if mean.dim() >= X.dim() else X.shape[:-2]
          mean = mean.view(view_shape)
          # noise_level = self.model.noise_model.posterior(X=X).mean
          # noise_level = noise_level.squeeze() / self.budget
          noise_level = torch.tensor([0.05]).to(X)
          # print(noise_level)
  
          sigma = posterior.variance.clamp_min(1e-9).sqrt().view(view_shape)
          sq = sigma.square() / (noise_level + sigma.square())
          # print('mean', mean, 'noise_level', noise_level, 'sigma', sigma)
          mq = mean + beta * ((noise_level * sigma.square()) / (noise_level + sigma.square())).sqrt()
          # print('mq', mq,'qmin', qmin,'sq', sq)
          t = (mq - qmin)/sq
          # print('t', t)
          normal = Normal(torch.zeros_like(mean), torch.ones_like(mean))
          eqi = t * sq * normal.cdf(t) + sq * torch.exp(normal.log_prob(t))
          # deal with batch evaluation and broadcasting
          # print(eqi)
          return eqi
  
  
  