from machine import Pin, SPI, ADC, PWM
import time


# ==========================================
# LCD CONNECTIONS
# ==========================================

# SPI communication
spi = SPI(
    0,
    baudrate=40000000,
    polarity=0,
    phase=0,
    sck=Pin(18),
    mosi=Pin(19)
)

# LCD control pins
cs = Pin(17, Pin.OUT)
rst = Pin(16, Pin.OUT)
dc = Pin(20, Pin.OUT)

# ==========================================
# TURBINE CONTROL
# ==========================================

# Turbine slider SIG is connected to GP26.
turbine_slider = ADC(Pin(26))

# ==========================================
# FISSION CONTROL
# ==========================================

# Fission slider SIG is connected to GP27.
fission_slider = ADC(Pin(27))

# ==========================================
# TURBINE GAUGE SERVO
# ==========================================

# Servo PWM is connected to GP14.
turbine_servo = PWM(Pin(14))
turbine_servo.freq(50)

# ==========================================
# FISSION GAUGE SERVO
# ==========================================

# Fission servo PWM is connected to GP13.
fission_servo = PWM(Pin(13))
fission_servo.freq(50)

def move_turbine_servo(angle):

    # Keep the angle between 0 and 180 degrees.
    angle = max(0, min(180, angle))

    # Our previously calibrated servo range.
    min_duty = 1638
    max_duty = 7864

    # Convert 0-180 degrees into PWM.
    duty = int(
        min_duty
        + (angle / 180)
        * (max_duty - min_duty)
    )

    turbine_servo.duty_u16(duty)

def move_fission_servo(angle):

    # Keep the angle between 0 and 180 degrees.
    angle = max(0, min(180, angle))

    # Use the same calibration as our
    # turbine servo for this Wokwi test.
    min_duty = 1638
    max_duty = 7864

    # Convert 0-180 degrees into PWM.
    duty = int(
        min_duty
        + (angle / 180)
        * (max_duty - min_duty)
    )

    fission_servo.duty_u16(duty)

# ==========================================
# LCD COMMAND FUNCTIONS
# ==========================================

def write_command(command):
    dc.value(0)
    cs.value(0)
    spi.write(bytes([command]))
    cs.value(1)


def write_data(data):
    dc.value(1)
    cs.value(0)
    spi.write(bytes(data))
    cs.value(1)


# ==========================================
# RESET LCD
# ==========================================

rst.value(0)
time.sleep(0.1)

rst.value(1)
time.sleep(0.1)


# ==========================================
# INITIALIZE ILI9341
# ==========================================

write_command(0x01)  # Software reset
time.sleep(0.15)

write_command(0x11)  # Sleep out
time.sleep(0.15)

# 16-bit RGB565 color mode
write_command(0x3A)
write_data([0x55])

# Set display orientation/color order.
write_command(0x36)
write_data([0x28])

# Turn display on
write_command(0x29)

print("LCD initialized!")

# ==========================================
# DRAW A SOLID COLOR
# ==========================================

def set_window(x0, y0, x1, y1):

    # Tell the LCD which columns we want to draw in.
    write_command(0x2A)
    write_data([
        x0 >> 8,
        x0 & 0xFF,
        x1 >> 8,
        x1 & 0xFF
    ])

    # Tell the LCD which rows we want to draw in.
    write_command(0x2B)
    write_data([
        y0 >> 8,
        y0 & 0xFF,
        y1 >> 8,
        y1 & 0xFF
    ])

    # Tell the LCD we're about to send pixel colors.
    write_command(0x2C)


def fill_screen(color):

    # ILI9341 resolution is 320 wide × 240 tall.
    set_window(0, 0, 319, 239)

    # RGB565 uses two bytes per pixel.
    high_byte = color >> 8
    low_byte = color & 0xFF

    pixel = bytes([high_byte, low_byte])

    dc.value(1)
    cs.value(0)

    # Send one row at a time.
    row = pixel * 320

    for _ in range(240):
        spi.write(row)

    cs.value(1)


# RGB565 red
RED = 0xF800

def fill_rect(x, y, width, height, color):

    # Select the rectangle we want to draw.
    set_window(
        x,
        y,
        x + width - 1,
        y + height - 1
    )

    # Convert our RGB565 color into two bytes.
    high_byte = color >> 8
    low_byte = color & 0xFF

    pixel = bytes([high_byte, low_byte])

    dc.value(1)
    cs.value(0)

    # Create one horizontal row of pixels.
    row = pixel * width

    # Draw that row repeatedly to make the rectangle.
    for _ in range(height):
        spi.write(row)

    cs.value(1)

# ==========================================
# TEST COLORS
# ==========================================

GREEN = 0x07E0
ORANGE = 0xFD20
BLUE = 0x001F

# ==========================================
# FIRST TURBINE GAUGE TEST
# ==========================================

BLACK = 0x0000

# Pretend the submarine currently needs
# about 60% turbine output.
target_turbine = 60

# Width of our gauge.
gauge_x = 20
gauge_y = 145
gauge_width = 280
gauge_height = 45

# Fission gauge position and size.
fission_gauge_x = 20
fission_gauge_y = 50
fission_gauge_width = 280
fission_gauge_height = 45

# Clear the screen.
fill_screen(BLACK)

# Draw the entire gauge orange first.
fill_rect(
    gauge_x,
    gauge_y,
    gauge_width,
    gauge_height,
    ORANGE
)

available_fuel = 80.0
degree_of_success = 0.5

optimal_turbine_tolerance = (
    2.5
    + (10.0 - 2.5) * degree_of_success
)

optimal_temperature_min = (
    40
    + (30 - 40) * degree_of_success
)

optimal_temperature_max = (
    60
    + (70 - 60) * degree_of_success
)

allowed_temperature_min = (
    30
    + (10 - 30) * degree_of_success
)

allowed_temperature_max = (
    70
    + (90 - 70) * degree_of_success
)

optimal_fission_min = (
    30
    + (20 - 30) * degree_of_success
)

optimal_fission_max = (
    (available_fuel - 20)
    + (
        (available_fuel - 10)
        - (available_fuel - 20)
    ) * degree_of_success
)

optimal_fission_min = min(
    optimal_fission_min,
    optimal_fission_max - 10
)

# Draw the fission gauge background.
fill_rect(
    fission_gauge_x,
    fission_gauge_y,
    fission_gauge_width,
    fission_gauge_height,
    ORANGE
)

# Use the calculated Barotrauma optimal fission range.

# Convert the minimum percentage into an LCD X position.
fission_green_x = int(
    fission_gauge_x
    + (optimal_fission_min / 100)
    * fission_gauge_width
)

# Calculate how wide the optimal range is.
fission_green_width_percent = (
    optimal_fission_max
    - optimal_fission_min
)

# Convert that percentage width into pixels.
fission_green_width = int(
    (fission_green_width_percent / 100)
    * fission_gauge_width
)

# Draw the optimal fission region.
fill_rect(
    fission_green_x,
    fission_gauge_y,
    fission_green_width,
    fission_gauge_height,
    GREEN
)

green_width_percent = (
    optimal_turbine_tolerance * 2
)

# Calculate where the green region starts.
green_start_percent = (
    target_turbine - green_width_percent / 2
)

# Convert percentages into screen pixels.
green_x = int(
    gauge_x
    + (green_start_percent / 100) * gauge_width
)

green_width = int(
    (green_width_percent / 100) * gauge_width
)

# Draw the recommended region.
fill_rect(
    green_x,
    gauge_y,
    green_width,
    gauge_height,
    GREEN
)

print("Target turbine:", target_turbine, "%")

# ==========================================
# FAKE SUBMARINE POWER SYSTEM
# ==========================================

# Maximum electrical output of our
# pretend reactor.
reactor_max_output = 5000

# Different loads our pretend submarine
# will experience.
ship_loads = [
    1000,
    2000,
    3000,
    4000,
    2500,
    1500
]

correct_turbine_output = 0.0

# Start with the first submarine load.
load_index = 0
ship_load = ship_loads[load_index]

# Remember when we last changed the load.
last_load_change = time.ticks_ms()

# ==========================================
# REACTOR TEMPERATURE
# ==========================================

# Starting temperature for our reactor simulation.
reactor_temperature = 50.0

while True:

    # ======================================
    # READ TURBINE CONTROL
    # ======================================

    raw_turbine = turbine_slider.read_u16()

    actual_turbine = (
        raw_turbine / 65535
    ) * 100

    # ======================================
    # READ FISSION CONTROL
    # ======================================

    # Read the physical fission slider.
    raw_fission = fission_slider.read_u16()

    # Convert the ADC reading into 0-100%.
    actual_fission = (
        raw_fission / 65535
    ) * 100

    # Convert 0-100% fission rate
    # into 0-180 degrees for the servo.
    fission_angle = (
        actual_fission / 100
    ) * 180

    move_fission_servo(fission_angle)


    # ======================================
    # REACTOR HEAT TEST
    # ======================================

    # Calculate generated heat using Barotrauma's
    # reactor heat formula.
    generated_heat = (
        actual_fission
        * (available_fuel / 100.0)
        * 2.0
    )

    # Find the temperature the reactor is
    # currently trying to reach.
    temperature_target = (
        generated_heat
        - actual_turbine
    )

    # Find the difference between our current
    # temperature and that target.
    temperature_difference = (
        temperature_target
        - reactor_temperature
    )

    # Barotrauma lets temperature move at
    # a maximum of 10 percentage points per second.
    max_temperature_change = 10.0 * 0.02

    # Move toward the target without overshooting it.
    if temperature_difference > max_temperature_change:
        reactor_temperature += max_temperature_change
    elif temperature_difference < -max_temperature_change:
        reactor_temperature -= max_temperature_change
    else:
        reactor_temperature = temperature_target

    # Keep simulated temperature within
    # our temporary 0-100 range.
    reactor_temperature = max(
        0,
        min(100, reactor_temperature)
    )

    if reactor_temperature > allowed_temperature_max:
        temperature_status = "CRITICAL"
    elif reactor_temperature > optimal_temperature_max:
        temperature_status = "HIGH"
    elif reactor_temperature < optimal_temperature_min:
        temperature_status = "LOW"
    else:
        temperature_status = "OPTIMAL"

    # Convert 0-100% turbine output
    # into 0-180 degrees for the servo.
    turbine_angle = (
        actual_turbine / 100
    ) * 180

    move_turbine_servo(turbine_angle)

    desired_turbine_output = (
    ship_load / reactor_max_output
    ) * 100

    desired_turbine_output = max(
        0,
        min(100, desired_turbine_output)
    )

    difference = (
        desired_turbine_output
        - correct_turbine_output
    )

    max_change = 20 * 0.02

    if difference > max_change:
        correct_turbine_output += max_change
    elif difference < -max_change:
        correct_turbine_output -= max_change
    else:
        correct_turbine_output = desired_turbine_output

    # ======================================
    # CHECK SUBMARINE LOAD TIMER
    # ======================================

    current_time = time.ticks_ms()

    if time.ticks_diff(
        current_time,
        last_load_change
    ) >= 2000:

        # Move to the next fake ship load.
        load_index += 1

        # If we reach the end of the list,
        # go back to the beginning.
        if load_index >= len(ship_loads):
            load_index = 0

        ship_load = ship_loads[load_index]

        # Remember when this change happened.
        last_load_change = current_time


        # ==================================
        # CALCULATE REQUIRED TURBINE
        # ==================================

        required_turbine = (
            ship_load / reactor_max_output
        ) * 100

        required_turbine = max(
            0,
            min(100, required_turbine)
        )


        # ==================================
        # UPDATE LCD TARGET
        # ==================================

        # Reset gauge to orange.
        fill_rect(
            gauge_x,
            gauge_y,
            gauge_width,
            gauge_height,
            ORANGE
        )

        green_start_percent = (
            correct_turbine_output
            - green_width_percent / 2
        )

        green_start_percent = max(
            0,
            min(
                100 - green_width_percent,
                green_start_percent
            )
        )

        green_x = int(
            gauge_x
            + (green_start_percent / 100)
            * gauge_width
        )

        fill_rect(
            green_x,
            gauge_y,
            green_width,
            gauge_height,
            GREEN
        )


        # ==================================
        # TERMINAL INFORMATION
        # ==================================

        print(
            "Optimal Fission Range:",
            round(optimal_fission_min, 1),
            "%",
            "to",
            round(optimal_fission_max, 1),
            "%"
        )

        print(
            "Optimal Temperature Range:",
            round(optimal_temperature_min, 1),
            "%",
            "to",
            round(optimal_temperature_max, 1),
            "%"
        )

        print(
            "Allowed Temperature Range:",
            round(allowed_temperature_min, 1),
            "%",
            "to",
            round(allowed_temperature_max, 1),
            "%"
        )
        
        print(
            "Ship Load:",
            ship_load,
            "kW",
            "| Required Turbine:",
            round(required_turbine, 1),
            "%",
            "| Correct Turbine:",
            round(correct_turbine_output, 1),
            "%",
            "| Your Turbine:",
            round(actual_turbine, 1),
            "%",
            "| Your Fission:",
            round(actual_fission, 1),
            "%",
            "| Temperature:",
            round(reactor_temperature, 1),
            "| Temp Status:",
            temperature_status
        )


    # Tiny delay so we don't hammer the simulator.
    time.sleep(0.02)