import pytest
from functions import Optimize, BalloonSystemType, TrackingType, PayloadOperation
from datetime import datetime
import numpy as np


def test_optimize():
    optimization = Optimize(
        # Balloon
        payload_mass_kg=250,
        maximal_altitude_m=35000,
        minimal_altitude_m=10000,
        balloon_1_film_density_kg_m2=0.01,
        balloon_1_tape_density_kg_m=0.0009,
        balloon_1_no_tapes=100,
        balloon_2_film_density_kg_m2=0.02,
        balloon_2_tape_density_kg_m=0.007,
        balloon_2_no_tapes=200,
        p_max_diff=200,
        p_diff_normal=5,
        maximal_upwind_speed=10,
        maximal_temp_difference=45,
        float_offset=1500,
        init_balloon_ratio_binb=1,
        init_balloon_ratio_tandem=0.5,
        balloon_system_type=BalloonSystemType.TANDEM_HE,
        SP_initial_theta=np.deg2rad(91),
        ZP_initial_theta=np.deg2rad(45),
        SP_incremental_parameter=1,
        ZP_incremental_parameter=5,
        SP_theta_end_error_tolerance=1e-5,
        ZP_theta_end_error_tolerance=1e-3,
        SP_max_iterations_balloon_shape=100000,
        ZP_max_iterations_balloon_shape=100000,
        SP_max_iterations_optimization=100,
        ZP_max_iterations_optimization=100,
        # Solar panel
        solar_panel_surface_density=2,
        solar_panel_efficiency=0.25,
        solar_panel_safety_coefficient=1,
        solar_panel_tracking_type=TrackingType.ONE_AXIS,
        solar_panel_power_fraction=0.05,
        # Battery
        battery_power_density=3000,
        battery_specific_energy=350,
        battery_max_discharge=0.25,
        battery_safety_coefficient=1,
        # Others
        latitude=20,
        longitude=20,
        flight_date=datetime(2024, 5, 12),
        flight_duration=25,
        desired_daily_altitude_change=15000,
        night_altitude_change=0.5,
        maximal_descend_rate=0.1,
        auxiliaries_weight=0.04,
        payload_power_requirement=10,
        payload_operation=PayloadOperation.DAY,
        power_heater=0.5,
        power_auxiliaries=2,
        # Stress
        balloon_1_max_stress_film=200,
        balloon_1_max_stress_tape=15000,
        balloon_2_max_stress_film=200,
        balloon_2_max_stress_tape=15000,
    )

    optimization.initialization()
    optimization.optimization()

    tolerance = 0.01  # General tolerance for floating-point comparisons

    assert optimization.battery.weight == pytest.approx(
        0.28329138250252256, abs=tolerance
    ), "Failed battery weight"
    assert optimization.battery.capacity == pytest.approx(
        99.15198387588289, abs=tolerance
    ), "Failed battery capacity"
    assert optimization.solar_panel.weight == pytest.approx(
        1.421020549482251, abs=tolerance
    ), "Failed solar panel weight"
    assert optimization.solar_panel.surface == pytest.approx(
        0.7105102747411255, abs=tolerance
    ), "Failed solar panel surface"
    assert optimization.solar_panel.total_energy == pytest.approx(
        1770.3838326166458, abs=tolerance
    ), "Failed solar panel total energy"
    assert optimization.total_accessories_weight == pytest.approx(
        261.77248440926417, abs=tolerance
    ), "Failed total accessories weight"
    assert optimization.balloon.balloon_1.volume == pytest.approx(
        32692.13148825687, abs=tolerance
    ), "Failed balloon 1 volume"
    assert optimization.balloon.balloon_2.volume == pytest.approx(
        42510.998769425976, abs=tolerance
    ), "Failed balloon 2 volume"


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