import math
import numpy as np

def total_horizontal_displacement(alpha0, speeds, widths, boundary_normal_angle):
    """
    Given an initial angle alpha0 (measured relative to the boundary normal),
    propagate angles through each layer via Snell's Law and compute
    the total horizontal displacement from start to finish.

    Parameters:
    -----------
    alpha0 : float
        Initial angle (radians) in the first layer relative to the boundary normal.
    speeds : list of float
        speeds[i] = travel speed in layer i
    widths : list of float
        widths[i] = thickness of layer i in the boundary-normal direction
    boundary_normal_angle : float
        Angle (radians) of the boundary normal relative to the horizontal (east direction).
        For example, if the boundary lines themselves are at 45° to the horizontal,
        the boundary normal might also be ±45°, depending on orientation.

    Returns:
    --------
    float
        The total horizontal (eastward) distance traveled.
        If Snell's Law fails (sin(...) out of range), returns a large number.
    """
    n = len(speeds) - 1  # Number of boundaries
    alpha = [0.0]*(n+1)
    alpha[0] = alpha0

    # Propagate angles using Snell's Law: v_i * sin(theta_i) = v_(i+1) * sin(theta_(i+1))
    for i in range(n):
        ratio = (speeds[i] / speeds[i+1]) * math.sin(alpha[i])
        if abs(ratio) > 1:
            # Invalid angle => no real solution
            return 1e9
        alpha[i+1] = math.asin(ratio)

    # Compute total horizontal displacement
    horiz = 0.0
    for i in range(n+1):
        # Path length in layer i = widths[i] / cos(alpha[i])
        # because alpha[i] is the angle to the normal, so the
        # "normal" component of the path is widths[i].
        seg_length = widths[i] / math.cos(alpha[i])

        # The angle of travel relative to the horizontal is
        # (boundary_normal_angle + alpha[i]).
        seg_angle = boundary_normal_angle + alpha[i]

        # Horizontal component
        horiz += seg_length * math.cos(seg_angle)

    return horiz

def total_time(alpha0, speeds, widths, boundary_normal_angle):
    """
    Given the same setup as total_horizontal_displacement, compute
    the total travel time if the initial angle is alpha0.
    """
    n = len(speeds) - 1
    alpha = [0.0]*(n+1)
    alpha[0] = alpha0

    # Again apply Snell's Law across boundaries
    for i in range(n):
        ratio = (speeds[i] / speeds[i+1]) * math.sin(alpha[0 + i])
        if abs(ratio) > 1:
            return 1e9
        alpha[i+1] = math.asin(ratio)

    # Sum travel times in each layer
    total_t = 0.0
    for i in range(n+1):
        seg_length = widths[i] / math.cos(alpha[i])
        total_t += seg_length / speeds[i]

    return total_t

def solve_snell_for_100_leagues(speeds, widths, boundary_normal_angle, target=100.0):
    """
    Uses a simple binary search on alpha0 to achieve a desired
    total horizontal displacement (default 100 leagues).
    Returns (alpha0_opt, min_time).
    """
    def f(alpha0):
        # difference from desired horizontal displacement
        return total_horizontal_displacement(alpha0, speeds, widths, boundary_normal_angle) - target

    # Search range for alpha0: from 0 to something less than pi/2
    # so that cos(alpha0) is positive.
    low, high = 0.0, math.radians(80.0)  # somewhat arbitrary upper bound
    for _ in range(100):
        mid = 0.5*(low + high)
        val = f(mid)
        if val > 0:
            # If we overshoot the horizontal distance, reduce angle
            high = mid
        else:
            # If we undershoot, increase angle
            low = mid

    alpha_opt = 0.5*(low + high)
    # Compute the time corresponding to that alpha
    t_opt = total_time(alpha_opt, speeds, widths, boundary_normal_angle)
    return alpha_opt, t_opt

def main():
    # Example speeds in each region:
    #   - Region 0: normal terrain (10 leagues/day)
    #   - Regions 1..5: inside the marsh, each with distinct speeds
    #   - Region 6: normal terrain again (10 leagues/day)
    speeds = [10, 9, 8, 7, 6, 5, 10]

    # Example widths (in leagues) in the boundary-normal direction:
    # Suppose each region is 10 leagues thick. Adjust as needed.
    widths = [10, 10, 10, 10, 10, 10, 10]

    # Suppose the marsh boundaries run at 45° to the horizontal,
    # so the boundary normal is also at ±45°. Let's pick +45°:
    boundary_normal_angle = math.radians(45.0)

    # Solve for the path that yields 100 leagues total horizontal displacement
    alpha_solution, time_solution = solve_snell_for_100_leagues(
        speeds, widths, boundary_normal_angle, target=100.0
    )

    print(f"Optimal initial angle (relative to boundary normal) = {math.degrees(alpha_solution):.3f} degrees")
    print(f"Total travel time = {time_solution:.5f} days")

if __name__ == "__main__":
    main()