import numpy as np
import random

def generate_layers_1d(grid_size, num_layers, min_size, transition_size, first_layer_minsize, first_layer_maxsize): 
    """Generate a 1D stratified model with transition zones."""
    usable_size = grid_size - (num_layers - 1) * transition_size
    grid = np.zeros(grid_size, dtype=int)
    remaining_size = usable_size
    layer_sizes = []

    # Determine layer sizes
    for i in range(num_layers - 1):
        if i == 0:
            size = random.randint(first_layer_minsize, first_layer_maxsize)
        else:
            size = random.randint(min_size, remaining_size - (num_layers - i - 1) * min_size)
        layer_sizes.append(size)
        remaining_size -= size
    layer_sizes.append(remaining_size)

    # Assign layers and transitions to grid
    current_position = 0
    for i, size in enumerate(layer_sizes):
        grid[current_position:current_position + size] = i + 1
        current_position += size
        if i < num_layers - 1:
            grid[current_position:current_position + transition_size] = 0  # 0 indicates transition zone
            current_position += transition_size

    return grid

def assign_permittivity(grid, num_layers, permittivity_range=(1, 81), transition_size=5):
    """Assign permittivity values to layers and interpolate values across transition zones."""
    permittivity_grid = np.zeros_like(grid, dtype=float)
    layer_permittivities = {}

    # Assign permittivity to each layer
    for layer in range(1, num_layers + 1):
        value = random.uniform(*permittivity_range)
        layer_permittivities[layer] = value
        permittivity_grid[grid == layer] = value

    # Smooth transitions between layers using linear interpolation
    for i in range(1, num_layers):
        start_val = layer_permittivities[i]
        end_val = layer_permittivities[i + 1]
        transition_start = np.where(grid == 0)[0][(i - 1) * transition_size]
        for j in range(transition_size):
            t = j / (transition_size - 1)
            permittivity_grid[transition_start + j] = (1 - t) * start_val + t * end_val

    return permittivity_grid

def assign_permittivity_with_smooth_transition(grid, num_layers, first_layer_eps_range, permittivity_range=(1, 81), transition_size=5):
    """Assign permittivity values using smooth cosine transition between layers."""
    permittivity_grid = np.zeros_like(grid, dtype=float)
    layer_permittivities = {}

    for layer in range(1, num_layers + 1):
        if layer == 1:
            value = random.randint(*first_layer_eps_range)
        else:
            value = random.uniform(*permittivity_range)
        layer_permittivities[layer] = value
        permittivity_grid[grid == layer] = value

    # Cosine-smooth transitions
    for i in range(1, num_layers):
        start_val = layer_permittivities[i]
        end_val = layer_permittivities[i + 1]
        transition_start = np.where(grid == 0)[0][(i - 1) * transition_size]
        for j in range(transition_size):
            t = j / (transition_size - 1)
            smooth = 0.5 * (1 - np.cos(np.pi * t))
            permittivity_grid[transition_start + j] = start_val + (end_val - start_val) * smooth

    return permittivity_grid

def assign_integer_values(permittivity_grid):
    """Convert permittivity values to discrete integers representing unique material zones."""
    layer_id = np.zeros_like(permittivity_grid, dtype=float)
    unique_permittivities = {}
    integer_value = 0

    for idx, val in enumerate(permittivity_grid):
        val = round(val, 5)
        if idx == 0 or val != round(permittivity_grid[idx - 1], 5):
            unique_permittivities[val] = integer_value
            integer_value += 1
        layer_id[idx] = unique_permittivities[val]

    return layer_id