# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Command-line script for simulating a fiber spectrograph.
"""
from __future__ import print_function, division

import warnings
import argparse

import numpy as np

import astropy.units as u

import specsim.config
import specsim.simulator

# This is a setup.py entry-point, not a standalone script.
# See http://astropy.readthedocs.io/en/latest/development/scripts.html

def main(args=None):
    # parse command-line arguments
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('-v', '--verbose', action='store_true',
        help='provide verbose output on progress')
    parser.add_argument('-c', '--config', default='test',
        help='name of the simulation configuration to use')
    parser.add_argument('--exposure-time', type=str, default='1000s',
        help='exposure time in to use (with units)')
    parser.add_argument('--sky-condition', type=str, default=None,
        help='sky condition to use (uses default if not set)')
    parser.add_argument('--airmass', type=float, default=1.,
        help='atmosphere airmass to use.')
    parser.add_argument('--moon-phase', type=float, default=None, metavar='P',
        help='moon phase between 0 (full) and 1 (new)')
    parser.add_argument('--moon-zenith', type=float, default=None, metavar='Z',
        help='zenith angle of the moon in degrees (>90 is below the horizon)')
    parser.add_argument('--moon-separation', type=float, default=None,
        metavar='S',
        help='opening angle between moon and this observation in degrees')
    parser.add_argument('--focal-x', type=str, default=None, metavar='X',
        help='Override x coordinate of source on focal plane (with units)')
    parser.add_argument('--focal-y', type=str, default=None, metavar='X',
        help='Override y coordinate of source on focal plane (with units)')
    parser.add_argument('--model', type=str, default=None,
        help='source fiberloss model to use (uses default if not set)')
    parser.add_argument('--z-in', type=float, default=None,
        help='redshift of input source data')
    parser.add_argument('--z-out', type=float, default=None,
        help='redshift that source should be transformed to')
    parser.add_argument('--filter', type=str, default=None,
        help='filter name to use for source flux normalization')
    parser.add_argument('--ab-mag', type=float, default=None,
        help='AB magnitude that source flux will be normalized to.')
    parser.add_argument('-o', '--output', type=str, default=None,
        help='optional output file name')
    parser.add_argument('--save-plot', type=str, default=None,
        help='save plot to the specified filename')
    parser.add_argument('--save-fiberloss', type=str, default=None,
        help='save fiberloss .fits and .ecsv files with this base filename')
    args = parser.parse_args(args)

    # Read the required configuration file.
    config = specsim.config.load_config(args.config)

    # Update configuration options from command-line options.
    config.verbose = args.verbose

    if args.sky_condition is not None:
        config.atmosphere.sky.condition = args.sky_condition
    config.atmosphere.airmass = args.airmass
    if (args.moon_phase is not None or args.moon_zenith is not None or
        args.moon_separation is not None):
        try:
            moon = config.atmosphere.moon.constants
        except AttributeError:
            print('Cannot set moon parameters when no moon defined in config.')
            return -1
        if args.moon_phase is not None:
            moon.moon_phase = args.moon_phase
        if args.moon_zenith is not None:
            moon.moon_zenith = '{0:f}deg'.format(args.moon_zenith)
        if args.moon_separation is not None:
            moon.separation_angle = '{0:f}deg'.format(args.moon_separation)

    if args.model is not None:
        config.source.type = args.model
    config.source.z_in = args.z_in
    config.source.z_out = args.z_out
    config.source.filter_name = args.filter
    config.source.ab_magnitude_out = args.ab_mag

    # Initialize the simulator.
    try:
        simulator = specsim.simulator.Simulator(config, verbose=args.verbose)
    except RuntimeError as e:
        print(e)
        return -1

    # Set parameters after configuration.
    try:
        simulator.observation.exposure_time = specsim.config.parse_quantity(
            args.exposure_time, u.s)
        if args.focal_x is not None:
            if args.focal_y is None:
                print('Must set both focal-x and focal-y.')
                return -1
            else:
                focal_x = specsim.config.parse_quantity(args.focal_x, u.mm)
                focal_y = specsim.config.parse_quantity(args.focal_y, u.mm)
                simulator.source.focal_xy = focal_x, focal_y
    except ValueError as e:
        print(e)
        return -1

    # Perform the simulation.
    simulator.simulate(save_fiberloss=args.save_fiberloss)

    # Summarize the results.
    print('Source at focal plane (x, y) = ({0:.1f}, {1:.1f}).'
          .format(simulator.focal_x[0], simulator.focal_y[0]))
    print('Observing airmass is {0:.3f}.'.format(simulator.atmosphere.airmass))
    for output in simulator.camera_output:
        camera_name = output.meta['name']
        pixel_size = output.meta['pixel_size']
        snr = (
            output['num_source_electrons'][:, 0] /
            np.sqrt(output['variance_electrons'][:, 0]))
        print('Median SNR in {0} camera = {1:.3f} / {2}'
              .format(camera_name, np.median(snr), pixel_size))

    # Save the results, if requested.
    if args.output:
        try:
            simulator.save(args.output)
        except Exception as e:
            print(e)
            return -1
        if args.verbose:
            print('Saved outputs to {0}'.format(args.output))

    # Plot the results if requested.
    if args.save_plot:
        # Defer these imports until now so that matplotlib is only required
        # if plots are requested.
        import matplotlib
        # Use a backend with minimal requirements (X11, etc).
        matplotlib.use('Agg')
        import matplotlib.pyplot as plt

        simulator.plot()

        with warnings.catch_warnings():
            # Silence expected matplotlib warnings.
            warnings.simplefilter('ignore', category=FutureWarning)
            plt.savefig(args.save_plot, facecolor='white', edgecolor='none')

        if args.verbose:
            print('Saved generated plot to {0}'.format(args.save_plot))


import ast
import community
import datetime
import lightgbm as lgb
import math
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import pandas as pd
import pickle
import plotly.express as px
import os
from sklearn.metrics import accuracy_score, confusion_matrix, roc_auc_score
from tqdm import tqdm
from make_boruta import *

class Zeus:
    """
    Class criada para construir modelo
    """

    def __init__(self, termo, user, treino_id, test_id):
        """
        Metodo construtor, aqui serão armazenadas as informações padrões do modelo:
        - User
        - run_id da base
        """
        self.term = termo
        self.data = datetime.date.today()
        self.user = str(user).upper()
        self.path_user = ''
        self.treino_id = treino_id
        self.test_id = test_id
        self.var_treino = ''
        self.var_teste = ''
        self.filtro_local = False
        self.filtro_data = False
        self.local = ''
        self.data_start = ''
        self.data_end = ''
        self.filtro = ''
        self.random_state = 101
        self.base_sintetica = ''
        self.data_active = False
        self.data_local = False
        self.mm = ''
        self.ids = ''
        self.train = ''
        self.numero_de_amostras_sinteticas_para_criar = ''
        self.porcentagem_para_criacao_de_amostras = ''
        self.df_cluster = ''
        self.clusters = ''
        self.var_teste_original = ''
        self.pega_variaveis()
        self.agregado = ''
        self.df_agregado = ''
        self.informacoes = ''
        self.sentimento = ''
        self.data_df = ''
        self.load_df = ''

    def pega_path_user(self):
        """
        Metodo que pega o path de acordo com o usuario que inicializou a class
        """
        os.chdir(os.path.dirname(
            r'C:\Users\wilgn\Desktop\Faculdade\3° Semestre\Insper Data\Projeto\Git projeto\Data_BCG_News\Model\\'))
        path_atual = os.getcwd()
        #print(os.listdir())
        if self.user == 'WILGNER':
            path_aux_funcs = path_atual.replace('Model', r'aux_funcs\\')
        else:
            path_aux_funcs = path_atual.replace('Model', r'aux_funcs/')

        os.chdir(os.path.dirname(path_aux_funcs))
        #print(os.listdir())
        with open('set_path.py', 'r') as arquivo_path:
            ler_arquivo = arquivo_path.read()
            dicionario = ast.literal_eval(ler_arquivo)

            lista_users = list(dicionario.keys())

            if self.user in lista_users:
                print('USUARIO VALIDO !')
                self.path_user = dicionario[self.user]
            else:
                raise TypeError(
                    'O USUARIO SELECIONADO NÃO TEM UM ENDEREÇO VALIDO CADASTRADO')

        os.chdir(os.path.dirname(
            r'C:\Users\wilgn\Desktop\Faculdade\3° Semestre\Insper Data\Projeto\Git projeto\Data_BCG_News\Model\\'))
        #print(os.listdir())
        # arquivo_path.close()

    def valida_acesso_path_user(self):
        self.pega_path_user()
        try:
            os.path.exists(self.path_user)
            print('PATH VALIDO PARA ACESSO')

        except:
            raise TypeError('IMPOSSIVEL ACESSAR O PATH')

    def pega_variaveis(self, teste=False, load_df=False, file_name_load=False):

        if teste:
            path_name = f'{self.path_user}Variables/{self.term}/{self.test_id}.parquet'
            os.chdir(os.path.dirname(path_name))
            self.var_teste = pd.read_parquet(os.path.basename(path_name))

        elif load_df:

            path_name = f'{self.path_user}Model/{file_name_load}'
            os.chdir(os.path.dirname(path_name))

            self.load_df = pd.read_parquet(os.path.basename(path_name))

        else:
            self.valida_acesso_path_user()
            path_name = f'{self.path_user}Variables/{self.term}/{self.treino_id}.parquet'
            os.chdir(os.path.dirname(path_name))

            self.var_treino = pd.read_parquet(os.path.basename(path_name))

    def seleciona_filtros(self, local=False, data_start=False, data_end=False):
        """
        Isinstance verifica se houve uma solicitação de filtro em alguma das variaveis
        """
        estado_local = isinstance(local, bool)
        estado_data_start = isinstance(data_start, bool)

        if not estado_local:
            self.local = local
            self.filtro_local = True
        if not estado_data_start:
            self.data_start = data_start
            self.data_end = data_end
            self.filtro_data = True

    def construir_filtro(self, teste=False):
        self.filtro = ''
        if not teste:
            if self.filtro_local and self.filtro_data:
                self.filtro = (self.var_treino.sigla == self.local.upper())
                self.data_local = True
                self.data_active = True

            elif self.filtro_data and not self.filtro_local:
                self.data_active = True

            elif self.filtro_local and not self.filtro_data:
                self.filtro = (self.var_treino.sigla == self.local.upper())
        else:
            if self.filtro_local and self.filtro_data:
                self.filtro = (self.var_teste.sigla == self.local.upper())
                self.data_local = True
                self.data_active = True

            elif self.filtro_data and not self.filtro_local:
                self.data_active = True

            elif self.filtro_local and not self.filtro_data:
                self.filtro = (self.var_teste.sigla == self.local.upper())

    def filtrar_treino(self, local=False, data_start=False, data_end=False):
        self.seleciona_filtros(
            local=local, data_start=data_start, data_end=data_end)
        self.construir_filtro()
        self.var_treino.data = pd.to_datetime(self.var_treino.data)

        if self.data_active and self.data_local:
            self.var_treino = self.var_treino[self.filtro]
            self.var_treino = self.var_treino[(self.var_treino.data > self.data_start) & (
                    self.var_treino.data < self.data_end)]
        elif self.data_active and not self.data_local:
            self.var_treino = self.var_treino[
                (self.var_treino.data > self.data_start) & (self.var_treino.data < self.data_end)]
        else:
            self.var_treino = self.var_treino[self.filtro]

    def filtrar_teste(self, local=False, data_start=False, data_end=False):
        self.seleciona_filtros(
            local=local, data_start=data_start, data_end=data_end)
        self.construir_filtro(teste=True)
        self.var_teste.data = pd.to_datetime(self.var_teste.data)

        if self.data_active and self.data_local:
            self.var_teste = self.var_teste[self.filtro]
            self.var_teste = self.var_teste[
                (self.var_teste.data > self.data_start) & (self.var_teste.data < self.data_end)]
        elif self.data_active and not self.data_local:
            self.var_teste = self.var_teste[
                (self.var_teste.data > self.data_start) & (self.var_teste.data < self.data_end)]
        else:
            self.var_teste = self.var_teste[self.filtro]

    def criar_base_sintetica(self, numero_de_amostras=3, porcentagem_para_criacao=.25):
        bases_sinteticas = []
        self.numero_de_amostras_sinteticas_para_criar = numero_de_amostras
        self.porcentagem_para_criacao_de_amostras = porcentagem_para_criacao
        colunas_pro_drop = ['unique_identifier', 'sigla', 'data']
        for i in range(numero_de_amostras):
            unique_identifier = self.var_treino['unique_identifier']
            df_com_drop = self.var_treino.drop(columns=colunas_pro_drop)
            df_com_colunas_sorteadas = df_com_drop.sample(
                frac=porcentagem_para_criacao, replace=True, random_state=self.random_state, axis=1)
            amostra = pd.concat(
                [unique_identifier, df_com_colunas_sorteadas], axis=1)
            amostra_sintetica = pd.DataFrame()
            amostra = amostra.loc[:, ~amostra.columns.duplicated()]

            for coluna in amostra.columns.tolist():
                amostra_sintetica[coluna] = amostra[coluna].sample(frac=1, replace=True,
                                                                   random_state=self.random_state).tolist()
            amostra_sintetica['label'] = 1
            amostra['label'] = 0

            amostra_concluida = pd.concat([amostra, amostra_sintetica])
            amostra_concluida.reset_index(inplace=True, drop=True)
            bases_sinteticas.append(amostra_concluida)

        self.base_sintetica = bases_sinteticas

    def treina_lightGBM(self, boruta_percs=[10], thr_bor_good=.5, thr_bor_ok=.9):
        numero_de_amostras = len(self.base_sintetica)
        x_list = []
        y_list = []
        col_lists = []
        model_list = []
        trained_models = []

        dfs = self.base_sintetica
        for i in range(numero_de_amostras):
            numero_de_colunas = dfs[i].shape[1]

            self.Y = dfs[i]['label']
            self.X = dfs[i].drop(columns=['unique_identifier', 'label'])
            self.take_out_cols_0 = []
            self.take_out_cols = []
            self.full_cols = self.X.columns.tolist()
            self.thr_bor_good = thr_bor_good
            self.thr_bor_ok = thr_bor_ok
            self.boruta_percs = boruta_percs
            self.boruta_res = boruta_select(
                X_df=self.X[[
                    col for col in self.full_cols if col not in self.take_out_cols]],
                Y=self.Y, perc_list=self.boruta_percs, allowed_perc_good=self.thr_bor_good,
                allowed_perc_med=self.thr_bor_ok)

            self.take_out_cols_irrelevant = self.boruta_res[0].loc[~self.boruta_res[0]['use']].index.tolist(
            )
            self.take_out_cols += self.take_out_cols_irrelevant

            self.use_cols = self.X[[col for col in self.X.columns.tolist(
            ) if col not in self.take_out_cols]].columns.tolist()

            y_list += [dfs[i]['label'].values]
            x_list += [dfs[i].drop(columns=['unique_identifier', 'label'])]
            col_lists += [self.use_cols]
            model_list += [{'type': 'LGBM',
                            'params': {'num_leaves': 25, 'n_estimators': 300, 'boosting_type': 'rf',
                                       'bagging_fraction': .8, 'bagging_freq': 1, 'random_state': self.random_state}}]

        # Treinando modelo
        for (model, x, y, cols) in zip(model_list, x_list, y_list, col_lists):
            X = x[cols]
            Y = y
            if model['type'] == 'LGBM':
                model_to_train = lgb.LGBMClassifier(**model['params'])

            trained_models += [model_to_train.fit(X=X.values, y=Y)]

        self.models = trained_models
        self.rf_models = self.models
        self.col_lists = col_lists

    def coleta_folhas(self, porcentagem_do_sample=0.1):

        self.df_random = self.var_treino.sample(
            frac=porcentagem_do_sample, replace=True, random_state=self.random_state, axis=0).copy()
        print(self.df_random.shape)
        # frame_list = []
        model_c = 0
        self.mm = set()
        self.ids = self.df_random['unique_identifier'].tolist()
        print('start with list values')

        # Pegando o resultado das folhas do model
        for (model, cols) in zip(self.rf_models, self.col_lists):
            if cols == 'label':
                continue
            else:
                raw_leafs = model.predict(
                    self.df_random[cols].values, pred_leaf=True)
                # return raw_leafs
                if model_c == 0:
                    full_leafs = raw_leafs
                else:
                    full_leafs = np.concatenate(
                        (full_leafs, raw_leafs), axis=1)

                model_c += 1
        self.raw = raw_leafs

    def criando_matriz_de_similaridade(self, porcentagem_do_sample=0.1):
        self.porcentagem_para_matriz = porcentagem_do_sample
        self.coleta_folhas(porcentagem_do_sample=porcentagem_do_sample)
        print('CRIANDO EDGES')
        edges = []
        # Criando matriz de similaridade
        for cc1, i in tqdm(enumerate(self.raw), 'FOLHAS:'):
            if cc1 % 100 == 0:
                print(cc1, datetime.datetime.now())
            for cc2_, j in enumerate(self.raw[cc1 + 1:]):
                cc2 = cc2_ + cc1 + 1
                if (cc1, cc2) not in self.mm and (cc2, cc1) not in self.mm:
                    leaf_count = sum(i == j)
                    # TODO: Fix similarity matrix with the square root
                    edges += [(self.ids[cc1], self.ids[cc2],
                               math.sqrt(leaf_count / len(self.raw[0])))]
                    self.mm.add((cc1, cc2))
        print('done with list values')

        # YOU ARE HERE
        G = nx.Graph()

        G.add_weighted_edges_from(edges)
        self.G = G

    def rodando_louvain(self, porcentagem_do_sample):
        self.criando_matriz_de_similaridade(
            porcentagem_do_sample=porcentagem_do_sample)
        self.clusters = (community.best_partition(
            self.G, weight='weight', randomize=True))

    def desenha_cluster_no_edges(self):
        plt.figure(figsize=(12, 8), dpi=150)
        plt.title('Louvain Tets', fontsize=20, loc='left', pad=15)
        self.pos = nx.spring_layout(self.G)
        nx.draw_networkx_nodes(self.G, self.pos, self.clusters.keys(), node_size=150,
                               node_color=list(self.clusters.values()))
        plt.show()

    def classifica_agrupamento(self, boruta_percs=[10], thr_bor_good=.5, thr_bor_ok=.9, take_out_cols=False):

        self.df_cluster = pd.DataFrame({'Rotulo': self.clusters.keys(),
                                        'Label': self.clusters.values()})
        print(f'Tamanho dos dados de cluster {self.df_cluster.shape}')
        self.train = self.var_treino[self.var_treino['unique_identifier'].isin(
            self.df_cluster.Rotulo.values.tolist())]
        self.train['label'] = self.df_cluster.Label.values.tolist()
        self.train.reset_index(drop=True)

        colunas_pro_drop = ['unique_identifier', 'sigla', 'data', 'artigo_original']
        self.var_teste_original = self.var_teste
        self.sentimento = self.var_teste_original['sentimento']
        self.data_df = self.var_teste_original['data']
        self.var_teste = self.var_teste.drop(columns=colunas_pro_drop)
        self.train = self.train.drop(columns=['sigla', 'data', 'artigo_original'])
        self.var_teste.reset_index(drop=True)
        self.train.reset_index(drop=True)

        print(f'Tamanho dos dados de treinamento {self.train.shape}')
        print(f'Tamanho dos dados de teste {self.var_teste.shape}')

        self.x_list = []
        self.y_list = []
        self.col_lists = []
        model_list = []
        trained_models = 0
        self.Y = self.train['label']
        self.X = self.train.drop(columns=['unique_identifier', 'label'])
        self.take_out_cols_0 = []
        self.take_out_cols = []
        self.full_cols = self.X.columns.tolist()
        self.thr_bor_good = thr_bor_good
        self.thr_bor_ok = thr_bor_ok
        self.boruta_percs = boruta_percs
        self.boruta_res = boruta_select(X_df=self.X[[col for col in self.full_cols if col not in self.take_out_cols]],
                                        Y=self.Y, perc_list=self.boruta_percs, allowed_perc_good=self.thr_bor_good,
                                        allowed_perc_med=self.thr_bor_ok)

        self.take_out_cols_irrelevant = self.boruta_res[0].loc[~self.boruta_res[0]['use']].index.tolist(
        )
        self.take_out_cols += self.take_out_cols_irrelevant

        self.use_cols = self.X[[col for col in self.X.columns.tolist(
        ) if col not in self.take_out_cols]].columns.tolist()

        if len(self.use_cols) < 1:
            self.use_cols = self.X.columns.tolist()

        self.y_list += [self.train['label'].values]
        self.x_list += [self.train.drop(columns=['unique_identifier', 'label'])]
        self.col_lists += [self.use_cols]
        model_list += [{'type': 'LGBM',
                        'params': {'num_leaves': 30, 'n_estimators': 500, 'boosting_type': 'rf',
                                   'bagging_fraction': .8, 'bagging_freq': 1, 'random_state': self.random_state}}]

        for (model, x, y, cols) in zip(model_list, self.x_list, self.y_list, self.col_lists):
            X = x[cols]
            print(X.shape)
            Y = y
            print(Y.shape)
            if model['type'] == 'LGBM':
                model_to_train = lgb.LGBMClassifier(**model['params'])

            trained_models = model_to_train.fit(X=X.values, y=Y)

        self.models = trained_models
        self.previsão = trained_models.predict(
            self.var_teste[self.col_lists[0]])
        self.resultado = self.previsão
        self.var_teste['label'] = self.resultado

        self.faz_agregacao()

        print('FREQUENCIA CLUSTER')
        print(self.df_cluster.Label.value_counts(sort=False))
        print('********************')
        print('FREQUENCIA CLASSIFICADO')
        print(self.var_teste.label.value_counts(sort=False))

    def plota_palavras_maiores(self, numero):
        for i in range(len(self.var_teste.label.unique())):
            df_data = pd.DataFrame({'word': self.var_teste[self.var_teste.label == i].drop(
                columns=['label']).sum(axis=0).nlargest(numero).index.tolist(),
                                    'value': self.var_teste[self.var_teste.label == i].drop(
                                        columns=['label']).sum(axis=0).nlargest(
                                        numero).values.tolist()})

            fig = px.bar(df_data, x='word', y='value', color='value', color_continuous_scale='Blues')
            fig.show()

    def salva_parametros(self):
        self.informacoes = {
            'user': self.user,
            'data': self.data,
            'run_id_treino': self.treino_id,
            'run_id_teste': self.test_id,
            'path_user': self.path_user,
            'filtro_nome': self.term,
            'filtro_data': self.data,
            'filtro_local': self.local,
            'numero_de_amostras_bases_sinteticas': self.numero_de_amostras_sinteticas_para_criar,
            'porcentagem_para_criacao_de_amostras': self.porcentagem_para_criacao_de_amostras,
            'porcentagem_para_matriz': self.porcentagem_para_matriz

        }

    def faz_agregacao(self):
        # Agrega os resultados
        self.var_teste_original['label'] = self.var_teste['label']
        self.agregado = self.var_teste_original[['unique_identifier', 'sigla', 'data', 'label']]
        self.df_agregado = pd.crosstab(self.agregado.sigla, self.agregado.label, normalize='index')

        return self.df_agregado


#!/usr/bin python3

import click

import numpy as np
import json

import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable

plt.rcParams["font.sans-serif"] = ["Arial", "sans-serif"]
plt.rcParams.update({"font.size": 12})

@click.command()
@click.option(
    "--group",
    default="1Lbb",
    type=click.Choice(
        [
            "1Lbb",
            "2L0J",
            "compressed",
            "3Loffshell",
            "stop1L",
            "3LRJR",
            "directstaus",
            "samesign",
            "sbottom",
        ]
    ),
)
@click.option(
    "--harvest",
    "-h",
    default=[],
    multiple=True,
    help="Harvest to plot",
)
@click.option(
    "--harvest-label",
    "-l",
    default=[],
    multiple=True,
    help="Label of the harvest",
)
@click.option("--width", type=int, default=1600, help="Width of canvas")
@click.option("--height", type=int, default=1200, help="Height of canvas")
@click.option(
    "--comEnergy", type=str, default="13 TeV", help="Center-of-mass energy for plot"
)
@click.option(
    "--lumi-Label", type=str, default="#sqrt{{s}} = {comEnergy}, {luminosity} fb^{{-1}}"
)
@click.option(
    "--xlabel",
    type=str,
    default="m(#tilde{#chi}^{#pm}_{1})/m(#tilde{#chi}^{0}_{2}) [GeV]",
)
@click.option(
    "--ylabel",
    type=str,
    default="m(#tilde{#chi}^{0}_{1}) [GeV]",
)
@click.option("--xmin", type=float, default=20, help="x-axis minimum")
@click.option("--xmax", type=float, default=1000, help="x-axis maximum")
@click.option("--ymin", type=float, default=20, help="y-axis minimum")
@click.option("--ymax", type=float, default=700, help="y-axis maximum")
@click.option(
    "--is-internal/--not-internal", default=False, help="Is this ATLAS Internal?"
)
@click.option(
    "--is-preliminary/--not-preliminary",
    default=False,
    help="Is this ATLAS Preliminary?",
)
@click.option(
    "--labels-left",
    type=float,
    default=0.2,
    help="x-position of labels drawn on the left side of the plot",
)
@click.option(
    "--labels-top",
    type=float,
    default=0.86,
    help="y-position of labels drawn on the top-left of the plot",
)
@click.option(
    "--log/--no-log",
    default=False,
    help="Draw logarithmic axes",
)
@click.option(
    "--draw-ATLAS/--no-ATLAS",
    default=False,
    help="Draw ATLAS label",
)
@click.option(
    "--process-label",
    type=str,
    default="pp #rightarrow #tilde{#chi}^{0}_{2} #tilde{#chi}^{#pm}_{1} (Wino) production ; #tilde{#chi}^{0}_{2} #rightarrow Z #tilde{#chi}^{0}_{1},#tilde{#chi}^{#pm}_{1} #rightarrow W #tilde{#chi}^{0}_{1} ; ",
    help="Label for the physics process drawn",
)
@click.option("--luminosity", type=float, default=139000)
def main(
    group,
    harvest,
    harvest_label,
    width,
    height,
    comenergy,
    lumi_label,
    xlabel,
    ylabel,
    xmin,
    xmax,
    ymin,
    ymax,
    is_internal,
    is_preliminary,
    labels_left,
    labels_top,
    log,
    process_label,
    draw_atlas,
    luminosity,
):

    colors = ["#e67e22", "#07a7ec", "#16a085", "#8e44ad", "#e74c3c"]
    i, yoffset = 0, 0

    g = {}

    for h in harvest:
        with open(f"analyses/{group}/harvests/{h}") as f:
            ds = json.load(f)

            for d in ds:
                if not (d["m1"], d["m2"]) in g.keys():
                    g[(d["m1"], d["m2"])] = [(d["CLsexp"], d["CLs"])]
                else:
                    g[(d["m1"], d["m2"])].append((d["CLsexp"], d["CLs"]))

    signal_grid = {}
    for k, i in g.items():
        if len(i) > 1:
            signal_grid[k] = i

    exp_x = np.array([i[0][0] for k, i in signal_grid.items()])
    exp_y = np.array([i[1][0] for k, i in signal_grid.items()])

    obs_x = np.array([i[0][1] for k, i in signal_grid.items()])
    obs_y = np.array([i[1][1] for k, i in signal_grid.items()])

    fig, ax = plt.subplots()
    fig.set_size_inches(5, 4)

    plt.locator_params(axis="x", nbins=8)

    ax.scatter(
        exp_x,
        exp_y,
        s=20,
        facecolors="none",
        label="Expected CLs",
        edgecolors="#686de0",
    )
    ax.scatter(
        obs_x,
        obs_y,
        s=20,
        facecolors="none",
        label="Observed CLs",
        edgecolors="#10ac84",
    )

    ax.plot((0, 1), linestyle="--", linewidth=1.2, c="#444444")

    ax.legend(
        loc="upper left", bbox_to_anchor=(-0.01, 0.90), frameon=False, fontsize=12
    )

    if log:
        tag = "log"
        ax.set_ylim([0.0000000001, 1])
        ax.set_xlim([0.0000000001, 1])
        ax.set_xscale("log")
        ax.set_yscale("log")
    else:
        tag = "lin"
        ax.set_ylim([0, 1])
        ax.set_xlim([0, 1])

    plt.ylabel("Full LH CLs", ha="right", y=1.0, fontsize=12)
    plt.xlabel("Simplified LH CLs", ha="right", x=1.0, fontsize=12)
    fig.tight_layout()

    ax.minorticks_on()
    plt.tick_params(direction="in", which="both", length=4)

    plt.text(
        0.04,
        0.92,
        r"$\sqrt{s}$=13 TeV, 139 fb$^{-1}$",
        transform=plt.gca().transAxes,
        alpha=1.0,
        fontsize=12,
    )
    plt.text(0.04, 0.86, f"", transform=plt.gca().transAxes, alpha=1.0, fontsize=12)

    plt.savefig(f"analyses/{group}/plots/cls_scatter_{tag}.pdf")

if __name__ == "__main__":
    # pylint: disable=no-value-for-parameter
    main()


# ------------------------------------------------------------------------
# Copyright (c) Hitachi, Ltd. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
# Modified from DETR (https://github.com/facebookresearch/detr)
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
# ------------------------------------------------------------------------
"""
Train and eval functions used in main.py
"""
import math
import os
import sys
from typing import Iterable
import numpy as np
import copy
import itertools

import torch

import util.misc as utils
from datasets.coco_eval import CocoEvaluator
from datasets.panoptic_eval import PanopticEvaluator
from datasets.hico_eval import HICOEvaluator
from datasets.vcoco_eval import VCOCOEvaluator

from util.misc import NestedTensor

def train_one_epoch(model: torch.nn.Module, criterion: torch.nn.Module,criterion_detr: torch.nn.Module,
                    criterion_verbs:torch.nn.Module,
                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
                    device: torch.device, epoch: int, args):
    max_norm = args.clip_max_norm
    model.train()
    criterion.train()
    criterion_detr.train()
    metric_logger = utils.MetricLogger(delimiter="  ")
    metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
    if hasattr(criterion, 'loss_labels'):
        metric_logger.add_meter('class_error', utils.SmoothedValue(window_size=1, fmt='{value:.2f}'))
    else:
        metric_logger.add_meter('obj_class_error', utils.SmoothedValue(window_size=1, fmt='{value:.2f}'))
    header = 'Epoch: [{}]'.format(epoch)
    print_freq = 100
    for samples, targets in metric_logger.log_every(data_loader, print_freq, header):
        samples = samples.to(device)
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

        outs = model(samples)    # pred_logits(100,classes+1), pred_boxes(100,4), aux_outputs(5 groups)
        output_detr = outs['detr']
        outputs = outs['hoi']
        verbs = outs['verbs']
        # cnn verbs loss
        losses_verbs = criterion_verbs(verbs,targets)
        # hoi
        loss_dict = criterion(outputs, targets)
        weight_dict = criterion.weight_dict
        losses_hoi = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)
        # detr
        loss_dict_detr = criterion_detr(output_detr, targets)
        weight_dict_detr = criterion_detr.weight_dict
        losses_detr = sum(loss_dict_detr[k] * weight_dict_detr[k] for k in loss_dict_detr.keys() if k in weight_dict_detr)
        # all
        losses = args.loss_hoi_weight*losses_hoi + args.loss_detr_weight*losses_detr + args.loss_verbs_weight*losses_verbs
        # reduce losses over all GPUs for logging purposes
        loss_dict_reduced = utils.reduce_dict(loss_dict)
        loss_dict_reduced_unscaled = {f'{k}_unscaled': v
                                      for k, v in loss_dict_reduced.items()}
        loss_dict_reduced_scaled = {f'{k}_scaled': v * weight_dict[k]
                                    for k, v in loss_dict_reduced.items() if k in weight_dict}
        losses_reduced_scaled = sum(loss_dict_reduced_scaled.values())

        loss_value = losses_reduced_scaled.item()
        # detr log
        loss_dict_reduced_detr = utils.reduce_dict(loss_dict_detr)
        loss_dict_reduced_unscaled_detr = {f'{k}_unscaled_detr': v
                                      for k, v in loss_dict_reduced_detr.items()}
        loss_dict_reduced_scaled_detr = {f'{k}_scaled_detr': v * weight_dict_detr[k]
                                    for k, v in loss_dict_reduced_detr.items() if k in weight_dict_detr}
        losses_reduced_scaled_detr = sum(loss_dict_reduced_scaled_detr.values())
        loss_value_detr = losses_reduced_scaled_detr.item()

        if not math.isfinite(loss_value):
            print("Loss is {}, stopping training".format(loss_value))
            print(loss_dict_reduced)
            sys.exit(1)

        optimizer.zero_grad()
        losses.backward()
        if max_norm > 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
        optimizer.step()

        metric_logger.update(loss=loss_value, **loss_dict_reduced_scaled, **loss_dict_reduced_unscaled)
        metric_logger.update(loss_detr=loss_value_detr, **loss_dict_reduced_scaled_detr, **loss_dict_reduced_unscaled_detr)
        metric_logger.update(loss_cnn_verbs = losses_verbs)

        if hasattr(criterion, 'loss_labels'):
            metric_logger.update(class_error=loss_dict_reduced['class_error'])
        else:
            metric_logger.update(obj_class_error=loss_dict_reduced['obj_class_error'])
        metric_logger.update(lr=optimizer.param_groups[0]["lr"])
    # gather the stats from all processes
    metric_logger.synchronize_between_processes()
    print("Averaged stats:", metric_logger)
    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}

def train_one_epoch_w_accum(model: torch.nn.Module, criterion: torch.nn.Module,criterion_detr: torch.nn.Module,
                    criterion_verbs:torch.nn.Module,
                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
                    device: torch.device, epoch: int, args):
    max_norm = args.clip_max_norm
    model.train()
    criterion.train()
    criterion_detr.train()
    metric_logger = utils.MetricLogger(delimiter="  ")
    metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
    if hasattr(criterion, 'loss_labels'):
        metric_logger.add_meter('class_error', utils.SmoothedValue(window_size=1, fmt='{value:.2f}'))
    else:
        metric_logger.add_meter('obj_class_error', utils.SmoothedValue(window_size=1, fmt='{value:.2f}'))
    header = 'Epoch: [{}]'.format(epoch)
    print_freq = 100
    for samples, targets in metric_logger.log_every(data_loader, print_freq, header):
        samples = samples.to(device)
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

        B = samples.tensors.shape[0]
        b = B // 2
        loss_dicts = list()
        optimizer.zero_grad()
        for i in range(2):
            b_samples = NestedTensor(samples.tensors[i*b:(i+1)*b], samples.mask[i*b:(i+1)*b])
            # b_samples = NestedTensor(samples.tendors[i*b:(i+1)*b], samples.tensors[i*b:(i+1)*b])
            b_targets = []
            for i in range(b*i, b* (i+1)):
                b_targets.append(targets[i])
            # b_targets = [{k: v for k, v in t.items()} for t in targets]
            b_outs = model(b_samples)
            b_output_detr = b_outs['detr']
            b_outputs = b_outs['hoi']
            b_verbs = b_outs['verbs']
            # cnn verbs loss
            b_losses_verbs = criterion_verbs(b_verbs, b_targets)
            # hoi
            b_loss_dict = criterion(b_outputs, b_targets)
            weight_dict = criterion.weight_dict
            b_losses_hoi = sum(b_loss_dict[k] * weight_dict[k] for k in b_loss_dict.keys() if k in weight_dict)
            # detr
            b_loss_dict_detr = criterion_detr(b_output_detr, b_targets)
            weight_dict_detr = criterion_detr.weight_dict
            b_losses_detr = sum(b_loss_dict_detr[k] * weight_dict_detr[k] for k in b_loss_dict_detr.keys() if k in weight_dict_detr)

            # all
            b_losses = (args.loss_hoi_weight * b_losses_hoi + args.loss_detr_weight * b_losses_detr + args.loss_verbs_weight * b_losses_verbs) / 2

            if torch.isnan(b_losses).sum() != 0:
                print('sum_losses',b_losses)
                print('b_losses_hoi',b_losses_hoi)
                print('b_losses_detr',b_losses_detr)
                print('b_losses_verbs',b_losses_verbs)
                sys.exit(1)
            b_losses.backward()
            loss_dicts.append(b_loss_dict)

        loss_dict = {k: (loss_dicts[0][k] + loss_dicts[1][k]) / 2 for k in loss_dicts[0].keys()}

        # reduce losses over all GPUs for logging purposes
        loss_dict_reduced = utils.reduce_dict(loss_dict)
        loss_dict_reduced_unscaled = {f'{k}_unscaled': v
                                      for k, v in loss_dict_reduced.items()}
        loss_dict_reduced_scaled = {f'{k}_scaled': v * weight_dict[k]
                                    for k, v in loss_dict_reduced.items() if k in weight_dict}
        losses_reduced_scaled = sum(loss_dict_reduced_scaled.values())

        loss_value = losses_reduced_scaled.item()
        # detr log
        loss_dict_reduced_detr = utils.reduce_dict(b_loss_dict_detr)
        loss_dict_reduced_unscaled_detr = {f'{k}_unscaled_detr': v
                                      for k, v in loss_dict_reduced_detr.items()}
        loss_dict_reduced_scaled_detr = {f'{k}_scaled_detr': v * weight_dict_detr[k]
                                    for k, v in loss_dict_reduced_detr.items() if k in weight_dict_detr}
        losses_reduced_scaled_detr = sum(loss_dict_reduced_scaled_detr.values())
        loss_value_detr = losses_reduced_scaled_detr.item()

        if not (math.isfinite(loss_value) and math.isfinite(loss_value_detr) and math.isfinite(b_losses_verbs)):
            print("Loss is {}, stopping training".format(loss_value))
            print("DETR Loss is {}, stopping training".format(loss_value_detr))
            print("Verb Loss is {}, stopping training".format(b_losses_verbs))
            print('HOI',loss_dict_reduced)
            print('DETR',loss_dict_reduced_detr)
            sys.exit(1)

        if max_norm > 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

        optimizer.step()

        metric_logger.update(loss=loss_value, **loss_dict_reduced_scaled, **loss_dict_reduced_unscaled)
        metric_logger.update(loss_detr=loss_value_detr, **loss_dict_reduced_scaled_detr, **loss_dict_reduced_unscaled_detr)
        metric_logger.update(loss_cnn_verbs = b_losses_verbs)

        if hasattr(criterion, 'loss_labels'):
            metric_logger.update(class_error=loss_dict_reduced['class_error'])
        else:
            metric_logger.update(obj_class_error=loss_dict_reduced['obj_class_error'])
        metric_logger.update(lr=optimizer.param_groups[0]["lr"])
    # gather the stats from all processes
    metric_logger.synchronize_between_processes()
    print("Averaged stats:", metric_logger)
    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}

@torch.no_grad()
def evaluate(model, criterion, postprocessors, data_loader, base_ds, device, output_dir):
    model.eval()
    criterion.eval()

    metric_logger = utils.MetricLogger(delimiter="  ")
    metric_logger.add_meter('class_error', utils.SmoothedValue(window_size=1, fmt='{value:.2f}'))
    header = 'Test:'

    iou_types = tuple(k for k in ('segm', 'bbox') if k in postprocessors.keys())
    coco_evaluator = CocoEvaluator(base_ds, iou_types)
    # coco_evaluator.coco_eval[iou_types[0]].params.iouThrs = [0, 0.1, 0.5, 0.75]

    panoptic_evaluator = None
    if 'panoptic' in postprocessors.keys():
        panoptic_evaluator = PanopticEvaluator(
            data_loader.dataset.ann_file,
            data_loader.dataset.ann_folder,
            output_dir=os.path.join(output_dir, "panoptic_eval"),
        )

    for samples, targets in metric_logger.log_every(data_loader, 100, header):
        samples = samples.to(device)
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

        outputs = model(samples)

        loss_dict = criterion(outputs, targets)
        weight_dict = criterion.weight_dict

        # reduce losses over all GPUs for logging purposes
        loss_dict_reduced = utils.reduce_dict(loss_dict)
        loss_dict_reduced_scaled = {k: v * weight_dict[k]
                                    for k, v in loss_dict_reduced.items() if k in weight_dict}
        loss_dict_reduced_unscaled = {f'{k}_unscaled': v
                                      for k, v in loss_dict_reduced.items()}
        metric_logger.update(loss=sum(loss_dict_reduced_scaled.values()),
                             **loss_dict_reduced_scaled,
                             **loss_dict_reduced_unscaled)
        metric_logger.update(class_error=loss_dict_reduced['class_error'])

        orig_target_sizes = torch.stack([t["orig_size"] for t in targets], dim=0)
        results = postprocessors['bbox'](outputs, orig_target_sizes)
        if 'segm' in postprocessors.keys():
            target_sizes = torch.stack([t["size"] for t in targets], dim=0)
            results = postprocessors['segm'](results, outputs, orig_target_sizes, target_sizes)
        res = {target['image_id'].item(): output for target, output in zip(targets, results)}
        if coco_evaluator is not None:
            coco_evaluator.update(res)

        if panoptic_evaluator is not None:
            res_pano = postprocessors["panoptic"](outputs, target_sizes, orig_target_sizes)
            for i, target in enumerate(targets):
                image_id = target["image_id"].item()
                file_name = f"{image_id:012d}.png"
                res_pano[i]["image_id"] = image_id
                res_pano[i]["file_name"] = file_name

            panoptic_evaluator.update(res_pano)

    # gather the stats from all processes
    metric_logger.synchronize_between_processes()
    print("Averaged stats:", metric_logger)
    if coco_evaluator is not None:
        coco_evaluator.synchronize_between_processes()
    if panoptic_evaluator is not None:
        panoptic_evaluator.synchronize_between_processes()

    # accumulate predictions from all images
    if coco_evaluator is not None:
        coco_evaluator.accumulate()
        coco_evaluator.summarize()
    panoptic_res = None
    if panoptic_evaluator is not None:
        panoptic_res = panoptic_evaluator.summarize()
    stats = {k: meter.global_avg for k, meter in metric_logger.meters.items()}
    if coco_evaluator is not None:
        if 'bbox' in postprocessors.keys():
            stats['coco_eval_bbox'] = coco_evaluator.coco_eval['bbox'].stats.tolist()
        if 'segm' in postprocessors.keys():
            stats['coco_eval_masks'] = coco_evaluator.coco_eval['segm'].stats.tolist()
    if panoptic_res is not None:
        stats['PQ_all'] = panoptic_res["All"]
        stats['PQ_th'] = panoptic_res["Things"]
        stats['PQ_st'] = panoptic_res["Stuff"]
    return stats, coco_evaluator

@torch.no_grad()
def evaluate_hoi(dataset_file, model, postprocessors, data_loader, subject_category_id, device):
    model.eval()

    metric_logger = utils.MetricLogger(delimiter="  ")
    header = 'Test:'

    preds = []
    gts = []
    indices = []
    for samples, targets in metric_logger.log_every(data_loader, 10, header):
        samples = samples.to(device)

        output_list = model(samples)
        outputs = output_list['hoi']
        orig_target_sizes = torch.stack([t["orig_size"] for t in targets], dim=0)
        results = postprocessors['hoi'](outputs, orig_target_sizes)

        preds.extend(list(itertools.chain.from_iterable(utils.all_gather(results))))
        # For avoiding a runtime error, the copy is used
        gts.extend(list(itertools.chain.from_iterable(utils.all_gather(copy.deepcopy(targets)))))

    # gather the stats from all processes
    metric_logger.synchronize_between_processes()

    img_ids = [img_gts['id'] for img_gts in gts]
    _, indices = np.unique(img_ids, return_index=True)
    preds = [img_preds for i, img_preds in enumerate(preds) if i in indices]
    gts = [img_gts for i, img_gts in enumerate(gts) if i in indices]

    if dataset_file == 'hico':
        evaluator = HICOEvaluator(preds, gts, subject_category_id, data_loader.dataset.rare_triplets,
                                  data_loader.dataset.non_rare_triplets, data_loader.dataset.correct_mat)
    elif dataset_file == 'vcoco':
        evaluator = VCOCOEvaluator(preds, gts, subject_category_id, data_loader.dataset.correct_mat)

    stats = evaluator.evaluate()

    return stats


"""MCEq Flux Models

This script implements the use of MCEq flux models via the IceCube standard
method of getFlux(ptype, energy, costheta). As such it may be used as a
drop-in replacement for other fluxes. Weighting in IceCube is performed
by multiplying the flux by the normalized one weight:

NuGen:
    (with generator)
        weight = p_int * (flux_val / unit) * generator(energy, ptype, costheta)
    (without generator)
        weight = flux_val * one_weight / (type_weight * n_events * n_files)

    with flux_val = flux_object.getFlux(ptype, energy, costheta)

It is recommended to cache the results of MCEq because these take a while
to produce. By default, the cache file is chosen
to be located in the 'resources' directory relative to the location of this
script. You may also set the environment variable 'MCEQ_CACHE_DIR' in order
to choose a different location for the cache file, or pass in an explicit
cache file when initializing the MCEQFlux object.

Environment Variables:

    'MCEQ_CACHE_DIR':
        If provided, the MCEq cache file will be written to this directory.

    'MKL_PATH':
        Path to the MKL libraries. If provided, these are passed on to MCEq.
        Note: the python package can be installed directly with MKL support
        via 'pip install MCEq[MKL]'.

Credit for the vast majority of code in this file goes to Mathis Boerner.
"""
import os
import logging
from copy import deepcopy
import os
import numpy as np
from scipy.interpolate import RectBivariateSpline

import ic3_labels

log = logging.getLogger('MCEqFlux')

# If cashier is available, set up directory for caching of MCEq results
try:
    from ic3_labels.weights.resources.cashier import cache
    got_cashier = True

    if 'MCEQ_CACHE_DIR' in os.environ:
        cache_dir = os.environ['MCEQ_CACHE_DIR']
        log.info("Found 'MCEQ_CACHE_DIR' in environment variables: {}".format(
            cache_dir))

        if not os.path.exists(cache_dir):
            log.info('Creating cache directory: {}'.format(cache_dir))
            os.makedirs(cache_dir)

        CACHE_FILE = os.path.join(cache_dir, 'mceq.cache')

    else:
        script_dir = os.path.dirname(os.path.abspath(__file__))
        CACHE_FILE = os.path.join(script_dir, 'resources', 'mceq.cache')

    log.info('Using MCEq cache file: {}'.format(CACHE_FILE))

except ImportError:
    got_cashier = False
    CACHE_FILE = None
    log.info("Could not import 'cashier'. MCEq results will not be cached!")

# Dictionary that converts ptype -> MCEq type string
PTYPE_CONVERTER = {
    12: 'nue',
    -12: 'antinue',
    14: 'numu',
    -14: 'antinumu',
    16: 'nutau',
    -16: 'antinutau',
}

def get_spline(
        interaction_model,
        primary_model,
        months,
        theta_grid,
        theta_grid_cos,
        cached=True,
        cache_file=CACHE_FILE,
        cache_read_only=False):
    """Get MCEq spline

    Solves the MCEq cascade equations for the given parameters. The equations
    are solved on the provided grid and interpolated.

    Parameters
    ----------
    interaction_model : str
        The interaction model. This is passed on to `MCEqRun`.
    primary_model : str
        The primary model to use. Must be one of:
            GST_3-gen, GST_4-gen, H3a, H4a, poly-gonato, TIG, ZS, ZSP, GH
    months : list of str
        The months for which to solve the cascade equations. These must be
        provided as a list of month names, e.g. ['January', 'August']. A list
        of splines will be returned of the same length as `months`.
    theta_grid : array_like
        The grid points in theta to evaluate on in degrees. If `theta_grid_cos`
        is True, this is instead cos(theta).
    theta_grid_cos : bool
        If True, `theta_grid` is interpreted as cos(theta), i.e. arccos() is
        applied first.
    cached : bool, optional
        If True, the result will be cached, or taken from cache if previously
        already computed. This is recommended, as MCEq takes a while to run.
    cache_file : str, optional
        The path to the cache file to use.
    cache_read_only : bool, optional
        If True, the cache is read only.

    Returns
    -------
    dict
        The result of MCEq together with the fitted splines. The structure is
        as follows:
        {
            # first month provided via `months`
            0: {
                'total_spline_dict': dict of RectBivariateSpline
                    A dictionary with the fitted splines for each particle
                    type for the 'total' flux. The dictionary keys are the
                    PDG particle encodings.
                'conv_spline_dict': dict of RectBivariateSpline
                    A dictionary with the fitted splines for each particle
                    type for the 'conv' flux. The dictionary keys are the
                    PDG particle encodings.
                'pr_spline_dict': dict of RectBivariateSpline
                    A dictionary with the fitted splines for each particle
                    type for the 'pr' flux. The dictionary keys are the
                    PDG particle encodings.
                'total_flux_dict': dict of array_like
                    A dictionary with the total flux for each grid point.
                    This is the result obtained from MCEq for the 'total' flux.
                'conv_flux_dict': dict of array_like
                    A dictionary with the conv flux for each grid point.
                    This is the result obtained from MCEq for the 'conv' flux.
                'pr_flux_dict': dict of array_like
                    A dictionary with the prompt flux for each grid point.
                    This is the result obtained from MCEq for the 'pr' flux.
                'config_updates':   dict
                    A dictionary of config updates that were applied to
                    mceq_config prior to solving the equations.
                'mceq_version' : str
                    The MCEq version that was used to create the splines.
                'ic3_labels_version' : str
                    The version of the ic3-labels package that was used to
                    create the splines.
                'e_grid' : array_like
                    The grid of energy points in log10.
                'theta_grid' : array_like
                    The grid of thetas.
            }

            # second month provided via `months`
            1: {
                ...
            }

            ...
        }
    """
    log.info('Getting Spline for {}; {} (cached={})'.format(
        interaction_model,
        primary_model,
        cached))

    def __solve_month__(
            mceq_run, e_grid, theta_grid, theta_grid_cos,
            ptype_converter=PTYPE_CONVERTER,
            eps=1e-128,
            ):
        """Solve MCEq equations for the provided mceq_run instance.

        Parameters
        ----------
        mceq_run : MCEqRun instance
            The MCEqRun instance. This instance must be configured to use
            the desired geometry and season.
        e_grid : array_like
            The grid of energy points in log10.
        theta_grid : array_like
            The grid points in theta to evaluate on in degrees.
            If `theta_grid_cos` is True, this is instead cos(theta).
        theta_grid_cos : bool
            If True, `theta_grid` is interpreted as cos(theta),
            i.e. arccos() is applied first.
        ptype_converter : dict, optional
            A dictionary that converts PDG encoding to MCEq type string.
        eps : float, optional
            A small float value > 0 that is used to clip the total flux
            prior to applying log10 for the spline fitting.

        Returns
        -------
        list of dict of RectBivariateSpline
            A list of dictionaries with the fitted splines for each particle
            type. The dictionary keys are the PDG particle encodings.
            The order of the dictionaries are: 'total', 'conv', 'pr'
        list of dict of array_like
            A list of dictionaries with the total flux for each grid point.
            This is the result obtained from MCEq.
            The order of the dictionaries are: 'total', 'conv', 'pr'
        """
        total_flux_dict = {}
        conv_flux_dict = {}
        pr_flux_dict = {}
        total_spline_dict = {}
        conv_spline_dict = {}
        pr_spline_dict = {}
        for key, value in ptype_converter.items():
            total_flux_dict[key] = np.ones((len(e_grid), len(theta_grid)))
            conv_flux_dict[key] = np.ones((len(e_grid), len(theta_grid)))
            pr_flux_dict[key] = np.ones((len(e_grid), len(theta_grid)))

        for i, theta_i in enumerate(theta_grid):
            if theta_grid_cos:
                theta_i = np.rad2deg(np.arccos(theta_i))
            mceq_run.set_theta_deg(theta_i)
            mceq_run.solve()

            # fill in flux totals
            for key, value in ptype_converter.items():
                total_flux_dict[key][:, i] = mceq_run.get_solution(
                    'total_{}'.format(value))
                conv_flux_dict[key][:, i] = mceq_run.get_solution(
                    'conv_{}'.format(value))
                pr_flux_dict[key][:, i] = mceq_run.get_solution(
                    'pr_{}'.format(value))

        # create splines
        for key, value in ptype_converter.items():
            total_spline_dict[key] = RectBivariateSpline(
                e_grid,
                theta_grid,
                np.log10(np.clip(total_flux_dict[key], eps, float('inf'))),
                s=0,
            )
            conv_spline_dict[key] = RectBivariateSpline(
                e_grid,
                theta_grid,
                np.log10(np.clip(conv_flux_dict[key], eps, float('inf'))),
                s=0,
            )
            pr_spline_dict[key] = RectBivariateSpline(
                e_grid,
                theta_grid,
                np.log10(np.clip(pr_flux_dict[key], eps, float('inf'))),
                s=0,
            )

        spline_dicts = [
            total_spline_dict, conv_spline_dict, pr_spline_dict
        ]

        flux_dicts = [
            total_flux_dict, conv_flux_dict, pr_flux_dict
        ]

        return spline_dicts, flux_dicts

    def __get_spline__(
            interaction_model,
            primary_model,
            months,
            theta_grid,
            theta_grid_cos,
            ):
        """Get MCEq spline for the provided settings

        Parameters
        ----------
        interaction_model : str
            The interaction model. This is passed on to `MCEqRun`.
        primary_model : str
            The primary model to use. Must be one of:
                GST_3-gen, GST_4-gen, H3a, H4a, poly-gonato, TIG, ZS, ZSP, GH
        months : list of str
            The months for which to solve the cascade equations. These must be
            provided as a list of month names, e.g. ['January', 'August']. A
            list of splines will be returned of the same length as `months`.
        theta_grid : array_like
            The grid points in theta to evaluate on in degrees.
            If `theta_grid_cos` is True, this is instead cos(theta).
        theta_grid_cos : bool
            If True, `theta_grid` is interpreted as cos(theta),
            i.e. arccos() is applied first.

        Returns
        -------
        dict
            The result of MCEq together with the fitted splines.
            See documentation of `get_spline()` for more details.

        Raises
        ------
        AttributeError
            If the provided `primary_model` is unknown.
        """
        log.info('\tCalculating \'{}\' \'{}\''.format(
            interaction_model, primary_model))

        import mceq_config
        from MCEq import version
        from MCEq.core import MCEqRun
        import crflux.models as pm

        config_updates = {
            'h_obs': 1000.,
            'debug_level': 1,
        }
        if 'MKL_PATH' in os.environ:
            config_updates['MKL_path'] = os.environ['MKL_PATH']

        splines = {}
        pmodels = {
            "GST_3-gen": (pm.GaisserStanevTilav, "3-gen"),
            "GST_4-gen": (pm.GaisserStanevTilav, "4-gen"),
            "H3a": (pm.HillasGaisser2012, "H3a"),
            "H4a": (pm.HillasGaisser2012, "H4a"),
            "poly-gonato": (pm.PolyGonato, False),
            "TIG": (pm.Thunman, None),
            "ZS": (pm.ZatsepinSokolskaya, 'default'),
            "ZSP": (pm.ZatsepinSokolskaya, 'pamela'),
            "GH": (pm.GaisserHonda, None),
        }

        for i, month in enumerate(months):
            config_updates['density_model'] = (
                'MSIS00_IC', ('SouthPole', month))

            # update settings in mceq_config
            # Previous method mceq_config.config is deprecated and resulted
            # in pickle errors for deepcopy.
            for name, value in config_updates.items():
                setattr(mceq_config, name, value)

            try:
                pmodel = pmodels[primary_model]
            except KeyError:
                raise AttributeError(
                    'primary_model {} unknown. options: {}'.format(
                        primary_model, pmodels.keys()))
            mceq_run = MCEqRun(
                interaction_model=interaction_model,
                primary_model=pmodel,
                theta_deg=0.0,
                **config_updates)
            e_grid = np.log10(deepcopy(mceq_run.e_grid))
            spline_dicts, flux_dicts = __solve_month__(
                mceq_run,
                e_grid,
                theta_grid,
                theta_grid_cos)

            splines[i] = {}
            splines[i]['total_spline_dict'] = spline_dicts[0]
            splines[i]['conv_spline_dict'] = spline_dicts[1]
            splines[i]['pr_spline_dict'] = spline_dicts[2]
            splines[i]['total_flux_dict'] = flux_dicts[0]
            splines[i]['conv_flux_dict'] = flux_dicts[1]
            splines[i]['pr_flux_dict'] = flux_dicts[2]
            splines[i]['config_updates'] = deepcopy(config_updates)
            splines[i]['mceq_version'] = version.__version__
            splines[i]['ic3_labels_version'] = ic3_labels.__version__
            splines[i]['e_grid'] = e_grid
            splines[i]['theta_grid'] = theta_grid
        return splines

    if got_cashier and cached:
        if cache_file is None:
            cache_f = 'mceq.cache'
        else:
            cache_f = cache_file
        log.info('\tUsing cache \'{}\''.format(cache_f))

        @cache(cache_file=cache_f, read_only=cache_read_only)
        def wrapped_get_spline(
                interaction_model,
                primary_model,
                months,
                theta_grid,
                theta_grid_cos,
                ):
            return __get_spline__(
                interaction_model=interaction_model,
                primary_model=primary_model,
                months=months,
                theta_grid=theta_grid,
                theta_grid_cos=theta_grid_cos,
            )

        return wrapped_get_spline(
            interaction_model,
            primary_model,
            months,
            theta_grid,
            theta_grid_cos,
        )
    else:
        return __get_spline__(
            interaction_model=interaction_model,
            primary_model=primary_model,
            months=months,
            theta_grid=theta_grid,
            theta_grid_cos=theta_grid_cos,
        )

class MCEQFlux(object):

    """MCQe Flux Wrapper

    Attributes
    ----------
    min_theta_deg : float, optional
        The minimum value of the theta grid in degrees.
        If `theta_grid_cos` is True, this is instead cos(theta).
    max_theta_deg : float, optional
        The maximum value of the theta grid in degrees.
        If `theta_grid_cos` is True, this is instead cos(theta).
    month_weights : array_like
        A list of probabilities for each month. These are used as weights
        to sample the corresponding month for each MC events.
        These weights can, for instance, be set to the relative livetime in
        each month of the year. This will then account for seasonal variations.
    months : list of str
        A list of months for which the interpolation splines are created.
    random_state : np.random.RandomState
        The random state that is used to draw the month for each MC event.
    splines : dict
        A dictionary containing the MCEq result and fitted splines.
        See documentation of `get_splines()` for more details.
    theta_grid : array_like
        The grid points in theta to evaluate on in degrees.
        If `theta_grid_cos` is True, this is instead cos(theta).
    theta_grid_cos : bool
        If True, `min_theta_deg` and `max_theta_deg` are interpreted as
        cos(theta), i.e. arccos() is applied first.
    """

    def __init__(
            self,
            min_theta_deg=0.,
            max_theta_deg=180.,
            theta_grid_cos=False,
            theta_steps=181,
            season='full_year',
            flux_type='total',
            random_state=None,
            **kwargs):
        """Initialize MCEQFlux Instance

        Parameters
        ----------
        min_theta_deg : float, optional
            The minimum value of the theta grid in degrees.
            If `theta_grid_cos` is True, this is instead cos(theta).
        max_theta_deg : float, optional
            The maximum value of the theta grid in degrees.
            If `theta_grid_cos` is True, this is instead cos(theta).
        theta_grid_cos : bool
            If True, `min_theta_deg` and `max_theta_deg` are interpreted as
            cos(theta), i.e. arccos() is applied first.
        theta_steps : int, optional
            The number of grid points between the specified min and max values.
        season : str, optional
            What season to use. This may either be a single month ,
            for example 'January', or 'full_year' may be used to run MCEq
            for every month of the year.
        flux_type : str, optional
            The flux type to compute. This must be one of
                'total': combined prompt and conv flux
                'pr': prompt neutrino flux
                'conv': conventional neutrino flux
            This will set the default flux type when calling `getFlux()`.
            You may, however, overwrite this defaul by passing an alternative
            flux type to `getFlux()`. Setting this default value allows for
            drop in replacement of other flux implementations in IceCube.
        random_state : np.random.Randomstate or int, optional
            An int or random state to set the seed.
        **kwargs
            Additional keyword arguments. (Not used!)
        """
        if not isinstance(random_state, np.random.RandomState):
            random_state = np.random.RandomState(random_state)
        self.random_state = random_state
        if season.lower() == 'full_year':
            self.months = [
                'January',
                'February',
                'March',
                'April',
                'May',
                'June',
                'July',
                'August',
                'September',
                'October',
                'November',
                'December',
            ]
        else:
            self.months = [season]

        if flux_type.lower() not in ['total', 'conv', 'pr']:
            raise ValueError('Flux type: {} must be on of {}'.format(
                flux_type.lower(), ['total', 'conv', 'pr']))

        self.flux_type = flux_type
        self.set_month_weights(np.ones_like(self.months, dtype=float))
        self.min_theta = min_theta_deg
        self.max_theta = max_theta_deg
        if theta_grid_cos:
            self.min_theta = np.cos(np.deg2rad(min_theta_deg))
            self.max_theta = np.cos(np.deg2rad(max_theta_deg))
        self.theta_grid_cos = theta_grid_cos
        self.theta_grid = np.linspace(
            self.min_theta, self.max_theta, theta_steps)

        self.splines = None

    def initialize(
            self,
            interaction_model='SIBYLL2.3c',
            primary_model='H3a',
            cached=True,
            cache_file=CACHE_FILE,
            cache_read_only=False,
            ):
        """Initialize MCEQFlux instance

        This will compute the splines or retrieve these from the cache
        if `cached` is True and if these have been previously computed.
        This method must be called prior to calls to `getFlux()`.

        Parameters
        ----------
        interaction_model : str
            The interaction model. This is passed on to `MCEqRun`.
        primary_model : str
            The primary model to use. Must be one of:
                GST_3-gen, GST_4-gen, H3a, H4a, poly-gonato, TIG, ZS, ZSP, GH
        cached : bool, optional
            If True, the result will be cached and if already computed, it will
            be retrieved from cache. This avoids recomputation of MCEq, which
            is recommended in order to reduce computation time.
        cache_file : str, optional
            The path to the cache file to use.
        cache_read_only : bool, optional
            If True, the cache is read only.
        """
        if cache_file is None:
            cache_file = CACHE_FILE

        self.splines = get_spline(
            interaction_model,
            primary_model,
            self.months,
            self.theta_grid,
            self.theta_grid_cos,
            cached=cached,
            cache_file=cache_file,
            cache_read_only=cache_read_only,
        )

        from MCEq import version

        # throw warning if there is a version mis-match.
        for key, spline in self.splines.items():
            msg = (
                'Cached file was created with {} version {}, '
                'but this is version {}!'
            )
            if version.__version__ != spline['mceq_version']:
                log.warning(msg.format(
                    'MCEq', spline['mceq_version'], version.__version__))

            if ic3_labels.__version__ != spline['ic3_labels_version']:
                log.warning(msg.format(
                    'ic3_labels',
                    spline['ic3_labels_version'],
                    ic3_labels.__version__,
                ))

    def set_month_weights(self, month_weights):
        """Summary

        Parameters
        ----------
        month_weights : array_like
            A list of probabilities for each month (these will be normalized
            internally). These are used as weights to sample the corresponding
            month for each MC events. These weights can, for instance, be set
            to the relative livetime in each month of the year.
            This will then account for seasonal variations.

        Raises
        ------
        AttributeError
            If the length of the provided `muon_weights` does not match the
            length of the specified months.
        """
        if len(month_weights) != len(self.months):
            raise AttributeError(
                'month_weights needs to be of the same '
                'length like self.months.'
            )
        self.month_weights = month_weights / np.sum(month_weights)

    def getFlux(
            self,
            ptype,
            energy,
            costheta,
            selected_month=None,
            random_state=None,
            flux_type=None,
            ):
        """Get flux for provided particle

        The flux is given in GeV^-1 cm^-2 s^-1 sr^-1 and may be used to
        weight NuGen events via the normalized `one_weight`:
            weight = flux * one_weight / (type_weight * n_events * n_files)

        Parameters
        ----------
        ptype : array_like or int
            The PDG encoding.
            For instance: I3MCWeightDict -> PrimaryNeutrinoType
        energy : array_like or float
            The energy of the primary particle.
            For instance: I3MCWeightDict -> PrimaryNeutrinoEnergy
        costheta : array_like or float
            The cos(zenith) angle of the primary particle.
            For instance: cos(I3MCWeightDict -> PrimaryNeutrinoZenith)
        selected_month : array_like, optional
            The month in which each event occurred. This must be given as
            an array of integer values between [0, 11] if `season` is
            'full_year'. If the `MCEQFlux` instance is initialized with only
            one month as the season, then `selected_month` must not be set.
            If None provided, the corresponding month of each event will be
            sampled via the defined `month_weights`.
        random_state : np.random.Randomstate or int, optional
            An int or random state to set the seed.
            If None provided, the random state will be used that was
            reated during initialization.
        flux_type : str, optional
            The flux type to compute. This must be one of
                'total': combined prompt and conv flux
                'pr': prompt neutrino flux
                'conv': conventional neutrino flux
            If None is provided, the specified default value at
            object instantiation time (__init__()) will be used.

        Returns
        -------
        array_like
            The flux for the given particle in GeV^-1 cm^-2 s^-1 sr^-1.

        Raises
        ------
        RuntimeError
            If MCEQFlux has not been initialized yet.
        ValueError
            If wrong `flux_type` is provided.
        """
        if self.splines is None:
            raise RuntimeError(
                'No splines calculated! Run \'initialize\' first')

        if len(self.months) == 1 and selected_month is not None:
            raise ValueError(
                'The months may not be set, since the MCEQFlux instance is '
                + 'initialized with only one month: {}'.format(self.months)
            )

        if flux_type is None:
            flux_type = self.flux_type
        elif flux_type.lower() not in ['total', 'conv', 'pr']:
            raise ValueError('Flux type: "{}" must be on of {}'.format(
                flux_type.lower(), ['total', 'conv', 'pr']))
            flux_type = flux_type.lower()

        if random_state is None:
            random_state = self.random_state
        elif not isinstance(random_state, np.random.RandomState):
            random_state = np.random.RandomState(random_state)

        # convert to numpy arrays and make sure these are at least 1D
        ptype = np.atleast_1d(ptype)
        energy = np.atleast_1d(energy)
        costheta = np.atleast_1d(costheta)

        if len(self.splines) > 1:
            if selected_month is None:
                int_months = np.arange(len(self.splines), dtype=int)
                selected_month = random_state.choice(
                    int_months,
                    replace=True,
                    size=len(energy),
                    p=self.month_weights,
                )
            else:
                selected_month = np.asarray(selected_month, dtype=int)
        else:
            selected_month = list(self.splines.keys())[0]
        flux = np.ones_like(energy)
        flux[:] = float('NaN')
        log10_energy = np.log10(energy)
        theta = np.rad2deg(np.arccos(costheta))

        for ptype_i in np.unique(ptype):
            mask_ptype = ptype == ptype_i
            for i in self.splines.keys():
                if isinstance(selected_month, int):
                    idx_ptype = mask_ptype
                else:
                    is_in_month = selected_month == i
                    idx_ptype = np.logical_and(mask_ptype, is_in_month)
                flux[idx_ptype] = self.splines[i][
                        flux_type + '_spline_dict'][ptype_i](
                    log10_energy[idx_ptype],
                    theta[idx_ptype],
                    grid=False)
        return np.power(10., flux)


import numpy as np
from collections import OrderedDict
import torch
import fire
from sklearn.metrics import confusion_matrix

def vote(**kwargs):
    dict =[]
    if kwargs['Flow'] == True:
        dict.append(torch.load("result/Flow.pth"))
    if kwargs['RGB'] == True:
        dict.append(torch.load("result/RGB.pth"))
    if kwargs['RGBDiff'] == True:
        dict.append(torch.load("result/RGBDiff.pth"))
    labels = []
    preds = []
    for key in dict[0].keys():
        temp = []
        for i in range(len(dict)):
            temp.append(dict[i][key]['pred'])
            label = dict[i][key]['label']
        pred = np.argmax(np.concatenate(temp,axis=0).sum(0),0)
        preds.append(pred)
        labels.append(label)
    cfs  = confusion_matrix(labels,preds)
    cfs_num = cfs.sum(1)
    cfs_hit = np.diag(cfs)
    pred_prob = (cfs_hit*1./cfs_num)
    print np.mean(pred_prob)
if __name__ == '__main__':
    fire.Fire()