"""MBT: Multimodal Bottleneck Transformers."""

import functools
from typing import Any, Callable, Dict, Iterable, Optional, Tuple

from absl import logging
import flax.linen as nn
from flax.linen.linear import default_kernel_init
from immutabledict import immutabledict
import jax
import jax.numpy as jnp
import ml_collections
import numpy as np
from scenic.model_lib.base_models import base_model
from scenic.model_lib.base_models import classification_model
from scenic.model_lib.base_models import model_utils as base_model_utils
from scenic.model_lib.base_models.classification_model import ClassificationModel
from scenic.model_lib.layers import attention_layers
from scenic.model_lib.layers import nn_layers
from scenic.projects.baselines import vit
from scenic.projects.mbt import model_utils

Initializer = Callable[[jnp.ndarray, Iterable[int], jnp.dtype], jnp.ndarray]

_MBT_CLASSIFICATION_METRICS = immutabledict({
    'accuracy': (base_model_utils.weighted_correctly_classified,
                 base_model_utils.num_examples),
    'accuracy_top_5': (functools.partial(
        base_model_utils.weighted_topk_correctly_classified,
        k=5), base_model_utils.num_examples),
    'loss': (base_model_utils.weighted_unnormalized_softmax_cross_entropy,
             base_model_utils.num_examples)
})

_MODALITIES = ['rgb', 'spectrogram']

def _reshape_to_time_space(x, temporal_dims):
  if x.ndim == 3:
    b, thw, d = x.shape
    assert thw % temporal_dims == 0
    hw = thw // temporal_dims
    x = jnp.reshape(x, [b, temporal_dims, hw, d])
  assert x.ndim == 4
  return x

def embed_2d_patch(x, patches, embedding_dim, name='embedding'):
  """Embedding input patches with 2D conv."""

  assert patches.get('size') is not None, ('patches.size is now the only way'
                                           'to define the patches')
  assert embedding_dim, 'embedding_dim must be specified'
  fh = patches.size[0]
  fw = patches.size[1]
  x = nn.Conv(
      embedding_dim, (fh, fw),
      strides=(fh, fw),
      padding='VALID',
      name=name)(x)

  return x

def embed_3d_patch(x,
                   patches,
                   embedding_dim,
                   kernel_init_method,
                   name='embedding'):
  """Embed 3D input patches into tokens."""

  assert patches.get('size') is not None, 'patches.size must be defined'
  assert len(patches.size) == 3, 'patches.size must have 3 elements'
  assert embedding_dim, 'embedding_dim must be specified'

  fh, fw, ft = patches.size

  if kernel_init_method == 'central_frame_initializer':
    kernel_initializer = model_utils.central_frame_initializer()
    logging.info('Using central frame initializer for input embedding')
  elif kernel_init_method == 'average_frame_initializer':
    kernel_initializer = model_utils.average_frame_initializer()
    logging.info('Using average frame initializer for input embedding')
  else:
    kernel_initializer = default_kernel_init
    logging.info('Using default initializer for input embedding')

  x = nn.Conv(
      embedding_dim, (ft, fh, fw),
      strides=(ft, fh, fw),
      padding='VALID',
      name=name,
      kernel_init=kernel_initializer)(x)

  return x

def temporal_encode(x,
                    modality,
                    temporal_encoding_config,
                    patches,
                    hidden_size,
                    return_1d=True):
  """Encode video for feeding into ViT."""
  if modality == 'spectrogram':
    # Spectrogram is treated as a big num_time_bins by num_mel_bins image.
    x = embed_2d_patch(x, patches, hidden_size, 'embedding_spectrogram')
    temporal_dims = 1
    if return_1d:
      n, h, w, c = x.shape
      x = jnp.reshape(x, [n, h * w, c])
  elif modality == 'rgb':
    if temporal_encoding_config.method == 'temporal_sampling':
      n, num_frames, in_h, in_w, c = x.shape
      n_sampled_frames = temporal_encoding_config.n_sampled_frames
      if n_sampled_frames < num_frames:
        t_start_idx = num_frames / (n_sampled_frames + 1)
        t_step = t_start_idx
      else:
        t_start_idx = 0
        t_step = 1
      t_end_idx = num_frames
      temporal_indices = jnp.arange(t_start_idx, t_end_idx, t_step)
      temporal_indices = jnp.round(temporal_indices).astype(jnp.int32)
      temporal_indices = jnp.minimum(temporal_indices, num_frames - 1)

      x = x[:, temporal_indices]  # [n, t_s, in_h, in_w, c]
      t_s = x.shape[1]
      x = jnp.reshape(x, [n, t_s * in_h, in_w, c])
      x = embed_2d_patch(x, patches, hidden_size)
      temporal_dims = t_s
      if return_1d:
        n, th, w, c = x.shape
        x = jnp.reshape(x, [n, th * w, c])
      else:
        n, th, w, c = x.shape
        x = jnp.reshape(x, [n, t_s, -1, w, c])
    if temporal_encoding_config.method == '3d_conv':
      kernel_init_method = temporal_encoding_config.get('kernel_init_method',
                                                        None)

      x = embed_3d_patch(x, patches, hidden_size, kernel_init_method)
      temporal_dims = x.shape[1]
      if return_1d:
        n, t, h, w, c = x.shape
        x = jnp.reshape(x, [n, t * h * w, c])
  else:
    raise AssertionError('Unknown temporal encoding method.')
  return x, temporal_dims

def add_positional_embed(x, feat_name):
  """Adds positional embedding."""
  assert x.ndim == 3  # (batch, len, emb)
  x = vit.AddPositionEmbs(
      posemb_init=nn.initializers.normal(stddev=0.02),  # from BERT.
      name=feat_name)(x)
  return x

class EncoderBlock(nn.Module):
  """Transformer encoder block.

  Attributes:
    mlp_dim: Dimension of the mlp on top of attention block.
    num_heads: Number of heads.
    dtype: The dtype of the computation (default: float32).
    dropout_rate: Dropout rate.
    attention_dropout_rate: Dropout for attention heads.
    attention_kernel_initializer: Initializer to use for attention
      layers.
    droplayer_p: Probability of dropping a layer.

  Returns:
    Output after transformer encoder block.
  """
  mlp_dim: int
  num_heads: int
  dtype: Any = jnp.float32
  dropout_rate: float = 0.1
  attention_dropout_rate: float = 0.1
  attention_kernel_initializer: Initializer = nn.initializers.xavier_uniform()
  droplayer_p: float = 0.0

  def get_drop_pattern(self, x, deterministic):
    if not deterministic and self.droplayer_p:
      shape = (x.shape[0],) + (1,) * (x.ndim - 1)
      return jax.random.bernoulli(
          self.make_rng('dropout'), self.droplayer_p, shape).astype('float32')
    else:
      return 0.0

  @nn.compact
  def __call__(self, inputs: jnp.ndarray, deterministic: bool) -> jnp.ndarray:
    """Applies Encoder1DBlock module."""

    # Attention block.
    x = nn.LayerNorm(dtype=self.dtype)(inputs)
    x = nn.MultiHeadDotProductAttention(
        num_heads=self.num_heads,
        kernel_init=self.attention_kernel_initializer,
        broadcast_dropout=False,
        dropout_rate=self.attention_dropout_rate,
        dtype=self.dtype)(x, x, deterministic=deterministic)
    x = nn.Dropout(rate=self.dropout_rate)(x, deterministic)

    drop_pattern = self.get_drop_pattern(x, deterministic)
    x = x * (1.0 - drop_pattern) + inputs

    # MLP block.
    y = nn.LayerNorm(dtype=self.dtype)(x)
    y = attention_layers.MlpBlock(
        mlp_dim=self.mlp_dim,
        dtype=self.dtype,
        dropout_rate=self.dropout_rate,
        activation_fn=nn.gelu,
        kernel_init=nn.initializers.xavier_uniform(),
        bias_init=nn.initializers.normal(stddev=1e-6))(
            y, deterministic=deterministic)

    drop_pattern = self.get_drop_pattern(x, deterministic)
    return y * (1.0 - drop_pattern) + x

class Encoder(nn.Module):
  """Transformer Encoder.

  Attributes:
    mlp_dim: Dimension of the mlp on top of attention block.
    num_layers: Number of layers.
    num_heads: Number of attention heads.
    dropout_rate: Dropout rate.
    attention_dropout_rate: Dropout for attention heads.
    stochastic_droplayer_rate: Probability of dropping a layer linearly
      grows from 0 to the provided value. Our implementation of stochastic
      depth follows timm library, which does per-example layer dropping and
      uses independent dropping patterns for each skip-connection.
    modality_fusion: Tuple with modalities to combine.
    fusion_layer: Which layer to fuse modalities. fusion_layer == 0 provides
      early fusion.
    use_bottleneck: If True, adds self-attention bottleneck.
    test_with_bottlenecks: Whether to use bottlenecks at test time.
    share_encoder: If True, different modalities share the same encoder weights
      for the layers before fusion.
    dtype: The dtype of the computation (default: float32).
  """
  mlp_dim: int
  num_layers: int
  num_heads: int
  dropout_rate: float = 0.1
  attention_dropout_rate: float = 0.1
  stochastic_droplayer_rate: float = 0.0
  modality_fusion: Tuple[str] = ('spectrogram',)
  fusion_layer: int = 0
  use_bottleneck: bool = False
  test_with_bottlenecks: bool = True
  share_encoder: bool = False
  dtype: Any = jnp.float32

  @nn.compact
  def __call__(self, x: Dict[str, Any],
               bottleneck: jnp.ndarray, *,
               train: bool):
    """Applies Transformer model on the inputs."""

    def get_encoder_block(encoder_block, droplayer_p, name):
      """Returns the encoder block for a single layer."""
      dtype = jax.dtypes.canonicalize_dtype(self.dtype)
      return encoder_block(
          mlp_dim=self.mlp_dim,
          num_heads=self.num_heads,
          dropout_rate=self.dropout_rate,
          attention_dropout_rate=self.attention_dropout_rate,
          droplayer_p=droplayer_p,
          name=name,
          dtype=dtype)

    def get_context(target_modality, modality_fusion, x):
      """Returns list of context modalities."""
      context = []
      for modality in _MODALITIES:
        if modality != target_modality and modality in modality_fusion:
          context.append(x[modality])
      return context

    def combine_context(x, other_modalities):
      """Combine x with a list of other modalities."""
      num_tokens = x.shape[1]
      # Append x to the end of the list
      other_modalities.append(x)
      x_combined = jnp.concatenate(other_modalities, axis=1)
      return x_combined, num_tokens

    assert self.modality_fusion

    # Add positional embeddings
    for modality in self.modality_fusion:
      if modality == 'rgb':
        name = ''
      else:
        name = '_' + modality
      x[modality] = add_positional_embed(x[modality], 'posembed_input' + name)

    use_bottlenecks = train or self.test_with_bottlenecks
    x_combined = None
    # Input Encoder
    for lyr in range(self.num_layers):
      droplayer_p = (
          lyr / max(self.num_layers - 1, 1)) * self.stochastic_droplayer_rate
      encoders = {}
      encoders['rgb'] = get_encoder_block(EncoderBlock, droplayer_p,
                                          f'encoderblock_{lyr}')

      for modality in self.modality_fusion:
        if modality != 'rgb':
          if self.share_encoder:
            encoders[modality] = encoders['rgb']
          else:
            encoders[modality] = get_encoder_block(
                EncoderBlock, droplayer_p,
                f'encoderblock_{lyr}_' + modality)

      if (lyr < self.fusion_layer or len(self.modality_fusion) == 1 or
          (self.use_bottleneck and not use_bottlenecks)):
        for modality in self.modality_fusion:
          x[modality] = encoders[modality](x[modality], deterministic=not train)
      else:
        if self.use_bottleneck:
          bottle = []
          for modality in self.modality_fusion:
            t_mod = x[modality].shape[1]
            in_mod = jnp.concatenate([x[modality], bottleneck], axis=1)
            out_mod = encoders[modality](in_mod, deterministic=not train)
            x[modality] = out_mod[:, :t_mod]
            bottle.append(out_mod[:, t_mod:])
          bottleneck = jnp.mean(jnp.stack(bottle, axis=-1), axis=-1)
        else:
          if not self.share_encoder and len(self.modality_fusion) > 1:
            x_new = {}
            for modality in self.modality_fusion:
              other_modalities = get_context(modality, self.modality_fusion, x)
              combined_mods, t = combine_context(x[modality], other_modalities)
              combined_mods = encoders[modality](
                  combined_mods, deterministic=not train)
              x_new[modality] = combined_mods[:, -t:]
            x = x_new

          elif self.share_encoder and len(self.modality_fusion) > 1:
            if x_combined is None:
              x_combined = []
              for modality in self.modality_fusion:
                x_combined.append(x[modality])
              x_combined = jnp.concatenate(x_combined, axis=1)
            x_combined = encoders['rgb'](x_combined, deterministic=not train)
    if x_combined is not None:
      x_out = x_combined
    else:
      x_out = []
      for modality in self.modality_fusion:
        x_out.append(x[modality])
      x_out = jnp.concatenate(x_out, axis=1)
    encoded = nn.LayerNorm(name='encoder_norm')(x_out)

    return encoded

class MBT(nn.Module):
  """Audio-Visual Fusion Transformer model for Video.

  Attributes:
    mlp_dim: Dimension of the mlp on top of attention block.
    num_layers: Number of layers.
    num_heads: Number of self-attention heads.
    num_classes: Number of output classes.
    patches: Configuration of the patches extracted in the stem of the model.
    hidden_size: Size of the hidden state of the output of model's stem.
      if None, we skip the extra projection + tanh activation at the end.
    temporal_encoding_config: ConfigDict which defines the type of input
      encoding when tokenising the video.
    attention_config: ConfigDict which defines the type of spatio-temporal
      attention applied in the model.
    representation_size: Size of the representation layer in the model's head.
    dropout_rate: Dropout rate.
    attention_dropout_rate: Dropout for attention heads.
    stochastic_droplayer_rate: Probability of dropping a layer. Linearly
      increases from 0 to the provided value..
    classifier: type of the classifier layer. Options are 'gap', 'gmp', 'gsp',
      'token'.
    modality_fusion: Tuple with modalities to combine.
    fusion_layer: Which layer to fuse modalities.
    return_prelogits: If true, return the final representation of the network
      before the classification head. Useful when using features for a
      downstream task.
    return_preclassifier: If true, return a dict of all token embeddings.
      Useful when using token embeddings for a downstream task.
    use_bottleneck: If True, adds self-attention bottleneck.
    n_bottlenecks: Number of bottleneck tokens.
    test_with_bottlenecks: Whether to use bottlenecks at test time.
    share_encoder: If True, different modalities share the same encoder weights
      for the layers before fusion.
    dtype: JAX data type for activations.
  """

  mlp_dim: int
  num_layers: int
  num_heads: int
  num_classes: int
  patches: ml_collections.ConfigDict
  hidden_size: int
  temporal_encoding_config: ml_collections.ConfigDict
  attention_config: ml_collections.ConfigDict
  representation_size: Optional[int] = None
  dropout_rate: float = 0.1
  attention_dropout_rate: float = 0.1
  stochastic_droplayer_rate: float = 0.
  classifier: str = 'gap'
  modality_fusion: Tuple[str] = ('spectrogram',)
  fusion_layer: int = 0
  return_prelogits: bool = False
  return_preclassifier: bool = False
  use_bottleneck: bool = False
  n_bottlenecks: int = 4
  test_with_bottlenecks: bool = True
  share_encoder: bool = False
  dtype: Any = jnp.float32

  @nn.compact
  def __call__(self,
               x,
               *,
               train: bool,
               debug: bool = False):
    assert self.fusion_layer <= self.num_layers and self.fusion_layer >= 0
    assert self.classifier in ['token', '0', 'gap', 'gmp', 'gsp']

    temporal_dims = {}
    for modality in self.modality_fusion:
      x[modality], _ = temporal_encode(
          x[modality], modality, self.temporal_encoding_config, self.patches,
          self.hidden_size)
      # If we want to add a class token, add it here.
      if self.classifier in ['token']:
        if modality == 'rgb' or len(self.modality_fusion) == 1:
          name = ''
        else:
          name = modality
        n, temporal_dims[modality], c = x[modality].shape
        cls = self.param('cls'+name, nn.initializers.zeros, (1, 1, c),
                         x[modality].dtype)
        cls = jnp.tile(cls, [n, 1, 1])
        x[modality] = jnp.concatenate([cls, x[modality]], axis=1)
        bottleneck_dtype = x[modality].dtype

    bottleneck = None
    if self.use_bottleneck:
      n_bottlenecks = self.n_bottlenecks
      if self.classifier in ['token']:
        n_bottlenecks += 1
      bottleneck = self.param('bottleneck',
                              nn.initializers.normal(stddev=0.02),  # From BERT.
                              (1, n_bottlenecks, c), bottleneck_dtype)
      bottleneck = jnp.tile(bottleneck, [n, 1, 1])

    x = Encoder(
        modality_fusion=self.modality_fusion,
        fusion_layer=self.fusion_layer,
        mlp_dim=self.mlp_dim,
        num_layers=self.num_layers,
        num_heads=self.num_heads,
        dropout_rate=self.dropout_rate,
        attention_dropout_rate=self.attention_dropout_rate,
        stochastic_droplayer_rate=self.stochastic_droplayer_rate,
        use_bottleneck=self.use_bottleneck,
        test_with_bottlenecks=self.test_with_bottlenecks,
        share_encoder=self.share_encoder,
        dtype=self.dtype,
        name='Transformer')(x, bottleneck, train=train)

    if self.return_preclassifier:
      return x

    if self.classifier in ['token', '0']:
      # Obtaining the CLS tokens for each modality.
      x_out = {}
      counter = 0
      for modality in self.modality_fusion:
        x_out[modality] = x[:, counter]
        counter += temporal_dims[modality] + 1
    elif self.classifier in ('gap', 'gmp', 'gsp'):
      fn = {'gap': jnp.mean, 'gmp': jnp.max, 'gsp': jnp.sum}[self.classifier]
      x_out = fn(x, axis=list(range(1, x.ndim - 1)))

    if self.representation_size is not None:
      pre_logits_fc = nn.Dense(self.representation_size, name='pre_logits')
      if isinstance(x_out, dict):
        for modality in x_out:
          x_out[modality] = pre_logits_fc(x_out[modality])
          x_out[modality] = nn.tanh(x_out[modality])
      else:
        x_out = nn.Dense(self.representation_size, name='pre_logits')(x_out)
        x_out = nn.tanh(x_out)
    else:
      if not isinstance(x_out, dict):
        x_out = nn_layers.IdentityLayer(name='pre_logits')(x_out)

    if self.return_prelogits:
      return x_out
    if isinstance(x_out, dict):
      output_projection_fc = nn.Dense(
          self.num_classes,
          kernel_init=nn.initializers.zeros,
          name='output_projection')
      x_pool = 0
      for modality in x_out:
        x_out[modality] = output_projection_fc(x_out[modality])
        x_pool += x_out[modality]
      x_pool /= len(x_out)
      if not train:
        return x_pool
    else:
      x_out = nn.Dense(
          self.num_classes,
          kernel_init=nn.initializers.zeros,
          name='output_projection')(
              x_out)
    return x_out

class MBTMultilabelClassificationModel(vit.ViTMultiLabelClassificationModel):
  """Video Transformer model for multi-class classification."""

  def build_flax_model(self) -> nn.Module:
    model_dtype = getattr(jnp, self.config.get('model_dtype_str', 'float32'))
    return MBT(
        num_classes=self.dataset_meta_data['num_classes'],
        modality_fusion=self.config.model.modality_fusion,
        fusion_layer=self.config.model.fusion_layer,
        use_bottleneck=self.config.model.get('use_bottleneck', False),
        test_with_bottlenecks=self.config.model.get(
            'test_with_bottlenecks', True),
        n_bottlenecks=self.config.model.get('n_bottlenecks', 4),
        share_encoder=self.config.model.get('share_encoder', False),
        mlp_dim=self.config.model.mlp_dim,
        num_layers=self.config.model.num_layers,
        num_heads=self.config.model.num_heads,
        representation_size=self.config.model.representation_size,
        patches=self.config.model.patches,
        hidden_size=self.config.model.hidden_size,
        temporal_encoding_config=self.config.model.temporal_encoding_config,
        attention_config=self.config.model.attention_config,
        classifier=self.config.model.classifier,
        dropout_rate=self.config.model.get('dropout_rate', 0.1),
        attention_dropout_rate=self.config.model.get('attention_dropout_rate',
                                                     0.1),
        stochastic_droplayer_rate=self.config.model.get(
            'stochastic_droplayer_rate', 0),
        return_prelogits=self.config.model.get('return_prelogits', False),
        dtype=model_dtype)

  def init_from_train_state(self,
                            train_state: Any,
                            restored_train_state: Any,
                            restored_model_cfg: ml_collections.ConfigDict,
                            restore_output_proj: bool = False) -> Any:
    """Updates the train_state with data from restored_train_state."""
    return model_utils.initialise_from_train_state(self.config, train_state,
                                                   restored_train_state,
                                                   restored_model_cfg,
                                                   restore_output_proj)

  def loss_function(self,
                    logits: jnp.array,
                    batch: base_model.Batch,
                    model_params: Optional[jnp.array] = None) -> float:
    """Returns sigmoid cross entropy loss with an L2 penalty on the weights.

    Args:
      logits: Output of model in shape [batch, length, num_classes]. Optionally,
        this can also be a dictionary with logits for individual modalities.
      batch: Batch of data that has 'label' and optionally 'batch_mask'.
      model_params: Parameters of the model, for optionally applying
        regularization.

    Returns:
      Total loss.
    """
    weights = batch.get('batch_mask')
    labels = batch['label']

    assert self.dataset_meta_data.get('target_is_onehot', False)

    label_weights = self.dataset_meta_data.get('class_weights', None)

    if isinstance(logits, dict):
      sig_ce_loss = []
      for modality in logits:
        sig_ce_loss.append(base_model_utils.weighted_sigmoid_cross_entropy(
            logits[modality],
            labels[modality],
            weights,
            label_weights=label_weights,
            label_smoothing=self.config.get('label_smoothing')))
      sig_ce_loss = jnp.mean(jnp.array(sig_ce_loss))
    else:
      if isinstance(labels, dict):
        assert 'all' in labels, 'mixmod must be turned off.'
        labels = labels['all']
      sig_ce_loss = base_model_utils.weighted_sigmoid_cross_entropy(
          logits,
          labels,
          weights,
          label_weights=label_weights,
          label_smoothing=self.config.get('label_smoothing'))
    if self.config.get('l2_decay_factor') is None:
      total_loss = sig_ce_loss
    else:
      l2_loss = base_model_utils.l2_regularization(model_params)
      total_loss = sig_ce_loss + 0.5 * self.config.l2_decay_factor * l2_loss
    return total_loss

class MBTClassificationModel(ClassificationModel):
  """Audio Video Transformer model for n-way classification."""

  def build_flax_model(self) -> nn.Module:
    assert (self.config.model.attention_config.get('type', 'spacetime') !=
            'factorized_encoder'), (
                'Other attention types not supported.')
    model_dtype = getattr(jnp, self.config.get('model_dtype_str', 'float32'))
    return MBT(
        num_classes=self.dataset_meta_data['num_classes'],
        modality_fusion=self.config.model.modality_fusion,
        fusion_layer=self.config.model.fusion_layer,
        use_bottleneck=self.config.model.get('use_bottleneck', False),
        test_with_bottlenecks=self.config.model.get(
            'test_with_bottlenecks', True),
        n_bottlenecks=self.config.model.get('n_bottlenecks', 4),
        share_encoder=self.config.model.get('share_encoder', False),
        mlp_dim=self.config.model.mlp_dim,
        num_layers=self.config.model.num_layers,
        num_heads=self.config.model.num_heads,
        representation_size=self.config.model.representation_size,
        patches=self.config.model.patches,
        hidden_size=self.config.model.hidden_size,
        temporal_encoding_config=self.config.model.temporal_encoding_config,
        attention_config=self.config.model.attention_config,
        classifier=self.config.model.classifier,
        dropout_rate=self.config.model.get('dropout_rate', 0.1),
        attention_dropout_rate=self.config.model.get('attention_dropout_rate',
                                                     0.1),
        stochastic_droplayer_rate=self.config.model.get(
            'stochastic_droplayer_rate', 0),
        return_prelogits=self.config.model.get('return_prelogits', False),
        dtype=model_dtype)

  def loss_function(self,
                    logits: jnp.ndarray,
                    batch: base_model.Batch,
                    model_params: Optional[jnp.ndarray] = None) -> float:
    """Returns softmax cross entropy loss with an L2 penalty on the weights.

    Args:
      logits: Output of model in shape [batch, length, num_classes]. Optionally,
        this can also be a dictionary with logits for individual modalities.
      batch: Batch of data that has 'label' and optionally 'batch_mask'.
      model_params: Parameters of the model, for optionally applying
        regularization.

    Returns:
      Total loss.
    """
    weights = batch.get('batch_mask')
    labels = batch['label']

    assert self.dataset_meta_data.get('target_is_onehot', False)

    if isinstance(logits, dict):
      sof_ce_loss = []
      for modality in logits:
        sof_ce_loss.append(base_model_utils.weighted_softmax_cross_entropy(
            logits[modality],
            labels[modality],
            weights,
            label_smoothing=self.config.get('label_smoothing')))
      sof_ce_loss = jnp.mean(jnp.array(sof_ce_loss))
    else:
      sof_ce_loss = base_model_utils.weighted_softmax_cross_entropy(
          logits,
          labels,
          weights,
          label_smoothing=self.config.get('label_smoothing'))
    if self.config.get('l2_decay_factor') is None:
      total_loss = sof_ce_loss
    else:
      l2_loss = base_model_utils.l2_regularization(model_params)
      total_loss = sof_ce_loss + 0.5 * self.config.l2_decay_factor * l2_loss
    return total_loss

  def get_metrics_fn(self, split: Optional[str] = None) -> base_model.MetricFn:
    """Returns a callable metric function for the model.

    Args:
      split: The split for which we calculate the metrics. It should be one
        of the ['train',  'validation', 'test'].
    Returns: A metric function with the following API: ```metrics_fn(logits,
      label, weights)```
    """
    del split  # for all splits, we return the same metric functions

    return functools.partial(
        classification_model.classification_metrics_function,
        target_is_onehot=self.dataset_meta_data.get('target_is_onehot', False),
        metrics=_MBT_CLASSIFICATION_METRICS)

  def init_from_train_state(self,
                            train_state: Any,
                            restored_train_state: Any,
                            restored_model_cfg: ml_collections.ConfigDict,
                            restore_output_proj: bool = False) -> Any:
    """Updates the train_state with data from restored_train_state."""
    return model_utils.initialise_from_train_state(self.config, train_state,
                                                   restored_train_state,
                                                   restored_model_cfg,
                                                   restore_output_proj)

class MBTMultiHeadClassificationModel(MBTClassificationModel):
  """Audio Visual Transformer model for multiple n-way classification."""

  def __init__(self, config, dataset_meta_data):
    super().__init__(config, dataset_meta_data)

    assert self.config.dataset_configs.get('class_splits'), (
        'dataset_configs.class_splits must be specified')
    self.class_splits = np.cumsum(self.config.dataset_configs.class_splits)
    if self.config.dataset_configs.get('split_names'):
      self.split_names = self.config.dataset_configs.split_names
    else:
      self.split_names = [str(x + 1) for x in range(len(self.class_splits))]

    assert not config.get('multicrop_softmax_logits', False), (
        'Returning softmaxed logits during multicrop evaluation is not '
        'supported for this model.')

  def loss_function(self,
                    logits: jnp.ndarray,
                    batch: base_model.Batch,
                    model_params: Optional[jnp.ndarray] = None) -> float:
    """Return softmax cross entropy loss with an L2 penalty on the weights."""
    weights = batch.get('batch_mask')
    labels = batch['label']

    assert self.dataset_meta_data.get('target_is_onehot', False)
    if not isinstance(logits, dict):
      all_logits = logits
      logits = {}
      logits['all'] = all_logits
      if isinstance(labels, dict):
        assert 'all' in labels, 'mixmod must be turned off.'
        labels = labels['all']
      else:
        all_labels = labels
        labels = {}
        labels['all'] = all_labels

    sof_ce_loss = []
    for modality in logits:
      if logits[modality].shape[-1] != self.class_splits[-1]:
        raise AssertionError(
            'Logit dimension must be equal to number of classes')

      logit_splits = jnp.split(logits[modality],
                               self.class_splits, axis=-1)[:-1]
      assert not isinstance(labels[modality], dict), labels.keys()
      labels_splits = jnp.split(
          labels[modality], self.class_splits, axis=-1)[:-1]
      label_smoothing = self.config.get('label_smoothing')

      sof_ce_losses = [
          base_model_utils.weighted_softmax_cross_entropy(
              logit_split, labels_split, weights, label_smoothing)
          for logit_split, labels_split in zip(logit_splits, labels_splits)
      ]
      sof_ce_loss.append(jnp.mean(jnp.array(sof_ce_losses)))
    sof_ce_loss = jnp.mean(jnp.array(sof_ce_loss))

    if self.config.get('l2_decay_factor') is None:
      total_loss = sof_ce_loss
    else:
      l2_loss = base_model_utils.l2_regularization(model_params)
      total_loss = sof_ce_loss + 0.5 * self.config.l2_decay_factor * l2_loss
    return total_loss

  def get_metrics_fn(self, split: Optional[str] = None) -> base_model.MetricFn:
    """Returns a callable metric function for the model.

    Args:
      split: The split for which we calculate the metrics. It should be one
        of the ['train',  'validation', 'test'].
    Returns: A metric function with the following API: ```metrics_fn(logits,
      label, weights)```
    """
    del split  # for all splits, we return the same metric functions

    def classification_metrics_function(logits, batch, metrics, class_splits,
                                        split_names):

      one_hot_targets = batch['label']
      weights = batch.get('batch_mask')  # batch_mask might not be defined

      logit_splits = jnp.split(logits, class_splits, axis=-1)[:-1]
      one_hot_target_splits = jnp.split(
          one_hot_targets, class_splits, axis=-1)[:-1]

      evaluated_metrics = {}
      total_loss = [0.0, 0.0]
      for logits_i, one_hot_targets_i, name in zip(logit_splits,
                                                   one_hot_target_splits,
                                                   split_names):
        for key, val in metrics.items():
          evaluated_metrics[
              f'{name}_{key}'] = base_model_utils.psum_metric_normalizer(
                  (val[0](logits_i, one_hot_targets_i,
                          weights), val[1](logits_i, one_hot_targets_i,
                                           weights)))
          if key == 'loss':
            total_loss[0] += evaluated_metrics[f'{name}_{key}'][0]
            total_loss[1] += evaluated_metrics[f'{name}_{key}'][1]
      evaluated_metrics['total_loss'] = total_loss

      if len(class_splits) == 2:
        pairwise_acc = base_model_utils.psum_metric_normalizer(
            (model_utils.joint_accuracy(logits, one_hot_targets, class_splits,
                                        weights),
             base_model_utils.num_examples(logits, one_hot_targets, weights)))
        pairwise_top_five = base_model_utils.psum_metric_normalizer(
            (model_utils.joint_top_k(
                logits, one_hot_targets, class_splits, k=5, weights=weights),
             base_model_utils.num_examples(logits, one_hot_targets, weights)))
        eval_name = f'{split_names[0]}-{split_names[1]}'
        evaluated_metrics[f'{eval_name}_accuracy'] = pairwise_acc
        evaluated_metrics[f'{eval_name}_accuracy_top_5'] = pairwise_top_five

      return evaluated_metrics

    return functools.partial(
        classification_metrics_function,
        metrics=_MBT_CLASSIFICATION_METRICS,
        class_splits=self.class_splits,
        split_names=self.split_names)


#!/usr/bin/env python
"""
 Author: Shane Bussmann, T. K. Daisy Leung

Similar to uvmcmcfit.py, but here we edited it to
    - throw out a pre-defined number of burn-in samples;
    - save acceptance_fraction + misc stuff as a separate file
    - only save samples every some number of samples (instead of every iteration)
    - email ourselves once a certain number of samples have been obtained, and so we can decide whether or not to stop sampling instead of interupting the code

 Last modified: 2016 Dec 15

 Note: This is experimental software that is in a very active stage of
 development.  If you are interested in using this for your research, please
 contact me first at tleung@astro.cornell.edu!  Thanks.

 Purpose: Fit a parametric model to interferometric data using Dan
 Foreman-Mackey's emcee routine.  Gravitationally lensed sources are accounted
 for using ray-tracing routines based on Adam Bolton's lensdemo_script and
 lensdemo_func python scripts.  Here is the copyright license from
 lensdemo_script.py:

 Copyright 2009 by Adam S. Bolton
 Creative Commons Attribution-Noncommercial-ShareAlike 3.0 license applies:
 http://creativecommons.org/licenses/by-nc-sa/3.0/
 All redistributions, modified or otherwise, must include this
 original copyright notice, licensing statement, and disclaimer.
 DISCLAIMER: ABSOLUTELY NO WARRANTY EXPRESS OR IMPLIED.
 AUTHOR ASSUMES NO LIABILITY IN CONNECTION WITH THIS COMPUTER CODE.

--------------------------
 USAGE

 python $PYSRC/uvmcmcfit2.py

--------------------------
 SETUP PROCEDURES

 1. Establish a directory that contains data for the specific target for which
 you wish to measure a lens model.  This is the directory from which you will
 run the software.

 I call this "uvfit00" for the first run on a given dataset, "uvfit01" for
 the second, etc.

 2. Inside this directory, you must ensure the following files are present:

 - "config.yaml": This is the configuration file that describes where the source
 of interest is located, what type of model to use for the lens and source, the
 name of the image of the target from your interferometric data, the name of
 the uvfits files containing the interferometric visibilities, and a few
 important processing options as well.  Syntax is yaml.

 - Image of the target from your interferometric data.  The spatial resolution
 of this image (arcseconds per pixel), modified by an optional oversampling
 parameter, defines the spatial resolution in both the unlensed and lensed
 surface brightness maps.

 - interferometric visibilities for every combination of array configuration,
 sideband, and date observed that you want to model.

 3. More info about the constraints and priors input files.

 - Lenses: The lenses are assumed to have singular isothermal ellipsoid
 profiles.

 - Sources: Sources are represented by Gaussian profiles.

--------
 OUTPUTS

 "posteriorpdf.fits": model parameters for every MCMC iteration, in fits
 format.

 "summary.txt": contains mean acceptance fraction

"""

from __future__ import print_function

# import the required modules
import os
import os.path
import sys
from astropy.io import fits
import numpy
from astropy.table import Table
import emcee
#import pyximport
#pyximport.install(setup_args={"include_dirs":numpy.get_include()})
import sample_vis
import lensutil
import uvutil
import setuputil
import yaml
from subprocess import call
import time

#cwd = os.getcwd()
#sys.path.append(cwd)
#import config

def lnprior(pzero_regions, paramSetup):

    """

    Function that computes the ln prior probabilities of the model parameters.

    """
    priorln = 0.0
    mu = 1

#     import pdb; pdb.set_trace()

    # ensure all parameters are finite
    if (pzero_regions * 0 != 0).any():
        priorln = -numpy.inf
        return priorln, mu

    # Uniform priors
    uniform_regions = paramSetup['PriorShape'] == 'Uniform'
    if uniform_regions.any():
        p_l_regions = paramSetup['p_l'][uniform_regions]
        p_u_regions = paramSetup['p_u'][uniform_regions]
        pzero_uniform = pzero_regions[uniform_regions]
        if (pzero_uniform > p_l_regions).all() and (pzero_uniform < p_u_regions).all():
            # log prior
            priorln += numpy.log(1.0/numpy.abs(p_l_regions - p_u_regions)).sum()
        else:
            priorln = -numpy.inf
            return priorln, mu

    # Gaussian priors
    gaussian_regions = paramSetup['PriorShape'] == 'Gaussian'
    if gaussian_regions.any():
        import scipy.stats as stats
        # initlized as [mean, blah, blah, sigma]
        mean_regions = paramSetup['p_l'][gaussian_regions]
        rms_regions = paramSetup['p_u'][gaussian_regions]
        pzero_gauss = pzero_regions[gaussian_regions]
        priorln += numpy.log(stats.norm(scale=rms_regions, loc=mean_regions).pdf(pzero_gauss)).sum()

    # Gaussian pos (for parameter that must be positive e.g. flux density)
    gaussPos_regions = paramSetup['PriorShape'] == 'GaussianPos'
    if gaussPos_regions.any():
        pzero_gaussPos = pzero_regions[gaussPos_regions]
        if pzero_gaussPos < 0.0:
            priorln = -numpy.inf
            return priorln, mu
        else:
            import scipy.stats as stats
            # initlized as [mean, blah, blah, sigma]
            mean_regions = paramSetup['p_l'][gaussPos_regions]
            rms_regions = paramSetup['p_u'][gaussPos_regions]
            priorln += numpy.log(stats.norm(scale=rms_regions, loc=mean_regions).pdf(pzero_gauss)).sum()

#     if not isinstance(priorln, float):
#         priorln = priorln.sum()
    return priorln, mu

def lnlike(pzero_regions, vis_complex, wgt, uuu, vvv, pcd,
           fixindx, paramSetup, computeamp=True, miriad=False):
    """ Function that computes the Ln likelihood of the data"""

    # search poff_models for parameters fixed relative to other parameters
    fixindx = numpy.array(fixindx)
    fixed = (numpy.where(fixindx >= 0))[0]
    nfixed = fixindx[fixed].size
    p_u_regions = paramSetup['p_u']
    poff_regions = p_u_regions.copy()
    poff_regions[:] = 0.
    #for ifix in range(nfixed):
    #    poff_regions[fixed[ifix]] = pzero_regions[fixindx[fixed[ifix]]]
    for ifix in range(nfixed):
        ifixed = int(fixed[ifix])
        subindx = int(fixindx[ifixed])
        par0 = 0
        if fixindx[subindx] > 0:
            par0 = pzero_regions[fixindx[subindx]]
        poff_regions[ifixed] = pzero_regions[subindx] + par0

    parameters_regions = pzero_regions + poff_regions

    npar_previous = 0

    amp = []  # Will contain the 'blobs' we compute
    g_image_all = 0.
    g_lensimage_all = 0.
    e_image_all = 0.
    e_lensimage_all = 0.

    nregions = paramSetup['nregions']
    for regioni in range(nregions):

        # get the model info for this model
        x = paramSetup['x'][regioni]
        y = paramSetup['y'][regioni]
        headmod = paramSetup['modelheader'][regioni]
        nlens = paramSetup['nlens_regions'][regioni]
        nsource = paramSetup['nsource_regions'][regioni]
        model_types = paramSetup['model_types'][regioni]

        # get pzero, p_u, and p_l for this specific model
        nparlens = 5 * nlens
        nparsource = 6 * nsource
        npar = nparlens + nparsource + npar_previous
        parameters = parameters_regions[npar_previous:npar]
        npar_previous = npar

        #-----------------------------------------------------------------
        # Create a surface brightness map of lensed emission for the given set
        # of foreground lens(es) and background source parameters.
        #-----------------------------------------------------------------

        g_image, g_lensimage, e_image, e_lensimage, amp_tot, amp_mask = \
                lensutil.sbmap(x, y, nlens, nsource, parameters, model_types,
                computeamp=computeamp)
        e_image_all += e_image
        e_lensimage_all += e_lensimage
        g_image_all += g_image
        g_lensimage_all += g_lensimage
        amp.extend(amp_tot)
        amp.extend(amp_mask)

        # --------------------------------------------------------------------
        # Python version of UVMODEL:
        # "Observe" the lensed emission with the interferometer
        # --------------------------------------------------------------------

        if nlens > 0:
            if computeamp:
                # Evaluate amplification for each region
                lensmask = e_lensimage != 0
                mask = e_image != 0
                numer = g_lensimage[lensmask].sum()
                denom = g_image[mask].sum()
                amp_mask = numer / denom
                numer = g_lensimage.sum()
                denom = g_image.sum()
                amp_tot = numer / denom
                if amp_tot > 1e2:
                    amp_tot = 1e2
                if amp_mask > 1e2:
                    amp_mask = 1e2
                amp.extend([amp_tot])
                amp.extend([amp_mask])
            else:
                amp.extend([1.0])
                amp.extend([1.0])

    if miriad:
        # save the fits image of the lensed source
        ptag = str(os.getpid())
        SBmapLoc = 'LensedSBmap' + ptag + '.fits'
        fits.writeto(SBmapLoc, g_lensimage_all, header=headmod, clobber=True)

        # convert fits format to miriad format
        SBmapMiriad = 'LensedSBmap' + ptag + '.miriad'
        os.system('rm -rf ' + SBmapMiriad)
        cmd = 'fits op=xyin in=' + SBmapLoc + ' out=' + SBmapMiriad
        call(cmd + ' > /dev/null 2>&1', shell=True)

        # compute simulated visibilities
        modelvisfile = 'SimulatedVisibilities' + ptag + '.miriad'
        call('rm -rf ' + modelvisfile, shell=True)
        cmd = 'uvmodel options=subtract vis=' + visfilemiriad + \
                ' model=' + SBmapMiriad + ' out=' + modelvisfile
        call(cmd + ' > /dev/null 2>&1', shell=True)

        # convert simulated visibilities to uvfits format
        mvuvfits = 'SimulatedVisibilities' + ptag + '.uvfits'
        call('rm -rf ' + mvuvfits, shell=True)
        cmd = 'fits op=uvout in=' + modelvisfile + ' out=' + mvuvfits
        call(cmd + ' > /dev/null 2>&1', shell=True)

        # read simulated visibilities
        mvuv = fits.open(mvuvfits)
        diff_real = mvuv[0].data['DATA'][:, 0, 0, 0, 0, 0]
        diff_imag = mvuv[0].data['DATA'][:, 0, 0, 0, 0, 1]
        wgt = mvuv[0].data['DATA'][:, 0, 0, 0, 0, 2]
        #model_complex = model_real[goodvis] + 1.0j * model_imag[goodvis]
        diff_all = numpy.append(diff_real, diff_imag)
        wgt = numpy.append(wgt, wgt)
        goodvis = wgt > 0
        diff_all = diff_all[goodvis]
        wgt = wgt[goodvis]
        chi2_all = wgt * diff_all * diff_all
    else:
        model_complex = sample_vis.uvmodel(g_lensimage_all, headmod,
                uuu, vvv, pcd)
        diff_all = numpy.abs(vis_complex - model_complex)
        chi2_all = wgt * diff_all * diff_all
    #model_real += numpy.real(model_complex)
    #model_imag += numpy.imag(model_complex)

    #fits.writeto('g_lensimage.fits', g_lensimage_all, headmod, clobber=True)
    #import matplotlib.pyplot as plt
    #print(pzero_regions)
    #plt.imshow(g_lensimage, origin='lower')
    #plt.colorbar()
    #plt.show()
    #plt.imshow(g_image, origin='lower')
    #plt.colorbar()
    #plt.show()

    # calculate chi^2 assuming natural weighting
    #fnuisance = 0.0
    #modvariance_real = 1 / wgt #+ fnuisance ** 2 * model_real ** 2
    #modvariance_imag = 1 / wgt #+ fnuisance ** 2 * model_imag ** 2
    #wgt = wgt / 4.
    #chi2_real_all = (real - model_real) ** 2. / modvariance_real
    #chi2_imag_all = (imag - model_imag) ** 2. / modvariance_imag
    #chi2_all = numpy.append(chi2_real_all, chi2_imag_all)

    # compute the sigma term
    #sigmaterm_real = numpy.log(2 * numpy.pi / wgt)
    #sigmaterm_imag = numpy.log(2 * numpy.pi * modvariance_imag)

    # compute the ln likelihood
    lnlikemethod = paramSetup['lnlikemethod']
    if lnlikemethod == 'chi2':
        lnlike = chi2_all
    else:
        # by definition, loglike = -n/2*ln(2pi sigma^2) - 1/(2sigma^2) sum of (data-model)^2 over i=1 to n; but the constant term doesn't matter
        sigmaterm_all = len(wgt) * numpy.log(2 * numpy.pi / wgt)
        lnlike = chi2_all   # + sigmaterm_all
        # * -1/2 factor in latter step

    # compute number of degrees of freedom
    #nmeasure = lnlike.size
    #nparam = (pzero != 0).size
    #ndof = nmeasure - nparam

    # assert that lnlike is equal to -1 * maximum likelihood estimate
    # use visibilities where weight is greater than 0
    #goodvis = wgt > 0
    #likeln = -0.5 * lnlike[goodvis].sum()
    likeln = -0.5 * lnlike.sum()
    #print(pcd, likeln)
    if likeln * 0 != 0:
        likeln = -numpy.inf

    return likeln, amp

def lnprob(pzero_regions, vis_complex, wgt, uuu, vvv, pcd,
           fixindx, paramSetup, computeamp=True, miriad=False):

    """

    Computes ln probabilities via ln prior + ln likelihood

    """

    lp, mu = lnprior(pzero_regions, paramSetup)

    if not numpy.isfinite(lp):
        probln = -numpy.inf
        mu = 1
        return probln, mu

    ll, mu = lnlike(pzero_regions, vis_complex, wgt, uuu, vvv, pcd,
           fixindx, paramSetup, computeamp=computeamp, miriad=miriad)

    normalization = 1.0#2 * real.size
    probln = lp * normalization + ll
#    print(probln, lp*normalization, ll)   # remove

    return probln, mu

configloc = 'config.yaml'
configfile = open(configloc, 'r')
config = yaml.load(configfile)

# Determine if we are going to compute the amplification of every model
if config.keys().count('ComputeAmp') > 0:
    computeamp = config['ComputeAmp']
else:
    computeamp = True

# Determine parallel processing options
if config.keys().count('MPI') > 0:
    mpi = config['MPI']
else:
    mpi = False

# multiple processors on a cluster using MPI
if mpi:

    from emcee.utils import MPIPool

    # One thread per slot
    Nthreads = 1

    # Initialize the pool object
    pool = MPIPool()

    # If this process is not running as master, wait for instructions, then exit
    if not pool.is_master():
        pool.wait()
        sys.exit(0)

# Single processor with Nthreads cores
else:

    if config.keys().count('Nthreads') > 0:
        # set the number of threads to use for parallel processing
        Nthreads = config['Nthreads']
    else:
        Nthreads = 1

    # Initialize the pool object
    pool = ''

#--------------------------------------------------------------------------
# Read in ALMA image and beam
#im = fits.getdata(config['ImageName'])
#im = im[0, 0, :, :].copy()
headim = fits.getheader(config['ImageName'])

# get resolution in ALMA image
#celldata = numpy.abs(headim['CDELT1'] * 3600)

#--------------------------------------------------------------------------
# read in visibility data
visfile = config['UVData']

# Determine if we will use miriad to compute simulated visibilities
if config.keys().count('UseMiriad') > 0:
    miriad = config['UseMiriad']

    if miriad:
        interactive = False
        index = visfile.index('uvfits')
        visfilemiriad = visfile[0:index] + 'miriad'

        # scale the weights
        newvisfile = visfile[0:index] + 'scaled.uvfits'
        uvutil.scalewt(visfile, newvisfile)
        visfile = newvisfile
    else:
        miriad = False
else:
    miriad = False

# attempt to process multiple visibility files.  This won't work if miriad=True
try:
    filetype = visfile[-6:]
    if filetype == 'uvfits':
        uvfits = True
    else:
        uvfits = False
    uuu, vvv, www = uvutil.uvload(visfile)
    pcd = uvutil.pcdload(visfile)
    vis_complex, wgt = uvutil.visload(visfile)
except:
    try:
        for i, ivisfile in enumerate(visfile):
            filetype = ivisfile[-6:]
            if filetype == 'uvfits':
                uvfits = True
            else:
                uvfits = False
            iuuu, ivvv, iwww = uvutil.uvload(ivisfile)
            ipcd = uvutil.pcdload(ivisfile)
            ivis_complex, iwgt = uvutil.visload(ivisfile)
            if i == 0:
                uuu = iuuu
                vvv = ivvv
                pcd = ipcd
                vis_complex = ivis_complex
                wgt = iwgt
            else:
                uuu = numpy.append(uuu, iuuu)
                vvv = numpy.append(vvv, ivvv)
                if ipcd != pcd:
                    data1 = visfile[0]
                    data2 = visfile[ivisfile]
                    msg = 'Phase centers in ' + data1 + ' and ' + data2 \
                            + ' do not match.  Please ensure phase ' \
                            + 'centers in all visibility datasets are equal.'
                    print(msg)
                    raise TypeError
                vis_complex = numpy.append(vis_complex, ivis_complex)
                wgt = numpy.append(wgt, iwgt)
    except:
        msg = "Visibility datasets must be specified as either a string or "\
                "a list of strings."
        print(msg)
        raise TypeError

# remove the data points with zero or negative weight
positive_definite = wgt > 0
assert len(positive_definite[positive_definite]) > 0, " --- Find no data to fit, check the weights --- "
vis_complex = vis_complex[positive_definite]
wgt = wgt[positive_definite]
uuu = uuu[positive_definite]
vvv = vvv[positive_definite]
#www = www[positive_definite]

npos = wgt.size

#----------------------------------------------------------------------------
# Load input parameters
paramSetup = setuputil.loadParams(config)
nwalkers = paramSetup['nwalkers']
nregions = paramSetup['nregions']
nparams = paramSetup['nparams']
pname = paramSetup['pname']
nsource_regions = paramSetup['nsource_regions']

# Use an intermediate posterior PDF to initialize the walkers if it exists
posteriorloc = 'posteriorpdf.fits'
if os.path.exists(posteriorloc):

    # read the latest posterior PDFs
    print("Found existing posterior PDF file: {:s}".format(posteriorloc))
    posteriordat = Table.read(posteriorloc)
    if len(posteriordat) > 1:

        # assign values to pzero
        nlnprob = 1
        pzero = numpy.zeros((nwalkers, nparams))
        startindx = nlnprob
        for j in range(nparams):
            namej = posteriordat.colnames[j + startindx]
            pzero[:, j] = posteriordat[namej][-nwalkers:]

        # number of mu measurements
        nmu = len(posteriordat.colnames) - nparams - nlnprob

        # output name is based on most recent burnin file name
        realpdf = True
    else:
        realpdf = False
else:
    realpdf = False

if not realpdf:
    extendedpname = ['lnprob']
    extendedpname.extend(pname)
    nmu = 0
    for regioni in range(nregions):
        ri = str(regioni)
        if paramSetup['nlens_regions'][regioni] > 0:
            nsource = nsource_regions[regioni]
            for i in range(nsource):
                si = '.Source' + str(i) + '.Region' + ri
                extendedpname.append('mu_tot' + si)
                nmu += 1
            for i in range(nsource):
                si = '.Source' + str(i) + '.Region' + ri
                extendedpname.append('mu_aper' + si)
                nmu += 1
            extendedpname.append('mu_tot.Region' + ri)
            extendedpname.append('mu_aper.Region' + ri)
            nmu += 2
    posteriordat = Table(names = extendedpname)
    pzero = numpy.array(paramSetup['pzero'])

# make sure no parts of pzero exceed p_u or p_l
#arrayp_u = numpy.array(p_u)
#arrayp_l = numpy.array(p_l)
#for j in range(nwalkers):
#    exceed = arraypzero[j] >= arrayp_u
#    arraypzero[j, exceed] = 2 * arrayp_u[exceed] - arraypzero[j, exceed]
#    exceed = arraypzero[j] <= arrayp_l
#    arraypzero[j, exceed] = 2 * arrayp_l[exceed] - arraypzero[j, exceed]
#pzero = arraypzero
#p_u = arrayp_u
#p_l = arrayp_l

# determine the indices for fixed parameters
fixindx = setuputil.fixParams(paramSetup)
fixindx = map(int, fixindx)

# Initialize the sampler with the chosen specs.
if mpi:
    sampler = emcee.EnsembleSampler(nwalkers, nparams, lnprob, pool=pool, \
        args=[vis_complex, wgt, uuu, vvv, pcd, \
        fixindx, paramSetup, computeamp, miriad])
else:
    sampler = emcee.EnsembleSampler(nwalkers, nparams, lnprob, \
        args=[vis_complex, wgt, uuu, vvv, pcd, \
        fixindx, paramSetup, computeamp, miriad], threads=Nthreads)

# Sample, outputting to a file
#os.system('date')
currenttime = time.time()

# do burn-in if posteriorpdf.fits doesn't exist or contains any samples
# But, it's difficult to judge how many steps is needed
# need to may sure later that we are sampling longer than the AC time
if not realpdf:
    burnin = 150
    print("*** Running Burn in phase of steps {:d} ***".format(burnin))
    try:
        pos0, lnprob0, rstate0 = sampler.run_mcmc(pzero, burnin)
    except ValueError:
        pos0, lnprob0, rstate0, _ = sampler.run_mcmc(pzero, burnin)
    sampler.reset()            # reset chain
else:
    pos0 = pzero

class AlarmException(Exception):
    pass

def alarmHandler(signum, frame):
    raise AlarmException

def nonBlockingRawInput(prompt='', timeout=20, response='yes'):
    '''

    '''
    import signal
    signal.signal(signal.SIGALRM, alarmHandler)
    signal.alarm(timeout)
    try:
        text = raw_input(prompt)
        signal.alarm(0)
        return text
    except AlarmException:
        print('\nPrompt timeout. Continuing...')
    signal.signal(signal.SIGALRM, signal.SIG_IGN)
    return response

def query_yes_no(question, default=None):
    """Ask a yes/no question via raw_input() and return their answer.

    "question" is a string that is presented to the user.
    "default" is the presumed answer if the user just hits <Enter>.
        It must be "yes" (the default), "no" or None (meaning
        an answer is required of the user).

    The "answer" return value is True for "yes" or False for "no".
    """

    import sys
    valid = {"yes": True, "y": True, "ye": True,
             "no": False, "n": False}
    if default is None:
        prompt = " [y/n] "
    elif default == "yes":
        prompt = " [Y/n] "
    elif default == "no":
        prompt = " [y/N] "
    else:
        raise ValueError("invalid default answer: '%s'" % default)

    while True:
        sys.stdout.write(question + prompt)
        choice = raw_input().lower()
        if default is not None and choice == '':
            return valid[default]
        elif choice in valid:
            return valid[choice]
        else:
            sys.stdout.write("Please respond with 'yes' or 'no' "
                             "(or 'y' or 'n').\n")
#        sys.stdout.flush()

def email_self(msg, receiver='tleung@astro.cornell.edu'):

    '''
    Parameters
    ----------
    msg: str
        in email

    '''

    import os

    #email
    SENDMAIL = "/usr/sbin/sendmail"
    p = os.popen("%s -t" % SENDMAIL, "w")
    p.write("To: "+receiver+"\n")
    p.write("Subject: uvmcmcfit needs a respond to continuue. \n")
    p.write("\n")    # blank line separating headers from body

    message = msg + "\n\n" + ' Continue?'

    p.write(message)
    sts = p.close()
    if sts != 0:
        print("Sendmail exit status {}".format(sts))

import cPickle as pickle
import os
# pos - A list of current positions of walkers in the parameter space; dim = (nwalkers, dim)
# prob - The list of log posterior probabilities for the walkers at positions given by pos . The shape of this object is (nwalkers, dim).
# state - the random number generator state
# amp - metadata 'blobs' associated with the current positon

# below for testing..
# nsamples = 1000
# nsessions = 2

# in general, we want many samples.
# niter & nsesions dep. on nwalkers
#
nsamples = 1e6
niter = int(round(nsamples/nwalkers))
nsessions = 10
saveint = niter/nsessions/3

valid = {"yes": True, "y": True, "ye": True,
         "no": False, "n": False}

for i in range(nsessions):
    saveidx = 0
    for pos, prob, state, amp in sampler.sample(pos0, iterations=int(niter/nsessions)):
    # using sampler.sample() will have pre-defined 0s in elements (cf. run_mcmc())
        walkers, steps, dim = sampler.chain.shape
        result = [
            "Mean Acceptance fraction across all walkers of this iteration: {:.2f}".format(numpy.mean(sampler.acceptance_fraction)),
            "Mean lnprob and Max lnprob values: {:f} {:f}".format(numpy.mean(prob), numpy.max(prob)),
            "Time to run previous set of walkers (seconds): {:f}".format(time.time() - currenttime)
                ]
        print('\n'.join(result))
        f = open('summary.txt', 'a')
        f.write('\n'.join(result))
        f.write('\n')
        f.close()

        currenttime = time.time()
        #ff.write(str(prob))

        superpos = numpy.zeros(1 + nparams + nmu)
        for wi in range(nwalkers):
            superpos[0] = prob[wi]
            superpos[1:nparams + 1] = pos[wi]
            superpos[nparams + 1:nparams + nmu + 1] = amp[wi]
            posteriordat.add_row(superpos)

        # only save if it has went through every saveint iterations or is the last sample
        if not sampler.chain[:, numpy.any(sampler.chain[0, :, :] != 0, axis=1), :].shape[1] % saveint or (sampler.chain[:, numpy.any(sampler.chain[0, :, :] != 0, axis=1), :].shape[1] == int(niter/nsessions)):
            print("Ran {:d} iterations in this session. Saving data".format(sampler.chain[:, numpy.any(sampler.chain[0, :, :] != 0, axis=1), :].shape[1]))
            posteriordat.write('posteriorpdf.fits', overwrite=True)
            #posteriordat.write('posteriorpdf.txt', format='ascii')

            saveidx = sampler.chain[:, numpy.any(sampler.chain[0, :, :] != 0, axis=1), :].shape[1]

    message = "We have finished {:d} iterations with {:d} walkers. ".format(sampler.chain[:, numpy.any(sampler.chain[0, :, :] != 0, axis=1), :].shape[1], nwalkers)

    if i < nsessions-1:
        email_self(message)
        print(message)
        ret = nonBlockingRawInput("Shall we continuue with next session? (Y/N)", timeout=600).lower()

        while not ret in valid:
            print("Please respond with 'yes' or 'no' (or 'y' or 'n').\n")
            ret = nonBlockingRawInput("Shall we continuue with next session? (Y/N)", timeout=600).lower()
        if not valid[ret]:
            import sys
            sys.exit("Quiting... ")
            if mpi: pool.close()

    sampler.reset()
    pos0 = pos

f = open('summary.txt', 'a')
f.write("Finish all {:d} sessions \n".format(nsessions))
f.write("Total number of samples: {:d} \n".format(niter/nsessions * nsessions * nwalkers))
f.write('\n')
f.close()

if mpi: pool.close()