# Copyright 2021 Calico LLC
# Copyright 2021 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import time
import os
import glob
import json
import functools
import inspect
from pathlib import Path

import tensorflow as tf
from tqdm import tqdm
import numpy as np
import pandas as pd
from typing import Any, Callable, Dict, Optional, Text, Union, Iterable, List, Sequence

import sonnet as snt
from sonnet.src import base, once, types, utils
from sonnet.src.optimizers import optimizer_utils

import tensorflow as tf
import wandb

# attribute

# Enformer tensorflow code was directly taken and modified for distributed training
# https://github.com/deepmind/deepmind-research/tree/master/enformer

# Genetic augmentation code was taken from
# https://github.com/calico/basenji/blob/84c681a4b02f592a3de90799cee7f17d96f81ef8/basenji/archive/augmentation.py

# constants

NUM_CORES_ENFORCE = 64  # using v3-64

SEQUENCE_LENGTH = 196_608
TARGET_LENGTH = 896
BIN_SIZE = 128

# assert TPUs

tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu = 'enformer')
tf.config.experimental_connect_to_cluster(tpu)
tf.tpu.experimental.initialize_tpu_system(tpu)
tpu_strategy = snt.distribute.TpuReplicator(tpu)

num_cores =  tpu_strategy.num_replicas_in_sync
assert num_cores == NUM_CORES_ENFORCE, f'must betraining on {num_cores} cores'

# optimizer

def adam_update(g, alpha, beta_1, beta_2, epsilon, t, m, v):
  """Implements 'Algorithm 1' from :cite:`kingma2014adam`."""
  m = beta_1 * m + (1. - beta_1) * g      # Biased first moment estimate.
  v = beta_2 * v + (1. - beta_2) * g * g  # Biased second raw moment estimate.
  m_hat = m / (1. - tf.pow(beta_1, t))    # Bias corrected 1st moment estimate.
  v_hat = v / (1. - tf.pow(beta_2, t))    # Bias corrected 2nd moment estimate.
  update = alpha * m_hat / (tf.sqrt(v_hat) + epsilon)
  return update, m, v

# https://github.com/deepmind/sonnet/blob/v2/sonnet/src/optimizers/adam.py
# modified for Adam with decoupled weight decay

class Adam(base.Optimizer):
  def __init__(self,
               learning_rate: Union[types.FloatLike, tf.Variable] = 0.001,
               beta1: Union[types.FloatLike, tf.Variable] = 0.9,
               beta2: Union[types.FloatLike, tf.Variable] = 0.999,
               epsilon: Union[types.FloatLike, tf.Variable] = 1e-8,
               weight_decay: Union[types.FloatLike, tf.Variable] = 1e-4,
               name: Optional[str] = None):
    super().__init__(name=name)
    self.learning_rate = learning_rate
    self.beta1 = beta1
    self.beta2 = beta2
    self.epsilon = epsilon
    self.weight_decay = weight_decay
    # TODO(petebu): Consider allowing the user to pass in a step.
    self.step = tf.Variable(0, trainable=False, name="t", dtype=tf.int64)
    self.m = []
    self.v = []

  @once.once
  def _initialize(self, parameters: Sequence[tf.Variable]):
    """First and second order moments are initialized to zero."""
    zero_var = lambda p: utils.variable_like(p, trainable=False)
    with tf.name_scope("m"):
      self.m.extend(zero_var(p) for p in parameters)
    with tf.name_scope("v"):
      self.v.extend(zero_var(p) for p in parameters)

  def apply(self, updates: Sequence[types.ParameterUpdate],
            parameters: Sequence[tf.Variable]):
    optimizer_utils.check_distribution_strategy()
    optimizer_utils.check_updates_parameters(updates, parameters)
    self._initialize(parameters)
    self.step.assign_add(1)
    for update, param, m_var, v_var in zip(updates, parameters, self.m, self.v):
      if update is None:
        continue

      optimizer_utils.check_same_dtype(update, param)
      learning_rate = tf.cast(self.learning_rate, update.dtype)
      beta_1 = tf.cast(self.beta1, update.dtype)
      beta_2 = tf.cast(self.beta2, update.dtype)
      epsilon = tf.cast(self.epsilon, update.dtype)
      step = tf.cast(self.step, update.dtype)

      update, m, v = adam_update(
        g=update, alpha=learning_rate, beta_1=beta_1, beta_2=beta_2,
        epsilon=epsilon, t=step, m=m_var, v=v_var)

      # decoupled weight decay
      # hack for now to exclude biases
      weight_decay_update = (param * self.weight_decay * learning_rate) if 'w:0' in param.name else tf.zeros_like(param)

      param.assign_sub(update)
      param.assign_sub(weight_decay_update)

      m_var.assign(m)
      v_var.assign(v)

# classes

class MultiheadAttention(snt.Module):
  """Multi-head attention."""

  def __init__(self,
               value_size: int,
               key_size: int,
               num_heads: int,
               scaling: bool = True,
               attention_dropout_rate: float = 0.1,
               relative_positions: bool = False,
               relative_position_symmetric: bool = False,
               relative_position_functions: Optional[List[str]] = None,
               num_relative_position_features: Optional[int] = None,
               positional_dropout_rate: float = 0.1,
               zero_initialize: bool = True,
               initializer: Optional[snt.initializers.Initializer] = None,
               name: str = None):
    """Creates a MultiheadAttention module.

    Args:.prefetch(2)
      value_size: The size of each value embedding per head.
      key_size: The size of each key and query embedding per head.
      num_heads: The number of independent queries per timestep.
      scaling: Whether to scale the attention logits.
      attention_dropout_rate: Dropout rate for attention logits.
      relative_positions: Whether to use TransformerXL style relative attention.
      relative_position_symmetric: If True, the symmetric version of basis
        functions will be used. If False, a symmetric and asymmetric versions
        will be use.
      relative_position_functions: List of function names used for relative
        positional biases.
      num_relative_position_features: Number of relative positional features
        to compute. If None, `value_size * num_heads` is used.
      positional_dropout_rate: Dropout rate for the positional encodings if
        relative positions are used.
      zero_initialize: if True, the final linear layer will be 0 initialized.
      initializer: Initializer for the projection layers. If unspecified,
        VarianceScaling is used with scale = 2.0.
      name: Name of module.
    """
    super().__init__(name=name)
    self._value_size = value_size
    self._key_size = key_size
    self._num_heads = num_heads
    self._attention_dropout_rate = attention_dropout_rate
    self._scaling = scaling
    self._relative_positions = relative_positions
    self._relative_position_symmetric = relative_position_symmetric
    self._relative_position_functions = relative_position_functions
    if num_relative_position_features is None:
      # num_relative_position_features needs to be divisible by the number of
      # relative positional functions *2 (for symmetric & asymmetric version).
      divisible_by = 2 * len(self._relative_position_functions)
      self._num_relative_position_features = (
          (self._value_size // divisible_by) * divisible_by)
    else:
      self._num_relative_position_features = num_relative_position_features
    self._positional_dropout_rate = positional_dropout_rate

    self._initializer = initializer
    if self._initializer is None:
      self._initializer = snt.initializers.VarianceScaling(scale=2.0)

    key_proj_size = self._key_size * self._num_heads
    embedding_size = self._value_size * self._num_heads

    self._q_layer = snt.Linear(
        key_proj_size,
        name='q_layer',
        with_bias=False,
        w_init=self._initializer)
    self._k_layer = snt.Linear(
        key_proj_size,
        name='k_layer',
        with_bias=False,
        w_init=self._initializer)
    self._v_layer = snt.Linear(
        embedding_size,
        name='v_layer',
        with_bias=False,
        w_init=self._initializer)
    w_init = snt.initializers.Constant(1e-8) if zero_initialize else self._initializer
    self._embedding_layer = snt.Linear(
        embedding_size,
        name='embedding_layer',
        w_init=w_init,
        b_init= snt.initializers.Constant(1e-8))

    # Create additional layers if using relative positions.
    if self._relative_positions:
      self._r_k_layer = snt.Linear(
          key_proj_size,
          name='r_k_layer',
          with_bias=False,
          w_init=self._initializer)
      self._r_w_bias = tf.Variable(
          self._initializer([1, self._num_heads, 1, self._key_size],
                            dtype=tf.float32),
          name='r_w_bias')
      self._r_r_bias = tf.Variable(
          self._initializer([1, self._num_heads, 1, self._key_size],
                            dtype=tf.float32),
          name='r_r_bias')

  def _multihead_output(self, linear, inputs):
    """Applies a standard linear to inputs and returns multihead output."""

    output = snt.BatchApply(linear)(inputs)  # [B, T, H * KV]
    num_kv_channels = output.shape[-1] // self._num_heads
    # Split H * Channels into separate axes.
    output = snt.reshape(output,
                         output_shape=[-1, self._num_heads, num_kv_channels])
    # [B, T, H, KV] -> [B, H, T, KV]
    return tf.transpose(output, [0, 2, 1, 3])

  def __call__(self,
               inputs,
               is_training=False):
    # Initialise the projection layers.
    embedding_size = self._value_size * self._num_heads
    seq_len = inputs.shape[1]

    # Compute q, k and v as multi-headed projections of the inputs.
    q = self._multihead_output(self._q_layer, inputs)  # [B, H, T, K]
    k = self._multihead_output(self._k_layer, inputs)  # [B, H, T, K]
    v = self._multihead_output(self._v_layer, inputs)  # [B, H, T, V]

    # Scale the query by the square-root of key size.
    if self._scaling:
      q *= self._key_size**-0.5

    if self._relative_positions:
      # For relative positions, we project positions to form relative keys.
      distances = tf.range(-seq_len + 1, seq_len, dtype=tf.float32)[tf.newaxis]
      positional_encodings = positional_features_all(
          positions=distances,
          feature_size=self._num_relative_position_features,
          seq_length=seq_len,
          feature_functions=self._relative_position_functions,
          symmetric=self._relative_position_symmetric)
      # [1, 2T-1, Cr]

      if is_training:
        positional_encodings = tf.nn.dropout(
            positional_encodings, rate=self._positional_dropout_rate)

      # [1, H, 2T-1, K]
      r_k = self._multihead_output(self._r_k_layer, positional_encodings)

      # Add shifted relative logits to content logits.
      # [B, H, T', T]
      content_logits = tf.matmul(q + self._r_w_bias, k, transpose_b=True)
      # [B, H, T', 2T-1]
      relative_logits = tf.matmul(
          q + self._r_r_bias, r_k, transpose_b=True)
      #  [B, H, T', T]
      relative_logits = relative_shift(relative_logits)
      logits = content_logits + relative_logits
    else:
      # [B, H, T', T]
      logits = tf.matmul(q, k, transpose_b=True)

    weights = tf.nn.softmax(logits)

    # Dropout on the attention weights.
    if is_training:
      weights = tf.nn.dropout(weights, rate=self._attention_dropout_rate)

    # Transpose and reshape the output.
    output = tf.matmul(weights, v)  # [B, H, T', V]
    output_transpose = tf.transpose(output, [0, 2, 1, 3])  # [B, T', H, V]

    # Final linear layer.
    attended_inputs = snt.reshape(
        output_transpose, output_shape=[embedding_size], preserve_dims=2)
    output = self._embedding_layer(attended_inputs)

    return output

def relative_shift(x):
  """Shift the relative logits like in TransformerXL."""
  # We prepend zeros on the final timescale dimension.
  to_pad = tf.zeros_like(x[..., :1])
  x = tf.concat([to_pad, x], -1)
  _, num_heads, t1, t2 = x.shape
  x = tf.reshape(x, [-1, num_heads, t2, t1])
  x = tf.slice(x, [0, 0, 1, 0], [-1, -1, -1, -1])
  x = tf.reshape(x, [-1, num_heads, t1, t2 - 1])
  x = tf.slice(x, [0, 0, 0, 0], [-1, -1, -1, (t2 + 1) // 2])
  return x

# Available feature functions:
def get_positional_feature_function(name):
  """Returns positional feature functions."""
  available = {
      'positional_features_exponential': positional_features_exponential,
      'positional_features_central_mask': positional_features_central_mask,
      'positional_features_gamma': positional_features_gamma
  }
  if name not in available:
    raise ValueError(f'Function {name} not available in {available.keys()}')
  return available[name]

def positional_features_all(positions: tf.Tensor,
                            feature_size: int,
                            seq_length: Optional[int] = None,
                            bin_size: Optional[int] = None,
                            feature_functions: Optional[List[str]] = None,
                            symmetric=False):
  """Compute relative positional encodings/features.

  Each positional feature function will compute/provide the same fraction of
  features, making up the total of feature_size.

  Args:
    positions: Tensor of relative positions of arbitrary shape.
    feature_size: Total number of basis functions.
    seq_length: Sequence length denoting the characteristic length that
      the individual positional features can use. This is required since the
      parametrization of the input features should be independent of `positions`
      while it could still require to use the total number of features.
    bin_size: Bin sized used to partition the sequence. This can be used to
      compute features on the absolute scale relative to the genome.
    feature_functions: List of different feature functions to use. Each function
      will take as argument: positions, sequence length and number of features
      to compute.
    symmetric: If True, the resulting features will be symmetric across the
      relative position of 0 (i.e. only absolute value of positions will
      matter). If false, then both the symmetric and asymmetric version
      (symmetric multiplied by sign(positions)) of the features will be used.

  Returns:
    Tensor of shape: `positions.shape + (feature_size,)`.
  """
  if feature_functions is None:
    feature_functions = ['positional_features_exponential',
                         'positional_features_central_mask',
                         'positional_features_gamma']
  num_components = len(feature_functions)  # 1 per each basis function
  if not symmetric:
    num_components = 2 * num_components

  # For now, we do not allow odd sized embeddings.
  if feature_size % num_components != 0:
    raise ValueError(
        f'feature_size has to be divisible by {num_components}')

  feature_functions = [get_positional_feature_function(f)
                       for f in feature_functions]
  num_basis_per_class = feature_size // num_components
  embeddings = tf.concat([f(tf.abs(positions), num_basis_per_class,
                            seq_length, bin_size)
                          for f in feature_functions],
                         axis=-1)
  if not symmetric:
    embeddings = tf.concat([embeddings,
                            tf.sign(positions)[..., tf.newaxis] * embeddings],
                           axis=-1)
  tf.TensorShape(embeddings.shape).assert_is_compatible_with(
      positions.shape + [feature_size])
  return embeddings

def _prepend_dims(x, num_dims):
  return tf.reshape(x, shape=[1] * num_dims + x.shape)

def positional_features_exponential(positions: tf.Tensor,
                                    feature_size: int,
                                    seq_length: Optional[int] = None,
                                    bin_size: Optional[int] = None,
                                    min_half_life: Optional[float] = 3.0):
  """Create exponentially decaying positional weights.

  Args:
    positions: Position tensor (arbitrary shape).
    feature_size: Number of basis functions to use.
    seq_length: Sequence length.
    bin_size: (unused). See `positional_features_all`.
    min_half_life: Smallest exponential half life in the grid of half lives.

  Returns:
    A Tensor with shape [2 * seq_length - 1, feature_size].
  """
  del bin_size  # Unused.
  if seq_length is None:
    seq_length = tf.reduce_max(tf.abs(positions)) + 1
  # Grid of half lifes from [3, seq_length / 2] with feature_size
  # distributed on the log scale.
  seq_length = tf.cast(seq_length, dtype=tf.float32)
  max_range = tf.math.log(seq_length) / tf.math.log(2.0)
  half_life = tf.pow(2.0, tf.linspace(min_half_life, max_range, feature_size))
  half_life = _prepend_dims(half_life, positions.shape.rank)
  positions = tf.abs(positions)
  outputs = tf.exp(-tf.math.log(2.0) / half_life * positions[..., tf.newaxis])
  tf.TensorShape(outputs.shape).assert_is_compatible_with(
      positions.shape + [feature_size])
  return outputs

def positional_features_central_mask(positions: tf.Tensor,
                                     feature_size: int,
                                     seq_length: Optional[int] = None,
                                     bin_size: Optional[int] = None):
  """Positional features using a central mask (allow only central features)."""
  del seq_length  # Unused.
  del bin_size  # Unused.
  center_widths = tf.pow(2.0, tf.range(1, feature_size + 1, dtype=tf.float32))
  center_widths = center_widths - 1
  center_widths = _prepend_dims(center_widths, positions.shape.rank)
  outputs = tf.cast(center_widths > tf.abs(positions)[..., tf.newaxis],
                    tf.float32)
  tf.TensorShape(outputs.shape).assert_is_compatible_with(
      positions.shape + [feature_size])
  return outputs

def gamma_pdf(x, concentration, rate):
  """Gamma probability distribution function: p(x|concentration, rate)."""
  log_unnormalized_prob = tf.math.xlogy(concentration - 1., x) - rate * x
  log_normalization = (tf.math.lgamma(concentration) -
                       concentration * tf.math.log(rate))
  return tf.exp(log_unnormalized_prob - log_normalization)

def positional_features_gamma(positions: tf.Tensor,
                              feature_size: int,
                              seq_length: Optional[int] = None,
                              bin_size: Optional[int] = None,
                              stddev=None,
                              start_mean=None):
  """Positional features computed using the gamma distributions."""
  del bin_size  # Unused.
  if seq_length is None:
    seq_length = tf.reduce_max(tf.abs(positions)) + 1
  if stddev is None:
    stddev = seq_length / (2 * feature_size)
  if start_mean is None:
    start_mean = seq_length / feature_size
  mean = tf.linspace(start_mean, seq_length, num=feature_size)
  mean = _prepend_dims(mean, positions.shape.rank)
  concentration = (mean / stddev)**2
  rate = mean / stddev**2
  probabilities = gamma_pdf(
      tf.abs(tf.cast(positions, dtype=tf.float32))[..., tf.newaxis],
      concentration, rate)
  probabilities += 1e-8  # To ensure numerical stability.
  outputs = probabilities / tf.reduce_max(probabilities)
  tf.TensorShape(outputs.shape).assert_is_compatible_with(
      positions.shape + [feature_size])
  return outputs

class Enformer(snt.Module):
  """Main model."""

  def __init__(self,
               channels: int = 1536,
               num_transformer_layers: int = 11,
               num_heads: int = 8,
               pooling_type: str = 'attention',
               use_convnext: bool = False,
               name: str = 'enformer'):
    """Enformer model.

    Args:
      channels: Number of convolutional filters and the overall 'width' of the
        model.
      num_transformer_layers: Number of transformer layers.
      num_heads: Number of attention heads.
      pooling_type: Which pooling function to use. Options: 'attention' or max'.
      name: Name of sonnet module.
    """
    super().__init__(name=name)
    # pylint: disable=g-complex-comprehension,g-long-lambda,cell-var-from-loop
    heads_channels = {'human': 5313, 'mouse': 1643}
    dropout_rate = 0.4
    assert channels % num_heads == 0, ('channels needs to be divisible '
                                       f'by {num_heads}')
    whole_attention_kwargs = {
        'attention_dropout_rate': 0.05,
        'initializer': None,
        'key_size': 64,
        'num_heads': num_heads,
        'num_relative_position_features': channels // num_heads,
        'positional_dropout_rate': 0.01,
        'relative_position_functions': [
            'positional_features_exponential',
            'positional_features_central_mask',
            'positional_features_gamma'
        ],
        'relative_positions': True,
        'scaling': True,
        'value_size': channels // num_heads,
        'zero_initialize': True
    }

    trunk_name_scope = tf.name_scope('trunk')
    trunk_name_scope.__enter__()
    from sonnet.src import moving_averages

    # lambda is used in Sequential to construct the module under tf.name_scope.
    def conv_block(filters, width=1, w_init=None, name='conv_block', **kwargs):
      with tf.name_scope(name or "batch_norm"):
        moving_mean = moving_averages.ExponentialMovingAverage(
            0.9, name="moving_mean")
        moving_variance = moving_averages.ExponentialMovingAverage(
            0.9, name="moving_variance")
      return Sequential(lambda: [
          snt.distribute.CrossReplicaBatchNorm(create_scale=True,
                        create_offset=True,
                        moving_mean = moving_mean,
                        moving_variance = moving_variance,
                        scale_init=snt.initializers.Ones()),
          gelu,
          snt.Conv1D(filters, width, w_init=w_init, **kwargs)
      ], name=name)

    def convnext_block(filters, width=1, mult = 4, ds_conv_kernel_size = 7, w_init=None, name='convnext_block', **kwargs):
      return Sequential(lambda: [
          ExpandDims(2),
          snt.DepthwiseConv2D((ds_conv_kernel_size, 1), name ='convnext_ds_conv'),
          Squeeze(2),
          snt.LayerNorm(axis=-1, create_scale=True, create_offset=True),
          snt.Linear(filters * mult, name='convnext_project_in'),
          tf.nn.relu,
          snt.Linear(filters, name='convnext_project_out')
      ], name=name)

    conv_block_fn = convnext_block if use_convnext else conv_block

    stem = Sequential(lambda: [
        snt.Conv1D(channels // 2, 15),
        Residual(conv_block(channels // 2, 1, name='pointwise_conv_block')),
        pooling_module(pooling_type, pool_size=2),
    ], name='stem')

    filter_list = exponential_linspace_int(start=channels // 2, end=channels,
                                           num=6, divisible_by=128)
    conv_tower = Sequential(lambda: [
        Sequential(lambda: [
            conv_block(num_filters, 5),
            Residual(conv_block(num_filters, 1, name='pointwise_conv_block')),
            pooling_module(pooling_type, pool_size=2),
            ],
                   name=f'conv_tower_block_{i}')
        for i, num_filters in enumerate(filter_list)], name='conv_tower')

    # Transformer.
    def transformer_mlp():
      return Sequential(lambda: [
          snt.LayerNorm(axis=-1, create_scale=True, create_offset=True),
          snt.Linear(channels * 2, name = 'project_in'),
          snt.Dropout(dropout_rate),
          tf.nn.relu,
          snt.Linear(channels, name = 'project_out'),
          snt.Dropout(dropout_rate)], name='mlp')

    transformer = Sequential(lambda: [
        Sequential(lambda: [
            Residual(Sequential(lambda: [
                snt.LayerNorm(axis=-1,
                              create_scale=True, create_offset=True,
                              scale_init=snt.initializers.Ones()),
                MultiheadAttention(**whole_attention_kwargs,
                                                    name=f'attention_{i}'),
                snt.Dropout(dropout_rate),
            ], name='mha')),
            Residual(transformer_mlp())], name=f'transformer_block_{i}')
        for i in range(num_transformer_layers)], name='transformer')

    crop_final = TargetLengthCrop1D(TARGET_LENGTH, name='target_input')

    final_pointwise = Sequential(lambda: [
        conv_block(channels * 2, 1),
        snt.Dropout(dropout_rate / 8),
        gelu], name='final_pointwise')

    self._trunk = Sequential([stem,
                              conv_tower,
                              transformer,
                              crop_final,
                              final_pointwise],
                             name='trunk')
    trunk_name_scope.__exit__(None, None, None)

    with tf.name_scope('heads'):
      self._heads = {
          head: Sequential(
              lambda: [snt.Linear(num_channels), tf.nn.softplus],
              name=f'head_{head}')
          for head, num_channels in heads_channels.items()
      }
    # pylint: enable=g-complex-comprehension,g-long-lambda,cell-var-from-loop

  @property
  def trunk(self):
    return self._trunk

  @property
  def heads(self):
    return self._heads

  def __call__(self, inputs: tf.Tensor,
               is_training: bool) -> Dict[str, tf.Tensor]:
    trunk_embedding = self.trunk(inputs, is_training=is_training)
    return {
        head: head_module(trunk_embedding, is_training=is_training)
        for head, head_module in self.heads.items()
    }

  @tf.function(input_signature=[
      tf.TensorSpec([None, SEQUENCE_LENGTH, 4], tf.float32)])
  def predict_on_batch(self, x):
    """Method for SavedModel."""
    return self(x, is_training=False)

class TargetLengthCrop1D(snt.Module):
  """Crop sequence to match the desired target length."""

  def __init__(self, target_length: int, name='target_length_crop'):
    super().__init__(name=name)
    self._target_length = target_length

  def __call__(self, inputs):
    trim = (inputs.shape[-2] - self._target_length) // 2
    if trim < 0:
      raise ValueError('inputs longer than target length')

    return inputs[..., trim:-trim, :]

class ExpandDims(snt.Module):

  def __init__(self, dim: int, name='expand_dims'):
    super().__init__(name=name)
    self._dim = dim

  def __call__(self, inputs):
    return tf.expand_dims(inputs, self._dim)

class Squeeze(snt.Module):

  def __init__(self, dim: int, name='squeeze'):
    super().__init__(name=name)
    self._dim = dim

  def __call__(self, inputs):
    return tf.squeeze(inputs, self._dim)

class Sequential(snt.Module):
  """snt.Sequential automatically passing is_training where it exists."""

  def __init__(self,
               layers: Optional[Union[Callable[[], Iterable[snt.Module]],
                                      Iterable[Callable[..., Any]]]] = None,
               name: Optional[Text] = None):
    super().__init__(name=name)
    if layers is None:
      self._layers = []
    else:
      # layers wrapped in a lambda function to have a common namespace.
      if hasattr(layers, '__call__'):
        with tf.name_scope(name):
          layers = layers()
      self._layers = [layer for layer in layers if layer is not None]

  def __call__(self, inputs: tf.Tensor, is_training: bool, **kwargs):
    outputs = inputs
    for _, mod in enumerate(self._layers):
      if accepts_is_training(mod):
        outputs = mod(outputs, is_training=is_training, **kwargs)
      else:
        outputs = mod(outputs, **kwargs)
    return outputs

def pooling_module(kind, pool_size):
  """Pooling module wrapper."""
  if kind == 'attention':
    return SoftmaxPooling1D(pool_size=pool_size, per_channel=True,
                            w_init_scale=2.0)
  elif kind == 'max':
    return tf.keras.layers.MaxPool1D(pool_size=pool_size, padding='same')
  else:
    raise ValueError(f'Invalid pooling kind: {kind}.')

class SoftmaxPooling1D(snt.Module):
  """Pooling operation with optional weights."""

  def __init__(self,
               pool_size: int = 2,
               per_channel: bool = False,
               w_init_scale: float = 0.0,
               name: str = 'softmax_pooling'):
    """Softmax pooling.

    Args:
      pool_size: Pooling size, same as in Max/AvgPooling.
      per_channel: If True, the logits/softmax weights will be computed for
        each channel separately. If False, same weights will be used across all
        channels.
      w_init_scale: When 0.0 is equivalent to avg pooling, and when
        ~2.0 and `per_channel=False` it's equivalent to max pooling.
      name: Module name.
    """
    super().__init__(name=name)
    self._pool_size = pool_size
    self._per_channel = per_channel
    self._w_init_scale = w_init_scale
    self._logit_linear = None

  @snt.once
  def _initialize(self, num_features):
    self._logit_linear = snt.Linear(
        output_size=num_features if self._per_channel else 1,
        with_bias=False,  # Softmax is agnostic to shifts.
        w_init=snt.initializers.Identity(self._w_init_scale))

  def __call__(self, inputs):
    _, length, num_features = inputs.shape
    self._initialize(num_features)
    inputs = tf.reshape(
        inputs,
        (-1, length // self._pool_size, self._pool_size, num_features))
    return tf.reduce_sum(
        inputs * tf.nn.softmax(self._logit_linear(inputs), axis=-2),
        axis=-2)

class Residual(snt.Module):
  """Residual block."""

  def __init__(self, module: snt.Module, name='residual'):
    super().__init__(name=name)
    self._module = module

  def __call__(self, inputs: tf.Tensor, is_training: bool, *args,
               **kwargs) -> tf.Tensor:
    return inputs + self._module(inputs, is_training, *args, **kwargs)

def gelu(x: tf.Tensor) -> tf.Tensor:
  """Applies the Gaussian error linear unit (GELU) activation function.

  Using approximiation in section 2 of the original paper:
  https://arxiv.org/abs/1606.08415

  Args:
    x: Input tensor to apply gelu activation.
  Returns:
    Tensor with gelu activation applied to it.
  """
  return tf.nn.sigmoid(1.702 * x) * x

def one_hot_encode(sequence: str,
                   alphabet: str = 'ACGT',
                   neutral_alphabet: str = 'N',
                   neutral_value: Any = 0,
                   dtype=np.float32) -> np.ndarray:
  """One-hot encode sequence."""
  def to_uint8(string):
    return np.frombuffer(string.encode('ascii'), dtype=np.uint8)
  hash_table = np.zeros((np.iinfo(np.uint8).max, len(alphabet)), dtype=dtype)
  hash_table[to_uint8(alphabet)] = np.eye(len(alphabet), dtype=dtype)
  hash_table[to_uint8(neutral_alphabet)] = neutral_value
  hash_table = hash_table.astype(dtype)
  return hash_table[to_uint8(sequence)]

def exponential_linspace_int(start, end, num, divisible_by=1):
  """Exponentially increasing values of integers."""
  def _round(x):
    return int(np.round(x / divisible_by) * divisible_by)

  base = np.exp(np.log(end / start) / (num - 1))
  return [_round(start * base**i) for i in range(num)]

def accepts_is_training(module):
  return 'is_training' in list(inspect.signature(module.__call__).parameters)

# data related functions

# @title `get_targets(organism)`
def get_targets(organism):
  targets_txt = f'https://raw.githubusercontent.com/calico/basenji/master/manuscripts/cross2020/targets_{organism}.txt'
  return pd.read_csv(targets_txt, sep='\t')

# @title `get_dataset(organism, subset, num_threads=8)`

def reverse_complement_transform(seq):
  """Reverse complement of batched onehot seq and corresponding label and na."""

  # reverse complement sequence
  seq_rc = tf.gather(seq, [3, 2, 1, 0], axis=-1)
  seq_rc = tf.reverse(seq_rc, axis=[0])
  return seq_rc

def shift_sequence(seq, shift_amount, pad_value=0.25):
  """Shift a sequence left or right by shift_amount.
  Args:
    seq: a [batch_size, sequence_length, sequence_depth] sequence to shift
    shift_amount: the signed amount to shift (tf.int32 or int)
    pad_value: value to fill the padding (primitive or scalar tf.Tensor)
  """
  input_shape = seq.shape

  pad = pad_value * tf.ones_like(seq[0:tf.abs(shift_amount), :])

  def _shift_right(_seq):
    sliced_seq = _seq[:-shift_amount:, :]
    return tf.concat([pad, sliced_seq], axis=0)

  def _shift_left(_seq):
    sliced_seq = _seq[-shift_amount:, :]
    return tf.concat([sliced_seq, pad], axis=0)

  output = tf.cond(
      tf.greater(shift_amount, 0), lambda: _shift_right(seq),
      lambda: _shift_left(seq))

  output.set_shape(input_shape)
  return output

def augment_stochastic_shifts(seq, augment_shifts):
  """Apply a stochastic shift augmentation.
  Args:
    seq: input sequence of size [batch_size, length, depth]
    augment_shifts: list of int offsets to sample from
  Returns:
    shifted and padded sequence of size [batch_size, length, depth]
  """
  shift_index = tf.random.uniform(shape=[], minval=0,
      maxval=len(augment_shifts), dtype=tf.int64)
  shift_value = tf.gather(tf.constant(augment_shifts), shift_index)

  seq = tf.cond(tf.not_equal(shift_value, 0),
                lambda: shift_sequence(seq, shift_value),
                lambda: seq)

  return seq

def augment_stochastic_shifts_map_fn(datum):
  augment_shifts = [-2, -1, 0, 1, 2]
  return dict(
    sequence = augment_stochastic_shifts(datum['sequence'], augment_shifts),
    target = datum['target']
  )

def augment_stochastic_rc_map_fn(datum):
  sequence, target = (datum['sequence'], datum['target'])
  augment = tf.random.uniform(shape=[]) > 0.5
  sequence, target = tf.cond(augment, lambda: (sequence[::-1, ::-1], target[::-1, :]),
                              lambda: (sequence, target))
  return dict(sequence = sequence, target = target)

def organism_path(organism):
    return os.path.join(f'gs://basenji_barnyard/data', organism)

def get_dataset(organism, subset, num_threads=8, shuffle=True, rotate = 0, augment = False):
  metadata = get_metadata(organism)
  files = tfrecord_files(organism, subset)
  files = files[rotate:] + files[:rotate]
  dataset = tf.data.TFRecordDataset(files,
                                    compression_type='ZLIB',
                                    num_parallel_reads=num_threads)
  if shuffle:
    dataset = dataset.repeat()
    dataset = dataset.shuffle(5000, seed = 42)

  dataset = dataset.map(functools.partial(deserialize, metadata=metadata),
                        num_parallel_calls=num_threads)
  if augment:
    dataset = dataset.map(augment_stochastic_shifts_map_fn, num_parallel_calls=num_threads)
    dataset = dataset.map(augment_stochastic_rc_map_fn, num_parallel_calls=num_threads)

  return dataset

def get_metadata(organism):
  # Keys:
  # num_targets, train_seqs, valid_seqs, test_seqs, seq_length,
  # pool_width, crop_bp, target_length
  path = os.path.join(organism_path(organism), 'statistics.json')
  with tf.io.gfile.GFile(path, 'r') as f:
    return json.load(f)

def tfrecord_files(organism, subset):
  # Sort the values by int(*).
  return sorted(tf.io.gfile.glob(os.path.join(
      organism_path(organism), 'tfrecords', f'{subset}-*.tfr'
  )), key=lambda x: int(x.split('-')[-1].split('.')[0]))

def deserialize(serialized_example, metadata):
  """Deserialize bytes stored in TFRecordFile."""
  feature_map = {
      'sequence': tf.io.FixedLenFeature([], tf.string),
      'target': tf.io.FixedLenFeature([], tf.string),
  }
  example = tf.io.parse_example(serialized_example, feature_map)
  sequence = tf.io.decode_raw(example['sequence'], tf.bool)
  sequence = tf.reshape(sequence, (metadata['seq_length'], 4))
  sequence = tf.cast(sequence, tf.float32)

  target = tf.io.decode_raw(example['target'], tf.float16)
  target = tf.reshape(target,
                      (metadata['target_length'], metadata['num_targets']))
  target = tf.cast(target, tf.float32)

  return {'sequence': sequence,
          'target': target}

# new get_dataset, for sequences that are actually 196_608

NEW_TFRECORD_LOCATIONS = dict(
  human = dict(
    train = 'gs://enformer-human-train/',
    valid = 'gs://enformer-human-valid/'
  ),
  mouse = dict(
    train = 'gs://enformer-mouse-train/',
    valid = 'gs://enformer-mouse-valid/'
  )
)

NUM_TRACKS_CONFIG = dict(human = 5313, mouse = 1643)

def new_dataset_map_seq_target(
  element,
  seq_len,
  species,  # 'human' or 'mouse'
  target_length = 896,
  shifts = None,
  augment_rc = False
):
  assert species in NUM_TRACKS_CONFIG, f'{species} not found in config'
  num_tracks = NUM_TRACKS_CONFIG[species]

  num_shifts = 0 if shifts is None else len(list(range(shifts[0], shifts[1] + 1)))

  data = {
    'seq': tf.io.FixedLenFeature([(seq_len + num_shifts) * 4], tf.float32),
    'target': tf.io.FixedLenFeature([target_length * num_tracks], tf.float32),
  }

  content = tf.io.parse_single_example(element, data)

  content['sequence'] = content.pop('seq')
  content['sequence'] = tf.reshape(content['sequence'], (-1, 4))
  content['target'] = tf.reshape(content['target'], (target_length, -1))

  # take care of shift augmentation

  shifts = tf.pad(tf.random.uniform(shape = [1], minval = 0, maxval = num_shifts, dtype = tf.int64), [[0, 1]])
  content['sequence'] = tf.slice(content['sequence'], shifts, (seq_len, -1))

  if augment_rc:
    content = augment_stochastic_rc_map_fn(content)

  content['sequence'].set_shape(tf.TensorShape([seq_len, 4]))
  content['target'].set_shape(tf.TensorShape([target_length, num_tracks]))

  return content

def get_dataset_new(
  organism,
  datatype,
  shifts = (-2, 2),
  augment_rc = False,
  num_threads = 8
):
  gcs_path = NEW_TFRECORD_LOCATIONS[organism][datatype]
  files = sorted(tf.io.gfile.glob(f'{gcs_path}*.tfrecord'))

  dataset = tf.data.TFRecordDataset(files, compression_type = 'ZLIB', num_parallel_reads = num_threads)
  map_element_fn = partial(new_dataset_map_seq_target, seq_len = SEQUENCE_LENGTH, species = organism, shifts = shifts, augment_rc = augment_rc)
  dataset = dataset.map(map_element_fn)
  return dataset

# training related functions

def corr_coef(x, y, eps = 0):
  x2 = tf.math.square(x)
  y2 = tf.math.square(y)
  xy = x * y
  ex = tf.reduce_mean(x, axis = 1)
  ey = tf.reduce_mean(y, axis = 1)
  exy = tf.reduce_mean(xy, axis = 1)
  ex2 = tf.reduce_mean(x2, axis = 1)
  ey2 = tf.reduce_mean(y2, axis = 1)
  r = (exy - ex * ey) / ((tf.math.sqrt(ex2 - tf.math.square(ex) + eps) * tf.math.sqrt(ey2 - tf.math.square(ey) + eps)) + eps)
  return tf.reduce_mean(r, axis = -1)

def create_eval_step(model, head):
  @tf.function
  def predict(seq, target):
    pred = model(seq, is_training=False)[head]
    return corr_coef(pred, target)
  return predict

def create_step_function(model, optimizer, head, clip_grad_norm = 1.0, weight_decay = 0.0001):

  @tf.function
  def train_step(batch_seq, batch_target):
    with tf.GradientTape() as tape:
      with snt.mixed_precision.scope(tf.float16):
        outputs = model(batch_seq, is_training=True)[head]

      corr_coef_loss = 1 - corr_coef(outputs, batch_target, eps = 1e-8)
      poisson = tf.reduce_mean(
          tf.keras.losses.poisson(batch_target, outputs))
      loss = poisson

    gradients = tape.gradient(loss, model.trainable_variables, unconnected_gradients=tf.UnconnectedGradients.ZERO)
    gradients = [tf.clip_by_norm(grad, clip_grad_norm) for grad in gradients]
    ctx = tf.distribute.get_replica_context()
    gradients = ctx.all_reduce("mean", gradients)
    optimizer.apply(gradients, model.trainable_variables)
    return loss

  return train_step

# instantiate model and training / eval functions

with tpu_strategy.scope():
  model = Enformer(channels=1536,
                   num_heads=8,
                   num_transformer_layers=11)

  learning_rate = tf.Variable(0., trainable=False, name='learning_rate')
  optimizer = snt.optimizers.Adam(learning_rate=learning_rate)

  train_step_human = create_step_function(model, optimizer, 'human')
  train_step_mouse = create_step_function(model, optimizer, 'mouse')

  eval_step_human = create_eval_step(model, 'human')
  eval_step_mouse = create_eval_step(model, 'mouse')

# experiment tracker

wandb.init(project='enformer')
wandb.run.save()

# Train the model

num_steps = int(2e6)
num_warmup_steps = 5000
target_learning_rate = 5e-4

checkpoint_every = 2500
max_eval_steps = 25
eval_every = 500

# Step variables

global_step = tf.Variable(0, name='global_step', trainable=False)

# checkpointing

checkpoint_root = "gs://enformer/"
checkpoint_name = "enformer"

save_prefix = os.path.join(checkpoint_root, checkpoint_name)

checkpoint = tf.train.Checkpoint(module = model,  step = global_step, optimizer = optimizer)

# load latest checkpoint if possible

latest = tf.train.latest_checkpoint(checkpoint_root)
if latest is not None:
  checkpoint.restore(latest)

@tf.function
def step():
  global_step.assign(global_step + 1)

  batch_human, batch_mouse = next(data_it)
  loss_human = tpu_strategy.run(train_step_human, args = (batch_human['sequence'], batch_human['target']))
  loss_mouse = tpu_strategy.run(train_step_mouse, args = (batch_mouse['sequence'], batch_mouse['target']))

  loss_human = tpu_strategy.reduce('mean', loss_human, axis = None)
  loss_mouse = tpu_strategy.reduce('mean', loss_mouse, axis = None)

  learning_rate_frac = tf.math.minimum(1.0, tf.cast(global_step, tf.float32) / tf.math.maximum(1.0, float(num_warmup_steps)))
  learning_rate.assign(target_learning_rate * learning_rate_frac)

  return loss_human, loss_mouse

@tf.function
def eval_step():
  batch_human = next(valid_human_data_it)
  batch_mouse = next(valid_mouse_data_it)
  human_r = tpu_strategy.run(eval_step_human, args = (batch_human['sequence'], batch_human['target']))
  mouse_r = tpu_strategy.run(eval_step_mouse, args = (batch_mouse['sequence'], batch_mouse['target']))

  human_r = tpu_strategy.reduce('mean', human_r, axis = 0)
  mouse_r = tpu_strategy.reduce('mean', mouse_r, axis = 0)
  return human_r, mouse_r

i = global_step.numpy()

total_mice = 114 * 256 + 111
total_human = 132 * 256 + 229
bucket_size = 256
num_seen = i * num_cores
human_file_skip = (num_seen % total_human) // bucket_size
mouse_file_skip = (num_seen % total_mice) // bucket_size

human_dataset = get_dataset('human', 'train', rotate = human_file_skip).batch(num_cores, drop_remainder = True)
mouse_dataset = get_dataset('mouse', 'train', rotate = mouse_file_skip).batch(num_cores, drop_remainder = True)
human_mouse_dataset = tf.data.Dataset.zip((human_dataset, mouse_dataset)).prefetch(2)

human_valid_dataset = get_dataset('human', 'valid', shuffle = False).repeat().batch(num_cores)
mouse_valid_dataset = get_dataset('mouse', 'valid', shuffle = False).repeat().batch(num_cores)

data_it = iter(tpu_strategy.experimental_distribute_dataset(human_mouse_dataset))
valid_human_data_it = iter(tpu_strategy.experimental_distribute_dataset(human_valid_dataset))
valid_mouse_data_it = iter(tpu_strategy.experimental_distribute_dataset(mouse_valid_dataset))

print(f'starting from {i}')

while i < num_steps:
  print(f'processing step {i}')
  loss_human, loss_mouse = step()
  loss_human = loss_human.numpy()
  loss_mouse = loss_mouse.numpy()
  learning_rate_numpy = learning_rate.numpy()
  print(f'completed step {i}')
  log = {
    'loss_human': loss_human,
    'loss_mouse': loss_mouse,
    'learning_rate': learning_rate_numpy
  }

  if i and not i % eval_every:
    print('evaluating')

    human_pearson_r, mouse_pearson_r = eval_step()
    human_pearson_r = human_pearson_r.numpy()
    mouse_pearson_r = mouse_pearson_r.numpy()

    log = {
      **log,
      'human_pearson_r': human_pearson_r,
      'mouse_pearson_r': mouse_pearson_r
    }

  wandb.log(log, step = i)

  if not i % checkpoint_every:
    print('checkpointing')
    checkpoint.save(save_prefix)

  i += 1


# -*- coding: utf-8 -*-
"""
为了适应中国雷达在不同仰角的探测距离不同以及前几层仰角 dop和ref分开扫描的问题
提出NuistRadar Object，以方便后续的算法及绘图
"""
import numpy as np
import xarray as xr
from ..configure.default_config import DEFAULT_METADATA, FILL_VALUE, CINRAD_field_mapping
from ..core.transforms import antenna_vectors_to_cartesian
from scipy import spatial
from ..interp.RadarInterp import radar_interp2d_var, radar_interp2d

class PRD(object):
    """
    Polarimetry Radar Data (PRD)
    A class for storing antenna coordinate radar data.
    Attributes
    ----------
    fields : dict
        Moment fields. with different variables
    scan_type : str
        Type of scan, one of 'ppi', 'rhi', 'sector' or 'other'. If the scan
        volume contains multiple sweep modes this should be 'other'.
    time : datetime object
        Time at the center of each ray.
    range : numpy array //m
        Range to the center of each gate (bin).
    latitude : scalar//units:degree
        Latitude of the instrument.
    longitude: scalar//units:degree
        Longitude of the instrument.
    altitude : scalar//units:m
        Altitude of the instrument, above sea level.
    fixed_angle : (nsweeps) units:degree
        Target angle for thr sweep. Azimuth angle in RHI modes, elevation
        angle in all other modes.
    azimuth : (nrays) units :degree
        Azimuth of antenna, relative to true North. Azimuth angles are
        recommended to be expressed in the range of [0, 360], but other
        representations are not forbidden.
    elevation : (nrays) units :degree
        Elevation of antenna, relative to the horizontal plane. Elevation
        angles are recommended to be expressed in the range of [-180, 180],
        but other representations are not forbidden.
    sweep_start_ray_index : numpy array(nsweeps)
        Index of the first ray in each sweep relative to the start of the
        volume, 0-based.
    sweep_end_ray_index : numpy array(nsweeps)
        Index of the last ray in each sweep relative to the start of the
        volume, 0-based.
    rays_per_sweep : numpy array (nsweeps)
        Number of rays in each sweep. The data key of this attribute is
        create upon first access from the data in the sweep_start_ray_index and
        sweep_end_ray_index attributes. If the sweep locations needs to be
        modified, do this prior to accessing this attribute or use
        :py:func:`init_rays_per_sweep` to reset the attribute.
    bins_per_sweep : numpy array (nsweeps)    !!!##added
        Number of bins in each sweep. The data key of this attribute is
        create upon first access from the data in the sweep_start_ray_index and
        sweep_end_ray_index attributes. If the sweep locations needs to be
        modified, do this prior to accessing this attribute or use
        :py:func:`init_rays_per_sweep` to reset the attribute.
    nyquist_velocity: numpy array (nsweeps) (m/s)
    unambiguous_range:numpy array (nsweeps) (m/s)
    frequency: constant (GHZ)
    nrays : int
        Number of rays in the volume.
    nsweeps : int
        Number of sweep in the volume.

    """

    def __init__(self, fields,  scan_type, time, range, azimuth, elevation,latitude,
                 longitude, altitude, sweep_start_ray_index, sweep_end_ray_index,
                 fixed_angle, bins_per_sweep, nyquist_velocity, frequency, unambiguous_range,
                 nrays, nsweeps, sitename):

        super(PRD, self).__init__()
        keys = fields.keys()
        self.fields = []
        for idx, (istart, iend) in enumerate(zip(sweep_start_ray_index, sweep_end_ray_index)):
            isweep_data = xr.Dataset(coords={'azimuth': (['time', ], azimuth[istart:iend+1]),
                                            'elevation': (['time',], elevation[istart:iend+1]),
                                            'range': range[:bins_per_sweep[idx]], 'time': time[istart:iend+1]})
            isweep_data.azimuth.attrs = DEFAULT_METADATA['azimuth']
            isweep_data.elevation.attrs = DEFAULT_METADATA['elevation']
            isweep_data.range.attrs = DEFAULT_METADATA['range']
            isweep_data.time.attrs = DEFAULT_METADATA['time']
            for ikey in keys:
                isweep_data[ikey] = (['time','range'], fields[ikey][istart:iend+1, :bins_per_sweep[idx]])
                isweep_data[ikey].attrs = DEFAULT_METADATA[CINRAD_field_mapping[ikey]]
            self.fields.append(isweep_data)
        self.scan_info = xr.Dataset(data_vars={"latitude":latitude,"longitude":longitude,
                        "altitude":altitude,"scan_type":scan_type,  "frequency":frequency,
                         "nyquist_velocity":(['sweep',], nyquist_velocity),
                        "unambiguous_range":(['sweep',], unambiguous_range),
                        "rays_per_sweep": (['sweep',], sweep_end_ray_index-sweep_start_ray_index+1),
                        "fixed_angle": (["sweep",], fixed_angle)},
                        coords={"sweep": np.arange(nsweeps, dtype=int)})

        self.scan_info['latitude'].attrs = DEFAULT_METADATA['latitude']
        self.scan_info['longitude'].attrs = DEFAULT_METADATA['longitude']
        self.scan_info['altitude'].attrs = DEFAULT_METADATA['altitude']
        self.scan_info['scan_type'].attrs = DEFAULT_METADATA['scan_type']
        self.scan_info['frequency'].attrs = DEFAULT_METADATA['frequency']
        self.scan_info['nyquist_velocity'].attrs = DEFAULT_METADATA['nyquist_velocity']
        self.scan_info['unambiguous_range'].attrs = DEFAULT_METADATA['unambiguous_range']
        self.scan_info['rays_per_sweep'].attrs = DEFAULT_METADATA['rays_per_sweep']
        self.scan_info['fixed_angle'].attrs = DEFAULT_METADATA['fixed_angle']
        self.nsweeps = nsweeps
        self.nrays = nrays
        self.sitename = sitename

    def get_vertical_section(self, start_point, end_point, field_name):
        """
        :param start_point: units:m
        :param end_point:  units:m
        :return:
        """
        start_x, start_y = start_point
        end_x, end_y = end_point
        bins_res = (self.fields[0].range[1] - self.fields[0].range[0]).values
        start_end_dis = np.sqrt((start_x-end_x)**2 + (start_y-end_y)**2)
        npoints = int(start_end_dis/(bins_res/2.) + 1)
        x_line = np.linspace(start_x, end_x, npoints)
        y_line = np.linspace(start_y, end_y, npoints)
        target = np.c_[x_line, y_line]
        xy_line_1d = np.linspace(0, start_end_dis, npoints)
        z_values = []
        field_values = []
        #先对剖线取最邻近点
        for ifield in self.fields:
            _x, _y, _z = antenna_vectors_to_cartesian(ifield.range.values, \
                                ifield.azimuth.values, ifield.elevation.values)
            kdtree = spatial.cKDTree(np.c_[_x.ravel(), _y.ravel()])
            _distance, _idx = kdtree.query(target, k=1, n_jobs=-1)
            _z_value = np.where(_distance > bins_res*2, np.nan, _z.ravel()[_idx])
            _field_value = np.where(_distance > bins_res*2, np.nan, ifield[field_name].values.ravel()[_idx])
            z_values.append(_z_value)
            field_values.append(_field_value)
        z_values = np.asarray(z_values)
        field_values = np.asarray(field_values)
        xy_line_values = np.stack([xy_line_1d,]*len(self.fields), axis=0)
        mask_flag = (np.isnan(z_values.ravel()) | np.isnan(field_values.ravel()))
        z_values_ravel = z_values.ravel()[~mask_flag]
        xy_line_values_ravel = xy_line_values.ravel()[~mask_flag]
        field_values_ravel = field_values.ravel()[~mask_flag]
        mesh_xy, mesh_z = np.mgrid[0:start_end_dis:npoints*1j, 0:np.nanmax(z_values):bins_res/2.] #生成剖面的网格
        grid_field = radar_interp2d_var(np.c_[xy_line_values_ravel, z_values_ravel], field_values_ravel,\
                              (mesh_xy, mesh_z), bandwidth=360/self.fields[0].azimuth.size)
        return mesh_xy, mesh_z, grid_field


import numpy as np
from astropy.io import fits
import matplotlib.pyplot as plt
import lightweaver.constants as Const
from lightweaver.rh_atoms import H_6_atom, H_6_CRD_atom, H_3_atom, C_atom, O_atom, OI_ord_atom, Si_atom, Al_atom, CaII_atom, Fe_atom, FeI_atom, He_9_atom, He_atom, He_large_atom, MgII_atom, N_atom, Na_atom, S_atom
from lightweaver.atmosphere import Atmosphere, ScaleType
from lightweaver.atomic_set import RadiativeSet
from lightweaver.atomic_table import get_global_atomic_table
from lightweaver.molecule import MolecularTable
from lightweaver.LwCompiled import LwContext
from lightweaver.utils import InitialSolution
from concurrent.futures import ProcessPoolExecutor, wait, as_completed
from tqdm import tqdm
from contextlib import redirect_stdout
import os
import pickle

def prep_atmos(data, xIdx, yIdx):
    height = data[xIdx, yIdx, :, 0].astype('<f8') / 1e2
    temp = data[xIdx, yIdx, :, 1].astype('<f8')
    vlos = data[xIdx, yIdx, :, 3].astype('<f8') / 1e2
    # pgasTop = data[xIdx, yIdx, 0, 2].astype('<f8') / (Const.CM_TO_M**2 / Const.G_TO_KG)
    pgas = data[xIdx, yIdx, :, 2].astype('<f8') / (Const.CM_TO_M**2 / Const.G_TO_KG)

    return {'height': height, 'temp': temp, 'vlos': vlos, 'pgas': pgas}

def iterate_ctx(ctx, prd=True, Nscatter=3, NmaxIter=1000):
    for i in range(NmaxIter):
        dJ = ctx.formal_sol_gamma_matrices()
        if i < Nscatter:
            continue
        delta = ctx.stat_equil()
        if prd:
            dRho = ctx.prd_redistribute(maxIter=5)

        if ctx.crswDone and dJ < 3e-3 and delta < 1e-3:
            print(i)
            print('----------')
            return

wave = np.linspace(853.9444, 854.9444, 1001)

data = fits.getdata('better_eb_310400.fits')
# atmosData = prep_atmos(data, 10,10)

def crsw_factory(initVal=1e3):
    val = initVal
    def callback():
        nonlocal val
        val = max(1.0, val * 0.1**(1/val))
        return val
    return callback

def cmo_synth(atmosData, crsw=None):
    with open(os.devnull, 'w') as f:
        with redirect_stdout(f):
            if crsw is not None:
                crsw = crsw()
            atmos = Atmosphere(ScaleType.Geometric, depthScale=atmosData['height'], temperature=atmosData['temp'], vlos=atmosData['vlos'], vturb=4000*np.ones_like(atmosData['height']))

            aSet = RadiativeSet([H_3_atom(), C_atom(), O_atom(), Si_atom(), Al_atom(), CaII_atom(), Fe_atom(), He_atom(), MgII_atom(), N_atom(), Na_atom(), S_atom()])
            aSet.set_active('H', 'Ca')

            spect = aSet.compute_wavelength_grid()

            atmos.convert_scales(Pgas=atmosData['pgas'])
            atmos.quadrature(5)

            mols = MolecularTable()
            eqPops = aSet.iterate_lte_ne_eq_pops(mols, atmos)
            ctx = LwContext(atmos, spect, eqPops, conserveCharge=True, initSol=InitialSolution.Lte, crswCallback=crsw)
            iterate_ctx(ctx, prd=False)
            eqPops.update_lte_atoms_Hmin_pops(atmos)
            Iwave = ctx.compute_rays(wave, [1.0])
            return Iwave

def cmo_synth_lte(atmosData):
    with open(os.devnull, 'w') as f:
        with redirect_stdout(f):
            atmos = Atmosphere(ScaleType.Geometric, depthScale=atmosData['height'], temperature=atmosData['temp'], vlos=atmosData['vlos'], vturb=4000*np.ones_like(atmosData['height']))

            aSet = RadiativeSet([H_3_atom(), C_atom(), O_atom(), Si_atom(), Al_atom(), CaII_atom(), Fe_atom(), He_atom(), MgII_atom(), N_atom(), Na_atom(), S_atom()])
            aSet.set_active('Ca')

            spect = aSet.compute_wavelength_grid()

            atmos.convert_scales(Pgas=atmosData['pgas'])
            atmos.quadrature(5)

            mols = MolecularTable()
            eqPops = aSet.iterate_lte_ne_eq_pops(mols, atmos)
            ctx = LwContext(atmos, spect, eqPops, conserveCharge=False)
            iterate_ctx(ctx, prd=False)
            eqPops.update_lte_atoms_Hmin_pops(atmos)
            Iwave = ctx.compute_rays(wave, [1.0])
            return Iwave

atmosData = []
with open('BrokenPixels.pickle', 'rb') as pkl:
    brokenPixels = pickle.load(pkl)

for x, y in brokenPixels:
    atmosData.append(prep_atmos(data, x, y))

with ProcessPoolExecutor() as executor:
    futures = [executor.submit(cmo_synth, d, crsw_factory) for d in atmosData]
    for f in tqdm(as_completed(futures), total=len(futures)):
        pass

spectra = []
for f in futures:
    try:
        spectra.append(f.result())
    except:
        spectra.append(None)

name = 'NlteRedo.pickle'
with open(name, 'wb') as f:
    pickle.dump(spectra, f)

# atmosHse = Atmosphere(ScaleType.Geometric, depthScale=atmosData['height'], temperature=atmosData['temp'], vlos=atmosData['vlos'], vturb=4000*np.ones_like(atmosData['height']))

# atmosHse.convert_scales(Ptop=atmosData['pgas'][0])
# atmosHse.quadrature(5)

# aSet = RadiativeSet([H_3_atom(), C_atom(), O_atom(), Si_atom(), Al_atom(), CaII_atom(), Fe_atom(), He_atom(), MgII_atom(), N_atom(), Na_atom(), S_atom()])
# aSet.set_active('Ca')
# spectHse = aSet.compute_wavelength_grid()
# eqPopsHse = aSet.iterate_lte_ne_eq_pops(mols, atmosHse)
# ctxHse = LwContext(atmosHse, spectHse, eqPopsHse, conserveCharge=False, initSol=InitialSolution.Lte)
# iterate_ctx(ctxHse, prd=False)
# eqPopsHse.update_lte_atoms_Hmin_pops(atmosHse)
# IwaveHse = ctxHse.compute_rays(wave, [1.0])

# plt.ion()
# plt.plot(wave, Iwave)
# plt.plot(wave, IwaveHse)
# plt.show()


#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Sky background related"""

import os
import warnings

import numpy as np

from astropy.table import Table

from scipy.stats import sigmaclip
from scipy.stats import binned_statistic_2d

from . import utils
from . import plotting

__all__ = ['SkyObjs', 'AperPhot', 'S18A_APER']

class AperPhot():
    """
    Class for aperture photometry in HSC.
    """
    PIX = 0.168 # arcsec/pixe

    def __init__(self, name, rad, rerun='s18a'):
        """Start a aperture photometry object."""
        self.aper_id = name
        self.name = "aper{0}".format(self.aper_id)
        self.r_pix = rad
        self.area_pix = np.pi * (rad ** 2.0)
        self.r_arcsec = rad * self.PIX
        self.area_arcsec = np.pi * (self.r_arcsec ** 2.0)
        self.rerun = rerun

        # Name of the columns for flux and flux error
        self.flux_col = self.flux(rerun=self.rerun)
        self.err_col = self.err(rerun=self.rerun)

    def flux(self, band=None, rerun='s18a'):
        """Aperture flux column name in S18A."""
        if rerun == 's18a':
            if band is not None:
                return "{0}_apertureflux_{1}_flux".format(band.strip(), self.aper_id)
            return "apertureflux_{0}_flux".format(self.aper_id)
        else:
            raise NotImplementedError("# Only S18A data are available.")

    def err(self, band=None, rerun='s18a'):
        """Aperture flux error column name in S18A."""
        if rerun == 's18a':
            if band is not None:
                return "{0}_{1}sigma".format(band.strip(), self.flux(rerun=rerun))
            return "{0}sigma".format(self.flux(rerun=rerun))
        else:
            raise NotImplementedError("# Only S18A data are available.")

# Aperture flux in S18A
S18A_APER_ID = ['10', '15', '20', '30', '40', '57', '84',
                '118', '168', '235']
S18A_APER_RAD = [3.0, 4.5, 6.0, 9.0, 12.0, 17.0, 25.0, 35.0, 50.0, 70.0]
S18A_APER = {}
for ii, rr in zip(S18A_APER_ID, S18A_APER_RAD):
    S18A_APER['aper{0}'.format(ii)] = AperPhot(ii, rr)

class SkyObjs():
    """
    Class for HSC sky objects.
    """
    # Convert the flux from erg/s/cm^2/Hz to HSC image value
    CGS_TO_IMG = 1.7378E30
    # Convert the flux from erg/s/cm^2/Hz to muJy
    CGS_TO_MUJY = 1.0E29
    # Convert the from from muJy to HSC image unit
    MUJY_TO_IMG = CGS_TO_IMG / CGS_TO_MUJY

    # List of filters
    FILTER_LIST = ['HSC-G', 'HSC-R', 'HSC-I', 'HSC-Z', 'HSC-Y']

    # Nicknames of filters
    FILTER_SHORT = ['g', 'r', 'i', 'z', 'y']

    def __init__(self, skyobjs, meas=False, nobj_min=5):
        """
        Initialize an object for HSC sky object catalog.
        """
        # Whether it is a forced photometry or a measurement catalog
        if meas:
            self.ra_col = 'i_ra'
            self.dec_col = 'i_dec'
            self.type = 'meas'
            self.meas = True
        else:
            self.ra_col = 'ra'
            self.dec_col = 'dec'
            self.type = 'force'
            self.meas = False

        # If skyobjs is a file name, read in the catalog
        if isinstance(skyobjs, str):
            _, file_ext = os.path.splitext(skyobjs)
            if file_ext == '.npy':
                self.skyobjs = np.load(skyobjs)
            elif file_ext == '.fits':
                self.skyobjs = Table.read(skyobjs).as_array().data
            else:
                raise TypeError("# Wrong file type: npy or fits!")
        elif isinstance(skyobjs, Table):
            try:
                self.skyobjs = skyobjs.as_array()
            except Exception:
                self.skyobjs = skyobjs.as_array()
        elif isinstance(skyobjs, np.ndarray) or isinstance(skyobjs, np.recarray):
            self.skyobjs = skyobjs

        # Minimum number of sky objects
        self.n_min = nobj_min

        # List of Tracts
        self.tract_list = list(np.unique(self.skyobjs['tract']))
        self.n_tract = len(self.tract_list)

        # List of Patches and Tracts
        self.tract_patch = np.unique(
            ["{0}_{1:03d}".format(t, p) for t, p in
             zip(self.skyobjs['tract'], self.skyobjs['patch'])])
        self.n_tract_patch = len(self.tract_patch)

    def select_tract(self, tract, patch=None, n_min=10, verbose=True) -> 'SkyObjs':
        """Select sky objects on one Tract (and Patch) from the catalog """
        tract_mask = self.skyobjs['tract'] == tract
        if tract_mask.sum() == 0:
            if verbose:
                warnings.warn("# Tract {0} is not available!".format(tract))
            return SkyObjs(self.skyobjs[self.skyobjs['tract'] < 0])

        if patch is not None:
            tract_mask = tract_mask & (self.skyobjs['patch'] == patch)
            if tract_mask.sum() == 0:
                if verbose:
                    warnings.warn(
                        "# Tract {0}-Patch {1} is not available!".format(tract, patch))
                return SkyObjs(self.skyobjs[self.skyobjs['tract'] < 0])

        # Number of sky objects available
        n_skyobj = tract_mask.sum()
        if n_skyobj <= n_min:
            if patch is None:
                if verbose:
                    warnings.warn("# Tract {0} has less than {1} skyobjs: {2}".format(
                        tract, n_min, n_skyobj))
            else:
                if verbose:
                    warnings.warn("# Tract {0}-Patch {1} has < {2} skyobjs: {3}".format(
                        tract, patch, n_min, n_skyobj))
            return SkyObjs(self.skyobjs[self.skyobjs['tract'] < 0])

        return SkyObjs(self.skyobjs[tract_mask])

    def select_box(self, ra1, ra2, dec1, dec2, n_min=5, verbose=True) -> 'SkyObjs':
        """Select sky objects in a box region."""
        # Order of the coordinates
        if ra1 >= ra2:
            ra1, ra2 = ra2, ra1
        if dec1 >= dec2:
            dec1, dec2 = dec2, dec1

        # Select sky objects in that region
        box_mask = ((self.skyobjs[self.ra_col] >= ra1) &
                    (self.skyobjs[self.ra_col] <= ra2) &
                    (self.skyobjs[self.dec_col] >= dec1) &
                    (self.skyobjs[self.dec_col] <= dec2))

        if box_mask.sum() == 0:
            if verbose:
                warnings.warn(
                    "# No sky object in this region: {0}:{1}-{2}:{3}".format(
                        ra1, ra2, dec1, dec2))
            return SkyObjs(self.skyobjs[self.skyobjs['tract'] < 0])

        if box_mask.sum() <= n_min:
            if verbose:
                warnings.warn("# Only find {0} sky object(s)".format(box_mask.sum()))

        return SkyObjs(self.skyobjs[box_mask])

    def select_circle(self, ra, dec, radius, n_min=5, verbose=True):
        """Select sky objects within a circle. Radius is in astropy.units."""
        from astropy.coordinates import SkyCoord
        import astropy.units as u
        if str(radius).replace('.', '', 1).isdigit():
            radius = radius * u.arcsec

        c = SkyCoord(ra, dec, frame='icrs', unit='deg')
        catalog = SkyCoord(self.skyobjs[self.ra_col],
                           self.skyobjs[self.dec_col], unit='deg', frame='icrs')
        circle_mask = (catalog.separation(c) < radius)

        if circle_mask.sum() == 0:
            if verbose:
                warnings.warn(
                    "# No sky object in this region: RA = {0}, DEC = {1}, r = {2} arcsec".format(
                        ra, dec, radius))
            return SkyObjs(self.skyobjs[self.skyobjs['tract'] < 0])

        if circle_mask.sum() <= n_min:
            if verbose:
                warnings.warn("# Only find {0} sky object(s)".format(circle_mask.sum()))

        return SkyObjs(self.skyobjs[circle_mask])

    def flux_stats(self, aper, band, rerun='s18a', sigma=3.5,
                   kde=False, bw=None, to_mujy=True, prefix=None):
        """Basic statistics of the flux."""
        u_factor = self.CGS_TO_MUJY if to_mujy else 1.0
        assert band in self.FILTER_SHORT, "# Wrong filter name: {}".format(band)

        flux_col = aper.flux(rerun=rerun, band=band)

        try:
            flux = self.skyobjs[flux_col] * u_factor
        except ValueError:
            raise Exception("# Wrong flux column name: {0}".format(flux_col))

        return utils.stats_summary(flux, sigma=sigma, n_min=self.n_min,
                                   kde=kde, bw=bw, prefix=prefix)

    def snr_stats(self, aper, band, rerun='s18a', sigma=3.5,
                  kde=False, bw=None, prefix=None):
        """Basic statistics of the S/N."""
        assert band in self.FILTER_SHORT, "# Wrong filter name: {}".format(band)

        flux_col = aper.flux(rerun=rerun, band=band)
        err_col = aper.err(rerun=rerun, band=band)

        try:
            snr = self.skyobjs[flux_col] / self.skyobjs[err_col]
        except ValueError:
            raise Exception("# Wrong column names: {0}/{1}".format(flux_col, err_col))

        return utils.stats_summary(snr, sigma=sigma, n_min=self.n_min,
                                   kde=kde, bw=bw, prefix=prefix)

    def mu_stats(self, aper, band, to_mujy=True, rerun='s18a', sigma=3.5,
                 kde=False, bw=None, prefix=None):
        """Basic statistics of the aperture flux density."""
        u_factor = self.CGS_TO_MUJY if to_mujy else 1.0
        assert band in self.FILTER_SHORT, "# Wrong filter name: {}".format(band)

        flux_col = aper.flux(rerun=rerun, band=band)

        try:
            mu = self.skyobjs[flux_col] * u_factor / aper.area_arcsec
        except ValueError:
            raise Exception("# Wrong flux column name: {0}".format(flux_col))

        return utils.stats_summary(mu, sigma=sigma, n_min=self.n_min,
                                   kde=kde, bw=bw, prefix=prefix)

    def sum_all_filters(self, aper, **kwargs):
        """Provide a summary of sky objects in all five bands."""
        aper_sum = {}
        for band in self.FILTER_SHORT:
            # Sky flux
            flux_pre = "{0}_{1}_flux".format(aper.name, band)
            flux_stats = self.flux_stats(aper, band, prefix=flux_pre, **kwargs)
            # S/N of sky flux
            snr_pre = "{0}_{1}_snr".format(aper.name, band)
            snr_stats = self.flux_stats(aper, band, prefix=snr_pre, **kwargs)
            # Surface flux density
            mu_pre = "{0}_{1}_mu".format(aper.name, band)
            mu_stats = self.flux_stats(aper, band, prefix=mu_pre, **kwargs)
            aper_sum = {**aper_sum, **flux_stats, **snr_stats, **mu_stats}

        return aper_sum

    def sum_aper_list(self, aper_list, **kwargs):
        """Summary of sky objects in all five bands for a list of apertures."""
        if isinstance(aper_list, list):
            return {key: value for stats in [
                self.sum_all_filters(aper, **kwargs) for aper in aper_list]
                    for key, value in stats.items()}
        else:
            raise TypeError("# Need a list of AperPhot objects!")

    def sum_all_tracts(self, aper_list, patch=False, verbose=True, **kwargs):
        """Provide summary for all the Tracts-(Patches) in the catalog."""
        result = []
        if not patch:
            for t in self.tract_list:
                if isinstance(aper_list, list):
                    t_sum = self.select_tract(t, verbose=verbose).sum_aper_list(
                        aper_list, **kwargs)
                    t_sum['tract'] = t
                    result.append(t_sum)
                elif isinstance(aper_list, AperPhot):
                    t_sum = self.select_tract(t, verbose=verbose).sum_all_filters(
                        aper_list, **kwargs)
                    t_sum['tract'] = t
                    result.append(t_sum)
        else:
            for t, p in [(int(tp.split('_')[0]), int(tp.split('_')[1]))
                         for tp in self.tract_patch]:
                if isinstance(aper_list, list):
                    t_sum = self.select_tract(t, patch=p, verbose=verbose).sum_aper_list(
                        aper_list, **kwargs)
                    t_sum['tract'] = t
                    t_sum['patch'] = p
                    result.append(t_sum)
                elif isinstance(aper_list, AperPhot):
                    t_sum = self.select_tract(t, patch=p, verbose=verbose).sum_all_filters(
                        aper_list, **kwargs)
                    t_sum['tract'] = t
                    t_sum['tract'] = p
                    result.append(t_sum)

        return result

    def get_summary(self, aper, band, prop, tract=None, patch=None,
                    rerun='s18a', kde=False, bw=0.2, sigma=3.0, to_mujy=True,
                    plot=False):
        """Show histogram of the properties of sky objects."""
        assert band in self.FILTER_SHORT, "# Wrong filter name: {}".format(band)
        u_factor = self.CGS_TO_MUJY if to_mujy else 1.0

        if tract is None:
            sky = self.skyobjs
        else:
            sky = self.select_tract(tract, patch=patch).skyobjs

        # Column names
        flux_col = aper.flux(rerun=rerun, band=band)
        err_col = aper.err(rerun=rerun, band=band)

        try:
            if prop == 'flux':
                values = sky[flux_col] * u_factor
            elif prop == 'snr':
                values = sky[flux_col] / sky[err_col]
            elif prop == 'mu':
                values = (sky[flux_col] * u_factor) / aper.area_arcsec
            else:
                raise Exception("# Wrong type of properties: flux/snr/mu")
        except ValueError:
            raise Exception("# Wrong flux column name: {0}".format(flux_col))

        clipped, summary = utils.stats_summary(
            values, sigma=sigma, n_min=self.n_min, kde=kde, bw=bw,
            return_clipped=True)

        if plot:
            if tract is None:
                region = None
            if tract is not None and patch is None:
                region = r'$\mathrm{Tract\ }%5d$' % tract
            elif tract is not None and patch is not None:
                region = r'${0}:{1}$'.format(tract, patch)

            aper_str = r"$\rm {0}$".format(aper.name[0].upper() + aper.name[1:])

            hist = plotting.plot_skyobj_hist(
                clipped, summary, band, prop, region=region, aper=aper_str, fontsize=20)

            return clipped, summary, hist

        return clipped, summary

    def plot_map(self, aper, band, prop, tract=None, patch=None, boxsize=0.19,
                 rerun='s18a', sigma=3.0, to_mujy=True, region=None, y_size=4,
                 margin=0.2, fontsize=30):
        """Show histogram of the properties of sky objects."""
        assert band in self.FILTER_SHORT, "# Wrong filter name: {}".format(band)
        u_factor = self.CGS_TO_MUJY if to_mujy else 1.0

        if tract is None:
            sky = self.skyobjs
        else:
            sky = self.select_tract(tract, patch=patch).skyobjs

        # Column names
        flux_col = aper.flux(rerun=rerun, band=band)
        err_col = aper.err(rerun=rerun, band=band)

        try:
            if prop == 'flux':
                values = sky[flux_col] * u_factor
            elif prop == 'snr':
                values = sky[flux_col] / sky[err_col]
            elif prop == 'mu':
                values = (sky[flux_col] * u_factor) / aper.area_arcsec
            else:
                raise Exception("# Wrong type of properties: flux/snr/mu")
        except ValueError:
            raise Exception("# Wrong flux column name: {0}".format(flux_col))

        flag = np.isfinite(values)
        values = values[flag]

        # RA, Dec
        ra, dec = sky['ra'][flag], sky['dec'][flag]

        # Number of bins
        x_bins = np.floor((np.max(ra) - np.min(ra)) / boxsize)
        y_bins = np.floor((np.max(dec) - np.min(dec)) / boxsize)

        _, low, upp = sigmaclip(values, low=sigma, high=sigma)
        mask = (values >= low) & (values <= upp)

        n_sky, x_edges, y_edges, _ = binned_statistic_2d(
            ra[mask], dec[mask], values[mask], 'count', bins=[x_bins, y_bins])

        mean_sky, _, _, _ = binned_statistic_2d(
            ra[mask], dec[mask], values[mask], 'mean', bins=[x_bins, y_bins])

        _, low_mean, upp_mean = sigmaclip(
            mean_sky[np.isfinite(mean_sky)].flatten(), low=sigma, high=sigma)

        v_edge = np.min(np.abs([low_mean, upp_mean]))

        if region is not None:
            region_str = r'$\rm {0}$'.format(region)
        else:
            region_str = ''

        band_str = r'$\ \ \ \rm {0}-band$'.format(band)
        aper_str = r"$\ \ \ \rm {0}$".format(aper.name[0].upper() + aper.name[1:])

        skyobj_map = plotting.map_skyobjs(
            x_edges, y_edges, n_sky, mean_sky,
            label=region_str + band_str + aper_str, n_min=10,
            vmin=-v_edge, vmax=v_edge, y_size=y_size, margin=margin, fontsize=fontsize)

        return x_edges, y_edges, n_sky, mean_sky, skyobj_map

import numpy,from scipy,import astropy,from astropy