"""
Generate 7000 new images and merge with existing 3000 into a new 10k dataset.
Copies existing data from the old file to avoid re-computing.
"""

import numpy as np
import h5py
import yaml
import os
import shutil
import concurrent.futures
from tqdm import tqdm
from data.generate_dataset import (
    compute_rs_from_solar_mass,
    generate_single_image,
    add_noise_and_fourier_artifacts,
    log1p_transform,
)


def worker_generate_image(args):
    idx, mass_solar, inclination_deg, disk_type, noise_std, fov_factor, config = args
    raw_image, rs = generate_single_image(mass_solar, inclination_deg, disk_type, config, fov_factor=fov_factor)
    noisy_image = add_noise_and_fourier_artifacts(raw_image, noise_std, float(config['data']['fourier_uv_coverage']))
    normalized_image = log1p_transform(noisy_image)
    return idx, normalized_image


def expand_dataset(num_new=7000, config_path='configs/config.yaml'):
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)

    old_path = config['data']['output_path']
    new_path = old_path.replace('.h5', '_10k.h5')
    image_size = int(config['data']['image_size'])
    mass_min = float(config['data']['mass_range_solar'][0])
    mass_max = float(config['data']['mass_range_solar'][1])
    incl_min = float(config['data']['inclination_range_deg'][0])
    incl_max = float(config['data']['inclination_range_deg'][1])
    noise_min = float(config['data']['noise_std_range'][0])
    noise_max = float(config['data']['noise_std_range'][1])

    # Read existing data
    with h5py.File(old_path, 'r') as hf_old:
        existing_count = hf_old['images'].shape[0]
        print(f'Existing dataset: {existing_count} images')

        total = existing_count + num_new
        print(f'Target: {total} images ({num_new} new)')

        # Sample parameters for new images
        masses_solar = np.exp(
            np.random.uniform(np.log(mass_min), np.log(mass_max), num_new)
        ).astype(np.float64)
        inclinations = np.random.uniform(incl_min, incl_max, num_new).astype(np.float32)
        disk_types = np.random.randint(0, 2, num_new)
        noise_stds = np.random.uniform(noise_min, noise_max, num_new).astype(np.float32)
        fov_factors = np.random.uniform(3.0, 15.0, num_new).astype(np.float32)

        rs_values = np.array([
            compute_rs_from_solar_mass(m, config) for m in masses_solar
        ], dtype=np.float64)

        # Create new file with combined data
        with h5py.File(new_path, 'w') as hf_new:
            # Create datasets at full size
            hf_new.create_dataset('images', shape=(total, 1, image_size, image_size), dtype=np.float32)
            hf_new.create_dataset('rs_meters', shape=(total,), dtype=np.float64)
            hf_new.create_dataset('mass_solar', shape=(total,), dtype=np.float64)
            hf_new.create_dataset('inclination_deg', shape=(total,), dtype=np.float32)
            hf_new.create_dataset('disk_type', shape=(total,), dtype=np.int32)
            hf_new.create_dataset('noise_std', shape=(total,), dtype=np.float32)
            hf_new.create_dataset('fov_factor', shape=(total,), dtype=np.float32)

            # Copy existing data
            print('Copying existing 3000 images...')
            hf_new['images'][:existing_count] = hf_old['images'][:]
            hf_new['rs_meters'][:existing_count] = hf_old['rs_meters'][:]
            hf_new['mass_solar'][:existing_count] = hf_old['mass_solar'][:]
            hf_new['inclination_deg'][:existing_count] = hf_old['inclination_deg'][:]
            hf_new['disk_type'][:existing_count] = hf_old['disk_type'][:]
            hf_new['noise_std'][:existing_count] = hf_old['noise_std'][:]
            hf_new['fov_factor'][:existing_count] = hf_old['fov_factor'][:]

            # Write new metadata
            hf_new['rs_meters'][existing_count:] = rs_values
            hf_new['mass_solar'][existing_count:] = masses_solar
            hf_new['inclination_deg'][existing_count:] = inclinations
            hf_new['disk_type'][existing_count:] = disk_types
            hf_new['noise_std'][existing_count:] = noise_stds
            hf_new['fov_factor'][existing_count:] = fov_factors

            # Generate new images in parallel
            tasks = [
                (idx, masses_solar[idx], inclinations[idx], disk_types[idx],
                 noise_stds[idx], fov_factors[idx], config)
                for idx in range(num_new)
            ]

            print(f'Generating {num_new} new images in parallel...')
            with concurrent.futures.ProcessPoolExecutor() as executor:
                futures = {executor.submit(worker_generate_image, task): task[0] for task in tasks}
                for future in tqdm(concurrent.futures.as_completed(futures), total=num_new, desc='Generating'):
                    idx = futures[future]
                    try:
                        _, normalized_image = future.result()
                        hf_new['images'][existing_count + idx, 0] = normalized_image
                    except Exception as exc:
                        print(f'Image {idx} generated an exception: {exc}')

    # Swap files
    backup_path = old_path.replace('.h5', '_3k_backup.h5')
    print(f'Backing up old dataset to {backup_path}')
    shutil.move(old_path, backup_path)
    shutil.move(new_path, old_path)

    print(f'\nDone! Dataset expanded: {existing_count} → {total} images')
    print(f'Saved to {old_path} (backup at {backup_path})')


if __name__ == '__main__':
    expand_dataset()