from collections import defaultdict
import logging
import os
from typing import Mapping, Tuple

import numpy as np
import pyro
from pyro.infer import SVI, TraceGraph_ELBO, Trace_ELBO
from pyro.nn import pyro_method
from pyro.optim import ExponentialLR, AdagradRMSProp  # noqa: F401
from pyro.distributions.torch_transform import ComposeTransformModule
from pyro.distributions.transforms import (
    ComposeTransform, AffineTransform, ExpTransform, Spline, Permute
)
from pyro.distributions.transforms import batchnorm, iterated
from pyro.distributions import (
    LowRankMultivariateNormal, MultivariateNormal, Normal, Laplace, TransformedDistribution  # noqa: F401
)
import torch
from torch.distributions import Independent
from torch.optim import AdamW

from counterfactualms.arch.medical import Decoder, Encoder
from counterfactualms.arch.nvae import Decoder as NDecoder
from counterfactualms.arch.nvae import Encoder as NEncoder
from counterfactualms.arch.thirdparty.neural_operations import Swish
from counterfactualms.distributions.transforms.reshape import ReshapeTransform
from counterfactualms.distributions.transforms.affine import LowerCholeskyAffine
from counterfactualms.utils.optim import OneCycleLR
from counterfactualms.utils.pyro_modifications import affine_autoregressive, affine_coupling
from counterfactualms.utils.pyro_modifications import spline_autoregressive, spline_coupling
from counterfactualms.distributions.deep import (
    DeepMultivariateNormal, DeepIndepNormal, DeepIndepMixtureNormal, Conv2dIndepNormal, DeepLowRankMultivariateNormal
)
from counterfactualms.experiments.calabresi.base_experiment import (
    BaseCovariateExperiment, BaseSEM, EXPERIMENT_REGISTRY  # noqa: F401
)

logger = logging.getLogger(__name__)

class StorageTraceGraph_ELBO(TraceGraph_ELBO):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.trace_storage = {'model': None, 'guide': None}

    def _get_trace(self, model, guide, args, kwargs):
        model_trace, guide_trace = super()._get_trace(model, guide, args, kwargs)
        self.trace_storage['model'] = model_trace
        self.trace_storage['guide'] = guide_trace
        return model_trace, guide_trace

class StorageTrace_ELBO(Trace_ELBO):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.trace_storage = {'model': None, 'guide': None}

    def _get_trace(self, model, guide, args, kwargs):
        model_trace, guide_trace = super()._get_trace(model, guide, args, kwargs)
        self.trace_storage['model'] = model_trace
        self.trace_storage['guide'] = guide_trace
        return model_trace, guide_trace

class Lambda(torch.nn.Module):
    def __init__(self, func):
        super().__init__()
        self.func = func

    def forward(self, x):
        return self.func(x)

class BaseVISEM(BaseSEM):
    context_dim = 0  # number of context dimensions for decoder

    def __init__(self, latent_dim:int, prior_components:int=1, posterior_components:int=1,
                 logstd_init:float=-5, enc_filters:Tuple[int]=(16,32,64,128),
                 dec_filters:Tuple[int]=(128,64,32,16), num_convolutions:int=3, use_upconv:bool=False,
                 decoder_type:str='fixed_var', decoder_cov_rank:int=10, img_shape:Tuple[int]=(128,128),
                 use_nvae=False, use_weight_norm=False, use_spectral_norm=False, laplace_likelihood=False,
                 eps=0.1, n_prior_flows=3, n_posterior_flows=3, use_autoregressive=False, use_swish=False,
                 use_spline=False, use_stable=False, pseudo3d=False, head_filters=(16,16), **kwargs):
        super().__init__(**kwargs)
        self.encoder_shape = ((3,) if pseudo3d else (1,)) + tuple(img_shape)
        self.decoder_shape = (head_filters[0],) + tuple(img_shape)
        self.img_shape = (1,) + tuple(img_shape)
        self.latent_dim = latent_dim
        self.prior_components = prior_components
        self.posterior_components = posterior_components
        self.logstd_init = logstd_init
        self.enc_filters = enc_filters
        self.dec_filters = dec_filters
        self.head_filters = head_filters
        self.num_convolutions = num_convolutions
        self.use_upconv = use_upconv
        self.decoder_type = decoder_type
        self.decoder_cov_rank = decoder_cov_rank
        self.use_nvae = use_nvae
        self.use_weight_norm = use_weight_norm
        self.use_spectral_norm = use_spectral_norm
        self.laplace_likelihood = laplace_likelihood
        self.eps = eps
        self.n_prior_flows = n_prior_flows
        self.n_posterior_flows = n_posterior_flows
        self.use_autoregressive = use_autoregressive
        self.use_spline = use_spline
        self.use_swish = use_swish
        self.use_stable = use_stable
        self.pseudo3d = pseudo3d
        self.annealing_factor = [1.]  # initialize here; will be changed during training
        self.n_levels = 0

        # decoder parts
        if use_nvae:
            decoder = NDecoder(
                num_convolutions=self.num_convolutions, filters=self.dec_filters,
                latent_dim=self.latent_dim + self.context_dim,
                output_size=self.decoder_shape
            )
        else:
            decoder = Decoder(
                num_convolutions=self.num_convolutions, filters=self.dec_filters,
                latent_dim=self.latent_dim + self.context_dim, upconv=self.use_upconv,
                output_size=self.decoder_shape,
                use_weight_norm=self.use_weight_norm,
                use_spectral_norm=self.use_spectral_norm,
            )

        self._create_decoder(decoder)

        # encoder parts
        if self.use_nvae:
            self.encoder = NEncoder(
                num_convolutions=self.num_convolutions,
                filters=self.enc_filters,
                latent_dim=self.latent_dim,
                input_size=self.encoder_shape
            )
        else:
            self.encoder = Encoder(
                num_convolutions=self.num_convolutions,
                filters=self.enc_filters,
                latent_dim=self.latent_dim,
                input_size=self.encoder_shape,
                use_weight_norm=self.use_weight_norm,
                use_spectral_norm = self.use_spectral_norm
            )

        nonlinearity = Swish() if self.use_swish else torch.nn.LeakyReLU(0.1)
        latent_layers = torch.nn.Sequential(
            torch.nn.Linear(self.latent_dim + self.context_dim, self.latent_dim),
            nonlinearity
        )

        if self.posterior_components > 1:
            self.latent_encoder = DeepIndepMixtureNormal(
                latent_layers, self.latent_dim, self.latent_dim, self.posterior_components)
        else:
            self.latent_encoder = DeepIndepNormal(latent_layers, self.latent_dim, self.latent_dim)

        if self.prior_components > 1:
            self.z_loc = torch.nn.Parameter(torch.randn([self.prior_components, self.latent_dim]))
            self.z_scale = torch.nn.Parameter(torch.randn([self.latent_dim]).clamp(min=-1.,max=None))  # log scale
            self.register_buffer('z_components',  # don't be bayesian about the mixture components
                ((1/self.prior_components)*torch.ones([self.prior_components], requires_grad=False)).log())
        else:
            self.register_buffer('z_loc', torch.zeros([latent_dim, ], requires_grad=False))
            self.register_buffer('z_scale', torch.ones([latent_dim, ], requires_grad=False))
            self.z_components = None

        # priors
        self.sex_logits = torch.nn.Parameter(torch.zeros([1, ]))
        self.register_buffer('slice_number_min', torch.zeros([1, ], requires_grad=False))
        self.register_buffer('slice_number_max', 241.*torch.ones([1, ], requires_grad=False)+1.)

        for k in self.required_data - {'sex', 'x', 'slice_number'}:
            self.register_buffer(f'{k}_base_loc', torch.zeros([1, ], requires_grad=False))
            self.register_buffer(f'{k}_base_scale', torch.ones([1, ], requires_grad=False))

        self.register_buffer('x_base_loc', torch.zeros(self.img_shape, requires_grad=False))
        self.register_buffer('x_base_scale', torch.ones(self.img_shape, requires_grad=False))

        for k in self.required_data - {'sex', 'x', 'slice_number'}:
            self.register_buffer(f'{k}_flow_lognorm_loc', torch.zeros([], requires_grad=False))
            self.register_buffer(f'{k}_flow_lognorm_scale', torch.ones([], requires_grad=False))

        perm = lambda: torch.randperm(self.latent_dim, dtype=torch.long, requires_grad=False)

        self.use_prior_flow = self.n_prior_flows > 0
        self.use_prior_permutations = self.n_prior_flows > 1
        if self.use_prior_permutations:
            for i in range(self.n_prior_flows):
                self.register_buffer(f'prior_flow_permutation_{i}', perm())

        self.use_posterior_flow = self.n_posterior_flows > 0
        self.use_posterior_permutations = self.n_posterior_flows > 1
        if self.use_posterior_permutations:
            for i in range(self.n_posterior_flows):
                self.register_buffer(f'posterior_flow_permutation_{i}', perm())

        # age flow
        self.age_flow_components = ComposeTransformModule([Spline(1)])
        self.age_flow_lognorm = AffineTransform(loc=self.age_flow_lognorm_loc.item(), scale=self.age_flow_lognorm_scale.item())
        self.age_flow_constraint_transforms = ComposeTransform([self.age_flow_lognorm, ExpTransform()])
        self.age_flow_transforms = ComposeTransform([self.age_flow_components, self.age_flow_constraint_transforms])

        # other flows shared components
        self.ventricle_volume_flow_lognorm = AffineTransform(loc=self.ventricle_volume_flow_lognorm_loc.item(), scale=self.ventricle_volume_flow_lognorm_scale.item())  # noqa: E501
        self.ventricle_volume_flow_constraint_transforms = ComposeTransform([self.ventricle_volume_flow_lognorm, ExpTransform()])

        self.brain_volume_flow_lognorm = AffineTransform(loc=self.brain_volume_flow_lognorm_loc.item(), scale=self.brain_volume_flow_lognorm_scale.item())
        self.brain_volume_flow_constraint_transforms = ComposeTransform([self.brain_volume_flow_lognorm, ExpTransform()])

        self.lesion_volume_flow_lognorm = AffineTransform(loc=self.lesion_volume_flow_lognorm_loc.item(), scale=self.lesion_volume_flow_lognorm_scale.item())
        self.lesion_volume_flow_eps = AffineTransform(loc=-eps, scale=1.)
        self.lesion_volume_flow_constraint_transforms = ComposeTransform([self.lesion_volume_flow_lognorm, ExpTransform(), self.lesion_volume_flow_eps])

        self.duration_flow_lognorm = AffineTransform(loc=self.duration_flow_lognorm_loc.item(), scale=self.duration_flow_lognorm_scale.item())
        self.duration_flow_eps = AffineTransform(loc=-eps, scale=1.)
        self.duration_flow_constraint_transforms = ComposeTransform([self.duration_flow_lognorm, ExpTransform(), self.duration_flow_eps])

        self.edss_flow_lognorm = AffineTransform(loc=self.edss_flow_lognorm_loc.item(), scale=self.edss_flow_lognorm_scale.item())
        self.edss_flow_eps = AffineTransform(loc=-eps, scale=1.)
        self.edss_flow_constraint_transforms = ComposeTransform([self.edss_flow_lognorm, ExpTransform(), self.edss_flow_eps])

        hidden_dims = (3 * self.latent_dim + 1,) if self.use_autoregressive else (2*self.latent_dim, 2*self.latent_dim)
        flow_kwargs = dict(hidden_dims=hidden_dims, nonlinearity=nonlinearity)
        if self.use_spline:
            flow_ = spline_autoregressive if self.use_autoregressive else spline_coupling
        else:
            flow_ = affine_autoregressive if self.use_autoregressive else affine_coupling
        if self.use_autoregressive:
            flow_kwargs['stable'] = self.use_stable

        if self.use_prior_permutations:
            self.prior_affine = iterated(self.n_prior_flows, batchnorm, self.latent_dim, momentum=0.05) if self.use_prior_flow else []
            self.prior_permutations = [Permute(getattr(self, f'prior_flow_permutation_{i}')) for i in range(self.n_prior_flows)]
            self.prior_flow_components = iterated(self.n_prior_flows, flow_, self.latent_dim, **flow_kwargs) if self.use_prior_flow else []
            self.prior_flow_transforms = [
                x for c in zip(self.prior_permutations, self.prior_affine, self.prior_flow_components) for x in c
            ]
        else:
            self.prior_affine = []
            self.prior_permutations = []
            self.prior_flow_components = flow_(self.latent_dim, **flow_kwargs) if self.use_prior_flow else []
            self.prior_flow_transforms = [self.prior_flow_components]

        if self.use_posterior_permutations:
            self.posterior_affine = iterated(self.n_posterior_flows, batchnorm, self.latent_dim, momentum=0.05)
            self.posterior_permutations = [Permute(getattr(self, f'posterior_flow_permutation_{i}')) for i in range(self.n_posterior_flows)]
            self.posterior_flow_components = iterated(self.n_posterior_flows, flow_, self.latent_dim, **flow_kwargs)
            self.posterior_flow_transforms = [
                x for c in zip(self.posterior_permutations, self.posterior_affine, self.posterior_flow_components) for x in c
            ]
        else:
            self.posterior_affine = []
            self.posterior_permutations = []
            self.posterior_flow_components = flow_(self.latent_dim, **flow_kwargs) if self.use_posterior_flow else []
            self.posterior_flow_transforms = [self.posterior_flow_components]

    def _create_decoder(self, decoder):
        co = 3 if self.pseudo3d else 1
        if self.decoder_type == 'fixed_var':
            self.decoder = Conv2dIndepNormal(decoder, self.head_filters, co,
                use_weight_norm=self.use_weight_norm, use_spectral_norm=self.use_spectral_norm)
            torch.nn.init.zeros_(self.decoder.logstd_head[-1].weight)
            self.decoder.logstd_head[-1].weight.requires_grad = False
            torch.nn.init.constant_(self.decoder.logstd_head[-1].bias, self.logstd_init)
            self.decoder.logstd_head[-1].bias.requires_grad = False

        elif self.decoder_type == 'learned_var':
            self.decoder = Conv2dIndepNormal(decoder, self.head_filters, co,
                use_weight_norm=self.use_weight_norm, use_spectral_norm=self.use_spectral_norm,
                logstd_ref=self.logstd_init)
            torch.nn.init.zeros_(self.decoder.logstd_head[-1].weight)
            self.decoder.logstd_head[-1].weight.requires_grad = False
            torch.nn.init.constant_(self.decoder.logstd_head[-1].bias, self.logstd_init)
            self.decoder.logstd_head[-1].bias.requires_grad = True

        elif self.decoder_type == 'independent_var':
            self.decoder = Conv2dIndepNormal(decoder, self.head_filters, co,
                use_weight_norm=self.use_weight_norm, use_spectral_norm=self.use_spectral_norm,
                logstd_ref=self.logstd_init)
            torch.nn.init.zeros_(self.decoder.logstd_head[-1].weight)
            self.decoder.logstd_head[-1].weight.requires_grad = True
            torch.nn.init.normal_(self.decoder.logstd_head[-1].bias, 0., 1e-1)
            self.decoder.logstd_head[-1].bias.requires_grad = True

        elif self.decoder_type == 'multivariate_gaussian':
            seq = torch.nn.Sequential(decoder, Lambda(lambda x: x.view(x.shape[0], -1)))
            self.decoder = DeepMultivariateNormal(seq, np.prod(self.decoder_shape), np.prod(self.img_shape))

        elif self.decoder_type == 'sharedvar_multivariate_gaussian':
            seq = torch.nn.Sequential(decoder, Lambda(lambda x: x.view(x.shape[0], -1)))
            self.decoder = DeepMultivariateNormal(seq, np.prod(self.decoder_shape), np.prod(self.img_shape))
            torch.nn.init.zeros_(self.decoder.logdiag_head.weight)
            self.decoder.logdiag_head.weight.requires_grad = False
            torch.nn.init.zeros_(self.decoder.lower_head.weight)
            self.decoder.lower_head.weight.requires_grad = False
            torch.nn.init.normal_(self.decoder.logdiag_head.bias, self.logstd_init, 1e-1)
            self.decoder.logdiag_head.bias.requires_grad = True

        elif self.decoder_type == 'lowrank_multivariate_gaussian':
            seq = torch.nn.Sequential(decoder, Lambda(lambda x: x.view(x.shape[0], -1)))
            self.decoder = DeepLowRankMultivariateNormal(
                seq, np.prod(self.decoder_shape), np.prod(self.img_shape), self.decoder_cov_rank
            )

        elif self.decoder_type == 'sharedvar_lowrank_multivariate_gaussian':
            seq = torch.nn.Sequential(decoder, Lambda(lambda x: x.view(x.shape[0], -1)))
            self.decoder = DeepLowRankMultivariateNormal(
                seq, np.prod(self.decoder_shape), np.prod(self.img_shape), self.decoder_cov_rank
            )
            torch.nn.init.zeros_(self.decoder.logdiag_head.weight)
            self.decoder.logdiag_head.weight.requires_grad = False
            torch.nn.init.zeros_(self.decoder.factor_head.weight)
            self.decoder.factor_head.weight.requires_grad = False
            torch.nn.init.normal_(self.decoder.logdiag_head.bias, self.logstd_init, 1e-1)
            self.decoder.logdiag_head.bias.requires_grad = True

        else:
            raise ValueError(f'unknown decoder type {self.decoder_type}.')

    def __setattr__(self, name, value):
        super().__setattr__(name, value)
        if 'flow_lognorm_loc' in name:
            name_ = name.replace('flow_lognorm_loc', '')
            getattr(self, name_ + 'flow_lognorm').loc = value.item()
        elif 'flow_lognorm_scale' in name:
            name_ = name.replace('flow_lognorm_scale', '')
            getattr(self, name_ + 'flow_lognorm').scale = value.item()
        elif 'flow_norm_loc' in name:
            name_ = name.replace('flow_norm_loc', '')
            getattr(self, name_ + 'flow_norm').loc = value.item()
        elif 'flow_norm_scale' in name:
            name_ = name.replace('flow_norm_scale', '')
            getattr(self, name_ + 'flow_norm').scale = value.item()
        elif 'prior_flow_permutation' in name:
            i = int(name[-1])
            self.prior_permutations[i].permutation = value
        elif 'posterior_flow_permutation' in name:
            i = int(name[-1])
            self.posterior_permutations[i].permutation = value

    def _get_preprocess_transforms(self):
        return super()._get_preprocess_transforms().inv

    def _get_transformed_x_dist(self, latent, ctx=None):
        x_pred_dist = self.decoder.predict(latent, ctx)  # returns a normal dist with mean of the predicted image
        if self.laplace_likelihood:
            x_base_dist = Laplace(self.x_base_loc, self.x_base_scale).to_event(3)
        else:
            x_base_dist = Normal(self.x_base_loc, self.x_base_scale).to_event(3)  # 3 dimensions starting from right dep.

        preprocess_transform = self._get_preprocess_transforms()

        if isinstance(x_pred_dist, MultivariateNormal) or isinstance(x_pred_dist, LowRankMultivariateNormal):
            chol_transform = LowerCholeskyAffine(x_pred_dist.loc, x_pred_dist.scale_tril)
            reshape_transform = ReshapeTransform(self.img_shape, (np.prod(self.img_shape), ))
            x_reparam_transform = ComposeTransform([reshape_transform, chol_transform, reshape_transform.inv])
        elif isinstance(x_pred_dist, Independent):
            x_pred_dist = x_pred_dist.base_dist
            x_reparam_transform = AffineTransform(x_pred_dist.loc, x_pred_dist.scale, 3)
        else:
            raise ValueError(f'{x_pred_dist} not valid.')

        return TransformedDistribution(x_base_dist, ComposeTransform([x_reparam_transform, preprocess_transform]))

    @pyro_method
    def guide(self, obs):
        raise NotImplementedError()

    @pyro_method
    def svi_guide(self, obs):
        self._check_observation(obs)
        self.guide(obs)

    @pyro_method
    def svi_model(self, obs):
        self._check_observation(obs)
        batch_size = obs['x'].shape[0]
        with pyro.plate('observations', batch_size):
            pyro.condition(self.model, data=obs)()

    @pyro_method
    def infer_z(self, *args, **kwargs):
        return self.guide(*args, **kwargs)

    @property
    def required_data(self):
        return {'x', 'sex', 'age', 'ventricle_volume', 'brain_volume', 'lesion_volume',
                'edss', 'duration', 'slice_number'}

    def _check_observation(self, obs):
        keys = obs.keys()
        assert self.required_data == set(keys), f'Incompatible observation: {tuple(keys)}'

    @pyro_method
    def infer(self, obs):
        self._check_observation(obs)
        obs_ = obs.copy()
        z = self.infer_z(obs_)
        obs_.update(dict(z=z))
        exogenous = self.infer_exogenous(obs_)
        exogenous['z'] = z
        return exogenous

    @pyro_method
    def reconstruct(self, obs, num_particles:int=1):
        self._check_observation(obs)
        z_dist = pyro.poutine.trace(self.guide).get_trace(obs).nodes['z']['fn']
        batch_size = obs['x'].shape[0]
        obs_ = {k: v for k, v in obs.items() if k != 'x'}
        recons = []
        for _ in range(num_particles):
            z = pyro.sample('z', z_dist)
            obs_.update({'z': z})
            recon = pyro.poutine.condition(
                self.sample, data=obs_)(batch_size)
            recons += [recon['x']]
        return torch.stack(recons).mean(0)

    def _cf_dict(self, counterfactuals):
        out = {k: [] for k in self.required_data}
        for cf in counterfactuals:
            for k in self.required_data:
                out[k].append(cf[k])
        out = {k: torch.stack(v).mean(0) for k, v in out.items()}
        return out

    @pyro_method
    def counterfactual(self, obs, condition:Mapping=None, num_particles:int=1):
        self._check_observation(obs)
        obs_ = obs.copy()
        z_dist = pyro.poutine.trace(self.guide).get_trace(obs_).nodes['z']['fn']  # variational posterior
        n = obs_['x'].shape[0]

        counterfactuals = []
        for _ in range(num_particles):
            z = pyro.sample('z', z_dist)
            obs_.update(dict(z=z))
            exogenous = self.infer_exogenous(obs_)
            exogenous['z'] = z
            # condition on these vars if they aren't included in 'do' as they are root nodes
            # and we don't have the exogenous noise for them yet
            if 'sex' not in condition.keys():
                exogenous['sex'] = obs_['sex']
            if 'slice_number' not in condition.keys():
                exogenous['slice_number'] = obs_['slice_number']

            cf = pyro.poutine.do(pyro.poutine.condition(self.sample_scm, data=exogenous), data=condition)(n)
            counterfactuals.append(cf)

        return self._cf_dict(counterfactuals)

    @classmethod
    def add_arguments(cls, parser):
        parser = super().add_arguments(parser)
        parser.add_argument('--latent-dim', default=100, type=int, help="latent dimension of model (default: %(default)s)")
        parser.add_argument('--prior-components', default=1, type=int, help="number of mixture components for prior (default: %(default)s)")
        parser.add_argument('--posterior-components', default=1, type=int, help="number of mixture components for posterior (default: %(default)s)")
        parser.add_argument('--logstd-init', default=-5, type=float, help="init/ref of logstd for fixed/learned (default: %(default)s)")
        parser.add_argument('--enc-filters', default=[16,32,64,128,256], nargs='+', type=int, help="number of filters in each layer of encoder (default: %(default)s)")
        parser.add_argument('--dec-filters', default=[256,128,64,32,16], nargs='+', type=int, help="number of filters in each layer of decoder (default: %(default)s)")
        parser.add_argument('--head-filters', default=[16], nargs='+', type=int, help="number of filters in each (mean/logstd) head (default: %(default)s)")
        parser.add_argument('--num-convolutions', default=3, type=int, help="number of convolutions in each layer (default: %(default)s)")
        parser.add_argument('--use-upconv', default=False, action='store_true', help="use upsample->conv instead of transpose conv (default: %(default)s)")
        parser.add_argument('--use-nvae', default=False, action='store_true', help="use nvae instead of standard vae (default: %(default)s)")
        parser.add_argument('--use-weight-norm', default=False, action='store_true', help="use weight norm in conv layers (not w/ nvae) (default: %(default)s)")
        parser.add_argument('--use-spectral-norm', default=False, action='store_true', help="use spectral norm in conv layers (not w/ nvae) (default: %(default)s)")
        parser.add_argument('--hierarchical-layers', default=(1,3,5), type=int, nargs='+', help="which filter layers are passed in hierarchical model (default: %(default)s)")
        parser.add_argument('--hierarchical-div', default=16, type=int, help="div factor in hierarchical model (default: %(default)s)")
        parser.add_argument('--temperature', default=2./3., type=float, help="temperature for ST Gumbel-softmax layers (default: %(default)s)")
        parser.add_argument('--laplace-likelihood', default=False, action='store_true', help="use laplace likelihood for image (default: %(default)s)")
        parser.add_argument('--n-prior-flows', default=3, type=int, help="use this number of flows for prior in flow net (default: %(default)s)")
        parser.add_argument('--n-posterior-flows', default=3, type=int, help="use this number of flows for posterior in flow net (default: %(default)s)")
        parser.add_argument('--use-autoregressive', default=False, action='store_true', help="use autoregressive flow for prior/post instead of coupling (default: %(default)s)")
        parser.add_argument('--use-spline', default=False, action='store_true', help="use spline flow for prior/post instead of affine (default: %(default)s)")
        parser.add_argument('--use-stable', default=False, action='store_true', help="use stable version of affine for prior/post instead (default: %(default)s)")
        parser.add_argument('--use-swish', default=False, action='store_true', help="use swish in flows for nonlinearity (default: %(default)s)")
        parser.add_argument('--pseudo3d', default=False, action='store_true', help="use pseudo3d images (default: %(default)s)")
        parser.add_argument(
            '--decoder-type', default='fixed_var', help="var type (default: %(default)s)",
            choices=['fixed_var', 'learned_var', 'independent_var', 'sharedvar_multivariate_gaussian',
                     'multivariate_gaussian', 'sharedvar_lowrank_multivariate_gaussian', 'lowrank_multivariate_gaussian'])
        parser.add_argument('--decoder-cov-rank', default=10, type=int, help="rank for lowrank cov approximation (requires lowrank decoder) (default: %(default)s)")  # noqa: E501
        return parser

class SVIExperiment(BaseCovariateExperiment):
    def __init__(self, hparams, pyro_model: BaseSEM):
        super().__init__(hparams, pyro_model)
        if hparams.tracegraph_elbo:
            self.svi_loss = StorageTraceGraph_ELBO(num_particles=hparams.num_svi_particles)
        else:
            self.svi_loss = StorageTrace_ELBO(num_particles=hparams.num_svi_particles)
        self._build_svi()

    def _build_svi(self, loss=None):
        def per_param_callable(module_name, param_name):
            if self.hparams.use_adagrad_rmsprop:
                params = {'eta': self.hparams.eta, 'delta': self.hparams.delta, 't': self.hparams.t}
            else:
                params = {'weight_decay': self.hparams.weight_decay,
                          'betas': self.hparams.betas, 'eps': 1e-5}
                if any([(pn in module_name) for pn in ('prior_flow', 'posterior_flow')]):
                    params['lr'] = self.hparams.lr
                elif 'affine' in module_name:
                    params['lr'] = self.hparams.lr
                    params['weight_decay'] = 0.
                elif 'flow_components' in module_name:
                    params['lr'] = self.hparams.pgm_lr
                elif 'sex_logits' in param_name:
                    params['lr'] = self.hparams.pgm_lr
                    params['weight_decay'] = 0.
                elif 'decoder' in module_name and 'logstd_head' in param_name:
                    params['weight_decay'] = self.hparams.logstd_weight_decay
                else:
                    params['lr'] = self.hparams.lr
                logger.info(f'building opt for {module_name} - {param_name} with p: {params}')
            return params

        def per_param_clip_args(module_name, param_name):
            clip_args = defaultdict(lambda: None)
            if any([(pn in module_name) for pn in ('prior_flow', 'posterior_flow')]):
                clip_args['clip_norm'] = self.hparams.flow_clip_norm
            elif any([(pn in param_name) for pn in ('affine', 'sex_logits', 'flow_components')]):
                clip_args['clip_norm'] = self.hparams.pgm_clip_norm
            else:
                clip_args['clip_norm'] = self.hparams.clip_norm
            logger.info(f'building clip args for {module_name} - {param_name} with p: {clip_args}')
            return clip_args

        if loss is None:
            loss = self.svi_loss

        optimizer = AdagradRMSProp if self.hparams.use_adagrad_rmsprop else AdamW
        verbose = self.hparams.verbosity > 1  # only print lr in debug mode
        if self.hparams.use_exponential_lr:
            self.scheduler = ExponentialLR({'optimizer': optimizer, 'optim_args': per_param_callable,
                                            'gamma': self.hparams.lrd, 'verbose': verbose},
                                            clip_args=per_param_clip_args)
        else:
            self.scheduler = OneCycleLR({'optimizer': optimizer, 'optim_args': per_param_callable,
                                         'epochs': self.hparams.n_epochs, 'steps_per_epoch': self._steps_per_epoch(),
                                         'pct_start': self.hparams.pct_start, 'div_factor': self.hparams.div_factor,
                                         'final_div_factor': self.hparams.final_div_factor, 'verbose': verbose},
                                         clip_args=per_param_clip_args)
        if self.hparams.use_cf_guide:
            def guide(*args, **kwargs):
                return self.pyro_model.counterfactual_guide(*args, **kwargs, counterfactual_type=self.hparams.cf_elbo_type)
            self.svi = SVI(self.pyro_model.svi_model, guide, self.scheduler, loss)
        else:
            self.svi = SVI(self.pyro_model.svi_model, self.pyro_model.svi_guide, self.scheduler, loss)
        self.svi.loss_class = loss

    def backward(self, *args, **kwargs):
        pass  # No loss to backpropagate since we're using Pyro's optimisation machinery

    def print_trace_updates(self, batch):
        with torch.no_grad():
            logger.info('Traces:\n' + ('#' * 10))

            guide_trace = pyro.poutine.trace(self.pyro_model.svi_guide).get_trace(batch)
            model_trace = pyro.poutine.trace(pyro.poutine.replay(self.pyro_model.svi_model, trace=guide_trace)).get_trace(batch)

            guide_trace = pyro.poutine.util.prune_subsample_sites(guide_trace)
            model_trace = pyro.poutine.util.prune_subsample_sites(model_trace)

            model_trace.compute_log_prob()
            guide_trace.compute_score_parts()

            logging.info(f'model: {model_trace.nodes.keys()}')
            for name, site in model_trace.nodes.items():
                if site["type"] == "sample":
                    fn = site['fn']
                    if isinstance(fn, Independent):
                        fn = fn.base_dist
                    try:
                        logging.info(f'{name}: {fn} - {fn.support}')
                    except NotImplementedError:
                        logging.info(f'{name}: {fn}')
                    log_prob_sum = site["log_prob_sum"]
                    is_obs = site["is_observed"]
                    logging.info(f'model - log p({name}) = {log_prob_sum} | obs={is_obs}')
                    if torch.isnan(log_prob_sum):
                        value = site['value'][0]
                        conc0 = fn.concentration0
                        conc1 = fn.concentration1
                        raise RuntimeError(f'Error: \n{value}\n{conc0}\n{conc1}')

            logging.info(f'guide: {guide_trace.nodes.keys()}')

            for name, site in guide_trace.nodes.items():
                if site["type"] == "sample":
                    fn = site['fn']
                    if isinstance(fn, Independent):
                        fn = fn.base_dist
                    try:
                        logging.info(f'{name}: {fn} - {fn.support}')
                    except NotImplementedError:
                        logging.info(f'{name}: {fn}')
                    entropy = site["score_parts"].entropy_term.sum()
                    is_obs = site["is_observed"]
                    logging.info(f'guide - log q({name}) = {entropy} | obs={is_obs}')

    def get_trace_metrics(self, batch):
        metrics = {}
        model = self.svi.loss_class.trace_storage['model']
        guide = self.svi.loss_class.trace_storage['guide']
        for k in self.required_data:
            metrics[f'log p({k})'] = model.nodes[k]['log_prob'].mean()
        if self.pyro_model.n_levels > 0:
            metrics['log p(z) - log q(z)'] = 0.
            for i in range(self.pyro_model.n_levels):
                metrics[f'log p(z{i})'] = model.nodes[f'z{i}']['log_prob'].mean()
                metrics[f'log q(z{i})'] = guide.nodes[f'z{i}']['log_prob'].mean()
                metrics['log p(z) - log q(z)'] += metrics[f'log p(z{i})'] - metrics[f'log q(z{i})']
        else:
            metrics['log p(z)'] = model.nodes['z']['log_prob'].mean()
            metrics['log q(z)'] = guide.nodes['z']['log_prob'].mean()
            metrics['log p(z) - log q(z)'] = metrics['log p(z)'] - metrics['log q(z)']
        return metrics

    def _theis_noise(self, obs):
        """ add noise to discrete variables per Theis 2016 """
        if self.training:
            obs['x'] += (torch.rand_like(obs['x']) - 0.5)
            obs['slice_number'] += (torch.rand_like(obs['slice_number']) - 0.5)
            obs['duration'] += torch.rand_like(obs['duration'] - 0.5)
            obs['duration'].clamp_(min=1e-4)
            obs['edss'] += ((torch.rand_like(obs['edss']) / 2.) - 0.25)
            obs['edss'].clamp_(min=1e-4)
        return obs

    @property
    def pseudo3d(self):
        return self.pyro_model.pseudo3d

    def prep_batch(self, batch):
        x = 255. * batch['image'].float()  # multiply by 255 b/c preprocess tfms
        out = dict(x=x)
        for k in self.required_data:
            if k in batch:
                out[k] = batch[k].unsqueeze(1).float()
        out = self._theis_noise(out)
        return out

    def _steps_per_epoch(self):
        return len(self.calabresi_train) // self.train_batch_size  # integer div b/c drop_last used

    def _set_annealing_factor(self, batch_idx=None):
        steps_per_epoch = self._steps_per_epoch()
        if batch_idx is None:
            batch_idx = steps_per_epoch
        not_in_sanity_check = self.hparams.annealing_epochs > 0
        in_annealing_epochs = self.current_epoch < self.hparams.annealing_epochs
        n_levels = max(self.pyro_model.n_levels, 1)
        self.pyro_model.annealing_factor = [1. for _ in range(n_levels)]
        for i in range(n_levels):
            if not_in_sanity_check and in_annealing_epochs and self.training:
                min_af = self.hparams.min_annealing_factor[i]
                max_af = self.hparams.max_annealing_factor[i]
                self.pyro_model.annealing_factor[i] = min_af + (max_af - min_af) * \
                                   (float(batch_idx + self.current_epoch * steps_per_epoch + 1) /
                                    float(self.hparams.annealing_epochs * steps_per_epoch))
            else:
                self.pyro_model.annealing_factor[i] = self.hparams.max_annealing_factor[i]
            if self.training:
                self.log(f'annealing_factor/af{i}', self.pyro_model.annealing_factor[i],
                         on_step=False, on_epoch=True)

    def training_step(self, batch, batch_idx):
        self._set_annealing_factor(batch_idx)
        batch = self.prep_batch(batch)
        if self.hparams.validate:
            logging.info('Validation:')
            self.print_trace_updates(batch)
        loss = self.svi.step(batch)
        self.scheduler.step()
        loss = torch.as_tensor(loss)
        self.log('train_loss', loss, on_step=False, on_epoch=True)
        metrics = self.get_trace_metrics(batch)
        if np.isnan(loss):
            self.logger.experiment.add_text('nan', f'nand at {self.current_epoch}:\n{metrics}')
            raise ValueError('loss went to nan with metrics:\n{}'.format(metrics))
        for k, v in metrics.items():
            self.log('train/' + k, v, on_step=False, on_epoch=True)
        return loss

    def validation_step(self, batch, batch_idx):
        self._set_annealing_factor()
        batch = self.prep_batch(batch)
        loss = self.svi.evaluate_loss(batch)
        self.log('val_loss', loss, on_step=False, on_epoch=True)
        metrics = self.get_trace_metrics(batch)
        for k, v in metrics.items():
            self.log('val/' + k, v, on_step=False, on_epoch=True)
        return metrics

    def test_step(self, batch, batch_idx):
        import nibabel as nib
        self._set_annealing_factor()
        subject = int(batch['subject'][0])
        scan = int(batch['scan'][0])
        batch = self.prep_batch(batch)
        loss = self.svi.evaluate_loss(batch)
        self.log('test_loss', loss, on_step=False, on_epoch=True)
        metrics = self.get_trace_metrics(batch)
        for k, v in metrics.items():
            self.log('test/' + k, v, on_step=False, on_epoch=True)
        samples = self.build_test_samples(batch)
        for intervention, data in samples.items():
            cf = data['x'].detach().cpu().numpy()
            if self.hparams.pseudo3d:
                cf = cf[:,1,...]  # get the middle slices
            cf = cf.squeeze()
            fn = os.path.join(self.hparams.test_dir, f'{subject}_{scan}_{intervention}.nii.gz')
            nib.Nifti1Image(cf,None).to_filename(fn)
        return {'samples': samples, 'metrics': metrics}

    @classmethod
    def add_arguments(cls, parser):
        parser = super().add_arguments(parser)
        parser.add_argument('--num-svi-particles', default=4, type=int, help="number of particles to use for ELBO (default: %(default)s)")
        parser.add_argument('--num-sample-particles', default=32, type=int, help="number of particles to use for MC sampling (default: %(default)s)")
        parser.add_argument('--use-cf-guide', default=False, action='store_true', help="whether to use counterfactual guide (default: %(default)s)")
        parser.add_argument(
            '--cf-elbo-type', default=-1, choices=[-1, 0, 1, 2],
            help="-1: randomly select per batch, 0: shuffle thickness, 1: shuffle intensity, 2: shuffle both (default: %(default)s)")
        parser.add_argument('--annealing-epochs', default=50, type=int, help="anneal kl div in z for this # epochs (default: %(default)s)")
        parser.add_argument('--min-annealing-factor', default=[0.2], type=float, nargs='+', help="anneal kl div in z starting here (per level for hierarchical) (default: %(default)s)")
        parser.add_argument('--max-annealing-factor', default=[1.0], type=float, nargs='+', help="anneal kl div in z ending here (per level for hierarchical) (default: %(default)s)")
        parser.add_argument('--tracegraph-elbo', default=False, action='store_true', help="use tracegraph elbo (much more computationally expensive) (default: %(default)s)")
        return parser

EXPERIMENT_REGISTRY[SVIExperiment.__name__] = SVIExperiment


# This script creates a yield table with CHEMPY (Rybizki et al. 2017) in numpy format and writes it in xdr file format.
# Yield tables are 4 dimensional depending on: stellar age, stellar (total) metallicity, source (AGB, SNia, SNii) and elemental species.
#
#
# table format:
#
# the general table format is as follows: A Header stating the structure of the table, the table body consisiting of N_z blocks
# of Ns * Ne lines for every stellar metallicity and columns for the stellar age of the SSP (although in xdr format we do not really have new lines...)
#
# Header:     nuber of steps in stellar metallicity Z (N_z), number of steps in stellar age t (N_t), number of elemental species Ne, number of sources Ns
# (no xdr    nSpecies, nSources
# format)    Z_min, Zmax, dZ
#             t_min=0, t_max=t_universe, dt (in log space to sample early stellar evolutionary phases better than late phases - see what Jan provides)
#            list of strings denoting the elements, e.g. 'Ag', 'Al', 'Ba', 'C', 'Fe', 'He', 'Mg', 'N', 'Na', 'Ne', 'O', 'S', 'Si', 'Ti', 'Zn' (should have same length as Ne)
#            write every element in a new line, makes it easier to read in in c.
#
# Table Body:     -------------> stellar age of SSP: N_t
# (xdr format)   |SSP of Z_min:       Mass Loss, source 1
#                |                    ...
#                |                    Mass Loss, source Ns
#                |                    Number of events, source 1
#                |                    ...
#                |                    Number of events, soure Ns
#                |                    Element 1, source 1
#                |                    Element 1, source 2
#                |                    ...
#                |                    Element 1, source Ns
#                |                    ...
#                |                    Element 2, source 1
#                |                    ...
#                |
#                |                    Element Ne, source Ns
#                |SSP of Z_min+dZ:    Mass loss, source 1
#                |                    ...
#                |                    Mass Loss, source Ns
#                |                    Number of events
#                |                    Element 1, source 1
#                |                    ...
#                |                    Element Ne, source Ns
#                |                    ...
#                |SSP of Z_max:       Mass Loss, source 1
#                |                    ....
#                |                    Element 1, source 1
#                |                    ...
#                |                    Element Ne, source Ns
#       Z of SSP |
#
#
# in fact, mass loss and unprocessed mass is the same.
# caveat: SNIA yields are not metallicity dependent and unprocessed mass loss is 0, thus we only store the yields once at the lowest metallicity.
# caveat: direct BH collapse returns only intial element fractions and happens before SNII so we combine the mass loss with SNII mass loss.

import numpy as np
import xdrlib
import numpy as np
import multiprocessing as mp
import os

####### SETTING THE CHEMPY PARAMETER #######################

from Chempy.parameter import ModelParameters
a = ModelParameters()

# Load solar abundances
from Chempy.solar_abundance import solar_abundances
basic_solar = solar_abundances()
getattr(basic_solar, 'Asplund09')()

# Load the yields
from Chempy.yields import SN2_feedback, AGB_feedback, SN1a_feedback
basic_sn2 = SN2_feedback()
getattr(basic_sn2, 'chieffi04_net')()
basic_1a = SN1a_feedback()
getattr(basic_1a, "Seitenzahl")()
basic_agb = AGB_feedback()
getattr(basic_agb, "Karakas16_net")()

# Print all supported elements
elements_to_trace = list(np.unique(basic_agb.elements+basic_sn2.elements+basic_1a.elements))
print(elements_to_trace)

# Producing the SSP birth elemental fractions (here we use solar)
solar_fractions = []
elements = np.hstack(basic_solar.all_elements)
for item in elements_to_trace:
    solar_fractions.append(float(basic_solar.fractions[np.where(elements==item)]))

# Initialise the SSP class with time-steps
time_steps = np.logspace(-2.44,1.139879,100) #np.linspace(0.,13.8,1024)
a.log_time = True

# yieldset
a.yield_table_name_sn2 = 'chieffi04_net'
a.yield_table_name_agb = 'Karakas16_net'
a.yield_table_name_1a = 'Seitenzahl'

# imf parameters
a.only_net_yields_in_process_tables = True
a.imf_type_name = 'Chabrier_1'
a.chabrier_para1 = 0.69
a.chabrier_para2 = 0.079
a.high_mass_slope = -2.3
a.imf_parameter = (a.chabrier_para1, a.chabrier_para2, a.high_mass_slope)
a.mmin = 0.1
a.mmax = 100
# 100,000,000 mass steps are smooth enough for 1000 time steps
a.mass_steps = 1000000 #100000000 #2000 # 200000
a.sn2mmin = 8.
a.sn2mmax = 40.
a.bhmmin = float(a.sn2mmax) ## maximum of hypernova
a.bhmmax = float(a.mmax) ## maximum of the IMF

# sn1a delay parameters for maoz
a.N_0 = np.power(10,-2.9)
a.sn1a_time_delay = np.power(10,-1.39794) #40 Myr
a.sn1a_exponent = 1.12
a.dummy = 0.0
a.sn1a_parameter = [a.N_0,a.sn1a_time_delay,a.sn1a_exponent,a.dummy]
######################## END OF SETTING CHEMPY PARAMETER ########################

######################## SETTING THE YIELDTABLE PARAMETERS ######################
list_of_metallicities = np.logspace(-5,-1.3,50)

from Chempy.wrapper import SSP_wrap

def create_one_SSP_table_old(parameters, source='SNII'):
    differential_table = True
    metallicity = parameters
    print(metallicity,a.yield_table_name_sn2)
    basic_ssp = SSP_wrap(a)
    basic_ssp.calculate_feedback(metallicity,list(elements_to_trace),list(solar_fractions),np.copy(time_steps),1)

    x = basic_ssp.agb_table
    y = basic_ssp.sn1a_table
    z = basic_ssp.sn2_table
    s = basic_ssp.bh_table
    d = basic_ssp.table

    u = np.zeros_like(x)
    names = list(u.dtype.names)

    for j,jtem in enumerate(names):
        if source == 'SNII':
            u[jtem] = z[jtem]
        if source == 'SNIA':
            u[jtem] = y[jtem]
        if source == 'AGB':
            u[jtem] = x[jtem]
        if source == 'BH':
            u[jtem] = s[jtem]
        if source == 'ALL':
            u[jtem] = x[jtem] + y[jtem] + z[jtem] + s[jtem]
    if differential_table:
        for el in elements_to_trace:
            d[el] = u[el]
    else:
        for el in elements_to_trace:
            d[el] = np.cumsum(u[el])
        for name in ['mass_of_ms_stars_dying', 'mass_in_remnants', 'sn2', 'sn1a', 'pn', 'bh', 'hydrogen_mass_accreted_onto_white_dwarfs', 'unprocessed_ejecta']:
            d[name] = np.cumsum(d[name])

    return(d)

def create_one_SSP_table(parameters, source='SNII'):
    metallicity = parameters
    print(metallicity,a.yield_table_name_sn2)
    basic_ssp = SSP_wrap(a)
    basic_ssp.calculate_feedback(metallicity,list(elements_to_trace),list(solar_fractions),np.copy(time_steps),1)

    x = basic_ssp.agb_table
    y = basic_ssp.sn1a_table
    z = basic_ssp.sn2_table
    s = basic_ssp.bh_table
    d = basic_ssp.table

    if source == 'SNII':
        return z
    if source == 'SNIA':
        return y
    if source == 'AGB':
        return x
    if source == 'BH':
        return s
    if source == 'ALL':
        u = np.zeros_like(x)
        names = list(u.dtype.names)
        for j,jtem in enumerate(names):
            u[jtem] = x[jtem] + y[jtem] + z[jtem] + s[jtem]
        return u

def my_wrap_AGB_table(parameters):
    return create_one_SSP_table(parameters, source='AGB')

def my_wrap_SNIA_table(parameters):
    return create_one_SSP_table(parameters, source='SNIA')

def my_wrap_SNII_table(parameters):
    return create_one_SSP_table(parameters, source='SNII')

def my_wrap_BH_table(parameters):
    return create_one_SSP_table(parameters, source='BH')

########## END OF SETTING YIELD TABLE PARAMETERS ################

# Call the SSP table creation routine
print('This python script reads a numpy file created with chempy to transform it to xdr file format to be read in by Gasoline!')

print("There are %d CPUs on this machine" % mp.cpu_count())
number_processes = max(1,20)# mp.cpu_count() - 1)
print("Using %d of them.", number_processes)

file = 'chempy_table_agb'
if not os.path.isfile(file+'.npy'):
    ############ CREATING THE ACTUAL TABLES ####################
    list_of_SSP_tables = []
    list_of_SSP_tables.append(list_of_metallicities)
    list_of_SSP_tables.append(time_steps)
    pool = mp.Pool(number_processes)
    results = pool.map(my_wrap_AGB_table, list_of_metallicities)
    pool.close()
    pool.join()
    list_of_SSP_tables.append(results)
    np.save(file, list_of_SSP_tables)

    list_of_SSP_tables = []
    list_of_SSP_tables.append(list_of_metallicities)
    list_of_SSP_tables.append(time_steps)
    pool = mp.Pool(number_processes)
    results = pool.map(my_wrap_SNIA_table, list_of_metallicities)
    pool.close()
    pool.join()
    list_of_SSP_tables.append(results)
    np.save('chempy_table_snia', list_of_SSP_tables)

    list_of_SSP_tables = []
    list_of_SSP_tables.append(list_of_metallicities)
    list_of_SSP_tables.append(time_steps)
    pool = mp.Pool(number_processes)
    results = pool.map(my_wrap_SNII_table, list_of_metallicities)
    pool.close()
    pool.join()
    list_of_SSP_tables.append(results)
    np.save('chempy_table_snii', list_of_SSP_tables)

    list_of_SSP_tables = []
    list_of_SSP_tables.append(list_of_metallicities)
    list_of_SSP_tables.append(time_steps)
    pool = mp.Pool(number_processes)
    results = pool.map(my_wrap_BH_table, list_of_metallicities)
    pool.close()
    pool.join()
    list_of_SSP_tables.append(results)
    np.save('chempy_table_bh', list_of_SSP_tables)

############### DOING THE EXPORT TO XDR FORMAT ######################
# import yield table created with chempy
print('Reading chempy yield table...')
yield_table_agb = np.load('chempy_table_agb.npy')
yield_table_snIa = np.load('chempy_table_snia.npy')
yield_table_snII = np.load('chempy_table_snii.npy')
yield_table_bh = np.load('chempy_table_bh.npy')

yield_sn2 = 'chieffi+04'
yield_snia = 'seitenzahl+13'
yield_agb = 'karakas+16'

# get parameters from yield table
N_Z = len(yield_table_snII[0])
N_t = len(yield_table_snII[1])
N_e = len(elements_to_trace) #len(yield_table_snII[2][1][0])-9 #change length of elements once final yield table is there

Zmin = yield_table_snII[0][0]
Zmax = yield_table_snII[0][len(yield_table_snII[0])-1]
dZ = np.log10(yield_table_snII[0][1])-np.log10(yield_table_snII[0][0])

tmin = yield_table_snII[1][0]*1e9
tmax = yield_table_snII[1][len(yield_table_snII[1])-1]*1e9
dt = np.log10(yield_table_snII[1][1])-np.log10(yield_table_snII[1][0])

# open new xdr file
xdr_table = open("yieldtable_xdr_high_Ia_norm","wb")

# initialize packing
p = xdrlib.Packer()

# wirte Header
print('Writing Header, stating information about metallicity steps, timesteps and number of elements and sources.')

xdr_table.write(b'###################################################################################\n')
#xdr_table.write(b'\n')
xdr_table.write(b'###   Yield table created with CHEMPY (https://github.com/jan-rybizki/Chempy)   ###\n')
#xdr_table.write(b'\n')
xdr_table.write(b'###               using the write_yield_lookup_table.py script.                 ###\n')
#xdr_table.write(b'\n')
xdr_table.write(b'###                Tobias Buck (tbuck@aip.de) in February 2021.                 ###\n')
#xdr_table.write(b'\n')
xdr_table.write(b'###################################################################################\n')

# here should go which yield sets we use.
# We should actually add all chempy paramters, like e.g. IMF and SNIA delay time parameters.
# In this way the yieldtable is self-consistent.

yieldsets = b'SNII, SNIA and AGB yieldsets used: %s\t %s\t %s\n'%(a.yield_table_name_sn2.encode('latin-1'),a.yield_table_name_1a.encode('latin-1'),a.yield_table_name_agb.encode('latin-1')) #(yield_sn2,yield_snia,yield_agb)
xdr_table.write(yieldsets)

imf = b'IMF parameters used: type: %s (%.3f,%.3f,%.3f)\t IMF min/max mass: %.2f/%.2f\t SN min/max mass: %.1f/%.1f\n'%(a.imf_type_name.encode('latin-1'),a.chabrier_para1,a.chabrier_para2,a.high_mass_slope,a.mmin,a.mmax,a.sn2mmin,a.sn2mmax)
xdr_table.write(imf)

snia = b'SNIA parameters used: type: Maoz+2012, normalization: %.3f\t delay time (Myr): %.2f\t exponent: %.2f\n'%(a.N_0,a.sn1a_time_delay,a.sn1a_exponent)
xdr_table.write(snia)

firstline = b'%i %i %i 3\n'%(N_Z, N_t, N_e)
xdr_table.write(firstline)

secondline = b'%E %E %E\n'%(Zmin, Zmax, dZ)
xdr_table.write(secondline)

thirdline = b'%E %E %E\n'%(tmin, tmax, dt)
xdr_table.write(thirdline)

#SNIA do not eject unprocessed material calculate the massloss from the newly synthesized fraction

for elem in elements_to_trace: #yield_table_snII[2][0].dtype.names[4:]:
    element = b'%s\t'%(elem.encode('latin-1'))
    xdr_table.write(element)
    yield_table_snIa[2][0]['unprocessed_ejecta'] += yield_table_snIa[2][0][elem]
xdr_table.write(b'\n')

print('Packing data to xdr format.')

for i in range(N_Z):
    #step through all metallicity bins
    #combine massloss from direct BH collapse with SNII mass loss
    yield_table_snII[2][i]["unprocessed_ejecta"] += yield_table_bh[2][i]["unprocessed_ejecta"]

    if i == 0:
        # at the lowest Z bin include the SNIA yields
        #p.pack_farray(N_t,np.cumsum(yield_table_snII[2][i]['mass_of_ms_stars_dying']) - np.cumsum(yield_table_snII[2][i]['mass_in_remnants']),p.pack_double)
        p.pack_farray(N_t,np.cumsum(yield_table_snII[2][i]['unprocessed_ejecta']),p.pack_double)
        p.pack_farray(N_t,np.cumsum(yield_table_snIa[2][i]['unprocessed_ejecta']),p.pack_double)
        p.pack_farray(N_t,np.cumsum(yield_table_agb[2][i]['unprocessed_ejecta']),p.pack_double)
        p.pack_farray(N_t,np.cumsum(yield_table_snII[2][i]['number_of_events']),p.pack_double)
        p.pack_farray(N_t,np.cumsum(yield_table_snIa[2][i]['number_of_events']),p.pack_double)
        p.pack_farray(N_t,np.cumsum(yield_table_agb[2][i]['number_of_events']),p.pack_double)
        for elem in elements_to_trace: #yield_table_snII[2][0].dtype.names[4:]:
            # step through all elements
            # and finally step also through all sources once we have them...
            p.pack_farray(N_t,np.cumsum(yield_table_snII[2][i][elem]),p.pack_double)
            p.pack_farray(N_t,np.cumsum(yield_table_snIa[2][i][elem]),p.pack_double)
            p.pack_farray(N_t,np.cumsum(yield_table_agb[2][i][elem]),p.pack_double)
    else:
        # now only SNII and AGB
        #p.pack_farray(N_t,np.cumsum(yield_table_snII[2][i]['mass_of_ms_stars_dying']) - np.cumsum(yield_table_snII[2][i]['mass_in_remnants']),p.pack_double)
        p.pack_farray(N_t,np.cumsum(yield_table_snII[2][i]['unprocessed_ejecta']),p.pack_double)
        p.pack_farray(N_t,np.cumsum(yield_table_agb[2][i]['unprocessed_ejecta']),p.pack_double)
        p.pack_farray(N_t,np.cumsum(yield_table_snII[2][i]['number_of_events']),p.pack_double)
        #p.pack_farray(N_t,np.cumsum(yield_table_snII[2][i]['sn1a']),p.pack_double)
        p.pack_farray(N_t,np.cumsum(yield_table_agb[2][i]['number_of_events']),p.pack_double)
        for elem in elements_to_trace: #yield_table_snII[2][0].dtype.names[4:]:
            # step through all elements
            # and finally step also through all sources once we have them...
            p.pack_farray(N_t,np.cumsum(yield_table_snII[2][i][elem]),p.pack_double)
            #p.pack_farray(N_t,np.cumsum(yield_table_snIa[2][i][elem]),p.pack_double)
            p.pack_farray(N_t,np.cumsum(yield_table_agb[2][i][elem]),p.pack_double)

print('Writing xdr part.')
#xdr_table.write(p.get_buffer())

xdr_table.write(p.get_buffer())

xdr_table.close()