
import pytorch_lightning as pl
import optuna
import xarray as xr

from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint

import os
import shutil
from argparse import ArgumentParser
from datetime import datetime

from project.fluxdata import FluxData
from models.hybrid import Q10Model

# Hardcoded `Trainer` args. Note that these cannot be passed via cli.
TRAINER_ARGS = dict(
    max_epochs=100,
    log_every_n_steps=1,
    weights_summary=None
)


class Objective(object):
    def __init__(self, args):
        self.args = args

    def __call__(self, trial: optuna.trial.Trial) -> float:
        q10_init = trial.suggest_float('q10_init', 0.0001, 1000.)
        seed = trial.suggest_int('seed', 0, 999999999999)
        use_ta = trial.suggest_categorical('use_ta', [True, False])
        dropout = trial.suggest_float('dropout', 0.0, 1.0)

        if use_ta:
            features = ['sw_pot', 'dsw_pot', 'ta']
        else:
            features = ['sw_pot', 'dsw_pot']

        pl.seed_everything(seed)

        # Further variables used in the hybrid model.
        physical = ['ta']

        # Target (multiple targets not possible currently).
        targets = ['reco']

        # Find variables that are only needed in physical model but not in NN.
        physical_exclusive = [v for v in physical if v not in features]

        # ------------
        # data
        # ------------
        ds = xr.open_dataset(self.args.data_path)

        fluxdata = FluxData(
            ds,
            features=features + physical_exclusive,
            targets=targets,
            context_size=1,
            train_time=slice('2003-01-01', '2006-12-31'),
            valid_time=slice('2007-01-01', '2007-12-31'),
            test_time=slice('2008-01-01', '2008-12-31'),
            batch_size=self.args.batch_size,
            data_loader_kwargs={'num_workers': 4})

        train_loader = fluxdata.train_dataloader()
        val_loader = fluxdata.val_dataloader()
        test_loader = fluxdata.test_dataloader()

        # Create empty xr.Dataset, will be used by the model to save predictions every epoch.
        max_epochs = TRAINER_ARGS['max_epochs']
        ds_pred = fluxdata.target_xr('valid', varnames=['reco', 'rb'], num_epochs=max_epochs)

        # ------------
        # model
        # ------------
        model = Q10Model(
            features=features,
            targets=targets,
            norm=fluxdata._norm,
            ds=ds_pred,
            q10_init=q10_init,
            hidden_dim=self.args.hidden_dim,
            num_layers=self.args.num_layers,
            learning_rate=self.args.learning_rate,
            dropout=dropout,
            weight_decay=self.args.weight_decay,
            num_steps=len(train_loader) * max_epochs)

        # ------------
        # training
        # ------------
        trainer = pl.Trainer.from_argparse_args(
            self.args,
            default_root_dir=self.args.log_dir,
            **TRAINER_ARGS,
            callbacks=[
                EarlyStopping(
                    monitor='valid_loss',
                    patience=10,
                    min_delta=0.00001),
                ModelCheckpoint(
                    filename='{epoch}-{val_loss:.2f}',
                    save_top_k=1,
                    verbose=False,
                    monitor='valid_loss',
                    mode='min',
                    prefix=model.__class__.__name__)
            ])
        trainer.fit(model, train_loader, val_loader)

        # ------------
        # testing
        # ------------
        # trainer.test(test_dataloaders=test_loader)

        # ------------
        # save results
        # ------------
        # Store predictions.
        ds = fluxdata.add_scalar_record(model.ds, varname='q10', x=model.q10_history)
        trial.set_user_attr('q10', ds.q10[-1].item())

        # Add some attributes that are required for analysis.
        ds.attrs = {
            'created': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            'author': 'bkraft@bgc-jena.mpg.de',
            'q10_init': q10_init,
            'dropout': dropout,
            'use_ta': int(use_ta),
            'loss': trainer.callback_metrics['valid_loss'].item()
        }

        ds = ds.isel(epoch=slice(0, trainer.current_epoch + 1))

        # Save data.
        save_dir = os.path.join(model.logger.log_dir, 'predictions.nc')
        print(f'Saving predictions to: {save_dir}')
        ds.to_netcdf(save_dir)

        return trainer.callback_metrics['valid_loss'].item()

    @staticmethod
    def add_project_specific_args(parent_parser: ArgumentParser) -> ArgumentParser:
        parser = ArgumentParser(parents=[parent_parser], add_help=False)
        parser.add_argument(
            '--batch_size', default=240, type=int)
        parser.add_argument(
            '--data_path', default='./data/Synthetic4BookChap.nc', type=str)
        parser.add_argument(
            '--log_dir', default='./logs/experiment_01/', type=str)
        return parser


def main(parser: ArgumentParser = None, **kwargs):
    """Use kwargs to overload argparse args."""

    # ------------
    # args
    # ------------
    if parser is None:
        parser = ArgumentParser()

    parser = Objective.add_project_specific_args(parser)
    parser = pl.Trainer.add_argparse_args(parser)
    parser = Q10Model.add_model_specific_args(parser)
    parser.add_argument('--create_study', action='store_true', help='create new study (deletes old) and exits')
    parser.add_argument('--single_seed', action='store_true', help='use only one seed instead of (1, ..., 10).')
    args = parser.parse_args()

    globargs = TRAINER_ARGS.copy()
    globargs.update(kwargs)

    for k, v in globargs.items():
        setattr(args, k, v)

    # ------------
    # study setup
    # ------------
    search_space = {
        'q10_init': [0.5, 1.5, 2.5],
        'seed': [0] if args.single_seed else [i for i in range(10)],
        'dropout': [0.0, 0.2, 0.4, 0.6],
        'use_ta': [True, False]
    }

    sql_file = os.path.abspath(os.path.join(args.log_dir, "optuna.db"))
    sql_path = f'sqlite:///{sql_file}'

    if args.create_study | (not os.path.isfile(sql_file)):
        if os.path.isdir(args.log_dir):
            shutil.rmtree(args.log_dir)
        os.makedirs(args.log_dir, exist_ok=True)
        study = optuna.create_study(
            study_name="q10hybrid",
            storage=sql_path,
            sampler=optuna.samplers.GridSampler(search_space),
            direction='minimize',
            load_if_exists=False)

        if args.create_study:
            return None

    if not os.path.isdir(args.log_dir):
        os.makedirs(args.log_dir)

    # ------------
    # run study
    # ------------
    n_trials = 1
    for _, v in search_space.items():
        n_trials *= len(v)
    study = optuna.load_study(
        study_name="q10hybrid",
        storage=sql_path,
        sampler=optuna.samplers.GridSampler(search_space))
    study.optimize(Objective(args), n_trials=n_trials)


if __name__ == '__main__':
    main()

[((166, 13, 166, 49), 'pytorch_lightning.Trainer.add_argparse_args', 'pl.Trainer.add_argparse_args', ({(166, 42, 166, 48): 'parser'}, {}), '(parser)', True, 'import pytorch_lightning as pl\n'), ((167, 13, 167, 53), 'models.hybrid.Q10Model.add_model_specific_args', 'Q10Model.add_model_specific_args', ({(167, 46, 167, 52): 'parser'}, {}), '(parser)', False, 'from models.hybrid import Q10Model\n'), ((40, 8, 40, 32), 'pytorch_lightning.seed_everything', 'pl.seed_everything', ({(40, 27, 40, 31): 'seed'}, {}), '(seed)', True, 'import pytorch_lightning as pl\n'), ((54, 13, 54, 49), 'xarray.open_dataset', 'xr.open_dataset', ({(54, 29, 54, 48): 'self.args.data_path'}, {}), '(self.args.data_path)', True, 'import xarray as xr\n'), ((138, 19, 138, 71), 'os.path.join', 'os.path.join', ({(138, 32, 138, 52): 'model.logger.log_dir', (138, 54, 138, 70): '"""predictions.nc"""'}, {}), "(model.logger.log_dir, 'predictions.nc')", False, 'import os\n'), ((146, 17, 146, 72), 'argparse.ArgumentParser', 'ArgumentParser', (), '', False, 'from argparse import ArgumentParser\n'), ((163, 17, 163, 33), 'argparse.ArgumentParser', 'ArgumentParser', ({}, {}), '()', False, 'from argparse import ArgumentParser\n'), ((188, 31, 188, 70), 'os.path.join', 'os.path.join', ({(188, 44, 188, 56): 'args.log_dir', (188, 58, 188, 69): '"""optuna.db"""'}, {}), "(args.log_dir, 'optuna.db')", False, 'import os\n'), ((192, 11, 192, 38), 'os.path.isdir', 'os.path.isdir', ({(192, 25, 192, 37): 'args.log_dir'}, {}), '(args.log_dir)', False, 'import os\n'), ((194, 8, 194, 48), 'os.makedirs', 'os.makedirs', (), '', False, 'import os\n'), ((205, 11, 205, 38), 'os.path.isdir', 'os.path.isdir', ({(205, 25, 205, 37): 'args.log_dir'}, {}), '(args.log_dir)', False, 'import os\n'), ((206, 8, 206, 33), 'os.makedirs', 'os.makedirs', ({(206, 20, 206, 32): 'args.log_dir'}, {}), '(args.log_dir)', False, 'import os\n'), ((191, 32, 191, 56), 'os.path.isfile', 'os.path.isfile', ({(191, 47, 191, 55): 'sql_file'}, {}), '(sql_file)', False, 'import os\n'), ((193, 12, 193, 39), 'shutil.rmtree', 'shutil.rmtree', ({(193, 26, 193, 38): 'args.log_dir'}, {}), '(args.log_dir)', False, 'import shutil\n'), ((217, 16, 217, 57), 'optuna.samplers.GridSampler', 'optuna.samplers.GridSampler', ({(217, 44, 217, 56): 'search_space'}, {}), '(search_space)', False, 'import optuna\n'), ((198, 20, 198, 61), 'optuna.samplers.GridSampler', 'optuna.samplers.GridSampler', ({(198, 48, 198, 60): 'search_space'}, {}), '(search_space)', False, 'import optuna\n'), ((99, 16, 102, 38), 'pytorch_lightning.callbacks.early_stopping.EarlyStopping', 'EarlyStopping', (), '', False, 'from pytorch_lightning.callbacks.early_stopping import EarlyStopping\n'), ((103, 16, 109, 52), 'pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint', 'ModelCheckpoint', (), '', False, 'from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint\n'), ((127, 23, 127, 37), 'datetime.datetime.now', 'datetime.now', ({}, {}), '()', False, 'from datetime import datetime\n')]