# 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
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
    # [NOTE] I got stuck on this for a bit, unsure if it was my sleep deprived brain but took me
    #   me a bit to get the efficiency calc.
    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]:
    
    water_mass_out = outlet_area * calc_speed_at_outlet(old_water_height) \
        * time_inc * WATER_DENSITY
    
    new_water_height = old_water_height - \
        (water_mass_out / (reservoir_area * WATER_DENSITY))

    return new_water_height, water_mass_out


# Task 6
def calc_heights_water_out(initial_height: float, final_height: float, 
                            reservoir_area: float, outlet_area: float, 
                            time_inc: float
                        ) -> tuple[tuple[float, ...], tuple[float, ...]]:
    
    # Track the water height & mass
    water_height, water_mass = [], []

    # Do the first iteration
    new_water_height, water_mass_out = \
        calc_new_water_height(initial_height, reservoir_area, 
                              outlet_area, time_inc)

    # Store the new values    
    water_height.append(new_water_height)
    water_mass.append(water_mass_out)


    # Do the loop
    while water_height[-1] > final_height:
        new_water_height, water_mass_out = \
            calc_new_water_height(water_height[-1], reservoir_area, 
                                outlet_area, time_inc)
        
        water_height.append(new_water_height)
        water_mass.append(water_mass_out)


    return tuple(water_height), tuple(water_mass)

# Task 7
def calc_energy_power(heights: tuple[float, ...], water_mass_outs: tuple[float, ...], 
                      relative_elevation: float, efficiency: float, 
                      time_inc: float) -> tuple[tuple[float, ...], tuple[float, ...]]:
    
    energy, power = (), ()

    for i in range(len(heights)):
        energy_jules = water_mass_outs[i] * GRAVITY_ACC * \
            (heights[i] + relative_elevation) * (efficiency/100) # note the comma

        energy += energy_jules / 3.6e6, # convert to kWh
        power += (energy_jules / time_inc) / 1e3, # convert to kW

    return energy, power

# Task 8
def calc_daily_profit(energy: tuple[float, ...], peak_tariff: float, 
                      off_peak_tariff: float, efficiency: float) -> float:
    
    total_profit = sum(energy) * (peak_tariff - off_peak_tariff / ((efficiency/100)**2))   

    return total_profit

# Task 9
def print_table(start_relative_elevation: float, step_size: float, 
                num_steps: int, initial_height: float, final_height: float, 
                reservoir_area: float, outlet_area: float, time_inc: float, 
                peak_tariff: float, off_peak_tariff: float, efficiency: float):

    # Print the header
    print("#"*79)
    print("# Relative elevation (m)  #    Daily Profit ($)     #   Total Energy (kWh)    #")
    print("#"*79)

    table_data = ()

    rel_elevation = start_relative_elevation
    for i in range(num_steps):
        
        heights, water_mass_outs = calc_heights_water_out(initial_height, final_height, reservoir_area, outlet_area, time_inc)
        energy, _ = calc_energy_power(heights, water_mass_outs, rel_elevation, efficiency, time_inc)
        
        total_energy = sum(energy)

        profit = calc_daily_profit(energy, peak_tariff, off_peak_tariff, efficiency)

        table_data += (rel_elevation, profit, total_energy), 
        
        rel_elevation += step_size

    # Print the table
    for elevation, profit, total_energy in table_data:
        print(f"#{elevation:^25d}#{profit:^25.2f}#{total_energy:^25.2f}#")

    print("#"*79)

def main():
    cmd, data = None, None
    while True:
        cmd = input("Please enter a command: ").lower()
        
        # Help CMD
        if cmd[0] == "h":
            print(HELP_MESSAGE)

        # Quit
        elif cmd[0] == "q":
            while True:
                cmd = input("Are you sure (y/n): ").lower()
                if cmd in ['y', 'n']:
                    break
                print("Please enter a valid reponse")
            if cmd == 'y':
                break
            
        # Read Data
        elif cmd[0] == "r":
            directory = input("Please specify the directory: ")
            filename = input("Please specify the filename: ")
            data = load_data(directory, filename)

        # Print Table
        elif cmd[0] == "p":
            if data is None:
                print("Please load data before using this command")
                continue
            
            cmds = cmd.split()
            if len(cmds) != 4:
                print("Please enter the correct number of arguments")
                continue
            
            # This makes ints & floats
            foo = ()
            for c in cmds:
                if "." in str(c):
                    foo += float(c),
                elif c.isnumeric():
                    foo += int(c),
                else:
                    foo += c,
            
            print_table(*foo[1:], *data)

        # Make Graph
        elif cmd[0] == "s":
            if data is None:
                print("Please load data before using this command")
                continue

            print("Simulating water heights...")
            heights, _ = calc_heights_water_out(*data[0:5])

            plot_water_height(heights, data[4])

        # Catch All
        else:
            print("Please enter a valid command")



##### End of Assisgnment #####
if __name__ == '__main__':
    # See if I get what the task sheet wants
    # [NOTE] It seems I am witing my own test suite :/ that wasn't the initention lol
    # Tests do equal what I get when i test it in the shell
    def sheet_tasks():
        from debug_lib import NO_EVAL, TUPLE_EVAL
        from debug_lib import compair_outputs, auto_input

        # from a1_alex import calc_daily_profit

        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)),
            (calc_heights_water_out, (40, 5, 40000, 1, 30), ((39.97898928844613, 39.95798409574207, 39.93698442188801), (838016.4324684857, 837796.3120398963, 837576.1916037091), TUPLE_EVAL)),
            (calc_energy_power, tuple(i for i in calc_heights_water_out(30, 20, 40000, 1, 30)) + (200, 85, 30), ((386.6002008976669, 386.4523727180418, 386.3045723539472), (46392.024107720026, 46374.28472616501, 46356.548682473665), TUPLE_EVAL)),
            (calc_daily_profit, (calc_energy_power(*(tuple(i for i in calc_heights_water_out(30, 20, 40000, 1, 30)) + (200, 85, 30)))[0], 0.02, 0.005, 85), 2718.700033709491)
        )

        compair_outputs(TASKS)
        
        print("\nTask 8")
        print_table(280, 20, 6, 30, 20, 40000, 1, 30, 0.02, 0.005, 85)


        print("\nTask 9")
        global input # This is some magic I wrote a while ago to be an automatic input func
        input = auto_input([
            "h", "s", "q", "n", 
            "r", "test", "test_data.txt", 
            "p 280 20 6",
            
            "r", "test", "test_data_2.txt",
            's',
            
            "q", "y"
            ])
        main()

    sheet_tasks()