import numpy as np
import random

def set_terrain(terrain, variation, difficulty):
    terrain_fns = [
        set_terrain_0,
        set_terrain_1,
        set_terrain_2,
        set_terrain_3,
        set_terrain_4,
        set_terrain_5,
        set_terrain_6,
        set_terrain_7,
        set_terrain_8,
        set_terrain_9,
        # INSERT TERRAIN FUNCTIONS HERE
    ]
    idx = int(variation * len(terrain_fns))
    height_field, goals = terrain_fns[idx](terrain.width * terrain.horizontal_scale, terrain.length * terrain.horizontal_scale, terrain.horizontal_scale, difficulty)
    terrain.height_field_raw = (height_field / terrain.vertical_scale).astype(np.int16)
    terrain.goals = goals
    return idx

def set_terrain_0(length, width, field_resolution, difficulty):
    """Mixed-height staggered platforms with interspersed lateral ramps for strategic agility and balance."""

    def m_to_idx(m):
        """Converts meters to quantized indices."""
        return np.round(m / field_resolution).astype(np.int16) if not (isinstance(m, list) or isinstance(m, tuple)) else [round(i / field_resolution) for i in m]

    height_field = np.zeros((m_to_idx(length), m_to_idx(width)))
    goals = np.zeros((8, 2))

    # Set up platform and ramp dimensions
    platform_length = 0.8 - 0.2 * difficulty  # Shorter platforms
    platform_length = m_to_idx(platform_length)
    platform_width = 1.0  # Constant width for simplicity
    platform_width = m_to_idx(platform_width)
    platform_height_min, platform_height_max = 0.0 + 0.1 * difficulty, 0.1 + 0.2 * difficulty
    lateral_shift = (0.2 + 0.3 * difficulty)  # Lateral shift of platforms increases with difficulty
    gap_length = 0.3 + 0.1 * difficulty
    gap_length = m_to_idx(gap_length)

    mid_y = m_to_idx(width) // 2

    def add_platform(start_x, end_x, shift_y):
        x1, x2 = start_x, end_x
        y1, y2 = mid_y - shift_y, mid_y + shift_y
        platform_height = np.random.uniform(platform_height_min, platform_height_max)
        height_field[x1:x2, y1:y2] = platform_height

    def add_lateral_ramp(start_x, end_x, shift_y, slope_up=True):
        half_width = platform_width // 2
        x1, x2 = start_x, end_x
        y1, y2 = mid_y - half_width, mid_y + half_width
        ramp_height = np.linspace(0, (platform_height_max-platform_height_min) if slope_up else (platform_height_min-platform_height_max), end_x-start_x)
        height_field[x1:x2, y1:y2] = ramp_height[:, None]

    # Set spawn area to flat ground
    spawn_length = m_to_idx(2)
    height_field[0:spawn_length, :] = 0
    goals[0] = [spawn_length - m_to_idx(0.5), mid_y]

    cur_x = spawn_length
    current_shift = 0

    for i in range(6):
        current_shift = m_to_idx(lateral_shift * ((-1) ** i))  # Alternating lateral shifts
        add_platform(cur_x, cur_x + platform_length, current_shift)
        goals[i + 1] = [cur_x + platform_length / 2, mid_y + current_shift / 2]
        
        cur_x += platform_length + gap_length

        if i < 5:  # Adding lateral ramps between platforms, except at the end
            add_lateral_ramp(cur_x, cur_x + gap_length, current_shift, slope_up=(i % 2 == 0))
        
    goals[-1] = [cur_x, mid_y]
    height_field[cur_x:, :] = 0

    return height_field, goals

def set_terrain_1(length, width, field_resolution, difficulty):
    """Zigzagging beams over varying heights for precision and agility improvement."""

    def m_to_idx(m):
        """Converts meters to quantized indices."""
        return np.round(m / field_resolution).astype(np.int16) if not isinstance(m, (list, tuple)) else [round(i / field_resolution) for i in m]

    height_field = np.zeros((m_to_idx(length), m_to_idx(width)))
    goals = np.zeros((8, 2))

    # Set up beam and gap dimensions
    beam_length = 1.2 - 0.2 * difficulty
    beam_width = 0.35  # Narrow for precision
    beam_height_min, beam_height_max = 0.0 + 0.4 * difficulty, 0.1 + 0.5 * difficulty
    gap_length = 0.3 + 0.5 * difficulty

    beam_length, beam_width, gap_length = map(m_to_idx, [beam_length, beam_width, gap_length])

    mid_y = m_to_idx(width) // 2

    def add_beam(start_x, end_x, mid_y, height):
        half_width = beam_width // 2
        x1, x2 = start_x, end_x
        y1, y2 = mid_y - half_width, mid_y + half_width
        height_field[x1:x2, y1:y2] = height

    # Set spawn area to flat ground
    spawn_length = m_to_idx(2)
    height_field[0:spawn_length, :] = 0
    goals[0] = [spawn_length - m_to_idx(0.5), mid_y]

    # Start setting zigzag beams
    cur_x = spawn_length
    for i in range(6):
        beam_height = random.uniform(beam_height_min, beam_height_max)
        lateral_shift = ((-1) ** i) * m_to_idx(0.5)
        
        add_beam(cur_x, cur_x + beam_length, mid_y + lateral_shift, beam_height)
        
        goals[i + 1] = [cur_x + beam_length / 2, mid_y + lateral_shift]

        # Add gap
        cur_x += beam_length + gap_length

    # Add final goal and fill the rest with flat ground
    goals[-1] = [cur_x + m_to_idx(0.5), mid_y]
    height_field[cur_x:, :] = 0

    return height_field, goals

def set_terrain_2(length, width, field_resolution, difficulty):
    """A course of staggered steps of varied widths and moderate heights for testing adaptability and coordination."""

    def m_to_idx(m):
        """Converts meters to quantized indices."""
        return np.round(m / field_resolution).astype(np.int16) if not (isinstance(m, list) or isinstance(m, tuple)) else [round(i / field_resolution) for i in m]

    height_field = np.zeros((m_to_idx(length), m_to_idx(width)))
    goals = np.zeros((8, 2))

    # Define staggered step characteristics
    step_length_min = 0.8
    step_length_max = 1.2
    step_width_min = 0.6
    step_width_max = 1.0
    step_heights = [0.05, 0.10, 0.15, 0.20]  # Varying step heights
    gap_length = 0.4
    gap_length = m_to_idx(gap_length)

    mid_y = m_to_idx(width) // 2

    def add_step(start_x, length, width, height, mid_y):
        half_width = m_to_idx(width) // 2
        x1, x2 = start_x, start_x + m_to_idx(length)
        y1, y2 = mid_y - half_width, mid_y + half_width
        height_field[x1:x2, y1:y2] = height

    # Set spawn area to flat ground
    spawn_length = m_to_idx(2)
    height_field[0:spawn_length, :] = 0
    # Put first goal at spawn
    goals[0] = [spawn_length - m_to_idx(0.5), mid_y]  
    
    # First step to provide easy beginning
    cur_x = spawn_length
    step_length = np.random.uniform(step_length_min, step_length_max)
    step_width = np.random.uniform(step_width_min, step_width_max)
    step_height = np.random.choice(step_heights)

    add_step(cur_x, step_length, step_width, step_height, mid_y)
    goals[1] = [cur_x + m_to_idx(step_length) / 2, mid_y]
    cur_x += m_to_idx(step_length) + gap_length

    # Add remaining steps
    for i in range(1, 6):
        step_length = np.random.uniform(step_length_min, step_length_max)
        step_width = np.random.uniform(step_width_min, step_width_max)
        step_height = np.random.choice(step_heights)

        add_step(cur_x, step_length, step_width, step_height, mid_y)
        goals[i+1] = [cur_x + m_to_idx(step_length) / 2, mid_y]

        # Add gap
        cur_x += m_to_idx(step_length) + gap_length
    
    # Add final goal behind the last step, fill in the remaining gap
    goals[-1] = [cur_x + m_to_idx(0.5), mid_y]
    height_field[cur_x:, :] = 0

    return height_field, goals

def set_terrain_3(length, width, field_resolution, difficulty):
    """Zigzagging platforms over pits with varying heights to test balance and precision."""

    def m_to_idx(m):
        """Converts meters to quantized indices."""
        return np.round(m / field_resolution).astype(np.int16) if not (isinstance(m, list) or isinstance(m, tuple)) else [round(i / field_resolution) for i in m]

    height_field = np.zeros((m_to_idx(length), m_to_idx(width)))
    goals = np.zeros((8, 2))

    # Set up platform dimensions with increased variability
    platform_length = 0.9 - 0.2 * difficulty
    platform_length = m_to_idx(platform_length)
    platform_width = np.random.uniform(1.0, 1.4)
    platform_width = m_to_idx(platform_width)
    platform_height_min, platform_height_max = 0.05 + 0.25 * difficulty, 0.1 + 0.35 * difficulty
    gap_length = 0.15 + 0.6 * difficulty
    gap_length = m_to_idx(gap_length)

    mid_y = m_to_idx(width) // 2

    def add_platform(start_x, end_x, mid_y, height):
        half_width = platform_width // 2
        height_field[start_x:end_x, max(mid_y-half_width, 0):min(mid_y+half_width, m_to_idx(width))] = height

    dx_min, dx_max = -0.1, 0.1
    dx_min, dx_max = m_to_idx(dx_min), m_to_idx(dx_max)
    dy_min, dy_max = -0.3, 0.3
    dy_min, dy_max = m_to_idx(dy_min), m_to_idx(dy_max)

    # Set spawn area to flat ground
    spawn_length = m_to_idx(2)
    height_field[0:spawn_length, :] = 0
    # First goal positioned at spawn
    goals[0] = [spawn_length - m_to_idx(0.5), mid_y]

    # Set remaining area to be a pit
    height_field[spawn_length:, :] = -1.0

    cur_x = spawn_length
    for i in range(1, 7):  # Setting up 6 platforms
        dx = np.random.randint(dx_min, dx_max)
        dy = np.random.randint(dy_min, dy_max)
        height = np.random.uniform(platform_height_min, platform_height_max)

        add_platform(cur_x, cur_x + platform_length + dx, mid_y + dy, height)

        # Place a goal at each platform center
        goals[i] = [cur_x + (platform_length + dx) / 2, mid_y + dy]

        # Add gap to the next platform
        cur_x += platform_length + dx + gap_length
    
    # Final goal at the end of the last platform, filling the remaining space
    goals[-1] = [cur_x + m_to_idx(0.5), mid_y]
    height_field[cur_x:, :] = 0

    return height_field, goals

def set_terrain_4(length, width, field_resolution, difficulty):
    """Convex and concave steps with varying widths for agility and balance navigation."""

    def m_to_idx(m):
        """Converts meters to quantized indices."""
        return np.round(m / field_resolution).astype(np.int16) if not (isinstance(m, list) or isinstance(m, tuple)) else [round(i / field_resolution) for i in m]

    height_field = np.zeros((m_to_idx(length), m_to_idx(width)))
    goals = np.zeros((8, 2))

    # Set up platform and step dimensions
    step_length = 0.8 - 0.2 * difficulty 
    step_length = m_to_idx(step_length)
    step_width_min = 0.4 
    step_width_max = 1.0 - 0.2 * difficulty 
    step_width_min = m_to_idx(step_width_min)
    step_width_max = m_to_idx(step_width_max)
    step_height_min, step_height_max = 0.05 + 0.1 * difficulty, 0.15 + 0.2 * difficulty

    mid_y = m_to_idx(width) // 2

    def add_step(start_x, end_x, mid_y, concave=False):
        half_width = np.random.uniform(step_width_min, step_width_max) // 2
        x1, x2 = start_x, end_x
        y1, y2 = int(mid_y - half_width), int(mid_y + half_width)
        step_height = np.random.uniform(step_height_min, step_height_max)
        if concave:
            slants = np.linspace(step_height, 0, num=x2-x1)
        else:
            slants = np.linspace(0, step_height, num=x2-x1)
        height_field[x1:x2, y1:y2] = slants[:, None]

    dx_min, dx_max = -0.05, 0.05
    dx_min, dx_max = m_to_idx(dx_min), m_to_idx(dx_max)
    dy_min, dy_max = -0.3, 0.3
    dy_min, dy_max = m_to_idx(dy_min), m_to_idx(dy_max)

    # Set spawn area to flat ground
    spawn_length = m_to_idx(2)
    height_field[0:spawn_length, :] = 0
    # Place the first goal at spawn
    goals[0] = [spawn_length - m_to_idx(0.5), mid_y]  

    # Define the pit level for depth below steps
    height_field[spawn_length:, :] = -1.0

    cur_x = spawn_length
    for i in range(1, 8):
        dx = np.random.randint(dx_min, dx_max)
        dy = np.random.randint(dy_min, dy_max)
        concave_if_even = (i % 2 == 0)  # Alternate between concave and convex steps
        
        add_step(cur_x, cur_x + step_length + dx, mid_y + dy, concave=concave_if_even)

        # Place goal at the center of each step
        goals[i] = [cur_x + (step_length + dx) / 2, mid_y + dy]

        # Move to next step
        cur_x += step_length + dx

    # End the course with ground level
    height_field[cur_x:, :] = 0

    return height_field, goals

def set_terrain_5(length, width, field_resolution, difficulty):
    """Serpentine elevated paths with alternating narrow, tilted, and ascending platforms to challenge navigation and balance."""

    def m_to_idx(m):
        """Converts meters to quantized indices."""
        return np.round(m / field_resolution).astype(np.int16) if not (isinstance(m, list) or isinstance(m, tuple)) else [round(i / field_resolution) for i in m]

    height_field = np.zeros((m_to_idx(length), m_to_idx(width)))
    goals = np.zeros((8, 2))

    # Define dimensions for platforms and paths
    platform_length = 0.8 - 0.2 * difficulty
    platform_length = m_to_idx(platform_length)
    path_width = np.random.uniform(1.0, 1.2) - 0.2 * difficulty
    path_width = m_to_idx(path_width)
    incline_height = 0.5 * difficulty
    incline_height_min, incline_height_max = incline_height, 0.05 + 0.5 * difficulty
    gap_length = 0.5 + 0.5 * difficulty
    gap_length = m_to_idx(gap_length)

    mid_y = m_to_idx(width) // 2

    def add_path(start_x, end_x, mid_y, height):
        half_width = path_width // 2
        x1, x2 = start_x, end_x
        y1, y2 = mid_y - half_width, mid_y + half_width
        height_field[x1:x2, y1:y2] = height

    def add_tilted_platform(start_x, end_x, mid_y, tilt_direction):
        half_width = path_width // 2
        x1, x2 = start_x, end_x
        y1, y2 = mid_y - half_width, mid_y + half_width
        incline = np.linspace(tilt_direction * incline_height, -tilt_direction * incline_height, num=y2-y1)
        incline = incline[None, :]  # Add a dimension for broadcasting to the x-range
        height_field[x1:x2, y1:y2] = incline

    # Initial setup
    spawn_length = m_to_idx(2)
    height_field[0:spawn_length, :] = 0
    goals[0] = [spawn_length - m_to_idx(0.5), mid_y] 

    cur_x = spawn_length
    for i in range(7):
        height = np.random.uniform(0.0, incline_height)

        if i % 2 == 0:
            # Regular path
            add_path(cur_x, cur_x + platform_length, mid_y, height)
        else:
            # Tilted platform
            tilt_direction = (-1)**i
            add_tilted_platform(cur_x, cur_x + platform_length, mid_y, tilt_direction)

        goals[i+1] = [cur_x + platform_length / 2, mid_y]

        cur_x += platform_length + gap_length

    # Final stretch
    goals[7] = [cur_x + m_to_idx(0.5), mid_y]
    height_field[cur_x:, :] = 0

    return height_field, goals

def set_terrain_6(length, width, field_resolution, difficulty):
    """Diagonal ascents, variable-width ledges, and jumping gaps for agility and balance challenges."""

    def m_to_idx(m):
        """Converts meters to quantized indices."""
        return np.round(m / field_resolution).astype(np.int16) if not (isinstance(m, list) or isinstance(m, tuple)) else [round(i / field_resolution) for i in m]

    height_field = np.zeros((m_to_idx(length), m_to_idx(width)))
    goals = np.zeros((8, 2))

    # Platform and obstacle dimensions
    platform_length = m_to_idx(0.8 - 0.2 * difficulty)
    platform_width_min, platform_width_max = m_to_idx(0.4), m_to_idx(1.0)
    platform_height_min, platform_height_max = 0.15 * difficulty, 0.45 * difficulty
    gap_length = m_to_idx(0.1 + 0.4 * difficulty)
    incline_length = m_to_idx(1.0)
    incline_height = np.linspace(0, 0.4 * difficulty, incline_length)

    mid_y = m_to_idx(width) // 2
    spawn_length = m_to_idx(2)
    
    # Ensure robot spawns in a flat area
    height_field[0:spawn_length, :] = 0
    goals[0] = [spawn_length - m_to_idx(0.5), mid_y]

    # Pit area
    height_field[spawn_length:, :] = -1.0
    cur_x = spawn_length

    for i in range(3):
        width = np.random.randint(platform_width_min, platform_width_max)
        height = np.random.uniform(platform_height_min, platform_height_max)
        
        # Add platform
        height_field[cur_x:cur_x + platform_length, mid_y - width//2:mid_y + width//2] = height
        goals[i + 1] = [cur_x + platform_length // 2, mid_y]
        cur_x += platform_length + gap_length

        # Add inclined path (diagonal ascent)
        height_field[cur_x:cur_x + incline_length, mid_y - width//2:mid_y + width//2] = incline_height[:, None]
        cur_x += incline_length

    # Final challenging platform
    final_width = np.random.randint(platform_width_min, platform_width_max)
    final_height = np.random.uniform(platform_height_min, platform_height_max)

    height_field[cur_x:cur_x + platform_length, mid_y - final_width//2:mid_y + final_width//2] = final_height
    goals[-2] = [cur_x + platform_length // 2, mid_y]

    # Add the last goal, restoring flat terrain
    cur_x += platform_length + gap_length
    goals[-1] = [cur_x + m_to_idx(0.3), mid_y]
    height_field[cur_x:, :] = 0

    return height_field, goals

def set_terrain_7(length, width, field_resolution, difficulty):
    """Terrain with alternating raised platforms and zigzag ramps testing agility and strategic navigation."""

    def m_to_idx(m):
        """Converts meters to quantized indices."""
        return np.round(m / field_resolution).astype(np.int16) if not (isinstance(m, list) or isinstance(m, tuple)) else [round(i / field_resolution) for i in m]

    height_field = np.zeros((m_to_idx(length), m_to_idx(width)))
    goals = np.zeros((8, 2))

    # Set up dimensions for platforms and ramps
    platform_length = 0.8 + 0.2 * difficulty
    platform_length = m_to_idx(platform_length)
    platform_width = np.random.uniform(1.0, 1.4)
    platform_width = m_to_idx(platform_width)
    platform_height_min, platform_height_max = 0.05 + 0.2 * difficulty, 0.3 + 0.3 * difficulty
    ramp_length = 0.8 + 0.3 * difficulty
    ramp_length = m_to_idx(ramp_length)
    ramp_height_min, ramp_height_max = 0.05 + 0.3 * difficulty, 0.15 + 0.4 * difficulty
    gap_length = 0.3 + 0.4 * difficulty
    gap_length = m_to_idx(gap_length)

    mid_y = m_to_idx(width) // 2

    def add_platform(start_x, end_x, mid_y):
        half_width = platform_width // 2
        x1, x2 = start_x, end_x
        y1, y2 = mid_y - half_width, mid_y + half_width
        platform_height = np.random.uniform(platform_height_min, platform_height_max)
        height_field[x1:x2, y1:y2] = platform_height

    def add_zigzag_ramp(start_x, end_x, mid_y, direction):
        half_width = platform_width // 2
        x1, x2 = start_x, end_x
        y1, y2 = mid_y - half_width, mid_y + half_width
        ramp_height = np.linspace(ramp_height_min, ramp_height_max, num=x2-x1)
        zigzag = np.tile(ramp_height, (2, 1)).flatten()[:y2-y1]
        zigzag = zigzag[None, :]  # Add a dimension for broadcasting to x
        if direction == -1:  # Reverse the zigzag order for alternating
            zigzag = zigzag[:, ::-1]
        height_field[x1:x2, y1:y2] = zigzag

    dx_min, dx_max = -0.1, 0.1
    dx_min, dx_max = m_to_idx(dx_min), m_to_idx(dx_max)
    dy_min, dy_max = -0.3, 0.3
    dy_min, dy_max = m_to_idx(dy_min), m_to_idx(dy_max)

    # Set flat ground for spawning area
    spawn_length = m_to_idx(2)
    height_field[0:spawn_length, :] = 0
    # Put first goal at spawn
    goals[0] = [spawn_length - m_to_idx(0.5), mid_y]

    # Set remaining area to have alternating platforms and ramps
    cur_x = spawn_length
    for i in range(3):  # Extend track with alternating ramps and platforms
        # Add platform
        add_platform(cur_x, cur_x + platform_length, mid_y)
        goals[2 * i + 1] = [cur_x + platform_length / 2, mid_y]
        cur_x += platform_length + gap_length

        # Add zigzag ramps
        direction = -1 if i % 2 else 1  # Alternate direction for zigzag effect
        add_zigzag_ramp(cur_x, cur_x + ramp_length, mid_y, direction)
        goals[2 * i + 2] = [cur_x + ramp_length / 2, mid_y]
        cur_x += ramp_length + gap_length

    # Add final goal at the end of the last element
    goals[-1] = [cur_x + m_to_idx(0.5), mid_y]
    height_field[cur_x:, :] = 0

    return height_field, goals

def set_terrain_8(length, width, field_resolution, difficulty):
    """Narrow beams and staggered platforms over pits to test balance and precision."""

    def m_to_idx(m):
        """Converts meters to quantized indices."""
        return np.round(m / field_resolution).astype(np.int16) if not (isinstance(m, list) or isinstance(m, tuple)) else [round(i / field_resolution) for i in m]

    height_field = np.zeros((m_to_idx(length), m_to_idx(width)))
    goals = np.zeros((8, 2))

    # Platform and beam setup
    platform_length = 0.8 - 0.2 * difficulty
    platform_length = m_to_idx(platform_length)
    platform_width = np.random.uniform(0.8, 1.0)
    platform_width = m_to_idx(platform_width)
    platform_height_range = (0.1 + 0.1 * difficulty, 0.15 + 0.15 * difficulty)

    beam_length = 1.0 - 0.5 * difficulty
    beam_length = m_to_idx(beam_length)
    beam_width = 0.4
    beam_width = m_to_idx(beam_width)

    gap_length = 0.6 + 0.3 * difficulty
    gap_length = m_to_idx(gap_length)

    mid_y = m_to_idx(width) // 2

    def add_platform(start_x, mid_y, height_range):
        half_width = platform_width // 2
        y1, y2 = mid_y - half_width, mid_y + half_width
        platform_height = np.random.uniform(*height_range)
        x2 = start_x + platform_length
        height_field[start_x:x2, y1:y2] = platform_height
        return x2

    def add_beam(start_x, mid_y, height_range):
        half_width = beam_width // 2
        y1, y2 = mid_y - half_width, mid_y + half_width
        beam_height = np.random.uniform(*height_range)
        x2 = start_x + beam_length
        height_field[start_x:x2, y1:y2] = beam_height
        return x2

    # Flat spawn area
    spawn_length = m_to_idx(2)
    height_field[0:spawn_length, :] = 0
    goals[0] = [spawn_length - m_to_idx(0.5), mid_y]

    # Set pit for challenge
    height_field[spawn_length:, :] = -1.0

    # Starting offset beyond spawn
    cur_x = spawn_length

    for i in range(4):  # Alternate between platforms and beams
        if i % 2 == 0:
            cur_x = add_platform(cur_x, mid_y, platform_height_range)
        else:
            cur_x = add_beam(cur_x, mid_y, platform_height_range)
        
        # Place goal in the center of each obstacle
        goals[i+1] = [cur_x - platform_length // 2, mid_y]

        # Add gap between elements
        cur_x += gap_length

    # Final platform at the end
    cur_x = add_platform(cur_x, mid_y, platform_height_range)
    goals[5] = [cur_x - platform_length // 2, mid_y]
    
    # Place the final goal beyond the last platform
    goals[6] = [cur_x + m_to_idx(0.25), mid_y]

    # Ensure the ground level at the end of the course
    height_field[cur_x:, :] = 0
    goals[7] = [cur_x + m_to_idx(1.0), mid_y]

    return height_field, goals

def set_terrain_9(length, width, field_resolution, difficulty):
    """Zigzag ramps with varied inclinations and strategic platforms for balanced navigation and jumping."""

    def m_to_idx(m):
        """Converts meters to quantized indices."""
        return np.round(m / field_resolution).astype(np.int16) if not (isinstance(m, list) or isinstance(m, tuple)) else [round(i / field_resolution) for i in m]

    height_field = np.zeros((m_to_idx(length), m_to_idx(width)))
    goals = np.zeros((8, 2))

    # Initialize platform and ramp dimensions
    platform_length = 0.9 - 0.2 * difficulty
    platform_length = m_to_idx(platform_length)
    platform_width = np.random.uniform(1.0, 1.2)
    platform_width = m_to_idx(platform_width)
    platform_height_min, platform_height_max = 0.05 + 0.15 * difficulty, 0.1 + 0.25 * difficulty
    ramp_height_min, ramp_height_max = 0.1 + 0.3 * difficulty, 0.15 + 0.4 * difficulty
    gap_length = 0.15 + 0.4 * difficulty
    gap_length = m_to_idx(gap_length)

    mid_y = m_to_idx(width) // 2

    def add_platform(start_x, end_x, mid_y, height_value):
        half_width = platform_width // 2
        x1, x2 = start_x, end_x
        y1, y2 = mid_y - half_width, mid_y + half_width
        height_field[x1:x2, y1:y2] = height_value

    def add_ramp(start_x, end_x, mid_y, start_height, end_height):
        half_width = platform_width // 2
        x1, x2 = start_x, end_x
        y1, y2 = mid_y - half_width, mid_y + half_width
        slant = np.linspace(start_height, end_height, x2-x1)[:, None]
        height_field[x1:x2, y1:y2] = slant

    # Set spawn area to flat ground
    spawn_length = m_to_idx(2)
    height_field[0:spawn_length, :] = 0
    goals[0] = [spawn_length - m_to_idx(0.5), mid_y]  

    # Add initial platform to transition into zigzag ramps
    cur_x = spawn_length
    platform_height = np.random.uniform(platform_height_min, platform_height_max)
    add_platform(cur_x, cur_x + platform_length, mid_y, platform_height)
    goals[1] = [cur_x + platform_length / 2, mid_y]
    cur_x += platform_length + gap_length

    for i in range(3):  # Introduce a series of ramps and platforms
        direction = (-1) ** i  # Alternate ramp direction for zigzag effect
        next_height = np.random.uniform(platform_height_min, platform_height_max)
        
        # Add ramp
        add_ramp(cur_x, cur_x + platform_length, mid_y + direction * platform_width // 2, platform_height, next_height)
        goals[2 * i + 2] = [cur_x + platform_length / 2, mid_y + direction * platform_width // 2]
        cur_x += platform_length + gap_length

        # Add connecting platform
        platform_height = next_height
        add_platform(cur_x, cur_x + platform_length, mid_y + direction * platform_width, platform_height)
        goals[2 * i + 3] = [cur_x + platform_length / 2, mid_y + direction * platform_width]
        cur_x += platform_length + gap_length

    # Add final goal after last platform
    goals[-1] = [cur_x + m_to_idx(0.5), mid_y]
    height_field[cur_x:, :] = 0

    return height_field, goals

# INSERT TERRAIN FUNCTION DEFINITIONS HERE