from .balloon_optimization import OptimizeBalloon
from .battery import Battery
from .panel import SolarPanel
from .enums import PayloadOperation

from scipy.optimize import newton
import numpy as np


class Optimize:
    def __init__(
        self,
        # Balloons # Missing strenghts + 4
        balloon_system_type,
        balloon_1_film_density_kg_m2,
        balloon_1_tape_density_kg_m,
        balloon_1_no_tapes,
        balloon_1_max_stress_film,
        balloon_1_max_stress_tape,
        balloon_2_film_density_kg_m2,
        balloon_2_tape_density_kg_m,
        balloon_2_no_tapes,
        balloon_2_max_stress_film,
        balloon_2_max_stress_tape,
        # Settings parameters
        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,
        # Solar panel
        solar_panel_surface_density,  #
        solar_panel_efficiency,  #
        solar_panel_safety_coefficient,  #
        solar_panel_tracking_type,  #
        solar_panel_power_fraction,
        # Battery
        battery_power_density,  #
        battery_specific_energy,  #
        battery_max_discharge,  #
        battery_safety_coefficient,  #
        # Mission
        latitude,  #
        longitude,  #
        maximal_altitude_m,  #
        minimal_altitude_m,
        flight_date,  #
        flight_duration,  #
        desired_daily_altitude_change,  #
        night_altitude_change,  #
        maximal_descend_rate,  #
        maximal_upwind_speed,  #
        maximal_temp_difference,  #
        # Payload
        payload_mass_kg,  #
        auxiliaries_weight,  #
        payload_power_requirement,  #
        payload_operation,  #
        # Others
        init_balloon_ratio_binb,
        init_balloon_ratio_tandem,
        p_max_diff,
        p_diff_normal,
        float_offset,
        power_heater,
        power_auxiliaries,
    ):
        # Initialize instance variables for balloon
        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
        self.balloon_1_max_stress_film = balloon_1_max_stress_film
        self.balloon_1_max_stress_tape = balloon_1_max_stress_tape
        self.balloon_2_max_stress_film = balloon_2_max_stress_film
        self.balloon_2_max_stress_tape = balloon_2_max_stress_tape
        # Initialize instance variables for solar panel
        self.solar_panel_surface_density = solar_panel_surface_density
        self.solar_panel_efficiency = solar_panel_efficiency
        self.solar_panel_safety_coefficient = solar_panel_safety_coefficient
        self.solar_panel_tracking_type = solar_panel_tracking_type
        self.solar_panel_power_fraction = solar_panel_power_fraction
        # Initialize instance variables for battery
        self.battery_power_density = battery_power_density
        self.battery_specific_energy = battery_specific_energy
        self.battery_max_discharge = battery_max_discharge
        self.battery_safety_coefficient = battery_safety_coefficient
        # Initialize other instance variables
        self.latitude = latitude
        self.longitude = longitude
        self.flight_date = flight_date
        self.flight_duration = flight_duration
        self.desired_daily_altitude_change = desired_daily_altitude_change
        self.night_altitude_change = night_altitude_change
        self.maximal_descend_rate = maximal_descend_rate
        self.auxiliaries_weight = auxiliaries_weight
        self.payload_power_requirement = payload_power_requirement
        self.payload_operation = payload_operation
        self.power_heater = power_heater
        self.power_auxiliaries = power_auxiliaries

        ## Initialization of output variables
        # For balloon iterations
        self.balloon1_radius_iterations = []
        self.balloon1_height_iterations = []
        self.balloon2_radius_iterations = []
        self.balloon2_height_iterations = []
        self.total_iterations = None
        self.balance_iterations = []
        self.night_energy_per_m = None
        self.day_energy_per_m = None
        self.altitude_array = None
        self.m_gas_balloon_1_array = None
        self.m_gas_balloon_2_array = None
        self.V_balloon_1_array = None
        self.V_balloon_2_array = None

    def initialization(self):
        # Initializing solar panel
        print(self.minimal_altitude_m)
        self.solar_panel = SolarPanel(
            efficiency=self.solar_panel_efficiency,
            safety_coefficient=self.solar_panel_safety_coefficient,
            surface_density=self.solar_panel_surface_density,
            latitude=self.latitude,
            longitude=self.longitude,
            min_altitude=self.minimal_altitude_m,
            tracking_type=self.solar_panel_tracking_type,
            date_departure=self.flight_date,
            flight_days=self.flight_duration,
        )

        self.hours_day = self.solar_panel.get_daylight_hours(
            self.solar_panel_power_fraction
        )
        self.hours_night = 24 - self.hours_day
        print(f"Self hours day: {self.hours_day}")
        print(f"Self hours night: {self.hours_night}")
        # Calculating energy requirements from auxiliaries and heater
        self.auxiliaries_energy = self.power_auxiliaries * 24
        self.heater_energy = self.power_heater * self.hours_night

        # Calculating initial energy requirement from payload and auxiliaries and heater
        match self.payload_operation:
            case PayloadOperation.DAY:
                self.payload_energy = self.payload_power_requirement * self.hours_day
                battery_energy_requirement = self.heater_energy
                battery_power_requirement = self.power_heater
            case PayloadOperation.NIGHT:
                self.payload_energy = self.payload_power_requirement * self.hours_night
                battery_energy_requirement = self.payload_energy + self.heater_energy
                battery_power_requirement = (
                    self.payload_power_requirement + self.power_heater
                )
            case PayloadOperation.BOTH:
                self.payload_energy = self.payload_power_requirement * 24
                battery_energy_requirement = (
                    self.payload_power_requirement * self.hours_night
                ) + self.heater_energy
                battery_power_requirement = (
                    self.payload_power_requirement + self.power_heater
                )
            case _:
                print("Unknown payload operation Day/Night/Both.")
        self.solar_panel.set_required_energy(
            self.payload_energy + self.auxiliaries_energy + self.heater_energy
        )
        self.solar_panel.set_required_power(
            self.payload_power_requirement + self.power_auxiliaries + self.power_heater
        )
        # Defining solar panel weight and surface area based on inputs
        self.solar_panel.define_output()

        # Calculating initial battery
        self.battery = Battery(
            power_density=self.battery_power_density,
            specific_energy=self.battery_specific_energy,
            max_discharge=self.battery_max_discharge,
            safety_coefficient=self.battery_safety_coefficient,
        )
        self.battery.calculate_total_weight(
            required_capacity=battery_energy_requirement,
            required_max_power=battery_power_requirement,
        )
        # Calculating initial payload and accessories weight
        self.total_accessories_weight = (
            self.solar_panel.weight + self.battery.weight + self.payload_mass_kg
        ) * (1 + self.auxiliaries_weight)
        # Initializing balloons
        self.balloon = OptimizeBalloon(
            payload_mass_kg=self.total_accessories_weight,
            maximal_altitude_m=self.maximal_altitude_m,
            minimal_altitude_m=self.minimal_altitude_m,
            balloon_1_film_density_kg_m2=self.balloon_1_film_density_kg_m2,
            balloon_1_tape_density_kg_m=self.balloon_1_tape_density_kg_m,
            balloon_1_no_tapes=self.balloon_1_no_tapes,
            balloon_2_film_density_kg_m2=self.balloon_2_film_density_kg_m2,
            balloon_2_tape_density_kg_m=self.balloon_2_tape_density_kg_m,
            balloon_2_no_tapes=self.balloon_2_no_tapes,
            p_max_diff=self.p_max_diff,
            p_diff_normal=self.p_diff_normal,
            maximal_upwind_speed=self.maximal_upwind_speed,
            maximal_temp_difference=self.maximal_temp_difference,
            float_offset=self.float_offset,
            init_balloon_ratio_binb=self.init_balloon_ratio_binb,
            init_balloon_ratio_tandem=self.init_balloon_ratio_tandem,
            balloon_system_type=self.balloon_system_type,
            SP_initial_theta=self.SP_initial_theta,
            ZP_initial_theta=self.ZP_initial_theta,
            SP_incremental_parameter=self.SP_incremental_parameter,
            ZP_incremental_parameter=self.ZP_incremental_parameter,
            SP_theta_end_error_tolerance=self.SP_theta_end_error_tolerance,
            ZP_theta_end_error_tolerance=self.ZP_theta_end_error_tolerance,
            SP_max_iterations_balloon_shape=self.SP_max_iterations_balloon_shape,
            ZP_max_iterations_balloon_shape=self.ZP_max_iterations_balloon_shape,
            SP_max_iterations_optimization=self.SP_max_iterations_optimization,
            ZP_max_iterations_optimization=self.ZP_max_iterations_optimization,
        )
        self.balloon.initial_balloons()
        self.balloon.optimization()

    def optimization(self):
        # Function to calculate optimal balloon, solar panels and battery

        # Initial balloon
        tolerance = 1e-3
        self.total_iterations = 0

        initial_guesses = [
            self.total_accessories_weight,
            self.total_accessories_weight * 1.01,
        ]  # Starting points

        def objective(payload):
            print("New minimization iteration")
            self.balloon.payload_mass_kg = payload
            self.balloon.recalculate()
            self.balloon.optimization()
            self.night_energy_per_m, self.day_energy_per_m = (
                self.balloon.energy_consumption()
            )
            self.__calculation_power_energy(
                self.night_energy_per_m, self.day_energy_per_m
            )
            self.total_accessories_weight = (
                self.solar_panel.weight + self.battery.weight + self.payload_mass_kg
            ) * (1 + self.auxiliaries_weight)
            balance = payload - self.total_accessories_weight
            print(
                f"Previous load: {payload}kg, New load: {self.total_accessories_weight}kg, Balance: {balance}"
            )
            self.total_iterations += 1

            self.balloon1_radius_iterations.append(self.balloon.balloon_1.radius)
            self.balloon1_height_iterations.append(self.balloon.balloon_1.height)
            self.balloon2_radius_iterations.append(self.balloon.balloon_2.radius)
            self.balloon2_height_iterations.append(self.balloon.balloon_2.height)
            self.balance_iterations.append(balance)
            return balance

        # Using the secant method to find the root of the objective function
        try:
            root = newton(
                func=objective,
                x0=initial_guesses[0],
                x1=initial_guesses[1],
                tol=tolerance,
                maxiter=100,
            )
            # print(f"Root found at payload: {root} kg")

            # After successful optimization recalculation for all altitudes
            (
                self.altitude_array,
                self.m_gas_balloon_1_array,
                self.m_gas_balloon_2_array,
                self.V_balloon_1_array,
                self.V_balloon_2_array,
            ) = self.balloon.recalculate_all_altitudes()
            # After successful optimization calculate all stress parameters
            self.__calculate_stress()
        except RuntimeError as e:
            print(
                f"No convergence of optimization of system after maximum iterations: {str(e)}"
            )

    def __calculation_power_energy(self, night_energy_per_m, day_energy_per_m):

        # Altitude control
        night_altitude_change = (
            self.desired_daily_altitude_change * self.night_altitude_change
        )
        day_altitude_change = self.desired_daily_altitude_change - night_altitude_change
        # Night and day and total daily altitude change energy in Wh
        night_altitude_control_energy = (
            night_altitude_change * night_energy_per_m / 3600
        )
        day_altitude_control_energy = day_altitude_change * day_energy_per_m / 3600

        total_altitude_control_energy = (
            night_altitude_control_energy + day_altitude_control_energy
        )
        # Altitude control power for night and day
        self.altitude_control_power_day = day_altitude_control_energy / (self.hours_day)
        self.altitude_control_power_night = night_altitude_control_energy / (
            self.hours_night
        )
        # Required power to move at given descend rate m/s
        altitude_control_power = self.maximal_descend_rate * day_energy_per_m

        # Total energy [Wh] and power [W]
        total_energy = (
            self.payload_energy
            + total_altitude_control_energy
            + self.auxiliaries_energy
            + self.heater_energy
        )
        total_power = (
            altitude_control_power
            + self.payload_power_requirement
            + self.power_auxiliaries
        )
        print(
            f"Total energy: {total_energy}, Payload energy: {self.payload_energy}, Altitude control energy: {total_altitude_control_energy}"
        )

        # Solar panels
        self.solar_panel.set_required_energy(total_energy)
        self.solar_panel.set_required_power(total_power)
        self.solar_panel.define_output()

        # Battery
        match self.payload_operation:
            case PayloadOperation.DAY:
                night_payload_energy_requirement = 0
                night_payload_power_requirement = 0
            case PayloadOperation.NIGHT:
                night_payload_energy_requirement = self.payload_energy
                night_payload_power_requirement = self.payload_power_requirement
            case PayloadOperation.BOTH:
                night_payload_energy_requirement = (
                    self.payload_power_requirement * self.hours_night
                )
                night_payload_power_requirement = self.payload_power_requirement
            case _:
                print("Unknown payload operation Day/Night/Both.")

        night_altitude_control_power = night_altitude_control_energy / (
            self.hours_night * 60 * 60
        )
        battery_energy_requirement = (
            night_payload_energy_requirement + night_altitude_control_energy
        ) + self.heater_energy
        battery_power_requirement = (
            night_payload_power_requirement + night_altitude_control_power
        ) + self.power_heater
        self.battery.calculate_total_weight(
            required_capacity=battery_energy_requirement,
            required_max_power=battery_power_requirement,
        )

    def get_energy_balance_plot_parameters(self):
        battery_capacity = self.battery.capacity
        min_discharge_level = self.battery_max_discharge * self.battery.capacity
        power_curve = self.solar_panel.power_curve_real

        # Equipment power requirements
        altitude_control_day = self.altitude_control_power_day
        altitude_control_night = self.altitude_control_power_night
        payload_power = self.payload_power_requirement

        # Operation mode of the payload
        payload_operation = self.payload_operation

        # Time parameters
        # sunrise = np.argmax(power_curve > 0)
        # sunset = len(power_curve) - np.argmax(power_curve[::-1] > 0) - 1
        sunrise = self.solar_panel.sunrise
        sunset = self.solar_panel.sunset
        # Minutes per day
        number_of_days = 3
        minutes_day = 24 * 60
        minutes_total = minutes_day * number_of_days

        # Calculate original total energy consumption for altitude control for one day
        original_day_energy = altitude_control_day * (sunset - sunrise)
        original_night_energy = altitude_control_night * (
            minutes_day - (sunset - sunrise)
        )

        # Scale altitude power based on solar power, maintaining total energy usage
        solar_power_scaled = np.where(
            power_curve > 0, power_curve / np.max(power_curve), 0
        )
        # Apply scaling to altitude power during day times
        scaled_day_power = altitude_control_day * solar_power_scaled

        # Ensure the total energy used in the scaled day matches the original
        scaled_day_energy = np.sum(scaled_day_power)
        scaling_factor = original_day_energy / scaled_day_energy
        scaled_day_power *= scaling_factor

        # Create complete altitude power array for one day
        altitude_power_day = np.where(
            (np.arange(minutes_day) >= sunrise) & (np.arange(minutes_day) <= sunset),
            scaled_day_power,
            altitude_control_night,
        )

        # Extend this pattern to the required number of days
        altitude_power_array = np.tile(altitude_power_day, number_of_days)

        # Extend power_curve to two days
        solar_power_array = np.tile(power_curve, number_of_days)

        # Determine day and night times for two days
        minutes = np.arange(minutes_total)  # Minutes for two days
        day_mask = (minutes % minutes_day >= sunrise) & (
            minutes % minutes_day <= sunset
        )
        night_mask = ~day_mask

        # Calculate power consumption arrays

        payload_power_array = np.zeros(minutes_total)
        heating_auxiliary_power_array = (
            np.ones(minutes_total) * self.power_auxiliaries
            + self.power_heater * night_mask
        )
        # Adjust payload power consumption based on operation mode
        if payload_operation == PayloadOperation.BOTH:
            payload_power_array[day_mask] = payload_power
            payload_power_array[night_mask] = payload_power
        elif payload_operation == PayloadOperation.DAY:
            payload_power_array[day_mask] = payload_power
        elif payload_operation == PayloadOperation.NIGHT:
            payload_power_array[night_mask] = payload_power

        # Calculate total power consumption and net power balance
        total_power_consumption = (
            altitude_power_array + payload_power_array + heating_auxiliary_power_array
        )
        power_balance = solar_power_array - total_power_consumption

        # Simulate battery charge state using a more controlled method
        charge_state = np.zeros(minutes_total)
        charge_state[0] = battery_capacity

        for i in range(1, minutes_total):
            potential_charge = charge_state[i - 1] + power_balance[i] / 60
            if potential_charge > battery_capacity:
                charge_state[i] = battery_capacity
            else:
                charge_state[i] = potential_charge
        # Return a dictionary containing all the parameters
        return {
            "charge_state": charge_state,
            "battery_capacity": battery_capacity,
            "min_discharge_level": min_discharge_level,
            "power_balance": power_balance,
            "solar_power_array": solar_power_array,
            "payload_power_array": payload_power_array,
            "altitude_power_array": altitude_power_array,
            "heating_auxiliary_power_array": heating_auxiliary_power_array,
        }

    def __calculate_stress(self):
        (
            self.balloon1_skin_length,
            self.balloon1_sigma_meridional,
            self.balloon1_sigma_film_calc,
            self.balloon1_sigma_tape_calc,
            self.balloon1_sigma_total,
        ) = self.balloon.balloon_1.get_meridional_stress(
            max_stress_film=self.balloon_1_max_stress_film,
            max_stress_tape=self.balloon_1_max_stress_tape,
        )
        # For balloon 2
        (
            self.balloon2_skin_length,
            self.balloon2_sigma_meridional,
            self.balloon2_sigma_film_calc,
            self.balloon2_sigma_tape_calc,
            self.balloon2_sigma_total,
        ) = self.balloon.balloon_2.get_meridional_stress(
            max_stress_film=self.balloon_2_max_stress_film,
            max_stress_tape=self.balloon_2_max_stress_tape,
        )

    def get_stress_plot_parameters(self):

        optimize_output_balloon_1_stress = {
            "skin_length": self.balloon1_skin_length,
            "sigma_meridional": self.balloon1_sigma_meridional,
            "sigma_film": self.balloon1_sigma_film_calc,
            "sigma_tape": self.balloon1_sigma_tape_calc,
            "sigma_total": self.balloon1_sigma_total,
        }

        optimize_output_balloon_2_stress = {
            "skin_length": self.balloon2_skin_length,
            "sigma_meridional": self.balloon2_sigma_meridional,
            "sigma_film": self.balloon2_sigma_film_calc,
            "sigma_tape": self.balloon2_sigma_tape_calc,
            "sigma_total": self.balloon2_sigma_total,
        }

        return optimize_output_balloon_1_stress, optimize_output_balloon_2_stress

    def get_output_massage(self):
        # Get output message
        # Validation of min altitudes
        min_altitude = self.balloon.real_min_altitude
        if self.minimal_altitude_m < min_altitude:
            message_1 = f"Minimal altitude is below the calculated minimal altitude. Difference is: {abs(self.minimal_altitude_m - min_altitude)}. Try different Balloon System Type if you want higher altitude range."
        else:
            message_1 = f"Minimal altitude is within the calculated minimal altitude."
        # Validation of float altitude
        float_altitude = self.balloon.float_altitude
        if float_altitude < min_altitude:
            message_2 = f"Float altitude is below the calculated minimal altitude. Difference is: {abs(float_altitude - min_altitude)}. Try lowering max altitude, the max temperature and upwind speed or use lighter materials."
        else:
            message_2 = f"Float altitude is valid."
        # Validation of max force balloon 2
        balloon1_calc_stress_median = np.median(self.balloon1_sigma_meridional)
        balloon1_max_stress_median = np.median(self.balloon1_sigma_total)

        if balloon1_calc_stress_median > balloon1_max_stress_median:
            message_3 = "The calculated stress of balloon 1 is higher than the maximum allowed stress. Try lowering the payload mass or use stronger materials."
        # If maximul allowed stress is more than 5 times more than allowed stress, thinner materials
        elif balloon1_max_stress_median / balloon1_calc_stress_median > 5:
            message_3 = "The calculated stress of balloon 1 is significantly lower than the maximum allowed stress. Try using thinner materials."
        else:
            message_3 = "The calculated stress of balloon 1 is within the allowed stress limits."

        # Validation of max force balloon 2
        balloon2_calc_stress_median = np.median(self.balloon2_sigma_meridional)
        balloon2_max_stress_median = np.median(self.balloon2_sigma_total)

        if balloon2_calc_stress_median > balloon2_max_stress_median:
            message_4 = " The calculated stress of balloon 2 is higher than the maximum allowed stress. Try lowering the payload mass or use stronger materials."
        elif balloon2_max_stress_median / balloon2_calc_stress_median > 5:
            message_4 = "The calculated stress of balloon 2 is significantly lower than the maximum allowed stress. Try using thinner materials."
        else:
            message_4 = "The calculated stress of balloon 2 is within the allowed stress limits."

        # Final message
        message = f"{message_1}\n{message_2}\n{message_3}\n{message_4}"
        return message
