import pytest
import numpy as np
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 load_npz(filepath):
    with np.load(filepath) as data:
        return {key: data[key] for key in data}


def test_gas_parameters_altitude():
    data = load_npz("tests/data/gas_parameters_altitude.npz")
    altitudes = data["altitudes"]
    reference_pressures = data["pressures"]
    reference_temperatures = data["temperatures"]
    pressures, temperatures = gas_parameters_altitude(altitudes)
    assert np.allclose(pressures, reference_pressures, rtol=1e-5)
    assert np.allclose(temperatures, reference_temperatures, rtol=1e-5)


def test_sphere_drag():
    data = load_npz("tests/data/sphere_drag.npz")
    reynolds_numbers = data["reynolds_numbers"]
    reference_drag_coeffs = data["reference_drag_coeffs"]
    for Re, ref_cd in zip(reynolds_numbers, reference_drag_coeffs):
        assert sphere_drag(Re) == pytest.approx(ref_cd, rel=1e-2)


def test_calculate_air_viscosity():
    data = load_npz("tests/data/air_viscosity.npz")
    temperatures = data["temperatures"]
    reference_viscosities = data["reference_viscosities"]
    for temp, ref_visc in zip(temperatures, reference_viscosities):
        assert calculate_air_viscosity(temp) == pytest.approx(ref_visc, rel=1e-2)


def test_calculate_gas_density():
    data = load_npz("tests/data/gas_density.npz")
    temperatures = data["temperatures"]
    pressures = data["pressures"]
    reference_densities = data["reference_densities"]
    specific_gas_constants = np.array([r_air, r_He])
    for temp, press, spec_gas_constant, ref_density in zip(
        temperatures, pressures, specific_gas_constants, reference_densities
    ):
        assert calculate_gas_density(temp, press, spec_gas_constant) == pytest.approx(
            ref_density, rel=1e-2
        )


def test_calculate_reynolds_number():
    data = load_npz("tests/data/reynolds_number.npz")
    temperatures = data["temperatures"]
    pressures = data["pressures"]
    velocities = data["velocities"]
    lengths = data["lengths"]
    reference_reynolds = data["reference_reynolds"]
    for temp, press, vel, length, ref_re in zip(
        temperatures, pressures, velocities, lengths, reference_reynolds
    ):
        assert calculate_reynolds_number(temp, press, vel, length) == pytest.approx(
            ref_re, rel=1e-2
        )


def test_calculate_specific_buoyancy():
    data = load_npz("tests/data/specific_buoyancy.npz")
    temperatures = data["temperatures"]
    pressures = data["pressures"]
    pressure_differences = data["pressure_differences"]
    reference_buoyancies = data["reference_buoyancies"]
    for temp, press, press_diff, ref_buoyancy in zip(
        temperatures, pressures, pressure_differences, reference_buoyancies
    ):
        assert calculate_specific_buoyancy(temp, press, press_diff) == pytest.approx(
            ref_buoyancy, rel=1e-2
        )


def test_calculate_radius_sphere():
    data = load_npz("tests/data/radius_sphere.npz")
    volumes = data["volumes"]
    reference_radii = data["reference_radii"]
    for vol, ref_radius in zip(volumes, reference_radii):
        assert calculate_radius_sphere(vol) == pytest.approx(ref_radius, rel=1e-5)


def test_calculate_area_circle():
    data = load_npz("tests/data/area_circle.npz")
    radii = data["radii"]
    reference_areas = data["reference_areas"]
    for rad, ref_area in zip(radii, reference_areas):
        assert calculate_area_circle(rad) == pytest.approx(ref_area, rel=1e-5)


def test_microhole_formation():
    data = load_npz("tests/data/microhole_formation.npz")
    formation_rates = data["formation_rate"]
    ages_of_film = data["age_of_film"]
    areas_holes_initial = data["area_holes_initial"]
    reference_hole_areas = data["reference_hole_areas"]
    for form_rate, age, area_hol_init, ref_hole_area in zip(
        formation_rates, ages_of_film, areas_holes_initial, reference_hole_areas
    ):
        assert microhole_formation(form_rate, age, area_hol_init) == pytest.approx(
            ref_hole_area, rel=1e-2
        )


def test_gas_leakage_micro_holes():
    data = load_npz("tests/data/gas_leakage_micro_holes.npz")
    discharge_coefficients = data["discharge_coefficients"]
    micro_hole_areas = data["micro_hole_areas"]
    pressure_differences = data["pressure_differences"]
    gas_densities = data["gas_densities"]
    time_steps = data["time_steps"]
    reference_leakages = data["reference_leakages"]
    for dc, mha, pd, gd, ts, ref_leakage in zip(
        discharge_coefficients,
        micro_hole_areas,
        pressure_differences,
        gas_densities,
        time_steps,
        reference_leakages,
    ):
        assert gas_leakage_micro_holes(dc, mha, pd, gd, ts) == pytest.approx(
            ref_leakage, rel=1e-2
        )


def test_gas_leakage_permeability():
    data = load_npz("tests/data/gas_leakage_permeability.npz")
    permeability_coefficients = data["permeability_coefficients"]
    pressure_differences = data["pressure_differences"]
    time_steps = data["time_steps"]
    balloon_skin_areas = data["balloon_skin_areas"]
    reference_leakages = data["reference_leakages"]
    for pc, pd, ts, bsa, ref_leakage in zip(
        permeability_coefficients,
        pressure_differences,
        time_steps,
        balloon_skin_areas,
        reference_leakages,
    ):
        assert gas_leakage_permeability(pc, pd, ts, bsa) == pytest.approx(
            ref_leakage, rel=1e-2
        )


if __name__ == "__main__":
    pytest.main()
