import numpy as np
import os
from math import atan, sin, cos, cbrt, pi, sqrt
from functions import (
    g0,
    r_air,
    r_He,
    H,
    K,
    P,
    T,
    gas_parameters_altitude,
    sphere_drag,
    calculate_air_viscosity,
    calculate_gas_density,
    calculate_reynolds_number,
    calculate_specific_buoyancy,
    calculate_radius_sphere,
    calculate_area_circle,
    microhole_formation,
    gas_leakage_micro_holes,
    gas_leakage_permeability,
)


def save_npz(directory, filename, **arrays):
    # Create the directory if it doesn't exist
    os.makedirs(directory, exist_ok=True)
    # Save the arrays to a .npz file
    np.savez_compressed(os.path.join(directory, filename), **arrays)


# Generate and save gas parameters
altitudes = np.linspace(0, 65000, 100)
pressures, temperatures = gas_parameters_altitude(altitudes)
save_npz(
    "tests/data",
    "gas_parameters_altitude.npz",
    altitudes=altitudes,
    pressures=pressures,
    temperatures=temperatures,
)

# Generate and save sphere drag coefficients
reynolds_numbers = np.logspace(-1, 7, 100)
reference_drag_coeffs = np.array([sphere_drag(Re) for Re in reynolds_numbers])
save_npz(
    "tests/data",
    "sphere_drag.npz",
    reynolds_numbers=reynolds_numbers,
    reference_drag_coeffs=reference_drag_coeffs,
)

# Generate and save air viscosities
temperatures = np.linspace(10, 300, 100)
reference_viscosities = np.array(
    [calculate_air_viscosity(temp) for temp in temperatures]
)
save_npz(
    "tests/data",
    "air_viscosity.npz",
    temperatures=temperatures,
    reference_viscosities=reference_viscosities,
)

# Generate and save gas densities
temperatures = np.linspace(200, 300, 100)
pressures = np.linspace(80000, 101325, 100)
specific_gas_constant = np.array([r_air, r_He])
reference_densities = np.array(
    [
        calculate_gas_density(temp, press, spec_gas_constant)
        for temp, press, spec_gas_constant in zip(
            temperatures, pressures, specific_gas_constant
        )
    ]
)
save_npz(
    "tests/data",
    "gas_density.npz",
    temperatures=temperatures,
    pressures=pressures,
    reference_densities=reference_densities,
)

# Generate and save Reynolds numbers
temperatures = np.linspace(200, 300, 100)
pressures = np.linspace(0.1, 100000, 100)
velocities = np.linspace(0.5, 50, 100)
lengths = np.linspace(0.1, 100, 100)
reference_reynolds = np.array(
    [
        calculate_reynolds_number(temp, press, vel, length)
        for temp, press, vel, length in zip(
            temperatures, pressures, velocities, lengths
        )
    ]
)
save_npz(
    "tests/data",
    "reynolds_number.npz",
    temperatures=temperatures,
    pressures=pressures,
    velocities=velocities,
    lengths=lengths,
    reference_reynolds=reference_reynolds,
)

# Generate and save specific buoyancies
temperatures = np.linspace(50, 300, 100)
pressures = np.linspace(0.1, 100000, 100)
pressure_differences = np.linspace(0, 500, 100)
reference_buoyancies = np.array(
    [
        calculate_specific_buoyancy(temp, press, press_diff)
        for temp, press, press_diff in zip(
            temperatures, pressures, pressure_differences
        )
    ]
)
save_npz(
    "tests/data",
    "specific_buoyancy.npz",
    temperatures=temperatures,
    pressures=pressures,
    pressure_differences=pressure_differences,
    reference_buoyancies=reference_buoyancies,
)

# Generate and save sphere radii
volumes = np.linspace(0.1, 500000, 100)
reference_radii = np.array([cbrt(3 * vol / (4 * pi)) for vol in volumes])
save_npz(
    "tests/data", "radius_sphere.npz", volumes=volumes, reference_radii=reference_radii
)

# Generate and save circle areas
radii = np.linspace(0.1, 200, 100)
reference_areas = np.array([2 * pi * rad for rad in radii])
save_npz("tests/data", "area_circle.npz", radii=radii, reference_areas=reference_areas)

# Generate and save microhole formation data
formation_rate = np.linspace(0, 10, 100)  # units per time
age_of_film = np.linspace(0, 100, 100)  # time units
area_holes_initial = np.linspace(0, 1, 10)  # units
reference_hole_areas = np.array(
    [
        microhole_formation(form_rate, age, area_hol_init)
        for form_rate, age, area_hol_init in zip(
            formation_rate, age_of_film, area_holes_initial
        )
    ]
)
save_npz(
    "tests/data",
    "microhole_formation.npz",
    formation_rate=formation_rate,
    age_of_film=age_of_film,
    area_holes_initial=area_holes_initial,
    reference_hole_areas=reference_hole_areas,
)

# Generate and save gas leakage through microholes data
discharge_coefficients = np.linspace(0.01, 1.2, 100)
micro_hole_areas = np.linspace(1e-12, 0.1, 100)  # m^2
pressure_differences = np.linspace(0.1, 500, 100)  # Pa
gas_densities = np.linspace(0.001, 2, 100)  # kg/m^3
time_steps = np.linspace(0.1, 100, 100)  # s
reference_leakages = np.array(
    [
        gas_leakage_micro_holes(dc, mha, pd, gd, ts)
        for dc, mha, pd, gd, ts in zip(
            discharge_coefficients,
            micro_hole_areas,
            pressure_differences,
            gas_densities,
            time_steps,
        )
    ]
)
save_npz(
    "tests/data",
    "gas_leakage_micro_holes.npz",
    discharge_coefficients=discharge_coefficients,
    micro_hole_areas=micro_hole_areas,
    pressure_differences=pressure_differences,
    gas_densities=gas_densities,
    time_steps=time_steps,
    reference_leakages=reference_leakages,
)

# Generate and save gas leakage through permeability data
permeability_coefficients = np.linspace(1e-12, 1e-5, 100)  # units
pressure_differences = np.linspace(0, 500, 100)  # Pa
time_steps = np.linspace(0.01, 10, 100)  # s
balloon_skin_areas = np.linspace(1, 1000, 100)  # m^2
reference_leakages = np.array(
    [
        gas_leakage_permeability(pc, pd, ts, bsa)
        for pc, pd, ts, bsa in zip(
            permeability_coefficients,
            pressure_differences,
            time_steps,
            balloon_skin_areas,
        )
    ]
)
save_npz(
    "tests/data",
    "gas_leakage_permeability.npz",
    permeability_coefficients=permeability_coefficients,
    pressure_differences=pressure_differences,
    time_steps=time_steps,
    balloon_skin_areas=balloon_skin_areas,
    reference_leakages=reference_leakages,
)
