from enum import Enum
from .balloon import Balloon
from .utils import (
    calculate_reynolds_number,
    calculate_radius_sphere,
    calculate_area_circle,
    sphere_drag,
    calculate_gas_density,
    gas_parameters_altitude,
)
from .constants import r_air, g0, r_He, lambda_air, lambda_He, Cp_Air, Cp_He
from .enums import BalloonSystemType
from scipy.optimize import brentq, newton, minimize, broyden2
from math import pi, ceil
import numpy as np
from logger import logger


class OptimizeBalloon:
    def __init__(
        self,
        payload_mass_kg,
        maximal_altitude_m,
        minimal_altitude_m,
        balloon_1_film_density_kg_m2,
        balloon_1_tape_density_kg_m,
        balloon_1_no_tapes,
        balloon_2_film_density_kg_m2,
        balloon_2_tape_density_kg_m,
        balloon_2_no_tapes,
        p_max_diff,
        p_diff_normal,
        maximal_upwind_speed,
        maximal_temp_difference,
        float_offset,
        init_balloon_ratio_binb,
        init_balloon_ratio_tandem,
        balloon_system_type,
        SP_initial_theta,
        ZP_initial_theta,
        SP_incremental_parameter,
        ZP_incremental_parameter,
        SP_theta_end_error_tolerance,
        ZP_theta_end_error_tolerance,
        SP_max_iterations_balloon_shape,
        ZP_max_iterations_balloon_shape,
        SP_max_iterations_optimization,
        ZP_max_iterations_optimization,
    ):
        # Verify that the balloon_system_type is an instance of the BalloonSystemType enum
        if not isinstance(balloon_system_type, BalloonSystemType):
            raise ValueError(
                "balloon_system_type must be an instance of BalloonSystemType Enum"
            )

        # Initialize instance variables
        self.payload_mass_kg = payload_mass_kg
        self.maximal_altitude_m = maximal_altitude_m
        self.minimal_altitude_m = minimal_altitude_m
        self.balloon_1_film_density_kg_m2 = balloon_1_film_density_kg_m2
        self.balloon_1_tape_density_kg_m = balloon_1_tape_density_kg_m
        self.balloon_1_no_tapes = balloon_1_no_tapes
        self.balloon_2_film_density_kg_m2 = balloon_2_film_density_kg_m2
        self.balloon_2_tape_density_kg_m = balloon_2_tape_density_kg_m
        self.balloon_2_no_tapes = balloon_2_no_tapes
        self.p_max_diff = p_max_diff
        self.p_diff_normal = p_diff_normal
        self.maximal_upwind_speed = maximal_upwind_speed
        self.maximal_temp_difference = maximal_temp_difference
        self.float_offset = float_offset
        self.init_balloon_ratio_binb = init_balloon_ratio_binb
        self.init_balloon_ratio_tandem = init_balloon_ratio_tandem
        self.balloon_system_type = balloon_system_type
        self.SP_initial_theta = SP_initial_theta
        self.ZP_initial_theta = ZP_initial_theta
        self.SP_incremental_parameter = SP_incremental_parameter
        self.ZP_incremental_parameter = ZP_incremental_parameter
        self.SP_theta_end_error_tolerance = SP_theta_end_error_tolerance
        self.ZP_theta_end_error_tolerance = ZP_theta_end_error_tolerance
        self.SP_max_iterations_balloon_shape = SP_max_iterations_balloon_shape
        self.ZP_max_iterations_balloon_shape = ZP_max_iterations_balloon_shape
        self.SP_max_iterations_optimization = SP_max_iterations_optimization
        self.ZP_max_iterations_optimization = ZP_max_iterations_optimization

    def initial_balloons(self):
        match self.balloon_system_type:
            case BalloonSystemType.TANDEM_HE:
                logger.info("Initializing Tandem HE balloon system.")
                self.balloon_ratio = self.init_balloon_ratio_tandem
                self.balloon_1 = Balloon(
                    payload=self.payload_mass_kg * self.balloon_ratio * g0,
                    altitude=self.maximal_altitude_m,
                    film_density=self.balloon_1_film_density_kg_m2 * g0,
                    tape_density=self.balloon_1_tape_density_kg_m * g0,
                    number_of_gores=self.balloon_1_no_tapes,
                    incremental_parameter=self.ZP_incremental_parameter,
                    initial_theta=self.ZP_initial_theta,
                    theta_end_error_tolerance=self.ZP_theta_end_error_tolerance,
                    max_iterations_balloon_shape=self.ZP_max_iterations_balloon_shape,
                    max_iterations_optimization=self.ZP_max_iterations_optimization,
                )
                self.balloon_2 = Balloon(
                    payload=self.payload_mass_kg * (1 - self.balloon_ratio) * g0,
                    altitude=self.maximal_altitude_m,
                    film_density=self.balloon_2_film_density_kg_m2 * g0,
                    tape_density=self.balloon_2_tape_density_kg_m * g0,
                    number_of_gores=self.balloon_2_no_tapes,
                    pressure_difference=self.p_max_diff,
                    pressure_difference_float=self.p_diff_normal,
                    incremental_parameter=self.SP_incremental_parameter,
                    initial_theta=self.SP_initial_theta,
                    theta_end_error_tolerance=self.SP_theta_end_error_tolerance,
                    max_iterations_balloon_shape=self.SP_max_iterations_balloon_shape,
                    max_iterations_optimization=self.SP_max_iterations_optimization,
                )

                self.total_balloon_weight = (
                    self.balloon_1.gores_weight
                    + self.balloon_1.skin_weight
                    + self.balloon_2.gores_weight
                    + self.balloon_2.skin_weight
                ) / g0
                self.fill_gas_weight = (
                    self.balloon_1.gas_weight + self.balloon_2.gas_weight
                )
                self.total_weight = (
                    self.total_balloon_weight
                    + self.fill_gas_weight
                    + self.payload_mass_kg
                )
                self.float_altitude = self.maximal_altitude_m - self.float_offset
            case BalloonSystemType.TANDEM_AIR:
                logger.exception("Tandem AIR balloon system not implemented yet.")
                raise Exception(
                    f"Tandem AIR balloon system not implemented yet. Use either Tandem HE or BINB AIR."
                )

            case BalloonSystemType.BINB_HE:
                logger.exception("BINB HE balloon system not implemented yet.")
                raise Exception(
                    f"BINB HE balloon system not implemented yet. Use either Tandem HE or BINB AIR."
                )

            case BalloonSystemType.BINB_AIR:
                logger.info(
                    "Initializing BINB AIR (Balloon in Balloon - Air Ballast) balloon system."
                )
                # Outer balloon
                self.balloon_ratio = self.init_balloon_ratio_binb
                # TODO: Balloon 1 film density and so on as sum of both balloon 1 and 2
                self.balloon_1 = Balloon(
                    payload=self.payload_mass_kg * self.balloon_ratio * g0,
                    altitude=self.maximal_altitude_m,
                    film_density=self.balloon_1_film_density_kg_m2 * g0,
                    tape_density=self.balloon_1_tape_density_kg_m * g0,
                    number_of_gores=self.balloon_1_no_tapes,
                    pressure_difference=self.p_max_diff,
                    pressure_difference_float=self.p_diff_normal,
                    incremental_parameter=self.SP_incremental_parameter,
                    initial_theta=self.SP_initial_theta,
                    theta_end_error_tolerance=self.SP_theta_end_error_tolerance,
                    max_iterations_balloon_shape=self.SP_max_iterations_balloon_shape,
                    max_iterations_optimization=self.SP_max_iterations_optimization,
                )
                # Inner balloon
                self.balloon_2 = Balloon(
                    payload=self.payload_mass_kg * self.balloon_ratio * g0,
                    altitude=self.maximal_altitude_m,
                    film_density=self.balloon_1_film_density_kg_m2 * g0,
                    tape_density=self.balloon_1_tape_density_kg_m * g0,
                    number_of_gores=self.balloon_1_no_tapes,
                    pressure_difference=self.p_max_diff,
                    pressure_difference_float=self.p_diff_normal,
                    incremental_parameter=self.SP_incremental_parameter,
                    initial_theta=self.SP_initial_theta,
                    theta_end_error_tolerance=self.SP_theta_end_error_tolerance,
                    max_iterations_balloon_shape=self.SP_max_iterations_balloon_shape,
                    max_iterations_optimization=self.SP_max_iterations_optimization,
                )

                self.balloon_2.film_density = self.balloon_2_film_density_kg_m2 * g0
                self.balloon_2.tape_density = self.balloon_2_tape_density_kg_m * g0
                self.balloon_2.number_of_gores = self.balloon_2_no_tapes
                self.balloon_2.pressure_difference = 0

                self.total_balloon_weight = (
                    self.balloon_1.gores_weight
                    + self.balloon_1.skin_weight
                    + self.balloon_2.gores_weight
                    + self.balloon_2.skin_weight
                ) / g0

                self.float_altitude = self.maximal_altitude_m

                # self.fill_gas_weight = (
                #     self.balloon_1.gas_weight + self.balloon_2.gas_weight
                # ) # Not correct
                # self.total_weight = (
                #     self.total_balloon_weight
                #     + self.fill_gas_weight
                #     + self.payload_mass_kg
                # )

                self.balloon_2.volume_rescale(
                    self.balloon_1.volume * self.balloon_ratio
                )
                pressure, temperature = gas_parameters_altitude(self.maximal_altitude_m)
                air_density = calculate_gas_density(temperature, pressure, r_air)

                def objective(volume):

                    outer_volume = volume[0]
                    inner_volume = outer_volume * self.balloon_ratio
                    self.balloon_1.volume_rescale(outer_volume)
                    self.balloon_2.volume_rescale(inner_volume)
                    # print(f"Outer volume: {self.balloon_1.volume}")
                    # print(f"Outer volume should be: {outer_volume}")
                    # print(f"Inner volume: {self.balloon_2.volume}")
                    # print(f"Inner volume should be: {inner_volume}")

                    # Boyle's law - same temperature and amount of gas
                    V_in = (
                        self.balloon_2.volume
                        * pressure
                        / (pressure + self.p_diff_normal)
                    )

                    m_He = V_in * (pressure + self.p_diff_normal) / (temperature * r_He)
                    self.balloon_2.gas_weight = m_He
                    m_air = (
                        (self.balloon_1.volume - V_in)
                        * (pressure + self.p_diff_normal)
                        / (temperature * r_air)
                    )
                    self.balloon_1.gas_weight = m_air
                    self.total_balloon_weight = (
                        self.balloon_1.gores_weight
                        + self.balloon_1.skin_weight
                        + self.balloon_2.gores_weight
                        + self.balloon_2.skin_weight
                    ) / g0
                    # print(f"m_He = {m_He}")
                    # print(f"m_air = {m_air}")
                    # print(volume)
                    balance = (
                        self.balloon_1.volume * air_density
                        - self.payload_mass_kg
                        - self.total_balloon_weight
                        - m_He
                        - m_air
                    )
                    # print(f"Balance: {balance}")
                    # rescale outer balloon to given volume
                    return abs(balance)

                tolerance = self.payload_mass_kg * 1e-7
                bounds = [(0.1 * self.balloon_1.volume, 10 * self.balloon_1.volume)]
                # Perform the optimization, TODO: maybe scant would be better without the abs in objective function
                logger.info(
                    "Optimizing the volume of the outer balloon for initialization."
                )
                minimize(
                    objective,
                    [self.balloon_1.volume],
                    bounds=bounds,
                    method="Nelder-Mead",
                    options={"fatol": tolerance, "xatol": tolerance},
                )

            case _:
                raise ValueError(
                    f"Unknown balloon system type. {self.balloon_system_type}"
                )

    def recalculate(self):
        match self.balloon_system_type:
            case BalloonSystemType.TANDEM_HE:
                logger.info("Recalculating Tandem HE balloon system.")
                # Balloon 1 - ZP balloon
                self.balloon_1.payload = self.payload_mass_kg * self.balloon_ratio * g0

                self.balloon_1.recalculate()
                logger.info("Balloon 1 recalculated.")

                # Balloon 2 - SP balloon
                self.balloon_2.payload = (
                    self.payload_mass_kg * (1 - self.balloon_ratio) * g0
                )

                pressure, temperature = gas_parameters_altitude(self.maximal_altitude_m)
                air_density = calculate_gas_density(temperature, pressure, r_air)

                def objective(balloon_2_volume):
                    volume = balloon_2_volume[0]

                    self.balloon_2.volume_rescale(volume)
                    balloon_weight = (
                        self.balloon_2.gores_weight + self.balloon_2.skin_weight
                    ) / g0
                    balance = (
                        volume * air_density
                        - self.balloon_2.payload / g0
                        - self.balloon_2.gas_weight
                        - balloon_weight
                    )
                    return abs(balance)

                tolerance = self.payload_mass_kg * 1e-7
                bounds = [(0.1 * self.balloon_1.volume, 10 * self.balloon_1.volume)]
                try:
                    minimize(
                        objective,
                        [self.balloon_2.volume],
                        bounds=bounds,
                        method="Nelder-Mead",
                        options={"fatol": tolerance, "xatol": tolerance},
                    )
                except Exception as e:
                    logger.error(f"Error during Balloon2 recalculation: {e}")
                    raise Exception(f"Error during Balloon2 recalculation: {e}")

                # self.balloon_2.recalculate()

                logger.info("Balloon 2 recalculated.")

                self.total_balloon_weight = (
                    self.balloon_1.gores_weight
                    + self.balloon_1.skin_weight
                    + self.balloon_2.gores_weight
                    + self.balloon_2.skin_weight
                ) / g0
                self.fill_gas_weight = (
                    self.balloon_1.gas_weight + self.balloon_2.gas_weight
                )
                self.total_weight = (
                    self.total_balloon_weight
                    + self.fill_gas_weight
                    + self.payload_mass_kg
                )
                self.float_altitude = self.maximal_altitude_m - self.float_offset
            case BalloonSystemType.TANDEM_AIR:
                print("Initializing Tandem AIR balloon system.")
            case BalloonSystemType.BINB_HE:
                print(
                    "Initializing BINB HE (Balloon in Balloon - Compressed Helium) balloon system."
                )
            case BalloonSystemType.BINB_AIR:
                logger.info(
                    "Recalculating BINB AIR (Balloon in Balloon - Air Ballast) balloon system."
                )

                self.balloon_1.pressure_difference = self.p_max_diff
                self.balloon_1.pressure_difference_float = self.p_diff_normal

                self.total_balloon_weight = (
                    self.balloon_1.gores_weight
                    + self.balloon_1.skin_weight
                    + self.balloon_2.gores_weight
                    + self.balloon_2.skin_weight
                ) / g0

                self.float_altitude = self.maximal_altitude_m

                self.balloon_2.volume_rescale(
                    self.balloon_1.volume * self.balloon_ratio
                )
                pressure, temperature = gas_parameters_altitude(self.maximal_altitude_m)
                air_density = calculate_gas_density(temperature, pressure, r_air)

                def objective(volume):

                    outer_volume = volume[0]
                    inner_volume = outer_volume * self.balloon_ratio
                    self.balloon_1.volume_rescale(outer_volume)
                    self.balloon_2.volume_rescale(inner_volume)
                    # print(f"Outer volume: {self.balloon_1.volume}")
                    # print(f"Outer volume should be: {outer_volume}")
                    # print(f"Inner volume: {self.balloon_2.volume}")
                    # print(f"Inner volume should be: {inner_volume}")

                    # Boyle's law - same temperature and amount of gas
                    V_in = (
                        self.balloon_2.volume
                        * pressure
                        / (pressure + self.p_diff_normal)
                    )

                    m_He = V_in * (pressure + self.p_diff_normal) / (temperature * r_He)
                    self.balloon_2.gas_weight = m_He
                    m_air = (
                        (self.balloon_1.volume - V_in)
                        * (pressure + self.p_diff_normal)
                        / (temperature * r_air)
                    )
                    self.balloon_1.gas_weight = m_air
                    self.total_balloon_weight = (
                        self.balloon_1.gores_weight
                        + self.balloon_1.skin_weight
                        + self.balloon_2.gores_weight
                        + self.balloon_2.skin_weight
                    ) / g0

                    balance = (
                        self.balloon_1.volume * air_density
                        - self.payload_mass_kg
                        - self.total_balloon_weight
                        - m_He
                        - m_air
                    )
                    logger.info(f"Balance during recalculation optimization: {balance}")

                    # rescale outer balloon to given volume
                    return abs(balance)

                tolerance = self.payload_mass_kg * 1e-7
                bounds = [(0.1 * self.balloon_1.volume, 10 * self.balloon_1.volume)]
                # Perform the optimization, TODO: maybe scant would be better without the abs in objective function
                minimize(
                    objective,
                    [self.balloon_1.volume],
                    bounds=bounds,
                    method="Nelder-Mead",
                    options={"fatol": tolerance, "xatol": tolerance},
                )
                logger.info("Recalculated BINB AIR balloon system.")
            case _:
                raise ValueError(
                    f"Unknown balloon system type. {self.balloon_system_type}"
                )

    def find_minimal_altitude(self):
        match self.balloon_system_type:
            case BalloonSystemType.TANDEM_HE:
                # TODO: For other cases
                ZP_balloon_weight = (
                    self.balloon_1.gores_weight + self.balloon_1.skin_weight
                ) / g0

                def target_function(altitude):
                    pressure, temperature = gas_parameters_altitude(altitude)
                    rho_air = calculate_gas_density(temperature, pressure, r_air)
                    rho_He = calculate_gas_density(temperature, pressure, r_He)
                    V_ZP = (
                        self.total_weight - self.balloon_2.volume * rho_air
                    ) / rho_air
                    balloon_gas_weight = V_ZP * rho_He
                    balloon_lift = V_ZP * rho_air

                    return balloon_lift - balloon_gas_weight - ZP_balloon_weight

                try:
                    minimal_altitude = newton(
                        target_function,
                        x0=self.minimal_altitude_m,
                        x1=self.maximal_altitude_m,
                        tol=1e-5,
                    )

                    minimal_altitude = float(minimal_altitude)

                    return minimal_altitude
                except RuntimeError as e:
                    logger.error(
                        f"Newton's method for minimal altitude calculation failed: {e}"
                    )
                    raise Exception(
                        f"Newton's method for minimal altitude calculation failed: {e}"
                    )

            case BalloonSystemType.BINB_AIR:
                return self.minimal_altitude_m
            case _:  # Catch-all for debugging
                raise ValueError(
                    f"Unknown balloon system type. {self.balloon_system_type}"
                )

    def recalculate_altitude(self, altitude, temperature_difference=None):
        # print(f"Received temperature_difference: {temperature_difference}")
        # TODO: For other cases
        pressure, temperature = gas_parameters_altitude(altitude)
        # print(f"Pressure: {pressure}, temperature: {temperature}")
        air_density = calculate_gas_density(temperature, pressure, r_air)
        match self.balloon_system_type:
            case BalloonSystemType.TANDEM_HE:
                # Calculation of volume of ZP balloon at float altitude
                V_balloon_1_float = (
                    self.total_weight - self.balloon_2.volume * air_density
                ) / air_density

                if temperature_difference == None:
                    m_gas_balloon_1_float = (
                        V_balloon_1_float * pressure / (temperature * r_He)
                    )

                else:
                    balloon_gas_temperature = temperature + temperature_difference
                    m_gas_balloon_1_float = (
                        V_balloon_1_float * pressure / (balloon_gas_temperature * r_He)
                    )

                m_gas_balloon_2_float = self.fill_gas_weight - m_gas_balloon_1_float
                # print("Recalculate altitude")
                # print(
                #     temperature_difference,
                #     m_gas_balloon_1_float,
                #     m_gas_balloon_2_float,
                #     V_balloon_1_float,
                # )
                return (
                    m_gas_balloon_1_float,
                    m_gas_balloon_2_float,
                    V_balloon_1_float,
                )

            case BalloonSystemType.TANDEM_AIR:
                pass
            case BalloonSystemType.BINB_HE:
                pass
            case BalloonSystemType.BINB_AIR:
                m_gas_balloon_2_float = self.balloon_2.gas_weight
                # 1st calculating needed air weight for balance
                m_gas_balloon_1_float = (
                    self.balloon_1.volume * air_density
                    - self.total_balloon_weight
                    - self.payload_mass_kg
                    - m_gas_balloon_2_float
                )

                # Without temp diff
                if temperature_difference == None:
                    # print("Without temp diff float altitude")
                    # 2nd calculating the total_pressure
                    pressure_total = (
                        (r_air * m_gas_balloon_1_float + r_He * m_gas_balloon_2_float)
                        * temperature
                        / self.balloon_1.volume
                    )
                    # overpressure = pressure_total - pressure
                    # 3rd calculating the internal balloon volume
                    V_balloon_2_float = (
                        m_gas_balloon_2_float * r_He * temperature / pressure_total
                    )

                    # print(
                    #     f"m_gas1: {m_gas_balloon_1_float}, m_gas2: {m_gas_balloon_2_float}, overpressure: {overpressure}, V_balloon1: {self.balloon_1.volume}, V_balloon2: {V_balloon_2_float}"
                    # )
                # With temp diff
                else:
                    # print("With temp diff float altitude")
                    balloon_gas_temperature = temperature + temperature_difference
                    pressure_total = (
                        (r_air * m_gas_balloon_1_float + r_He * m_gas_balloon_2_float)
                        * balloon_gas_temperature
                        / self.balloon_1.volume
                    )
                    # overpressure = pressure_total - pressure
                    # 3rd calculating the internal balloon volume
                    V_balloon_2_float = (
                        m_gas_balloon_2_float
                        * r_He
                        * balloon_gas_temperature
                        / pressure_total
                    )

                    # print(
                    #     f"m_gas1: {m_gas_balloon_1_float}, m_gas2: {m_gas_balloon_2_float}, overpressure: {overpressure}, V_balloon1: {self.balloon_1.volume}, V_balloon2: {V_balloon_2_float}"
                    # )
                return m_gas_balloon_1_float, m_gas_balloon_2_float, V_balloon_2_float
            case _:
                print("Unknown balloon system type.")

    def recalculate_all_altitudes(self):
        # TODO: For other cases
        # Round min and max altitude
        min_altitude = round(self.real_min_altitude)
        max_altitude = round(self.maximal_altitude_m)
        # Create altitude array
        altitude_array = np.arange(max_altitude, min_altitude - 1, -1)

        # Calculate the number of altitudes between min and max (step = 1)
        num_altitudes = max_altitude - min_altitude + 1

        # Create pressure, temperature and air density arrays
        pressure_array, temperature_array = gas_parameters_altitude(altitude_array)
        air_density_array = calculate_gas_density(
            temperature_array, pressure_array, r_air
        )

        # Initialize arrays to store results for all altitudes
        m_gas_balloon_1_array = np.zeros(num_altitudes)
        m_gas_balloon_2_array = np.zeros(num_altitudes)
        V_balloon_1_array = np.zeros(num_altitudes)
        V_balloon_2_array = np.zeros(num_altitudes)

        match self.balloon_system_type:
            case BalloonSystemType.TANDEM_HE:
                V_balloon_2_array = np.full_like(altitude_array, self.balloon_2.volume)
                # Calculation of volume of ZP balloon at float altitude
                V_balloon_1_array = (
                    self.total_weight - self.balloon_2.volume * air_density_array
                ) / air_density_array

                m_gas_balloon_1_array = (
                    V_balloon_1_array * pressure_array / (temperature_array * r_He)
                )
                m_gas_balloon_2_array = self.fill_gas_weight - m_gas_balloon_1_array
                return (
                    altitude_array,
                    m_gas_balloon_1_array,
                    m_gas_balloon_2_array,
                    V_balloon_1_array,
                    V_balloon_2_array,
                )

            case BalloonSystemType.TANDEM_AIR:
                pass
            case BalloonSystemType.BINB_HE:
                pass
            case BalloonSystemType.BINB_AIR:
                V_balloon_1_array = np.full_like(altitude_array, self.balloon_1.volume)
                m_gas_balloon_2_array = np.full_like(
                    altitude_array, self.balloon_2.gas_weight
                )
                # 1st calculating needed air weight for balance
                m_gas_balloon_1_array = (
                    self.balloon_1.volume * air_density_array
                    - self.total_balloon_weight
                    - self.payload_mass_kg
                    - self.balloon_2.gas_weight
                )
                gas_mass_total_array = self.balloon_2.gas_weight + m_gas_balloon_1_array
                # Density of gas inside the balloon
                rho_array = gas_mass_total_array / self.balloon_1.volume
                # Mean specific gas constant for the mixture
                r_mean_array = (
                    r_air * m_gas_balloon_1_array + r_He * m_gas_balloon_2_array
                ) / gas_mass_total_array
                pressure_total_array = rho_array * r_mean_array * temperature_array
                # overpressure_array = pressure_total_array - pressure_array
                # 3rd calculating the internal balloon volume
                V_balloon_2_array = (
                    m_gas_balloon_2_array
                    * r_He
                    * temperature_array
                    / pressure_total_array
                )
                return (
                    altitude_array,
                    m_gas_balloon_1_array,
                    m_gas_balloon_2_array,
                    V_balloon_1_array,
                    V_balloon_2_array,
                )
            case _:
                print("Unknown balloon system type.")

    def find_maximal_stabilizing_force_altitude(self):
        # Maximal stabilizing force with maximal overpressure and the corresponding altitude
        # TODO: For other cases

        bounds = [(0.1, 60000.0)]
        tolerance = 1e-5
        # Setting float altitude gas weights
        # Dont want to have temp difference here, if yes it would calculate balloon with already good balance of weights of gasses, this way it calculates it as sudden change of temperature
        m_gas_balloon_1_float, m_gas_balloon_2_float, _ = self.recalculate_altitude(
            altitude=self.float_altitude,
            # temperature_difference=self.maximal_temp_difference,
        )

        # Preparing the objective function that calculates the absolute value of difference between overpressure and maximal overpressure
        match self.balloon_system_type:
            case BalloonSystemType.TANDEM_HE:

                def objective(altitude):
                    pressure, temperature = gas_parameters_altitude(altitude)
                    pressure_total = (
                        m_gas_balloon_2_float
                        * r_He
                        * (temperature + self.maximal_temp_difference)
                        / self.balloon_2.volume
                    )
                    overpressure = pressure_total - pressure
                    return abs(self.p_max_diff - overpressure)

            case BalloonSystemType.BINB_AIR:

                def objective(altitude):
                    pressure, temperature = gas_parameters_altitude(altitude)
                    pressure_total = (
                        (temperature + self.maximal_temp_difference)
                        * (m_gas_balloon_2_float * r_He + m_gas_balloon_1_float * r_air)
                        / self.balloon_1.volume
                    )
                    overpressure = pressure_total - pressure
                    # print(
                    #     f"Altitude: {altitude}, Pressure Toal: {pressure_total}, Overpressure: {overpressure}"
                    # )
                    return abs(self.p_max_diff - overpressure)

        # Optimization
        altitude_max_overpressure = minimize(
            objective,
            [self.float_altitude],
            bounds=bounds,
            method="Nelder-Mead",
            options={"fatol": tolerance, "xatol": tolerance},
        )

        # Max altitude and altitude parameters
        altitude_max_overpressure = altitude_max_overpressure.x[0]
        pressure, temperature = gas_parameters_altitude(altitude_max_overpressure)
        air_density = calculate_gas_density(temperature, pressure, r_air)

        # Calculating force based on max altitude for each balloon type
        match self.balloon_system_type:
            case BalloonSystemType.TANDEM_HE:
                force = (
                    self.total_balloon_weight
                    + self.fill_gas_weight
                    + self.payload_mass_kg
                    - self.balloon_2.volume * air_density
                    - m_gas_balloon_1_float
                    * r_He
                    * (temperature + self.maximal_temp_difference)
                    * air_density
                    / pressure
                ) * g0
                return force, altitude_max_overpressure
            case BalloonSystemType.BINB_AIR:
                force = (
                    self.total_balloon_weight
                    + self.payload_mass_kg
                    + m_gas_balloon_1_float
                    + m_gas_balloon_2_float
                    - self.balloon_1.volume * air_density
                ) * g0
                return force, altitude_max_overpressure

    def calculate_force_overpressure_altitude(self, altitude):
        # Function to calculate force at maximal altitude together with overpressure
        # TODO: For other cases
        pressure, temperature = gas_parameters_altitude(altitude)
        air_density = calculate_gas_density(temperature, pressure, r_air)
        m_gas_balloon_1_float, m_gas_balloon_2_float, _ = self.recalculate_altitude(
            altitude=self.float_altitude,
        )
        # Calculating for balloon type
        match self.balloon_system_type:
            case BalloonSystemType.TANDEM_HE:
                force = (
                    self.total_balloon_weight
                    + self.fill_gas_weight
                    + self.payload_mass_kg
                    - self.balloon_2.volume * air_density
                    - m_gas_balloon_1_float
                    * r_He
                    * (temperature + self.maximal_temp_difference)
                    * air_density
                    / pressure
                ) * g0

                pressure_total = (
                    m_gas_balloon_2_float
                    * r_He
                    * (temperature + self.maximal_temp_difference)
                    / self.balloon_2.volume
                )
                overpressure = pressure_total - pressure
                return force, overpressure
            case BalloonSystemType.BINB_AIR:
                force = (
                    self.total_balloon_weight
                    + self.payload_mass_kg
                    + m_gas_balloon_1_float
                    + m_gas_balloon_2_float
                    - self.balloon_1.volume * air_density
                ) * g0
                # Correct
                pressure_total = (
                    (temperature + self.maximal_temp_difference)
                    * (m_gas_balloon_2_float * r_He + m_gas_balloon_1_float * r_air)
                    / self.balloon_1.volume
                )
                overpressure = pressure_total - pressure
                return force, overpressure

    def temperature_change_force(self):
        # Calculating force from temperature difference between balloon and environment
        # TODO: For other cases

        # Pressure, temperature and air density at float altitude
        pressure, temperature = gas_parameters_altitude(self.float_altitude)
        air_density = calculate_gas_density(temperature, pressure, r_air)
        m_gas_balloon_1_float, m_gas_balloon_2_float, _ = self.recalculate_altitude(
            altitude=self.float_altitude,
            # temperature_difference=self.maximal_temp_difference,
        )
        match self.balloon_system_type:
            case BalloonSystemType.TANDEM_HE:
                # Increase of force, expension of ZP balloon, but also overpressure due to SP balloon
                force = (
                    self.total_balloon_weight
                    + self.fill_gas_weight
                    + self.payload_mass_kg
                    - self.balloon_2.volume * air_density
                    - m_gas_balloon_1_float
                    * r_He
                    * (temperature + self.maximal_temp_difference)
                    * air_density
                    / pressure
                ) * g0
                return force
            case BalloonSystemType.BINB_AIR:
                # No increase in force, but increase in pressure due to SP balloon
                return 0

    def upwind_force(self):

        # TODO: For other cases
        # TODO: What if float altitude below minimal calculated altitude
        # Float altitude
        # Pressure, temperature at float altitude
        pressure, temperature = gas_parameters_altitude(self.float_altitude)
        air_density = calculate_gas_density(temperature, pressure, r_air)
        match self.balloon_system_type:
            case BalloonSystemType.TANDEM_HE:
                _, _, V_balloon_1_float = self.recalculate_altitude(self.float_altitude)
                # Radius of sphere with same volume as balloon
                r_1 = calculate_radius_sphere(V_balloon_1_float)
                r_2 = calculate_radius_sphere(self.balloon_2.volume)
                # Reynolds numbers
                Re_1 = calculate_reynolds_number(
                    temperature, pressure, self.maximal_upwind_speed, r_1 * 2
                )

                Re_2 = calculate_reynolds_number(
                    temperature, pressure, self.maximal_upwind_speed, r_2 * 2
                )
                # Drag coefficient
                Cd_1 = sphere_drag(Re_1)
                Cd_2 = sphere_drag(Re_2)
                # Drag surface area
                A_1 = calculate_area_circle(r_1)
                A_2 = calculate_area_circle(r_2)

                force = (
                    air_density
                    * -self.maximal_upwind_speed**2
                    * (Cd_1 * A_1 + 0.5 * Cd_2 * A_2)
                    / 2
                )
                return force
            case BalloonSystemType.BINB_AIR:
                # Radius of sphere with same volume as balloon
                r = calculate_radius_sphere(self.balloon_1.volume)
                # Reynolds number
                Re = calculate_reynolds_number(
                    temperature, pressure, self.maximal_upwind_speed, r * 2
                )
                # Drag coefficient
                Cd = sphere_drag(Re)
                # Drag surface area
                A = calculate_area_circle(r)

                force = air_density * -self.maximal_upwind_speed**2 * (Cd * A) / 2
                return force

    def check_altitude(self):
        """
        Function to check that the minimal possible altitude as below float altitude
        """
        min_altitude = self.find_minimal_altitude()
        if min_altitude < self.float_altitude:
            return True
        else:
            return False

    def force_balance(self):
        # TODO: Add condition with minimal altitude, if float is below that.
        # TODO: What if max force altitude is lower than max altitude, add condition - not sure what it means, probably working
        logger.info("Calculating force balance.")
        force_maximal, altitude_max_force = (
            self.find_maximal_stabilizing_force_altitude()
        )
        force_altitude, overpressure_altitude = (
            self.calculate_force_overpressure_altitude(self.maximal_altitude_m)
        )
        self.wind_force = self.upwind_force()
        # temperature_force = self.temperature_change_force()
        match self.balloon_system_type:
            case BalloonSystemType.TANDEM_HE:
                # ZP care about maximal altitude -> gas leakage
                if self.maximal_altitude_m <= altitude_max_force:
                    # Force is equal to force at maximal altitude
                    self.force = force_altitude
                    self.max_force_altitude = self.maximal_altitude_m
                else:
                    # Force is equal to force at maximal overpressure
                    self.force = force_maximal
                    self.max_force_altitude = altitude_max_force
            case BalloonSystemType.BINB_AIR:
                # Superpressure doesnt care about maximal altitude -> no gas leakage, but about maximal pressure
                self.force = force_maximal
                self.max_force_altitude = altitude_max_force

        self.force_balance_total = self.force + self.wind_force
        return self.force_balance_total

    def optimization(self):
        match self.balloon_system_type:
            case BalloonSystemType.TANDEM_HE:
                """
                For tandem helium balloons, optimization finds the lowest ratio,
                meaning it can descend to lowest possible altitude,
                while satisfying the uplift drag and temperature difference forces.
                """
                tolerance = 1e-5

                def objective(balloon_ratio):
                    print(f"Balloon ratio: {balloon_ratio}")
                    if not (0 < balloon_ratio):
                        # Return a large penalty value if the ratio is out of bounds
                        return -np.sign(balloon_ratio) * 1e6
                    self.balloon_ratio = balloon_ratio
                    self.recalculate()
                    force_balance = self.force_balance()
                    print(f"Force balance: {force_balance}")
                    maximal_force_overpressure, altitude_max_force = (
                        self.find_maximal_stabilizing_force_altitude()
                    )

                    altitude_difference = altitude_max_force - self.maximal_altitude_m

                    def objective2(float_offset):
                        logger.info(f"    Float offset: {float_offset}")
                        self.float_offset = float_offset
                        self.recalculate()
                        _, altitude_max_force = (
                            self.find_maximal_stabilizing_force_altitude()
                        )
                        altitude_difference = (
                            altitude_max_force - self.maximal_altitude_m
                        )
                        logger.info(f"    Altitude Difference: {altitude_difference}")

                        return altitude_difference

                    if abs(altitude_difference) > 1 and altitude_max_force != 60000:
                        # Initial guesses
                        initial_guess = self.float_offset
                        next_guess = self.float_offset + 1
                        logger.info("Optimizing float offset.")
                        sol = newton(
                            objective2, x0=initial_guess, x1=next_guess, tol=tolerance
                        )
                        logger.info(f"Optimized float offset.{sol}")
                        force_balance = self.force_balance()

                    logger.info(f"Balloon ratio: {self.balloon_ratio}")
                    logger.info(f"Force Balance: {force_balance}")
                    return force_balance

                # Initial guesses
                initial_guess = (
                    self.balloon_ratio
                )  # First guess as current balloon_ratio
                next_guess = (
                    self.balloon_ratio + 0.1
                )  # Second guess as current balloon_ratio + 0.1

                try:
                    # Using the secant method to find the root of the objective function.
                    logger.info("Optimizing tandem helium balloon system.")
                    optimal_ratio = newton(
                        objective, x0=initial_guess, x1=next_guess, tol=tolerance
                    )
                    logger.info("Optimized tandem helium balloon system.")
                    # print(f"Optimized balloon ratio: {optimal_ratio}")
                    return optimal_ratio
                except ValueError as e:
                    raise Exception(
                        f"Failed to find a root of tandem balloon optimization. {e}"
                    )

            case BalloonSystemType.BINB_AIR:
                # For BINB AIR balloons, optimization finds the lowest overpressure, while satisfying the uplift drag and temperature difference forces.
                def objective(overpressure):
                    self.p_max_diff = overpressure
                    self.recalculate()
                    force_balance = self.force_balance()
                    return force_balance

                # Tolerance for optimization
                tolerance = 1e-5

                # Initial guesses
                initial_guess = self.p_max_diff
                next_guess = self.p_max_diff + 1

                try:
                    # Using the secant method to find the root of the objective function.
                    optimal_overpressure = newton(
                        objective, x0=initial_guess, x1=next_guess, tol=tolerance
                    )

                    return optimal_overpressure
                except ValueError as e:
                    raise Exception(f"Failed to find a root of BinB optimization. {e}")

    def energy_consumption(self):
        self.real_min_altitude = round(self.find_minimal_altitude())
        max_altitude = round(self.maximal_altitude_m)

        # When minimal altitude is below required minimal_altitude
        if self.real_min_altitude < self.minimal_altitude_m:
            self.real_min_altitude = self.minimal_altitude_m

        altitude_array = np.arange(max_altitude, self.real_min_altitude - 1, -1)
        pressure_array, temperature_array = gas_parameters_altitude(altitude_array)
        air_density_array = calculate_gas_density(
            temperature_array, pressure_array, r_air
        )
        # TODO: Check how to calculate temperature difference energy requirements during day
        day_temp_difference = 0.66 * self.maximal_temp_difference
        day_balloon_gas_temperature_array = temperature_array + day_temp_difference
        altitude_difference = max_altitude - self.real_min_altitude

        match self.balloon_system_type:
            case BalloonSystemType.TANDEM_HE:
                # TODO: Calc for this category
                # Limit volume of ZP balloon, to lift its weight
                volume_balloon1_array = (
                    self.total_weight - self.balloon_2.volume * air_density_array
                ) / air_density_array

                ## Night
                # Calculation of lift gas fill in ZP balloon
                night_m_ZP_He_array = (
                    volume_balloon1_array * pressure_array / (temperature_array * r_He)
                )
                night_m_SP_He_array = self.fill_gas_weight - night_m_ZP_He_array

                # Calculation of overpressure in SP balloon
                night_overpressure_array = (
                    night_m_SP_He_array * r_He * (temperature_array)
                ) / self.balloon_2.volume - pressure_array

                # Calculate p_out
                night_total_pressure_array = pressure_array + night_overpressure_array
                # Calculate delta_m as the difference between consecutive masses
                night_delta_m = np.diff(night_m_SP_He_array)
                night_delta_m = np.insert(
                    night_delta_m, 0, 0
                )  # Insert a zero at the beginning

                # Compute the term for the energy calculation
                night_term = (
                    Cp_He
                    * temperature_array
                    * (
                        (night_total_pressure_array / pressure_array)
                        ** ((lambda_He - 1) / lambda_He)
                        - 1
                    )
                )
                # Compute the energy changes for each altitude step
                night_delta_E_pump = night_term * night_delta_m
                # Calculate cumulative energy change
                self.night_cumulative_delta_E_pump = np.cumsum(night_delta_E_pump)

                ## Day
                # Calculation of lift gas fill in ZP balloon
                day_m_ZP_He_array = (
                    volume_balloon1_array
                    * pressure_array
                    / (day_balloon_gas_temperature_array * r_He)
                )
                day_m_SP_He_array = self.fill_gas_weight - day_m_ZP_He_array

                # Calculation of overpressure in SP balloon
                day_overpressure_array = (
                    day_m_SP_He_array * r_He * day_balloon_gas_temperature_array
                ) / self.balloon_2.volume - pressure_array

                # Calculate p_out
                day_total_pressure_array = pressure_array + day_overpressure_array
                # Calculate delta_m as the difference between consecutive masses
                day_delta_m = np.diff(day_m_SP_He_array)
                day_delta_m = np.insert(
                    day_delta_m, 0, 0
                )  # Insert a zero at the beginning

                # Compute the term for the energy calculation
                day_term = (
                    Cp_He
                    * temperature_array
                    * (
                        (day_total_pressure_array / pressure_array)
                        ** ((lambda_He - 1) / lambda_He)
                        - 1
                    )
                )
                # Compute the energy changes for each altitude step
                day_delta_E_pump = day_term * day_delta_m
                # Calculate cumulative energy change
                self.day_cumulative_delta_E_pump = np.cumsum(day_delta_E_pump)
            case BalloonSystemType.BINB_AIR:
                # TODO: DOnt use the max temp, but maybe like 0.6 of it.
                # Calculation of energy requirements for day and night
                # 1st calculating needed air weight for balance
                m_gas_balloon_2 = self.balloon_2.gas_weight
                m_gas_balloon_1_array = (
                    self.balloon_1.volume * air_density_array
                    - self.total_balloon_weight
                    - self.payload_mass_kg
                    - m_gas_balloon_2
                )
                gas_mass_total_array = m_gas_balloon_2 + m_gas_balloon_1_array

                # Difference of air mass between each altitude step
                delta_m = np.diff(m_gas_balloon_1_array)
                # Insert a zero at the beginning
                delta_m = np.insert(delta_m, 0, 0)

                # Density of gas inside the balloon
                rho_array = gas_mass_total_array / self.balloon_1.volume
                # Mean specific gas constant for the mixture
                r_mean_array = (
                    r_air * m_gas_balloon_1_array + r_He * m_gas_balloon_2
                ) / gas_mass_total_array

                ## Night
                # 2nd calculating the total_pressure from ideal gas law
                night_total_pressure_array = (
                    rho_array * r_mean_array * temperature_array
                )
                night_overpressure_array = night_total_pressure_array - pressure_array

                night_term = (
                    Cp_Air
                    * temperature_array
                    * (
                        (night_total_pressure_array / pressure_array)
                        ** ((lambda_air - 1) / lambda_air)
                        - 1
                    )
                )
                night_delta_E_pump = night_term * delta_m
                self.night_cumulative_delta_E_pump = np.cumsum(night_delta_E_pump)

                ## Day
                # Equation of force equilibrium - overpressure calculation for day high balloon temp
                day_total_pressure_array = (
                    rho_array * r_mean_array * day_balloon_gas_temperature_array
                )
                day_overpressure_array = day_total_pressure_array - pressure_array

                day_term = (
                    Cp_Air
                    * temperature_array
                    * (
                        (day_total_pressure_array / pressure_array)
                        ** ((lambda_air - 1) / lambda_air)
                        - 1
                    )
                )
                day_delta_E_pump = day_term * delta_m
                self.day_cumulative_delta_E_pump = np.cumsum(day_delta_E_pump)

        night_energy_per_m = (
            self.night_cumulative_delta_E_pump[-1] / altitude_difference
        )
        day_energy_per_m = self.day_cumulative_delta_E_pump[-1] / altitude_difference

        return night_energy_per_m, day_energy_per_m
