#!/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, 0, 21, 17), 'visual.misc.set_plot_config', 'set_plot_config', ({}, {}), '()', False, 'from visual.misc import set_plot_config\n'), ((84, 15, 84, 58), 'archi.classic.L4', 'ARCHI', (), '', True, 'from archi.classic import L4 as ARCHI\n'), ((85, 21, 85, 40), 'utils.model.get_optimizer', 'get_optimizer', ({(85, 35, 85, 39): 'args'}, {}), '(args)', False, 'from utils.model import get_optimizer\n'), ((86, 12, 86, 50), 'utils.model.get_model', 'get_model', ({(86, 22, 86, 26): 'args', (86, 28, 86, 49): 'TangentPropClassifier'}, {}), '(args, TangentPropClassifier)', False, 'from utils.model import get_model\n'), ((96, 13, 96, 25), 'utils.log.set_logger', 'set_logger', ({}, {}), '()', False, 'from utils.log import set_logger\n'), ((99, 4, 99, 17), 'utils.log.flush', 'flush', ({(99, 10, 99, 16): 'logger'}, {}), '(logger)', False, 'from utils.log import flush\n'), ((102, 4, 102, 55), 'os.makedirs', 'os.makedirs', (), '', False, 'import os\n'), ((106, 14, 106, 51), 'pandas.concat', 'pd.concat', (), '', True, 'import pandas as pd\n'), ((114, 13, 114, 32), 'logging.getLogger', 'logging.getLogger', ({}, {}), '()', False, 'import logging\n'), ((115, 4, 115, 16), 'utils.log.print_line', 'print_line', ({}, {}), '()', False, 'from utils.log import print_line\n'), ((117, 4, 117, 16), 'utils.log.print_line', 'print_line', ({}, {}), '()', False, 'from utils.log import print_line\n'), ((122, 13, 122, 21), 'problem.higgs.HiggsConfigTesOnly', 'Config', ({}, {}), '()', True, 'from problem.higgs import HiggsConfigTesOnly as Config\n'), ((124, 55, 124, 97), 'problem.higgs.get_generators_torch', 'get_generators_torch', (), '', False, 'from problem.higgs import get_generators_torch\n'), ((126, 22, 126, 51), 'problem.higgs.GeneratorCPU', 'GeneratorCPU', ({(126, 35, 126, 50): 'valid_generator'}, {}), '(valid_generator)', False, 'from problem.higgs import GeneratorCPU\n'), ((127, 21, 127, 49), 'problem.higgs.GeneratorCPU', 'GeneratorCPU', ({(127, 34, 127, 48): 'test_generator'}, {}), '(test_generator)', False, 'from problem.higgs import GeneratorCPU\n'), ((132, 4, 132, 50), 'os.makedirs', 'os.makedirs', (), '', False, 'import os\n'), ((133, 4, 133, 17), 'utils.log.flush', 'flush', ({(133, 10, 133, 16): 'logger'}, {}), '(logger)', False, 'from utils.log import flush\n'), ((136, 4, 136, 74), 'utils.model.train_or_load_neural_net', 'train_or_load_neural_net', (), '', False, 'from utils.model import train_or_load_neural_net\n'), ((66, 31, 70, 34), 'collections.OrderedDict', 'OrderedDict', ({(66, 43, 70, 33): "[('tes', self.tes), ('jes', self.jes), ('les', self.les)]"}, {}), "([('tes', self.tes), ('jes', self.jes), ('les', self.les)])", False, 'from collections import OrderedDict\n'), ((108, 19, 108, 73), 'os.path.join', 'os.path.join', ({(108, 32, 108, 55): 'model.results_directory', (108, 57, 108, 72): '"""threshold.csv"""'}, {}), "(model.results_directory, 'threshold.csv')", False, 'import os\n')]