"""Generate dCDB9 collocation points and grids, with no CFD training data.

The filename follows the requested package interface. The generated data are
physical coordinates and exact initial conditions, not fabricated CFD fields.
"""
import numpy as np

from common import arguments, load_config, output_path, save_json, sha256


def initial_values(points, physics):
    target = np.zeros((len(points), 3), dtype=np.float32)
    inside = np.sum(points[:, :2] ** 2, axis=1) <= physics['dam_radius'] ** 2
    target[:, 0] = np.where(inside, physics['depth_inside'], physics['depth_outside'])
    return target


def sample_points(counts, physics, seed):
    rng = np.random.default_rng(seed)
    lower = np.array([physics['x_range'][0], physics['y_range'][0], physics['t_range'][0]])
    upper = np.array([physics['x_range'][1], physics['y_range'][1], physics['t_range'][1]])
    interior = rng.uniform(lower, upper, (counts['interior'], 3)).astype(np.float32)
    initial = rng.uniform(lower, upper, (counts['initial'], 3)).astype(np.float32)
    initial[:, 2] = physics['t_range'][0]
    if counts['boundary'] % 4:
        raise ValueError('Boundary count must be divisible by four for per-side losses.')
    per_side = counts['boundary'] // 4
    boundary = rng.uniform(lower, upper, (counts['boundary'], 3)).astype(np.float32)
    pair = boundary.copy()
    for side in range(4):
        axis, high = side // 2, side % 2
        sl = slice(side * per_side, (side + 1) * per_side)
        boundary[sl, axis] = upper[axis] if high else lower[axis]
        pair[sl, axis] = lower[axis] if high else upper[axis]
    return {'interior': interior, 'initial': initial,
            'initial_target': initial_values(initial, physics),
            'boundary': boundary, 'boundary_pair': pair,
            'boundary_side': np.repeat(np.arange(4, dtype=np.int64), per_side)}


def generate(config, root):
    samples_path = output_path(root, config['paths']['samples'])
    grid_path = output_path(root, config['paths']['grid'])
    samples_path.parent.mkdir(parents=True, exist_ok=True)
    train = sample_points(config['sampling'], config['physics'], config['seed'])
    val = sample_points(config['validation'], config['physics'], config['seed'] + 1)
    arrays = {**train, **{'val_' + k: v for k, v in val.items()}}
    np.savez_compressed(samples_path, **arrays)
    spacing = config['evaluation']['grid_spacing']
    xmin, xmax = config['physics']['x_range']
    ymin, ymax = config['physics']['y_range']
    nx, ny = round((xmax - xmin) / spacing), round((ymax - ymin) / spacing)
    x = xmin + (np.arange(nx) + 0.5) * spacing
    y = ymin + (np.arange(ny) + 0.5) * spacing
    xx, yy = np.meshgrid(x, y, indexing='xy')
    np.savez_compressed(grid_path, x=x, y=y, xy=np.column_stack((xx.ravel(), yy.ravel())),
                        times=np.array(config['evaluation']['times']), spacing=spacing)
    manifest = {'source': 'paper Eq.55 and periodic boundary geometry only',
                'external_training_data': False, 'seed': config['seed'],
                'train': {k: list(v.shape) for k, v in train.items()},
                'validation': {k: list(v.shape) for k, v in val.items()},
                'grid_shape': [ny, nx], 'samples_sha256': sha256(samples_path),
                'grid_sha256': sha256(grid_path), 'units': ['m', 'm', 's'],
                'initial_target_units': ['m', 'm/s', 'm/s']}
    save_json(samples_path.parent / 'sampling_manifest.json', manifest)
    print('SAMPLING', manifest, flush=True)
    return samples_path


if __name__ == '__main__':
    args = arguments(__doc__).parse_args()
    config, root = load_config(args.config)
    generate(config, root)