#!/usr/bin/env python
# coding: utf-8
from __future__ import print_function
from __future__ import division
from __future__ import absolute_import
from __future__ import unicode_literals

# Command line :
# python -m benchmark.VAR.GG.TP

import os
import logging
from config import SEED
from config import _ERROR
from config import _TRUTH

import numpy as np
import pandas as pd

from visual.misc import set_plot_config
set_plot_config()

from utils.log import set_logger
from utils.log import flush
from utils.log import print_line
from utils.model import get_model
from utils.model import get_optimizer
from utils.model import train_or_load_neural_net
from utils.evaluation import evaluate_summary_computer
from utils.images import gather_images

from visual.misc import plot_params

from problem.higgs import HiggsConfigTesOnly as Config
from problem.higgs import get_generators_torch
from problem.higgs import GeneratorCPU
from problem.higgs import GeneratorTorch
from problem.higgs import HiggsNLL as NLLComputer

from model.tangent_prop import TangentPropClassifier
from archi.classic import L4 as ARCHI
from ...my_argparser import TP_parse_args
from collections import OrderedDict

from .common import measurement


DATA_NAME = 'HIGGSTES'
BENCHMARK_NAME = 'VAR-'+DATA_NAME
N_ITER = 30


class TrainGenerator:
    def __init__(self, data_generator, cuda=False):
        self.data_generator = data_generator
        if cuda:
            self.data_generator.cuda()
        else:
            self.data_generator.cpu()

        self.mu  = self.tensor(Config.CALIBRATED.mu, requires_grad=True)
        self.tes = self.tensor(Config.CALIBRATED.tes, requires_grad=True)
        self.jes = self.tensor(Config.CALIBRATED.jes, requires_grad=True)
        self.les = self.tensor(Config.CALIBRATED.les, requires_grad=True)
        self.params = (self.tes, self.jes, self.tes, self.mu)
        self.nuisance_params = OrderedDict([
                                ('tes', self.tes),
                                ('jes', self.jes),
                                ('les', self.les),
                                ])

    def generate(self, n_samples=None):
            X, y, w = self.data_generator.diff_generate(*self.params, n_samples=n_samples)
            return X, y, w

    def reset(self):
        self.data_generator.reset()

    def tensor(self, data, requires_grad=False, dtype=None):
        return self.data_generator.tensor(data, requires_grad=requires_grad, dtype=dtype)


def build_model(args, i_cv):
    args.net = ARCHI(n_in=29, n_out=2, n_unit=args.n_unit)
    args.optimizer = get_optimizer(args)
    model = get_model(args, TangentPropClassifier)
    model.set_info(DATA_NAME, BENCHMARK_NAME, i_cv)
    return model


# =====================================================================
# MAIN
# =====================================================================
def main():
    # BASIC SETUP
    logger = set_logger()
    args = TP_parse_args(main_description="Training launcher for INFERNO on GG benchmark")
    logger.info(args)
    flush(logger)
    # INFO
    model = build_model(args, -1)
    os.makedirs(model.results_directory, exist_ok=True)
    # RUN
    logger.info(f'Running runs [{args.start_cv},{args.end_cv}[')
    results = [run(args, i_cv) for i_cv in range(args.start_cv, args.end_cv)]
    results = pd.concat(results, ignore_index=True)
    # EVALUATION
    results.to_csv(os.path.join(model.results_directory, 'threshold.csv'))
    print(results)
    print("DONE !")


def run(args, i_cv):
    logger = logging.getLogger()
    print_line()
    logger.info('Running iter n°{}'.format(i_cv))
    print_line()


    # LOAD/GENERATE DATA
    logger.info('Set up data generator')
    config = Config()
    seed = SEED + i_cv * 5
    train_generator, valid_generator, test_generator = get_generators_torch(seed, cuda=args.cuda)
    train_generator = TrainGenerator(train_generator, cuda=args.cuda)
    valid_generator = GeneratorCPU(valid_generator)
    test_generator = GeneratorCPU(test_generator)

    # SET MODEL
    logger.info('Set up classifier')
    model = build_model(args, i_cv)
    os.makedirs(model.results_path, exist_ok=True)
    flush(logger)

    # TRAINING / LOADING
    train_or_load_neural_net(model, train_generator, retrain=args.retrain)

    # MEASUREMENT
    results = measurement(model, i_cv, config, valid_generator, test_generator)
    print(results)
    return results

if __name__ == '__main__':
    main()

[(21, 'visual.misc.set_plot_config', 'set_plot_config', False, 'from visual.misc import set_plot_config\n'), (84, 'archi.classic.L4', 'ARCHI', True, 'from archi.classic import L4 as ARCHI\n'), (85, 'utils.model.get_optimizer', 'get_optimizer', False, 'from utils.model import get_optimizer\n'), (86, 'utils.model.get_model', 'get_model', False, 'from utils.model import get_model\n'), (96, 'utils.log.set_logger', 'set_logger', False, 'from utils.log import set_logger\n'), (99, 'utils.log.flush', 'flush', False, 'from utils.log import flush\n'), (102, 'os.makedirs', 'os.makedirs', False, 'import os\n'), (106, 'pandas.concat', 'pd.concat', True, 'import pandas as pd\n'), (114, 'logging.getLogger', 'logging.getLogger', False, 'import logging\n'), (115, 'utils.log.print_line', 'print_line', False, 'from utils.log import print_line\n'), (117, 'utils.log.print_line', 'print_line', False, 'from utils.log import print_line\n'), (122, 'problem.higgs.HiggsConfigTesOnly', 'Config', True, 'from problem.higgs import HiggsConfigTesOnly as Config\n'), (124, 'problem.higgs.get_generators_torch', 'get_generators_torch', False, 'from problem.higgs import get_generators_torch\n'), (126, 'problem.higgs.GeneratorCPU', 'GeneratorCPU', False, 'from problem.higgs import GeneratorCPU\n'), (127, 'problem.higgs.GeneratorCPU', 'GeneratorCPU', False, 'from problem.higgs import GeneratorCPU\n'), (132, 'os.makedirs', 'os.makedirs', False, 'import os\n'), (133, 'utils.log.flush', 'flush', False, 'from utils.log import flush\n'), (136, 'utils.model.train_or_load_neural_net', 'train_or_load_neural_net', False, 'from utils.model import train_or_load_neural_net\n'), (66, 'collections.OrderedDict', 'OrderedDict', False, 'from collections import OrderedDict\n'), (108, 'os.path.join', 'os.path.join', False, 'import os\n')]