# -*- coding: future_fstrings -*-

import numpy as np
import torch
from torch.distributions import Normal, kl_divergence, Bernoulli, RelaxedBernoulli
from torch import nn
import torch.nn.functional as F
from .utils import spatial_transform, calc_kl_z_pres_bernoulli
from .common import *
from .modules import NumericalRelaxedBernoulli

class TrackerRNN(nn.Module):

    def __init__(self, hid_dim):
        super(TrackerRNN, self).__init__()

        self.cell = nn.GRUCell(temporal_rnn_inp_dim, hid_dim)

    def forward(self, h_pre, c_pre, temporal_rnn_inp):
        h = self.cell(temporal_rnn_inp, c_pre)

        # output and hidden, for vanilla rnn, output == hidden
        return h, h

class AttEncoder(nn.Module):

    def __init__(self, args):
        self.args = args
        super(AttEncoder, self).__init__()

        self.temporal_img_conv_net = nn.Sequential(
            nn.Conv2d(img_encode_dim, temporal_img_enc_hid_dim, 1),
            nn.ELU(),
            nn.GroupNorm(8, temporal_img_enc_hid_dim)
        )

        self.temporal_img_enc_net = nn.Linear(
            temporal_img_enc_hid_dim * self.args.num_cell_h // 2 * self.args.num_cell_w // 2, temporal_img_enc_dim)

    def forward(self, img_enc):
        """

        :param x: (bs, dim, img_h, img_w)
        """
        bs = img_enc.size(0)
        x = self.temporal_img_enc_net(self.temporal_img_conv_net(img_enc).view(bs, -1))

        # bs, dim
        return x

class PropagationCell(nn.Module):

    def __init__(self, args, z_what_net, glimpse_dec_net):
        super(PropagationCell, self).__init__()
        self.args = args

        self.z_pres_logits_bias = 2.
        self.where_update_scale = where_update_scale
        self.z_where_std_bias = -2

        # self.z_what_gate_bias = 2
        self.register_buffer('z_pres_stop_threshold', torch.tensor(0.6))

        z_where_transit_bias_net_input_dim = temporal_rnn_out_dim + z_what_dim + z_where_scale_dim + \
                                             z_where_shift_dim + z_where_bias_dim + temporal_img_enc_dim

        self.z_where_transit_bias_net = nn.Sequential(
            nn.Linear(z_where_transit_bias_net_input_dim, z_where_transit_bias_net_hid_dim),
            nn.ELU(),
            nn.Linear(z_where_transit_bias_net_hid_dim, (z_where_scale_dim + z_where_shift_dim) * 2)
        )

        z_depth_transit_net_input_dim = temporal_rnn_out_dim + z_what_dim + temporal_img_enc_dim

        self.z_depth_transit_net = nn.Sequential(
            nn.Linear(z_depth_transit_net_input_dim, z_depth_transit_net_hid_dim),
            nn.ELU(),
            nn.Linear(z_depth_transit_net_hid_dim, z_depth_dim * 2)
        )

        self.z_what_from_temporal_net = nn.Sequential(
            nn.Linear(temporal_rnn_out_dim, z_what_from_temporal_hid_dim),
            nn.ELU(),
            nn.Linear(z_what_from_temporal_hid_dim, z_what_dim * 2)
        )

        z_what_gate_net_inp_dim = temporal_rnn_out_dim + temporal_img_enc_dim

        self.z_what_gate_net = nn.Sequential(
            nn.Linear(z_what_gate_net_inp_dim, 64),
            nn.ELU(),
            nn.Linear(64, 2),
            nn.Sigmoid(),
        )

        z_pres_transit_input_dim = temporal_rnn_out_dim + z_where_scale_dim + \
                                   z_where_shift_dim + z_where_bias_dim + z_what_dim

        self.z_pres_transit = nn.Sequential(
            nn.Linear(z_pres_transit_input_dim, z_pres_hid_dim),
            nn.ELU(),
            nn.Linear(z_pres_hid_dim, z_pres_dim),
        )

        temporal_rnn_inp_net_inp_dim = z_where_scale_dim + z_where_shift_dim + z_pres_dim + \
                                       z_what_dim + z_where_bias_dim + temporal_img_enc_dim

        self.temporal_rnn_inp_net = nn.Linear(temporal_rnn_inp_net_inp_dim, temporal_rnn_inp_dim)

        self.temporal_rnn = TrackerRNN(temporal_rnn_hid_dim)
        self.attention_encoding = AttEncoder(self.args)

        self.glimpse_dec_net = glimpse_dec_net
        self.z_what_net = z_what_net
        self.prior_cell = PropagatePrior()

    def update_z_where(self, z_where_pre, z_where_bias):
        z_where_shift = z_where_pre[:, 2:] + self.where_update_scale * z_where_bias[:, 2:].tanh()

        scale, ratio = z_where_bias[:, :2].tanh().chunk(2, 1)
        scale = self.args.size_anc + self.args.var_s * scale  # add bias to let masking do its job
        ratio = self.args.ratio_anc + self.args.var_anc * ratio
        ratio_sqrt = ratio.sqrt()

        z_where = torch.cat((scale / ratio_sqrt, scale * ratio_sqrt, z_where_shift), dim=1)
        # # always within the image
        z_where = torch.cat((z_where[:, :2], z_where[:, 2:].clamp(-1.05, 1.05)), dim=1)
        return z_where

    def forward(self, x, img_enc, temporal_rnn_out_pre, temporal_rnn_hid_pre, prior_rnn_out_pre,
                prior_rnn_hid_pre, z_what_pre, z_where_pre, z_where_bias_pre, z_depth_pre, z_pres_pre,
                cumsum_one_minus_z_pres, ids_pre, lengths, max_length, t, no_disc=False, eps=1e-15):
        """

        :param x: input image (bs, c, h, w)
        :param img_enc: input image encode (bs, c, num_cell_h, num_cell_w)
        :param temporal_rnn_out_pre: (bs, max_num_obj, dim)
        :param temporal_rnn_hid_pre: (bs, max_num_obj, dim)
        :param z_what_pre: (bs, max_num_obj, dim)
        :param z_where_pre: (bs, max_num_obj, dim)
        :param z_depth_pre: (bs, max_num_obj, dim)
        :param z_pres_pre: (bs, max_num_obj, dim)
        :param cumsum_one_minus_z_pres: (bs, max_num_obj, dim)
        :param lengths: (bs)
        :return:
        """
        bs = x.size(0)
        device = x.device
        max_num_obj = max_length
        bns = bs * max_num_obj
        obj_mask = (z_pres_pre.view(bs, max_num_obj) != 0).float()
        temporal_rnn_out_pre, temporal_rnn_hid_pre, prior_rnn_out_pre, \
        prior_rnn_hid_pre, z_what_pre, z_where_pre, z_where_bias_pre, z_depth_pre, \
        z_pres_pre, cumsum_one_minus_z_pres = \
            temporal_rnn_out_pre.view(bns, -1), temporal_rnn_hid_pre.view(bns, -1), \
            prior_rnn_out_pre.view(bns, -1), prior_rnn_hid_pre.view(bns, -1), \
            z_what_pre.view(bns, -1), z_where_pre.view(bns, -1), z_where_bias_pre.view(bns, -1), \
            z_depth_pre.view(bns, -1), z_pres_pre.view(bns, -1), \
            cumsum_one_minus_z_pres.view(bns, -1)

        prior_rnn_out, prior_rnn_hid, prior_what_mean, prior_what_std, prior_where_bias_mean, \
        prior_where_bias_std, prior_depth_mean, prior_depth_std, prior_pres_prob = \
            self.prior_cell(prior_rnn_out_pre, prior_rnn_hid_pre, z_what_pre,
                            z_where_pre, z_where_bias_pre, z_depth_pre, z_pres_pre)

        z_where_att = x.new_ones(z_where_pre.size()) * .5
        z_where_att[:, 2:] = z_where_pre[:, 2:].detach()
        img_enc_att = spatial_transform(
            img_enc.unsqueeze(1).expand(-1, max_num_obj, -1, -1, -1).contiguous().
                view(bns, img_encode_dim, self.args.num_cell_h, self.args.num_cell_w), z_where_att,
            (bns, img_encode_dim, self.args.num_cell_h // 2, self.args.num_cell_w // 2), inverse=False
        )
        # bns, dim
        temporal_img_enc = self.attention_encoding(img_enc_att).view(-1, temporal_img_enc_dim)
        temporal_img_enc = \
            temporal_img_enc.view(bs, -1, temporal_img_enc_dim).contiguous().view(-1, temporal_img_enc_dim)

        temporal_rnn_inp_net_inp = torch.cat(
            [z_where_pre, z_pres_pre, z_what_pre, z_where_bias_pre, temporal_img_enc],
            dim=1
        )
        temporal_rnn_inp = self.temporal_rnn_inp_net(temporal_rnn_inp_net_inp)
        # bns, dim
        temporal_rnn_out, temporal_rnn_hid = self.temporal_rnn(
            temporal_rnn_out_pre, temporal_rnn_hid_pre, temporal_rnn_inp
        )

        # z_where transition
        z_where_transit_bias_net_inp = torch.cat(
            [temporal_rnn_out, z_what_pre, z_where_pre, z_where_bias_pre, temporal_img_enc], dim=1
        )
        z_where_bias_mean, z_where_bias_std = \
            self.z_where_transit_bias_net(z_where_transit_bias_net_inp).chunk(2, -1)
        z_where_bias_std = F.softplus(z_where_bias_std + self.z_where_std_bias)
        if self.args.phase_generate and t >= self.args.observe_frames:
            z_where_bias_dist = Normal(prior_where_bias_mean, prior_where_bias_std)
        else:
            z_where_bias_dist = Normal(z_where_bias_mean, z_where_bias_std)

        z_where_bias = z_where_bias_dist.rsample()
        z_where = self.update_z_where(z_where_pre, z_where_bias)
        z_where_mean = self.update_z_where(z_where_pre, z_where_bias_mean)

        # get glimpse encode
        x_att = \
            spatial_transform(
                x.unsqueeze(1).expand(-1, max_num_obj, -1, -1, -1).contiguous().view(bns, 3, img_h, img_w), z_where,
                (bns, 3, glimpse_size, glimpse_size), inverse=False
            )

        z_what_from_enc_mean, z_what_from_enc_std = self.z_what_net(
            x_att
        )
        z_what_from_enc_std = F.softplus(z_what_from_enc_std)

        # z_what transit
        z_what_from_temporal_mean, z_what_from_temporal_std = \
            self.z_what_from_temporal_net(temporal_rnn_out).chunk(2, -1)

        z_what_from_temporal_std = F.softplus(z_what_from_temporal_std)

        z_what_gate_net_inp = torch.cat((temporal_rnn_out, temporal_img_enc), dim=1)
        forget_gate, input_gate = self.z_what_gate_net(z_what_gate_net_inp).chunk(2, -1)

        z_what_mean = input_gate * z_what_from_enc_mean + \
                      forget_gate * z_what_from_temporal_mean

        z_what_std = F.softplus(input_gate * z_what_from_enc_std + \
                                forget_gate * z_what_from_temporal_std)

        if self.args.phase_generate and t >= self.args.observe_frames:
            z_what_dist = Normal(prior_what_mean, prior_what_std)
        else:
            z_what_dist = Normal(z_what_mean, z_what_std)

        z_what = z_what_dist.rsample()

        z_depth_transit_net_inp = torch.cat(
            [temporal_rnn_out, z_what, temporal_img_enc],
            dim=1
        )
        z_depth_mean, z_depth_std = self.z_depth_transit_net(z_depth_transit_net_inp).chunk(2, -1)
        z_depth_std = F.softplus(z_depth_std)

        if self.args.phase_generate and t >= self.args.observe_frames:
            z_depth_dist = Normal(prior_depth_mean, prior_depth_std)
        else:
            z_depth_dist = Normal(z_depth_mean, z_depth_std)

        z_depth = z_depth_dist.rsample()

        # z_pres bns, dim
        z_pres_transit_inp = torch.cat(
            [temporal_rnn_out, z_where, z_where_bias, z_what],
            dim=1
        )
        z_pres_logits = pres_logit_factor * torch.tanh(self.z_pres_transit(z_pres_transit_inp) +
                                                       self.z_pres_logits_bias)
        if self.args.phase_generate and t >= self.args.observe_frames:
            q_z_pres = NumericalRelaxedBernoulli(probs=prior_pres_prob, temperature=self.args.tau)
        else:
            q_z_pres = NumericalRelaxedBernoulli(logits=z_pres_logits, temperature=self.args.tau)

        # for z_pres, we end up setting this to one during generation
        z_pres_y = q_z_pres.rsample()
        z_pres = torch.sigmoid(z_pres_y)
        if no_disc:
            z_pres = torch.ones_like(z_pres)
        cumsum_one_minus_z_pres += (1 - z_pres) * obj_mask.view(bns, 1)
        z_pres = z_pres * (cumsum_one_minus_z_pres < self.z_pres_stop_threshold).float()

        # (bs, dim, glimpse_size, glimpse_size)
        o_att, alpha_att = self.glimpse_dec_net(z_what)

        alpha_att_hat = alpha_att * z_pres.view(-1, 1, 1, 1)
        y_att = alpha_att_hat * o_att

        # (bs, 3, img_h, img_w)
        y_each_obj = spatial_transform(y_att, z_where, (bns, 3, img_h, img_w), inverse=True)

        # (batch_size_t, 1, glimpse_size, glimpse_size)
        importance_map = alpha_att_hat * torch.sigmoid(-z_depth).view(-1, 1, 1, 1)

        # (batch_size_t, 1, img_h, img_w)
        importance_map_full_res = spatial_transform(importance_map, z_where, (bns, 1, img_h, img_w),
                                                    inverse=True)

        # (batch_size_t, 1, img_h, img_w)
        alpha_map = spatial_transform(alpha_att_hat, z_where, (bns, 1, img_h, img_w), inverse=True)

        kl_z_pres = \
            (calc_kl_z_pres_bernoulli(z_pres_logits, prior_pres_prob) *
             obj_mask.view(bns)).view(bs, max_num_obj).sum(1)

        prior_what_dist = Normal(prior_what_mean, prior_what_std)
        prior_where_bias_dist = Normal(prior_where_bias_mean, prior_where_bias_std)
        prior_depth_dist = Normal(prior_depth_mean, prior_depth_std)

        kl_z_what = \
            (kl_divergence(z_what_dist, prior_what_dist).sum(1) * \
             z_pres.squeeze() * obj_mask.view(bns)).view(bs, max_num_obj).sum(1)
        kl_z_where = \
            (kl_divergence(z_where_bias_dist, prior_where_bias_dist).sum(1) * \
             z_pres.squeeze() * obj_mask.view(bns)).view(bs, max_num_obj).sum(1)
        kl_z_depth = \
            (kl_divergence(z_depth_dist, prior_depth_dist).sum(1) * \
             z_pres.squeeze() * obj_mask.view(bns)).view(bs, max_num_obj).sum(1)

        ########################################### Compute log importance ############################################
        log_imp = x.new_zeros(bs, 1)
        if not self.training and self.args.phase_nll:
            z_pres_binary = (z_pres > 0.5).float()
            # (bns, dim)
            log_imp_what = (prior_what_dist.log_prob(z_what) - z_what_dist.log_prob(z_what)) * \
                           z_pres_binary * obj_mask.view(bns, 1)
            log_imp_depth = (prior_depth_dist.log_prob(z_depth) - z_depth_dist.log_prob(z_depth)) * \
                            z_pres_binary * obj_mask.view(bns, 1)
            log_imp_where = (prior_where_bias_dist.log_prob(z_where_bias) - z_where_bias_dist.log_prob(z_where_bias)) * \
                            z_pres_binary * obj_mask.view(bns, 1)

            log_pres_prior = z_pres_binary * torch.log(prior_pres_prob + eps) + \
                             (1 - z_pres_binary) * torch.log(1 - prior_pres_prob + eps)

            log_pres_pos = z_pres_binary * torch.log(torch.sigmoid(z_pres_logits) + eps) + \
                           (1 - z_pres_binary) * torch.log(1 - torch.sigmoid(z_pres_logits) + eps)

            log_imp_pres = (log_pres_prior - log_pres_pos) * obj_mask.view(bns, 1)

            log_imp = log_imp_what.view(bs, -1).sum(1) + log_imp_depth.view(bs, -1).sum(1) + \
                      log_imp_where.view(bs, -1).sum(1) + log_imp_pres.view(bs, -1).sum(1)

        ######################################## End of Compute log importance #########################################
        z_what_all = z_what.view(bs, max_num_obj, -1) * obj_mask.view(bs, max_num_obj, 1)
        z_where_dummy = x.new_ones(bs, max_num_obj, (z_where_scale_dim + z_where_shift_dim)) * .5
        z_where_dummy[:, :, z_where_scale_dim:] = 2
        z_where_all = z_where.view(bs, max_num_obj, -1) * obj_mask.view(bs, max_num_obj, 1) + \
                      z_where_dummy * (1 - obj_mask.view(bs, max_num_obj, 1))
        z_where_bias_all = z_where_bias.view(bs, max_num_obj, -1) * obj_mask.view(bs, max_num_obj, 1)
        z_pres_all = z_pres.view(bs, max_num_obj, -1) * obj_mask.view(bs, max_num_obj, 1)
        temporal_rnn_hid_all = \
            temporal_rnn_hid.view(bs, max_num_obj, -1) * obj_mask.view(bs, max_num_obj, 1)
        temporal_rnn_out_all = \
            temporal_rnn_out.view(bs, max_num_obj, -1) * obj_mask.view(bs, max_num_obj, 1)
        z_depth_all = z_depth.view(bs, max_num_obj, -1) * obj_mask.view(bs, max_num_obj, 1)
        y_each_obj_all = \
            y_each_obj.view(bs, max_num_obj, 3, img_h, img_w) * obj_mask.view(bs, max_num_obj, 1, 1, 1)
        alpha_map_all = \
            alpha_map.view(bs, max_num_obj, 1, img_h, img_w) * obj_mask.view(bs, max_num_obj, 1, 1, 1)
        importance_map_all = \
            importance_map_full_res.view(bs, max_num_obj, 1, img_h, img_w) * \
            obj_mask.view(bs, max_num_obj, 1, 1, 1)

        cumsum_one_minus_z_pres = cumsum_one_minus_z_pres.view(bs, max_num_obj, -1)
        prior_rnn_out = prior_rnn_out.view(bs, max_num_obj, -1)
        prior_rnn_hid = prior_rnn_hid.view(bs, max_num_obj, -1)

        if self.args.log_phase:
            self.log = {
                'z_what': z_what_all,
                'z_where': z_where_all,
                'z_pres': z_pres_all,
                'z_what_std': z_what_std.view(bs, max_num_obj, -1),
                'z_what_mean': z_what_mean.view(bs, max_num_obj, -1),
                'z_where_bias_std': z_where_bias_std.view(bs, max_num_obj, -1),
                'z_where_bias_mean': z_where_bias_mean.view(bs, max_num_obj, -1),
                'glimpse': x_att.view(bs, max_num_obj, 3, glimpse_size, glimpse_size),
                'glimpse_recon': y_att.view(bs, max_num_obj, 3, glimpse_size, glimpse_size),
                'prior_z_pres_prob': prior_pres_prob.view(bs, max_num_obj, -1),
                'prior_where_bias_std': prior_where_bias_std.view(bs, max_num_obj, -1),
                'prior_where_bias_mean': prior_where_bias_mean.view(bs, max_num_obj, -1),
                'prior_what_mean': prior_what_mean.view(bs, max_num_obj, -1),
                'prior_what_std': prior_what_std.view(bs, max_num_obj, -1),
                'lengths': lengths,
                'z_depth': z_depth_all,
                'z_depth_std': z_depth_std.view(bs, max_num_obj, -1),
                'z_depth_mean': z_depth_mean.view(bs, max_num_obj, -1),
                'y_each_obj': y_each_obj_all.view(bs, max_num_obj, 3, img_h, img_w),
                'alpha_map': alpha_map_all.view(bs, max_num_obj, 1, img_h, img_w),
                'importance_map': importance_map_all.view(bs, max_num_obj, 1, img_h, img_w),
                'z_pres_logits': z_pres_logits.view(bs, max_num_obj, -1),
                'z_pres_y': z_pres_y.view(bs, max_num_obj, -1),
                'o_att': o_att.view(bs, max_num_obj, 3, glimpse_size, glimpse_size),
                'z_where_bias': z_where_bias_all,
                'ids': ids_pre
            }
        else:
            self.log = {}
        representation = {"z_where": z_where_mean.view(bs, max_num_obj, -1) * obj_mask.view(bs, max_num_obj, 1)+z_where_dummy * (1 - obj_mask.view(bs, max_num_obj, 1)),
                          "z_what": z_what_mean.view(bs, max_num_obj, -1) * obj_mask.view(bs, max_num_obj, 1),
                          "z_depth": z_depth_mean.view(bs, max_num_obj, -1) * obj_mask.view(bs, max_num_obj, 1)
                          }
        only_prop_representation = {"z_where": z_where_mean,
                          "z_what": z_what_mean,
                          "z_depth": z_depth_mean
                          }
        return y_each_obj_all, alpha_map_all, importance_map_all, z_what_all, z_where_all, \
               z_where_bias_all, z_depth_all, z_pres_all, ids_pre, kl_z_what, kl_z_where, kl_z_depth, \
               kl_z_pres, temporal_rnn_out_all, temporal_rnn_hid_all, prior_rnn_out, \
               prior_rnn_hid, cumsum_one_minus_z_pres, log_imp, self.log, representation, only_prop_representation

class PropagatePrior(nn.Module):
    """Attention, initial state of rnn is learnable"""

    def __init__(self):
        super(PropagatePrior, self).__init__()

        prior_rnn_inp_net_inp_dim = z_what_dim + z_where_scale_dim + z_where_shift_dim + \
                                    z_where_bias_dim + z_depth_dim + z_pres_dim

        self.prior_rnn_inp_net = nn.Linear(prior_rnn_inp_net_inp_dim, prior_rnn_inp_dim)

        self.prior_rnn = nn.LSTMCell(prior_rnn_inp_dim, prior_rnn_hid_dim)

        self.prior_what_net = nn.Linear(prior_rnn_out_dim, z_what_dim * 2)
        self.prior_where_bias_net = nn.Linear(
            prior_rnn_out_dim, (z_where_scale_dim + z_where_shift_dim) * 2)

        self.prior_depth_net = nn.Linear(prior_rnn_out_dim, z_depth_dim * 2)
        self.prior_pres_net = nn.Linear(prior_rnn_out_dim, z_pres_dim)
        self.prior_z_pres_logits_bias = 5.
        self.where_update_scale = where_update_scale

    def forward(self, prior_rnn_out_pre, prior_rnn_hid_pre, z_what_pre,
                z_where_pre, z_where_bias_pre, z_depth_pre, z_pres_pre, eps=1e-15):
        bns = z_what_pre.size(0)

        z_what_pre_flat = z_what_pre.view(-1, z_what_dim)
        z_where_pre_flat = z_where_pre.view(-1, z_where_scale_dim + z_where_shift_dim)
        z_where_bias_pre_flat = z_where_bias_pre.view(-1, z_where_bias_dim)
        z_depth_pre_flat = z_depth_pre.view(-1, z_depth_dim)
        z_pres_pre_flat = z_pres_pre.view(-1, z_pres_dim)

        prior_rnn_out_pre_flat = prior_rnn_out_pre.view(bns, -1)
        prior_rnn_hid_pre_flat = prior_rnn_hid_pre.view(bns, -1)

        # prior_rnn
        prior_rnn_inp_net_inp = torch.cat((z_what_pre_flat, z_where_pre_flat, z_where_bias_pre_flat,
                                           z_depth_pre_flat, z_pres_pre_flat), dim=1)
        prior_rnn_inp = self.prior_rnn_inp_net(prior_rnn_inp_net_inp)

        prior_rnn_out, prior_rnn_hid = self.prior_rnn(prior_rnn_inp, (prior_rnn_out_pre_flat,
                                                                      prior_rnn_hid_pre_flat))

        prior_what_mean, prior_what_std = self.prior_what_net(prior_rnn_out).chunk(2, -1)
        prior_depth_mean, prior_depth_std = self.prior_depth_net(prior_rnn_out).chunk(2, -1)
        prior_where_bias_mean, prior_where_bias_std = self.prior_where_bias_net(prior_rnn_out).chunk(2, -1)

        prior_pres_probs = torch.sigmoid(self.prior_pres_net(prior_rnn_out) + \
                                         self.prior_z_pres_logits_bias)

        prior_rnn_out = prior_rnn_out.view(bns, -1)
        prior_rnn_hid = prior_rnn_hid.view(bns, -1)
        prior_what_mean = prior_what_mean.view(bns, -1)
        prior_what_std = prior_what_std.view(bns, -1)
        prior_where_bias_mean = prior_where_bias_mean.view(bns, -1)
        prior_where_bias_std = prior_where_bias_std.view(bns, -1)
        prior_depth_mean = prior_depth_mean.view(bns, -1)
        prior_depth_std = prior_depth_std.view(bns, -1)
        prior_pres_probs = prior_pres_probs.view(bns, -1)

        return prior_rnn_out, prior_rnn_hid, prior_what_mean, F.softplus(prior_what_std), \
               prior_where_bias_mean, F.softplus(prior_where_bias_std), prior_depth_mean, \
               F.softplus(prior_depth_std), prior_pres_probs


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

import enum
from abc import ABC, abstractmethod
from typing import Optional, Tuple, Union, cast

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

from . import covariances, kernels, mean_functions
from .base import Module, Parameter, TensorType
from .conditionals.util import (
    base_conditional,
    base_conditional_with_lm,
    expand_independent_outputs,
    fully_correlated_conditional,
    independent_interdomain_conditional,
    mix_latent_gp,
    separate_independent_conditional_implementation,
)
from .config import default_float, default_jitter
from .inducing_variables import (
    FallbackSeparateIndependentInducingVariables,
    FallbackSharedIndependentInducingVariables,
    InducingPoints,
    InducingVariables,
    SeparateIndependentInducingVariables,
    SharedIndependentInducingVariables,
)
from .types import MeanAndVariance
from .utilities import Dispatcher, add_noise_cov

class _QDistribution(Module):
    """
    Base class for our parametrization of q(u) in the `AbstractPosterior`.
    Internal - do not rely on this outside of GPflow.
    """

class _DeltaDist(_QDistribution):
    def __init__(self, q_mu):
        self.q_mu = q_mu  # [M, L]

    @property
    def q_sqrt(self):
        return None

class _DiagNormal(_QDistribution):
    def __init__(self, q_mu, q_sqrt):
        self.q_mu = q_mu  # [M, L]
        self.q_sqrt = q_sqrt  # [M, L]

class _MvNormal(_QDistribution):
    def __init__(self, q_mu, q_sqrt):
        self.q_mu = q_mu  # [M, L]
        self.q_sqrt = q_sqrt  # [L, M, M], lower-triangular

class PrecomputeCacheType(enum.Enum):
    """
    - `PrecomputeCacheType.TENSOR` (or `"tensor"`): Precomputes the cached
      quantities and stores them as tensors (which allows differentiating
      through the prediction). This is the default.
    - `PrecomputeCacheType.VARIABLE` (or `"variable"`): Precomputes the cached
      quantities and stores them as variables, which allows for updating
      their values without changing the compute graph (relevant for AOT
      compilation).
    - `PrecomputeCacheType.NOCACHE` (or `"nocache"` or `None`): Avoids
      immediate cache computation. This is useful for avoiding extraneous
      computations when you only want to call the posterior's
      `fused_predict_f` method.
    """

    TENSOR = "tensor"
    VARIABLE = "variable"
    NOCACHE = "nocache"

def _validate_precompute_cache_type(value) -> PrecomputeCacheType:
    if value is None:
        return PrecomputeCacheType.NOCACHE
    elif isinstance(value, PrecomputeCacheType):
        return value
    elif isinstance(value, str):
        return PrecomputeCacheType(value.lower())
    else:
        raise ValueError(
            f"{value} is not a valid PrecomputeCacheType. Valid options: 'tensor', 'variable', 'nocache' (or None)."
        )

class AbstractPosterior(Module, ABC):
    def __init__(
        self,
        kernel,
        X_data: Union[tf.Tensor, InducingVariables],
        alpha: Optional[TensorType] = None,
        Qinv: Optional[TensorType] = None,
        mean_function: Optional[mean_functions.MeanFunction] = None,
    ):
        """
        Users should use `create_posterior` to create instances of concrete
        subclasses of this AbstractPosterior class instead of calling this
        constructor directly. For `create_posterior` to be able to correctly
        instantiate subclasses, developers need to ensure their subclasses
        don't change the constructor signature.
        """
        super().__init__()

        self.kernel = kernel
        self.X_data = X_data
        self.alpha = alpha
        self.Qinv = Qinv
        self.mean_function = mean_function

    def _add_mean_function(self, Xnew, mean):
        if self.mean_function is None:
            return mean
        else:
            return mean + self.mean_function(Xnew)

    def fused_predict_f(
        self, Xnew, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, including mean_function
        Does not make use of caching
        """
        mean, cov = self._conditional_fused(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )
        return self._add_mean_function(Xnew, mean), cov

    @abstractmethod
    def _conditional_fused(
        self, Xnew, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, *excluding* mean_function
        Does not make use of caching
        """

    def predict_f(
        self, Xnew, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, including mean_function.
        Relies on precomputed alpha and Qinv (see _precompute method)
        """
        if self.alpha is None or self.Qinv is None:
            raise ValueError(
                "Cache has not been precomputed yet. Call update_cache first or use fused_predict_f"
            )
        mean, cov = self._conditional_with_precompute(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )
        return self._add_mean_function(Xnew, mean), cov

    @abstractmethod
    def _conditional_with_precompute(
        self, Xnew, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, *excluding* mean_function.
        Relies on cached alpha and Qinv.
        """

    def update_cache(self, precompute_cache: Optional[PrecomputeCacheType] = None) -> None:
        """
        Sets the cache depending on the value of `precompute_cache` to a
        `tf.Tensor`, `tf.Variable`, or clears the cache. If `precompute_cache`
        is not given, the setting defaults to the most-recently-used one.
        """
        if precompute_cache is None:
            try:
                precompute_cache = cast(
                    PrecomputeCacheType,
                    self._precompute_cache,  # type: ignore
                )
            except AttributeError:
                raise ValueError(
                    "You must pass precompute_cache explicitly (the cache had not been updated before)."
                )
        else:
            self._precompute_cache = precompute_cache

        if precompute_cache is PrecomputeCacheType.NOCACHE:
            self.alpha = self.Qinv = None

        elif precompute_cache is PrecomputeCacheType.TENSOR:
            self.alpha, self.Qinv = self._precompute()

        elif precompute_cache is PrecomputeCacheType.VARIABLE:
            alpha, Qinv = self._precompute()
            if isinstance(self.alpha, tf.Variable) and isinstance(self.Qinv, tf.Variable):
                # re-use existing variables
                self.alpha.assign(alpha)
                self.Qinv.assign(Qinv)
            else:  # create variables
                self.alpha = tf.Variable(alpha, trainable=False)
                self.Qinv = tf.Variable(Qinv, trainable=False)

class GPRPosterior(AbstractPosterior):
    def __init__(
        self,
        kernel,
        X_data: tf.Tensor,
        Y_data: tf.Tensor,
        likelihood_variance: Parameter,
        mean_function: Optional[mean_functions.MeanFunction] = None,
        *,
        precompute_cache: Optional[PrecomputeCacheType],
    ):

        super().__init__(kernel, X_data, mean_function=mean_function)
        self.mean_function = mean_function
        self.Y_data = Y_data
        self.likelihood_variance = likelihood_variance

        if precompute_cache is not None:
            self.update_cache(precompute_cache)

    def _conditional_with_precompute(
        self, Xnew, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, *excluding* mean_function.
        Relies on cached alpha and Qinv.
        """
        Kmn = self.kernel(self.X_data, Xnew)
        Knn = self.kernel(Xnew, full_cov=full_cov)

        return base_conditional_with_lm(Kmn, self.Qinv, Knn, self.alpha, full_cov=full_cov)

    def _precompute(self) -> Tuple[tf.Tensor, tf.Tensor]:

        """
        Precomputes the cholesky decomposition of Kmm_plus_s for later reuse we will call
        base_conditional_with_lm implementation ('Qinv' in the Abstract Posterior class). We also
        precompute the less compute intensive error term ('alpha' in the Abstract Posterior class)
        """

        Kmm = self.kernel(self.X_data)
        Kmm_plus_s = add_noise_cov(Kmm, self.likelihood_variance)

        # obtain the cholesky decomposition of Kmm_plus_s
        Lm = tf.linalg.cholesky(Kmm_plus_s)

        alpha = self.Y_data - self.mean_function(self.X_data)  # type: ignore
        tf.debugging.assert_shapes(
            [
                (Lm, ["M", "M"]),
                (Kmm, ["M", "M"]),
            ]
        )
        return alpha, Lm

    def _conditional_fused(
        self, Xnew, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, *excluding* mean_function
        Does not make use of caching
        """

        # taken directly from the deprecated GPR implementation
        err = self.Y_data - self.mean_function(self.X_data)  # type: ignore

        Kmm = self.kernel(self.X_data)
        Knn = self.kernel(Xnew, full_cov=full_cov)
        Kmn = self.kernel(self.X_data, Xnew)
        Kmm_plus_s = add_noise_cov(Kmm, self.likelihood_variance)

        return base_conditional(
            Kmn, Kmm_plus_s, Knn, err, full_cov=full_cov, white=False
        )  # [N, P], [N, P] or [P, N, N]

class BasePosterior(AbstractPosterior):
    def __init__(
        self,
        kernel,
        inducing_variable,
        q_mu: tf.Tensor,
        q_sqrt: tf.Tensor,
        whiten: bool = True,
        mean_function: Optional[mean_functions.MeanFunction] = None,
        *,
        precompute_cache: Optional[PrecomputeCacheType],
    ):

        super().__init__(kernel, inducing_variable, mean_function=mean_function)
        self.whiten = whiten
        self._set_qdist(q_mu, q_sqrt)

        if precompute_cache is not None:
            self.update_cache(precompute_cache)

    @property
    def q_mu(self):
        return self._q_dist.q_mu

    @property
    def q_sqrt(self):
        return self._q_dist.q_sqrt

    def _set_qdist(self, q_mu, q_sqrt):
        if q_sqrt is None:
            self._q_dist = _DeltaDist(q_mu)
        elif len(q_sqrt.shape) == 2:  # q_diag
            self._q_dist = _DiagNormal(q_mu, q_sqrt)
        else:
            self._q_dist = _MvNormal(q_mu, q_sqrt)

    def _precompute(self):
        Kuu = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [(R), M, M]
        q_mu = self._q_dist.q_mu

        if Kuu.shape.ndims == 4:
            ML = tf.reduce_prod(tf.shape(Kuu)[:2])
            Kuu = tf.reshape(Kuu, [ML, ML])
        if Kuu.shape.ndims == 3:
            q_mu = tf.linalg.adjoint(self._q_dist.q_mu)[..., None]  # [..., R, M, 1]
        L = tf.linalg.cholesky(Kuu)

        if not self.whiten:
            # alpha = Kuu⁻¹ q_mu
            alpha = tf.linalg.cholesky_solve(L, q_mu)
        else:
            # alpha = L⁻ᵀ q_mu
            alpha = tf.linalg.triangular_solve(L, q_mu, adjoint=True)
        # predictive mean = Kfu alpha
        # predictive variance = Kff - Kfu Qinv Kuf
        # S = q_sqrt q_sqrtᵀ
        I = tf.eye(tf.shape(L)[-1], dtype=L.dtype)
        if isinstance(self._q_dist, _DeltaDist):
            B = I
        else:
            if not self.whiten:
                # Qinv = Kuu⁻¹ - Kuu⁻¹ S Kuu⁻¹
                #      = Kuu⁻¹ - L⁻ᵀ L⁻¹ S L⁻ᵀ L⁻¹
                #      = L⁻ᵀ (I - L⁻¹ S L⁻ᵀ) L⁻¹
                #      = L⁻ᵀ B L⁻¹
                if isinstance(self._q_dist, _DiagNormal):
                    q_sqrt = tf.linalg.diag(tf.linalg.adjoint(self._q_dist.q_sqrt))
                elif isinstance(self._q_dist, _MvNormal):
                    q_sqrt = self._q_dist.q_sqrt
                Linv_qsqrt = tf.linalg.triangular_solve(L, q_sqrt)
                Linv_cov_u_LinvT = tf.matmul(Linv_qsqrt, Linv_qsqrt, transpose_b=True)
            else:
                if isinstance(self._q_dist, _DiagNormal):
                    Linv_cov_u_LinvT = tf.linalg.diag(tf.linalg.adjoint(self._q_dist.q_sqrt ** 2))
                elif isinstance(self._q_dist, _MvNormal):
                    q_sqrt = self._q_dist.q_sqrt
                    Linv_cov_u_LinvT = tf.matmul(q_sqrt, q_sqrt, transpose_b=True)
                # Qinv = Kuu⁻¹ - L⁻ᵀ S L⁻¹
                # Linv = (L⁻¹ I) = solve(L, I)
                # Kinv = Linvᵀ @ Linv
            B = I - Linv_cov_u_LinvT
        LinvT_B = tf.linalg.triangular_solve(L, B, adjoint=True)
        B_Linv = tf.linalg.adjoint(LinvT_B)
        Qinv = tf.linalg.triangular_solve(L, B_Linv, adjoint=True)

        M, L = tf.unstack(tf.shape(self._q_dist.q_mu), num=2)
        Qinv = tf.broadcast_to(Qinv, [L, M, M])

        tf.debugging.assert_shapes(
            [
                (Qinv, ["L", "M", "M"]),
            ]
        )

        return alpha, Qinv

class IndependentPosterior(BasePosterior):
    def _post_process_mean_and_cov(self, mean, cov, full_cov, full_output_cov):
        return mean, expand_independent_outputs(cov, full_cov, full_output_cov)

    def _get_Kff(self, Xnew, full_cov):

        # TODO: this assumes that Xnew has shape [N, D] and no leading dims

        if isinstance(self.kernel, (kernels.SeparateIndependent, kernels.IndependentLatent)):
            # NOTE calling kernel(Xnew, full_cov=full_cov, full_output_cov=False) directly would return
            # if full_cov: [P, N, N] -- this is what we want
            # else: [N, P] instead of [P, N] as we get from the explicit stack below
            Kff = tf.stack([k(Xnew, full_cov=full_cov) for k in self.kernel.kernels], axis=0)
        elif isinstance(self.kernel, kernels.MultioutputKernel):
            # effectively, SharedIndependent path
            Kff = self.kernel.kernel(Xnew, full_cov=full_cov)
            # NOTE calling kernel(Xnew, full_cov=full_cov, full_output_cov=False) directly would return
            # if full_cov: [P, N, N] instead of [N, N]
            # else: [N, P] instead of [N]
        else:
            # standard ("single-output") kernels
            Kff = self.kernel(Xnew, full_cov=full_cov)  # [N, N] if full_cov else [N]

        return Kff

    def _conditional_with_precompute(
        self, Xnew, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        # Qinv: [L, M, M]
        # alpha: [M, L]

        Kuf = covariances.Kuf(self.X_data, self.kernel, Xnew)  # [(R), M, N]
        Kff = self._get_Kff(Xnew, full_cov)

        mean = tf.matmul(Kuf, self.alpha, transpose_a=True)
        if Kuf.shape.ndims == 3:
            mean = tf.linalg.adjoint(tf.squeeze(mean, axis=-1))

        if full_cov:
            Kfu_Qinv_Kuf = tf.matmul(Kuf, self.Qinv @ Kuf, transpose_a=True)
            cov = Kff - Kfu_Qinv_Kuf
        else:
            # [Aᵀ B]_ij = Aᵀ_ik B_kj = A_ki B_kj
            # TODO check whether einsum is faster now?
            Kfu_Qinv_Kuf = tf.reduce_sum(Kuf * tf.matmul(self.Qinv, Kuf), axis=-2)
            cov = Kff - Kfu_Qinv_Kuf
            cov = tf.linalg.adjoint(cov)

        return self._post_process_mean_and_cov(mean, cov, full_cov, full_output_cov)

class IndependentPosteriorSingleOutput(IndependentPosterior):
    # could almost be the same as IndependentPosteriorMultiOutput ...
    def _conditional_fused(
        self, Xnew, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        # same as IndependentPosteriorMultiOutput, Shared~/Shared~ branch, except for following line:
        Knn = self.kernel(Xnew, full_cov=full_cov)

        Kmm = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [M, M]
        Kmn = covariances.Kuf(self.X_data, self.kernel, Xnew)  # [M, N]

        fmean, fvar = base_conditional(
            Kmn, Kmm, Knn, self.q_mu, full_cov=full_cov, q_sqrt=self.q_sqrt, white=self.whiten
        )  # [N, P],  [P, N, N] or [N, P]
        return self._post_process_mean_and_cov(fmean, fvar, full_cov, full_output_cov)

class IndependentPosteriorMultiOutput(IndependentPosterior):
    def _conditional_fused(
        self, Xnew, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        if isinstance(self.X_data, SharedIndependentInducingVariables) and isinstance(
            self.kernel, kernels.SharedIndependent
        ):
            # same as IndependentPosteriorSingleOutput except for following line
            Knn = self.kernel.kernel(Xnew, full_cov=full_cov)
            # we don't call self.kernel() directly as that would do unnecessary tiling

            Kmm = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [M, M]
            Kmn = covariances.Kuf(self.X_data, self.kernel, Xnew)  # [M, N]

            fmean, fvar = base_conditional(
                Kmn, Kmm, Knn, self.q_mu, full_cov=full_cov, q_sqrt=self.q_sqrt, white=self.whiten
            )  # [N, P],  [P, N, N] or [N, P]
        else:
            # this is the messy thing with tf.map_fn, cleaned up by the st/clean_up_broadcasting_conditionals branch

            # Following are: [P, M, M]  -  [P, M, N]  -  [P, N](x N)
            Kmms = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [P, M, M]
            Kmns = covariances.Kuf(self.X_data, self.kernel, Xnew)  # [P, M, N]
            if isinstance(self.kernel, kernels.Combination):
                kernel_list = self.kernel.kernels
            else:
                kernel_list = [self.kernel.kernel] * len(self.X_data.inducing_variable_list)
            Knns = tf.stack(
                [k.K(Xnew) if full_cov else k.K_diag(Xnew) for k in kernel_list], axis=0
            )

            fmean, fvar = separate_independent_conditional_implementation(
                Kmns,
                Kmms,
                Knns,
                self.q_mu,
                q_sqrt=self.q_sqrt,
                full_cov=full_cov,
                white=self.whiten,
            )

        return self._post_process_mean_and_cov(fmean, fvar, full_cov, full_output_cov)

class LinearCoregionalizationPosterior(IndependentPosteriorMultiOutput):
    def _post_process_mean_and_cov(self, mean, cov, full_cov, full_output_cov):
        """
        mean: [N, L]
        cov: [L, N, N] or [N, L]
        """
        cov = expand_independent_outputs(cov, full_cov, full_output_cov=False)
        mean, cov = mix_latent_gp(self.kernel.W, mean, cov, full_cov, full_output_cov)
        return mean, cov

class FullyCorrelatedPosterior(BasePosterior):
    def _conditional_with_precompute(
        self, Xnew, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        # TODO: this assumes that Xnew has shape [N, D] and no leading dims

        # Qinv: [L, M, M]
        # alpha: [M, L]

        Kuf = covariances.Kuf(self.X_data, self.kernel, Xnew)
        assert Kuf.shape.ndims == 4
        M, L, N, K = tf.unstack(tf.shape(Kuf), num=Kuf.shape.ndims, axis=0)
        Kuf = tf.reshape(Kuf, (M * L, N * K))

        Kff = self.kernel(Xnew, full_cov=full_cov, full_output_cov=full_output_cov)
        # full_cov=True and full_output_cov=True: [N, P, N, P]
        # full_cov=True and full_output_cov=False: [P, N, N]
        # full_cov=False and full_output_cov=True: [N, P, P]
        # full_cov=False and full_output_cov=False: [N, P]
        if full_cov == full_output_cov:
            new_shape = (N * K, N * K) if full_cov else (N * K,)
            Kff = tf.reshape(Kff, new_shape)

        N = tf.shape(Xnew)[0]
        K = tf.shape(Kuf)[-1] // N

        mean = tf.matmul(Kuf, self.alpha, transpose_a=True)
        if Kuf.shape.ndims == 3:
            mean = tf.linalg.adjoint(tf.squeeze(mean, axis=-1))

        if not full_cov and not full_output_cov:
            # fully diagonal case in both inputs and outputs
            # [Aᵀ B]_ij = Aᵀ_ik B_kj = A_ki B_kj
            # TODO check whether einsum is faster now?
            Kfu_Qinv_Kuf = tf.reduce_sum(Kuf * tf.matmul(self.Qinv, Kuf), axis=-2)
        else:
            Kfu_Qinv_Kuf = tf.matmul(Kuf, self.Qinv @ Kuf, transpose_a=True)
            if not (full_cov and full_output_cov):
                # diagonal in either inputs or outputs
                new_shape = tf.concat([tf.shape(Kfu_Qinv_Kuf)[:-2], (N, K, N, K)], axis=0)
                Kfu_Qinv_Kuf = tf.reshape(Kfu_Qinv_Kuf, new_shape)
                if full_cov:
                    # diagonal in outputs: move outputs to end
                    tmp = tf.linalg.diag_part(tf.einsum("...ijkl->...ikjl", Kfu_Qinv_Kuf))
                elif full_output_cov:
                    # diagonal in inputs: move inputs to end
                    tmp = tf.linalg.diag_part(tf.einsum("...ijkl->...jlik", Kfu_Qinv_Kuf))
                Kfu_Qinv_Kuf = tf.einsum("...ijk->...kij", tmp)  # move diagonal dim to [-3]
        cov = Kff - Kfu_Qinv_Kuf

        if not full_cov and not full_output_cov:
            cov = tf.linalg.adjoint(cov)

        mean = tf.reshape(mean, (N, K))
        if full_cov == full_output_cov:
            cov_shape = (N, K, N, K) if full_cov else (N, K)
        else:
            cov_shape = (K, N, N) if full_cov else (N, K, K)
        cov = tf.reshape(cov, cov_shape)

        return mean, cov

    def _conditional_fused(
        self, Xnew, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        Kmm = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [M, L, M, L]
        Kmn = covariances.Kuf(self.X_data, self.kernel, Xnew)  # [M, L, N, P]
        Knn = self.kernel(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )  # [N, P](x N)x P  or  [N, P](x P)

        M, L, N, K = tf.unstack(tf.shape(Kmn), num=Kmn.shape.ndims, axis=0)
        Kmm = tf.reshape(Kmm, (M * L, M * L))

        if full_cov == full_output_cov:
            Kmn = tf.reshape(Kmn, (M * L, N * K))
            Knn = tf.reshape(Knn, (N * K, N * K)) if full_cov else tf.reshape(Knn, (N * K,))
            mean, cov = base_conditional(
                Kmn, Kmm, Knn, self.q_mu, full_cov=full_cov, q_sqrt=self.q_sqrt, white=self.whiten
            )  # [K, 1], [1, K](x NK)
            mean = tf.reshape(mean, (N, K))
            cov = tf.reshape(cov, (N, K, N, K) if full_cov else (N, K))
        else:
            Kmn = tf.reshape(Kmn, (M * L, N, K))
            mean, cov = fully_correlated_conditional(
                Kmn,
                Kmm,
                Knn,
                self.q_mu,
                full_cov=full_cov,
                full_output_cov=full_output_cov,
                q_sqrt=self.q_sqrt,
                white=self.whiten,
            )
        return mean, cov

class FallbackIndependentLatentPosterior(FullyCorrelatedPosterior):  # XXX
    def _conditional_fused(
        self, Xnew, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        Kmm = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [L, M, M]
        Kmn = covariances.Kuf(self.X_data, self.kernel, Xnew)  # [M, L, N, P]
        Knn = self.kernel(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )  # [N, P](x N)x P  or  [N, P](x P)

        return independent_interdomain_conditional(
            Kmn,
            Kmm,
            Knn,
            self.q_mu,
            full_cov=full_cov,
            full_output_cov=full_output_cov,
            q_sqrt=self.q_sqrt,
            white=self.whiten,
        )

get_posterior_class = Dispatcher("get_posterior_class")

@get_posterior_class.register(kernels.Kernel, InducingVariables)
def _get_posterior_base_case(kernel, inducing_variable):
    # independent single output
    return IndependentPosteriorSingleOutput

@get_posterior_class.register(kernels.MultioutputKernel, InducingPoints)
def _get_posterior_fully_correlated_mo(kernel, inducing_variable):
    return FullyCorrelatedPosterior

@get_posterior_class.register(
    (kernels.SharedIndependent, kernels.SeparateIndependent),
    (SeparateIndependentInducingVariables, SharedIndependentInducingVariables),
)
def _get_posterior_independent_mo(kernel, inducing_variable):
    # independent multi-output
    return IndependentPosteriorMultiOutput

@get_posterior_class.register(
    kernels.IndependentLatent,
    (FallbackSeparateIndependentInducingVariables, FallbackSharedIndependentInducingVariables),
)
def _get_posterior_independentlatent_mo_fallback(kernel, inducing_variable):
    return FallbackIndependentLatentPosterior

@get_posterior_class.register(
    kernels.LinearCoregionalization,
    (SeparateIndependentInducingVariables, SharedIndependentInducingVariables),
)
def _get_posterior_linearcoregionalization_mo_efficient(kernel, inducing_variable):
    # Linear mixing---efficient multi-output
    return LinearCoregionalizationPosterior

def create_posterior(
    kernel,
    inducing_variable,
    q_mu,
    q_sqrt,
    whiten,
    mean_function=None,
    precompute_cache: Union[PrecomputeCacheType, str, None] = PrecomputeCacheType.TENSOR,
):
    posterior_class = get_posterior_class(kernel, inducing_variable)
    precompute_cache = _validate_precompute_cache_type(precompute_cache)
    return posterior_class(
        kernel,
        inducing_variable,
        q_mu,
        q_sqrt,
        whiten,
        mean_function,
        precompute_cache=precompute_cache,
    )


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

"""This file is used for generator training under reinforcement learning framework.

It is implemented by integrating exploration strategy into REINFORCE algorithm.
The deep learning code is implemented by PyTorch ( >= version 1.0)
"""

import os

import click
import numpy as np
import torch
from rdkit import rdBase
from torch.utils.data import DataLoader, TensorDataset
from tqdm import trange

from drugex import model, util

def policy_gradient(agent, environ, explore=None, *, batch_size, mc, epsilon, baseline):
    """Training generator under reinforcement learning framework,
    The rewoard is only the final reward given by environment (predictor).

    agent (model.Generator): the exploitation network for SMILES string generation
    environ (util.Activity): the environment provide the final reward for each SMILES
    explore (model.Generator): the exploration network for SMILES string generation,
        it has the same architecture with the agent.
    """
    seqs = []

    # repeated sampling with MC times
    for _ in range(mc):
        seq = agent.sample(batch_size, explore=explore, epsilon=epsilon)
        seqs.append(seq)
    seqs = torch.cat(seqs, dim=0)
    ix = util.unique(seqs)
    seqs = seqs[ix]
    smiles, valids = util.check_smiles(seqs, agent.voc)

    # obtaining the reward
    preds = environ(smiles)
    preds[valids == False] = 0
    preds -= baseline
    preds = torch.Tensor(preds.reshape(-1, 1)).to(util.dev)

    ds = TensorDataset(seqs, preds)
    loader = DataLoader(ds, batch_size=batch_size)

    # Training Loop
    for seq, pred in loader:
        score = agent.likelihood(seq)
        agent.optim.zero_grad()
        loss = agent.PGLoss(score, pred)
        loss.backward()
        agent.optim.step()

def rollout_pg(agent, environ, explore=None, *, batch_size, baseline, mc, epsilon):
    """Training generator under reinforcement learning framework.

    The reward is given for each token in the SMILES, which is generated by
    Monte Carlo Tree Search based on final reward given by the environment.

    Arguments:

        agent (model.Generator): the exploitation network for SMILES string generation
        environ (util.Activity): the environment provide the final reward for each SMILES
        explore (model.Generator): the exploration network for SMILES string generation,
            it has the same architecture with the agent.
    """

    agent.optim.zero_grad()
    seqs = agent.sample(batch_size, explore=explore, epsilon=epsilon)
    batch_size = seqs.size(0)
    seq_len = seqs.size(1)
    rewards = np.zeros((batch_size, seq_len))
    smiles, valids = util.check_smiles(seqs, agent.voc)
    preds = environ(smiles) - baseline
    preds[valids == False] = - baseline
    scores, hiddens = agent.likelihood(seqs)

    # Monte Carlo Tree Search for step rewards generation
    for _ in trange(mc):
        for i in range(seq_len):
            if (seqs[:, i] != 0).any():
                h = hiddens[:, :, i, :]
                subseqs = agent.sample(batch_size, inits=(seqs[:, i], h, i + 1, None))
                subseqs = torch.cat([seqs[:, :i+1], subseqs], dim=1)
                subsmile, subvalid = util.check_smiles(subseqs, voc=agent.voc)
                subpred = environ(subsmile) - baseline
                subpred[1 - subvalid] = -baseline
            else:
                subpred = preds
            rewards[:, i] += subpred
    loss = agent.PGLoss(scores, seqs, torch.FloatTensor(rewards / mc))
    loss.backward()
    agent.optim.step()
    return 0, valids.mean(), smiles, preds

def _main_helper(*, epsilon, baseline, batch_size, mc, vocabulary_path, output_dir):
    #: Vocabulary containing all of the tokens for SMILES construction
    voc = util.Voc(vocabulary_path)
    #: File path of predictor in the environment
    environ_path = os.path.join(output_dir, 'RF_cls_ecfp6.pkg')
    #: file path of hidden states in RNN for initialization
    initial_path = os.path.join(output_dir, 'net_p.pkg')
    #: file path of hidden states of optimal exploitation network
    agent_path = os.path.join(output_dir, 'net_e_%.2f_%.1f_%dx%d' % (epsilon, baseline, batch_size, mc))
    #: file path of hidden states of exploration network
    explore_path = os.path.join(output_dir, 'net_p.pkg')

    # Environment (predictor)
    environ = util.Environment(environ_path)
    # Agent (generator, exploitation network)
    agent = model.Generator(voc)
    agent.load_state_dict(torch.load(initial_path))

    # exploration network
    explore = model.Generator(voc)
    explore.load_state_dict(torch.load(explore_path))

    best_score = 0
    log_file = open(agent_path + '.log', 'w')

    it = trange(1000)
    for epoch in it:
        it.write('\n--------\nEPOCH %d\n--------' % (epoch + 1))
        it.write('\nForward Policy Gradient Training Generator : ')
        policy_gradient(agent, environ, explore=explore,
                        baseline=baseline, batch_size=batch_size, mc=mc, epsilon=epsilon)
        seqs = agent.sample(1000)
        ix = util.unique(seqs)
        smiles, valids = util.check_smiles(seqs[ix], agent.voc)
        scores = environ(smiles)
        scores[valids == False] = 0
        unique = (scores >= 0.5).sum() / 1000
        # The model with best percentage of unique desired SMILES will be persisted on the hard drive.
        if best_score < unique:
            torch.save(agent.state_dict(), agent_path + '.pkg')
            best_score = unique
        print("Epoch+: %d average: %.4f valid: %.4f unique: %.4f" % (epoch, scores.mean(), valids.mean(), unique), file=log_file)
        for i, smile in enumerate(smiles):
            print('%f\t%s' % (scores[i], smile), file=log_file)

        # Learing rate exponential decay
        for param_group in agent.optim.param_groups:
            param_group['lr'] *= (1 - 0.01)
    log_file.close()

@click.command()
@click.option('-d', '--input-directory', type=click.Path(file_okay=False, dir_okay=True), show_default=True)
@click.option('-o', '--output-directory', type=click.Path(file_okay=False, dir_okay=True), show_default=True)
@click.option('--mc', type=int, default=10, show_default=True)
@click.option('--batch-size', type=int, default=500, show_default=True)
@click.option('--num-threads', type=int, default=1, show_default=True)
@click.option('-e', '--epsilon', type=float, default=0.1, show_default=True)
@click.option('-b', '--baseline', type=float, default=0.1, show_default=True)
@click.option('-g', '--cuda-visible-devices')
def main(input_directory, output_directory, mc, batch_size, num_threads, epsilon, baseline, cuda_visible_devices):
    rdBase.DisableLog('rdApp.error')
    torch.set_num_threads(num_threads)
    if cuda_visible_devices:
        os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices
    _main_helper(
        baseline=baseline,
        batch_size=batch_size,
        mc=mc,
        epsilon=epsilon,
        vocabulary_path=os.path.join(input_directory, "voc.txt"),
        output_dir=output_directory,
    )

if __name__ == "__main__":
    main()


# Derive optimal tolerance by ESM-SSP and  well-characterized errors
# for an audience seeking to emulate from the full CMIP6 archive,
# for a selection of ESMs.

# For each ESM, loop over various tolerances and generate Ndraws = 500
# GSAT trajectories. Archives with and without each target to characterize
# error. (Reproducible mode off so draws are different).
# Compare to the target ensemble via 4 metrics (E1, E2 on both Tgavs and
# jumps) to select the optimal tol for each ESM-SSP-archive combo.
# Use the resultant optimal tol for a final, reproducible mode set of recipes,
# Tgavs, and gridded data sets of monthly + daily for ??? variables.
# Start with tas, pr, psl - must subset archive up front to only include
# ensemble members that have all of these files available.
# look at CanESM5, MIROC6 and ACCESS-ESM1-5

## TODO functionalize at least some part of the analysis or at least use a for-loop
## over the different SSP targets so that the code isn't so long and repetitive
# making a table of avail runs X planned archives and for looping over that would
# trim things down (see approach for max tol runs). And rewrite the tolerance
# iteration to be a while loop, comparing current to prev instead of calculating
# and saving it all? Update writes and reads to be subdir so things are tidier

# would be better to functionalize this script with ESM, tol and Ndraws as arguments
# and then have the .sh just call the function and dispatch to diff nodes for each run I guess.

# #############################################################################
# General setup
# #############################################################################
# Import packages
import pandas as pd
import numpy as np
import stitches as stitches
import pkg_resources
import os
from  pathlib import Path

pd.set_option('display.max_columns', None)

OUTPUT_DIR = pkg_resources.resource_filename('stitches', 'data/created_data')
# OUTPUT_DIR = '/pic/projects/GCAM/stitches_pic/paper1_outputs'

# #############################################################################
# Experiment  setup
# #############################################################################
# experiment parameters
tolerances = np.round(np.arange(0.07, 0.225, 0.005), 3)
Ndraws = 10
error_threshold = 0.1

# pangeo table of ESMs for reference
pangeo_path = pkg_resources.resource_filename('stitches', 'data/pangeo_table.csv')
pangeo_data = pd.read_csv(pangeo_path)
pangeo_data = pangeo_data[((pangeo_data['variable'] == 'tas') | (pangeo_data['variable'] == 'pr') | (pangeo_data['variable'] == 'psl'))
                          & ((pangeo_data['domain'] == 'Amon') | (pangeo_data['domain'] == 'day')) ].copy()

# Keep only the runs that have data for all vars X all timesteps:
pangeo_good_ensembles =[]
for name, group in pangeo_data.groupby(['model', 'experiment', 'ensemble']):
    df = group.drop_duplicates().copy()
    if len(df) == 6:
        pangeo_good_ensembles.append(df)
    del(df)
pangeo_good_ensembles = pd.concat(pangeo_good_ensembles)
pangeo_good_ensembles  = pangeo_good_ensembles[['model', 'experiment', 'ensemble']].drop_duplicates().copy()
pangeo_good_ensembles = pangeo_good_ensembles.reset_index(drop=True).copy()

# won't use idealized runs
pangeo_good_ensembles = pangeo_good_ensembles[~((pangeo_good_ensembles['experiment'] == '1pctCO2') |
                                                (pangeo_good_ensembles['experiment'] == 'abrupt-4xCO2')) ].reset_index(drop=True).copy()

esms = ['ACCESS-ESM1-5', 'CanESM5', 'MIROC6']
# ['ACCESS-CM2', 'ACCESS-ESM1-5', 'AWI-CM-1-1-MR', 'BCC-CSM2-MR',
#        'BCC-ESM1', 'CESM2', 'CESM2-FV2', 'CESM2-WACCM', 'CMCC-CM2-HR4',
#        'CMCC-CM2-SR5', 'CMCC-ESM2', 'CanESM5', 'HadGEM3-GC31-LL',
#        'HadGEM3-GC31-MM', 'IITM-ESM', 'MIROC-ES2L', 'MIROC6',
#        'MPI-ESM-1-2-HAM', 'MPI-ESM1-2-HR', 'MPI-ESM1-2-LR', 'MRI-ESM2-0',
#        'NorESM2-LM', 'NorESM2-MM', 'SAM0-UNICON', 'TaiESM1',
#        'UKESM1-0-LL']

# #############################################################################
# Load full archive and target data
# #############################################################################

# Load the full archive of all staggered windows, which we will be matching on
full_archive_path = pkg_resources.resource_filename('stitches', 'data/matching_archive.csv')
full_archive_data = pd.read_csv(full_archive_path)

# Keep only the entries that appeared in pangeo_good_ensembles:
keys =['model', 'experiment', 'ensemble']
i1 = full_archive_data.set_index(keys).index
i2 = pangeo_good_ensembles.set_index(keys).index
full_archive_data= full_archive_data[i1.isin(i2)].copy()
del(i1)
del(i2)

# Load the original archive without staggered windows, which we will draw
# the target trajectories from for matching
full_target_path = pkg_resources.resource_filename('stitches', 'data/matching_archive.csv')
full_target_data = pd.read_csv(full_target_path)

# Keep only the entries that appeared in pangeo_good_ensembles:
keys =['model', 'experiment', 'ensemble']
i1 = full_target_data.set_index(keys).index
i2 = pangeo_good_ensembles.set_index(keys).index
full_target_data = full_target_data[i1.isin(i2)].copy()
del(i1)
del(i2)
del(keys)

# #############################################################################
# Some helper functions
# #############################################################################
def prep_target_data(target_df):
    if not target_df.empty:
        grped = target_df.groupby(['experiment', 'variable', 'ensemble', 'model'])
        for name, group in grped:
            df1 = group.copy()
            # if it isn't a complete time series (defined as going to 2099 or 2100),
            # remove it from the target data frame:
            if max(df1.end_yr) < 2099:
                target_df = target_df.loc[(target_df['ensemble'] != df1.ensemble.unique()[0])].copy().reset_index(
                    drop=True)
            del (df1)
        del (grped)

        target_df = target_df.reset_index(drop=True).copy()
        return(target_df)

def get_orig_data(target_df):
    if not target_df.empty:
        esm_name = target_df.model.unique()[0]
        scn_name = target_df.experiment.unique()[0]

        full_rawtarget_path = pkg_resources.resource_filename('stitches', ('data/tas-data/' + esm_name + '_tas.csv'))
        full_rawtarget_data = pd.read_csv(full_rawtarget_path)

        orig_data = full_rawtarget_data[(full_rawtarget_data['experiment'] == scn_name)].copy()
        keys = ['experiment', 'ensemble', 'model']
        i1 = orig_data.set_index(keys).index
        i2 = target_df.set_index(keys).index
        orig_data = orig_data[i1.isin(i2)].copy()
        del (i1)
        del (i2)
        del (keys)
        del (full_rawtarget_data)
        del (full_rawtarget_path)

        orig_data = orig_data.reset_index(drop=True).copy()
        return (orig_data)

def match_draw_stitchTgav(target_df, archive_df, toler, num_draws, TGAV_OUTPUT_DIR, reproducible):

    esm_name = archive_df.model.unique()[0]

    if not target_df.empty:
        # Use the match_neighborhood function to generate all of the matches between the target and
        # archive data points.
        match_df = stitches.match_neighborhood(target_df, archive_df, tol=toler)

        if target_df.experiment.unique() in archive_df.experiment.unique():
            archive_id = 'w_target'
        else:
            archive_id = 'wo_target'

        scn_name = target_df.experiment.unique()[0]

        for draw in range(0, num_draws):
            # Do the random draw of recipes
            if reproducible:
                unformatted_recipe = stitches.permute_stitching_recipes(N_matches=10000,
                                                                        matched_data=match_df,
                                                                        archive=archive_df,
                                                                        testing=True)
            else:
                unformatted_recipe = stitches.permute_stitching_recipes(N_matches=10000,
                                                                        matched_data=match_df,
                                                                        archive=archive_df,
                                                                        testing=False)

            new_ids = ('tol' + str(toler) + '~draw' + str(draw) + '~' + archive_id + '~'+
                       unformatted_recipe['stitching_id'].astype(str)).copy()
            unformatted_recipe = unformatted_recipe.drop(columns=['stitching_id']).copy()
            unformatted_recipe['stitching_id'] = new_ids
            del (new_ids)

            # format the recipe
            recipe = stitches.generate_gridded_recipe(unformatted_recipe)
            recipe.columns = ['target_start_yr', 'target_end_yr', 'archive_experiment', 'archive_variable',
                              'archive_model', 'archive_ensemble', 'stitching_id', 'archive_start_yr',
                              'archive_end_yr', 'tas_file']
            recipe['tolerance'] = toler
            recipe['draw'] = draw
            recipe['archive'] = archive_id
            recipe.to_csv((OUTPUT_DIR + '/' + esm_name + '/experiment_CMIP6/' +
                           'gridded_recipes_' + esm_name + '_target_' + scn_name +
                           '_tol' + str(toler) +
                           '_draw' + str(draw) +
                           '_archive_' + archive_id + '.csv'), index=False)
            del (unformatted_recipe)

            # stitch the GSAT values and save as csv
            try:

                gsat = stitches.gmat_stitching(recipe)
                gsat['tolerance'] = toler
                gsat['draw'] = draw
                gsat['archive'] = archive_id
                for id in gsat.stitching_id.unique():
                    ds = gsat[gsat['stitching_id'] == id].copy()
                    fname = (TGAV_OUTPUT_DIR +
                             'stitched_' + esm_name + '_GSAT_' + id + '.csv')
                    ds.to_csv(fname, index=False)
                    del (ds)

                del (gsat)

            except:
                print(("Some issue stitching GMAT for " + esm_name + ". Skipping and moving on"))

    else:
        recipe = []
        print('Some missing target data for ' + esm_name + '. Analysis will be skipped')

    return(recipe)

def get_jumps(tgav_df):
    tgav_jump = []
    for name, group in tgav_df.groupby(['variable', 'experiment', 'ensemble', 'model']):
        ds = group.copy()
        ds['jump'] = ds.value.diff().copy()
        ds = ds.dropna().copy()
        tgav_jump.append(ds)
        del (ds)
    tgav_jump = pd.concat(tgav_jump)
    tgav_jump = tgav_jump.drop(columns=['value']).copy()
    tgav_jump = tgav_jump.drop_duplicates().reset_index(drop=True).copy()
    return(tgav_jump)

def four_errors(gen_data, orig_data):
    gen_data_jump = get_jumps(gen_data)
    orig_data_jump = get_jumps(orig_data)

    orig_stats = []
    for name, group in orig_data.groupby(['model', 'variable', 'experiment']):
        ds = group.copy()
        ds1 = ds[['model', 'variable', 'experiment']].drop_duplicates().copy()
        ds1['mean_orig_tgav'] = np.mean(ds.value.values)
        ds1['sd_orig_tgav'] = np.std(ds.value.values)
        orig_stats.append(ds1)
        del (ds)
        del (ds1)
    orig_stats = pd.concat(orig_stats).reset_index(drop=True).copy()

    orig_stats_jump = []
    for name, group in orig_data_jump.groupby(['model', 'variable', 'experiment']):
        ds = group.copy()
        ds1 = ds[['model', 'variable', 'experiment']].drop_duplicates().copy()
        ds1['mean_orig_jump'] = np.mean(ds.jump.values)
        ds1['sd_orig_jump'] = np.std(ds.jump.values)
        orig_stats_jump.append(ds1)
        del (ds)
        del (ds1)
    orig_stats_jump = pd.concat(orig_stats_jump).reset_index(drop=True).copy()

    orig_stats = orig_stats.merge(orig_stats_jump, how='left', on=['model', 'variable', 'experiment']).copy()
    del (orig_stats_jump)

    gen_stats = []
    for name, group in gen_data.groupby(['model', 'variable', 'experiment', 'tolerance', 'draw', 'archive']):
        ds = group.copy()
        ds1 = ds[['model', 'variable', 'experiment', 'tolerance', 'draw', 'archive']].drop_duplicates().copy()
        ds1['mean_gen_tgav'] = np.mean(ds.value.values)
        ds1['sd_gen_tgav'] = np.std(ds.value.values)
        gen_stats.append(ds1)
        del (ds)
        del (ds1)
    gen_stats = pd.concat(gen_stats).reset_index(drop=True).copy()

    gen_stats_jump = []
    for name, group in gen_data_jump.groupby(['model', 'variable', 'experiment', 'tolerance', 'draw', 'archive']):
        ds = group.copy()
        ds1 = ds[['model', 'variable', 'experiment', 'tolerance', 'draw', 'archive']].drop_duplicates().copy()
        ds1['mean_gen_jump'] = np.mean(ds.jump.values)
        ds1['sd_gen_jump'] = np.std(ds.jump.values)
        gen_stats_jump.append(ds1)
        del (ds)
        del (ds1)
    gen_stats_jump = pd.concat(gen_stats_jump).reset_index(drop=True).copy()

    gen_stats = gen_stats.merge(gen_stats_jump, how='left',
                                on=['model', 'variable', 'experiment', 'tolerance', 'draw', 'archive']).copy()
    del (gen_stats_jump)

    compare = gen_stats.merge(orig_stats, how='left', on=['model', 'variable', 'experiment']).copy()
    del (gen_stats)
    del (orig_stats)

    compare['E1_tgav'] = abs(compare.mean_orig_tgav - compare.mean_gen_tgav) / compare.sd_orig_tgav
    compare['E2_tgav'] = compare.sd_gen_tgav / compare.sd_orig_tgav

    compare['E1_jump'] = abs(compare.mean_orig_jump - compare.mean_gen_jump) / compare.sd_orig_jump
    compare['E2_jump'] = compare.sd_gen_jump / compare.sd_orig_jump

    compare = compare[['model', 'variable', 'experiment', 'tolerance', 'draw', 'archive',
                       'E1_tgav', 'E2_tgav', 'E1_jump', 'E2_jump']].copy()

    four_values = []
    for name, group in compare.groupby(['model', 'variable', 'experiment', 'tolerance', 'draw', 'archive']):
        ds = group.copy()
        ds['max_metric'] = np.max(
            [ds.E1_tgav.values, abs(1 - ds.E2_tgav.values), ds.E1_jump.values, abs(1 - ds.E2_jump.values)])
        four_values.append(ds)
        del (ds)
    four_values = pd.concat(four_values).reset_index(drop=True).copy()
    del (compare)

    return(four_values)

# #############################################################################
# The experiment
# #############################################################################

# for each of the esms in the experiment, subset to what we want
# to work with and run the experiment.
for esm in esms:
    print(esm)

    # subset the archive and the targets to this ESM
    archive_w_all = full_archive_data[(full_archive_data['model'] == esm)].copy()

    archive_wo245 = full_archive_data[(full_archive_data['model'] == esm) &
                                      (full_archive_data['experiment'] != 'ssp245')].copy()

    archive_wo370 = full_archive_data[(full_archive_data['model'] == esm) &
                                      (full_archive_data['experiment'] != 'ssp370')].copy()

    target_245 = full_target_data[(full_target_data['model'] == esm) &
                                  (full_target_data['experiment'] == 'ssp245')].copy()

    target_370 = full_target_data[(full_target_data['model'] == esm) &
                                  (full_target_data['experiment'] == 'ssp370')].copy()

    # Clean up target data and pull corresponding original/raw data
    if not target_245.empty:
        # clean up
        target_245 = prep_target_data(target_245).copy()

        # and pull corresponding original/raw data for later comparison
        orig_245 = get_orig_data(target_245).copy()

    if not target_370.empty:
        # clean up
        target_370 = prep_target_data(target_370).copy()

        # and pull corresponding original/raw data for later comparison
        orig_370 = get_orig_data(target_370).copy()

    # loop over tolerances:
    for tolerance in tolerances:

        rp_245_w = match_draw_stitchTgav(target_245, archive_w_all,
                                         toler=tolerance, num_draws=Ndraws,
                                         TGAV_OUTPUT_DIR=(OUTPUT_DIR + '/' + esm + '/experiment_CMIP6/tolerance-generated-tgavs/' ),
                                         reproducible=False)
        rp_245_wo = match_draw_stitchTgav(target_245, archive_wo245,
                                         toler=tolerance, num_draws=Ndraws,
                                         TGAV_OUTPUT_DIR=(OUTPUT_DIR + '/' + esm + '/experiment_CMIP6/tolerance-generated-tgavs/' ),
                                         reproducible=False)
        rp_370_w = match_draw_stitchTgav(target_370, archive_w_all,
                                         toler=tolerance, num_draws=Ndraws,
                                         TGAV_OUTPUT_DIR=(OUTPUT_DIR + '/' + esm + '/experiment_CMIP6/tolerance-generated-tgavs/' ),
                                         reproducible=False)
        rp_370_wo = match_draw_stitchTgav(target_370, archive_wo370,
                                         toler=tolerance, num_draws=Ndraws,
                                         TGAV_OUTPUT_DIR=(OUTPUT_DIR + '/' + esm + '/experiment_CMIP6/tolerance-generated-tgavs/' ),
                                         reproducible=False)

    #########################################################
    # Now we've generated all the GSAT files we're going to for each target.
    # It's time to compare to the raw ensemble statistics for both Tgav and jumps.

    if (((not orig_245.empty) & (not target_245.empty)) |
            ((not orig_370.empty) & (not target_370.empty))):

        # form the raw ensemble data frame
        if orig_245.empty:
            orig = orig_370.reset_index(drop=True).copy()
        elif orig_370.empty:
            orig = orig_245.reset_index(drop=True).copy()
        else:
            orig = orig_245.append(orig_370).reset_index(drop=True).copy()

        # Read in all generated GSAT files and format so error metrics can
        # be calculated.
        gen = []
        entries = Path((OUTPUT_DIR + '/' + esm + '/experiment_CMIP6/tolerance-generated-tgavs/'))
        for entry in entries.iterdir():
            if (('stitched' in entry.name) & ('GSAT' in entry.name)):
                data = pd.read_csv((OUTPUT_DIR + '/' + esm + '/experiment_CMIP6/tolerance-generated-tgavs/') + entry.name)
                data['model'] = esm

                if ('ssp245' in entry.name):
                    data['experiment'] = 'ssp245'

                if ('ssp370' in entry.name):
                    data['experiment'] = 'ssp370'

                gen.append(data)
                del (data)
        gen = pd.concat(gen).reset_index(drop=True).copy()
        gen = gen.rename(columns={"stitching_id": "ensemble"}).copy()

        compared_data = four_errors(gen_data=gen, orig_data=orig)
        compared_data.to_csv((OUTPUT_DIR + '/' + esm + '/experiment_CMIP6/all_metrics.csv'), index=False)

        # average over draws
        aggregate_metrics = []
        for name, group in compared_data.groupby(['model', 'variable', 'experiment', 'tolerance', 'archive']):
            ds = group.copy()
            ds1 = ds[['model', 'variable', 'experiment', 'tolerance', 'archive']].drop_duplicates().copy()
            ds1['aggregate_E1_tgav'] = np.mean(ds.E1_tgav.values)
            ds1['aggregate_E2_tgav'] = np.mean(ds.E2_tgav.values)
            ds1['aggregate_E1_jump'] = np.mean(ds.E1_jump.values)
            ds1['aggregate_E2_jump'] = np.mean(ds.E2_jump.values)
            ds1['max_metric'] = np.max([ds1.aggregrate_E1_tgav.values,
                                        abs(1 - ds1.aggregate_E2_tgav.values),
                                        ds1.aggregate_E1_jump.values,
                                        abs(1 - ds1.aggregate_E2_jump.values)])
            aggregate_metrics.append(ds1)
            del (ds)
            del (ds1)
        aggregate_metrics = pd.concat(aggregate_metrics).reset_index(drop=True).copy()

        # filter to the largest tolerance that keeps max_metric<error_threshold
        # for each model, variable, experiment, archive.
        max_tol = []
        for name, group in aggregate_metrics.groupby(['model', 'variable', 'experiment', 'archive']):
            ds = group.copy()
            ds = ds[ds["max_metric"] < error_threshold].copy()
            ds = ds[ds['tolerance'] == ds.tolerance.max()].copy()
            ds = ds.rename(columns={"tolerance": "max_tol"}).copy()
            ds = ds[['model', 'variable', 'experiment', 'archive', 'max_tol']].copy()
            max_tol.append(ds)
            del (ds)
        max_tol = pd.concat(max_tol).reset_index(drop=True).copy()
        max_tol.to_csv((OUTPUT_DIR + '/' + esm + '/experiment_CMIP6/max_tol_by_ESM_SSP.csv'), index=False)

    if not max_tol.empty:
        # now that we have max tolerances, we will use those to do a final,
        # reproducible draw and construction of gridded data
        for name, group in max_tol.groupby(['model', 'variable', 'experiment', 'archive']):
            tolerance = group.max_tol.unique()[0]
            arch_id = group.archive.unique()[0]
            targ_id = group.experiment.unique()[0]

            if targ_id == 'ssp245':
                target = target_245.copy()
                if group.archive.unique == 'wo_target':
                    archive = archive_wo245.copy()
                else:
                    archive = archive_w_all.copy()

            if targ_id == 'ssp370':
                target = target_370.copy()
                if group.archive.unique == 'wo_target':
                    archive = archive_wo370.copy()
                else:
                    archive = archive_w_all.copy()

            match_df = stitches.match_neighborhood(target, archive,
                                                   tol=tolerance)

            unformatted_recipe = stitches.permute_stitching_recipes(N_matches=10000,
                                                                    matched_data=match_df,
                                                                    archive=archive, testing=True)
            new_ids = ('~tol' + str(tolerance) + '~archive_' + arch_id + '~' +
                       unformatted_recipe['stitching_id'].astype(str)).copy()
            unformatted_recipe = unformatted_recipe.drop(columns=['stitching_id']).copy()
            unformatted_recipe['stitching_id'] = new_ids
            del (new_ids)

            recipe = stitches.generate_gridded_recipe(unformatted_recipe)
            recipe.columns = ['target_start_yr', 'target_end_yr', 'archive_experiment', 'archive_variable',
                              'archive_model', 'archive_ensemble', 'stitching_id', 'archive_start_yr',
                              'archive_end_yr', 'tas_file']
            recipe.to_csv((OUTPUT_DIR + '/' + esm + '/experiment_CMIP6/max-tol-runs/' +
                               'gridded_recipes_' + esm + '_target'+
                           targ_id +'_archive_' + arch_id +'.csv'), index=False)

            gsat = stitches.gmat_stitching(recipe)
            gsat['tolerance'] = tolerance
            gsat['archive'] = arch_id
            for id in gsat.stitching_id.unique():
                ds = gsat[gsat['stitching_id'] == id].copy()
                fname = (OUTPUT_DIR + '/' + esm + '/experiment_CMIP6/max-tol-runs/' +
                         'stitched_' + esm + '_GSAT_' + id + '.csv')
                ds.to_csv(fname, index=False)

                # for single_id in recipe['stitching_id'].unique():
                #     single_rp = recipe.loc[recipe['stitching_id'] == single_id].copy()
                #     outputs = stitches.gridded_stitching((OUTPUT_DIR + '/' + esm + '/experiment_CMIP6'),
                #                                          single_rp)
                #     del (single_rp)
                #     del(outputs)

# end for loop over ESMs


# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
ssobj --- Solar System objects.
===============================

   Classes
   -------
   SolarSysObject

   Functions
   ---------
   getgeom
   getspiceobj
   getxyz
   summarizegeom

"""

import numpy as np
import astropy.units as u
from astropy.time import Time

from . import core

__all__ = [
    'SolarSysObject',
    'getgeom',
    'getspiceobj',
    'getxyz',
    'summarizegeom',
]

class SolarSysObject:
    """A star, planet, comet, etc. in the Solar System.

    Parameters
    ----------
    state : State
      The object from which to retrieve positions and velocities.
    M : float or Quantity, optional
      Mass of the object.  [float: kg]
    GM : float or Quantity, optional
      Gravitational constant times the mass of the object. [float:
      km**3/s**2]
    name : string, optional
      The name of the object.

    Methods
    -------
    ephemeris : Ephemeris for an observer.
    fluxd : Total flux density as seen by an observer.
    lightcurve : An ephemeris table with fluxes.
    observe : Distance, phase angle, etc. to another object.
    orbit : Osculating orbital parameters at date.
    r : Position vector
    v : Velocity vector
    h : Angular momentum integral, `r × v`

    Notes
    -----
    Inheriting classes should override `fluxd`.

    """

    def __init__(self, state, M=None, GM=None, name=None):
        from .state import State
        assert isinstance(state, State)
        self.state = state
        self.name = name

        if M is None and GM is None:
            self.M = 0
        elif M is not None:
            self.M = M
        else:
            self.GM = GM

    @property
    def M(self):
        """Mass.  [kg]"""
        return self._GM / 6.67384e-20

    @M.setter
    def M(self, m):
        m = u.Quantity(m, u.kg)
        self._GM = 6.67384e-20 * m.value

    @property
    def GM(self):
        """Gravitational constant times mass. [kg**3 / s**2]"""
        return self._GM

    @GM.setter
    def GM(self, gm):
        gm = u.Quantity(gm, u.km**3 / u.s**2)
        self._GM = gm.value

    def __repr__(self):
        return '<SolarSysObject name="{}" M="{:.3g} kg">'.format(
            self.name, self.M)

    def r(self, date):
        """Position vector.

        Parameters
        ----------
        date : string, float, astropy Time, datetime, or array
          Processed via `util.date2time`.

        Returns
        -------
        r : ndarray
          Position vector (3-element or Nx3 element array). [km]

        """
        return self.state.r(date)

    def v(self, date):
        """Velocity vector.

        Parameters
        ----------
        date : string, float, astropy Time, datetime, or array
          Processed via `util.date2time`.

        Returns
        -------
        v : ndarray
          Velocity vector (3-element or Nx3 element array). [km/s]

        """
        return self.state.v(date)

    def rv(self, date):
        """Position and velocity vectors.

        Parameters
        ----------
        date : string, float, astropy Time, datetime, or array
          Processed via `util.date2time`.

        Returns
        -------
        r : ndarray
          Position vector (3-element or Nx3 element array). [km]
        v : ndarray
          Velocity vector (3-element or Nx3 element array). [km/s]

        """
        return self.state.rv(date)

    def h(self, date):
        """Angular momentum integral, `r × v`.

        Parameters
        ----------
        date : string, float, astropy Time, datetime, or array
          Processed via `util.date2time`.

        Returns
        -------
        h : ndarray
          3-element or Nx3 element array.  [km2/s]

        """
        return self.state.h(date)

    def ephemeris(self, observer, dates, num=None, columns=None,
                  cformats=None, ra_unit='hourangle', date_format=None,
                  ltt=False, **kwargs):
        """Ephemeris for an observer.

        Parameters
        ----------
        observer : SolarSysObject
          The observer.
        dates : array
          If `num` is `None`, set `dates` to a list of exact times for
          the ephemeris, otherwise, let `dates` be a start and stop
          time, and `num` be the number of dates to generate.  `dates`
          may be in any format that `observe` accepts.
        num : int, optional
          If not `None`, generate this many time steps between
          `min(dates)` and `max(dates)`.
        columns : array, optional
          A list of `Geom` keywords to use as table columns, or `None`
          for the default list.
        cformats : dict or list, optional
          A dictionary of formats with keys corresponding to
          `columns`.
        ra_unit : str or astropy Unit, optional
          The unit of Right Ascention output, e.g., hourangle or deg.
        date_format : function
          A function to format the `date` column before creating the
          table.
        ltt : bool, optional
          Set to `True` to account for light travel time.

        Returns
        -------
        eph : astropy Table

        Notes
        -----
        `date_format` should be removed when astropy `Table` can
        handle Time objects.

        Override `_add_lc_columns` to add additional columns to
        lightcurve output.

        """

        from astropy.table import Table, Column
        from astropy.units import Quantity
        from ..util import date2time, dh2hms

        dates = date2time(dates)
        if num is not None:
            if num <= 0:
                time = []
            elif num == 1:
                time = [dates[0], dates[-1]]
            else:
                step = (dates[-1] - dates[0]) / float(num - 1)
                time = []
                for i in range(num):
                    time += [dates[0] + step * i]
        else:
            time = dates

        g = observer.observe(self, time, ltt=ltt)

        if columns is None:
            columns = ['date', 'ra', 'dec', 'rh', 'delta', 'phase', 'selong']

        if cformats is None:
            cformats = dict()

        _cformats = dict(
            date = '{:s}',
            ra = lambda x: dh2hms(x, "{:2d}:{:02d}"),
            dec = lambda x: dh2hms(x, "{:2d}:{:02d}"),
            lam = '{:.0f}',
            bet = '{:+.0f}',
            rh = '{:.3f}',
            delta = '{:.3f}',
            phase = '{:.0f}',
            selong = '{:.0f}',
            lelong = '{:.0f}')

        for k, v in cformats:
            _cformats[k] = v

        if date_format is None:
            date_format = lambda d: d.iso[:-7]

        eph = Table()
        eph.meta['ltt'] = ltt
        #eph.meta['observer'] = str(self)
        #eph.meta['target'] = str(target)

        for c in columns:
            if c == 'date':
                data = [date_format(d) for d in g[c]]
            elif c == 'ra':
                data = g[c].to(ra_unit)
            else:
                data = g[c]

            if c in _cformats:
                cf = _cformats[c]
            else:
                cf = None

            eph.add_column(Column(data=data, name=c, format=cf))

        return eph

    def fluxd(self, observer, date, wave, ltt=False,
              unit=u.Unit('W / (m2 um)'), **kwargs):
        """Total flux density as seen by an observer.

        Parameters
        ----------
        observer : SolarSysObject
          The observer.
        date : string, float, astropy Time, datetime
          The time of the observation in any format acceptable to
          `observer`.
        wave : Quantity
          The wavelengths at which to compute `fluxd`.
        ltt : bool, optional
          Set to `True` to correct the object's position for light
          travel time.
        unit : astropy Unit
          The return unit, must be flux density.

        Returns
        -------
        fluxd : Quantity

        """
        raise NotImplemented('This class has not implemented fluxd.')

    def lightcurve(self, observer, dates, wave, wformat='f{:.1f}',
                   verbose=True, **kwargs):
        """An ephemeris table with fluxes.

        Parameters
        ----------
        observer : SolarSysObject
          The observer.
        dates : string, float, astropy Time, datetime
          The dates of the observation in any format acceptable to
          `SolarSysObservable.ephemeris`.
        wave : Quantity
          The wavelengths at which to compute `fluxd`.
        wformat : string, optional
          The flux density columns will are labed with their
          wavelengths using this format string.
        verbose : bool, optional
          If `True`, give some visual feedback.
        **kwargs
          Additional `fluxd` and `ephemeris` keywords.

        Returns
        -------
        lc : astropy Table

        Notes
        -----
        `date` must be in the table returned by `ephemeris`.

        The `cformat` keyword `fluxd` will be used to format all flux
        density columns.

        """

        from astropy.table import Column

        lc = self.ephemeris(observer, dates, **kwargs)
        if 'date' not in lc.columns:
            raise KeyError("Ephemeris must return a date column."
                           "  Update columns.")

        if not np.iterable(wave):
            wave = [wave.value] * wave.unit

        if verbose:
            print('Computing date:')

        fluxd = np.zeros((len(lc), len(wave)))
        for i, d in enumerate(lc['date']):
            if verbose:
                print(d)
            f = self.fluxd(observer, d, wave, **kwargs)
            fluxd[i] = f.value
            unit = f.unit

        if verbose:
            print()

        lc = self._add_lc_columns(lc)

        cf = kwargs.get('cformat', dict()).get('fluxd', '{:9.3g}')
        unit = str(unit)
        for i in range(fluxd.shape[1]):
            lc.add_column(Column(data=fluxd[:, i],
                                 name=wformat.format(wave.value[i]),
                                 format=cf, unit=unit))
        return lc

    def _add_lc_columns(self, lc):
        """Add additional columns to a lightcurve table.

        Parameters
        ----------
        lc : Table
          The current lightcurve table.

        """
        return lc

    def observe(self, target, date, ltt=False):
        """Distance, phase angle, etc. to another object.

        Parameters
        ----------
        target : SolarSysObject
          The target to observe.
        date : string, float, astropy Time, datetime, or array
          Processed via `util.date2time`.
        ltt : bool, optional
          Account for light travel time when `True`.

        Returns
        -------
        geom : Geom
          The geometric parameters of the observation.

        """

        from astropy.time import TimeDelta
        import astropy.constants as const
        from . import Geom
        from ..util import date2time

        date = date2time(date)

        rt = target.r(date) # target postion
        ro = self.r(date)   # observer position
        vt = target.v(date) # target velocity
        vo = self.v(date)   # observer velocity

        g = Geom(ro * u.km, rt * u.km,
                 vo=vo * u.km / u.s, vt=vt * u.km / u.s,
                 date=date)

        if ltt:
            dt = (g['delta'] / const.c.si).decompose().value
            date -= TimeDelta(dt, format='sec')
            g = self.observe(target, date, ltt=False)

        return g

    def orbit(self, date):
        """Osculating orbital elements.

        Parameters
        ----------
        date : string, float, astropy Time, datetime, or array
          Processed via `util.date2time`.

        Returns
        -------
        orbit : dict
          See `util.state2orbit`.

        """

        from ..util import state2orbit, date2time
        r = self.r(date)
        v = self.v(date)
        jd = date2time(date).jd
        return state2orbit(r, v)

def getgeom(target, observer, date=None, ltt=False, kernel=None):
    """Moving target geometry parameters for an observer and date.

    Parameters
    ----------
    target : string, SolarSysObject
      The object's name or NAIF ID, as found in the relevant SPICE
      kernel, or a `SolarSysObject`.
    observer : string, array, SolarSysObject
      A valid built-in observer name, set of heliocentric rectangular
      ecliptic J2000 coordinates, or a `SolarSysObject`.  See the
      `ephem` package documentation for built-in moving objects that
      can be used as observers.
    date : string, float or array, optional
      The date(s) for which to compute the target's geometry.
    ltt : bool, optional
      Set to true to correct parameters for light travel time
      (currently, only one ltt iteration is implemented).
    kernel : string, optional
      If the target or observer is a string, use this kernel.

    Returns
    -------
    geom : Geom
      The geometric parameters of the observation.

    Raises
    ------
    ValueError on invalid input.

    """

    from . import _loaded_objects
    from .state import SpiceState

    if isinstance(target, str):
        target = SolarSysObject(SpiceState(target, kernel=kernel))
    elif not isinstance(target, SolarSysObject):
        raise ValueError("target must be a string or SolarSysObject.")

    if isinstance(observer, str):
        if observer.lower() in _loaded_objects:
            observer = _loaded_objects[observer.lower()]
        else:
            observer = SolarSysObject(SpiceState(target, kernel=kernel))
    elif np.iterable(observer):
        observer = FixedObject(observer)
    elif not isinstance(observer, SolarSysObject):
        raise ValueError("observer must be a string, array, or SolarSysObject")

    return observer.observe(target, date, ltt=ltt)

def getspiceobj(obj, kernel=None, name=None, **kwarg):
    """Create a new SolarSysObject with a SPICE kernel, for your convenience.

    Parameters
    ----------
    obj : string or int
      The object's name or NAIF ID, as found in the relevant SPICE
      kernel.
    kernel : string, optional
      The name of a specific SPICE planetary ephemeris kernel (SPK) to
      use for this object, or `None` to automatically search for a
      kernel through `find_kernel`.
    name : string
      The name of the object, or `None` to use `obj`.
    **kwarg
      Any other `SolarSysObject` keyword argument.

    Returns
    -------
    ssobj : SolarSysObject
      A `SolarSysObject` loaded with the requested SPICE ephemeris
      file.

    """

    from .state import SpiceState
    name = str(obj) if name is None else name
    return SolarSysObject(SpiceState(obj, kernel=kernel), name=name, **kwarg)

def getxyz(obj, date=None, kernel=None):
    """Coordinates and velocity from an ephemeris kernel.

    Coordinates are heliocentric rectangular ecliptic J2000.

    Parameters
    ----------
    obj : string or int
      The object's name or NAIF ID, as found in the relevant SPICE
      kernel.
    date : string, float, astropy Time, datetime, or array, optional
      Processed via `util.date2time`.
    kernel : string, optional
      The name of a specific SPICE planetary ephemeris kernel (SPK) to
      use for this object, or `None` to automatically search for a
      kernel through `find_kernel`.

    Returns
    -------
    r, v: array
      The position and veloctiy vectors as 3-element arrays, or, if
      multiple dates are requested, Nx3-element arrays. [km and km/s]

    Raises
    ------
    ValueError for invalid `date`.
    ObjectError for `obj` not in `kernel`.

    """

    from .state import SpiceState
    obj = SpiceState(obj, kernel=kernel)
    return obj.r(date), obj.v(date)

def summarizegeom(*args, **kwargs):
    """Pretty print a summary of the observing geometry.

    See `getgeom` for a description of the parameters.

    """

    getgeom(*args, **kwargs).summary()

# update module docstring
from ..util import autodoc
autodoc(globals())
del autodoc

import numpy,import astropy,from astropy