import argparse
import json
from copy import deepcopy

from .Model import Model, model_types
from .DataReader import HistoricalDataReader, AggregatorReaderWrapper
from .DataReader import RadiationDataWrapper, WeatherDataWrapper, WeatherDataReader
from .SlidingWindow import SlidingWindowAdapter


def add_args_group(parser: argparse.ArgumentParser, uses_model=True, uses_future=False):
    group = parser.add_argument_group('Input data')
    group.add_argument('--file', required=True, help='Data file with one minute values')
    group.add_argument('--aggregate', type=int,
                       help='How many values to aggregate into one data set (default: do not aggregate)')
    group.add_argument('--weather', default='0_data/weather_export.json', help='JSON file with weather data to use')
    if uses_model:
        group_model = parser.add_argument_group('Model parameters')
        group_model_file = group_model.add_mutually_exclusive_group()
        group_model_file.add_argument('--model-load', type=str, default=None,
                                      help='File name where to load a model from')
        group_model_file.add_argument('--model-save', type=str, default=None,
                                      help='File name where to save a model to')
        group_model.add_argument('--model-type', choices=model_types.keys(), default='linear')
        group_model.add_argument('--model-epochs', type=int, default=100,
                                 help='Number of epochs of neural network training (default: 100)')
        group_model.add_argument('--model-past', type=int, default=5,
                                 help='Number of past values to use to create thermal model (default: 5)')
        group_model.add_argument('--model-future', type=int, default=0,
                                 help='Number of forecast values to use to create thermal model (default: 0)')
    if uses_future:
        group.add_argument('--future', type=int, default=None, required=True,
                           help='For how many steps should environment be evaluated before performing action')


def get_config_from_args(args: argparse.Namespace):
    if args.model_load:
        with open(args.model_load, 'r') as fp:
            config = json.load(fp)
    else:
        config = {
            'aggregate': args.aggregate,
            'model_type': args.model_type,
            'model_epochs': args.model_epochs,
            'model_past_values': args.model_past,
            'model_past_fields': ('temp_in', 'temp_out', 'mode'),
            'model_future_values': args.model_future,
            'model_future_fields': ('temperature', 'wind_speed', 'radiation', 'humid'),
        }

    config['file'] = args.file
    config['model_load'] = args.model_load
    config['model_save'] = args.model_save
    config['weather'] = args.weather
    if 'future' in args.__dict__:
        config['future_values'] = args.future
    return config


def get_config(**kwargs):
    parser = argparse.ArgumentParser()
    add_args_group(parser, **kwargs)
    args = parser.parse_args()
    return get_config_from_args(args)


def prepare_config_save(config: dict) -> dict:
    config = deepcopy(config)
    del config['file']
    del config['model_load']
    del config['model_save']
    del config['weather']
    if 'future_values' in config:
        del config['future_values']
    return config


def get_reader_from_config(config: dict):
    reader = HistoricalDataReader(config['file'])
    if config['aggregate']:
        reader = AggregatorReaderWrapper(reader, aggregate=config['aggregate'])
    if config['weather']:
        weather = WeatherDataReader('0_data/weather_export.json')
        reader = WeatherDataWrapper(reader, weather=weather)
        reader = RadiationDataWrapper(reader)
    return reader


def get_sliding_window_from_config(config: dict, **kwargs):
    reader = get_reader_from_config(config)
    return SlidingWindowAdapter(reader, **config, **kwargs)


def get_sliding_window_from_args(args: argparse.Namespace, **kwargs):
    config = get_config_from_args(args)
    return get_sliding_window_from_config(config, **kwargs)


def get_model_from_config(config: dict):
    if config['model_load']:
        model = Model.load(config, config['model_load'] + '.model')
    else:
        model = Model.generate(config)

    if config['model_save']:
        with open(config['model_save'], 'w') as fp:
            json.dump(prepare_config_save(config), fp)
        model.save(config['model_save'] + '.model')

    return model