"""Load the selected checkpoint and query the three paper snapshot grids."""
from pathlib import Path
import sys
import time

import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from model import PINNShallowWater
from common import arguments, device_for, errors, load_config, output_path, runtime_info, save_json, sha256
from fake_data import initial_values
from train import to_device, validation, weights


def main():
    args = arguments(__doc__).parse_args()
    config, root = load_config(args.config)
    device = device_for(config)
    torch.set_num_threads(8)
    checkpoint_path = output_path(root, config['paths']['best_checkpoint'])
    checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
    if checkpoint.get('completed_iterations') != 5000 or not checkpoint.get('formal_run_complete'):
        raise ValueError('Formal 5000-iteration training has not completed.')
    if checkpoint['config'] != config or checkpoint['model_class'] != 'PINNShallowWater':
        raise ValueError('Checkpoint/config/model identity mismatch.')
    model = PINNShallowWater(**config['model']).to(device)
    model.load_state_dict(checkpoint['model_state_dict'], strict=True)
    model.eval()
    with np.load(output_path(root, config['paths']['grid'])) as grid:
        x, y, xy, times = [grid[k] for k in ['x', 'y', 'xy', 'times']]
    batch_size = config['evaluation']['inference_batch_size']
    snapshots = []
    started = time.monotonic()
    with torch.no_grad():
        for t in times:
            coordinates = np.column_stack((xy, np.full(len(xy), t))).astype(np.float32)
            batches = []
            for begin in range(0, len(coordinates), batch_size):
                prediction = model(torch.as_tensor(coordinates[begin:begin + batch_size], device=device))
                batches.append(prediction.cpu().numpy())
            field = np.concatenate(batches).reshape(len(y), len(x), 3)
            if not np.isfinite(field).all():
                raise FloatingPointError(f'Nonfinite prediction at t={t}')
            snapshots.append(field)
            print('INFERENCE_SNAPSHOT', {'t': float(t), 'shape': field.shape,
                  'min_huv': field.reshape(-1, 3).min(axis=0).tolist(),
                  'max_huv': field.reshape(-1, 3).max(axis=0).tolist()}, flush=True)
    fields = np.stack(snapshots)
    path = output_path(root, config['paths']['predictions'])
    np.savez_compressed(path, x=x, y=y, times=times, fields=fields,
                        variables=np.array(['h', 'u', 'v']), units=np.array(['m', 'm/s', 'm/s']))
    initial_errors = errors(fields[0].reshape(-1, 3), initial_values(xy, config['physics']))
    samples_path = output_path(root, config['paths']['samples'])
    if checkpoint['sampling_sha256'] != sha256(samples_path):
        raise ValueError('Checkpoint sampling hash mismatch.')
    with np.load(samples_path) as arrays:
        val = to_device(arrays, 'val_', device)
    check = validation(model, val, config, weights(config, device))
    metadata = {'status': 'completed', 'checkpoint': str(checkpoint_path),
                'checkpoint_sha256': sha256(checkpoint_path), 'selected_iteration': checkpoint['iteration'],
                'completed_training_iterations': 5000, 'state_dict_strict_load': True,
                'fields_shape': list(fields.shape), 'all_finite': True, 'initial_grid_errors': initial_errors,
                'best_checkpoint_validation': check, 'runtime': runtime_info(device),
                'elapsed_seconds': time.monotonic() - started, 'predictions_sha256': sha256(path)}
    save_json(path.parent / 'inference_summary.json', metadata)
    print('INFERENCE_COMPLETE', metadata, flush=True)


if __name__ == '__main__':
    main()