import numpy as np
from torch import Tensor
from torch.autograd import Variable

class ReplayBuffer(object):
    """
    Replay Buffer for multi-agent RL with parallel rollouts
    """

    def __init__(self,
                 max_steps,
                 num_agents,
                 obs_dims,
                 ac_dims,
                 self_play=True):
        """
        Inputs:
            max_steps (int): Maximum number of timepoints to store in buffer
            num_agents (int): Number of agents in environment
            obs_dims (list of ints): number of obervation dimensions for each
                                     agent
            ac_dims (list of ints): number of action dimensions for each agent
        """
        self.max_steps = max_steps
        self.self_play = self_play
        self.num_agents = num_agents * 2 if self_play else num_agents
        obs_dims = obs_dims * 2 if self_play else obs_dims
        ac_dims = ac_dims * 2 if self_play else ac_dims

        # self.num_agents = num_agents
        self.obs_buffs = []
        self.ac_buffs = []
        self.rew_buffs = []
        self.next_obs_buffs = []
        self.done_buffs = []
        for odim, adim in zip(obs_dims, ac_dims):
            self.obs_buffs.append(np.zeros((max_steps, odim),
                                           dtype=np.float32))
            self.ac_buffs.append(np.zeros((max_steps, adim), dtype=np.float32))
            self.rew_buffs.append(np.zeros(max_steps, dtype=np.float32))
            self.next_obs_buffs.append(
                np.zeros((max_steps, odim), dtype=np.float32))
            self.done_buffs.append(np.zeros(max_steps, dtype=np.uint8))

        # index of first empty location in buffer (last index when full)
        self.filled_i = 0
        self.curr_i = 0  # current index to write to (ovewrite oldest data)

        return

    def __len__(self):
        return self.filled_i

    def push(self, observations, actions, rewards, next_observations, dones):
        # handle multiple parallel environments
        nentries = observations.shape[1]
        observations, actions, rewards, next_observations, dones = self.pack_self_play(
            observations, actions, rewards, next_observations, dones)

        if self.curr_i + nentries > self.max_steps:
            # num of indices to roll over
            rollover = self.max_steps - self.curr_i
            for agent_i in range(self.num_agents):
                self.obs_buffs[agent_i] = np.roll(self.obs_buffs[agent_i],
                                                  rollover,
                                                  axis=0)
                self.ac_buffs[agent_i] = np.roll(self.ac_buffs[agent_i],
                                                 rollover,
                                                 axis=0)
                self.rew_buffs[agent_i] = np.roll(self.rew_buffs[agent_i],
                                                  rollover)
                self.next_obs_buffs[agent_i] = np.roll(
                    self.next_obs_buffs[agent_i], rollover, axis=0)
                self.done_buffs[agent_i] = np.roll(self.done_buffs[agent_i],
                                                   rollover)
            self.curr_i = 0
            self.filled_i = self.max_steps

        for agent_i in range(self.num_agents):

            self.obs_buffs[agent_i][self.curr_i:self.curr_i +
                                    nentries] = observations[agent_i]
            # actions are already batched by agent, so they are indexed differently
            self.ac_buffs[agent_i][self.curr_i:self.curr_i +
                                   nentries] = actions[agent_i]
            self.rew_buffs[agent_i][self.curr_i:self.curr_i +
                                    nentries] = rewards[agent_i]
            self.next_obs_buffs[agent_i][self.curr_i:self.curr_i +
                                         nentries] = next_observations[agent_i]
            self.done_buffs[agent_i][self.curr_i:self.curr_i +
                                     nentries] = dones[agent_i]

        self.curr_i += nentries
        if self.filled_i < self.max_steps:
            self.filled_i += nentries
        if self.curr_i == self.max_steps:
            self.curr_i = 0

        return

    def sample(self, N, to_gpu=False, norm_rews=True):
        inds = np.random.choice(np.arange(self.filled_i), size=N, replace=True)
        if to_gpu:
            cast = lambda x: Variable(Tensor(x), requires_grad=False).cuda()
        else:
            cast = lambda x: Variable(Tensor(x), requires_grad=False)
        if norm_rews:
            ret_rews = [
                cast((self.rew_buffs[i][inds] -
                      self.rew_buffs[i][:self.filled_i].mean()) /
                     (self.rew_buffs[i][:self.filled_i].std() + 1e-7))
                for i in range(self.num_agents)
            ]
        else:
            ret_rews = [
                cast(self.rew_buffs[i][inds]) for i in range(self.num_agents)
            ]

        obs_buffs = [
            cast(self.obs_buffs[i][inds]) for i in range(self.num_agents)
        ]
        ac_buffs = [
            cast(self.ac_buffs[i][inds]) for i in range(self.num_agents)
        ]
        next_obs_buffs = [
            cast(self.next_obs_buffs[i][inds]) for i in range(self.num_agents)
        ]
        done_buffs = [
            cast(self.done_buffs[i][inds]) for i in range(self.num_agents)
        ]

        return (obs_buffs, ac_buffs, ret_rews, next_obs_buffs, done_buffs)

    def pack_self_play(self, observations, actions, rewards, next_observations,
                       dones):
        swap = np.roll(np.arange(16), 8)
        obs_packed = np.zeros((self.num_agents, 16, 112))
        nex_obs_packed = np.zeros((self.num_agents, 16, 112))

        ac_packed = actions + [actions[0][swap], actions[1][swap]]
        rew_packed = rewards + [rewards[0][swap], rewards[1][swap]]
        done_packed = dones + [dones[0][swap], dones[1][swap]]
        obs_packed[:2], obs_packed[2:] = observations, observations[:, swap, :]
        nex_obs_packed[:2] = next_observations
        nex_obs_packed[2:] = [
            next_observations[0][swap], next_observations[1][swap]
        ]

        return obs_packed, ac_packed, rew_packed, nex_obs_packed, done_packed

    def get_average_rewards(self, N):
        if self.filled_i == self.max_steps:
            inds = np.arange(self.curr_i - N,
                             self.curr_i)  # allow for negative indexing
        else:
            inds = np.arange(max(0, self.curr_i - N), self.curr_i)
        return [self.rew_buffs[i][inds].mean() for i in range(self.num_agents)]


from rlgym.utils.state_setters.wrappers import PhysicsWrapper
from rlgym.utils.gamestates import PlayerData
import numpy as np

class CarWrapper(PhysicsWrapper):

    def __init__(self, team_num: int = -1, id: int = -1, player_data: PlayerData = None):
        """
        CarWrapper constructor. Under most circumstances, users should not expect to instantiate their own CarWrapper objects.

        :param team_num: Integer indicating 0 for blue and 1 for orange.
        :param id: Integer indicating the spectator ID assigned to the car.
        :param player_data: PlayerData object for values to be copied from.

        NOTE: team_num and id will be ignored if a PlayerData object is passed.
        """
        if player_data is None:
            super().__init__()
            self.position = np.asarray([id * 100, 0, 17])
            self.rotation: np.ndarray = np.zeros(3)
            self.team_num: int = team_num
            self.id: int = id
            self.boost: float = 0
        else:
            super().__init__(phys_obj=player_data.car_data)
            self._read_from_player_data(player_data)

    def _read_from_player_data(self, player_data: PlayerData):
        """
        A function to modify CarWrapper values from values in a PlayerData object.
        """
        self.rotation = player_data.car_data.euler_angles()
        self.team_num = player_data.team_num
        self.id = player_data.car_id
        self.boost = player_data.boost_amount

    def set_rot(self, pitch: float = None, yaw: float = None, roll: float = None):
        """
        Sets rotation in terms of euler angles.
        """
        if pitch is not None:
            self.rotation[0] = pitch
        if yaw is not None:
            self.rotation[1] = yaw
        if roll is not None:
            self.rotation[2] = roll

    def _encode(self) -> list:
        """
        Function called by a StateWrapper to produce a state string.

        :return: String containing value data.
        """
        encoded = np.concatenate(((self.id,), self.position, self.linear_velocity,
                                    self.angular_velocity, self.rotation, (self.boost,)))

        return encoded.tolist()


# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
# Modifications Copyright 2017 Abigail See
#
# 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.
# ==============================================================================

"""This file contains code to build and run the tensorflow graph for the sequence-to-sequence model"""

import os
import time
import numpy as np
import tensorflow as tf
from tensorflow.contrib.tensorboard.plugins import projector
import data

FLAGS = tf.app.flags.FLAGS

class Discriminator(object):
    """A class to represent a sequence-to-sequence model for text summarization. Supports both baseline mode, pointer-generator mode, and coverage"""

    def __init__(self, hps, vocab):
        self._hps = hps
        self._vocab = vocab

    def _add_placeholders(self):
        """Add placeholders to the graph. These are entry points for any input data."""
        hps = self._hps

        # encoder part
        self._target_batch = tf.placeholder(tf.int32, [hps.batch_size* hps.max_enc_sen_num, hps.max_enc_seq_len], name='enc_batch')
        #self._target_lens = tf.placeholder(tf.int32, [hps.batch_size* hps.max_enc_sen_num], name='enc_lens')

        self._dec_batch = tf.placeholder(tf.int32, [hps.batch_size * hps.max_enc_sen_num, hps.max_enc_seq_len], name='enc_batch')
        self._dec_lens = tf.placeholder(tf.int32, [hps.batch_size * hps.max_enc_sen_num], name='enc_lens')
        #self._enc_sen_lens = tf.placeholder(tf.int32, [hps.batch_size * hps.], name='enc_sen_lens')
        self._target_mask = tf.placeholder(tf.float32,
                                            [hps.batch_size* hps.max_enc_sen_num, hps.max_enc_seq_len],
                                            name='target_mask')
        #self._enc_padding_mask = tf.placeholder(tf.float32, [hps.batch_size, None], name='enc_padding_mask')
        self._decay = tf.placeholder(tf.float32, name="decay_learning_rate")
        self.label = tf.placeholder(tf.float32, [hps.batch_size * hps.max_enc_sen_num, hps.max_enc_seq_len], name="positive_negtive")

        #self._target_batch = tf.placeholder(tf.int32,
        #                                    [hps.batch_size* hps.max_enc_sen_num],
        #                                    name='target_batch')

    def _make_feed_dict(self, batch):
        feed_dict = {}
        feed_dict[self._target_batch] = batch.target_batch
        feed_dict[self._dec_batch] = batch.dec_batch
        feed_dict[self._dec_lens] = batch.dec_sen_lens
        feed_dict[self.label] = batch.labels
        #feed_dict[self._enc_sen_lens] = batch.enc_sen_lens
        #feed_dict[self._enc_padding_mask] = batch.enc_padding_mask
        feed_dict[self._target_mask] = batch.dec_padding_mask
        #feed_dict[self.label] = batch.labels
        return feed_dict

    def _build_model(self):
        """Add the whole sequence-to-sequence model to the graph."""
        hps = self._hps
        vsize = self._vocab.size()  # size of the vocabulary

        with tf.variable_scope('discriminator'):
            # Some initializers
            self.rand_unif_init = tf.random_uniform_initializer(-hps.rand_unif_init_mag, hps.rand_unif_init_mag,
                                                                seed=123)
            self.trunc_norm_init = tf.truncated_normal_initializer(stddev=hps.trunc_norm_init_std)

            # Add embedding matrix (shared by the encoder and decoder inputs)
            with tf.variable_scope('embedding'):
                embedding = tf.get_variable('embedding', [vsize, hps.emb_dim], dtype=tf.float32,
                                            initializer=self.trunc_norm_init)

                emb_dec_inputs = tf.nn.embedding_lookup(embedding,
                                                        self._dec_batch)  # tensor with shape (batch_size, max_enc_steps, emb_size)
                self.emb_enc_inputs = emb_dec_inputs

            ## Add the encoder.
            #encoder_vector = self._add_encoder(emb_enc_inputs, self._enc_lens, hps)

            with tf.variable_scope('output_projection'):
                w = tf.get_variable('w_output', [hps.hidden_dim, vsize], dtype=tf.float32,
                                    initializer=self.trunc_norm_init)
                v = tf.get_variable('v_output', [vsize], dtype=tf.float32, initializer=self.trunc_norm_init)

            with tf.variable_scope('decoder'):
                # When decoding, use model output from the previous step
                # for the next step.
                loop_function = None

                cell = tf.contrib.rnn.LSTMCell(
                    hps.hidden_dim,
                    initializer=tf.random_uniform_initializer(-0.1, 0.1, seed=113),
                    state_is_tuple=False)

                #tf.logging.info(emb_dec_inputs)
                emb_dec_inputs = tf.unstack(emb_dec_inputs, axis=1)
                self._dec_in_state = cell.zero_state(FLAGS.batch_size* hps.max_enc_sen_num, tf.float32)
                # tf.logging.info(self._dec_in_state)
                # tf.logging.info(emb_dec_inputs)
                decoder_outputs, self._dec_out_state = tf.contrib.legacy_seq2seq.rnn_decoder(
                    emb_dec_inputs,self._dec_in_state,
                    cell, loop_function=None
                )
                decoder_outputs = tf.transpose(decoder_outputs, [1, 0, 2])

            decoder_outputs = tf.reshape(decoder_outputs,
                                         [-1,
                                          hps.hidden_dim])
            decoder_outputs = tf.nn.xw_plus_b(decoder_outputs, w, v)

            decoder_outputs = tf.reshape(decoder_outputs,
                                             [hps.batch_size * hps.max_enc_sen_num, hps.max_enc_seq_len,
                                              FLAGS.vocab_size])

            '''crossent = tf.nn.sparse_softmax_cross_entropy_with_logits(
                    labels=self._target_batch, logits=decoder_outputs)
            self.out_put = tf.argmax(crossent, 1)
            self.out_put = tf.reshape(self.out_put, [hps.batch_size, hps.max_enc_sen_num])'''
            '''weights = self._target_mask * self.label
            self.train_loss = tf.contrib.seq2seq.sequence_loss(
                decoder_outputs,
                self._target_batch,
                weights,
                average_across_timesteps=True,
                average_across_batch=True)'''

            weights = self._target_mask * self.label
            self.train_loss = tf.contrib.seq2seq.sequence_loss(
                decoder_outputs,
                self._target_batch,
                weights,
                average_across_timesteps=True,
                average_across_batch=True)

            self.out_loss = tf.contrib.seq2seq.sequence_loss(
                decoder_outputs,
                self._target_batch,
                self._target_mask,
                average_across_timesteps=False,
                average_across_batch=False)
            self.out_loss=tf.reshape(self.out_loss, [-1])
            #label=tf.reshape(self.label, [-1])
            #self.train_loss = tf.reduce_mean(self.out_loss)/(hps.batch_size*hps.max_enc_sen_num*hps.max_enc_seq_len)
            self.out_loss = tf.reshape(self.out_loss, [hps.batch_size, hps.max_enc_sen_num, hps.max_enc_seq_len])
            self.out_loss_sentence = tf.reduce_mean(self.out_loss,axis = -1)

    def _add_train_op(self):
        """Sets self._train_op, the op to run for training."""
        # Take gradients of the trainable variables w.r.t. the loss function to minimize
        loss_to_minimize =  self.train_loss
        tvars = tf.trainable_variables()
        gradients = tf.gradients(loss_to_minimize, tvars, aggregation_method=tf.AggregationMethod.EXPERIMENTAL_TREE)

        grads, global_norm = tf.clip_by_global_norm(gradients, self._hps.max_grad_norm)

        # Add a summary
        tf.summary.scalar('global_norm', global_norm)

        # Apply adagrad optimizer
        optimizer = tf.train.AdagradOptimizer(self._hps.lr, initial_accumulator_value=self._hps.adagrad_init_acc)

        self._train_op = optimizer.apply_gradients(zip(grads, tvars), global_step=self.global_step, name='train_step')

    def build_graph(self):

        """Add the placeholders, model, global step, train_op and summaries to the graph"""
        with tf.device("/gpu:" + str(FLAGS.gpuid)):
            tf.logging.info('Building graph...')
            t0 = time.time()
            self._add_placeholders()
            self._build_model()
            self.global_step = tf.Variable(0, name='global_step', trainable=False)
            self._add_train_op()
            t1 = time.time()
            tf.logging.info('Time to build graph: %i seconds', t1 - t0)

    def run_train_step(self, sess, batch, decay=False):
        """Runs one training iteration. Returns a dictionary containing train op, summaries, loss, global_step and (optionally) coverage loss."""

        feed_dict = self._make_feed_dict(batch)
        feed_dict[self._decay] = 1.0
        if decay:
            feed_dict[self._decay] = 0.001

        to_return = {
            'train_op': self._train_op,
            'loss': self.train_loss,
            'out_loss': self.out_loss,
            'global_step': self.global_step,
        }

        return sess.run(to_return, feed_dict)
    def run_pre_train_step(self, sess, batch):
        """Runs one training iteration. Returns a dictionary containing train op, summaries, loss, global_step and (optionally) coverage loss."""
        feed_dict = self._make_feed_dict(batch)
        feed_dict[self._decay] = 1.0
        to_return = {
            'train_op': self._train_op,
            'loss': self.train_loss,
            'out_loss': self.out_loss,
            'global_step': self.global_step,
        }

        return sess.run(to_return, feed_dict)

    def run_ypred_auc(self, sess, batch):
        """Runs one training iteration. Returns a dictionary containing train op, summaries, loss, global_step and (optionally) coverage loss."""
        feed_dict = self._make_feed_dict(batch)
        to_return = {
            'y_pred_auc': self.out_loss,
            'y_pred_auc_sentence': self.out_loss_sentence
        }

        return sess.run(to_return, feed_dict)

    '''def run_eval_step(self, sess, batch):
        """Runs one evaluation iteration. Returns a dictionary containing summaries, loss, global_step and (optionally) coverage loss."""
        feed_dict = self._make_feed_dict(batch)
        error_list =[]
        error_label = []
        to_return = {
            'predictions': self.out_put,
        }
        results = sess.run(to_return, feed_dict)
        right =0
        number =0
        output = results['predictions']

        for i in range(len(batch.labels)):
            if batch.target_mask[i] == 1:
                if results['predictions'][i] == batch.labels[i]:
                    right +=1
                else:
                    error_label.append(results['predictions'][i])
                    error_list.append(batch.original_reviews[i])
                number+=1
        print (batch.labels)
        print (batch.target_mask)
        print (results['predictions'])
        print (right)
        print (number)
        print (error_label)
        print (error_list)
        return right, number,error_list,error_label'''


#!/usr/bin/python

'''
Use synthetic partial columns from GEOS-Chem to obtain cloud-sliced
NO2 in the upper troposphere and compare this to UT NO2 obtained if
simply average the NO2 mixing ratios from the model over the same
pressure range (the "truth"). Both are obtained by Gaussian
weighting toward the pressure center.

GEOS-Chem partial columns are obtained over Europe, North America,
and China at the GEOS-FP meteorology native resolution (0.25x0.3125)
(latxlon) for June-August 2016-2017.

Input options to process the data include the region, the horizontal
resolution, and model simulation years.

.. code-block:: bash

    usage: ut_no2_gc_test.py [-h] [--gc_dir GC_DIR] [--out_dir OUT_DIR]
                         [--resolution RESOLUTION] [--region REGION]
                         [--strat_filter_threshold STRAT_FILTER_THRESHOLD]
                         [--start_date START_DATE]
                         [--end_date END_DATE] [-p PLOT]
                         [--do_temp_correct DO_TEMP_CORRECT]
                         [--apply_cld_frac_filter APPLY_CLD_FRAC_FILTER]
                         [--do_cld_hght_test DO_CLD_HGHT_TEST]

    optional arguments:
      -h, --help            show this help message and exit
      --gc_dir GC_DIR
      --out_dir OUT_DIR
      --resolution RESOLUTION
                            Can be 8x10, 4x5, 2x25 or 1x1
      --region REGION       Can be EU, NA, or CH
      --strat_filter_threshold STRAT_FILTER_THRESHOLD
      --start_date START_DATE
      --end_date END_DATE
      -p PLOT, --plot PLOT
      --do_temp_correct DO_TEMP_CORRECT
      --apply_cld_frac_filter APPLY_CLD_FRAC_FILTER
      --do_cld_hght_test DO_CLD_HGHT_TEST

'''

# Note for the future; most of this can probably be replaced with a modification of the GridAggregator class from
# tropomi_ut_no2

# Import relevant packages:
import numpy as np
from netCDF4 import Dataset
from scipy import stats
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
import argparse
from sklearn.linear_model import LinearRegression
import sys
import os
import datetime as dt
from dateutil import rrule as rr

# Import hack
sys.path.append(
    os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        '..'))

from uptrop.date_file_utils import get_gc_file_list
from uptrop.constants import AVOGADRO
from uptrop.constants import G
from uptrop.constants import MW_AIR
from uptrop.bootstrap import rma
from uptrop.cloud_slice_no2 import cldslice, CLOUD_SLICE_ERROR_ENUM
from uptrop.height_pressure_converter import alt2pres, pres2alt

# Turn off warnings:
np.warnings.filterwarnings('ignore')

# Define pressure range:
# These are fixed, so can be constants, as opposed to inputs.
P_MIN=180
P_MAX=450

class ProcessingException(Exception):
    pass

class CloudSliceException(Exception):
    pass

class InvalidRegionException(Exception):
    pass

class InvalidResolutionException(Exception):
    pass

class DomainIssueException(Exception):
    pass

class ProcessedData:
    """A class for comparing geoschem and tropomi data"""
    # Note for anyone reading this; this performs very similar functions to the GridAggregator class, but in a different
    # way. I prefer the method in GridAggregator, but this is fine for now. -John Roberts
    def __init__(self, region, str_res, strat_thld, do_temperature_correction=False, do_cld_frac_filter=False, do_cld_hght_test=False):
        """Creates and returns an instance of ProcessedData for a given region

        This class aggregates geoschem data over a grid defined by the :ref:`define_grid method`.

        :param region: The region to aggregate geoschem over; can be NA, EU or CH.
        :type region: str
        :param str_res: The resolution of the grid around teh region to analyse; can be 8x10, 4x5, 2x25 or 1x1
        :type str_res: str
        :param do_temperature_correction: Whether to perform the temperature correction step
        :type do_temperature_correction: bool
        :param do_cld_frac_filter: Whether to perform fractional filtering of clouds
        :type do_cld_frac_filter: bool

        :returns: A ProcessedData isntance ready to recieve data
        :rtype: ProcessedData
        """

        self.temperature_correction = do_temperature_correction
        self.cld_frac_filter = do_cld_frac_filter
        self.cloud_height_test = do_cld_hght_test

        # Filtering threshold for stratosphere:
        self.strat_filt = strat_thld

        self.define_grid(region, str_res)
        grid_shape = self.X.shape

        # Define output arrays:
        self.g_no2_vcd = np.zeros(grid_shape)
        self.g_no2_vmr = np.zeros(grid_shape)
        self.g_cld_fr = np.zeros(grid_shape)
        self.g_cld_p = np.zeros(grid_shape)
        self.g_slope_err = np.zeros(grid_shape)
        self.g_gaus_wgt = np.zeros(grid_shape)
        self.true_o3 = np.zeros(grid_shape)  # ozone mixing ratio coincident with cloud-sliced ut self
        self.true_no2 = np.zeros(grid_shape)  # "true" cloudy UT NO2
        self.g_askut_no2 = np.zeros(grid_shape)  # "true" all-sky UT NO2
        self.g_ask_gaus_wgt = np.zeros(grid_shape)  # "true" all-sky UT NO2 weights
        self.g_as_cnt = np.zeros(grid_shape)  # Count all-sky
        self.g_cnt = np.zeros(grid_shape)

        # Type of data loss in the cloud-slicing retrieval:
        self.loss_count = {
            "too_few_points": 0,
            "low_cloud_height_range": 0,
            "low_cloud_height_std": 0,
            "large_error": 0,
            "sig_diff_from_zero": 0,
            "no2_outlier": 0,
            "non_uni_strat": 0,
        }

        # Initialize:
        self.cloud_slice_count = 0
        self.maxcnt = 0
        self.grad_retain = 0
        self.grad_remove = 0
        # Print out min and max cloud-sliced UT NO2 error size:
        self.minerr = 1   # Set to maximum possible error (100%)
        self.maxerr = 0

        # Define string to represent the layer range:
        self.prange = str(P_MIN) + '-' + str(P_MAX)

        # Define factor to convert slope of NO2 mixing ratio versus pressure
        # to VMR:
        self.den2mr = np.divide((np.multiply(G, MW_AIR)), AVOGADRO)

    def define_grid(self, region, str_res):
        """Defines a grid based on a set of regional bounding boxes and a resolution

        :param region: The region to aggregate geoschem over; can be NA, EU or CH.
        :type region: str
        :param str_res: The resolution of the grid around the region to analyse; can be 8x10, 4x5, 2x25 or 1x1
        :type str_res: str
        """
        # Define target grid:
        if region == 'NA':
            self.minlat = 2.
            self.maxlat = 62.
            self.minlon = -140.
            self.maxlon = -55.
        elif region == 'EU':
            self.minlat = 26.
            self.maxlat = 66.
            self.minlon = -25.
            self.maxlon = 45.
        elif region == 'CH':
            self.minlat = 6.
            self.maxlat = 62.
            self.minlon = 60.
            self.maxlon = 140.
        else:
            print("Invalid region; valid regions are 'NA','EU','CH'.")
            raise InvalidRegionException
        # NOTE FOR LATER: This snippet turns up in fresco_cld_err; a candidate for the library.
        # Define grid information:
        if str_res == '8x10':
            self.dellat, self.dellon = 8, 10
        elif str_res == '4x5':
            self.dellat, self.dellon = 4, 5
        elif str_res == '2x25':
            self.dellat, self.dellon = 2, 2.5
        elif str_res == '1x1':
            self.dellat, self.dellon = 1, 1
        else:
            print("Invalid resolution: valid resolutions are 8x10, 4x5, 2x25 (two pt 5) and 1x1")
            raise InvalidResolutionException
        self.out_lon = np.arange(self.minlon, self.maxlon + self.dellon, self.dellon)
        self.out_lat = np.arange(self.minlat, self.maxlat + self.dellat, self.dellat)
        # Convert output lats and long to 2D:
        self.X, self.Y = np.meshgrid(self.out_lon, self.out_lat, indexing='ij')
        # Dimensions of output data:
        self.xdim = len(self.out_lon)
        self.ydim = len(self.out_lat)
        # Get maximum possible number of native resolution model gridsquares
        # in the coarser grid (to use in process_grid_square to check that
        # don't exceed this):
        self.max_limit=(self.dellon/0.3125)*(self.dellat/0.25)

    # ----Processing methods----
    def process_geoschem_day(self, file_path):
        """Aggregates geoschem data into the processing grid for a file on a given day.

        Calls :ref:`regrid_and_process` for each pixel in the geoschem file, then calls
        :ref:`process_grid_square` for each square of the grid

        :param file_path: Path to the geoschem file to load into the processing grid
        :type file_path: str"""
        # Define output data for this day:
        out_shape = (self.xdim, self.ydim)  # This feel gross. A 3-d list of appendable lists.
        self.g_no2 = [[[] for n in range(self.ydim)] for m in range(self.xdim)]
        self.g_grad = [[[] for n in range(self.ydim)] for m in range(self.xdim)]
        self.g_o3 = [[[] for n in range(self.ydim)] for m in range(self.xdim)]
        self.all_cld_fr = [[[] for n in range(self.ydim)] for m in range(self.xdim)]
        self.strat_no2 = [[[] for n in range(self.ydim)] for m in range(self.xdim)]
        self.all_cld_p = [[[] for n in range(self.ydim)] for m in range(self.xdim)]
        self.g_true_no2 = [[[] for n in range(self.ydim)] for m in range(self.xdim)]

        this_geoschem_day = GeosChemDay(file_path,
                                        temperature_correction=self.temperature_correction,
                                        cloud_height_test=self.cloud_height_test)
        # Get column values:
        for y in range(len(this_geoschem_day.t_lat)):
            for x in range(len(this_geoschem_day.t_lon)):
                this_geoschem_day.prepare_no2_pixel(x, y)
                self.regrid_and_process(x, y, this_geoschem_day)
        for i in range(self.xdim):
            for j in range(self.ydim):
                if any(value != 0 for value in self.g_no2[i][j]):
                    self.process_grid_square(i, j)

    def regrid_and_process(self, x, y, gc_data):
        """Inserts data from pixel x,y from geoschem data into the appropriate grid cell

        :param x: X index of geoschem grid cell
        :type x: int
        :param y: Y index of geoschem grid cell
        :type y: int
        :param gc_data: An instance of GeosChemDay data
        :type gc_data: GeosChemDay
        """
        # If no valid data, skip:
        if ( len(gc_data.askind) == 0):
            return

        # Find nearest gridsquare in output grid:
        lon = int(np.argmin(abs(self.out_lon - gc_data.t_lon[x])))
        lat = int(np.argmin(abs(self.out_lat - gc_data.t_lat[y])))

        self.g_ask_gaus_wgt[lon, lat] += np.sum(gc_data.twgt)

        self.g_askut_no2[lon, lat] += np.sum(gc_data.t_gc_no2[gc_data.askind, y, x] * gc_data.twgt * 1e3)
        self.g_as_cnt[lon, lat] += 1.0

        # Add relevant data to array if there are clouds at 450-180 hPa:
        if (gc_data.lcld < gc_data.level_min) or (gc_data.lcld > gc_data.level_max):
            # Check for clouds between min and max pressure. If none, move to next pixel.
            #print("Cloud top outside pressure range in pixel {},{}".format(x,y))
            return

        # Find where cloud fraction in UT exceeds 0.7 after calculating
        # true all-sky NO2:
        # (Keep for testing effect of thick clouds on cloud-sliced UT NO2):
        if ( self.cld_frac_filter ):
            if np.sum(gc_data.t_cld_fr[gc_data.level_min:gc_data.level_max + 1, y, x]) < 0.7:
                return

        self.g_no2[lon][lat].append(gc_data.no2_2d)
        self.g_grad[lon][lat].append(gc_data.grad)
        self.strat_no2[lon][lat].append(gc_data.strat_col)
        self.all_cld_p[lon][lat].append(gc_data.t_cld_hgt[y, x])
        self.g_o3[lon][lat].append(
            np.mean(gc_data.t_gc_o3[gc_data.level_min:gc_data.level_max + 1, y, x]))
        self.g_true_no2[lon][lat].append(
            np.mean(gc_data.t_gc_no2[gc_data.level_min:gc_data.level_max + 1, y, x]))
        self.all_cld_fr[lon][lat].append(
            np.sum(gc_data.t_cld_fr[gc_data.level_min:gc_data.level_max + 1, y, x]))
        pass

    def process_grid_square(self, i, j):

        # Define vectors of relevant data:
        # These should all have the same length
        t_col_no2 = np.asarray(self.g_no2[i][j],dtype=np.float)
        t_strat_no2 = np.asarray(self.strat_no2[i][j],dtype=np.float)
        t_fr_c = np.asarray(self.all_cld_fr[i][j],dtype=np.float)
        t_cld = np.asarray(self.all_cld_p[i][j],dtype=np.float)
        t_mr_no2 = np.asarray(self.g_true_no2[i][j],dtype=np.float)
        t_o3 = np.asarray(self.g_o3[i][j],dtype=np.float)
        t_grad_no2 = np.asarray(self.g_grad[i][j],dtype=np.float)

        # Skip if fewer than 10 points:
        if len(t_col_no2) < 10:
            #print("Grid square {}, {} failed with condition too_few_points".format(i, j))
            self.loss_count["too_few_points"] += 1
            return

        # Remove non-uniform stratosphere:
        if (np.std(t_strat_no2) / np.mean(t_strat_no2)) > self.strat_filt:
            self.loss_count["non_uni_strat"] += 1
            return  # continue

        # Get number of points:
        n_pnts = len(t_cld)
        if n_pnts > self.maxcnt:
            self.maxcnt = n_pnts
            print('Max no. of points in a cluster: ', self.maxcnt, flush=True)

        # Check if number of gridsquares exceeds max possible
        # (indicates error in coarser grid domain):
        if ( self.maxcnt > self.max_limit ):
            print('No. of model grids exceeds max possible')
            raise DomainIssueException

        # Use cloud_slice_ut_no2 function to get cloud-sliced UT NO2.
        # Treat data differently depending on whether there are 10-99 points:
        if n_pnts >= 10 and n_pnts <100:
            self.add_slice(i,j,t_cld,t_col_no2, t_mr_no2, t_fr_c, t_grad_no2, t_o3)
        # Or at least 100 points:
        elif n_pnts >= 100:
            num_slices = 40
            stride = round(n_pnts / num_slices)
            nloop = list(range(stride))
            for w in nloop:
                subset_t_col_no2 = t_col_no2[w::stride]
                subset_t_cld = t_cld[w::stride]
                subset_t_mr_no2 = t_mr_no2[w::stride]
                subset_t_fr_c = t_fr_c[w::stride]
                subset_t_grad_no2 = t_grad_no2[w::stride]
                subset_t_o3 = t_o3[w::stride]
                self.add_slice(i, j, subset_t_cld, subset_t_col_no2, subset_t_mr_no2, subset_t_fr_c, subset_t_grad_no2, subset_t_o3)

    def add_slice(self, i, j, t_cld, t_col_no2, t_mr_no2, t_fr_c, t_grad_no2, t_o3):
        """Extracts the upper troposphere gc_data, gc_data error, ozone data, ozone error,
         and mean cloud pressure for grid square [i,j]

        This method uses the cloud-slicing function :ref:`uptrop.cloud_slice_ut_no2.cldslice`
        Once calculated, the a weighting is derived from cloud pressure.
        The weighted upper tropospheric gc_data and error is added to the rolling total for this season.
        If the cloud slicing fails, then the reason is added to loss_count for the end report.

        :param i: X-index of grid square
        :type i: int
        :param j: Y-index of grid square
        :type j: int
        :param t_cld: Tropospheric cloud
        :type t_cld: float
        :param t_col_no2: Total (stratosphere+troposphere) column no2
        :type t_col_no2: float
        :param t_mr_no2: NO2 mixing ratio between pmax and pmin
        :type t_mr_no2: float
        :param t_fr_c: Cloud fraction between pmax and pmin
        :type t_fr_c: float
        :param t_grad_no2: NO2 gradient between pmax and pmin in pptv/hPa
        :type t_grad_no2: float
        :param t_o3: O3 mixing ratio to test influence from stratosphere
        :type t_o3: float
        """
        utmrno2, utmrno2err, stage_reached, mean_cld_pres = cldslice(t_col_no2, t_cld,140)

        # Skip if approach didn't work (i.e. cloud-sliced UT NO2 is NaN):
        # Drop out after the reason for data loss is added to loss_count.
        if np.isnan(utmrno2) or np.isnan(utmrno2err):
            # Cloud-slicing led to nan, but due to rma regression rather
            # than data filtering (these are rare):
            if ( stage_reached==0 ):
                print("Cloud-sliced NO2 NAN for pixel i:{} j:{}".format(i, j))
                return
            self.loss_count[CLOUD_SLICE_ERROR_ENUM[stage_reached]] += 1
            # Track non-uniform NO2 removed:
            grad_ind=np.where(np.abs(t_grad_no2) >= 0.33)[0]
            self.grad_remove += len(grad_ind)
            #print("Cloud-slice exception {} in pixel i:{} j:{}".format(
            #    CLOUD_SLICE_ERROR_ENUM[stage_reached], i, j))
        else:
            # Calculate weights:
            #err_wgt = 1.0 / (utmrno2err ** 2)  # not use, but preserve anyway
            gaus_wgt = np.exp((-(mean_cld_pres - 315) ** 2) / (2 * 135 ** 2))

            # Get and track error range:
            nerr = utmrno2err / utmrno2
            if nerr > self.maxerr:
                self.maxerr = nerr
                print('Max rel. cloud-sliced NO2 error: ', self.maxerr, flush=True)
            if nerr < self.minerr:
                self.minerr = nerr
                print('Min rel. cloud-sliced NO2 error: ', self.minerr, flush=True)

            # Track non-uniform NO2 retained:
            grad_ind=np.where(np.abs(t_grad_no2) >= 0.33)[0]
            self.grad_retain += len(grad_ind)

            # Weighted mean for each pass of cldslice:
            weight_mean = np.mean(t_mr_no2 * 1e3) * gaus_wgt
            self.true_no2[i, j] += weight_mean
            self.true_o3[i, j] += (np.mean(t_o3) * gaus_wgt)
            self.g_no2_vmr[i, j] += utmrno2 * gaus_wgt
            self.g_slope_err[i, j] += utmrno2err * gaus_wgt
            self.g_gaus_wgt[i, j] += gaus_wgt
            self.g_cnt[i, j] += 1
            self.g_cld_fr[i, j] += np.mean(t_fr_c)
            self.cloud_slice_count += 1

    def get_weighted_mean(self):
        """Applies weighting to the aggregated means.
        """
        # Mean physical parameters:
        self.g_no2_vmr = np.divide(self.g_no2_vmr, self.g_gaus_wgt, where=self.g_cnt != 0)
        self.g_cld_fr = np.divide(self.g_cld_fr, self.g_cnt, where=self.g_cnt != 0)
        self.true_no2 = np.divide(self.true_no2, self.g_gaus_wgt, where=self.g_cnt != 0)
        self.g_cld_p = np.divide(self.g_cld_p, self.g_cnt, where=self.g_cnt != 0)
        self.true_o3 = np.divide(self.true_o3, self.g_gaus_wgt, where=self.g_cnt != 0)
        self.g_askut_no2 = np.divide(self.g_askut_no2, self.g_ask_gaus_wgt, where=self.g_as_cnt != 0)
        # Mean cloud-sliced UT NO2 error:
        self.g_slope_err = np.divide(self.g_slope_err, self.g_gaus_wgt, where=self.g_cnt != 0)
        # Mean weights:
        self.g_gaus_wgt = np.divide(self.g_gaus_wgt, self.g_cnt, where=self.g_cnt != 0)
        # No data (nan):
        self.true_no2[self.g_cnt == 0] = np.nan
        self.true_o3[self.g_cnt == 0] = np.nan
        self.g_slope_err[self.g_cnt == 0] = np.nan
        self.g_gaus_wgt[self.g_cnt == 0] = np.nan
        self.g_askut_no2[self.g_as_cnt == 0] = np.nan
        self.g_no2_vmr[self.g_cnt == 0] = np.nan
        self.g_cld_fr[self.g_cnt == 0] = np.nan
        self.g_cld_p[self.g_cnt == 0] = np.nan

    # ----Reporting and saving methods----
    def print_data_report(self):
        """Prints a set of reasons for missing data
        """
        # Output code diagnostics:
        # No. of data points:
        # print('No. of valid data points: ',cloud_slice_count,flush=True)
        print('Max no. of data points in a gridsquare: ', np.amax(self.g_cnt), flush=True)
        # Track reasons for data loss:
        print('(1) Too few points: ', self.loss_count["too_few_points"], flush=True)
        print('(2) Low cloud height range: ', self.loss_count["low_cloud_height_range"], flush=True)
        print('(3) Low cloud height std dev: ', self.loss_count["low_cloud_height_std"], flush=True)
        print('(4) Large error: ', self.loss_count["large_error"], flush=True)
        print('(5) Significantly less than zero: ', self.loss_count["sig_diff_from_zero"], flush=True)
        print('(6) Outlier (NO2 > 200 pptv): ', self.loss_count["no2_outlier"], flush=True)
        print('(7) Non-uniform stratosphere: ', self.loss_count["non_uni_strat"], flush=True)
        print('(8) Successful retrievals: ', self.cloud_slice_count, flush=True)
        print('(9) Total possible points: ', (sum(self.loss_count.values()) + self.cloud_slice_count), flush=True)
        print('Non-uniform NO2 retained = ', self.grad_retain)
        print('Non-uniform NO2 removed = ', self.grad_remove)
        print('Error range of cloud-sliced UT NO2: ', np.nanmin(self.g_slope_err), np.nanmax(self.g_slope_err))

    def plot_data(self):
        """Plots the present state of the gridded data
        """
        # ===> FUTURE UPDATE : change from basemap to cartopy <===
        # Plot the data:
        m = Basemap(resolution='l', projection='merc',
                    lat_0=0, lon_0=0, llcrnrlon=self.minlon,
                    llcrnrlat=self.minlat, urcrnrlon=self.maxlon, urcrnrlat=self.maxlat)
        xi, yi = m(self.X, self.Y)
        plt.subplot(2, 3, 1)
        cs = m.pcolor(xi, yi, np.squeeze(self.g_no2_vmr), vmin=0, vmax=80, cmap='jet')
        m.drawparallels(np.arange(-80., 81., 45.), labels=[1, 0, 0, 0], fontsize=8)
        m.drawmeridians(np.arange(-180., 181., 45.), labels=[0, 0, 0, 1], fontsize=8)
        m.drawcoastlines()
        m.drawcountries()
        cbar = m.colorbar(cs, location='bottom', pad="10%")
        plt.title('Cloud-sliced NO2 VMRs')
        plt.subplot(2, 3, 2)
        cs = m.pcolor(xi, yi, np.squeeze(self.true_no2), vmin=0.0, vmax=80, cmap='jet')
        m.drawparallels(np.arange(-80., 81., 45.), labels=[1, 0, 0, 0], fontsize=8)
        m.drawmeridians(np.arange(-180., 181., 45.), labels=[0, 0, 0, 1], fontsize=8)
        m.drawcoastlines()
        m.drawcountries()
        cbar = m.colorbar(cs, location='bottom', pad="10%")
        plt.title('True cloudy NO2')
        plt.subplot(2, 3, 3)
        cs = m.pcolor(xi, yi, np.squeeze(self.g_askut_no2), vmin=0, vmax=80, cmap='jet')
        m.drawparallels(np.arange(-80., 81., 45.), labels=[1, 0, 0, 0], fontsize=8)
        m.drawmeridians(np.arange(-180., 181., 45.), labels=[0, 0, 0, 1], fontsize=8)
        m.drawcoastlines()
        m.drawcountries()
        cbar = m.colorbar(cs, location='bottom', pad="10%")
        plt.title('True NO2 VMRs under all-sky conditions')
        plt.subplot(2, 3, 4)
        plt.plot(self.true_no2, self.g_no2_vmr, 'o', color='black', markersize=6)
        r = stats.pearsonr(self.true_no2[~np.isnan(self.g_no2_vmr)],
                           self.g_no2_vmr[~np.isnan(self.g_no2_vmr)])
        # print('Correlation = ', r[0])
        result = rma(self.true_no2[~np.isnan(self.g_no2_vmr)], self.g_no2_vmr[~np.isnan(self.g_no2_vmr)],
                     len(self.true_no2[~np.isnan(self.g_no2_vmr)]), 1000)
        print(result, flush=True)
        xvals = np.arange(0, 100, 5)
        yvals = result[1] + xvals * result[0]
        plt.plot(xvals, yvals, '-')
        plt.xlim(-4, 80)
        plt.ylim(-4, 80)
        plt.xlabel('True NO2 (cloudy)')
        plt.ylabel('Cloud-sliced NO2')
        print('===== True (cloudy) vs cloud-sliced UT NO2 ====')
        print('R = ', r[0], flush=True)
        print('Slope = ', result[0])
        print('Slope Err = ', result[2], flush=True)
        print('Intercept = ', result[1], flush=True)
        print('Intercept Err = ', result[3], flush=True)
        add2plt = ("y = {a:6.2f}x + {b:6.3f}".format(a=result[0], b=result[1]))
        plt.text(2, 75, add2plt, fontsize=8,
                 ha='left', va='center')  # , transform=ax.transAxes)
        add2plt = ("R = {a:6.2f}".format(a=r[0]))
        plt.text(2, 65, add2plt, fontsize=8,
                 ha='left', va='center')  # , transform=ax.transAxes)
        plt.subplot(2, 3, 5)
        plt.plot(self.g_askut_no2, self.g_no2_vmr, 'o', color='black', markersize=6)
        r = stats.pearsonr(self.g_askut_no2[(~np.isnan(self.g_no2_vmr)) & (~np.isnan(self.g_askut_no2))],
                           self.g_no2_vmr[(~np.isnan(self.g_no2_vmr)) & (~np.isnan(self.g_askut_no2))])
        result = rma(self.g_askut_no2[(~np.isnan(self.g_no2_vmr)) & (~np.isnan(self.g_askut_no2))],
                     self.g_no2_vmr[(~np.isnan(self.g_no2_vmr)) & (~np.isnan(self.g_askut_no2))],
                     len(self.g_askut_no2[(~np.isnan(self.g_no2_vmr)) & (~np.isnan(self.g_askut_no2))]),
                     1000)
        xvals = np.arange(0, 100, 5)
        yvals = result[1] + xvals * result[0]
        plt.plot(xvals, yvals, '-')
        plt.xlim(-4, 80)
        plt.ylim(-4, 80)
        plt.xlabel('True NO2 (all-sky)')
        plt.ylabel('Cloud-sliced NO2')
        add2plt = ("y = {a:6.2f}x + {b:6.3f}". \
                   format(a=result[0], b=result[1]))
        plt.text(2, 75, add2plt, fontsize=8, \
                 ha='left', va='center')  # , transform=ax.transAxes)
        add2plt = ("R = {a:6.2f}".format(a=r[0]))
        plt.text(2, 65, add2plt, fontsize=8, \
                 ha='left', va='center')  # , transform=ax.transAxes)
        plt.subplot(2, 3, 6)
        plt.plot(self.g_askut_no2, self.true_no2, 'o', color='black', markersize=6)
        r = stats.pearsonr(self.g_askut_no2[(~np.isnan(self.true_no2)) & (~np.isnan(self.g_askut_no2))],
                           self.true_no2[(~np.isnan(self.true_no2)) & (~np.isnan(self.g_askut_no2))])
        result = rma(self.g_askut_no2[(~np.isnan(self.true_no2)) & (~np.isnan(self.g_askut_no2))],
                     self.true_no2[(~np.isnan(self.true_no2)) & (~np.isnan(self.g_askut_no2))],
                     len(self.g_askut_no2[(~np.isnan(self.true_no2)) & (~np.isnan(self.g_askut_no2))]),
                     1000)
        xvals = np.arange(0, 100, 5)
        yvals = result[1] + xvals * result[0]
        plt.plot(xvals, yvals, '-')
        plt.xlim(-4, 80)
        plt.ylim(-4, 80)
        plt.xlabel('True NO2 (all-sky)')
        plt.ylabel('True NO2 (cloudy)')
        add2plt = ("y = {a:6.2f}x + {b:6.3f}".format(a=result[0], b=result[1]))
        plt.text(2, 75, add2plt, fontsize=8, ha='left', va='center')
        add2plt = ("R = {a:6.2f}".format(a=r[0]))
        plt.text(2, 65, add2plt, fontsize=8, ha='left', va='center')
        plt.show()

    def save_to_netcdf(self, out_path):
        """Saves the present state of the grid to a netCDF4 file

        :param out_path: Path to the output file
        :type out_path: str
        """
        # Save the data to NetCDF:
        ncout = Dataset(out_path, mode='w', format='NETCDF4')
        # Create data dimensions:
        ncout.createDimension('lat', self.ydim)
        ncout.createDimension('lon', self.xdim)
        # create lon axis:
        lon = ncout.createVariable('lon', np.float32, ('lon',))
        lon.units = 'degrees_east'
        lon.long_name = 'longitude'
        lon[:] = self.out_lon
        # create lat axis:
        lat = ncout.createVariable('lat', np.float32, ('lat',))
        lat.units = 'degrees_north'
        lat.long_name = 'latgitude'
        lat[:] = self.out_lat
        # create data axes:
        # (1) Cloud-sliced NO2:
        csutno2 = ncout.createVariable('csutno2', np.float32, ('lon', 'lat'))
        csutno2.units = 'pptv'
        csutno2.long_name = 'UT NO2 mixing ratio (180-450 hPa) obtained using cloud-slicing'
        csutno2[:] = self.g_no2_vmr
        # (2a) Double-weighting error:
        utdblerr = ncout.createVariable('utdblerr', np.float32, ('lon', 'lat'))
        utdblerr.units = 'pptv'
        utdblerr.long_name = 'Standard error of the NO2 mixing ratios in the UT (180-450 hPa) obtained using cloud-slicing'
        utdblerr[:] = self.g_slope_err
        # (2b) Gaussian-weighting error:
        utgauserr = ncout.createVariable('utgauserr', np.float32, ('lon', 'lat'))
        utgauserr.units = 'pptv'
        utgauserr.long_name = 'Standard error of the NO2 mixing ratios in the UT (180-450 hPa) obtained using cloud-slicing'
        utgauserr[:] = self.g_gaus_wgt
        # (3) Number of observations in each gridsquare:
        nobs = ncout.createVariable('nobs', np.float32, ('lon', 'lat'))
        nobs.units = 'unitless'
        nobs.long_name = 'Number of observations in each gridsquare used to obtain cloud-sliced UT NO2 mixing ratios'
        nobs[:] = self.g_cnt
        # (4) Mean cloud pressure for season between 450-180 hPa:
        utcld = ncout.createVariable('utcld', np.float32, ('lon', 'lat'))
        utcld.units = 'hPa'
        utcld.long_name = 'Mean cloud pressure between 450 and 180 hPa'
        utcld[:] = self.g_cld_p
        # (5) Mean NO2 mixing ratio at 450-180 hPa for scenes with clouds:
        cldutno2 = ncout.createVariable('cldutno2', np.float32, ('lon', 'lat'))
        cldutno2.units = 'pptv'
        cldutno2.long_name = 'UT NO2 mixing ratio (180-450 hPa) obtained if clouds are present'
        cldutno2[:] = self.true_no2
        # (6) Mean NO2 mixing ratio at 450-180 hPa under all conditions (all-sky):
        askutno2 = ncout.createVariable('askutno2', np.float32, ('lon', 'lat'))
        askutno2.units = 'pptv'
        askutno2.long_name = 'UT NO2 mixing ratio (180-450 hPa) obtained under all conditions (all-sky)'
        askutno2[:] = self.g_askut_no2
        # (7) Cloud fraction:
        utcldfrc = ncout.createVariable('utcldfrc', np.float32, ('lon', 'lat'))
        utcldfrc.units = 'unitless'
        utcldfrc.long_name = 'GEOS-FP cloud fraction obtained as sum of 3D cloud fractions across range of interest (180-450 hPa)'
        utcldfrc[:] = self.g_cld_fr
        # (8) O3 sampled coincident with cloud-slicing retrieval:
        uto3 = ncout.createVariable('uto3', np.float32, ('lon', 'lat'))
        uto3.units = 'ppbv'
        uto3.long_name = 'GEOS-Chem ozone obtained coincident with cloud-sliced NO2'
        uto3[:] = self.true_o3
        # Close the file:
        ncout.close()

class GeosChemDay:
    """A class for reading, preprocessing and accessing Geoschem data on a given day
    """
    def __init__(self, file_path, temperature_correction=False, cloud_height_test=False):
        """Reads the data at file_path and returns a GeosChemDay object containing that data

        :param file_path: Path to the netcdf4 file containing the GeosChem data
        :type file_path: str
        :param temperature_correction: Whether to apply temperature correction
        :type temperature_correction: bool
        :param cloud_height_test: Whether to test effect of systematic underestimate in cloud height
        :type cloud_height_test: bool

        :returns: A GeosChemDay class
        :rtype: GeosChemDay
        """
        print('File path: ',file_path, flush=True)

        self.temperature_correction = temperature_correction
        self.cloud_height_test = cloud_height_test

        # Read dataset:
        fh = Dataset(file_path, mode='r')
        # Extract data of interest:
        tlon, tlat, tgcno2, tcldfr, tcldhgt, tadn, tbxhgt, tpedge, tpause, tgco3, tdegk = \
            fh.variables['LON'], fh.variables['LAT'], \
            fh.variables['IJ-AVG-S__NO2'], fh.variables['TIME-SER__CF'], \
            fh.variables['TIME-SER__CThgt'], fh.variables['TIME-SER__AIRDEN'], \
            fh.variables['BXHGHT-S__BXHEIGHT'], fh.variables['PEDGE-S__PSURF'], \
            fh.variables['TR-PAUSE__TP-PRESS'], fh.variables['IJ-AVG-S__O3'], \
            fh.variables['DAO-3D-S__TMPU']
        self.t_lon = tlon[:]
        self.t_lat = tlat[:]
        self.t_gc_no2 = tgcno2[:]
        self.t_cld_fr = tcldfr[:]
        self.t_cld_hgt = tcldhgt[0, :, :]
        self.t_adn = tadn[:]  # in molec/cm3
        self.t_bx_hgt = tbxhgt[:]
        self.t_p_edge = tpedge[:]
        self.t_pause = tpause[0, :, :]
        self.t_gc_o3 = tgco3[:]
        self.t_deg_k = tdegk[:]
        # Convert box height from m to cm:
        self.t_bx_hgt = self.t_bx_hgt * 1e2

        if self.cloud_height_test:
            # Lower cloud heights by 1 km to roughly mimic lower altitude
            # clouds retrieved for TROPOMI assuming clouds are reflective
            # boundaries with uniform reflectivity:
            # Calculate cloud top height in m:
            t_cld_hgt = pres2alt(self.t_cld_hgt*1e2)
            # Lower the clouds by 1 km (1000 m) (this won't work for low-altitude
            # clouds):
            t_cld_hgt = t_cld_hgt - 1e3
            # Convert back to Pa and convert that to hPa:
            self.t_cld_hgt = alt2pres(t_cld_hgt)*1e-2

        #Get outputs ready here for tidyness:
        self.no2_2d = None
        self.trop_col = None
        self.strat_col = None
        self.gcutno2 = None
        self.gascnt = None
        self.grad = None

        self.level_min = None
        self.level_max = None
        self.askind = None

    def prepare_no2_pixel(self, x, y):
        """Extracts preprocesed no2 from the geoschem pixel at x,y

        :param x: The x index of the pixel
        :type x: int
        :param y: The y index of the pixel
        :type y: int
        """
        # Calculate corresponding mid-pressure values:
        tp_mid = np.zeros(len(self.t_p_edge[:, y, x]))
        # Get mid-pressure values, except for highest layer:
        for k in range(len(self.t_p_edge[:, y, x]) - 1):
            tp_mid[k] = np.multiply(0.5, (self.t_p_edge[k, y, x] + self.t_p_edge[k + 1, y, x]))
        # Estimate mid-pressure for highest value (doesn't need to
        # be accurate, as surpassing the range of interset):
        # Data output from the model includes 47 vertical layers. This means that only 46 pressure centres can be calculated as the calculation requires pressure edges.
        tp_mid[46] = np.multiply(0.5, (self.t_p_edge[46, y, x] + (self.t_p_edge[46, y, x] - 0.1)))
        # Get model layer of tropopause:
        tppind = np.argmin(abs(tp_mid - self.t_pause[y, x]))
        # Get indices that fall between 450 and 180 hPa for estimating
        # "true' all-sky UT NO2 and partial columns:
        lind = np.where((tp_mid >= P_MIN) & (tp_mid <= P_MAX))[0]
        # Get UT NO2 under "true" all-sky conditions:
        # Make sure this is below the tropopause:
        # If below tropopause, use full extent (180-450 hPa):
        if lind[len(lind) - 1] <= tppind:
            self.askind = lind
        # If above tropopause, trim to tropopause-450 hPa:
        if lind[len(lind) - 1] > tppind:
            self.askind = lind[np.where(lind <= tppind)[0]]
        # If tropopause below 450 hPa, skip entirely:
        if self.t_pause[y, x] > P_MAX:
            #print("Tropopause less than P_MAX in geoschem pixel x:{}, y:{}".format(x,y))
            return  # continue

        # Get Guassian weights that allocate higher weights to points
        # closest to the pressure centre (315 hPa):
        # Equation is:
        #   w = exp(-(p-315)^2/2*135^2 ) where 315 hPa is the centre and
        #         135 hPa is the standard deviation.
        self.twgt = np.exp((-(tp_mid[self.askind] - 315) ** 2) / (2 * 135 ** 2))

        # Get model level of cloud top height closest to lowest
        # pressure-altitude of interest (P_MIN):
        self.lcld = np.argmin(abs(self.t_cld_hgt[y, x] - tp_mid))
        # Skip if cloud top height ouside pressure range of interest:
        self.level_min, self.level_max = np.amin(lind), np.amax(lind)

        if (self.temperature_correction):

            # Equation is from the TROPOMI product ATBD (p. 32, Eqn 18)
            # (product document abbrevation: S5P-KNMI-L2-0005-RP)
            self.temp_corr = 1 - (3.16e-3 * (self.t_deg_k[self.level_min:, y, x] - 220.)) + \
                        (3.39e-6 * ((self.t_deg_k[self.level_min:, y, x] - 220) ** 2))
        else:
            # Set to 1 so that no scaling is applied:
            # (might be a more eloquent way to do this)
            self.temp_corr = np.ones(len(self.t_gc_no2[self.level_min:, y, x]))

        # Calculate NO2 gradient:
        regr = LinearRegression()
        # Pressure (hPa)
        x_vals=tp_mid[self.level_min:self.level_max].reshape(-1,1)
        # NO2 (ppbv)
        y_vals=self.t_gc_no2[self.level_min:self.level_max, y, x].reshape(-1,1)
        # NO2 (pptv)
        y_vals=y_vals*1e3
        # Perform regression:
        regr.fit(x_vals,y_vals)
        # Define gradient from regression slope (pptv/hPa):
        self.grad = regr.coef_

        # Get partial NO2 column in molec/m2 from cloud top height
        # to highest model level (output up to level 47):
        # print(t_gc_no2[self.level_min:tppind,y,x])
        # print(t_gc_no2[self.level_min:tppind,y,x]*1.5)
        self.no2_2d = np.sum(self.t_gc_no2[self.level_min:, y, x]
                             * 1e-5
                             * self.temp_corr
                             * self.t_adn[self.level_min:, y, x]
                             * self.t_bx_hgt[self.level_min:, y, x])
        # Get stratospheric column from 180 hPa aloft:
        # Previous approach (remove when model simulations done):
        # tppind=np.where(tpmid<180.)[0]
        self.strat_col = np.sum(self.t_gc_no2[tppind:, y, x]
                                * 1e-5 * self.t_adn[tppind:, y, x]
                                * self.t_bx_hgt[tppind:, y, x])

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    # Shorten directory name to up to "GC/", then define the subdirectory
    # as 'geosfp' + dirreg + 'iccw/' in get_file_list.
    # This is now done in get_gc_file_list
    parser.add_argument("--gc_dir")
    parser.add_argument("--out_dir")
    parser.add_argument("--resolution", default="4x5", help="Can be 8x10, 4x5, 2x25 or 1x1")
    parser.add_argument("--region", default="EU", help="Can be EU, NA, or CH")
    parser.add_argument("--strat_filter_threshold", default="002", help="")
    #parser.add_argument("--start_date", default="2016-06-01")
    #parser.add_argument("--end_date", default="2017-08-31")
    parser.add_argument("-p", "--plot", type=bool, default=False)
    parser.add_argument("--do_temp_correct", type=bool)
    parser.add_argument("--apply_cld_frac_filter", type=bool)
    parser.add_argument("--do_cld_hght_test", type=bool)
    args = parser.parse_args()

    # Get files:
    gc_dir = args.gc_dir
    STR_RES = args.resolution
    REGION = args.region

    # Hard code for GEOS-Chem synthetic experiment:
    yrrange = '2016-2017'

    # Define output file depending on input arguments:
    out_file = os.path.join(args.out_dir, 'gc-v12-1-0-ut-no2'
                            + '-' + args.region
                            + '-' + args.resolution
                            + '-jja-' + yrrange
                            + '-' + "gaus-wgt" )

    if args.do_cld_hght_test:
        out_file += "-cldtop"
    if args.strat_filter_threshold != "002":
        out_file += '-' + args.strat_filter_threshold +'strat'
    if args.do_temp_correct:
        out_file += "-temp-corr"
    if args.apply_cld_frac_filter:
        out_file += "-cld-filt"

    out_file += "-v1.nc"

    strat_filter_threshold = float(args.strat_filter_threshold)/100
    files = get_gc_file_list(gc_dir, args.region)
    print('Number of files:', len(files), flush=True)

    rolling_total = ProcessedData(REGION, STR_RES, strat_filter_threshold,
                                  do_temperature_correction=args.do_temp_correct,
                                  do_cld_frac_filter=args.apply_cld_frac_filter,
                                  do_cld_hght_test=args.do_cld_hght_test)

    for file_path in files:
        rolling_total.process_geoschem_day(file_path)

    rolling_total.get_weighted_mean()
    rolling_total.print_data_report()
    if args.plot==True: rolling_total.plot_data()
    rolling_total.save_to_netcdf(out_file)


# -*- coding: utf-8 -*-
# @Time    : 2021-12-10 8:47 p.m.
# @Author  : young wang
# @FileName: tympanotomy.py
# @Software: PyCharm

import os
from scipy import ndimage
from skimage import exposure
import numpy as np
from pydicom.uid import generate_uid
import pydicom
from os.path import join, isfile
import matplotlib.pyplot as plt
import glob

def oct_to_dicom(data, resolutionx, resolutiony, PatientName, seriesdescription,
                 dicom_folder, dicom_prefix):
    """
    convert pixel array [512,512,330] to DICOM format
    using MRI template, this template will be deprecated
    in the future once OCT template is received
    """
    # data = load_from_oct_file(oct_file)

    dss = []

    template_file = '../template data/template.dcm'
    ds = pydicom.dcmread(template_file)

    # SeriesInstanceUID refers to each series, and should be
    # unqie for each sesssion, and generate_uid() provides an unique
    # identifier
    ds.SeriesInstanceUID = generate_uid()

    all_files_exist = False

    # looping through all 330 slice of images with [512(row) x 512(column)]
    for i in range(data.shape[2]):
        # UID used for indexing slices
        ds.SOPInstanceUID = generate_uid()

        # update row and column numbers to be 512
        ds.Rows = data.shape[0]
        ds.Columns = data.shape[1]

        # define the bottom(assuming the middle plane to be zero,
        # that -165 * 30um(axial resolution) = -4.95 mm)
        # DICOM assumes physical dimension to be in mm unit
        bottom = -4.95

        # elevate the z by its axial resolution at a time
        z = bottom + (i * 0.03)

        # update meta properties

        # 1cm / 512 = 0.02 mm, needs to check with rob
        # this spacing should be calculated as radiant/pixel then mm to pixel
        #
        ds.PixelSpacing = [resolutionx, resolutiony]  # pixel spacing in x, y planes [mm]
        ds.SliceThickness = 0.03  # slice thickness in axial(z) direction [mm]
        ds.SpacingBetweenSlices = 0.03  # slice spacing in axial(z) direction [mm]
        ds.SliceLocation = '%0.2f' % z  # slice location in axial(z) direction
        ds.InstanceNumber = '%0d' % (i + 1,)  # instance number, 330 in total
        ds.ImagePositionPatient = [z, 0, 0]  # patient physical location
        ds.Manufacturer = 'Audioptics Medical Inc'
        ds.InstitutionName = 'Audioptics Medical'
        ds.InstitutionAddress = '1344 Summer St., #55, Halifax, NS, Canada'
        ds.StudyDescription = 'Example DICOM export'
        ds.StationName = 'Unit 1'
        ds.SeriesDescription = seriesdescription
        ds.PhysiciansOfRecord = ''
        ds.PerformingPhysicianName = ''
        ds.InstitutionalDepartmentName = ''
        ds.ManufacturerModelName = 'Mark II'
        ds.PatientName = PatientName
        ds.PatientBirthDate = '20201123'
        ds.PatientAddress = ''

        # setting the dynamic range with WindowCenter and WindowWidth
        # lowest_visible_value = window_center — window_width / 2
        # highest_visible_value = window_center + window_width / 2

        ds.WindowCenter = '248'
        ds.WindowWidth = '396'

        # # set highest and lowest pixel values
        ds.LargestImagePixelValue = 446
        ds.SmallestImagePixelValue = 50
        dicom_file = join(dicom_folder, "%s%05d.dcm" % (dicom_prefix, i))

        pixel_data = data[:, :, i]
        # pixel_data[pixel_data <= 50] = 0
        ds.PixelData = pixel_data.tobytes()
        ds.save_as(dicom_file)
        dss.append(ds)
        all_files_exist = all_files_exist and isfile(dicom_file)
    return all_files_exist

''
if __name__ == '__main__':

    desktop_loc = os.path.expanduser('~/Desktop/GeoCorrection')
    patient_uid = 'patient'
    # study_uid = 'patient'
    # input_dir = join(desktop_loc, patient_uid, study_uid)

    study_uid = 'patient'
    input_dir = join(desktop_loc, patient_uid)
    file_extension = '*.npy'

    data_path = []
    for filename in glob.glob(join(input_dir, file_extension)):
        data_path.append(filename)

    volume_data = np.load(data_path[-1])

    umbom  = []
    tip = [ ]
    from scipy.signal import find_peaks, medfilt
    from scipy.ndimage import gaussian_filter1d,gaussian_filter

    slice = np.flip(volume_data[256,:,:].T)
    slice = gaussian_filter(slice, sigma=0.25)

    for i in range(slice.shape[0]):

        line = gaussian_filter1d(slice[i],sigma = 5 )

        peak_umbom, _ = find_peaks(line, height= 45 * np.mean(slice),distance = 20)

        if len(peak_umbom) >= 1:
            umboml = all(number > 200 for number in list(peak_umbom))
            umbomh = all(number <400 for number in list(peak_umbom))

            if umboml and umbomh:
                x = peak_umbom[0]
                y = line[x]
                umbom.append((x,y))

            else:
                pass

        else:
            pass

    for i in range(slice.shape[0]):

        line = gaussian_filter1d(slice[i], sigma=5)

        peak_tip, _ = find_peaks(line, height=25 * np.mean(slice), distance=20)

        if len(peak_tip) >= 1:

            tipl = all(number < 150 for number in list(peak_umbom))
            tiph = all(number > 450 for number in list(peak_umbom))

            if tipl and tiph:
                x = peak_tip[0]
                y = line[x]
                tip.append((i, x, y))
            else:
                pass

        else:
            pass

    print(len(umbom))
    plt.imshow(slice,'gray')
    pint = np.asarray(umbom)
    tip_array = np.asarray(tip)
    # x_min = min(tip_array[:,0])
    # ymin = slice[x_min]
    tip_cor = tip_array[0]
    plt.scatter(np.median(pint[:,0]),np.median(pint[:,1]))
    plt.scatter(tip_cor[0],tip_cor[1])
    plt.show()

    x1 = tip_cor[0]
    y1 = tip_cor[1]

    x2 = np.median(pint[:,0])
    y2 = np.median(pint[:,1])

    dy = y2 -y1
    dx = x2 - x1

    import math
    rad = math.atan2(dy,dx)
    degrees = math.degrees(rad)

    points = []

    portion = int(y2) - int(y1)
    for i in range (int(y1),int(y2), 1):

        scale = (i-int(y1))/portion
        y = y1 + scale/math.cos(rad)
        points.append((i,y))


from keras.utils import Sequence

import os
import cv2
import copy
import threading
import numpy as np
import imgaug as ia
from imgaug import augmenters as iaa
from imgaug.augmentables.segmaps import SegmentationMapsOnImage
import imgaug as ia
from keras.utils import Sequence
import xml.etree.ElementTree as ET
import matplotlib.pyplot as plt
from rhs_utils import KeyPointPair, draw_segmap
import math

import cv2

def read_annotations(img_dir, num_exposures):
    all_imgs = []
    seen_labels = {}
    i_base_file = 0
    i_bmp_file = 0

    all_file_names = sorted(os.listdir(img_dir))
    img = {'image_file_names': [], 'segmap_file_name': []}

    i_in_tuple = 0  # image in tuple counter
    for file_name in all_file_names:
        ext = os.path.splitext(file_name)[1]

        if ext == '.bmp':
            img['image_file_names'].append(img_dir + file_name)
            i_in_tuple += 1

        elif ext == '.cmp':
            img['segmap_file_name'] = img_dir + file_name
            i_in_tuple += 1

        if i_in_tuple > num_exposures:
            all_imgs += [img]
            i_in_tuple = 0  # reset image in tuple counter

            for i in img['image_file_names']:
                print("bmp_file=", i)
            print("cmp_file=", img['segmap_file_name'])
            img = {'image_file_names': [], 'segmap_file_name': []}

    return all_imgs

class YoloBatchGenerator(Sequence):
    def __init__(self,
                 all_imgs,
                 config,
                 shuffle=True,
                 jitter=True,
                 norm=None):
        self.generator = None

        self.all_imgs = all_imgs
        self.config = config

        print("self.config=", self.config)

        self.shuffle = shuffle
        self.jitter = jitter
        self.norm = norm
        self.image_counter = 0
        self.lock = threading.Lock()

        ia.seed(1)

        # augmentors by https://github.com/aleju/imgaug

        def sometimes(aug):
            return iaa.Sometimes(0.3, aug)

        self.aug_pipe = iaa.Sequential(
            [
                sometimes(iaa.Affine(
                    scale={"x": (0.8, 1.2), "y": (0.8, 1.2)},
                    translate_percent={"x": (-0.05, 0.05),
                                       "y": (-0.05, 0.05)},
                    rotate=(-5, 5),
                    shear=(-0.1, 0.1),
                    mode="edge")
                ),
                iaa.SomeOf(1,
                           [iaa.Add((-10, 10), per_channel=1),
                            iaa.Multiply((0.7, 1.3), per_channel=0)
                            ]
                           )
            ],
            random_order=True
        )

        if shuffle:
            np.random.shuffle(self.all_imgs)

    # scale augmented exposure-images and segmap downto size of CNN input gate
    def scale_imgs_and_segmap(self, imgs_array_aug, segmap_aug):
        # classmap section
        # print( np.shape( segmap_aug ) )

        scale = self.config['SCALE']
        width_scaled = int(self.config['IMAGE_WDT'] * scale)
        height_scaled = int(self.config['IMAGE_HGT'] * scale)

        segmap_aug_scaled = cv2.resize(
            segmap_aug, dim, interpolation=cv2.INTER_NEAREST)

        plt.imshow(segmap_aug_scaled, interpolation='nearest')
        plt.show()

        imgs_array_aug_scaled = np.zeros(
            (width_scaled, height_scaled, self.config['NUM_CHANNELS']))

        for i_expo in range(self.config['NUM_EXPOSURES']):
            img = cv2.GaussianBlur(imgs_array_aug[:, :, i_expo*self.config['NUM_CHANNELS']:(i_expo+1)*self.config['NUM_CHANNELS']],
                                   (self.config['GAUSS_KERNEL_RADIUS'] * 2+1,
                                    self.config['GAUSS_KERNEL_RADIUS']*2+1),
                                   self.config['GAUSS_SIGMA'] *
                                   self.config['GAUSS_KERNEL_RADIUS'],
                                   cv2.BORDER_DEFAULT)
            img_aug_scaled = cv2.resize(
                img, (width_scaled, height_scaled), interpolation=cv2.INTER_LINEAR)
            plt.imshow(img_aug_scaled, interpolation='nearest')
            plt.show()

            # concatenate images to get a train array with shape( IMAGE_H, IMAGE_W, 3*num_exposures )
            imgs_array_aug_scaled = np.concatenate(
                (imgs_array_aug_scaled, img_aug_scaled), axis=2)

    def __len__(self):
        return int(np.ceil(float(len(self.all_imgs))/self.config['BATCH_SIZE']))

    def size(self):
        return len(self.all_imgs)

    # load original sized segmap file
    def load_segmap(self, segmap_file_name):
        file = open(segmap_file_name, "rb")  # open binary file

        # get file size
        file.seek(0, 2)  # set pointer to end of file
        nb_bytes = file.tell()
        file.seek(0, 0)  # set pointer to begin of file
        buf = file.read(nb_bytes)
        file.close()
        # convert from byte stream to numpy array
        segmap_orig = np.asarray(list(buf), dtype=np.byte)
        segmap_orig = segmap_orig.reshape(
            (self.config['IMAGE_H'], self.config['IMAGE_W']))

        if self.config['CRATES_ONLY'] == 1:
            segmap_orig = np.where(segmap_orig > 0, 1, 0).astype(np.byte)

        return segmap_orig

    def load_images(self, i_images_tuple):
        # read all exposures of this i_images_tuple
        images = []
        for iExposure in range(self.config['NUM_EXPOSURES']):
            images.append(cv2.imread(
                self.all_imgs[i_images_tuple]['image_file_names'][iExposure]))
        return images

    def y_from_segmap(self, segmap):
        nb_cells_x = self.config['GRID_H']
        nb_cells_y = self.config['GRID_W']
        input_wdt = self.config['INPUT_W']
        input_hgt = self.config['INPUT_H']
        y = np.zeros((nb_cells_x, nb_cells_y, self.config['NUM_CLASSES']))

        # old coarse grid
        if False:
            for i_cell_y in range(nb_cells_y):
                start_pix_y = int(1.0*i_cell_y*input_hgt/nb_cells_y)
                end_pix_y = int(1.0*(i_cell_y+1)*input_hgt/nb_cells_y)
                for i_cell_x in range(nb_cells_x):
                    start_pix_x = int(1.0*i_cell_x*input_wdt/nb_cells_x)
                    end_pix_x = int(1.0*(i_cell_x+1)*input_wdt/nb_cells_x)
                    v_nb_class = np.zeros(self.config['NUM_CLASSES'])
                    nb_cell_pix = (end_pix_x-start_pix_x) * \
                        (end_pix_y - start_pix_y)

                    for i_pix_y in range(start_pix_y, end_pix_y):
                        for i_pix_x in range(start_pix_x, end_pix_x):
                            val = segmap[i_pix_y, i_pix_x]
                            v_nb_class[val] += 1

                    for i_class in range(self.config['NUM_CLASSES']):
                        y[i_cell_y, i_cell_x, i_class] = 1.0 * \
                            v_nb_class[i_class]/nb_cell_pix

        # new pixel grid
        for i_cell_y in range(nb_cells_y):
            for i_cell_x in range(nb_cells_x):
                i_class = segmap[i_cell_y, i_cell_x]
                y[i_cell_y, i_cell_x, i_class] = 1

        return y

    def save_augimage(self, nb, path, tag, img):  # save augmented image for debugging
        fullName = path + "img_" + str(nb) + "_"+tag+".bmp"
        cv2.imwrite(fullName, img)

    # save augmented segmap for debugging
    def save_augsegmap(self, nb, path, tag, segmap):
        fullName = "..\\imgaug\\img_" + str(nb) + "_"+tag + ".cmp"
        buf = segmap.tobytes()
        file = open(fullName, "wb")
        file.write(buf)

    def save_y(self, nb, path, y):  # save augmented y for debugging
        fullName = "..\\imgaug\\img_" + str(nb) + ".txt"
        file = open(fullName, "w")
        file.write(np.array_str(y))

    def __getitem__(self, idx):  # returns a complete batch pair x_batch and y_batch
        l_bound = idx*self.config['BATCH_SIZE']
        r_bound = (idx+1)*self.config['BATCH_SIZE']

        x_batch = np.zeros((0))
        y_batch = np.zeros((0))
        x_batch = x_batch.reshape(
            (0, self.config['IMAGE_H'], self.config['IMAGE_W'], 3*self.config['NUM_EXPOSURES']))
        y_batch = y_batch.reshape(
            (0, self.config['GRID_H'], self.config['GRID_W'], self.config['NUM_CLASSES']))
        # old y_batch = np.zeros((r_bound - l_bound, self.config['GRID_H'],  self.config['GRID_W'], 1 + self.config['NUM_CLASSES']))

        i_img = 0
        images_batch = []
        segmaps_batch = []

        self.lock.acquire()  # to ensure synchronization of image and segmap

        # keypoints-list aufbauen
        num_images = len(self.all_imgs)
        for instance_count in range(r_bound - l_bound):
            instance_src_index = (l_bound + instance_count) % num_images
            # augment all exposure instances in all_imgs
            imgs = self.load_images(instance_src_index)

            # construct output grid fro segment map
            segmap_file_name = self.all_imgs[instance_src_index]['segmap_file_name']
            segmap = self.load_segmap(segmap_file_name)

            # for synchronization, but maybe it´s not necessary
            self.aug_pipe_det = self.aug_pipe.to_deterministic()

            # combine images to a single stacked array
            imgs_array = np.copy(imgs[0])
            for i_img in range(1, self.config['NUM_EXPOSURES']):
                # concatenate images to get a train array with shape( IMAGE_H, IMAGE_W, 3*num_exposures )
                imgs_array = np.concatenate((imgs_array, imgs[i_img]), axis=2)

            segmap = SegmentationMapsOnImage(segmap, shape=imgs_array.shape)
            imgs_array_aug, segmap_aug = self.aug_pipe_det(
                image=imgs_array, segmentation_maps=segmap)  # augment images with 3*num_exposures

            # scale augmented exposure-images and segmap downto size of CNN input gate
            imgs_array_aug_scaled, segmap_aug_scaled = self.scale_imgs_and_segmap(
                imgs_array_aug, segmap_aug)

            y = self.y_from_segmap(segmap_aug_scaled.get_arr())

            #self.save_augimage( self.image_counter, "..\\imgaug\\", "a", imgs_array_aug[...,0:3].reshape((self.config['IMAGE_H'],self.config['IMAGE_W'],3 )) )
            #self.save_augimage( self.image_counter, "..\\imgaug\\", "b", imgs_array_aug[...,3:6].reshape((self.config['IMAGE_H'],self.config['IMAGE_W'],3 )) )
            #self.save_augimage( self.image_counter, "..\\imgaug\\", "c", imgs_array_aug[...,6:9].reshape((self.config['IMAGE_H'],self.config['IMAGE_W'],3 )) )
            #self.save_augsegmap( self.image_counter, "..\\imgaug\\","a",  segmap_aug.get_arr() )
            #self.save_augsegmap( self.image_counter, "..\\imgaug\\","b",  segmap_aug.get_arr() )
            #self.save_augsegmap( self.image_counter, "..\\imgaug\\","c",  segmap_aug.get_arr() )

#            segmap = SegmentationMapsOnImage(segmap, shape=img.shape)

#            self.aug_pipe_det = self.aug_pipe.to_deterministic() # for synchronization, but maybe it´s not necessary

            # augment all images in tuple with same augmentation mask
#            images_aug_i = []

#            for i_img in self.config['num_exposures']
#                image_aug_i, segmap_aug_i = self.aug_pipe_det(image=img, segmentation_maps=segmap) #augment one image in tuple
#                images_aug_i.append( image_aug_i )
#                old_segmap_aug_i = np.copy( segmap_aug_i )
#                if i_img == 0:
#                    train_array = np.copy( image_aug_i )
#                elif i_img > 0:
#                    train_array = np.concatenate( (train_array, image_aug_i), axis=2 ) # concatenate images to get a train array with shape( IMAGE_H, IMAGE_W, 3*num_exposures )
#                    diff = np.sum( segmap_aug_i - old_segmap_aug_i )  # check if all images in tuple are augmented in same manner
#                    print( "augdiff=", diff )

            # merge augmented images in tuple

#            y = self.y_from_segmap( segmaps_aug_i.get_arr() )

            #self.save_augimage( self.image_counter, "..\\imgaug\\", images_aug_i.reshape((self.config['IMAGE_H'],self.config['IMAGE_W'],3 )) )
            #self.save_y( self.image_counter, "..\\imgaug\\", np.argmax( y, axis=2 ) )
            #self.save_augsegmap( self.image_counter, "..\\imgaug\\",  segmaps_aug_i.get_arr() )

            self.image_counter += 1
            i_img += 1

#            x_batch = np.append( x_batch, images_aug_i.reshape((1,self.config['IMAGE_H'],self.config['IMAGE_W'],3 )),0)
            x_batch = np.append(x_batch, imgs_array_aug_scaled.reshape(
                (1, self.config['INPUT_H'], self.config['INPUT_W'], 3*self.config['NUM_EXPOSURES'])), 0)
            y_batch = np.append(y_batch, y.reshape(
                (1, self.config['GRID_H'], self.config['GRID_W'], self.config['NUM_CLASSES'])), 0)

        self.lock.release()
        x_batch = self.norm(x_batch)

        np.set_printoptions(threshold=np.inf, linewidth=500)

        # for i in range( x_batch.shape[0] ):
        #    x = x_batch[i,:]*255.0
        #    self.save_augimage( i, "..\\imgaug\\", x )
        #    self.save_y( i, "..\\imgaug\\", y_batch[i, :] )#np.argmax( y_batch[i,:], axis=2 ) )
        # self.save_y( i, "..\\imgaug\\", x )#np.argmax( y_batch[i,:], axis=2 ) )

        return x_batch, y_batch  # x_batch=image normalized and y_batch=one_hot_classes

    def on_epoch_end(self):
        if self.shuffle:
            # shuffle along the first axis only
            np.random.shuffle(self.all_imgs)


# Copyright (c) 2017 Sony Corporation. All Rights Reserved.
#
# 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 argparse
import os

import nnabla as nn
from nnabla.ext_utils import get_extension_context
import numpy as np
from scipy.io import wavfile

from hparams import hparams as hp
from model.model import Tacotron

import sys
sys.path.append(str(Path().cwd().parents[2] / 'utils'))

from neu.tts.audio import synthesize_from_spec
from neu.tts.text import text_normalize

def synthesize(args):
    char2idx = {ch: i for i, ch in enumerate(hp.vocab)}
    with open(args.f_text, 'r') as file:
        text = ''.join(file.readlines())
    # normalize the text
    text = text_normalize(text, hp.vocab)
    if len(text) >= hp.text_len - 1:
        text = text[:hp.text_len-1]
    text += '~' * (hp.text_len - len(text))
    text = np.array([char2idx[ch] for ch in text]).reshape(-1)

    hp.batch_size = 1
    # load the model
    model = Tacotron(hp)
    model.training = False
    model.load_parameters(args.f_model)

    x_txt = nn.Variable([hp.batch_size, hp.text_len])
    _, mag, _ = model(x_txt)
    x_txt.d = text[np.newaxis, :]

    mag.forward(clear_buffer=True)
    wave = synthesize_from_spec(mag.d[0].copy(), hp)
    wavfile.write(args.f_output, rate=hp.sr, data=wave)  # write a sample

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--context', '-c', type=str, default='cudnn',
                        help="Extension module. 'cudnn' is highly recommended.")
    parser.add_argument("--device-id", "-d", type=str, default='-1',
                        help='A list of device ids to use.\
                        This is only valid if you specify `-c cudnn`. \
                        Defaults to use all available GPUs.')
    parser.add_argument("--f-model", "-m", type=str,
                        help='File path to the trained model.')
    parser.add_argument("--f-text", "-f", type=str,
                        help='File path to the text file.')
    parser.add_argument("--f-output", "-o", type=str, default='sample.wav',
                        help='File path to the synthetic output waveform.')
    args = parser.parse_args()

    # setup context for nnabla
    if args.device_id != '-1':
        os.environ["CUDA_VISIBLE_DEVICES"] = args.device_id

    # setup nnabla context
    ctx = get_extension_context(args.context, device_id='0')
    nn.set_default_context(ctx)

    synthesize(args)


#!/usr/bin/env python
import argparse
import sys
import networkx as nx
import re
import numpy as np
import itertools

_s = re.compile('\s+')
_p = re.compile('(\d+)\s+(\d+)')
_rings_str = """* place dummy atoms at the center of each rings
*

set ncycles = %s

read sequence POL @NCYCLES
generate DUM warn

%s

scalar wmain set 0
%s

return
"""

def build_topology(psffile):
    g = nx.Graph()
    flag = 0

    for line in open(psffile).readlines():
        if flag == 0 and line.strip().endswith('NATOM'):
            flag = 1
            continue
        if flag == 0 and line.strip().endswith('bonds'):
            flag = 2
            continue
        if flag == 1 and not line.strip(): flag = 0
        if flag == 2 and not line.strip(): break

        if flag == 1:
            num, segid, resid, resname, name = _s.split(line)[1:6]
            if resname.startswith('TIP3'): continue
            if name.startswith('H'): continue
            g.add_node(int(num), {'segid': segid, 'resname': resname, 'name': name, 'resid': resid})

        if flag == 2:
            for pair in _p.findall(line):
                num1, num2 = map(int, pair)
                if g.has_node(num1) and g.has_node(num2): g.add_edge(num1, num2)
    return g

def build_atomtable(psf, crdfile):
    crds = {}
    flag = 0
    for line in open(crdfile).readlines():
        entries = _s.split(line)
        if flag == 0 and entries[0] != '*':
            flag = 1
            continue
        if flag == 1 and not line.strip(): break

        if flag == 1:
            num, resid, resname, name, x, y, z, segid = _s.split(line.strip())[:8]
            if resname.startswith('TIP3'): continue
            if name.startswith('H'): continue
            if psf.node[int(num)]['name'] != name: raise AtomMismatch("%d %s != %d %s" % (int(num), psf.node[int(num)]['name'], int(num), name))
            crds[int(num)] = np.array((float(x), float(y), float(z)))
    return crds

def build_exclusionlist(penstr):
    exclud_list = []
    for line in open(penstr).readlines():
        entries = _s.split(line)
        if entries[0].startswith('%'): continue
        if entries[0].startswith('-'):
            exclud_list.append(entries[1])
    return exclud_list

class AtomMismatch(Exception):
    pass

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('psf', metavar='psffile', help='PSF file')
    parser.add_argument('crd', metavar='crdfile', help='CRD file')
    parser.add_argument('penstr', metavar='penstrfile', help='pentest.py output STR file')
    args = parser.parse_args()

    # build connectivity of atoms
    psf = build_topology(args.psf)
    crd = build_atomtable(psf, args.crd)
    exclud_list = build_exclusionlist(args.penstr)

    # sanity check
    if len(psf.nodes()) != len(crd): raise AtomMismatch('Number of atom does not match')

    # print out ring atoms
    molecules = nx.connected_component_subgraphs(psf)
    dummys = []
    fixes = []
    i = 1
    _coor = 'coor set xdir %8.3f ydir %8.3f zdir %8.3f sele segid dum .and. resi %d end'
    _fix = 'scalar wmain set 1.0 sele segid %s .and. resi %s end'
    for m in molecules:
        cycles = nx.cycle_basis(m)
        if not cycles: continue
        for cycle in cycles:
            segid = psf.node[cycle[0]]['segid']
            if segid != 'MEMB' and segid in exclud_list: continue
            atoms = np.array([crd[num] for num in cycle])
            com = np.sum(atoms, axis=0)/len(atoms)
            dummys.append(_coor % (com[0], com[1], com[2], i))
            if _fix % (psf.node[cycle[0]]['segid'], psf.node[cycle[0]]['resid']) not in fixes:
                fixes.append(_fix % (psf.node[cycle[0]]['segid'], psf.node[cycle[0]]['resid']))
            i += 1
    print _rings_str % (len(dummys), "\n".join(dummys), "\n".join(fixes))

if __name__ == '__main__':
    main()