import numpy as np
import h5py
import yaml
import os
from tqdm import tqdm
from scipy.ndimage import gaussian_filter
from data.ray_tracer import RayTracer
from data.accretion_disk import ThinDisk, ThickDisk


def compute_rs_from_solar_mass(mass_solar, config):
    G = float(config['physics']['G'])
    c = float(config['physics']['c'])
    M_sun = float(config['physics']['M_sun'])
    rs = 2 * G * mass_solar * M_sun / (c ** 2)
    return rs


def add_noise_and_fourier_artifacts(image, noise_std, uv_coverage):
    noisy = image + np.random.normal(0, noise_std, image.shape).astype(np.float32)
    noisy = np.clip(noisy, 0, None)
    fft = np.fft.fft2(noisy)
    mask = np.zeros_like(fft, dtype=bool)
    n = image.shape[0]
    num_sampled = int(uv_coverage * n * n)
    indices = np.random.choice(n * n, num_sampled, replace=False)
    flat_mask = np.zeros(n * n, dtype=bool)
    flat_mask[indices] = True
    mask = flat_mask.reshape(n, n)
    masked_fft = fft * mask
    reconstructed = np.abs(np.fft.ifft2(masked_fft)).astype(np.float32)
    smoothed = gaussian_filter(reconstructed, sigma=0.5)
    return smoothed


def normalize_image(image):
    min_val = image.min()
    max_val = image.max()
    if max_val - min_val < 1e-10:
        return np.zeros_like(image)
    return ((image - min_val) / (max_val - min_val)).astype(np.float32)


def generate_single_image(mass_solar, inclination_deg, disk_type, config):
    rs = compute_rs_from_solar_mass(mass_solar, config)
    observer_distance = float(config['data']['observer_distance_rs'])
    image_size = int(config['data']['image_size'])
    inner = float(config['data']['disk_inner_radius_rs'])
    outer = float(config['data']['disk_outer_radius_rs'])
    scale = float(config['data']['disk_brightness_scale'])

    if disk_type == 0:
        disk = ThinDisk(rs, inner, outer, scale)
    else:
        disk = ThickDisk(rs, inner, outer, scale)

    tracer = RayTracer(rs, observer_distance, image_size, inclination_deg)
    image = tracer.render(disk)
    return image, rs


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


def generate_dataset(config_path='configs/config.yaml'):
    import concurrent.futures
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)

    num_images = int(config['data']['num_images'])
    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])
    output_path = config['data']['output_path']

    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    masses_solar = np.exp(
        np.random.uniform(np.log(mass_min), np.log(mass_max), num_images)
    ).astype(np.float64)
    inclinations = np.random.uniform(incl_min, incl_max, num_images).astype(np.float32)
    disk_types = np.random.randint(0, 2, num_images)
    noise_stds = np.random.uniform(noise_min, noise_max, num_images).astype(np.float32)

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

    tasks = [
        (idx, masses_solar[idx], inclinations[idx], disk_types[idx], noise_stds[idx], config)
        for idx in range(num_images)
    ]

    with h5py.File(output_path, 'w') as hf:
        hf.create_dataset('images', shape=(num_images, 1, image_size, image_size), dtype=np.float32)
        hf.create_dataset('rs_meters', data=rs_values, dtype=np.float64)
        hf.create_dataset('mass_solar', data=masses_solar, dtype=np.float64)
        hf.create_dataset('inclination_deg', data=inclinations, dtype=np.float32)
        hf.create_dataset('disk_type', data=disk_types, dtype=np.int32)
        hf.create_dataset('noise_std', data=noise_stds, dtype=np.float32)

        print(f"Generating dataset in parallel using up to {os.cpu_count()} CPU cores...")
        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_images, desc='Generating dataset'):
                idx = futures[future]
                try:
                    _, normalized_image = future.result()
                    hf['images'][idx, 0] = normalized_image
                except Exception as exc:
                    print(f'Image {idx} generated an exception: {exc}')

    print(f'Dataset saved to {output_path}')
    print(f'Total images: {num_images}')
    print(f'Mass range: {mass_min:.2e} to {mass_max:.2e} solar masses')
    print(f'RS range: {rs_values.min():.4e} to {rs_values.max():.4e} meters')


if __name__ == '__main__':
    generate_dataset()