# Copyright (c) 2018, Henrique Miranda
# All rights reserved.
#
# This file is part of the yambopy project
#
from __future__ import print_function, division
from yambopy import *
from itertools import product
from netCDF4 import Dataset

class YamboExcitonWeight(YamboSaveDB):
    """
    Class to read the excitonic weight writen by ypp
    """

    def __init__(self,filename,save='SAVE',path='.'):
        #read save database

        # ! I have commented line 20 because is was broken. Pls check

        ysave = YamboSaveDB.from_db_file(folder=save,filename='ns.db1')
        self.sym_car  = ysave.sym_car
        self.kpts_car = ysave.car_kpoints #kpts_car
        self.lat = ysave.lat

        #super(YamboSaveDB,self).from_db_file(folder=save,filename='ns.db1')

        #read excitons file
        self.excitons = np.loadtxt(filename)

        self.weights = None

    def write_irr(self,filename="irr.dat"):
        """ write the list of kpoints from the irreducible brillouin zone
        """
        f = open("irr.dat",'w')
        for k in self.kpts_car:
            f.write(("%12.8lf"*3)%(k[1],k[0],k[2])+"\n")
        f.close()

    def write_full(self,filename="full.dat"):
        """ write the list of kpoints in the full brillouin zone
        """
        #generate all the possible points
        kpoints = self.kpts_car
        kpts = []
        for k in self.kpts_car:
            for sym in self.sym_car:
                kpts.append(np.dot(sym,k))

        f = open("full.dat",'w')
        for q in kpts:
            f.write(("%12.8lf "*3)%tuple(q)+"\n")
        f.close()

    def get_data(self):
        qpts, weights, transitions = self.calc_kpts_weights()
        return { "qpts": qpts,
                 "weights": weights,
                 "lattice": self.lat,
                 "reciprocal_lattice": self.rlat,
                 "transitions": transitions
                 }

    def calc_kpts_weights(self,repx=list(range(-1,2)),repy=list(range(-1,2)),repz=list(range(-1,2))):
        """ Calculate the weights and kpoints of the excitons
        """
        self.weights     = dict()

        #first run set everything to zero
        for line in self.excitons:
            v,c,k,sym,w,e = line
            self.weights[(int(k),int(sym))] = 0

        #add weights
        for line in self.excitons:
            v,c,k,sym,w,e = line
            self.weights[(int(k),int(sym))] += w

        #rename symmetries and kpoints
        sym = self.sym_car
        kpoints = self.kpts_car

        qpts     = []
        kidx     = []
        weights  = []

        for r in product(repx,repy,repz):
          for k,s in list(self.weights.keys()):
            w   = self.weights[(k,s)]
            weights.append( w )
            qpt = np.dot(sym[s-1],kpoints[k-1])+red_car([r],self.rlat)[0]
            qpts.append( qpt )
            kidx.append( k )

        return np.array(qpts), np.array(weights)

    def calc_kpts_transitions(self,repx=list(range(-1,2)),repy=list(range(-1,2)),repz=list(range(-1,2)),debug=False):
        """ Calculate the transitions and kpoints of the excitons
        """
        self.weights     = dict()
        self.transitions = dict()
        self.transitions_v_to_c = dict()

        #first run set everything to zero
        for line in self.excitons:
            v,c,k,sym,w,e = line
            self.weights[(int(k),int(sym))] = 0
            self.transitions[(int(v),int(c),int(k),int(sym))] = 0
            self.transitions_v_to_c[(int(v),int(c))] = 0

        #add weights
        for line in self.excitons:
            v,c,k,sym,w,e = line
            self.weights[(int(k),int(sym))] += w

        #add transitions
        for line in self.excitons:
            v,c,k,sym,w,e = line
            self.transitions[(int(v),int(c),int(k),int(sym))] += w

        #add percentage of a given v => c transition
        norm = sum(self.excitons[:,4])
        for v,c,k,s in list(self.transitions.keys()):
          self.transitions_v_to_c[(int(v),int(c))] += self.transitions[(v,c,k,s)]
        if debug: print('transitions (valence > condution):')
        for v,c in self.transitions_v_to_c:
          self.transitions_v_to_c[(v,c)] = self.transitions_v_to_c[(v,c)]/norm
          if debug: print('%3d > %3d'%(v,c))

        #rename symmetries and kpoints
        sym = self.sym_car
        kpoints = self.kpts_car

        qpts     = []
        kidx     = []
        t_v_c    = []

        for r in product(repx,repy,repz):
          for k,s in list(self.weights.keys()):
            qpt = np.dot(sym[s-1],kpoints[k-1])+red_car([r],self.rlat)[0]
            qpts.append( qpt )
            kidx.append( k )
            #print (v_ref,c_ref,k,s)
            #aux.append(self.transitions[(v_ref,c_ref,k,s)])

        for v_ref,c_ref in list(self.transitions_v_to_c.keys()):
          aux = []
          for r in product(repx,repy,repz):
            for k,s in list(self.weights.keys()):
              aux.append(self.transitions[(v_ref,c_ref,k,s)])
          t_v_c.append(np.array(aux))

        return np.array(qpts), t_v_c, np.array(kidx)

    def plot_contour(self,resX=500,resY=500):
        """ plot a contour
            idea taken from http://stackoverflow.com/questions/18764814/make-contour-of-scatter
        """
        kpts, z = self.calc_kpts_weights()
        x,y = kpts[:,0],kpts[:,1]
        xi = np.linspace(min(x), max(x), resX)
        yi = np.linspace(min(y), max(y), resY)
        Z = griddata(x, y, z, xi, yi, interp='cubic')
        X, Y = np.meshgrid(xi, yi)

        plt.contourf(X, Y, Z, cmap='gist_heat_r')
        plt.show()

    def plot_weights(self,ax,size=20,marker='H',set_origin=0.0,lim=0.2,cmap='viridis',log_scale=False,set_maximum=1.0):
        """

        Plot the weights in a scatter plot of this exciton (1st version tuned by A. Molina-Sanchez)
        Options:
        cmap : colormap. Default viridis
        log_scale : Logarithmic scale for the intensity (True or False)
        set_maximum : Only applied for linear scale. Apply a cut for a selected intensity (values between 0 and 1)
        Further development: Option for the colorbar

        """
        from numpy import sqrt
        import matplotlib.pyplot as plt
        import matplotlib.colors as colors

        """
        These options can be decided by the user. In this first version we just:
        remove axis
        """
        ax.set_aspect('equal')
        ax.set_xlim(-lim,lim)
        ax.set_ylim(-lim,lim)
        ax.axes.get_xaxis().set_visible(False)
        ax.axes.get_yaxis().set_visible(False)

        kpts, weights = self.calc_kpts_weights()

        if log_scale == True:
           norm = colors.LogNorm(vmin=weights.min(),vmax=weights.max())
        else:
           if abs(set_maximum)>1.:
              set_maximum = 1.
           norm = colors.Normalize(vmin=weights.min(),vmax=abs(set_maximum)*weights.max())

        cmap = plt.get_cmap(cmap)

        ax.scatter(kpts[:,0]-set_origin, kpts[:,1]-set_origin, s=size, marker=marker, color=cmap(norm(weights)))

    def __str__(self):
        s = ""
        s += "reciprocal lattice:\n"
        s += "\n".join([("%12.8lf "*3)%tuple(r) for r in self.rlat])+"\n"
        s += "lattice:\n"
        s += "\n".join([("%12.8lf "*3)%tuple(r) for r in self.lat])+"\n"
        s += "alat:\n"
        s += ("%12.8lf "*3)%tuple(self.alat)+"\n"
        return s

    def plot_transitions(self,size=30,lim=0.2):
        """
        Plot the weight of a given transition in a scatter plot of this exciton.
        My idea is to associate for each transition a color and to plot in a different plot
        """

        from numpy import sqrt
        cmap = plt.get_cmap("gist_heat_r")

        fig = plt.figure(figsize=(10,10))
        kpts, t_v_c, _ = self.calc_kpts_transitions()
        for individual in t_v_c:
          plt.scatter(kpts[:,0], kpts[:,1], s=size, marker='H', color=[cmap(sqrt(c)) for c in individual])

        plt.xlim([-lim,lim])
        plt.ylim([-lim,lim])
        ax = plt.axes()
        ax.set_aspect('equal')
        plt.show()

    def plot_exciton_bs(self,ax,path,nbands='all',space='transition',color='#1f77b4'):
        """
        Plot the excitonic weights of a given transition in the band-structure
        """
        kpts, t_v_c, kidx = self.calc_kpts_transitions(repx=list(range(1)),repy=list(range(1)),repz=list(range(1)))
        t_v_c = np.array(t_v_c)

        #get_path is provided by savedb
        bands_kpoints, bands_indexes, path_car = self.get_path(path,kpts=kpts)

        #calculate distances
        bands_distances = [0]
        distance = 0
        for nk in range(1,len(bands_kpoints)):
            distance += np.linalg.norm(bands_kpoints[nk-1]-bands_kpoints[nk])
            bands_distances.append(distance)

        #get energies at these k-points
        eig = self.eigenvalues[kidx-1]
        eig = eig[bands_indexes]
        transition_weight  = t_v_c[:,bands_indexes]
        if nbands == 'all': nbands = self.nbands
        for tw,t in zip(transition_weight,list(self.transitions_v_to_c.keys())):
            v,c = t
            if space == 'transition':
                ax.plot(bands_distances,eig[:,c-1]-eig[:,v-1])
                ax.scatter(bands_distances,eig[:,c-1]-eig[:,v-1],s=tw*1e4)
            else:
                ax.plot(bands_distances,eig[:,c-1],c=color)
                ax.plot(bands_distances,eig[:,v-1],c=color)
                ax.scatter(bands_distances,eig[:,c-1],s=tw*1e4,c=color)
                ax.scatter(bands_distances,eig[:,v-1],s=tw*1e4,c=color)

    def __str__(self):
        s = ""
        s += "reciprocal lattice:\n"
        s += "\n".join([("%12.8lf "*3)%tuple(r) for r in self.rlat])+"\n"
        s += "lattice:\n"
        s += "\n".join([("%12.8lf "*3)%tuple(r) for r in self.lat])+"\n"
        s += "alat:\n"
        s += ("%12.8lf "*3)%tuple(self.alat)+"\n"
        return s

if __name__ == "__main__":
    ye = YamboExciton('o-yambo.exc_weights_at_1_02')
    print(ye)
    ye.write_irr()
    ye.write_full()
    #ye.plot_contour()
    ye.plot_weights()


# coding: utf-8
import copy

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim

from flearn.client.datasets.get_data import DictDataset, get_dataloader
from flearn.common.distiller import KDLoss
from flearn.common.strategy import ParentStrategy
from flearn.common.strategy.df import DF, DFDistiller
from flearn.common.trainer import Trainer

class MyDFDistiller(DFDistiller):
    def multi(
        self, teacher_lst, student, method="avg_logits", weight_lst=None, **kwargs
    ):
        self._init_kd(teacher_lst, student, **kwargs)
        if weight_lst == None:
            self.weight_lst = [1 / len(teacher_lst)] * len(teacher_lst)
        else:
            self.weight_lst = weight_lst

        for _ in range(self.epoch):
            for _, (x, _) in enumerate(self.kd_loader):
                x = x.to(self.device)
                _, _, output = self.student(x)

                soft_target_lst = []
                for teacher in self.teacher_lst:
                    with torch.no_grad():
                        _, _, soft_target = teacher(x)
                        soft_target_lst.append(soft_target)

                loss = self.multi_loss(method, soft_target_lst, output)

                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()
                # print("train_loss_fine_tuning", loss.data)
        return student.state_dict()

class MyDF(DF):
    def server_post_processing(self, ensemble_params_lst, ensemble_params, **kwargs):
        w_glob = ensemble_params["w_glob"]
        agg_weight_lst, w_local_lst = self.server_pre_processing(ensemble_params_lst)

        teacher_lst = []
        for w_local in w_local_lst:
            self.model_base.load_state_dict(w_local)
            teacher_lst.append(copy.deepcopy(self.model_base))

        self.model_base.load_state_dict(w_glob)
        student = copy.deepcopy(self.model_base)

        kd_loader, device = kwargs.pop("kd_loader"), kwargs.pop("device")
        temperature = kwargs.pop("T")
        distiller = MyDFDistiller(
            kd_loader,
            device,
            kd_loss=KDLoss(temperature),
        )

        molecular = np.sum(agg_weight_lst)
        weight_lst = [w / molecular for w in agg_weight_lst]
        # agg_weight_lst：应该依照每个模型在验证集上的性能来进行分配
        ensemble_params["w_glob"] = distiller.multi(
            teacher_lst, student, kwargs.pop("method"), weight_lst=weight_lst, **kwargs
        )
        return ensemble_params

class CCVR(ParentStrategy):
    """
    Federated learning via Classifier Calibration with Virtual Representations

    [1] Luo M, Chen F, Hu D, et al. No Fear of Heterogeneity: Classifier Calibration for Federated Learning with Non-IID Data[J]. arXiv preprint arXiv:2106.05001, 2021.
    """

    def __init__(self, head_model_base, strategy):
        super().__init__(strategy)
        self.head_model_base = head_model_base

    @staticmethod
    def client_mean_feat(feat_lst, label_lst):
        sum_ = 0
        # 按照类别提取特征
        d = {}
        for h_l, label_l in zip(feat_lst, label_lst):
            sum_ += len(h_l)
            for h, label in zip(h_l, label_l):
                label = int(label.cpu())
                if label not in d.keys():
                    d[label] = [h]
                else:
                    d[label].append(h)
        # label_len = len(d.keys())

        # 计算mu, sigma
        upload_d = {}
        for k, v in d.items():
            v_item = torch.stack(v).detach().cpu()
            # 考虑样本数量过少不上传的情况
            if len(v_item) < 10:
                continue
            # if len(v_item) * label_len * 2 < sum_:
            #     continue
            mu, sigma = v_item.mean(dim=0), v_item.var(dim=0)
            upload_d[k] = {"mu": mu, "sigma": sigma, "N": len(v)}
        return upload_d

    def client(self, trainer, agg_weight=1.0):
        w_shared = super().client(trainer, agg_weight)
        w_shared["fd"] = self.client_mean_feat(trainer.feat_lst, trainer.label_lst)
        return w_shared

    @staticmethod
    def load_model(model_base, new_model_dict):
        model_base_dict = model_base.state_dict()
        model_base_dict.update(new_model_dict)
        model_base.load_state_dict(model_base_dict)
        return model_base

    def server_mean_feat(self, fd_lst):
        # 获取每个客户端的标签，集成一起
        label_lst = []
        for x in fd_lst:
            label_lst += list(x.keys())
        label_lst = list(set(label_lst))
        print("labels: ", label_lst)

        fd_d = {}
        # 统计每个标签的特征分布
        for l in label_lst:
            # 客户端不一定具备所有的标签，异质
            labeled_fd_lst = [x for x in fd_lst if l in x.keys()]
            sum_n = sum(x[l]["N"] for x in labeled_fd_lst)

            mu_lst = [fd[l]["mu"] * fd[l]["N"] / sum_n for fd in labeled_fd_lst]
            mu = torch.stack(mu_lst).sum(dim=0)

            sigma1 = torch.stack(
                [fd[l]["mu"] * (fd[l]["N"] - 1) / (sum_n - 1) for fd in labeled_fd_lst]
            ).sum(dim=0)
            sigma2 = torch.stack(
                [
                    fd[l]["mu"] * fd[l]["mu"].T * fd[l]["N"] / (sum_n - 1)
                    for fd in labeled_fd_lst
                ]
            ).sum(dim=0)

            sigma = sigma1 + sigma2 - sum_n / (sum_n - 1) * mu * mu.T

            # 生成batchsize为200的数据样本，总共有10类，所以就是2k个样本
            dist_c = np.random.normal(mu, sigma, size=(200, mu.size()[0]))
            fd_d[l] = torch.tensor(dist_c)

        return fd_d

    def server_post_processing(self, ensemble_params_lst, ensemble_params, **kwargs):
        # 特征参数提取
        fd_lst = self.extract_lst(ensemble_params_lst, "fd")
        fd_d = self.server_mean_feat(fd_lst)

        # 准备好数据集，模型、优化器等等
        trainset = DictDataset(fd_d)
        trainloader, _ = get_dataloader(trainset, trainset, batch_size=64)
        optimizer = optim.SGD(
            self.head_model_base.parameters(), lr=1e-2, momentum=0.9, weight_decay=0.05
        )
        criterion = nn.CrossEntropyLoss()
        self.glob_model_base = self.load_model(
            self.head_model_base, ensemble_params["w_glob"]
        )

        # 重新训练分类器
        trainer = Trainer(
            self.head_model_base, optimizer, criterion, kwargs["device"], False
        )
        trainer.train(trainloader, epochs=1)
        w_train = trainer.weight

        for k in w_train.keys():
            ensemble_params["w_glob"][k] = w_train[k].cpu()

        return ensemble_params

    def server(self, ensemble_params_lst, round_, **kwargs):
        ensemble_params = super().server(ensemble_params_lst, round_)
        return self.server_post_processing(
            ensemble_params_lst, ensemble_params, **kwargs
        )

class DFCCVR(CCVR):
    def __init__(self, model_base, head_model_base, strategy):
        super().__init__(head_model_base, strategy)
        self.model_base = model_base
        self.df = MyDF(model_base, strategy)

    def server(self, ensemble_params_lst, round_, **kwargs):
        # 先DF后CCVR
        ensemble_params = self.df.server(ensemble_params_lst, round_, **kwargs)
        return self.server_post_processing(
            ensemble_params_lst, ensemble_params, **kwargs
        )


# -*- coding: utf-8 -*-
"""
References
----------
-   :cite:`Meng2015c` : Meng, J., Simon, F., Hanika, J., & Dachsbacher, C.
    (2015). Physically Meaningful Rendering using Tristimulus Colours. Computer
    Graphics Forum, 34(4), 31-40. doi:10.1111/cgf.12676
-   :cite:`Smits1999a` : Smits, B. (1999). An RGB-to-Spectrum Conversion for
    Reflectances. Journal of Graphics Tools, 4(4), 11-22.
    doi:10.1080/10867651.1999.10487511
"""

from __future__ import absolute_import

from colour.utilities import (CaseInsensitiveMapping, as_float_array,
                              filter_kwargs)

from .dataset import *  # noqa
from . import dataset
from .meng2015 import XYZ_to_sd_Meng2015
from .smits1999 import RGB_to_sd_Smits1999

__all__ = []
__all__ += dataset.__all__
__all__ += ['XYZ_to_sd_Meng2015']
__all__ += ['RGB_to_sd_Smits1999']

XYZ_TO_SD_METHODS = CaseInsensitiveMapping({
    'Meng 2015': XYZ_to_sd_Meng2015,
    'Smits 1999': RGB_to_sd_Smits1999,
})
XYZ_TO_SD_METHODS.__doc__ = """
Supported spectral distribution recovery methods.

References
----------
:cite:`Meng2015c`, :cite:`Smits1999a`

XYZ_TO_SD_METHODS : CaseInsensitiveMapping
    **{'Meng 2015', 'Smits 1999'}**
"""

def XYZ_to_sd(XYZ, method='Meng 2015', **kwargs):
    """
    Recovers the spectral distribution of given *CIE XYZ* tristimulus
    values using given method.

    Parameters
    ----------
    XYZ : array_like
        *CIE XYZ* tristimulus values to recover the spectral distribution
        from.
    method : unicode, optional
        **{'Meng 2015', 'Smits 1999'}**,
        Computation method.

    Other Parameters
    ----------------
    cmfs : XYZ_ColourMatchingFunctions
        {:func:`colour.recovery.XYZ_to_sd_Meng2015`},
        Standard observer colour matching functions.
    interval : numeric, optional
        {:func:`colour.recovery.XYZ_to_sd_Meng2015`},
        Wavelength :math:`\\lambda_{i}` range interval in nm. The smaller
        ``interval`` is, the longer the computations will be.
    optimisation_parameters : dict_like, optional
        {:func:`colour.recovery.XYZ_to_sd_Meng2015`},
        Parameters for :func:`scipy.optimize.minimize` definition.

    Returns
    -------
    SpectralDistribution
        Recovered spectral distribution.

    Notes
    -----

    +------------+-----------------------+---------------+
    | **Domain** | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``XYZ``    | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    -   *Smits (1999)* method will internally convert given *CIE XYZ*
        tristimulus values to *RGB* colourspace array assuming equal energy
        illuminant *E*.

    References
    ----------
    :cite:`Meng2015c`, :cite:`Smits1999a`

    Examples
    --------

    *Meng (2015)* reflectance recovery:

    >>> import numpy as np
    >>> from colour.utilities import numpy_print_options
    >>> from colour.colorimetry import (
    ...     STANDARD_OBSERVERS_CMFS, SpectralShape, sd_to_XYZ_integration)
    >>> XYZ = np.array([0.21781186, 0.12541048, 0.04697113])
    >>> cmfs = (
    ...     STANDARD_OBSERVERS_CMFS['CIE 1931 2 Degree Standard Observer'].
    ...     copy().align(SpectralShape(360, 780, 10))
    ... )
    >>> sd = XYZ_to_sd(XYZ, cmfs=cmfs)
    >>> with numpy_print_options(suppress=True):
    ...     # Doctests skip for Python 2.x compatibility.
    ...     sd  # doctest: +SKIP
    SpectralDistribution([[ 360.        ,    0.0780114...],
                          [ 370.        ,    0.0780316...],
                          [ 380.        ,    0.0780471...],
                          [ 390.        ,    0.0780351...],
                          [ 400.        ,    0.0779702...],
                          [ 410.        ,    0.0778033...],
                          [ 420.        ,    0.0770958...],
                          [ 430.        ,    0.0748008...],
                          [ 440.        ,    0.0693230...],
                          [ 450.        ,    0.0601136...],
                          [ 460.        ,    0.0477407...],
                          [ 470.        ,    0.0334964...],
                          [ 480.        ,    0.0193352...],
                          [ 490.        ,    0.0074858...],
                          [ 500.        ,    0.0001225...],
                          [ 510.        ,    0.       ...],
                          [ 520.        ,    0.       ...],
                          [ 530.        ,    0.       ...],
                          [ 540.        ,    0.0124896...],
                          [ 550.        ,    0.0389831...],
                          [ 560.        ,    0.0775105...],
                          [ 570.        ,    0.1247947...],
                          [ 580.        ,    0.1765339...],
                          [ 590.        ,    0.2281918...],
                          [ 600.        ,    0.2751347...],
                          [ 610.        ,    0.3140115...],
                          [ 620.        ,    0.3433561...],
                          [ 630.        ,    0.3635777...],
                          [ 640.        ,    0.3765428...],
                          [ 650.        ,    0.3841726...],
                          [ 660.        ,    0.3883633...],
                          [ 670.        ,    0.3905415...],
                          [ 680.        ,    0.3916742...],
                          [ 690.        ,    0.3922554...],
                          [ 700.        ,    0.3925427...],
                          [ 710.        ,    0.3926783...],
                          [ 720.        ,    0.3927330...],
                          [ 730.        ,    0.3927586...],
                          [ 740.        ,    0.3927548...],
                          [ 750.        ,    0.3927681...],
                          [ 760.        ,    0.3927813...],
                          [ 770.        ,    0.3927840...],
                          [ 780.        ,    0.3927536...]],
                         interpolator=SpragueInterpolator,
                         interpolator_args={},
                         extrapolator=Extrapolator,
                         extrapolator_args={...})
    >>> sd_to_XYZ_integration(sd) / 100  # doctest: +ELLIPSIS
    array([ 0.2178545...,  0.1254141...,  0.0470095...])

    *Smits (1999)* reflectance recovery:

    >>> sd = XYZ_to_sd(XYZ, method='Smits 1999')
    >>> with numpy_print_options(suppress=True):
    ...     sd  # doctest: +ELLIPSIS
    SpectralDistribution([[ 380.        ,    0.07691923],
                          [ 417.7778    ,    0.0587005 ],
                          [ 455.5556    ,    0.03943195],
                          [ 493.3333    ,    0.03024978],
                          [ 531.1111    ,    0.02750692],
                          [ 568.8889    ,    0.02808645],
                          [ 606.6667    ,    0.34298985],
                          [ 644.4444    ,    0.41185795],
                          [ 682.2222    ,    0.41185795],
                          [ 720.        ,    0.41180754]],
                         interpolator=LinearInterpolator,
                         interpolator_args={},
                         extrapolator=Extrapolator,
                         extrapolator_args={...})
    >>> sd_to_XYZ_integration(sd) / 100  # doctest: +ELLIPSIS
    array([ 0.2004523...,  0.1105627...,  0.0420964...])
    """

    a = as_float_array(XYZ)

    function = XYZ_TO_SD_METHODS[method]

    if function is RGB_to_sd_Smits1999:
        from colour.recovery.smits1999 import XYZ_to_RGB_Smits1999

        a = XYZ_to_RGB_Smits1999(XYZ)

    return function(a, **filter_kwargs(function, **kwargs))

__all__ += ['XYZ_TO_SD_METHODS', 'XYZ_to_sd']


import torch
import numpy as np

def torch_nms(tlbr, scores, classes=None, thresh=.5, bias=0, fast=False):
    """
    Non maximum suppression implemented with pytorch tensors

    CURRENTLY NOT WORKING

    Args:
        tlbr (Tensor): Bounding boxes of one image in the format (tlbr)
        scores (Tensor): Scores of each box
        classes (Tensor, optional): the classes of each box. If specified nms is applied to each class separately.
        thresh (float): iou threshold

    Returns:
        ByteTensor: keep: boolean array indicating which boxes were not pruned.

    Example:
        >>> # DISABLE_DOCTEST
        >>> import torch
        >>> import numpy as np
        >>> tlbr = torch.FloatTensor(np.array([
        >>>     [0, 0, 100, 100],
        >>>     [100, 100, 10, 10],
        >>>     [10, 10, 100, 100],
        >>>     [50, 50, 100, 100],
        >>>     [100, 100, 130, 130],
        >>>     [100, 100, 130, 130],
        >>>     [100, 100, 130, 130],
        >>> ], dtype=np.float32))
        >>> scores = torch.FloatTensor(np.array([.1, .5, .9, .1, .3, .5, .4]))
        >>> classes = torch.FloatTensor(np.array([0, 0, 0, 0, 0, 0]))
        >>> thresh = .5
        >>> keep = torch_nms(tlbr, scores, classes, thresh)
        >>> bboxes[keep]

    Example:
        >>> # DISABLE_DOCTEST
        >>> import torch
        >>> import numpy as np
        >>> # Test to check that conflicts are correctly resolved
        >>> tlbr = torch.FloatTensor(np.array([
        >>>     [100, 100, 150, 101],
        >>>     [120, 100, 180, 101],
        >>>     [150, 100, 200, 101],
        >>> ], dtype=np.float32))
        >>> scores = torch.FloatTensor(np.linspace(.8, .9, len(tlbr)))
        >>> classes = None
        >>> thresh = .3
        >>> keep = torch_nms(tlbr, scores, classes, thresh, fast=False)
        >>> bboxes[keep]
    """
    if tlbr.numel() == 0:
        return []

    # Sort coordinates by descending score
    ordered_scores, order = scores.sort(0, descending=True)

    from netharn import util
    boxes = util.Boxes(tlbr[order], 'tlbr')
    ious = boxes.ious(boxes, bias=bias)

    # if False:
    #     x1, y1, x2, y2 = tlbr[order].split(1, 1)

    #     # Compute dx and dy between each pair of boxes (these mat contain every pair twice...)
    #     dx = (x2.min(x2.t()) - x1.max(x1.t())).clamp_(min=0)
    #     dy = (y2.min(y2.t()) - y1.max(y1.t())).clamp_(min=0)

    #     # Compute iou
    #     intersections = dx * dy
    #     areas = (x2 - x1) * (y2 - y1)
    #     unions = (areas + areas.t()) - intersections
    #     ious = intersections / unions

    # Filter based on iou (and class)
    conflicting = (ious > thresh).triu(1)

    if classes is not None:
        ordered_classes = classes[order]
        same_class = (ordered_classes.unsqueeze(0) == ordered_classes.unsqueeze(1))
        conflicting = (conflicting & same_class)
    # Now we have a 2D matrix where conflicting[i, j] indicates if box[i]
    # conflicts with box[j]. For each box[i] we want to only keep the first
    # one that does not conflict with any other box[j].

    # Find out how many conflicts each ordered box has with other boxes that
    # have higher scores than it does. In other words...
    # n_conflicts[i] is the number of conflicts box[i] has with other boxes
    # that have a **higher score** than box[i] does. We will definately
    # keep any box where n_conflicts is 0, but we need to postprocess because
    # we might actually keep some boxes currently marked as conflicted.
    n_conflicts = conflicting.sum(0).byte()

    if not fast:
        # It is not enought to simply use all places where there are no
        # conflicts. Say we have boxes A, B, and C, where A conflicts with B,
        # B conflicts with C but A does not conflict with C. The fact that we
        # use A should mean that C is not longer conflicted.

        if True:
            # Marginally faster. best=618.2 us
            ordered_keep = np.zeros(len(conflicting), dtype=np.uint8)
            supress = np.zeros(len(conflicting), dtype=np.bool)
            for i, row in enumerate(conflicting.cpu().numpy() > 0):
                if not supress[i]:
                    ordered_keep[i] = 1
                    supress[row] = 1
            ordered_keep = torch.ByteTensor(ordered_keep).to(tlbr.device)
        else:
            # Marginally slower: best=1.382 ms,
            n_conflicts_post = n_conflicts.cpu()
            conflicting = conflicting.cpu()

            keep_len = len(n_conflicts_post) - 1
            for i in range(1, keep_len):
                if n_conflicts_post[i] > 0:
                    n_conflicts_post -= conflicting[i]

            n_conflicts = n_conflicts_post.to(n_conflicts.device)
            ordered_keep = (n_conflicts == 0)
    else:
        # Now we can simply keep any box that has no conflicts.
        ordered_keep = (n_conflicts == 0)

    # Unsort, so keep is aligned with input boxes
    keep = ordered_keep.new(*ordered_keep.size())
    keep.scatter_(0, order, ordered_keep)
    return keep

def test_class_torch():
    import numpy as np
    import torch
    import netharn as nh
    import ubelt as ub
    # from netharn.util.nms.torch_nms import torch_nms
    # from netharn.util import non_max_supression

    thresh = .5

    num = 500
    rng = nh.util.ensure_rng(0)
    cpu_boxes = nh.util.Boxes.random(num, scale=400.0, rng=rng, format='tlbr', tensor=True)
    cpu_tlbr = cpu_boxes.to_tlbr().data
    # cpu_scores = torch.Tensor(rng.rand(len(cpu_tlbr)))
    # make all scores unique to ensure comparability
    cpu_scores = torch.Tensor(np.linspace(0, 1, len(cpu_tlbr)))
    cpu_cls = torch.LongTensor(rng.randint(0, 10, len(cpu_tlbr)))

    tlbr = cpu_boxes.to_tlbr().data.to('cuda')
    scores = cpu_scores.to('cuda')
    classes = cpu_cls.to('cuda')

    keep1 = []
    for idxs in ub.group_items(range(len(classes)), classes.cpu().numpy()).values():
        # cls_tlbr = tlbr.take(idxs, axis=0)
        # cls_scores = scores.take(idxs, axis=0)
        cls_tlbr = tlbr[idxs]
        cls_scores = scores[idxs]
        cls_keep = torch_nms(cls_tlbr, cls_scores, thresh=thresh, bias=0)
        keep1.extend(list(ub.compress(idxs, cls_keep.cpu().numpy())))
    keep1 = sorted(keep1)

    keep_ = torch_nms(tlbr, scores, classes=classes, thresh=thresh, bias=0)
    keep2 = np.where(keep_.cpu().numpy())[0].tolist()

    keep3 = nh.util.non_max_supression(tlbr.cpu().numpy(),
                                       scores.cpu().numpy(),
                                       classes=classes.cpu().numpy(),
                                       thresh=thresh, bias=0, impl='gpu')

    print(len(keep1))
    print(len(keep2))
    print(len(keep3))

    print(set(keep1) - set(keep2))
    print(set(keep2) - set(keep1))

def _benchmark():
    """
    python -m netharn.util.nms.torch_nms _benchmark --show

    SeeAlso:
        PJR Darknet NonMax supression
        https://github.com/pjreddie/darknet/blob/master/src/box.c

        Lightnet NMS
        https://gitlab.com/EAVISE/lightnet/blob/master/lightnet/data/transform/_postprocess.py#L116

    """
    import torch
    import numpy as np
    import netharn as nh
    from netharn.util.nms.torch_nms import torch_nms
    from netharn.util import non_max_supression
    import ubelt as ub
    import itertools as it

    N = 100
    bestof = 10

    ydata = ub.ddict(list)
    # xdata = [10, 20, 40, 80, 100, 200, 300, 400, 500, 600, 700, 1000, 1500, 2000]

    # max number of boxes yolo will spit out at a time
    max_boxes = 19 * 19 * 5

    xdata = [10, 20, 40, 80, 100, 200, 300, 400, 500, 600, 700, 1000, 1500, max_boxes]
    # xdata = [10, 20, 40, 80, 100, 200, 300, 400, 500]
    xdata = [10, 100, 500]

    rng = nh.util.ensure_rng(0)

    thresh = 0.5

    for num in xdata:
        print('\n\n---- number of boxes = {} ----\n'.format(num))

        outputs = {}

        # Build random test boxes and scores
        cpu_boxes = nh.util.Boxes.random(num, scale=10.0, rng=rng, format='tlbr', tensor=True)
        cpu_tlbr = cpu_boxes.to_tlbr().data
        # cpu_scores = torch.Tensor(rng.rand(len(cpu_tlbr)))
        # make all scores unique to ensure comparability
        cpu_scores = torch.Tensor(np.linspace(0, 1, len(cpu_tlbr)))
        cpu_cls = torch.LongTensor(rng.randint(0, 10, len(cpu_tlbr)))

        # Format boxes in lightnet format
        cpu_ln_boxes = torch.cat([cpu_boxes.to_cxywh().data, cpu_scores[:, None], cpu_cls.float()[:, None]], dim=-1)

        # Move boxes to numpy
        np_tlbr = cpu_tlbr.numpy()
        np_scores = cpu_scores.numpy()
        np_cls = cpu_cls.numpy()  # NOQA

        gpu = torch.device('cuda', 0)

        measure_gpu = torch.cuda.is_available()
        measure_cpu = False or not torch.cuda.is_available()

        def _ln_output_to_keep(ln_output, ln_boxes):
            keep = []
            for row in ln_output:
                # Find the index that we kept
                idxs = np.where(np.all(np.isclose(ln_boxes, row), axis=1))[0]
                assert len(idxs) == 1
                keep.append(idxs[0])
            assert np.all(np.isclose(ln_boxes[keep], ln_output))
            return keep

        if measure_gpu:
            # Move boxes to the GPU
            gpu_tlbr = cpu_tlbr.to(gpu)
            gpu_scores = cpu_scores.to(gpu)
            gpu_cls = cpu_cls.to(gpu)  # NOQA
            gpu_ln_boxes = cpu_ln_boxes.to(gpu)

            t1 = ub.Timerit(N, bestof=bestof, label='torch(gpu)')
            for timer in t1:
                with timer:
                    keep = torch_nms(gpu_tlbr, gpu_scores, thresh=thresh)
                    torch.cuda.synchronize()
            ydata[t1.label].append(t1.min())
            outputs[t1.label] = np.where(keep.cpu().numpy())[0]

            t1 = ub.Timerit(N, bestof=bestof, label='cython(gpu)')
            for timer in t1:
                with timer:
                    keep = non_max_supression(np_tlbr, np_scores, thresh=thresh, impl='gpu')
                    torch.cuda.synchronize()
            ydata[t1.label].append(t1.min())
            outputs[t1.label] = sorted(keep)

            from lightnet.data.transform._postprocess import NonMaxSupression
            t1 = ub.Timerit(N, bestof=bestof, label='lightnet-slow(gpu)')
            for timer in t1:
                with timer:
                    ln_output = NonMaxSupression._nms(gpu_ln_boxes, nms_thresh=thresh, class_nms=False, fast=False)
                    torch.cuda.synchronize()
            # convert lightnet NMS output to keep for consistency
            keep = _ln_output_to_keep(ln_output, gpu_ln_boxes)
            ydata[t1.label].append(t1.min())
            outputs[t1.label] = sorted(keep)

            if False:
                t1 = ub.Timerit(N, bestof=bestof, label='lightnet-fast(gpu)')
                for timer in t1:
                    with timer:
                        ln_output = NonMaxSupression._nms(gpu_ln_boxes, nms_thresh=thresh, class_nms=False, fast=True)
                        torch.cuda.synchronize()
                # convert lightnet NMS output to keep for consistency
                keep = _ln_output_to_keep(ln_output, gpu_ln_boxes)
                ydata[t1.label].append(t1.min())
                outputs[t1.label] = sorted(keep)

        if measure_cpu:
            t1 = ub.Timerit(N, bestof=bestof, label='torch(cpu)')
            for timer in t1:
                with timer:
                    keep = torch_nms(cpu_tlbr, cpu_scores, thresh=thresh)
            ydata[t1.label].append(t1.min())
            outputs[t1.label] = np.where(keep.cpu().numpy())[0]

        if True:
            t1 = ub.Timerit(N, bestof=bestof, label='cython(cpu)')
            for timer in t1:
                with timer:
                    keep = non_max_supression(np_tlbr, np_scores, thresh=thresh, impl='cpu')
            ydata[t1.label].append(t1.min())
            outputs[t1.label] = sorted(keep)

            t1 = ub.Timerit(N, bestof=bestof, label='numpy(cpu)')
            for timer in t1:
                with timer:
                    keep = non_max_supression(np_tlbr, np_scores, thresh=thresh, impl='py')
            ydata[t1.label].append(t1.min())
            outputs[t1.label] = sorted(keep)

        # Check that all kept boxes do not have more than `threshold` ious
        for key, idxs in outputs.items():
            ious = nh.util.box_ious(np_tlbr[idxs], np_tlbr[idxs])
            max_iou = (np.tril(ious) - np.eye(len(ious))).max()
            if max_iou > thresh:
                print('{} produced a bad result with max_iou={}'.format(key, max_iou))

        # Check result consistency:
        print('\nResult stats:')
        for key in sorted(outputs.keys()):
            print('    * {:<20}: num={}'.format(key, len(outputs[key])))

        print('\nResult overlaps (method1, method2: jaccard):')
        datas = []
        for k1, k2 in it.combinations(sorted(outputs.keys()), 2):
            idxs1 = set(outputs[k1])
            idxs2 = set(outputs[k2])
            jaccard = len(idxs1 & idxs2) / len(idxs1 | idxs2)
            datas.append((k1, k2, jaccard))
        datas = sorted(datas, key=lambda x: -x[2])
        for k1, k2, jaccard in datas:
            print('    * {:<20}, {:<20}: {:0.4f}'.format(k1, k2, jaccard))

    nh.util.mplutil.autompl()
    nh.util.mplutil.multi_plot(xdata, ydata, xlabel='num boxes', ylabel='seconds')
    nh.util.show_if_requested()

if __name__ == '__main__':
    """
    CommandLine:
        python -m netharn.util.nms.torch_nms all
    """
    import xdoctest
    xdoctest.doctest_module(__file__)


import numpy as np
import netCDF4 as nc
import csv
from hurp.RDConverter import * #RDWGSConverter
conv = RDWGSConverter()
import matplotlib.pyplot as plt
from scipy.interpolate import griddata
from sys import platform
#from mpl_toolkits.basemap import interp as mapinterp
from itertools import product
from os.path import exists
import cPickle

def msg(string,type=2):
    import sys

    if type == 1: print(string),
    if type == 2: print(string)

    sys.stdout.flush()
    return

def nei2wrf(NEIpath, wrfchemipath, wrfoutpath, Times, month):
    '''
    The purpose of this script is to write WRF-Chem emission files for HURP.
    There are three input sources:
    - NEI - Netherlands Emission Inventory: a list of area sources of CO2, NOx and PM10, with RD coordinates
            This data needs to be gridded.
    - wrfout - contains the information of the domains in HURP, one wrfout file for each domain. The coordinates are in WGS84 lat/lon, and spaced at fixed distance (km) intervals.

    Four wrfchemi output files will be written, one for each domain, containing the gridded emission data projected on the WRF grids.
    '''
    #--- settings
    FS               = 8 # FontSize
    domains          = ['d04','d03','d02','d01']
    kg2mol           = dict()
    kg2mol['CO2']    = 1000./44. # 1 kmol = 44 kg       | 1/44 kmol = 1 kg | 1 kg = 1000/44 mol
    kg2mol['NOx']    = 1000./30.8 # 1 kmol = 44 kg       | 1/44 kmol = 1 kg | 1 kg = 1000/44 mol
    kg2mol['PM10']   = 1.
    yr2hr            = 1./(365.*24.)
    Nsnaps           = 14 # 13 SNAPS + 1 for road traffic
    zTimes           = ['00z','12z']
    tracers          = ['NOx','PM10','CO2']
    selectedtracers  = tracers[:2]
    dows             = ('Mon','Tue','Wed','Thu','Fri','Sat','Sun')
    moys             = ('Jan','Feb','Mar','Apr','May','Jun','Jul','Aug','Sep','Oct','Nov','Dec')

                                    # in NER output file                                 # In Ingrid Super's file
    SNAPnames                    = [ [ '01: Afvalverwijdering                       '  ],   #SNAP9='waste treatment and disposal';
                                    [ '02: Bouw                                    '  ],   #
                                    [ '03: Chemische Industrie                     '  ],   #SNAP6='solvents and other product use';
                                    [ '04: Consumenten                             '  ],   #
                                    [ '05: Drinkwatervoorziening                   '  ],   #
                                    [ '06: Energiesector                           '  ],   #SNAP1='combustion in energy and transformation industries';
                                    [ '07: Handel, Diensten en Overheid (HDO)      '  ],   # SNAP2='non-industrial combustion plants';
                                    [ '08: Landbouw                                '  ],   #SNAP10='agriculture'
                                    [ '09: Natuur                                  '  ],   #
                                    [ '10: Overige industrie                       '  ],   #SNAP3='combustion in manufacturing industry'; SNAP4='production processes'
                                    [ '11: Raffinaderijen                          '  ],   #SNAP5='extraction and distribution of fossil fuels and geothermal energy';
                                    [ '12: Riolering en waterzuiveringsinstallaties'  ],   #
                                    [ '13: Verkeer en vervoer - Except road traffic'  ],   #SNAP8='other mobile sources and machinery'
                                    [ '14: Verkeer en vervoer - Wegverkeer         '  ] ]  #SNAP7='road transport';

    # 01: Afvalverwerking                          = SNAP  9 'waste treatment and disposal';
    # 02: Bouw                                     = SNAP  3 'combustion in manufacturing industry'; SNAP4='production processes' : SNAP 3 & 4 combined
    # 03: Chemische industrie                      = SNAP  6 'solvents and other product use';
    # 04: Consumenten                              = SNAP  3 'combustion in manufacturing industry'; SNAP4='production processes' : SNAP 3 & 4 combined
    # 05: Drinkwatervoorziening                    = Constant genomen
    # 06: Energiesector                            = SNAP  1 'combustion in energy and transformation industries';
    # 07: Handel, Diensten en Overheid (HDO)       = SNAP 7 'road transport'
    # 08: Landbouw                                 = SNAP 10 'agriculture'
    # 09: Natuur                                   = Constant genomen
    # 10: Overige industrie                        = SNAP3='combustion in manufacturing industry'; SNAP4='production processes' : SNAP 3 & 4 combined
    # 11: Raffinaderijen                           = SNAP  5 'extraction and distribution of fossil fuels and geothermal energy';
    # 12: Riolering en waterzuiveringsinstallaties = Constant genomen
    # 13: Verkeer en vervoer - Except road traffic = SNAP  8 'other mobile sources and machinery'
    # 14: Verkeer en vervoer - Wegverkeer          = SNAP  7 'road transport'
    TP_moy = dict() # Month of Year
    TP_dow = dict() # Day of Week
    TP_hod = dict() # Hour of Day

    # SNAP                      1      2      3      4      5      6      7      8      9     10     11     12     13     14
    TP_moy['Jan'] = np.array([1.000, 1.060, 0.950, 1.060, 1.000, 1.200, 0.880, 0.450, 1.000, 1.060, 1.200, 1.000, 0.880, 0.880])
    TP_moy['Feb'] = np.array([1.000, 1.047, 0.960, 1.047, 1.000, 1.150, 0.920, 1.300, 1.000, 1.047, 1.200, 1.000, 0.920, 0.920])
    TP_moy['Mar'] = np.array([1.000, 1.035, 1.020, 1.035, 1.000, 1.050, 0.980, 2.350, 1.000, 1.035, 1.200, 1.000, 0.980, 0.980])
    TP_moy['Apr'] = np.array([1.000, 1.010, 1.000, 1.010, 1.000, 1.000, 1.030, 1.700, 1.000, 1.010, 0.800, 1.000, 1.030, 1.030])
    TP_moy['May'] = np.array([1.000, 0.985, 1.010, 0.985, 1.000, 0.900, 1.050, 0.850, 1.000, 0.985, 0.800, 1.000, 1.050, 1.050])
    TP_moy['Jun'] = np.array([1.000, 0.960, 1.030, 0.960, 1.000, 0.850, 1.060, 0.850, 1.000, 0.960, 0.800, 1.000, 1.060, 1.060])
    TP_moy['Jul'] = np.array([1.000, 0.965, 1.030, 0.965, 1.000, 0.800, 1.010, 0.850, 1.000, 0.965, 0.800, 1.000, 1.010, 1.010])
    TP_moy['Aug'] = np.array([1.000, 0.935, 1.010, 0.935, 1.000, 0.875, 1.020, 1.000, 1.000, 0.935, 0.800, 1.000, 1.020, 1.020])
    TP_moy['Sep'] = np.array([1.000, 0.995, 1.040, 0.995, 1.000, 0.950, 1.060, 1.100, 1.000, 0.995, 0.800, 1.000, 1.060, 1.060])
    TP_moy['Oct'] = np.array([1.000, 1.010, 1.030, 1.010, 1.000, 1.000, 1.050, 0.650, 1.000, 1.010, 1.200, 1.000, 1.050, 1.050])
    TP_moy['Nov'] = np.array([1.000, 1.023, 1.010, 1.023, 1.000, 1.075, 1.010, 0.450, 1.000, 1.023, 1.200, 1.000, 1.010, 1.010])
    TP_moy['Dec'] = np.array([1.000, 0.975, 0.910, 0.975, 1.000, 1.150, 0.930, 0.450, 1.000, 0.975, 1.200, 1.000, 0.930, 0.930])

    TP_dow['Mon'] = np.array([1.000, 1.050, 1.200, 1.050, 1.000, 1.060, 1.020, 1.000, 1.000, 1.050, 1.000, 1.000, 1.000, 1.020])
    TP_dow['Tue'] = np.array([1.000, 1.050, 1.200, 1.050, 1.000, 1.060, 1.060, 1.000, 1.000, 1.050, 1.000, 1.000, 1.000, 1.060])
    TP_dow['Wed'] = np.array([1.000, 1.050, 1.200, 1.050, 1.000, 1.060, 1.080, 1.000, 1.000, 1.050, 1.000, 1.000, 1.000, 1.080])
    TP_dow['Thu'] = np.array([1.000, 1.050, 1.200, 1.050, 1.000, 1.060, 1.100, 1.000, 1.000, 1.050, 1.000, 1.000, 1.000, 1.100])
    TP_dow['Fri'] = np.array([1.000, 1.050, 1.200, 1.050, 1.000, 1.060, 1.140, 1.000, 1.000, 1.050, 1.000, 1.000, 1.000, 1.140])
    TP_dow['Sat'] = np.array([1.000, 0.910, 0.500, 0.910, 1.000, 0.850, 0.810, 1.000, 1.000, 0.910, 1.000, 1.000, 1.000, 0.810])
    TP_dow['Sun'] = np.array([1.000, 0.840, 0.500, 0.840, 1.000, 0.850, 0.790, 1.000, 1.000, 0.840, 1.000, 1.000, 1.000, 0.790])

    TP_hod['00']  = np.array([1.000, 0.875, 0.500, 0.875, 1.000, 0.790, 0.190, 0.600, 1.000, 0.875, 1.000, 1.000, 1.000, 0.190])
    TP_hod['01']  = np.array([1.000, 0.875, 0.350, 0.875, 1.000, 0.720, 0.090, 0.600, 1.000, 0.875, 1.000, 1.000, 1.000, 0.090])
    TP_hod['02']  = np.array([1.000, 0.890, 0.200, 0.890, 1.000, 0.720, 0.060, 0.600, 1.000, 0.890, 1.000, 1.000, 1.000, 0.060])
    TP_hod['03']  = np.array([1.000, 0.910, 0.100, 0.910, 1.000, 0.710, 0.050, 0.600, 1.000, 0.910, 1.000, 1.000, 1.000, 0.050])
    TP_hod['04']  = np.array([1.000, 0.940, 0.100, 0.940, 1.000, 0.740, 0.090, 0.600, 1.000, 0.940, 1.000, 1.000, 1.000, 0.090])
    TP_hod['05']  = np.array([1.000, 0.975, 0.200, 0.975, 1.000, 0.800, 0.220, 0.650, 1.000, 0.975, 1.000, 1.000, 1.000, 0.220])
    TP_hod['06']  = np.array([1.000, 1.010, 0.750, 1.010, 1.000, 0.920, 0.860, 0.750, 1.000, 1.010, 1.000, 1.000, 1.000, 0.860])
    TP_hod['07']  = np.array([1.000, 1.045, 1.250, 1.045, 1.000, 1.080, 1.840, 0.900, 1.000, 1.045, 1.000, 1.000, 1.000, 1.840])
    TP_hod['08']  = np.array([1.000, 1.080, 1.400, 1.080, 1.000, 1.190, 1.860, 1.100, 1.000, 1.080, 1.000, 1.000, 1.000, 1.860])
    TP_hod['09']  = np.array([1.000, 1.110, 1.500, 1.110, 1.000, 1.220, 1.410, 1.350, 1.000, 1.110, 1.000, 1.000, 1.000, 1.410])
    TP_hod['10']  = np.array([1.000, 1.140, 1.500, 1.140, 1.000, 1.210, 1.240, 1.450, 1.000, 1.140, 1.000, 1.000, 1.000, 1.240])
    TP_hod['11']  = np.array([1.000, 1.150, 1.500, 1.150, 1.000, 1.210, 1.200, 1.600, 1.000, 1.150, 1.000, 1.000, 1.000, 1.200])
    TP_hod['12']  = np.array([1.000, 1.110, 1.500, 1.110, 1.000, 1.170, 1.320, 1.650, 1.000, 1.110, 1.000, 1.000, 1.000, 1.320])
    TP_hod['13']  = np.array([1.000, 1.120, 1.500, 1.120, 1.000, 1.150, 1.440, 1.750, 1.000, 1.120, 1.000, 1.000, 1.000, 1.440])
    TP_hod['14']  = np.array([1.000, 1.125, 1.500, 1.125, 1.000, 1.140, 1.450, 1.700, 1.000, 1.125, 1.000, 1.000, 1.000, 1.450])
    TP_hod['15']  = np.array([1.000, 1.080, 1.500, 1.080, 1.000, 1.130, 1.590, 1.550, 1.000, 1.080, 1.000, 1.000, 1.000, 1.590])
    TP_hod['16']  = np.array([1.000, 1.040, 1.500, 1.040, 1.000, 1.100, 2.030, 1.350, 1.000, 1.040, 1.000, 1.000, 1.000, 2.030])
    TP_hod['17']  = np.array([1.000, 1.005, 1.400, 1.005, 1.000, 1.070, 2.080, 1.100, 1.000, 1.005, 1.000, 1.000, 1.000, 2.080])
    TP_hod['18']  = np.array([1.000, 0.975, 1.250, 0.975, 1.000, 1.040, 1.510, 0.900, 1.000, 0.975, 1.000, 1.000, 1.000, 1.510])
    TP_hod['19']  = np.array([1.000, 0.950, 1.100, 0.950, 1.000, 1.020, 1.060, 0.750, 1.000, 0.950, 1.000, 1.000, 1.000, 1.060])
    TP_hod['20']  = np.array([1.000, 0.925, 1.000, 0.925, 1.000, 1.020, 0.740, 0.650, 1.000, 0.925, 1.000, 1.000, 1.000, 0.740])
    TP_hod['21']  = np.array([1.000, 0.905, 0.900, 0.905, 1.000, 1.010, 0.620, 0.600, 1.000, 0.905, 1.000, 1.000, 1.000, 0.620])
    TP_hod['22']  = np.array([1.000, 0.890, 0.800, 0.890, 1.000, 0.960, 0.610, 0.600, 1.000, 0.890, 1.000, 1.000, 1.000, 0.610])
    TP_hod['23']  = np.array([1.000, 0.875, 0.700, 0.875, 1.000, 0.880, 0.440, 0.600, 1.000, 0.875, 1.000, 1.000, 1.000, 0.440])

    #--- input
    if not 'progress' in globals(): progress = list()
    if not 'NEIloaded' in progress:
        #--- init
        msg('Loading gridded NEI files: ',type=1)
        E_nei             = dict()
        for tracer in selectedtracers:
            NEIncfilename = '%s/NEI_%s.nc'%(NEIpath,tracer)
            msg('%s...'%NEIncfilename,type=1)
            NEIncfile     = nc.Dataset(NEIncfilename,'r')
            E_nei[tracer] = NEIncfile.variables['E_%s'%tracer][:] # kg per area
            if not 'xnei' in globals():
                xnei      = NEIncfile.variables['Longitude'][:]
                ynei      = NEIncfile.variables['Latitude' ][:]
                Xnei,Ynei = np.meshgrid(xnei,ynei)
                ny,nx     = Xnei.shape
            NEIncfile.close()
        progress.append('NEIloaded')

    if not 'WRFOUTloaded' in progress:
        msg('Loading lat/lon from wrfout...',type=1)
        xlat = dict()
        xlon = dict()
        for domain in domains:
            wrfoutfilename = '%s/wrfinput_%s'%(wrfoutpath,domain)
            wrfoutfile     = nc.Dataset(wrfoutfilename,'r')
            xlat[domain]   = wrfoutfile.variables['XLAT' ][0,:] # centerpoints
            xlon[domain]   = wrfoutfile.variables['XLONG'][0,:] # centerpoints
            wrfoutfile.close()
        progress.append('WRFOUTloaded')
        msg('done.')

    if not 'NEIinterpolatedToWRF' in progress:
        Xwrf             = dict()
        Ywrf             = dict()
        E_wrf            = dict()
        for tracer in selectedtracers:
            picklefilename = '%s/E_wrf_%s.pickle'%(wrfoutpath,tracer)
            if exists(picklefilename):
                msg('Loading pre-cooked %s'%picklefilename,type=1)
                pickledata = cPickle.load(open(picklefilename,'r'))
                Xwrf,Ywrf,E_wrf_tmp,domains,tracers = pickledata
                for domain in domains:
                    E_wrf[domain,tracer] = E_wrf_tmp[domain,tracer]
                del(E_wrf_tmp)
            else:

                # Note: d04 has a higher resolution than NEI --? interpolate or nearest neighbor
                #       d03 has the same resolution as NEI, but may be shifted in x or y direction --> nearest neighbour regridding
                #       d02 has a coarser resolution than NEI (but still fits in the NEI coverage?) --> mass conservative interpolation
                #       d01 has a coarser resolution than NEI and also extends beyond the NEI domain --> Use MACC (how about MACC resolution?)
                msg('Interpolating...',type=1)
                for domain in domains:
                    nlat,nlon    = xlat[domain].shape
                    if (tracer == selectedtracers[0]):
                        msg('Converting (lon,lat) in (x,y) for %s'%domain,type=1)
                        Xwrf[domain] = np.empty((nlat,nlon))
                        Ywrf[domain] = np.empty((nlat,nlon))
                        for ilat,ilon in product(range(nlat), range(nlon)):
                            x,y   = conv.fromWgsToRd([xlat[domain][ilat,ilon],xlon[domain][ilat,ilon]])
            #               msg xlon[domain][ilat,ilon], xlat[domain][ilat,ilon], x, y
                            Xwrf[domain][ilat,ilon] = x
                            Ywrf[domain][ilat,ilon] = y

                    msg('Interpolating %s...'%tracer,type=1)
                    E_wrf[domain,tracer] = np.zeros((Nsnaps,nlat,nlon))
                    for isnap in range(Nsnaps):
                        msg('snap %02d...'%isnap,type=1)
                        E_wrf[domain,tracer][isnap,:,:] = griddata((Xnei.flatten(),Ynei.flatten()),E_nei[tracer][isnap].flatten() , (Xwrf[domain],Ywrf[domain]),method='nearest')
                    # of
                    # sst_smooth = basemap.interp(sst_coarse, lons_sub[0,:], lats_sub[:,0], *np.meshgrid(lons, lats), order=1)
                    #E_NOx_wrf [domain] = mapinterp(E_NOx_nei [Isnaps].sum(axis=0),xnei,ynei, Xwrf[domain],Ywrf[domain],order=0) # order 0:Nearest
                    #E_PM10_wrf[domain] = mapinterp(E_PM10_nei[Isnaps].sum(axis=0),xnei,ynei, Xwrf[domain],Ywrf[domain],order=0) # order 1: bilinear

                E_wrf_tmp = {}
                for domain in domains:
                    E_wrf_tmp[domain,tracer] = E_wrf[domain,tracer]
                    pickledata     = (Xwrf,Ywrf,E_wrf_tmp,domains,tracers)
                    cPickle.dump(pickledata,open(picklefilename,'wb'))
                del(E_wrf_tmp)

        progress.append('NEIinterpolatedToWRF')
        msg('done.')

    moy = moys[month-1]
    if not 'wrfchemi_written' in progress:
        for domain,zTime,dow in product(domains,zTimes,dows):
            print('Writing wrfchemifile for %s %s %s %s'%(domain,zTime,moy,dow))
            #its = range(12) if zTime == '00z' else range(12,24)
            zTimeOffset = 12 if zTime == '12z' else 0

            wrfoutfilename            = '%s/wrfinput_%s'%(wrfoutpath,domain)
            wrfoutfile                = nc.Dataset(wrfoutfilename,'r')
            wrfchemifilename          = '%s/NEI_wrfchemi_%s_%s_%s_%s'%(wrfchemipath,zTime,domain,moy,dow)
            wrfchemifile              = nc.Dataset(wrfchemifilename, 'w',format = 'NETCDF3_CLASSIC' )

            nz                        =  1
            nlat,nlon                 = xlat[domain].shape

            wrfchemifile.createDimension('Time',        None) # will be 12
            wrfchemifile.createDimension('bottom_top' , nz)
            wrfchemifile.createDimension('south_north', nlat)
            wrfchemifile.createDimension('west_east'  , nlon)
            wrfchemifile.createDimension('DateStrLen' , 19)

            ncTimes                   = wrfchemifile.createVariable('Times'               ,'c',dimensions=('Time','DateStrLen'))
            ncXLAT                    = wrfchemifile.createVariable('XLAT'                ,'d',dimensions=('south_north','west_east'))
            ncXLONG                   = wrfchemifile.createVariable('XLONG'               ,'d',dimensions=('south_north','west_east'))
            ncE_NOx_stat              = wrfchemifile.createVariable('E_NOx_stat'          ,'d',dimensions=('Time','bottom_top','south_north','west_east'),fill_value = -1.e34)
            ncE_NOx_traf              = wrfchemifile.createVariable('E_NOx_traf'          ,'d',dimensions=('Time','bottom_top','south_north','west_east'),fill_value = -1.e34)
            ncE_PM10_stat             = wrfchemifile.createVariable('E_PM10_stat'         ,'d',dimensions=('Time','bottom_top','south_north','west_east'),fill_value = -1.e34)
            ncE_PM10_traf             = wrfchemifile.createVariable('E_PM10_traf'         ,'d',dimensions=('Time','bottom_top','south_north','west_east'),fill_value = -1.e34)

            ncTimes                   = Times[zTime]
            ncXLAT.FieldType          = 104.
            ncXLAT.MemoryOrder        = 'XY'
            ncXLAT.description        = 'LATITUDE, SOUTH IS NEGATIVE'
            ncXLAT.units              = 'degree_north'
            ncXLAT.stagger            = ' '
            ncXLAT.coordinates        = 'XLONG XLAT'
            ncXLAT[:]                 = xlat[domain]
            ncXLONG.FieldType         = 104.
            ncXLONG.MemoryOrder       = 'XY'
            ncXLONG.description       = 'LONGGITUDE, WEST IS NEGATIVE'
            ncXLONG.units             = 'degree_north'
            ncXLONG.stagger           = ' '
            ncXLONG.coordinates       = 'XLONG XLAT'
            ncXLONG[:]                = xlon[domain]

            ncE_NOx_stat.FieldType    = 104.
            ncE_NOx_stat.MemoryOrder  = 'XYZ'
            ncE_NOx_stat.description  = 'NOx emissions from stationary sources'
            ncE_NOx_stat.units        = 'mole/km2/hr'
            ncE_NOx_stat.stagger      = ' '
            ncE_NOx_stat.coordinates  = 'XLONG XLAT'
            Isnap                     = range(13)
            for it in range(12):
                TP                    = TP_hod['%02d'%(it+zTimeOffset)][Isnap] * TP_dow[dow][Isnap] * TP_moy[moy][Isnap]
                TP2                   = np.tile(TP[...,np.newaxis,np.newaxis],(1,nlat,nlon))
                ncE_NOx_stat[it,0,:,:]= (TP2 * E_wrf[domain,'NOx'][Isnap,:,:]).sum(axis=0)
    #           ncE_NOx_stat[it,:]    = (TP2 * E_wrf[domain,'NOx'][Isnap]).sum(axis=0)

            ncE_NOx_traf.FieldType    = 104.
            ncE_NOx_traf.MemoryOrder  = 'XYZ'
            ncE_NOx_traf.description  = 'NOx emissions from road traffic'
            ncE_NOx_traf.units        = 'mole/km2/hr'
            ncE_NOx_traf.stagger      = ' '
            ncE_NOx_traf.coordinates  = 'XLONG XLAT'
            Isnap                     = 13
            for it in range(12):
                TP                    = TP_hod['%02d'%(it+zTimeOffset)][Isnap] * TP_dow[dow][Isnap] * TP_moy[moy][Isnap]
                ncE_NOx_traf[it,0,:,:] = TP*E_wrf[domain,'NOx'][Isnap,:,:]
    #           TP2                   = np.tile(TP[...,np.newaxis,np.newaxis],(1,nlat,nlon))

            ncE_PM10_stat.FieldType   = 104.
            ncE_PM10_stat.MemoryOrder = 'XYZ'
            ncE_PM10_stat.description = 'PM10 emissions from stationary sources'
            ncE_PM10_stat.units       = 'kg/km2/hr'
            ncE_PM10_stat.stagger     = ' '
            ncE_PM10_stat.coordinates = 'XLONG XLAT'
            Isnap                     = range(13)
            for it in range(12):
                TP                    = TP_hod['%02d'%(it+zTimeOffset)][Isnap] * TP_dow[dow][Isnap] * TP_moy[moy][Isnap]
                TP2                   = np.tile(TP[...,np.newaxis,np.newaxis],(1,nlat,nlon))
                ncE_PM10_stat[it,0,:,:]= (TP2 * E_wrf[domain,'PM10'][Isnap,:,:]).sum(axis=0)

            ncE_PM10_traf.FieldType   = 104.
            ncE_PM10_traf.MemoryOrder = 'XYZ'
            ncE_PM10_traf.description = 'PM10 emissions from road traffic'
            ncE_PM10_traf.units       = 'kg/km2/hr'
            ncE_PM10_traf.stagger     = ' '
            ncE_PM10_traf.coordinates = 'XLONG XLAT'
            Isnap                     = 13
            for it in range(12):
                TP                    = TP_hod['%02d'%(it+zTimeOffset)][Isnap] * TP_dow[dow][Isnap] * TP_moy[moy][Isnap]
    #           TP2                   = np.tile(TP[...,np.newaxis,np.newaxis],(1,nlat,nlon))
                ncE_PM10_traf[it,0,:,:]   = TP*E_wrf[domain,'PM10'][Isnap,:,:]

            wrfchemifile.CEN_LAT      = wrfoutfile.getncattr('CEN_LAT')
            wrfchemifile.CEN_LON      = wrfoutfile.getncattr('CEN_LON')
            wrfchemifile.TRUELAT1     = wrfoutfile.getncattr('TRUELAT1')
            wrfchemifile.TRUELAT2     = wrfoutfile.getncattr('TRUELAT2')
            wrfchemifile.MOAD_CEN_LAT = wrfoutfile.getncattr('MOAD_CEN_LAT')
            wrfchemifile.STAND_LON    = wrfoutfile.getncattr('STAND_LON')
            wrfchemifile.POLE_LAT     = wrfoutfile.getncattr('POLE_LAT')
            wrfchemifile.POLE_LON     = wrfoutfile.getncattr('POLE_LON')
            wrfchemifile.GMT          = wrfoutfile.getncattr('GMT')
            wrfchemifile.JULYR        = wrfoutfile.getncattr('JULYR')
            wrfchemifile.JULDAY       = wrfoutfile.getncattr('JULDAY')
            wrfchemifile.MAP_PROJ     = wrfoutfile.getncattr('MAP_PROJ')
            wrfchemifile.MMINLU       = wrfoutfile.getncattr('MMINLU')
            wrfchemifile.NUM_LAND_CAT = wrfoutfile.getncattr('NUM_LAND_CAT')
            wrfchemifile.ISWATER      = wrfoutfile.getncattr('ISWATER')
            wrfchemifile.ISLAKE       = wrfoutfile.getncattr('ISLAKE')
            wrfchemifile.ISICE        = wrfoutfile.getncattr('ISICE')
            wrfchemifile.ISURBAN      = wrfoutfile.getncattr('ISURBAN')
            wrfchemifile.ISOILWATER   = wrfoutfile.getncattr('ISOILWATER')

            wrfchemifile.close()
            progress.append('wrfchemi_written')

    #--- output
    if not 'dofig' in globals(): dofig = [None, False, True, False]
    dofig = 4*(False,)
    #E_CO2_nei = np.ma.masked_where(E_CO2_nei == 0, E_CO2_nei)
    #E_NOx_nei = np.ma.masked_where(E_NOx_nei == 0, E_NOx_nei)
    if dofig[1]:
        tracer = 'NOx'
        f = plt.figure(1)
        f.clf()
        ax = f.add_subplot(111)
        ax.scatter(data[tracer][:,0],data[tracer][:,1],s=20,c=np.log10(data[tracer][:,5]),marker='.',edgecolors=None)
        msg('fig 1 done.')
    plt.show()

    if dofig[2]:
        Isnap = range(14)
        tracer = 'NOx'
        f = plt.figure(2)
        f.clf()
        ax = f.add_subplot(111)
        Z = E_nei[tracer][Isnap,:].sum(axis=0); Z = np.ma.masked_where(Z == 0, Z)
        h=ax.pcolor(Xnei,Ynei,np.log10(Z))
        for domain in domains[:3]:
            xwrf = Xwrf[domain]
            ywrf = Ywrf[domain]
            ax.plot([xwrf.min(), xwrf.max(), xwrf.max(),xwrf.min(),xwrf.min()],[ywrf.min(),ywrf.min(),ywrf.max(),ywrf.max(),ywrf.min()],'k-')
            cb = f.colorbar(h,orientation='vertical')
            pow = list()
            for YTL in cb.ax.get_ymajorticklabels(): pow.append(np.double(YTL.get_text().replace(u'\N{MINUS SIGN}', '-')))
            cb.ax.set_yticklabels(np.power(10,pow),fontsize=FS)

        msg('fig 2 done.')
        dofig[2] = False

    if dofig[3]:
        msg('Doing figure 3...',type=1)
        f = plt.figure(3)
        f.clf()
        for domain in domains[::-1]:
            tracer = 'NOx'

            ax= f.add_subplot(111)
            Z = E_wrf_stat[domain,tracer]+E_wrf_traf[domain,tracer]; Z = np.ma.masked_where(Z == 0, Z)
            h=ax.pcolor(Xwrf[domain],Ywrf[domain],np.log10(Z),vmin=0,vmax=np.log10(2000))
            xwrf = Xwrf[domain]
            ywrf = Ywrf[domain]
            ax.plot([xwrf.min(), xwrf.max(), xwrf.max(),xwrf.min(),xwrf.min()],[ywrf.min(),ywrf.min(),ywrf.max(),ywrf.max(),ywrf.min()],'k-')

        cb = f.colorbar(h,orientation='vertical')
        pow = list()
        for YTL in cb.ax.get_ymajorticklabels(): pow.append(np.double(YTL.get_text().replace(u'\N{MINUS SIGN}', '-')))
        cb.ax.set_yticklabels(np.power(10,pow),fontsize=FS)

        msg('Done figure 3.')

if __name__=="__main__":
    if platform == 'win32':
        NEIpath      = 'U:/AMS_Stimulus_2016/Data_share/Emissions/NEI'
        wrfchemipath = 'U:/AMS_Stimulus_2016/Data_share/Emissions/wrfchemi'
        wrfoutpath   = 'S:/data/HURP/Adam_wrfout'
    elif (platform == 'darwin' and exists('/Users/molen050/mnt/promise')): # CapeGrim
        NEIpath      = '/Users/molen050/mnt/promise/WRF/michiel/HURP/Data_share/Emissions/NEI'
        wrfchemipath = '/Users/molen050/mnt/promise/WRF/michiel/HURP/Data_share/Emissions/wrfchemi'
        wrfoutpath   = '/Users/molen050/mnt/promise/WRF/michiel/HURP/Data_share/Domains'
    elif (platform == 'linux2' ): # Cartesius
        print('Setting env for Cartesius.')
        NEIpath      = '/projects/0/aams/wrfv3/Data_share/Emissions/NEI'
        wrfchemipath = '/projects/0/aams/wrfv3/Data_share/Emissions/wrfchemi'
        wrfoutpath   = '/projects/0/aams/wrfv3/OutputMichiel/HURP_20150630-20150704'
    else:  # maunaloa
        NEIpath      = '/Storage/WRF/michiel/HURP/Data_share/Emissions/NEI'
        wrfchemipath = '/Storage/WRF/michiel/HURP/Data_share/Emissions/wrfchemi'
        wrfoutpath   = '/Storage/WRF/michiel/HURP/Data_share/Domains'

    Times            = dict()
    Times['00z']     = [['2014-01-01_00:00:00'],['2014-01-01_01:00:00'],['2014-01-01_02:00:00'],['2014-01-01_03:00:00'],
                        ['2014-01-01_04:00:00'],['2014-01-01_05:00:00'],['2014-01-01_06:00:00'],['2014-01-01_07:00:00'],
                        ['2014-01-01_08:00:00'],['2014-01-01_09:00:00'],['2014-01-01_10:00:00'],['2014-01-01_11:00:00']]

    Times['12z']     = [['2014-01-01_12:00:00'],['2014-01-01_13:00:00'],['2014-01-01_14:00:00'],['2014-01-01_15:00:00'],
                        ['2014-01-01_16:00:00'],['2014-01-01_17:00:00'],['2014-01-01_18:00:00'],['2014-01-01_19:00:00'],
                        ['2014-01-01_20:00:00'],['2014-01-01_21:00:00'],['2014-01-01_22:00:00'],['2014-01-01_23:00:00']]
    nei2wrf(NEIpath, wrfchemipath, wrfoutpath, Times)


# Licensed under a 3-clause BSD style license - see LICENSE.rst
import numpy as np
from astropy.coordinates import SkyOffsetFrame
from gammapy.data import FixedPointingInfo
from gammapy.irf import EDispMap, PSFMap
from gammapy.maps import Map, WcsNDMap
from gammapy.modeling.models import PowerLawSpectralModel
from gammapy.utils.coordinates import sky_to_fov

__all__ = [
    "make_map_background_irf",
    "make_edisp_map",
    "make_edisp_kernel_map",
    "make_psf_map",
    "make_map_exposure_true_energy",
]

def make_map_exposure_true_energy(pointing, livetime, aeff, geom):
    """Compute exposure map.

    This map has a true energy axis, the exposure is not combined
    with energy dispersion.

    Parameters
    ----------
    pointing : `~astropy.coordinates.SkyCoord`
        Pointing direction
    livetime : `~astropy.units.Quantity`
        Livetime
    aeff : `~gammapy.irf.EffectiveAreaTable2D`
        Effective area
    geom : `~gammapy.maps.WcsGeom`
        Map geometry (must have an energy axis)

    Returns
    -------
    map : `~gammapy.maps.WcsNDMap`
        Exposure map
    """
    offset = geom.separation(pointing)
    energy = geom.get_axis_by_name("energy_true").center

    exposure = aeff.data.evaluate(
        offset=offset, energy_true=energy[:, np.newaxis, np.newaxis]
    )
    # TODO: Improve IRF evaluate to preserve energy axis if length 1
    # For now, we handle that case via this hack:
    if len(exposure.shape) < 3:
        exposure = np.expand_dims(exposure.value, 0) * exposure.unit

    exposure = (exposure * livetime).to("m2 s")

    return WcsNDMap(geom, exposure.value.reshape(geom.data_shape), unit=exposure.unit)

def _map_spectrum_weight(map, spectrum=None):
    """Weight a map with a spectrum.

    This requires map to have an "energy" axis.
    The weights are normalised so that they sum to 1.
    The mean and unit of the output image is the same as of the input cube.

    At the moment this is used to get a weighted exposure image.

    Parameters
    ----------
    map : `~gammapy.maps.Map`
        Input map with an "energy" axis.
    spectrum : `~gammapy.modeling.models.SpectralModel`
        Spectral model to compute the weights.
        Default is power-law with spectral index of 2.

    Returns
    -------
    map_weighted : `~gammapy.maps.Map`
        Weighted image
    """
    if spectrum is None:
        spectrum = PowerLawSpectralModel(index=2.0)

    # Compute weights vector
    energy_edges = map.geom.get_axis_by_name("energy_true").edges
    weights = spectrum.integral(
        emin=energy_edges[:-1], emax=energy_edges[1:], intervals=True
    )
    weights /= weights.sum()
    shape = np.ones(len(map.geom.data_shape))
    shape[0] = -1
    return map * weights.reshape(shape.astype(int))

def make_map_background_irf(pointing, ontime, bkg, geom, oversampling=None):
    """Compute background map from background IRFs.

    Parameters
    ----------
    pointing : `~gammapy.data.FixedPointingInfo` or `~astropy.coordinates.SkyCoord`
        Observation pointing

        - If a ``FixedPointingInfo`` is passed, FOV coordinates are properly computed.
        - If a ``SkyCoord`` is passed, FOV frame rotation is not taken into account.
    ontime : `~astropy.units.Quantity`
        Observation ontime. i.e. not corrected for deadtime
        see https://gamma-astro-data-formats.readthedocs.io/en/stable/irfs/full_enclosure/bkg/index.html#notes)
    bkg : `~gammapy.irf.Background3D`
        Background rate model
    geom : `~gammapy.maps.WcsGeom`
        Reference geometry
    oversampling: int
        Oversampling factor in energy, used for the background model evaluation.

    Returns
    -------
    background : `~gammapy.maps.WcsNDMap`
        Background predicted counts sky cube in reco energy
    """
    # TODO:
    #  This implementation can be improved in two ways:
    #  1. Create equal time intervals between TSTART and TSTOP and sum up the
    #  background IRF for each interval. This is instead of multiplying by
    #  the total ontime. This then handles the rotation of the FoV.
    #  2. Use the pointing table (does not currently exist in CTA files) to
    #  obtain the RA DEC and time for each interval. This then considers that
    #  the pointing might change slightly over the observation duration

    # Get altaz coords for map
    if oversampling is not None:
        geom = geom.upsample(factor=oversampling, axis="energy")

    map_coord = geom.to_image().get_coord()
    sky_coord = map_coord.skycoord

    if isinstance(pointing, FixedPointingInfo):
        altaz_coord = sky_coord.transform_to(pointing.altaz_frame)

        # Compute FOV coordinates of map relative to pointing
        fov_lon, fov_lat = sky_to_fov(
            altaz_coord.az, altaz_coord.alt, pointing.altaz.az, pointing.altaz.alt
        )
    else:
        # Create OffsetFrame
        frame = SkyOffsetFrame(origin=pointing)
        pseudo_fov_coord = sky_coord.transform_to(frame)
        fov_lon = pseudo_fov_coord.lon
        fov_lat = pseudo_fov_coord.lat

    energies = geom.get_axis_by_name("energy").edges

    bkg_de = bkg.evaluate_integrate(
        fov_lon=fov_lon,
        fov_lat=fov_lat,
        energy_reco=energies[:, np.newaxis, np.newaxis],
    )

    d_omega = geom.to_image().solid_angle()
    data = (bkg_de * d_omega * ontime).to_value("")
    bkg_map = WcsNDMap(geom, data=data)

    if oversampling is not None:
        bkg_map = bkg_map.downsample(factor=oversampling, axis="energy")

    return bkg_map

def make_psf_map(psf, pointing, geom, exposure_map=None):
    """Make a psf map for a single observation

    Expected axes : rad and true energy in this specific order
    The name of the rad MapAxis is expected to be 'rad'

    Parameters
    ----------
    psf : `~gammapy.irf.PSF3D`
        the PSF IRF
    pointing : `~astropy.coordinates.SkyCoord`
        the pointing direction
    geom : `~gammapy.maps.Geom`
        the map geom to be used. It provides the target geometry.
        rad and true energy axes should be given in this specific order.
    exposure_map : `~gammapy.maps.Map`, optional
        the associated exposure map.
        default is None

    Returns
    -------
    psfmap : `~gammapy.cube.PSFMap`
        the resulting PSF map
    """
    energy_axis = geom.get_axis_by_name("energy_true")
    energy = energy_axis.center

    rad_axis = geom.get_axis_by_name("theta")
    rad = rad_axis.center

    # Compute separations with pointing position
    offset = geom.separation(pointing)

    # Compute PSF values
    # TODO: allow broadcasting in PSF3D.evaluate()
    psf_values = psf._interpolate(
        (
            rad[:, np.newaxis, np.newaxis],
            offset,
            energy[:, np.newaxis, np.newaxis, np.newaxis],
        )
    )

    # TODO: this probably does not ensure that probability is properly normalized in the PSFMap
    # Create Map and fill relevant entries
    data = psf_values.to_value("sr-1")
    psfmap = Map.from_geom(geom, data=data, unit="sr-1")
    return PSFMap(psfmap, exposure_map)

def make_edisp_map(edisp, pointing, geom, exposure_map=None):
    """Make a edisp map for a single observation

    Expected axes : migra and true energy in this specific order
    The name of the migra MapAxis is expected to be 'migra'

    Parameters
    ----------
    edisp : `~gammapy.irf.EnergyDispersion2D`
        the 2D Energy Dispersion IRF
    pointing : `~astropy.coordinates.SkyCoord`
        the pointing direction
    geom : `~gammapy.maps.Geom`
        the map geom to be used. It provides the target geometry.
        migra and true energy axes should be given in this specific order.
    exposure_map : `~gammapy.maps.Map`, optional
        the associated exposure map.
        default is None

    Returns
    -------
    edispmap : `~gammapy.cube.EDispMap`
        the resulting EDisp map
    """
    energy_axis = geom.get_axis_by_name("energy_true")
    energy = energy_axis.center

    migra_axis = geom.get_axis_by_name("migra")
    migra = migra_axis.center

    # Compute separations with pointing position
    offset = geom.separation(pointing)

    # Compute EDisp values
    edisp_values = edisp.data.evaluate(
        offset=offset,
        energy_true=energy[:, np.newaxis, np.newaxis, np.newaxis],
        migra=migra[:, np.newaxis, np.newaxis],
    )

    # Create Map and fill relevant entries
    data = edisp_values.to_value("")
    edispmap = Map.from_geom(geom, data=data, unit="")
    return EDispMap(edispmap, exposure_map)

def make_edisp_kernel_map(edisp, pointing, geom, exposure_map=None):
    """Make a edisp kernel map for a single observation

    Expected axes : (reco) energy and true energy in this specific order
    The name of the reco energy MapAxis is expected to be 'energy'.
    The name of the true energy MapAxis is expected to be 'energy_true'.

    Parameters
    ----------
    edisp : `~gammapy.irf.EnergyDispersion2D`
        the 2D Energy Dispersion IRF
    pointing : `~astropy.coordinates.SkyCoord`
        the pointing direction
    geom : `~gammapy.maps.Geom`
        the map geom to be used. It provides the target geometry.
        energy and true energy axes should be given in this specific order.
    exposure_map : `~gammapy.maps.Map`, optional
        the associated exposure map.
        default is None

    Returns
    -------
    edispmap : `~gammapy.cube.EDispKernelMap`
        the resulting EDispKernel map
    """
    # Use EnergyDispersion2D migra axis.
    migra_axis = edisp.data.axis("migra")

    # Create temporary EDispMap Geom
    new_geom = geom.to_image().to_cube(
        [migra_axis, geom.get_axis_by_name("energy_true")]
    )

    edisp_map = make_edisp_map(edisp, pointing, new_geom, exposure_map)

    return edisp_map.to_edisp_kernel_map(geom.get_axis_by_name("energy"))


import torch
import torch.nn.functional as F
from torch.distributions.normal import Normal
import numpy as np

class PolicyNet(torch.nn.Module):
    """
    Gaussian Policy
    """
    def __init__(self, state_dim, hidden_dim, action_dim, action_bound):
        super(PolicyNet, self).__init__()
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
        self.fc_mu = torch.nn.Linear(hidden_dim, action_dim)
        self.fc_std = torch.nn.Linear(hidden_dim, action_dim)
        self.action_bound = action_bound

    def forward(self, x):
        x = F.relu(self.fc1(x))
        mu = self.fc_mu(x)
        std = F.softplus(self.fc_std(x))
        dist = Normal(mu, std)
        normal_sample = dist.rsample()  # rsample()是重参数化采样函数
        log_prob = dist.log_prob(normal_sample)
        action = torch.tanh(normal_sample)  # 计算tanh_normal分布的对数概率密度
        log_prob = log_prob - torch.log(1 - torch.tanh(action).pow(2) + 1e-7)
        action = action * self.action_bound
        return action, log_prob

class QValueNet(torch.nn.Module):
    def __init__(self, state_dim, hidden_dim, action_dim):
        super(QValueNet, self).__init__()
        self.fc1 = torch.nn.Linear(state_dim + action_dim, hidden_dim)
        self.fc2 = torch.nn.Linear(hidden_dim, 1)

    def forward(self, x, a):
        cat = torch.cat([x, a], dim=1)  # 拼接状态和动作
        x = F.relu(self.fc1(cat))
        return self.fc2(x)

class SAC:
    """处理连续动作的SAC算法"""
    def __init__(self, state_dim, hidden_dim, action_dim, action_bound,
                 actor_lr, critic_lr, alpha_lr, target_entropy, tau, gamma, device):
        self.actor = PolicyNet(state_dim, hidden_dim, action_dim,
                               action_bound).to(device)  # 策略网络
        # 第一个Q网络
        self.critic_1 = QValueNet(state_dim, hidden_dim, action_dim).to(device)
        # 第二个Q网络
        self.critic_2 = QValueNet(state_dim, hidden_dim, action_dim).to(device)
        self.target_critic_1 = QValueNet(state_dim, hidden_dim,
                                         action_dim).to(device)  # 第一个目标Q网络
        self.target_critic_2 = QValueNet(state_dim, hidden_dim,
                                         action_dim).to(device)  # 第二个目标Q网络
        # 令目标Q网络的初始参数和Q网络一样
        self.target_critic_1.load_state_dict(self.critic_1.state_dict())
        self.target_critic_2.load_state_dict(self.critic_2.state_dict())
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(),
                                                lr=actor_lr)
        self.critic_1_optimizer = torch.optim.Adam(self.critic_1.parameters(),
                                                   lr=critic_lr)
        self.critic_2_optimizer = torch.optim.Adam(self.critic_2.parameters(),
                                                   lr=critic_lr)
        # 使用alpha的log值,可以使训练结果比较稳定
        self.log_alpha = torch.tensor(np.log(0.01), dtype=torch.float)
        self.log_alpha.requires_grad = True  # 可以对alpha求梯度
        self.log_alpha_optimizer = torch.optim.Adam([self.log_alpha],
                                                    lr=alpha_lr)
        self.target_entropy = target_entropy  # 目标熵的大小
        self.gamma = gamma
        self.tau = tau
        self.device = device

    def take_action(self, state):
        state = torch.tensor([state], dtype=torch.float).to(self.device)
        action = self.actor(state)[0]
        return [action.item()]

    def calc_target(self, rewards, next_states, dones):  # 计算目标Q值
        next_actions, log_prob = self.actor(next_states)
        entropy = -log_prob
        q1_value = self.target_critic_1(next_states, next_actions)
        q2_value = self.target_critic_2(next_states, next_actions)
        next_value = torch.min(q1_value,
                               q2_value) + self.log_alpha.exp() * entropy
        td_target = rewards + self.gamma * next_value * (1 - dones)
        return td_target

    def soft_update(self, net, target_net):
        for param_target, param in zip(target_net.parameters(),
                                       net.parameters()):
            param_target.data.copy_(param_target.data * (1.0 - self.tau) +
                                    param.data * self.tau)

    def update(self, transition_dict):
        states = torch.tensor(transition_dict['states'],
                              dtype=torch.float).to(self.device)
        actions = torch.tensor(transition_dict['actions'],
                               dtype=torch.float).view(-1, 1).to(self.device)
        rewards = torch.tensor(transition_dict['rewards'],
                               dtype=torch.float).view(-1, 1).to(self.device)
        next_states = torch.tensor(transition_dict['next_states'],
                                   dtype=torch.float).to(self.device)
        dones = torch.tensor(transition_dict['dones'],
                             dtype=torch.float).view(-1, 1).to(self.device)
        rewards = (rewards + 8.0) / 8.0  # 对倒立摆环境的奖励进行重塑

        # 更新两个Q网络
        td_target = self.calc_target(rewards, next_states, dones)
        critic_1_loss = torch.mean(
            F.mse_loss(self.critic_1(states, actions), td_target.detach()))
        critic_2_loss = torch.mean(
            F.mse_loss(self.critic_2(states, actions), td_target.detach()))
        self.critic_1_optimizer.zero_grad()
        critic_1_loss.backward()
        self.critic_1_optimizer.step()
        self.critic_2_optimizer.zero_grad()
        critic_2_loss.backward()
        self.critic_2_optimizer.step()

        # 更新策略网络
        new_actions, log_prob = self.actor(states)
        entropy = -log_prob
        q1_value = self.critic_1(states, new_actions)
        q2_value = self.critic_2(states, new_actions)
        actor_loss = torch.mean(-self.log_alpha.exp() * entropy -
                                torch.min(q1_value, q2_value))
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

        # 更新alpha值
        alpha_loss = torch.mean(
            (entropy - self.target_entropy).detach() * self.log_alpha.exp())
        self.log_alpha_optimizer.zero_grad()
        alpha_loss.backward()
        self.log_alpha_optimizer.step()

        self.soft_update(self.critic_1, self.target_critic_1)
        self.soft_update(self.critic_2, self.target_critic_2)