# These could be from the std math lib, but I like the numpy ones better personally
from numpy import sin, cos, pi, sqrt

# Pull a1 support
from a1_support import GRAVITY_ACC, WATER_DENSITY, HELP_MESSAGE
from a1_support import load_data, plot_water_height

# Fill these in with your details
__author__ = "Cal Wing"
__email__ = "cal@wing.id.au"
__date__ = "14/02/2025"
__version__ = "1.0.0"

# Task 1 [Broken]
def determine_power_used(water_mass: float, elevation: float, 
                         pumping_time: float, efficiency: float) -> float:
    """
    Calculates the power required to pump a certain mass of water a certain height
    taking into account the pumping_time and efficiency.
    
    Parameters: 
        water_mass (float): the mass of the water pumped [kg]
        elevation (float): the height difference [m]
        pumping_time (float): the amount of time the pump is running for [hrs]
        efficiency (float): the conversion efficiency [%]
        
    Returns:
        (float): the power required by the pump [kW]
    """
    # Ideal energy needed to lift the water
    potenital_energy = water_mass * GRAVITY_ACC * elevation # J = kg * m/s^2 * m = mgh

    # Actual amount of energy needed based on eff
    #   Here an 85% eff means an extra 15% is needed to pump the water up
    electrical_energy = potenital_energy / (efficiency/100) # J = J * SCALER
    
    # Actual electrial power used to pump the water
    power_used = electrical_energy / (pumping_time*60*60) # W = J / S = J / (hr*60*60)
    
    return power_used / 1e3 # W -> kW

# Task 2
def determine_water_released(gen_power: float, elevation: float, 
                             pumping_time: float, efficiency: float) -> float:
    """
    Calculates the mass of water released required to generate a specified power
    
    Parameters:
    gen_power (float): the specified power to be generated [kW]
    elevation (float): the height difference [m]
    pumping_time (float): the time the pump is running for [hrs]
    efficiency (float): the conversion efficiency [%]
    
    Returns:
        (float): the mass of the water required [kg]
    """
    
    # How much electrical engery was generated?
    electrical_energy = (gen_power * 1e3) * (pumping_time*60*60) # J = W * S = (kW * 10^3) * (hrs*60*60)
    
    # Need more potential energy to get electrial energy
    #   Again here for 85% eff need 15% more potenitial enegry to get pot-eng
    potential_energy = electrical_energy * (efficiency/100) # J = J * Scaler

    water_mass = potential_energy / (GRAVITY_ACC * elevation) # kg = J / (m/s^2 * m)

    return water_mass

# Task 3
def determine_cost_to_pump(gen_power: float, pumping_time: float, 
                           off_peak_tariff: float) -> float:
    """
    Calculates the cost of using the pump during off peak period

    Parameters:
        gen_power (float): the amount of power generated by the pump [kW]
        pumping_time (float): the amount of time the pump is running for [hrs]
        off_peak_tariff (float): the off-peak tarriff cost for electricity [$/kWh]

    Returns:
        (float): the cost of running the pump ($)
    """

    cost = gen_power * pumping_time * off_peak_tariff

    return cost

# Task 4
def calc_speed_at_outlet(water_height: float) -> float:
    """
    Calculates the speed of the water at the outlet of the dam

    Parameters:
        water_height (float): the height of the water in the pipe [m]

    Returns:
        (float): the speed of the water at the outlet [m/s]
    """

    speed = sqrt(2 * water_height * GRAVITY_ACC)

    return speed

# Task 5
def calc_new_water_height(old_water_height: float, reservoir_area: float, 
                    outlet_area: float, time_inc: float) -> tuple[float, float]:
    
    

    return (0.0, 0.0)


# See if I get what the task sheet wants
def sheet_tasks():
    import numpy as np
    TASKS = (
        (determine_power_used, (5e6, 250, 8, 85), 500.9191176470588),
        (determine_water_released, (300, 250, 8, 85), 2994495.4128440367),
        (determine_cost_to_pump, (300, 8, 0.02), 48),
        (calc_speed_at_outlet, (30), 24.261079942986875),
        (calc_new_water_height, (20, 40000, 1, 30), (19.985143183382704, 592567.1021442247)),
    )

    for i, (task, args, expected) in enumerate(TASKS):
        print(f'Task {i+1}')
        result = task(*args)

        # Number comparison
        if isinstance(expected, (float, int)):
            foo = np.isclose(result, expected, 0.001)
            print(f'{task.__qualname__}{args} -> {result} {"~=" if foo else "!="} {expected}')
            print(f'Task is{"" if foo else " NOT"} close enough to expected result.')        
        
        # Literal Comparison
        elif expected is not None:                        
            print(f'{task.__qualname__}{args} -> {result} {"==" if result == expected else "!="} {expected}')
            print(f'Task{"" if result == expected else " NOT"} equal to expected result.')
        
        # Just run the task
        else:
            print(f'{task.__qualname__}{args} -> {result}')
        
        # Newline
        print()

if __name__ == '__main__':
    sheet_tasks()