"""
@Fire
https://github.com/fire717
"""
import sys

import torch
import math
import numpy as np

import torch.nn.functional as F
import cv2

_img_size = 192
_feature_map_size = _img_size//4

_center_weight_path = 'lib/data/center_weight_origin.npy'

class JointBoneLoss(torch.nn.Module):
    def __init__(self, joint_num):
        super(JointBoneLoss, self).__init__()
        id_i, id_j = [], []
        for i in range(joint_num):
            for j in range(i+1, joint_num):
                id_i.append(i)
                id_j.append(j)
        self.id_i = id_i
        self.id_j = id_j

        # self.id_i = [0,1,2,3,4,5,2]
        # self.id_j = [1,2,3,4,5,6,4]

    def forward(self, joint_out, joint_gt):
        J = torch.norm(joint_out[:,self.id_i,:] - joint_out[:,self.id_j,:], p=2, dim=-1, keepdim=False)
        Y = torch.norm(joint_gt[:,self.id_i,:] - joint_gt[:,self.id_j,:], p=2, dim=-1, keepdim=False)
        loss = torch.abs(J-Y)
        # loss = loss.mean()
        loss = torch.sum(loss)/joint_out.shape[0]/len(self.id_i)
        return loss

class MovenetLoss(torch.nn.Module):
    def __init__(self, use_target_weight=False, target_weight=[1]):
        super(MovenetLoss, self).__init__()
        self.mse = torch.nn.MSELoss(size_average=True)
        self.use_target_weight = use_target_weight
        self.target_weight=target_weight

        self.center_weight = torch.from_numpy(np.load(_center_weight_path))
        self.make_center_w = False

        # self.range_weight_x = torch.from_numpy(np.array([[x for x in range(48)] for _ in range(48)]))
        # self.range_weight_y = self.range_weight_x.T

        self.boneloss = JointBoneLoss(17)

    def l1(self, pre, target,kps_mask):
        # print("1 ",pre.shape, pre.device)
        # print("2 ",target.shape, target.device)
        # b

        # return torch.mean(torch.abs(pre - target)*kps_mask)
        return torch.sum(torch.abs(pre - target)*kps_mask)/ (kps_mask.sum() + 1e-4)

    def l2_loss(self, pre, target):
        loss = (pre - target)
        loss = (loss * loss) / 2 / pre.shape[0]

        return loss.sum()

    def centernetfocalLoss(self, pred, gt):
        ''' Modified focal loss. Exactly the same as CornerNet.
          Runs faster and costs a little bit more memory
        Arguments:
          pred (batch x c x h x w)
          gt_regr (batch x c x h x w)
        '''
        pos_inds = gt.eq(1).float()
        neg_inds = gt.lt(1).float()

        neg_weights = torch.pow(1 - gt, 4)

        loss = 0

        pos_loss = torch.log(pred) * torch.pow(1 - pred, 2) * pos_inds
        neg_loss = torch.log(1 - pred) * torch.pow(pred, 2) * neg_weights * neg_inds

        num_pos  = pos_inds.float().sum()
        pos_loss = pos_loss.sum()
        neg_loss = neg_loss.sum()

        if num_pos == 0:
            loss = loss - neg_loss
        else:
            loss = loss - (pos_loss + neg_loss) / num_pos
        return loss

    def myMSEwithWeight(self, pre, target):
        #target 0-1
        # pre = torch.sigmoid(pre)
        # print(torch.max(pre), torch.min(pre))
        # b
        loss = torch.pow((pre-target),2)
        # loss = torch.abs(pre-target)

        #weight_mask = (target+0.1)/1.1
        weight_mask = target*8+1
        # weight_mask = torch.pow(target,2)*8+1

        #gamma from focal loss
        #gamma = torch.pow(torch.abs(target-pre), 2)

        loss = loss*weight_mask#*gamma

        loss = torch.sum(loss)/target.shape[0]/target.shape[1]

        # bg_loss = self.bgLoss(pre, target)
        return loss

    def heatmapL1(self, pre, target):
        #target 0-1
        # pre = torch.sigmoid(pre)
        # print(torch.max(pre), torch.min(pre))
        # b
        loss = torch.abs(pre-target)

        #weight_mask = (target+0.1)/1.1
        weight_mask = target*4+1

        #gamma from focal loss
        #gamma = torch.pow(torch.abs(target-pre), 2)

        loss = loss*weight_mask#*gamma

        loss = torch.sum(loss)/target.shape[0]/target.shape[1]
        return loss

    ###############
    def boneLoss(self, pred, target):
        #[64, 7, 48, 48]
        def _Frobenius(mat1, mat2):
            return torch.pow(torch.sum(torch.pow(mat1-mat2,2)),0.5)
            # return torch.sum(torch.pow(mat1-mat2,2))

        _bone_idx = [[0,1],[1,2],[2,3],[3,4],[4,5],[5,6],[2,4]]

        loss = 0
        for bone_id in _bone_idx:
            bone_pre = pred[:,bone_id[0],:,:]-pred[:,bone_id[1],:,:]
            bone_gt = target[:,bone_id[0],:,:]-target[:,bone_id[1],:,:]

            f = _Frobenius(bone_pre,bone_gt)
            loss+=f

        loss = loss/len(_bone_idx)/pred.shape[0]
        return loss

    def bgLoss(self, pre, target):
        ##[64, 7, 48, 48]

        bg_pre = torch.sum(pre, axis=1)
        bg_pre = 1-torch.clamp(bg_pre, 0, 1)

        bg_gt = torch.sum(target, axis=1)
        bg_gt = 1-torch.clamp(bg_gt, 0, 1)

        #weight_mask = (1-bg_gt)*4+1

        loss = torch.sum(torch.pow((bg_pre-bg_gt),2))/pre.shape[0]

        return loss

    def heatmapLoss(self, pred, target, batch_size):
        #[64, 7, 48, 48]
        # print(pred.shape, target.shape)

        # heatmaps_pred = pred.reshape((batch_size, pred.shape[1], -1)).split(1, 1)
        # #对tensor在某一dim维度下，根据指定的大小split_size=int，或者list(int)来分割数据，返回tuple元组
        # #print(len(heatmaps_pred), heatmaps_pred[0].shape)#7 torch.Size([64, 1, 48*48]
        # heatmaps_gt = target.reshape((batch_size, pred.shape[1], -1)).split(1, 1)

        # loss = 0

        # for idx in range(pred.shape[1]):
        #     heatmap_pred = heatmaps_pred[idx].squeeze()#[64, 40*40]
        #     heatmap_gt = heatmaps_gt[idx].squeeze()
        #     if self.use_target_weight:
        #         loss += self.centernetfocalLoss(
        #                         heatmap_pred.mul(self.target_weight[idx//2]),
        #                         heatmap_gt.mul(self.target_weight[idx//2])
        #                     )
        #     else:

        #         loss += self.centernetfocalLoss(heatmap_pred, heatmap_gt)
        # loss /= pred.shape[1]

        return self.myMSEwithWeight(pred,target)

    def centerLoss(self, pred, target, batch_size):
        # heatmaps_pred = pred.reshape((batch_size, -1))
        # heatmaps_gt = target.reshape((batch_size, -1))
        return self.myMSEwithWeight(pred, target)

    def regsLoss(self, pred, target, cx0, cy0,  kps_mask, batch_size, num_joints):
        #[64, 14, 48, 48]
        # print(target.shape, cx0.shape, cy0.shape)#torch.Size([64, 14, 48, 48]) torch.Size([64]) torch.Size([64])

        _dim0 = torch.arange(0,batch_size).long()
        _dim1 = torch.zeros(batch_size).long()

        #print("regsLoss: " , cx0,cy0)
        # print(target.shape)#torch.Size([1, 14, 48, 48])
        # print(torch.max(target[0][2]), torch.min(target[0][2]))
        # print(torch.max(target[0][3]), torch.min(target[0][3]))

        # cv2.imwrite("t.jpg", target[0][2].cpu().numpy()*255)
        loss = 0
        for idx in range(num_joints):

            gt_x = target[_dim0,_dim1+idx*2,cy0,cx0]
            gt_y = target[_dim0,_dim1+idx*2+1,cy0,cx0]

            pre_x = pred[_dim0,_dim1+idx*2,cy0,cx0]
            pre_y = pred[_dim0,_dim1+idx*2+1,cy0,cx0]

            # print(torch.max(target[_dim0,_dim1+idx*2,:,:]),torch.min(target[_dim0,_dim1+idx*2,:,:]))
            # print(gt_x,pre_x)
            # print(gt_y,pre_y)

            # print(kps_mask[:,idx])
            # print(gt_x,pre_x)
            # print(self.l1(gt_x,pre_x,kps_mask[:,idx]))
            # print('---')
            #

            loss+=self.l1(gt_x,pre_x,kps_mask[:,idx])
            loss+=self.l1(gt_y,pre_y,kps_mask[:,idx])
        #b
        # offset_x_pre = torch.clip(pre_x,0,_feature_map_size-1).long()
        # offset_y_pre = torch.clip(pre_y,0,_feature_map_size-1).long()
        # offset_x_gt = torch.clip(gt_x+cx0,0,_feature_map_size-1).long()
        # offset_y_gt = torch.clip(gt_y+cy0,0,_feature_map_size-1).long()

        return loss / num_joints

    def offsetLoss(self, pred, target,  cx0, cy0, regs, kps_mask, batch_size, num_joints):
        _dim0 = torch.arange(0,batch_size).long()
        _dim1 = torch.zeros(batch_size).long()
        loss = 0
        # print(gt_y,gt_x)
        for idx in range(num_joints):
            gt_x = regs[_dim0,_dim1+idx*2,cy0,cx0].long()+cx0
            gt_y = regs[_dim0,_dim1+idx*2+1,cy0,cx0].long()+cy0

            gt_x[gt_x>47]=47
            gt_x[gt_x<0]=0
            gt_y[gt_y>47]=47
            gt_y[gt_y<0]=0

            gt_offset_x = target[_dim0,_dim1+idx*2,gt_y,gt_x]
            gt_offset_y = target[_dim0,_dim1+idx*2+1,gt_y,gt_x]

            pre_offset_x = pred[_dim0,_dim1+idx*2,gt_y,gt_x]
            pre_offset_y = pred[_dim0,_dim1+idx*2+1,gt_y,gt_x]

            # print(gt_offset_x, torch.max(target[_dim0,_dim1+idx*2,...]),torch.min(target[_dim0,_dim1+idx*2,...]))
            # print(gt_offset_y, torch.max(target[_dim0,_dim1+idx*2+1,...]),torch.min(target[_dim0,_dim1+idx*2+1,...]))
            loss+=self.l1(gt_offset_x,pre_offset_x,kps_mask[:,idx])
            loss+=self.l1(gt_offset_y,pre_offset_y,kps_mask[:,idx])
        #     print(gt_y,gt_x)
        # b
        return loss / num_joints

        """
        0.0 0.5
        0.0 0.75
        0.75 0.25
        0.0 0.75
        0.0 0.5
        """

    def maxPointPth(self, heatmap, center=True):
        #pytorch version
        # n,1,h,w
        # 计算center heatmap的最大值得到中心点
        if center:
            heatmap = heatmap*self.center_weight[:heatmap.shape[0],...]
            #加权取最靠近中间的

        n,c,h,w = heatmap.shape
        heatmap = heatmap.reshape((n, -1)) #64, 48x48
        # print(heatmap[0])
        # max_id = torch.argmax(heatmap, 1)#64, 1
        # print(max_id)
        max_v,max_id = torch.max(heatmap, 1)#64, 1
        # print(max_v)
        # print("max_i: ",max_i)

        # mask0 = torch.zeros(max_v.shape).to(heatmap.device)
        # mask1 = torch.ones(max_v.shape).to(heatmap.device)
        # mask = torch.where(torch.gt(max_v,th), mask1, mask0)
        # print(mask)
        # b
        y = max_id//w
        x = max_id%w

        return x,y

    def forward(self, output, target, kps_mask):
        batch_size = output[0].size(0)
        num_joints = output[0].size(1)
        #print("output: ", [x.shape for x in output])
        #[64, 7, 48, 48] [64, 1, 48, 48] [64, 14, 48, 48] [64, 14, 48, 48]
        # print("target: ", [x.shape for x in target])#[64, 36, 48, 48]
        #print(weights.shape)# [14,]
        heatmaps = target[:,:17,:,:]
        centers = target[:,17:18,:,:]
        regs = target[:,18:52,:,:]
        offsets = target[:,52:,:,:]

        heatmap_loss = self.heatmapLoss(output[0], heatmaps, batch_size)

        # bg_loss = self.bgLoss(output[0], heatmaps)
        #bone_loss = self.boneloss(output[0], heatmaps)
        bone_loss = self.boneLoss(output[0], heatmaps)
        #print(heatmap_loss)
        center_loss = self.centerLoss(output[1], centers, batch_size)

        if not self.make_center_w:
            self.center_weight = torch.reshape(self.center_weight,(1,1,48,48))
            self.center_weight = self.center_weight.repeat((output[1].shape[0],output[1].shape[1],1,1))
            # print(self.center_weight.shape)
            # b
            self.center_weight = self.center_weight.to(target.device)
            self.make_center_w = True
            self.center_weight.requires_grad_(False)

            # self.range_weight_x = self.range_weight_x.to(target.device)
            # self.range_weight_y = self.range_weight_y.to(target.device)
            # self.range_weight_x.requires_grad_(False)
            # self.range_weight_y.requires_grad_(False)
        #print(self.center_weight)

        cx0, cy0 = self.maxPointPth(centers)
        # cx1, cy1 = self.maxPointPth(pre_centers)
        cx0 = torch.clip(cx0,0,_feature_map_size-1).long()
        cy0 = torch.clip(cy0,0,_feature_map_size-1).long()
        # cx1 = torch.clip(cx1,0,_feature_map_size-1).long()
        # cy1 = torch.clip(cy1,0,_feature_map_size-1).long()

        # print(cx0, cy0)
        # bbb
        # cv2.imwrite("_centers.jpg", centers[0][0].cpu().numpy()*255)
        # b

        regs_loss = self.regsLoss(output[2], regs, cx0, cy0, kps_mask,batch_size, num_joints)
        offset_loss = self.offsetLoss(output[3], offsets,
                            cx0, cy0,regs,
                            kps_mask,batch_size, num_joints)

        # total_loss = heatmap_loss+center_loss+0.1*regs_loss+offset_loss
        # print(heatmap_loss,center_loss,regs_loss,offset_loss)
        # b

        """

        """
        # boneloss = self.boneLoss(output[3], offsets,
        #                     cx0, cy0,regs,
        #                     kps_mask,batch_size, num_joints)

        return [heatmap_loss,bone_loss,center_loss,regs_loss,offset_loss]

movenetLoss = MovenetLoss(use_target_weight=False)

def calculate_loss(predict, label):
    loss = movenetLoss(predict, label)
    return loss


from scipy.interpolate import interp1d
from pyHalo.defaults import *
from pyHalo.Cosmology.cosmology import Cosmology
from pyHalo.Halos.lens_cosmo import LensCosmo
from pyHalo.Halos.HaloModels.NFW import NFWSubhhalo, NFWFieldHalo
from pyHalo.Halos.HaloModels.TNFW import TNFWFieldHalo, TNFWSubhalo
from pyHalo.Halos.HaloModels.PsuedoJaffe import PJaffeSubhalo
from pyHalo.Halos.HaloModels.PTMass import PTMass
from pyHalo.Halos.HaloModels.coreTNFW import coreTNFWFieldHalo, coreTNFWSubhalo
from pyHalo.Halos.HaloModels.ULDM import ULDMFieldHalo, ULDMSubhalo
import numpy as np
from copy import deepcopy

def realization_at_z(realization, z, angular_coordinate_x=None, angular_coordinate_y=None, max_range=None):
    """
    :param realization: an instance of Realization
    :param z: the redshift where we want to extract halos
    :param angular_coordinate_x: if max_range is specified, will only keep halos within
    max_range of (angular_coordinate_x, angular_coordinate_y)
    :param angular_coordinate_y:
    :param max_range: radius in arcseconds where we want to keep halos. If None, will return a new realization class
     that contains all halos at redshift z contained in the input realization class
    :return: a new instance of Realization
    """

    _halos, _indexes = realization.halos_at_z(z)
    halos = []
    indexes = []

    if max_range is not None:
        for i, halo in enumerate(_halos):
            dx, dy = halo.x - angular_coordinate_x, halo.y - angular_coordinate_y
            dr = (dx ** 2 + dy ** 2) ** 0.5
            if dr < max_range:
                halos.append(halo)
                indexes.append(i)
    else:
        halos = _halos
        indexes = _indexes

    centerx, centery = realization.rendering_center

    return Realization.from_halos(halos, realization.lens_cosmo,
                                  realization._prof_params,
                                  realization._mass_sheet_correction,
                                  realization.rendering_classes,
                                  centerx, centery), indexes

class Realization(object):

    """
    This is the main class for storing a population of dark matter halos, both in the main lens plane and along the
    line of sight. This class is created by the main pyhalo module.
    """

    def __init__(self, masses, x, y, r3d, mdefs, z, subhalo_flag, lens_cosmo,
                 halos=None, kwargs_realization={}, mass_sheet_correction=True,
                 rendering_classes=None, rendering_center_x=None, rendering_center_y=None,
                 geometry=None):

        """

        This class is the main class that stores information regarding realizations of dark matter halos. It is not
        intended to be created directly by the user. Instances of this class are created through the class
        pyHalo or pyHalo_dynamic.

        :param masses: an array of halo masses (units solar mass)
        :param x: an array of halo x-coordinates (units arcsec)
        :param y: an array of halo y-coordinates (units arcsec)
        :param r2d: an array of halo 2-d distances from lens center (units kpc / (kpc/arsec), or arcsec,
        at halo redshift)
        :param r3d: an array of halo 2-d distances from lens center (units kpc / (kpc/arsec), or arcsec,
        at halo redshift)
        :param mdefs: mass definition of each halo
        :param z: halo redshift
        :param subhalo_flag: whether each halo is a subhalo or a regular halo
        :param lens_cosmo: an instance of LensCosmo (see Halos.lens_cosmo)
        :param halos: a list of halo class instances
        :param kwargs_realization: kwargs for the realiztion
        :param mass_sheet_correction: whether to apply a mass sheet correction
        :param rendering_classes: a list of rendering class instances
        :param rendering_center_x: an instance of scipy.interp1d that returns an angular position given a comoving distance.
        The angular coordinate defines the center of the rendering volume, and halos will be distributed symmetrically around it.
        The value defaults to 0, but is overridden when the method shift_background_to_source is called
        :param rendering_center_y: same as rendering_center_x, but for the y angular coordinate
        :param geometry: (optional, only relevant is subtract_exact_mass_sheets=True is specified in kwargs_realization)
        an instance of Geometry (pyHalo.Cosmology.geometry) that defines the rendering volume
        """

        self.apply_mass_sheet_correction = mass_sheet_correction
        self.geometry = geometry
        self.lens_cosmo = lens_cosmo
        self._zlens, self._zsource = self.lens_cosmo.z_lens, self.lens_cosmo.z_source
        self.astropy_instance = self.lens_cosmo.cosmo.astropy
        self.halos = []
        self._loaded_models = {}
        self._has_been_shifted = False
        self._prof_params = set_default_kwargs(kwargs_realization, self._zsource)

        if halos is None:

            for mi, xi, yi, r3di, mdefi, zi, sub_flag in zip(masses, x, y, r3d,
                           mdefs, z, subhalo_flag):

                unique_tag = np.random.rand()
                model = self._load_halo_model(mi, xi, yi, r3di, mdefi, zi, sub_flag, self.lens_cosmo,
                                              self._prof_params, unique_tag)
                self.halos.append(model)

        else:

            self.halos = halos

        self._reset()

        self.set_rendering_classes(rendering_classes)

        if rendering_center_x is None or rendering_center_y is None:
            _z = np.linspace(0, self._zsource, 100)
            d = [self.lens_cosmo.cosmo.D_C_transverse(zi) for zi in _z]
            angle = np.zeros_like(d)
            rendering_center_x = interp1d(d, angle)
            rendering_center_y = interp1d(d, angle)

        self._rendering_center_x = rendering_center_x
        self._rendering_center_y = rendering_center_y

    @classmethod
    def from_halos(cls, halos, lens_cosmo, prof_params, msheet_correction, rendering_classes,
                   rendering_center_x=None, rendering_center_y=None, geometry=None):

        """

        :param halos: a list of halo class instances
        :param lens_cosmo: an instance of LensCosmo (see Halos.lens_cosmo)
        :param prof_params: keyword arguments for the realization
        :param msheet_correction: whether or not to apply a mass sheet correction
        :param rendering_classes: a list of rendering classes
        :param rendering_center_x: an instance of scipy.interp1d that returns an angular position given a comoving distance.
        The angular coordinate defines the center of the rendering volume, and halos will be distributed symmetrically around it.
        The value defaults to 0, but is overridden when the method shift_background_to_source is called
        :param rendering_center_y: same as rendering_center_x, but for the y angular coordinate
        :param geometry: (optional, only relevant is subtract_exact_mass_sheets=True is specified in kwargs_realization)
        an instance of Geometry (pyHalo.Cosmology.geometry) that defines the rendering volume
        :return: an instance of Realization created directly from the halo class instances
        """

        realization = Realization(None, None, None, None, None, None, None, lens_cosmo,
                                  halos=halos, kwargs_realization=prof_params,
                                  mass_sheet_correction=msheet_correction,
                                  rendering_classes=rendering_classes,
                                  rendering_center_x=rendering_center_x,
                                  rendering_center_y=rendering_center_y,
                                  geometry=geometry)

        return realization

    @property
    def rendering_center(self):

        """
        Returns the instances of scipy.interp1d that compute the coordinate center of the lensing volume given a comoving
        distance.
        """
        return self._rendering_center_x, self._rendering_center_y

    def filter(self, aperture_radius_front,
               aperture_radius_back,
               log_mass_allowed_in_aperture_front,
               log_mass_allowed_in_aperture_back,
               log_mass_allowed_global_front,
               log_mass_allowed_global_back,
               interpolated_x_angle, interpolated_y_angle,
               zmin=None, zmax=None, aperture_units='ANGLES'):

        """

        :param aperture_radius_front: the radius of a circular window around each light ray where halos are halo kept
        if they are more massive than log_mass_allowed_in_aperture_front (applied for z < z_lens)
        :param aperture_radius_back: the radius of a circular window around each light ray where halos are halo kept
        if they are more massive than log_mass_allowed_in_aperture_back (applied for z < z_lens)
        :param log_mass_allowed_in_aperture_front: the minimum halo mass to be kept inside the tube around each light ray
        in the foreground
        :param log_mass_allowed_in_aperture_back: the minimum halo mass to be kept inside the tube around each light ray
        in the background
        :param log_mass_allowed_global_front: The minimum mass to be kept everywhere in the foreground (if this is smaller
        than log_mass_allowed_in_aperture_front, then the argument aperture_radius_front will have no effect)
        :param log_mass_allowed_global_back: The minimum mass to be kept everywhere in the background (if this is smaller
        than log_mass_allowed_in_aperture_back, then the argument aperture_radius_back will have no effect)
        :param interpolated_x_angle: a list of scipy.interp1d that retuns the x angular position of a ray in
        arcsec given a comoving distance
        :param interpolated_y_angle: a list of scipy.interp1d that retuns the y angular position of a ray in
        arcsec given a comoving distance
        :param zmin: only keep halos at z > zmin
        :param zmax: only keep halos at z < zmax
        :param aperture_units: either 'ANGLES' or 'MPC'

        - If 'ANGLES', then halos are kept inside angular apertures
        around each light ray with size aperture_radius_front/aperture_radius_back.
        - If 'MPC', then halos are kept inside circular apertures with radius
        R = aperture_radius_front/back * mpc_per_arcsec(0.5)
        where D_C(0.5) is the comoving transverse distance at z = 0.5. The unit of R is arcsec * Mpc

        'ANGLES' is more conservative in that it keeps more halos in the lens model; the rendering area is basically a cone
        since the aperture size is a fixed angle at every redshift, whereas 'MPC' distributes halos in cylindrical
        tubes around each light ray along the line of sight.

        :return: A new instance of Realization with the cuts on position and mass applied
        """
        halos = []

        if zmax is None:
            zmax = self._zsource
        if zmin is None:
            zmin = 0

        for plane_index, zi in enumerate(self.unique_redshifts):

            plane_halos, _ = self.halos_at_z(zi)
            inds_at_z = np.where(self.redshifts == zi)[0]
            x_at_z = self.x[inds_at_z]
            y_at_z = self.y[inds_at_z]
            masses_at_z = self.masses[inds_at_z]

            if zi < zmin:
                continue
            if zi > zmax:
                continue

            comoving_distance_z = self.lens_cosmo.cosmo.D_C_z(zi)

            if zi <= self._zlens:

                minimum_mass_everywhere = deepcopy(log_mass_allowed_global_front)
                minimum_mass_in_window = deepcopy(log_mass_allowed_in_aperture_front)
                aperture_radius_arcsec = deepcopy(aperture_radius_front)

            else:

                minimum_mass_everywhere = deepcopy(log_mass_allowed_global_back)
                minimum_mass_in_window = deepcopy(log_mass_allowed_in_aperture_back)
                aperture_radius_arcsec = deepcopy(aperture_radius_back)

            keep_inds_mass = np.where(masses_at_z >= 10 ** minimum_mass_everywhere)[0]

            inds_m_low = np.where(masses_at_z < 10 ** minimum_mass_everywhere)[0]

            keep_inds_dr = []
            for idx in inds_m_low:
                for k, (interp_x, interp_y) in enumerate(zip(interpolated_x_angle, interpolated_y_angle)):

                    dx = x_at_z[idx] - interp_x(comoving_distance_z)
                    dy = y_at_z[idx] - interp_y(comoving_distance_z)

                    if aperture_units == 'ANGLES':
                        dr_cut = aperture_radius_arcsec

                    elif aperture_units == 'MPC':
                        dx *= comoving_distance_z
                        dy *= comoving_distance_z
                        dr_cut = aperture_radius_arcsec * self.lens_cosmo.cosmo.D_C_z(0.5)
                    else:
                        raise Exception('aperture units must be either MPC or ANGLES')

                    dr = np.sqrt(dx ** 2 + dy ** 2)

                    if dr <= dr_cut:
                        keep_inds_dr.append(idx)
                        break

            keep_inds = np.append(keep_inds_mass, np.array(keep_inds_dr)).astype(int)

            tempmasses = masses_at_z[keep_inds]
            keep_inds = keep_inds[np.where(tempmasses >= 10 ** minimum_mass_in_window)[0]]

            for halo_index in keep_inds:
                halos.append(plane_halos[halo_index])

        return Realization.from_halos(halos, self.lens_cosmo, self._prof_params,
                                      self.apply_mass_sheet_correction, self.rendering_classes,
                                      self._rendering_center_x, self._rendering_center_y, self.geometry)

    def set_rendering_classes(self, rendering_classes):

        """
        This method sets the rendering classes for the realization, which are used to apply the negative convergence sheet
        corrections after adding halos. The properties of the convergence sheets you need to add depend on the form of the
        mass function you have specified, so this information is stored in the classes used to render halos
        (LOSPowerLaw, MainLensPowerLaw, etc, see classes in pyHalo/Rendering)

        If the rendering classes are not specified for whatever reason, the code will still run but no negative convergence
        sheets will be included in your lens models. This could potentially bias results as you've effectively made every
        light cone overdense relative to the mean matter density in the Universe.

        :param rendering_classes: a list or an instance of a rendering class (LOSPowerLaw, MainLensPowerLaw)
        """

        if not isinstance(rendering_classes, list):
            rendering_classes = [rendering_classes]
        self.rendering_classes = rendering_classes

    def join(self, real, join_rendering_classes=False):

        """
        This routine combines one realization with another realization, keeping only the unqiue halos
        (as identified by their .unique_tag attribute) between them.

        :param real: another realization, possibly a filtered version of self
        :param join_rendering_classes: If True, the rendering classes associated with the new
        realization will include both the rendering class associated with self and that of real.

        :return: a new realization that contains all unique halos from self and real
        """

        halos = []

        tags = self._tags(self.halos)
        real_tags = self._tags(real.halos)
        if len(tags) >= len(real_tags):
            long, short = tags, real_tags
            halos_long, halos_short = self.halos, real.halos
        else:
            long, short = real_tags, tags
            halos_long, halos_short = real.halos, self.halos

        for halo in halos_short:
            halos.append(halo)

        for i, tag in enumerate(long):

            if tag not in short:
                halos.append(halos_long[i])

        if join_rendering_classes:
            rendering_class_self = self.rendering_classes
            rendering_class_new = real.rendering_classes
            rendering_classes = rendering_class_self + rendering_class_new
        else:
            rendering_classes = self.rendering_classes

        centerx, centery = self.rendering_center
        return Realization.from_halos(halos, self.lens_cosmo, self._prof_params,
                                      self.apply_mass_sheet_correction, rendering_classes,
                                      centerx, centery, self.geometry)

    def shift_background_to_source(self, ray_interp_x, ray_interp_y):

        """
        This routine shifts the entire relation along a path specified by ray_interp_x/y. This routine is intended
        to be used in situations where the source is significantly offset from the origin, and you want to align the
        center of the rendering volume such that tracks the path of the light.

        :param ray_interp_x: instance of scipy.interp1d, returns the angular position of a ray
        fired through the lens center given a comoving distance
        :param ray_interp_y: same but for the y coordinate
        :return:
        """

        if self._has_been_shifted:
            return self

        halos = []

        for halo in self.halos:

            comoving_distance_z = self.lens_cosmo.cosmo.D_C_z(halo.z)

            xshift, yshift = ray_interp_x(comoving_distance_z), ray_interp_y(comoving_distance_z)

            halo.x += xshift
            halo.y += yshift
            halos.append(halo)

        new_realization = Realization.from_halos(halos, self.lens_cosmo, self._prof_params, self.apply_mass_sheet_correction,
                                                 self.rendering_classes, ray_interp_x, ray_interp_y, self.geometry)

        new_realization._has_been_shifted = True

        return new_realization

    def lensing_quantities(self, add_mass_sheet_correction=True, z_mass_sheet_max=None,
                           kwargs_mass_sheet_correction=None):

        """
        :param add_mass_sheet_correction: include sheets of negative convergence to correct for mass added subhalos/field halos
        :param z_mass_sheet_max: don't include negative convergence sheets at z>z_mass_sheet_max (this does nothing
        if the previous argument is False
        :pararm kwargs_mass_sheet_correction: additional keyword arguments for the mass sheet correction
        (changes the default setting)
        :return: the lens_model_list, redshift_list, kwargs_lens, and numerical_alpha_class keywords that can be plugged
        directly into a lenstronomy LensModel class
        """

        kwargs_lens = []
        lens_model_list = []
        redshift_array = []
        numerical_interp = None

        for i, halo in enumerate(self.halos):

            lens_model_name = halo.lenstronomy_ID
            kwargs_halo, interp_class = halo.lenstronomy_params
            lens_model_list += lens_model_name
            kwargs_lens += kwargs_halo
            redshift_array += [halo.z] * len(lens_model_name)

            if interp_class is not None:
                numerical_interp = interp_class

        if self.apply_mass_sheet_correction and add_mass_sheet_correction:

            if self.rendering_classes is None:
                raise Exception('if applying a convergence sheet correction, must specify '
                                'the rendering classes.')

            kwargs_mass_sheets, profile_list, z_sheets = self._mass_sheet_correction(self.rendering_classes,
                                                                                z_mass_sheet_max,
                                                                                kwargs_mass_sheet_correction)
            kwargs_lens += kwargs_mass_sheets
            lens_model_list += profile_list
            redshift_array = np.append(redshift_array, z_sheets)

        return lens_model_list, redshift_array, kwargs_lens, numerical_interp

    def split_at_z(self, z):
        """
        Splits the realization at redshift z, returning one instance at Realization containing all halos with
        redshift < zlens and another with all halos at redshift >= z. Be careful with the mass sheet corrections contained
        in the rendering_classes, as both new realizations will get all rendering classes from the parent realization.

        :param z: the redshift at which to split the realization
        :return: two instances at Realization divided at redshift z
        """

        halos_1, halos_2 = [], []
        for halo in self.halos:
            if halo.z <= z:
                halos_1.append(halo)
            else:
                halos_2.append(halo)

        centerx, centery = self.rendering_center
        realization_1 = Realization.from_halos(halos_1, self.lens_cosmo,
                                               self._prof_params, self.apply_mass_sheet_correction, self.rendering_classes,
                                               centerx, centery, self.geometry)
        realization_2 = Realization.from_halos(halos_2, self.lens_cosmo,
                                               self._prof_params, self.apply_mass_sheet_correction, self.rendering_classes,
                                               centerx, centery, self.geometry)

        return realization_1, realization_2

    def halo_comoving_coordinates(self):

        """
        :param halos: a list of halos
        :return: the comoving (x, y) position, mass, and redshift of each halo in the realization
        """
        xcoords, ycoords, masses, redshifts = [], [], [], []

        for halo in self.halos:
            D = self.lens_cosmo.cosmo.D_C_z(halo.z)
            x_arcsec, y_arcsec = halo.x, halo.y
            x_comoving, y_comoving = D * x_arcsec, D * y_arcsec
            xcoords.append(x_comoving)
            ycoords.append(y_comoving)
            masses.append(halo.mass)
            redshifts.append(halo.z)

        return np.array(xcoords), np.array(ycoords), np.log10(masses), np.array(redshifts)

    def halos_at_z(self, z):
        """

        :param z: redshift
        :return: all halos in the realization that are at redshift z
        """
        halos = []
        index = []
        for i, halo in enumerate(self.halos):
            if halo.z != z:
                continue
            halos.append(halo)
            index.append(i)

        return halos, index

    def mass_at_z_exact(self, z):

        """
        Computes the total mass rendered at each redshift z
        :param z: redshift
        :return: total mass rendered at z
        """

        inds = np.where(self.redshifts == z)
        m_exact = np.sum(self.masses[inds])
        return m_exact

    def number_of_halos_before_redshift(self, z):

        """
        Computes the number of halos with redshifts < z
        :param z: redshift
        :return: number of halos with redshift < z
        """
        n = 0
        for halo in self.halos:
            if halo.z < z:
                n += 1
        return n

    def number_of_halos_after_redshift(self, z):

        """
        Computes the number of halos with redshifts > z
        :param z: redshift
        :return: number of halos with redshift > z
        """
        n = 0
        for halo in self.halos:
            if halo.z > z:
                n += 1
        return n

    def number_of_halos_at_redshift(self, z):

        """
        Computes the number of halos with redshifts == z
        :param z: redshift
        :return: number of halos at z
        """

        n = 0
        for halo in self.halos:
            if halo.z == z:
                n += 1
        return n

    def _mass_sheet_correction(self, rendering_classes, z_mass_sheet_max, kwargs_mass_sheet_correction):

        """
        This routine adds the negative mass sheet corrections along the LOS and in the main lens plane.
        The actual physics that determines the amount of negative convergence to add is encoded in the rendering_classes
        (see for example Rendering.Field.PowerLaw.powerlaw_base.py)

        :param rendering_classes: the rendering class associated with each realization
        :param z_mass_sheet_max: don't add mass sheets at lens planes with redshift > z_mass_sheet_max
        :param kwargs_mass_sheet_correction: keyword arguments for the convergence sheet correction
        :return: the kwargs_lens, lens_model_list, and redshift_list of the mass sheets that can be plugged into lenstronomy
        """

        kwargs_mass_sheets = []

        redshifts = []

        profiles = []

        if self._prof_params['subtract_exact_mass_sheets']:

            for zi in self.unique_redshifts:
                area = self.geometry.angle_to_physical_area(0.5 * self.geometry.cone_opening_angle, zi)
                kwargs_mass_sheets += [{'kappa_ext': -self.mass_at_z_exact(zi) / self.lens_cosmo.sigma_crit_mass(zi, area)}]

            redshifts = self.unique_redshifts

            profiles = ['CONVERGENCE'] * len(kwargs_mass_sheets)

        else:

            for rendering_class in rendering_classes:

                if rendering_class is None:
                    continue

                kwargs_new, profiles_new, redshifts_new = \
                    rendering_class.convergence_sheet_correction(kwargs_mass_sheet_correction)

                kwargs_mass_sheets += kwargs_new
                redshifts += redshifts_new
                profiles += profiles_new

        if z_mass_sheet_max is not None:
            kwargs_mass_sheets_out = []
            profiles_out = []
            redshifts_out = []

            inds_keep = np.where(np.array(redshifts) <= z_mass_sheet_max)[0]

            for i in range(0, len(kwargs_mass_sheets)):
                if i in inds_keep:
                    kwargs_mass_sheets_out.append(kwargs_mass_sheets[i])
                    profiles_out.append(profiles[i])
                    redshifts_out.append(redshifts[i])
        else:
            kwargs_mass_sheets_out, profiles_out, redshifts_out = kwargs_mass_sheets, profiles, redshifts

        # define the center of mass sheet to be the center of the rendering volume
        centerx_interp, centery_interp = self.rendering_center

        for i, (zi, profile_name) in enumerate(zip(redshifts_out, profiles_out)):
            di = self.lens_cosmo.cosmo.D_C_z(zi)
            x_center, y_center = centerx_interp(di), centery_interp(di)
            if profile_name == 'CONVERGENCE':
                kwargs_mass_sheets_out[i]['ra_0'] = float(x_center)
                kwargs_mass_sheets_out[i]['dec_0'] = float(y_center)
            else:
                kwargs_mass_sheets_out[i]['center_x'] = float(x_center)
                kwargs_mass_sheets_out[i]['center_y'] = float(y_center)

        return kwargs_mass_sheets_out, profiles_out, redshifts_out

    def _load_halo_model(self, mass, x, y, r3d, mdef, z, is_subhalo,
                         lens_cosmo_instance, args, unique_tag):

        """
        Loads the halo model for each object based on the mass definition
        :param halo: an instance of Halo
        :return: the specific Halo class corresponding to mass definition mdef
        """

        if mdef == 'NFW':

            if is_subhalo:
                model = NFWSubhhalo(mass, x, y, r3d, mdef, z, is_subhalo,
                                    lens_cosmo_instance, args, unique_tag)
            else:
                model = NFWFieldHalo(mass, x, y, r3d, mdef, z, is_subhalo,
                                    lens_cosmo_instance, args, unique_tag)

        elif mdef == 'TNFW':

            if is_subhalo:
                model = TNFWSubhalo(mass, x, y, r3d, mdef, z, is_subhalo,
                                    lens_cosmo_instance, args, unique_tag)

            else:
                model = TNFWFieldHalo(mass, x, y, r3d, mdef, z, is_subhalo,
                                      lens_cosmo_instance, args, unique_tag)

        elif mdef == 'PT_MASS':

            model = PTMass(mass, x, y, r3d, mdef, z, is_subhalo,
                           lens_cosmo_instance, args, unique_tag)

        elif mdef == 'PJAFFE':

            model = PJaffeSubhalo(mass, x, y, r3d, mdef, z, is_subhalo,
                                  lens_cosmo_instance, args, unique_tag)

        elif mdef == 'coreTNFW':

            if is_subhalo:
                model = coreTNFWSubhalo(mass, x, y, r3d, mdef, z, is_subhalo,
                                  lens_cosmo_instance, args, unique_tag)
            else:
                model = coreTNFWFieldHalo(mass, x, y, r3d, mdef, z, is_subhalo,
                                  lens_cosmo_instance, args, unique_tag)

        elif mdef == 'ULDM':

            if is_subhalo:
                model = ULDMSubhalo(mass, x, y, r3d, mdef, z, is_subhalo,
                                  lens_cosmo_instance, args, unique_tag)
            else:
                model = ULDMFieldHalo(mass, x, y, r3d, mdef, z, is_subhalo,
                                  lens_cosmo_instance, args, unique_tag)

        else:
            raise ValueError('halo profile ' + str(mdef) + ' not recongnized.')

        return model

    def _tags(self, halos=None):

        """

        :param halos: a list of halos
        :return: the unique tag for each halo in halos; if halos is not specified, returns the unique tag for each
        halo in the realization
        """
        if halos is None:
            halos = self.halos
        tags = []

        for halo in halos:

            tags.append(halo.unique_tag)

        return tags

    def _reset(self):
        """
        Resets all class attributes to the current set of halos contained in the realization
        :return:
        """

        self.x = []
        self.y = []
        self.masses = []
        self.redshifts = []
        self.r3d = []
        self.mdefs = []
        self._halo_tags = []
        self.subhalo_flags = []

        for halo in self.halos:
            self.masses.append(halo.mass)
            self.x.append(halo.x)
            self.y.append(halo.y)
            self.redshifts.append(halo.z)
            self.r3d.append(halo.r3d)
            self.mdefs.append(halo.mdef)
            self._halo_tags.append(halo.unique_tag)
            self.subhalo_flags.append(halo.is_subhalo)

        self.masses = np.array(self.masses)
        self.x = np.array(self.x)
        self.y = np.array(self.y)
        self.r3d = np.array(self.r3d)
        self.redshifts = np.array(self.redshifts)

        self.unique_redshifts = np.unique(self.redshifts)

    def __eq__(self, other_reealization):

        """
        Defintes equality between two realizations if they contain the same halos with the same unique tags
        :param other_reealization:
        :return:
        """
        tags = self._tags(self.halos)
        other_tags = other_reealization._tags()
        for tag in other_tags:
            if tag not in tags:
                return False
        else:
            return True

class SingleHalo(Realization):

    def __init__(self, halo_mass, x, y, mdef, z, zlens, zsource, r3d=None, subhalo_flag=False,
                 kwargs_halo={}, cosmo=None):

        """
       Useful for generating a realization with a single or a few
        user-specified halos.
        :param halo_mass: mass of the halo in M_sun
        :param x: halo x coordinate in arcsec
        :param y: halo y coordinate in arcsec
        :param mdef: halo mass definition
        :param z: halo redshift
        :param zlens: main deflector redshift
        :param zsource: source redshift
        :param r3d: three dimensional coordinate of halo inside the host in kpc
        (only relevant for tidally-truncated subhalos, for field halos this can be None)
        :param subhalo_flag: bool, sets whether or not a halo is a subhalo
        :param kwargs_halo: keyword arguments for the halo
        :param cosmo: an instance of Cosmology(); if none is provided a default cosmology will be used
        """
        if cosmo is None:
            cosmo = Cosmology()

        lens_cosmo = LensCosmo(zlens, zsource, cosmo)

        # these are redundant keywords for a single halo, but we need to specify them anyways
        kwargs_halo.update({'cone_opening_angle': 6., 'log_mlow': 6., 'log_mhigh': 10.})
        super(SingleHalo, self).__init__([halo_mass], [x], [y],
                                         [r3d], [mdef], [z], [subhalo_flag], lens_cosmo,
                                         kwargs_realization=kwargs_halo, mass_sheet_correction=False)


from .utils import PyKEArgumentHelpFormatter
from . import kepio, kepmsg, kepkey, kepfunc, kepstat
from astropy.io import fits as pyfits
from matplotlib import pyplot as plt
import numpy as np
from tqdm import tqdm

__all__ = ['kepfilter']

def kepfilter(infile, passband, outfile=None, datacol='SAP_FLUX', function='boxcar',
              cutoff=1.0, plot=False, overwrite=False, verbose=False,
              logfile='kepfilter.log'):
    """
    kepfilter -- bandpass filtering of Kepler light curve data

    ``kepfilter`` applies a bandpass filter to Kepler light curve data. In the
    low bandpass option, the data is convolved with a function of
    user-specified width. Choices of convolution function are **boxcar**,
    **Gaussian** or **sinc**. In the high bandpass option the convolution minus
    the median of the convolution is subtracted from the original data. The
    filtered data is copied to a new FITS file with the same structure as the
    input file.

    Parameters
    ----------
    infile : str
        The name of a MAST standard format FITS file containing Kepler light
        curve data within the first data extension.
    passband : str
        The type of filter to be applied. A low bandpass filter will suppress
        high-frequency signal shorter than the cutoff. A high bandpass filter
        will suppress low-frequency signal longer than the cutoff.
        The options are:

        * low
        * high
    outfile : str
        The name of the output FITS file. The output file is identical in
        format to the input file. The data to be filtered will be overwritten
        in the output file by its filtered version.
    datacol : str
        The name of the data column in the input FITS file to be filtered, e.g.
        SAP_FLUX, PDCSAP_FLUX, MOM_CENTR1 etc. A full list of
        archived data columns is provided in the Kepler Archive Manual.
    function : string
        The functional form of the bandpass convolution function.
        The options are:

        * boxcar
        * gauss
        * sinc
    cutoff : float
        The frequency of the bandpass cutoff in units of days-1.
    plot : bool
        Plot the original light curve and the result of the filter?
    overwrite : bool
        Overwrite the output file? if overwrite is **False** and an existing
        file has the same name as outfile then the task will stop with an
        error.
    verbose : bool
        Print informative messages and warnings to the shell and logfile?
    logfile : str
        Name of the logfile containing error and warning messages.

    Examples
    --------

    .. code-block :: bash

        $ kepfilter kplr002436324-2009259160929_llc.fits --datacol 'SAP_FLUX' --function 'boxcar'
        --plot --verbose --overwrite

    .. image :: ../_static/images/api/kepfilter.png
        :align: center
    """
    if outfile is None:
        outfile = infile.split('.')[0] + "-{}.fits".format(__all__[0])
    ## log the call
    hashline = '--------------------------------------------------------------'
    kepmsg.log(logfile, hashline, verbose)
    call = ('KEPFILTER -- '
            + ' infile={}'.format(infile)
            + ' outfile={}'.format(outfile)
            + ' datacol={}'.format(datacol)
            + ' function={}'.format(function)
            + ' cutoff={}'.format(cutoff)
            + ' passband={}'.format(passband)
            + ' plot={}'.format(plot)
            + ' overwrite={}'.format(overwrite)
            + ' verbose={}'.format(verbose)
            + ' logfile={}'.format(logfile))
    kepmsg.log(logfile, call+'\n', verbose)
    ## start time
    kepmsg.clock('KEPFILTER started at',logfile,verbose)
    ## overwrite output file
    if overwrite:
        kepio.overwrite(outfile, logfile, verbose)
    if kepio.fileexists(outfile):
        errmsg = 'ERROR -- KEPFILTER: {} exists. Use --overwrite'.format(outfile)
        kepmsg.err(logfile, message, verbose)

    ## open input file
    instr = pyfits.open(infile, 'readonly')
    tstart, tstop, bjdref, cadence = kepio.timekeys(instr, infile,
                                                    logfile, verbose)
    try:
        work = instr[0].header['FILEVER']
        cadenom = 1.0
    except:
        cadenom = cadence

    ## fudge non-compliant FITS keywords with no values
    instr = kepkey.emptykeys(instr, infile, logfile, verbose)
    ## read table structure
    table = kepio.readfitstab(infile, instr[1], logfile, verbose)
    # read time and flux columns
    barytime = kepio.readtimecol(infile, table, logfile, verbose)
    flux= kepio.readsapcol(infile, table, logfile, verbose)
    # filter input data table
    try:
        nanclean = instr[1].header['NANCLEAN']
    except:
        naxis2 = 0
        for i in range(len(table.field(0))):
            if (np.isfinite(barytime[i]) and np.isfinite(flux[i])
                and flux[i] != 0.0):
                table[naxis2] = table[i]
                naxis2 += 1
        instr[1].data = table[:naxis2]
        kepkey.new('NANCLEAN', True, 'NaN cadences removed from data',
                   instr[1], outfile, logfile, verbose)

    ## read table columns
    intime = (kepio.readtimecol(infile, instr[1].data, logfile, verbose)
              + bjdref)
    indata = kepio.readfitscol(infile, instr[1].data, datacol, logfile,
                               verbose) / cadenom
    ## define data sampling
    tr = 1.0 / (cadence / 86400)
    timescale = 1.0 / (cutoff / tr)
    ## define convolution function
    if function == 'boxcar':
        filtfunc = np.ones(int(np.ceil(timescale)))
    elif function == 'gauss':
        timescale /= 2
        dx = np.ceil(timescale * 10 + 1)
        filtfunc = kepfunc.gauss([1.0, dx / 2 - 1.0, timescale],
                                 np.linspace(0, dx - 1, dx))
    elif function == 'sinc':
        dx = np.ceil(timescale * 12 + 1)
        fx = (np.linspace(0, dx - 1, dx) - dx / 2 + 0.5) / timescale
        filtfunc = np.sinc(fx)

    filtfunc /= np.sum(filtfunc)
    ## pad time series at both ends with noise model
    ave, sigma = (np.mean(indata[:len(filtfunc)]),
                  np.std(indata[:len(filtfunc)]))
    padded = np.append(kepstat.randarray(np.ones(len(filtfunc)) * ave,
                       np.ones(len(filtfunc)) * sigma), indata)
    ave, sigma = (np.mean(indata[-len(filtfunc):]),
                  np.std(indata[-len(filtfunc):]))
    padded = np.append(padded, kepstat.randarray(np.ones(len(filtfunc)) * ave,
                       np.ones(len(filtfunc)) * sigma))
    ## convolve data
    convolved = np.convolve(padded,filtfunc,'same')
    ## remove padding from the output array
    if function == 'boxcar':
        outdata = convolved[len(filtfunc):-len(filtfunc)]
    else:
        outdata = convolved[len(filtfunc):-len(filtfunc)]
    ## subtract low frequencies
    if passband == 'high':
        outmedian = np.median(outdata)
        outdata = indata - outdata + outmedian
    ## comment keyword in output file
    kepkey.history(call, instr[0], outfile, logfile, verbose)
    ## clean up x-axis unit
    intime0 = float(int(tstart / 100) * 100.0)
    if intime0 < 2.4e6: intime0 += 2.4e6
    ptime = intime - intime0
    xlab = 'BJD $-$ {}'.format(intime0)
    ## clean up y-axis units
    pout = indata * 1.0
    pout2 = outdata * 1.0
    nrm = len(str(int(np.nanmax(pout)))) - 1
    pout = pout / 10 ** nrm
    pout2 = pout2 / 10 ** nrm
    ylab = '10$^{}$ {}'.format(nrm, 'e$^-$ s$^{-1}$')
    ## data limits
    xmin = ptime.min()
    xmax = ptime.max()
    ymin = np.nanmin(pout)
    ymax = np.nanmax(pout)
    xr = xmax - xmin
    yr = ymax - ymin
    ptime = np.insert(ptime, [0], [ptime[0]])
    ptime = np.append(ptime, [ptime[-1]])
    pout = np.insert(pout, [0], [0.0])
    pout = np.append(pout, 0.0)
    pout2 = np.insert(pout2, [0], [0.0])
    pout2 = np.append(pout2, 0.0)
    ## plot light curve
    if plot:
        plt.figure()
        plt.clf()

        ## plot filtered data
        ax = plt.axes([0.06, 0.1, 0.93, 0.87])
        plt.gca().xaxis.set_major_formatter(plt.ScalarFormatter(useOffset=False))
        plt.gca().yaxis.set_major_formatter(plt.ScalarFormatter(useOffset=False))
        plt.plot(ptime, pout, color='#ff9900', linestyle='-', linewidth=1.0)
        plt.fill(ptime, pout, color='#ffff00', linewidth=0.0, alpha=0.2)
        if passband == 'low':
            plt.plot(ptime[1:-1], pout2[1:-1], color='#0000ff', linestyle='-',
                     linewidth=1.0)
        else:
            plt.plot(ptime, pout2, color='#0000ff', linestyle='-',
                     linewidth=1.0)
            plt.fill(ptime, pout2, color='#0000ff', linewidth=0.0, alpha=0.2)
        plt.xlabel(xlab, {'color' : 'k'})
        plt.ylabel(ylab, {'color' : 'k'})
        plt.xlim(xmin-xr*0.01,xmax+xr*0.01)
        if ymin >= 0.0:
            plt.ylim(ymin-yr*0.01,ymax+yr*0.01)
        else:
            plt.ylim(1.0e-10,ymax+yr*0.01)
        plt.grid()
        # render plot
        plt.show()
    ## write output file
    print("Writing output file {}...".format(outfile))
    for i in tqdm(range(len(outdata))):
        instr[1].data.field(datacol)[i] = outdata[i]
    instr.writeto(outfile)
    ## close input file
    instr.close()
    ## end time
    kepmsg.clock('KEPFILTER completed at', logfile, verbose)

def kepfilter_main():
    import argparse

    parser = argparse.ArgumentParser(
             description='Low bandpass or high bandpass signal filtering',
             formatter_class=PyKEArgumentHelpFormatter)
    parser.add_argument('infile', help='Name of input file', type=str)
    parser.add_argument('--passband', help='low- or high-bandpass filter',
                        type=str, choices=['low','high'])
    parser.add_argument('--outfile',
                        help=('Name of FITS file to output.'
                              ' If None, outfile is infile-kepfilter.'),
                        default=None)
    parser.add_argument('--datacol', default='SAP_FLUX',
                        help='Name of data column', type=str)
    parser.add_argument('--function', default='boxcar',
                        help='The bandpass convolution function', type=str,
                        choices=['boxcar','gauss','sinc'])
    parser.add_argument('--cutoff', default=1.0,
                        help='Characteristic frequency cutoff of filter [1/days]',
                        type=float)
    parser.add_argument('--plot', action='store_true',
                        help='Plot result?')
    parser.add_argument('--overwrite', action='store_true',
                        help='Overwrite output file?')
    parser.add_argument('--verbose', action='store_true',
                        help='Write to a log file?')
    parser.add_argument('--logfile', help='Name of ascii log file',
                        default='kepfilter.log', type=str)
    args = parser.parse_args()
    kepfilter(args.infile, args.passband, args.outfile, args.datacol,
              args.function, args.cutoff, args.plot, args.overwrite,
              args.verbose, args.logfile)


# -*- coding: utf-8 -*-
"""Class holding a set of TwoDSpectra

    Class Details
    -------------

"""
import numbers

#import h5py
#import matplotlib.pyplot as plt
import numpy

from ..core.time import TimeAxis
from ..core.valueaxis import ValueAxis
from ..core.frequency import FrequencyAxis
from ..core.dfunction import DFunction
#from .twod2 import TwoDResponse
from .twod import TwoDSpectrum

from ..core.managers import Manager, energy_units

#from ..core.managers import energy_units
from .. import COMPLEX
from .. import REAL

from ..core.saveable import Saveable

from .. import part_REAL, part_IMAGINARY, part_COMPLEX, part_ABS
from .. import signal_TOTL #, signal_REPH, signal_NONR

class TwoDResponseContainer(Saveable):
    """Class holding a set of TwoDSpectra

    Parameters
    ----------

    t2axis: TimeAxis
       object holding waiting times at which spectra are calculated

    keep_pathways: bool
       if set True, the container will keep all types of Liouville pathways
       stored separately

    keep_stypes: bool
       if se t True, the container will keep rephasing and non-rephasing
       spectra stored separately

    """

    def __init__(self, t2axis=None, keep_pathways=False, keep_stypes=True):

        self.t2axis = t2axis
        self.keep_pathways = keep_pathways
        self.keep_stypes = keep_stypes

        self.axis = None

        self.itype = None
        self.index = 0
        self.tags = []

        if self.keep_pathways:
            raise Exception("Container keeping pathways not available yet")

        self.spectra = {}
        self._which = None

        if t2axis is not None:
            self.use_indexing_type(itype=t2axis)

    def use_indexing_type(self, itype):
        """Sets the type of indices used to identify spectra

        Parameters
        ----------

        itype : string, ValueAxis, TimeAxis, FrequencyAxis
            Type of indexig. If itype is a string, it should have values
            either 'integer' or 'string' in which case the specra will be
            stored by integer index or by string (as in dictionary). If
            itype is a ValueAxis (TimeAxis, FrequencyAxis), spectra will
            be indexed by the values in the axis object.

        """

        if isinstance(itype, str):
            if itype == "integer":
                self.itype = "integer"
            elif itype == "string":
                self.itype = "string"
            else:
                raise Exception("Unknown indexing type")
        elif isinstance(itype, ValueAxis):
            if isinstance(itype, TimeAxis):
                self.itype = "TimeAxis"
                self.axis = itype
                # This axis must FFT in the "standard" way
                # -- it must of the "complete" type
                self.axis.atype = "complete"
            elif isinstance(itype, FrequencyAxis):
                self.itype = "FrequencyAxis"
                self.axis = itype
            else:
                self.itype = "ValueAxis"
                self.axis = itype
        else:
            raise Exception("Unknown indexing type")

    def set_spectrum(self, spect, tag=None):
        """Stores spectrum with a tag (time, index, etc.)

        Stores the spectrum according to present indexing scheme

        Parameters
        ----------

        spect : TwoDSpectrum
            Object holding the spectrum; when not tag is specified for a
            spectrum which as its t2 time set, the tag is set to t2 time.

        tag : {int, string, ValuesAxis, TimeAxis, FrequencyAxis}
            Tag which will be used for retrieval of the spectrum from
            the container.

        """

        if self.itype == "integer":

            if tag is None:
                self.spectra[self.index] = spect
                self.index += 1
                return self.index
            else:
                if isinstance(tag, numbers.Integral):
                    self.spectra[tag] = spect
                else:
                    raise Exception("The spectrum has to be tagged by an integer")
                return tag

        elif self.itype in ["ValueAxis", "TimeAxis", "FrequencyAxis"]:

            if tag is None:
                # we will read the spectrum intrinsic t2 time and set it as tag
                if spect.t2 >= 0.0:
                    tag = spect.t2

            if tag is not None:
                if tag in self.axis.data:
                    self.spectra[tag] = spect
                    self.tags.append(tag)
                    self.index += 1
                else:
                    raise Exception("Tag not compatible with the ValueAxis")
            else:
                raise Exception("No tag specified, and spectrum"
                                +" does not have t2 time set"
                                +" - cannot store spectrum")
            return self.index

        elif self.itype == "string":
            if tag is not None:
                stag = str(tag)
                self.spectra[stag] = spect
                self.tags.append(stag)
                self.index += 1
            else:
                raise Exception("No tag specified - cannot store spectrum")
            return self.index

        else:

            raise Exception("Unknown type of indexing")

    def _lousy_equal(self, x1, x2, dx, frac=0.25):
        """Equals up to fraction of dx

        This function returns True if x1 is closer to x2 than `frac` of
        a specified interval. In addition it saves the value of x2 to which
        x1 is equal in the attribute _which of the present class.

        """
        if abs(x1-x2) < dx*frac:
            self._which = x2
            return True

        self._which = None
        return False

    def get_spectrum_by_index(self, indx):
        """Returns spectrum by integet index

        The integer index is assigned to all spectra in the order they were
        saved to the container. They can be retrieved in this order

        Parameters
        ----------

        indx : int
            Index of the spectrum to be retrieved

        """

        if self.itype == "integer":

            return self.get_spectrum(indx)

        else:

            return self.spectra[self.tags[indx]]

    def get_response(self, tag):
        """Same as get_spectrum, but the name more sense for a response container

        """
        return self.get_spectrum(tag)

    def get_spectrum(self, tag):
        """Returns spectrum corresponing to time t2

        Checks if the time t2 is present in the t2axis

        Parameters
        ----------

        t2 : float
            Waiting time for which spectrum should be returned

        """
        if self.itype in ["integer"]:

            return self.spectra[tag]

        elif self.itype in ["string"]:

            return self.spectra[str(tag)]

        elif self.itype in ["ValueAxis", "TimeAxis", "FrequencyAxis"]:

            with energy_units("int"):
                if any(self._lousy_equal(tag, li, self.axis.step)
                   for li in self.axis.data):

                    try:
                        return self.spectra[self._which]
                    except KeyError:
                        print(self.spectra)
                        raise Exception()
                else:
                    raise Exception("Tag not compatible with the ValueAxis")

        else:

            raise Exception("Unknown type of indexing")

    def set_data_flag(self, flag):
        """Sets data flag for all spectra in the container

        """

        for tag in self.spectra:

            sp = self.spectra[tag]
            sp.set_data_flag(flag)

    def get_TwoDSpectrumContainer(self, stype=signal_TOTL):
        """Returns a container with specific spectra

        """
        if self.itype in ["ValueAxis", "TimeAxis", "FrequencyAxis"]:
            axis = self.axis.deepcopy()

            cont = TwoDSpectrumContainer(axis)

            for val in self.axis.data:
                sp = self.get_spectrum(val)
                nsp = sp.get_TwoDSpectrum(dtype=stype)
                cont.set_spectrum(nsp, tag=val)

            return cont

        else:

            raise Exception("")

    def get_nearest(self, val):

        if self.itype == "FrequencyAxis":
            #print(Manager().current_units["frequency"])
            #print(Manager().current_units["energy"])
            nval = Manager().convert_energy_2_internal_u(val)
            # get tags and convert them to numbers
            ntags = numpy.zeros(len(self.tags), dtype=REAL)
            k = 0
            for stag in self.tags:
                #print(stag)
                ntag = float(stag)
                ntags[k] = ntag
                k += 1
            dtags = numpy.abs(ntags - nval)
            imin = numpy.argmin(dtags)
            #print("Returning spectrum at: ",
            #      Manager().convert_energy_2_current_u(self.tags[imin]))
            return self.spectra[self.tags[imin]], imin

    def length(self):
        """Returns the length of the container

        """
        return len(self.spectra.keys())

    def get_spectra(self, start=None, end=None):
        """Returns a list of the calculated spectra

        Returns all spectra or an interval of spectra when `start` and `end`
        are specified

        Parameters
        ----------

        start : int
            Index of the first spectrum to be returned

        end : int
            Index of the last spectrum to be returned

        """

        ven = [value for (key, value) in sorted(self.spectra.items())]

        if (start is None) and (end is None):
            return ven
        else:
            ven2 = []
            vkeys = [key for (key, value) in sorted(self.spectra.items())]
            for k in vkeys:
                if k >= start and k <= end:
                    ven2.append(self.spectra[k])
            return ven2

    def get_PumpProbeSpectrumContainer(self, skip=0):
        """Converts this container into PumpProbeSpectrumContainer

        """

        from .pumpprobe import PumpProbeSpectrumContainer

        k = 0
        ppc = []
        ttc = []
        ii = 0
        for sp in self.get_spectra():
            if k == 0:
                pp = sp.get_PumpProbeSpectrum()
                ppc.append(pp)
                ttc.append(self.axis.data[ii])
            k += 1
            if k > skip:
                k = 0
            ii += 1

        length = len(ppc)
        start = ppc[0].get_t2()
        step = ppc[1].get_t2()-start

        naxis = TimeAxis(start,length,step)
        ppcont = PumpProbeSpectrumContainer(t2axis=naxis)
        ppcont.itype = self.itype

        ii = 0
        for sp in ppc:
            tt = ttc[ii]
            ppcont.set_spectrum(sp, tt)
            ii += 1

        return ppcont

    def get_integrated_area_evolution(self, times, area, dpart=part_REAL):
        """Returns the integrated area of the 2D spectra as a function of their index

        """
        vals = numpy.zeros(times.length, dtype=COMPLEX)
        k = 0

        # this only acts on Frequency axis
        with energy_units("int"):
            tms = times.data

        for t2 in tms:

            sp = self.get_spectrum(t2)
            vals[k] = sp.get_area_integral(area, dpart=part_REAL)
            k +=1

        return DFunction(times, vals)

    def get_point_evolution(self, x, y, times):
        """Tracks an evolution of a single point on the 2D spectrum

        Parameters
        ----------

        x : float
            x coordinate in the 2D spectrum (usually omega_1 axis)

        y : float
            y coordinate in the 2D spectrum (usually omega_3 axis)

        times : ValueAxis
            Times (usually waiting t_2 times) in which spectra are taken

        """

        vals = numpy.zeros(times.length, dtype=COMPLEX)
        k = 0

        # this only acts on Frequency axis
        with energy_units("int"):
            tms = times.data

        for t2 in tms:

            sp = self.get_spectrum(t2)
            vals[k] = sp.get_value_at(x, y)
            k +=1

        return DFunction(times, vals)

    def global_fit_exponential(self, guess=None):
        """Global fit of the data with exponentials

        """
        from scipy.optimize import least_squares
        from functools import partial

        if guess is None:
            guess = [1.0, 1.0/100.0, 0.0]

        _exp_2D_fcion = partial(_exp_2D_data0, times=self.axis.data, cont=self)

        params = least_squares(_exp_2D_fcion, guess)

        return params

    def fft(self, ffttype="complex-positive", window=None, offset=0.0,
            dtype=None, dpart=part_COMPLEX, tag=None):
        """Fourier transform in t2 time

        This method performs FFT on the container data determined by the
        value of the `dtype` argument. The new container is created and
        the storage resolution of its components is set `off`. This means
        that the container and its spectra have no idea about what data they
        store. Even when plotting the spectra, one has to set plotting of
        the `total` spectrum.

        Parameters
        ----------

        ffttype : string
            Specifies the type Fourier transform we perform

        window : DFunction
            Windowing function for the data. Default is None

        """

        if dtype is None:
            raise Exception("Type of the data for FFT has to be specified")

        if self.itype not in ["ValueAxis", "TimeAxis", "FrequencyAxis"]:
            raise Exception("FFT cannot be performed for"+
                            " this type of indexing")

        # even when no window function is supplied, we create one with
        # all elements equal to one
        if window is None:
            winfce = DFunction(self.axis,
                               numpy.ones(self.axis.length, dtype=REAL))
        else:
            winfce = window

        if isinstance(self.axis, TimeAxis):
            # restrict the time axis by the off-set
            tlist = []
            for tt in self.axis.data:
                if tt >= offset:
                    tlist.append(tt)
            if len(tlist) > 1:
                dt = tlist[1]-tlist[0]
                Nt = len(tlist)
                t0 = tlist[0]
                # effective time axis for fft
                eff_axis = TimeAxis(t0, Nt, dt, atype="complete")
            else:
                raise Exception("Offset too large")
        else:
            eff_axis = self.axis

        # put all data into one array

        #raise Exception()
        self.set_data_flag(dtype)

        tags = eff_axis.data #self.axis.data
        Nos = self.length()

        if len(tags) <= Nos:
            tag1 = eff_axis.data[0] #self.axis.data[0]
            sp1 = self.get_spectrum(tag1)

            N1, N2 = sp1.d__data.shape
            data = numpy.zeros((N1, N2, len(tags)), dtype=sp1.d__data.dtype)

            for k_n in range(len(tags)):
                tag = eff_axis.data[k_n] #self.axis.data[k_n]

                spect = self.get_spectrum(tag)
                spect.set_data_flag(dtype)
                if dpart == part_COMPLEX:
                    data[:,:,k_n] = spect.d__data
                elif dpart == part_REAL:
                    data[:,:,k_n] = numpy.real(spect.d__data)
                elif dpart == part_IMAGINARY:
                    data[:,:,k_n] = numpy.imag(spect.d__data)
                elif dpart == part_ABS:
                    data[:,:,k_n] = numpy.abs(spect.d__data)

        else:
            raise Exception("Number of spectra not consistent"+
                            " with ValueAxis object")

        #
        # FFT of the axis
        #
        axis = eff_axis # = self.axis

        if isinstance(axis, TimeAxis):
            axis.shift_to_zero()
            new_axis = axis.get_FrequencyAxis()

        elif isinstance(axis, FrequencyAxis):
            new_axis = axis.get_TimeAxis()

        else:
            # this must be ValueAxis

            ftaxdata = (2.0*numpy.pi)*numpy.fft.fftfreq(axis.length,
                                                        axis.step)
            ftaxdata = numpy.fft.fftshift(ftaxdata)

            start = ftaxdata[0]
            length = len(ftaxdata)
            step = ftaxdata[1]-ftaxdata[0]

            new_axis = ValueAxis(start, length, step)

        #
        # FFT of the data
        #

        # window function
        ftdata = numpy.zeros(data.shape, dtype=data.dtype)
        Nwin = len(winfce.data)
        Ndat = data.shape[2]
        for i_n in range(data.shape[0]):
            for j_n in range(data.shape[1]):
                ftdata[i_n,j_n,:] = data[i_n,j_n,:]*winfce.data[Nwin-Ndat:Nwin]

        ftdata = numpy.fft.ifft(ftdata, axis=2)
        ftdata = numpy.fft.fftshift(ftdata, axes=2)

        # save it to a new container
        new_container = TwoDSpectrumContainer()
        new_container.use_indexing_type(new_axis)

        for k_n in range(Ndat):
            tag = new_axis.data[k_n]
            spect = TwoDSpectrum()
            spect.set_axis_1(sp1.xaxis)
            spect.set_axis_3(sp1.yaxis)

            spect.set_data(ftdata[:, :, k_n], dtype=signal_TOTL)

            new_container.set_spectrum(spect, tag=tag)

        return new_container

    def _create_root_group(self, start, name):
        return start.create_group(name)

    def _save_axis(self, rt, name, ax):
        axdir = rt.create_group(name)
        axdir.attrs.create("start",ax.start)
        axdir.attrs.create("length",ax.length)
        axdir.attrs.create("step",ax.step)

    def _load_axis(self, rt, name):
        axdir = rt[name]
        start = axdir.attrs["start"]
        length = axdir.attrs["length"]
        step = axdir.attrs["step"]
        return TimeAxis(start, length, step)

    def trimall_to(self, window=None):
        """Trims all spectra in the container

        Parameters
        ----------

        window: list of floats
            Window, specified by four float number, to which all spectra
            in the container should be trimmed

        """
        if window is not None:
            axes = window
            for s in self.get_spectra():
                s.trim_to(window=axes)

    # FIXME: this is still legacy version
    def amax(self, spart=part_REAL):
        """Returns maximum amplitude of the spectra in the container

        """

        mxs = []
        for s in self.get_spectra():
            #spect2D = numpy.real(s.d__data)
            if spart == part_REAL:
                spect2D = numpy.real(s.data)
            elif spart == part_IMAGINARY:
                spect2D = numpy.imag(s.data)
            elif spart == part_ABS:
                spect2D = numpy.abs(s.data)
            else:
                raise Exception("Unknow part of the spectrum:", spart)
            mx = numpy.amax(spect2D)
            mxs.append(mx)
        return numpy.amax(numpy.array(mxs))

    # Print iterations progress
    def _printProgressBar(self, iteration, total,
                          prefix = '', suffix = '',
                          decimals = 1, length = 100,
                          fill='*'):
        """
        Call in a loop to create terminal progress bar
        @params:
            iteration   - Required  : current iteration (Int)
            total       - Required  : total iterations (Int)
            prefix      - Optional  : prefix string (Str)
            suffix      - Optional  : suffix string (Str)
            decimals    - Optional  : positive number of decimals in percent complete (Int)
            length      - Optional  : character length of bar (Int)
            fill        - Optional  : bar fill character (Str)

        Based on:
        https://stackoverflow.com/questions/3173320/text-progress-bar-in-the-console
        """
#                          fill = '█'):
        percent = ("{0:." + str(decimals) + "f}").format(100 * (iteration / float(total)))
        filledLength = int(length * iteration // total)
        bar = fill * filledLength + '-' * (length - filledLength)
        print('\r%s |%s| %s%% %s' % (prefix, bar, percent, suffix), end = '\r')
        # Print New Line on Complete
        if iteration == total:
            print()

    def make_movie(self, filename, window=None,
                   stype=signal_TOTL, spart=part_REAL,
                   cmap=None,
                   Npos_contours=10,
                   vmax=None, vmin_ratio=0.5,
                   xlabel=None,
                   ylabel=None,
                   axis_label_font=None,
                   start=None, end=None,
                   frate=20, dpi=100,
                   show_states=None,
                   show_states_func=None,
                   label=None,
                   label_func=None,
                   text_loc=None,
                   progressbar=False,
                   use_t2=True,
                   title="Quantarhei movie",
                   comment="Created with Quantarhei"):
        """Creates a movie out of the spectra in the container

        Parameters
        ----------

        Npos_contours : int
            Nomber of positive value contours in the plot

        """

        import matplotlib.pyplot as plt
        import matplotlib.animation as manimation

        FFMpegWriter = manimation.writers["ffmpeg"]

        metadata = dict(title=title, artist='Quantarhei',
                comment=comment)
        writer = FFMpegWriter(fps=frate, metadata=metadata)

        fig = plt.figure()

        spctr = self.get_spectra()
        l = len(spctr)

        if use_t2 and (spctr[0].get_t2() < 0.0):
            #print("Warning: switching off usage of t2"
            #      +" information obtained from the spectrum object (t2 < 0)")
            use_t2 = False

        if use_t2:
            last_t2 = spctr[l-1].get_t2()
            first_t2 = spctr[0].get_t2()

            if start is None:
                start = first_t2
            if end is None:
                end = last_t2

        if vmax is None:
            mx = self.amax(spart=spart)
        else:
            mx = vmax

        with writer.saving(fig, filename, dpi):
            k = 0
            # Initial call to print 0% progress
            if use_t2:
                sp2write = self.get_spectra(start=start, end=end)
            else:
                sp2write = self.get_spectra()
            l = len(sp2write)

            if progressbar:
                self._printProgressBar(0, l, prefix = 'Progress:',
                                       suffix = 'Complete', length = 50)

            for sp in sp2write: #self.get_spectra(start=start, end=end):

                if label_func is not None:
                    (label, text_loc) = label_func(sp)
                if show_states_func is not None:
                    show_states = show_states_func(sp)

                sp.plot(fig=fig, window=window, cmap=cmap, vmax=mx,
                        vmin_ratio=vmin_ratio,
                        Npos_contours=Npos_contours,
                        stype=stype,spart=spart,
                        show_states=show_states,
                        xlabel=xlabel, ylabel=ylabel,
                        axis_label_font=axis_label_font,
                        label=label, text_loc=text_loc) #"T="+str(sp.get_t2())+"fs")
                writer.grab_frame()
                if progressbar:
                    self._printProgressBar(k + 1, l, prefix = 'Progress:',
                                           suffix = 'Complete', length = 50)

                k += 1
#                if k == 20:
#                    return

def _exp_2D_data0(params, times=None, cont=None):
    """Returns a residue between time dependent matrix data and a matrix
       multipled by a sum of exponentials

    """

    if times is None:
        raise Exception("Times have to be supplied")
    if cont is None:
        raise Exception("Spectra container has to be supplied")

    np = len(params)

    data0 = cont.get_spectrum(0.0)

    residues = numpy.zeros((times.shape[0], data0.data.shape[0],
                            data0.data.shape[1]), dtype=data0.dtype)

    N = times.shape[0]*data0.data.shape[0]*data0.data.shape[0]

    nexp = int((np - 1)/2)

    ii = 0
    for t in times:

        ret = 0.0
        kp = 0

        for kk in range(nexp):
            ret += params[kp]*numpy.exp(-params[kp+1]*t)
            kp += 2
        ret += params[kp]

        spect = cont.get_spectrum(t)

        residues[ii,:,:] = numpy.abs(spect.data - ret*data0.data)

        ii += 1

    return residues.reshape(N)

class TwoDSpectrumContainer(TwoDResponseContainer):

    def __init__(self, t2axis=None, dtype=signal_TOTL):

        self.t2axis = t2axis

        self.axis = None

        self.itype = None
        self.index = 0
        self.tags = []
        self.dtype = dtype

        self.spectra = {}
        self._which = None

        if t2axis is not None:
            self.use_indexing_type(itype=t2axis)

    def set_data_flag(self, flag):
        """Sets data flag for all spectra in the container

        """

        if flag != self.dtype:
            raise Exception("Cannot change spectra type")

    def get_TwoDSpectrumContainer(self, stype=signal_TOTL):
        """Returns a container with specific spectra

        """
        if stype == self.dtype:
            return self
        else:
            raise Exception("Cannot change spectra type in this container")

    # FIXME: This needs to be reimplemented
    def get_PumpProbeSpectrumContainer(self, skip=0):
        """Converts this container into PumpProbeSpectrumContainer

        """

        if self.dtype == signal_TOTL:

            from .pumpprobe import PumpProbeSpectrumContainer

            k = 0
            ppc = []
            ttc = []
            ii = 0
            for sp in self.get_spectra():
                if k == 0:
                    pp = sp.get_PumpProbeSpectrum()
                    ppc.append(pp)
                    ttc.append(self.axis.data[ii])
                k += 1
                if k > skip:
                    k = 0
                ii += 1

            length = len(ppc)
            start = ppc[0].get_t2()
            step = ppc[1].get_t2()-start

            naxis = TimeAxis(start,length,step)
            ppcont = PumpProbeSpectrumContainer(t2axis=naxis)

            ii = 0
            for sp in ppc:
                tt = ttc[ii]
                ppcont.set_spectrum(sp, tt)
                ii += 1

            return ppcont

        else:

            raise Exception("Cannot calculate Pump-probe from 2D"+
                            " spectra of type"+self.dtype)

    def normalize2(self, norm=1.0, each=False, dpart=part_REAL):
        """Normalize the whole container of spectra

        Normalization of the whole container so that the maximum
        value of the spectrum is equal to the `norm`.

        Parameters
        ----------

        norm : float
            Value to which we normalize the spectra

        each: bool
            If False, we normalize the global maximum of the container, i.e.
            a maximum accross whole spectra. if True, each spectrum is
            normalized individually against its own maximum.

        dpart: string
            Part of the spectrum from which the maximum is calculated, it
            can be part_REAL, part_IMAGINARY or part_ABS. The values of these
            constants are defined in the highest level of namespace in
            Quantarhei.

        """

        nsp = len(self.spectra)
        mxs = numpy.zeros(nsp, dtype=REAL)
        ii = 0
        for tag in self.spectra.keys():
            sp = self.get_spectrum(tag)
            mxs[ii] = sp.get_max_value(dpart=dpart)
            ii += 1

        mx = numpy.amax(mxs)

        nmax = [mx]
        for tag in self.spectra.keys():
            sp = self.get_spectrum(tag)
            if each:
                nmax = [sp.get_max_value(dpart=dpart)]
            sp.normalize2(norm, dpart=dpart, nmax=nmax)

        return nmax

    def fft(self, ffttype="complex-positive", window=None, offset=0.0,
            dtype=None, dpart=part_COMPLEX, tag=None):
        """Fourier transform in t2 time

        This method performs FFT on the container data determined by the
        value of the `dtype` argument. The new container is created and
        the storage resolution of its components is set `off`. This means
        that the container and its spectra have no idea about what data they
        store. Even when plotting the spectra, one has to set plotting of
        the `total` spectrum.

        Parameters
        ----------

        ffttype : string
            Specifies the type Fourier transform we perform

        window : DFunction
            Windowing function for the data. Default is None

        """

        if dtype is not None:
            if dtype != self.dtype:
                raise Exception("Cannot change spectra type"+
                                " in TwoDSpectrumContainer")

        if self.itype not in ["ValueAxis", "TimeAxis", "FrequencyAxis"]:
            raise Exception("FFT cannot be performed for"+
                            " this type of indexing")

        # even when no window function is supplied, we create one with
        # all elements equal to one
        if window is None:
            winfce = DFunction(self.axis,
                               numpy.ones(self.axis.length, dtype=REAL))
        else:
            winfce = window

        if isinstance(self.axis, TimeAxis):
            # restrict the time axis by the off-set
            tlist = []
            for tt in self.axis.data:
                if tt >= offset:
                    tlist.append(tt)
            if len(tlist) > 1:
                dt = tlist[1]-tlist[0]
                Nt = len(tlist)
                t0 = tlist[0]
                # effective time axis for fft
                eff_axis = TimeAxis(t0, Nt, dt, atype="complete")
            else:
                raise Exception("Offset too large")
        else:
            eff_axis = self.axis

        # put all data into one array

        #raise Exception()

        tags = eff_axis.data #self.axis.data
        Nos = self.length()

        if len(tags) <= Nos:
            tag1 = eff_axis.data[0] #self.axis.data[0]
            sp1 = self.get_spectrum(tag1)

            N1, N2 = sp1.data.shape
            data = numpy.zeros((N1, N2, len(tags)), dtype=sp1.data.dtype)

            for k_n in range(len(tags)):
                tag = eff_axis.data[k_n] #self.axis.data[k_n]

                spect = self.get_spectrum(tag)

                if dpart == part_COMPLEX:
                    data[:,:,k_n] = spect.data
                elif dpart == part_REAL:
                    data[:,:,k_n] = numpy.real(spect.data)
                elif dpart == part_IMAGINARY:
                    data[:,:,k_n] = numpy.imag(spect.data)
                elif dpart == part_ABS:
                    data[:,:,k_n] = numpy.abs(spect.data)

        else:
            raise Exception("Number of spectra not consistent"+
                            " with ValueAxis object")

        #
        # FFT of the axis
        #
        axis = eff_axis # = self.axis

        if isinstance(axis, TimeAxis):
            axis.shift_to_zero()
            new_axis = axis.get_FrequencyAxis()

        elif isinstance(axis, FrequencyAxis):
            new_axis = axis.get_TimeAxis()

        else:
            # this must be ValueAxis

            ftaxdata = (2.0*numpy.pi)*numpy.fft.fftfreq(axis.length,
                                                        axis.step)
            ftaxdata = numpy.fft.fftshift(ftaxdata)

            start = ftaxdata[0]
            length = len(ftaxdata)
            step = ftaxdata[1]-ftaxdata[0]

            new_axis = ValueAxis(start, length, step)

        #
        # FFT of the data
        #

        # window function
        ftdata = numpy.zeros(data.shape, dtype=data.dtype)
        Nwin = len(winfce.data)
        Ndat = data.shape[2]
        for i_n in range(data.shape[0]):
            for j_n in range(data.shape[1]):
                ftdata[i_n,j_n,:] = data[i_n,j_n,:]*winfce.data[Nwin-Ndat:Nwin]

        ftdata = numpy.fft.ifft(ftdata, axis=2)
        ftdata = numpy.fft.fftshift(ftdata, axes=2)

        # save it to a new container
        new_container = TwoDSpectrumContainer()
        new_container.use_indexing_type(new_axis)

        for k_n in range(Ndat):
            tag = new_axis.data[k_n]
            spect = TwoDSpectrum()
            spect.set_axis_1(sp1.xaxis)
            spect.set_axis_3(sp1.yaxis)

            spect.set_data(ftdata[:, :, k_n], dtype=self.dtype)

            new_container.set_spectrum(spect, tag=tag)

        return new_container

    def unitedir(self, dname):

        clist = self.loaddir(dname)

        n = len(clist)

        cont = clist[1]

        for k in range(1, n):

            ctn = clist[k+1]

            for tag in ctn.spectra:
                cont.set_spectrum(ctn.spectra[tag], tag=tag)

        return cont