"""
Spectral data representation.

"""

from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

from .core import Data
from .atomic import get_atomdat

from astropy.units import (erg, km, cm, s, angstrom, spectral,
                           spectral_density, Quantity, UnitsError)
from astropy.constants import m_e, c, e
from astropy.table import Table, Column

from math import pi, sqrt, exp, log10
from warnings import warn

import numpy as np

__all__ = ['Spectrum2D', 'Spectrum1D', 'Absorber', 'EmissionLine']

e2_me_c = (e.esu ** 2 / (m_e.cgs * c.cgs)).to(cm ** 2 / s)
c_kms = c.to(km / s)

atomdat = None

def find_bin_edges(bin_centres):
    """
    Find the bin edges given the bin centres.

    Parameters
    ----------
    bin_centres : array, shape (N,)
        The bin centres.

    Returns
    -------
    bins : array, shape (N + 1,)
        The bin edges.

    """

    if not isinstance(bin_centres, np.ndarray):
        bin_centres = np.asarray(bin_centres)

    edges = bin_centres[:-1] + 0.5 * (bin_centres[1:] - bin_centres[:-1])
    bins = np.concatenate(([2 * bin_centres[0] - edges[0]], edges,
                           [2 * bin_centres[-1] - edges[-1]]))

    return bins

class Spectrum2D(object):
    """
    A 2D spectrum.

    Parameters
    ----------
    dispersion : `astropy.units.Quantity` or array, shape (N,)
        Spectral dispersion axis.

    data : array, shape (N, M)
        The spectral data.

    unit : `astropy.units.UnitBase` or str, optional
        Unit for the dispersion axis.

    """

    def __init__(self, dispersion, data, unit=None):

        self.dispersion = Quantity(dispersion, unit=unit)

        if unit is not None:
            self.wavelength = self.dispersion.to(angstrom)

        else:
            self.wavelength = self.dispersion

        self.data = data

class Spectrum1D(object):
    """
    A 1D spectrum. Assumes wavelength units unless otherwise specified.

    Parameters
    ----------
    dispersion : `astropy.units.Quantity` or array
        Spectral dispersion axis.

    flux : `igmtools.data.Data`, `astropy.units.Quantity` or array
        Spectral flux. Should have the same length as `dispersion`.

    error : `astropy.units.Quantity` or array, optional
        Error on each flux value.

    continuum : `astropy.units.Quantity` or array, optional
        An estimate of the continuum flux.

    mask : array, optional
        Mask for the spectrum. The values must be False where valid and True
        where not.

    unit : `astropy.units.UnitBase` or str, optional
        Spectral unit.

    dispersion_unit : `astropy.units.UnitBase` or str, optional
        Unit for the dispersion axis.

    meta : dict, optional
        Meta data for the spectrum.

    """

    def __init__(self, dispersion, flux, error=None, continuum=None,
                 mask=None, unit=None, dispersion_unit=None, meta=None):

        _unit = flux.unit if unit is None and hasattr(flux, 'unit') else unit

        if isinstance(error, (Quantity, Data)):
            if error.unit != _unit:
                raise UnitsError('The error unit must be the same as the '
                                 'flux unit.')
            error = error.value

        elif isinstance(error, Column):
            if error.unit != _unit:
                raise UnitsError('The error unit must be the same as the '
                                 'flux unit.')
            error = error.data

        # Set zero error elements to NaN:
        if error is not None:
            zero = error == 0
            error[zero] = np.nan

            # Mask these elements:
            if mask is not None:
                self.mask = mask | np.isnan(error)

            else:
                self.mask = np.isnan(error)

        # If dispersion is a `Quantity`, `Data`, or `Column` instance with the
        # unit attribute set, that unit is preserved if `dispersion_unit` is
        # None, but overriden otherwise
        self.dispersion = Quantity(dispersion, unit=dispersion_unit)

        if dispersion_unit is not None:
            self.wavelength = self.dispersion.to(angstrom)

        else:
            # Assume wavelength units:
            self.wavelength = self.dispersion

        self.flux = Data(flux, error, unit)

        if continuum is not None:
            self.continuum = Quantity(continuum, unit=unit)

        else:
            self.continuum = None

        self.meta = meta

    @classmethod
    def from_table(cls, table, dispersion_column, flux_column,
                   error_column=None, continuum_column=None, unit=None,
                   dispersion_unit=None):
        """
        Initialises a `Spectrum1D` object from an `astropy.table.Table`
        instance.

        Parameters
        ----------
        table : `astropy.table.Table`
            Contains information used to construct the spectrum. Must have
            columns for the dispersion axis and the spectral flux.

        dispersion_column : str
            Name for the dispersion column.

        flux_column : str
            Name for the flux column.

        error_column : str, optional
            Name for the error column.

        continuum_column : str, optional
            Name for the continuum column.

        unit : `astropy.units.UnitBase` or str, optional
            Spectral unit.

        dispersion_unit : `astropy.units.UnitBase` or str, optional
            Unit for the dispersion axis.

        """

        dispersion = Quantity(table[dispersion_column])
        flux = Quantity(table[flux_column])

        if error_column is not None:
            error = Quantity(table[error_column])
        else:
            error = None

        if continuum_column is not None:
            continuum = Quantity(table[continuum_column])
        else:
            continuum = None

        meta = table.meta
        mask = table.mask

        return cls(dispersion, flux, error, continuum, mask, unit,
                   dispersion_unit, meta)

    def write(self, *args, **kwargs):
        """
        Write the spectrum to a file. Accepts the same arguments as
        `astropy.table.Table.write`

        """

        if self.dispersion.unit is None:
            label_string = 'WAVELENGTH'

        else:
            if self.dispersion.unit.physical_type == 'length':
                label_string = 'WAVELENGTH'

            elif self.dispersion.unit.physical_type == 'frequency':
                label_string = 'FREQUENCY'

            elif self.dispersion.unit.physical_type == 'energy':
                label_string = 'ENERGY'

            else:
                raise ValueError('unrecognised unit type')

        t = Table([self.dispersion, self.flux, self.flux.uncertainty.value],
                  names=[label_string, 'FLUX', 'ERROR'])
        t['ERROR'].unit = t['FLUX'].unit

        if self.continuum is not None:
            t['CONTINUUM'] = self.continuum

        t.write(*args, **kwargs)

    def plot(self, **kwargs):
        """
        Plot the spectrum. Accepts the same arguments as
        `igmtools.plot.Plot`.

        """

        from ..plot import Plot

        p = Plot(1, 1, 1, **kwargs)

        p.axes[0].plot(self.dispersion.value, self.flux.value,
                       drawstyle='steps-mid')

        if self.flux.uncertainty is not None:
            p.axes[0].plot(self.dispersion.value, self.flux.uncertainty.value,
                           drawstyle='steps-mid')

        p.tidy()
        p.display()

    def normalise_to_magnitude(self, magnitude, band):
        """
        Normalises the spectrum to match the flux equivalent to the
        given AB magnitude in the given passband.

        Parameters
        ----------
        magnitude : float
            AB magnitude.

        band : `igmtools.photometry.Passband`
            The passband.

        """

        from ..photometry import mag2flux

        mag_flux = mag2flux(magnitude, band)
        spec_flux = self.calculate_flux(band)
        norm = mag_flux / spec_flux
        self.flux *= norm

    def calculate_flux(self, band):
        """
        Calculate the mean flux for a passband, weighted by the response
        and wavelength in the given passband.

        Parameters
        ----------
        band : `igmtools.photometry.Passband`
            The passband.

        Returns
        -------
        flux : `astropy.units.Quantity`
            The mean flux in erg / s / cm^2 / Angstrom.

        Notes
        -----
        This function does not calculate an uncertainty.

        """

        if (self.wavelength[0] > band.wavelength[0] or
                self.wavelength[-1] < band.wavelength[-1]):

            warn('Spectrum does not cover the whole bandpass, '
                 'extrapolating...')
            dw = np.median(np.diff(self.wavelength.value))
            spec_wavelength = np.arange(
                band.wavelength.value[0],
                band.wavelength.value[-1] + dw, dw) * angstrom
            spec_flux = np.interp(spec_wavelength, self.wavelength,
                                  self.flux.value)

        else:
            spec_wavelength = self.wavelength
            spec_flux = self.flux.value

        i, j = spec_wavelength.searchsorted(
            Quantity([band.wavelength[0], band.wavelength[-1]]))
        wavelength = spec_wavelength[i:j]
        flux = spec_flux[i:j]

        dw_band = np.median(np.diff(band.wavelength))
        dw_spec = np.median(np.diff(wavelength))

        if dw_spec.value > dw_band.value > 20:

            warn('Spectrum wavelength sampling interval {0:.2f}, but bandpass'
                 'sampling interval {1:.2f}'.format(dw_spec, dw_band))

            # Interpolate the spectrum to the passband wavelengths:
            flux = np.interp(band.wavelength, wavelength, flux)
            band_transmission = band.transmission
            wavelength = band.wavelength

        else:
            # Interpolate the band transmission to the spectrum wavelengths:
            band_transmission = np.interp(
                wavelength, band.wavelength, band.transmission)

        # Weight by the response and wavelength, appropriate when we're
        # counting the number of photons within the band:
        flux = (np.trapz(band_transmission * flux * wavelength, wavelength) /
                np.trapz(band_transmission * wavelength, wavelength))
        flux *= erg / s / cm ** 2 / angstrom

        return flux

    def calculate_magnitude(self, band, system='AB'):
        """
        Calculates the magnitude in a given passband.

        band : `igmtools.photometry.Passband`
            The passband.

        system : {`AB`, `Vega`}
            Magnitude system.

        Returns
        -------
        magnitude : float
            Magnitude in the given system.

        """

        if system not in ('AB', 'Vega'):
            raise ValueError('`system` must be one of `AB` or `Vega`')

        f1 = self.calculate_flux(band)

        if f1 > 0:
            magnitude = -2.5 * log10(f1 / band.flux[system])

            if system == 'Vega':
                # Add 0.026 because Vega has V = 0.026:
                magnitude += 0.026

        else:
            magnitude = np.inf

        return magnitude

    def apply_extinction(self, EBmV):
        """
        Apply Milky Way extinction.

        Parameters
        ----------
        EBmV : float
            Colour excess.

        """

        from astro.extinction import MWCardelli89

        tau = MWCardelli89(self.wavelength, EBmV=EBmV).tau
        self.flux *= np.exp(-tau)

        if self.continuum is not None:
            self.continuum *= np.exp(-tau)

    def rebin(self, dispersion):
        """
        Rebin the spectrum onto a new dispersion axis.

        Parameters
        ----------
        dispersion : float, `astropy.units.Quantity` or array
            The dispersion for the rebinned spectrum. If a float, assumes a
            linear scale with that bin size.

        """

        if isinstance(dispersion, float):
            dispersion = np.arange(
                self.dispersion.value[0], self.dispersion.value[-1],
                dispersion)

        old_bins = find_bin_edges(self.dispersion.value)
        new_bins = find_bin_edges(dispersion)

        widths = np.diff(old_bins)

        old_length = len(self.dispersion)
        new_length = len(dispersion)

        i = 0  # index of old array
        j = 0  # index of new array

        # Variables used for rebinning:
        df = 0.0
        de2 = 0.0
        nbins = 0.0

        flux = np.zeros_like(dispersion)
        error = np.zeros_like(dispersion)

        # Sanity check:
        if old_bins[-1] < new_bins[0] or new_bins[-1] < old_bins[0]:
            raise ValueError('Dispersion scales do not overlap!')

        # Find the first contributing old pixel to the rebinned spectrum:
        if old_bins[i + 1] < new_bins[0]:

            # Old dispersion scale extends lower than the new one. Find the
            # first old bin that overlaps with the new scale:
            while old_bins[i + 1] < new_bins[0]:
                i += 1

            i -= 1

        elif old_bins[0] > new_bins[j + 1]:

            # New dispersion scale extends lower than the old one. Find the
            # first new bin that overlaps with the old scale:
            while old_bins[0] > new_bins[j + 1]:
                flux = np.nan
                error = np.nan
                j += 1

            j -= 1

        l0 = old_bins[i]  # lower edge of contributing old bin

        while True:

            h0 = old_bins[i + 1]  # upper edge of contributing old bin
            h1 = new_bins[j + 1]  # upper edge of jth new bin

            if h0 < h1:
                # Count up the decimal number of old bins that contribute to
                # the new one and start adding up fractional flux values:
                if self.flux.uncertainty.value[i] > 0:
                    bin_fraction = (h0 - l0) / widths[i]
                    nbins += bin_fraction

                    # We don't let `Data` handle the error propagation here
                    # because a sum of squares will not give us what we
                    # want, i.e. 0.25**2 + 0.75**2 != 0.5**2 + 0.5**2 != 1**2
                    df += self.flux.value[i] * bin_fraction
                    de2 += self.flux.uncertainty.value[i] ** 2 * bin_fraction

                l0 = h0
                i += 1

                if i == old_length:
                    break

            else:
                # We have all but one of the old bins that contribute to the
                # new one, so now just add the remaining fraction of the new
                # bin to the decimal bin count and add the remaining
                # fractional flux value to the sum:
                if self.flux.uncertainty.value[i] > 0:
                    bin_fraction = (h1 - l0) / widths[i]
                    nbins += bin_fraction
                    df += self.flux.value[i] * bin_fraction
                    de2 += self.flux.uncertainty.value[i] ** 2 * bin_fraction

                if nbins > 0:
                    # Divide by the decimal bin count to conserve flux density:
                    flux[j] = df / nbins
                    error[j] = sqrt(de2) / nbins

                else:
                    flux[j] = 0.0
                    error[j] = 0.0

                df = 0.0
                de2 = 0.0
                nbins = 0.0

                l0 = h1
                j += 1

                if j == new_length:
                    break

        if hasattr(self.dispersion, 'unit'):
            dispersion = Quantity(dispersion, self.dispersion.unit)

        if hasattr(self.flux, 'unit'):
            flux = Data(flux, error, self.flux.unit)

        # Linearly interpolate the continuum onto the new dispersion scale:
        if self.continuum is not None:
            continuum = np.interp(dispersion, self.dispersion, self.continuum)
        else:
            continuum = None

        return self.__class__(dispersion, flux, continuum=continuum)

class Absorber(object):
    """
    Class representation of an absorber.

    Parameters
    ----------
    identifier : str
        Name of the ion, molecule or isotope, e.g. `HI`.

    redshift : float, optional
        Redshift of the absorber.

    logn : float, optional
        Log10 column density (cm^-2).

    b : float, optional
        Doppler broadening parameter (km/s).

    covering_fraction : float, optional
        Covering fraction.

    atom : `igmtools.data.AtomDat`, optional
        Atomic data.

    """

    def __init__(self, identifier, redshift=None, logn=None, b=None,
                 covering_fraction=1, atom=None):

        if atom is None:
            atom = get_atomdat()

        self.identifier = identifier
        self.transitions = atom[identifier]
        self.redshift = redshift
        self.logn = logn
        self.b = b
        self.covering_fraction = covering_fraction

    def __repr__(self):

        return 'Absorber({0}, z={1:.2f}, logN={2:.2f}, b={3})'.format(
            self.identifier, self.redshift, self.logn, int(self.b))

    @classmethod
    def from_tau_peak(cls, transition, tau, b):
        """
        Initialise an absorber from the optical depth at line centre and
        Doppler broadining parameter of a given transition.

        Parameters
        ----------
        transition : str
            Name of the transition, e.g. `HI 1215'

        tau : float
            Optical depth at the line centre.

        b : float
            Doppler broadening parameter (km/s).

        """

        atom = get_atomdat()
        transition = atom.get_transition(transition)

        if isinstance(b, Quantity):
            b = b.to(cm / s)
        else:
            b = (b * km / s).to(cm / s)

        wavelength = transition.wavelength.to(cm)
        osc = transition.osc

        column = tau * b / (sqrt(pi) * e2_me_c * osc * wavelength)
        logn = log10(column.value)

        return cls(identifier=transition.parent, logn=logn, b=b)

    def optical_depth(self, dispersion):
        """
        Calculates the optical depth profile for a given spectral
        dispersion array.

        Parameters
        ----------
        dispersion : array
            Spectral dispersion.

        Returns
        -------
        tau : array
            The optical depth profile.

        """

        from ..calculations import optical_depth, tau_peak

        if isinstance(dispersion, Quantity):
            dispersion = dispersion.to(angstrom)

        elif hasattr(dispersion, 'unit'):
            if dispersion.unit is not None:
                dispersion = dispersion.to(angstrom)

        else:
            dispersion = Quantity(dispersion, unit=angstrom)

        velocity_range = ([-20000, 20000] * km / s if self.logn > 18
                          else [-1000, 1000] * km / s)

        # Select only transitions with redshifted central wavelengths inside
        # `dispersion` +/- 500 km/s:
        rest_wavelengths = Quantity([t.wavelength for t in self.transitions])
        observed_wavelengths = rest_wavelengths * (1 + self.redshift)

        wmin = dispersion[0] * (1 - 500 * km / s / c_kms)
        wmax = dispersion[-1] * (1 - 500 * km / s / c_kms)

        in_range = ((observed_wavelengths >= wmin) &
                    (observed_wavelengths <= wmax))
        transitions = np.array(self.transitions)[in_range]

        tau = np.zeros_like(dispersion.value)

        for i, transition in enumerate(transitions):

            tau_max = tau_peak(transition, self.logn, self.b)

            if 1 - exp(-tau_max) < 1e-3:
                continue

            observed_wavelength = transition.wavelength * (1 + self.redshift)
            dv = ((dispersion - observed_wavelength) /
                  observed_wavelength * c_kms)

            i0, i1 = dv.searchsorted(velocity_range)
            tau0 = optical_depth(dv[i0:i1], transition, self.logn, self.b)
            tau[i0:i1] += tau0

        return tau

class EmissionLine(object):
    """
    Class representation of an emission line and its properties.

    Parameters
    ----------
    wavelength : float
        Rest frame wavelength of the line in Angstrom.

    redshift : float
        Redshift of the emission line.

    flux : `igmtools.data.Data`, optional
        Integrated line flux.

    cont : `igmtools.data.Data`, optional
        Continuum flux at the line centre.

    eqw : `igmtools.data.Data`, optional
        Equivalent width of the line.

    """

    def __init__(self, wavelength, redshift, flux=None, cont=None, eqw=None):

        from ..calculations import comoving_distance

        if flux and not isinstance(flux, Data):
            raise ValueError('flux must be an instance of a Data object')

        if cont and not isinstance(cont, Data):
            raise ValueError('cont must be an instance of a Data object')

        if eqw and not isinstance(eqw, Data):
            raise ValueError('eqw must be an instance of a Data object')

        if isinstance(wavelength, Quantity):
            self.wavelength = wavelength.to(angstrom)

        else:
            self.wavelength = wavelength * angstrom

        self.redshift = redshift
        self.wavelength_observed = self.wavelength * (1 + self.redshift)

        if flux:

            self._flux = flux.to(erg / cm ** 2 / s, equivalencies=spectral())
            self.rflux = self._flux * (1 + redshift) ** 2

            distance = comoving_distance(self.redshift).cgs
            self.luminosity = 4 * pi * distance ** 2 * self._flux

            if cont and eqw:

                self._cont = cont.to(
                    erg / cm ** 2 / s / angstrom,
                    equivalencies=spectral_density(self.wavelength_observed))
                self.rcont = self._cont * (1 + redshift) ** 3

                self._eqw = eqw.to(angstrom)
                self.reqw = self._eqw / (1 + redshift)

            elif cont and not eqw:

                self._cont = cont.to(
                    erg / cm ** 2 / s / angstrom,
                    equivalencies=spectral_density(self.wavelength_observed))
                self.rcont = self._cont * (1 + redshift) ** 3

                self._eqw = self._flux / self._cont
                self.reqw = self._eqw / (1 + redshift)

            elif eqw and not cont:

                self._eqw = eqw.to(angstrom)
                self.reqw = self._eqw / (1 + redshift)

                self._cont = self._flux / self._eqw
                self.rcont = self._cont * (1 + redshift) ** 3

            else:

                self._eqw = eqw
                self.reqw = None
                self._cont = cont
                self.rcont = None

        elif cont:

            self._cont = cont.to(
                erg / cm ** 2 / s / angstrom,
                equivalencies=spectral_density(self.wavelength_observed))
            self.rcont = self._cont * (1 + redshift) ** 3

            if eqw:

                self._eqw = eqw.to(angstrom)
                self.reqw = self._eqw / (1 + redshift)

                self._flux = self._cont * self._eqw
                self.rflux = self._flux * (1 + redshift) ** 2

                distance = comoving_distance(self.redshift).cgs
                self.luminosity = 4 * pi * distance ** 2 * self._flux

            else:

                self._eqw = eqw
                self.reqw = None
                self._flux = flux
                self.rflux = None
                self.luminosity = None

        elif eqw:

            self._eqw = eqw.to(angstrom)
            self.reqw = self._eqw / (1 + redshift)

            self._flux = flux
            self.rflux = None
            self._cont = cont
            self.rcont = None
            self.luminosity = None

        else:

            self._flux = flux
            self.rflux = None
            self._cont = cont
            self.rcont = None
            self._eqw = eqw
            self.reqw = None
            self.luminosity = None

    @property
    def flux(self):
        return self._flux

    @flux.setter
    def flux(self, value):

        from ..calculations import comoving_distance

        if not isinstance(value, Data):
            raise ValueError('flux must be an instance of a Data object')

        self._flux = value.to(erg / cm ** 2 / s, equivalencies=spectral())
        self.rflux = self._flux * (1 + self.redshift) ** 2

        distance = comoving_distance(self.redshift).cgs
        self.luminosity = 4 * pi * distance ** 2 * self._flux

    @property
    def cont(self):
        return self._cont

    @cont.setter
    def cont(self, value):

        if not isinstance(value, Data):
            raise ValueError('cont must be an instance of a Data object')

        self._cont = value.to(
            erg / cm ** 2 / s / angstrom,
            equivalencies=spectral_density(self.wavelength_observed))
        self.rcont = self._cont * (1 + self.redshift) ** 3

    @property
    def eqw(self):
        return self._eqw

    @eqw.setter
    def eqw(self, value):

        if not isinstance(value, Data):
            raise ValueError('eqw must be an instance of a Data object')

        self._eqw = value.to(angstrom)
        self.reqw = self._eqw / (1 + self.redshift)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys

import argparse

import multiprocessing
import numpy as np

import torch
from torch import nn, Tensor
import torch.nn.functional as F

from torch_geometric.data import Data as GeometricData, Batch

from ctp.util import make_batches
from ctp.clutrr import Data, Instance

from ctp.geometric import GraphAttentionNetwork
from ctp.geometric import GraphConvolutionalNetwork

from ctp.geometric import VecBaselineNetworkV1
from ctp.geometric import VecBaselineNetworkV2
from ctp.geometric import Seq2VecEncoderFactory

from typing import Dict, List, Tuple, Optional

import logging

logger = logging.getLogger(os.path.basename(sys.argv[0]))
np.set_printoptions(linewidth=256, precision=4, suppress=True, threshold=sys.maxsize)

torch.set_num_threads(multiprocessing.cpu_count())

# PYTHONPATH=. python3 ./bin/geometric-cli.py
#  --train data/clutrr-emnlp/data_089907f8/*train*
#  --test data/clutrr-emnlp/data_089907f8/*test*

def to_data(instance: Instance,
            relation_to_idx: Dict[str, int],
            test_relation_to_idx: Dict[str, int],
            nb_entities: int,

            is_predicate: bool,
            predicate_to_idx: Dict[str, int],
            relation_to_predicate: Dict[str, str],
            test_predicate_to_idx: Dict[str, int],

            device: Optional[torch.device] = None) -> Tuple[GeometricData, Tuple[int, int]]:
    entity_lst = sorted({x for t in instance.story for x in {t[0], t[2]}})
    entity_to_idx = {e: i for i, e in enumerate(entity_lst)}

    x = torch.arange(nb_entities, device=device).view(-1, 1)

    edge_list = [(entity_to_idx[s], entity_to_idx[o]) for (s, _, o) in instance.story]
    edge_index = torch.tensor(list(zip(*edge_list)), dtype=torch.long, device=device)

    if is_predicate is True:
        edge_types = [predicate_to_idx[relation_to_predicate[p]] for (_, p, _) in instance.story]
        y = torch.tensor([test_predicate_to_idx[relation_to_predicate[instance.target[1]]]], device=device)
    else:
        edge_types = [relation_to_idx[p] for (_, p, _) in instance.story]
        y = torch.tensor([test_relation_to_idx[instance.target[1]]], device=device)

    edge_attr = torch.tensor(edge_types, dtype=torch.long, device=device).view(-1, 1)

    res = GeometricData(x=x, edge_index=edge_index, edge_attr=edge_attr, y=y)

    target_pair = (entity_to_idx[instance.target[0]], entity_to_idx[instance.target[2]])
    return res, target_pair

def to_batches(instances: List[Instance],
               batch_size: int,
               relation_to_idx: Dict[str, int],
               test_relation_to_idx: Dict[str, int],

               is_predicate: bool,
               predicate_to_idx: Dict[str, int],
               relation_to_predicate: Dict[str, str],
               test_predicate_to_idx: Dict[str, int],

               device: Optional[torch.device] = None) -> List[Tuple[Batch, List[int], Tensor, List[Instance]]]:
    nb_instances, res = len(instances), []
    batches = make_batches(nb_instances, batch_size)

    for batch_start, batch_end in batches:
        batch_instances = instances[batch_start:batch_end]
        max_nb_entities = max(i.nb_nodes for i in batch_instances)
        this_batch_size = len(batch_instances)

        batch_pairs = [
            to_data(i, relation_to_idx, test_relation_to_idx, max_nb_entities,
                    is_predicate, predicate_to_idx, relation_to_predicate, test_predicate_to_idx, device=device)
            for i in batch_instances
        ]

        batch_data: List[GeometricData] = [d for d, _ in batch_pairs]
        batch_targets: List[List[int]] = [[p[0], p[1]] for _, p in batch_pairs]

        max_node = max(i + 1 for b in batch_data for i in b.x[:, 0].cpu().numpy())

        batch = Batch.from_data_list(batch_data)
        slices = [max_node for _ in batch_data]

        targets = torch.tensor(batch_targets, dtype=torch.long, device=device).view(this_batch_size, 1, 2)

        res += [(batch, slices, targets, batch_instances)]
    return res

def main(argv):
    argparser = argparse.ArgumentParser('Geometric CLUTRR', formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    train_path = "data/clutrr-emnlp/data_test/64.csv"

    argparser.add_argument('--train', action='store', type=str, default=train_path)
    argparser.add_argument('--test', nargs='+', type=str, default=[])

    argparser.add_argument('--model', '-m', action='store', type=str, default='gat')

    # training params
    argparser.add_argument('--epochs', '-e', action='store', type=int, default=100)
    argparser.add_argument('--learning-rate', '-l', action='store', type=float, default=0.001)
    argparser.add_argument('--batch-size', '-b', action='store', type=int, default=100)

    argparser.add_argument('--embedding-size', '-k', action='store', type=int, default=100)
    argparser.add_argument('--edge-embedding-size', '-K', action='store', type=int, default=20)
    argparser.add_argument('--hidden-size', action='store', type=int, default=100)
    argparser.add_argument('--nb-filters', action='store', type=int, default=100)
    argparser.add_argument('--nb-heads', action='store', type=int, default=3)
    argparser.add_argument('--nb-rounds', action='store', type=int, default=3)
    argparser.add_argument('--nb-highway', action='store', type=int, default=2)

    argparser.add_argument('--seed', action='store', type=int, default=0)

    argparser.add_argument('--evaluate-every', '-V', action='store', type=int, default=1)

    argparser.add_argument('--v2', action='store_true', default=False)
    argparser.add_argument('--predicate', action='store_true', default=False)

    args = argparser.parse_args(argv)

    train_path = args.train
    test_paths = args.test

    model_name = args.model

    nb_epochs = args.epochs
    learning_rate = args.learning_rate
    batch_size = args.batch_size

    embedding_size = args.embedding_size
    edge_embedding_size = args.edge_embedding_size
    hidden_size = args.hidden_size
    nb_filters = args.nb_filters
    nb_heads = args.nb_heads
    nb_rounds = args.nb_rounds
    nb_highway = args.nb_highway

    seed = args.seed

    evaluate_every = args.evaluate_every

    is_v2 = args.v2
    is_predicate = args.predicate

    np.random.seed(seed)
    torch.manual_seed(seed)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    logger.info(f'Device: {device}')

    if torch.cuda.is_available():
        torch.set_default_tensor_type(torch.cuda.FloatTensor)

    data = Data(train_path=train_path, test_paths=test_paths)

    entity_lst, _, relation_lst = data.entity_lst, data.predicate_lst, data.relation_lst
    predicate_lst = data.predicate_lst

    relation_to_predicate = data.relation_to_predicate

    test_relation_lst = ["aunt", "brother", "daughter", "daughter-in-law", "father", "father-in-law", "granddaughter",
                         "grandfather", "grandmother", "grandson", "mother", "mother-in-law", "nephew", "niece",
                         "sister", "son", "son-in-law", "uncle"]

    test_predicate_lst = sorted({relation_to_predicate[r] for r in test_relation_lst})

    relation_to_idx = {r: i for i, r in enumerate(relation_lst)}
    test_relation_to_idx = {r: i for i, r in enumerate(test_relation_lst)}

    predicate_to_idx = {p: i for i, p in enumerate(predicate_lst)}
    test_predicate_to_idx = {p: i for i, p in enumerate(test_predicate_lst)}

    nb_nodes = len(entity_lst)
    nb_edge_types = len(relation_lst)
    nb_targets = len(test_relation_lst)

    if is_predicate is True:
        nb_edge_types = len(predicate_lst)
        nb_targets = len(test_predicate_lst)

    nb_instances = len(data.train)
    batches = to_batches(data.train, batch_size=batch_size,
                         relation_to_idx=relation_to_idx,
                         test_relation_to_idx=test_relation_to_idx,

                         is_predicate=is_predicate,
                         predicate_to_idx=predicate_to_idx,
                         relation_to_predicate=relation_to_predicate,
                         test_predicate_to_idx=test_predicate_to_idx,

                         device=device)

    if model_name in {'gat'}:
        model = GraphAttentionNetwork(nb_nodes=nb_nodes, nb_edge_types=nb_edge_types, target_size=nb_targets,
                                      nb_heads=nb_heads, embedding_size=embedding_size,
                                      edge_embedding_size=edge_embedding_size, nb_rounds=nb_rounds)
    elif model_name in {'gcn'}:
        model = GraphConvolutionalNetwork(nb_nodes=nb_nodes, nb_edge_types=nb_edge_types, target_size=nb_targets,
                                          embedding_size=embedding_size, edge_embedding_size=edge_embedding_size,
                                          nb_rounds=nb_rounds)
    else:
        encoder_factory = Seq2VecEncoderFactory()
        encoder = encoder_factory.build(name=model_name, embedding_dim=embedding_size, hidden_size=hidden_size,
                                        num_filters=nb_filters, num_heads=nb_heads, num_highway=nb_highway)

        if is_v2 is False:
            model = VecBaselineNetworkV1(nb_nodes=nb_nodes, nb_edge_types=nb_targets, relation_lst=relation_lst,
                                         encoder=encoder, embedding_size=embedding_size)
        else:
            model = VecBaselineNetworkV2(nb_nodes=nb_nodes, nb_edge_types=nb_targets, relation_lst=relation_lst,
                                         encoder=encoder, embedding_size=embedding_size)

    model = model.to(device)

    params_lst = nn.ParameterList([p for p in model.parameters()])
    optimizer = torch.optim.Adam(params_lst, lr=learning_rate)

    def test(test_set) -> float:
        correct = 0
        model.eval()

        test_batches = to_batches(test_set, batch_size=batch_size, relation_to_idx=relation_to_idx,
                                  test_relation_to_idx=test_relation_to_idx,

                                  is_predicate=is_predicate,
                                  predicate_to_idx=predicate_to_idx,
                                  relation_to_predicate=relation_to_predicate,
                                  test_predicate_to_idx=test_predicate_to_idx,

                                  device=device)

        for test_batch, test_slices, test_targets, test_instances in test_batches:
            test_logits = model(test_batch, test_slices, test_targets, test_instances)
            test_predictions = test_logits.max(dim=1)[1]
            correct += test_predictions.eq(test_batch.y).sum().item()
        return correct / len(test_set)

    for epoch in range(1, nb_epochs + 1):
        loss_total = 0.0
        model.train()

        for batch, slices, targets, instances in batches:
            logits = model(batch, slices, targets, instances)

            assert logits.shape[1] == len(test_relation_lst if not is_predicate else test_predicate_lst)

            loss = F.cross_entropy(logits, batch.y, reduction='sum')
            loss_total += loss.item()

            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

        # train_accuracy = test(data.train)
        print(f'Epoch: {epoch:03d}, Train Loss: {loss_total / nb_instances:.7f}')

        if epoch % evaluate_every == 0:
            for name in data.test:
                test_accuracy = test(data.test[name])
                print(f'Epoch: {epoch:03d}, Test Set: {name}, Accuracy: {test_accuracy:.7f}')

    logger.info("Training finished")

if __name__ == '__main__':
    logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
    print(' '.join(sys.argv))
    main(sys.argv[1:])


from pathlib import Path

import math
import os
import numpy as np
import torch
from tqdm import tqdm
from PIL import Image

from vp_suite.base.base_dataset import VPDataset, VPData
import vp_suite.constants as constants

class MovingMNISTDataset(VPDataset):
    r"""
    Dataset class for the dataset "Moving MNIST", as firstly encountered in
    "Unsupervised Learning of Video Representations using LSTMs" by Srivastava et al.
    (https://arxiv.org/pdf/1502.04681v3.pdf).

    Each sequence depicts two digits from the MNIST dataset moving linearly in front of a black background,
    occasionally bouncing off the wall and overlapping each other.

    For downloading and preparing the dataset, scripts have been developed by Tencia Lee,
    ported to python 3 by Praateek Mahajan (https://gist.github.com/praateekmahajan/b42ef0d295f528c986e2b3a0b31ec1fe)
    and further modified here.
    """

    NAME = "Moving MNIST"
    DEFAULT_DATA_DIR = constants.DATA_PATH / "moving_mnist"
    MIN_SEQ_LEN = 20
    ACTION_SIZE = 0
    DATASET_FRAME_SHAPE = (64, 64, 3)

    train_to_val_ratio = 0.96

    def __init__(self, split, **dataset_kwargs):
        super(MovingMNISTDataset, self).__init__(split, **dataset_kwargs)
        self.NON_CONFIG_VARS.extend(["data_ids", "data_fps"])

        self.data_dir = str((Path(self.data_dir) / split).resolve())
        self.data_ids = sorted(os.listdir(self.data_dir))
        self.data_fps = [os.path.join(self.data_dir, image_id) for image_id in self.data_ids]

    def __len__(self):
        return len(self.data_fps)

    def __getitem__(self, i) -> VPData:
        if not self.ready_for_usage:
            raise RuntimeError("Dataset is not yet ready for usage (maybe you forgot to call set_seq_len()).")

        rgb_raw = np.load(self.data_fps[i])  # [t', h, w]
        rgb_raw = np.expand_dims(rgb_raw, axis=-1).repeat(3, axis=-1) # [t', h, w, c]
        rgb_raw = rgb_raw[:self.seq_len:self.seq_step]  # [t, h, w, c]
        rgb = self.preprocess(rgb_raw)

        actions = torch.zeros((self.total_frames, 1))  # [t, a], actions should be disregarded in training logic

        data = { "frames": rgb, "actions": actions }
        return data

    def download_and_prepare_dataset(self):
        frame_size = (64, 64)
        num_frames = 20  # length of each sequence
        digit_size = 28  # size of mnist digit within frame
        digits_per_image = 2  # number of digits in each frame
        d_path = self.DEFAULT_DATA_DIR
        d_path.mkdir(parents=True)

        # training sequences
        train_seqs = 60000
        print("generating training set...")
        train_data = generate_moving_mnist(d_path, training=True, shape=frame_size, num_frames=num_frames,
                                           num_images=train_seqs, original_size=digit_size,
                                           nums_per_image=digits_per_image)
        print("saving training set...")
        save_generated_mmnist(train_data, train_seqs, frame_size, d_path / "train")

        # testing sequences
        test_seqs = 10000
        print("generating test set...")
        test_data = generate_moving_mnist(d_path, training=False, shape=frame_size, num_frames=num_frames,
                                          num_images=test_seqs, original_size=digit_size,
                                          nums_per_image=digits_per_image)
        print("saving test set...")
        save_generated_mmnist(test_data, test_seqs, frame_size, d_path / "test")

# === MMNIST data preparation tools ============================================

def save_generated_mmnist(data: np.ndarray, seqs: int, frame_size: (int, int), out_path: Path):
    r"""
    Save generated data per-sequence to specified out path.

    Args:
        data (np.ndarray): The generated data to save.
        seqs (int): The number of generated sequences.
        frame_size ((int, int)): The frame size.
        out_path (Path): The path where the data should be saved.
    """
    out_path.mkdir()
    num_frames = data.shape[0] // seqs
    data = data.reshape(seqs, num_frames, *frame_size)
    for i in tqdm(range(data.shape[0])):
        cur_out_fp = out_path / f"seq_{i:05d}.npy"
        np.save(str(cur_out_fp), data[i])

# helper functions
def arr_from_img(im, mean: float = 0, std: float = 1):
    r"""
    Convert image to array.

    Args:
        im(): Image.
        mean(float): Mean to subtract.
        std(float): Standard Deviation to subtract.

    Returns:
        Image in np.float32 format, in width height channel format. With values in range 0,1
        Shift means subtract by certain value. Could be used for mean subtraction.

    """
    width, height = im.size
    arr = im.getdata()
    c = int(np.product(arr.size) / (width * height))

    return (np.asarray(arr, dtype=np.float32).reshape((height, width, c)).transpose(2, 1, 0) / 255. - mean) / std

def img_from_arr(X: np.ndarray, index: int, mean: float = 0, std: float = 1):
    r"""
    Convert array to image.

    Args:
        X(np.ndarray): Dataset of shape N x C x W x H.
        index(int): Index of image we want to fetch.
        mean(float): Mean to add.
        std(float): Standard Deviation to add.

    Returns:
        Image with dimensions H x W x C or H x W if it's a single channel image.
    """
    ch, w, h = X.shape[1], X.shape[2], X.shape[3]
    ret = (((X[index] + mean) * 255.) * std).reshape(ch, w, h).transpose(2, 1, 0).clip(0, 255).astype(np.uint8)
    if ch == 1:
        ret = ret.reshape(h, w)
    return ret

def load_dataset(d_path: Path, training: bool):
    r"""
    Loads MNIST from the web on demand.

    Args:
        d_path (Path): The path where the downloaded digits should be stored.
        training (bool): Whether to use the training images (True) or the test images (False).

    Returns: The loaded MNIST images.

    """
    from vp_suite.utils.utils import download_from_url
    import gzip

    def load_mnist_images(filename):
        if not os.path.exists(filename):
            mnist_source = 'http://yann.lecun.com/exdb/mnist/'
            fname_ = filename.split("/")[-1]
            download_from_url(mnist_source + fname_, filename)
        with gzip.open(filename, 'rb') as f:
            data = np.frombuffer(f.read(), np.uint8, offset=16)
        data = data.reshape(-1, 1, 28, 28).transpose(0, 1, 3, 2)
        return data / np.float32(255)

    if training:
        return load_mnist_images(str(d_path / 'train-images-idx3-ubyte.gz'))
    return load_mnist_images(str(d_path / 't10k-images-idx3-ubyte.gz'))

def generate_moving_mnist(d_path: Path, training: bool = False, shape: (int, int) = (64, 64),
                          num_frames: int = 30, num_images: int = 100, original_size: int = 28,
                          nums_per_image: int = 2):
    r"""
    Generate sequences of moving MNIST digits by moving them around between frames.

    Args:
        training (bool): Used to decide if downloading/generating training set or test set.
        shape ((int, int)): Shape we want for our moving images (new_width and new_height).
        num_frames (int): Number of frames in a particular movement/animation/gif.
        num_images (int): Number of movement/animations/gif to generate.
        original_size (int): Real size of the images (eg: MNIST is 28x28).
        nums_per_image (int): Digits per movement/animation/gif.

    Returns:
        Dataset of np.uint8 type with dimensions num_frames * num_images x 1 x new_width x new_height

    """
    mnist = load_dataset(d_path, training)
    width, height = shape

    # Get how many pixels can we move around a single image
    lims = (x_lim, y_lim) = width - original_size, height - original_size

    # Create a dataset of shape of num_frames * num_images x 1 x new_width x new_height
    # Eg : 3000000 x 1 x 64 x 64
    dataset = np.empty((num_frames * num_images, 1, width, height), dtype=np.uint8)

    for img_idx in tqdm(range(num_images)):
        # Randomly generate direction, speed and velocity for both images
        direcs = np.pi * (np.random.rand(nums_per_image) * 2 - 1)
        speeds = np.random.randint(5, size=nums_per_image) + 2
        veloc = np.asarray([(speed * math.cos(direc), speed * math.sin(direc)) for direc, speed in zip(direcs, speeds)])
        # Get a list containing two PIL images randomly sampled from the database
        mnist_images = [Image.fromarray(img_from_arr(mnist, r, mean=0)).resize((original_size, original_size),
                                                                               Image.ANTIALIAS) \
                        for r in np.random.randint(0, mnist.shape[0], nums_per_image)]
        # Generate tuples of (x,y) i.e initial positions for nums_per_image (default : 2)
        positions = np.asarray([(np.random.rand() * x_lim, np.random.rand() * y_lim) for _ in range(nums_per_image)])

        # Generate new frames for the entire num_framesgth
        for frame_idx in range(num_frames):

            canvases = [Image.new('L', (width, height)) for _ in range(nums_per_image)]
            canvas = np.zeros((1, width, height), dtype=np.float32)

            # In canv (i.e Image object) place the image at the respective positions
            # Super impose both images on the canvas (i.e empty np array)
            for i, canv in enumerate(canvases):
                canv.paste(mnist_images[i], tuple(positions[i].astype(int)))
                canvas += arr_from_img(canv, mean=0)

            # Get the next position by adding velocity
            next_pos = positions + veloc

            # Iterate over velocity and see if we hit the wall
            # If we do then change the  (change direction)
            for i, pos in enumerate(next_pos):
                for j, coord in enumerate(pos):
                    if coord < -2 or coord > lims[j] + 2:
                        veloc[i] = list(list(veloc[i][:j]) + [-1 * veloc[i][j]] + list(veloc[i][j + 1:]))

            # Make the permanent change to position by adding updated velocity
            positions = positions + veloc

            # Add the canvas to the dataset array
            dataset[img_idx * num_frames + frame_idx] = (canvas * 255).clip(0, 255).astype(np.uint8)

    return dataset


'''

Authors: Renke Huang, Qiuhua Huang
Contact: qiuhua.huang@pnnl.gov

'''

import sys, os, time, parser, math
import numpy as np
import gym, ray
import logz, optimizers, utils
from obserfilter import RunningStat
from policy_LSTM import *
import socket
from shared_noise import *
from GridPackPowerDynSimEnvDef_v2 import GridPackPowerDynSimEnv
from statistics import *

folder_dir = os.path.abspath(os.path.join(os.getcwd(), ".."))
print('-----------------root path of the rlgc:', folder_dir)

FAULT_BUS_CANDIDATES = [0, 1, 2, 3, 4, 5, 6, 7, 8]
FAULT_START_TIME = 1.0
FTD_CANDIDATES = [0.00, 0.05, 0.1]
FAULT_CASES = [(FAULT_BUS_CANDIDATES[i], FAULT_START_TIME, FTD_CANDIDATES[j]) for i in range(len(FAULT_BUS_CANDIDATES))
               for j in range(len(FTD_CANDIDATES))]

GKE = True # whether the training is running on GKE
EIOC = True # set this to be true if running on GKE

b_debug_timing = True
NUM_CORES = 18

OB_DIM = 301
AC_DIM = 81

if GKE:
    simu_input_file = folder_dir + '/testData/tamu2000/input_tamu2000_step005_rmgensmall30_withob_2acloads_v33_nocsv_superlu.xml'
else:
    simu_input_file = folder_dir + '/testData/tamu2000/input_tamu2000_step005_rmgensmall30_withob_2acloads_v33_nocsv_eioc.xml'

rl_config_file = folder_dir + '/testData/tamu2000/json/RLGC_RL_tamu2000_loadShedding_zone3_gp.json'

# ======================================================================

@ray.remote
class SingleRolloutWorker(object):
    '''
    responsible for doing one single rollout in the env
    '''

    def __init__(self, rollout_length, policy_params):

        print('\n\n\n -----------------------Set Env=GridPackPowerDynSimEnv------------------------\n\n\n')
        self.env = GridPackPowerDynSimEnv(simu_input_file, rl_config_file, force_symmetric_continuous_action=True)
        print('\n\n\n ----------------------------Set Env Done-----------------------------\n\n\n')
        self.rollout_length = rollout_length
        self.policy_type = policy_params['type']
        if policy_params['type'] == 'linear':
            self.policy = LinearPolicy(policy_params)
        elif policy_params['type'] == 'nonlinear':
            self.policy = FullyConnectedNeuralNetworkPolicy(policy_params)
        elif policy_params['type'] == 'LSTM':
            self.policy = LSTMPolicy(policy_params)
        else:
            raise NotImplementedError
        time.sleep(5)

    def single_rollout(self, fault_tuple, weights, ob_mean, ob_std):
        # one SingleRolloutWorker is only doing one fault case in an iter
        # fault_tuple = FAULT_CASES[fault_case_id]
        total_reward = 0.
        steps = 0
        self.policy.update_weights(weights)

        # us RS for collection of observation states; restart every rollout
        self.RS = RunningStat(shape=(OB_DIM,))

        t1 = time.time()
        ob = self.env.validate(case_Idx=0, fault_bus_idx=fault_tuple[0],
                               fault_start_time=fault_tuple[1], fault_duration_time=fault_tuple[2])
        t_validate = time.time() - t1

        if self.policy_type == 'LSTM':
            self.policy.reset()

        t_act = 0
        t_step = 0
        for _ in range(self.rollout_length):
            ob = np.asarray(ob, dtype=np.float64)
            self.RS.push(ob)
            normal_ob = (ob - ob_mean) / (ob_std + 1e-8)
            t3 = time.time()
            action_org = self.policy.act(normal_ob)
            t4 = time.time()
            ob, reward, done, _ = self.env.step(action_org)
            t5 = time.time()
            t_act += t4 - t3
            t_step += t5 - t4
            total_reward += reward
            steps += 1
            if done:
                break
        t2 = time.time()

        if b_debug_timing:
            return {'reward': total_reward, 'step': steps, 'time': [{'Total Time': t2 - t1},
                                                                    {'Reset Time': t_validate},
                                                                    {'Action Time': t_act},
                                                                    {'Step Time': t_step} ]}
        else:
            return {'reward': total_reward, 'step': steps, 'time': [t2 - t1, t_validate, t_act, t_step]}

    def return_RS(self):
        return self.RS

    def close_env(self):
        self.env.close_env()
        return

def run_ars_test(params):

    # set policy parameters.
    if params['policy_file'] != "":
        nonlin_policy_org = np.load(params['policy_file'], allow_pickle=True)
        nonlin_policy = nonlin_policy_org['arr_0']

        w_M = nonlin_policy[0].copy()
        w_mean = nonlin_policy[1].copy()
        w_std = nonlin_policy[2].copy()

        policy_params = {'type': params['policy_type'],
                         'policy_network_size': params['policy_network_size'],
                         'ob_dim': OB_DIM,
                         'ac_dim': AC_DIM,
                         'weights': w_M,
                         'w_mean': w_mean,
                         'w_std': w_std,
                         }
    else:
        # set policy parameters.
        policy_params = {'type': params['policy_type'],
                         'policy_network_size': params['policy_network_size'],
                         'ob_dim': OB_DIM,
                         'ac_dim': AC_DIM}

    single_rollout_workers = []
    nsingleworks = min( len(FAULT_CASES), params['onedirection_numofcasestorun'])

    #-----------just test one single worker----------------------------------------

    for i in range(nsingleworks):
        single_rollout_workers.append(SingleRolloutWorker.remote(rollout_length=150,
                                                                      policy_params=policy_params))
        time.sleep(2)

    weights_id = ray.put(w_M)
    ob_mean = ray.put(w_mean)
    ob_std = ray.put(w_std)

    # -----------just test one single worker----------------------------------------
    icase = 1
    tmp_id = single_rollout_workers[icase].single_rollout.remote(
                fault_tuple=FAULT_CASES[icase], weights=weights_id, ob_mean=ob_mean,
                ob_std=ob_std)
    tmp_result = ray.get(tmp_id)

    print ('-----------testing just one single roll out for fault tuple: ', FAULT_CASES[icase])
    print (tmp_result)
    print('----------- finished testing for just one single roll out for fault tuple: ', FAULT_CASES[icase])
    print ('\n')

    for itr in range (5):

        results_ids = []

        # -----------just test different number of works for rolling out fully distributed for each itration-----------------
        nworkers = nsingleworks//(itr+1)

        print ('---------start multiple single roll outs, itr: ', itr, ', n-workers: ', nworkers)
        print('\n')
        for i in range(nworkers):
            results_ids.append(single_rollout_workers[i].single_rollout.remote(
                fault_tuple=FAULT_CASES[i], weights=weights_id, ob_mean=ob_mean,
                ob_std=ob_std))

        results = ray.get(results_ids)
        reward_list = []
        step_list = []
        time_list = []
        for result in results:
            reward_list.append(result['reward'])
            step_list.append(result['step'])
            time_list.append(result['time'])
        reward_ave = mean(reward_list)

        print ('--------------n-works: ', nworkers, ', average reward: ', reward_ave)
        for itmp in range(len(reward_list)):
            print('Reward: ', reward_list[itmp], ', Steps: ', step_list[itmp], ', Timing: ', time_list[itmp])
        # ----- delete the weight in the global table in ray before return

        print('\n')

    print ('-------------------finished all testing!!!!!!!!!!!!!--------------')
    del weights_id
    del ob_mean
    del ob_std

    return

# Setting
if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('--n_iter', '-n', type=int, default=500)  # Number of iterations
    if EIOC:  # EIOC test
        parser.add_argument('--n_directions', '-nd', type=int, default=1)
    else:  # Constance training
        parser.add_argument('--n_directions', '-nd', type=int, default=128)
    if EIOC:  # EIOC test
        parser.add_argument('--deltas_used', '-du', type=int, default=2)
    else:  # Constance training
        parser.add_argument('--deltas_used', '-du', type=int, default=64)
    parser.add_argument('--step_size', '-s', type=float, default=1)  # default=0.02
    parser.add_argument('--delta_std', '-std', type=float, default=2)  # default=.03
    parser.add_argument('--decay', type=float, default=0.995)  # decay of step_size and delta_std
    parser.add_argument('--rollout_length', '-r', type=int, default=150)
    parser.add_argument('--seed', type=int, default=589)  # Seed Number for randomization
    parser.add_argument('--policy_type', type=str, default='LSTM')
    parser.add_argument('--dir_path', type=str,
                        default='ars_tamu2000_testrayandgridpack')  # Folder Name for outputs
    if EIOC:  # EIOC test
        parser.add_argument('--save_per_iter', type=int, default=1)  # save the .npz file per x iterations
    else:  # Constance training
        parser.add_argument('--save_per_iter', type=int, default=10)  # save the .npz file per x iterations
    parser.add_argument('--policy_network_size', type=list, default=[64, 64])

    parser.add_argument('--onedirection_numofcasestorun', type=int,
                        default=18)  # len(FAULT_CASES))  # For each direction, how many cases to run, default is to run all the fault
    # please set onedirection_numofcasestorun to be the multiplie of 3, such as 6,9,12,15,18, etc

    # please specify number of cores here; it is to help determine the number of workers
    if EIOC:  # EIOC test
        parser.add_argument('--cores', type=int, default=NUM_CORES)  # how many cores available as of hardware, EIOC test
    else:  # Constance training
        parser.add_argument('--cores', type=int, default=NUM_CORES)  # how many cores available as of hardware
    parser.add_argument('--policy_file', type=str, default="training_results/1_3faultcases/nonlinear_policy_plus490.npz")

    # the parameters controlling iterations for convergence
    # n_iter is maximum number of iterations
    # tol_p is tolerance for the percentage of average reward change between iterations
    # tol_steps is the total iterations for maintaining tol_p
    parser.add_argument('--tol_p', type=float, default=0.001)
    parser.add_argument('--tol_steps', type=int, default=100)
    local_ip = socket.gethostbyname(socket.gethostname())

    # Init Ray in Cluster, use log_to_driver=True for printing messages from remote nodes
    ##!!#####################################################################################################
    # !! if you are using EIOC, use this line of code below
    if EIOC:  # EIOC test
        ray.init(address="localhost:6379", log_to_driver=False)
    else:
        # !! if you are using Constance, use this line of code below
        ray.init(temp_dir=os.environ["tmpfolder"], redis_address=os.environ["ip_head"], log_to_driver=False)
    ##!!######################################################################################################

    args = parser.parse_args()
    params = vars(args)
    t1 = time.time()
    run_ars_test(params)
    t2 = time.time()
    print('Total Time run_ars:', t2 - t1)


# -*- coding: utf-8 -*-
import numpy as np
from scipy import interpolate
from .BaseDataProtocol.SABProtocol import dtype_sab
from .util import _prepare_for_read, _unpack_from_buf, julian2date, \
    get_radar_info, make_time_unit_str, get_radar_sitename
from netCDF4 import date2num
from ..core.NRadar import PRD
from ..configure.pyart_config import get_metadata, get_fillvalue
from ..configure.default_config import CINRAD_field_mapping, _LIGHT_SPEED
from ..core.PyartRadar import Radar

class SABBaseData(object):
    """
    解码SA/SB/CB/SC2.0的雷达数据，仅仅对数据（dBZ, V, W）做了转换
    """

    def __init__(self, filename, station_lon=None, station_lat=None, station_alt=None):
        """
        :param filename:  radar basedata filename
        :param station_lon:  radar station longitude //units: degree east
        :param station_lat:  radar station latitude //units:degree north
        :param station_alt:  radar station altitude //units: meters
        """
        super(SABBaseData, self).__init__()
        self.filename = filename
        self.station_lon = station_lon
        self.station_lat = station_lat
        self.station_alt = station_alt
        self.fid = _prepare_for_read(self.filename)
        self.RadialNum, self.nrays = self._RadialNum_SAB_CB()  ##检查文件有无问题
        self.radial = self._parse_radial()
        status = np.array([istatus['RadialStatus'] for istatus in self.radial[:]])
        self.sweep_start_ray_index = np.where((status == 0) | (status == 3))[0]
        self.sweep_end_ray_index = np.where((status == 2) | (status == 4))[0]
        self.nsweeps = len(self.sweep_start_ray_index)
        self.fid.close()

    def _RadialNum_SAB_CB(self):
        """f: a file-like object was provided, 确定雷达数据的径向字节长度"""
        assert self.fid.read(28)[14:16] == b'\x01\x00', 'file in not a valid SA/SB/CB file!'
        self.fid.seek(0, 0)  ##移动到开头
        data_len = len(self.fid.read())
        assert (data_len % 2432 == 0) | (data_len % 4132 == 0) | (data_len % 3132 == 0), "file size has problems!"
        ###判断雷达数据类型SA/SB 或者 CB
        if data_len % 2432 == 0:
            RadialNum = 2432
            self.Type = "SAB"
        elif data_len % 4132 == 0:
            RadialNum = 4132
            self.Type = 'CB'
        else:
            RadialNum = 3132
            self.Type = 'SC'
        self.fid.seek(0, 0)  ##移动到开头
        return RadialNum, int(data_len / RadialNum)

    def _parse_radial(self):
        """
        循环读取所有径向数据
        :param fid:
        :return:
        """
        radial = []
        for _ in range(self.nrays):
            radial.append(self._parse_radial_single(self.fid.read(self.RadialNum)))
        return radial

    def _parse_radial_single(self, radial_buf):
        Radial = {}
        RadialHeader, size_tmp = _unpack_from_buf(radial_buf, 0, dtype_sab.RadialHeader())
        Radial.update(RadialHeader)
        dBZ = np.frombuffer(radial_buf[RadialHeader['PtrOfReflectivity'] + dtype_sab.InfSize:\
                                                     RadialHeader['PtrOfReflectivity'] + dtype_sab.InfSize +\
                                                     RadialHeader['GatesNumberOfReflectivity']],\
                                                    dtype="u1")
        V = np.frombuffer(radial_buf[RadialHeader['PtrOfVelocity'] + dtype_sab.InfSize: \
                                                     RadialHeader['PtrOfVelocity'] + dtype_sab.InfSize + \
                                                     RadialHeader['GatesNumberOfDoppler']], \
                                                    dtype="u1")
        W = np.frombuffer(radial_buf[RadialHeader['PtrOfSpectrumWidth'] + dtype_sab.InfSize: \
                                                   RadialHeader['PtrOfSpectrumWidth']  + dtype_sab.InfSize+ \
                                                   RadialHeader['GatesNumberOfDoppler']], \
                                                    dtype="u1")
        Radial['fields'] = {}
        Radial['fields']['dBZ'] = np.where(dBZ > 1, (dBZ.astype(int) - 2) / 2. - 32, np.nan).astype(np.float32)
        Radial['fields']['V'] = np.where(V > 1, (V.astype(int) - 2) / 2. - 63.5, np.nan).astype(np.float32)
        Radial['fields']['W'] = np.where(W > 1, (W.astype(int) - 2) / 2. - 63.5, np.nan).astype(np.float32)
        return Radial

    def get_nyquist_velocity(self):
        """get nyquist vel per ray
        获取每根径向的不模糊速度
        :return:(nRays)
        """
        return np.array([iradial['Nyquist'] / 100. for iradial in self.radial])

    def get_unambiguous_range(self):
        """
        获取每根径向的不模糊距离 units:km
        :return:(nRays)
        """
        return np.array([iradial['URange'] / 10. for iradial in self.radial])

    def get_scan_time(self):
        """
        获取每根径向的扫描时间
        :return:(nRays)
        """
        return np.array([julian2date(iradial['JulianDate'], iradial['mSends']) for iradial in self.radial])

    def get_sweep_end_ray_index(self):
        """
        获取每个sweep的结束的index，包含在内
        :return:(nsweep)
        """
        return self.sweep_end_ray_index

    def get_sweep_start_ray_index(self):
        """
        获取每个sweep的开始的index
        :return:(nsweep)
        """
        return self.sweep_start_ray_index

    def get_rays_per_sweep(self):
        """
        获取每个sweep的径向数
        :return:(nsweep)
        """
        return self.sweep_end_ray_index - self.sweep_start_ray_index + 1

    def get_azimuth(self):
        """
        获取每根径向的方位角
        :return:(nRays)
        """
        return np.array([iradial['AZ'] / 8. * 180. / 4096. for iradial in self.radial])

    def get_elevation(self):
        """
        获取每根径向的仰角
        :return: (nRays)
        """
        return np.array([iradial['El'] / 8. * 180. / 4096. for iradial in self.radial])

    def get_latitude_longitude_altitude_frequency(self):
        """
        获取经纬度高度，雷达频率
        :return:lat, lon, alt, frequency
        """
        lat, lon, alt, frequency = get_radar_info(self.filename)
        if self.station_lon is not None:
            lon = self.station_lon
        if self.station_lat is not None:
            lat = self.station_lat
        if self.station_alt is not None:
            alt = self.station_alt
        return lat, lon, alt, frequency

    def get_scan_type(self):
        """
        获取扫描的类型
        :return:
        """
        return "ppi"

    def get_sitename(self):
        return get_radar_sitename(self.filename)

class SAB2NRadar(object):
    """到NusitRadar object 的桥梁"""

    def __init__(self, SAB):
        self.SAB = SAB
        self.v_index_alone = self.get_v_idx()
        self.dBZ_index_alone = self.get_dbz_idx()
        self.dBZ_Res = self.SAB.radial[0]["GateSizeOfReflectivity"] ##反射率因子的分辨率
        for index_with_dbz, index_with_v in zip(self.dBZ_index_alone, self.v_index_alone):
            assert abs(self.SAB.get_elevation()[index_with_v] - \
                       self.SAB.get_elevation()[index_with_dbz]) < 0.5, "warning! maybe it is a problem."
            self.interp_dBZ(index_with_dbz, index_with_v)
        ind_remove = self.get_reomve_radial_num()
        self.radial = [iray for ind, iray in enumerate(self.SAB.radial) if ind not in ind_remove]
        self.nrays = len(self.radial)
        self.nsweeps = self.SAB.nsweeps - self.dBZ_index_alone.size
        status = np.array([istatus['RadialStatus'] for istatus in self.radial[:]])
        self.sweep_start_ray_index = np.where((status == 0) | (status == 3))[0]
        self.sweep_end_ray_index = np.where((status == 2) | (status == 4))[0]
        self.scan_type = self.SAB.get_scan_type()
        self.latitude, self.longitude, self.altitude, self.frequency = \
            self.SAB.get_latitude_longitude_altitude_frequency()
        self.bins_per_sweep = self.get_nbins_per_sweep()
        self.max_bins = self.bins_per_sweep.max()
        self.range = self.get_range_per_radial(self.max_bins)  ##所有的数据向多普勒数据对齐
        self.azimuth = self.get_azimuth()
        self.elevation = self.get_elevation()
        self.fields = self._get_fields()
        self.sitename = self.SAB.get_sitename()

    def get_reomve_radial_num(self):
        """获得需要remove的radial的index"""
        """获得需要remove的radial的index"""
        dBZ_alone = self.get_dbz_idx()
        index_romove = []
        for isweep in dBZ_alone:
            index_romove.extend(range(self.SAB.sweep_start_ray_index[isweep], \
                                      self.SAB.sweep_end_ray_index[isweep] + 1))
        return index_romove

    def get_v_idx(self):
        """获取需要插值的sweep, 插值到有径向速度仰角"""
        flag = np.array([((self.SAB.radial[idx]['fields']["V"].size != 0) and
                          (self.SAB.radial[idx]['fields']["dBZ"].size == 0)) \
                         for idx in self.SAB.sweep_start_ray_index])
        return np.where(flag == 1)[0]

    def get_dbz_idx(self):
        """获取含有dbz的sweep"""
        flag = np.array([((self.SAB.radial[idx]['fields']["V"].size == 0) and
                          (self.SAB.radial[idx]['fields']["dBZ"].size != 0)) \
                         for idx in self.SAB.sweep_start_ray_index])
        return np.where(flag == 1)[0]

    def interp_dBZ(self, field_with_dBZ_num, field_without_dBZ_num):
        """
        将dBZ插值到不含dBZ的仰角
        :param field_with_dBZ_num: 要插值的sweep num, （从0开始）
        :param field_without_dBZ_num: 要插值到的sweep num, (从0开始)  which to evaluate the interpolated values
        :return:
        """
        azimuth = self.SAB.get_azimuth()  ##
        assert (field_with_dBZ_num + 1) == field_without_dBZ_num, "check interp sweep!"
        dbz_az = azimuth[self.SAB.sweep_start_ray_index[field_with_dBZ_num]: \
                         self.SAB.sweep_end_ray_index[field_with_dBZ_num] + 1]
        v_az = azimuth[self.SAB.sweep_start_ray_index[field_without_dBZ_num]: \
                       self.SAB.sweep_end_ray_index[field_without_dBZ_num] + 1]
        dbz_idx = np.argmin(np.abs(dbz_az.reshape(-1, 1) - v_az.reshape(1, -1)), axis=0) + \
                  self.SAB.sweep_start_ray_index[field_with_dBZ_num]
        v_idx = np.arange(self.SAB.sweep_start_ray_index[field_without_dBZ_num], \
                          self.SAB.sweep_end_ray_index[field_without_dBZ_num] + 1)
        for ind_dbz, ind_v in zip(dbz_idx, v_idx):
            self.SAB.radial[ind_v]["fields"]['dBZ'] = self.SAB.radial[ind_dbz]["fields"]['dBZ']

    def get_azimuth(self):
        """
        获取每根径向的方位角
        :return:(nRays)
        """
        return np.array([iradial['AZ'] / 8. * 180. / 4096. for iradial in self.radial])

    def get_elevation(self):
        """
        获取每根径向的仰角
        :return: (nRays)
        """
        return np.array([iradial['El'] / 8. * 180. / 4096. for iradial in self.radial])

    def get_rays_per_sweep(self):
        """
        获取每个sweep的径向数
        :return:(nsweep)
        """
        return self.sweep_end_ray_index - self.sweep_start_ray_index + 1

    def get_scan_time(self):
        """
        获取每根径向的扫描时间
        :return:(nRays)
        """
        return np.array([julian2date(iradial['JulianDate'], iradial['mSends']) for iradial in self.radial])

    def get_nyquist_velocity(self):
        """get nyquist vel per ray
        获取每根径向的不模糊速度
        :return:(nRays)
        """
        return np.array([iradial['Nyquist'] / 100. for iradial in self.radial])

    def get_unambiguous_range(self):
        """
        获取每根径向的不模糊距离
        :return:(nRays)
        """
        return np.array([iradial['URange'] / 10. for iradial in self.radial])

    def get_sweep_end_ray_index(self):
        """
        获取每个sweep的结束的index，包含在内
        :return:(nsweep)
        """
        return self.sweep_end_ray_index

    def get_sweep_start_ray_index(self):
        """
        获取每个sweep的开始的index
        :return:(nsweep)
        """
        return self.sweep_start_ray_index

    def get_nbins_per_sweep(self):
        """
        确定每个sweep V探测的库数
        :return:
        """
        return np.array([self.radial[idx]['fields']['V'].size for idx in self.sweep_start_ray_index])

    def get_range_per_radial(self, length):
        """
        确定径向每个库的距离
        :param length:
        :return:
        """
        Resolution = self.radial[0]["GateSizeOfDoppler"]
        return np.linspace(Resolution, Resolution * length, length)

    def get_dbz_range_per_radial(self, length):
        """
        确定径向每个库的距离
        :param length:
        :return:
        """
        Resolution = self.dBZ_Res
        start_range = self.radial[0]["GateSizeOfDoppler"]
        return np.linspace(start_range, start_range + Resolution * (length - 1), length)

    def _get_fields(self):
        """将所有的field的数据提取出来"""
        fields = {}
        field_keys = self.radial[0]['fields'].keys()
        for ikey in field_keys:
            fields[ikey] = np.array([self._add_or_del_field(iray['fields'], ikey) for iray in self.radial])
        return fields

    def _add_or_del_field(self, dat_fields, key):
        """
        根据fields的key提取数据, 将dbz的数据和dop的数据分辨率统一
        :param dat_fields: fields的数据
        :param key: key words
        :return:
        """
        length = self.max_bins
        if key == "dBZ":
            dbz_range = self.get_dbz_range_per_radial(dat_fields[key].size)
            dop_range = self.range
            match_data = interpolate.interp1d(dbz_range, dat_fields[key], kind="nearest",
                                              bounds_error=False, fill_value=np.nan)
            dat_ray = match_data(dop_range)
            #print(dop_range, dbz_range)
            #print(dat_ray)
            return dat_ray.ravel()
        else:
            dat_ray = dat_fields[key]
        if dat_ray.size >= length:
            return (dat_ray[:length]).ravel()
        else:
            out = np.full((length,), np.nan)
            out[:dat_ray.size] = dat_ray
            return out.ravel()

    def get_NRadar_nyquist_speed(self):
        """array shape (nsweeps)"""
        return np.array([self.radial[idx]['Nyquist'] / 100. for idx in self.sweep_start_ray_index])

    def get_NRadar_unambiguous_range(self):
        """array shape (nsweeps)"""
        return np.array([self.radial[idx]['URange'] / 10. for idx in self.sweep_start_ray_index])

    def get_fixed_angle(self):
        if self.nsweeps == 9:
            fixed_angle = np.array([0.50, 1.45, 2.40, 3.35, 4.30, 6.00, 9.00, 14.6, 19.5])
        elif self.nsweeps == 14:
            fixed_angle = np.array([0.50, 1.45, 2.40, 3.35, 4.30, 5.25, 6.2, 7.5, 8.7, 10, 12, 14, 16.7, 19.5])
        elif self.nsweeps == 6:
            fixed_angle = np.array([0.50, 1.50, 2.50, 2.50, 3.50, 4.50])
        elif self.nsweeps == 4:
            fixed_angle = np.array([0.50, 2.50, 3.50, 4.50])
        else:
            fixed_angle = np.array([self.radial[idx]['El'] / 8. * 180. / 4096. for idx in self.sweep_start_ray_index])
        return fixed_angle

    def ToPRD(self):
        """将WSR98D数据转为PRD的数据格式"""
        return PRD(fields=self.fields, scan_type=self.scan_type, time=self.get_scan_time(), \
                          range=self.range, azimuth=self.azimuth, elevation=self.elevation, latitude=self.latitude, \
                          longitude=self.longitude, altitude=self.altitude,
                          sweep_start_ray_index=self.sweep_start_ray_index, \
                          sweep_end_ray_index=self.sweep_end_ray_index, fixed_angle=self.get_fixed_angle(), \
                          bins_per_sweep=self.bins_per_sweep, nyquist_velocity=self.get_NRadar_nyquist_speed(), \
                          frequency=self.frequency, unambiguous_range=self.get_NRadar_unambiguous_range(), \
                          nrays=self.nrays, nsweeps=self.nsweeps, sitename = self.sitename, pyart_radar=self.ToPyartRadar())

    def ToPyartRadar(self):
        """转化为Pyart Radar的对象"""
        dts = self.get_scan_time()
        units = make_time_unit_str(min(dts))
        time = get_metadata('time')
        time['units'] = units
        time['data'] = date2num(dts, units).astype('float32')

        # range
        _range = get_metadata('range')
        # assume that the number of gates and spacing from the first ray is
        # representative of the entire volume
        _range['data'] = self.range
        _range['meters_to_center_of_first_gate'] = self.radial[0]["GateSizeOfDoppler"]
        _range['meters_between_gates'] = self.radial[0]["GateSizeOfDoppler"]

        latitude = get_metadata('latitude')
        longitude = get_metadata('longitude')
        altitude = get_metadata('altitude')
        latitude['data'] = np.array([self.latitude], dtype='float64')
        longitude['data'] = np.array([self.longitude], dtype='float64')
        altitude['data'] = np.array([self.altitude], dtype='float64')

        metadata = get_metadata('metadata')
        metadata['original_container'] = 'CINRAD/SAB'
        metadata['site_name'] = self.sitename
        metadata['radar_name'] = "CINRAD/SA/SB/CB/SC"

        sweep_start_ray_index = get_metadata('sweep_start_ray_index')
        sweep_end_ray_index = get_metadata('sweep_end_ray_index')
        sweep_start_ray_index['data'] = self.sweep_start_ray_index
        sweep_end_ray_index['data'] = self.sweep_end_ray_index

        sweep_number = get_metadata('sweep_number')
        sweep_number['data'] = np.arange(self.nsweeps, dtype='int32')

        scan_type = self.scan_type

        sweep_mode = get_metadata('sweep_mode')
        sweep_mode['data'] = np.array(self.nsweeps * ['azimuth_surveillance'], dtype='S')

        # elevation
        elevation = get_metadata('elevation')
        elevation['data'] = self.elevation

        # azimuth
        azimuth = get_metadata('azimuth')
        azimuth['data'] = self.azimuth

        # fixed_angle
        fixed_angle = get_metadata('fixed_angle')
        fixed_angle['data'] = self.get_fixed_angle()

        # instrument_parameters
        instrument_parameters = self._get_instrument_parameters()

        # fields
        fields = {}
        for field_name_abbr in self.fields.keys():
            field_name = CINRAD_field_mapping[field_name_abbr]
            if field_name is None:
                continue
            field_dic = get_metadata(field_name)
            field_dic['data'] = np.ma.masked_array(self.fields[field_name_abbr],\
                                mask=np.isnan(self.fields[field_name_abbr]), fill_value=get_fillvalue())
            field_dic['_FillValue'] = get_fillvalue()
            fields[field_name] = field_dic

        return Radar(time, _range, fields, metadata, scan_type,
                     latitude, longitude, altitude,
                     sweep_number, sweep_mode, fixed_angle, sweep_start_ray_index,
                     sweep_end_ray_index,
                     azimuth, elevation,
                     instrument_parameters=instrument_parameters)

    def _get_instrument_parameters(self):
        """ Return a dictionary containing instrument parameters. """

        # pulse width
        pulse_width = get_metadata('pulse_width')
        pulse_width['data'] = np.array([self.radial[0]["GateSizeOfDoppler"] / _LIGHT_SPEED,], dtype='float32')  # m->sec

        # assume that the parameters in the first ray represent the beam widths,
        # bandwidth and frequency in the entire volume

        wavelength_hz = self.frequency * 10 ** 9

        # radar_beam_width_h
        radar_beam_width_h = get_metadata('radar_beam_width_h')
        radar_beam_width_h['data'] = np.array([1, ], dtype='float32')

        # radar_beam_width_v
        radar_beam_width_v = get_metadata('radar_beam_width_v')
        radar_beam_width_v['data'] = np.array([1, ], dtype='float32')

        # frequency
        frequency = get_metadata('frequency')
        frequency['data'] = np.array([wavelength_hz], dtype='float32')

        instrument_parameters = {
            'pulse_width': pulse_width,
            'radar_beam_width_h': radar_beam_width_h,
            'radar_beam_width_v': radar_beam_width_v,
            'frequency': frequency, }

        # nyquist velocity if defined
        nyquist_velocity = get_metadata('nyquist_velocity')
        nyquist_velocity['data'] = self.get_nyquist_velocity()
        instrument_parameters['nyquist_velocity'] = nyquist_velocity
        return instrument_parameters