"""Compute various values for THE Henry Pump."""
from math import sqrt
import matplotlib.pyplot as plt

# GLOBAL PARAMETERS
g_dt = 0.050  # seconds
g_thread_pitch = 0.5906

g_accel_time = 1.0  # seconds
g_decel_time = 1.0  # seconds

g_stroke_length = 60.0  # inches
g_upper_offset = 9.0  # inches
g_lower_offset = 16.0  # inches

g_motor_full_speed_hz = 120.0  # Hz
g_motor_full_speed_rpm = 1200.0  # RPM

g_run_speed = 120.0  # Hz

MAX_NUT_ENDLOAD = 11000.0  # pounds
FLUID_GRADIENT = 0.43733  # PSI/ft

pump_constants = {
    0.625: 0.046,
    0.75: 0.066,
    0.9375: 0.117,
    1.0625: 0.132,
    1.125: 0.148,
    1.25: 0.182,
    1.5: 0.262,
    1.75: 0.357,
    1.78125: 0.370,
    2.0: 0.466,
    2.25: 0.590,
    2.5: 0.728,
    2.75: 0.861,
    3.75: 1.640,
    4.75: 2.630
}

fluid_load_constants = {
    1.0625: 0.384,
    1.25: 0.531,
    1.5: 0.765,
    1.75: 1.041,
    2.0: 1.36,
    2.25: 1.721,
    2.5: 2.125,
    2.75: 2.571
}


def _ftlb_to_nm(ftlb):
    """Convert ft-lbs to n-m."""
    return ftlb / 0.73756


def _nm_to_ftlb(nm):
    """Convert n-m to ft-lbs."""
    return nm * 0.73756


def max_production(pump_size, strokes_per_minute, stroke_length):
    """Calculate the maximum production capabilities of a pump at a speed."""
    try:
        return pump_constants[pump_size] * strokes_per_minute * stroke_length
    except KeyError:
        print("No pump constant for that size pump.")


def end_load(pump_size, depth, strokes_per_minute, stroke_length):
    """Calculate the end load on the nut."""
    try:
        fluid_load_constant = fluid_load_constants[pump_size]
        return depth * fluid_load_constant * (FLUID_GRADIENT / 0.433)
    except KeyError:
        print("No fluid load constant for that size pump.")


def max_producing_depth(pump_size):
    """Calculate the maximum producing depth."""
    try:
        fluid_load_constant = fluid_load_constants[pump_size]
        return MAX_NUT_ENDLOAD / (fluid_load_constant * (FLUID_GRADIENT / 0.433))
    except KeyError:
        print("No fluid load constant for that size pump.")


def voltage_drop(length, current, temperature, circular_mils, alpha):
    """Calculate the voltage drop in the wire."""
    k20 = 10.37  # 1.26 / 1000.0  # ohms of 1 circular mil foot of conductor
    return sqrt(3) * k20 * (1.0 + alpha * ((5.0 / 9.0) * (temperature - 32.0) - 20.0)) * length * current / circular_mils


def calc_voltage_needed(input_voltage, voltage_drop):
    """Calculate the voltage needed given an input voltage and a voltage drop."""
    return input_voltage + voltage_drop


def print_voltage_drops():
    """Print values for voltage drop at distances."""
    current = 10.4  # Amps
    temperature = 250.0  # degress F
    circular_mils = 10384.0  # cross-sectional area in circular mils of conductor
    alpha_copper = 0.00393
    motor_input_voltage = 380.0
    print("---")
    for x in [400.0, 1000.0, 2000.0, 3000.0, 4000.0, 5000.0, 6000.0, 7000.0, 8000.0, 9000.0, 10000.0]:
        v_dropped = voltage_drop(x, current, temperature, circular_mils, alpha_copper)
        v_output = calc_voltage_needed(motor_input_voltage, v_dropped)
        print("Length: {}".format(x))
        print("Voltage Drop: {} V".format(round(v_dropped, 3)))
        print("Output Voltage: {} V".format(round(v_output, 3)))
        print("Motor Input Voltage {} V".format(round(motor_input_voltage, 3)))
        print("---")


def time_for_freq_change(freq_delta, np_hz, ramp_time):
    """Calculate the time required for a frequency change."""
    return (freq_delta / np_hz) * ramp_time


def dist_for_freq_change(freq_delta, ramp_time, np_hz, np_rpm, thread_pitch):
    """Calculate the distance required for a frequency change."""
    time_required = time_for_freq_change(freq_delta, np_hz, ramp_time)
    return 0.5 * time_required * freq_delta * (np_rpm / np_hz) * (1.0 / 60.0) * thread_pitch


def time_for_uniform_freq(freq, distance, np_hz, np_rpm, thread_pitch):
    """Calculate the time required to move a distance at a constant frequency."""
    return distance / (freq * (np_rpm / np_hz) * (1.0 / 60.0) * thread_pitch)


def calculate_stroke_time(stroke_length, upper_offset, lower_offset, run_frequency, accel_time, decel_time, nameplate_hz=120.0, nameplate_rpm=1200.0, thread_pitch=0.5406):
    """Calculate the time it takes to make a full stroke."""
    ramp_up_time = time_for_freq_change(run_frequency, nameplate_hz, accel_time)
    ramp_up_dist = dist_for_freq_change(run_frequency, accel_time, nameplate_hz, nameplate_rpm, thread_pitch)

    ramp_down_time = time_for_freq_change(run_frequency, nameplate_hz, decel_time)
    ramp_down_dist = dist_for_freq_change(run_frequency, decel_time, nameplate_hz, nameplate_rpm, thread_pitch)

    run_dist = stroke_length - (lower_offset + upper_offset + ramp_up_dist + ramp_down_dist)
    run_time = time_for_uniform_freq(run_frequency, run_dist, nameplate_hz, nameplate_rpm, thread_pitch)

    stroke_time = 2.0 * (ramp_up_time + run_time + ramp_down_time)
    return stroke_time


def print_stroke_times():
    """Print the required frequency for different SPM's with different ramp times."""
    for ramp_i in [0.2, 0.4, 0.6, 0.8, 1.0, 1.2, 1.4, 1.6, 1.8, 2.0]:
        print("-----")
        print("Ramp Rate: {}".format(ramp_i))
        for hz_i in [20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0, 110.0, 120.0]:
            spm = 60 / calculate_stroke_time(g_stroke_length, g_upper_offset, g_lower_offset, hz_i, ramp_i, ramp_i)
            print("Freq: {} Hz., SPM: {}".format(hz_i, spm))


def sim(run_freq, dt, num_strokes, stroke_length, accel_time, decel_time, upper_offset, lower_offset, motor_full_speed_hz, motor_full_speed_rpm, thread_pitch, use_average=False):
    """Run the main simulation loop."""
    strokes = 0

    time_to_turnaround = time_for_freq_change(run_freq, motor_full_speed_hz, decel_time)
    # time_to_turnaround = decel_time - (motor_full_speed_hz - run_freq) * decel_time / motor_full_speed_hz
    print("Turn Time: {} s.".format(time_to_turnaround))

    distance_to_turnaround = dist_for_freq_change(run_freq, decel_time, motor_full_speed_hz, motor_full_speed_rpm, thread_pitch)
    # distance_to_turnaround = 0.5 * time_to_turnaround * run_freq * (motor_full_speed_rpm / motor_full_speed_hz) * (1.0 / 60.0) * thread_pitch
    print("Turn Distance: {} in.".format(distance_to_turnaround))

    upper_turnaround_target = stroke_length - (upper_offset + distance_to_turnaround)
    upper_offset_target = stroke_length - upper_offset
    lower_turnaround_target = lower_offset + distance_to_turnaround

    print("--")
    print("Targets")
    print("-")
    print("Lower Offset: {} in.".format(lower_offset))
    print("Turnaround Target: {} in.".format(upper_turnaround_target))
    print("Upper Offset: {} in.".format(upper_offset_target))
    print("Lower Turnaround Target: {} in.".format(lower_turnaround_target))
    print("--")

    position = lower_offset
    time = 0.0

    direction = 1
    starting_speed = 0.0  # Hz
    speed = starting_speed
    last_speed = speed

    speed_array = []
    position_array = []
    time_array = []
    stroke_part = -1

    max_position = 0
    min_position = stroke_length

    while strokes != num_strokes:
        last_speed = speed
        if direction == 1:
            if position < upper_turnaround_target:
                # below offset and turnaround, ramp up to setpoint, then run at setpoint
                if stroke_part != 0:
                    stroke_part = 0
                    print("{} - {}".format(stroke_part, position))
                if speed < run_freq:
                    speed += (motor_full_speed_hz / accel_time) * dt
                if speed > run_freq:
                    speed = run_freq
            elif position >= upper_turnaround_target:
                # above turnaround distance, ramp to 0 and change direction
                if stroke_part != 1:
                    stroke_part = 1
                    print("{} - {}".format(stroke_part, position))
                if speed > 0:
                    speed += -1.0 * (motor_full_speed_hz / decel_time) * dt
                if speed <= 0:
                    speed = 0
                    direction = -1

        elif direction == -1:
            if position > lower_turnaround_target:
                # above offset and rampdown distance, ramp up to setpoint (negative), then run at setpoint
                if stroke_part != 2:
                    stroke_part = 2
                    print("{} - {}".format(stroke_part, position))
                if speed > (-1.0 * run_freq):
                    speed += -1.0 * (motor_full_speed_hz / accel_time) * dt
                if speed < (-1.0 * run_freq):
                    speed = -1.0 * run_freq
            elif position <= lower_turnaround_target:
                # below turnaround distance, ramp to 0 and change direction
                if stroke_part != 3:
                    stroke_part = 3
                    print("{} - {}".format(stroke_part, position))
                if speed < 0:
                    speed += 1.0 * (motor_full_speed_hz / decel_time) * dt
                if speed >= 0:
                    speed = 0
                    direction = 1
                    strokes += 1

        if use_average:
            delta_x = (speed + last_speed) / 2.0 * (motor_full_speed_rpm / motor_full_speed_hz) * (dt / 60.0) * thread_pitch
        else:
            delta_x = speed * (motor_full_speed_rpm / motor_full_speed_hz) * (dt / 60.0) * thread_pitch

        position += delta_x
        time += dt
        speed_array.append(speed * (motor_full_speed_rpm / motor_full_speed_hz))
        position_array.append(position)
        time_array.append(time)
        # print("Time: {} sec., Position: {} in., Speed: {}".format(time, position, speed * (motor_full_speed_rpm / motor_full_speed_hz)))
        # sleep(dt)

        if stroke_part == 1:
            max_position = max([max_position, position])

        if stroke_part == 3:
            min_position = min([min_position, position])

    SPM = 60.0 / (time / strokes)

    # max_position = max(position_array)
    # min_position = min(position_array)
    print("Max: {}, Min: {}".format(max_position, min_position))
    print("SPM: {}".format(SPM))

    fig, ax1 = plt.subplots()
    ax1.plot(time_array, speed_array, "b-", linewidth=2)
    # ax1.plot(position_array, speed_array, "b-", linewidth=2)
    ax1.set_ylabel("RPM", color="b")
    ax2 = ax1.twinx()
    ax2.plot(time_array, position_array, "r-", linewidth=2)
    ax2.set_ylabel("in.", color="r")
    fig.tight_layout()
    ax1.grid(color='black', linestyle='--', linewidth=1)
    plt.show()


def calc_travel_time(distance, speed_rpm, thread_pitch):
    """Calculate the time it takes to travel a distatnce at a speed."""
    rotations_needed = distance / thread_pitch
    minutes_for_full_travel = rotations_needed / speed_rpm
    seconds_for_full_travel = minutes_for_full_travel / 60.0
    return seconds_for_full_travel


def calc_theoretical_production(stroke_length, spm, pump_diameter):
    """Calculate the maximum production for a given stroke length and speed."""
    try:
        return stroke_length * pump_constants[pump_diameter] * spm
    except IndexError:
        print("Cannot find pump constant for that diameter.")
        return False


if __name__ == '__main__':
    sim(g_run_speed, g_dt, 5, g_stroke_length, g_accel_time, g_decel_time, g_upper_offset,
        g_lower_offset, g_motor_full_speed_hz, g_motor_full_speed_rpm, g_thread_pitch, use_average=False)
