import pytest
from functions import OptimizeBalloon, BalloonSystemType
import numpy as np

# Define the array of test cases for OptimizeBalloon
test_cases_optimize_balloon = [
    {
        "params": {
            "payload_mass_kg": 100,
            "maximal_altitude_m": 40000,
            "minimal_altitude_m": 30000,
            "balloon_1_film_density_kg_m2": 0.01,
            "balloon_1_tape_density_kg_m": 0.0007,
            "balloon_1_no_tapes": 20,
            "balloon_2_film_density_kg_m2": 0.01,
            "balloon_2_tape_density_kg_m": 0.0007,
            "balloon_2_no_tapes": 200,
            "p_max_diff": 200,
            "p_diff_normal": 5,
            "maximal_upwind_speed": 10,
            "maximal_temp_difference": 45,
            "float_offset": 500,
            "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(55),
            "SP_incremental_parameter": 1,
            "ZP_incremental_parameter": 5,
            "SP_theta_end_error_tolerance": 1e-5,
            "ZP_theta_end_error_tolerance": 0.001,
            "SP_max_iterations_balloon_shape": 100000,
            "ZP_max_iterations_balloon_shape": 100000,
            "SP_max_iterations_optimization": 100,
            "ZP_max_iterations_optimization": 100,
        },
        # Expected values
        # General
        "balloon_ratio": 0.849144080804889,  # instance.balloon_ratio
        "float_altitude": 37552.19869954808,  # instance.float_altitude
        "force": 21.22489848686857,  # instance.force
        "minimal_altitude": 33838.508192299196,
        "energy_day": 262.5773950180138,
        "energy_night": 11.965087014349427,
        # Balloon 1
        "b_1_expected_radius_max": 23.737976096306994,
        "b_1_expected_height": 41.24704079765522,
        "b_1_expected_volume": 44471.00824079631,
        "b_1_gas_weight": 23.665731107599033,
        # Balloon 2
        "b_2_expected_radius_max": 18.163122602496856,
        "b_2_expected_height": 21.618516804292327,
        "b_2_expected_volume": 16404.926670250403,
        "b_2_gas_weight": 8.887347103778154,
    },
    {
        "params": {
            "payload_mass_kg": 100,
            "maximal_altitude_m": 35000,
            "minimal_altitude_m": 15000,
            "balloon_1_film_density_kg_m2": 0.01,
            "balloon_1_tape_density_kg_m": 0.01,
            "balloon_1_no_tapes": 0,
            "balloon_2_film_density_kg_m2": 0.02,
            "balloon_2_tape_density_kg_m": 0.02,
            "balloon_2_no_tapes": 0,
            "p_max_diff": 200,
            "p_diff_normal": 10,
            "maximal_upwind_speed": 25,
            "maximal_temp_difference": 15,
            "float_offset": 100,
            "init_balloon_ratio_binb": 1,
            "init_balloon_ratio_tandem": 0.5,
            "balloon_system_type": BalloonSystemType.BINB_AIR,
            "SP_initial_theta": np.deg2rad(91),
            "ZP_initial_theta": np.deg2rad(55),
            "SP_incremental_parameter": 0.5,
            "ZP_incremental_parameter": 2,
            "SP_theta_end_error_tolerance": 1e-5,
            "ZP_theta_end_error_tolerance": 1e-5,
            "SP_max_iterations_balloon_shape": 150000,
            "ZP_max_iterations_balloon_shape": 150000,
            "SP_max_iterations_optimization": 100,
            "ZP_max_iterations_optimization": 100,
        },
        # Expected values
        # General
        "balloon_ratio": 1,
        "float_altitude": 35000,
        "force": 158.6248638318734,
        "minimal_altitude": 15000,
        "energy_day": 1091.859548638294,
        "energy_night": 58.944458235270716,
        # Balloon 1
        "b_1_expected_radius_max": 24.398581517075158,
        "b_1_expected_height": 29.1712955640616,
        "b_1_expected_volume": 39842.848622925354,
        "b_1_gas_weight": 5.855291519285194,
        # Balloon 2
        "b_2_expected_radius_max": 24.398581517075158,
        "b_2_expected_height": 29.1712955640616,
        "b_2_expected_volume": 39842.848622925354,
        "b_2_gas_weight": 45.22405352182044,
    },
]


@pytest.mark.parametrize("case", test_cases_optimize_balloon)
def test_optimize_balloon(case):
    params = case["params"]

    optimization_instance = OptimizeBalloon(
        payload_mass_kg=params["payload_mass_kg"],
        maximal_altitude_m=params["maximal_altitude_m"],
        minimal_altitude_m=params["minimal_altitude_m"],
        balloon_1_film_density_kg_m2=params["balloon_1_film_density_kg_m2"],
        balloon_1_tape_density_kg_m=params["balloon_1_tape_density_kg_m"],
        balloon_1_no_tapes=params["balloon_1_no_tapes"],
        balloon_2_film_density_kg_m2=params["balloon_2_film_density_kg_m2"],
        balloon_2_tape_density_kg_m=params["balloon_2_tape_density_kg_m"],
        balloon_2_no_tapes=params["balloon_2_no_tapes"],
        p_max_diff=params["p_max_diff"],
        p_diff_normal=params["p_diff_normal"],
        maximal_upwind_speed=params["maximal_upwind_speed"],
        maximal_temp_difference=params["maximal_temp_difference"],
        float_offset=params["float_offset"],
        init_balloon_ratio_binb=params["init_balloon_ratio_binb"],
        init_balloon_ratio_tandem=params["init_balloon_ratio_tandem"],
        balloon_system_type=params["balloon_system_type"],
        SP_initial_theta=params["SP_initial_theta"],
        ZP_initial_theta=params["ZP_initial_theta"],
        SP_incremental_parameter=params["SP_incremental_parameter"],
        ZP_incremental_parameter=params["ZP_incremental_parameter"],
        SP_theta_end_error_tolerance=params["SP_theta_end_error_tolerance"],
        ZP_theta_end_error_tolerance=params["ZP_theta_end_error_tolerance"],
        SP_max_iterations_balloon_shape=params["SP_max_iterations_balloon_shape"],
        ZP_max_iterations_balloon_shape=params["ZP_max_iterations_balloon_shape"],
        SP_max_iterations_optimization=params["SP_max_iterations_optimization"],
        ZP_max_iterations_optimization=params["ZP_max_iterations_optimization"],
    )

    optimization_instance.initial_balloons()
    optimization_instance.optimization()
    min_altitude = optimization_instance.find_minimal_altitude()
    night_energy, day_energy = optimization_instance.energy_consumption()
    tolerance_length = 0.01  # Tolerance of 0.01 meters
    tolerance_volume = 0.5  # Tolerance of 0.5 cubic meters
    tolerance_weight = 0.1  # Tolerance of 0.1 N
    tolerance_general = 0.01  # General tolerance

    # General assertions
    assert optimization_instance.balloon_ratio == pytest.approx(
        case["balloon_ratio"], abs=tolerance_general
    ), f"Failed balloon ratio for params: {params}"
    assert optimization_instance.float_altitude == pytest.approx(
        case["float_altitude"], abs=tolerance_length
    ), f"Failed float altitude for params: {params}"
    assert optimization_instance.force == pytest.approx(
        case["force"], abs=tolerance_weight
    ), f"Failed force for params: {params}"
    assert min_altitude == pytest.approx(
        case["minimal_altitude"], abs=tolerance_length
    ), f"Failed minimal altitude for params: {params}"
    assert day_energy == pytest.approx(
        case["energy_day"], abs=tolerance_general
    ), f"Failed minimal altitude for params: {params}"
    assert night_energy == pytest.approx(
        case["energy_night"], abs=tolerance_general
    ), f"Failed minimal altitude for params: {params}"

    # Balloon 1 assertions
    assert max(optimization_instance.balloon_1.radius) == pytest.approx(
        case["b_1_expected_radius_max"], abs=tolerance_length
    ), f"Failed Balloon 1 radius for params: {params}"
    assert optimization_instance.balloon_1.height[-1] == pytest.approx(
        case["b_1_expected_height"], abs=tolerance_length
    ), f"Failed Balloon 1 height for params: {params}"
    assert optimization_instance.balloon_1.volume == pytest.approx(
        case["b_1_expected_volume"], abs=tolerance_volume
    ), f"Failed Balloon 1 volume for params: {params}"
    assert optimization_instance.balloon_1.gas_weight == pytest.approx(
        case["b_1_gas_weight"], abs=tolerance_weight
    ), f"Failed Balloon 1 gas weight for params: {params}"

    # Balloon 2 assertions
    assert max(optimization_instance.balloon_2.radius) == pytest.approx(
        case["b_2_expected_radius_max"], abs=tolerance_length
    ), f"Failed Balloon 2 radius for params: {params}"
    assert optimization_instance.balloon_2.height[-1] == pytest.approx(
        case["b_2_expected_height"], abs=tolerance_length
    ), f"Failed Balloon 2 height for params: {params}"
    assert optimization_instance.balloon_2.volume == pytest.approx(
        case["b_2_expected_volume"], abs=tolerance_volume
    ), f"Failed Balloon 2 volume for params: {params}"
    assert optimization_instance.balloon_2.gas_weight == pytest.approx(
        case["b_2_gas_weight"], abs=tolerance_weight
    ), f"Failed Balloon 2 gas weight for params: {params}"
