import json
import uuid

from utils.types import JSON

base_data: str = 'PATH/samples/imagenet-64x64'
base_results: str = 'results'
num_sample_steps: int = 40
num_sample_steps_train: int = 20

model_load_path: str = 'PATH/results/lines/lines-imagenet-64x64-3-loss-lpips-vgg-64-bs-64-br-2-6e73977c-3d1a-47bc-8ba5-98127a2ab937/checkpoints/last/index.pth'

def get_config(run_id: str) -> JSON:
    return {
        'log_filepath': f'{base_results}/{run_id}/log/log.log',
        'tensorboard_log_dir': f'{base_results}/{run_id}/tensorboard/logs',
        'train_steps': 100_000_000_000,
        'batch_size': 32,
        'batch_repeats': 64,
        'data_loader_workers': 4,
        'dataset': {'dataset_name': 'imagenet-64x64'},
        'report': {'train_steps': 25, 'test_steps': None},
        'checkpoint': {
            'folder': f'{base_results}/{run_id}/checkpoints',
            'save_steps': 15_000,
            'last_steps': 5_000,
            'save_time': None,
            'last_time': 6 * 60 * 60
        },
        'edm_scheduler': {},
        'edm_sampler': {'num_steps': num_sample_steps},
        'model': {
            'name': 'edm-imagenet-64x64-cond-adm',
            'load_path': model_load_path,
            'load_keys': ['model']
        },
        'model_optimizer': {'name': 'r-adam', 'learning_rate': 0.000008},
        'ema': [
            {
                'beta': 0.999,
                'update_every': 1,
                'load_path': model_load_path,
                'load_keys': ['ema', 'beta_0.999_update_every_1_update_after_step_100_inv_gamma_1.0_power_0.66667']
            },
            {
                'beta': 0.9999,
                'update_every': 1,
                'load_path': model_load_path,
                'load_keys': ['ema', 'beta_0.9999_update_every_1_update_after_step_100_inv_gamma_1.0_power_0.66667']
            }
        ],
        'fid': {
            'steps': 5_000,
            'batch_size': 25,
            'reference_path': 'PATH/fid-refs/imagenet-64x64.npz'
        },
        'reconstruction_loss': {'name': 'lpips-vgg', 'params': {'size': 64}},
        'amp': {
            'use_fp16': True,
            'use_autocast': False,
            'use_gard_scaler': True
        },
        'train_dataset': [
            {
                'noise': {
                    'folder': f'{base_data}/noises',
                    'num_samples': 4_000_000
                },
                'image': {
                    'folder': f'{base_data}/edm-imagenet-64x64-cond-adm',
                    'num_samples': 4_000_000
                },
                'label': {
                    'folder': f'{base_data}/labels',
                    'num_samples': 4_000_000
                },
                'time_step': i
            } for i in range(num_sample_steps_train)
        ],
        'test_dataset': [
            {
                'noise': {
                    'folder': f'{base_data}/noises',
                    'num_samples': 10_000,
                    'start_index': 5_000_000
                },
                'image': {
                    'folder': f'{base_data}/edm-imagenet-64x64-cond-adm',
                    'num_samples': 10_000,
                    'start_index': 5_000_000
                },
                'label': {
                    'folder': f'{base_data}/labels',
                    'num_samples': 10_000,
                    'start_index': 5_000_000
                },
                'time_step': 0
            }
        ],
        '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': 5_000_000
                },
                'label': {
                    'folder': f'{base_data}/labels',
                    'num_samples': 50_000,
                    'start_index': 5_000_000
                }
            }
        ]
    }


def main() -> None:
    name: str = 'lines-imagenet-64x64-5-loss-lpips-vgg-64-bs-32-br-64'
    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()
