# TODO: Calculation for caps is starting from wrong initial angle (theta 0), needs more iterations and takes more time. Possible fix would be to take the last theta 0 from previous calculation and check two points close by
# TODO: When recalculating, starting from previous angle, but doing large steps and still requires many iterations to recalculate.

from matplotlib import pyplot as plt
from math import atan, tan, sin, cos, pi, atan2, sqrt, cbrt
import numpy as np
from scipy.optimize import minimize, dual_annealing
import time
from functions import (
    gas_parameters_altitude,
    calculate_gas_density,
    calculate_air_viscosity,
    calculate_specific_buoyancy,
)
from functions import r_He, g0, r_air

# Simplified balloon calculation


class Balloon:

    iteration = 0
    iteration_shape = 0

    # Modifiable constants
    pressure_head = 0  # a

    def __init__(
        self,
        payload,
        altitude,
        film_density,
        tape_density,
        number_of_gores,
        incremental_parameter,
        initial_theta,
        theta_end_error_tolerance,
        pressure_difference=0,
        pressure_difference_float=0,
        max_iterations_optimization=100,
        max_iterations_balloon_shape=10000,
    ):
        """
        Initialize the Balloon object with given parameters.

        Parameters:
        payload (float): The payload of the balloon in Newtons.
        altitude (float): The geopotential altitude in meters.
        film_density (float or list of float): The weight density of the film in N/m^2.
        tape_density (float): The weight density of the tape in N/m.
        number_of_gores (int): The initial number of gores.
        incremental_parameter (float): The length of one section.
        initial_theta (float): The initial estimation of the angle in radians.
        theta_end_error_tolerance (float): The angular precision at the end.
        pressure_difference (float, optional): The maximal pressure difference of the SP balloon in Pascals. Default is 0.
        pressure_difference_float (float, optional): The normal pressure difference of the SP balloon in Pascals. Default is 0.
        max_iterations_optimization (int, optional): The maximum number of iterations for optimization. Default is 100.
        max_iterations_balloon_shape (int, optional): The maximum number of iterations for balloon shape calculation. Default is 10000.
        """

        # Assign values to instance variables
        self.payload = payload  # Payload [N]
        self.altitude = altitude  # Geopotential altitude [m]
        self.film_density = (
            film_density  # Film weight density [N/m^2] if caps -> array of len(caps)+1
        )
        self.tape_density = tape_density  # Tape weight density [N/m]
        self.number_of_gores = number_of_gores  # Initial umber of gores
        self.pressure_difference = (
            pressure_difference  # Maximal pressure difference of SP balloon [Pa]
        )
        self.pressure_difference_float = (
            pressure_difference_float  # Normal pressure difference of SP balloon [Pa]
        )
        # Calculation parameters
        self.incremental_parameter = (
            incremental_parameter  # Incremental parameter - length of one section
        )
        self.initial_theta = initial_theta  # Estimation of initial angle [rad]
        self.theta_end_error_tolerance = theta_end_error_tolerance  # Incremental parameter - angular precision at the end
        self.max_iterations_optimization = max_iterations_optimization
        self.max_iterations_balloon_shape = max_iterations_balloon_shape

        self.total_iter_balloon_calc = 0
        self.total_iter_optimization_calc = 0

        # Scaling factor to adjust for lift / load balance
        self.scale_factor = 1
        # Initialization of parameters for storing optimization values

        self.theta_0 = [0, self.initial_theta]
        self.theta_end = [0]

        self.skin_elements = [0]
        self.theta = []
        self.radius = []
        self.height = []
        self.all_radius = []
        self.all_heights = []
        # Gas Parameters

        self.pressure, self.temperature = gas_parameters_altitude(self.altitude)
        self.gas_density = calculate_gas_density(
            self.temperature, self.pressure + self.pressure_difference_float, r_He
        )
        self.buoyancy = calculate_specific_buoyancy(
            self.temperature, self.pressure, self.pressure_difference_float
        )

        # Calculation
        # Parametrized
        self.__calculation_balloon_parametrized()

    def recalculate(self):
        """
        Recalculate the balloon shape and parameters based on the updated optimal initial angle.
        """
        self.iteration = 0
        self.iteration_shape = 0

        self.initial_theta = self.optimal_theta_first
        self.theta_0 = [0, self.initial_theta]
        self.theta_end = [0]

        self.skin_elements = [0]
        self.all_radius = []
        self.all_heights = []
        # Gas Parameters

        self.pressure, self.temperature = gas_parameters_altitude(self.altitude)
        self.gas_density = calculate_gas_density(
            self.temperature, self.pressure + self.pressure_difference_float, r_He
        )
        self.buoyancy = calculate_specific_buoyancy(
            self.temperature, self.pressure, self.pressure_difference_float
        )

        # Calculation
        self.__calculation_balloon_parametrized()

    # Calculate with parametrization
    def __parametrization(self):
        """
        Calculate the parametrization factor for balloon calculations.
        """
        k = 0.0002
        self.parameter_lambda = np.cbrt(self.payload / self.buoyancy)
        self.calc_incremental_parameter = (
            k * self.incremental_parameter * self.parameter_lambda
        )

    # Main function to calculate balloon shape with or without caps
    def __calculate(self):
        """
        Main function to calculate balloon shape
        """

        # Function to calculate balloon shape
        def calculate_balloon_shape(theta_first):
            """
            Calculate the balloon shape based on the initial angle.

            Parameters:
            theta_first (float): The initial angle in radians.
            """
            # Precalculation of common calculations in loop to save time
            changed = False
            if theta_first > pi / 2 or self.pressure_difference > 5:
                changed = True
                theta_first = theta_first - pi / 2

            tan_theta_first = tan(theta_first)
            pi_2 = 2 * pi

            # Initialize lists to store balloon parameters
            if changed:
                theta = [theta_first + pi / 2]
            else:
                theta = [theta_first]
            r = [0]  # Array of distance to center line
            z = [0]  # Vertical distance from bottom of the balloon
            I_r = [0]
            I_z = [0]
            I_s = [0]
            iter = 0  # Number of balloon skin parts = iterations

            # Main loop for balloon shape calculation
            while iter == 0 or (r[-1] > 0):

                # Calculate the new theta angle based on the balloon parameters
                # Numerator = dr
                numerator = self.payload * tan_theta_first - I_r[iter]
                # Denominator = dz
                denominator = self.payload + I_z[iter] + I_s[iter]

                sin_theta_n1 = numerator / sqrt(numerator**2 + denominator**2)
                cos_theta_n1 = denominator / sqrt(numerator**2 + denominator**2)
                r_n = r[iter]
                z_n = z[iter]
                if changed:
                    theta_calc = pi / 2 + atan2(numerator, denominator)
                    if theta_calc > pi:
                        theta.append(theta_calc - 2 * pi)
                    else:
                        theta.append(theta_calc)
                    r_n1 = r_n + cos_theta_n1 * self.calc_incremental_parameter
                    z_n1 = z_n - sin_theta_n1 * self.calc_incremental_parameter
                else:
                    theta.append(atan2(numerator, denominator))
                    r_n1 = r_n + sin_theta_n1 * self.calc_incremental_parameter
                    z_n1 = z_n + cos_theta_n1 * self.calc_incremental_parameter

                surface = pi_2 * r_n1 * self.calc_incremental_parameter
                buoyancy_adjusted = self.buoyancy * (z_n1 - self.pressure_head)
                r.append(r_n1)
                z.append(z_n1)
                # Inertia
                I_r.append(
                    I_r[iter]
                    + ((buoyancy_adjusted + self.pressure_difference) * surface)
                    * cos_theta_n1
                )
                I_z.append(
                    I_z[iter]
                    + ((buoyancy_adjusted + self.pressure_difference) * surface)
                    * sin_theta_n1
                )
                I_s.append(
                    I_s[iter]
                    + surface * self.film_density
                    + self.tape_density
                    * self.number_of_gores
                    * self.calc_incremental_parameter
                )

                # Increment the iteration index
                iter += 1

                # print("")
                # print(
                #     "theta: {:<5.5f}     sin_theta: {:<10.5f}     cos theta:  {:<10.5f}     r: {:<10.5f}     z: {:<10.5f}     numerator: {:<10.5f}     denominator:  {:<10.5f}".format(
                #         theta[-1] * 180 / pi,
                #         sin_theta_n1,
                #         cos_theta_n1,
                #         r_n1,
                #         z_n1,
                #         numerator,
                #         denominator,
                #     )
                # )

                if iter >= self.max_iterations_balloon_shape:
                    raise ValueError(
                        "Exceeded maximum iterations for balloon shape calculation"
                    )

            # Return balloon parameters based on whether it's an ending calculation or not

            # Store shapes, final angle and number of elements
            self.theta = np.array(theta)
            self.radius = np.array(r)
            self.height = np.array(z)
            self.all_radius.append(self.radius)
            self.all_heights.append(self.height)
            self.theta_end.append(theta[-1])
            self.skin_elements.append(iter)
            self.total_iter_balloon_calc += iter

        # Function to optimize the initial angle for the balloon shape calculation
        def theta_first_optimization():
            """
            Optimize the initial angle for the balloon shape calculation using a PID controller.
            """

            iter = 1  # Number of iteration

            # theta_0 = self.theta_0
            # theta_end = self.theta_end
            theta_0 = [0, self.initial_theta]
            theta_end = [0]
            # print(f"theta_0 init: {theta_0}")
            # print(f"theta_end init: {theta_end}")
            kp = 0.6  # Proportional gain
            ki = 0.00  # Integral gain
            kd = 0.0  # Derivative gain
            integral_correction = 0
            previous_correction = None
            fallback_step_size = 0.01
            iter_to_PID = 10
            while (
                iter == 1
                or abs(theta_end[iter - 1] + pi / 2) > self.theta_end_error_tolerance
            ):

                calculate_balloon_shape(theta_0[iter])
                # self.balloon_plot()
                theta_end.append(self.theta_end[-1])
                delta_theta_0 = theta_0[iter] - theta_0[iter - 1]
                delta_theta_end = theta_end[iter] - theta_end[iter - 1]
                error = theta_end[iter] + pi / 2
                # print("")
                # print(f"Iter: {iter}")
                if delta_theta_end == 0:
                    theta_0.append(theta_0[iter] + np.sign(error) * fallback_step_size)
                else:
                    correction = delta_theta_0 / delta_theta_end * error
                    if iter < iter_to_PID:
                        correction_factor = correction
                    else:
                        integral_correction += correction
                        derivative_correction = (
                            (correction - previous_correction)
                            if previous_correction is not None
                            else 0
                        )
                        correction_factor = (
                            kp * correction
                            + ki * integral_correction
                            + kd * derivative_correction
                        )
                        # print(
                        #     "Proportional correction: {:<5.5f}     Integral correction: {:<5.5f}     Derivative correction:  {:<5.5f}".format(
                        #         correction * kp,
                        #         integral_correction * ki,
                        #         derivative_correction * kd,
                        #     )
                        # )
                    previous_correction = correction_factor
                    theta_0.append(theta_0[iter] - correction_factor)

                # print(
                #     "theta_0: {:<5.5f}     theta_0 previous: {:<10.5f}     theta_0 new:  {:<10.5f}     theta_end: {:<15}     theta_end previous: {:<15}     theta_end error: {:<15}".format(
                #         theta_0[iter] * 180 / pi,
                #         theta_0[iter - 1] * 180 / pi,
                #         theta_0[iter + 1] * 180 / pi,
                #         theta_end[iter] * 180 / pi,
                #         theta_end[iter - 1] * 180 / pi,
                #         (theta_end[iter] + pi / 2) * 180 / pi,
                #     )
                # )

                iter += 1
                if iter >= self.max_iterations_optimization:
                    raise ValueError(
                        "Exceeded maximum iterations for theta first optimization"
                    )
            calculate_balloon_shape(theta_0[-1])
            self.total_iter_optimization_calc = iter
            self.theta_0 = theta_0

        # Main calculation function
        theta_first_optimization()
        # minimize_optimize()
        self.optimal_theta_first = self.theta_0[-1]

    # Function to reshape balloon to adjust for lift/weight offset scale factor
    def __lift_weight_rescale(self):
        """
        Rescale the balloon shape to adjust for lift/weight offset scale factor.

        Returns:
        float: The optimized scale factor.
        """
        # Function to reshape balloon to adjust for lift/weight offset scale factor
        init_scale_factor = abs(self.payload / (self.lift_balance + self.payload))
        # init_scale_factor = sqrt(abs(self.payload / (self.lift_balance + self.payload)))
        # TODO: The bounds should be also adjusted based on the accuracy of initial solution similarly as init_scale_factor
        bounds = [(1 / 10, 100)]
        # print(f"Init scale factor:  {init_scale_factor}")
        # print(f"Bound:              {bounds}")
        original_radius = np.array(self.radius)
        original_height = np.array(self.height)
        original_incremental_parameter = self.calc_incremental_parameter
        air_density = calculate_gas_density(self.temperature, self.pressure, r_air)
        self.gas_density = calculate_gas_density(
            self.temperature, self.pressure + self.pressure_difference_float, r_He
        )

        def objective(scale_factor):
            scale_factor = scale_factor[0]
            radius = original_radius * scale_factor
            height = original_height * scale_factor
            incremental_parameter = original_incremental_parameter * scale_factor

            # Volume approximation through conical frustum
            volume = self.__calculate_volume(radius, height)

            # Calculate length of skin
            length = len(radius) * incremental_parameter

            # Surface approximation through conical frustum

            surface = self.__calculate_surface(radius, incremental_parameter)
            lift = volume * air_density * g0

            skin_weight = surface * self.film_density

            gores_weight = self.tape_density * self.number_of_gores * length
            gas_weight = volume * self.gas_density
            lift_balance = lift - (
                self.payload + skin_weight + gores_weight + gas_weight * g0
            )
            # print("Balloon parameters:")
            # print(f"    Height of balloon:          {height[-1]:.5f} m")
            # print(f"    Max radius of balloon:      {max(radius):.5f} m")
            # print(f"    Volume of balloon:          {volume:.5f} m3")
            # print(f"    Surface of balloon:         {surface:.5f} m2")
            # print(f"    Length of skin of balloon:  {length:.5f} m")
            # print(f"    Lift of the balloon:        {lift:.5f} N")
            # print(f"    Weight of skin:             {skin_weight:.5f} N")
            # print(f"    Weight of gores             {gores_weight:.5f} N")
            # print(f"    Weight of balloon gas:      {gas_weight * g0:.5f} N")
            # print(f"    Balance of lift and load    {(lift_balance):.12f} N")
            # print(f"    Scale Factor:               {scale_factor:.5f}")

            return abs(lift_balance)

        # Tolerance
        tolerance = self.payload * 1e-7
        # # Start timing
        # start_time = time.time()
        # Perform the optimization
        if init_scale_factor**2 < 4:
            result = minimize(
                objective,
                [init_scale_factor],
                bounds=bounds,
                method="Nelder-Mead",
                options={"fatol": tolerance, "xatol": tolerance},
            )
        else:
            # Perform the optimization using dual_annealing
            result = dual_annealing(
                objective,
                bounds=bounds,
                x0=[init_scale_factor],
                maxiter=1000,  # This is just an example, adjust as needed
                seed=123,
            )

        # # End timing
        # end_time = time.time()

        # # Calculate elapsed time
        # elapsed_time = end_time - start_time

        # # Print the time taken
        # print(f"Optimization took {elapsed_time:.2f} seconds.")

        # Check if the optimization was successful and return the result
        if result.success:
            optimized_scale = result.x[0]
            return optimized_scale
        else:
            raise Exception("Optimization failed: " + result.message)

    # Function to reshape balloon to adjust for lift/weight offset scale factor
    def volume_rescale(self, volume):
        """
        Rescale the balloon shape to match the given volume.

        Parameters:
        volume (float): The target volume for the balloon in cubic meters.
        """
        # Function to reshape balloon to adjust for lift/weight offset scale factor
        init_scale_factor = sqrt(volume / (self.volume))

        bounds = [(1 / 25, 25)]
        # print(f"Init scale factor:  {init_scale_factor}")
        # print(f"Bound:              {bounds}")
        original_radius = np.array(self.radius)
        original_height = np.array(self.height)

        def objective(scale_factor):
            scale_factor = scale_factor[0]
            radius = original_radius * scale_factor
            height = original_height * scale_factor

            # Volume approximation through conical frustum
            volume_new = self.__calculate_volume(radius, height)
            volume_difference = volume - volume_new
            # print("Balloon parameters:")
            # print(f"    Height of balloon:          {height[-1]:.5f} m")
            # print(f"    Max radius of balloon:      {max(radius):.5f} m")
            # print(f"    Volume of balloon:          {volume:.5f} m3")
            return abs(volume_difference)

        # Tolerance
        tolerance = self.volume * 1e-7
        # # Start timing
        # start_time = time.time()
        # Perform the optimization
        if init_scale_factor**2 < 4:
            result = minimize(
                objective,
                [init_scale_factor],
                bounds=bounds,
                method="Nelder-Mead",
                options={"fatol": tolerance, "xatol": tolerance},
            )
        else:
            # Perform the optimization using dual_annealing
            result = dual_annealing(
                objective,
                bounds=bounds,
                x0=[init_scale_factor],
                maxiter=1000,  # This is just an example, adjust as needed
                seed=123,
            )

        # # End timing
        # end_time = time.time()

        # # Calculate elapsed time
        # elapsed_time = end_time - start_time
        self.scale_factor = result.x[0]
        # Print the time taken
        # print(f"Optimization took {elapsed_time:.2f} seconds.")
        # print(f"Scale factor:{self.scale_factor}")

        # Recalculating balloons shape
        self.radius = self.radius * self.scale_factor
        self.height = self.height * self.scale_factor
        self.calc_incremental_parameter = (
            self.calc_incremental_parameter * self.scale_factor
        )
        # All radius and height data from iterations
        self.all_radius = [
            radius_array * self.scale_factor for radius_array in self.all_radius
        ]
        self.all_heights = [
            heights_array * self.scale_factor for heights_array in self.all_heights
        ]
        # Recalculating other parameters
        self.__additional_parameters()

    def __calculation_balloon_parametrized(self):
        # Calculation of parametrization factor
        self.__parametrization()
        # Main calculation
        self.__calculate()
        # Calculation of additional parameters
        self.__additional_parameters()
        # Rescale if needed
        if abs(self.lift_balance) > 1e-3:
            # Calculating scale factor
            self.scale_factor = self.__lift_weight_rescale()
            # Recalculating balloons shape
            self.radius = self.radius * self.scale_factor
            self.height = self.height * self.scale_factor
            self.calc_incremental_parameter = (
                self.calc_incremental_parameter * self.scale_factor
            )
            # All radius and height data from iterations
            self.all_radius = [
                radius_array * self.scale_factor for radius_array in self.all_radius
            ]
            self.all_heights = [
                heights_array * self.scale_factor for heights_array in self.all_heights
            ]
            # Recalculating other parameters
            self.__additional_parameters()

    def plot_optimization(self):
        fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(10, 10))

        # Convert lists to numpy arrays for easier manipulation
        theta_first = np.array(self.theta_0) * 180 / np.pi  # Convert radians to degrees
        theta_end_errors = (
            (-np.array(self.theta_end) - np.pi / 2) * 180 / np.pi
        )  # Convert radians to degrees
        skin_elements = np.array(self.skin_elements)

        # Exclude the first element of each array
        theta_first = theta_first[1:]
        theta_end_errors = theta_end_errors[1:]
        skin_elements = skin_elements[1:]

        # Create an x-axis array starting from 1
        iterations = np.arange(1, len(theta_first) + 1)

        # Plot for Theta First
        ax1.set_xlabel("Iteration")
        ax1.set_ylabel("Theta First (degrees)")
        ax1.plot(
            iterations,
            theta_first,
            label="Theta First",
            color="red",
            linestyle="-",
            marker="o",
        )
        ax1.grid(True, which="both", linestyle="--", linewidth=0.5)

        # Plot for Theta End Error
        ax2.set_xlabel("Iteration")
        ax2.set_ylabel("Theta End Error (degrees)")
        ax2.plot(
            iterations,
            theta_end_errors,
            label="Theta End Error",
            color="blue",
            linestyle="--",
            marker="x",
        )
        ax2.grid(True, which="both", linestyle="--", linewidth=0.5)

        # Plot for Skin Elements
        ax3.set_xlabel("Iteration")
        ax3.set_ylabel("Skin Elements")
        ax3.plot(
            iterations,
            skin_elements,
            label="Skin Elements",
            color="green",
            linestyle="-.",
            marker="s",
        )
        ax3.grid(True, which="both", linestyle="--", linewidth=0.5)

        # Enhance layout to accommodate three subplots
        fig.tight_layout()

        plt.show()

    def plot_balloon_shapes(self):
        # Create a figure and axis for the plot
        fig, ax = plt.subplots(figsize=(10, 6))
        colors = plt.cm.viridis(np.linspace(0, 1, len(self.all_radius)))

        for i, (radius, height) in enumerate(zip(self.all_radius, self.all_heights)):
            ax.plot(radius, height, color=colors[i])
            ax.plot(-radius, height, color=colors[i])  # Assuming symmetry

        ax.set_title("Balloon Shape Evolution During Optimization")
        ax.set_xlabel("Radius (m)")
        ax.set_ylabel("Height (m)")
        ax.axis("equal")
        ax.grid(True)

        # Create a ScalarMappable with the same colormap and normalization as the line plot
        sm = plt.cm.ScalarMappable(
            cmap="viridis", norm=plt.Normalize(vmin=0, vmax=len(self.all_radius))
        )
        sm.set_array(
            []
        )  # You can set_array to an empty list as the ScalarMappable is not used for a direct plotting
        cbar = plt.colorbar(sm, ax=ax, orientation="vertical")
        cbar.set_label("Iteration Step")

        plt.show()

    def get_meridional_stress(self, max_stress_film, max_stress_tape):
        # Create docstring
        """
        Calculate the meridional stress of the balloon skin and tapes.

        Parameters
        ----------
        max_stress_film : float
            Maximum stress of the film in N/m.
        max_stress_tape : float
            Maximum force of the tapes in N.

        Returns
        -------
        tuple
            A tuple containing the following values:
            - skin_length_cumulative : numpy.ndarray
                Cumulative length of the skin.
            - sigma_meridional : numpy.ndarray
                Meridional stress of the balloon skin.
            - sigma_film_calc : numpy.ndarray
                Calculated film stress.
            - sigma_tape_calc : numpy.ndarray
                Calculated tape stress.
            - sigma_total : numpy.ndarray
                Total stress of the balloon skin and tapes.
        """

        # Calculate the meridional stress
        # max stress film in N/m - calculated as tensile strenght of film MPa times thickness in mm times 1000 to get N/m
        # max stress tape in N

        # Clip theta to avoid division by zero and other numerical issues
        theta = np.clip(self.theta, -np.pi / 2, np.pi / 2)

        surface_array = 2 * pi * self.radius * self.calc_incremental_parameter
        pressure_force_cumulative = np.cumsum(
            (
                self.buoyancy * (self.height - self.pressure_head)
                + self.pressure_difference
            )
            * surface_array
            * np.sin(theta)
        )
        skin_weight_cumulative = np.cumsum(self.film_density * surface_array)
        skin_length_cumulative = np.cumsum(
            self.calc_incremental_parameter * np.ones_like(self.radius)
        )
        tapes_weight_cumulative = np.cumsum(
            self.tape_density
            * self.number_of_gores
            * self.calc_incremental_parameter
            * np.ones_like(self.radius)
        )
        meridional_force = (
            self.payload
            + pressure_force_cumulative
            + skin_weight_cumulative
            + tapes_weight_cumulative
        )
        sigma_meridional = abs(meridional_force[1:-1]) / (
            2 * np.pi * self.radius[1:-1] * abs(np.cos(theta[1:-1]))
        )

        sigma_film_calc = max_stress_film * np.ones_like(self.radius[1:-1])
        sigma_tape_calc = (
            max_stress_tape * self.number_of_gores / (2 * np.pi * self.radius[1:-1])
        )
        sigma_total = sigma_film_calc + sigma_tape_calc
        return (
            skin_length_cumulative[1:-1],
            sigma_meridional,
            sigma_film_calc,
            sigma_tape_calc,
            sigma_total,
        )

    def data_print(self):
        # Inputs
        print("Inputs:")
        print(f"    Payload: {self.payload:.2f} N")
        print(f"    Altitude: {self.altitude:.2f} m")
        print(f"    Overpressure: {self.pressure_difference_float:.2f} Pa")

        # Gas parameters
        print("Gas parameters:")
        print(f"    Buoyancy: {self.buoyancy:.4f} N/m3")
        print(f"    Pressure: {self.pressure:.2f} Pa")
        print(f"    Temperature: {self.temperature:.2f} K")

        print("Balloon parameters:")
        print(f"    Height of balloon:          {self.max_height :.5f} m")
        print(f"    Max radius of balloon:      {self.max_radius:.5f} m")
        print(f"    Volume of balloon:          {self.volume:.5f} m3")
        print(f"    Surface of balloon:         {self.surface:.5f} m2")
        print(f"    Length of skin of balloon:  {self.length:.5f} m")
        print(f"    Lift of the balloon:        {self.lift:.5f} N")
        print(f"    Weight of skin:             {self.skin_weight:.5f} N")
        print(f"    Weight of gores             {self.gores_weight:.5f} N")
        print(f"    Weight of balloon gas:      {self.gas_weight * g0:.5f} N")
        print(f"    Balance of lift and load    {self.lift_balance:.5f} N")
        print(
            f"    Width to Height Ratio       {2 * max(self.radius) / self.height[-1]:.5f}"
        )
        print("Other parameters:")
        print(
            f"    Number of iterations of angle optimization:   {self.total_iter_optimization_calc}"
        )
        print(
            f"    Number of total shape calculation steps:      {self.total_iter_balloon_calc}"
        )
        print(
            f"    Optimal theta 0:                              {self.optimal_theta_first * 180 / pi:.5f}°"
        )
        print(
            f"    Scale factor:                                 {self.scale_factor:.5f}"
        )
        print(
            f"    Incremental parameter:                        {self.calc_incremental_parameter:.5f}"
        )

    def __calculate_volume(self, radius, height):
        # volume = sum(
        #     [
        #         pi
        #         * (height[i + 1] - height[i])
        #         * (((radius[i] + radius[i + 1]) / 2) ** 2)
        #         for i in range(len(radius) - 1)
        #     ]
        # )
        # volume = sum(
        #     [
        #         1
        #         / 3
        #         * pi
        #         * (height[i + 1] - height[i])
        #         * (radius[i] ** 2 + radius[i] * radius[i + 1] + radius[i + 1] ** 2)
        #         for i in range(len(radius) - 1)
        #     ]
        # )

        differences = height[1:] - height[:-1]
        radii_sqr = radius[:-1] ** 2 + radius[:-1] * radius[1:] + radius[1:] ** 2
        volumes = (1 / 3) * pi * differences * radii_sqr
        volume = np.sum(volumes)

        return volume

    def __calculate_surface(self, radius, calc_incremental_parameter):
        # Surface approximation of balloons skin

        # surface = sum(
        #     pi * (radius[i] + radius[i + 1]) * calc_incremental_parameter
        #     for i in range(len(radius) - 1)
        # )
        surface = np.sum(2 * pi * radius * calc_incremental_parameter)
        return surface

    # Calculating additional parameters
    def __additional_parameters(self):
        # Balloon parameters

        # Volume approximation through conical frustum
        self.volume = self.__calculate_volume(self.radius, self.height)

        # Calculate length of skin
        self.length = len(self.radius) * self.calc_incremental_parameter

        # Surface approximation through conical frustum

        self.surface = self.__calculate_surface(
            self.radius, self.calc_incremental_parameter
        )
        self.lift = (
            self.volume
            * calculate_gas_density(self.temperature, self.pressure, r_air)
            * g0
        )

        self.skin_weight = self.surface * self.film_density

        self.gores_weight = self.tape_density * self.number_of_gores * self.length
        self.gas_weight = self.volume * self.gas_density
        self.lift_balance = self.lift - (
            self.payload + self.skin_weight + self.gores_weight + self.gas_weight * g0
        )
        self.max_height = self.height[-1]
        self.max_radius = max(self.radius)

    def balloon_plot(self):
        plt.figure()
        plt.axis("equal")
        plt.grid()
        plt.xlabel("Radius (m)")
        plt.ylabel("Height (m)")
        plt.title("Balloon Shape")
        plt.plot(self.radius, self.height, "b")
        plt.plot([-x for x in self.radius], self.height, "b")
        plt.show()
