import json
import uuid

import numpy as np

from utils.types import JSON

base_data: str = 'PATH/cifar10-32x32/second-phase'
base_real_data: str = 'PATH'
base_results: str = 'results'
num_sample_steps: int = 18

base_load_path: str = 'PATH/results/lines-gan-cifar10-32x32-2-bs-32-br-16-pretrained-discriminator-a1f6c7f6-e641-4299-9f06-93da7ba955aa/checkpoints/last'
train_steps: list[int] = [0, 1]


def get_config(run_id: str) -> JSON:
    return {
        'base_folder': f'{base_results}/{run_id}',
        'train_steps': 100_000_000_000,
        'batch_size': 32,
        'batch_repeats': 16,
        'data_loader_workers': 4,
        'validate_ddp_consistency_steps': 200,
        'dataset': {'dataset_name': 'cifar10-32x32'},
        'report': {'train_steps': 20, 'test_steps': None},
        'checkpoint': {
            'save_steps': 20_000,
            'last_steps': 1_000,
        },
        'edm_scheduler': {},
        'edm_sampler': {'num_steps': num_sample_steps},
        'model': {
            'name': 'edm-cifar10-32x32-cond-vp',
            'load_path': f'{base_load_path}/model.pth.tmp',
        },
        'model_optimizer': {
            'name': 'r-adam',
            'learning_rate': 0.0004
        },
        'discriminator': {
            'discriminator_load_path': f'{base_load_path}/discriminator.pth.tmp'
        },
        'discriminator_optimizer': {
            'name': 'r-adam',
            'learning_rate': 0.002,
            'params': {'betas': (0.5, 0.9)},
        },
        'ema': [
            {
                'beta': 0.999,
                'update_every': 1,
                'load_path': f'{base_load_path}/ema.tmp/beta_0.999_update_every_1_update_after_step_100_inv_gamma_1.0_power_0.66667.pth',
            },
            {
                'beta': 0.9999,
                'update_every': 1,
                'load_path': f'{base_load_path}/ema.tmp/beta_0.9999_update_every_1_update_after_step_100_inv_gamma_1.0_power_0.66667.pth',
            },
            {
                'beta': 0.999,
                'update_every': 1,
                'power': 0.75,
                'load_path': f'{base_load_path}/ema.tmp/beta_0.999_update_every_1_update_after_step_100_inv_gamma_1.0_power_0.75.pth',
            },
            {
                'beta': 0.9999,
                'update_every': 1,
                'power': 0.75,
                'load_path': f'{base_load_path}/ema.tmp/beta_0.9999_update_every_1_update_after_step_100_inv_gamma_1.0_power_0.75.pth',
            }
        ],
        'fid': {
            'steps': 2_500,
            'batch_size': 50,
            'reference_path': 'PATH/cifar10-32x32-train.npz'
        },
        'reconstruction_loss': {'name': 'lpips-vgg', 'params': {'size': 224}},
        'reconstruction_lambda': 0.5,
        'amp': {},
        'distributed_sampler_seed': np.random.randint(0, 2 ** 31),
        'distributed_sampler_real_seed': np.random.randint(0, 2 ** 31),
        'train_dataset': [
            {
                'noise': {
                    'folder': f'{base_data}/noises',
                    'num_samples': 1_000_000
                },
                'image': {
                    'folder': f'{base_data}/20caa972-e407-4143-abb7-dc0dcad730d1_[0,1,18]',
                    'num_samples': 1_000_000
                },
                'label': {
                    'folder': f'{base_data}/labels',
                    'num_samples': 1_000_000
                },
                'time_step': i
            } for i in train_steps
        ],
        'test_dataset': [
            {
                'noise': {
                    'folder': f'{base_data}/noises',
                    'num_samples': 10_000,
                    'start_index': 1_000_000
                },
                'image': {
                    'folder': f'{base_data}/20caa972-e407-4143-abb7-dc0dcad730d1_[0,1,18]',
                    'num_samples': 10_000,
                    'start_index': 1_000_000
                },
                'label': {
                    'folder': f'{base_data}/labels',
                    'num_samples': 10_000,
                    'start_index': 1_000_000
                },
                'time_step': 0
            }
        ],
        'train_real_dataset': [
            {
                'image': {
                    'folder': f'{base_real_data}/train/images',
                    'num_samples': 50_000
                },
                'label': {
                    'folder': f'{base_real_data}/train/labels',
                    'num_samples': 50_000
                }
            }
        ],
        'test_real_dataset': [
            {
                'image': {
                    'folder': f'{base_real_data}/test/images',
                    'num_samples': 10_000
                },
                'label': {
                    'folder': f'{base_real_data}/test/labels',
                    'num_samples': 10_000
                }
            }
        ],
        'fid_train_dataset': [
            {
                'noise': {
                    'folder': f'{base_data}/noises',
                    'num_samples': 50_000
                },
                'label': {
                    'folder': f'{base_data}/labels',
                    'num_samples': 50_000
                }
            }
        ],
        'fid_test_dataset': [
            {
                'noise': {
                    'folder': f'{base_data}/noises',
                    'num_samples': 50_000,
                    'start_index': 1_000_000
                },
                'label': {
                    'folder': f'{base_data}/labels',
                    'num_samples': 50_000,
                    'start_index': 1_000_000
                }
            }
        ]
    }


def main() -> None:
    name: str = 'lines-gan-cifar10-32x32-second-phase-middle-1-20caa972-e407-4143-abb7-dc0dcad730d1_[0,1,18]-6-bs-32-br-16'
    config: JSON = get_config(f'{name}-{uuid.uuid4()}')
    print(config)
    with open(f'configs/{name}-{uuid.uuid4()}.json', 'w') as file:
        json.dump(config, file)


if __name__ == '__main__':
    main()
