"""
Functions are modified on top of GFLA.
GFLA's license: https://github.com/RenYurui/Global-Flow-Local-Attention/blob/master/LICENSE.md

"""

import torch
import torch.nn as nn
import torchvision.models as models
import torch.nn.functional as F
import os
import torchvision.transforms as transforms
import numpy as np

class GANLoss(nn.Module):
    """Define different GAN objectives.
    The GANLoss class abstracts away the need to create the target label tensor
    that has the same size as the input.
    """

    def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0):
        """ Initialize the GANLoss class.
        Parameters:
            gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp.
            target_real_label (bool) - - label for a real image
            target_fake_label (bool) - - label of a fake image
        Note: Do not use sigmoid as the last layer of Discriminator.
        LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss.
        """
        super(GANLoss, self).__init__()
        self.register_buffer('real_label', torch.tensor(target_real_label))
        self.register_buffer('fake_label', torch.tensor(target_fake_label))
        self.gan_mode = gan_mode
        if gan_mode == 'lsgan':
            self.loss = nn.MSELoss()
        elif gan_mode == 'vanilla':
            self.loss = nn.BCEWithLogitsLoss()
        elif gan_mode in ['wgangp']:
            self.loss = None
        else:
            raise NotImplementedError('gan mode %s not implemented' % gan_mode)

    def get_target_tensor(self, prediction, target_is_real):
        """Create label tensors with the same size as the input.
        Parameters:
            prediction (tensor) - - tpyically the prediction from a discriminator
            target_is_real (bool) - - if the ground truth label is for real images or fake images
        Returns:
            A label tensor filled with ground truth label, and with the size of the input
        """

        if target_is_real:
            target_tensor = self.real_label
        else:
            target_tensor = self.fake_label
        return target_tensor.expand_as(prediction)

    def __call__(self, prediction, target_is_real):
        """Calculate loss given Discriminator's output and grount truth labels.
        Parameters:
            prediction (tensor) - - tpyically the prediction output from a discriminator
            target_is_real (bool) - - if the ground truth label is for real images or fake images
        Returns:
            the calculated loss.
        """
        if self.gan_mode in ['lsgan', 'vanilla']:
            target_tensor = self.get_target_tensor(prediction, target_is_real)
            loss = self.loss(prediction, target_tensor)
        elif self.gan_mode == 'wgangp':
            if target_is_real:
                loss = -prediction.mean()
            else:
                loss = prediction.mean()
        return loss

def cal_gradient_penalty(netD, real_data, fake_data, device, type='mixed', constant=1.0, lambda_gp=10.0):
    """Calculate the gradient penalty loss, used in WGAN-GP paper https://arxiv.org/abs/1704.00028
    Arguments:
        netD (network)              -- discriminator network
        real_data (tensor array)    -- real images
        fake_data (tensor array)    -- generated images from the generator
        device (str)                -- GPU / CPU: from torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')
        type (str)                  -- if we mix real and fake data or not [real | fake | mixed].
        constant (float)            -- the constant used in formula ( | |gradient||_2 - constant)^2
        lambda_gp (float)           -- weight for this loss
    Returns the gradient penalty loss
    """
    if lambda_gp > 0.0:
        if type == 'real':   # either use real images, fake images, or a linear interpolation of two.
            interpolatesv = real_data
        elif type == 'fake':
            interpolatesv = fake_data
        elif type == 'mixed':
            alpha = torch.rand(real_data.shape[0], 1, device=device)
            alpha = alpha.expand(real_data.shape[0], real_data.nelement() // real_data.shape[0]).contiguous().view(*real_data.shape)
            interpolatesv = alpha * real_data + ((1 - alpha) * fake_data)
        else:
            raise NotImplementedError('{} not implemented'.format(type))
        interpolatesv.requires_grad_(True)
        disc_interpolates = netD(interpolatesv)
        gradients = torch.autograd.grad(outputs=disc_interpolates, inputs=interpolatesv,
                                        grad_outputs=torch.ones(disc_interpolates.size()).to(device),
                                        create_graph=True, retain_graph=True, only_inputs=True)
        gradients = gradients[0].view(real_data.size(0), -1)  # flat the data
        gradient_penalty = (((gradients + 1e-16).norm(2, dim=1) - constant) ** 2).mean() * lambda_gp        # added eps
        return gradient_penalty, gradients
    else:
        return 0.0, None

class MultiAffineRegularizationLoss(nn.Module):
    def __init__(self, kz_dic):
        super(MultiAffineRegularizationLoss, self).__init__()
        self.kz_dic=kz_dic
        self.method_dic={}
        for key in kz_dic:
            instance = AffineRegularizationLoss(kz_dic[key])
            self.method_dic[key] = instance
        self.layers = sorted(kz_dic, reverse=True)

    def __call__(self, flow_fields):
        loss=0
        for i in range(len(flow_fields)):
            method = self.method_dic[self.layers[i]]
            loss += method(flow_fields[i])
        return loss

class AffineRegularizationLoss(nn.Module):
    """docstring for AffineRegularizationLoss"""
    # kernel_size: kz
    def __init__(self, kz):
        super(AffineRegularizationLoss, self).__init__()
        self.kz = kz
        self.criterion = torch.nn.L1Loss()
        from models.networks.block_extractor.block_extractor   import BlockExtractor
        from models.networks.local_attn_reshape.local_attn_reshape   import LocalAttnReshape

        self.extractor = BlockExtractor(kernel_size=kz)
        self.reshape = LocalAttnReshape()

        temp = np.arange(kz)
        A = np.ones([kz*kz, 3])
        A[:, 0] = temp.repeat(kz)
        A[:, 1] = temp.repeat(kz).reshape((kz,kz)).transpose().reshape(kz**2)
        AH = A.transpose()
        k = np.dot(A, np.dot(np.linalg.inv(np.dot(AH, A)), AH)) - np.identity(kz**2) #K = (A((AH A)^-1)AH - I)
        self.kernel = np.dot(k.transpose(), k)
        self.kernel = torch.from_numpy(self.kernel).unsqueeze(1).view(kz**2, kz, kz).unsqueeze(1)

    def __call__(self, flow_fields):
        grid = self.flow2grid(flow_fields)

        grid_x = grid[:,0,:,:].unsqueeze(1)
        grid_y = grid[:,1,:,:].unsqueeze(1)
        weights = self.kernel.type_as(flow_fields)
        #import pdb; pdb.set_trace()
        loss_x = self.calculate_loss(grid_x, weights)
        loss_y = self.calculate_loss(grid_y, weights)
        return loss_x+loss_y

    def calculate_loss(self, grid, weights):
        results = nn.functional.conv2d(grid, weights)   # KH K B [b, kz*kz, w, h]
        b, c, h, w = results.size()
        kernels_new = self.reshape(results, self.kz)
        f = torch.zeros(b, 2, h, w).type_as(kernels_new) + float(int(self.kz/2))
        grid_H = self.extractor(grid, f)
        result = torch.nn.functional.avg_pool2d(grid_H*kernels_new, self.kz, self.kz)
        loss = torch.mean(result)*self.kz**2
        return loss

    def flow2grid(self, flow_field):
        b,c,h,w = flow_field.size()
        x = torch.arange(w).view(1, -1).expand(h, -1).type_as(flow_field).float()
        y = torch.arange(h).view(-1, 1).expand(-1, w).type_as(flow_field).float()
        grid = torch.stack([x,y], dim=0)
        grid = grid.unsqueeze(0).expand(b, -1, -1, -1)
        return flow_field+grid

class VGGLoss(nn.Module):
    r"""
    Perceptual loss, VGG-based
    https://arxiv.org/abs/1603.08155
    https://github.com/dxyang/StyleTransfer/blob/master/utils.py
    """

    def __init__(self, weights=[1.0, 1.0, 1.0, 1.0, 1.0]):
        super(VGGLoss, self).__init__()
        self.add_module('vgg', VGG19())
        self.criterion = torch.nn.L1Loss()
        self.weights = weights

    def compute_gram(self, x):
        b, ch, h, w = x.size()
        f = x.view(b, ch, w * h)
        f_T = f.transpose(1, 2)
        G = f.bmm(f_T) / (h * w * ch)
        return G

    def __call__(self, x, y, last_only=False, content_only=False):
        # Compute features
        x_vgg, y_vgg = self.vgg(x), self.vgg(y)
        if not last_only:
            content_loss = 0.0
            content_loss += self.weights[0] * self.criterion(x_vgg['relu1_1'], y_vgg['relu1_1'])
            content_loss += self.weights[1] * self.criterion(x_vgg['relu2_1'], y_vgg['relu2_1'])
            content_loss += self.weights[2] * self.criterion(x_vgg['relu3_1'], y_vgg['relu3_1'])
            content_loss += self.weights[3] * self.criterion(x_vgg['relu4_1'], y_vgg['relu4_1'])
            content_loss += self.weights[4] * self.criterion(x_vgg['relu5_1'], y_vgg['relu5_1'])
            if content_only:
                return content_loss

            # Compute loss
            style_loss = 0.0
            style_loss += self.criterion(self.compute_gram(x_vgg['relu2_2']), self.compute_gram(y_vgg['relu2_2']))
            style_loss += self.criterion(self.compute_gram(x_vgg['relu3_4']), self.compute_gram(y_vgg['relu3_4']))
            style_loss += self.criterion(self.compute_gram(x_vgg['relu4_4']), self.compute_gram(y_vgg['relu4_4']))
            style_loss += self.criterion(self.compute_gram(x_vgg['relu5_2']), self.compute_gram(y_vgg['relu5_2']))
        else:
            content_loss = self.criterion(x_vgg['relu5_1'], y_vgg['relu5_1'])
            if content_only:
                return content_loss
            style_loss = self.criterion(self.compute_gram(x_vgg['relu5_2']), self.compute_gram(y_vgg['relu5_2']))

        return content_loss, style_loss

class PerceptualCorrectness(nn.Module):
    r"""

    """

    def __init__(self, layer=['rel1_1','relu2_1','relu3_1','relu4_1']):
        super(PerceptualCorrectness, self).__init__()
        self.add_module('vgg', VGG19())
        self.layer = layer
        self.eps=1e-8
        from models.networks.resample2d_package.resample2d import Resample2d
        self.resample = Resample2d(4, 1, sigma=2)

    def __call__(self, target, source, flow_list, used_layers, mask=None, use_bilinear_sampling=False):
        used_layers=sorted(used_layers, reverse=True)
        # self.target=target
        # self.source=source
        self.target_vgg, self.source_vgg = self.vgg(target), self.vgg(source)
        loss = 0
        for i in range(len(flow_list)):
            loss += self.calculate_loss(flow_list[i], self.layer[used_layers[i]], mask, use_bilinear_sampling)

        return loss

    def calculate_loss(self, flow, layer, mask=None, use_bilinear_sampling=False):
        target_vgg = self.target_vgg[layer]
        source_vgg = self.source_vgg[layer]
        [b, c, h, w] = target_vgg.shape
        # maps = F.interpolate(maps, [h,w]).view(b,-1)
        flow = F.interpolate(flow, [h,w])

        target_all = target_vgg.view(b, c, -1)                      #[b C N2]
        source_all = source_vgg.view(b, c, -1).transpose(1,2)       #[b N2 C]

        source_norm = source_all/(source_all.norm(dim=2, keepdim=True)+self.eps)
        target_norm = target_all/(target_all.norm(dim=1, keepdim=True)+self.eps)

        correction = torch.bmm(source_norm, target_norm) #[b N2 N2]
        (correction_max,max_indices) = torch.max(correction, dim=1)

        # interple with bilinear sampling
        if use_bilinear_sampling:
            input_sample = self.bilinear_warp(source_vgg, flow).view(b, c, -1)
        else:
            input_sample = self.resample(source_vgg, flow).view(b, c, -1)

        correction_sample = F.cosine_similarity(input_sample, target_all)    #[b 1 N2]
        loss_map = torch.exp(-correction_sample/(correction_max+self.eps))
        if mask is None:
            loss = torch.mean(loss_map) - torch.exp(torch.tensor(-1).type_as(loss_map))
        else:
            mask=F.interpolate(mask, size=(target_vgg.size(2), target_vgg.size(3)))
            mask=mask.view(-1, target_vgg.size(2)*target_vgg.size(3))
            loss_map = loss_map - torch.exp(torch.tensor(-1).type_as(loss_map))
            loss = torch.sum(mask * loss_map)/(torch.sum(mask)+self.eps)

        # print(correction_sample[0,2076:2082])
        # print(correction_max[0,2076:2082])
        # coor_x = [32,32]
        # coor = max_indices[0,32+32*64]
        # coor_y = [int(coor%64), int(coor/64)]
        # source = F.interpolate(self.source, [64,64])
        # target = F.interpolate(self.target, [64,64])
        # source_i = source[0]
        # target_i = target[0]

        # source_i = source_i.view(3, -1)
        # source_i[:,coor]=-1
        # source_i[0,coor]=1
        # source_i = source_i.view(3,64,64)
        # target_i[:,32,32]=-1
        # target_i[0,32,32]=1
        # lists = str(int(torch.rand(1)*100))
        # img_numpy = util.tensor2im(source_i.data)
        # util.save_image(img_numpy, 'source'+lists+'.png')
        # img_numpy = util.tensor2im(target_i.data)
        # util.save_image(img_numpy, 'target'+lists+'.png')
        return loss

    def bilinear_warp(self, source, flow):
        [b, c, h, w] = source.shape
        x = torch.arange(w).view(1, -1).expand(h, -1).type_as(source).float() / (w-1)
        y = torch.arange(h).view(-1, 1).expand(-1, w).type_as(source).float() / (h-1)
        grid = torch.stack([x,y], dim=0)
        grid = grid.unsqueeze(0).expand(b, -1, -1, -1)
        grid = 2*grid - 1
        flow = 2*flow/torch.tensor([w, h]).view(1, 2, 1, 1).expand(b, -1, h, w).type_as(flow)
        grid = (grid+flow).permute(0, 2, 3, 1)
        input_sample = F.grid_sample(source, grid).view(b, c, -1)
        return input_sample

class VGG19(torch.nn.Module):
    def __init__(self):
        super(VGG19, self).__init__()
        features = models.vgg19(pretrained=True).features
        self.relu1_1 = torch.nn.Sequential()
        self.relu1_2 = torch.nn.Sequential()

        self.relu2_1 = torch.nn.Sequential()
        self.relu2_2 = torch.nn.Sequential()

        self.relu3_1 = torch.nn.Sequential()
        self.relu3_2 = torch.nn.Sequential()
        self.relu3_3 = torch.nn.Sequential()
        self.relu3_4 = torch.nn.Sequential()

        self.relu4_1 = torch.nn.Sequential()
        self.relu4_2 = torch.nn.Sequential()
        self.relu4_3 = torch.nn.Sequential()
        self.relu4_4 = torch.nn.Sequential()

        self.relu5_1 = torch.nn.Sequential()
        self.relu5_2 = torch.nn.Sequential()
        self.relu5_3 = torch.nn.Sequential()
        self.relu5_4 = torch.nn.Sequential()

        for x in range(2):
            self.relu1_1.add_module(str(x), features[x])

        for x in range(2, 4):
            self.relu1_2.add_module(str(x), features[x])

        for x in range(4, 7):
            self.relu2_1.add_module(str(x), features[x])

        for x in range(7, 9):
            self.relu2_2.add_module(str(x), features[x])

        for x in range(9, 12):
            self.relu3_1.add_module(str(x), features[x])

        for x in range(12, 14):
            self.relu3_2.add_module(str(x), features[x])

        for x in range(14, 16):
            self.relu3_2.add_module(str(x), features[x])

        for x in range(16, 18):
            self.relu3_4.add_module(str(x), features[x])

        for x in range(18, 21):
            self.relu4_1.add_module(str(x), features[x])

        for x in range(21, 23):
            self.relu4_2.add_module(str(x), features[x])

        for x in range(23, 25):
            self.relu4_3.add_module(str(x), features[x])

        for x in range(25, 27):
            self.relu4_4.add_module(str(x), features[x])

        for x in range(27, 30):
            self.relu5_1.add_module(str(x), features[x])

        for x in range(30, 32):
            self.relu5_2.add_module(str(x), features[x])

        for x in range(32, 34):
            self.relu5_3.add_module(str(x), features[x])

        for x in range(34, 36):
            self.relu5_4.add_module(str(x), features[x])

        # don't need the gradients, just want the features
        for param in self.parameters():
            param.requires_grad = False

    def forward(self, x):
        relu1_1 = self.relu1_1(x)
        relu1_2 = self.relu1_2(relu1_1)

        relu2_1 = self.relu2_1(relu1_2)
        relu2_2 = self.relu2_2(relu2_1)

        relu3_1 = self.relu3_1(relu2_2)
        relu3_2 = self.relu3_2(relu3_1)
        relu3_3 = self.relu3_3(relu3_2)
        relu3_4 = self.relu3_4(relu3_3)

        relu4_1 = self.relu4_1(relu3_4)
        relu4_2 = self.relu4_2(relu4_1)
        relu4_3 = self.relu4_3(relu4_2)
        relu4_4 = self.relu4_4(relu4_3)

        relu5_1 = self.relu5_1(relu4_4)
        relu5_2 = self.relu5_2(relu5_1)
        relu5_3 = self.relu5_3(relu5_2)
        relu5_4 = self.relu5_4(relu5_3)

        out = {
            'relu1_1': relu1_1,
            'relu1_2': relu1_2,

            'relu2_1': relu2_1,
            'relu2_2': relu2_2,

            'relu3_1': relu3_1,
            'relu3_2': relu3_2,
            'relu3_3': relu3_3,
            'relu3_4': relu3_4,

            'relu4_1': relu4_1,
            'relu4_2': relu4_2,
            'relu4_3': relu4_3,
            'relu4_4': relu4_4,

            'relu5_1': relu5_1,
            'relu5_2': relu5_2,
            'relu5_3': relu5_3,
            'relu5_4': relu5_4,
        }
        return out


"""
Version 8.0
***********
This version is object oriented and uses 5 methods via
SEIQRDPSolver to solve the SEIQRDP model.
It supports multiprocessing (with experiment module or any other method
                             by 'adding' experiments together: exp1 + exp2).
This version is not optimized.

In this version (8.0.alpha):
    - The dataset is now handled properly. No more global variable LOCATIONS.
    - Only the data for the population/region in question is loaded.
    - Dataset is downloaded with user's permission if no data file is found.
    - Corrected the fitness function (nRMSE as described in the paper).
    - The fitness function has been re-written to be readable.
    - Got rid of Comparator class. Fitness is handled in the GA class.
    - Got rid of Simulator class. SEIQRDPSolver does the same job.
    - *MAJOR* Dictionaries are now the data structure for the model's
        parameters. This will allow easily integration of other models.
    - Got rid of Configurator class and other parameter's handling functions.
        The parameters are accessed by their names in the new dictionary
        structure.
    - The model's search space is no longer hardcoded in this file. It was
        moved to the model's file.
    - Corrected the 95% confidence interval computation. (was off by 1%)
    - GeneticFit no longer fits at initialization. Requires a call to run().
    - Data is loaded more properly and much faster.
    - Renamed a lot of stuff (name conventions were not respected before)
    - Some code refactoring. It runs significantly faster.
"""
__version__ = '8.0.a'

# Math
import numpy as np
import scipy.stats as st
from math import sqrt
from random import random
from random import choice as rand_choice
from itertools import combinations
# Results
import matplotlib.pyplot as plt
from datetime import timedelta
# Data
import pandas as pd
from os.path import isfile, isdir
from os import mkdir
# Utils
from copy import deepcopy  # Is needed to create hard copies
# Internal
from seiqrdp_model.region_data import Region
import seiqrdp_model.seiqrdp_solver as slv

# ###############    DATA    ################
def update_data():
    """
    Downloads data from
    https://raw.githubusercontent.com/datasets/
    """
    dataset = pd.read_csv('https://raw.githubusercontent.com/datasets/'
                          'covid-19/master/data/'
                          'time-series-19-covid-combined.csv')
    if not isdir('Data/'):
        mkdir('Data/')
    dataset.to_csv(r'Data/time-series-19-covid-combined.csv',
                   index=False)

def load_data(region_name, pop_size):
    """
        Loading data from csv file.
        Returns the Region object.
    """
    print("Dataset is loading...")

    # If data is missing, prompt for data download.
    if not isfile('Data/time-series-19-covid-combined.csv'):
        print("Dataset not found.")
        usr_input = input('Download data from github repo? (y/any) :\n')
        if usr_input == 'y' or usr_input == 'Y':
            print("Dataset is being updated...")
            update_data()
    # Load data from file.
    dataframe = pd.read_csv('Data/time-series-19-covid-combined.csv')

    region = Region.from_dataframe(dataframe, region_name, pop_size)
    return region
# ###############    DATA    ################

class GeneticFit:
    """
        This class implements an evolutionary genetic algorithm that
        finds the best fitting SEIR parameters for a 'region' (pop + cases).
        At each generation we keep the best 'elite_size' parameters and breed
        them with a 'mutation_rate'.
        We stop at 'max_gen' generations.
        The best result (parameters) is 'self.elite'.
        Set 'verbose' = True to see the progress.
    """

    def __init__(self, region, elite_size, max_gen, mutation_rate=0.4,
                 search_space=slv.search_space, method='LSODA', verbose=False):
        self.region = region
        self.elite_size = elite_size
        self.max_gen = max_gen
        self.mutation_rate = mutation_rate
        self.search_space = search_space

        # I0 and E0 first guesses.
        self.search_space['I0'] = [region.rcQ[0], region.rcQ[0]*5]
        self.search_space['E0'] = [region.rcQ[0], region.rcQ[0]*5]

        self.method = method
        self.verbose = verbose

        (self.lembda, self.kappa) = self.compute_lambda_kappa()

        self.result = []
        self.elite = []

    def run(self):
        self.result = self.evolve()
        self.elite = self.result[0]

    def compute_lambda_kappa(self):
        """
            Determine Lambda and Kappa from real data.
            Note: 'lembda' is used instead of the Python reserved name lambda.
        """
        lembda = []
        kappa = []
        cQ = np.array(self.region.rcQ)
        R = np.array(self.region.rR)
        D = np.array(self.region.rD)
        Q = cQ - R - D

        for n in range(len(Q)-1):
            if (R[n] == 0) and (D[n] == 0):  # v007.4
                continue
            lembda.append((R[n+1]-R[n])/Q[n])
            kappa.append((D[n+1]-D[n])/Q[n])

        Lembda = np.mean(np.array(lembda))
        Kappa = np.mean(np.array(kappa))
        return (Lembda, Kappa)

    def evolve(self):
        """
        Evolves the initial population of parameters over a number of
        generations to fit the epidemic model.

        Returns:
        (parameters, fitness)
            A tuple of the best parameters set and its fitness value.

        """
        # Computing the total population size from the elite size.
        pop_size = self.elite_size*(self.elite_size-1)

        # Computing the search space's widths (in which params can vary)
        width = {}
        for param in self.search_space:
            width[param] = self.search_space[param][1]\
                - self.search_space[param][0]

        population = []
        elite = []

        # Generating the initial population
        for i in range(pop_size):
            parameters = {'lambda': self.lembda, 'kappa': self.kappa}
            for param in self.search_space:
                parameters[param] = random()*width[param]\
                    + self.search_space[param][0]

            # Adding the set of parameters to the population along with
            # its fitness.
            population.append([parameters, self.fitness(parameters)])

        # Selecting the elite
        elite = self.generate_elite(population, self.elite_size)

        # Main loop: iterates up to max_gen generations.
        for g in range(self.max_gen):
            # Passing the elite into the next population
            new_population = self.generate_next_population(elite, width)
            # Generate the next elite.
            elite = self.generate_elite(new_population, self.elite_size)

            if self.verbose:
                print(f'Generation {g+1} of {self.max_gen}.'
                      f'Best: {elite[0][1]}\n')
        # Return a tuple (params, fitness) of the best elite.
        return ([elite[0][0], elite[0][1]])

    def generate_next_population(self, elite, width):
        """
        Generates the next generation's population from the previous ones'
        elite.

        Parameters:
            elite :
                The elite parameters which survive and breed.
            width :
                List of widths of the search space; by how much the parameters
                can mutate.

        Returns:
            new_population : list
                The population of the next generation.
        """
        # Passing the elite directly into the next population
        new_population = deepcopy(elite)

        # Breed and mutate, and add to the new population
        for par1, par2 in combinations(elite, 2):
            offspring = self.cross(par1[0], par2[0],
                                   width, self.search_space)
            # Add the offspring to the new population along with
            # its fitness.
            new_population.append([offspring, self.fitness(offspring)])

        return new_population

    def sort_population(self, pop):
        """
            Sorts the population of the genetic pool by
            their fitting score (2nd element in population list)
            Note: get rid of lambda and replace it with a externally
                    declared function (performance issue)
        """
        # reverse = None (Sorts in Ascending order)
        pop.sort(key=lambda x: x[1])
        return pop

    def cross(self, par1, par2, width, search_space):
        """
            Generates an offspring by breeding 2 parents, with
            a 'mutation_chance' = 0-1.
        """
        offspring = {'lambda': self.lembda, 'kappa': self.kappa}
        for param in search_space:
            if random() < self.mutation_rate:
                #   Mutation: the gene is random
                offspring[param] = random()*width[param]+search_space[param][0]
            else:
                #   Breeding: the gene is from a single random parent
                offspring[param] = rand_choice([par1[param], par2[param]])
        return offspring

    def generate_elite(self, population, elite_size):
        """
            Sorts the population by their scores and
            returns 'elite_size' best elements.
        """
        # Selecting the elite
        self.sort_population(population)
        return population[:elite_size]

    def fitness(self, genome):
        """ Computes the fitness of a genome.
            genome = [6 values], no kappa/lambda.
        """
        simulation = slv.SEIQRDPSolver(genome, self.region, mode='fit')

        simulation.solve_model(0, len(self.region.rcQ), 1, self.method)

        sqr_sum = 0
        for a, b in zip(self.region.rcQ, simulation.cQ):
            sqr_sum += (a-b)**2
        n_rmse = sqrt(np.mean(sqr_sum))/np.mean(self.region.rcQ)
        return n_rmse

class Analyzer:
    """
        This class measures BRN, peak_time, Peak values for a given 'Region'
        and a set of SEIR parameters.
        It gives also the data of simulation for n_sim_days period after
        the first day.
    """

    def __init__(self, region, parameters, n_sim_days):
        self.region = region
        self.parameters = parameters
        self.n_sim_days = n_sim_days

        simulation = slv.SEIQRDPSolver(parameters, region)
        simulation.solve_model(0, self.n_sim_days)

        active_cases = simulation.Q

        peak_active_infections = max(simulation.Q)
        peak_day = simulation.Q.index(peak_active_infections)
        self.peak_is_outofrange = False
        if peak_day == len(simulation.Q)-1:
            self.peak_is_outofrange = True

        self.peak = active_cases[peak_day-1]
        self.peak_time = peak_day
        self.brn = parameters['beta']/parameters['delta']
        self.fit = simulation

    def get_epidemic_startdate(self):
        """
        This can be used to seek the probable first day when the
        first exposition occured.
        """

        solver = slv.SEIQRDPSolver(self.parameters, self.region)

        step = -30
        solver.solve_model(0, step, 1, 'LSODA')
        E = solver.E

        while E[0] > 1:
            step += -30
            # Simulates the previous 30 days
            solver.solve_model(step+30, step, 1, 'LSODA')
            E = solver.E
        starting_point = 0

        while E[starting_point] < 0:
            starting_point += 1

        starting_point += step
        starting_date = self.region.first_day\
            + timedelta(days=starting_point)
        return(starting_date)

class BaseExperiment:
    """
        A BaseExperiment perfoms a genetic fitting + simulation
        and computes the model variables.
        It is suited for parallel computing since 2 instences of
        BaseExeriment can be added together to form a 3rd one:
            exp3 = exp1 + exp2
        Example:
            >> import Algeria_SEIR_COVID2019_Object as SEIR
            >> SEIR.load_data()
            >> ita_exp = SEIR.BaseExperiment(region=SEIR.LOCATIONS['Italy'],
                                             n_sim_days=100, n_experiments=4)
            >> ita_exp.run()
            >> ita_exp.compute_results()
            >> ita_exp.show_results('all')
    """

    def __init__(self, region, n_sim_days, n_experiments,
                 elite_size=10, max_gen=50, mutation_rate=0.4,
                 search_space=slv.search_space, method='LSODA', verbose=False):

        self.region = region
        self.n_sim_days = n_sim_days
        self.n_experiments = n_experiments
        self.elite_size = elite_size
        self.max_gen = max_gen
        self.mutation_rate = mutation_rate
        self.search_space = search_space
        self.method = method
        self.verbose = verbose

    def init_states(self):

        # typical values
        self.typical_brn = 0
        self.typical_experiment_index = 0
        self.typical_parameters = {}
        self.typical_peak = 0
        self.typical_peak_time = 0
        self.typical_peak_found = True

        # observables mean values
        self.brn = 0
        self.peak = 0
        self.peak_time = 0
        self.peak_found = True

        # observables covariances
        self.err_brn = 0
        self.err_peak = 0
        self.err_peak_time = 0

        # Elite parameters
        self.elites = []

        # All exps. data.
        self.raw_brn, self.raw_Rn, self.raw_peak, self.raw_peak_time,\
            self.raw_S, self.raw_E, self.raw_I, self.raw_Q, self.raw_cQ,\
            self.raw_R, self.raw_D, self.raw_P\
            = [], [], [], [], [], [], [], [], [], [], [], []

        # mean compartmental variables
        self.S, self.E, self.I, self.Q, self.cQ, self.R, self.D,\
            self.P, self.Rn\
            = [], [], [], [], [], [], [], [], []

        # compartmental variables covariances
        self.err_S, self.err_E, self.err_I, self.err_Q, self.err_cQ,\
            self.err_R, self.err_D, self.err_P, self.err_Rn\
            = [], [], [], [], [], [], [], [], []

        # days list of simulation
        self.days = []

    def run(self):
        """
            Runs the simulations WITHOUT computing the final results.
        """
        # Here we perform n_experiments of experiments
        # (fitting + calculating observables)
        self.init_states()

        for n in range(self.n_experiments):
            # Fit real data and obtaining the n'th elite parameter
            if self.verbose:
                print("Fitting data: ", n+1, " out of ",
                      self.n_experiments, ' experiments.')
            fit = GeneticFit(self.region, self.elite_size, self.max_gen,
                             self.mutation_rate, self.search_space,
                             self.method, self.verbose)
            fit.run()

            # Calculate observables of n'th elite parameter
            if self.verbose:
                print("Calculating epidemic data.")
            observables = Analyzer(self.region, fit.elite, self.n_sim_days)
            # Assign what you have obtained to lists
            self.elites.append(fit.elite)
            self.raw_brn.append(observables.brn)
            self.raw_peak.append(observables.peak)
            self.raw_peak_time.append(observables.peak_time)
            self.raw_S.append(observables.fit.S)
            self.raw_E.append(observables.fit.E)
            self.raw_I.append(observables.fit.I)
            self.raw_Q.append(observables.fit.Q)
            self.raw_cQ.append(observables.fit.cQ)
            self.raw_R.append(observables.fit.R)
            self.raw_D.append(observables.fit.D)
            self.raw_P.append(observables.fit.P)
            # R(t) = brn * S/N  v007.4
            self.raw_Rn.append([observables.brn * s/self.region.N
                                for s in self.raw_S[n]])
            if observables.peak_is_outofrange:
                self.peak_found = False

    def compute_results(self):
        """
            Computes the final results of the simulations.
        """
        # Calculate the {mean values} and {typical values}
        # of observables from n_experiments of experiments
        if self.verbose:
            print("I am calculating expected values")

        self.brn = np.mean(np.array(self.raw_brn))

        self.typical_brn = self.raw_brn[0]
        for n in range(1, self.n_experiments):
            if (self.raw_brn[n]-self.brn)**2 < (self.typical_brn-self.brn)**2:
                self.typical_brn = self.raw_brn[n]
                self.typical_experiment_index = n
            else:
                continue

        self.typical_parameters = self.elites[self.typical_experiment_index]

        self.peak = np.mean(np.array(self.raw_peak))
        self.typical_peak = self.raw_peak[self.typical_experiment_index]

        self.peak_time = np.mean(np.array(self.raw_peak_time))
        self.typical_peak_time\
            = self.raw_peak_time[self.typical_experiment_index]

        # Calculate the mean values of compartmental variables
        for i in range(self.n_sim_days):
            self.S.append(np.mean(np.array(self.raw_S)[:, i]))
            self.E.append(np.mean(np.array(self.raw_E)[:, i]))
            self.I.append(np.mean(np.array(self.raw_I)[:, i]))
            self.Q.append(np.mean(np.array(self.raw_Q)[:, i]))
            self.cQ.append(np.mean(np.array(self.raw_cQ)[:, i]))
            self.R.append(np.mean(np.array(self.raw_R)[:, i]))
            self.D.append(np.mean(np.array(self.raw_D)[:, i]))
            self.P.append(np.mean(np.array(self.raw_P)[:, i]))
            self.Rn.append(np.mean(np.array(self.raw_Rn)[:, i]))

        # Calculate the error of observables
        # Critical value for 95% confidence interval.
        crit_value = st.t.ppf((1 + 0.95) / 2, self.n_sim_days - 1)
        if self.verbose:
            print("Calculating errors")
        self.err_brn = st.sem(np.array(self.raw_brn)) * crit_value
        self.err_peak = st.sem(np.array(self.raw_peak)) * crit_value
        self.err_peak_time = st.sem(np.array(self.raw_peak_time)) * crit_value

        # Calculate the std error on comartmental variables  v007.4
        #  * crit_value means sigma2 : 95% confidence
        for i in range(self.n_sim_days):
            self.days.append(i)
            self.err_S.append(st.sem(np.array(self.raw_S)[:, i]) * crit_value)
            self.err_E.append(st.sem(np.array(self.raw_E)[:, i]) * crit_value)
            self.err_I.append(st.sem(np.array(self.raw_I)[:, i]) * crit_value)
            self.err_Q.append(st.sem(np.array(self.raw_Q)[:, i]) * crit_value)
            self.err_cQ.append(st.sem(np.array(self.raw_cQ)[:, i])*crit_value)
            self.err_R.append(st.sem(np.array(self.raw_R)[:, i]) * crit_value)
            self.err_D.append(st.sem(np.array(self.raw_D)[:, i]) * crit_value)
            self.err_P.append(st.sem(np.array(self.raw_P)[:, i]) * crit_value)
            self.err_Rn.append(st.sem(np.array(self.raw_Rn)[:, i])*crit_value)

    def show_results(self, *curves):
        """
            Shows the curves and prints the results of
            n_experiments experiments.
        """
        print('___________________ Plots ___________________')
        self.show_curves(*curves)

        self.show_result_values()
        return

    def show_curves(self, *curves):
        """
            Shows the curves only.
            Available curves:
                cq: Cumulative quarantined.
                q: active quarantined.
                rn: Rn(t)
            Choosing 'all' or no arguments shows all the curves.
        """
        # Errors handling
        args = ['cq', 'q', 'rn', 'all']
        n_plots = 0
        for arg in curves:
            if arg not in args:
                print(f'Warning: \'{arg}\' is not a valid curve.')
            else:
                n_plots += 1

        draw_all = (len(curves) == 0) or ('all' in curves)
        if draw_all:
            n_plots = 3

        # Plot cQ and Q with error bars
        fig, plots = plt.subplots(nrows=n_plots, ncols=1, sharex=False,
                                  figsize=[8, (13/3)*n_plots])

        if n_plots == 1:
            plots = (plots,)

        last_plot = -1
        if 'cq' in curves or draw_all:
            last_plot += 1
            cq = plots[last_plot]
            cq.set_title('Cumulative Number of cases')
            cq.plot(self.region.rcQ, color='red', marker="o",
                    label='Real data (fitting)')
            cq.errorbar(self.days, self.cQ, yerr=self.err_cQ, errorevery=6,
                        ecolor='black', label='Model')
            cq.legend(loc="lower right")
            cq.set_xlabel(f'Days since {self.region.first_day}')
            cq.set_ylabel('cQ(t)')
            cq.grid(True)

        if 'q' in curves or draw_all:
            last_plot += 1
            q = plots[last_plot]
            q.set_title('Number of active cases')
            q.errorbar(self.days, self.Q, yerr=self.err_Q, errorevery=6,
                       ecolor='black')
            q.set_xlabel(f'Days since {self.region.first_day}')
            q.set_ylabel('Q(t)')
            q.grid(True)

        if 'rn' in curves or draw_all:
            last_plot += 1
            rn = plots[last_plot]
            rn.set_title('Reproduction Number')
            rn.errorbar(self.days, self.Rn, yerr=self.err_Rn, errorevery=6,
                        ecolor='black', label='R(t)')
            # Drawing Rn = 1
            rn.plot(self.days, [1]*len(self.Rn), label='R=1', color='red',
                    linestyle='dashed')
            rn.legend(loc="upper right")
            rn.set_xlabel(f'Days since {self.region.first_day}')
            rn.set_ylabel('R(t)')
            rn.grid(True)

        fig.tight_layout()
        plt.show()

    def show_result_values(self):
        """
            Shows the numerical values of the experiment.
        """
        # Showing mean values
        print('_______________________ Mean Values _______________________')
        print('BRN: ', round(self.brn, 3), "\u00B1", round(self.err_brn, 3))

        if self.peak_found:
            print('peak_time: ',
                  self.region.first_day + timedelta(round(self.peak_time)),
                  ' (', round(self.peak_time, 2), ") \u00B1",
                  round(self.err_peak_time, 2))
            print('Peak:', int(self.peak), "\u00B1", int(self.err_peak))
        else:
            print("A peak could not be found in the chosen time period.")
        # Totals (mean)
        print('')
        print(f'Total infected (t={self.n_sim_days}):',
              int(self.cQ[-1]), "\u00B1", int(self.err_cQ[-1]))
        print(f'Total recovered (t={self.n_sim_days}):',
              int(self.R[-1]), "\u00B1", int(self.err_R[-1]))
        print(f'Total deaths (t={self.n_sim_days}):',
              int(self.D[-1]), "\u00B1", int(self.err_D[-1]))

        # Showing typical values
        print('_______________________ Typical Values _______________________')

        print('Typical values are chosen to give the closest BRN '
              'to the mean BRN.\n')

        print('BRN: ', round(self.typical_brn, 3))
        if self.peak_found:
            print('peak_time: ', self.region.first_day
                  + timedelta(round(self.typical_peak_time)),
                  ' (', round(self.typical_peak_time, 2), ')')
            print('Peak:', int(self.typical_peak))
        else:
            print("A peak could not be found in the chosen time period.")

        # Totals (typical)
        print('')
        print(f'Total infected (t={self.n_sim_days}):',
              int(self.raw_cQ[self.typical_experiment_index][-1]))
        print(f'Total recovered (t={self.n_sim_days}):',
              int(self.raw_R[self.typical_experiment_index][-1]))
        print(f'Total deaths (t={self.n_sim_days}):',
              int(self.raw_D[self.typical_experiment_index][-1]))

        print('')
        print('Protection rate: ',
              round(self.typical_parameters['alpha'], 3))
        print('Transmission rate: ',
              round(self.typical_parameters['beta'], 3))
        print('Latent period: ',
              round(1/self.typical_parameters['gamma'], 3))
        print('Infectious period: ',
              round(1/self.typical_parameters['delta'], 3))
        print('Recovery rate: ',
              round(self.typical_parameters['lambda'], 3))
        print('Fatality rate: ',
              round(self.typical_parameters['kappa'], 3))
        print('Number of exposed people in (', self.region.first_day, ') = ',
              round(self.typical_parameters['E0']))
        print('Number of infectious people in (', self.region.first_day,
              ') = ', round(self.typical_parameters['I0']))

    def __add__(self, other):
        """
            Defines the + operator for BaseExperiment.
            NOTE : This is just a hack for now!
            To Do: Make this more robust
        """
        newExp = deepcopy(self)
        newExp.elites += other.elites
        newExp.raw_brn += other.raw_brn
        newExp.raw_peak += other.raw_peak
        newExp.raw_peak_time += other.raw_peak_time
        newExp.raw_S += other.raw_S
        newExp.raw_E += other.raw_E
        newExp.raw_I += other.raw_I
        newExp.raw_Q += other.raw_Q
        newExp.raw_cQ += other.raw_cQ
        newExp.raw_R += other.raw_R
        newExp.raw_D += other.raw_D
        newExp.raw_P += other.raw_P
        newExp.raw_Rn += other.raw_Rn
        newExp.n_experiments += other.n_experiments

        # observables mean values
        newExp.brn = 0
        newExp.peak = 0
        newExp.peak_time = 0
        newExp.peak_found = newExp.peak_found and other.peak_found

        # observables covariances
        newExp.err_brn = 0
        newExp.err_peak = 0
        newExp.err_peak_time = 0

        # mean compartmental variables
        newExp.S, newExp.E, newExp.I, newExp.Q, newExp.cQ, newExp.R,\
            newExp.D, newExp.P, newExp.Rn = [], [], [], [], [], [], [], [], []

        # compartmental variables covariances
        newExp.err_S, newExp.err_E, newExp.err_I, newExp.err_Q, newExp.err_cQ,\
            newExp.err_R, newExp.err_D, newExp.err_P, newExp.err_Rn = [], [],\
            [], [], [], [], [], [], []

        newExp.days = []

        return newExp


#!/usr/bin/env python3.7

##################################################################################################
# This script is used for setting the velocity-dependent transfer fucntion (TF) to 1e-8 for TF   #
# file generated by CAMB(planck_2018_transfer_out.dat) and axsionCAMB(planck_2018_axion_transfer #
# _out.dat). Be sure to prepare the input transfer function files and specify their name!! The   #
# output file planck_2018_transfer_out_no_vel.dat/planck_2018_axion_transfer_out.dat will be     #
# produced accordingly(corresponding to CAMB/axionCAMB input).                                   #
##################################################################################################

import sys
import numpy as np
import pandas as pd

if len(sys.argv) != 2:
    raise RuntimeError("Input Error: Please select do axion or not (T/F)!")

if sys.argv[1] == "F":
    # no axion
    df_trans = pd.read_csv('planck_2018_transfer_out.dat', delimiter='\s+', index_col=False)
    col = list(df_trans.columns[1:])
    col.append("None")
    df_trans.columns = col
    df_trans = df_trans.iloc[:,:-1]
    df_trans['v_CDM'] = 1e-8*np.ones(df_trans['v_CDM'].shape)
    df_trans['v_b'] = 1e-8*np.ones(df_trans['v_b'].shape)
    df_trans['v_b-v_c'] = 1e-8*np.ones(df_trans['v_b-v_c'].shape)

    df_trans.to_csv('planck_2018_transfer_out_no_vel.dat',header=None, sep='\t', float_format = "%.6e", index=None)
####################################################################################################################
elif sys.argv[1] == "T":
    # with axion
    df_trans_axion = pd.read_csv('planck_2018_axion_transfer_out.dat', delimiter='\s+', index_col=False, header=None)

    # switch the CDM column with axion column
    temp = df_trans_axion.iloc[:,1].copy()
    df_trans_axion.iloc[:,1] = df_trans_axion.iloc[:,6].copy()
    df_trans_axion.iloc[:,6] = temp.copy()

    # switch the CDM column with total matter column
    temp = df_trans_axion.iloc[:,6].copy()
    df_trans_axion.iloc[:,6] = df_trans_axion.iloc[:,-1].copy()
    df_trans_axion.iloc[:,-1] = temp.copy()

    # subtitute the data with negative TF by its absolute value
    df_trans_axion = abs(df_trans_axion)

    # subtitute the data with negative TF by 1e-8
#    cri = df_trans_axion<0.0
#    print("Total %.2f%% of tranfer function is negative term!"%(100*cri.sum().sum()/cri.shape[0]/cri.shape[1]))
#    df_trans_axion[cri] = 1e-8

    for i in range(9,13):
        df_trans_axion[i] = 1e-8*np.ones(len(df_trans_axion.iloc[:,0]))

    df_trans_axion.iloc[:,:100].to_csv('planck_2018_axion_transfer_out_no_vel.dat', \
                                        header=None, sep='\t', float_format = "%.6e", index=None)
    #with open('planck_2018_axion_transfer_out_no_vel.dat','w') as f:
    #    np.savetxt(f, df_trans_axion.iloc[:,:],delimiter=' ', fmt = '%.6e')
####################################################################################################################
else:
    raise RuntimeError("Input Error: Please select do axion or not (T/F)!")


from __future__ import division
import sys,os

if sys.version_info.major>=3:
	from io import StringIO
else:
	from StringIO import StringIO

from .nbody import NbodySnapshot
from .. import extern as ext
from .settings import LTSettings

import numpy as np
import astropy.units as u
from astropy.cosmology import w0waCDM,LambdaCDM

############################################################
################Gadget2Settings class#######################
############################################################

class Gadget2Settings(LTSettings):

	"""
	Class handler of the tunable settings in a Gadget2 run

	"""

	file_names = ["InitCondFile","OutputDir","EnergyFile","InfoFile","TimingsFile","CpuFile","RestartFile","SnapshotFileBase","OutputListFilename"]
	cpu_timings = ["TimeLimitCPU","ResubmitOn","ResubmitCommand"]
	code_options = ["ICFormat","SnapFormat","ComovingIntegrationOn","TypeOfTimestepCriterion","OutputListOn","PeriodicBoundariesOn"]
	characteristics_of_run = ["TimeMax"]
	output_frequency = ["TimeBetSnapshot","TimeOfFirstSnapshot","CpuTimeBetRestartFile","TimeBetStatistics","NumFilesPerSnapshot","NumFilesWrittenInParallel"]
	accuracy_time_integration = ["ErrTolIntAccuracy","MaxRMSDisplacementFac","CourantFac","MaxSizeTimestep","MinSizeTimestep"]
	tree_algorithm = ["ErrTolTheta","TypeOfOpeningCriterion","ErrTolForceAcc","TreeDomainUpdateFrequency"]
	sph = ["DesNumNgb","MaxNumNgbDeviation","ArtBulkViscConst","InitGasTemp","MinGasTemp"]
	memory_allocation = ["PartAllocFactor","TreeAllocFactor","BufferSize"]
	system_of_units = ["UnitLength_in_cm","UnitMass_in_g","UnitVelocity_in_cm_per_s","GravityConstantInternal"]
	softening = ["MinGasHsmlFractional","SofteningGas","SofteningHalo","SofteningDisk","SofteningBulge","SofteningStars","SofteningBndry","SofteningGasMaxPhys","SofteningHaloMaxPhys","SofteningDiskMaxPhys","SofteningBulgeMaxPhys","SofteningStarsMaxPhys","SofteningBndryMaxPhys"]

	def __init__(self,**kwargs):

		#Default outputs
		self.OutputScaleFactor = np.array([0.2463687286034,0.25331915596808,0.2603792960781,0.26755062217624,0.27483471475984,0.282233268286,0.28974809827109,0.29738114881152,0.30513450055509,0.31301037915502,0.32101116424033,0.32913939894726,0.3373978000372,0.34578926866836,0.35431690185511,0.36298400467612,0.37179410329025,0.38075095882651,0.38985858222059,0.39912125007794,0.40854352165158,0.4181302570317,0.42788663665433,0.43781818224776,0.4479307793476,0.45823070152614,0.46872463649679,0.47941971427274,0.49032353757851,0.50144421473605,0.51279039527177,0.52437130852031,0.53619680553253,0.54827740463258,0.56062434101017,0.57324962078195,0.58616608000982,0.59938744922579,0.61292842408364,0.62680474283878,0.64103327145057,0.65563209720868,0.67062063190864,0.68601972574487,0.70185179325524,0.71814095284413,0.73491318163551,0.75219648767026,0.77002110176951,0.78841969174729,0.80742760208188,0.82708312265889,0.84742779079644,0.86850673147378,0.89036904153293,0.91306822464037,0.9366626850185,0.96121628943351,0.98679900871668,1.0])

		#File names
		self.InitCondFile = "gadget_ic"
		self.OutputDir = "snapshots"
		self.EnergyFile = "energy.txt"
		self.InfoFile = "info.txt"
		self.TimingsFile = "timings.txt"
		self.CpuFile = "cpu.txt"
		self.RestartFile = "restart"
		self.SnapshotFileBase = "snapshot"
		self.OutputListFilename = "outputs.txt"

		#CPU Timings
		self.TimeLimitCPU = 1.0*u.day
		self.ResubmitOn = 0
		self.ResubmitCommand = "my-scriptfile"

		#Code options
		self.ICFormat  = 1
		self.SnapFormat = 1
		self.ComovingIntegrationOn = 1
		self.TypeOfTimestepCriterion = 0
		self.OutputListOn = 1
		self.PeriodicBoundariesOn = 1

		#Caracteristics of run
		self.TimeMax = 1.0

		#Output frequency
		self.TimeBetSnapshot = 0.5
		self.TimeOfFirstSnapshot = 0
		self.CpuTimeBetRestartFile = 12.5*u.hour
		self.TimeBetStatistics = 0.05
		self.NumFilesPerSnapshot = 16
		self.NumFilesWrittenInParallel = 8

		#Accuracy of time integration
		self.ErrTolIntAccuracy = 0.025
		self.MaxRMSDisplacementFac = 0.2
		self.CourantFac = 0.15
		self.MaxSizeTimestep = 0.02
		self.MinSizeTimestep = 0.0

		#Tree algorithm, force accuracy, domain update frequency
		self.ErrTolTheta = 0.45
		self.TypeOfOpeningCriterion = 1
		self.ErrTolForceAcc = 0.005
		self.TreeDomainUpdateFrequency = 0.025

		#Further parameters of SPH
		self.DesNumNgb = 33
		self.MaxNumNgbDeviation = 2
		self.ArtBulkViscConst = 0.8
		self.InitGasTemp = 1000.0
		self.MinGasTemp = 50.0

		#Memory allocation
		self.PartAllocFactor = 1.3
		self.TreeAllocFactor = 0.7
		self.BufferSize = 20*u.Mbyte

		#System of units
		self.UnitLength_in_cm = 3.085678e21       # ;  1.0 kpc
		self.UnitMass_in_g = 1.989e43    #;  1.0e10 solar masses
		self.UnitVelocity_in_cm_per_s = 1.0e5  # ;  1 km/sec
		self.GravityConstantInternal = 0

		#Softening lengths
		self.MinGasHsmlFractional = 0.25
		self.SofteningGas = 0
		self.SofteningHalo = 9.000000
		self.SofteningDisk = 0
		self.SofteningBulge = 0
		self.SofteningStars = 0
		self.SofteningBndry = 0

		self.SofteningGasMaxPhys = 0
		self.SofteningHaloMaxPhys = 9.000000
		self.SofteningDiskMaxPhys = 0
		self.SofteningBulgeMaxPhys = 0
		self.SofteningStarsMaxPhys = 0
		self.SofteningBndryMaxPhys = 0

		#Allow for kwargs override
		for key in kwargs.keys():
			setattr(self,key,kwargs[key])

	def sections(self):

		return [ "file_names","cpu_timings","code_options","characteristics_of_run","output_frequency","accuracy_time_integration","tree_algorithm","sph","memory_allocation","system_of_units","softening" ]

	def showSection(self,section):

		if section not in self.sections():
			raise ValueError("Parameter file does not admit a section named {0}".format(section))

		for option in getattr(self,section):
			print("{0} = {1}".format(option,getattr(self,option)))

	def show(self):

		for section in self.sections():
			print(section+":\n")
			self.showSection(section)
			print("\n")

	def writeSection(self,section):

		"""
		Writes the corresponding section of the Gadget2 parameter file

		"""

		output = StringIO()

		#Write preamble
		output.write("% {0}\n\n".format(section))

		#Cycle through options
		for option in getattr(self,section):

			#Read the corresponding value
			value = getattr(self,option)

			#Convert units as necessary
			if type(value)==u.quantity.Quantity:

				if value.unit.physical_type=="time":
					value = value.to(u.s).value
				elif value.unit.physical_type=="speed":
					value = value.to(u.cm/u.s).value
				elif "byte" in value.unit.to_string():
					value = value.to(u.Mbyte).value

			#Write the line
			output.write("{0}		{1}\n".format(option,value))

		#Finish
		output.write("\n\n")
		output.seek(0)

		return output.read()

	@classmethod
	def default(cls):

		"""
		Generate default settings
		"""

		return cls()

############################################################
################Gadget2Header class#########################
############################################################

class Gadget2Header(dict):

	"""
	Class handler of a Gadget2 snapshot header

	"""

	def __init__(self,HeaderDict=dict()):

		super(Gadget2Header,self).__init__()
		for key in HeaderDict.keys():
			self[key] = HeaderDict[key]

	def __repr__(self):

		keys = self.keys()
		keys.sort()

		return "\n".join([ "{0} : {1}".format(key,self[key]) for key in keys ])

	def __add__(self,rhs):

		assert isinstance(rhs,Gadget2Header),"addition not defined if rhs is not a Gadget2Header!"

		#Check that it makes sense to add the snapshots (cosmological parameters, box size, time and redshift must agree)
		fields_to_match = ["Ode0","Om0","h","w0","wa","box_size","endianness","flag_cooling","flag_feedback","flag_sfr","num_files"]
		fields_to_match += ["num_particles_total","num_particles_total_gas","num_particles_total_side","num_particles_total_with_mass","redshift","scale_factor"]

		for field in fields_to_match:
			assert self[field] == rhs[field],"{0} fields do not match!".format(field)

		assert np.all(self["masses"]==rhs["masses"])
		assert np.all(self["num_particles_total_of_type"]==rhs["num_particles_total_of_type"])

		#Construct the header of the merged snapshot
		merged_header = self.copy()
		merged_header["files"] += rhs["files"]
		merged_header["num_particles_file"] += rhs["num_particles_file"]
		merged_header["num_particles_file_gas"] += rhs["num_particles_file_gas"]
		merged_header["num_particles_file_of_type"] += rhs["num_particles_file_of_type"]
		merged_header["num_particles_file_with_mass"] += rhs["num_particles_file_with_mass"]

		return merged_header

##############################################################
#################Gadget2Snapshot class######################
##############################################################

class Gadget2Snapshot(NbodySnapshot):

	"""
	A class that handles Gadget2 snapshots, mainly I/O from the binary format and spatial information statistics.Inherits from the abstract NbodySnapshot

	"""

	###############################################################################################
	#########################Abstract methods implementation#######################################
	###############################################################################################

	@classmethod
	def buildFilename(cls,root,pool,**kwargs):

		if pool is not None:
			return root+".{0}".format(pool.rank)
		else:
			return root

	@classmethod
	def int2root(cls,name,n):
		return name + "_{0:03d}".format(n)

	############################################################################################

	def getHeader(self):

		header = Gadget2Header(ext._gadget2.getHeader(self.fp))
		header["files"] = [self.fp.name]

		header["w0"] = -1.0
		header["wa"] = 0.0
		header["comoving_distance"] = LambdaCDM(H0=header["h"]*100,Om0=header["Om0"],Ode0=header["Ode0"]).comoving_distance(header["redshift"]).to(u.kpc).value * header["h"]

		return header

	############################################################################################

	def setLimits(self):
		self._first = None
		self._last = None

	############################################################################################

	def getPositions(self,first=None,last=None,save=True):

		"""
		Reads in the particles positions (read in of a subset is allowed): when first and last are specified, the numpy array convention is followed (i.e. getPositions(first=a,last=b)=getPositions()[a:b])

		:param first: first particle in the file to be read, if None 0 is assumed
		:type first: int. or None

		:param last: last particle in the file to be read, if None the total number of particles is assumed
		:type last: int. or None

		:param save: if True saves the particles positions as attribute
		:type save: bool.

		:returns: numpy array with the particle positions

		"""

		#The file must not be closed
		assert not self.fp.closed

		#Particles do not have structure
		self.weights = None
		self.virial_radius = None
		self.concentration = None

		numPart = self._header["num_particles_file"]

		#Calculate the offset from the beginning of the file: 4 bytes (endianness) + 256 bytes (header) + 8 bytes (void)
		offset = 4 + 256 + 8

		#If first is specified, offset the file pointer by that amount
		if first is not None:

			assert first>=0
			offset += 4 * 3 * first
			numPart -= first

		if last is not None:

			if first is not None:

				assert last>=first and last<=self._header["num_particles_file"]
				numPart = last - first

			else:

				assert last<=self._header["num_particles_file"]
				numPart = last

		#Read in the particles positions and return the corresponding array
		try:
			positions = (ext._gadget2.getPosVel(self.fp,offset,numPart) * self.kpc_over_h).to(self.Mpc_over_h)
		except AttributeError:
			positions = ext._gadget2.getPosVel(self.fp,offset,numPart) * u.kpc

		if save:
			self.positions = positions
			return self.positions

		#Return
		return positions

	############################################################################################

	def getVelocities(self,first=None,last=None,save=True):

		"""
		Reads in the particles velocities (read in of a subset is allowed): when first and last are specified, the numpy array convention is followed (i.e. getVelocities(first=a,last=b)=getVelocities()[a:b])

		:param first: first particle in the file to be read, if None 0 is assumed
		:type first: int. or None

		:param last: last particle in the file to be read, if None the total number of particles is assumed
		:type last: int. or None

		:param save: if True saves the particles velocities as attrubute
		:type save: bool.

		:returns: numpy array with the particle velocities

		"""

		assert not self.fp.closed

		numPart = self._header["num_particles_file"]

		#Calculate the offset from the beginning of the file: 4 bytes (endianness) + 256 bytes (header) + 8 bytes (void)
		offset = 4 + 256 + 8

		#Skip all the particle positions
		offset += 4 * 3 * numPart

		#Skip other 8 void bytes
		offset += 8

		#If first is specified, offset the file pointer by that amount
		if first is not None:

			assert first>=0
			offset += 4 * 3 * first
			numPart -= first

		if last is not None:

			if first is not None:

				assert last>=first and last<=self._header["num_particles_file"]
				numPart = last - first

			else:

				assert last<=self._header["num_particles_file"]
				numPart = last

		#Read in the particles positions and return the corresponding array
		velocities = ext._gadget2.getPosVel(self.fp,offset,numPart)

		#Scale units
		velocities *= self._velocity_unit
		velocities *= u.cm / u.s

		if save:
			self.velocities = velocities
			return self.velocities

		#Return
		return velocities

	############################################################################################

	def getID(self,first=None,last=None,save=True):

		"""
		Reads in the particles IDs, 4 byte ints, (read in of a subset is allowed): when first and last are specified, the numpy array convention is followed (i.e. getID(first=a,last=b)=getID()[a:b])

		:param first: first particle in the file to be read, if None 0 is assumed
		:type first: int. or None

		:param last: last particle in the file to be read, if None the total number of particles is assumed
		:type last: int. or None

		:param save: if True saves the particles IDs as attribute
		:type save: bool.

		:returns: numpy array with the particle IDs

		"""

		assert not self.fp.closed

		numPart = self._header["num_particles_file"]

		#Calculate the offset from the beginning of the file: 4 bytes (endianness) + 256 bytes (header) + 8 bytes (void)
		offset = 4 + 256 + 8

		#Skip all the particle positions
		offset += 4 * 3 * numPart

		#Skip other 8 void bytes
		offset += 8

		#Skip all the particle velocities
		offset += 4 * 3 * numPart

		#Skip other 8 void bytes
		offset += 8

		#If first is specified, offset the file pointer by that amount
		if first is not None:

			assert first>=0
			offset += 4 * first
			numPart -= first

		if last is not None:

			if first is not None:

				assert last>=first and last<=self._header["num_particles_file"]
				numPart = last - first

			else:

				assert last<=self._header["num_particles_file"]
				numPart = last

		#Read in the particles positions and return the corresponding array
		ids = ext._gadget2.getID(self.fp,offset,numPart)
		if save:
			self.id = ids
			return self.id

		#Return
		return ids

	############################################################################################

	def write(self,filename,files=1):

		"""
		Writes particles information (positions, velocities, etc...) to a properly formatter Gadget snapshot

		:param filename: name of the file to which to write the snapshot
		:type filename: str.

		:param files: number of files on which to split the writing of the snapshot (useful if the number of particles is large); if > 1 the extension ".n" is appended to the filename
		:type files: int.

		"""

		#Sanity checks
		assert hasattr(self,"positions"),"Positions must be specified!!"
		assert self.positions.shape[1]==3

		if not hasattr(self,"_header"):
			self.setHeaderInfo()

		#Build a bare header based on the available info (need to convert units back to the Gadget ones)
		_header_bare = self._header.copy()
		_header_bare["box_size"] = _header_bare["box_size"].to(self.kpc_over_h).value
		_header_bare["masses"] = _header_bare["masses"].to(u.g).value * _header_bare["h"] / self._mass_unit
		_header_bare["num_particles_file_of_type"] = _header_bare["num_particles_file_of_type"].astype(np.int32)
		_header_bare["num_particles_total_of_type"] = _header_bare["num_particles_total_of_type"].astype(np.int32)
		_header_bare["comoving_distance"] = _header_bare["comoving_distance"].to(self.Mpc_over_h).value * 1.0e3

		#Convert units for positions and velocities
		_positions_converted = self.positions.to(self.kpc_over_h).value.astype(np.float32)

		if hasattr(self,"velocities"):
			assert self.positions.shape==self.velocities.shape
			_velocities_converted = (self.velocities.to(u.cm/u.s).value / self._velocity_unit).astype(np.float32)
			writeVel = 1
		else:
			_velocities_converted = np.zeros((1,3),dtype=np.float32)
			writeVel = 0

		#Check if we want to split on multiple files (only DM particles supported so far for this feature)
		if files>1:

			#Number of files the snapshot is split into
			_header_bare["num_files"] = files

			#Update the header with the file names
			self.header["files"] = [ "{0}.{1}".format(filename,n) for n in range(files) ]

			#Distribute particles among files
			particles_per_file = _header_bare["num_particles_total"] // files

			#Write each file
			for n in range(files - 1):

				#Update header
				_header_bare["num_particles_file"] = particles_per_file
				#TODO all particles are DM, fix distribution in the future
				_header_bare["num_particles_file_of_type"] = np.array([0,particles_per_file,0,0,0,0],dtype=np.int32)

				#Write it!
				filename_with_extension = "{0}.{1}".format(filename,n)
				ext._gadget2.write(_header_bare,_positions_converted[n*particles_per_file:(n+1)*particles_per_file],_velocities_converted[n*particles_per_file:(n+1)*particles_per_file],n*particles_per_file+1,filename_with_extension,writeVel)

			#The last file might have a different number of particles
			particles_last_file = len(_positions_converted[particles_per_file*(files-1):])

			#Update header
			_header_bare["num_particles_file"] = particles_last_file
			#TODO all particles are DM, fix distribution in the future
			_header_bare["num_particles_file_of_type"] = np.array([0,particles_last_file,0,0,0,0],dtype=np.int32)

			#Write it!
			filename_with_extension = "{0}.{1}".format(filename,files-1)
			ext._gadget2.write(_header_bare,_positions_converted[particles_per_file*(files-1):],_velocities_converted[particles_per_file*(files-1):],(files-1)*particles_per_file+1,filename_with_extension,writeVel)

		else:

			#Update the num_files key only if not present already
			if "num_files" not in _header_bare.keys():
				_header_bare["num_files"] = 1

			#Update the header with the file names
			self.header["files"] = [ filename ]

			#Write it!!
			ext._gadget2.write(_header_bare,_positions_converted,_velocities_converted,1,filename,writeVel)

	############################################################################################
	###########################Extra methods####################################################
	############################################################################################

	def setHeaderInfo(self,Om0=0.26,Ode0=0.74,w0=-1.0,wa=0.0,h=0.72,redshift=100.0,box_size=15.0*u.Mpc/0.72,flag_cooling=0,flag_sfr=0,flag_feedback=0,flag_stellarage=0,flag_metals=0,flag_entropy_instead_u=0,masses=np.array([0,1.03e10,0,0,0,0])*u.Msun,num_particles_file_of_type=None,npartTotalHighWord=np.zeros(6,dtype=np.uint32)):

		"""
		Sets the header info in the snapshot to write

		"""

		if num_particles_file_of_type is None:
			num_particles_file_of_type = np.array([0,1,0,0,0,0],dtype=np.int32) * self.positions.shape[0]

		assert num_particles_file_of_type.sum()==self.positions.shape[0],"The total number of particles must match!!"
		assert box_size.unit.physical_type=="length"
		assert masses.unit.physical_type=="mass"

		#Create the header
		self._header = Gadget2Header()

		#Fill in
		self._header["Om0"] = Om0
		self._header["Ode0"] = Ode0
		self._header["w0"] = w0
		self._header["wa"] = wa
		self._header["h"] = h
		self._header["H0"] = 100.0*h*u.km/(u.s*u.Mpc)
		self._header["redshift"] = redshift
		self._header["scale_factor"] = 1.0 / (1.0 + redshift)
		self._header["box_size"] = box_size
		self._header["flag_cooling"] = flag_cooling
		self._header["flag_sfr"] = flag_sfr
		self._header["flag_feedback"] = flag_feedback
		self._header["flag_stellarage"] = flag_stellarage
		self._header["flag_metals"] = flag_metals
		self._header["flag_entropy_instead_u"] = flag_entropy_instead_u
		self._header["masses"] = masses
		self._header["num_particles_file_of_type"] = num_particles_file_of_type
		self._header["num_particles_file"] = num_particles_file_of_type.sum()
		self._header["num_particles_total_of_type"] = num_particles_file_of_type
		self._header["num_particles_total"] = num_particles_file_of_type.sum()
		self._header["npartTotalHighWord"] = npartTotalHighWord

		#Define the kpc/h and Mpc/h units for convenience
		self.kpc_over_h = u.def_unit("kpc/h",u.kpc/self._header["h"])
		self.Mpc_over_h = u.def_unit("Mpc/h",u.Mpc/self._header["h"])

		#Compute the comoving distance according to the model
		cosmo = w0waCDM(H0=100.0*h,Om0=Om0,Ode0=Ode0,w0=w0,wa=wa)
		self._header["comoving_distance"] = cosmo.comoving_distance(redshift).to(self.Mpc_over_h)

	def writeParameterFile(self,filename,settings):

		"""
		Writes a Gadget2 parameter file to evolve the current snapshot using Gadget2

		:param filename: name of the file to which to write the parameters
		:type filename: str.

		:param settings: tunable settings of Gadget2 (see Gadget2 manual)
		:type settings: Gadget2Settings

		"""

		#Create output directory if not existent already
		outputdir = settings.OutputDir
		if not(os.path.isdir(outputdir)):
			os.mkdir(outputdir)

		#Update the settings according to the physical units of the current snapshot
		settings.UnitLength_in_cm = self._length_unit
		settings.UnitMass_in_g = self._mass_unit
		settings.UnitVelocity_in_cm_per_s = self._velocity_unit

		#Set the appropriate name for the initial condition file
		if "files" in self.header.keys():
			suffix = self.header["files"][0].split(".")[-1]
			settings.InitCondFile = os.path.abspath(self.header["files"][0].rstrip(".{0}".format(suffix)))

		#Write the options
		with open(filename,"w") as paramfile:

			#Filenames section
			paramfile.write(settings.writeSection("file_names"))

			#CPU time limit section
			paramfile.write(settings.writeSection("cpu_timings"))

			#Code options section
			paramfile.write(settings.writeSection("code_options"))

			#Initial scale factor time
			paramfile.write("{0}			{1}\n".format("TimeBegin",self.header["scale_factor"]))

			#Characteristics of run section
			paramfile.write(settings.writeSection("characteristics_of_run"))

			#Cosmological parameters
			paramfile.write("{0}			{1}\n".format("Omega0",self.header["Om0"]))
			paramfile.write("{0}			{1}\n".format("OmegaLambda",self.header["Ode0"]))
			paramfile.write("{0}			{1}\n".format("OmegaBaryon",0.046))
			paramfile.write("{0}			{1}\n".format("HubbleParam",self.header["h"]))
			paramfile.write("{0}			{1}\n".format("BoxSize",self.header["box_size"].to(self.kpc_over_h).value))
			paramfile.write("{0}			{1}\n".format("w0",self.header["w0"]))
			paramfile.write("{0}			{1}\n\n".format("wa",self.header["wa"]))

			#Output frequency section
			paramfile.write(settings.writeSection("output_frequency"))

			#Accuracy of time integration section
			paramfile.write(settings.writeSection("accuracy_time_integration"))

			#Tree algorithm section
			paramfile.write(settings.writeSection("tree_algorithm"))

			#SPH section
			paramfile.write(settings.writeSection("sph"))

			#Memory allocation section
			paramfile.write(settings.writeSection("memory_allocation"))

			#System of units section
			paramfile.write(settings.writeSection("system_of_units"))

			#Softening lengths section
			paramfile.write(settings.writeSection("softening"))

##############################################################
#################Gadget2SnapshotDE class######################
##############################################################

class Gadget2SnapshotDE(Gadget2Snapshot):

	"""
	A class that handles Gadget2 snapshots, mainly I/O from the binary format and spatial information statistics.Inherits from Gadget2Snapshot; assumes that the header includes Dark Energy information

	"""

	def getHeader(self):
		header = Gadget2Header(ext._gadget2.getHeader(self.fp))
		header["files"] = [self.fp.name]

		return header

##############################################################
#################Gadget2SnapshotNu class######################
##############################################################

class Gadget2SnapshotNu(Gadget2Snapshot):

	"""
	A class that handles Gadget2 snapshots with neutrino effects

	"""

	def getHeader(self):

		header = Gadget2Header(ext._gadget2.getHeader(self.fp))
		header["files"] = [self.fp.name]

		del header["comoving_distance"]
		header["w0"] = -1.
		header["wa"] = 0.

		return header

##################################################################
#################Gadget2SnapshotPipe class########################
##################################################################

class Gadget2SnapshotPipe(Gadget2SnapshotDE):

	"""
	Read in the particle positions when calling the constructor, without calling fseek

	"""

	def __init__(self,*args,**kwargs):

		#Call parent constructor
		super(Gadget2SnapshotPipe,self).__init__(*args,**kwargs)

		#Read in the positions
		npart = self.header["num_particles_file"]
		self.fp.read(8)

		try:
			self.positions = (np.fromstring(self.fp.read(4*3*npart),dtype=np.float32).reshape(npart,3) * self.kpc_over_h).to(self.Mpc_over_h)
		except AttributeError:
			pass

		#Read the rest
		self.fp.read()

		#Particles do not have structure
		self.weights = None
		self.virial_radius = None
		self.concentration = None

import numpy,import astropy,from astropy