"""TODO."""

from __future__ import print_function, absolute_import
import numba
from numba import jit
import numpy as np
import math
import random
from numbskull_extend.inference import draw_sample, eval_factor

@jit(cache=True, nogil=True)
def learnthread(shardID, nshards, step, regularization, reg_param,
                truncation,var_copy, weight_copy, weight,variable,
                factor, fmap,vmap, factor_index, Z,
                fids,var_value, var_value_evid,weight_value, learn_non_evidence,
                poential_weight,alpha_bound,tau_bound,sample_list=None, wmap=None,wfactor=None):
    """TODO."""
    # Identify start and end variable
    nvar = variable.shape[0]
    start = (shardID * nvar) // nshards
    end = ((shardID + 1) * nvar) // nshards
    if sample_list == None:     #sample_list为None表示不需要平衡化
        for var_samp in range(start, end):
            if variable[var_samp]["isEvidence"] == 4:
                # This variable is not owned by this machine
                continue
            sample_and_sgd(var_samp, step, regularization, reg_param, truncation,
                           var_copy, weight_copy, weight, variable,factor,
                           fmap, vmap,factor_index, Z[shardID], fids[shardID],
                           var_value,var_value_evid, weight_value, learn_non_evidence,poential_weight,
                           alpha_bound,tau_bound)
    else:      #需要平衡化
        sample_num = sample_list.shape[0]
        start = (shardID * sample_num) // nshards
        end = ((shardID + 1) * sample_num) // nshards
        for i in range(0,sample_num) :
            var_samp = sample_list[i]['vid']
            if variable[var_samp]["isEvidence"] == 4:
                # This variable is not owned by this machine
                continue
            sample_and_sgd(var_samp, step, regularization, reg_param, truncation,
                           var_copy, weight_copy, weight, variable,factor,
                           fmap, vmap,factor_index, Z[shardID],fids[shardID],
                           var_value,var_value_evid,weight_value, learn_non_evidence,poential_weight,
                           alpha_bound,tau_bound)

@jit(cache=True, nogil=True)
def learnthread_bgd(shardID, nshards, step, regularization, reg_param,
                truncation,var_copy, weight_copy, weight,variable,
                factor, fmap,vmap, factor_index, Z,
                fids,var_value, var_value_evid,weight_value, learn_non_evidence,
                poential_weight,alpha_bound,tau_bound,sample_list=None,wmap=None,wfactor=None):
    nweight = weight.shape[0]
    start = (shardID * nweight) // nshards
    end = ((shardID + 1) * nweight) // nshards
    for wid in range(start,end):
        if weight[wid]["isFixed"]:
            continue
        else:
            sample_and_bgd(wid,step, regularization, reg_param, truncation,
                           var_copy, weight_copy, weight, variable,factor,
                           fmap,vmap, factor_index, Z[shardID],fids[shardID],
                           var_value, var_value_evid,weight_value, learn_non_evidence,poential_weight,
                           alpha_bound,tau_bound,wmap,wfactor)

@jit(nopython=True, cache=True, nogil=True)
def sample_and_bgd(wid,step, regularization, reg_param, truncation,
                   var_copy, weight_copy, weight, variable,factor,
                   fmap,vmap, factor_index, Z,fids,
                   var_value, var_value_evid,weight_value, learn_non_evidence,poential_weight,
                   alpha_bound,tau_bound,wmap,wfactor):    #批量梯度下降不需要考虑poential_weight
    #1.计算梯度和
    weight_id = wmap[wid]["weightId"]
    weight_index_offset = wmap[wid]["weight_index_offset"]
    weight_index_length = wmap[wid]["weight_index_length"]
    #找到此权重相关的每一个factor
    factor_count = weight_index_length   #此权重拥有的因子个数
    for fIndex in range(weight_index_offset,weight_index_offset+weight_index_length):
        factor_id = wfactor[fIndex]["factorId"]
        ftv_offset = factor[factor_id]["ftv_offset"]
        ftv_length = factor[factor_id]["arity"]
        var_count = ftv_length          #此因子拥有的变量个数
        gradient_sum = 0   #不需要参数化时，所有梯度的和
        gradient1_sum = 0  #需要参数化时，参数1的梯度和
        gradient2_sum = 0  # 需要参数化时，参数2的梯度和
        # 找到每一个factor相关的每一个变量
        for vIndex in range(ftv_offset,ftv_offset+ftv_length):
            var_samp = fmap[vIndex]["vid"]
            if variable[var_samp]["isEvidence"] != 1:
                evidence = draw_sample(var_samp, var_copy, weight_copy,
                                       weight, variable, factor,
                                       fmap, vmap, factor_index, Z,
                                       var_value_evid, weight_value)
                # If evidence then store the initial value in a tmp variable
            # then sample and compute the gradient.
            else:
                evidence = variable[var_samp]["initialValue"]
            var_value_evid[var_copy][var_samp] = evidence
            # Sample the variable
            proposal = draw_sample(var_samp, var_copy, weight_copy, weight,
                                   variable, factor, fmap, vmap,
                                   factor_index, Z, var_value, weight_value)
            var_value[var_copy][var_samp] = proposal
            if not learn_non_evidence and variable[var_samp]["isEvidence"] != 1:
                return
            truncate = random.random() < 1.0 / truncation if regularization == 1 else False
            p0 = eval_factor(factor_id, var_samp,
                             evidence, var_copy,
                             variable, factor, fmap,
                             var_value_evid)
            p1 = eval_factor(factor_id, var_samp,
                             proposal, var_copy,
                             variable, factor, fmap,
                             var_value)
            # if need parameterize
            if weight[factor[factor_id]['weightId']]['parameterize']:
                x = fmap[factor[factor_id]["ftv_offset"]]['x']
                theta = fmap[factor[factor_id]["ftv_offset"]]['theta']
                a = weight[factor[factor_id]['weightId']]['a']
                b = weight[factor[factor_id]['weightId']]['b']
                gradient1 = (p1 - p0) * theta * factor[factor_id]["featureValue"] * (x - b)
                gradient2 = (p1 - p0) * theta * factor[factor_id]["featureValue"] * (-a)
                gradient1_sum += gradient1
                gradient2_sum += gradient2
            # if not need parameterize
            else:
                gradient = (p1 - p0) * factor[factor_id]["featureValue"]
                gradient_sum += gradient
        #求平均值
        # print("权重开始平均")
        gradient1_sum /= var_count
        gradient2_sum /= var_count
        gradient_sum /= var_count

    #2.更新参数，分为需要参数化的和不需要参数化的
    #if need parameterize
    if weight[factor[factor_id]['weightId']]['parameterize'] == 1:
            if regularization == 2:  # 是否需要正则化
                a *= (1.0 / (1.0 + reg_param * step))
                a -= step * gradient1_sum/factor_count
                b *= (1.0 / (1.0 + reg_param * step))
                b -= step * gradient2_sum/factor_count
            elif regularization == 1:
                # Truncated Gradient
                # "Sparse Online Learning via Truncated Gradient"
                #  Langford et al. 2009
                a -= step * gradient1_sum/factor_count
                b -= step * gradient2_sum/factor_count
                if truncate:
                    l1delta = reg_param * step * truncation
                    a = max(0, a - l1delta) if a > 0 else min(0, a + l1delta)
                    b = max(0, b - l1delta) if b > 0 else min(0, b + l1delta)
            else:
                a -= step * gradient1_sum/factor_count
                b -= step * gradient2_sum/factor_count
            if a < tau_bound[factor[factor_id]['weightId']]['lowerBound']:
                a = tau_bound[factor[factor_id]['weightId']]['lowerBound']
            elif a > tau_bound[factor[factor_id]['weightId']]['upperBound']:
                a = tau_bound[factor[factor_id]['weightId']]['upperBound']
            if b > alpha_bound[factor[factor_id]['weightId']]['upperBound']:
                b = alpha_bound[factor[factor_id]['weightId']]['upperBound']
            elif b < alpha_bound[factor[factor_id]['weightId']]['lowerBound']:
                b = alpha_bound[factor[factor_id]['weightId']]['lowerBound']
            w = theta * a * (x - b)
            weight[factor[factor_id]['weightId']]['a'] = a
            weight[factor[factor_id]['weightId']]['b'] = b
    # if not need parameterize
    elif weight[factor[factor_id]['weightId']]['parameterize'] == 0:
        w = weight_value[weight_copy][weight_id]
        if regularization == 2:
            w *= (1.0 / (1.0 + reg_param * step))
            w -= step * gradient_sum /factor_count
        elif regularization == 1:
            # Truncated Gradient
            # "Sparse Online Learning via Truncated Gradient"
            #  Langford et al. 2009
            w -= step * gradient_sum /factor_count
            if truncate:
                l1delta = reg_param * step * truncation
                w = max(0, w - l1delta) if w > 0 else min(0, w + l1delta)
        else:
            w -= step * gradient_sum /factor_count
    weight_value[weight_copy][weight_id] = w
    weight[factor[factor_id]['weightId']]['initialValue'] = w
    if variable[var_samp]["isEvidence"] != 1:
        poential_weight[factor[factor_id]['weightId']] = w

@jit(nopython=True, cache=True, nogil=True)
def get_factor_id_range(variable, vmap, var_samp, val):
    """TODO."""
    varval_off = val
    if variable[var_samp]["dataType"] == 0:
        varval_off = 0
    vtf = vmap[variable[var_samp]["vtf_offset"] + varval_off]
    start = vtf["factor_index_offset"]
    end = start + vtf["factor_index_length"]
    return (start, end)

@jit(nopython=True, cache=True, nogil=True)
def sample_and_sgd(var_samp, step, regularization, reg_param, truncation,
                   var_copy, weight_copy, weight, variable,
                   factor, fmap,vmap, factor_index, Z,
                   fids, var_value, var_value_evid,
                   weight_value, learn_non_evidence,
                   poential_weight,alpha_bound,tau_bound):
    """TODO."""
    # If learn_non_evidence sample twice.
    # The method corresponds to expectation-conjugate descent.
    if variable[var_samp]["isEvidence"] != 1:
        evidence = draw_sample(var_samp, var_copy, weight_copy,
                               weight, variable, factor,
                               fmap, vmap, factor_index, Z,
                               var_value_evid, weight_value)
        # If evidence then store the initial value in a tmp variable
    # then sample and compute the gradient.
    else:
        evidence = variable[var_samp]["initialValue"]

    var_value_evid[var_copy][var_samp] = evidence
    # Sample the variabl e
    proposal = draw_sample(var_samp, var_copy, weight_copy, weight,
                           variable, factor, fmap, vmap,
                           factor_index, Z, var_value, weight_value)

    var_value[var_copy][var_samp] = proposal
    if not learn_non_evidence and variable[var_samp]["isEvidence"] != 1:
        return
    # Compute the gradient and update the weights
    # Iterate over corresponding factors

    range_fids = get_factor_id_range(variable, vmap, var_samp, evidence)
    # TODO: is it possible to avoid copying around fids
    if evidence != proposal:
        range_prop = get_factor_id_range(variable, vmap, var_samp, proposal)
        s1 = range_fids[1] - range_fids[0]
        s2 = range_prop[1] - range_prop[0]
        s = s1 + s2
        fids[:s1] = factor_index[range_fids[0]:range_fids[1]]
        fids[s1:s] = factor_index[range_prop[0]:range_prop[1]]
        fids[:s].sort()
    else:
        s = range_fids[1] - range_fids[0]
        fids[:s] = factor_index[range_fids[0]:range_fids[1]]

    truncate = random.random() < 1.0 / truncation if regularization == 1 else False
    # go over all factor ids, ignoring dupes
    last_fid = -1  # numba 0.28 would complain if this were None
    for factor_id in fids[:s]:
        if factor_id == last_fid:
            continue
        last_fid = factor_id
        weight_id = factor[factor_id]["weightId"]
        if weight[weight_id]["isFixed"]:
            continue
        # Compute Gradient
        p0 = eval_factor(factor_id, var_samp,
                         evidence, var_copy,
                         variable, factor, fmap,
                         var_value_evid)
        p1 = eval_factor(factor_id, var_samp,
                         proposal, var_copy,
                         variable, factor, fmap,
                         var_value)
        #if need parameterize
        if weight[factor[factor_id]['weightId']]['parameterize'] == 1:
            x = fmap[factor[factor_id]["ftv_offset"]]['x']
            theta = fmap[factor[factor_id]["ftv_offset"]]['theta']
            a = weight[factor[factor_id]['weightId']]['a']
            b = weight[factor[factor_id]['weightId']]['b']
            gradient1 = (p1 - p0) * theta * factor[factor_id]["featureValue"] * (x - b)
            gradient2 = (p1 - p0) * theta * factor[factor_id]["featureValue"] * (-a)
            if regularization == 2:  # 是否需要正则化
                a *= (1.0 / (1.0 + reg_param * step))
                a -= step * gradient1
                b *= (1.0 / (1.0 + reg_param * step))
                b -= step * gradient2
            elif regularization == 1:
            # Truncated Gradient
            # "Sparse Online Learning via Truncated Gradient"
            #  Langford et al. 2009
                a -= step * gradient1
                b -= step * gradient2
                if truncate:
                    l1delta = reg_param * step * truncation
                    a = max(0, a - l1delta) if a > 0 else min(0, a + l1delta)
                    b = max(0, b - l1delta) if b > 0 else min(0, b + l1delta)
            else:
                a -= step * gradient1
                b -= step * gradient2

            # if alpha_bound != None and tau_bound != None:
            if a < tau_bound[factor[factor_id]['weightId']]['lowerBound']:
                a = tau_bound[factor[factor_id]['weightId']]['lowerBound']
            elif a > tau_bound[factor[factor_id]['weightId']]['upperBound']:
                a = tau_bound[factor[factor_id]['weightId']]['upperBound']
            if b > alpha_bound[factor[factor_id]['weightId']]['upperBound']:
                b = alpha_bound[factor[factor_id]['weightId']]['upperBound']
            elif  b < alpha_bound[factor[factor_id]['weightId']]['lowerBound']:
                b = alpha_bound[factor[factor_id]['weightId']]['lowerBound']
            w = theta * a * (x - b)
            weight[factor[factor_id]['weightId']]['a'] = a
            weight[factor[factor_id]['weightId']]['b'] = b
        #如果不需要参数化
        elif weight[factor[factor_id]['weightId']]['parameterize'] == 0:
            gradient = (p1 - p0) * factor[factor_id]["featureValue"]
        # Update weight
            w = weight_value[weight_copy][weight_id]
            if regularization == 2:
                w *= (1.0 / (1.0 + reg_param * step))
                w -= step * gradient
            elif regularization == 1:
            # Truncated Gradient
            # "Sparse Online Learning via Truncated Gradient"
            #  Langford et al. 2009
                w -= step * gradient
                if truncate:
                    l1delta = reg_param * step * truncation
                    w = max(0, w - l1delta) if w > 0 else min(0, w + l1delta)
            else:
                w -= step * gradient
        weight_value[weight_copy][weight_id] = w
        weight[factor[factor_id]['weightId']]['initialValue'] = w
        if variable[var_samp]["isEvidence"] != 1:
            poential_weight[factor[factor_id]['weightId']] = w


from __future__ import division
from __future__ import print_function
from __future__ import absolute_import

import sys
import numpy as np

import torch
import torch.nn as nn
from torch.nn import functional as F
import torch.nn.utils.rnn as rnn_utils

from .loss import (
    gaussian_log_pdf,
    unit_gaussian_log_pdf,
    bernoulli_log_pdf,
    categorical_program_log_pdf,
)
from .utils import log_mean_exp

class ProgramRNN(nn.Module):
    r"""Supervised recurrent neural network for predicting feedback labels.
    Parameterizes p(label|program).

    @param z_dim: integer
                  size of latent vector
    @param label_dim: integer
                      number of (binary) labels
    @param vocab_size: integer
                       size of vocabulary
    @param embedding_dim: integer [default: 300]
                          size of learned embedding
    @param hidden_dim: integer [default: 256]
                       size of hidden layer
    @param word_dropout: float [default: 0.5]
                         probability to dropout input sequences to decoder
    @param num_layers: integer [default: 2]
                       number of hidden layers in GRU
    """
    def __init__(   self, z_dim, label_dim, vocab_size, embedding_dim=300,
                    hidden_dim=256, num_layers=2):
        super(ProgramRNN, self).__init__()

        self.z_dim = z_dim
        self.label_dim = label_dim
        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers

        self.embedding_module = nn.Embedding(self.vocab_size, self.embedding_dim)
        self.encoder = ProgramEncoder(  self.embedding_module, self.z_dim,
                                        hidden_dim=self.hidden_dim,
                                        num_layers=self.num_layers)
        self.decoder = LabelDecoder(self.z_dim, self.label_dim,
                                    hidden_dim=self.hidden_dim)

    def forward(self, seq, length):
        z_mu, _  = self.encoder(seq, length)
        return self.decoder(z_mu)

class ProgramMVAE(nn.Module):
    r"""Multimodal Variational Autoencoder (MVAE) with Expert Supervision.
    The MVAE is the equivalent of a soft transformation of the expert pCFG graph.

    @param z_dim: integer
                  size of latent vector
    @param label_dim: integer
                      number of (binary) labels
    @param vocab_size: integer
                       size of vocabulary
    @param sutilos_idx: integer
                    index for start-of-sentence
    @param eos_idx: integer
                    index for end-of-sentence
    @param pad_idx: integer
                    index for padding
    @param unk_idx: integer
                    index for unknown tokens
    @param embedding_dim: integer [default: 300]
                          size of learned embedding
    @param hidden_dim: integer [default: 256]
                       size of hidden layer
    @param word_dropout: float [default: 0.5]
                         probability to dropout input sequences to decoder
    @param num_layers: integer [default: 2]
                       number of hidden layers in GRU
    """
    def __init__(self, z_dim, label_dim, vocab_size, sos_idx, eos_idx, pad_idx, unk_idx,
                 embedding_dim=300, hidden_dim=256, word_dropout=0.5, num_layers=2):
        super(ProgramMVAE, self).__init__()

        self.z_dim = z_dim
        self.label_dim = label_dim
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        self.vocab_size = vocab_size
        self.sos_idx = sos_idx
        self.eos_idx = eos_idx
        self.pad_idx = pad_idx
        self.unk_idx = unk_idx
        self.word_dropout = word_dropout
        self.num_layers = num_layers

        self.product_experts = ProductOfExperts()
        self.embedding_module = nn.Embedding(self.vocab_size, self.embedding_dim)

        self.program_encoder = ProgramEncoder(  self.embedding_module, self.z_dim,
                                                hidden_dim=self.hidden_dim,
                                                num_layers=self.num_layers)
        self.program_decoder = ProgramDecoder(
            self.embedding_module, self.z_dim, self.sos_idx, self.eos_idx, self.pad_idx, self.unk_idx,
            hidden_dim=self.hidden_dim, word_dropout=self.word_dropout, num_layers=self.num_layers)

        self.label_encoder = LabelEncoder(
            self.z_dim, self.label_dim, hidden_dim=self.hidden_dim)
        self.label_decoder = LabelDecoder(
            self.z_dim, self.label_dim, hidden_dim=self.hidden_dim)

    def reparameterize(self, z_mu, z_logvar):
        std = torch.exp(0.5 * z_logvar)
        eps = torch.randn_like(std)
        return eps.mul(std).add_(z_mu)

    def forward(self, seq, length, label, hide_seq=False):
        if hide_seq:
            z_mu, z_logvar = self.inference(None, None, label)
        else:
            z_mu, z_logvar = self.inference(seq, length, label)

        z = self.reparameterize(z_mu, z_logvar)
        seq_logits = self.program_decoder(z, seq, length)
        label_out = self.label_decoder(z)

        return seq_logits, label_out, z, z_mu, z_logvar

    def prior_expert(self, mu):
        p_mu = torch.zeros_like(mu)
        p_logvar = torch.zeros_like(mu)
        return p_mu, p_logvar

    def inference(self, seq, length, label):
        z_mu, z_logvar = [], []

        if seq is not None:
            program_mu, program_logvar = self.program_encoder(seq, length)
            z_mu.append(program_mu.unsqueeze(0))
            z_logvar.append(program_logvar.unsqueeze(0))

        if label is not None:
            label_mu, label_logvar = self.label_encoder(label)
            z_mu.append(label_mu.unsqueeze(0))
            z_logvar.append(label_logvar.unsqueeze(0))

        prior_mu, prior_logvar = self.prior_expert(z_mu[0])
        z_mu.append(prior_mu)
        z_logvar.append(prior_logvar)

        z_mu = torch.cat(z_mu, dim=0)
        z_logvar = torch.cat(z_logvar, dim=0)

        z_mu, z_logvar = self.product_experts(z_mu, z_logvar)

        return z_mu, z_logvar

    def get_joint_marginal(self, seq, length, label, n_samples=100):
        z_mu, z_logvar = self.inference(seq, length, label)

        log_w = []
        for i in xrange(n_samples):
            z_i = self.reparameterize(z_mu, z_logvar)
            x_logits_i = self.program_decoder(z_i, seq, length)
            y_out_i = self.label_decoder(z_i)

            log_p_x_given_z_i = categorical_program_log_pdf(seq[:, 1:], x_logits_i[:, :-1])
            log_p_y_given_z_i = bernoulli_log_pdf(label, y_out_i)
            log_q_z_given_x_y_i = gaussian_log_pdf(z_i, z_mu, z_logvar)
            log_p_z_i = unit_gaussian_log_pdf(z_i)

            log_w_i = log_p_x_given_z_i + log_p_y_given_z_i + log_p_z_i - log_q_z_given_x_y_i
            log_w.append(log_w_i.unsqueeze(1))

        log_w = torch.cat(log_w, dim=1)
        log_p_x_y = log_mean_exp(log_w, dim=1)
        log_p_x_y = -torch.mean(log_p_x_y)

        return log_p_x_y

    def get_program_marginal(self, seq, length, n_samples=100):
        z_mu, z_logvar = self.inference(seq, length, None)

        log_w = []
        for i in xrange(n_samples):
            z_i = self.reparameterize(z_mu, z_logvar)
            seq_logits_i = self.program_decoder(z_i, seq, length)

            # probability of text is product of probabilities of each word
            log_p_x_given_z_i = categorical_program_log_pdf(seq[:, 1:], seq_logits_i[:, :-1])
            log_q_z_given_x_i = gaussian_log_pdf(z_i, z_mu, z_logvar)
            log_p_z_i = unit_gaussian_log_pdf(z_i)

            log_w_i = log_p_x_given_z_i + log_p_z_i - log_q_z_given_x_i
            log_w.append(log_w_i.unsqueeze(1))

        log_w = torch.cat(log_w, dim=1)
        log_p_x = log_mean_exp(log_w, dim=1)
        log_p_x = -torch.mean(log_p_x)

        return log_p_x

    def get_label_marginal(self, y, n_samples=100):
        z_mu, z_logvar = self.inference(None, None, y)

        log_w = []
        for i in xrange(n_samples):
            z_i = self.reparameterize(z_mu, z_logvar)
            y_out_i = self.label_decoder(z_i)

            log_p_y_given_z_i = bernoulli_log_pdf(y, y_out_i)
            log_q_z_given_y_i = gaussian_log_pdf(z_i, z_mu, z_logvar)
            log_p_z_i = unit_gaussian_log_pdf(z_i)

            log_w_i = log_p_y_given_z_i + log_p_z_i - log_q_z_given_y_i
            log_w.append(log_w_i.unsqueeze(1))

        log_w = torch.cat(log_w, dim=1)
        log_p_y = log_mean_exp(log_w, dim=1)
        log_p_y = -torch.mean(log_p_y)

        return log_p_y

class ProgramEncoder(nn.Module):
    r"""Parameterizes q(z|program) with RNN.

    Inspired by Bowman et. al. (https://arxiv.org/abs/1511.06349).

    @param embedding_module: nn.Embedding
                             we initialize this separately from the encoder
    @param z_dim: integer
                  size of latent vector
    @param hidden_dim: integer [default: 256]
                       size of hidden layer
    @param num_layers: integer [default: 2]
                       number of hidden layers in GRU
    """
    def __init__(self, embedding_module, z_dim, hidden_dim=256, num_layers=2):
        super(ProgramEncoder, self).__init__()

        self.embedding = embedding_module
        self.embedding_dim = self.embedding.embedding_dim
        self.z_dim = z_dim
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers

        self.gru = nn.GRU(self.embedding_dim, self.hidden_dim, num_layers=num_layers)
        self.h2mu = nn.Linear(self.hidden_dim * self.num_layers, self.z_dim)
        self.h2logvar = nn.Linear(self.hidden_dim * self.num_layers, self.z_dim)

    def forward(self, seq, length):
        batch_size = seq.size(0)

        if batch_size > 1:
            # sort in decreasing order of length in order to pack
            # sequence; if only 1 element in batch, nothing to do.
            sorted_lengths, sorted_idx = torch.sort(length, descending=True)
            seq = seq[sorted_idx]

        embed_seq = self.embedding(seq)
        # reorder from (B,L,D) to (L,B,D)
        embed_seq = embed_seq.transpose(0, 1)

        packed = rnn_utils.pack_padded_sequence(
            embed_seq,
            sorted_lengths.data.tolist() if batch_size > 1 else length.data.tolist())

        _, hidden = self.gru(packed)
        hidden = hidden.permute(1, 0, 2).contiguous()
        hidden = hidden.view(batch_size, self.hidden_dim * self.num_layers)

        if batch_size > 1:
            _, reversed_idx = torch.sort(sorted_idx)
            hidden = hidden[reversed_idx]

        z_mu = self.h2mu(hidden)
        z_logvar = self.h2logvar(hidden)

        return z_mu, z_logvar

class ProgramDecoder(nn.Module):
    r"""Parameterizes p(program|z) with RNN to generate a distribution of a
    sequence of tokens. Assumes a maximum sequence length and a fixed vocabulary.

    We return logits to a categorical so please use
        nn.CrossEntropy
    instead of
        nn.NLLLoss

    Inspired by Bowman et. al. (https://arxiv.org/abs/1511.06349).

    @param embedding_module: nn.Embedding
                             pass the embedding module (share with encoder)
    @param z_dim: integer
                  size of latent vector
    @param sos_idx: integer
                    index for start-of-sentence
    @param eos_idx: integer
                    index for end-of-sentence
    @param pad_idx: integer
                    index for padding
    @param unk_idx: integer
                    index for unknown tokens
    @param hidden_dim: integer [default: 256]
                       size of hidden layer
    @param word_dropout: float [default: 0.5]
                         probability to dropout input sequences to decoder
    @param num_layers: integer [default: 2]
                       number of hidden layers in GRU
    """
    def __init__(   self, embedding_module, z_dim, sos_idx, eos_idx, pad_idx, unk_idx,
                    hidden_dim=256, word_dropout=0.5, num_layers=2):
        super(ProgramDecoder, self).__init__()

        self.embedding = embedding_module
        self.embedding_dim = embedding_module.embedding_dim
        self.vocab_size = embedding_module.num_embeddings
        self.z_dim = z_dim
        self.hidden_dim = hidden_dim
        self.sos_idx = sos_idx
        self.eos_idx = eos_idx
        self.pad_idx = pad_idx
        self.unk_idx = unk_idx
        self.word_dropout = word_dropout
        self.num_layers = num_layers
        self.gru = nn.GRU(  self.embedding_dim, self.hidden_dim,
                            num_layers=self.num_layers)
        self.z2h = nn.Linear(self.z_dim, self.hidden_dim * self.num_layers)
        self.outputs2vocab = nn.Linear(self.hidden_dim, self.vocab_size)

    def forward(self, z, seq, length):
        batch_size = z.size(0)

        if batch_size > 1:
            sorted_lengths, sorted_idx = torch.sort(length, descending=True)

            z = z[sorted_idx]
            seq = seq[sorted_idx]

        if self.word_dropout > 0:
            # randomly replace with unknown tokens
            prob = torch.rand(seq.size())
            prob[(seq.cpu().data - self.sos_idx) & \
                 (seq.cpu().data - self.pad_idx) == 0] = 1
            mask_seq = seq.clone()
            mask_seq[(prob < self.word_dropout).to(z.device)] = self.unk_idx
            seq = mask_seq

        embed_seq = self.embedding(seq)
        # reorder from (B,L,D) to (L,B,D)
        embed_seq = embed_seq.transpose(0, 1)

        packed = rnn_utils.pack_padded_sequence(
            embed_seq,
            sorted_lengths.data.tolist() if batch_size > 1 else length.data.tolist())

        hidden = self.z2h(z)
        hidden = hidden.view(batch_size, self.num_layers, self.hidden_dim)
        hidden = hidden.permute(1, 0, 2).contiguous()

        outputs, _ = self.gru(packed, hidden)
        outputs = rnn_utils.pad_packed_sequence(outputs, batch_first=True)[0]
        outputs = outputs.contiguous()

        # reorder from (L,B,D) to (B,L,D)
        outputs = outputs.transpose(0, 1)

        if batch_size > 1:
            _, reversed_idx = torch.sort(sorted_idx)
            outputs = outputs[reversed_idx]

        max_length = outputs.size(1)
        outputs_2d = outputs.view(batch_size * max_length, self.hidden_dim)
        outputs_2d = self.outputs2vocab(outputs_2d)
        outputs = outputs_2d.view(batch_size, max_length, self.vocab_size)

        return outputs.contiguous()

    def sample(self, z, max_seq_len, greedy=False):
        r"""Sample tokens in an auto-regressive framework.

        @param z: torch.Tensor
                  sample of latent variables
        @param max_seq_len: integer
                            maximum size of sequence
        @param greedy: boolean [default: False]
                       pick most likely token or sample token?
        """
        with torch.no_grad():
            batch_size = z.size(0)

            # initialize hidden state
            hidden = self.z2h(z)
            hidden = hidden.view(batch_size, self.num_layers, self.hidden_dim)
            hidden = hidden.permute(1, 0, 2).contiguous()

            # first input is SOS token
            inputs = np.array([self.sos_idx for _ in xrange(batch_size)])
            inputs = torch.from_numpy(inputs)
            inputs = inputs.unsqueeze(1)
            inputs = inputs.to(z.device)

            # save SOS as first generated token
            inputs_npy = inputs.squeeze(1).cpu().numpy()
            sampled_ids = [[w] for w in inputs_npy]

            # (B,L,D) to (L,B,D)
            inputs = inputs.transpose(0, 1)

            # compute embeddings
            inputs = self.embedding(inputs)

            for i in xrange(max_seq_len):
                outputs, hidden = self.gru(inputs, hidden)  # outputs: (L=1,B,H)
                outputs = outputs.squeeze(0)                # outputs: (B,H)
                outputs = self.outputs2vocab(outputs)       # outputs: (B,V)

                if greedy:
                    predicted = outputs.max(1)[1]
                    predicted = predicted.unsqueeze(1)
                else:
                    outputs = F.softmax(outputs, dim=1)
                    predicted = torch.multinomial(outputs, 1)

                predicted_npy = predicted.squeeze(1).cpu().numpy()
                predicted_lst = predicted_npy.tolist()

                for w, so_far in zip(predicted_lst, sampled_ids):
                    if so_far[-1] != self.eos_idx:
                        so_far.append(w)

                inputs = predicted.transpose(0, 1)          # inputs: (L=1,B)
                inputs = self.embedding(inputs)             # inputs: (L=1,B,E)

            sampled_lengths = [len(text) for text in sampled_ids]
            sampled_lengths = np.array(sampled_lengths)

            max_length = max(sampled_lengths)
            padded_ids = np.ones((batch_size, max_length)) * self.pad_idx

            for i in xrange(batch_size):
                padded_ids[i, :sampled_lengths[i]] = sampled_ids[i]

            sampled_lengths = torch.from_numpy(sampled_lengths).long()
            sampled_ids = torch.from_numpy(padded_ids).long()

        return sampled_ids, sampled_lengths

class LabelEncoder(nn.Module):
    r"""Parameterizes q(z|label).

    @param z_dim: integer
                  number of latent dimensions
    @param label_dim: integer
                      number of label dimensions
    @param hidden_dim: integer [default: 256]
                       number of hidden dimensions
    """
    def __init__(self, z_dim, label_dim, hidden_dim=256):
        super(LabelEncoder, self).__init__()
        self.z_dim = z_dim
        self.label_dim = label_dim
        self.hidden_dim = hidden_dim

        self.net = nn.Sequential(
            nn.Linear(self.label_dim, self.hidden_dim),
            Swish(),
            nn.Linear(self.hidden_dim, self.hidden_dim),
            Swish(),
            # here we don't explicitly separate into (mu, logvar)
            # but equivalent thing
            nn.Linear(self.hidden_dim, self.z_dim * 2))

    def forward(self, x):
        h = self.net(x)
        z_mu, z_logvar = torch.chunk(h, 2, dim=1)

        return z_mu, z_logvar

class LabelDecoder(nn.Module):
    r"""Parameterizes p(label|z).

    @param z_dim: integer
                  number of latent dimensions
    @param label_dim: integer
                      number of label dimensions
    @param hidden_dim: integer [default: 256]
                       number of hidden dimensions
    """
    def __init__(self, z_dim, label_dim, hidden_dim=256):
        super(LabelDecoder, self).__init__()

        self.z_dim = z_dim
        self.label_dim = label_dim
        self.hidden_dim = hidden_dim

        self.net = nn.Sequential(
            nn.Linear(self.z_dim, self.hidden_dim),
            Swish(),
            nn.Linear(self.hidden_dim, self.hidden_dim),
            Swish(),
            nn.Linear(self.hidden_dim, self.hidden_dim),
            Swish(),
            nn.Linear(self.hidden_dim, self.label_dim))

    def forward(self, z):
        # we assume binary labels
        return torch.sigmoid(self.net(z))

class ProductOfExperts(nn.Module):
    r"""Return parameters for product of independent experts.
    See https://arxiv.org/pdf/1410.7827.pdf for equations.

    @param mu: M x D for M experts
    @param logvar: M x D for M experts
    """
    def forward(self, mu, logvar, eps=1e-8):
        var = torch.exp(logvar) + eps
        T = 1 / var  # precision of i-th Gaussian expert at point x
        pd_mu = torch.sum(mu * T, dim=0) / torch.sum(T, dim=0)
        pd_var = 1 / torch.sum(T, dim=0)
        pd_logvar = torch.log(pd_var)

        return pd_mu, pd_logvar

class Swish(nn.Module):
    def forward(self, x):
        return x * torch.sigmoid(x)

def swish(x):
    return x * torch.sigmoid(x)


#!/usr/bin/env python
# -*- coding:utf8 -*-
"""
@author: ruohua
@file: eval.py
@time: 2021/2/22 8:03 PM
"""
import numpy as np
import argparse
from multiprocessing.dummy import Pool as ThreadPool
import os
import cv2
from functools import partial
import glob

def single_eval_boundary(fg_boundary, gt_boundary, bound_pix=0):
    """
    Compute mean,recall and decay from per-frame evaluation.
    Calculates precision/recall for boundaries between foreground_mask and
    gt_mask using morphological operators to speed it up.

    Arguments:
        fg_boundary (ndarray): binary boundary prediction image.     Shape: H x W     Value: 0 and 255
        gt_boundary (ndarray): binary annotated boundary image.      Shape: H x W     Value: 0 and 255
        bound_pix: half of the thickness of boundary.
                  A morphology dilation will make the thickness of edge to what we set here
                   this is the radius of disk

    Returns:
        F (float): boundaries F-measure
        P (float): boundaries precision
        R (float): boundaries recall
    """
    assert np.atleast_3d(fg_boundary).shape[2] == 1

    from skimage.morphology import binary_dilation, disk

    fg_dil = binary_dilation(fg_boundary, disk(bound_pix))
    gt_dil = binary_dilation(gt_boundary, disk(bound_pix))

    # Get the intersection
    gt_match = gt_boundary * fg_dil
    fg_match = fg_boundary * gt_dil

    # Area of the intersection
    n_fg = np.sum(fg_boundary)
    n_gt = np.sum(gt_boundary)

    # % Compute precision and recall
    if n_fg == 0 and n_gt > 0:
        precision = 1
        recall = 0
    elif n_fg > 0 and n_gt == 0:
        precision = 0
        recall = 1
    elif n_fg == 0 and n_gt == 0:
        precision = 1
        recall = 1
    else:
        precision = np.sum(fg_match) / float(n_fg)
        recall = np.sum(gt_match) / float(n_gt)

    # Compute F meas
    # ure
    if precision + recall == 0:
        F = 0
    else:
        F = 2 * precision * recall / (precision + recall)

    return F, precision, recall

def reverse_black_and_white(img):
    """
    reverse black and white color
    :param img: binary img
    :return:
    """
    img = 255 - img

    return img

def check_if_white_back_black_edge(pred):
    """
        check if the prediction image is binary
            1. 0 for membrane; 255 for not manbrane
            2. check if binary
    """
    values = np.unique(pred)
    # print(values)

    # check if binary
    if len(values) > 2:
        print("Your prediction result has not been binarized, please prompt them to choose the appropriate threshold for binarization.")
        raise ValueError

    white_pos = np.where(pred == 255)
    # print(len(white_pos[0]))
    white_count = len(white_pos[0])
    black_pos = np.where(pred == 0)
    # print(len(black_pos[0]))
    black_count = len(black_pos[0])
    # print(black_count / white_count)
    rate = black_count / white_count
    if rate < 5:
        print("The results must be submitted with white background and black edge. Please submit after correction.")
        raise ValueError

def single_eval_wrapper(bound_pix, result_matrix, data_list_line):
    id_, pred_path, gt_path = data_list_line

    print("Now evaluating the %s prediction, path = %s" % (id_, pred_path))

    pred_ = cv2.imread(pred_path, 0)
    gt_ = cv2.imread(gt_path, 0)

    pred_ = reverse_black_and_white(pred_)
    gt_ = reverse_black_and_white(gt_)

    check_if_white_back_black_edge(pred_)

    F, precision, recall = single_eval_boundary(pred_, gt_, bound_pix)

    result_matrix[int(id_)] = [F, precision, recall]

def eval_on_whole_dataset(pred_folder_path, gt_folder_path, bound_pix, thread_num):
    pred_list_all =  glob.glob(pred_folder_path + '/*.png')
    print(pred_list_all)

    gt_list_all = []
    for pre in pred_list_all:
        gt_list_all.append(pre.replace(pred_folder_path,gt_folder_path))
    print(gt_list_all)

    id_list = [x for x in range(len(gt_list_all))]

    gt_list = np.asarray(gt_list_all)
    pred_list = np.asarray(pred_list_all)
    id_list = np.asarray(id_list)

    gt_list = np.expand_dims(gt_list, 1)
    pred_list = np.expand_dims(pred_list, 1)
    id_list = np.expand_dims(id_list, 1)

    data_list = np.concatenate([id_list, pred_list, gt_list], axis=1)  # data list for pool

    result_matrix = np.zeros_like(data_list)  # each line contains: F-score, precision, recall

    pool = ThreadPool(thread_num)

    pool.map(partial(single_eval_wrapper, bound_pix, result_matrix), data_list)

    pool.close()
    pool.join()

    result_matrix = np.asarray(result_matrix, dtype=np.float)
    F_score = result_matrix[:, 0]
    print(F_score)
    mean_F = np.mean(F_score)

    return mean_F

def eval_on_onepic(pred_path, gt_path, bound_pix=0):
    pred_ = cv2.imread(pred_path, 0)
    gt_ = cv2.imread(gt_path, 0)

    pred_ = 255 - pred_
    gt_ = 255 - gt_

    check_if_white_back_black_edge(pred_)
    print("Check format: OK!")

    F, precision, recall = single_eval_boundary(pred_, gt_, bound_pix=0)
    return F

def calc_ap(prec, rec):
    mrec = np.array([0, rec, 1])
    mpre = np.array([0, prec, 0])

    for i in range(mrec.size - 2, -1, -1):
        mpre[i] = max(mpre[i], mpre[i+1]);
        #print mpre[i]

    idx1 = np.where(mrec[1 : ]   != mrec[0 : 2])
    idx2 = [x + 1 for x in idx1]

    ap = sum((mrec.take(idx2) - mrec.take(idx1)) * mpre.take(idx2))
    print("ap = " + str(ap[0]))
    return ap[0]

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--evalon1pic', type=str,
                    default='True', help='if evaluate on one image, True or False')
    parser.add_argument('--pre_path', type=str,
                    default='./pre.png', help='prediction path')
    parser.add_argument('--gt_path', type=str,
                    default='./gt.png', help='ground truth path')
    args = parser.parse_args()

    pre_path = args.pre_path
    gt_path = args.gt_path
    bound_pix = 0  # the thickness of boundary
    thread_num = 4  # number of parallel threads

    print("Evaluate on one image? {}".format(args.evalon1pic))

    if args.evalon1pic == 'True':
        # eval on single image
        print("evaluate {} and {}".format(pre_path,gt_path))
        mean_F = eval_on_onepic(pre_path, gt_path, bound_pix)
        print("The F1-score of prediction is : {}".format(mean_F))

    if args.evalon1pic == 'False':
        # eval on folder
        mean_F = eval_on_whole_dataset(pre_path, gt_path, bound_pix, thread_num)
        print("Mean F_score of all prediction is {}".format(mean_F))


# -*- coding: utf-8 -*-
# Author: Daniel Ryan <ryand5@tcd.ie>

"""Some IRIS instrument tools."""

import datetime
import warnings
import os.path

import numpy as np
import astropy.units as u
from astropy.units.quantity import Quantity
from astropy.modeling import fitting
from astropy.modeling.models import custom_model
from astropy import constants
import scipy.io
from scipy import ndimage
from scipy import interpolate
from sunpy.time import parse_time
import sunpy.util.config
from sunpy.util.net import check_download_file
from ndcube import NDCube

# Define some properties of IRIS detectors.  Source: IRIS instrument
# paper.
DETECTOR_GAIN = {"NUV": 18., "FUV": 6., "SJI": 18.}
DETECTOR_YIELD = {"NUV": 1., "FUV": 1.5, "SJI": 1.}
SJI_DEFAULT_BSCALE = 0.25
SJI_DEFAULT_BZERO = 7992.0
DN_UNIT = {
    "NUV": u.def_unit("DN_IRIS_NUV",
                      DETECTOR_GAIN["NUV"] / DETECTOR_YIELD["NUV"]*u.photon),
    "FUV": u.def_unit("DN_IRIS_FUV",
                      DETECTOR_GAIN["FUV"]/DETECTOR_YIELD["FUV"]*u.photon),
    "SJI": u.def_unit("DN_IRIS_SJI",
                      DETECTOR_GAIN["SJI"]/DETECTOR_YIELD["SJI"]*u.photon),
    "SJI_UNSCALED": u.def_unit("DN_IRIS_SJI_UNSCALED", u.ct)}
# Define an equivalency between SJI and SJI_UNSCALED units
SJI_SCALING = [(DN_UNIT["SJI"],
                DN_UNIT["SJI_UNSCALED"],
                lambda x: (x - SJI_DEFAULT_BZERO) / SJI_DEFAULT_BSCALE,
                lambda x: x * SJI_DEFAULT_BSCALE + SJI_DEFAULT_BZERO)]

READOUT_NOISE = {"NUV": 1.2*DN_UNIT["NUV"],
                 "FUV": 3.1*DN_UNIT["FUV"],
                 "SJI": 1.2*DN_UNIT["SJI"]}
RADIANCE_UNIT = u.erg / u.cm ** 2 / u.s / u.steradian / u.Angstrom
SLIT_WIDTH = 0.33*u.arcsec

IRIS_RESPONSE_REMOTE_PATH = "https://sohowww.nascom.nasa.gov/solarsoft/iris/response/"
RESPONSE_VERSION_FILENAMES = {"1": "iris_sra_20130211.geny",
                              "2": "iris_sra_20130715.geny",
                              "3": "iris_sra_c_20150331.geny",
                              "4": "iris_sra_c_20161022.geny"}

# Define some custom error messages.
APPLY_EXPOSURE_TIME_ERROR = ("Exposure time correction has probably already "
                             "been applied since the unit already includes "
                             "inverse time. To apply exposure time correction "
                             "anyway, set 'force' kwarg to True.")
UNDO_EXPOSURE_TIME_ERROR = ("Exposure time correction has probably already "
                            "been undone since the unit does not include "
                            "inverse time. To undo exposure time correction "
                            "anyway, set 'force' kwarg to True.")

# Define whether IRIS WCS is 0 or 1 origin based.
WCS_ORIGIN = 1

def get_iris_response(pre_launch=False, response_file=None, response_version=None,
                      force_download=False):
    """Returns IRIS response structure.

    One and only one of pre_launch, response_file and response_version must be set.

    Parameters
    ----------
    pre_launch: `bool`
        Equivalent to setting response_version=2.  Cannot be set
        simultaneously with response_file kwarg. Default=False
    response_file: `str`
        Version number of effective area file to be used.  Cannot be set
        simultaneously with pre_launch kwarg.  Default=latest
    response_version : `int`
        Version number of effective area file to be used. Cannot be set
        simultaneously with response_file or pre_launch kwarg. Default=latest

    Returns
    -------
    iris_response: `dict`
        Various parameters regarding IRIS response.  The following keys:
        date_obs: `datetime.datetime`
        lambda: `astropy.units.Quantity`
        area_sg: `astropy.units.Quantity`
        name_sg: `str`
        dn2phot_sg: `tuple` of length 2
        area_sji: `astropy.units.Quantity`
        name_sji: `str`
        dn2phot_sji:  `tuple` of length 4
        comment: `str`
        version: `int`
        version_date: `datetime.datetime`

    Notes
    -----
    This routine does not calculate time dependent effective areas using
    version 3 and above of the response functions as is done in the SSW version
    of this code.  Therefore, asking it to read a version 3 or above response
    function will result in an error.  This code should be updated in future
    versions to calculate time dependent effective areas.

    """
    # Ensures the file exits in the path given.
    if response_file is not None:
        if not(os.path.isfile(response_file)):
            raise KeyError("Not a valid file path")

    # Ensure conflicting kwargs are not set.
    if response_file:
        response_file_set = True
    else:
        response_file_set = False
    if response_version:
        response_version_set = True
    else:
        response_version_set = False
    if response_file_set+pre_launch+response_version_set != 1:
        raise ValueError("One and only one of kwargs pre_launch, response_file "
                         "and response_version must be set.")
    # If pre_launch set, define response_version to 2.
    if pre_launch:
        response_version = 2
    # If response_file not set, define appropriate response file
    # based on version.
    if not response_file:
        try:
            response_filename = RESPONSE_VERSION_FILENAMES[str(response_version)]
        except KeyError:
            raise KeyError("Version number not recognized.")
        if response_version > 2:
            warnings.warn("Effective areas are not available (i.e. set  to zero).  "
                  "For response file versions > 2 time dependent effective "
                  "areas must be calculated via fitting, which is not supported "
                  "by this function at this time. "
                  "Version of this response file = {0}".format(response_version))
        # Define the directory in which the response file should exist
        # to be the sunpy download directory.
        config = sunpy.util.config.load_config()
        download_dir = config.get('downloads', 'download_dir')
        # Check response file exists in download_dir.  If not, download it.
        check_download_file(response_filename, IRIS_RESPONSE_REMOTE_PATH, download_dir,
                            replace=force_download)
        # Define response file as path + filename.
        response_file = os.path.join(download_dir, response_filename)

    # Read response file and store in a dictionary.
    raw_response_data = scipy.io.readsav(response_file)
    iris_response = dict([(name, raw_response_data["p0"][name][0])
                          for name in raw_response_data["p0"].dtype.names])
    # Convert some properties to more convenient types.
    iris_response["LAMBDA"] = Quantity(iris_response["LAMBDA"], unit=u.nm)
    iris_response["AREA_SG"] = Quantity(iris_response["AREA_SG"], unit=u.cm**2)
    iris_response["AREA_SJI"] = Quantity(iris_response["AREA_SJI"], unit=u.cm**2)
    iris_response["GEOM_AREA"] = Quantity(iris_response["GEOM_AREA"], unit=u.cm**2)
    iris_response["VERSION"] = int(iris_response["VERSION"])
    # Convert some properties not found in version below version 3 to
    # more convenient types.
    if iris_response["VERSION"] > 2:
        # If DATE_OBS has a value, convert to datetime, else set to
        # None.
        try:
            iris_response["DATE_OBS"] = parse_time(iris_response["DATE_OBS"])
        except:
            iris_response["DATE_OBS"] = None
        # Convert C_F_TIME to array of datetime objects while
        # conserving shape.
        c_f_time = np.empty(iris_response["C_F_TIME"].shape, dtype=object)
        for i, row in enumerate(iris_response["C_F_TIME"]):
            for j, t in enumerate(row):
                c_f_time[i][j] = parse_time(float(t))
        iris_response["C_F_TIME"] = c_f_time
        # Convert C_F_LAMBDA to Quantity.
        iris_response["C_F_LAMBDA"] = Quantity(iris_response["C_F_LAMBDA"], unit="nm")
        # Convert C_N_TIME to array of datetime objects while
        # conserving shape.
        c_n_time = np.empty(iris_response["C_N_TIME"].shape, dtype=object)
        for i, row in enumerate(iris_response["C_N_TIME"]):
            for j, t in enumerate(row):
                c_n_time[i][j] = parse_time(float(t))
        iris_response["C_N_TIME"] = c_n_time
        # Convert C_N_LAMBDA to Quantity.
        iris_response["C_N_LAMBDA"] = Quantity(iris_response["C_N_LAMBDA"], unit="nm")
        # Convert C_S_TIME to array of datetime objects while
        # conserving shape.
        c_s_time = np.empty(iris_response["C_S_TIME"].shape, dtype=object)
        for i, row in enumerate(iris_response["C_S_TIME"]):
            for j, column in enumerate(row):
                for k, t in enumerate(column):
                    c_s_time[i][j][k] = parse_time(float(t))
        iris_response["C_S_TIME"] = c_s_time
        # Convert DATE in ELEMENTS array to array of datetime objects.
        for i, t in enumerate(iris_response["ELEMENTS"]["DATE"]):
            iris_response["ELEMENTS"]["DATE"][i] = parse_time(t.decode())
        # Convert VERSION_DATE to datetime object.
        iris_response["VERSION_DATE"] = parse_time(iris_response["VERSION_DATE"].decode())
    else:
        # Change DATE tag in data with version < 2 to VERSION_DATE to
        # be consistent with more recent versions.
        iris_response["VERSION_DATE"] = datetime.datetime(int(iris_response["DATE"][0:4]),
                                                          int(iris_response["DATE"][4:6]),
                                                          int(iris_response["DATE"][6:8]))
        del(iris_response["DATE"])
    return iris_response

@custom_model
def _gaussian1d_on_linear_bg(x, amplitude=None, mean=None, standard_deviation=None,
                             constant_term=None, linear_term=None):
    return amplitude * np.exp(-((x - mean) / standard_deviation) ** 2) + constant_term + linear_term * x

def _calculate_orbital_wavelength_variation(data_array, date_data_created, slit_pixel_range=None,
                                            spline_smoothing=False, fit_individual_profiles=False,
                                            spacecraft_velocity=None, orbital_phase=None, roll_angle=None):
    """Calculates orbital corrections of spectral line positions using level 2 files.

    For data generated from the April 2014 pipeline, thermal and spacecraft velocity components
    have both been subtracted in the level 2 files.  Therefore, this routine calculates the
    residual orbital (thermal) variation.  For data generated from the Oct 2013 pipeline,
    this routine calculates the total of thermal and spacecraft velocity components.

    Parameters
    ----------
    data_array: `xarray.DataArray`
        IRIS spectrograph data from spectral window Mg II k 2796 as generated by
        `sunpy.spectra.sources.IRISRaster.`
    date_data_created: `datetime.datetime`
        Date the data was created by IRIS pipeline.  Used to determine where spacecraft
        velocity etc. needs to be accounted for.
    spacecraft_velocity: `astropy.units.quantity.Quantity`
        Velocity of spacecraft at each exposure in data_array.
        Must be set if date_data_created < 1 April 2014.
    orbital_phase: `numpy.array`
        Orbital phase of spacecraft at each exposure in data_array.  Available from
        auxiliary data in IRIS spectrograph fits files.
        Must be set if date_data_created < 1 April 2014.
    roll_angle: `astropy.units.quantity.Quantity`
        Roll angle of spacecraft. Must be set if date_data_created < 1 April 2014.

    Returns
    -------
    orbital_wavelength_variation: `astropy.table.Table`
        Contains the following columns:
        time: `datetime.datetime` objects
            Observation times of wavelength variations.
        FUV: `astropy.quantity.Quantity`
            Wavelength variation in the FUV.
        NUV: `astropy.quantity.Quantity`
            Wavelength variation in the NUV.

    """
    # Define vacuum rest wavelength of Ni I 2799 line.
    wavelength_nii = 2799.474 * u.Angstrom
    # Define factor converting NUV spectral pixel size to Angstrom
    specsize = 0.0255
    # Define date of new pipeline.
    date_new_pipeline = datetime.datetime(2014, 4, 1)
    if date_data_created < date_new_pipeline:
        # Check that there are measurement times with good values of
        # spacecraft velocity and orbital phase.
        bad_aux = np.asarray(np.isfinite(spacecraft_velocity) * np.isfinite(orbital_phase) * (-1), dtype=bool)
    # Generate wavelength vector containing only Ni I line.
    wavelength_window = Quantity(data_array.coords["wavelength"].values,
                                 unit=data_array.attrs["units"]["wavelength"])
    wavelength_roi_index = np.arange(len(wavelength_window))[
        np.logical_and(wavelength_window >= 2799.3 * u.Angstrom, wavelength_window <= 2799.8 * u.Angstrom)]
    # Check that there are at least 5 points in wavelength region.
    # Must have at least this many for a gaussian fit.
    if len(wavelength_roi_index) < 5:
        wavelength_roi_index = np.arange(5) + wavelength_roi_index[0]
    # Extract wavelength of region around Ni I line as array in units
    # of Angstroms.
    wavelength_roi = wavelength_window.to(u.Angstrom).value[wavelength_roi_index]
    # Keep only data within wavelength region of interest.
    data_array = data_array.isel(spectral_axis=slice(wavelength_roi_index[0], wavelength_roi_index[-1] + 1))
    # If user selected a sub-region of the slit, reduce data to just
    # that region.
    if slit_pixel_range:
        if len(slit_pixel_range) == 2:
            data_array = data_array.isel(slit_axis, slice(slit_pixel_range[0], slit_pixel_range[1]))
        else:
            raise TypeError("slit_pixel_range must be tuple of length 2 giving lower and " +
                            "upper bounds of section of slit over which to average line fits.")

    # Derive residual orbital variation.
    # Define array to hold averaged position of Ni I line at different
    # times.
    mean_line_wavelengths = np.empty(len(data_array.time)) * np.nan
    # Define initial guess for gaussian model.
    g_init = _gaussian1d_on_linear_bg(amplitude=-2., mean=wavelength_nii.value,
                                      standard_deviation=2., constant_term=50., linear_term=1.5)
    # Define fitting method.
    fit_g = fitting.LevMarLSQFitter()
    # Depending on user choice, either fit line as measured by each
    # pixel then average line position, or fit average line spectrum
    # from all slit pixels.
    if fit_individual_profiles:
        pixels_in_slit = len(raster.slit_axis)
        for k in range(len(raster.time)):
            pixel_line_wavelengths = np.empty(pixels_in_slit)*np.nan
            data_single_time = raster.isel(raster_axis=k)
            # Iterate through each pixel along slit and perform fit to
            # Ni I line.
            for j in range(2, pixels_in_slit-2):
                # Average over 5 pixels to improve signal-to-noise.
                intensity_mean_5pix = data_single_time.isel(slit_axis=slice(j-2, j+3)).mean(axis=0)
                # Fit gaussian to Ni I line.
                g = fit_g(g_init, wavelength_roi, intensity_mean_5pix)
                # Check that fit is within physically reasonable
                # limits.  If so, store line center wavelength in
                # mean_line_wavelengths array. Else leave element as
                # defined, i.e. NaN.
                if np.isfinite(g.amplitude) and g.amplitude < 0. and \
                            wavelength_roi[0] < g.mean < wavelength_roi[-1]:
                    pixel_line_wavelengths[j] = g.mean
            # Take average of Ni I line position from fits in each
            # pixel.
            mean_line_wavelengths[k] = np.nanmean(pixel_line_wavelengths)
    else:
        # Else average all line profiles then perform fit.
        # Iterate through each measurement time and fit a gaussian to
        # Ni I line.
        for k in range(len(raster.time)):
            # Get data averaged over slit.
            data_single_time = raster.isel(raster_axis=k)
            data_slit_averaged = data_single_time.to_masked_array().mean(axis=0).data
            # Fit Ni I line with a gaussian.
            # Perform fit.
            g = fit_g(g_init, wavelength_roi, data_slit_averaged)
            # Check that fit is within physically reasonable limits.
            # If so, store line center wavelength in
            # mean_line_wavelengths array. Else leave element as
            # defined, i.e. NaN.
            if np.isfinite(g.amplitude) and g.amplitude < 0. and \
                        wavelength_roi[0] < g.mean < wavelength_roi[-1]:
                mean_line_wavelengths[k] = g.mean
            # If data produced by old pipeline, subtract spacecraft velocity
            # from the line position.
            if date_created < date_new_pipeline:
                mean_line_wavelengths[k] = \
                    mean_line_wavelengths[k]-spacecraft_velocity[k]/3e8*wavelength_nii.to(u.Angstrom).value

    # Mark abnormal values.  Thermal drift is of the order of 2
    # unsummed wavelength pixels peak-to-peak.
    w_abnormal = np.where(np.abs(mean_line_wavelengths-np.nanmedian(mean_line_wavelengths)) >= specsize*2)[0]
    if len(w_abnormal) > 0:
        mean_line_wavelengths[w_abnormal] = np.nan
    # Further data reduction required for files from old pipeline.
    if date_created < date_new_pipeline:
        dw_th_A = mean_line_wavelengths - np.nanmean(mean_line_wavelengths)
        # Change the unit from Angstrom into unsummed wavelength pixel.
        dw_th_p = dw_th_A/specsize
        # Adjust reference wavelength using orbital phase information.
        if not(np.isfinite(orbital_phase)).all():
            warnings.warn("Orbital phase values are invalid.  Thermal drift may be offset by at most one pixel.")
            dw_th = dw_th
            # For absolute wavelength calibration of NUV, the
            # following amount (unit Angstrom) has to be
            # subtracted from the wavelengths.
            abswvl_nuv = np.nanmean(mean_line_wavelengths)-wavelength_nii.to(u.Angstrom).value
        else:
            # Define empirical sine fitting at 0 roll angle shifted by
            # different phase.
            sine_params = [-0.66615146, -1.0, 53.106583-roll_angle/360.*2*np.pi]
            phase_adj=np.nanmean(sine_params[0]*np.sin(sine_params[1]*orbital_phase+sine_params[2]))
            # thermal component of the orbital variation, in the unit of unsummed wavelength pixel
            dw_th=dw_th_p+phase_adj
            # For absolute wavelength calibration of NUV the following
            # amount (unit Angstrom) has to be subtracted from the
            # wavelengths.
            abswvl_nuv = np.nanmean(mean_line_wavelengths)-wavelength_nii.to(u.Angstrom).value-phase_adj*specsize
    else:
        # Calculate relative variation of the line position.
        dw_th = mean_line_wavelengths-np.nanmean(mean_line_wavelengths)

    # If spline_smoothing=True, perform spline fit a smoothing to
    # eliminate the 5 minute photospheric oscillation.
    if spline_smoothing:
        # Define spacing of spline knots in seconds.
        spline_knot_spacing = 300.
        # Create array of time in seconds from first time and
        # calculate duration of fitting period.
        time_s = np.asarray(x.coords["time"]-x.coords["time"][0], dtype=float)/1e9
        duration = time_s[-1]-time_s[0]
        # Check whether there is enough good data for a spline fit.
        if duration < spline_knot_spacing:
            raise ValueError("Not enough data for spline fit.")
        # Check whether there is enough good data for a spline fit.
        wgood = np.isfinite(mean_line_wavelengths)
        ngood = float(sum(wgood))
        wbad = not(np.isfinite(mean_line_wavelengths))
        nbad = float(sum(wbad))
        if nbad/ngood > 0.25:
            raise ValuError("Not enough good data for spline fit.")
        # Smooth residual thermal variation curve to eliminate the
        # 5-min photospheric oscillation.
        # Determine number of smoothing point using 3 point
        # lagrangian derivative.
        deriv_time = np.array([(time_s[i+1]-time_s[i-1])/2. for i in range(1,len(time_s)-1)])
        deriv_time = np.insert(deriv_time, 0, (-3*time_s[0]+4*time_s[1]-time_s[2])/2)
        deriv_time = np.insert(deriv_time, -1, (3*time_s[-1]-4*time_s[-2]+time_s[-3])/2)
        n_smooth = int(spline_knot_spacing/deriv_time.mean())
        if n_smooth < len(wgood):
            dw_good = convolve(dw_th[good], Box1DKernel(n_smooth))
        else:
            dw_good = dw_th[good]
        time_good = time_s[good]
        # Fit spline.
        tck = interpolate.splrep(time_good, dw_good, s=0)
        dw_th = interpolate.splev(time_s, tck)

    # Derive residual orbital curves in FUV and NUV and store
    # in a table.
    times = [datetime.datetime.utcfromtimestamp(t/1e9) for t in raster.coords["time"].values.tolist()]
    # Depeding on which pipeline produced the files...
    if date_created < date_new_pipeline:
        dw_orb_fuv = dw_th * (-0.013) + spacecraft_velocity.to(u.km/u.s).value / (3.e5) * 1370. * u.Angstrom
        dw_orb_nuv = dw_th * 0.0255 + spacecraft_velocity.to(u.km/u.s).value / (3.e5) * 2800. * u.Angstrom
    else:
        dw_orb_fuv = dw_th*(-1)*u.Angstrom
        dw_orb_nuv = dw_th*u.Angstrom

    orbital_wavelength_variation = Table([times, dw_orb_fuv, dw_orb_nuv],
                                         names=("time", "wavelength variation FUV", "wavelength variation NUV"))
    return orbital_wavelength_variation

def get_detector_type(meta):
    """
    Gets the IRIS detector type from a meta dictionary.

    In this function, FUV1 and FUV2 are just assigned as FUV.

    Parameters
    ----------
    meta: dict-like
        Dictionary-like object containing entry for "detector type"

    Returns
    -------
    detector_type: `str`
       Detector type.

    """
    if "FUV" in meta["detector type"]:
        detector_type = "FUV"
    else:
        detector_type = meta["detector type"]
    return detector_type

def convert_between_DN_and_photons(old_data_arrays, old_unit, new_unit):
    """Converts arrays from IRIS DN to photons or vice versa.

    In this function, an inverse time component due to exposure time
    correction is ignored during calculations but preserved in final unit.

    Parameters
    ----------
    old_data_arrays: iterable of `numpy.ndarray`s
        Arrays of data to be converted.

    old_unit: `astropy.unit.Unit`
        Unit of data arrays.

    new_unit: `astropy.unit.Unit`
        Unit to convert data arrays to.

    Returns
    -------
    new_data_arrays: `list` of `numpy.ndarray`s
        Data arrays converted to new_unit.

    new_unit_time_accounted: `astropy.unit.Unit`
        Unit of new data arrays with any inverse time component preserved.

    """
    if old_unit == new_unit or old_unit == new_unit / u.s:
        new_data_arrays = [data for data in old_data_arrays]
        new_unit_time_accounted = old_unit
    else:
        # During calculations, the time component due to exposure
        # time correction, if it has been applied, is ignored.
        # Check here whether the time correction is present in the
        # original unit so that is carried through to new unit.
        if u.s not in (old_unit * u.s).decompose().bases:
            old_unit_without_time = old_unit * u.s
            new_unit_time_accounted = new_unit / u.s
        else:
            old_unit_without_time = old_unit
            new_unit_time_accounted = new_unit
        # Convert data and uncertainty to new unit.
        new_data_arrays = [(data * old_unit_without_time).to(new_unit).value
                           for data in old_data_arrays]
    return new_data_arrays, new_unit_time_accounted

def calculate_exposure_time_correction(old_data_arrays, old_unit, exposure_time,
                                       force=False):
    """
    Applies exposure time correction to data arrays.

    Parameters
    ----------
    old_data_arrays: iterable of `numpy.ndarray`s
        Arrays of data to be converted.

    old_unit: `astropy.unit.Unit`
        Unit of data arrays.

    exposure_time: `numpy.ndarray`
        Exposure time in seconds for each exposure in data arrays.

    Returns
    -------
    new_data_arrays: `list` of `numpy.ndarray`s
        Data arrays with exposure time corrected for.

    new_unit_time_accounted: `astropy.unit.Unit`
        Unit of new data arrays after exposure time correction.

    """
    # If force is not set to True and unit already includes inverse time,
    # raise error as exposure time correction has probably already been
    # applied and should not be applied again.
    if force is not True and u.s in old_unit.decompose().bases:
        raise ValueError(APPLY_EXPOSURE_TIME_ERROR)
    else:
        # Else, either unit does not include inverse time and so
        # exposure does need to be applied, or
        # user has set force=True and wants the correction applied
        # regardless of the unit.
        new_data_arrays = [old_data/exposure_time for old_data in old_data_arrays]
        new_unit = old_unit/u.s
    return new_data_arrays, new_unit

def uncalculate_exposure_time_correction(old_data_arrays, old_unit, exposure_time,
                                         force=False):
    """
    Removes exposure time correction from data arrays.

    Parameters
    ----------
    old_data_arrays: iterable of `numpy.ndarray`s
        Arrays of data to be converted.

    old_unit: `astropy.unit.Unit`
        Unit of data arrays.

    exposure_time: `numpy.ndarray`
        Exposure time in seconds for each exposure in data arrays.

    Returns
    -------
    new_data_arrays: `list` of `numpy.ndarray`s
        Data arrays with exposure time correction removed.

    new_unit_time_accounted: `astropy.unit.Unit`
        Unit of new data arrays after exposure time correction removed.

    """
    # If force is not set to True and unit does not include inverse time,
    # raise error as exposure time correction has probably already been
    # undone and should not be undone again.
    if force is not True and u.s in (old_unit*u.s).decompose().bases:
        raise ValueError(UNDO_EXPOSURE_TIME_ERROR)
    else:
        # Else, either unit does include inverse time and so
        # exposure does need to be removed, or
        # user has set force=True and wants the correction removed
        # regardless of the unit.
        new_data_arrays = [old_data * exposure_time for old_data in old_data_arrays]
        new_unit = old_unit*u.s
    return new_data_arrays, new_unit

def convert_or_undo_photons_per_sec_to_radiance(
        data_quantities, obs_wavelength, detector_type,
        spectral_dispersion_per_pixel, solid_angle, undo=False):
    """
    Converts data quantities from counts/s to radiance (or vice versa).

    Parameters
    ----------
    data_quantities: iterable of `astropy.units.Quantity`s
        Quantities to be converted.  Must have units of counts/s or
        radiance equivalent counts, e.g. erg / cm**2 / s / sr / Angstrom.

    obs_wavelength: `astropy.units.Quantity`
        Wavelength at each element along spectral axis of data quantities.

    detector_type: `str`
        Detector type: 'FUV', 'NUV', or 'SJI'.

    spectral_dispersion_per_pixel: scalar `astropy.units.Quantity`
        spectral dispersion (wavelength width) of a pixel.

    solid_angle: scalar `astropy.units.Quantity`
        Solid angle corresponding to a pixel.

    undo: `bool`
        If False, converts counts/s to radiance.
        If True, converts radiance to counts/s.
        Default=False

    Returns
    -------
    new_data_quantities: `list` of `astropy.units.Quantity`s
        Data quantities converted to radiance or counts/s
        depending on value of undo kwarg.

    """
    # Check data quantities are in the right units.
    if undo is True:
        for i, data in enumerate(data_quantities):
            if not data.unit.is_equivalent(RADIANCE_UNIT):
                raise ValueError(
                    "Invalid unit provided.  As kwarg undo=True, "
                    "unit must be equivalent to {0}.  Error found for {1}th element "
                    "of data_quantities. Unit: {2}".format(RADIANCE_UNIT, i, data.unit))
    else:
        for data in data_quantities:
            if data.unit != u.photon/u.s:
                raise ValueError(
                    "Invalid unit provided.  As kwarg undo=False, "
                    "unit must be equivalent to {0}.  Error found for {1}th element "
                    "of data_quantities. Unit: {2}".format(u.photon/u.s, i, data.unit))
    photons_per_sec_to_radiance_factor = calculate_photons_per_sec_to_radiance_factor(
        obs_wavelength, detector_type, spectral_dispersion_per_pixel, solid_angle)
    # Change shape of arrays so they are compatible for broadcasting
    # with data and uncertainty arrays.
    photons_per_sec_to_radiance_factor = \
        _reshape_1D_wavelength_dimensions_for_broadcast(photons_per_sec_to_radiance_factor,
                                                        data_quantities[0].ndim)
    # Perform (or undo) radiometric conversion.
    if undo is True:
        new_data_quantities = [(data / photons_per_sec_to_radiance_factor).to(u.photon/u.s)
                               for data in data_quantities]
    else:
        new_data_quantities = [(data*photons_per_sec_to_radiance_factor).to(RADIANCE_UNIT)
                               for data in data_quantities]
    return new_data_quantities

def calculate_photons_per_sec_to_radiance_factor(
        wavelength, detector_type, spectral_dispersion_per_pixel, solid_angle):
    """
    Calculates multiplicative factor that converts counts/s to radiance for given wavelengths.

    Parameters
    ----------
    wavelength: `astropy.units.Quantity`
        Wavelengths for which counts/s-to-radiance factor is to be calculated

    detector_type: `str`
        Detector type: 'FUV' or 'NUV'.

    spectral_dispersion_per_pixel: scalar `astropy.units.Quantity`
        spectral dispersion (wavelength width) of a pixel.

    solid_angle: scalar `astropy.units.Quantity`
        Solid angle corresponding to a pixel.

    Returns
    -------
    radiance_factor: `astropy.units.Quantity`
        Mutliplicative conversion factor from counts/s to radiance units
        for input wavelengths.

    """
    # Get effective area and interpolate to observed wavelength grid.
    eff_area_interp = _get_interpolated_effective_area(detector_type, wavelength)
    # Return radiometric conversed data assuming input data is in units of photons/s.
    return constants.h * constants.c / wavelength / u.photon / \
           spectral_dispersion_per_pixel / eff_area_interp / solid_angle

def _get_interpolated_effective_area(detector_type, obs_wavelength):
    # Get effective area
    ########### This needs to be generalized to the time of OBS once that functionality is written #########
    iris_response = get_iris_response(pre_launch=True)
    if detector_type == "FUV":
        detector_type_index = 0
    elif detector_type == "NUV":
        detector_type_index = 1
    else:
        raise ValueError("Detector type not recognized.")
    eff_area = iris_response["AREA_SG"][detector_type_index, :]
    response_wavelength = iris_response["LAMBDA"]
    # Interpolate the effective areas to cover the wavelengths
    # at which the data is recorded:
    eff_area_interp_base_unit = u.Angstrom
    tck = interpolate.splrep(response_wavelength.to(eff_area_interp_base_unit).value,
                             eff_area.to(eff_area_interp_base_unit ** 2).value, s=0)
    eff_area_interp = interpolate.splev(
        obs_wavelength.to(eff_area_interp_base_unit).value, tck) * eff_area_interp_base_unit ** 2
    return eff_area_interp

def _reshape_1D_wavelength_dimensions_for_broadcast(wavelength, n_data_dim):
    if n_data_dim == 1:
        pass
    elif n_data_dim == 2:
        wavelength = wavelength[np.newaxis, :]
    elif n_data_dim == 3:
        wavelength = wavelength[np.newaxis, np.newaxis, :]
    else:
        raise ValueError("IRISSpectrogram dimensions must be 2 or 3.")
    return wavelength

def _convert_iris_sequence(sequence, new_unit):
    """Converts data and uncertainty in an IRISSpectrogramSequence between units.

    Parameters
    ----------
    sequence: `NDCubeSequence`, `SpectrogramSequence` or `IRISSpectrogramSequence`
        Sequence whose constituent NDCubes are be converted to new units.

    new_unit: `astropy.units.Unit` or `str`
       Unit to which the data is to be converted.

    Returns
    -------
    converted_data_list: `list` of `NDCube`s.
       List of NDCubes with data and uncertainty attributes converted to new_unit.

    """
    # Define empty list to hold NDCubes with converted data and uncertainty.
    converted_data_list = []
    # Cycle through each NDCube, convert data and uncertainty to new
    # units, and append to list.
    for i, cube in enumerate(sequence.data):
        # Determine what type of DN unit is needed based on detector type.
        detector_type = _get_detector_type(cube.meta)
        if new_unit == "DN":
            new_unit = DN_UNIT[detector_type]
        # If NDCube is already in new unit, add NDCube as is to list.
        if cube.unit is new_unit or cube.unit is new_unit / u.s:
            converted_data_list.append(cube)
        # Else convert data and uncertainty to new unit.
        if cube.unit != new_unit or cube.unit != new_unit / u.s:
            # During calculations, the time component due to exposure
            # time correction, if it has been applied, is ignored.
            # Check here whether the time correction is present in the
            # original unit so that is carried through to new unit.
            if u.s not in (cube.unit.decompose() * u.s).bases:
                new_unit_time_accounted = new_unit / u.s
            else:
                new_unit_time_accounted = new_unit
            # Convert data and uncertainty to new unit.
            data = (cube.data * cube.unit).to(new_unit).value
            uncertainty = (cube.uncertainty.array * cube.unit).to(new_unit).value
            # Append new instance of NDCube in new unit to list.
            converted_data_list.append(NDCube(
                data, wcs=cube.wcs, meta=cube.meta, mask=cube.mask,
                unit=new_unit_time_accounted, uncertainty=uncertainty,
                extra_coords=_extra_coords_to_input_format(cube._extra_coords)))
    return converted_data_list

def _apply_or_undo_exposure_time_correction(sequence, correction_function):
    """Applies or undoes exposure time correction to a sequence of NDCubes.

    Correction is applied (or undone) to both data and uncertainty attributes of NDCubes.

    Parameters
    ----------
    sequence: `NDCubeSequence`, `SpectrogramSequence` or `IRISSpectrogramSequence`
        Sequence whose constituent NDCubes are be converted to new units.
        NDCubes with sequence must have an 'exposure time' entry in its extra
        coords attribute.

    correction_function: function
        Function applying or undoing exposure time correction.

    Returns
    -------
    converted_data_list: `list` of `NDCube`s.
       List of NDCubes with data and uncertainty corrected (or uncorrected)
       for exposure time.

    """
    converted_data_list = []
    for i, cube in enumerate(sequence.data):
        if u.s not in cube.unit.decompose().bases:
            exposure_time_s = cube._extra_coords["exposure time"]["value"].to(u.s).value
            if len(cube.dimensions.shape) == 1:
                pass
            elif len(cube.dimensions.shape) == 2:
                exposure_time_s = exposure_time_s[:, np.newaxis]
            elif len(cube.dimensions.shape) == 3:
                exposure_time_s = exposure_time_s[:, np.newaxis, np.newaxis]
            else:
                raise ValueError("NDCube dimensions must be 2 or 3. Dimensions={0}".format(
                    len(cube.dimensions.shape)))
            data = correction_function(cube.data, exposure_time_s)
            uncertainty = correction_function(cube.uncertainty.array, exposure_time_s)
            converted_data_list.append(NDCube(
                data, wcs=cube.wcs, meta=cube.meta, mask=cube.mask, unit=cube.unit / u.s,
                uncertainty=uncertainty,
                extra_coords=_extra_coords_to_input_format(cube._extra_coords)))
        else:
            converted_data_list.append(cube)
    return converted_data_list

def calculate_dust_mask(data_array):
    """Calculate a mask with the dust positions in a given arrayself.

    Parameters
    ----------
    data_array : `numpy.ndarray`
        This array contains some dust poisition that will be calculated. The array
        must have scaled values.

    Returns
    -------
    dust : `numpy.ndarray` of `bool`
        This array has the same shape than data_array and contains the dust positions
        when the value is True.

    """
    # Creating a mask with the same shape than the inputed data array.
    mask = np.zeros_like(data_array, dtype=bool)
    # Set the pixel value to True is the pixel is recognized as a dust pixel.
    mask[(data_array < 0.5) & (data_array > -200)] = True
    # Extending the mask to avoid the neighbours pixel influenced by the dust pixels.
    struct = np.array([np.zeros((3, 3)), np.ones((3, 3)), np.zeros((3, 3))], dtype=bool)
    mask = ndimage.binary_dilation(mask, structure=struct).astype(mask.dtype)
    return mask

import numpy,import scipy,from scipy,import astropy,from astropy


from __future__ import division, print_function, absolute_import

from tmm.tmm_core import (coh_tmm, unpolarized_RT, ellips,
                       position_resolved, find_in_structure_with_inf)

import tmm.tmm_core as tmm

import numpy
from numpy import pi, linspace, inf, array
import pandas
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt

try:
    import colorpy.illuminants
    import colorpy.colormodels
    from . import color
    colors_were_imported = True
except ImportError:
    # without colorpy, you can't run sample5(), but everything else is fine.
    colors_were_imported = False

from matplotlib.pyplot import show
from argparse import ArgumentParser
import lowtran
from lowtran.lowtran.plots import plottrans

import xarray

#def main():
p = ArgumentParser(description='Lowtran 7 interface')
p.add_argument('-z', '--obsalt', help='altitude of observer [km]', type=float, default=0.)
p.add_argument('-a', '--zenang', help='observer zenith angle [deg]', type=float, nargs='+', default=[0, 45, 60, 80])
p.add_argument('-s', '--short', help='shortest wavelength nm ', type=float, default=200)
p.add_argument('-l', '--long', help='longest wavelength cm^-1 ', type=float, default=30000)
p.add_argument('-step', help='wavelength step size cm^-1', type=float, default=20)
p.add_argument('--model', help='0-6, see Card1 "model" reference. 6 = 1976 US Standard', type=int, default=6)
P = p.parse_args()

c1 = {'model': P.model,
      'h1': P.obsalt,
      'angle': P.zenang,
      'wlshort': P.short,
      'wllong': P.long,
      'wlstep': P.step,
      }

tr = lowtran.transmittance(c1)
wl = tr.wavelength_nm.data # This gives the wavelength range
#a_transmission_data = tr.transmission.data.squeeze() # Transmission data as a funciton of angle.

#a = tr.transmission # make a DataArray from a DataSet from Xarray package
#a2 = a.values # make
#a3 = a.to_dataset
#a4 = a.to_dataframe
#a5 = a.to_series
#a6 = a.to_pandas
# to convert to pandas series to tr.transmission.to_series()

plottrans(tr, c1)
#
show()

"""
Here's a thin non-absorbing layer, on top of a thick absorbing layer, with
air on both sides. Plotting reflected intensity versus wavenumber, at two
different incident angles.
"""
degree = pi/180

# list of layer thicknesses in nm
d_list = [inf, 100, 300, inf]
# list of refractive indices
n_list = [1, 2.2, 3.3+0.3j, 1]
# list of wavenumbers to plot in nm^-1
ks = linspace(0.0001, .01, num=400)
# initialize lists of y-values to plot
rnorm = []
r45 = []
for k in ks:
		# For normal incidence, s and p polarizations are identical.
		# I arbitrarily decided to use 's'.
    rnorm.append(tmm.coh_tmm('s', n_list, d_list, 0, 1/k)['R'])
    r45.append(tmm.unpolarized_RT(n_list, d_list, 45*degree, 1/k)['R'])
kcm = ks * 1e7 #ks in cm^-1 rather than nm^-1
plt.figure()
plt.plot(kcm, rnorm, 'blue', kcm, r45, 'purple')
plt.xlabel('k (cm$^{-1}$)')
plt.ylabel('Fraction reflected')
plt.title('Reflection of unpolarized light at 0$^\circ$ incidence (blue), '
          '45$^\circ$ (purple)')
plt.show()

#if __name__ == '__main__':
#    main()


#!/usr/bin/python
# -*- coding: latin-1 -*-
"""
Materials define the surface properties and specify how light rays get coloured and reflected. All materials are callables - they implement the __call__ method
and can be used like functions.

.. moduleauthor:: Adrian Köring
"""

import numpy as np

from padvinder.util import normalize
from padvinder.util import check_finite

class Material(object):
    """
    An emission material consists of an emitted colour only.
    Without gradients, lighting or anything.

    Parameters
    ----------
    color : numpy.ndarray_like
        of three dimensions and contains colors as (Red, Green, Blue) where
        (0,0,0) is black and (1,1,1) is white

    Raises
    ------
    ValueError
        if the color contains any non-finite (inf, nan) values

    Examples
    --------
    >>> Material((0.8, 0.8, 0.8))
    Material(color=[.8 .8 .8])
    """
    def __init__(self, color = (.5, .5, .5)):
        check_finite(color)
        self._color = np.array(color).astype(np.float64)

    @property
    def color(self):
        """
        Returns the color of the material.
        """
        return self._color

    def __call__(self, surface_normal,
                       incoming_color,
                       incoming_direction,
                       outgoing_direction):
        """
        Calculate light reflected from the material toward the outgoing
        direction. Keep in mind, while pathtracing starts at the camera and
        heads into the scene, the rays contribution is accumulated 'backwards'.
        Therefore the incoming direction is further down the path and
        outgoing_direction is closer towards the camera.

        Parameters
        ----------
        surface_normal : numpy.ndarray_like
            normal vector at the geometries surface
        incoming_color : numpy.ndarray_like
            the color the ray has accumulated up to this point
        incoming_direction : numpy.ndarray_like
            the direction from where the 'light shines' onto the surface
        outgoing_direction : numpy.ndarray_like
            the direction into which the 'light gets reflected' from the surface

        Returns
        -------
        color : numpy.ndarray_like
            the light color 'getting reflected' from the surface
        """
        return self._color

    def outgoing_direction(self, normal, incoming_direction):
        """
        Given a surface normal and an incoming direction, determine the
        direction in which the path continues.

        normal : numpy.ndarray_like of shape (3, )
            the surface normal at the intersection point

        incoming_direction : numpy.ndarray_like of shape (3, )
            the direction from which light hits the surface

        Returns
        -------
        outgoing direction : numpy.ndarray_like of shape (3, 0)
            the direction in which light is reflected from the surface
        """
        # BaseClass randomly (not uniformly) samples the hemisphere
        point_on_sphere = normalize(np.random.uniform(0, 1, size=(3, )))
        # ensure it is in the same hemisphere as the normal
        if np.dot(normal, point_on_sphere) < 0:
            point_on_sphere = -point_on_sphere
        return normalize(normal + point_on_sphere)

    def __repr__(self):
        return "Material(color={})".format(self._color)

class Emission(Material):
    """
    Emission is equivalent to the abstract base class Material. Due to semantics
    this class exists and merely inherits without modifications.

    Parameters
    ----------
    color : numpy.ndarray_like
        of three dimensions and contains colors as (Red, Green, Blue) where
        (0,0,0) is black and (1,1,1) is white

    Raises
    ------
    ValueError
        if the color contains any non-finite (inf, nan) values

    Examples
    --------
    >>> Emission()
    Emission(color=[10.0, 10.0, 10.0])
    """
    def __init__(self, color = (10, 10, 10)):
        super().__init__(color)

    def __repr__(self):
        return "Emission(color={})".format(self._color)

class Lambert(Material):
    def __init__(self, color = (0.5, 0.5, 0.5), diffuse = 1):
        """
        A lambert material consists of a colour value and a diffuse coefficient.

        Parameters
        ----------
        color : numpy.ndarray_like
            of three dimensions and contains colors as (Red, Green, Blue) where
            (0,0,0) is black and (1,1,1) is white
        diffuse : number in [0, 1]
            percentage of incoming light that is reflected again

        Raises
        ------
        ValueError
            if the color contains any non-finite (inf, nan) values

        Examples
        --------
        >>> Lambert((0.8, 0.8, 0.8), 1)
        Lambert(color=[0.8, 0.8, 0.8], diffuse=1)
        """
        super().__init__(color)
        self._diffuse = diffuse

    @property
    def diffuse(self):
        """
        Returns the diffuse value of the material.
        """
        return self._diffuse

    # def __call__(self, surface_normal,
    #                    incoming_light,
    #                    incoming_direction,
    #                    outgoing_direction):
    #     """
    #     Calculate light reflected from the material toward the outgoing
    #     direction. Keep in mind, while pathtracing starts at the camera and
    #     heads into the scene, the rays contribution is accumulated 'backwards'.
    #     Therefore the incoming direction is further down the path and
    #     outgoing_direction is closer towards the camera.
    #
    #     Parameters
    #     ----------
    #     surface_normal : numpy.ndarray_like
    #         normal vector at the geometries surface
    #     incoming_color : numpy.ndarray_like
    #         the color the ray has accumulated up to this point
    #     incoming_direction : numpy.ndarray_like
    #         the direction from where the 'light shines' onto the surface
    #     outgoing_direction : numpy.ndarray_like
    #         the direction into which the 'light gets reflected' from the surface
    #
    #     Returns
    #     -------
    #     color : numpy.ndarray_like
    #         the light color 'getting reflected' from the surface
    #     """
    #     raise NotImplemented()

    def __repr__(self):
        c, d = self._color, self._diffuse
        return "Lambert(color={0}, diffuse={1})".format(c, d)