import pytest
import io
import sys
from omegaconf import DictConfig

from src.utils.utils import _check_if_in_config, REQUIRED_CONFIGS, check_config, print_config


@pytest.fixture
def get_dict():
    return DictConfig({'plugins': {
        'ddp_plugin': {'_target_': 'pytorch_lightning.plugins.DDPPlugin', 'find_unused_parameters': False}}, 'task': {
        '_target_': 'src.tasks.semantic_segmentation.semantic_segmentation.SemanticSegmentation',
        'confusion_matrix_log_every_n_epoch': 1, 'confusion_matrix_val': True, 'confusion_matrix_test': True},
        'loss': {'_target_': 'torch.nn.CrossEntropyLoss'},
        'metric': {'_target_': 'src.metrics.divahisdb.HisDBIoU',
                   'num_classes': '5'}, 'model': {
            'backbone': {'_target_': 'pl_bolts.models.vision.UNet', 'num_classes': '5',
                         'num_layers': 2, 'features_start': 32}, 'header': {'_target_': 'torch.nn.Identity'}},
        'optimizer': {'_target_': 'torch.optim.Adam', 'lr': 0.001, 'betas': [0.9, 0.999], 'eps': 1e-08,
                      'weight_decay': 0, 'amsgrad': False}, 'callbacks': {
            'check_backbone_header_compatibility': {
                '_target_': 'src.callbacks.model_callbacks.CheckBackboneHeaderCompatibility'},
            'model_checkpoint': {'_target_': 'src.callbacks.model_callbacks.SaveModelStateDictAndTaskCheckpoint',
                                 'monitor': 'val/crossentropyloss', 'save_top_k': 1, 'save_last': True, 'mode': 'min',
                                 'verbose': False, 'dirpath': 'checkpoints/',
                                 'filename': 'dev-baby-unet-cb55-10',
                                 'backbone_filename': 'backbone',
                                 'header_filename': 'header'},
            'watch_model': {'_target_': 'src.callbacks.wandb_callbacks.WatchModelWithWandb', 'log': 'all',
                            'log_freq': 1}}, 'logger': {
            'wandb': {'_target_': 'pytorch_lightning.loggers.wandb.WandbLogger', 'project': 'unsupervised',
                      'name': 'dev-baby-unet-cb55-10', 'offline': False, 'job_type': 'train', 'group': 'dev-runs',
                      'tags': ['best_model', 'USL'], 'save_dir': '.', 'log_model': False, 'notes': 'Testing'},
            'csv': {'_target_': 'pytorch_lightning.loggers.csv_logs.CSVLogger', 'save_dir': '.', 'name': 'csv/'}},
        'seed': 42, 'train': True, 'test': True,
        'trainer': {'_target_': 'pytorch_lightning.Trainer', 'gpus': -1, 'accelerator': 'ddp',
                    'min_epochs': 1, 'max_epochs': 3, 'enable_model_summary': 'full', 'precision': 16},
        'datamodule': {
            '_target_': 'src.datamodules.hisDBDataModule.DIVAHisDBDataModule.DIVAHisDBDataModuleCropped',
            'data_dir': '/net/research-hisdoc/datasets/semantic_segmentation/datasets_cropped/CB55-10-segmentation',
            'crop_size': 256, 'num_workers': 4, 'batch_size': 16, 'shuffle': True, 'drop_last': True},
        'save_config': True, 'checkpoint_folder_name': '{epoch}/', 'work_dir': '.',
        'debug': False, 'print_config': True, 'disable_warnings': True})


def test_check_config_everything_good(get_dict):
    check_config(get_dict)
    assert not get_dict['debug']
    assert get_dict['metric'] == {'_target_': 'src.metrics.divahisdb.HisDBIoU',
                                  'num_classes': '5'}
    assert get_dict['train']
    assert get_dict['test']
    assert 'predict' not in get_dict
    assert get_dict['seed'] == 42


def test_check_config_no_seed(get_dict, caplog):
    del get_dict['seed']
    check_config(get_dict)
    assert get_dict['seed'] > -1
    assert f'No seed specified! Seed set to {get_dict["seed"]}'


def test_check_config_fast_dev_run(get_dict, caplog):
    get_dict['trainer']['fast_dev_run'] = True
    check_config(get_dict)
    assert get_dict['trainer']['gpus'] == 0
    assert get_dict['datamodule']['num_workers'] == 0
    assert "Forcing debugger friendly configuration! <config.trainer.fast_dev_run=True>" in caplog.text


def test_check_config_ddp_cpu_and_precision(get_dict, caplog):
    get_dict['trainer']['accelerator'] = 'cpu'
    get_dict['trainer']['precision'] = 16
    check_config(get_dict)
    assert 'You are using ddp_cpu without precision=16. This can lead to a crash! Use 64 or 32!' in caplog.text


def test__check_if_in_config_good_config(get_dict):
    for cf in REQUIRED_CONFIGS:
        _check_if_in_config(config=get_dict, name=cf)


def test__check_if_in_config_bad_config(get_dict):
    del get_dict['datamodule']
    with pytest.raises(ValueError):
        for cf in REQUIRED_CONFIGS:
            _check_if_in_config(config=get_dict, name=cf)


def test_print_config(get_dict, capsys):
    print_config(config=get_dict, fields=(
        "trainer",
        "task",
        "model",
        "optimizer",
        "datamodule",
        "callbacks",
        "loss",
        "metric",
        "logger",
        "seed",
        "train",
        "test",
        "predict"
    ))
    expected_result = \
        "⚙ CONFIG\n" \
        "├── trainer\n" \
        "│   └── _target_: pytorch_lightning.Trainer\n" \
        "│       gpus: -1\n" \
        "│       accelerator: ddp\n" \
        "│       min_epochs: 1\n" \
        "│       max_epochs: 3\n" \
        "│       enable_model_summary: full\n" \
        "│       precision: 16\n" \
        "│       \n" \
        "├── task\n" \
        "│   └── _target_: src.tasks.semantic_segmentation.semantic_segmentation.Semantic\n" \
        "│       confusion_matrix_log_every_n_epoch: 1\n" \
        "│       confusion_matrix_val: true\n" \
        "│       confusion_matrix_test: true\n" \
        "│       \n" \
        "├── model\n" \
        "│   └── backbone:\n" \
        "│         _target_: pl_bolts.models.vision.UNet\n" \
        "│         num_classes: '5'\n" \
        "│         num_layers: 2\n" \
        "│         features_start: 32\n" \
        "│       header:\n" \
        "│         _target_: torch.nn.Identity\n" \
        "│       \n" \
        "├── optimizer\n" \
        "│   └── _target_: torch.optim.Adam\n" \
        "│       lr: 0.001\n" \
        "│       betas:\n" \
        "│       - 0.9\n" \
        "│       - 0.999\n" \
        "│       eps: 1.0e-08\n" \
        "│       weight_decay: 0\n" \
        "│       amsgrad: false\n" \
        "│       \n" \
        "├── datamodule\n" \
        "│   └── _target_: src.datamodules.hisDBDataModule.DIVAHisDBDataModule.DIVAHisDBD\n" \
        "│       data_dir: /net/research-hisdoc/datasets/semantic_segmentation/datasets_c\n" \
        "│       crop_size: 256\n" \
        "│       num_workers: 4\n" \
        "│       batch_size: 16\n" \
        "│       shuffle: true\n" \
        "│       drop_last: true\n" \
        "│       \n" \
        "├── callbacks\n" \
        "│   └── check_backbone_header_compatibility:\n" \
        "│         _target_: src.callbacks.model_callbacks.CheckBackboneHeaderCompatibili\n" \
        "│       model_checkpoint:\n" \
        "│         _target_: src.callbacks.model_callbacks.SaveModelStateDictAndTaskCheck\n" \
        "│         monitor: val/crossentropyloss\n" \
        "│         save_top_k: 1\n" \
        "│         save_last: true\n" \
        "│         mode: min\n" \
        "│         verbose: false\n" \
        "│         dirpath: checkpoints/\n" \
        "│         filename: dev-baby-unet-cb55-10\n" \
        "│         backbone_filename: backbone\n" \
        "│         header_filename: header\n" \
        "│       watch_model:\n" \
        "│         _target_: src.callbacks.wandb_callbacks.WatchModelWithWandb\n" \
        "│         log: all\n" \
        "│         log_freq: 1\n" \
        "│       \n" \
        "├── loss\n" \
        "│   └── _target_: torch.nn.CrossEntropyLoss\n" \
        "│       \n" \
        "├── metric\n" \
        "│   └── _target_: src.metrics.divahisdb.HisDBIoU\n" \
        "│       num_classes: '5'\n" \
        "│       \n" \
        "├── logger\n" \
        "│   └── wandb:\n" \
        "│         _target_: pytorch_lightning.loggers.wandb.WandbLogger\n" \
        "│         project: unsupervised\n" \
        "│         name: dev-baby-unet-cb55-10\n" \
        "│         offline: false\n" \
        "│         job_type: train\n" \
        "│         group: dev-runs\n" \
        "│         tags:\n" \
        "│         - best_model\n" \
        "│         - USL\n" \
        "│         save_dir: .\n" \
        "│         log_model: false\n" \
        "│         notes: Testing\n" \
        "│       csv:\n" \
        "│         _target_: pytorch_lightning.loggers.csv_logs.CSVLogger\n" \
        "│         save_dir: .\n" \
        "│         name: csv/\n" \
        "│       \n" \
        "├── seed\n" \
        "│   └── 42\n" \
        "├── train\n" \
        "│   └── True\n" \
        "├── test\n" \
        "│   └── True\n" \
        "├── predict\n" \
        "│   └── None\n" \
        "├── checkpoint_folder_name\n" \
        "│   └── {epoch}/\n" \
        "├── debug\n" \
        "│   └── False\n" \
        "├── disable_warnings\n" \
        "│   └── True\n" \
        "├── plugins\n" \
        "│   └── ddp_plugin:\n" \
        "│         _target_: pytorch_lightning.plugins.DDPPlugin\n" \
        "│         find_unused_parameters: false\n" \
        "│       \n" \
        "├── print_config\n" \
        "│   └── True\n" \
        "├── save_config\n" \
        "│   └── True\n" \
        "└── work_dir\n" \
        "    └── .\n"
    caps = capsys.readouterr()
    assert caps.out == expected_result
