from functions import g0, r_air, r_He, H, K, P, T
import numpy as np
from math import atan, sin, cos, cbrt, pi

# ==================
# Calculation of Gas Parameters
# ==================
# region Gas Parameters


# def gas_parameters_altitude(Altitude):
#     """
#     Calculate atmospheric pressure and temperature at a specified altitude
#     following the ICAO Standard Atmosphere 1973 model.

#     Parameters:
#     - Altitude (float): Geopotential altitude in meters.

#     Returns:
#     - Press (float): Atmospheric pressure in Pascals (Pa).
#     - Temp (float): Temperature in Kelvin (K).

#     Note: This function relies on predefined constants for the Earth's atmosphere
#     and assumes a piecewise linear temperature gradient based on altitude.

#     Author: David Hofman (hofmada7@fel.cvut.cz), 2024
#     """
#     layer = np.searchsorted(H, Altitude, side="right") - 1

#     Temp = T[layer] + (Altitude - H[layer]) * K[layer]
#     if K[layer] == 0:  # Isothermal layer
#         Press = P[layer] * np.exp(-g0 * (Altitude - H[layer]) / (r_air * T[layer]))
#     else:
#         Press = P[layer] * (1 + K[layer] * (Altitude - H[layer]) / T[layer]) ** (
#             -g0 / (K[layer] * r_air)
#         )


#     return Press, Temp
# For arrays
def gas_parameters_altitude(Altitude):
    """
    Calculate atmospheric pressure and temperature at specified altitudes
    following the ICAO Standard Atmosphere 1973 model. Accepts both numpy arrays
    and single float values.

    Parameters:
    - Altitude (float or np.array): Geopotential altitudes in meters.

    Returns:
    - Press (float or np.array): Atmospheric pressures in Pascals (Pa).
    - Temp (float or np.array): Temperatures in Kelvin (K).
    """
    # Ensure input is array-like for uniform processing
    scalar_input = np.isscalar(Altitude)
    if scalar_input:
        Altitude = np.array([Altitude], dtype=float)

    layer = np.searchsorted(H, Altitude, side="right") - 1

    Temp = T[layer] + (Altitude - H[layer]) * K[layer]
    Press = np.zeros_like(Altitude, dtype=float)

    # Isothermal layer calculation
    isothermal = K[layer] == 0
    non_isothermal = ~isothermal

    # For non-isothermal layers
    if np.any(non_isothermal):
        temp_ratio = (
            1
            + K[layer][non_isothermal]
            * (Altitude[non_isothermal] - H[layer][non_isothermal])
            / T[layer][non_isothermal]
        )
        Press[non_isothermal] = P[layer][non_isothermal] * (
            temp_ratio ** (-g0 / (K[layer][non_isothermal] * r_air))
        )

    # For isothermal layers
    if np.any(isothermal):
        Press[isothermal] = P[layer][isothermal] * np.exp(
            -g0
            * (Altitude[isothermal] - H[layer][isothermal])
            / (r_air * T[layer][isothermal])
        )

    # If the input was a scalar, return scalar values instead of one-element arrays
    if scalar_input:
        return Press[0], Temp[0]
    return Press, Temp


# endregion

# ==================
# Calculation of Drag Coefficient for Spheres
# ==================
# region Drag Coefficient Calculation


def sphere_drag(Re):
    """
    Calculates the drag coefficient for a sphere as a function of the Reynolds number,
    using predefined coefficients for different ranges of Reynolds numbers.

    Parameters:
    - Re (float): Reynolds number, a dimensionless quantity in fluid mechanics.

    Returns:
    - float: Drag coefficient for the sphere at the given Reynolds number.

    Note: This function uses piecewise defined polynomial approximations for the drag
    coefficient based on the Reynolds number. It handles both interpolation within the
    provided ranges and extrapolation beyond the defined limits (Re > 10^7 and Re < 0.1).
    Sphere drag data based on: Carlson, Leland A., and Walter J. Horn. "New thermal and trajectory model for high-altitude balloons." Journal of Aircraft 20.6 (1983): 500-507.

    Author: David Hofman (hofmada7@fel.cvut.cz), 2024
    """
    coefficients = [
        (0.1, 40, 14.966487175135969, -0.5883081156967886),
        (40, 170, 11.588961547473467, -0.5189752838410412),
        (170, 1000, 3.9948716140022236, -0.311599230175758),
        (1000, 1000000, 0.4642, 0.0),
        (1000000, 2000000, 0.015797725394546776, 0.24468509595490195),
        (2000000, 3345180, 1.4472765531252076e-07, 1.0442437455029776),
        (3345180, 10000000, 2.226066742665555e-21, 3.1613690104808154),
    ]

    for x1, x2, a, b in coefficients:
        if x1 <= Re <= x2:
            return a * (Re**b)

    # Handle extrapolation
    if Re < coefficients[0][0]:
        a, b = coefficients[0][2], coefficients[0][3]
        return a * (Re**b)
    else:
        a, b = coefficients[-1][2], coefficients[-1][3]
        return a * (Re**b)


# endregion


# ==================
# Calculation of Air Viscosity
# ==================


def calculate_air_viscosity(temperature_K):
    """
    Calculate the dynamic viscosity of air using Sutherland's Law.

    Parameters:
    temperature_K : float
        Temperature in Kelvin at which to calculate viscosity.

    Returns:
    float
        Dynamic viscosity in Pa·s.
    """
    # Constants for air
    mu_0 = 1.716e-5  # Pa·s, viscosity at reference temperature
    T_0 = 273.15  # K, reference temperature (0°C)
    C = 110.4  # K, Sutherland's constant

    # Sutherland's Law formula
    mu = mu_0 * ((temperature_K + C) / (T_0 + C)) * ((T_0 / temperature_K) ** 1.5)

    return mu


# ==================
# Calculation of Air Density
# ==================


def calculate_gas_density(temperature_K, pressure_Pa, specific_gas_constants):
    """
    Calculate the density of a gas using the ideal gas law for given temperature and pressure conditions.

    Parameters:
    - temperature_K: The absolute temperature of the gas in Kelvin.
    - pressure_Pa: The pressure of the gas in Pascals.
    - specific_gas_constant: The specific gas constant for the gas in question.

    Returns:
    - Gas density in kilograms per cubic meter (kg/m³).

    The function computes the density of the gas by applying the ideal gas law, which relates the pressure, volume, temperature, and the amount of gas. The density is calculated as the ratio of pressure to the product of the specific gas constant and the temperature.
    """

    # Calculate air density using the ideal gas law
    air_density = pressure_Pa / (specific_gas_constants * temperature_K)

    return air_density


# ==================
# Calculation of Reynolds number
# ==================


def calculate_reynolds_number(temperature_K, pressure_Pa, velocity_m_s, length_m):
    """
    Calculate the Reynolds number for air at a given temperature, pressure, velocity, and length.

    Parameters:
    - temperature_K: Temperature in Kelvins
    - pressure_Pa: Pressure in Pascals
    - velocity_m_s: Velocity of the flow in meters per second
    - length_m: Characteristic length in meters (sphere - diameter)

    Returns:
    - Reynolds number (dimensionless)
    """

    # Calculate air density using the ideal gas law
    density = calculate_gas_density(temperature_K, pressure_Pa, r_air)

    # Calculate air viscosity using Sutherland's law
    viscosity = calculate_air_viscosity(temperature_K)

    # Calculate Reynolds number
    reynolds_number = (density * velocity_m_s * length_m) / viscosity

    return reynolds_number


# ==================
# Calculation of Specific Buoyancy
# ==================
# def calculate_specific_buoyancy(temperature_K, pressure_Pa):
#     # Density of air at external conditions
#     density_air = calculate_gas_density(temperature_K, pressure_Pa, r_air)

#     # Buoyant force calculation based solely on external air density
#     b = g0 * density_air
#     return b


def calculate_specific_buoyancy(temperature_K, pressure_Pa, pressure_difference_Pa):
    """
    Calculate the specific buoyancy in Newtons per cubic meter for a given set of environmental conditions, where the medium is a mixture of air and helium. The specific buoyancy is affected by temperature and pressure, and pressure difference.

    Parameters:
    - temperature_K: Temperature in Kelvin, which influences the density and buoyancy of gases.
    - pressure_Pa: Initial atmospheric pressure in Pascals, representing the environmental baseline pressure.
    - pressure_difference_Pa: Difference in pressure in Pascals, indicating an increase  from the baseline pressure.

    Returns:
    - Specific buoyancy (b) in Newtons per cubic meter (N/m³).

    The function uses the gravity constant (g0), specific gas constants for helium (r_He) and air (r_air), to compute the specific buoyancy based on the change in pressure due to the presence of helium compared to air.
    """

    b = (
        g0
        / temperature_K
        * (pressure_Pa * r_He - (pressure_Pa + pressure_difference_Pa) * r_air)
        / (r_air * r_He)
    )
    return b


# ==================
# Calculation of Sin and Cos from Tan
# ==================


def trig_via_arctan(t):
    """
    Compute sin(alpha) and cos(alpha) by first calculating the angle alpha using the arctan function from the tan(alpha) value.

    Parameters:
    - t (float): The value of tan(alpha), where alpha is the angle in radians derived from the arctangent of t.

    Returns:
    - sin(alpha) (float): Value of sin(alpha)
    - cos(alpha) (float): Value of cos(alpha)

    This method involves calculating the angle alpha using math.atan(t), which is the inverse tangent function, and then using this angle to calculate sin(alpha) and cos(alpha) directly. This approach is traditional and commonly used in scenarios where the angle itself might be needed for further calculations or when using libraries that may optimize these trigonometric computations.
    """
    alpha = atan(t)
    sin_alpha = sin(alpha)
    cos_alpha = cos(alpha)
    return sin_alpha, cos_alpha


# ==================
# Calculation of Sphere Radius
# ==================
def calculate_radius_sphere(Volume):
    """
    Calculate the radius of a sphere given its volume.

    Parameters:
    - Volume (float): The volume of the sphere.

    Returns:
    - radius (float): The radius of the sphere.
    """
    radius = cbrt(3 * Volume / (4 * pi))
    return radius


# ==================
# Calculation of  Circle Area
# ==================
def calculate_area_circle(radius):
    """
    Calculate the area of a circle given its radius.

    Parameters:
    - radius (float): The radius of the circle.

    Returns:
    - area (float): The area of the circle.
    """
    area = 2 * pi * radius
    return area
