import torch
from torchvision.utils import save_image
import torchvision.models as models
import os
import time
import ipdb
import sys
sys.path.append("..")
import utils, models
import matplotlib.pyplot as plt
plt.switch_backend('agg')
import numpy as np

from model import Generator, DCGAN
from model import CifarDiscriminator, Discriminator
from model import Inverter
from classifier import LeNet

class Solver(object):
    def __init__(self, train_loader, test_loader, config):
        # Data loader.
        self.train_loader = train_loader
        self.test_loader = test_loader

        self.config = config

        # Model configurations.
        self.z_dim = config.z_dim
        self.image_dim = config.image_dim
        self.conv_dim = config.conv_dim
        self.lambda_gp = config.lambda_gp
        self.lambda_i = config.lambda_i

        # Training configurations.
        self.dataset = config.dataset
        self.batch_size = config.batch_size
        self.num_iters = config.num_iters
        self.dis_iters = config.dis_iters # discriminator iteration
        self.resume_iters = config.resume_iters
        self.g_lr = config.g_lr
        self.i_lr = config.i_lr
        self.d_lr = config.d_lr
        self.beta1 = config.beta1
        self.beta2 = config.beta2
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        # Testing configurations.
        self.classifier = config.classifier
        self.cla_iters = config.cla_iters
        self.search = config.search
        self.n_samples = config.n_samples
        self.step = config.step

        # Directories.
        self.model_dir = config.model_dir
        self.sample_dir = config.sample_dir
        self.cla_dir = config.cla_dir
        self.adversary_dir = config.adversary_dir

        # Step size.
        self.log_step = config.log_step
        self.model_step = config.model_step
        self.sample_step = config.sample_step

        # Build the model.
        self.build_model()

    def build_model(self):
        self.G = DCGAN(self.config.batch_size,nz=100)
        self.I = Inverter(self.z_dim, self.image_dim, self.conv_dim)
        self.D = CifarDiscriminator()
        self.g_optimizer = torch.optim.Adam(self.G.parameters(), self.g_lr, [self.beta1, self.beta2])
        self.i_optimizer = torch.optim.Adam(self.I.parameters(), self.i_lr, [self.beta1, self.beta2])
        self.d_optimizer = torch.optim.Adam(self.D.parameters(), self.d_lr, [self.beta1, self.beta2])

        self.G.to(self.device)
        self.I.to(self.device)
        self.D.to(self.device)

    def load_model(self, resume_iters):
        print('Loading the trained model from step {}...'.format(resume_iters))
        G_path = os.path.join(self.model_dir, '{}_G.ckpt'.format(resume_iters))
        I_path = os.path.join(self.model_dir, '{}_I.ckpt'.format(resume_iters))
        D_path = os.path.join(self.model_dir, '{}_D.ckpt'.format(resume_iters))
        self.G.load_state_dict(torch.load(G_path, map_location=lambda storage, loc: storage))
        self.I.load_state_dict(torch.load(I_path, map_location=lambda storage, loc: storage))
        self.D.load_state_dict(torch.load(D_path, map_location=lambda storage, loc: storage))

    def reset_grad(self):
        self.d_optimizer.zero_grad()
        self.i_optimizer.zero_grad()
        self.g_optimizer.zero_grad()

    def denorm(self, x):
        """Convert the range from [-1, 1] to [0, 1]."""
        out = (x + 1) / 2
        return out.clamp_(0, 1)

    def gradient_penalty(self, y, x):
        """Compute gradient penalty: (L2_norm(dy/dx)-1)**2."""
        weight = torch.ones(y.size()).to(self.device)
        dydx = torch.autograd.grad(outputs=y,
                                   inputs=x,
                                   grad_outputs=weight,
                                   retain_graph=True,
                                   create_graph=True,
                                   only_inputs=True)[0]
        dydx = dydx.view(dydx.size(0), -1)
        dydx_l2norm = torch.sqrt(torch.sum(dydx**2, dim=1))
        gradient_penalty = torch.mean((dydx_l2norm - 1.)**2)
        return gradient_penalty

    def train(self):
        # Start training from scratch or resume training.
        start_iters = 0
        if self.resume_iters:
            start_iters = self.resume_iters
            self.load_model(self.resume_iters)

        # Fixed inputs for sampling.
        fixed_noise = torch.randn(self.batch_size, self.z_dim,1,1).to(self.device)
        fixed_images = next(iter(self.train_loader))[0].to(self.device)

        # Start training.
        print('Start training...')
        start_time = time.time()
        for i in range(start_iters, self.num_iters):
            for j, (images, _) in enumerate(self.train_loader):

                # =================== Process input data. =================== #
                noise = torch.randn(len(images), self.z_dim,1,1).to(self.device)
                images = images.to(self.device)

                # ================ Train the discriminator. ================= #
                # Compute loss (Wasserstein-1 distance).
                d_real_loss = self.D(images)
                fake_images = self.G(noise)
                d_fake_loss = self.D(fake_images)

                # Compute loss for gradient penalty.
                alpha = torch.rand(images.size(0), 1, 1, 1).to(self.device)
                x_hat = alpha * fake_images + (1-alpha) * images
                d_gp_loss = self.gradient_penalty(self.D(x_hat), x_hat)

                d_loss = torch.mean(d_fake_loss) - torch.mean(d_real_loss) + self.lambda_gp * d_gp_loss

                # Backprop and optimize.
                self.reset_grad()
                d_loss.backward()
                self.d_optimizer.step()

                # =================== Train the inverter. =================== #
                # Compute loss.
                reconst_images = self.G(self.I(images))
                reconst_noise = self.I(self.G(noise))

                i_loss = torch.mean((images-reconst_images)**2) \
                         + self.lambda_i * torch.mean((noise-reconst_noise)**2)

                # Backprop and optimize.
                self.reset_grad()
                i_loss.backward()
                self.i_optimizer.step()

                # ================== Train the generator. =================== #
                # Compute loss.
                g_loss = -torch.mean(self.D(self.G(noise)))

                # Backprop and optimize.
                self.reset_grad()
                g_loss.backward()
                self.g_optimizer.step()

            # Print out training information.
            if (i+1) % self.log_step == 0:
                elapsed_time = time.time() - start_time
                print('Elapsed time [{:.4f}], Iteration [{}/{}], '
                      'D Loss: {:.4f}, I Loss: {:.4f}, G Loss: {:.4f}'.format(
                    elapsed_time, i+1, self.num_iters,
                    d_loss.item(), i_loss.item(), g_loss.item()
                ))

            # Save sampled images.
            if (i+1) % self.sample_step == 0:
                # Generate from noise.
                output = self.G(fixed_noise)
                sample_path = os.path.join(self.sample_dir, '{}_samples.jpg'.format(i+1))
                save_image(output.data.cpu(), sample_path)
                print('Saved generated images into {}...'.format(sample_path))

                # Reconstruct images.
                reconst_images = self.G(self.I(fixed_images))
                comparison = torch.zeros((fixed_images.size(0) * 2,
                                          fixed_images.size(1),
                                          fixed_images.size(2),
                                          fixed_images.size(3)),
                                         dtype=torch.float).to(self.device)
                for k in range(fixed_images.size(0)):
                    comparison[2*k] = fixed_images[k]
                    comparison[2*k+1] = reconst_images[k]

                sample_path = os.path.join(self.sample_dir, '{}_reconstructions.jpg'.format(i+1))
                save_image(comparison.data.cpu(), sample_path)

            # Save model checkpoints.
            if (i+1) % self.model_step == 0:
                G_path = os.path.join(self.model_dir, '{}_G.ckpt'.format(i+1))
                I_path = os.path.join(self.model_dir, '{}_I.ckpt'.format(i+1))
                D_path = os.path.join(self.model_dir, '{}_D.ckpt'.format(i+1))
                torch.save(self.G.state_dict(), G_path)
                torch.save(self.I.state_dict(), I_path)
                torch.save(self.D.state_dict(), D_path)
                print('Saved model checkpoints into {}...'.format(self.model_dir))

    def load_unk_model(self, args,retrain=False):
        """
        Load an unknown model. Used for convenience to easily swap unk model
        """
        if args.mnist:
            if os.path.exists("../saved_models/mnist_cnn.pt"):
                model = Net().to(args.device)
                model.load_state_dict(torch.load("../saved_models/mnist_cnn.pt"))
                model.eval()
            else:
                model = main_mnist(args)
        if args.cifar:
            if os.path.exists("../saved_models/cifar_VGG16.pt") and retrain == False:
                model = VGG().to(args.device)
                model = nn.DataParallel(model)
                model.load_state_dict(torch.load("../saved_models/cifar_VGG16.pt"))
            else:
                model = utils.main_cifar(args,normalize=False)
        else:
            # load pre-trained ResNet50
            model = resnet50(pretrained=True).to(args.device)
        return model

    def generate_adversary(self,args):
        # Load the trained models.
        self.load_model(self.resume_iters)

        # Choose search method.
        print('Search method:', self.search)

        # Choose the classifier to generate adversary examples against.
        print('Classifier:', self.classifier)
        C = load_unk_model(args,retrain=False)
        # if self.classifier == 'lenet':
            # C = LeNet().to(self.device)
            # cla_path = os.path.join(self.cla_dir, self.classifier, '{}_lenet.ckpt'.format(self.cla_iters))
            # C.load_state_dict(torch.load(cla_path, map_location=lambda storage, loc: storage))

        # Generate adversary examples.
        for j, (images, labels) in enumerate(self.test_loader):
            for i in range(32):
                x = images[i].unsqueeze(0).to(self.device)
                y = labels[i].to(self.device)

                adversary = self.iterative_search(self.G, self.I, C, x, y,
                                                  n_samples=self.n_samples, step=self.step)
                adversary_path = os.path.join(self.adversary_dir,
                                 '{}_{}_{}.jpg'.format(self.classifier, j+1, i+1))
                self.save_adversary(adversary, adversary_path)
                print('Saved natural adversary example...'.format(adversary_path))

    def save_adversary(self, adversary, filename):
        fig, ax = plt.subplots(1, 2, figsize=(7, 3))

        ax[0].imshow(adversary['x'],
                     interpolation='none', cmap=plt.get_cmap('gray'))
        ax[0].text(1, 5, str(adversary['y']), color='white', fontsize=20)
        ax[0].axis('off')

        ax[1].imshow(adversary['x_adv'],
                     interpolation='none', cmap=plt.get_cmap('gray'))
        ax[1].text(1, 5, str(adversary['y_adv']), color='white', fontsize=20)
        ax[1].axis('off')

        fig.savefig(filename)
        plt.close()

    def iterative_search(self, G, I, C, x, y, y_target=None, z=None,
                         n_samples=5000, step=0.01, l=0., h=10., ord=2):
        """
        :param G: function of generator
        :param I: function of inverter
        :param C: function of classifier
        :param x: input instance
        :param y: label
        :param y_target: target label for adversary
        :param z: latent vector corresponding to x
        :param n_samples: number of samples in each search iteration
        :param step: delta r for search step size
        :param l: lower bound of search range
        :param h: upper bound of search range
        :param ord: indicating norm order
        :return: adversary for x against the classifier (d_adv is delta z between z and z_adv)
        """

        x_adv, y_adv, z_adv, d_adv = None, None, None, None
        h = l + step

        if z is None:
            z = I(x)

        while True:
            delta_z = np.random.randn(n_samples, z.size(1))
            d = np.random.rand(n_samples) * (h - l) + l        # random values between the search range [l, h)
            norm_p = np.linalg.norm(delta_z, ord=ord, axis=1)  # L2 norm of delta z along axis=1
            d_norm = np.divide(d, norm_p).reshape(-1, 1)       # rescale/normalize d
            delta_z = np.multiply(delta_z, d_norm)             # random noise vectors of norms within (r, r + delta r]
            delta_z = torch.from_numpy(delta_z).float().to(self.device)
            z_tilde = z + delta_z
            x_tilde = G(z_tilde)
            y_tilde = torch.argmax(C(x_tilde), dim=1)

            if y_target is None:
                indices_adv = np.where(y_tilde != y)[0]
            else:
                indices_adv = np.where(y_tilde == y_target)[0]

            # No candidate generated.
            if len(indices_adv) == 0:
                print('No candidate generated, increasing search range...')
                l = h
                h = l + step

            # Certain candidates generated.
            else:
                # Choose the data index with the least perturbation.
                idx_adv = indices_adv[np.argmin(d[indices_adv])]

                # For debugging.
                if y_target is None:
                    assert (y_tilde[idx_adv] != y)

                else:
                    assert (y_tilde[idx_adv] == y_target)

                # Save natural adversary example.
                if d_adv is None or d[idx_adv] < d_adv:
                    x_adv = x_tilde[idx_adv]
                    y_adv = y_tilde[idx_adv]
                    z_adv = z_tilde[idx_adv]
                    d_adv = d[idx_adv]

                    if y_target is None:
                        print("Untarget y=%d y_adv=%d d_adv=%.4f l=%.4f h=%.4f" % (y, y_adv, d_adv, l, h))
                    else:
                        print("Targeted y=%d y_adv=%d d_adv=%.4f l=%.4f h=%.4f" % (y, y_adv, d_adv, l, h))

                    break

        adversary = {'x': x.squeeze(0).squeeze(0).data.cpu().numpy(),
                     'y': y.data.cpu().numpy(),
                     'z': z.squeeze(0).data.cpu().numpy(),
                     'x_adv': x_adv.squeeze(0).data.cpu().numpy(),
                     'y_adv': y_adv.data.cpu().numpy(),
                     'z_adv': z_adv.data.cpu().numpy(),
                     'd_adv': d_adv}

        return adversary


# -*- coding: utf-8 -*-
"""
Created on Wed Jan 27 19:32:21 2021

@author: zju
"""
import numpy as np

def preprocess(pair_file, seq_file):

    with open(pair_file, 'r') as f:

        lines = f.readlines()

    proteins_1 = [line.strip().split('\t')[0] for line in lines]
    proteins_2 = [line.strip().split('\t')[1] for line in lines]
    labels = [line.strip().split('\t')[2] for line in lines]

    protein_list = list(set(proteins_1 + proteins_2))

    protein_seq = {}

    with open(seq_file, 'r') as f:

        lines = f.readlines()

    for i in range(len(lines)):

        line = lines[i].strip().split('\t')
        protein_seq[line[0]] = line[1]

#    return proteins_1, proteins_2, labels, protein_seq
    amino_acid ={'A':1,'C':2,'D':3,'E':4,'F':5,
                'G':6,'H':7,'I':8,'K':9,'L':10,
                'M':11,'N':12,'P':13,'Q':14,'R':15,'S':16,
                'T':17,'V':18,'W':19,'Y':20,'U':21,'X':22,'B':0}

# positive and negative setting
    k1 = []
    k2 = []
    k3 = []
    k_h = []

    for i in range(len(labels)):

        protein_1 = proteins_1[i]
        protein_2 = proteins_2[i]

        label = labels[i]

        seq_1 = protein_seq[protein_1]
        seq_2 = protein_seq[protein_2]

        a1 = np.zeros([1500,], dtype = int)
        a2 = np.zeros([1500,], dtype = int)
        a3 = np.zeros([3,], dtype = float)

        k = 0
        for AA in seq_1:
            a1[k] = amino_acid[AA]
            k += 1
        k1.append(a1)

        k = 0
        for AA in seq_2:
            a2[k] = amino_acid[AA]
            k += 1
        k2.append(a2)

        if int(label) == 0:
            a3[1] = 1
        elif int(label) == 1:
            a3[0] = 1
        else:
            print('error')
            break
        k3.append(a3)

        k_h.append(np.array([protein_1, protein_2]))

    m1 = np.stack(k1, axis=0)
    m2 = np.stack(k2, axis=0)
    m3 = np.stack(k3, axis=0)
    m_h = np.stack(k_h, axis=0)

# single protein setting
    k1 = []
    k2 = []
    k3 = []

    for protein in protein_list:

        seq_1 = protein_seq[protein]
        seq_2 = 'B'
        label = 2

        a1 = np.zeros([1500,], dtype = int)
        a2 = np.zeros([1500,], dtype = int)
        a3 = np.zeros([3,], dtype = float)

        k = 0
        for AA in seq_1:
            a1[k] = amino_acid[AA]
            k += 1
        k1.append(a1)

        k = 0
        for AA in seq_2:
            a2[k] = amino_acid[AA]
            k += 1
        k2.append(a2)

        if int(label) == 2:
            a3[2] = 1
        else:
            print('error')
            break
        k3.append(a3)

    n1 = np.stack(k1, axis=0)
    n2 = np.stack(k2, axis=0)
    n3 = np.stack(k3, axis=0)

    return m1, m2, m3, n1, n2, n3, m_h


from typing import List, Tuple
from functools import reduce
from time import time
from copy import deepcopy
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

from regym.rl_algorithms.replay_buffers import Storage
from regym.networks import random_sample

from regym.rl_algorithms.PPO.ppo_loss import compute_loss

class PPOAlgorithm():

    def __init__(self, kwargs, model):
        '''
        Refer to original paper for further explanation: https://arxiv.org/pdf/1707.06347.pdf
        horizon: (0, infinity) Number of timesteps that will elapse in between optimization calls.
        discount: (0,1) Reward discount factor
        use_gae: Flag, wether to use Generalized Advantage Estimation (GAE) (instead of return base estimation)
        gae_tau: (0,1) GAE hyperparameter.
        use_cuda: Flag, to specify whether to use CUDA tensors in Pytorch calculations
        entropy_weight: (0,1) Coefficient for (regularatization) entropy based loss
        gradient_clip: float, Clips gradients to reduce the chance of destructive updates
        optimization_epochs: int, Number of epochs per optimization step.
        mini_batch_size: int, Mini batch size to use to calculate losses (Use power of 2 for efficciency)
        ppo_ratio_clip: float, clip boundaries (1 - clip, 1 + clip) used in clipping loss function.
        learning_rate: float, optimizer learning rate.
        adam_eps: (float), Small Epsilon value used for ADAM optimizer. Prevents numerical instability when v^{hat} (Second momentum estimator) is near 0.
        model: (Pytorch nn.Module) Used to represent BOTH policy network and value network
        '''
        self.kwargs = deepcopy(kwargs)
        self.horizon = kwargs['horizon']
        self.discount = kwargs['discount']
        self.use_gae = kwargs['use_gae']
        self.gae_tau = kwargs['gae_tau']
        self.ppo_ratio_clip = kwargs['ppo_ratio_clip']
        self.entropy_weight = kwargs['entropy_weight']
        self.optimization_epochs = kwargs['optimization_epochs']
        self.learning_rate = kwargs['learning_rate']
        self.gradient_clip = kwargs['gradient_clip']
        self.adam_eps = kwargs['adam_eps']
        self.use_cuda = kwargs['use_cuda']
        self.batch_size = kwargs['mini_batch_size']
        self.model = model
        if self.use_cuda:
            self.model = self.model.cuda()

        self.optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate, eps=self.adam_eps)

        self.recurrent = False
        self.rnn_keys = [key for key,value in self.kwargs.items() if isinstance(value, str) and 'RNN' in value]
        if len(self.rnn_keys):
            self.recurrent = True

        # As a default, create a single storage
        self.storages: List[Storage] = self.create_storages(num_storages=1)

        # Number of times train was called
        self.num_updates = 0

        # Number of policy updates (Number of times compute_loss() was called)
        self.num_optimizer_steps = 0

        # To be set by a PPOAgent
        self.summary_writer: torch.util.tensorboard.SummaryWritter = None

    def create_storages(self, num_storages: int, size=-1) -> List[Storage]:
        if size == -1: size = self.horizon
        storages = [Storage(size) for _ in range(num_storages)]
        if self.recurrent:
            for storage in storages:
                storage.add_key('rnn_states')
                storage.add_key('next_rnn_states')
        return storages

    def train(self):
        start_time = time()
        self.num_updates += 1
        for storage in self.storages:
            storage.placeholder(num_elements=len(storage.s))
            self.compute_advantages_and_returns(storage)

        (states, actions, log_probs_old, returns,
         advantages, rnn_states) = self.retrieve_values_from_storages()

        if self.recurrent: rnn_states = self.reformat_rnn_states(rnn_states)

        total_loss = 0
        for _ in range(self.optimization_epochs):
            epoch_loss = self.optimize_model(states, actions, log_probs_old, returns,
                                             advantages, rnn_states)
            total_loss += epoch_loss

        for storage in self.storages: storage.reset()

        if self.summary_writer:
            self.summary_writer.add_scalar('Timing/Policy_update', time() - start_time,
                                           self.num_updates)
            self.summary_writer.add_scalar('Training/Avg_loss_per_epoch',
                                           (total_loss / self.optimization_epochs).cpu(),
                                           self.num_updates)

    def reformat_rnn_states(self, rnn_states):
        reformated_rnn_states = { k: ( [list()], [list()] ) for k in self.rnn_keys }
        for rnn_state in rnn_states:
            for k, (hstates, cstates) in rnn_state.items():
                for idx_layer, (h,c) in enumerate(zip(hstates, cstates)):
                    reformated_rnn_states[k][0][0].append(h)
                    reformated_rnn_states[k][1][0].append(c)
        for k, (hstates, cstates) in reformated_rnn_states.items():
            hstates = torch.cat(hstates[0], dim=0)
            cstates = torch.cat(cstates[0], dim=0)
            reformated_rnn_states[k] = ([hstates], [cstates])
        return reformated_rnn_states

    def compute_advantages_and_returns(self, storage):
        advantages = torch.from_numpy(np.zeros((1, 1), dtype=np.float32)) # TODO explain (used to be number of workers)
        returns = storage.V[-1].detach()
        num_observed_rewards = len(storage.r)
        for i in reversed(range(num_observed_rewards)):
            returns = storage.r[i] + self.discount * storage.non_terminal[i] * returns
            if not self.use_gae:
                advantages = returns - storage.V[i].detach()
            else:
                td_error = storage.r[i] + self.discount * storage.non_terminal[i] * storage.V[i + 1] - storage.V[i]
                advantages = advantages * self.gae_tau * self.discount * storage.non_terminal[i] + td_error
            storage.adv[i] = advantages.detach()
            storage.ret[i] = returns.detach()

    def retrieve_values_from_storages(self):
        all_states = []
        all_actions = []
        all_log_probs_old = []
        all_returns = []
        all_advantages = []
        all_rnn_states = [] if self.recurrent else None

        for storage in self.storages:
            (states, actions, log_probs_old, returns, advantages,
             rnn_states) = self.retrieve_values_from_single_storage(storage)
            all_states += states
            all_actions += actions
            all_log_probs_old += log_probs_old
            all_returns += returns
            all_advantages += advantages
            if self.recurrent: all_rnn_states += rnn_states

        all_states = torch.stack(all_states, dim=0)
        all_actions = torch.stack(all_actions, dim=0)
        all_log_probs_old = torch.stack(all_log_probs_old, dim=0)
        all_returns = torch.stack(all_returns, dim=0)
        all_advantages = torch.stack(all_advantages, dim=0)
        if self.recurrent: all_rnn_states = torch.stack(all_rnn_states, dim=0)
        return (all_states, all_actions, all_log_probs_old, all_returns,
               all_advantages, all_rnn_states)

    def retrieve_values_from_single_storage(self, storage):
        if self.recurrent:
            cat = storage.cat(['s', 'a', 'log_pi_a', 'ret', 'adv', 'rnn_states'])
            rnn_states = cat[-1]
            states, actions, log_probs_old, returns, advantages = map(lambda x: torch.cat(x, dim=1), cat[:-1])
        else:
            states, actions, log_probs_old, returns, advantages = map(lambda x: torch.cat(x, dim=0),
                                                                     storage.cat(['s', 'a', 'log_pi_a', 'ret', 'adv']) )
            rnn_states = None

        # Turning actions into dim: self.horizon x 1
        # (Might break in environments where actions are not a single number)
        assert storage.s != self.horizon, (f'Number of states in storage ({len(storage.s)}) '
                                           f"should be equal to PPO's horizon ({self.horizon})")
        #states = states.view(len(storage.s), -1)
        actions = actions.view(-1, 1)
        log_probs_old = log_probs_old.view(-1, 1)
        returns = returns.view(-1, 1)
        advantages = advantages.view(-1, 1)
        advantages = self.standardize(advantages)
        return states, actions, log_probs_old, returns, advantages, rnn_states

    def standardize(self, x):
        return (x - x.mean()) / x.std()

    def sample_batch_from_indices(self, batch_indices, states, actions,
                                  log_probs_old, returns, advantages,
                                  rnn_states) -> Tuple[torch.Tensor]:
        sampled_states = states[batch_indices].cuda() if self.use_cuda else states[batch_indices]
        sampled_actions = actions[batch_indices].cuda() if self.use_cuda else actions[batch_indices]
        sampled_log_probs_old = log_probs_old[batch_indices].cuda() if self.use_cuda else log_probs_old[batch_indices]
        sampled_returns = returns[batch_indices].cuda() if self.use_cuda else returns[batch_indices]
        sampled_advantages = advantages[batch_indices].cuda() if self.use_cuda else advantages[batch_indices]

        sampled_rnn_states = None
        if self.recurrent:
            sampled_rnn_states = { k: ([None]*nbr_layers_per_rnn[k] , [None]*nbr_layers_per_rnn[k]) for k in self.rnn_keys}
            for k in sampled_rnn_states:
                for idx in range(nbr_layers_per_rnn[k]):
                    sampled_rnn_states[k][0][idx] = rnn_states[k][0][idx][batch_indices].cuda() if self.use_cuda else rnn_states[k][0][idx][batch_indices]
                    sampled_rnn_states[k][1][idx] = rnn_states[k][1][idx][batch_indices].cuda() if self.use_cuda else rnn_states[k][1][idx][batch_indices]
        return (sampled_states, sampled_actions, sampled_log_probs_old,
                sampled_returns, sampled_advantages, sampled_rnn_states)

    def optimize_model(self, states, actions, log_probs_old, returns, advantages, rnn_states=None):
        sampler = random_sample(np.arange(states.size(0)), self.batch_size)
        assert( (self.recurrent and rnn_states is not None) or not(self.recurrent or rnn_states is not None) )
        if self.recurrent:
            nbr_layers_per_rnn = { k:len(rnn_states[k][0] ) for k in self.rnn_keys}

        total_loss = 0.
        for batch_indices in sampler:
            self.num_optimizer_steps += 1
            batch_indices = torch.from_numpy(batch_indices).long()

            (sampled_states, sampled_actions, sampled_log_probs_old,
            sampled_returns, sampled_advantages,
            sampled_rnn_states) = self.sample_batch_from_indices(batch_indices,
                    states, actions, log_probs_old, returns, advantages,
                    rnn_states)

            batch_loss = compute_loss(states=sampled_states,
                                      actions=sampled_actions,
                                      log_probs_old=sampled_log_probs_old,
                                      returns=sampled_returns,
                                      advantages=sampled_advantages,
                                      model=self.model,
                                      rnn_states=sampled_rnn_states,
                                      ratio_clip=self.ppo_ratio_clip,
                                      entropy_weight=self.entropy_weight,
                                      iteration_count=self.num_optimizer_steps,
                                      summary_writer=self.summary_writer)

            self.optimizer.zero_grad()
            batch_loss.backward(retain_graph=False)
            nn.utils.clip_grad_norm_(self.model.parameters(), self.gradient_clip)
            self.optimizer.step()
            total_loss += batch_loss.cpu().detach()
        return total_loss

    def __getstate__(self):
        '''
        Function invoked when pickling.

        torch.utils.SummaryWriters are not pickable.
        '''
        to_pickle_dict = self.__dict__
        if 'summary_writer' in to_pickle_dict:
            to_pickle_dict = self.__dict__.copy()
            to_pickle_dict['summary_writer'] = None
        return to_pickle_dict

    def __repr__(self):
        return (f'Num updates: {self.num_updates}\n'
         f'Num optimizer steps: {self.num_optimizer_steps}\n'
         f'Horizon: {self.horizon}\n'
         f'Discount: {self.discount}\n'
         f'Use_gae: {self.use_gae}\n'
         f'Gae tau: {self.gae_tau}\n'
         f'Ppo ratio clip: {self.ppo_ratio_clip}\n'
         f'Entropy weight: {self.entropy_weight}\n'
         f'Optimization epochs: {self.optimization_epochs}\n'
         f'Learning rate: {self.learning_rate}\n'
         f'Gradient clip: {self.gradient_clip}\n'
         f'Adam eps: {self.adam_eps}\n'
         f'Batch size: {self.batch_size}\n'
         f'Use cuda: {self.use_cuda}\n'
         f'Storages: {self.storages}'
         )


import argparse
import os

import numpy as np
import torch as t

from utils.batch_loader import BatchLoader
from utils.parameters import Parameters
from model.rvae_dilated import RVAE_dilated

if __name__ == '__main__':

    assert os.path.exists('trained_RVAE'), \
        'trained model not found'

    parser = argparse.ArgumentParser(description='Sampler')
    parser.add_argument('--use-cuda', type=bool, default=True, metavar='CUDA',
                        help='use cuda (default: True)')
    parser.add_argument('--num-sample', type=int, default=10, metavar='NS',
                        help='num samplings (default: 10)')

    args = parser.parse_args()

    batch_loader = BatchLoader('')
    parameters = Parameters(batch_loader.max_word_len,
                            batch_loader.max_seq_len,
                            batch_loader.words_vocab_size,
                            batch_loader.chars_vocab_size)

    rvae = RVAE_dilated(parameters)
    rvae.load_state_dict(t.load('trained_RVAE'))
    if args.use_cuda:
        rvae = rvae.cuda()

    for iteration in range(args.num_sample):
        seed = np.random.normal(size=[1, parameters.latent_variable_size])
        result = rvae.sample(batch_loader, 50, seed, args.use_cuda)
        print(result)
        print()


#!/usr/bin/env python
"""
radarData.py
This file holds the RadarData class that hold the radar data and processes it.

@author: John Swoboda
"""
import ipdb
import scipy as sp
# My modules
from isrutilities.physConstants import v_C_0, v_Boltz
from SimISR import Path
from SimISR.IonoContainer import IonoContainer
from SimISR.utilFunctions import CenteredLagProduct, MakePulseDataRepLPC,dict2h5,h52dict,readconfigfile, BarkerLag, update_progress
#from SimISR import specfunctions
#from SimISR.analysisplots import plotspecsgen

class RadarDataFile(object):
    """
        This class will will take the ionosphere class and create radar data both
        at the IQ and fitted level.

        Variables
        simparams - A dictionary that holds simulation parameters the keys are the
            following
            'angles': Angles (in degrees) for the simulation.  The az and el angles
                are stored in a list of tuples [(az_0,el_0),(az_1,el_1)...]
            'IPP' : Interpulse period in seconds, this is the IPP from position to
                position.  In other words if a IPP is 10 ms and there are 10 beams it will
                take 100 ms to go through all of the beams.
            'Tint'  This is the integration time in seconds.
            'Timevec': This is the time vector for each frame.  It is referenced to
                the begining of the frame.  Also in seconds
            'Pulse': The shape of the pulse.  This is a numpy array.
            'TimeLim': The length of time that the experiment will run, which cooresponds
                to how many pulses per position are used.
        sensdict - A dictionary that holds the sensor parameters.
        rawdata: This is a NbxNpxNr numpy array that holds the raw IQ data.
        rawnoise: This is a NbxNpxNr numpy array that holds the raw noise IQ data.
    """

    def __init__(self, Ionodict, inifile, outdir, outfilelist=None):
        """
            This function will create an instance of the RadarData class.  It will
            take in the values and create the class and make raw IQ data.
            Inputs:
            sensdict - A dictionary of sensor parameters
            angles - A list of tuples which the first position is the az angle
                and the second position is the el angle.
            IPP - The interpulse period in seconds represented as a float.
            Tint - The integration time in seconds as a float.  This will be the
            integration time of all of the beams.
            time_lim - The length of time of the simulation the number of time points
                will be calculated.
            pulse - A numpy array that represents the pulse shape.
            rng_lims - A numpy array of length 2 that holds the min and max range
                that the radar will cover.
        """
        (sensdict, simparams) = readconfigfile(inifile)
        self.simparams = simparams
        N_angles = len(self.simparams['angles'])

        NNs = int(self.simparams['NNs'])
        self.sensdict = sensdict
        Npall = sp.floor(self.simparams['TimeLim']/self.simparams['IPP'])
        Npall = int(sp.floor(Npall/N_angles)*N_angles)
        Np = Npall/N_angles

        print("All spectrums created already")
        filetimes = Ionodict.keys()
        filetimes.sort()
        ftimes = sp.array(filetimes)
        simdtype = self.simparams['dtype']
        pulsetimes = sp.arange(Npall)*self.simparams['IPP'] +ftimes.min()
        pulsefile = sp.array([sp.where(itimes-ftimes >= 0)[0][-1] for itimes in pulsetimes])
        # differentiate between phased arrays and dish antennas
        if sensdict['Name'].lower() in ['risr', 'pfisr', 'risr-n']:

            beams = sp.tile(sp.arange(N_angles), Npall/N_angles)
        else:

            # for dish arrays
            brate = simparams['beamrate']
            beams2 = sp.repeat(sp.arange(N_angles), brate)
            beam3 = sp.concatenate((beams2, beams2[::-1]))
            ntile = int(sp.ceil(Npall/len(beam3)))
            leftover = int(Npall-ntile*len(beam3))
            if ntile > 0:
                beams = sp.tile(beam3, ntile)
                beams = sp.concatenate((beams, beam3[:leftover]))
            else:
                beams = beam3[:leftover]

        pulsen = sp.repeat(sp.arange(Np), N_angles)
        pt_list = []
        pb_list = []
        pn_list = []
        fname_list = []
        self.datadir = outdir
        self.maindir = outdir.parent
        self.procdir = self.maindir/'ACF'
        Nf = len(filetimes)
        progstr = 'Data from {:d} of {:d} being processed Name: {:s}.'
        if outfilelist is None:
            print('\nData Now being created.')

            Noisepwr = v_Boltz*sensdict['Tsys']*sensdict['BandWidth']
            self.outfilelist = []
            for ifn, ifilet in enumerate(filetimes):

                outdict = {}
                ifile = Ionodict[ifilet]
                ifilename = Path(ifile).name
                update_progress(float(ifn)/Nf, progstr.format(ifn, Nf, ifilename))
                curcontainer = IonoContainer.readh5(ifile)
                if ifn == 0:
                    self.timeoffset = curcontainer.Time_Vector[0, 0]
                pnts = pulsefile == ifn
                pt = pulsetimes[pnts]
                pb = beams[pnts]
                pn = pulsen[pnts].astype(int)
                rawdata = self.__makeTime__(pt, curcontainer.Time_Vector,
                                            curcontainer.Sphere_Coords,
                                            curcontainer.Param_List, pb)
                d_shape = rawdata.shape
                n_tempr = sp.random.randn(*d_shape).astype(simdtype)
                n_tempi = 1j*sp.random.randn(*d_shape).astype(simdtype)
                noise = sp.sqrt(Noisepwr/2)*(n_tempr+n_tempi)
                outdict['AddedNoise'] = noise
                outdict['RawData'] = rawdata+noise
                outdict['RawDatanonoise'] = rawdata
                outdict['NoiseData'] = sp.sqrt(Noisepwr/2)*(sp.random.randn(len(pn), NNs).astype(simdtype)+
                                                            1j*sp.random.randn(len(pn), NNs).astype(simdtype))
                outdict['Pulses'] = pn
                outdict['Beams'] = pb
                outdict['Time'] = pt
                fname = '{0:d} RawData.h5'.format(ifn)
                newfn = self.datadir/fname
                self.outfilelist.append(str(newfn))
                dict2h5(str(newfn), outdict)

                #Listing info
                pt_list.append(pt)
                pb_list.append(pb)
                pn_list.append(pn)
                fname_list.append(fname)
            infodict = {'Files':fname_list, 'Time':pt_list, 'Beams':pb_list, 'Pulses':pn_list}
            dict2h5(str(outdir.joinpath('INFO.h5')), infodict)

        else:
            infodict = h52dict(str(outdir.joinpath('INFO.h5')))
            alltime = sp.hstack(infodict['Time'])
            self.timeoffset = alltime.min()
            self.outfilelist = outfilelist

#%% Make functions
    def __makeTime__(self,pulsetimes,spectime,Sphere_Coords,allspecs,beamcodes):
        """This is will make raw radar data for a given time period and set of
        spectrums. This is an internal function called by __init__.
        Inputs-
        self - RadarDataFile object.
        pulsetimes - The time the pulses are sent out in reference to the spectrums.
        spectime - The times for the spectrums.
        Sphere_Coords - An Nlx3 array that holds the spherical coordinates of the spectrums.
        allspecs - An NlxNdtimexNspec array that holds the spectrums.
        beamcodes - A NBx4 array that holds the beam codes along with the beam location
        in az el along with a system constant in the array. """

        range_gates = self.simparams['Rangegates']
        pulse = self.simparams['Pulse']
        sensdict = self.sensdict
        pulse2spec = sp.array([sp.where(itimes-spectime >= 0)[0][-1] for itimes in pulsetimes])
        Np = len(pulse2spec)
        lp_pnts = len(pulse)
        samp_num = sp.arange(lp_pnts)
        #N_rg = len(range_gates)# take the size
        N_samps = len(range_gates)
        angles = self.simparams['angles']
        Nbeams = len(angles)
        rho = Sphere_Coords[:, 0]
        Az = Sphere_Coords[:, 1]
        El = Sphere_Coords[:, 2]
        rng_len = self.sensdict['t_s']*v_C_0*1e-3/2.
        speclen = allspecs.shape[-1]
        simdtype = self.simparams['dtype']
        out_data = sp.zeros((Np, N_samps), dtype=simdtype)
        weights = {ibn:self.sensdict['ArrayFunc'](Az, El, ib[0], ib[1], sensdict['Angleoffset'])
                   for ibn, ib in enumerate(angles)}
        for istn in range(len(spectime)):
            for ibn in range(Nbeams):
                #print('\t\t Making Beam {0:d} of {1:d}'.format(ibn, Nbeams))
                weight = weights[ibn]
                for isamp in range(N_samps):
                    range_g = range_gates[isamp]

                    if range_g <= 0:
                        continue
                    range_m = range_g*1e3
                    rnglims = [range_g-rng_len/2.0, range_g+rng_len/2.0]
                    rangelog = (rho >= rnglims[0])&(rho < rnglims[1])
                    # Get the number of points covered
                    cur_pnts = samp_num+isamp
                    keep_pnt = sp.logical_and(cur_pnts >= 0, cur_pnts < N_samps)
                    cur_pnts = cur_pnts[keep_pnt]
                    # This is a nearest neighbors interpolation for the
                    # spectrums in the range domain
                    if sp.sum(rangelog) == 0:
                        minrng = sp.argmin(sp.absolute(range_g-rho))
                        rangelog[minrng] = True

                    #create the weights and weight location based on the beams pattern.
                    weight_cur = weight[rangelog]
                    weight_cur = weight_cur/weight_cur.sum()
                    specsinrng = allspecs[rangelog]
                    if specsinrng.ndim == 3:
                        specsinrng = specsinrng[:, istn]
                    elif specsinrng.ndim == 2:
                        specsinrng = specsinrng[istn]
                    specsinrng = specsinrng*sp.tile(weight_cur[:, sp.newaxis], (1, speclen))
                    cur_spec = specsinrng.sum(0)
                    # based off new way of calculating
                    pow_num = sensdict['Pt']*sensdict['Ksys'][ibn]*sensdict['t_s']
                    pow_den = range_m**2
                    curdataloc = sp.where(sp.logical_and((pulse2spec == istn),
                                                         (beamcodes == ibn)))[0]
                    # create data
                    if len(curdataloc) == 0:
                        print('\t\t No data for {0:d} of {1:d} in this time period'.format(ibn, Nbeams))
                        continue
                    cur_pulse_data = MakePulseDataRepLPC(pulse, cur_spec, 20,
                                                         len(curdataloc), numtype=simdtype)
                    cur_pulse_data = cur_pulse_data*sp.sqrt(pow_num/pow_den)
                    for idatn, idat in enumerate(curdataloc):
                        d_pnts = cur_pnts[::-1]-isamp
                        cursum = cur_pulse_data[idatn, d_pnts]+out_data[idat, cur_pnts]
                        out_data[idat, cur_pnts] = cursum
        return out_data
        #%% Processing
    def processdataiono(self):
        """ This will perform the the data processing and create the ACF estimates
        for both the data and noise but put it in an Ionocontainer.
        Inputs:

        Outputs:
        Ionocontainer- This is an instance of the ionocontainer class that will hold the acfs.
        """
        (datalags, noiselags) = self.processdata()
        return lagdict2ionocont(datalags, noiselags, self.sensdict, self.simparams,
                                datalags['Time'])

    def processdata(self):
        """ This will perform the the data processing and create the ACF estimates
        for both the data and noise.
        Inputs:
        timevec - A numpy array of times in seconds where the integration will begin.
        inttime - The integration time in seconds.
        lagfunc - A function that will make the desired lag products.
        Outputs:
        DataLags: A dictionary with keys 'Power' 'ACF','RG','Pulses' that holds
        the numpy arrays of the data.
        NoiseLags: A dictionary with keys 'Power' 'ACF','RG','Pulses' that holds
        the numpy arrays of the data.
        """
        timevec = self.simparams['Timevec'] +self.timeoffset
        inttime = self.simparams['Tint']
        # Get array sizes

        NNs = int(self.simparams['NNs'])
        range_gates = self.simparams['Rangegates']
        N_samps = len(range_gates)# take the size
        pulse = self.simparams['Pulse']
        Pulselen = len(pulse)
        N_rg = N_samps - Pulselen+1
        simdtype = self.simparams['dtype']
        Ntime = len(timevec)

        if 'outangles' in self.simparams.keys():
            Nbeams = len(self.simparams['outangles'])
            inttime = inttime
        else:
            Nbeams = len(self.simparams['angles'])

        # Choose type of processing
        if self.simparams['Pulsetype'].lower() == 'barker':
            lagfunc = BarkerLag
            Nlag = 1
        else:
            lagfunc = CenteredLagProduct
            Nlag = Pulselen
        # initialize output arrays
        outdata = sp.zeros((Ntime, Nbeams, N_rg, Nlag), dtype=simdtype)
        outaddednoise = sp.zeros((Ntime, Nbeams, N_rg, Nlag), dtype=simdtype)
        outnoise = sp.zeros((Ntime, Nbeams, NNs-Pulselen+1, Nlag), dtype=simdtype)
        pulses = sp.zeros((Ntime, Nbeams))
        pulsesN = sp.zeros((Ntime, Nbeams))
        timemat = sp.zeros((Ntime, 2))
        Ksysvec = self.sensdict['Ksys']
        # set up arrays that hold the location of pulses that are to be processed together
        infoname = self.datadir / 'INFO.h5'
        # Just going to assume that the info file is in the directory
        infodict = h52dict(str(infoname))
        flist = infodict['Files']
        file_list = [str(self.datadir/i) for i in flist]
        pulsen_list = infodict['Pulses']
        beamn_list = infodict['Beams']
        time_list = infodict['Time']
        file_loclist = [ifn*sp.ones(len(ifl)) for ifn, ifl in enumerate(beamn_list)]

        pulsen = sp.hstack(pulsen_list).astype(int)# pulse number
        beamn = sp.hstack(beamn_list).astype(int)# beam numbers
        ptimevec = sp.hstack(time_list).astype(float)# time of each pulse
        file_loc = sp.hstack(file_loclist).astype(int)# location in the file

        # run the time loop
        print("Forming ACF estimates")

        # For each time go through and read only the necisary files
        for itn, it in enumerate(timevec):
            update_progress(float(itn)/Ntime, "Time {0:d} of {1:d}".format(itn, Ntime))
            # do the book keeping to determine locations of data within the files
            cur_tlim = (it, it+inttime)
            curcases = sp.logical_and(ptimevec >= cur_tlim[0], ptimevec < cur_tlim[1])
            if  not sp.any(curcases):
                prog_str = "No pulses for time {0:d} of {1:d}, lagdata adjusted accordinly"
                update_progress(float(itn)/Ntime, prog_str.format(itn, Ntime))
                outdata = outdata[:itn]
                outnoise = outnoise[:itn]
                pulses = pulses[:itn]
                pulsesN = pulsesN[:itn]
                timemat = timemat[:itn]
                continue
            pulseset = set(pulsen[curcases])
            poslist = [sp.where(pulsen == item)[0] for item in pulseset]
            pos_all = sp.hstack(poslist)
            try:
                pos_all = sp.hstack(poslist)
                curfileloc = file_loc[pos_all]
            except:
                ipdb.set_trace()
            # Find the needed files and beam numbers
            curfiles = set(curfileloc)
            beamlocs = beamn[pos_all]
            timemat[itn, 0] = ptimevec[pos_all].min()
            timemat[itn, 1] = ptimevec[pos_all].max()
            # cur data pulls out all data from all of the beams and posisions
            curdata = sp.zeros((len(pos_all), N_samps), dtype=simdtype)
            curaddednoise = sp.zeros((len(pos_all), N_samps), dtype=simdtype)
            curnoise = sp.zeros((len(pos_all), NNs), dtype=simdtype)
            # Open files and get required data
            # XXX come up with way to get open up new files not have to reread in data that is already in memory
            for ifn in curfiles:
                curfileit = [sp.where(pulsen_list[ifn] == item)[0] for item in pulseset]
                curfileitvec = sp.hstack(curfileit)
                ifile = file_list[ifn]
                curh5data = h52dict(ifile)
                file_arlocs = sp.where(curfileloc == ifn)[0]
                curdata[file_arlocs] = curh5data['RawData'][curfileitvec]

                curaddednoise[file_arlocs] = curh5data['AddedNoise'].astype(simdtype)[curfileitvec]
                # Read in noise data when you have don't have ACFs

                curnoise[file_arlocs] = curh5data['NoiseData'].astype(simdtype)[curfileitvec]

            # differentiate between phased arrays and dish antennas
            if self.sensdict['Name'].lower() in ['risr', 'pfisr', 'risr-n']:
                # After data is read in form lags for each beam
                for ibeam in range(Nbeams):
                    prog_num = float(itn)/Ntime+float(ibeam)/Ntime/Nbeams
                    update_progress(prog_num, "Beam {0:d} of {1:d}".format(ibeam, Nbeams))
                    beamlocstmp = sp.where(beamlocs == ibeam)[0]
                    pulses[itn, ibeam] = len(beamlocstmp)

                    outdata[itn, ibeam] = lagfunc(curdata[beamlocstmp].copy(),
                                                  numtype=self.simparams['dtype'],
                                                  pulse=pulse, lagtype=self.simparams['lagtype'])

                    pulsesN[itn, ibeam] = len(beamlocstmp)
                    outnoise[itn, ibeam] = lagfunc(curnoise[beamlocstmp].copy(),
                                                   numtype=self.simparams['dtype'],
                                                   pulse=pulse, lagtype=self.simparams['lagtype'])
                    outaddednoise[itn, ibeam] = lagfunc(curaddednoise[beamlocstmp].copy(),
                                                        numtype=self.simparams['dtype'],
                                                        pulse=pulse, lagtype=self.simparams['lagtype'])
            else:
                for ibeam, ibeamlist in enumerate(self.simparams['outangles']):
                    prog_num = float(itn)/Ntime+float(ibeam)/Ntime/Nbeams
                    update_progress(prog_num, "Beam {0:d} of {1:d}".format(ibeam, Nbeams))
                    beamlocstmp = sp.where(sp.in1d(beamlocs, ibeamlist))[0]
                    inputdata = curdata[beamlocstmp].copy()
                    noisedata = curnoise[beamlocstmp].copy()
                    noisedataadd = curaddednoise[beamlocstmp].copy()
                    pulses[itn, ibeam] = len(beamlocstmp)
                    pulsesN[itn, ibeam] = len(beamlocstmp)
                    outdata[itn, ibeam] = lagfunc(inputdata, numtype=self.simparams['dtype'],
                                                  pulse=pulse, lagtype=self.simparams['lagtype'])
                    outnoise[itn, ibeam] = lagfunc(noisedata, numtype=self.simparams['dtype'],
                                                   pulse=pulse, lagtype=self.simparams['lagtype'])
                    outaddednoise[itn, ibeam] = lagfunc(noisedataadd, numtype=self.simparams['dtype'],
                                                        pulse=pulse, lagtype=self.simparams['lagtype'])
        # Create output dictionaries and output data
        DataLags = {'ACF':outdata, 'Pow':outdata[:, :, :, 0].real, 'Pulses':pulses,
                    'Time':timemat, 'AddedNoiseACF':outaddednoise}
        NoiseLags = {'ACF':outnoise, 'Pow':outnoise[:, :, :, 0].real, 'Pulses':pulsesN,
                     'Time':timemat}
        return(DataLags, NoiseLags)

#%% Make Lag dict to an iono container
def lagdict2ionocont(DataLags, NoiseLags, sensdict, simparams, time_vec):
    """This function will take the data and noise lags and create an instance of the
    Ionocontanier class. This function will also apply the summation rule to the lags.
    Inputs
    DataLags - A dictionary """
    # Pull in Location Data
    angles = simparams['angles']
    ang_data = sp.array([[iout[0], iout[1]] for iout in angles])
    rng_vec = simparams['Rangegates']
    n_samps = len(rng_vec)
    # pull in other data
    pulse = simparams['Pulse']
    p_samps = len(pulse)
    pulsewidth = p_samps*sensdict['t_s']
    txpower = sensdict['Pt']
    if sensdict['Name'].lower() in ['risr', 'pfisr', 'risr-n']:
        Ksysvec = sensdict['Ksys']
    else:

        beamlistlist = sp.array(simparams['outangles']).astype(int)
        inplist = sp.array([i[0] for i in beamlistlist])
        Ksysvec = sensdict['Ksys'][inplist]
        ang_data_temp = ang_data.copy()
        ang_data = sp.array([ang_data_temp[i].mean(axis=0) for i in beamlistlist])

    sumrule = simparams['SUMRULE']
    rng_vec2 = simparams['Rangegatesfinal']
    Nrng2 = len(rng_vec2)
    minrg = p_samps-1+sumrule[0].min()
    maxrg = Nrng2+minrg

    # Copy the lags
    lagsData = DataLags['ACF'].copy()
    # Set up the constants for the lags so they are now
    # in terms of density fluxtuations.
    angtile = sp.tile(ang_data, (Nrng2, 1))
    rng_rep = sp.repeat(rng_vec2, ang_data.shape[0], axis=0)
    coordlist = sp.zeros((len(rng_rep), 3))
    [coordlist[:, 0], coordlist[:, 1:]] = [rng_rep, angtile]
    (Nt, Nbeams, Nrng, Nlags) = lagsData.shape

    # make a range average to equalize out the conntributions from each gate

    plen2 = int(sp.floor(float(p_samps-1)/2))
    samps = sp.arange(0, p_samps, dtype=int)

    rng_ave = sp.zeros((Nrng, p_samps))

    for isamp in range(plen2, Nrng+plen2):
        for ilag in range(p_samps):
            toplag = int(sp.floor(float(ilag)/2))
            blag = int(sp.ceil(float(ilag)/2))
            if toplag == 0:
                sampsred = samps[blag:]
            else:
                sampsred = samps[blag:-toplag]
            cursamps = isamp-sampsred

            keepsamps = sp.logical_and(cursamps >= 0, cursamps < Nrng)
            cursamps = cursamps[keepsamps]
            rng_samps = rng_vec[cursamps]**2*1e6
            keepsamps2 = rng_samps > 0
            if keepsamps2.sum() == 0:
                continue
            rng_samps = rng_samps[keepsamps2]
            rng_ave[isamp-plen2, ilag] = 1./(sp.mean(1./(rng_samps)))
    rng_ave_temp = rng_ave.copy()
    if simparams['Pulsetype'].lower() is 'barker':
        rng_ave_temp = rng_ave[:,0][:,sp.newaxis]
    # rng_ave = rng_ave[int(sp.floor(plen2)):-int(sp.ceil(plen2))]
    # rng_ave = rng_ave[minrg:maxrg]
    rng3d = sp.tile(rng_ave_temp[sp.newaxis, sp.newaxis, :, :], (Nt, Nbeams, 1, 1))
    ksys3d = sp.tile(Ksysvec[sp.newaxis,:, sp.newaxis, sp.newaxis], (Nt, 1, Nrng, Nlags))

    # rng3d = sp.tile(rng_ave[:, sp.newaxis, sp.newaxis, sp.newaxis], (1, Nlags, Nt, Nbeams))
    # ksys3d = sp.tile(Ksysvec[sp.newaxis, sp.newaxis, sp.newaxis, :], (Nrng2, Nlags, Nt, 1))
    radar2acfmult = rng3d/(pulsewidth*txpower*ksys3d)
    pulses = sp.tile(DataLags['Pulses'][:, :, sp.newaxis, sp.newaxis], (1, 1, Nrng, Nlags))
    time_vec = time_vec[:Nt]
    # Divid lags by number of pulses
    lagsData = lagsData/pulses
    # Set up the noise lags and divid out the noise.
    lagsNoise = NoiseLags['ACF'].copy()
    lagsNoise = sp.mean(lagsNoise, axis=2)
    pulsesnoise = sp.tile(NoiseLags['Pulses'][:, :, sp.newaxis], (1, 1, Nlags))
    lagsNoise = lagsNoise/pulsesnoise
    lagsNoise = sp.tile(lagsNoise[:, :, sp.newaxis, :], (1, 1, Nrng, 1))

    # multiply the data and the sigma by inverse of the scaling from the radar
    lagsData = lagsData*radar2acfmult
    lagsNoise = lagsNoise*radar2acfmult
    # Apply summation rule
    # lags transposed from (time,beams,range,lag)to (range,lag,time,beams)
    lagsData = sp.transpose(lagsData, axes=(2, 3, 0, 1))
    lagsNoise = sp.transpose(lagsNoise, axes=(2, 3, 0, 1))
    lagsDatasum = sp.zeros((Nrng2, Nlags, Nt, Nbeams), dtype=lagsData.dtype)
    lagsNoisesum = sp.zeros((Nrng2, Nlags, Nt, Nbeams), dtype=lagsNoise.dtype)

    for irngnew, irng in enumerate(sp.arange(minrg, maxrg)):
        for ilag in range(Nlags):
            lsumtemp = lagsData[irng+sumrule[0, ilag]:irng+sumrule[1, ilag]+1, ilag].sum(axis=0)
            lagsDatasum[irngnew, ilag] = lsumtemp
            nsumtemp = lagsNoise[irng+sumrule[0, ilag]:irng+sumrule[1, ilag]+1, ilag].sum(axis=0)
            lagsNoisesum[irngnew, ilag] = nsumtemp
    # subtract out noise lags
    lagsDatasum = lagsDatasum-lagsNoisesum

    # Put everything in a parameter list
    Paramdata = sp.zeros((Nbeams*Nrng2, Nt, Nlags), dtype=lagsData.dtype)
    # Put everything in a parameter list
    # transpose from (range,lag,time,beams) to (beams,range,time,lag)
    # lagsDatasum = lagsDatasum*radar2acfmult
    # lagsNoisesum = lagsNoisesum*radar2acfmult
    lagsDatasum = sp.transpose(lagsDatasum, axes=(3, 0, 2, 1))
    lagsNoisesum = sp.transpose(lagsNoisesum, axes=(3, 0, 2, 1))

    # multiply the data and the sigma by inverse of the scaling from the radar
    # lagsDatasum = lagsDatasum*radar2acfmult
    # lagsNoisesum = lagsNoisesum*radar2acfmult

    # Calculate a variance using equation 2 from Hysell's 2008 paper. Done use full covariance matrix because assuming nearly diagonal.
    # Get the covariance matrix
    pulses_s = sp.transpose(pulses, axes=(1, 2 ,0, 3))[:, :Nrng2]
    Cttout = makeCovmat(lagsDatasum, lagsNoisesum, pulses_s, Nlags)

    Paramdatasig = sp.zeros((Nbeams*Nrng2 ,Nt, Nlags, Nlags), dtype=Cttout.dtype)

    curloc = 0
    for irng in range(Nrng2):
        for ibeam in range(Nbeams):
            Paramdata[curloc] = lagsDatasum[ibeam, irng].copy()
            Paramdatasig[curloc] = Cttout[ibeam, irng].copy()
            curloc += 1
    ionodata = IonoContainer(coordlist, Paramdata, times=time_vec, ver=1,
                             paramnames=sp.arange(Nlags)*sensdict['t_s'])
    ionosigs = IonoContainer(coordlist, Paramdatasig, times=time_vec, ver=1,
                             paramnames=sp.arange(Nlags**2).reshape(Nlags, Nlags)*sensdict['t_s'])
    return (ionodata, ionosigs)

def makeCovmat(lagsDatasum, lagsNoisesum, pulses_s, Nlags):
    """
        Makes the covariance matrix for the lags given the noise acf and number
        of pulses.
    """
    axvec = sp.roll(sp.arange(lagsDatasum.ndim), 1)
    # Get the covariance matrix
    R = sp.transpose(lagsDatasum/sp.sqrt(2.*pulses_s), axes=axvec)
    Rw = sp.transpose(lagsNoisesum/sp.sqrt(2.*pulses_s), axes=axvec)
    l = sp.arange(Nlags)
    T1, T2 = sp.meshgrid(l, l)
    R0 = R[sp.zeros_like(T1)]
    Rw0 = Rw[sp.zeros_like(T1)]
    Td = sp.absolute(T1-T2)
    Tl = T1>T2
    R12 = R[Td]
    R12[Tl] = sp.conjugate(R12[Tl])
    Rw12 = Rw[Td]
    Rw12[Tl] = sp.conjugate(Rw12[Tl])
    Ctt = R0*R12+R[T1]*sp.conjugate(R[T2])+Rw0*Rw12+Rw[T1]*sp.conjugate(Rw[T2])
    avecback = sp.roll(sp.arange(Ctt.ndim), -2)
    Cttout = sp.transpose(Ctt, avecback)
    return Cttout


# -*- coding: utf-8 -*-
"""
@author: Johannes Elfner <johannes.elfner@googlemail.com>
Date: Aug 2018
"""

import numpy as np

from .pipe import Pipe
from ..precomp_funs import pipe1D_branched_diff

class PipeBranched(Pipe):
    """
    type: Single Pipe

    Can be added to the simulation environment by using the following method:
        .add_part(Pipe, name, volume=..., grid_points=..., outer_diameter=...,
        shell_thickness=...)
    Part creation parameters:
    -------------------------
    Pipe : class name
        This parameter can\'t be changed and needs to be passed exactly like
        this.
    name : string
        Pipe identifier as string. Needs to be unique.
    length : integer, float
        Thermal energy storage volume in [m^3].
    grid_points : integer
        Number of grid points to discretize the pipe with.

    For part initialization the following additional parameters need to be
    passed:
        .init_part(insulation_thickness=..., insulation_lambda=...)

    Part initialization parameters:
    -------------------------------
    insulation_thickness : float
        in [m]
    insulation_lambda : float
        in [W/(m*K)]

    """

    def __init__(self, name, master_cls, **kwargs):
        self.constr_type = 'PipeBranched'  # define construction type
        # since this part is a subclass of Pipe, initialize Pipe:
        super().__init__(
            name, master_cls, **kwargs, constr_type=self.constr_type
        )

        # preallocate mass flow grids:
        self.dm = np.zeros_like(self.T)
        self._dm_top = np.zeros_like(self.T)
        self._dm_bot = np.zeros_like(self.T)
        # set dm char so that all ports face inwards (more convenient when
        # adding new ports):
        self.dm_char = tuple(('in', 'in'))

        # if ports can be added to this part:
        self.can_add_ports = True

        # add ports
        self._add_ports(**kwargs)

        # construct all port shape dependent vars:
        dummy_var = list(self.port_names)
        for i in range(self.port_num):
            dummy_var[i] = self.name + ';' + dummy_var[i]
        self._own_ports = tuple(dummy_var)
        # preallocate port values:
        self._port_vals = np.zeros(self.port_num)
        # preallocate grids for port connection parameters:
        # cross section area of wall of connected pipe, fluid cross section
        # area of, gridspacing and lambda of wall of connected pipe
        Tpshp = self._T_port.shape
        self._A_wll_conn_p = np.zeros(Tpshp)
        self._A_fld_conn_p = np.zeros(Tpshp)
        self._A_p_wll_mean = np.zeros(Tpshp)
        self._A_p_fld_mean = np.zeros(Tpshp)
        self._port_gsp = np.zeros(Tpshp)
        self._lam_wll_conn_p = np.full_like(self._T_port, 1e-2)
        self._lam_port_fld = np.zeros(Tpshp)
        self._lam_fld_own_p = np.zeros(Tpshp)
        self._UA_port_fld = np.zeros(Tpshp)
        self._UA_port_wll = np.zeros(Tpshp)

        # if the topology construction method has to stop when it reaches the
        # part to solve more ports from other sides before completely solving
        # the massflow of it. This will be set to false as soon as only one
        # port to solve is remaining:
        self.break_topology = True
        # count how many ports are still open to be solved by topology. If
        # break topology is True, this is used to set it to False if 1 is
        # reached.
        self._cnt_open_prts = self.port_num
        # determine if part has the capability to affect massflow (dm) by
        # diverting flow through ports or adding flow through ports:
        self.affect_dm = False
        # if the massflow (dm) has the same value in all cells of the part
        # (respectively in each flow channel for parts with multiple flows):
        self.dm_invariant = False

        # IMPORTANT: THIS VARIABLE **MUST NOT BE INHERITED BY SUB-CLASSES**!!
        # If sub-classes are inherited from this part, this bool checker AND
        # the following variables MUST BE OVERWRITTEN!
        # ist the diff function fully njitted AND are all input-variables
        # stored in a container?
        self._diff_fully_njit = False
        # self._diff_njit = pipe1D_diff  # handle to njitted diff function
        # input args are created in simenv _create_diff_inputs method

    def init_part(self, **kwargs):
        # since this part is a subclass of Pipe, call init_part of Pipe:
        super().init_part(**kwargs)

        # add massflow grid argument to input args at correct position:
        self._input_arg_names_sorted.insert(6, 'dm')

    def _get_flow_routine(  # overwrite pipe routine, since new ports exist
        self, port, parent_port=None, subnet=False, **kwargs
    ):
        """
        Returns the massflow calculation routine for the port of the current
        part to the topology construction. The massflow calculation routine has
        to look like:

        routine = (memory_view_to_target_port,
                   operation_id,
                   memory_view_to_port1, memory_view_to_port2, ...)

        with target_port being the port which has to be calculated and port1
        and port2 being the other/source ports which **don't** have to be
        calculated with this routine! These source ports **must be given**
        when the routine is called.

        Parameters:
        -----------
        port : string
            Port name of the port which shall be calculated (target port).

        """

        # if more than one port is open to be solved, it will always be solved
        # by passing on the value from the connected parent part
        if self._cnt_open_prts != 1:
            # get topology connection conditions (target index, source
            # part/port identifiers, source index and algebraic sign for passed
            # massflow):
            (
                trgt_idx,
                src_part,
                src_port,
                src_idx,
                alg_sign,
            ) = self._get_topo_cond(port, parent_port)
            # the operation id is always 0 or -1, depending on the
            # algebraic sign of the ports which are connected, when there is
            # more than one port remaining to be solved. This means the value
            # of connected port is either passed on (0) or the negative value
            # of it is used (-1):
            if alg_sign == 'positive':
                operation_id = 0
            else:
                operation_id = -1
            # add operation instructions to tuple (memory view to target
            # massflow array cell, operation id and memory view source port's
            # massflow array cells)
            op_routine = (
                self._dm_io.reshape(-1)[trgt_idx],
                operation_id,
                self._models.parts[src_part]._dm_io.reshape(-1)[src_idx],
            )
        # only one port remaining to be solved:
        else:
            # There are two possibilities to get the last port:
            # 1. - Get it from a connected port -> parent_port must be given
            # 2. - Calculate it from the missing massflow through the other
            #      already solved own ports. - > no parent_port needed
            # Method 1 is highly preferred, since it only involves accessing
            # the connected port's massflow array cell, while method 2 requires
            # summing up all other ports.
            if parent_port is not None:
                # Method 1!
                # get topology connection conditions (target index, source
                # part/port identifiers, source index and algebraic sign for
                # passed massflow):
                (
                    trgt_idx,
                    src_part,
                    src_port,
                    src_idx,
                    alg_sign,
                ) = self._get_topo_cond(port, parent_port)
                # passing values is always 0 or -1, depending on pos.
                # massflow directions of connected ports
                if alg_sign == 'positive':
                    # in/out connected
                    operation_id = 0
                else:
                    # in/in or out/out connected
                    operation_id = -1
                # add operation instructions to tuple (memory view to target
                # massflow array cell, operation id and memory view source
                # port's massflow array cells)
                op_routine = (
                    self._dm_io.reshape(-1)[trgt_idx],
                    operation_id,
                    self._models.parts[src_part]._dm_io.reshape(-1)[src_idx],
                )
            else:
                # Method 2!
                # get port index slice of target port to create memory view:
                trgt_idx = self._get_topo_cond(port)
                # operation id for this case is always -1, since the positive
                # massflow direction of all ports is facing inwards
                operation_id = -1
                # add operation instructions to tuple (memory view to target
                # massflow array cell, operation id and memory view to ALL
                # source/other port's massflow array cells of the TES)
                op_routine = (
                    self._dm_io.reshape(-1)[trgt_idx],
                    operation_id,
                )
                # add memory view to source massflow array cell(s)
                # therefore loop over already solved ports
                for slvd_prt in self._solved_ports:
                    # get solved port index as slice to create memory views
                    slvd_prt_idx = self._get_topo_cond(slvd_prt)
                    # add memory view to this ports massflow
                    op_routine += (self._dm_io.reshape(-1)[slvd_prt_idx],)

        # update solved ports list and counter stop break:
        self._solved_ports.append(port)
        self._cnt_open_prts = self.port_num - len(self._solved_ports)
        # set break topology to False if enough ports have been solved:
        self.break_topology = True if self._cnt_open_prts > 0 else False
        # remove part from hydr_comps if completely solved:
        if self._cnt_open_prts == 0:
            self._models._hydr_comps.remove(self.name)

        # save topology parameters to dict for easy information lookup:
        net = 'Subnet' if subnet else 'Flownet'
        operation_routine = (
            'Negative of sum'
            if operation_id == -1
            else 'Pass on value'
            if operation_id == 0
            else 'Error'
        )
        src_part = src_part if self._cnt_open_prts > 1 else self.name
        source_ports = (
            parent_port
            if operation_id == 0
            else tuple(set(self.port_names) - set(port))
        )
        kwargs['pump_side'] = (
            'decoupled from pump by hydraulic compensator'
            if kwargs['pump_side'] == 'sub net - undefined'
            else kwargs['pump_side']
        )
        # add port dict for current port and fill it:
        if port not in self.info_topology:
            self.info_topology[port] = dict()
        self.info_topology[port].update(
            {
                'Net': net,
                'Massflow': self._dm_io.reshape(-1)[trgt_idx],
                'Calculation routine': operation_routine,
                'Source part': src_part,
                'Source port(s)': source_ports,
                'Connected part': (
                    self._models.port_links[self.name + ';' + port].split(';')[
                        0
                    ]
                ),
                'Connected port': (
                    self._models.port_links[self.name + ';' + port].split(';')[
                        1
                    ]
                ),
                'Parent pump/part': kwargs['parent_pump'],
                'Pump side': kwargs['pump_side'],
            }
        )

        return op_routine

    def get_diff(self, timestep):
        """
        This function just calls a jitted calculation function.

        """

        pipe1D_branched_diff(*self._input_args, timestep)

        return self.dT_total


#!/usr/bin/env python
'''
mcu: Modeling and Crystallographic Utilities
Copyright (C) 2019 Hung Q. Pham. 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.

Email: Hung Q. Pham <pqh3.14@gmail.com>
'''

import numpy as np
from ..utils import plot, str_format
from ..utils.misc import check_exist
from ..cell import utils as cell_utils
from ..cell import cell
from ..vasp import const
from . import crystal_io
from . import utils as crystal_utils

class main(cell.main, plot.main):
    def __init__(self,  prefix="outfile"):
        assert prefix is not None, "Provide a prefix name for your project, for example, prefix.scf.out, prefix.band.out, etc."
        self.prefix = prefix
        self.get_info()

############ General #################
    def get_info(self, filename=None):
        '''Extract basis information from the vasprun.xml'''
        if filename is None:
            if check_exist(self.prefix + ".out"):
                filename = self.prefix + ".out"
            else:
                assert 0, "Cannot find any prefix.out file"

        data = crystal_io.read_out(filename)
        self.nelec = data['nelec']
        self.nelec_core = data['nelec core']
        self.nao = data['nao']
        self.basis = data['basis']
        self.atom  = data['atom']
        self.natom  = len(self.atom)
        self.element = list(dict.fromkeys(self.atom ))
        self.kpts = data['kpts']
        if data['total energy - final'] is None:
            self.e_tot = data['total energy - sp']
        else:
            self.e_tot = data['total energy - final']

        # Make a cell object in the spglib format
        lattice = data['lattice']
        positions = data['atom_position']
        numbers = cell_utils.convert_atomtype(self.atom)
        self.cell_init = (lattice, positions, numbers)
        self.cell = self.cell_init

############ Preprocessing #################
    def make_DOS_input(self, lm='None', filename=None):
        '''Giving a lm string, export a d3 file for DOS calculation.
        '''
        if lm is None: lm = 'spdf'
        if filename is None: filename = self.prefix + '.d3'
        formatted_atom, formatted_lm = str_format.general_lm(lm)

        pDOS_orb_idx = []
        for i, atoms in enumerate(formatted_atom):
            proj_orbs = []
            for j, atom in enumerate(atoms):
                # Locate the atom
                if atom is None:
                    idx_atom = np.arange(len(self.atom))
                else:
                    atom_, id = str_format.format_atom(atom)
                    assert atom_ in self.element, "This is wrong: " + atom + ". Check the lm string/list. Atom is must be in the element list: " + " ".join(self.element)
                    available_atom = [(n, atm) for n, atm in enumerate(self.atom) if atm == atom_]
                    natom = len(available_atom)
                    if id is not None:
                        assert id <= natom, "This is wrong: " + atom + ". Check the lm string/list. Atom id is must be <= " + str(natom) + " for: " + atom_

                    idx_atom = []
                    nspecies = self.atom.count(atom_)
                    nwfc = nspecies // natom
                    count = 0
                    for n, atm in enumerate(self.atom):
                        if atm == atom_:
                            if id is None:
                                idx_atom.append(n)
                            elif count // nwfc == id - 1:
                                idx_atom.append(n)
                            count += 1

                if atom is None:
                    atom_ = ""
                    available_orbs = crystal_utils.basis_short2long(self.basis)
                    available_orbs = sum(list(available_orbs.values()),[])
                else:
                    atom_, id = str_format.format_atom(atom)
                    available_orbs = crystal_utils.basis_short2long(self.basis)[atom_]

                available_orbs = list(dict.fromkeys(['None'] + available_orbs))

                for atm_idx in idx_atom:
                    for orb in formatted_lm[i][j]:
                        assert str(orb) in available_orbs, "This is wrong: " + str(orb) + ". Check the lm string/list. Available basis functions " + atom_ +  " are: " +  ", ".join(available_orbs)
                        idx = crystal_utils.orb_index(self.atom, self.basis, atm_idx, orb)
                        proj_orbs.append(idx)
            proj_orbs = sum(proj_orbs,[])
            proj_orbs.sort()
            pDOS_orb_idx.append(proj_orbs)

        # Export d3 file using pDOS_orb_idx list
        nDOS = len(pDOS_orb_idx)
        npoints = 1000
        first_band = self.nelec_core//2 + 1
        last_band = self.nao
        store_dos = 1
        legendre = 14
        printing = 0
        with open(filename, 'w') as f:
            f.write('NEWK\n')
            f.write('12 12\n')
            f.write('1 0\n')
            f.write('DOSS\n')
            f.write('%d %d %d %d %d %d %d\n' % (nDOS, npoints, first_band, last_band, store_dos, legendre, printing))
            for dos in pDOS_orb_idx:
                norb = len(dos)
                orb_string = [str(orb) for orb in dos]
                f.write('%d %s\n' % (norb, ' '.join(orb_string)))
            f.write('END\n')

############ Plotting #################
    def get_band(self, filename=None, phonon=False, gamma_correct=False, threshold=8.06554):
        '''Make a band from from
           NOTE: the proj_kpath is computed from the dk and nkp. Due to the round out error in f25, the computed high symmetric k-point coordinates won't be exactly the same as values obtained from *.BAND file.

           if phonon == True: the phonon band is read. Three small negative acoustic modes at Gamma can be removed by using the gamma_correct keyword and threshold 1 meV = 8.06554 cm^-1
        '''
        if filename is None:
            if check_exist(self.prefix + ".f25"):
                filename = self.prefix + ".f25"
            elif check_exist("fort.25"):
                filename = "fort.25"
                print('Found fort.25. Band structure is extracted from fort.25')
            else:
                assert 0, "Cannot find " + self.prefix + ".f25" + " or fort.25. Check if you has band structure file"

        data = crystal_io.read_f25(filename)[0]
        assert data is not None, "Cannot find BAND information, check if you have generated fort.25: " + filename
        temp = []
        ihferm_list = []
        proj_kpath = []
        sym_kpoint_coor = [0.0]
        shift = 0
        for block in data:
            ihferm = block['ihferm']
            nband = block['nband']
            nkp = block['nkp']
            dum = block['dum']
            dk = block['dk']
            efermi = block['efermi']
            emin = block['emin']
            emax = block['emax']
            ivalues = block['values']
            eigenvals = block['eigenvals']

            # Compute hight symmetric k-pts coordinates
            temp.append(eigenvals.reshape(-1, nband))
            path = np.arange(nkp)*dk + shift
            proj_kpath.append(path)
            shift = path[-1] + dk
            sym_kpoint_coor.append(path[-1])

        if ihferm == 0 or ihferm == 2:
            band = np.float64([np.vstack(temp)])
            proj_kpath = np.hstack(proj_kpath)
        elif ihferm == 1 or ihferm == 3:
            nblock = len(temp) // 2
            band_up = np.vstack(temp[:nblock])
            band_down = np.vstack(temp[nblock:])
            band = np.float64([band_up, band_down])
            proj_kpath = np.hstack(proj_kpath[:nblock])

        sym_kpoint_coor = np.float64(sym_kpoint_coor)

        if phonon:
            if gamma_correct:
                # Should be checked first before using the correction to make sure it is really acoustic phonon modes
                nspin, nkpts, nband = band.shape
                band = band.flatten()
                imag_mode_idx = band < 0.0
                imag_mode = band[imag_mode_idx]
                imag_mode_idx2 = imag_mode > -threshold
                imag_mode[imag_mode_idx2] = 0.0
                band[imag_mode_idx] = imag_mode
                band = band.reshape(nspin, nkpts, nband)
        else:
            band = const.AUTOEV * band
            efermi = const.AUTOEV * efermi

        return band, proj_kpath, sym_kpoint_coor, efermi

    def get_bandgap(self, filename=None, efermi=None):
        '''Get the bandgap'''
        if filename is None:
            if check_exist(self.prefix + ".f25"):
                filename = self.prefix + ".f25"
            elif check_exist("fort.25"):
                filename = "fort.25"
                print('Found fort.25. Band structure is extracted from fort.25')
            else:
                assert 0, "Cannot find " + self.prefix + ".f25" + " or fort.25. Check if you has band structure file"

        band, proj_kpath, sym_kpoint_coor, efermi_ = self.get_band(filename)
        if efermi is None: efermi = efermi_
        nspin, nkpts, nbands = band.shape
        for spin in range(nspin):
            print('Spin:', spin)
            CBM = None
            for bandth in range(nbands):
                shifted_band = band[spin,:,bandth] - efermi
                if (shifted_band > 0.0).all() == True:
                    CBM = band[spin,:, bandth]
                    VBM = band[spin,:, bandth -1]
                    break
                elif ((shifted_band < 0.0).any() == True) and ((shifted_band > 0.0).any() == True):
                    print("This is a metal")
                    break

            if CBM is not None:
                vbm_idx = np.argmax(VBM)
                cbm_idx = np.argmin(CBM)
                bandgap = CBM[cbm_idx] - VBM[vbm_idx]
                direct = False
                if vbm_idx == cbm_idx: direct = True

                # TODO: the kpath_frac currently cannot be obtained form f25.
                # Other outputs are needed
                # kpath_frac = self.cp2k_io.kpath_frac[set_block]
                # print('  E(VBM) = %7.4f at k = [%6.4f,%6.4f,%6.4f]' % (VBM[vbm_idx],
                                                                # kpath_frac[vbm_idx,0], kpath_frac[vbm_idx,1], kpath_frac[vbm_idx,2]))
                # print('  E(CBM) = %7.4f at k = [%6.4f,%6.4f,%6.4f]' % (CBM[cbm_idx],
                                                                # kpath_frac[cbm_idx,0], kpath_frac[cbm_idx,1], kpath_frac[cbm_idx,2]))
                if direct == True:
                    print('  Direct bandgap   : %6.3f' % (bandgap))
                else:
                    print('  Indirect bandgap : %6.3f' % (bandgap))
                    gap1 = CBM[vbm_idx] - VBM[vbm_idx]
                    gap2 = CBM[cbm_idx] - VBM[cbm_idx]
                    direct_gap = min(gap1, gap2)
                    print('  Direct bandgap   : %6.3f' % (direct_gap))

    def _generate_band(self, filename=None, efermi=None, spin=0, klabel=None):
        '''Processing/collecting the band data before the plotting function
        '''
        if filename is None:
            if check_exist(self.prefix + ".f25"):
                filename = self.prefix + ".f25"
            elif check_exist("fort.25"):
                filename = "fort.25"
                print('Found fort.25. Band structure is extracted from fort.25')
            else:
                assert 0, "Cannot find " + self.prefix + ".f25" + " or fort.25. Check if you has band structure file"

        band, proj_kpath, sym_kpoint_coor, efermi_ = self.get_band(filename)
        if efermi is None: efermi = efermi_
        band = band[spin] - efermi

        # Process the labels for k-point
        if klabel is not None:
            klabel, coor_kpts = str_format.format_klabel(klabel)
            assert len(klabel) == len(sym_kpoint_coor), "The number of k label must be " + str(len(sym_kpoint_coor))

        return band, proj_kpath, sym_kpoint_coor, klabel

    def _generate_phononband(self, unit="CM", gamma_correct=False, threshold=8.06554, spin=0, klabel=None):
        '''Processing/collecting the phnon band (cm^-1) data before the plotting function
           Unit: CM = CM^-1, THZ = THz, MEV = meV
        '''
        band, proj_kpath, sym_kpoint_coor, efermi_ = self.get_band(phonon=True, gamma_correct=gamma_correct, threshold=threshold)
        if unit.lower() == "thz":
            band = const.CMTOTHZ * band
        elif unit.lower() == "mev":
            band = const.CMTOMEV * band

        # Process the labels for k-point
        if klabel is not None:
            klabel, coor_kpts = str_format.format_klabel(klabel)
            assert len(klabel) == len(sym_kpoint_coor), "The number of k label must be " + str(len(sym_kpoint_coor))

        return band[spin], proj_kpath, sym_kpoint_coor, klabel

    def _generate_dos(self, filename=None, efermi=None, spin=0, lm=None):
        '''Processing/collecting the DOS data before the plotting function

            TDOS dimensions: [spin , [E(eV), tdos(E)]]

            spin            : spin of DOS.
            lm              : string or a list of string, e.g. 'Ni:s' or ['Ni:s','C:s,px,pz']
        '''
        if lm is None:
            lm = [atom+':s,p,d' for atom in self.element]

        if filename is None:
            if check_exist(self.prefix + ".f25"):
                filename = self.prefix + ".f25"
            elif check_exist("fort.25"):
                filename = "fort.25"
                print('Found fort.25. DOS is extracted from fort.25')
            else:
                assert 0, "Cannot find " + self.prefix + ".f25" + " or fort.25. Check if you has DOS file"

        # Compute pDOS
        data = crystal_io.read_f25(filename)[1]
        assert data is not None, "Cannot find DOS information, check if you have generated fort.25: " + filename

        dos_data = []
        for block in data:
            ihferm = block['ihferm']
            nrow = block['nrow']
            ncol = block['ncol']
            dx = block['dx']
            dy = block['dy']
            efermi = const.AUTOEV * block['efermi']
            emin = block['emin']
            emax = block['emax']
            ivalues = block['values']
            dos_data.append(block['dos'])

        if ihferm == 0 or ihferm == 2:
            band = np.float64([np.vstack(temp)])
            proj_kpath = np.hstack(proj_kpath)
        elif ihferm == 1 or ihferm == 3:
            nblock = len(temp) // 2
            band_up = np.vstack(temp[:nblock])
            band_down = np.vstack(temp[nblock:])
            band = np.float64([band_up, band_down])
            proj_kpath = np.hstack(proj_kpath[:nblock])

        # Compute energy points:
        energy_path = const.AUTOEV * (np.arange(ncol)*dy + emax) - efermi
        tdos = np.float64([energy_path, dos_data[-1]]).T
        pdos = None

        if len(dos_data) != 1:
            pdos = np.asarray(dos_data[:-1]).T

        return tdos, pdos